pymc
pymc copied to clipboard
pymc-devs
BUG: NaNs in logp mess up resampling step in SMC
Describe the issue:
The model I am working with sometimes returns NaN for the logp (it is an ODE model, and for certain nasty parameter combinations the solver gives up). When using SMC as a sampler, this leads to unexpected behavior in the resampling step: All the weights will be NaN, and as a result the resampling indices will all be zero.
Arguably, the logp should not be NaN, and it is unclear (at least to me) what behavior would be correct in this situation. But the current behavior will lead to a collapse of the samples without any warning or explanation of what is happening. My pragmatic solution would be to ignore any samples with a NaN logp during resampling by setting their logp to -inf.
I don't know how to create a complete example model to reproduce this. But the following code demonstrates the behavior of the current SMC implementation:
Reproduceable code example:
from scipy.special import logsumexp
from pymc.smc.kernels import systematic_resampling
import numpy as np
# Create random logp values
probs = np.random.rand(20)
probs[3] = np.NaN
log_weights_un = np.log(probs)
# Set NaNs to -inf to resolve the problem
# log_weights_un = np.where(np.isnan(log_weights_un), -np.inf, log_weights_un)
# Compute the weights as in the SMC kernel
log_weights = log_weights_un - logsumexp(log_weights_un)
weights = np.exp(log_weights)
# Get the resampling indices
rng = np.random.default_rng(seed=5)
systematic_resampling(weights, rng=rng)
Error message:
No response
PyMC version information:
PyMC version: 5.13.1 (local installation of main branch)
Context for the issue:
No response
Can't the handling of nan be done on your/user side?
If that is possible that might be the preferred solution. But how would I do that inside the PyMC model?
I think that the sampler should at least give a warning about NaN values and how they influence the results. What I currently end up with is a “posterior” that consists of a single sample and it is absolutely not clear that this could be related to NaN values.
@ricardoV94 Nuts for instance also treats nan as -inf. Without this many more models would diverge. Especially early on during sampling nans are not that uncommon, when we might be using unrealistic step sizes for a couple iterations.
Doing the same thing in smc makes sense to me.
But then, I guess a big difference is that in NUTS we still record that this happend in the form of a divergence. I don't think there is currently anything like this in smc...
If that is possible that might be the preferred solution. But how would I do that inside the PyMC model?
How are you plugging the ODE into PyMC? Wherever the logp is defined there can be a switch if it comes out nan.
I use sunode to solve the ODE and then simply define a Normal distribution with observed values around the solution in PyMC. Something like this:
# Solve ODE with sunode
mu = solve_ivp(x0, params, rhs, tvals, t0)
# Data follow Normal distribution around the mean
y = pm.Normal(mu=mu, sigma=sigma, observed=data)
where mu can be NaN.