def test_logp_scalar_ode():
"""Test the computation of the log probability for these models"""
# Differential equation
def system_1(y, t, p):
return np.exp(-t) - p[0] * y[0]
# Parameters and inital condition
alpha = 0.4
y0 = 0.0
times = np.arange(0.5, 8, 0.5)
yobs = np.array(
[0.30, 0.56, 0.51, 0.55, 0.47, 0.42, 0.38, 0.30, 0.26, 0.21, 0.22, 0.13, 0.13, 0.09, 0.09]
)[:, np.newaxis]
ode_model = DifferentialEquation(func=system_1, t0=0, times=times, n_theta=1, n_states=1)
integrated_solution, *_ = ode_model._simulate([y0], [alpha])
assert integrated_solution.shape == yobs.shape
# compare automatic and manual logp values
manual_logp = norm.logpdf(x=np.ravel(yobs), loc=np.ravel(integrated_solution), scale=1).sum()
with pm.Model() as model_1:
forward = ode_model(theta=[alpha], y0=[y0])
y = pm.Normal("y", mu=forward, sd=1, observed=yobs)
pymc3_logp = model_1.logp()
np.testing.assert_allclose(manual_logp, pymc3_logp)
def test_op_equality(self):
"""Tests that the equality of mathematically identical Ops evaluates True"""
# Create ODE to test with
def ode_func(y, t, p):
return np.exp(-t) - p[0] * y[0]
t = np.linspace(0, 2, 12)
# Instantiate two Ops
op_1 = DifferentialEquation(func=ode_func,
t0=0,
times=t,
n_states=1,
n_theta=1)
op_2 = DifferentialEquation(func=ode_func,
t0=0,
times=t,
n_states=1,
n_theta=1)
op_other = DifferentialEquation(func=ode_func,
t0=0,
times=np.linspace(0, 2, 16),
n_states=1,
n_theta=1)
assert op_1 == op_2
assert op_1 != op_other
return
def test_simulate():
"""Tests the integration in DifferentialEquation"""
# Create an ODe to integrate
def ode_func(y, t, p):
return np.exp(-t) - p[0] * y[0]
# Evaluate exact solution
y0 = 0
t = np.arange(0, 12, 0.25).reshape(-1, 1)
a = 0.472
y = 1.0 / (a - 1) * (np.exp(-t) - np.exp(-a * t))
# Instantiate ODE model
ode_model = DifferentialEquation(func=ode_func,
t0=0,
times=t,
n_states=1,
n_theta=1)
simulated_y, sens = ode_model._simulate([y0], [a])
assert simulated_y.shape == (len(t), 1)
assert sens.shape == (len(t), 1, 1 + 1)
np.testing.assert_allclose(y, simulated_y, rtol=1e-5)
def test_number_of_params(self):
with pytest.raises(ValueError):
DifferentialEquation(func=self.system,
t0=0,
times=self.times,
n_states=1,
n_theta=0)
def test_func_callable(self):
with pytest.raises(ValueError):
DifferentialEquation(func=1,
t0=0,
times=self.times,
n_states=1,
n_theta=1)
class TestErrors:
"""Test running model for a scalar ODE with 1 parameter"""
def system(y, t, p):
return np.exp(-t) - p[0] * y[0]
times = np.arange(0, 9)
ode_model = DifferentialEquation(func=system,
t0=0,
times=times,
n_states=1,
n_theta=1)
@pytest.mark.xfail(condition=(aesara.config.floatX == "float32"),
reason="Fails on float32")
def test_too_many_params(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[1, 1], y0=[0])
@pytest.mark.xfail(condition=(aesara.config.floatX == "float32"),
reason="Fails on float32")
def test_too_many_y0(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[1], y0=[0, 0])
@pytest.mark.xfail(condition=(aesara.config.floatX == "float32"),
reason="Fails on float32")
def test_too_few_params(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[], y0=[1])
@pytest.mark.xfail(condition=(aesara.config.floatX == "float32"),
reason="Fails on float32")
def test_too_few_y0(self):
with pytest.raises(pm.ShapeError):
self.ode_model(theta=[1], y0=[])
def test_func_callable(self):
with pytest.raises(ValueError):
DifferentialEquation(func=1,
t0=0,
times=self.times,
n_states=1,
n_theta=1)
def test_number_of_states(self):
with pytest.raises(ValueError):
DifferentialEquation(func=self.system,
t0=0,
times=self.times,
n_states=0,
n_theta=1)
def test_number_of_params(self):
with pytest.raises(ValueError):
DifferentialEquation(func=self.system,
t0=0,
times=self.times,
n_states=1,
n_theta=0)
def test_scalar_ode_1_param(self):
"""Test running model for a scalar ODE with 1 parameter"""
def system(y, t, p):
return np.exp(-t) - p[0] * y[0]
times = np.array([
0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5,
7.0, 7.5
])
yobs = np.array([
0.31, 0.57, 0.51, 0.55, 0.47, 0.42, 0.38, 0.3, 0.26, 0.22, 0.22,
0.14, 0.14, 0.09, 0.1
])[:, np.newaxis]
ode_model = DifferentialEquation(func=system,
t0=0,
times=times,
n_states=1,
n_theta=1)
with pm.Model() as model:
alpha = pm.HalfCauchy("alpha", 1)
y0 = pm.Lognormal("y0", 0, 1)
sigma = pm.HalfCauchy("sigma", 1)
forward = ode_model(theta=[alpha], y0=[y0])
y = pm.Lognormal("y",
mu=pm.math.log(forward),
sd=sigma,
observed=yobs)
idata = pm.sample(100, tune=0, chains=1)
assert idata.posterior["alpha"].shape == (1, 100)
assert idata.posterior["y0"].shape == (1, 100)
assert idata.posterior["sigma"].shape == (1, 100)
def inferenceCases(self,S0,I0,R0,observedCumConfirmedCases):
import pymc3 as pm
import theano
import theano.tensor as tt # <-
from theano import printing
from pymc3.ode import DifferentialEquation
from scipy.integrate import odeint
import arviz as az
N = S0+I0+R0
def odem(y,t,p):
# y= [S,I,R,c]
# p= [beta,gamma,delta]
dS = -1.*y[0]*y[1]*p[0]
dI = y[0]*y[1]*p[0] - ( p[1]*y[1] ) # infected
dR = p[1]*y[1]
dc = y[0]*y[1]*p[0] # cumulative cases
return [dS,dI,dR,dc] # last state is cumulative cases
dta = observedCumConfirmedCases.reshape(-1,)
#dta = O[:,-1].reshape(-1,)
timesteps = len(dta)
times = np.arange(0,timesteps)
sir_model = DifferentialEquation(
func = odem,
times = times,
n_states = 4,
n_theta = 2,
t0 = 0,
)
with pm.Model() as model:
beta = pm.Normal('beta' ,2.0,1)
gamma = pm.Normal('gamma' ,2.0,1)
tt.printing.Print('beta')(beta)
tt.printing.Print('gamma')(gamma)
sigmas = pm.Gamma('sigmas', 0.5,0.5, shape=1)
sir_curves = sir_model(y0=[S0,I0,R0,0], theta=[beta,gamma]) # maybe??
obs = pm.Normal("obs", observed = dta , mu = sir_curves[:,-1], sd=sigmas) # maybe?
with model:
MAP = pm.find_MAP()
self.MAP = MAP
self.beta = MAP['beta']
self.gamma = MAP['gamma']
def inferenceCasesAndDeaths(self, S0, I0, R0, D0,
observedNewConfirmedCases, observedNewDeaths):
import pymc3 as pm
import theano
import theano.tensor as tt
from theano import printing
from pymc3.ode import DifferentialEquation
from scipy.integrate import odeint
import arviz as az
def odem(y, t, p):
# y= [S,I,R,D,c,d,r]
# p= [beta,gamma,delta]
dS = -1. * y[0] * y[1] * p[0]
dI = y[0] * y[1] * p[0] - (p[1] * y[1] + p[2] * y[1])
dR = p[1] * y[1]
dD = p[2] * y[1]
dc = y[0] * y[1] * p[0]
return [dS, dI, dR, dD, dc]
dta = np.hstack((observedNewDeaths.reshape(-1, 1),
observedNewConfirmedCases.reshape(-1, 1)))
timesteps = len(dta)
times = np.arange(0, timesteps)
sir_model = DifferentialEquation(
func=odem,
times=times,
n_states=5,
n_theta=3,
t0=0,
)
with pm.Model() as model:
beta = pm.Uniform('beta', 0.0, 10.0)
gamma = pm.Uniform('gamma', 0.0, 10.0)
delta = pm.Uniform('delta', 0.0, 10.0)
sigmas = pm.Gamma('sigmas', 0.5, 0.5, shape=2)
sir_curves = sir_model(y0=[S0, I0, R0, D0, 0],
theta=[beta, gamma, delta])
obs = pm.Normal("obs",
observed=dta,
mu=sir_curves[:, [3, 4]],
sd=sigmas)
with model:
MAP = pm.find_MAP()
self.MAP = MAP
def test_sens_ic_scalar_2_param(self):
# Scalar ODE 2 Param
def ode_func_2(y, t, p):
return p[0] * np.exp(-p[0] * t) - p[1] * y[0]
# Instantiate ODE model
model2 = DifferentialEquation(func=ode_func_2, t0=0, times=self.t, n_states=1, n_theta=2)
model2_sens_ic = np.array([1, 0, 0])
np.testing.assert_array_equal(model2_sens_ic, model2._sens_ic)
def test_sens_ic_vector_2_param(self):
# Vector ODE 2 Param
def ode_func_4(y, t, p):
ds = -p[0] * y[0] * y[1]
di = p[0] * y[0] * y[1] - p[1] * y[1]
return [ds, di]
# Instantiate ODE model
model4 = DifferentialEquation(func=ode_func_4, t0=0, times=self.t, n_states=2, n_theta=2)
model4_sens_ic = np.array([1, 0, 0, 0, 0, 1, 0, 0])
np.testing.assert_array_equal(model4_sens_ic, model4._sens_ic)
def test_sens_ic_scalar_1_param(self):
"""Tests the creation of the initial condition for the sensitivities"""
# Scalar ODE 1 Param
# Create an ODe to integrate
def ode_func_1(y, t, p):
return np.exp(-t) - p[0] * y[0]
# Instantiate ODE model
# Instantiate ODE model
model1 = DifferentialEquation(func=ode_func_1, t0=0, times=self.t, n_states=1, n_theta=1)
# Sensitivity initial condition for this model should be 1 by 2
model1_sens_ic = np.array([1, 0])
np.testing.assert_array_equal(model1_sens_ic, model1._sens_ic)
def test_vector_ode_2_param(self):
"""Test running model for a vector ODE with 2 parameters"""
def system(y, t, p):
ds = -p[0] * y[0] * y[1]
di = p[0] * y[0] * y[1] - p[1] * y[1]
return [ds, di]
times = np.array(
[0.0, 0.8, 1.6, 2.4, 3.2, 4.0, 4.8, 5.6, 6.4, 7.2, 8.0])
yobs = np.array([
[1.02, 0.02],
[0.86, 0.12],
[0.43, 0.37],
[0.14, 0.42],
[0.05, 0.43],
[0.03, 0.14],
[0.02, 0.08],
[0.02, 0.04],
[0.02, 0.01],
[0.02, 0.01],
[0.02, 0.01],
])
ode_model = DifferentialEquation(func=system,
t0=0,
times=times,
n_states=2,
n_theta=2)
with pm.Model() as model:
beta = pm.HalfCauchy("beta", 1)
gamma = pm.HalfCauchy("gamma", 1)
sigma = pm.HalfCauchy("sigma", 1, shape=2)
forward = ode_model(theta=[beta, gamma], y0=[0.99, 0.01])
y = pm.Lognormal("y",
mu=pm.math.log(forward),
sd=sigma,
observed=yobs)
idata = pm.sample(100, tune=0, chains=1)
assert idata.posterior["beta"].shape == (1, 100)
assert idata.posterior["gamma"].shape == (1, 100)
assert idata.posterior["sigma"].shape == (1, 100, 2)
def test_sens_ic_vector_3_params(self):
# Big System with Many Parameters
def ode_func_5(y, t, p):
dx = p[0] * (y[1] - y[0])
ds = y[0] * (p[1] - y[2]) - y[1]
dz = y[0] * y[1] - p[2] * y[2]
return [dx, ds, dz]
# Instantiate ODE model
model5 = DifferentialEquation(func=ode_func_5, t0=0, times=self.t, n_states=3, n_theta=3)
# First three columns are derivatives with respect to ode parameters
# Last three coluimns are derivatives with repsect to initial condition
# So identity matrix should appear in last 3 columns
model5_sens_ic = np.array([[1, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0]])
np.testing.assert_array_equal(np.ravel(model5_sens_ic), model5._sens_ic)
def test_vector_ode_1_param(self):
"""Test running model for a vector ODE with 1 parameter"""
def system(y, t, p):
ds = -p[0] * y[0] * y[1]
di = p[0] * y[0] * y[1] - y[1]
return [ds, di]
times = np.array([0.0, 0.8, 1.6, 2.4, 3.2, 4.0, 4.8, 5.6, 6.4, 7.2, 8.0])
yobs = np.array(
[
[1.02, 0.02],
[0.86, 0.12],
[0.43, 0.37],
[0.14, 0.42],
[0.05, 0.43],
[0.03, 0.14],
[0.02, 0.08],
[0.02, 0.04],
[0.02, 0.01],
[0.02, 0.01],
[0.02, 0.01],
]
)
ode_model = DifferentialEquation(func=system, t0=0, times=times, n_states=2, n_theta=1)
with pm.Model() as model:
R = pm.Lognormal("R", 1, 5)
sigma = pm.HalfCauchy("sigma", 1, shape=2)
forward = ode_model(theta=[R], y0=[0.99, 0.01])
y = pm.Lognormal("y", mu=pm.math.log(forward), sd=sigma, observed=yobs)
trace = pm.sample(100, tune=0, chains=1)
assert trace["R"].size > 0
assert trace["sigma"].size > 0
def get_SIR(x, y, y0, country, forecast_len=0, load_post=False):
'''
If 'forecast_len' is nonzero, attempts to load a trace corresponding to the
country of interest from the directory 'traces' and retrieves predicted numbers
of infected and susceptible patients 'forecast_len' days into the future after the
1st case is detected in the country.
'''
# If in 'prediction mode', modify x, y to reflect forecast length
if forecast_len != 0:
ext = np.arange(1, forecast_len + 1).astype(float)
ext += x[-1]
x = np.append(x, ext)
y = np.empty((x.shape[0], y.shape[1]))
# SIR Model
# p[0]: beta, p[1]: lambda
def SIR(y, t, p):
ds = -p[0] * y[0] * y[1] # Susceptible differential
di = p[0] * y[0] * y[1] - p[1] * y[1] # Infected differential
return [ds, di]
# Initialize ODE
sir_ode = DifferentialEquation(func=SIR,
times=x,
n_states=2,
n_theta=2,
t0=0)
load_dir = osp.join('traces', country.lower())
with pm.Model() as model:
sigma = pm.HalfNormal('sigma', 3, shape=2)
# R0 is bounded below by 1 because we see an epidemic has occured
R0 = pm.Normal('R0', 2, 3)
lmbda = pm.Normal('lambda', 0.1, 0.1)
beta = pm.Deterministic('beta', lmbda * R0)
print('Setting up model for ' + country)
sir_curves = sir_ode(y0=y0, theta=[beta, lmbda])
y_obs = pm.Normal('y_obs', mu=sir_curves, sigma=sigma, observed=y)
if forecast_len == 0:
trace = pm.sample(2000,
tune=1000,
cores=2,
chains=2,
progressbar=True)
# Save trace
pm.save_trace(trace, load_dir, overwrite=True)
# Get the posterior
post = pm.sample_posterior_predictive(trace, progressbar=True)
out_post = post
else:
# Load trace
print('Loading trace')
trace = pm.load_trace(load_dir)
print('Computing posterior')
#Get posterior
if not load_post:
post = pm.sample_posterior_predictive(trace[500:],
progressbar=True)
out_post = post
with open(country + '_post.pkl', 'wb') as buff:
pickle.dump({'post': post}, buff)
else:
with open(country + '_post.pkl', 'rb') as buff:
data = pickle.load(buff)
out_post = data['post']
print('Done')
return trace, out_post, x
# Define ODE
def freefall(y, t, p):
return 2.0 * p[1] - p[0] * y[0]
# Get data to fit
times = np.arange(0, 10, 0.5)
gamma, g, y0, sigma = 0.4, 9.8, -2, 2
y = odeint(freefall, t=times, y0=y0, args=tuple([[gamma, g]]))
yobs = np.random.normal(y, sigma)
# Define inference model
ode_model = DifferentialEquation(func=freefall,
times=times,
n_states=1,
n_theta=2,
t0=0)
with pm.Model() as model:
# priors
sigma = pm.HalfCauchy('sigma', 1)
g = pm.HalfCauchy('g', 1)
gamma = pm.Lognormal('gamma', 0, 1)
# observed
ode_solution = ode_model(y0=[0], theta=[gamma, g])
'''The ode_solution has a shape of (n_times, n_states)'''
Y = pm.Normal('Y', mu=ode_solution, sd=sigma, observed=yobs)
# inference
trace = pm.sample(1000, tune=1000, cores=1)
# prior = pm.sample_prior_predictive()
# posterior_predictive = pm.sample_posterior_predictive(trace)
# s: y[0], e: y[1], i_1: y[2], i_2: y[3], a: y[4]
def SEIR(y, t, p):
ds = -p[0] * y[0] * (y[3] + alpha_1 * y[2] + alpha_2 * y[4]) / N
de = p[0] * y[0] * (y[3] + alpha_1 * y[2] + alpha_2 * y[4]) / N - (1-r) * y[1] / D_e - r * y[1] / D_e
di_1 = (1-r) * y[1] / D_e - y[2] / D_i
di_2 = y[2] / D_i - y[3] / D_r
da = r * y[1] / D_e - y[4] / D_r
return [ds, de, di_1, di_2, da]
seir_model = DifferentialEquation(
func=SEIR,
times=one_estimation_period,
n_states=5, # S, E, I, A, H. (no need to have R)
n_theta=1,
t0=0,
)
with pm.Model() as modelSEIR:
beta1 = pm.Uniform('beta1', 0, 5)
beta2 = pm.Uniform('beta2', 0, 5)
beta3 = pm.Uniform('beta3', 0, 5)
beta4 = pm.Uniform('beta4', 0, 5)
beta5 = pm.Uniform('beta5', 0, 5)
beta6 = pm.Uniform('beta6', 0, 5)
seir_curves1 = seir_model(y0 = [sus0, exp0, inf0, asy0, hos0], theta=[beta1,])
otypes=[t.dmatrix])
def seirdaq_ode_wrapper(time_exp, initial_conditions, beta, gamma_I):
time_span = (time_exp.min(), time_exp.max())
args = [beta, gamma_I]
y_model = seirdaq_ode_solver(initial_conditions, time_span, time_exp,
*args)
simulated_time = y_model.t
simulated_ode_solution = y_model.y
return simulated_ode_solution
seir_jia_model_diffeq = DifferentialEquation(func=seir_jia_model_pymc3,
times=time_range,
n_states=1,
n_theta=13,
t0=0)
# %%
# beta_deterministic, gamma_I_deterministic, gamma_A_deterministic, gamma_D_deterministic, d_I_deterministic, d_D_deterministic = result_seirdaq.x
beta_deterministic, gamma_I_deterministic = result_seirdaq.x
print(
f"parameters: beta = {beta_deterministic}, gamma_I = {gamma_I_deterministic}"
)
# %%
number_of_cores = 24
population_uncertain_variance = 0.05 * np.max(dead_individuals)
variance = population_uncertain_variance * population_uncertain_variance
# Observed system response
Xobs = np.random.normal(X, 0.02)
# Plot the actual and observed system response
plt.plot(timerange, Xobs, marker='o', linestyle='none')
plt.plot(timerange, X, color='C0', alpha=0.5)
plt.legend()
plt.show()
# Define the diferential equation for the pymc3odint module
lmm_model = DifferentialEquation(
func=X_dot, # The DE function, first order
times=np.linspace(0, 0.1,
10), # time range (why not use already defined time?
n_states=4, # Degrees of freedom
n_theta=1, # Model parameters
t0=0, # Start at t=0
)
with pm.Model():
sigma = pm.HalfNormal('sigma', 0.05, shape=4)
#m1 is bounded because mass cannot take negative value
m1 = pm.Bound(pm.Normal, lower=0)('m1', 0.5, 3)
# Set intitial conditions to known values, only parameter is mass
lmm_sol = lmm_model(y0=X0, theta=[m1])
Y = pm.Lognormal('Y', mu=pm.math.log(lmm_sol), sd=sigma, observed=Xobs)
X=np.array([np.array(data['s'])[7:14],
np.array(data['i'])[7:14]]
)
#X=X/pop #code to normalize data to 1
time_range = np.arange(0,len(X[0]))
#%%
# Create differential equation models
model1 = DifferentialEquation(
func = SIR_diffeq, # what is the differential equation?
times = time_range, # what is the time grid for numerical integration?
n_states = 3, # what is the dimensionality of the system?
n_theta = 2, # how many parameters are there?
# t0 = 0 # are we starting at the beginning?
)
model2 = DifferentialEquation(
func = SIR_wDR_diffeq, # what is the differential equation?
times = time_range, # what is the time grid for numerical integration?
n_states = 2, # what is the dimensionality of the system?
n_theta = 3, # how many parameters are there?
# t0 = 0 # are we starting at the beginning?
)
model3=DifferentialEquation(
func = SI_diffeq, # what is the differential equation?
plt.legend()
# plt.show()
# exit()
survivor = ([0] + observations[ObsEnum.RSURVIVOR.value]) - (
observations[ObsEnum.RSURVIVOR.value] + [0])
print(observations[ObsEnum.RSURVIVOR.value])
print("Making the model...")
# This is a DiffEq definition for pymc3.
pm_ode_model = DifferentialEquation(
func=ode_model,
times=[x for x in range(STEPS)],
n_states=10, # dimension of the return value of 'func'
n_theta=10, # dimension of p (additional parameters)
t0=0)
# Now we use some declarations to se up the random
# variables and how they are tied together
basic_model = pm.Model()
with basic_model:
# See : https://docs.pymc.io/notebooks/getting_started.html
# Prior belief (because it's not based on data).
# Bound = put bounds around the parameter we're modelling
# Normal = the parameter is ditributed around some normal distribution
gamma1 = pm.Bound(pm.Normal, lower=1 / 10,
# Observed system response
yobs = np.random.lognormal(mean=np.log(y[1::]), sigma=[0.2, 0.3])
# Plot the actual and observed system response
plt.plot(times[1::], yobs, marker='o', linestyle='none')
plt.plot(times, y[:, 0], color='C0', alpha=0.5, label=f'$S(t)$')
plt.plot(times, y[:, 1], color='C1', alpha=0.5, label=f'$I(t)$')
plt.legend()
plt.show()
# Define the diferential equation for the pymc3odint module
sir_model = DifferentialEquation(
func=SIR, # The DE function, first order
times=np.arange(0.25, 5,
0.25), # time range (why not use already defined time?
n_states=2, # Degrees of freedom
n_theta=2, # Model parameters
t0=0, # Start at t=0
)
with pm.Model():
sigma = pm.HalfCauchy('sigma', 1, shape=2)
# R0 is bounded below by 1 because we see an epidemic has occured
R0 = pm.Bound(pm.Normal, lower=1)('R0', 2, 3)
lam = pm.Lognormal('lambda', pm.math.log(2), 2)
beta = pm.Deterministic('beta', lam * R0)
sir_curves = sir_model(y0=[0.99, 0.01], theta=[beta, lam])
Y = pm.Lognormal('Y',
ds = -p[0] * y[0] * y[1]
di = p[0] * y[0] * y[1] - p[1] * y[1]
return [ds, di]
# Get data to fit
times = np.arange(0, 5, 0.25)
beta, gamma = 4, 1.0
y = odeint(SIR, t=times, y0=[0.99, 0.01], args=((beta, gamma), ), rtol=1e-8)
yobs = np.random.lognormal(mean=np.log(y[1::]), sigma=[0.2, 0.3])
# Define inference model
sir_model = DifferentialEquation(
func=SIR,
times=np.arange(0.25, 5, 0.25),
n_states=2,
n_theta=2,
t0=0,
)
with pm.Model() as model:
# priors
sigma = pm.HalfCauchy('sigma', 1, shape=2)
R0 = pm.Bound(pm.Normal, lower=1)(
'R0', 2,
3) # R0 is bounded below by 1 because we see an epidemic has occured
lam = pm.Lognormal('lambda', pm.math.log(2), 2)
beta = pm.Deterministic('beta', lam * R0)
# observed
sir_curves = sir_model(y0=[0.99, 0.01], theta=[beta, lam])
Y = pm.Lognormal('Y', mu=pm.math.log(sir_curves), sd=sigma, observed=yobs)
# inference
