Euler Maruyama issue

Questions v5
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