def build_model(params: dict) -> CompartmentalModel:
"""
Build the master function to run a simple SIR model
"""
# Define model compartments.
compartments = ["S", "I", "R"]
time = params["time"]
model = CompartmentalModel(
times=[time["start"], time["end"]],
compartments=compartments,
infectious_compartments=["I"],
timestep=time["step"],
)
model.set_initial_population({
"I": 1,
"S": 999999,
})
# Add flows
model.add_infection_frequency_flow("infection", params["contact_rate"],
"S", "I")
model.add_fractional_flow("recovery", params["recovery_rate"], "I", "R")
# Request derived outputs
model.request_output_for_compartments("prevalence_susceptible",
compartments=["S"])
model.request_output_for_compartments("prevalence_infectious",
compartments=["I"])
return model
def request_victorian_outputs(model: CompartmentalModel, params: Parameters):
# Victorian cluster model outputs.
clusters = [
Region.to_filename(region) for region in Region.VICTORIA_SUBREGIONS
]
# Track incidence of disease (transition from exposed to active) - overall and for each cluster
model.request_output_for_flow(name="incidence", flow_name="incidence")
for cluster in clusters:
model.request_output_for_flow(
name=f"incidence_for_cluster_{cluster}",
flow_name="incidence",
dest_strata={"cluster": cluster},
)
# Track progress of disease (transition from early to late active)
# - overall, by cluster, by clinical stratum and by both
model.request_output_for_flow(name="progress", flow_name="progress")
for cluster in clusters:
model.request_output_for_flow(
name=f"progress_for_cluster_{cluster}",
flow_name="progress",
dest_strata={"cluster": cluster},
)
# Notifications.
notification_sources = []
for clinical in NOTIFICATION_STRATA:
name = f"progressX{clinical}"
notification_sources.append(name)
model.request_output_for_flow(
name=name,
flow_name="progress",
dest_strata={"clinical": clinical},
save_results=False,
)
model.request_aggregate_output(name="notifications",
sources=notification_sources)
# Cluster-specific notifications.
for cluster in clusters:
cluster_notification_sources = []
for clinical in NOTIFICATION_STRATA:
name = f"progress_for_cluster_{cluster}X{clinical}"
cluster_notification_sources.append(name)
model.request_output_for_flow(
name=name,
flow_name="progress",
dest_strata={
"cluster": cluster,
"clinical": clinical
},
save_results=False,
)
model.request_aggregate_output(
name=f"notifications_for_cluster_{cluster}",
sources=cluster_notification_sources)
model.request_cumulative_output(
name=f"accum_notifications_for_cluster_{cluster}",
source=f"notifications_for_cluster_{cluster}",
)
# Track non-ICU hospital admissions (transition from early to late active in hospital, non-ICU stratum)
model.request_output_for_flow(
name="non_icu_admissions",
flow_name="progress",
source_strata={"clinical": Clinical.HOSPITAL_NON_ICU},
dest_strata={"clinical": Clinical.HOSPITAL_NON_ICU},
)
for cluster in clusters:
model.request_output_for_flow(
name=f"non_icu_admissions_for_cluster_{cluster}",
flow_name="progress",
source_strata={
"clinical": Clinical.HOSPITAL_NON_ICU,
"cluster": cluster
},
dest_strata={
"clinical": Clinical.HOSPITAL_NON_ICU,
"cluster": cluster
},
)
# Track ICU admissions (transition from early to late active in ICU stratum)
model.request_output_for_flow(
name="icu_admissions",
flow_name="progress",
source_strata={"clinical": Clinical.ICU},
dest_strata={"clinical": Clinical.ICU},
)
for cluster in clusters:
model.request_output_for_flow(
name=f"icu_admissions_for_cluster_{cluster}",
flow_name="progress",
source_strata={
"clinical": Clinical.ICU,
"cluster": cluster
},
dest_strata={
"clinical": Clinical.ICU,
"cluster": cluster
},
)
model.request_cumulative_output(
name=f"accum_icu_admissions_for_cluster_{cluster}",
source=f"icu_admissions_for_cluster_{cluster}",
)
# Create hospitalisation functions as sum of hospital non-ICU and ICU
model.request_aggregate_output(
"hospital_admissions",
sources=["icu_admissions", "non_icu_admissions"])
for cluster in clusters:
model.request_aggregate_output(
f"hospital_admissions_for_cluster_{cluster}",
sources=[
f"icu_admissions_for_cluster_{cluster}",
f"non_icu_admissions_for_cluster_{cluster}",
],
)
model.request_cumulative_output(
name=f"accum_hospital_admissions_for_cluster_{cluster}",
source=f"hospital_admissions_for_cluster_{cluster}",
)
# Hospital occupancy
# We count all ICU and hospital late active compartments and a proportion of early active ICU cases.
compartment_periods = params.sojourn.compartment_periods
icu_early_period = compartment_periods["icu_early"]
hospital_early_period = compartment_periods["hospital_early"]
period_icu_patients_in_hospital = max(
icu_early_period - hospital_early_period, 0.0)
proportion_icu_patients_in_hospital = period_icu_patients_in_hospital / icu_early_period
model.request_output_for_compartments(
"_late_active_hospital",
compartments=[Compartment.LATE_ACTIVE],
strata={"clinical": Clinical.HOSPITAL_NON_ICU},
save_results=False,
)
model.request_output_for_compartments(
"icu_occupancy",
compartments=[Compartment.LATE_ACTIVE],
strata={"clinical": Clinical.ICU},
)
model.request_output_for_compartments(
"_early_active_icu",
compartments=[Compartment.EARLY_ACTIVE],
strata={"clinical": Clinical.ICU},
save_results=False,
)
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",
],
)
for cluster in clusters:
model.request_output_for_compartments(
f"_late_active_hospital_for_cluster_{cluster}",
compartments=[Compartment.LATE_ACTIVE],
strata={
"clinical": Clinical.HOSPITAL_NON_ICU,
"cluster": cluster
},
save_results=False,
)
model.request_output_for_compartments(
f"icu_occupancy_for_cluster_{cluster}",
compartments=[Compartment.LATE_ACTIVE],
strata={
"clinical": Clinical.ICU,
"cluster": cluster
},
)
model.request_output_for_compartments(
f"_early_active_icu_for_cluster_{cluster}",
compartments=[Compartment.EARLY_ACTIVE],
strata={
"clinical": Clinical.ICU,
"cluster": cluster
},
save_results=False,
)
model.request_function_output(
name=f"_early_active_icu_proportion_for_cluster_{cluster}",
func=lambda patients: patients *
proportion_icu_patients_in_hospital,
sources=["_early_active_icu"],
save_results=False,
)
model.request_aggregate_output(
name=f"hospital_occupancy_for_cluster_{cluster}",
sources=[
f"_late_active_hospital_for_cluster_{cluster}",
f"icu_occupancy_for_cluster_{cluster}",
f"_early_active_icu_proportion_for_cluster_{cluster}",
],
)
# Infection deaths.
model.request_output_for_flow(name="infection_deaths",
flow_name="infect_death")
model.request_cumulative_output(name="accum_deaths",
source="infection_deaths")
for cluster in clusters:
model.request_output_for_flow(
name=f"infection_deaths_for_cluster_{cluster}",
flow_name="infect_death",
source_strata={"cluster": cluster},
)
model.request_cumulative_output(
name=f"accum_infection_deaths_for_cluster_{cluster}",
source="infection_deaths")
# Proportion seropositive
model.request_output_for_compartments(name="_total_population",
compartments=COMPARTMENTS,
save_results=False)
model.request_output_for_compartments(name="_recovered",
compartments=[Compartment.RECOVERED],
save_results=False)
model.request_function_output(
name="proportion_seropositive",
sources=["_recovered", "_total_population"],
func=lambda recovered, total: recovered / total,
)
for agegroup in AGEGROUP_STRATA:
total_name = f"_total_populationXagegroup_{agegroup}"
recovered_name = f"_recoveredXagegroup_{agegroup}"
model.request_output_for_compartments(
name=total_name,
compartments=COMPARTMENTS,
strata={"agegroup": agegroup},
save_results=False,
)
model.request_output_for_compartments(
name=recovered_name,
compartments=[Compartment.RECOVERED],
strata={"agegroup": agegroup},
save_results=False,
)
model.request_function_output(
name=f"proportion_seropositiveXagegroup_{agegroup}",
sources=[recovered_name, total_name],
func=lambda recovered, total: recovered / total,
)
for cluster in clusters:
total_name = f"_total_populationXcluster_{cluster}"
recovered_name = f"_recoveredXcluster_{cluster}"
model.request_output_for_compartments(
name=total_name,
compartments=COMPARTMENTS,
strata={"cluster": cluster},
save_results=False,
)
model.request_output_for_compartments(
name=recovered_name,
compartments=[Compartment.RECOVERED],
strata={"cluster": cluster},
save_results=False,
)
model.request_function_output(
name=f"proportion_seropositiveXcluster_{cluster}",
sources=[recovered_name, total_name],
func=lambda recovered, total: recovered / total,
)
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 build_model(params: dict) -> CompartmentalModel:
time = params["time"]
model = CompartmentalModel(
times=[time["start"], time["end"]],
compartments=COMPARTMENTS,
infectious_compartments=INFECTIOUS_COMPS,
timestep=time["step"],
)
# Add initial population
init_pop = {
Compartment.EARLY_LATENT: params["initial_early_latent_population"],
Compartment.LATE_LATENT: params["initial_late_latent_population"],
Compartment.INFECTIOUS: params["initial_infectious_population"],
Compartment.DETECTED: params["initial_detected_population"],
Compartment.ON_TREATMENT: params["initial_on_treatment_population"],
Compartment.RECOVERED: 0,
}
sum_init_pop = sum(init_pop.values())
init_pop[Compartment.
SUSCEPTIBLE] = params["start_population_size"] - sum_init_pop
model.set_initial_population(init_pop)
# Add inter-compartmental flows
params = _get_derived_params(params)
# Entry flows
model.add_crude_birth_flow(
"birth",
params["crude_birth_rate"],
Compartment.SUSCEPTIBLE,
)
# Infection flows.
model.add_infection_frequency_flow(
"infection",
params["contact_rate"],
Compartment.SUSCEPTIBLE,
Compartment.EARLY_LATENT,
)
model.add_infection_frequency_flow(
"infection_from_latent",
params["contact_rate_from_latent"],
Compartment.LATE_LATENT,
Compartment.EARLY_LATENT,
)
model.add_infection_frequency_flow(
"infection_from_recovered",
params["contact_rate_from_recovered"],
Compartment.RECOVERED,
Compartment.EARLY_LATENT,
)
# Transition flows.
model.add_fractional_flow(
"treatment_early",
params["preventive_treatment_rate"],
Compartment.EARLY_LATENT,
Compartment.RECOVERED,
)
model.add_fractional_flow(
"treatment_late",
params["preventive_treatment_rate"],
Compartment.LATE_LATENT,
Compartment.RECOVERED,
)
model.add_fractional_flow(
"stabilisation",
params["stabilisation_rate"],
Compartment.EARLY_LATENT,
Compartment.LATE_LATENT,
)
model.add_fractional_flow(
"early_activation",
params["early_activation_rate"],
Compartment.EARLY_LATENT,
Compartment.INFECTIOUS,
)
model.add_fractional_flow(
"late_activation",
params["late_activation_rate"],
Compartment.LATE_LATENT,
Compartment.INFECTIOUS,
)
# Post-active-disease flows
model.add_fractional_flow(
"detection",
params["detection_rate"],
Compartment.INFECTIOUS,
Compartment.DETECTED,
)
model.add_fractional_flow(
"treatment_commencement",
params["treatment_commencement_rate"],
Compartment.DETECTED,
Compartment.ON_TREATMENT,
)
model.add_fractional_flow(
"missed_to_active",
params["missed_to_active_rate"],
Compartment.DETECTED,
Compartment.INFECTIOUS,
)
model.add_fractional_flow(
"self_recovery_infectious",
params["self_recovery_rate"],
Compartment.INFECTIOUS,
Compartment.LATE_LATENT,
)
model.add_fractional_flow(
"self_recovery_detected",
params["self_recovery_rate"],
Compartment.DETECTED,
Compartment.LATE_LATENT,
)
model.add_fractional_flow(
"treatment_recovery",
params["treatment_recovery_rate"],
Compartment.ON_TREATMENT,
Compartment.RECOVERED,
)
model.add_fractional_flow(
"treatment_default",
params["treatment_default_rate"],
Compartment.ON_TREATMENT,
Compartment.INFECTIOUS,
)
model.add_fractional_flow(
"failure_retreatment",
params["failure_retreatment_rate"],
Compartment.ON_TREATMENT,
Compartment.DETECTED,
)
model.add_fractional_flow(
"spontaneous_recovery",
params["spontaneous_recovery_rate"],
Compartment.ON_TREATMENT,
Compartment.LATE_LATENT,
)
# Death flows
# Universal death rate to be overriden by a multiply in age stratification.
uni_death_flow_names = model.add_universal_death_flows("universal_death",
death_rate=1)
model.add_death_flow(
"infectious_death",
params["infect_death_rate"],
Compartment.INFECTIOUS,
)
model.add_death_flow(
"detected_death",
params["infect_death_rate"],
Compartment.DETECTED,
)
model.add_death_flow(
"treatment_death",
params["treatment_death_rate"],
Compartment.ON_TREATMENT,
)
# Apply age-stratification
age_strat = _build_age_strat(params, uni_death_flow_names)
model.stratify_with(age_strat)
# Add vaccination stratification.
vac_strat = _build_vac_strat(params)
model.stratify_with(vac_strat)
# Apply organ stratification
organ_strat = _build_organ_strat(params)
model.stratify_with(organ_strat)
# Apply strain stratification
strain_strat = _build_strain_strat(params)
model.stratify_with(strain_strat)
# Add amplification flow
model.add_fractional_flow(
name="amplification",
fractional_rate=params["amplification_rate"],
source=Compartment.ON_TREATMENT,
dest=Compartment.ON_TREATMENT,
source_strata={"strain": "ds"},
dest_strata={"strain": "mdr"},
expected_flow_count=9,
)
# Add cross-strain reinfection flows
model.add_infection_frequency_flow(
name="reinfection_ds_to_mdr",
contact_rate=params["reinfection_rate"],
source=Compartment.EARLY_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "ds"},
dest_strata={"strain": "mdr"},
expected_flow_count=3,
)
model.add_infection_frequency_flow(
name="reinfection_mdr_to_ds",
contact_rate=params["reinfection_rate"],
source=Compartment.EARLY_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "mdr"},
dest_strata={"strain": "ds"},
expected_flow_count=3,
)
model.add_infection_frequency_flow(
name="reinfection_late_ds_to_mdr",
contact_rate=params["reinfection_rate"],
source=Compartment.LATE_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "ds"},
dest_strata={"strain": "mdr"},
expected_flow_count=3,
)
model.add_infection_frequency_flow(
name="reinfection_late_mdr_to_ds",
contact_rate=params["reinfection_rate"],
source=Compartment.LATE_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "mdr"},
dest_strata={"strain": "ds"},
expected_flow_count=3,
)
# Apply classification stratification
class_strat = _build_class_strat(params)
model.stratify_with(class_strat)
# Apply retention stratification
retention_strat = _build_retention_strat(params)
model.stratify_with(retention_strat)
# Register derived output functions, which are calculations based on the model's compartment values or flows.
# These are calculated after the model is run.
model.request_output_for_flow("notifications", flow_name="detection")
model.request_output_for_flow("early_activation",
flow_name="early_activation")
model.request_output_for_flow("late_activation",
flow_name="late_activation")
model.request_output_for_flow("infectious_deaths",
flow_name="infectious_death")
model.request_output_for_flow("detected_deaths",
flow_name="detected_death")
model.request_output_for_flow("treatment_deaths",
flow_name="treatment_death")
model.request_output_for_flow("progression_early",
flow_name="early_activation")
model.request_output_for_flow("progression_late",
flow_name="late_activation")
model.request_aggregate_output("progression",
["progression_early", "progression_late"])
model.request_output_for_compartments("population_size", COMPARTMENTS)
model.request_aggregate_output(
"_incidence",
sources=["early_activation", "late_activation"],
save_results=False)
model.request_function_output("incidence",
sources=["_incidence", "population_size"],
func=lambda i, p: 1e5 * i / p)
model.request_aggregate_output(
"disease_deaths",
sources=["infectious_deaths", "detected_deaths", "treatment_deaths"])
cum_start_time = params["cumulative_output_start_time"]
model.request_cumulative_output("cumulative_diseased",
source="_incidence",
start_time=cum_start_time)
model.request_cumulative_output("cumulative_deaths",
source="disease_deaths",
start_time=cum_start_time)
model.request_output_for_compartments("_count_infectious",
INFECTIOUS_COMPS,
save_results=False)
model.request_function_output(
"prevalence_infectious",
sources=["_count_infectious", "population_size"],
func=lambda c, p: 1e5 * c / p,
)
model.request_output_for_compartments(
"_count_latent", [Compartment.EARLY_LATENT, Compartment.LATE_LATENT],
save_results=False)
model.request_function_output(
"percentage_latent",
sources=["_count_latent", "population_size"],
func=lambda c, p: 100 * c / p,
)
return model
def request_standard_outputs(
model: CompartmentalModel,
params: Parameters,
):
country = params.country
pop = params.population
is_region_vic = pop.region and Region.to_name(
pop.region) in Region.VICTORIA_SUBREGIONS
# Disease incidence
model.request_output_for_flow(name="incidence", flow_name="incidence")
notification_at_sympt_onset_sources = []
for agegroup in AGEGROUP_STRATA:
# Track incidence for each agegroup.
model.request_output_for_flow(
name=f"incidenceXagegroup_{agegroup}",
flow_name="incidence",
dest_strata={"agegroup": agegroup},
)
for clinical in CLINICAL_STRATA:
# Track incidence for each agegroup and clinical status.
name = f"incidenceXagegroup_{agegroup}Xclinical_{clinical}"
model.request_output_for_flow(
name=name,
flow_name="incidence",
dest_strata={
"agegroup": agegroup,
"clinical": clinical
},
)
if clinical in NOTIFICATION_STRATA:
notification_at_sympt_onset_sources.append(name)
# Notifications at symptom onset.
model.request_aggregate_output(name="notifications_at_sympt_onset",
sources=notification_at_sympt_onset_sources)
# Disease progresssion
model.request_output_for_flow(name="progress", flow_name="progress")
for agegroup in AGEGROUP_STRATA:
for clinical in NOTIFICATION_STRATA:
model.request_output_for_flow(
name=f"progressXagegroup_{agegroup}Xclinical_{clinical}",
flow_name="progress",
dest_strata={
"agegroup": agegroup,
"clinical": clinical
},
)
# New hospital admissions
hospital_sources = []
icu_sources = []
for agegroup in AGEGROUP_STRATA:
icu_sources.append(
f"progressXagegroup_{agegroup}Xclinical_{Clinical.ICU}")
hospital_sources += [
f"progressXagegroup_{agegroup}Xclinical_{Clinical.ICU}",
f"progressXagegroup_{agegroup}Xclinical_{Clinical.HOSPITAL_NON_ICU}",
]
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"progressXagegroup_{a}Xclinical_{c}" for a in AGEGROUP_STRATA
for c in NOTIFICATION_STRATA
]
model.request_aggregate_output(name="local_notifications",
sources=notification_sources)
model.request_aggregate_output(
name="notifications", sources=notification_sources
) # Used to be different coz we had imports.
# Notification by age group
for agegroup in AGEGROUP_STRATA:
sympt_isolate_name = f"progressXagegroup_{agegroup}Xclinical_sympt_isolate"
hospital_non_icu_name = f"progressXagegroup_{agegroup}Xclinical_hospital_non_icu"
icu_name = f"progressXagegroup_{agegroup}Xclinical_icu"
model.request_function_output(
name=f"notificationsXagegroup_{agegroup}",
sources=[sympt_isolate_name, hospital_non_icu_name, icu_name],
func=lambda sympt, hosp, icu: sympt + hosp + icu,
)
# Infection deaths.
model.request_output_for_flow(name="infection_deaths",
flow_name="infect_death")
for agegroup in AGEGROUP_STRATA:
model.request_output_for_flow(
name=f"infection_deathsXagegroup_{agegroup}",
flow_name="infect_death",
source_strata={"agegroup": agegroup},
save_results=False,
)
for clinical in CLINICAL_STRATA:
model.request_output_for_flow(
name=
f"infection_deathsXagegroup_{agegroup}Xclinical_{clinical}",
flow_name="infect_death",
source_strata={
"agegroup": agegroup,
"clinical": clinical
},
)
model.request_cumulative_output(name="accum_deaths",
source="infection_deaths")
if not is_region_vic:
for agegroup in AGEGROUP_STRATA:
model.request_cumulative_output(
name=f"accum_deathsXagegroup_{agegroup}",
source=f"infection_deathsXagegroup_{agegroup}",
)
# Track years of life lost per year.
life_expectancy = inputs.get_life_expectancy_by_agegroup(
AGEGROUP_STRATA, country.iso3)[0]
life_expectancy_latest = [
life_expectancy[agegroup][-1] for agegroup in life_expectancy
]
yoll_sources = []
for idx, agegroup in enumerate(AGEGROUP_STRATA):
# Use default parameter to bind loop variable to function.
l = life_expectancy_latest[idx]
def get_yoll(deaths, life_exp=l):
return deaths * life_exp
yoll_source = f"_yoll_{agegroup}"
yoll_sources.append(yoll_source)
model.request_function_output(
name=yoll_source,
func=get_yoll,
sources=[f"infection_deathsXagegroup_{agegroup}"],
save_results=False,
)
model.request_aggregate_output(name="years_of_life_lost",
sources=yoll_sources)
if params.country.iso3 in OPTI_ISO3S:
# Derived outputs for the optimization project.
model.request_cumulative_output(name="accum_years_of_life_lost",
source="years_of_life_lost")
# Track hospital occupancy.
# We count all ICU and hospital late active compartments and a proportion of early active ICU cases.
compartment_periods = params.sojourn.compartment_periods
icu_early_period = compartment_periods["icu_early"]
hospital_early_period = compartment_periods["hospital_early"]
period_icu_patients_in_hospital = max(
icu_early_period - hospital_early_period, 0.0)
proportion_icu_patients_in_hospital = period_icu_patients_in_hospital / icu_early_period
model.request_output_for_compartments(
"_late_active_hospital",
compartments=[Compartment.LATE_ACTIVE],
strata={"clinical": Clinical.HOSPITAL_NON_ICU},
save_results=False,
)
model.request_output_for_compartments(
"icu_occupancy",
compartments=[Compartment.LATE_ACTIVE],
strata={"clinical": Clinical.ICU},
)
model.request_output_for_compartments(
"_early_active_icu",
compartments=[Compartment.EARLY_ACTIVE],
strata={"clinical": Clinical.ICU},
save_results=False,
)
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=COMPARTMENTS,
save_results=False)
model.request_output_for_compartments(name="_recovered",
compartments=[Compartment.RECOVERED],
save_results=False)
model.request_function_output(
name="proportion_seropositive",
sources=["_recovered", "_total_population"],
func=lambda recovered, total: recovered / total,
)
if not is_region_vic:
for agegroup in AGEGROUP_STRATA:
total_name = f"_total_populationXagegroup_{agegroup}"
recovered_name = f"_recoveredXagegroup_{agegroup}"
model.request_output_for_compartments(
name=total_name,
compartments=COMPARTMENTS,
strata={"agegroup": agegroup},
save_results=False,
)
model.request_output_for_compartments(
name=recovered_name,
compartments=[Compartment.RECOVERED],
strata={"agegroup": agegroup},
save_results=False,
)
model.request_function_output(
name=f"proportion_seropositiveXagegroup_{agegroup}",
sources=[recovered_name, total_name],
func=lambda recovered, total: recovered / total,
)
if params.stratify_by_immunity and params.vaccination:
model.request_output_for_flow(name="vaccination",
flow_name="vaccination")