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_sojourn_flow__expect_flows_applied( inf_pop, exp_flow): """ Expect a flow to occur proportional to the compartment size and parameter. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population(distribution={"I": inf_pop}) model.add_sojourn_flow("recovery", 10, "I", "R") model._prepare_to_run() actual_flow_rates = model._get_compartment_rates(model.initial_population, 0) # Expect exp_flow = inf_pop * () exp_flow expected_flow_rates = np.array([0, -exp_flow, exp_flow]) assert_array_equal(actual_flow_rates, expected_flow_rates)