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_get_flow_adjustment(param_name, stratum, comp, adjustment):
"""
Ensure we can create a stratification that parses flow adjustments.
"""
strat = Stratification(
name="foo",
strata=["1", "2", "3"],
compartments=["sus", "inf", "rec"],
comp_split_props={},
flow_adjustments={},
)
strat.flow_adjustments = {
"infect_death": [
{
"strata": {},
"adjustments": {
"2": (FlowAdjustment.MULTIPLY, 0.5)
}
},
{
"strata": {
"location": "work"
},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.9)
}
},
],
"contact_rate": [
{
"strata": {
"age": "10"
},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.5),
"2": (FlowAdjustment.OVERWRITE, 0.2),
},
},
{
"strata": {
"age": "20"
},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.6),
"3": (FlowAdjustment.COMPOSE, "some_function"),
},
},
],
}
actual_adj = strat.get_flow_adjustment(comp, stratum, param_name)
assert actual_adj == adjustment
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_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_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_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 create(
strat: Stratification,
compartments: List[Compartment],
param_func: Callable[[str, float], float],
):
"""
Create inter-compartmental flows for ageing from one stratum to the next.
The ageing rate is proportional to the width of the age bracket.
It's assumed that both ages and model timesteps are in years.
"""
assert strat.is_ageing()
ageing_flows = []
ageing_params = {}
ages = list(sorted(map(int, strat.strata)))
for age_idx in range(len(ages) - 1):
start_age = int(ages[age_idx])
end_age = int(ages[age_idx + 1])
param_name = f"ageing{start_age}to{end_age}"
ageing_rate = 1.0 / (end_age - start_age)
ageing_params[param_name] = ageing_rate
for comp in compartments:
if not comp.has_name_in_list(strat.compartments):
# Don't include unstratified compartments
continue
flow = AgeingFlow(
source=comp.stratify(strat.name, str(start_age)),
dest=comp.stratify(strat.name, str(end_age)),
param_name=param_name,
param_func=param_func,
)
ageing_flows.append(flow)
return ageing_flows, ageing_params
def stratify(self, strat: Stratification) -> List[BaseFlow]:
"""
Returns a list of new, stratified exit flows to replace the current flow.
"""
if not self.source.has_name_in_list(strat.compartments):
# Flow source is not stratified, do not stratify this flow.
return [self]
new_flows = []
for stratum in strat.strata:
adjustment = strat.get_flow_adjustment(self.source, stratum,
self.param_name)
if adjustment:
new_adjustments = [*self.adjustments, adjustment]
else:
new_adjustments = self.adjustments
new_source = self.source.stratify(strat.name, stratum)
new_flow = self.copy(
source=new_source,
param_name=self.param_name,
param_func=self.param_func,
adjustments=new_adjustments,
)
new_flows.append(new_flow)
return new_flows
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 stratify(self, strat: Stratification) -> List[BaseFlow]:
"""
Returns a list of new, stratified entry flows to replace the current flow.
"""
if not self.dest.has_name_in_list(strat.compartments):
# Flow destination is not stratified, do not stratify this flow.
return [self]
new_flows = []
for stratum in strat.strata:
adjustment = None
if strat.is_ageing():
# Use special rules for ageing.
if stratum == "0":
# Babies get born at age 0, add a null op to prevent default behaviour.
adjustment = (FlowAdjustment.MULTIPLY, 1)
else:
# Babies do not get born at any other age.
adjustment = (FlowAdjustment.OVERWRITE, 0)
else:
# Not an ageing stratification, check for user-specified flow adjustments.
adjustment = strat.get_flow_adjustment(self.dest, stratum, self.param_name)
if not adjustment:
# Default to equally dividing entry population between all strata.
num_strata = len(strat.strata)
entry_fraction = 1.0 / num_strata
adjustment = (FlowAdjustment.MULTIPLY, entry_fraction)
new_adjustments = [*self.adjustments, adjustment]
new_dest = self.dest.stratify(strat.name, stratum)
new_flow = self.copy(
dest=new_dest,
param_name=self.param_name,
param_func=self.param_func,
adjustments=new_adjustments,
)
new_flows.append(new_flow)
return new_flows
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_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 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_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_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_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 test_stratification_with_compartment_split_proportions():
"""
Ensure we can create a stratification that parses compartment splits.
"""
strat = Stratification(
name="foo",
strata=["1", "2", "3"],
compartments=["sus", "inf", "rec"],
comp_split_props={
"1": 0.2,
"2": 0.2
},
flow_adjustments={},
)
assert strat.comp_split_props == {
"1": 0.2,
"2": 0.2,
"3": 0.6,
}
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"})
def stratify(self, strat: Stratification) -> List[BaseFlow]:
"""
Returns a list of new, stratified flows to replace the current flow.
"""
is_source_strat = self.source.has_name_in_list(strat.compartments)
is_dest_strat = self.dest.has_name_in_list(strat.compartments)
if not (is_dest_strat or is_source_strat):
# Flow is not stratified, do not stratify this flow.
return [self]
new_flows = []
for stratum in strat.strata:
# Find new compartments
if is_source_strat:
new_source = self.source.stratify(strat.name, stratum)
else:
new_source = self.source
if is_dest_strat:
new_dest = self.dest.stratify(strat.name, stratum)
else:
new_dest = self.dest
# Find flow adjustments to apply to the new stratified flows.
# First, we try to find an adjustment for the source compartment.
# This is for when the source has the required stratifications and the destination does not.
# For example - people recovering from I -> R with multiple I strata, all with different recovery rates.
adjustment = strat.get_flow_adjustment(self.source, stratum,
self.param_name)
if not adjustment:
# Otherwise, try find an adjustment for the destination compartment.
# This is for when the destination has the required stratifications and the source does not.
# For example - people recovering from I -> R with multiple R strata, with different recovery proportions.
adjustment = strat.get_flow_adjustment(self.dest, stratum,
self.param_name)
# Should we apply an adjustment to conserve the number of people?
should_apply_conservation_split = ((not strat.is_strain())
and (not adjustment)
and (is_dest_strat
and not is_source_strat))
if should_apply_conservation_split:
# If the source is stratified but not the destination, then we need to account
# for the resulting fan-out of flows by reducing the flow rate.
# We don't do this for strains because this effect is already
# captured by the infecitousness multiplier.
num_strata = len(strat.strata)
entry_fraction = 1.0 / num_strata
adjustment = (FlowAdjustment.MULTIPLY, entry_fraction)
if adjustment:
new_adjustments = [*self.adjustments, adjustment]
else:
new_adjustments = self.adjustments
new_flow = self.copy(
source=new_source,
dest=new_dest,
param_name=self.param_name,
param_func=self.param_func,
adjustments=new_adjustments,
)
new_flows.append(new_flow)
return new_flows
def test_stratification_with_flow_adjustments():
"""
Ensure we can create a stratification that parses flow adjustments.
"""
strat = Stratification(
name="foo",
strata=["1", "2", "3"],
compartments=["sus", "inf", "rec"],
comp_split_props={},
flow_adjustments={
"infect_death": {
"2": 0.5
},
"infect_deathXlocation_work": {
"1": 0.9
},
"contact_rateXage_10": {
"1": 0.5,
"2W": 0.2
},
"contact_rateXage_20": {
"1": 0.5,
"3": "some_function"
},
},
)
assert strat.flow_adjustments == {
"infect_death": [
{
"strata": {},
"adjustments": {
"2": (FlowAdjustment.MULTIPLY, 0.5)
}
},
{
"strata": {
"location": "work"
},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.9)
}
},
],
"contact_rate": [
{
"strata": {
"age": "10"
},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.5),
"2": (FlowAdjustment.OVERWRITE, 0.2),
},
},
{
"strata": {
"age": "20"
},
"adjustments": {
"1": (FlowAdjustment.MULTIPLY, 0.5),
"3": (FlowAdjustment.COMPOSE, "some_function"),
},
},
],
}
def test_get_stratified_compartment_names__with_extisting_strat():
age_strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S", "I", "R"],
comp_split_props={},
flow_adjustments={},
)
comp_names = [Compartment("S"), Compartment("I"), Compartment("R")]
age_comp_names = get_stratified_compartment_names(age_strat, comp_names)
loc_strat = Stratification(
name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"],
comp_split_props={},
flow_adjustments={},
)
loc_comp_names = get_stratified_compartment_names(loc_strat,
age_comp_names)
assert loc_comp_names == [
Compartment("S", ["age", "location"], {
"age": "0",
"location": "rural"
}),
Compartment("S", ["age", "location"], {
"age": "0",
"location": "urban"
}),
Compartment("S", ["age", "location"], {
"age": "10",
"location": "rural"
}),
Compartment("S", ["age", "location"], {
"age": "10",
"location": "urban"
}),
Compartment("S", ["age", "location"], {
"age": "20",
"location": "rural"
}),
Compartment("S", ["age", "location"], {
"age": "20",
"location": "urban"
}),
Compartment("I", ["age", "location"], {
"age": "0",
"location": "rural"
}),
Compartment("I", ["age", "location"], {
"age": "0",
"location": "urban"
}),
Compartment("I", ["age", "location"], {
"age": "10",
"location": "rural"
}),
Compartment("I", ["age", "location"], {
"age": "10",
"location": "urban"
}),
Compartment("I", ["age", "location"], {
"age": "20",
"location": "rural"
}),
Compartment("I", ["age", "location"], {
"age": "20",
"location": "urban"
}),
Compartment("R", ["age", "location"], {
"age": "0",
"location": "rural"
}),
Compartment("R", ["age", "location"], {
"age": "0",
"location": "urban"
}),
Compartment("R", ["age", "location"], {
"age": "10",
"location": "rural"
}),
Compartment("R", ["age", "location"], {
"age": "10",
"location": "urban"
}),
Compartment("R", ["age", "location"], {
"age": "20",
"location": "rural"
}),
Compartment("R", ["age", "location"], {
"age": "20",
"location": "urban"
}),
]
def test_create_ageing_flows():
strat = Stratification(
name="age",
strata=["0", "10", "30", "40"],
compartments=["S", "I"],
comp_split_props={},
flow_adjustments={},
)
previous_compartments = [
Compartment("S", strat_names=["location"], strat_values={"location": "home"}),
Compartment("S", strat_names=["location"], strat_values={"location": "work"}),
Compartment("I", strat_names=["location"], strat_values={"location": "home"}),
Compartment("I", strat_names=["location"], strat_values={"location": "work"}),
Compartment("R", strat_names=["location"], strat_values={"location": "home"}),
Compartment("R", strat_names=["location"], strat_values={"location": "work"}),
]
ageing_flows, ageing_params = AgeingFlow.create(strat, previous_compartments, _get_param_value)
assert ageing_params == {"ageing0to10": 0.1, "ageing10to30": 0.05, "ageing30to40": 0.1}
flow = ageing_flows[0]
assert flow.source == Compartment.deserialize("SXlocation_homeXage_0")
assert flow.dest == Compartment.deserialize("SXlocation_homeXage_10")
assert flow.param_name == "ageing0to10"
assert flow.param_func == _get_param_value
flow = ageing_flows[1]
assert flow.source == Compartment.deserialize("SXlocation_workXage_0")
assert flow.dest == Compartment.deserialize("SXlocation_workXage_10")
assert flow.param_name == "ageing0to10"
assert flow.param_func == _get_param_value
flow = ageing_flows[2]
assert flow.source == Compartment.deserialize("IXlocation_homeXage_0")
assert flow.dest == Compartment.deserialize("IXlocation_homeXage_10")
assert flow.param_name == "ageing0to10"
assert flow.param_func == _get_param_value
flow = ageing_flows[3]
assert flow.source == Compartment.deserialize("IXlocation_workXage_0")
assert flow.dest == Compartment.deserialize("IXlocation_workXage_10")
assert flow.param_name == "ageing0to10"
assert flow.param_func == _get_param_value
flow = ageing_flows[4]
assert flow.source == Compartment.deserialize("SXlocation_homeXage_10")
assert flow.dest == Compartment.deserialize("SXlocation_homeXage_30")
assert flow.param_name == "ageing10to30"
assert flow.param_func == _get_param_value
flow = ageing_flows[5]
assert flow.source == Compartment.deserialize("SXlocation_workXage_10")
assert flow.dest == Compartment.deserialize("SXlocation_workXage_30")
assert flow.param_name == "ageing10to30"
assert flow.param_func == _get_param_value
flow = ageing_flows[6]
assert flow.source == Compartment.deserialize("IXlocation_homeXage_10")
assert flow.dest == Compartment.deserialize("IXlocation_homeXage_30")
assert flow.param_name == "ageing10to30"
assert flow.param_func == _get_param_value
flow = ageing_flows[7]
assert flow.source == Compartment.deserialize("IXlocation_workXage_10")
assert flow.dest == Compartment.deserialize("IXlocation_workXage_30")
assert flow.param_name == "ageing10to30"
assert flow.param_func == _get_param_value
flow = ageing_flows[8]
assert flow.source == Compartment.deserialize("SXlocation_homeXage_30")
assert flow.dest == Compartment.deserialize("SXlocation_homeXage_40")
assert flow.param_name == "ageing30to40"
assert flow.param_func == _get_param_value
flow = ageing_flows[9]
assert flow.source == Compartment.deserialize("SXlocation_workXage_30")
assert flow.dest == Compartment.deserialize("SXlocation_workXage_40")
assert flow.param_name == "ageing30to40"
assert flow.param_func == _get_param_value
flow = ageing_flows[10]
assert flow.source == Compartment.deserialize("IXlocation_homeXage_30")
assert flow.dest == Compartment.deserialize("IXlocation_homeXage_40")
assert flow.param_name == "ageing30to40"
assert flow.param_func == _get_param_value
flow = ageing_flows[11]
assert flow.source == Compartment.deserialize("IXlocation_workXage_30")
assert flow.dest == Compartment.deserialize("IXlocation_workXage_40")
assert flow.param_name == "ageing30to40"
assert flow.param_func == _get_param_value
def stratify(
self,
stratification_name: str,
strata_request: List[str],
compartments_to_stratify: List[str],
comp_split_props: Dict[str, float] = {},
flow_adjustments: Dict[str, Dict[str, float]] = {},
infectiousness_adjustments: Dict[str, float] = {},
mixing_matrix: Union[np.ndarray, Callable[[float], np.ndarray]] = None,
):
"""
Apply a stratification to the model's compartments.
stratification_name: The name of the stratification
strata_names: The names of the strata to apply
compartments_to_stratify: The compartments that will have the stratification applied. Falsey args interpreted as "all".
comp_split_props: Request to split existing population in the compartments according to specific proportions
"""
validate_stratify(
self,
stratification_name,
strata_request,
compartments_to_stratify,
comp_split_props,
flow_adjustments,
infectiousness_adjustments,
mixing_matrix,
)
strat = Stratification(
name=stratification_name,
strata=strata_request,
compartments=compartments_to_stratify,
comp_split_props=comp_split_props,
flow_adjustments=flow_adjustments,
)
self.stratifications.append(strat)
# Add this stratification's mixing matrix if a new one is provided.
if mixing_matrix is not None:
assert not strat.is_strain(), "Strains cannot have a mixing matrix."
# Only allow mixing matrix to be supplied if there is a complete stratification.
assert (
compartments_to_stratify == self.original_compartment_names
), "Mixing matrices only allowed for full stratification."
self.mixing_matrices.append(mixing_matrix)
self.mixing_categories = strat.update_mixing_categories(self.mixing_categories)
# Prepare infectiousness levels for force of infection adjustments.
self.infectiousness_levels[strat.name] = infectiousness_adjustments
if strat.is_strain():
# Track disease strain names, overriding default values.
self.disease_strains = strat.strata
# Stratify compartments, split according to split_proportions
prev_compartment_names = copy.copy(self.compartment_names)
self.compartment_names = get_stratified_compartment_names(strat, self.compartment_names)
self.compartment_values = get_stratified_compartment_values(
strat, prev_compartment_names, self.compartment_values
)
for idx, c in enumerate(self.compartment_names):
c.idx = idx
# Stratify flows
prev_flows = self.flows
self.flows = []
for flow in prev_flows:
self.flows += flow.stratify(strat)
if strat.is_ageing():
ageing_flows, ageing_params = AgeingFlow.create(
strat, prev_compartment_names, self.get_parameter_value
)
self.flows += ageing_flows
self.parameters.update(ageing_params)
# Update indicies used by flows to lookup compartment values.
compartment_idx_lookup = {name: idx for idx, name in enumerate(self.compartment_names)}
for flow in self.flows:
flow.update_compartment_indices(compartment_idx_lookup)
