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 _build_model(self, params):
model = CompartmentalModel(
times=[params["time"]["start"], 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"],
)
model.set_initial_population(distribution={"S": 1000, "I": 1000})
model.add_crude_birth_flow("birth", params["birth_rate"], "S")
model.add_fractional_flow("recovery", params["recovery_rate"], "I",
"R")
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 build_model(params: dict) -> CompartmentalModel:
"""
Build the compartmental model from the provided parameters.
"""
params = Parameters(**params)
model = CompartmentalModel(
times=[params.time.start, params.time.end],
compartments=COMPARTMENTS,
infectious_compartments=INFECTIOUS_COMPARTMENTS,
timestep=params.time.step,
)
# Population distribution
country = params.country
pop = params.population
# Time periods calculated from periods (ie "sojourn times")
compartment_periods = preprocess.compartments.calc_compartment_periods(
params.sojourn)
# Get country population by age-group
total_pops = inputs.get_population_by_agegroup(AGEGROUP_STRATA,
country.iso3,
pop.region,
year=pop.year)
# Distribute infectious seed across infectious split sub-compartments
total_disease_time = sum(
[compartment_periods[c] for c in DISEASE_COMPARTMENTS])
init_pop = {
c: params.infectious_seed * compartment_periods[c] / total_disease_time
for c in DISEASE_COMPARTMENTS
}
# Assign the remainder starting population to the S compartment
init_pop[Compartment.SUSCEPTIBLE] = sum(total_pops) - sum(
init_pop.values())
model.set_initial_population(init_pop)
# Add intercompartmental flows
if params.seasonal_force:
# Use a time-varying, sinusoidal seasonal forcing function for contact rate.
contact_rate = get_seasonal_forcing(365.0, 173.0,
params.seasonal_force,
params.contact_rate)
else:
# Use a static contact rate.
contact_rate = params.contact_rate
# Adjust contact rate for Variant of Concerns
if params.voc_emmergence:
voc_multiplier = scale_up_function(
x=[
params.voc_emmergence.start_time,
params.voc_emmergence.end_time
],
y=[
1.0,
1.0 + params.voc_emmergence.final_proportion *
(params.voc_emmergence.contact_rate_multiplier - 1.0),
],
method=4,
)
raw_contact_rate = contact_rate
if isinstance(contact_rate, float):
def contact_rate(t):
return raw_contact_rate * voc_multiplier(t)
else:
def contact_rate(t):
return raw_contact_rate(t) * voc_multiplier(t)
model.add_infection_frequency_flow(
name="infection",
contact_rate=contact_rate,
source=Compartment.SUSCEPTIBLE,
dest=Compartment.EARLY_EXPOSED,
)
# Infection progress flows.
model.add_fractional_flow(
name="infect_onset",
fractional_rate=1.0 / compartment_periods[Compartment.EARLY_EXPOSED],
source=Compartment.EARLY_EXPOSED,
dest=Compartment.LATE_EXPOSED,
)
model.add_fractional_flow(
name="incidence",
fractional_rate=1.0 / compartment_periods[Compartment.LATE_EXPOSED],
source=Compartment.LATE_EXPOSED,
dest=Compartment.EARLY_ACTIVE,
)
model.add_fractional_flow(
name="progress",
fractional_rate=1.0 / compartment_periods[Compartment.EARLY_ACTIVE],
source=Compartment.EARLY_ACTIVE,
dest=Compartment.LATE_ACTIVE,
)
# Recovery flows
model.add_fractional_flow(
name="recovery",
fractional_rate=1.0 / compartment_periods[Compartment.LATE_ACTIVE],
source=Compartment.LATE_ACTIVE,
dest=Compartment.RECOVERED,
)
# Infection death
model.add_death_flow(
name="infect_death",
death_rate=0, # Will be overwritten later in clinical stratification.
source=Compartment.LATE_ACTIVE,
)
if params.waning_immunity_duration is not None:
# Waning immunity (if requested)
model.add_fractional_flow(
name="warning_immunity",
fractional_rate=1.0 / params.waning_immunity_duration,
source=Compartment.RECOVERED,
dest=Compartment.SUSCEPTIBLE,
)
# Stratify the model by age group.
age_strat = get_agegroup_strat(params, total_pops)
model.stratify_with(age_strat)
# Stratify the model by clinical status
clinical_strat = get_clinical_strat(params)
model.stratify_with(clinical_strat)
# Stratify by immunity - which will include vaccination and infection history
if params.stratify_by_immunity:
immunity_strat = get_immunity_strat(params)
model.stratify_with(immunity_strat)
if params.vaccination:
vacc_params = params.vaccination
for roll_out_component in vacc_params.roll_out_components:
add_vaccination_flows(model, roll_out_component,
age_strat.strata)
# Infection history stratification
if params.stratify_by_infection_history:
history_strat = get_history_strat(params, compartment_periods)
model.stratify_with(history_strat)
# Stratify model by Victorian subregion (used for Victorian cluster model).
if params.victorian_clusters:
cluster_strat = get_cluster_strat(params)
model.stratify_with(cluster_strat)
apply_post_cluster_strat_hacks(params, model)
# Set up derived output functions
if not params.victorian_clusters:
request_standard_outputs(model, params)
else:
request_victorian_outputs(model, params)
return model