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 _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 test_create_stratification__with_infectiousness_adjustments():
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
assert strat.infectiousness_adjustments == {}
# Fail coz not all strata specified
with pytest.raises(AssertionError):
strat.add_infectiousness_adjustments(
compartment_name="S",
adjustments={"rural": Multiply(1.2)},
)
# Fail coz an incorrect strata specified
with pytest.raises(AssertionError):
strat.add_infectiousness_adjustments(
compartment_name="S",
adjustments={
"rural": Multiply(1.2),
"urban": Multiply(0.8),
"alpine": Multiply(1.1),
},
)
# Fail coz a time-varying function was used (not allowed!)
with pytest.raises(AssertionError):
strat.add_infectiousness_adjustments(
compartment_name="S",
adjustments={
"rural": Multiply(1.2),
"urban": Multiply(lambda t: 2),
},
)
strat.add_infectiousness_adjustments(
compartment_name="S",
adjustments={
"rural": Multiply(1.2),
"urban": Multiply(2),
},
)
assert strat.infectiousness_adjustments["S"]["rural"]._is_equal(
Multiply(1.2))
assert strat.infectiousness_adjustments["S"]["urban"]._is_equal(
Multiply(2))
# Fail coz we just did this
with pytest.raises(AssertionError):
strat.add_infectiousness_adjustments(
compartment_name="S",
adjustments={
"rural": Multiply(1.2),
"urban": Multiply(2),
},
)
strat.add_infectiousness_adjustments(
compartment_name="I",
adjustments={
"rural": Multiply(1.2),
"urban": None,
},
)
assert strat.infectiousness_adjustments["I"]["rural"]._is_equal(
Multiply(1.2))
assert strat.infectiousness_adjustments["I"]["urban"] is None
def test_strat_infectiousness__with_adjustments():
"""
Ensure multiply infectiousness adjustment is applied.
"""
# Create a model
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 900, "I": 100})
strat = Stratification("age", ["baby", "child", "adult"], ["S", "I", "R"])
strat.set_population_split({"baby": 0.1, "child": 0.3, "adult": 0.6})
strat.add_infectiousness_adjustments(
"I", {"child": adjust.Multiply(3), "adult": adjust.Multiply(0.5), "baby": None}
)
model.stratify_with(strat)
assert_array_equal(
model.initial_population,
np.array([90, 270, 540, 10, 30, 60, 0, 0, 0]),
)
# Do pre-run force of infection calcs.
model._prepare_to_run()
assert_array_equal(
model._compartment_infectiousness["default"],
np.array([0, 0, 0, 1, 3, 0.5, 0, 0, 0]),
)
# Do pre-iteration force of infection calcs
model._prepare_time_step(0, model.initial_population)
# Get multipliers
infectees = model.compartments[0:3]
infectors = model.compartments[3:6]
expected_density = 10 * 1 + 30 * 3 + 60 * 0.5
expected_frequency = expected_density / 1000
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_density_multiplier(infectee, infector) == expected_density
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency
# Stratify again, now with overwrites
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
strat.add_infectiousness_adjustments(
"I", {"urban": adjust.Overwrite(1), "rural": adjust.Multiply(7)}
)
model.stratify_with(strat)
assert_array_equal(
model.initial_population,
np.array([45, 45, 135, 135, 270.0, 270, 5, 5, 15, 15, 30, 30, 0, 0, 0, 0, 0, 0]),
)
# Do pre-run force of infection calcs.
model._prepare_to_run()
assert_array_equal(
model._compartment_infectiousness["default"],
np.array([0, 0, 0, 0, 0, 0, 1, 7, 1, 21, 1, 3.5, 0, 0, 0, 0, 0, 0]),
)
# Do pre-iteration force of infection calcs
model._prepare_time_step(0, model.initial_population)
# Get multipliers
infectees = model.compartments[0:6]
infectors = model.compartments[6:12]
expected_density = 5 * 1 + 5 * 7 + 15 * 1 + 15 * 21 + 30 * 1 + 30 * 3.5
expected_frequency = expected_density / 1000
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_density_multiplier(infectee, infector) == expected_density
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency