def test_add_exit_flows_post_stratification(): """ Ensure user can add exit flows post stratification. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) assert len(model._flows) == 0 # Apply partial stratification strat = Stratification("location", ["urban", "rural"], ["S", "I"]) model.stratify_with(strat) assert len(model._flows) == 0 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", {"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)])
def test_stochastic_exitinction(death_rate, recovery_rate, contact_rate): """ A smokey test to make sure the disease goes extinct sometimes, because the infectious person recovers before they can infect someone else. Calculations similar to test_stochastic_death_exitinction """ pr_death_or_recovery = 1 - np.exp(-1 * (death_rate + recovery_rate)) pr_infected = 1 - np.exp(-contact_rate / 1000) pr_noone_infected = binom.pmf(0, 1000, pr_infected) pr_extinction = pr_death_or_recovery * pr_noone_infected expected_extinctions = _find_num_successes(pr_extinction, TRIALS, ERROR_RATE) count_extinctions = 0 for _ in range(TRIALS): model = CompartmentalModel( times=[0, 1], compartments=["S", "I", "R"], infectious_compartments=["I"], ) model.set_initial_population(distribution={"S": 999, "I": 1}) model.add_death_flow("infect_death", death_rate, "I") model.add_transition_flow("recovery", recovery_rate, "I", "R") model.add_infection_frequency_flow("infection", contact_rate, "S", "I") model.run_stochastic() is_extinct = model.outputs[1, 1] == 0 if is_extinct: count_extinctions += 1 assert count_extinctions >= expected_extinctions
def test_stochastic_death_exitinction(death_rate, contact_rate): """ A smokey test to make sure the disease goes extinct around the right amount, because the infectious person dies before they can infect someone else. See here for how this stuff is calculated https://autumn-files.s3-ap-southeast-2.amazonaws.com/Switching_to_stochastic_mode.pdf Consider the following flow rates: - 0.5 infected deaths timestep - 2 people infected per timestep - infection frequency force of infection of 1 inf / 1000 pop - sus pop of 999 - contact rate of 2 - flow rate of 2 * 999 / 1000 = 1.998 ~= 2 Based on stochastic model (per person) - P(infect_death) ~=40%(1 - e^(-0.5/1)) - P(infected) ~= 0.2% (1 - e^(-2/1000)) Using a binomial calculator, we get - ~86% chance of 1 or more people getting infected - ~14% chance of noone getting infected Death and infection are independent processes within the model. So then we expect a ~6% chance of exctinction (infected person dies, no one infected) (40% * 14%) Given this there is a > 0.999999 chance that we see at least 25 disease exctinctions in 1000 runs (using binomial calculation) """ pr_death = 1 - np.exp(-death_rate) pr_infected = 1 - np.exp(-contact_rate / 1000) pr_noone_infected = binom.pmf(0, 1000, pr_infected) pr_extinction = pr_death * pr_noone_infected expected_extinctions = _find_num_successes(pr_extinction, TRIALS, ERROR_RATE) count_extinctions = 0 for _ in range(TRIALS): model = CompartmentalModel( times=[0, 1], compartments=["S", "I", "R"], infectious_compartments=["I"], ) model.set_initial_population(distribution={"S": 999, "I": 1}) model.add_death_flow("infect_death", death_rate, "I") model.add_infection_frequency_flow("infection", contact_rate, "S", "I") model.run_stochastic() is_extinct = model.outputs[1, 1] == 0 if is_extinct: count_extinctions += 1 assert count_extinctions >= expected_extinctions
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 test_model__with_complex_dynamics__expect_correct_outputs(): """ Ensure that a model with the "full suite" of TB dynamics produces correct results: - 5 compartments - birth rate + universal death rate - standard inter-compartment flows """ model = CompartmentalModel( times=[0, 5], compartments=["S", "EL", "LL", "I", "R"], infectious_compartments=["I"] ) model.set_initial_population(distribution={"S": 900, "I": 100}) model.add_crude_birth_flow("births", 0.02, "S") model.add_universal_death_flows("universal_deaths", 0.02) model.add_infection_frequency_flow("infection", 14, "S", "EL") model.add_infection_frequency_flow("reinfection", 14, "R", "EL") model.add_infection_frequency_flow("regression", 3, "LL", "EL") model.add_transition_flow("early_progression", 2, "EL", "I") model.add_transition_flow("stabilisation", 3, "EL", "LL") model.add_transition_flow("early_progression", 1, "LL", "I") model.add_death_flow("infect_death", 0.4, "I") model.add_transition_flow("recovery", 0.2, "I", "R") model.add_transition_flow("case_detection", 1, "I", "R") model.run() # Expect that the results are consistent, nothing crazy happens. # These results were not independently calculated, so this is more of an "acceptance test". expected_outputs = np.array( [ [900.0, 0.0, 0.0, 100.0, 0.0], [66.1, 203.8, 274.2, 307.2, 75.3], [2.9, 150.9, 220.5, 345.3, 69.4], [2.2, 127.3, 175.6, 297.0, 58.1], [1.8, 106.4, 145.6, 248.8, 48.5], [1.5, 88.8, 121.5, 207.8, 40.5], ] ) assert_allclose(model.outputs, expected_outputs, atol=0.2, verbose=True)
def test_model__with_infect_death_rate__expect_infected_pop_decrease(): """ Ensure that a model with two compartments and only infected death rate dynamics results in fewer infected people, but no change to susceptible pop. """ # Set up a model with 100 people, all infectious. model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.set_initial_population(distribution={"S": 50, "I": 50}) # Add some dying at ~2 people / 100 / year. model.add_death_flow("infect_death", 0.02, "I") model.run() expected_outputs = np.array( [ [50.00, 50.00, 0], # Initial conditions [50.00, 49.01, 0], [50.00, 48.04, 0], [50.00, 47.09, 0], [50.00, 46.16, 0], [50.00, 45.24, 0], ] ) assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
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 test_solve_stochastic(monkeypatch): """ Test that _solve_stochastic glue code works. Don't test the actual flow rate calculations or stochastic sampling bits. """ model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"], ) # Add some people to the model, expect initial conditions of [990, 10, 0] model.set_initial_population(distribution={"S": 990, "I": 10}) # Add flows - the parameters add here will be overidden by `mock_get_rates` # but the flow directions will be used. model.add_crude_birth_flow("birth", 8, "S") model.add_infection_frequency_flow("infection", 6, "S", "I") model.add_death_flow("infect_death", 3, "I") model.add_transition_flow("recovery", 2, "I", "R") # Mock out flow rate calculation - tested elsewhere and tricky to predict. def mock_get_rates(comp_vals, time): # Return the flow rates that will be used to solve the model return None, np.array([float(f.param) for f in model._flows]) monkeypatch.setattr(model, "_get_rates", mock_get_rates) # Mock out stochastic flow sampling - tested elsewhere. def mock_sample_entry_flows(seed, entry_flow_rates, timestep): assert not seed assert 0 < timestep <= 5 expected_flow_rates = np.array([8, 0, 0]) assert_array_equal(entry_flow_rates, expected_flow_rates) return np.array([8, 0, 0]) def mock_sample_transistion_flows(seed, flow_rates, flow_map, comp_vals, timestep): assert not seed assert 0 < timestep <= 5 # Flows get re-arranged by setup process expected_flow_map = np.array([[0, 1, -1], [2, 0, 1], [3, 1, 2]]) assert_array_equal(flow_map, expected_flow_map) expected_flow_rates = np.array( [[0.0, 3.0, 0.0], [0.0, 0.0, 0.0], [6.0, 0.0, 0.0], [0.0, 2.0, 0.0]] ) assert_array_equal(flow_rates, expected_flow_rates) return np.array([-6, 1, 2]) monkeypatch.setattr(stochastic, "sample_entry_flows", mock_sample_entry_flows) monkeypatch.setattr(stochastic, "sample_transistion_flows", mock_sample_transistion_flows) model.run_stochastic() expected_outputs = np.array( [ [990, 10, 0], [992, 11, 2], [994, 12, 4], [996, 13, 6], [998, 14, 8], [1000, 15, 10], ] ) assert_array_equal(model.outputs, expected_outputs)
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