Euler Maruyama issue

Questions v5
1

Hi all. I’m trying to run the following example (2nd in the link): Inferring parameters of SDEs using a Euler-Maruyama scheme — PyMC3 3.11.4 documentation . However I run into and index error.

import numpy as np
import pymc as pm
import pytensor.tensor as pt

dt=0.1
N, τ, a, m, σ2 = 200, 3.0, 1.05, 0.2, 1e-1
xs, ys = [0.0], [1.0]
for i in range(N):
    x, y = xs[-1], ys[-1]
    dx = τ * (x - x ** 3.0 / 3.0 + y)
    dy = (1.0 / τ) * (a - x)
    xs.append(x + dt * dx + np.sqrt(dt) * σ2 * np.random.randn())
    ys.append(y + dt * dy + np.sqrt(dt) * σ2 * np.random.randn())
xs, ys = np.array(xs), np.array(ys)
zs = m * xs + (1 - m) * ys + np.random.randn(xs.size) * 0.1

def osc_sde(xy, τ, a):
    x, y = xy[:, 0], xy[:, 1]
    dx = τ * (x - x ** 3.0 / 3.0 + y)
    dy = (1.0 / τ) * (a - x)
    dxy = pt.stack([dx, dy], axis=0).T
    return dxy, σ2


xys = np.array([xs, ys]).T
initd = pm.Normal.dist(0, 1, shape=xys.shape)
coords = {"time":np.arange(xys.shape[0]), "vars":np.arange(xys.shape[1])}
with pm.Model(coords=coords) as model:
    τh = pm.Uniform("τh", lower=0.1, upper=5.0)
    ah = pm.Uniform("ah", lower=0.5, upper=1.5)
    mh = pm.Uniform("mh", lower=0.0, upper=1.0)
    xyh = pm.EulerMaruyama("xyh", dt, osc_sde, (τh, ah), dims=("time", "vars"), init_dist=initd)
    zh = pm.Normal("zh", mu=mh * xyh[:, 0] + (1 - mh) * xyh[:, 1], sigma=0.1, observed=zs)

...
step
    f, g = sde_fn(prev_y, *prev_sde_pars)

  Cell In[1], line 18 in osc_sde
    x, y = xy[:, 0], xy[:, 1]

IndexError: too many indices for array

The example runs fine in PyMC3. I cannot figure out the issue, nor find any updated tutorials or documents on how to appropriately implement this in PyMC v5. Any clue would be really appreciated. Many thanks in advance.

1 Like
2

Any chance you could use something more modern like sunode? GitHub - pymc-devs/sunode: Solve ODEs fast, with support for PyMC

1 Like
3

Thanks for the reply. I hadn’t thought of it. Maybe I should give a try to sunode, I guess just adding the parameters for the SDEs as priors may work. I’m trying with scan, but I tend to get confused when trying to recover the rngs. I couldn’t figure out how to use “pytensorf.collect_default_updates”. I’m looking at this thread: Simple markov chain with aesara scan - #10 by ricardoV94 . So maybe that’ll help to solve the issue.

4

Have you seen these AR/MA notebooks? They show the more updated way to work with scan. I can also help debug code if you post what you have so far.

1 Like
5

Thanks so much. I’m trying to sample some stochastic simulations of compartment models. My idea was simulating via an Euler-Maruyama or Euler-Binomial using PyMC to get MCMC samples or to sample the simulated output to test parameters recoverability. There are several models of increasing complexity, so I made an example with very simple simulations:

import pymc as pm
import numpy as np
import pytensor
import pytensor.tensor as pt
import matplotlib.pyplot as plt
from pymc.pytensorf import collect_default_updates
import arviz as az

# rng = np.random.default_rng()
# rvs = sp.stats.norm(loc=1000, scale=500).rvs
# M = sp.sparse.random(10, 10, density=0.75, random_state=rng, data_rvs=rvs).A
# M[M<0] = 0

N = pm.draw( pm.TruncatedNormal.dist(mu=20e6, sigma=10e6), 1).round(0)

#### Simulate SIR epidemic development

# initial parameters
initial_infected = 1

#epidemic length
n_days = 100 #days
times = np.arange(n_days)  #days array index

#defined parameters
beta = 0.5 #infection rate
gamma = 0.1 #recovery rate

# Initialize arrays
susceptible = np.zeros(n_days)
infected = np.zeros(n_days)
recovered = np.zeros(n_days)

# Set initial conditions
susceptible[0] = N - initial_infected
infected[0] = initial_infected

I_begin = infected[0] 
S_begin = N - infected[0] 
R_begin = I_begin*0 

dt = 1

not_take_off = True
while not_take_off:
    for t in range(1, n_days):
        # Compute new infections and recoveries
        # new_infections = beta * susceptible[t - 1] * infected[t - 1] / N
        # new_recoveries = gamma * infected[t - 1]
        
        new_infections = np.random.binomial(susceptible[t - 1], 1 - np.exp(-beta*dt * infected[t - 1] / N))
        new_recoveries = np.random.binomial(infected[t - 1], 1 - np.exp(-gamma*dt))
        
        # Update compartments
        susceptible[t] = susceptible[t - 1] - new_infections
        infected[t] = infected[t - 1] + new_infections - new_recoveries
        recovered[t] = recovered[t - 1] + new_recoveries
        
    if np.array([infected, recovered]).sum() > 0:
        not_take_off = False
    
