I think this specific error is because you’re trying to mix automatic logp inference with an “ordinary” logp.
Logp inference happens in the context of a pm.CustomDist. This takes the function with .dist RVs and works out the associated logp.
In your case, no automatic inference is going on, because you’re giving the .dist RVs to pm.Poisson, who doesn’t know what to do with them. You can confirm that this model (which is more obviously wrong) throws the same error:
with pm.Model() as m:
sigma_mu = pm.Exponential('sigma_mu', lam=1)
mu = pm.LogNormal.dist(0, sigma_mu)
y = pm.Poisson('y', mu=x * mu, observed=np.random.poisson(lam=3, size=(100)))
pm.sample()
The scan you’ve written already implicitly defines a distribution over the S, I, and R, so you don’t need to resort to pm.Poisson. You should be using a CustomDist directly. This has some limitations. In particular, you can only have a single output from the dist function. Here’s what I came up with:
def sir_model(S0, I0, R0, beta, gamma, N, dt, n_steps, size):
def sir_step(*args):
S_prev, I_prev, R_prev, beta, gamma, N, dt = args
new_I = pm.Binomial.dist(S_prev, 1 - pt.exp(-beta*dt * I_prev / N))
new_R = pm.Binomial.dist(I_prev, 1 - pt.exp(-gamma*dt))
S = S_prev - new_I
I = I_prev + new_I - new_R
R = R_prev + new_R
return (S, I, R), collect_default_updates([S, I, R])
(S, I, R), updates = pytensor.scan(
fn=sir_step,
outputs_info=[S0, I0, R0],
non_sequences=[beta, gamma, N, dt],
n_steps=n_steps,
strict=True
)
# We're only allowed a single output in the dist function of CustomDist
return I
coords = {'days':np.arange(n_days, dtype='int')}
with pm.Model(coords=coords) as mod:
data = pm.MutableData('infected', infected, dims=['days'])
dt = pt.as_tensor_variable(0.1)
# Priors
beta = pm.HalfNormal("beta", 1)
gamma = pm.HalfNormal("gamma", 1)
# Define a single distribution over the entire sequence using the scan
y_hat = pm.CustomDist('SIR_process', S_begin, I_begin, R_begin,
beta,
gamma,
N,
dt,
n_days,
dist=sir_model,
observed=data,
dims=['days'])
idata_prior = pm.sample_prior_predictive()
idata = pm.sample()
Unfortunately, this raises a new error:
Long error
ValueError Traceback (most recent call last)
Cell In[52], line 49
39 y_hat = pm.CustomDist('SIR_process', S_begin, I_begin, R_begin,
40 beta,
41 gamma,
(...)
46 observed=data,
47 dims=['days'])
48 idata_prior = pm.sample_prior_predictive()
---> 49 idata = pm.sample()
File ~/mambaforge/envs/cge-dev/lib/python3.11/site-packages/pymc/sampling/mcmc.py:651, in sample(draws, tune, chains, cores, random_seed, progressbar, step, nuts_sampler, initvals, init, jitter_max_retries, n_init, trace, discard_tuned_samples, compute_convergence_checks, keep_warning_stat, return_inferencedata, idata_kwargs, nuts_sampler_kwargs, callback, mp_ctx, model, **kwargs)
648 auto_nuts_init = False
650 initial_points = None
--> 651 step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
653 if nuts_sampler != "pymc":
654 if not isinstance(step, NUTS):
File ~/mambaforge/envs/cge-dev/lib/python3.11/site-packages/pymc/sampling/mcmc.py:211, in assign_step_methods(model, step, methods, step_kwargs)
209 methods_list: List[Type[BlockedStep]] = list(methods or pm.STEP_METHODS)
210 selected_steps: Dict[Type[BlockedStep], List] = {}
--> 211 model_logp = model.logp()
213 for var in model.value_vars:
214 if var not in assigned_vars:
215 # determine if a gradient can be computed
File ~/mambaforge/envs/cge-dev/lib/python3.11/site-packages/pymc/model/core.py:720, in Model.logp(self, vars, jacobian, sum)
718 rv_logps: List[TensorVariable] = []
719 if rvs:
--> 720 rv_logps = transformed_conditional_logp(
721 rvs=rvs,
722 rvs_to_values=self.rvs_to_values,
723 rvs_to_transforms=self.rvs_to_transforms,
724 jacobian=jacobian,
725 )
726 assert isinstance(rv_logps, list)
728 # Replace random variables by their value variables in potential terms
File ~/mambaforge/envs/cge-dev/lib/python3.11/site-packages/pymc/logprob/basic.py:612, in transformed_conditional_logp(rvs, rvs_to_values, rvs_to_transforms, jacobian, **kwargs)
609 transform_rewrite = TransformValuesRewrite(values_to_transforms) # type: ignore
611 kwargs.setdefault("warn_rvs", False)
--> 612 temp_logp_terms = conditional_logp(
613 rvs_to_values,
614 extra_rewrites=transform_rewrite,
615 use_jacobian=jacobian,
616 **kwargs,
617 )
619 # The function returns the logp for every single value term we provided to it.
620 # This includes the extra values we plugged in above, so we filter those we
621 # actually wanted in the same order they were given in.
622 logp_terms = {}
File ~/mambaforge/envs/cge-dev/lib/python3.11/site-packages/pymc/logprob/basic.py:541, in conditional_logp(rv_values, warn_rvs, ir_rewriter, extra_rewrites, **kwargs)
538 q_values = remapped_vars[: len(q_values)]
539 q_rv_inputs = remapped_vars[len(q_values) :]
--> 541 q_logprob_vars = _logprob(
542 node.op,
543 q_values,
544 *q_rv_inputs,
545 **kwargs,
546 )
548 if not isinstance(q_logprob_vars, (list, tuple)):
549 q_logprob_vars = [q_logprob_vars]
File ~/mambaforge/envs/cge-dev/lib/python3.11/functools.py:909, in singledispatch.<locals>.wrapper(*args, **kw)
905 if not args:
906 raise TypeError(f'{funcname} requires at least '
907 '1 positional argument')
--> 909 return dispatch(args[0].__class__)(*args, **kw)
File ~/mambaforge/envs/cge-dev/lib/python3.11/site-packages/pymc/logprob/scan.py:308, in logprob_ScanRV(op, values, name, *inputs, **kwargs)
305 scan_args = ScanArgs.from_node(new_node)
306 rv_outer_outs = get_random_outer_outputs(scan_args)
--> 308 var_indices, rv_vars, io_vars = zip(*rv_outer_outs)
309 value_map = {_rv: _val for _rv, _val in zip(rv_vars, values)}
311 def create_inner_out_logp(value_map: Dict[TensorVariable, TensorVariable]) -> TensorVariable:
ValueError: not enough values to unpack (expected 3, got 0)
I’m pretty sure the way I wrote it should work, but it doesn’t so I have to call in the @ricardoV94 cavalry. I tried it every which way – with vector-valued returns, with non-discrete distributions, with manual updates, with inputs as tensor/non-tensors/random variables… so I’m good and stumped. Hopefully you see an obvious error I’m missing.
In principal, I think this next model should be the “most correct” way to do it. It takes an SIR vector as initial value and returns a (n_days, 3) matrix. Here, the scan defines a joint distribution over the whole system, not just the bit you observe. This is what the model actually does, after all:
Best model?
def sir_model(SIR0, beta, gamma, N, dt, n_steps, size):
# Transition matrix to update SIR variables. SIR + A @ IR_shocks leads to the 3 update equations
A = pt.as_tensor_variable(np.array([[-1, 0],
[1, -1],
[0, 1]]))
def sir_step(*args):
# Same as yours, but I changed it to be vector valued
SIR_prev, beta, gamma, N, dt = args
probs = pt.stack([1 - pt.exp(-beta*dt * SIR_prev[1] / N),
1 - pt.exp(-gamma*dt)])
IR_shocks = pm.Binomial.dist(n=SIR_prev[:-1], p=probs)
SIR_next = SIR_prev + A @ IR_shocks
return SIR_next, collect_default_updates([SIR_next])
SIR, updates = pytensor.scan(
fn=sir_step,
outputs_info=[SIR0],
non_sequences=[beta, gamma, N, dt],
n_steps=n_steps,
strict=True
)
# We're only allowed a single output in the dist function of CustomDist, that's why
# I made it work with vectors.
return SIR
coords = {'days':np.arange(n_days, dtype='int'),
'SIR':['Susceptible', 'Infected', 'Recovered']}
with pm.Model(coords=coords) as mod:
# Data and constants. Set unobserved processes (Susceptible, Recovered) to all NaN
# and let PyMC impute them
sir_data = np.full((n_days, 3), np.nan)
sir_data[:, 1] = infected
dt = 0.1
SIR0 = pt.as_tensor(np.r_[S_begin, I_begin, R_begin], name='initial_SIR', dtype='int64')
# Priors
beta = pm.HalfNormal("beta", 3)
gamma = pm.HalfNormal("gamma", 1)
# Define a single distribution over the entire sequence using the scan
y_hat = pm.CustomDist('SIR_process', SIR0,
beta,
gamma,
N,
dt,
n_days,
dist=sir_model,
observed=sir_data,
dims=['days', 'SIR'])
idata_prior = pm.sample_prior_predictive()
But the automatic inference fails on this one, and in a very opaque way. Hopefully there’s an obvious bug?