def test_stochastic_well_mixed(self):
S, E, I, R = list("SEIR")
N = 75000
tmax = 100
model = EpiModel([S, E, I, R], N)
model.set_processes([
(S, I, 2, E, I),
(I, 1, R),
(E, 1, I),
])
model.set_initial_conditions({S: N - 100, I: 100})
tt = np.linspace(0, tmax, 10000)
result_int = model.integrate(tt)
t, result_sim = model.simulate(tmax,
sampling_dt=1,
return_compartments=[S, R])
model = StochasticEpiModel([S, E, I, R], N)
model.set_link_transmission_processes([
(I, S, 2, I, E),
])
model.set_node_transition_processes([
(I, 1, R),
(E, 1, I),
])
model.set_random_initial_conditions({S: N - 100, I: 100})
t, result_sim2 = model.simulate(tmax,
sampling_dt=1,
return_compartments=[S, R])
for c, res in result_sim2.items():
#print(c, np.abs(1-res[-1]/result_int[c][-1]))
#print(c, np.abs(1-res[-1]/result_sim[c][-1]))
assert (np.abs(1 - res[-1] / result_int[c][-1]) < 0.05)
assert (np.abs(1 - res[-1] / result_sim[c][-1]) < 0.05)
t, result_sim = model.simulate(tmax,sampling_dt=1)
end = time()
print("numeric model needed", end-start, "s")
for c, res in result_sim.items():
pl.plot(t, res, '--')
model = StochasticEpiModel([S,E,I,R],N)
model.set_link_transmission_processes([
( I, S, 2, I, E ),
])
model.set_node_transition_processes([
( I, 1, R),
( E, 1, I),
])
model.set_random_initial_conditions({S: N-100, I: 100})
start = time()
t, result_sim = model.simulate(tmax,sampling_dt=1)
end = time()
for c, res in result_sim.items():
pl.plot(t, res, ':')
print("stochastic model needed", end-start, "s")
pl.show()