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_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_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_sojourn_flow_get_net_flow():
flow = SojournFlow(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=lambda t: 2 * t,
adjustments=[],
)
flow.source.idx = 1
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 3 / (2 * 7)
def test_sojourn_flow_get_net_flow_with_adjust():
flow = SojournFlow(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=lambda t: 2 * t,
adjustments=[adjust.Multiply(13)],
)
flow.source.idx = 2
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
np.testing.assert_almost_equal(net_flow, 5 / (2 * 13 * 7))
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__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_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)])