Can pm.draw() be used inside a pymc model or will it break links between RVs and observed/likelihood variable?

Questions v5
1

I’m fitting for the means and covariance matrix of a 3D MvNormal: the 3 means, the 3 standard deviations, and only one of the covariances (the other covariances are assumed to be zero). All follow Uniform priors. An older version of the code works well but relied on a gigantic pt.tensordot that involved evaluating pm.logp(mvnormal, points) at 1e6 pre-tabulated points. I’m exploring an alternative that involves drawing a smaller number (~1e3) of points given the current sample of mvnormal:

with model: 
    mvn = pm.MvNormal.dist(mu=mu_vector,cov=cov_matrix)
    random_draws = pm.draw(mvn,1000)
   
    [... model code that makes use of random_draws ...] 

    phi = pm.draw(pm.Uniform.dist(lower=0,upper=2*np.pi),1000)

    [... model code that uses both phi and random_draws to predict hist_mean_pdf ...]

    z_obs = pm.Poisson('z_obs', mu=hist_mean_pdf, observed=hist_observed)

This alternative model isn’t working yet and I noticed the pm.model_to_graphviz(model=model) doesn’t connect the RV boxes to the observed/likelihood variable box.

Are those pm.draw allowed inside a pymc model, or do I need something like pm.Deterministic? Note that I use pm.draw twice in my model – the first time is drawing from the mvn object which is completely determined by the values of the stochastic RVs mu_vector and cov_matrix, and the second time is just randomly drawing phi from a Uniform dist – phi is unconnected to any of the stochastic RVs but does contribute to the prediction that enters the z_obs likelihood variable.

2

You should just use the dists directly with the size you want. draw is a utility to materialize draws which returns a (one time) numpy array.

with model: 
    mvn = pm.MvNormal.dist(mu=mu_vector,cov=cov_matrix, size=1000)
   
    [... model code that makes use of mvn ...] 

    phi = pm.Uniform.dist(lower=0,upper=2*np.pi, size=1000)

    [... model code that uses both phi and random_draws to predict hist_mean_pdf ...]

    z_obs = pm.Poisson('z_obs', mu=hist_mean_pdf, observed=hist_observed)

However PyMC might complain that your logp is stochastic in which case you may need to hide the random draws inside a custom sampler:

3

Thanks @ricardoV94 ! I tried your size=1000 dists approach first and the model graphviz is successfully showing arrows connecting my RVs to the likelihood variable, but I got the following error – is this PyMC complaining that my logp is stochastic as you expected?

While trying logp_fn = model.compile_logp():

