def test_stratify_transition_flows__with_dest_only_stratified__with_adjustments():
"""
Ensure transition flows are stratified correctly when only the flow destination is stratified.
Expect adjustments to override the automatic person-conserving adjustment.
"""
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"])
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 _build_organ_strat(params):
organ_strat = Stratification(
"organ", ["smear_positive", "smear_negative", "extra_pulmonary"],
INFECTIOUS_COMPS)
organ_strat.set_population_split(params["organ"]["props"])
# Add infectiousness adjustments
for comp in INFECTIOUS_COMPS:
organ_strat.add_infectiousness_adjustments(
comp, _adjust_all_multiply(params["organ"]["foi"]))
organ_strat.add_flow_adjustments(
"early_activation", _adjust_all_multiply(params["organ"]["props"]))
organ_strat.add_flow_adjustments(
"late_activation", _adjust_all_multiply(params["organ"]["props"]))
organ_strat.add_flow_adjustments(
"detection",
_adjust_all_multiply(params["detection_rate_stratified"]["organ"]),
)
organ_strat.add_flow_adjustments(
"self_recovery_infectious",
_adjust_all_multiply(params["self_recovery_rate_stratified"]["organ"]),
)
organ_strat.add_flow_adjustments(
"self_recovery_detected",
_adjust_all_multiply(params["self_recovery_rate_stratified"]["organ"]),
)
return organ_strat
def test_stratify_exit_flows():
"""
Ensure exit flows are stratified correctly.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_death_flow("d_S", 3, "S")
model.add_death_flow("d_I", 5, "I")
model.add_death_flow("d_R", 7, "R")
expected_flows = [
DeathFlow("d_S", C("S"), 3),
DeathFlow("d_I", C("I"), 5),
DeathFlow("d_R", 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 partial stratification
strat = Stratification("location", ["urban", "rural"], ["S", "I"])
model.stratify_with(strat)
expected_flows = [
DeathFlow("d_S", C("S", {"location": "urban"}), 3),
DeathFlow("d_S", C("S", {"location": "rural"}), 3),
DeathFlow("d_I", C("I", {"location": "urban"}), 5),
DeathFlow("d_I", C("I", {"location": "rural"}), 5),
DeathFlow("d_R", 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 partial stratification with flow adjustments
strat = Stratification("age", ["young", "old"], ["I", "R"])
strat.add_flow_adjustments("d_I", {"young": Multiply(0.5), "old": Multiply(2)})
strat.add_flow_adjustments("d_R", {"young": Multiply(0.5), "old": Multiply(2)})
model.stratify_with(strat)
expected_flows = [
DeathFlow("d_S", C("S", {"location": "urban"}), 3),
DeathFlow("d_S", C("S", {"location": "rural"}), 3),
DeathFlow("d_I", C("I", {"location": "urban", "age": "young"}), 5, [Multiply(0.5)]),
DeathFlow("d_I", C("I", {"location": "urban", "age": "old"}), 5, [Multiply(2)]),
DeathFlow("d_I", C("I", {"location": "rural", "age": "young"}), 5, [Multiply(0.5)]),
DeathFlow("d_I", C("I", {"location": "rural", "age": "old"}), 5, [Multiply(2)]),
DeathFlow("d_R", C("R", {"age": "young"}), 7, [Multiply(0.5)]),
DeathFlow("d_R", C("R", {"age": "old"}), 7, [Multiply(2)]),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def _build_retention_strat(params):
strat = Stratification("retained", ["yes", "no"], [Compartment.DETECTED])
strat.add_flow_adjustments(
"detection",
{
"yes": Multiply(params["retention_prop"]),
"no": Multiply(1 - params["retention_prop"]),
},
)
strat.add_flow_adjustments(
"missed_to_active",
{
"yes": Multiply(0),
"no": Multiply(1)
},
)
strat.add_flow_adjustments(
"treatment_commencement",
{
"yes": Multiply(1),
"no": Multiply(0)
},
)
strat.add_flow_adjustments(
"failure_retreatment",
{
"yes": Multiply(1),
"no": Multiply(0)
},
)
return strat
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 get_cluster_strat(params: Parameters) -> Stratification:
cluster_strat = Stratification("cluster", CLUSTER_STRATA, COMPARTMENTS)
country = params.country
vic = params.victorian_clusters
# Determine how to split up population by cluster
# There is -0.5% to +4% difference per age group between sum of region population in 2018 and
# total VIC population in 2020
region_pops = {
region: sum(
inputs.get_population_by_agegroup(AGEGROUP_STRATA,
country.iso3,
region.upper(),
year=2018))
for region in CLUSTER_STRATA
}
sum_region_props = sum(region_pops.values())
cluster_split_props = {
region: pop / sum_region_props
for region, pop in region_pops.items()
}
cluster_strat.set_population_split(cluster_split_props)
# Adjust contact rate multipliers
contact_rate_adjustments = {}
for cluster in Region.VICTORIA_METRO + [Region.BARWON_SOUTH_WEST]:
cluster_name = cluster.replace("-", "_")
contact_rate_multiplier = getattr(
vic, f"contact_rate_multiplier_{cluster_name}")
contact_rate_adjustments[cluster_name] = Multiply(
contact_rate_multiplier)
for cluster in Region.VICTORIA_RURAL:
if cluster != Region.BARWON_SOUTH_WEST:
cluster_name = cluster.replace("-", "_")
contact_rate_multiplier = getattr(
vic, "contact_rate_multiplier_regional")
contact_rate_adjustments[cluster_name] = Multiply(
contact_rate_multiplier)
# Add in flow adjustments per-region so we can calibrate the contact rate for each region.
cluster_strat.add_flow_adjustments("infection", contact_rate_adjustments)
# Use an identity mixing matrix to temporarily declare no inter-cluster mixing, which will then be over-written
cluster_mixing_matrix = np.eye(len(CLUSTER_STRATA))
cluster_strat.set_mixing_matrix(cluster_mixing_matrix)
return cluster_strat
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 get_agegroup_strat(params: Parameters,
total_pops: List[int]) -> Stratification:
"""
Age stratification
"""
# We use "Stratification" instead of "AgeStratification" for this model, to avoid triggering
# automatic demography features (which work on the assumption that the time unit is years, so would be totally wrong)
age_strat = Stratification("agegroup", AGEGROUP_STRATA, COMPARTMENTS)
country = params.country
# Dynamic heterogeneous mixing by age
if params.elderly_mixing_reduction and not params.mobility.age_mixing:
# Apply eldery protection to the age mixing parameters
params.mobility.age_mixing = preprocess.elderly_protection.get_elderly_protection_mixing(
params.elderly_mixing_reduction)
static_mixing_matrix = inputs.get_country_mixing_matrix(
"all_locations", country.iso3)
dynamic_mixing_matrix = preprocess.mixing_matrix.build_dynamic_mixing_matrix(
static_mixing_matrix,
params.mobility,
country,
)
age_strat.set_mixing_matrix(dynamic_mixing_matrix)
# Set distribution of starting population
age_split_props = {
agegroup: prop
for agegroup, prop in zip(AGEGROUP_STRATA,
normalise_sequence(total_pops))
}
age_strat.set_population_split(age_split_props)
# Adjust flows based on age group.
age_strat.add_flow_adjustments(
"infection", {
s: Multiply(v)
for s, v in params.age_stratification.susceptibility.items()
})
return age_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 _build_vac_strat(params):
vac_strat = Stratification("vac", ["unvaccinated", "vaccinated"],
[Compartment.SUSCEPTIBLE])
vac_strat.set_population_split({"unvaccinated": 1.0, "vaccinated": 0.0})
# Apply flow adjustments
vac_strat.add_flow_adjustments(
"infection",
{
"unvaccinated": Multiply(1.0),
"vaccinated": Multiply(params["bcg"]["rr_infection_vaccinated"]),
},
)
vac_strat.add_flow_adjustments(
"treatment_early",
{
"unvaccinated": Multiply(0.0),
"vaccinated": Multiply(1.0)
},
)
vac_strat.add_flow_adjustments("treatment_late", {
"unvaccinated": Multiply(0.0),
"vaccinated": Multiply(1.0)
})
def time_varying_vaccination_coverage(t):
return params["bcg"][
"coverage"] if t > params["bcg"]["start_time"] else 0.0
def time_varying_unvaccinated_coverage(t):
return 1 - params["bcg"]["coverage"] if t > params["bcg"][
"start_time"] else 1.0
vac_strat.add_flow_adjustments(
"birth",
{
"unvaccinated": Multiply(time_varying_unvaccinated_coverage),
"vaccinated": Multiply(time_varying_vaccination_coverage),
},
)
return vac_strat
def get_history_strat(
params: Parameters,
compartment_periods: Dict[str, float],
) -> Stratification:
"""
Infection history stratification
"""
history_strat = Stratification(
"history",
["naive", "experienced"],
[Compartment.SUSCEPTIBLE, Compartment.EARLY_EXPOSED],
)
# Everyone starts out naive.
history_strat.set_population_split({"naive": 1.0, "experienced": 0.0})
# Waning immunity makes recovered individuals transition to the 'experienced' stratum
history_strat.add_flow_adjustments("warning_immunity", {
"naive": Multiply(0.0),
"experienced": Multiply(1.0)
})
# Get the proportion of people in each clinical stratum, relative to total people in compartment.
clinical_params = params.clinical_stratification
symptomatic_props = get_proportion_symptomatic(params)
abs_props = get_absolute_strata_proportions(
symptomatic_props=symptomatic_props,
icu_props=clinical_params.icu_prop,
hospital_props=clinical_params.props.hospital.props,
symptomatic_props_multiplier=clinical_params.props.symptomatic.
multiplier,
hospital_props_multiplier=clinical_params.props.hospital.multiplier,
)
# Adjust parameters defining progression from early exposed to late exposed to obtain the requested proportion
for age_idx, agegroup in enumerate(AGEGROUP_STRATA):
# Collect existing rates of progressions for symptomatic vs non-symptomatic
rate_non_sympt = (abs_props[Clinical.NON_SYMPT][age_idx] /
compartment_periods[Compartment.EARLY_EXPOSED])
total_progression_rate = 1.0 / compartment_periods[
Compartment.EARLY_EXPOSED]
rate_sympt = total_progression_rate - rate_non_sympt
# Multiplier for symptomatic is rel_prop_symptomatic_experienced
sympt_multiplier = params.rel_prop_symptomatic_experienced
# Multiplier for asymptomatic rate is 1 + rate_sympt / rate_non_sympt * (1 - sympt_multiplier) in order to preserve aggregated exit flow.
non_sympt_multiplier = 1 + rate_sympt / rate_non_sympt * (
1.0 - sympt_multiplier)
# Create adjustment requests
for clinical in CLINICAL_STRATA:
experienced_multiplier = (non_sympt_multiplier
if clinical == Clinical.NON_SYMPT else
sympt_multiplier)
adjustments = {
"naive": Multiply(1.0),
# Sojourn flow, divide by proportion.
"experienced": Multiply(1.0 / experienced_multiplier),
}
history_strat.add_flow_adjustments(
"infect_onset",
adjustments,
dest_strata={
"agegroup": agegroup,
"clinical": clinical
},
)
return history_strat
def _get_test_model(timestep=1, times=[0, 150]):
comps = ["S", "EE", "LE", "EA", "LA", "R"]
infectious_comps = ["LE", "EA", "LA"]
model = CompartmentalModel(
times=times,
compartments=comps,
infectious_compartments=infectious_comps,
timestep=timestep,
)
model.set_initial_population({"S": int(20e6), "LA": 100})
# Add flows
model.add_infection_frequency_flow(name="infection", contact_rate=0.03, source="S", dest="EE")
model.add_sojourn_flow(name="infect_onset", sojourn_time=7, source="EE", dest="LE")
model.add_sojourn_flow(name="incidence", sojourn_time=7, source="LE", dest="EA")
model.add_sojourn_flow(name="progress", sojourn_time=7, source="EA", dest="LA")
model.add_sojourn_flow(name="recovery", sojourn_time=7, source="LA", dest="R")
model.add_death_flow(name="infect_death", death_rate=0.005, source="LA")
model.add_transition_flow(name="warning_immunity", fractional_rate=0.01, source="R", dest="S")
# Stratify by age
age_strat = Stratification("age", AGE_STRATA, comps)
age_strat.set_population_split(AGE_SPLIT_PROPORTIONS)
age_strat.set_mixing_matrix(AGE_MIXING_MATRIX)
age_strat.add_flow_adjustments(
"infection", {s: Multiply(v) for s, v in AGE_SUSCEPTIBILITY.items()}
)
model.stratify_with(age_strat)
# Stratify by clinical status
clinical_strat = Stratification("clinical", CLINICAL_STRATA, infectious_comps)
clinical_strat.add_infectiousness_adjustments("LE", {**ADJ_BASE, "non_sympt": Overwrite(0.25)})
clinical_strat.add_infectiousness_adjustments("EA", {**ADJ_BASE, "non_sympt": Overwrite(0.25)})
clinical_strat.add_infectiousness_adjustments(
"LA",
{
**ADJ_BASE,
"non_sympt": Overwrite(0.25),
"sympt_isolate": Overwrite(0.2),
"hospital": Overwrite(0.2),
"icu": Overwrite(0.2),
},
)
clinical_strat.add_flow_adjustments(
"infect_onset",
{
"non_sympt": Multiply(0.26),
"icu": Multiply(0.01),
"hospital": Multiply(0.04),
"sympt_public": Multiply(0.66),
"sympt_isolate": Multiply(0.03),
},
)
model.stratify_with(clinical_strat)
# Request derived outputs.
model.request_output_for_flow(name="incidence", flow_name="incidence")
model.request_output_for_flow(name="progress", flow_name="progress")
for age in AGE_STRATA:
for clinical in NOTIFICATION_STRATA:
model.request_output_for_flow(
name=f"progressXage_{age}Xclinical_{clinical}",
flow_name="progress",
dest_strata={"age": age, "clinical": clinical},
)
hospital_sources = []
icu_sources = []
for age in AGE_STRATA:
icu_sources.append(f"progressXage_{age}Xclinical_icu")
hospital_sources += [
f"progressXage_{age}Xclinical_icu",
f"progressXage_{age}Xclinical_hospital",
]
model.request_aggregate_output(
name="new_hospital_admissions",
sources=hospital_sources,
)
model.request_aggregate_output(name="new_icu_admissions", sources=icu_sources)
# Get notifications, which may included people detected in-country as they progress, or imported cases which are detected.
notification_sources = [
f"progressXage_{a}Xclinical_{c}" for a in AGE_STRATA for c in NOTIFICATION_STRATA
]
model.request_aggregate_output(name="notifications", sources=notification_sources)
# Infection deaths.
model.request_output_for_flow(name="infection_deaths", flow_name="infect_death")
model.request_cumulative_output(name="accum_deaths", source="infection_deaths")
# Track hospital occupancy.
# We count all ICU and hospital late active compartments and a proportion of early active ICU cases.
model.request_output_for_compartments(
"_late_active_hospital",
compartments=["LA"],
strata={"clinical": "hospital"},
save_results=False,
)
model.request_output_for_compartments(
"icu_occupancy",
compartments=["LA"],
strata={"clinical": "icu"},
)
model.request_output_for_compartments(
"_early_active_icu",
compartments=["EA"],
strata={"clinical": "icu"},
save_results=False,
)
proportion_icu_patients_in_hospital = 0.25
model.request_function_output(
name="_early_active_icu_proportion",
func=lambda patients: patients * proportion_icu_patients_in_hospital,
sources=["_early_active_icu"],
save_results=False,
)
model.request_aggregate_output(
name="hospital_occupancy",
sources=[
"_late_active_hospital",
"icu_occupancy",
"_early_active_icu_proportion",
],
)
# Proportion seropositive
model.request_output_for_compartments(
name="_total_population", compartments=comps, save_results=False
)
model.request_output_for_compartments(name="_recovered", compartments=["R"], save_results=False)
model.request_function_output(
name="proportion_seropositive",
sources=["_recovered", "_total_population"],
func=lambda recovered, total: recovered / total,
)
return model
def get_clinical_strat(params: Parameters) -> Stratification:
"""
Stratify the model by clinical status
Stratify the infectious compartments of the covid model (not including the pre-symptomatic compartments, which are
actually infectious)
- notifications are derived from progress from early to late for some strata
- the proportion of people moving from presympt to early infectious, conditioned on age group
- rate of which people flow through these compartments (reciprocal of time, using within_* which is a rate of ppl / day)
- infectiousness levels adjusted by early/late and for clinical strata
- we start with an age stratified infection fatality rate
- 50% of deaths for each age bracket die in ICU
- the other deaths go to hospital, assume no-one else can die from COVID
- should we ditch this?
"""
clinical_strat = Stratification("clinical", CLINICAL_STRATA, INFECTIOUS_COMPARTMENTS)
clinical_params = params.clinical_stratification
country = params.country
pop = params.population
"""
Infectiousness adjustments for clinical stratification
"""
# Add infectiousness reduction multiplier for all non-symptomatic infectious people.
# These people are less infectious because of biology.
non_sympt_adjust = Overwrite(clinical_params.non_sympt_infect_multiplier)
clinical_strat.add_infectiousness_adjustments(
Compartment.LATE_EXPOSED,
{
Clinical.NON_SYMPT: non_sympt_adjust,
Clinical.SYMPT_NON_HOSPITAL: None,
Clinical.SYMPT_ISOLATE: None,
Clinical.HOSPITAL_NON_ICU: None,
Clinical.ICU: None,
},
)
clinical_strat.add_infectiousness_adjustments(
Compartment.EARLY_ACTIVE,
{
Clinical.NON_SYMPT: non_sympt_adjust,
Clinical.SYMPT_NON_HOSPITAL: None,
Clinical.SYMPT_ISOLATE: None,
Clinical.HOSPITAL_NON_ICU: None,
Clinical.ICU: None,
},
)
# Add infectiousness reduction for people who are late active and in isolation or hospital/icu.
# These people are less infectious because of physical distancing/isolation/PPE precautions.
late_infect_multiplier = clinical_params.late_infect_multiplier
clinical_strat.add_infectiousness_adjustments(
Compartment.LATE_ACTIVE,
{
Clinical.NON_SYMPT: non_sympt_adjust,
Clinical.SYMPT_ISOLATE: Overwrite(late_infect_multiplier[Clinical.SYMPT_ISOLATE]),
Clinical.SYMPT_NON_HOSPITAL: None,
Clinical.HOSPITAL_NON_ICU: Overwrite(late_infect_multiplier[Clinical.HOSPITAL_NON_ICU]),
Clinical.ICU: Overwrite(late_infect_multiplier[Clinical.ICU]),
},
)
"""
Adjust infection death rates for hospital patients (ICU and non-ICU)
"""
symptomatic_adjuster = params.clinical_stratification.props.symptomatic.multiplier
hospital_adjuster = params.clinical_stratification.props.hospital.multiplier
ifr_adjuster = params.infection_fatality.multiplier
# Get all the adjustments in the same way as we will do if the immunity stratification is implemented
entry_adjustments, death_adjs, progress_adjs, recovery_adjs, _, _ = get_all_adjs(
clinical_params,
country,
pop,
params.infection_fatality.props,
params.sojourn,
params.testing_to_detection,
params.case_detection,
ifr_adjuster,
symptomatic_adjuster,
hospital_adjuster,
)
# Assign all the adjustments to the model
for agegroup in AGEGROUP_STRATA:
source = {"agegroup": agegroup}
clinical_strat.add_flow_adjustments(
"infect_onset", entry_adjustments[agegroup], source_strata=source
)
clinical_strat.add_flow_adjustments(
"infect_death", death_adjs[agegroup], source_strata=source
)
clinical_strat.add_flow_adjustments(
"progress",
progress_adjs,
source_strata=source,
)
clinical_strat.add_flow_adjustments(
"recovery",
recovery_adjs[agegroup],
source_strata=source,
)
return clinical_strat
def _build_class_strat(params):
strat = Stratification("classified", ["correctly", "incorrectly"],
[Compartment.DETECTED, Compartment.ON_TREATMENT])
# Apply only to ds flows
strat.add_flow_adjustments(
"detection",
{
"correctly": Multiply(1.0),
"incorrectly": Multiply(0.0)
},
source_strata={"strain": "ds"},
)
# Apply only to mdr flows
strat.add_flow_adjustments(
"detection",
{
"correctly": Multiply(params["frontline_xpert_prop"]),
"incorrectly": Multiply(1.0 - params["frontline_xpert_prop"]),
},
source_strata={"strain": "mdr"},
)
# Apply only to mdr flows that have extra_pulmonary organ.
strat.add_flow_adjustments(
"detection",
{
"correctly": Multiply(0),
"incorrectly": Multiply(1)
},
source_strata={
"strain": "mdr",
"organ": "extra_pulmonary"
},
)
strat.add_flow_adjustments(
"treatment_default",
_adjust_all_multiply(
params["treatment_default_rate_stratified"]["classified"]),
source_strata={"strain": "mdr"},
)
strat.add_flow_adjustments(
"spontaneous_recovery",
_adjust_all_multiply(
params["spontaneous_recovery_rate_stratified"]["classified"]),
)
strat.add_flow_adjustments(
"failure_retreatment",
_adjust_all_multiply(
params["failure_retreatment_rate_stratified"]["classified"]),
)
return strat
def test_create_stratification__with_flow_adjustments():
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
assert strat.flow_adjustments == {}
# Fail coz not all strata specified
with pytest.raises(AssertionError):
strat.add_flow_adjustments(
flow_name="recovery",
adjustments={"rural": Multiply(1.2)},
)
# Fail coz an incorrect strata specified
with pytest.raises(AssertionError):
strat.add_flow_adjustments(
flow_name="recovery",
adjustments={
"rural": Multiply(1.2),
"urban": Multiply(0.8),
"alpine": Multiply(1.1),
},
)
strat.add_flow_adjustments(
flow_name="recovery",
adjustments={
"rural": Multiply(1.2),
"urban": Multiply(0.8)
},
)
assert len(strat.flow_adjustments["recovery"]) == 1
adj, src, dst = strat.flow_adjustments["recovery"][0]
assert adj["rural"]._is_equal(Multiply(1.2))
assert adj["urban"]._is_equal(Multiply(0.8))
assert not (src or dst)
# Add another adjustment for the same flow.
strat.add_flow_adjustments(
flow_name="recovery",
adjustments={
"rural": Multiply(1.3),
"urban": Multiply(0.9)
},
source_strata={"age": "10"},
dest_strata={"work": "office"},
)
assert len(strat.flow_adjustments["recovery"]) == 2
adj, src, dst = strat.flow_adjustments["recovery"][0]
assert adj["rural"]._is_equal(Multiply(1.2))
assert adj["urban"]._is_equal(Multiply(0.8))
assert not (src or dst)
adj, src, dst = strat.flow_adjustments["recovery"][1]
assert adj["rural"]._is_equal(Multiply(1.3))
assert adj["urban"]._is_equal(Multiply(0.9))
assert src == {"age": "10"}
assert dst == {"work": "office"}
def urban_infection_adjustment(t):
return 2 * t
strat.add_flow_adjustments(
flow_name="infection",
adjustments={
"rural": Multiply(urban_infection_adjustment),
"urban": None,
},
)
assert len(strat.flow_adjustments["infection"]) == 1
adj, src, dst = strat.flow_adjustments["infection"][0]
assert adj["rural"]._is_equal(Multiply(urban_infection_adjustment))
assert adj["urban"] is None
assert not (src or dst)
def get_immunity_strat(params: Parameters) -> Stratification:
immunity_strat = Stratification("immunity", IMMUNITY_STRATA, COMPARTMENTS)
relative_severity_effect = 1.0
# Everyone starts out unvaccinated
immunity_strat.set_population_split({
"unvaccinated": 1.0,
"vaccinated": 0.0
})
if params.vaccination:
# Apply vaccination effect against severe disease given infection
relative_severity_effect -= params.vaccination.severity_efficacy
# Apply vaccination effect against infection/transmission
immunity_strat.add_flow_adjustments(
"infection",
{
"vaccinated":
Multiply(1.0 - params.vaccination.infection_efficacy),
"unvaccinated": None,
},
)
symptomatic_adjuster = (
params.clinical_stratification.props.symptomatic.multiplier *
relative_severity_effect)
hospital_adjuster = (
params.clinical_stratification.props.hospital.multiplier *
relative_severity_effect)
ifr_adjuster = params.infection_fatality.multiplier * relative_severity_effect
# Get all the adjustments in the same way as we did for the clinical stratification
entry_adjustments, death_adjs, progress_adjs, recovery_adjs, _, _ = get_all_adjs(
params.clinical_stratification,
params.country,
params.population,
params.infection_fatality.props,
params.sojourn,
params.testing_to_detection,
params.case_detection,
ifr_adjuster,
symptomatic_adjuster,
hospital_adjuster,
)
for i_age, agegroup in enumerate(AGEGROUP_STRATA):
for clinical_stratum in CLINICAL_STRATA:
relevant_strata = {
"agegroup": agegroup,
"clinical": clinical_stratum,
}
immunity_strat.add_flow_adjustments(
"infect_onset",
{
"unvaccinated":
None,
"vaccinated":
entry_adjustments[agegroup][clinical_stratum],
},
dest_strata=relevant_strata, # Must be dest
)
immunity_strat.add_flow_adjustments(
"infect_death",
{
"unvaccinated": None,
"vaccinated": death_adjs[agegroup][clinical_stratum]
},
source_strata=relevant_strata, # Must be source
)
immunity_strat.add_flow_adjustments(
"progress",
{
"unvaccinated": None,
"vaccinated": progress_adjs[clinical_stratum]
},
source_strata=
relevant_strata, # Either source or dest or both
)
immunity_strat.add_flow_adjustments(
"recovery",
{
"unvaccinated": None,
"vaccinated": recovery_adjs[agegroup][clinical_stratum]
},
source_strata=relevant_strata, # Must be source
)
return immunity_strat
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_stratify_entry_flows__with_explicit_adjustments():
"""
Ensure entry flows are stratified correctly when adjustments are requested.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_importation_flow("imports", 10, "S")
expected_flows = [ImportFlow("imports", C("S"), 10)]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
strat.add_flow_adjustments("imports", {"urban": Multiply(0.9), "rural": None})
model.stratify_with(strat)
expected_flows = [
ImportFlow(
"imports",
C("S", {"location": "urban"}),
10,
[Multiply(0.9)],
),
ImportFlow(
"imports",
C("S", {"location": "rural"}),
10,
[],
),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
strat = Stratification("age", ["young", "old"], ["S", "I", "R"])
strat.add_flow_adjustments("imports", {"young": Multiply(0.8), "old": Overwrite(1)})
model.stratify_with(strat)
expected_flows = [
ImportFlow(
"imports",
C("S", {"location": "urban", "age": "young"}),
10,
[Multiply(0.9), Multiply(0.8)],
),
ImportFlow(
"imports",
C("S", {"location": "urban", "age": "old"}),
10,
[Multiply(0.9), Overwrite(1)],
),
ImportFlow(
"imports",
C("S", {"location": "rural", "age": "young"}),
10,
[Multiply(0.8)],
),
ImportFlow(
"imports",
C("S", {"location": "rural", "age": "old"}),
10,
[Overwrite(1)],
),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])