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_compartment_values__with_subset_stratified():
strat = Stratification(
name="age",
strata=["0", "10", "20"],
compartments=["S"],
)
strat.set_population_split({"0": 0.25, "10": 0.5, "20": 0.25})
comps = [Compartment("S"), Compartment("I"), Compartment("R")]
comp_values = np.array([1000.0, 100.0, 0.0])
new_comp_values = strat._stratify_compartment_values(comps, 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 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_stratify_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])
comps = [
Compartment("S", {"age": "0"}),
Compartment("S", {"age": "10"}),
Compartment("S", {"age": "20"}),
Compartment("I", {"age": "0"}),
Compartment("I", {"age": "10"}),
Compartment("I", {"age": "20"}),
Compartment("R", {"age": "0"}),
Compartment("R", {"age": "10"}),
Compartment("R", {"age": "20"}),
]
strat = Stratification(
name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"],
)
strat.set_population_split({"rural": 0.1, "urban": 0.9})
new_comp_values = strat._stratify_compartment_values(comps, 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_stratify__single_with_pop_split__validate_compartments():
"""
Ensure stratifying a model correctly adjusts the model compartments.
Also the population split should be applied.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population({"S": 900, "I": 90, "R": 10})
# Compartments exist
assert model.compartments == [
Compartment("S"), Compartment("I"),
Compartment("R")
]
# Each compartment knows its index
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
# Compartments have the correct population
assert_array_equal(model.initial_population, np.array([900, 90, 10]))
# Stratify the model
strat = Stratification(name="age",
strata=["child", "adult"],
compartments=["S", "I", "R"])
strat.set_population_split({"child": 0.8, "adult": 0.2})
model.stratify_with(strat)
assert model._stratifications == [strat]
# Ensure compartments are stratified correctly
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
assert model.compartments == [
Compartment("S", {"age": "child"}),
Compartment("S", {"age": "adult"}),
Compartment("I", {"age": "child"}),
Compartment("I", {"age": "adult"}),
Compartment("R", {"age": "child"}),
Compartment("R", {"age": "adult"}),
]
expected_pop_arr = np.array([720, 180, 72, 18, 8, 2])
assert_array_equal(model.initial_population, expected_pop_arr)
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_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_create_stratification__with_pop_split():
strat = Stratification(name="location",
strata=["rural", "urban"],
compartments=["S", "I", "R"])
assert strat.population_split == {"rural": 0.5, "urban": 0.5}
# Works
strat.set_population_split({"rural": 0.2, "urban": 0.8})
assert strat.population_split == {"rural": 0.2, "urban": 0.8}
# Fails coz missing a key
with pytest.raises(AssertionError):
strat.set_population_split({"urban": 1})
# Fails coz doesn't sum to 1
with pytest.raises(AssertionError):
strat.set_population_split({"urban": 0.2, "rural": 0.3})
# Fails coz contains negative number
with pytest.raises(AssertionError):
strat.set_population_split({"urban": -2, "rural": 3})
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 test_stratify__double_with_split_and_partial__validate_compartments():
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population({"S": 900, "I": 90, "R": 10})
# Compartments exist
assert model.compartments == [
Compartment("S"), Compartment("I"),
Compartment("R")
]
# Each compartment knows its index
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
# Compartments have the correct population
assert_array_equal(model.initial_population, np.array([900, 90, 10]))
# Stratify the model
age_strat = Stratification(name="age",
strata=["child", "adult"],
compartments=["S", "R"])
age_strat.set_population_split({"child": 0.8, "adult": 0.2})
model.stratify_with(age_strat)
assert model._stratifications == [age_strat]
# Ensure compartments are stratified correctly
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
assert model.compartments == [
Compartment("S", {"age": "child"}),
Compartment("S", {"age": "adult"}),
Compartment("I"),
Compartment("R", {"age": "child"}),
Compartment("R", {"age": "adult"}),
]
expected_pop_arr = np.array([720, 180, 90, 8, 2])
assert_array_equal(model.initial_population, expected_pop_arr)
# Stratify the model again!
loc_strat = Stratification(name="location",
strata=["urban", "rural", "alpine"],
compartments=["S", "I"])
loc_strat.set_population_split({"urban": 0.7, "rural": 0.2, "alpine": 0.1})
model.stratify_with(loc_strat)
assert model._stratifications == [age_strat, loc_strat]
# Ensure compartments are stratified correctly
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
assert model.compartments == [
Compartment("S", {
"age": "child",
"location": "urban"
}),
Compartment("S", {
"age": "child",
"location": "rural"
}),
Compartment("S", {
"age": "child",
"location": "alpine"
}),
Compartment("S", {
"age": "adult",
"location": "urban"
}),
Compartment("S", {
"age": "adult",
"location": "rural"
}),
Compartment("S", {
"age": "adult",
"location": "alpine"
}),
Compartment("I", {"location": "urban"}),
Compartment("I", {"location": "rural"}),
Compartment("I", {"location": "alpine"}),
Compartment("R", {"age": "child"}),
Compartment("R", {"age": "adult"}),
]
expected_pop_arr = np.array([504, 144, 72, 126, 36, 18, 63, 18, 9, 8, 2])
assert_allclose(model.initial_population,
expected_pop_arr,
atol=1e-9,
rtol=0)
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_strain__with_infectious_multipliers_and_heterogeneous_mixing():
"""
Test infectious multiplier and flow rate calculations for
3 strains which have different infectiousness levels plus a seperate
stratification which has a mixing matrix.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
contact_rate = 0.2
model.add_infection_frequency_flow("infection", contact_rate, "S", "I")
age_strat = Stratification("age", ["child", "adult"], ["S", "I", "R"])
age_strat.set_population_split({
"child": 0.6, # 600 people
"adult": 0.4, # 400 people
})
# Higher mixing among adults or children,
# than between adults or children.
age_strat.set_mixing_matrix(np.array([[1.5, 0.5], [0.5, 1.5]]))
model.stratify_with(age_strat)
strain_strat = StrainStratification("strain", ["a", "b", "c"], ["I"])
strain_strat.set_population_split({
"a": 0.7, # 70 people
"b": 0.2, # 20 people
"c": 0.1, # 10 people
})
strain_strat.add_infectiousness_adjustments(
"I",
{
"a": adjust.Multiply(0.5), # 0.5x as susceptible
"b": adjust.Multiply(3), # 3x as susceptible
"c": adjust.Multiply(2), # 2x as susceptible
},
)
model.stratify_with(strain_strat)
# Do pre-run force of infection calcs.
model._prepare_to_run()
assert_array_equal(model._compartment_infectiousness["a"],
np.array([0, 0, 0.5, 0, 0, 0.5, 0, 0, 0, 0]))
assert_array_equal(model._compartment_infectiousness["b"],
np.array([0, 0, 0, 3, 0, 0, 3, 0, 0, 0]))
assert_array_equal(model._compartment_infectiousness["c"],
np.array([0, 0, 0, 0, 2, 0, 0, 2, 0, 0]))
# 0 for child, 1 for adult
assert model._category_lookup == {
0: 0,
1: 1,
2: 0,
3: 0,
4: 0,
5: 1,
6: 1,
7: 1,
8: 0,
9: 1
}
assert_array_equal(
model._category_matrix,
np.array([
[1, 0, 1, 1, 1, 0, 0, 0, 1, 0],
[0, 1, 0, 0, 0, 1, 1, 1, 0, 1],
]),
)
# Do pre-iteration force of infection calcs
model._prepare_time_step(0, model.initial_population)
assert_array_equal(model._category_populations, np.array([600, 400]))
assert_array_equal(
model._infection_density["a"],
np.array([0.5 * (42 * 1.5 + 28 * 0.5), 0.5 * (42 * 0.5 + 28 * 1.5)]),
)
assert_array_equal(
model._infection_density["b"],
np.array([
3 * (12 * 1.5 + 8 * 0.5),
3 * (8 * 1.5 + 12 * 0.5),
]),
)
assert_array_equal(
model._infection_density["c"],
np.array([2 * (6 * 1.5 + 4 * 0.5), 2 * (4 * 1.5 + 6 * 0.5)]),
)
assert_array_equal(
model._infection_frequency["a"],
np.array([
0.5 * ((42 / 600) * 1.5 + (28 / 400) * 0.5),
0.5 * ((42 / 600) * 0.5 + (28 / 400) * 1.5),
]),
)
assert_array_equal(
model._infection_frequency["b"],
np.array([
3 * ((12 / 600) * 1.5 + (8 / 400) * 0.5),
3 * ((8 / 400) * 1.5 + (12 / 600) * 0.5),
]),
)
assert_array_equal(
model._infection_frequency["c"],
np.array([
2 * ((6 / 600) * 1.5 + (4 / 400) * 0.5),
2 * ((4 / 400) * 1.5 + (6 / 600) * 0.5)
]),
)
# Get multipliers
sus_child = model.compartments[0]
sus_adult = model.compartments[1]
inf_child_a = model.compartments[2]
inf_child_b = model.compartments[3]
inf_child_c = model.compartments[4]
inf_adult_a = model.compartments[5]
inf_adult_b = model.compartments[6]
inf_adult_c = model.compartments[7]
density = model._get_infection_density_multiplier
freq = model._get_infection_frequency_multiplier
assert density(sus_child, inf_child_a) == 0.5 * (42 * 1.5 + 28 * 0.5)
assert density(sus_adult, inf_adult_a) == 0.5 * (42 * 0.5 + 28 * 1.5)
assert density(sus_child, inf_child_b) == 3 * (12 * 1.5 + 8 * 0.5)
assert density(sus_adult, inf_adult_b) == 3 * (8 * 1.5 + 12 * 0.5)
assert density(sus_child, inf_child_c) == 2 * (6 * 1.5 + 4 * 0.5)
assert density(sus_adult, inf_adult_c) == 2 * (4 * 1.5 + 6 * 0.5)
assert freq(sus_child,
inf_child_a) == 0.5 * ((42 / 600) * 1.5 + (28 / 400) * 0.5)
assert freq(sus_adult,
inf_adult_a) == 0.5 * ((42 / 600) * 0.5 + (28 / 400) * 1.5)
assert freq(sus_child,
inf_child_b) == 3 * ((12 / 600) * 1.5 + (8 / 400) * 0.5)
assert freq(sus_adult,
inf_adult_b) == 3 * ((8 / 400) * 1.5 + (12 / 600) * 0.5)
assert freq(sus_child,
inf_child_c) == 2 * ((6 / 600) * 1.5 + (4 / 400) * 0.5)
assert freq(sus_adult,
inf_adult_c) == 2 * ((4 / 400) * 1.5 + (6 / 600) * 0.5)
# Get infection flow rates
flow_to_inf_child_a = 540 * contact_rate * freq(sus_child, inf_child_a)
flow_to_inf_adult_a = 360 * contact_rate * freq(sus_adult, inf_adult_a)
flow_to_inf_child_b = 540 * contact_rate * freq(sus_child, inf_child_b)
flow_to_inf_adult_b = 360 * contact_rate * freq(sus_adult, inf_adult_b)
flow_to_inf_child_c = 540 * contact_rate * freq(sus_child, inf_child_c)
flow_to_inf_adult_c = 360 * contact_rate * freq(sus_adult, inf_adult_c)
expected_flow_rates = np.array([
-flow_to_inf_child_a - flow_to_inf_child_b - flow_to_inf_child_c,
-flow_to_inf_adult_a - flow_to_inf_adult_b - flow_to_inf_adult_c,
flow_to_inf_child_a,
flow_to_inf_child_b,
flow_to_inf_child_c,
flow_to_inf_adult_a,
flow_to_inf_adult_b,
flow_to_inf_adult_c,
0.0,
0.0,
])
flow_rates = model._get_compartment_rates(model.initial_population, 0)
assert_allclose(expected_flow_rates, flow_rates, verbose=True)
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 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