def test_adding_quadratic_processes(self):
epi = MatrixEpiModel(list("SEIAR"))
Q = [ np.zeros((5,5)) for C in epi.compartments ]
Q[0][0,2] = -1
Q[0][0,3] = -1
Q[1][0,2] = +1
Q[1][0,3] = +1
epi.set_processes([
("S", "I", 1.0, "I", "E"),
])
epi.add_transmission_processes([
("S", "A", 1.0, "A", "E"),
])
for iM, M in enumerate(epi.quadratic_rates):
assert(np.all(M.toarray().flatten()==Q[iM].flatten()))
def test_adding_linear_rates(self):
epi = MatrixEpiModel(list("SEIR"))
epi.set_processes([
("E", 1.0, "I"),
])
epi.add_transition_processes([
("I", 1.0, "R"),
])
linear_rates = np.zeros((4,4))
linear_rates[1,1] = -1
linear_rates[2,1] = +1
linear_rates[2,2] = -1
linear_rates[3,2] = +1
assert(np.all(epi.linear_rates.toarray().flatten()==linear_rates.flatten()))
def test_birth_death(self):
epi = MatrixEpiModel(list("SIR"))
R0 = 2
rho = 1
mu = 0.2
eta = R0 * rho
with self.assertWarns(UserWarning):
epi.set_processes([
("S", "I", eta, "I", "I"),
("I", rho, "R"),
(None, mu, "S"),
("S", mu, None),
("R", mu, None),
("I", mu, None),
])
epi.set_initial_conditions({'S': 0.8, 'I':0.2 })
t = [0,1000]
res = epi.integrate(t)
assert(np.isclose(res['S'][-1],(mu+rho)/eta))
assert(np.isclose(res['I'][-1],mu/eta*(eta-mu-rho)/(mu+rho)))
def test_R0(self):
"""
R0 computation test based on
"The construction of next-generation matrices for compartmental epidemic models"
by Diekmann, Heesterbeek, Roberts, J. R. Soc. Interface (2010) 7, 873–885
doi.org/10.1098/rsif.2009.0386
(Section 2.2)
"""
b11 = 1
b12 = 2
b21 = 3
b22 = 6
#we'll use b11*b22 = b12*b21 because it makes
#our lives easier in computing R0 analytically
mu = 4
g1 = 5
g2 = 7
nu1 = 8
nu2 = 9
p = 0.2
N = 100
base_comps = list("SEIR")
cats = [1,2]
comps = []
for C in base_comps:
for c in cats:
comps.append(C+str(c))
model = MatrixEpiModel(comps,initial_population_size=N)
transition_processes = [
("S1", mu, None),
("E1", mu, None),
("I1", mu, None),
("R1", mu, None),
("S2", mu, None),
("E2", mu, None),
("I2", mu, None),
("R2", mu, None),
("E1", nu1, "I1"),
("E2", nu2, "I2"),
("I1", g1, "R1"),
("I2", g2, "R2"),
]
transmission_processes = [
( "S1", "I1", b11, "E1", "I1"),
( "S1", "I2", b12, "E1", "I2"),
( "S2", "I1", b21, "E2", "I1"),
( "S2", "I2", b22, "E2", "I2"),
]
birth_processes =[
(None, p*mu*N, "S1"),
(None, (1-p)*mu*N, "S2"),
]
model.set_processes(transition_processes + \
transmission_processes + \
birth_processes,
allow_nonzero_column_sums=True,
)
model.set_initial_conditions({
'S1': p*N,
'S2': (1-p)*N,
})
R0 = model.get_next_generation_matrix_leading_eigenvalue()
R0theory = p*b11*nu1/(nu1 + mu)/(g1+mu) + (1-p)*b22*nu2/(nu2+mu)/(g2+mu)
assert(np.isclose(R0,R0theory))