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