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)])