def test_will_update_population_on_timestep(self):
population = FixedPopulationModel()
susceptable = population.create_compartment('susceptable',
initial=1000)
infected = population.create_compartment('infected', initial=1)
model = TemporalDiseaseModel(population, 1)
model.add_transition(susceptable, infected, Constant(1))
model.trace(1, 1)
self.assertEqual(susceptable.size.value, 999)
self.assertEqual(infected.size.value, 2)
def test_can_trace_several_timesteps(self):
population = FixedPopulationModel()
susceptable = population.create_compartment('susceptable',
initial=1000)
infected = population.create_compartment('infected', initial=1)
model = TemporalDiseaseModel(population, 1)
dt = model.dt_expression
model.add_transition(susceptable, infected, dt)
self.assertEqual(model.trace(3, 1),
[['t', 'susceptable', 'infected'], [0, 1000, 1],
[1, 999, 2], [2, 998, 3]])
def test_can_update_population_on_based_on_dt(self):
population = FixedPopulationModel()
susceptable = population.create_compartment('susceptable',
initial=1000)
infected = population.create_compartment('infected', initial=1)
model = TemporalDiseaseModel(population, 4)
dt = model.dt_expression
model.add_transition(susceptable, infected, Constant(2) * dt)
model.trace(1, 1)
self.assertEqual(susceptable.size.value, 992)
self.assertEqual(infected.size.value, 9)
def test_will_error_gracefully_when_not_parameterised(self):
population = FixedPopulationModel()
susceptable = population.create_compartment('susceptable',
initial=1000)
infected = population.create_compartment('infected', initial=1)
model = TemporalDiseaseModel(population, 1)
alpha = model.create_parameter('alpha')
model.add_transition(susceptable, infected, alpha)
with self.assertRaises(ValueError):
model.trace(1, 1)
alpha.set_value(1)
model.trace(1, 1)
def test_will_transition_correctly_at_zero(self):
population = FixedPopulationModel()
s = population.create_compartment('susceptable', initial=200)
i = population.create_compartment('infected', initial=200)
population.initialise_expressions(1)
population.move(s, i, 300)
self.assertEqual(s.size.value, 0)
self.assertEqual(i.size.value, 400)
from connectors.cli import cli_executor
from core.models.disease import TemporalDiseaseModel
from core.models.population import FixedPopulationModel
from core.expressions import Binomial
# set up our fixed population of S I R
population = FixedPopulationModel()
susceptable = population.create_compartment('susceptable', initial=1000)
infected = population.create_compartment('infected', initial=1)
recovered = population.create_compartment('recovered')
# set up our stochastic transformation of the infecious disease model:
# dS/dt = beta * S * I / N
# dI/dt = beta * S * I / N - sigma * I
# dR/dt = sigma * I
model = TemporalDiseaseModel(population, .1)
beta = model.create_parameter('beta')
sigma = model.create_parameter('sigma')
dt = model.dt_expression
infection_expression = Binomial(susceptable.size,
beta * (infected.size / population.size) * dt)
recovery_expression = Binomial(infected.size, sigma * dt)
model.add_transition(susceptable, infected, infection_expression)
model.add_transition(infected, recovered, recovery_expression)
if __name__ == '__main__':
cli_executor(model)
def test_cannot_create_invalid_compartments(self):
population = FixedPopulationModel()
population.initialise_expressions(1)
with self.assertRaises(ValueError):
population.create_compartment('susceptable', initial=-200)
def test_reports_the_correct_size(self):
population = FixedPopulationModel()
s = population.create_compartment('susceptable', initial=200)
i = population.create_compartment('infected', initial=100)
population.initialise_expressions(1)
self.assertEqual(population.size.evaluate(), 300)