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?