def test_entry_flow_stratify__with_no_flow_adjustments():
strat = Stratification(
name="location",
strata=["1", "2"],
compartments=["infect", "recovery"],
comp_split_props={},
flow_adjustments={},
)
flow = EntryFlow(
dest=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0].param_name == "foo"
assert new_flows[0].param_func == _get_param_value
assert new_flows[0].adjustments == [(FlowAdjustment.MULTIPLY, 0.5)]
assert new_flows[0].dest == Compartment("infect",
strat_names=["location"],
strat_values={"location": "1"})
assert new_flows[1].param_name == "foo"
assert new_flows[1].param_func == _get_param_value
assert new_flows[1].adjustments == [(FlowAdjustment.MULTIPLY, 0.5)]
assert new_flows[1].dest == Compartment("infect",
strat_names=["location"],
strat_values={"location": "2"})
def test_exit_flow_stratify__with_no_flow_adjustments():
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["infect", "recovery"],
comp_split_props={},
flow_adjustments={},
)
flow = ExitFlow(
source=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0].param_name == "foo"
assert new_flows[0].param_func == _get_param_value
assert new_flows[0].adjustments == []
assert new_flows[0].source == Compartment("infect",
strat_names=["age"],
strat_values={"age": "1"})
assert new_flows[1].param_name == "foo"
assert new_flows[1].param_func == _get_param_value
assert new_flows[1].adjustments == []
assert new_flows[1].source == Compartment("infect",
strat_names=["age"],
strat_values={"age": "2"})
def test_stratification_with_basic_setup():
"""
Ensure we can create a simple stratification.
"""
strat = Stratification(
name="foo",
strata=[1, "2", "bar", "baz"],
compartments=["sus", "inf", "rec"],
comp_split_props={},
flow_adjustments={},
)
assert strat.name == "foo"
assert strat.compartments == [
Compartment("sus"),
Compartment("inf"),
Compartment("rec")
]
assert strat.strata == ["1", "2", "bar", "baz"]
assert strat.comp_split_props == {
"1": 0.25,
"2": 0.25,
"bar": 0.25,
"baz": 0.25
}
assert strat.flow_adjustments == {}
def test_get_infection_density_multiplier():
model = EpiModel(**MODEL_KWARGS)
model.prepare_to_run()
model.prepare_time_step(0, model.compartment_values)
c_src = Compartment("S")
c_dst = Compartment("I")
multiplier = model.get_infection_density_multiplier(c_src, c_dst)
assert model.population_infectious == 10
assert model.population_total == 1000
assert multiplier == 10
def test_serialize():
c = Compartment(
"infected",
strat_names=["location", "age"],
strat_values={
"location": "hawaii",
"age": "15"
},
)
assert c.serialize() == "infectedXlocation_hawaiiXage_15"
def test_standard_flow_get_net_flow():
flow = StandardFlow(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
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_entry_flow_stratify_with_ageing():
strat = Stratification(
name="age",
strata=["0", "1", "2"],
compartments=["infect", "recovery"],
comp_split_props={},
flow_adjustments={},
)
# Expect these to be ignored in favour of birth specific adjustments.
strat.flow_adjustments = {
"foo": [{
"strata": {},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.1)
}
}]
}
flow = EntryFlow(
dest=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[(FlowAdjustment.OVERWRITE, 0.2)],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 3
assert new_flows[0].param_name == "foo"
assert new_flows[0].param_func == _get_param_value
assert new_flows[0].adjustments == [
(FlowAdjustment.OVERWRITE, 0.2),
(FlowAdjustment.MULTIPLY, 1),
]
assert new_flows[0].dest == Compartment("infect",
strat_names=["age"],
strat_values={"age": "0"})
assert new_flows[1].param_name == "foo"
assert new_flows[1].param_func == _get_param_value
assert new_flows[1].adjustments == [
(FlowAdjustment.OVERWRITE, 0.2),
(FlowAdjustment.OVERWRITE, 0),
]
assert new_flows[1].dest == Compartment("infect",
strat_names=["age"],
strat_values={"age": "1"})
assert new_flows[2].param_name == "foo"
assert new_flows[2].param_func == _get_param_value
assert new_flows[2].adjustments == [
(FlowAdjustment.OVERWRITE, 0.2),
(FlowAdjustment.OVERWRITE, 0),
]
assert new_flows[2].dest == Compartment("infect",
strat_names=["age"],
strat_values={"age": "2"})
def test_standard_flow_get_net_flow_with_adjust():
flow = StandardFlow(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
adjustments=[(FlowAdjustment.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_infect_density_get_net_flow_with_adjust(FlowClass):
flow = FlowClass(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
find_infectious_multiplier=find_infectious_multiplier,
adjustments=[(FlowAdjustment.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 * 23
def test_infect_get_net_flow(FlowClass):
flow = FlowClass(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
find_infectious_multiplier=find_infectious_multiplier,
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 * 23
def test_stratify():
c = Compartment("infected")
assert c._name == "infected"
assert c._strat_names == tuple()
assert c._strat_values == {}
assert c.serialize() == "infected"
c_age = c.stratify("age", "15")
assert c_age._name == "infected"
assert c_age._strat_names == ("age", )
assert c_age._strat_values == {"age": "15"}
assert c_age.serialize() == "infectedXage_15"
c_loc = c_age.stratify("location", "work")
assert c_loc._name == "infected"
assert c_loc._strat_names == ("age", "location")
assert c_loc._strat_values == {"age": "15", "location": "work"}
assert c_loc.serialize() == "infectedXage_15Xlocation_work"
def _add_flow(self, requested_flow):
flow = None
f_source = requested_flow.get("origin")
f_source = Compartment.deserialize(f_source) if f_source else None
f_dest = requested_flow.get("to")
f_dest = Compartment.deserialize(f_dest) if f_dest else None
f_param = requested_flow["parameter"]
f_type = requested_flow["type"]
if f_type == Flow.STANDARD:
flow = StandardFlow(
source=f_source,
dest=f_dest,
param_name=f_param,
param_func=self.get_parameter_value,
)
elif f_type == Flow.INFECTION_FREQUENCY:
flow = InfectionFrequencyFlow(
source=f_source,
dest=f_dest,
param_name=f_param,
param_func=self.get_parameter_value,
find_infectious_multiplier=self.get_infection_frequency_multiplier,
)
elif f_type == Flow.INFECTION_DENSITY:
flow = InfectionDensityFlow(
source=f_source,
dest=f_dest,
param_name=f_param,
param_func=self.get_parameter_value,
find_infectious_multiplier=self.get_infection_density_multiplier,
)
elif f_type == Flow.DEATH:
flow = InfectionDeathFlow(
source=f_source, param_name=f_param, param_func=self.get_parameter_value
)
elif f_type == Flow.IMPORT:
flow = ImportFlow(
dest=self.entry_compartment,
param_name=f_param,
param_func=self.get_parameter_value,
)
if flow:
self.flows.append(flow)
return flow
def test_transition_flow_stratify_with_no_matching_compartments():
strat = Stratification(
name="age",
strata=["1", "2", "3"],
compartments=["recovery"], # Flow compartments not included
comp_split_props={},
flow_adjustments={},
)
flow = TransitionFlow(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert new_flows == [flow]
def test_get_stratified_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])
comp_names = [
Compartment("S", strat_names=["age"], strat_values={"age": "0"}),
Compartment("S", strat_names=["age"], strat_values={"age": "10"}),
Compartment("S", strat_names=["age"], strat_values={"age": "20"}),
Compartment("I", strat_names=["age"], strat_values={"age": "0"}),
Compartment("I", strat_names=["age"], strat_values={"age": "10"}),
Compartment("I", strat_names=["age"], strat_values={"age": "20"}),
Compartment("R", strat_names=["age"], strat_values={"age": "0"}),
Compartment("R", strat_names=["age"], strat_values={"age": "10"}),
Compartment("R", strat_names=["age"], strat_values={"age": "20"}),
]
strat = Stratification(
name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"],
comp_split_props={
"rural": 0.1,
"urban": 0.9
},
flow_adjustments={},
)
new_comp_values = get_stratified_compartment_values(
strat, comp_names, 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_replace_deaths_birth_flow_get_net_flow():
flow = ReplacementBirthFlow(
dest=Compartment("susceptible"),
get_total_deaths=_get_total_deaths,
adjustments=[],
)
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 23
def test_get_stratified_compartment_values__with_subset_stratified():
strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S"],
comp_split_props={
"0": 0.25,
"10": 0.5,
"20": 0.25
},
flow_adjustments={},
)
comp_names = [Compartment("S"), Compartment("I"), Compartment("R")]
comp_values = np.array([1000.0, 100.0, 0.0])
new_comp_values = get_stratified_compartment_values(
strat, comp_names, comp_values)
expected_arr = np.array([250.0, 500.0, 250.0, 100.0, 0.0])
assert_array_equal(expected_arr, new_comp_values)
def test_replace_deaths_birth_flow_get_net_flow_with_adjust():
flow = ReplacementBirthFlow(
dest=Compartment("susceptible"),
get_total_deaths=_get_total_deaths,
adjustments=[(FlowAdjustment.MULTIPLY, 13)],
)
flow.dest.idx = 2
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():
flow = CrudeBirthFlow(
dest=Compartment("susceptible"),
param_name="crude_birth_rate",
param_func=_get_param_value,
adjustments=[],
)
vals = np.array([1, 3, 5])
net_flow = flow.get_net_flow(vals, 7)
assert net_flow == 0.1 * 7 * (1 + 3 + 5)
def test_setup__with_double_strat():
c = Compartment(
"infected",
strat_names=["location", "age"],
strat_values={
"location": "hawaii",
"age": "15"
},
)
assert str(c) == "infectedXlocation_hawaiiXage_15"
assert hash(c) == hash("infectedXlocation_hawaiiXage_15")
def test_death_flow_get_net_flow(FlowClass):
flow = FlowClass(
source=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
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_death_flow_get_net_flow_with_adjust(FlowClass):
flow = FlowClass(
source=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[(FlowAdjustment.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_deserialzie():
s = "infectedXlocation_hawaiiXage_15"
c = Compartment.deserialize(s)
assert c == Compartment(
"infected",
strat_names=["location", "age"],
strat_values={
"location": "hawaii",
"age": "15"
},
)
assert c != Compartment(
"infected",
strat_names=["age", "location"],
strat_values={
"location": "hawaii",
"age": "15"
},
)
assert str(c) == "infectedXlocation_hawaiiXage_15"
assert hash(c) == hash("infectedXlocation_hawaiiXage_15")
def test_exit_flow_stratify_with_flow_adjustments():
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["infect", "recovery"],
comp_split_props={},
flow_adjustments={},
)
strat.flow_adjustments = {
"foo": [{
"strata": {},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.1)
}
}]
}
flow = ExitFlow(
source=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[(FlowAdjustment.OVERWRITE, 0.2)],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0].param_name == "foo"
assert new_flows[0].param_func == _get_param_value
assert new_flows[0].adjustments == [
(FlowAdjustment.OVERWRITE, 0.2),
(FlowAdjustment.MULTIPLY, 0.1),
]
assert new_flows[0].source == Compartment("infect",
strat_names=["age"],
strat_values={"age": "1"})
assert new_flows[1].param_name == "foo"
assert new_flows[1].param_func == _get_param_value
assert new_flows[1].adjustments == [(FlowAdjustment.OVERWRITE, 0.2)]
assert new_flows[1].source == Compartment("infect",
strat_names=["age"],
strat_values={"age": "2"})
def test_entry_flow_stratify__when_not_applicable():
strat = Stratification(
name="location",
strata=["1", "2", "3"],
compartments=["recovery"], # infect compartment not included
comp_split_props={},
flow_adjustments={},
)
flow = EntryFlow(
dest=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert new_flows == [flow]
def test_exit_flow_stratify__when_no_compartment_match():
strat = Stratification(
name="age",
strata=["1", "2", "3"],
compartments=["recovery"], # infect compartment not included
comp_split_props={},
flow_adjustments={},
)
flow = ExitFlow(
source=Compartment("infect"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert new_flows == [flow]
def __init__(
self,
name: str,
strata: List[str],
compartments: List[str],
comp_split_props: Dict[str, float],
flow_adjustments: Dict[str, Dict[str, float]],
):
self.name = name
self.strata = list(map(str, strata))
self.comp_split_props = get_all_proportions(self.strata,
comp_split_props)
self.flow_adjustments = parse_flow_adjustments(flow_adjustments)
self.compartments = [
Compartment(c) if type(c) is str else c for c in compartments
]
def test_get_stratified_compartment_names__with_no_extisting_strat_and_subset_only(
):
strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S"],
comp_split_props={},
flow_adjustments={},
)
comp_names = [Compartment("S"), Compartment("I"), Compartment("R")]
strat_comp_names = get_stratified_compartment_names(strat, comp_names)
assert strat_comp_names == [
Compartment("S", strat_names=["age"], strat_values={"age": "0"}),
Compartment("S", strat_names=["age"], strat_values={"age": "10"}),
Compartment("S", strat_names=["age"], strat_values={"age": "20"}),
Compartment("I"),
Compartment("R"),
]
def __init__(
self,
times: List[float],
compartment_names: List[str],
initial_conditions: Dict[str, float],
parameters: Dict[str, float],
requested_flows: List[dict],
infectious_compartments: List[str],
birth_approach: str,
entry_compartment: str,
starting_population: int,
):
super().__init__(
times=times,
compartment_names=compartment_names,
initial_conditions=initial_conditions,
parameters=parameters,
requested_flows=requested_flows,
infectious_compartments=infectious_compartments,
birth_approach=birth_approach,
entry_compartment=entry_compartment,
starting_population=starting_population,
)
# Keeps track of Stratifications that have been applied.
self.stratifications = []
# Keeps track of original, pre-stratified compartment names.
self.original_compartment_names = [Compartment.deserialize(n) for n in compartment_names]
# A list of the different sub-categories ('strains') of the diease that we are modelling.
self.disease_strains = [self.DEFAULT_DISEASE_STRAIN]
# Strata-based multipliers for compartmental infectiousness levels.
self.infectiousness_levels = {}
# Compartment-based overwrite values for compartmental infectiousness levels.
self.individual_infectiousness_adjustments = []
# Mixing matrices a list of NxN arrays used to calculate force of infection.
self.mixing_matrices = []
# A list of dicts that has the strata required to match a row in the mixing matrix.
self.mixing_categories = [{}]
def test_transition_flow_stratify_dest_but_not_source__with_flow_adjustments():
"""
Ensure flow is adjusted to account for fan out.
"""
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["happy", "recovery"],
comp_split_props={},
flow_adjustments={},
)
strat.flow_adjustments = {
"foo": [{
"strata": {},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.1)
}
}]
}
flow = TransitionFlow(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0].param_name == "foo"
assert new_flows[0].param_func == _get_param_value
assert new_flows[0].adjustments == [(FlowAdjustment.MULTIPLY, 0.1)]
assert new_flows[0].source == Compartment("infect",
strat_names=[],
strat_values={})
assert new_flows[0].dest == Compartment("happy",
strat_names=["age"],
strat_values={"age": "1"})
assert new_flows[1].param_name == "foo"
assert new_flows[1].param_func == _get_param_value
assert new_flows[1].adjustments == [(FlowAdjustment.MULTIPLY, 0.5)]
assert new_flows[1].source == Compartment("infect",
strat_names=[],
strat_values={})
assert new_flows[1].dest == Compartment("happy",
strat_names=["age"],
strat_values={"age": "2"})
def test_transition_flow_stratify_source_and_dest():
strat = Stratification(
name="age",
strata=["1", "2"],
compartments=["infect", "happy", "recovery"],
comp_split_props={},
flow_adjustments={},
)
flow = TransitionFlow(
source=Compartment("infect"),
dest=Compartment("happy"),
param_name="foo",
param_func=_get_param_value,
adjustments=[],
)
new_flows = flow.stratify(strat)
assert len(new_flows) == 2
assert new_flows[0].param_name == "foo"
assert new_flows[0].param_func == _get_param_value
assert new_flows[0].adjustments == []
assert new_flows[0].source == Compartment("infect",
strat_names=["age"],
strat_values={"age": "1"})
assert new_flows[0].dest == Compartment("happy",
strat_names=["age"],
strat_values={"age": "1"})
assert new_flows[1].param_name == "foo"
assert new_flows[1].param_func == _get_param_value
assert new_flows[1].adjustments == []
assert new_flows[1].source == Compartment("infect",
strat_names=["age"],
strat_values={"age": "2"})
assert new_flows[1].dest == Compartment("happy",
strat_names=["age"],
strat_values={"age": "2"})
