def test_stratify_transition_flows__with_dest_only_stratified(): """ Ensure transition flows are stratified correctly when only the flow destination is stratified. Expect an person-conserving adjustment of 1/N to be applied to each flow - N being the number of new strata. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_sojourn_flow("recovery", 7, "I", "R") expected_flows = [ SojournFlow("recovery", C("I"), C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) # Apply stratification strat = Stratification("location", ["urban", "rural"], ["R"]) model.stratify_with(strat) expected_flows = [ SojournFlow("recovery", C("I"), C("R", {"location": "urban"}), 7, [Multiply(0.5)]), SojournFlow("recovery", C("I"), C("R", {"location": "rural"}), 7, [Multiply(0.5)]), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_strat_model__with_age_and_starting_proportion__expect_ageing(): """ Ensure that a module with age stratification and starting proporptions produces ageing flows, and the correct output. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I"], infectious_compartments=["I"]) model.set_initial_population(distribution={"S": 1000, "I": 0}) strat = AgeStratification("age", [0, 5, 15, 60], ["S", "I"]) strat.set_population_split({"0": 0.8, "5": 0.1, "15": 0.1, "60": 0}) model.stratify_with(strat) # Run the model for 5 years. model.run() # Expect everyone to generally get older, but no one should die or get sick. # Expect initial distribution of ages to be set according to "requested_proportions". expected_arr = np.array( [ [800.0, 100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0], [655.0, 228.3, 114.4, 2.4, 0.0, 0.0, 0.0, 0.0], [536.2, 319.4, 139.2, 5.2, 0.0, 0.0, 0.0, 0.0], [439.0, 381.3, 171.1, 8.6, 0.0, 0.0, 0.0, 0.0], [359.5, 420.6, 207.1, 12.8, 0.0, 0.0, 0.0, 0.0], [294.3, 442.5, 245.4, 17.8, 0.0, 0.0, 0.0, 0.0], ] ) assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_add_entry_flows_post_stratification(): """ Ensure we can add flows after a model is stratified. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) assert len(model._flows) == 0 strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) model.stratify_with(strat) with pytest.raises(AssertionError): model.add_importation_flow("imports", 10, "S", expected_flow_count=1) assert len(model._flows) == 0 model.add_importation_flow("imports", 10, "S", expected_flow_count=2) assert len(model._flows) == 2 expected_flows = [ ImportFlow("imports", C("S", {"location": "urban"}), 10, []), ImportFlow("imports", C("S", {"location": "rural"}), 10, []), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_stratify_transition_flows__with_dest_only_stratified__with_adjustments_and_strains(): """ Ensure transition flows are stratified correctly when only the flow destination is stratified. Expect adjustments to override the automatic person-conserving adjustment when using a strain strat. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_sojourn_flow("recovery", 7, "I", "R") expected_flows = [ SojournFlow("recovery", C("I"), C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) # Apply stratification strat = StrainStratification("location", ["urban", "rural"], ["R"]) strat.add_flow_adjustments("recovery", {"urban": Overwrite(0.7), "rural": Overwrite(0.1)}) model.stratify_with(strat) expected_flows = [ SojournFlow("recovery", C("I"), C("R", {"location": "urban"}), 7, [Overwrite(0.7)]), SojournFlow("recovery", C("I"), C("R", {"location": "rural"}), 7, [Overwrite(0.1)]), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_stratify__age__validate_ageing_flows_added(): """ Ensure, when using an age stratification, that ageing flows are automatically added and that birth flows are all sent to age 0. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) assert len(model._flows) == 0 model.add_crude_birth_flow("births", 0.02, "S") assert len(model._flows) == 1 strat = AgeStratification("age", ["0", "5", "15"], ["S", "I", "R"]) model.stratify_with(strat) # Expect ageing flows amongst age group and a birth flow that only goes to age 0. expected_flows = [ CrudeBirthFlow("births", C("S", {"age": "0"}), 0.02), SojournFlow("ageing_SXage_0_to_SXage_5", C("S", {"age": "0"}), C("S", {"age": "5"}), 5), SojournFlow("ageing_IXage_0_to_IXage_5", C("I", {"age": "0"}), C("I", {"age": "5"}), 5), SojournFlow("ageing_RXage_0_to_RXage_5", C("R", {"age": "0"}), C("R", {"age": "5"}), 5), SojournFlow("ageing_SXage_5_to_SXage_15", C("S", {"age": "5"}), C("S", {"age": "15"}), 10), SojournFlow("ageing_IXage_5_to_IXage_15", C("I", {"age": "5"}), C("I", {"age": "15"}), 10), SojournFlow("ageing_RXage_5_to_RXage_15", C("R", {"age": "5"}), C("R", {"age": "15"}), 10), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_single_static_mixing_matrix(): """ Test that we are using the correct mixing matrix when we have a single static mixing matrix """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) # Apply first stratification with a mixing matrix. strat = Stratification(name="agegroup", strata=["child", "adult"], compartments=["S", "I", "R"]) mixing_matrix = np.array([[2, 3], [5, 7]]) strat.set_mixing_matrix(mixing_matrix) model.stratify_with(strat) # We should get the default mixing matrix actual_mixing = model._get_mixing_matrix(0) assert_array_equal(actual_mixing, mixing_matrix) # Static matrices shouldn't change over time actual_mixing = model._get_mixing_matrix(123) assert_array_equal(actual_mixing, mixing_matrix) # Agegroup mixing categories have been added. assert model._mixing_categories == [{ "agegroup": "child" }, { "agegroup": "adult" }]
def test_stratify_transition_flows__with_source_only_stratified(): """ Ensure transition flows are stratified correctly when only the flow source is stratified. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_sojourn_flow("recovery", 7, "I", "R") expected_flows = [ SojournFlow("recovery", C("I"), C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) # Apply stratification strat = Stratification("location", ["urban", "rural"], ["S", "I"]) model.stratify_with(strat) expected_flows = [ SojournFlow("recovery", C("I", {"location": "urban"}), C("R"), 7), SojournFlow("recovery", C("I", {"location": "rural"}), C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_strat_model__with_age__expect_ageing(): """ Ensure that a module with age stratification produces ageing flows, and the correct output. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I"], infectious_compartments=["I"]) model.set_initial_population(distribution={"S": 1000, "I": 0}) strat = AgeStratification("age", [0, 5, 15, 60], ["S", "I"]) model.stratify_with(strat) # Run the model for 5 years. model.run() # Expect everyone to generally get older, but no one should die or get sick expected_arr = np.array( [ [250.0, 250.0, 250.0, 250.0, 0.0, 0.0, 0.0, 0.0], [204.7, 269.3, 270.3, 255.8, 0.0, 0.0, 0.0, 0.0], [167.6, 278.9, 291.5, 262.0, 0.0, 0.0, 0.0, 0.0], [137.2, 281.2, 312.8, 268.7, 0.0, 0.0, 0.0, 0.0], [112.3, 278.1, 333.7, 275.9, 0.0, 0.0, 0.0, 0.0], [92.0, 270.9, 353.5, 283.6, 0.0, 0.0, 0.0, 0.0], ] ) assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_add_exit_flows_post_stratification(): """ Ensure user can add exit flows post stratification. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) assert len(model._flows) == 0 # Apply partial stratification strat = Stratification("location", ["urban", "rural"], ["S", "I"]) model.stratify_with(strat) assert len(model._flows) == 0 model.add_death_flow("d_S", 3, "S") model.add_death_flow("d_I", 5, "I") model.add_death_flow("d_R", 7, "R") expected_flows = [ DeathFlow("d_S", C("S", {"location": "urban"}), 3), DeathFlow("d_S", C("S", {"location": "rural"}), 3), DeathFlow("d_I", C("I", {"location": "urban"}), 5), DeathFlow("d_I", C("I", {"location": "rural"}), 5), DeathFlow("d_R", C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_full_age_strat(): model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) age_strat = AgeStratification(name="age", strata=["0", "5", "10"], compartments=["S", "I", "R"]) model.stratify_with(age_strat)
def test_repeat_age_strat_fails(): model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) age_strat = AgeStratification(name="age", strata=["0", "5", "10"], compartments=["S", "I", "R"]) model.stratify_with(age_strat) with pytest.raises(AssertionError): model.stratify_with(age_strat)
def test_stratify_entry_flows__with_no_explicit_adjustments(): """ Ensure entry flows are stratified correctly when no adjustments are requested. Expect flow to be conserved, split evenly over the new strata. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_importation_flow("imports", 10, "S") expected_flows = [ImportFlow("imports", C("S"), 10)] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) model.stratify_with(strat) expected_flows = [ ImportFlow("imports", C("S", {"location": "urban"}), 10, [Multiply(0.5)]), ImportFlow("imports", C("S", {"location": "rural"}), 10, [Multiply(0.5)]), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) strat = Stratification("age", ["young", "old"], ["S", "I", "R"]) model.stratify_with(strat) expected_flows = [ ImportFlow( "imports", C("S", {"location": "urban", "age": "young"}), 10, [Multiply(0.5), Multiply(0.5)], ), ImportFlow( "imports", C("S", {"location": "urban", "age": "old"}), 10, [Multiply(0.5), Multiply(0.5)], ), ImportFlow( "imports", C("S", {"location": "rural", "age": "young"}), 10, [Multiply(0.5), Multiply(0.5)], ), ImportFlow( "imports", C("S", {"location": "rural", "age": "old"}), 10, [Multiply(0.5), Multiply(0.5)], ), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_repeat_strat_including_age(): model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) age_strat = AgeStratification(name="age", strata=["0", "5", "10"], compartments=["S", "I", "R"]) other_strat = Stratification(name="gender", strata=["female", "male"], compartments=["S", "I", "R"]) model.stratify_with(age_strat) model.stratify_with(other_strat)
def test_stratify_exit_flows(): """ Ensure exit flows are stratified correctly. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_death_flow("d_S", 3, "S") model.add_death_flow("d_I", 5, "I") model.add_death_flow("d_R", 7, "R") expected_flows = [ DeathFlow("d_S", C("S"), 3), DeathFlow("d_I", C("I"), 5), DeathFlow("d_R", C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) # Apply partial stratification strat = Stratification("location", ["urban", "rural"], ["S", "I"]) model.stratify_with(strat) expected_flows = [ DeathFlow("d_S", C("S", {"location": "urban"}), 3), DeathFlow("d_S", C("S", {"location": "rural"}), 3), DeathFlow("d_I", C("I", {"location": "urban"}), 5), DeathFlow("d_I", C("I", {"location": "rural"}), 5), DeathFlow("d_R", C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) # Apply partial stratification with flow adjustments strat = Stratification("age", ["young", "old"], ["I", "R"]) strat.add_flow_adjustments("d_I", {"young": Multiply(0.5), "old": Multiply(2)}) strat.add_flow_adjustments("d_R", {"young": Multiply(0.5), "old": Multiply(2)}) model.stratify_with(strat) expected_flows = [ DeathFlow("d_S", C("S", {"location": "urban"}), 3), DeathFlow("d_S", C("S", {"location": "rural"}), 3), DeathFlow("d_I", C("I", {"location": "urban", "age": "young"}), 5, [Multiply(0.5)]), DeathFlow("d_I", C("I", {"location": "urban", "age": "old"}), 5, [Multiply(2)]), DeathFlow("d_I", C("I", {"location": "rural", "age": "young"}), 5, [Multiply(0.5)]), DeathFlow("d_I", C("I", {"location": "rural", "age": "old"}), 5, [Multiply(2)]), DeathFlow("d_R", C("R", {"age": "young"}), 7, [Multiply(0.5)]), DeathFlow("d_R", C("R", {"age": "old"}), 7, [Multiply(2)]), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_model__with_two_symmetric_stratifications(): """ Adding two strata with the same properties should yield the exact same infection dynamics and outputs as having no strata at all. This does not test strains directly, but if this doesn't work then further testing is pointless. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population(distribution={"S": 900, "I": 100}) model.add_infection_frequency_flow("infection", 0.2, "S", "I") model.add_sojourn_flow("recovery", 10, "I", "R") # Do pre-run force of infection calcs. model._prepare_to_run() model._prepare_time_step(0, model.initial_population) # Check infectiousness multipliers susceptible = model.compartments[0] infectious = model.compartments[1] assert model._get_infection_density_multiplier(susceptible, infectious) == 100.0 assert model._get_infection_frequency_multiplier(susceptible, infectious) == 0.1 model.run() # Create a stratified model where the two non-strain strata are symmetric stratified_model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) stratified_model.set_initial_population(distribution={"S": 900, "I": 100}) stratified_model.add_infection_frequency_flow("infection", 0.2, "S", "I") stratified_model.add_sojourn_flow("recovery", 10, "I", "R") strat = Stratification("clinical", ["home", "hospital"], ["I"]) stratified_model.stratify_with(strat) stratified_model.run() # Ensure stratified model has the same results as the unstratified model. merged_outputs = np.zeros_like(model.outputs) merged_outputs[:, 0] = stratified_model.outputs[:, 0] merged_outputs[:, 1] = stratified_model.outputs[:, 1] + stratified_model.outputs[:, 2] merged_outputs[:, 2] = stratified_model.outputs[:, 3] assert_allclose(merged_outputs, model.outputs, atol=0.01, rtol=0.01, verbose=True)
def test_add_mixing_matrix_fails(): """ Ensure validation works when trying to add a mixing matrix. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) strat = Stratification(name="agegroup", strata=["child", "adult"], compartments=["S", "R"]) mixing_matrix = np.array([[2, 3], [5, 7]]) strat.set_mixing_matrix(mixing_matrix) # Expect this to fail because it's not a full stratification (no I compartment). with pytest.raises(AssertionError): model.stratify_with(strat)
def test_strains__with_two_symmetric_strains(): """ Adding two strains with the same properties should yield the same infection dynamics and outputs as having no strains at all. We expect the force of infection for each strain to be 1/2 of the unstratified model, but the stratification process will not apply the usual conservation fraction to the fan out flows. """ # Create an unstratified model model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population(distribution={"S": 900, "I": 100}) model.add_infection_frequency_flow("infection", 0.2, "S", "I") model.add_sojourn_flow("recovery", 10, "I", "R") # Do pre-run force of infection calcs. model._prepare_to_run() model._prepare_time_step(0, model.initial_population) # Check infectiousness multipliers susceptible = model.compartments[0] infectious = model.compartments[1] assert model._get_infection_density_multiplier(susceptible, infectious) == 100.0 assert model._get_infection_frequency_multiplier(susceptible, infectious) == 0.1 model.run() # Create a stratified model where the two strain strata are symmetric strain_model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) strain_model.set_initial_population(distribution={"S": 900, "I": 100}) strain_model.add_infection_frequency_flow("infection", 0.2, "S", "I") strain_model.add_sojourn_flow("recovery", 10, "I", "R") strat = StrainStratification("strain", ["a", "b"], ["I"]) strain_model.stratify_with(strat) strain_model.run() # Ensure stratified model has the same results as the unstratified model. merged_outputs = np.zeros_like(model.outputs) merged_outputs[:, 0] = strain_model.outputs[:, 0] merged_outputs[:, 1] = strain_model.outputs[:, 1] + strain_model.outputs[:, 2] merged_outputs[:, 2] = strain_model.outputs[:, 3] assert_allclose(merged_outputs, model.outputs, atol=0.01, rtol=0.01, verbose=True)
def test_no_mixing_matrix__with_previous_strat(): """ Test that we are using the default 'null-op' mixing matrix when we have a no user-specified mixing matrix and a stratification has already been applied """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) # Apply first stratification with a mixing matrix. strat = Stratification(name="agegroup", strata=["child", "adult"], compartments=["S", "I", "R"]) first_strat_matrix = np.array([[2, 3], [5, 7]]) strat.set_mixing_matrix(first_strat_matrix) model.stratify_with(strat) # We should get the default mixing matrix actual_mixing = model._get_mixing_matrix(0) assert_array_equal(actual_mixing, first_strat_matrix) # Static matrices shouldn't change over time actual_mixing = model._get_mixing_matrix(123) assert_array_equal(actual_mixing, first_strat_matrix) # Agegroup mixing categories have been added. assert model._mixing_categories == [{ "agegroup": "child" }, { "agegroup": "adult" }] # Apply second stratification with no mixing matrix. strat = Stratification(name="location", strata=["work", "home"], compartments=["S", "I", "R"]) model.stratify_with(strat) # We should get the same results as before. actual_mixing = model._get_mixing_matrix(0) assert_array_equal(actual_mixing, first_strat_matrix) actual_mixing = model._get_mixing_matrix(123) assert_array_equal(actual_mixing, first_strat_matrix) assert model._mixing_categories == [{ "agegroup": "child" }, { "agegroup": "adult" }]
def test_stratify_transition_flows__with_source_and_dest_stratified(): """ Ensure transition flows are stratified correctly when both the flow source and dest are stratified. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_infection_frequency_flow("infection", 0.03, "S", "I") model.add_sojourn_flow("recovery", 7, "I", "R") expected_flows = [ InfectionFrequencyFlow( "infection", C("S"), C("I"), 0.03, model._get_infection_frequency_multiplier ), SojournFlow("recovery", C("I"), C("R"), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) # Apply stratification strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) model.stratify_with(strat) expected_flows = [ InfectionFrequencyFlow( "infection", C("S", {"location": "urban"}), C("I", {"location": "urban"}), 0.03, model._get_infection_frequency_multiplier, ), InfectionFrequencyFlow( "infection", C("S", {"location": "rural"}), C("I", {"location": "rural"}), 0.03, model._get_infection_frequency_multiplier, ), SojournFlow("recovery", C("I", {"location": "urban"}), C("R", {"location": "urban"}), 7), SojournFlow("recovery", C("I", {"location": "rural"}), C("R", {"location": "rural"}), 7), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_stratify__single_with_pop_split__validate_compartments(): """ Ensure stratifying a model correctly adjusts the model compartments. Also the population split should be applied. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population({"S": 900, "I": 90, "R": 10}) # Compartments exist assert model.compartments == [ Compartment("S"), Compartment("I"), Compartment("R") ] # Each compartment knows its index assert [c.idx for c in model.compartments ] == list(range(len(model.compartments))) # Compartments have the correct population assert_array_equal(model.initial_population, np.array([900, 90, 10])) # Stratify the model strat = Stratification(name="age", strata=["child", "adult"], compartments=["S", "I", "R"]) strat.set_population_split({"child": 0.8, "adult": 0.2}) model.stratify_with(strat) assert model._stratifications == [strat] # Ensure compartments are stratified correctly assert [c.idx for c in model.compartments ] == list(range(len(model.compartments))) assert model.compartments == [ Compartment("S", {"age": "child"}), Compartment("S", {"age": "adult"}), Compartment("I", {"age": "child"}), Compartment("I", {"age": "adult"}), Compartment("R", {"age": "child"}), Compartment("R", {"age": "adult"}), ] expected_pop_arr = np.array([720, 180, 72, 18, 8, 2]) assert_array_equal(model.initial_population, expected_pop_arr)
def test_no_mixing_matrix(): """ Test that we are using the default 'null-op' mixing matrix when we have a no user-specified mixing matrix """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) strat = Stratification(name="agegroup", strata=["child", "adult"], compartments=["S", "I", "R"]) model.stratify_with(strat) # We should get the default mixing matrix default_matrix = np.array([[1]]) actual_mixing = model._get_mixing_matrix(0) assert_array_equal(actual_mixing, default_matrix) # Static matrices shouldn't change over time actual_mixing = model._get_mixing_matrix(123) assert_array_equal(actual_mixing, default_matrix) # No mixing categories have been added. assert model._mixing_categories == [{}]
def test_add_entry_flows_post_stratification__with_filter(): """ Ensure we can add flows after a model is stratified when a strata filter is applied """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) assert len(model._flows) == 0 strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) model.stratify_with(strat) assert len(model._flows) == 0 model.add_importation_flow( "imports", 10, "S", dest_strata={"location": "urban"}, expected_flow_count=1 ) assert len(model._flows) == 1 expected_flows = [ ImportFlow("imports", C("S", {"location": "urban"}), 10, []), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def _get_test_model(timestep=1, times=[0, 150]): comps = ["S", "EE", "LE", "EA", "LA", "R"] infectious_comps = ["LE", "EA", "LA"] model = CompartmentalModel( times=times, compartments=comps, infectious_compartments=infectious_comps, timestep=timestep, ) model.set_initial_population({"S": int(20e6), "LA": 100}) # Add flows model.add_infection_frequency_flow(name="infection", contact_rate=0.03, source="S", dest="EE") model.add_sojourn_flow(name="infect_onset", sojourn_time=7, source="EE", dest="LE") model.add_sojourn_flow(name="incidence", sojourn_time=7, source="LE", dest="EA") model.add_sojourn_flow(name="progress", sojourn_time=7, source="EA", dest="LA") model.add_sojourn_flow(name="recovery", sojourn_time=7, source="LA", dest="R") model.add_death_flow(name="infect_death", death_rate=0.005, source="LA") model.add_transition_flow(name="warning_immunity", fractional_rate=0.01, source="R", dest="S") # Stratify by age age_strat = Stratification("age", AGE_STRATA, comps) age_strat.set_population_split(AGE_SPLIT_PROPORTIONS) age_strat.set_mixing_matrix(AGE_MIXING_MATRIX) age_strat.add_flow_adjustments( "infection", {s: Multiply(v) for s, v in AGE_SUSCEPTIBILITY.items()} ) model.stratify_with(age_strat) # Stratify by clinical status clinical_strat = Stratification("clinical", CLINICAL_STRATA, infectious_comps) clinical_strat.add_infectiousness_adjustments("LE", {**ADJ_BASE, "non_sympt": Overwrite(0.25)}) clinical_strat.add_infectiousness_adjustments("EA", {**ADJ_BASE, "non_sympt": Overwrite(0.25)}) clinical_strat.add_infectiousness_adjustments( "LA", { **ADJ_BASE, "non_sympt": Overwrite(0.25), "sympt_isolate": Overwrite(0.2), "hospital": Overwrite(0.2), "icu": Overwrite(0.2), }, ) clinical_strat.add_flow_adjustments( "infect_onset", { "non_sympt": Multiply(0.26), "icu": Multiply(0.01), "hospital": Multiply(0.04), "sympt_public": Multiply(0.66), "sympt_isolate": Multiply(0.03), }, ) model.stratify_with(clinical_strat) # Request derived outputs. model.request_output_for_flow(name="incidence", flow_name="incidence") model.request_output_for_flow(name="progress", flow_name="progress") for age in AGE_STRATA: for clinical in NOTIFICATION_STRATA: model.request_output_for_flow( name=f"progressXage_{age}Xclinical_{clinical}", flow_name="progress", dest_strata={"age": age, "clinical": clinical}, ) hospital_sources = [] icu_sources = [] for age in AGE_STRATA: icu_sources.append(f"progressXage_{age}Xclinical_icu") hospital_sources += [ f"progressXage_{age}Xclinical_icu", f"progressXage_{age}Xclinical_hospital", ] model.request_aggregate_output( name="new_hospital_admissions", sources=hospital_sources, ) model.request_aggregate_output(name="new_icu_admissions", sources=icu_sources) # Get notifications, which may included people detected in-country as they progress, or imported cases which are detected. notification_sources = [ f"progressXage_{a}Xclinical_{c}" for a in AGE_STRATA for c in NOTIFICATION_STRATA ] model.request_aggregate_output(name="notifications", sources=notification_sources) # Infection deaths. model.request_output_for_flow(name="infection_deaths", flow_name="infect_death") model.request_cumulative_output(name="accum_deaths", source="infection_deaths") # Track hospital occupancy. # We count all ICU and hospital late active compartments and a proportion of early active ICU cases. model.request_output_for_compartments( "_late_active_hospital", compartments=["LA"], strata={"clinical": "hospital"}, save_results=False, ) model.request_output_for_compartments( "icu_occupancy", compartments=["LA"], strata={"clinical": "icu"}, ) model.request_output_for_compartments( "_early_active_icu", compartments=["EA"], strata={"clinical": "icu"}, save_results=False, ) proportion_icu_patients_in_hospital = 0.25 model.request_function_output( name="_early_active_icu_proportion", func=lambda patients: patients * proportion_icu_patients_in_hospital, sources=["_early_active_icu"], save_results=False, ) model.request_aggregate_output( name="hospital_occupancy", sources=[ "_late_active_hospital", "icu_occupancy", "_early_active_icu_proportion", ], ) # Proportion seropositive model.request_output_for_compartments( name="_total_population", compartments=comps, save_results=False ) model.request_output_for_compartments(name="_recovered", compartments=["R"], save_results=False) model.request_function_output( name="proportion_seropositive", sources=["_recovered", "_total_population"], func=lambda recovered, total: recovered / total, ) return model
def test_stratify__double_with_split_and_partial__validate_compartments(): model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population({"S": 900, "I": 90, "R": 10}) # Compartments exist assert model.compartments == [ Compartment("S"), Compartment("I"), Compartment("R") ] # Each compartment knows its index assert [c.idx for c in model.compartments ] == list(range(len(model.compartments))) # Compartments have the correct population assert_array_equal(model.initial_population, np.array([900, 90, 10])) # Stratify the model age_strat = Stratification(name="age", strata=["child", "adult"], compartments=["S", "R"]) age_strat.set_population_split({"child": 0.8, "adult": 0.2}) model.stratify_with(age_strat) assert model._stratifications == [age_strat] # Ensure compartments are stratified correctly assert [c.idx for c in model.compartments ] == list(range(len(model.compartments))) assert model.compartments == [ Compartment("S", {"age": "child"}), Compartment("S", {"age": "adult"}), Compartment("I"), Compartment("R", {"age": "child"}), Compartment("R", {"age": "adult"}), ] expected_pop_arr = np.array([720, 180, 90, 8, 2]) assert_array_equal(model.initial_population, expected_pop_arr) # Stratify the model again! loc_strat = Stratification(name="location", strata=["urban", "rural", "alpine"], compartments=["S", "I"]) loc_strat.set_population_split({"urban": 0.7, "rural": 0.2, "alpine": 0.1}) model.stratify_with(loc_strat) assert model._stratifications == [age_strat, loc_strat] # Ensure compartments are stratified correctly assert [c.idx for c in model.compartments ] == list(range(len(model.compartments))) assert model.compartments == [ Compartment("S", { "age": "child", "location": "urban" }), Compartment("S", { "age": "child", "location": "rural" }), Compartment("S", { "age": "child", "location": "alpine" }), Compartment("S", { "age": "adult", "location": "urban" }), Compartment("S", { "age": "adult", "location": "rural" }), Compartment("S", { "age": "adult", "location": "alpine" }), Compartment("I", {"location": "urban"}), Compartment("I", {"location": "rural"}), Compartment("I", {"location": "alpine"}), Compartment("R", {"age": "child"}), Compartment("R", {"age": "adult"}), ] expected_pop_arr = np.array([504, 144, 72, 126, 36, 18, 63, 18, 9, 8, 2]) assert_allclose(model.initial_population, expected_pop_arr, atol=1e-9, rtol=0)
def test_strat_infectiousness__with_adjustments(): """ Ensure multiply infectiousness adjustment is applied. """ # Create a model model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.set_initial_population(distribution={"S": 900, "I": 100}) strat = Stratification("age", ["baby", "child", "adult"], ["S", "I", "R"]) strat.set_population_split({"baby": 0.1, "child": 0.3, "adult": 0.6}) strat.add_infectiousness_adjustments( "I", {"child": adjust.Multiply(3), "adult": adjust.Multiply(0.5), "baby": None} ) model.stratify_with(strat) assert_array_equal( model.initial_population, np.array([90, 270, 540, 10, 30, 60, 0, 0, 0]), ) # Do pre-run force of infection calcs. model._prepare_to_run() assert_array_equal( model._compartment_infectiousness["default"], np.array([0, 0, 0, 1, 3, 0.5, 0, 0, 0]), ) # Do pre-iteration force of infection calcs model._prepare_time_step(0, model.initial_population) # Get multipliers infectees = model.compartments[0:3] infectors = model.compartments[3:6] expected_density = 10 * 1 + 30 * 3 + 60 * 0.5 expected_frequency = expected_density / 1000 for infectee, infector in zip(infectees, infectors): assert model._get_infection_density_multiplier(infectee, infector) == expected_density for infectee, infector in zip(infectees, infectors): assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency # Stratify again, now with overwrites strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) strat.add_infectiousness_adjustments( "I", {"urban": adjust.Overwrite(1), "rural": adjust.Multiply(7)} ) model.stratify_with(strat) assert_array_equal( model.initial_population, np.array([45, 45, 135, 135, 270.0, 270, 5, 5, 15, 15, 30, 30, 0, 0, 0, 0, 0, 0]), ) # Do pre-run force of infection calcs. model._prepare_to_run() assert_array_equal( model._compartment_infectiousness["default"], np.array([0, 0, 0, 0, 0, 0, 1, 7, 1, 21, 1, 3.5, 0, 0, 0, 0, 0, 0]), ) # Do pre-iteration force of infection calcs model._prepare_time_step(0, model.initial_population) # Get multipliers infectees = model.compartments[0:6] infectors = model.compartments[6:12] expected_density = 5 * 1 + 5 * 7 + 15 * 1 + 15 * 21 + 30 * 1 + 30 * 3.5 expected_frequency = expected_density / 1000 for infectee, infector in zip(infectees, infectors): assert model._get_infection_density_multiplier(infectee, infector) == expected_density for infectee, infector in zip(infectees, infectors): assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency
def build_model(params: dict) -> CompartmentalModel: time = params["time"] model = CompartmentalModel( times=[time["start"], time["end"]], compartments=COMPARTMENTS, infectious_compartments=INFECTIOUS_COMPS, timestep=time["step"], ) # Add initial population init_pop = { Compartment.EARLY_LATENT: params["initial_early_latent_population"], Compartment.LATE_LATENT: params["initial_late_latent_population"], Compartment.INFECTIOUS: params["initial_infectious_population"], Compartment.DETECTED: params["initial_detected_population"], Compartment.ON_TREATMENT: params["initial_on_treatment_population"], Compartment.RECOVERED: 0, } sum_init_pop = sum(init_pop.values()) init_pop[Compartment. SUSCEPTIBLE] = params["start_population_size"] - sum_init_pop model.set_initial_population(init_pop) # Add inter-compartmental flows params = _get_derived_params(params) # Entry flows model.add_crude_birth_flow( "birth", params["crude_birth_rate"], Compartment.SUSCEPTIBLE, ) # Infection flows. model.add_infection_frequency_flow( "infection", params["contact_rate"], Compartment.SUSCEPTIBLE, Compartment.EARLY_LATENT, ) model.add_infection_frequency_flow( "infection_from_latent", params["contact_rate_from_latent"], Compartment.LATE_LATENT, Compartment.EARLY_LATENT, ) model.add_infection_frequency_flow( "infection_from_recovered", params["contact_rate_from_recovered"], Compartment.RECOVERED, Compartment.EARLY_LATENT, ) # Transition flows. model.add_fractional_flow( "treatment_early", params["preventive_treatment_rate"], Compartment.EARLY_LATENT, Compartment.RECOVERED, ) model.add_fractional_flow( "treatment_late", params["preventive_treatment_rate"], Compartment.LATE_LATENT, Compartment.RECOVERED, ) model.add_fractional_flow( "stabilisation", params["stabilisation_rate"], Compartment.EARLY_LATENT, Compartment.LATE_LATENT, ) model.add_fractional_flow( "early_activation", params["early_activation_rate"], Compartment.EARLY_LATENT, Compartment.INFECTIOUS, ) model.add_fractional_flow( "late_activation", params["late_activation_rate"], Compartment.LATE_LATENT, Compartment.INFECTIOUS, ) # Post-active-disease flows model.add_fractional_flow( "detection", params["detection_rate"], Compartment.INFECTIOUS, Compartment.DETECTED, ) model.add_fractional_flow( "treatment_commencement", params["treatment_commencement_rate"], Compartment.DETECTED, Compartment.ON_TREATMENT, ) model.add_fractional_flow( "missed_to_active", params["missed_to_active_rate"], Compartment.DETECTED, Compartment.INFECTIOUS, ) model.add_fractional_flow( "self_recovery_infectious", params["self_recovery_rate"], Compartment.INFECTIOUS, Compartment.LATE_LATENT, ) model.add_fractional_flow( "self_recovery_detected", params["self_recovery_rate"], Compartment.DETECTED, Compartment.LATE_LATENT, ) model.add_fractional_flow( "treatment_recovery", params["treatment_recovery_rate"], Compartment.ON_TREATMENT, Compartment.RECOVERED, ) model.add_fractional_flow( "treatment_default", params["treatment_default_rate"], Compartment.ON_TREATMENT, Compartment.INFECTIOUS, ) model.add_fractional_flow( "failure_retreatment", params["failure_retreatment_rate"], Compartment.ON_TREATMENT, Compartment.DETECTED, ) model.add_fractional_flow( "spontaneous_recovery", params["spontaneous_recovery_rate"], Compartment.ON_TREATMENT, Compartment.LATE_LATENT, ) # Death flows # Universal death rate to be overriden by a multiply in age stratification. uni_death_flow_names = model.add_universal_death_flows("universal_death", death_rate=1) model.add_death_flow( "infectious_death", params["infect_death_rate"], Compartment.INFECTIOUS, ) model.add_death_flow( "detected_death", params["infect_death_rate"], Compartment.DETECTED, ) model.add_death_flow( "treatment_death", params["treatment_death_rate"], Compartment.ON_TREATMENT, ) # Apply age-stratification age_strat = _build_age_strat(params, uni_death_flow_names) model.stratify_with(age_strat) # Add vaccination stratification. vac_strat = _build_vac_strat(params) model.stratify_with(vac_strat) # Apply organ stratification organ_strat = _build_organ_strat(params) model.stratify_with(organ_strat) # Apply strain stratification strain_strat = _build_strain_strat(params) model.stratify_with(strain_strat) # Add amplification flow model.add_fractional_flow( name="amplification", fractional_rate=params["amplification_rate"], source=Compartment.ON_TREATMENT, dest=Compartment.ON_TREATMENT, source_strata={"strain": "ds"}, dest_strata={"strain": "mdr"}, expected_flow_count=9, ) # Add cross-strain reinfection flows model.add_infection_frequency_flow( name="reinfection_ds_to_mdr", contact_rate=params["reinfection_rate"], source=Compartment.EARLY_LATENT, dest=Compartment.EARLY_LATENT, source_strata={"strain": "ds"}, dest_strata={"strain": "mdr"}, expected_flow_count=3, ) model.add_infection_frequency_flow( name="reinfection_mdr_to_ds", contact_rate=params["reinfection_rate"], source=Compartment.EARLY_LATENT, dest=Compartment.EARLY_LATENT, source_strata={"strain": "mdr"}, dest_strata={"strain": "ds"}, expected_flow_count=3, ) model.add_infection_frequency_flow( name="reinfection_late_ds_to_mdr", contact_rate=params["reinfection_rate"], source=Compartment.LATE_LATENT, dest=Compartment.EARLY_LATENT, source_strata={"strain": "ds"}, dest_strata={"strain": "mdr"}, expected_flow_count=3, ) model.add_infection_frequency_flow( name="reinfection_late_mdr_to_ds", contact_rate=params["reinfection_rate"], source=Compartment.LATE_LATENT, dest=Compartment.EARLY_LATENT, source_strata={"strain": "mdr"}, dest_strata={"strain": "ds"}, expected_flow_count=3, ) # Apply classification stratification class_strat = _build_class_strat(params) model.stratify_with(class_strat) # Apply retention stratification retention_strat = _build_retention_strat(params) model.stratify_with(retention_strat) # Register derived output functions, which are calculations based on the model's compartment values or flows. # These are calculated after the model is run. model.request_output_for_flow("notifications", flow_name="detection") model.request_output_for_flow("early_activation", flow_name="early_activation") model.request_output_for_flow("late_activation", flow_name="late_activation") model.request_output_for_flow("infectious_deaths", flow_name="infectious_death") model.request_output_for_flow("detected_deaths", flow_name="detected_death") model.request_output_for_flow("treatment_deaths", flow_name="treatment_death") model.request_output_for_flow("progression_early", flow_name="early_activation") model.request_output_for_flow("progression_late", flow_name="late_activation") model.request_aggregate_output("progression", ["progression_early", "progression_late"]) model.request_output_for_compartments("population_size", COMPARTMENTS) model.request_aggregate_output( "_incidence", sources=["early_activation", "late_activation"], save_results=False) model.request_function_output("incidence", sources=["_incidence", "population_size"], func=lambda i, p: 1e5 * i / p) model.request_aggregate_output( "disease_deaths", sources=["infectious_deaths", "detected_deaths", "treatment_deaths"]) cum_start_time = params["cumulative_output_start_time"] model.request_cumulative_output("cumulative_diseased", source="_incidence", start_time=cum_start_time) model.request_cumulative_output("cumulative_deaths", source="disease_deaths", start_time=cum_start_time) model.request_output_for_compartments("_count_infectious", INFECTIOUS_COMPS, save_results=False) model.request_function_output( "prevalence_infectious", sources=["_count_infectious", "population_size"], func=lambda c, p: 1e5 * c / p, ) model.request_output_for_compartments( "_count_latent", [Compartment.EARLY_LATENT, Compartment.LATE_LATENT], save_results=False) model.request_function_output( "percentage_latent", sources=["_count_latent", "population_size"], func=lambda c, p: 100 * c / p, ) return model
def test_stratify_entry_flows__with_explicit_adjustments(): """ Ensure entry flows are stratified correctly when adjustments are requested. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.add_importation_flow("imports", 10, "S") expected_flows = [ImportFlow("imports", C("S"), 10)] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) strat.add_flow_adjustments("imports", {"urban": Multiply(0.9), "rural": None}) model.stratify_with(strat) expected_flows = [ ImportFlow( "imports", C("S", {"location": "urban"}), 10, [Multiply(0.9)], ), ImportFlow( "imports", C("S", {"location": "rural"}), 10, [], ), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)]) strat = Stratification("age", ["young", "old"], ["S", "I", "R"]) strat.add_flow_adjustments("imports", {"young": Multiply(0.8), "old": Overwrite(1)}) model.stratify_with(strat) expected_flows = [ ImportFlow( "imports", C("S", {"location": "urban", "age": "young"}), 10, [Multiply(0.9), Multiply(0.8)], ), ImportFlow( "imports", C("S", {"location": "urban", "age": "old"}), 10, [Multiply(0.9), Overwrite(1)], ), ImportFlow( "imports", C("S", {"location": "rural", "age": "young"}), 10, [Multiply(0.8)], ), ImportFlow( "imports", C("S", {"location": "rural", "age": "old"}), 10, [Overwrite(1)], ), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_strain__with_flow_adjustments(): """ Test infectious multiplier and flow rate calculations for 3 strains which have different flow adjustments. These flow adjustments would correspond to some physical process that we're modelling, and they should be effectively the same as applying infectiousness multipliers. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population(distribution={"S": 900, "I": 100}) contact_rate = 0.2 model.add_infection_frequency_flow("infection", contact_rate, "S", "I") strat = StrainStratification("strain", ["a", "b", "c"], ["I"]) strat.set_population_split({ "a": 0.7, # 70 people "b": 0.2, # 20 people "c": 0.1, # 10 people }) strat.add_flow_adjustments( "infection", { "a": adjust.Multiply(0.5), # 0.5x as susceptible "b": adjust.Multiply(3), # 3x as susceptible "c": adjust.Multiply(2), # 2x as susceptible }, ) model.stratify_with(strat) # Do pre-run force of infection calcs. model._prepare_to_run() assert_array_equal(model._compartment_infectiousness["a"], np.array([0, 1, 0, 0, 0])) assert_array_equal(model._compartment_infectiousness["b"], np.array([0, 0, 1, 0, 0])) assert_array_equal(model._compartment_infectiousness["c"], np.array([0, 0, 0, 1, 0])) assert model._category_lookup == {0: 0, 1: 0, 2: 0, 3: 0, 4: 0} assert_array_equal(model._category_matrix, np.array([[1, 1, 1, 1, 1]])) # Do pre-iteration force of infection calcs model._prepare_time_step(0, model.initial_population) assert_array_equal(model._category_populations, np.array([1000])) assert_array_equal(model._infection_density["a"], np.array([70])) assert_array_equal(model._infection_density["b"], np.array([20])) assert_array_equal(model._infection_density["c"], np.array([10])) assert_array_equal(model._infection_frequency["a"], np.array([70 / 1000])) assert_array_equal(model._infection_frequency["b"], np.array([20 / 1000])) assert_array_equal(model._infection_frequency["c"], np.array([10 / 1000])) # Get multipliers susceptible = model.compartments[0] infectious_a = model.compartments[1] infectious_b = model.compartments[2] infectious_c = model.compartments[3] assert model._get_infection_density_multiplier(susceptible, infectious_a) == 70 assert model._get_infection_density_multiplier(susceptible, infectious_b) == 20 assert model._get_infection_density_multiplier(susceptible, infectious_c) == 10 assert model._get_infection_frequency_multiplier(susceptible, infectious_a) == 70 / 1000 assert model._get_infection_frequency_multiplier(susceptible, infectious_b) == 20 / 1000 assert model._get_infection_frequency_multiplier(susceptible, infectious_c) == 10 / 1000 # Get infection flow rates flow_rates = model._get_compartment_rates(model.initial_population, 0) sus_pop = 900 flow_to_a = sus_pop * contact_rate * (70 * 0.5 / 1000) flow_to_b = sus_pop * contact_rate * (20 * 3 / 1000) flow_to_c = sus_pop * contact_rate * (10 * 2 / 1000) expected_flow_rates = np.array([ -flow_to_a - flow_to_b - flow_to_c, flow_to_a, flow_to_b, flow_to_c, 0.0 ]) assert_allclose(expected_flow_rates, flow_rates, verbose=True)
def test_strain__with_infectious_multipliers_and_heterogeneous_mixing(): """ Test infectious multiplier and flow rate calculations for 3 strains which have different infectiousness levels plus a seperate stratification which has a mixing matrix. """ model = CompartmentalModel(times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]) model.set_initial_population(distribution={"S": 900, "I": 100}) contact_rate = 0.2 model.add_infection_frequency_flow("infection", contact_rate, "S", "I") age_strat = Stratification("age", ["child", "adult"], ["S", "I", "R"]) age_strat.set_population_split({ "child": 0.6, # 600 people "adult": 0.4, # 400 people }) # Higher mixing among adults or children, # than between adults or children. age_strat.set_mixing_matrix(np.array([[1.5, 0.5], [0.5, 1.5]])) model.stratify_with(age_strat) strain_strat = StrainStratification("strain", ["a", "b", "c"], ["I"]) strain_strat.set_population_split({ "a": 0.7, # 70 people "b": 0.2, # 20 people "c": 0.1, # 10 people }) strain_strat.add_infectiousness_adjustments( "I", { "a": adjust.Multiply(0.5), # 0.5x as susceptible "b": adjust.Multiply(3), # 3x as susceptible "c": adjust.Multiply(2), # 2x as susceptible }, ) model.stratify_with(strain_strat) # Do pre-run force of infection calcs. model._prepare_to_run() assert_array_equal(model._compartment_infectiousness["a"], np.array([0, 0, 0.5, 0, 0, 0.5, 0, 0, 0, 0])) assert_array_equal(model._compartment_infectiousness["b"], np.array([0, 0, 0, 3, 0, 0, 3, 0, 0, 0])) assert_array_equal(model._compartment_infectiousness["c"], np.array([0, 0, 0, 0, 2, 0, 0, 2, 0, 0])) # 0 for child, 1 for adult assert model._category_lookup == { 0: 0, 1: 1, 2: 0, 3: 0, 4: 0, 5: 1, 6: 1, 7: 1, 8: 0, 9: 1 } assert_array_equal( model._category_matrix, np.array([ [1, 0, 1, 1, 1, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 1, 1, 1, 0, 1], ]), ) # Do pre-iteration force of infection calcs model._prepare_time_step(0, model.initial_population) assert_array_equal(model._category_populations, np.array([600, 400])) assert_array_equal( model._infection_density["a"], np.array([0.5 * (42 * 1.5 + 28 * 0.5), 0.5 * (42 * 0.5 + 28 * 1.5)]), ) assert_array_equal( model._infection_density["b"], np.array([ 3 * (12 * 1.5 + 8 * 0.5), 3 * (8 * 1.5 + 12 * 0.5), ]), ) assert_array_equal( model._infection_density["c"], np.array([2 * (6 * 1.5 + 4 * 0.5), 2 * (4 * 1.5 + 6 * 0.5)]), ) assert_array_equal( model._infection_frequency["a"], np.array([ 0.5 * ((42 / 600) * 1.5 + (28 / 400) * 0.5), 0.5 * ((42 / 600) * 0.5 + (28 / 400) * 1.5), ]), ) assert_array_equal( model._infection_frequency["b"], np.array([ 3 * ((12 / 600) * 1.5 + (8 / 400) * 0.5), 3 * ((8 / 400) * 1.5 + (12 / 600) * 0.5), ]), ) assert_array_equal( model._infection_frequency["c"], np.array([ 2 * ((6 / 600) * 1.5 + (4 / 400) * 0.5), 2 * ((4 / 400) * 1.5 + (6 / 600) * 0.5) ]), ) # Get multipliers sus_child = model.compartments[0] sus_adult = model.compartments[1] inf_child_a = model.compartments[2] inf_child_b = model.compartments[3] inf_child_c = model.compartments[4] inf_adult_a = model.compartments[5] inf_adult_b = model.compartments[6] inf_adult_c = model.compartments[7] density = model._get_infection_density_multiplier freq = model._get_infection_frequency_multiplier assert density(sus_child, inf_child_a) == 0.5 * (42 * 1.5 + 28 * 0.5) assert density(sus_adult, inf_adult_a) == 0.5 * (42 * 0.5 + 28 * 1.5) assert density(sus_child, inf_child_b) == 3 * (12 * 1.5 + 8 * 0.5) assert density(sus_adult, inf_adult_b) == 3 * (8 * 1.5 + 12 * 0.5) assert density(sus_child, inf_child_c) == 2 * (6 * 1.5 + 4 * 0.5) assert density(sus_adult, inf_adult_c) == 2 * (4 * 1.5 + 6 * 0.5) assert freq(sus_child, inf_child_a) == 0.5 * ((42 / 600) * 1.5 + (28 / 400) * 0.5) assert freq(sus_adult, inf_adult_a) == 0.5 * ((42 / 600) * 0.5 + (28 / 400) * 1.5) assert freq(sus_child, inf_child_b) == 3 * ((12 / 600) * 1.5 + (8 / 400) * 0.5) assert freq(sus_adult, inf_adult_b) == 3 * ((8 / 400) * 1.5 + (12 / 600) * 0.5) assert freq(sus_child, inf_child_c) == 2 * ((6 / 600) * 1.5 + (4 / 400) * 0.5) assert freq(sus_adult, inf_adult_c) == 2 * ((4 / 400) * 1.5 + (6 / 600) * 0.5) # Get infection flow rates flow_to_inf_child_a = 540 * contact_rate * freq(sus_child, inf_child_a) flow_to_inf_adult_a = 360 * contact_rate * freq(sus_adult, inf_adult_a) flow_to_inf_child_b = 540 * contact_rate * freq(sus_child, inf_child_b) flow_to_inf_adult_b = 360 * contact_rate * freq(sus_adult, inf_adult_b) flow_to_inf_child_c = 540 * contact_rate * freq(sus_child, inf_child_c) flow_to_inf_adult_c = 360 * contact_rate * freq(sus_adult, inf_adult_c) expected_flow_rates = np.array([ -flow_to_inf_child_a - flow_to_inf_child_b - flow_to_inf_child_c, -flow_to_inf_adult_a - flow_to_inf_adult_b - flow_to_inf_adult_c, flow_to_inf_child_a, flow_to_inf_child_b, flow_to_inf_child_c, flow_to_inf_adult_a, flow_to_inf_adult_b, flow_to_inf_adult_c, 0.0, 0.0, ]) flow_rates = model._get_compartment_rates(model.initial_population, 0) assert_allclose(expected_flow_rates, flow_rates, verbose=True)
def test_stratify__age__validate_ageing_flows_added_second(): """ Ensure that age stratification works when applied after a previous stratification. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) assert len(model._flows) == 0 model.add_crude_birth_flow("births", 0.02, "S") assert len(model._flows) == 1 strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) model.stratify_with(strat) strat = AgeStratification("age", ["0", "5", "15"], ["S", "I", "R"]) model.stratify_with(strat) # Expect ageing flows amongst age group and a birth flow that only goes to age 0. expected_flows = [ CrudeBirthFlow("births", C("S", {"location": "urban", "age": "0"}), 0.02, [Multiply(0.5)]), CrudeBirthFlow("births", C("S", {"location": "rural", "age": "0"}), 0.02, [Multiply(0.5)]), SojournFlow( "ageing_SXlocation_urbanXage_0_to_SXlocation_urbanXage_5", C("S", {"location": "urban", "age": "0"}), C("S", {"location": "urban", "age": "5"}), 5, ), SojournFlow( "ageing_SXlocation_ruralXage_0_to_SXlocation_ruralXage_5", C("S", {"location": "rural", "age": "0"}), C("S", {"location": "rural", "age": "5"}), 5, ), SojournFlow( "ageing_IXlocation_urbanXage_0_to_IXlocation_urbanXage_5", C("I", {"location": "urban", "age": "0"}), C("I", {"location": "urban", "age": "5"}), 5, ), SojournFlow( "ageing_IXlocation_ruralXage_0_to_IXlocation_ruralXage_5", C("I", {"location": "rural", "age": "0"}), C("I", {"location": "rural", "age": "5"}), 5, ), SojournFlow( "ageing_RXlocation_urbanXage_0_to_RXlocation_urbanXage_5", C("R", {"location": "urban", "age": "0"}), C("R", {"location": "urban", "age": "5"}), 5, ), SojournFlow( "ageing_RXlocation_ruralXage_0_to_RXlocation_ruralXage_5", C("R", {"location": "rural", "age": "0"}), C("R", {"location": "rural", "age": "5"}), 5, ), SojournFlow( "ageing_SXlocation_urbanXage_5_to_SXlocation_urbanXage_15", C("S", {"location": "urban", "age": "5"}), C("S", {"location": "urban", "age": "15"}), 10, ), SojournFlow( "ageing_SXlocation_ruralXage_5_to_SXlocation_ruralXage_15", C("S", {"location": "rural", "age": "5"}), C("S", {"location": "rural", "age": "15"}), 10, ), SojournFlow( "ageing_IXlocation_urbanXage_5_to_IXlocation_urbanXage_15", C("I", {"location": "urban", "age": "5"}), C("I", {"location": "urban", "age": "15"}), 10, ), SojournFlow( "ageing_IXlocation_ruralXage_5_to_IXlocation_ruralXage_15", C("I", {"location": "rural", "age": "5"}), C("I", {"location": "rural", "age": "15"}), 10, ), SojournFlow( "ageing_RXlocation_urbanXage_5_to_RXlocation_urbanXage_15", C("R", {"location": "urban", "age": "5"}), C("R", {"location": "urban", "age": "15"}), 10, ), SojournFlow( "ageing_RXlocation_ruralXage_5_to_RXlocation_ruralXage_15", C("R", {"location": "rural", "age": "5"}), C("R", {"location": "rural", "age": "15"}), 10, ), ] assert len(expected_flows) == len(model._flows) assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])