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_flows__with_function_flow__expect_flows_applied(
inf_pop, sus_pop, exp_flow):
"""
Expect a flow to occur proportional to the result of `get_flow_rate`.
"""
def get_flow_rate(flow, comps, comp_vals, flows, flow_rates, time):
_, i_pop, _ = comp_vals
i_flow, _ = flow_rates
return i_pop + i_flow
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": sus_pop, "I": inf_pop})
model.add_transition_flow("infection", 0.1, "S", "I")
model.add_function_flow("treatment", get_flow_rate, "I", "S")
model._prepare_to_run()
actual_flow_rates = model._get_compartment_rates(model.initial_population,
0)
expected_infected = sus_pop * 0.1
expected_flow_rates = np.array([
exp_flow - expected_infected,
expected_infected - exp_flow,
0,
])
assert_array_equal(actual_flow_rates, expected_flow_rates)
def test_apply_flows__with_transition_flow__expect_flows_applied(
inf_pop, sus_pop, exp_flow):
"""
Expect a flow to occur proportional to the compartment size and parameter.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": sus_pop, "I": inf_pop})
model.add_transition_flow("deliberately_infected", 0.1, "S", "I")
model._prepare_to_run()
actual_flow_rates = model._get_compartment_rates(model.initial_population,
0)
# Expect sus_pop * 0.1 = exp_flow
expected_flow_rates = np.array([-exp_flow, exp_flow])
assert_array_equal(actual_flow_rates, expected_flow_rates)
def test_flow_derived_outputs():
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 900, "I": 100})
# Constant entry.
model.add_importation_flow("imports", num_imported=2, dest="S")
model.request_output_for_flow(name="importation", flow_name="imports")
model.request_output_for_flow(name="importation_raw", flow_name="imports", raw_results=True)
# Linear entry.
model.add_importation_flow("imports_land", num_imported=lambda t: 3 * t, dest="S")
model.request_output_for_flow(name="importation_land", flow_name="imports_land")
# Quadratic entry.
model.add_importation_flow("imports_air", num_imported=lambda t: t ** 2, dest="S")
model.request_output_for_flow(name="importation_air", flow_name="imports_air")
# Fractional transition flow
model.add_transition_flow("recovery", 0.1, "I", "R")
model.request_output_for_flow(name="recovery_raw", flow_name="recovery", raw_results=True)
model.request_output_for_flow(name="recovery_delta", flow_name="recovery", raw_results=False)
model.run()
dos = model.derived_outputs
# Raw outputs are the instantaneous flow rate at a given time.
assert_allclose(dos["recovery_raw"], model.outputs[:, 1] * 0.1, rtol=0.01)
# Post-processed outputs reflect changes in compartment size.
recovered_count = np.zeros_like(model.outputs[:, 2])
recovered_count[1:] = np.diff(model.outputs[:, 2])
assert_allclose(dos["recovery_delta"][1:], recovered_count[1:], rtol=0.01)
# Good match for constant
assert_array_equal(dos["importation"][1:], np.array([2, 2, 2, 2, 2]))
assert_array_equal(dos["importation_raw"], np.array([2, 2, 2, 2, 2, 2]))
# Good match for linear
assert_allclose(dos["importation_land"][1:], np.array([1.5, 4.5, 7.5, 10.5, 13.5]))
# So-so match for quadratic
assert_allclose(dos["importation_air"][1:], np.array([0.5, 2.5, 6.5, 12.5, 20.5]), rtol=0.1)
