def test_stratify_compartment_values__with_subset_stratified():
strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S"],
)
strat.set_population_split({"0": 0.25, "10": 0.5, "20": 0.25})
comps = [Compartment("S"), Compartment("I"), Compartment("R")]
comp_values = np.array([1000.0, 100.0, 0.0])
new_comp_values = strat._stratify_compartment_values(comps, comp_values)
expected_arr = np.array([250.0, 500.0, 250.0, 100.0, 0.0])
assert_array_equal(expected_arr, new_comp_values)
def test_stratify_compartment_values__with_extisting_strat():
"""
Stratify compartments for the second time, expect that compartments
are are split according to proportions and old compartments are removed.
"""
comp_values = np.array(
[250.0, 500.0, 250.0, 25.0, 50.0, 25.0, 0.0, 0.0, 0.0])
comps = [
Compartment("S", {"age": "0"}),
Compartment("S", {"age": "10"}),
Compartment("S", {"age": "20"}),
Compartment("I", {"age": "0"}),
Compartment("I", {"age": "10"}),
Compartment("I", {"age": "20"}),
Compartment("R", {"age": "0"}),
Compartment("R", {"age": "10"}),
Compartment("R", {"age": "20"}),
]
strat = Stratification(
name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"],
)
strat.set_population_split({"rural": 0.1, "urban": 0.9})
new_comp_values = strat._stratify_compartment_values(comps, comp_values)
expected_arr = np.array([
25,
225.0,
50.0,
450.0,
25.0,
225.0,
2.5,
22.5,
5.0,
45.0,
2.5,
22.5,
0.0,
0.0,
0.0,
0.0,
0.0,
0.0,
])
assert_array_equal(expected_arr, new_comp_values)