def test_transition_flow_get_net_flow_with_adjust():
flow = TransitionFlow(
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)
assert net_flow == 2 * 5 * 7 * 13
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_transition_flow_get_net_flow():
flow = TransitionFlow(
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 == 2 * 3 * 7
def test_infection_get_net_flow_with_adjust(FlowClass):
flow = FlowClass(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=lambda t: 2 * t,
find_infectious_multiplier=lambda s, d: 23,
adjustments=[adjust.Multiply(13)],
)
flow.source.idx = 1
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 2 * 3 * 7 * 23 * 13
def test_stratify_compartment_values__with_no_extisting_strat():
strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S", "I", "R"],
)
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, 500.0, 250.0, 25.0, 50.0, 25.0, 0.0, 0.0, 0.0])
assert_array_equal(expected_arr, new_comp_values)
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_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 test_replace_deaths_birth_flow_get_net_flow_with_adjust():
flow = ReplacementBirthFlow(
name="flow",
dest=Compartment("S"),
param=lambda t: 23,
adjustments=[adjust.Multiply(13)],
)
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 23 * 13
def test_crude_birth_flow_get_net_flow_with_adjust():
flow = CrudeBirthFlow(
name="flow",
dest=Compartment("S"),
param=lambda t: 0.1 * t,
adjustments=[adjust.Multiply(13)],
)
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 0.1 * 7 * (1 + 3 + 5) * 13
def test_import_flow_get_net_flow_with_adjust():
flow = ImportFlow(
name="flow",
dest=Compartment("S"),
param=lambda t: 0.1 * t,
adjustments=[adjust.Multiply(13)],
)
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 0.1 * 7 * 13
def test_create_model():
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert_array_equal(model.times, np.array([0, 1, 2, 3, 4, 5]))
assert model.compartments == [Compartment("S"), Compartment("I"), Compartment("R")]
assert model._infectious_compartments == [Compartment("I")]
assert_array_equal(model.initial_population, np.array([0, 0, 0]))
# Times out of order
with pytest.raises(AssertionError):
CompartmentalModel(
times=[5, 0], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
# Infectious compartment not a compartment
with pytest.raises(AssertionError):
CompartmentalModel(
times=[-1, 5], compartments=["S", "I", "R"], infectious_compartments=["E"]
)
def test_death_flow_get_net_flow():
flow = DeathFlow(
name="flow",
source=Compartment("I"),
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 == 2 * 3 * 7
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_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_function_flow_get_net_flow_with_adjust():
def get_flow_rate(flow, comps, comp_vals, flows, flow_rates, time):
i_pop = sum([comp_vals[c.idx] for c in comps if c.name == "I"])
s_flows = sum([flow_rates[c.idx] for c in comps if c.name == "S"])
return s_flows * i_pop * time
flow = FunctionFlow(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=get_flow_rate,
adjustments=[adjust.Multiply(13)],
)
flow.source.idx = 1
compartments = [Compartment("S"), Compartment("I"), Compartment("R")]
for i, c in enumerate(compartments):
c.idx = i
compartment_values = np.array([1, 3, 5])
flow_rates = np.array([2, 7, 13])
net_flow = flow.get_net_flow(compartments, compartment_values, [flow],
flow_rates, 7)
assert net_flow == 2 * 3 * 7 * 13
def test_entry_flow_stratify_with_ageing():
strat = AgeStratification(
name="age",
strata=["0", "1", "2"],
compartments=["I", "R"],
)
flow = EntryFlow(
name="birth",
dest=Compartment("I"),
param=2,
adjustments=[adjust.Overwrite(0.2)],
)
flow._is_birth_flow = False # Not marked as a birth flow!
new_flows = flow.stratify(strat)
assert len(new_flows) == 3 # So the birth flow rules don't apply.
flow._is_birth_flow = True # Marked as a birth flow.
new_flows = flow.stratify(strat)
assert len(new_flows) == 1 # So the birth flow rules apply.
# Only age 0 babies get born.
assert new_flows[0]._is_equal(
EntryFlow(
name="birth",
param=2,
dest=Compartment("I", {"age": "0"}),
adjustments=[adjust.Overwrite(0.2)],
))
# Expect this to fail coz you can't adjust birth flows for age stratifications.
strat.add_flow_adjustments("birth", {
"0": adjust.Multiply(0.1),
"1": None,
"2": None
})
with pytest.raises(AssertionError):
flow.stratify(strat)
def test_entry_flow_stratify_with_adjustments():
flow = EntryFlow(
name="flow",
dest=Compartment("I"),
param=2,
adjustments=[adjust.Overwrite(0.2)],
)
strat = Stratification(
name="location",
strata=["1", "2"],
compartments=["I", "R"],
)
strat.add_flow_adjustments("flow", {
"1": adjust.Multiply(0.1),
"2": adjust.Multiply(0.3)
})
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0]._is_equal(
EntryFlow(
name="flow",
param=2,
dest=Compartment("I", {"location": "1"}),
adjustments=[adjust.Overwrite(0.2),
adjust.Multiply(0.1)],
))
assert new_flows[1]._is_equal(
EntryFlow(
name="flow",
param=2,
dest=Compartment("I", {"location": "2"}),
adjustments=[adjust.Overwrite(0.2),
adjust.Multiply(0.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 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_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_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_transition_flow_stratify_source_but_not_dest():
"""
Ensure two parallel flows created, no adjustments required.
"""
flow = TransitionFlow(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=2,
adjustments=[],
)
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["S", "I"],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0]._is_equal(
TransitionFlow(
name="flow",
param=2,
source=Compartment("I", {"age": "1"}),
dest=Compartment("R"),
adjustments=[],
)
)
assert new_flows[1]._is_equal(
TransitionFlow(
name="flow",
param=2,
source=Compartment("I", {"age": "2"}),
dest=Compartment("R"),
adjustments=[],
)
)
def test_transition_flow_stratify_dest_but_not_source():
"""
Ensure two new flows are created and automatically adjusted to account for fan out.
"""
flow = TransitionFlow(
name="flow",
source=Compartment("I"),
dest=Compartment("R"),
param=2,
adjustments=[],
)
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["S", "R"],
)
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.5)],
)
)
assert new_flows[1]._is_equal(
TransitionFlow(
name="flow",
param=2,
source=Compartment("I"),
dest=Compartment("R", {"age": "2"}),
adjustments=[adjust.Multiply(0.5)],
)
)
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_get_flow_adjustments__with_strata_whitelist():
# Latest matching flow adjustment should always win.
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
strat.add_flow_adjustments("flow", {"rural": Multiply(1), "urban": None})
strat.add_flow_adjustments("flow", {
"rural": Multiply(3),
"urban": Overwrite(2)
})
strat.add_flow_adjustments("flow", {
"rural": Multiply(2),
"urban": Overwrite(1)
},
source_strata={"age": "20"})
# No source strata
entry_flow = EntryFlow("flow", Compartment("S"), 1)
with pytest.raises(AssertionError):
strat.get_flow_adjustment(entry_flow)
other_flow = TransitionFlow("other", Compartment("S"), Compartment("I"), 1)
assert strat.get_flow_adjustment(other_flow) is None
trans_flow = TransitionFlow("flow", Compartment("S"), Compartment("I"), 1)
exit_flow = ExitFlow("flow", Compartment("I"), 1)
for flow in [trans_flow, exit_flow]:
adj = strat.get_flow_adjustment(flow)
assert adj["rural"]._is_equal(Multiply(3))
assert adj["urban"]._is_equal(Overwrite(2))
# Only flows with matching strata should get the adjustment
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
strat.add_flow_adjustments("flow", {"rural": Multiply(1), "urban": None})
strat.add_flow_adjustments("flow", {
"rural": Multiply(3),
"urban": Overwrite(2)
})
strat.add_flow_adjustments("flow", {
"rural": Multiply(2),
"urban": Overwrite(1)
},
dest_strata={"age": "20"})
# No dest strata
exit_flow = ExitFlow("flow", Compartment("I"), 1)
with pytest.raises(AssertionError):
strat.get_flow_adjustment(exit_flow)
# No matching dest strata
other_flow = TransitionFlow("other", Compartment("S"), Compartment("I"), 1)
other_flow_strat = TransitionFlow("other", Compartment("S"),
Compartment("I", {"age": "20"}), 1)
assert strat.get_flow_adjustment(other_flow) is None
assert strat.get_flow_adjustment(other_flow_strat) is None
# Flows without age 20 get the last match.
trans_flow = TransitionFlow("flow", Compartment("S"), Compartment("I"), 1)
entry_flow = EntryFlow("flow", Compartment("S"), 1)
trans_flow_strat_wrong = TransitionFlow("flow", Compartment("S"),
Compartment("I", {"age": "10"}), 1)
entry_flow_strat_wrong = EntryFlow("flow", Compartment("S", {"age": "10"}),
1)
trans_flow_strat_wrong_2 = TransitionFlow("flow",
Compartment("S", {"age": "20"}),
Compartment("I"), 1)
for flow in [
trans_flow,
entry_flow,
trans_flow_strat_wrong,
entry_flow_strat_wrong,
trans_flow_strat_wrong_2,
]:
adj = strat.get_flow_adjustment(flow)
assert adj["rural"]._is_equal(Multiply(3))
assert adj["urban"]._is_equal(Overwrite(2))
trans_flow_strat = TransitionFlow("flow", Compartment("S"),
Compartment("I", {"age": "20"}), 1)
entry_flow_strat = EntryFlow("flow", Compartment("S", {"age": "20"}), 1)
for flow in [trans_flow_strat, entry_flow_strat]:
adj = strat.get_flow_adjustment(flow)
assert adj["rural"]._is_equal(Multiply(2))
assert adj["urban"]._is_equal(Overwrite(1))
# The last created matching flow adjustment will win, also include both source and dest.
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
strat.add_flow_adjustments("flow", {"rural": Multiply(1), "urban": None})
strat.add_flow_adjustments(
"flow",
{
"rural": Multiply(5),
"urban": Overwrite(7)
},
source_strata={"age": "20"},
dest_strata={
"age": "30",
"work": "home"
},
)
strat.add_flow_adjustments(
"flow",
{
"rural": Multiply(2),
"urban": Overwrite(1)
},
source_strata={"age": "20"},
dest_strata={"age": "30"},
)
strat.add_flow_adjustments("flow", {
"rural": Multiply(3),
"urban": Overwrite(2)
})
# Missing source strata
trans_flow_strat_wrong = TransitionFlow(
"flow", Compartment("S"),
Compartment("I", {
"age": "30",
"work": "home"
}), 1)
# Missing dest strata
trans_flow_strat_wrong_2 = TransitionFlow("flow",
Compartment("S", {"age": "20"}),
Compartment("I"), 1)
# Incomplete dest strata - less specific still wins because of ordering.
trans_flow_strat_wrong_3 = TransitionFlow("flow",
Compartment("S", {"age": "20"}),
Compartment("I", {"age": "30"}),
1)
for flow in [
trans_flow_strat_wrong, trans_flow_strat_wrong_2,
trans_flow_strat_wrong_3
]:
adj = strat.get_flow_adjustment(flow)
assert adj["rural"]._is_equal(Multiply(3))
assert adj["urban"]._is_equal(Overwrite(2))
# Match to the last created stratification
trans_flow_strat = TransitionFlow(
"flow", Compartment("S", {"age": "20"}),
Compartment("I", {
"age": "30",
"work": "home"
}), 1)
for flow in [trans_flow_strat]:
adj = strat.get_flow_adjustment(flow)
assert adj["rural"]._is_equal(Multiply(3))
assert adj["urban"]._is_equal(Overwrite(2))
def test_is_match():
"""
Ensure is_match is behaving sensibly.
TODO: Come up with a more rigorous, parametrized check.
"""
source = Compartment("S")
dest = Compartment("I")
trans_flow = TransitionFlow("x", source, dest, None, None)
assert trans_flow.is_match("x", source_strata={}, dest_strata={})
assert not trans_flow.is_match("y", source_strata={}, dest_strata={})
assert not trans_flow.is_match(
"x", source_strata={"thing": "good"}, dest_strata={})
assert not trans_flow.is_match(
"x", source_strata={}, dest_strata={"thing": "good"})
entry_flow = EntryFlow("x", dest, None, None)
assert entry_flow.is_match("x", source_strata={}, dest_strata={})
assert not entry_flow.is_match("y", source_strata={}, dest_strata={})
assert entry_flow.is_match("x",
source_strata={"thing": "good"},
dest_strata={})
assert not entry_flow.is_match(
"x", source_strata={}, dest_strata={"thing": "good"})
exit_flow = ExitFlow("x", source, None, None)
assert exit_flow.is_match("x", source_strata={}, dest_strata={})
assert not exit_flow.is_match("y", source_strata={}, dest_strata={})
assert not exit_flow.is_match(
"x", source_strata={"thing": "good"}, dest_strata={})
assert exit_flow.is_match("x",
source_strata={},
dest_strata={"thing": "good"})
source = Compartment("S", {"thing": "good", "color": "red"})
dest = Compartment("I", {"thing": "good", "color": "red"})
trans_flow = TransitionFlow("x", source, dest, None, None)
assert trans_flow.is_match("x", source_strata={}, dest_strata={})
assert not trans_flow.is_match("y", source_strata={}, dest_strata={})
assert trans_flow.is_match("x",
source_strata={"thing": "good"},
dest_strata={})
assert trans_flow.is_match("x",
source_strata={},
dest_strata={"thing": "good"})
assert trans_flow.is_match("x",
source_strata={"thing": "good"},
dest_strata={"color": "red"})
assert trans_flow.is_match("x",
source_strata={
"thing": "good",
"color": "red"
},
dest_strata={})
assert trans_flow.is_match("x",
source_strata={},
dest_strata={
"thing": "good",
"color": "red"
})
assert not trans_flow.is_match(
"x", source_strata={"thing": "good"}, dest_strata={"color": "green"})
assert not trans_flow.is_match(
"x", source_strata={"thing": "bad"}, dest_strata={})
assert not trans_flow.is_match(
"x", source_strata={}, dest_strata={"thing": "bad"})
entry_flow = EntryFlow("x", dest, None, None)
assert entry_flow.is_match("x", source_strata={}, dest_strata={})
assert not entry_flow.is_match("y", source_strata={}, dest_strata={})
assert entry_flow.is_match("x",
source_strata={"thing": "good"},
dest_strata={})
assert entry_flow.is_match("x",
source_strata={},
dest_strata={"thing": "good"})
assert entry_flow.is_match("x",
source_strata={"thing": "bad"},
dest_strata={})
assert not entry_flow.is_match(
"x", source_strata={}, dest_strata={"thing": "bad"})
exit_flow = ExitFlow("x", source, None, None)
assert exit_flow.is_match("x", source_strata={}, dest_strata={})
assert not exit_flow.is_match("y", source_strata={}, dest_strata={})
assert exit_flow.is_match("x",
source_strata={"thing": "good"},
dest_strata={})
assert exit_flow.is_match("x",
source_strata={},
dest_strata={"thing": "good"})
assert not exit_flow.is_match(
"x", source_strata={"thing": "bad"}, dest_strata={})
assert exit_flow.is_match("x",
source_strata={},
dest_strata={"thing": "bad"})
assert exit_flow.is_match("x", source_strata={}, dest_strata={})
assert not exit_flow.is_match("y", source_strata={}, dest_strata={})
assert exit_flow.is_match("x",
source_strata={"thing": "good"},
dest_strata={})
assert exit_flow.is_match("x",
source_strata={},
dest_strata={"thing": "good"})
assert not exit_flow.is_match(
"x", source_strata={"thing": "bad"}, dest_strata={})
assert exit_flow.is_match("x",
source_strata={},
dest_strata={"thing": "bad"})
SOURCE = Compartment("source")
DEST = Compartment("dest")
def test_update_compartment_indices():
mapping = {"source": 2, "dest": 7}
trans_flow = TransitionFlow("x", SOURCE, DEST, None, None)
trans_flow.update_compartment_indices(mapping)
assert trans_flow.source.idx == 2
assert trans_flow.dest.idx == 7
entry_flow = EntryFlow("x", DEST, None, None)
entry_flow.update_compartment_indices(mapping)
assert entry_flow.dest.idx == 7
def test_stratify_compartments__with_extisting_strat():
age_strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S", "I", "R"],
)
comps = [Compartment("S"), Compartment("I"), Compartment("R")]
age_comps = age_strat._stratify_compartments(comps)
loc_strat = Stratification(
name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"],
)
loc_comps = loc_strat._stratify_compartments(age_comps)
assert loc_comps == [
Compartment("S", {
"age": "0",
"location": "rural"
}),
Compartment("S", {
"age": "0",
"location": "urban"
}),
Compartment("S", {
"age": "10",
"location": "rural"
}),
Compartment("S", {
"age": "10",
"location": "urban"
}),
Compartment("S", {
"age": "20",
"location": "rural"
}),
Compartment("S", {
"age": "20",
"location": "urban"
}),
Compartment("I", {
"age": "0",
"location": "rural"
}),
Compartment("I", {
"age": "0",
"location": "urban"
}),
Compartment("I", {
"age": "10",
"location": "rural"
}),
Compartment("I", {
"age": "10",
"location": "urban"
}),
Compartment("I", {
"age": "20",
"location": "rural"
}),
Compartment("I", {
"age": "20",
"location": "urban"
}),
Compartment("R", {
"age": "0",
"location": "rural"
}),
Compartment("R", {
"age": "0",
"location": "urban"
}),
Compartment("R", {
"age": "10",
"location": "rural"
}),
Compartment("R", {
"age": "10",
"location": "urban"
}),
Compartment("R", {
"age": "20",
"location": "rural"
}),
Compartment("R", {
"age": "20",
"location": "urban"
}),
]
def test_optimize_adjustments(adjustments, optimized):
flow = EntryFlow("flow", Compartment("S"), 1, adjustments)
flow.optimize_adjustments()
assert len(flow.adjustments) == len(optimized)
assert all([a._is_equal(o) for a, o in zip(flow.adjustments, optimized)])
