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_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_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_entry_flow_stratify__with_no_flow_adjustments():
flow = EntryFlow(
name="flow",
dest=Compartment("I"),
param=2,
adjustments=[],
)
strat = Stratification(
name="location",
strata=["1", "2"],
compartments=["I", "R"],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
# Both flows has 50% flow adjustment applied to conserve inflows of people.
assert new_flows[0]._is_equal(
EntryFlow(
name="flow",
param=2,
dest=Compartment("I", {"location": "1"}),
adjustments=[adjust.Multiply(0.5)],
))
assert new_flows[1]._is_equal(
EntryFlow(
name="flow",
param=2,
dest=Compartment("I", {"location": "2"}),
adjustments=[adjust.Multiply(0.5)],
))
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_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_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 _build_organ_strat(params):
organ_strat = Stratification(
"organ", ["smear_positive", "smear_negative", "extra_pulmonary"],
INFECTIOUS_COMPS)
organ_strat.set_population_split(params["organ"]["props"])
# Add infectiousness adjustments
for comp in INFECTIOUS_COMPS:
organ_strat.add_infectiousness_adjustments(
comp, _adjust_all_multiply(params["organ"]["foi"]))
organ_strat.add_flow_adjustments(
"early_activation", _adjust_all_multiply(params["organ"]["props"]))
organ_strat.add_flow_adjustments(
"late_activation", _adjust_all_multiply(params["organ"]["props"]))
organ_strat.add_flow_adjustments(
"detection",
_adjust_all_multiply(params["detection_rate_stratified"]["organ"]),
)
organ_strat.add_flow_adjustments(
"self_recovery_infectious",
_adjust_all_multiply(params["self_recovery_rate_stratified"]["organ"]),
)
organ_strat.add_flow_adjustments(
"self_recovery_detected",
_adjust_all_multiply(params["self_recovery_rate_stratified"]["organ"]),
)
return organ_strat
def test_stratify_transition_flows__with_dest_only_stratified__with_adjustments():
"""
Ensure transition flows are stratified correctly when only the flow destination is stratified.
Expect adjustments to override the automatic person-conserving adjustment.
"""
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"])
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 _build_retention_strat(params):
strat = Stratification("retained", ["yes", "no"], [Compartment.DETECTED])
strat.add_flow_adjustments(
"detection",
{
"yes": Multiply(params["retention_prop"]),
"no": Multiply(1 - params["retention_prop"]),
},
)
strat.add_flow_adjustments(
"missed_to_active",
{
"yes": Multiply(0),
"no": Multiply(1)
},
)
strat.add_flow_adjustments(
"treatment_commencement",
{
"yes": Multiply(1),
"no": Multiply(0)
},
)
strat.add_flow_adjustments(
"failure_retreatment",
{
"yes": Multiply(1),
"no": Multiply(0)
},
)
return strat
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_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 _build_class_strat(params):
strat = Stratification("classified", ["correctly", "incorrectly"],
[Compartment.DETECTED, Compartment.ON_TREATMENT])
# Apply only to ds flows
strat.add_flow_adjustments(
"detection",
{
"correctly": Multiply(1.0),
"incorrectly": Multiply(0.0)
},
source_strata={"strain": "ds"},
)
# Apply only to mdr flows
strat.add_flow_adjustments(
"detection",
{
"correctly": Multiply(params["frontline_xpert_prop"]),
"incorrectly": Multiply(1.0 - params["frontline_xpert_prop"]),
},
source_strata={"strain": "mdr"},
)
# Apply only to mdr flows that have extra_pulmonary organ.
strat.add_flow_adjustments(
"detection",
{
"correctly": Multiply(0),
"incorrectly": Multiply(1)
},
source_strata={
"strain": "mdr",
"organ": "extra_pulmonary"
},
)
strat.add_flow_adjustments(
"treatment_default",
_adjust_all_multiply(
params["treatment_default_rate_stratified"]["classified"]),
source_strata={"strain": "mdr"},
)
strat.add_flow_adjustments(
"spontaneous_recovery",
_adjust_all_multiply(
params["spontaneous_recovery_rate_stratified"]["classified"]),
)
strat.add_flow_adjustments(
"failure_retreatment",
_adjust_all_multiply(
params["failure_retreatment_rate_stratified"]["classified"]),
)
return strat
def test_get_flow_adjustments__with_no_adjustments():
trans_flow = TransitionFlow("flow", Compartment("S"), Compartment("I"), 1)
entry_flow = EntryFlow("flow", Compartment("S"), 1)
exit_flow = ExitFlow("flow", Compartment("I"), 1)
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
for flow in [trans_flow, entry_flow, exit_flow]:
assert strat.get_flow_adjustment(flow) is None
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_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_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_stratify_compartments__with_no_extisting_strat_and_subset_only():
strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S"],
)
comps = [Compartment("S"), Compartment("I"), Compartment("R")]
strat_comps = strat._stratify_compartments(comps)
assert strat_comps == [
Compartment("S", {"age": "0"}),
Compartment("S", {"age": "10"}),
Compartment("S", {"age": "20"}),
Compartment("I"),
Compartment("R"),
]
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_transition_flow_stratify_with_no_matching_compartments():
flow = TransitionFlow(
name="flow",
source=Compartment("S"),
dest=Compartment("I"),
param=2,
adjustments=[],
)
strat = Stratification(
name="location",
strata=["1", "2", "3"],
compartments=["R"],
)
new_flows = flow.stratify(strat)
assert new_flows == [flow]
def test_get_flow_adjustments__with_one_adjustment():
other_flow = TransitionFlow("other", Compartment("S"), Compartment("I"), 1)
trans_flow = TransitionFlow("flow", Compartment("S"), Compartment("I"), 1)
entry_flow = EntryFlow("flow", Compartment("S"), 1)
exit_flow = ExitFlow("flow", Compartment("I"), 1)
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
strat.add_flow_adjustments("flow", {"rural": Multiply(1), "urban": None})
assert strat.get_flow_adjustment(other_flow) is None
for flow in [trans_flow, entry_flow, exit_flow]:
adj = strat.get_flow_adjustment(flow)
assert adj["urban"] is None
assert adj["rural"]._is_equal(Multiply(1))
def test_create_stratification():
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
assert strat.name == "location"
assert strat.compartments == [
Compartment("S"), Compartment("I"),
Compartment("R")
]
assert strat.strata == ["rural", "urban"]
assert strat.population_split == {"rural": 0.5, "urban": 0.5}
assert strat.flow_adjustments == {}
assert strat.infectiousness_adjustments == {}
assert strat.mixing_matrix is None
assert not strat.is_ageing()
assert not strat.is_strain()
def get_cluster_strat(params: Parameters) -> Stratification:
cluster_strat = Stratification("cluster", CLUSTER_STRATA, COMPARTMENTS)
country = params.country
vic = params.victorian_clusters
# Determine how to split up population by cluster
# There is -0.5% to +4% difference per age group between sum of region population in 2018 and
# total VIC population in 2020
region_pops = {
region: sum(
inputs.get_population_by_agegroup(AGEGROUP_STRATA,
country.iso3,
region.upper(),
year=2018))
for region in CLUSTER_STRATA
}
sum_region_props = sum(region_pops.values())
cluster_split_props = {
region: pop / sum_region_props
for region, pop in region_pops.items()
}
cluster_strat.set_population_split(cluster_split_props)
# Adjust contact rate multipliers
contact_rate_adjustments = {}
for cluster in Region.VICTORIA_METRO + [Region.BARWON_SOUTH_WEST]:
cluster_name = cluster.replace("-", "_")
contact_rate_multiplier = getattr(
vic, f"contact_rate_multiplier_{cluster_name}")
contact_rate_adjustments[cluster_name] = Multiply(
contact_rate_multiplier)
for cluster in Region.VICTORIA_RURAL:
if cluster != Region.BARWON_SOUTH_WEST:
cluster_name = cluster.replace("-", "_")
contact_rate_multiplier = getattr(
vic, "contact_rate_multiplier_regional")
contact_rate_adjustments[cluster_name] = Multiply(
contact_rate_multiplier)
# Add in flow adjustments per-region so we can calibrate the contact rate for each region.
cluster_strat.add_flow_adjustments("infection", contact_rate_adjustments)
# Use an identity mixing matrix to temporarily declare no inter-cluster mixing, which will then be over-written
cluster_mixing_matrix = np.eye(len(CLUSTER_STRATA))
cluster_strat.set_mixing_matrix(cluster_mixing_matrix)
return cluster_strat
def test_entry_flow_stratify__when_not_applicable():
flow = EntryFlow(
name="flow",
dest=Compartment("I"),
param=2,
adjustments=[],
)
strat = Stratification(
name="location",
strata=["1", "2", "3"],
compartments=["R"],
)
# Expect no stratification because compartment not being stratified.
new_flows = flow.stratify(strat)
assert new_flows == [flow]
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)
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_transition_flow_stratify_dest_but_not_source__with_flow_adjustments():
"""
Ensure flow is adjusted to account for fan out.
"""
flow = TransitionFlow(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=2,
adjustments=[adjust.Multiply(0.1)],
)
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["S", "R"],
)
strat.add_flow_adjustments(
"flow",
{
"1": adjust.Multiply(0.2),
"2": adjust.Multiply(0.3),
},
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0]._is_equal(
TransitionFlow(
name="flow",
param=2,
source=Compartment("I"),
dest=Compartment("R", {"age": "1"}),
adjustments=[adjust.Multiply(0.1), adjust.Multiply(0.2)],
)
)
assert new_flows[1]._is_equal(
TransitionFlow(
name="flow",
param=2,
source=Compartment("I"),
dest=Compartment("R", {"age": "2"}),
adjustments=[adjust.Multiply(0.1), adjust.Multiply(0.3)],
)
)
def test_create_stratification__with_pop_split():
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
assert strat.population_split == {"rural": 0.5, "urban": 0.5}
# Works
strat.set_population_split({"rural": 0.2, "urban": 0.8})
assert strat.population_split == {"rural": 0.2, "urban": 0.8}
# Fails coz missing a key
with pytest.raises(AssertionError):
strat.set_population_split({"urban": 1})
# Fails coz doesn't sum to 1
with pytest.raises(AssertionError):
strat.set_population_split({"urban": 0.2, "rural": 0.3})
# Fails coz contains negative number
with pytest.raises(AssertionError):
strat.set_population_split({"urban": -2, "rural": 3})
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 _build_vac_strat(params):
vac_strat = Stratification("vac", ["unvaccinated", "vaccinated"],
[Compartment.SUSCEPTIBLE])
vac_strat.set_population_split({"unvaccinated": 1.0, "vaccinated": 0.0})
# Apply flow adjustments
vac_strat.add_flow_adjustments(
"infection",
{
"unvaccinated": Multiply(1.0),
"vaccinated": Multiply(params["bcg"]["rr_infection_vaccinated"]),
},
)
vac_strat.add_flow_adjustments(
"treatment_early",
{
"unvaccinated": Multiply(0.0),
"vaccinated": Multiply(1.0)
},
)
vac_strat.add_flow_adjustments("treatment_late", {
"unvaccinated": Multiply(0.0),
"vaccinated": Multiply(1.0)
})
def time_varying_vaccination_coverage(t):
return params["bcg"][
"coverage"] if t > params["bcg"]["start_time"] else 0.0
def time_varying_unvaccinated_coverage(t):
return 1 - params["bcg"]["coverage"] if t > params["bcg"][
"start_time"] else 1.0
vac_strat.add_flow_adjustments(
"birth",
{
"unvaccinated": Multiply(time_varying_unvaccinated_coverage),
"vaccinated": Multiply(time_varying_vaccination_coverage),
},
)
return vac_strat