def __init__(self,
compartments,
initial_population_size=1,
correct_for_dynamical_population_size=False):
"""
"""
MatrixEpiModel.__init__(self, compartments, initial_population_size,
correct_for_dynamical_population_size)
self.t = sympy.symbols("t")
if self.t in self.compartments:
raise ValueError(
"Don't use `t` as a compartment symbol, as it is reserved for time."
)
self.has_functional_rates = False
self.birth_rates = sympy.zeros(self.N_comp, 1)
self.linear_rates = sympy.zeros(self.N_comp, self.N_comp)
self.quadratic_rates = [ sympy.zeros(self.N_comp, self.N_comp)\
for c in range(self.N_comp) ]
self.birth_events = sympy.zeros(self.N_comp, 1)
self.linear_events = sympy.zeros(self.N_comp, self.N_comp)
self.quadratic_events = [ sympy.zeros(self.N_comp, self.N_comp)\
for c in range(self.N_comp) ]
self.parameter_values = {}
def test_fusion_and_adding_rates(self):
A, B, C = list("ABC")
epi = MatrixEpiModel(list("ABC"))
# this should not raise a warning that rates do not sum to zero
# as it will be actively suppressed
epi.add_fusion_processes([
(A, B, 1, C),
])
with self.assertWarns(UserWarning):
# this should raise a warning that rates do not sum to zero
epi.add_quadratic_rates([
(A, B, C, -1),
(A, B, A, +1),
])
# now rates should sum to zero
epi.add_quadratic_rates([
(A, B, B, +1),
])
with self.assertWarns(UserWarning):
# this should raise a warning that rates do not sum to zero
epi.add_linear_rates([
(A, B, -1)
])
def test_quadratic_processes(self):
epi = MatrixEpiModel(list("SEIR"))
Q = [ np.zeros((4,4)) for C in epi.compartments ]
Q[0][0,2] = -1
Q[1][0,2] = +1
epi.add_transmission_processes([
("S", "I", 1.0, "I", "E"),
])
for iM, M in enumerate(epi.quadratic_rates):
assert(np.all(M.toarray().flatten()==Q[iM].flatten()))
def test_linear_rates(self):
epi = MatrixEpiModel(list("SEIR"))
epi.add_transition_processes([
("E", 1.0, "I"),
("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_initial_condition_warnings(self):
A, B, C = list("ABC")
epi = MatrixEpiModel(list("ABC"))
with self.assertWarns(UserWarning):
# this should raise a warning that rates do not sum to zero
epi.set_initial_conditions({A:0.1,B:0.2})
with self.assertWarns(UserWarning):
# this should raise a warning that initial conditions were set twice
epi.set_initial_conditions([(A,0.1),(A,0.2)])
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))
def test_compartments(self):
epi = MatrixEpiModel(list("SEIR"))
assert(all([ i == epi.get_compartment_id(C) for i, C in enumerate("SEIR") ]))
assert(epi.get_compartment_id("E") == 1)
assert(epi.get_compartment(1) == "E")