def test_apply_many_flows():
"""
Expect multiple flows to operate independently and produce the correct final flow rate.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
model.add_death_flow("infection_death", 0.1, "I")
model.add_universal_death_flows("universal_death", 0.1)
model.add_infection_frequency_flow("infection", 0.2, "S", "I")
model.add_sojourn_flow("recovery", 10, "I", "R")
model.add_transition_flow("vaccination", 0.1, "S", "R")
model.add_crude_birth_flow("births", 0.1, "S")
model._prepare_to_run()
actual_flow_rates = model._get_compartment_rates(model.initial_population,
0)
# Expect the effects of all these flows to be linearly superimposed.
infect_death_flows = np.array([0, -10, 0])
universal_death_flows = np.array([-90, -10, 0])
infected = 900 * 0.2 * (100 / 1000)
infection_flows = np.array([-infected, infected, 0])
recovery_flows = np.array([0, -10, 10])
vaccination_flows = np.array([-90, 0, 90])
birth_flows = np.array([100, 0, 0])
expected_flow_rates = (infect_death_flows + universal_death_flows +
infection_flows + recovery_flows +
vaccination_flows + birth_flows)
assert_array_equal(actual_flow_rates, expected_flow_rates)
def test_apply_universal_death_flow():
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 990, "I": 10})
model.add_universal_death_flows("universal_death", 0.1)
model._prepare_to_run()
actual_flow_rates = model._get_compartment_rates(model.initial_population,
0)
expected_flow_rates = np.array([-99, -1])
assert_array_equal(actual_flow_rates, expected_flow_rates)
def test_apply_replace_death_birth_flow():
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
model.add_death_flow("infection_death", 0.1, "I")
model.add_universal_death_flows("universal_death", 0.05)
model.add_replacement_birth_flow("births", "S")
model._prepare_to_run()
actual_flow_rates = model._get_compartment_rates(model.initial_population,
0)
# Expect 10 people to die and 10 to be born
exp_i_flow_rate = -0.1 * 100 - 0.05 * 100
exp_s_flow_rate = -exp_i_flow_rate # N.B births + deaths in the S compartment should balance.
expected_flow_rates = np.array([exp_s_flow_rate, exp_i_flow_rate])
assert_array_equal(actual_flow_rates, expected_flow_rates)