S = susceptible
I = infected
R = recovered

plt.plot(range(n_days), S, linestyle=":", label='Susceptible')
plt.plot(range(n_days), I, linestyle="--", label='Infected')
plt.plot(range(n_days), R, label='Recovered')
plt.xlabel('Days')
plt.ylabel('Number of Individuals')
plt.title('SIR Model Simulation')
plt.rc('axes.spines',top=False, right=False)
plt.tight_layout()
plt.legend()
plt.grid(alpha=0.2)
plt.show()
plt.close()


rngI = pytensor.shared(np.random.default_rng())
rngR = pytensor.shared(np.random.default_rng())

def sir_func(S_prev, I_prev, R_prev, beta, gamma, N, dt):
    
    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))
    
    # new_I = beta * S_prev * I_prev / N
    # new_R = gamma * I_prev 
    
    S = S_prev - new_I
    I = I_prev + new_I - new_R
    R = R_prev + new_R
    
    return (S, I, R), collect_default_updates([new_I, new_R])


with pm.Model() as mod:
    beta = pm.HalfNormal("beta", 1)
    gamma = pm.HalfNormal("gamma", 1)
    dt = 0.1

    outputs, _ = pytensor.scan(
        fn=sir_func,
        outputs_info=[S_begin, I_begin, R_begin],
        non_sequences=[beta, gamma, N, dt],
        n_steps=n_days
    )
    
    Send = pm.Deterministic("S", outputs[0])
    Iend = pm.Deterministic("I", outputs[1])
    Rend = pm.Deterministic("R", outputs[2])
    
    y = pm.Poisson("y", mu=Iend, observed=infected)
    
with mod:
    idata = pm.sample()
    

S = az.extract(idata)["S"].values.mean(axis=1)
I = az.extract(idata)["I"].values.mean(axis=1)
R = az.extract(idata)["R"].values.mean(axis=1)

plt.plot(range(n_days), S, linestyle=":", label='Susceptible')
plt.plot(range(n_days), I, linestyle="--", label='Infected')
plt.plot(range(n_days), R, label='Recovered')
plt.xlabel('Days')
plt.ylabel('Number of Individuals')
plt.title('SIR Model Simulation\n(origin country)')
plt.rc('axes.spines',top=False, right=False)
plt.tight_layout()
plt.legend()
plt.grid(alpha=0.2)
plt.show()
plt.close()

The notebook you linked was of great help to get a relatively functioning model. There’s one issue, however, that I cannot figure out from the notebook. If I try to sample without the “y = pm.Poisson(“y”, mu=Iend, observed=infected)” likelihood the model samples well, though with a bit of a clumsy output. But when I add the aforementioned likelihood, I get the following error:

ValueError: Random variables detected in the logp graph: {binomial_rv{0, (0, 0), int64, False}.out, binomial_rv{0, (0, 0), int64, False}.0, halfnormal_rv{0, (0, 0), floatX, False}.0, binomial_rv{0, (0, 0), int64, False}.out, binomial_rv{0, (0, 0), int64, False}.0, halfnormal_rv{0, (0, 0), floatX, False}.0}.
This can happen when DensityDist logp or Interval transform functions reference nonlocal variables,
or when not all rvs have a corresponding value variable.
1 Like
6

Regarding the lack of examples, sometimes the best source comes from the tests in the codebase: https://github.com/pymc-devs/pymc/blob/602234b98d7c2c1e4382a56c16e7f71c1241bc8a/tests/distributions/test_timeseries.py#L848

1 Like
7

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?

1 Like
8

Thanks for the reply. Sadly, I cannot find an example on that link which could help me with systems of 2 or more differential equations. I find the same errors described below by @jessegrabowski when trying to implement his code.

9

Many thanks for taking this throughout look to the problem. I have tried to run both models your proposed and I get the same error: ValueError: not enough values to unpack (expected 3, got 0) .

Indeed, a model like the second one is ideal, as the three SIR outputs are required. Hopefully there’s a way around the issue.

10

Isn’t this the same old convolution problem (just looking at that @ there). PyMC doesn’t currently infer any logps for operations that are not obviously invertible.

1 Like
11

For the simple SIR logprob, PyMC inference fails for two quaint reasons:

  1. We don’t support transformations of discrete variables yet. So stuff like Binomial + const won’t work, although Normal + const works. We are working on expanding this functionality
  2. We don’t support logp that are conditioned on the transformation of other inferred RVs. We can boil this limitation to the following example:
import pymc as pm

prev_I = 0.5
prev_R = 0.5

new_I = pm.Normal.dist()
new_R = pm.Normal.dist()

I = prev_I + new_I - new_R
R = prev_R + new_R

I.name = "I"
R.name = "R"
        
pm.logprob.conditional_logp({I: I.clone(), R: R.clone()})
# RuntimeError: The logprob terms of the following value variables could not be derived: {I}

There is nothing challenging about the logp of I, but we don’t have the mechanism to tell it that instead of new_R it can count on R - prev_R. All it sees is an expression that depends on two “RVs” which is not supported.

If the EulerMuryama doesn’t support multivariate processes, I think you have to fallback to implementing your own logp by hand. You can still use CustomDist and pass a logp=logp_fn kwarg.

2 Likes
12

Many thanks, I’ll give it a try.