/opt/anaconda3/envs/pymc/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:190: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: [multivariate_normal_rv{1, (1, 2), floatX, False}.0, multivariate_normal_rv{1, (1, 2), floatX, False}.out]
  warnings.warn(
/opt/anaconda3/envs/pymc/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:190: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: [uniform_rv{0, (0, 0), floatX, False}.0, uniform_rv{0, (0, 0), floatX, False}.out]
  warnings.warn(

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[45], line 1
----> 1 logp_fn = model.compile_logp()

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/pymc/model.py:671, in Model.compile_logp(self, vars, jacobian, sum)
    652 def compile_logp(
    653     self,
    654     vars: Optional[Union[Variable, Sequence[Variable]]] = None,
    655     jacobian: bool = True,
    656     sum: bool = True,
    657 ) -> PointFunc:
    658     """Compiled log probability density function.
    659 
    660     Parameters
   (...)
    669         Defaults to True.
    670     """
--> 671     return self.model.compile_fn(self.logp(vars=vars, jacobian=jacobian, sum=sum))

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/pymc/model.py:759, in Model.logp(self, vars, jacobian, sum)
    757 rv_logps: List[TensorVariable] = []
    758 if rvs:
--> 759     rv_logps = joint_logp(
    760         rvs=rvs,
    761         rvs_to_values=self.rvs_to_values,
    762         rvs_to_transforms=self.rvs_to_transforms,
    763         jacobian=jacobian,
    764     )
    765     assert isinstance(rv_logps, list)
    767 # Replace random variables by their value variables in potential terms

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:308, in joint_logp(rvs, rvs_to_values, rvs_to_transforms, jacobian, **kwargs)
    305     value_var = rvs_to_values[rv]
    306     logp_terms[value_var] = temp_logp_terms[value_var]
--> 308 _check_no_rvs(list(logp_terms.values()))
    309 return list(logp_terms.values())

File /opt/anaconda3/envs/pymc/lib/python3.11/site-packages/pymc/logprob/joint_logprob.py:265, in _check_no_rvs(logp_terms)
    255 unexpected_rv_nodes = [
    256     node
    257     for node in pytensor.graph.ancestors(logp_terms)
   (...)
    262     )
    263 ]
    264 if unexpected_rv_nodes:
--> 265     raise ValueError(
    266         f"Random variables detected in the logp graph: {unexpected_rv_nodes}.\n"
    267         "This can happen when DensityDist logp or Interval transform functions "
    268         "reference nonlocal variables."
    269     )

ValueError: Random variables detected in the logp graph: [uniform_rv{0, (0, 0), floatX, False}.out, multivariate_normal_rv{1, (1, 2), floatX, False}.out, uniform_rv{0, (0, 0), floatX, False}.out].
This can happen when DensityDist logp or Interval transform functions reference nonlocal variables.

I am not familiar with BlockedStep – why is it that much more complicated machinery necessary? I noticed that wrapping the two pm.draw or .dists(size=1000) inside pm.Deterministic didn’t work either – is there a way to make pm.Deterministic work? In fact it might be nice to have these sets of 1000 “nuisance” parameters saved in the final trace (but only if it is not expensive).

4

We just decided to forbid it because we don’t know how samplers handle stochastic logp or gradients. It would be uncharted territory.

The extra machinery with the sampler has the benefit that the random component is “frozen” across logp and gradient evaluations (of which NUTS could do 1000s per step). Other samplers like Slice explicitly assume you can always contract or expand your proposal window in a way that you comeback to something as probable as before, which might not be the case with a stochastic logp.

If you have an informed opinion there’s a discussion here to at least provide this step sampler to users: Add FixedDistribution helper · pymc-devs/pymc · Discussion #6275 · GitHub

5

How would I adapt your custom sampling class UpdateFakeVariablesStep(BlockedStep) for my use case of needing N random draws from mvn = pm.MvNormal.dist(mu=mu_vector,cov=cov_matrix) and N random draws from phi = pm.Uniform.dist(lower=0,upper=2*np.pi)? Do I need two BlockedStep classes or just one?

I think maybe I can replace your i = self.rng.integers(100) with something like phi = self.rng.uniform(low=0,high=2*np.pi,size=1000) but I am not sure how to assign it back to the phi inside the with model context.

And it looks like similarly I can use mvn = self.rng.multivariate_normal(mean=mu_vector,cov=cov_matrix,size=1000) but again I am totally lost on how to connect this back into the with model context so that I can then use those resulting random numbers throughout the rest of my model to compute hist_mean_pdf for the Poisson likelihood.

6

A sampler can be responsible for multiple variables, my example in the other thread actually samples both a and aT0. Reproduced here:

import arviz as az
import numpy as np
import pymc as pm

from pymc.util import get_value_vars_from_user_vars
from pymc.step_methods.arraystep import BlockedStep


class UpdateFakeVariablesStep(BlockedStep):

    def __init__(self, vars, seed=None, model=None):
        model = pm.modelcontext(model)
        self.vars = get_value_vars_from_user_vars(vars, model)
        self.rng = np.random.default_rng(seed=seed)
        print(self.vars)

    def step(self, point: dict):
        # a1, a2
        new_point = point.copy()
        new_point[self.vars[0].name] = np.asarray(self.rng.integers(100))
        new_point[self.vars[1].name] = np.asarray(self.rng.integers(100) + 100)
        stats = [{}]
        return new_point, stats


with pm.Model() as local_model:
    a1 = pm.Flat('a1')
    a2 = pm.Flat('a2', initval=100)

    pm.Normal('y', a1, a2, observed=[0, 1, 2])

    update_fake_step = UpdateFakeVariablesStep([a1, a2])
    idata = pm.sample(step=[update_fake_step], chains=2)

az.summary(idata)