def _build_organ_strat(params): organ_strat = Stratification( "organ", ["smear_positive", "smear_negative", "extra_pulmonary"], INFECTIOUS_COMPS) organ_strat.set_population_split(params["organ"]["props"]) # Add infectiousness adjustments for comp in INFECTIOUS_COMPS: organ_strat.add_infectiousness_adjustments( comp, _adjust_all_multiply(params["organ"]["foi"])) organ_strat.add_flow_adjustments( "early_activation", _adjust_all_multiply(params["organ"]["props"])) organ_strat.add_flow_adjustments( "late_activation", _adjust_all_multiply(params["organ"]["props"])) organ_strat.add_flow_adjustments( "detection", _adjust_all_multiply(params["detection_rate_stratified"]["organ"]), ) organ_strat.add_flow_adjustments( "self_recovery_infectious", _adjust_all_multiply(params["self_recovery_rate_stratified"]["organ"]), ) organ_strat.add_flow_adjustments( "self_recovery_detected", _adjust_all_multiply(params["self_recovery_rate_stratified"]["organ"]), ) return organ_strat
def _get_test_model(timestep=1, times=[0, 150]): comps = ["S", "EE", "LE", "EA", "LA", "R"] infectious_comps = ["LE", "EA", "LA"] model = CompartmentalModel( times=times, compartments=comps, infectious_compartments=infectious_comps, timestep=timestep, ) model.set_initial_population({"S": int(20e6), "LA": 100}) # Add flows model.add_infection_frequency_flow(name="infection", contact_rate=0.03, source="S", dest="EE") model.add_sojourn_flow(name="infect_onset", sojourn_time=7, source="EE", dest="LE") model.add_sojourn_flow(name="incidence", sojourn_time=7, source="LE", dest="EA") model.add_sojourn_flow(name="progress", sojourn_time=7, source="EA", dest="LA") model.add_sojourn_flow(name="recovery", sojourn_time=7, source="LA", dest="R") model.add_death_flow(name="infect_death", death_rate=0.005, source="LA") model.add_transition_flow(name="warning_immunity", fractional_rate=0.01, source="R", dest="S") # Stratify by age age_strat = Stratification("age", AGE_STRATA, comps) age_strat.set_population_split(AGE_SPLIT_PROPORTIONS) age_strat.set_mixing_matrix(AGE_MIXING_MATRIX) age_strat.add_flow_adjustments( "infection", {s: Multiply(v) for s, v in AGE_SUSCEPTIBILITY.items()} ) model.stratify_with(age_strat) # Stratify by clinical status clinical_strat = Stratification("clinical", CLINICAL_STRATA, infectious_comps) clinical_strat.add_infectiousness_adjustments("LE", {**ADJ_BASE, "non_sympt": Overwrite(0.25)}) clinical_strat.add_infectiousness_adjustments("EA", {**ADJ_BASE, "non_sympt": Overwrite(0.25)}) clinical_strat.add_infectiousness_adjustments( "LA", { **ADJ_BASE, "non_sympt": Overwrite(0.25), "sympt_isolate": Overwrite(0.2), "hospital": Overwrite(0.2), "icu": Overwrite(0.2), }, ) clinical_strat.add_flow_adjustments( "infect_onset", { "non_sympt": Multiply(0.26), "icu": Multiply(0.01), "hospital": Multiply(0.04), "sympt_public": Multiply(0.66), "sympt_isolate": Multiply(0.03), }, ) model.stratify_with(clinical_strat) # Request derived outputs. model.request_output_for_flow(name="incidence", flow_name="incidence") model.request_output_for_flow(name="progress", flow_name="progress") for age in AGE_STRATA: for clinical in NOTIFICATION_STRATA: model.request_output_for_flow( name=f"progressXage_{age}Xclinical_{clinical}", flow_name="progress", dest_strata={"age": age, "clinical": clinical}, ) hospital_sources = [] icu_sources = [] for age in AGE_STRATA: icu_sources.append(f"progressXage_{age}Xclinical_icu") hospital_sources += [ f"progressXage_{age}Xclinical_icu", f"progressXage_{age}Xclinical_hospital", ] model.request_aggregate_output( name="new_hospital_admissions", sources=hospital_sources, ) model.request_aggregate_output(name="new_icu_admissions", sources=icu_sources) # Get notifications, which may included people detected in-country as they progress, or imported cases which are detected. notification_sources = [ f"progressXage_{a}Xclinical_{c}" for a in AGE_STRATA for c in NOTIFICATION_STRATA ] model.request_aggregate_output(name="notifications", sources=notification_sources) # Infection deaths. model.request_output_for_flow(name="infection_deaths", flow_name="infect_death") model.request_cumulative_output(name="accum_deaths", source="infection_deaths") # Track hospital occupancy. # We count all ICU and hospital late active compartments and a proportion of early active ICU cases. model.request_output_for_compartments( "_late_active_hospital", compartments=["LA"], strata={"clinical": "hospital"}, save_results=False, ) model.request_output_for_compartments( "icu_occupancy", compartments=["LA"], strata={"clinical": "icu"}, ) model.request_output_for_compartments( "_early_active_icu", compartments=["EA"], strata={"clinical": "icu"}, save_results=False, ) proportion_icu_patients_in_hospital = 0.25 model.request_function_output( name="_early_active_icu_proportion", func=lambda patients: patients * proportion_icu_patients_in_hospital, sources=["_early_active_icu"], save_results=False, ) model.request_aggregate_output( name="hospital_occupancy", sources=[ "_late_active_hospital", "icu_occupancy", "_early_active_icu_proportion", ], ) # Proportion seropositive model.request_output_for_compartments( name="_total_population", compartments=comps, save_results=False ) model.request_output_for_compartments(name="_recovered", compartments=["R"], save_results=False) model.request_function_output( name="proportion_seropositive", sources=["_recovered", "_total_population"], func=lambda recovered, total: recovered / total, ) return model
def get_clinical_strat(params: Parameters) -> Stratification: """ Stratify the model by clinical status Stratify the infectious compartments of the covid model (not including the pre-symptomatic compartments, which are actually infectious) - notifications are derived from progress from early to late for some strata - the proportion of people moving from presympt to early infectious, conditioned on age group - rate of which people flow through these compartments (reciprocal of time, using within_* which is a rate of ppl / day) - infectiousness levels adjusted by early/late and for clinical strata - we start with an age stratified infection fatality rate - 50% of deaths for each age bracket die in ICU - the other deaths go to hospital, assume no-one else can die from COVID - should we ditch this? """ clinical_strat = Stratification("clinical", CLINICAL_STRATA, INFECTIOUS_COMPARTMENTS) clinical_params = params.clinical_stratification country = params.country pop = params.population """ Infectiousness adjustments for clinical stratification """ # Add infectiousness reduction multiplier for all non-symptomatic infectious people. # These people are less infectious because of biology. non_sympt_adjust = Overwrite(clinical_params.non_sympt_infect_multiplier) clinical_strat.add_infectiousness_adjustments( Compartment.LATE_EXPOSED, { Clinical.NON_SYMPT: non_sympt_adjust, Clinical.SYMPT_NON_HOSPITAL: None, Clinical.SYMPT_ISOLATE: None, Clinical.HOSPITAL_NON_ICU: None, Clinical.ICU: None, }, ) clinical_strat.add_infectiousness_adjustments( Compartment.EARLY_ACTIVE, { Clinical.NON_SYMPT: non_sympt_adjust, Clinical.SYMPT_NON_HOSPITAL: None, Clinical.SYMPT_ISOLATE: None, Clinical.HOSPITAL_NON_ICU: None, Clinical.ICU: None, }, ) # Add infectiousness reduction for people who are late active and in isolation or hospital/icu. # These people are less infectious because of physical distancing/isolation/PPE precautions. late_infect_multiplier = clinical_params.late_infect_multiplier clinical_strat.add_infectiousness_adjustments( Compartment.LATE_ACTIVE, { Clinical.NON_SYMPT: non_sympt_adjust, Clinical.SYMPT_ISOLATE: Overwrite(late_infect_multiplier[Clinical.SYMPT_ISOLATE]), Clinical.SYMPT_NON_HOSPITAL: None, Clinical.HOSPITAL_NON_ICU: Overwrite(late_infect_multiplier[Clinical.HOSPITAL_NON_ICU]), Clinical.ICU: Overwrite(late_infect_multiplier[Clinical.ICU]), }, ) """ Adjust infection death rates for hospital patients (ICU and non-ICU) """ symptomatic_adjuster = params.clinical_stratification.props.symptomatic.multiplier hospital_adjuster = params.clinical_stratification.props.hospital.multiplier ifr_adjuster = params.infection_fatality.multiplier # Get all the adjustments in the same way as we will do if the immunity stratification is implemented entry_adjustments, death_adjs, progress_adjs, recovery_adjs, _, _ = get_all_adjs( clinical_params, country, pop, params.infection_fatality.props, params.sojourn, params.testing_to_detection, params.case_detection, ifr_adjuster, symptomatic_adjuster, hospital_adjuster, ) # Assign all the adjustments to the model for agegroup in AGEGROUP_STRATA: source = {"agegroup": agegroup} clinical_strat.add_flow_adjustments( "infect_onset", entry_adjustments[agegroup], source_strata=source ) clinical_strat.add_flow_adjustments( "infect_death", death_adjs[agegroup], source_strata=source ) clinical_strat.add_flow_adjustments( "progress", progress_adjs, source_strata=source, ) clinical_strat.add_flow_adjustments( "recovery", recovery_adjs[agegroup], source_strata=source, ) return clinical_strat
def test_create_stratification__with_infectiousness_adjustments(): strat = Stratification(name="location", strata=["rural", "urban"], compartments=["S", "I", "R"]) assert strat.infectiousness_adjustments == {} # Fail coz not all strata specified with pytest.raises(AssertionError): strat.add_infectiousness_adjustments( compartment_name="S", adjustments={"rural": Multiply(1.2)}, ) # Fail coz an incorrect strata specified with pytest.raises(AssertionError): strat.add_infectiousness_adjustments( compartment_name="S", adjustments={ "rural": Multiply(1.2), "urban": Multiply(0.8), "alpine": Multiply(1.1), }, ) # Fail coz a time-varying function was used (not allowed!) with pytest.raises(AssertionError): strat.add_infectiousness_adjustments( compartment_name="S", adjustments={ "rural": Multiply(1.2), "urban": Multiply(lambda t: 2), }, ) strat.add_infectiousness_adjustments( compartment_name="S", adjustments={ "rural": Multiply(1.2), "urban": Multiply(2), }, ) assert strat.infectiousness_adjustments["S"]["rural"]._is_equal( Multiply(1.2)) assert strat.infectiousness_adjustments["S"]["urban"]._is_equal( Multiply(2)) # Fail coz we just did this with pytest.raises(AssertionError): strat.add_infectiousness_adjustments( compartment_name="S", adjustments={ "rural": Multiply(1.2), "urban": Multiply(2), }, ) strat.add_infectiousness_adjustments( compartment_name="I", adjustments={ "rural": Multiply(1.2), "urban": None, }, ) assert strat.infectiousness_adjustments["I"]["rural"]._is_equal( Multiply(1.2)) assert strat.infectiousness_adjustments["I"]["urban"] is None
def test_strat_infectiousness__with_adjustments(): """ Ensure multiply infectiousness adjustment is applied. """ # Create a model model = CompartmentalModel( times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"] ) model.set_initial_population(distribution={"S": 900, "I": 100}) strat = Stratification("age", ["baby", "child", "adult"], ["S", "I", "R"]) strat.set_population_split({"baby": 0.1, "child": 0.3, "adult": 0.6}) strat.add_infectiousness_adjustments( "I", {"child": adjust.Multiply(3), "adult": adjust.Multiply(0.5), "baby": None} ) model.stratify_with(strat) assert_array_equal( model.initial_population, np.array([90, 270, 540, 10, 30, 60, 0, 0, 0]), ) # Do pre-run force of infection calcs. model._prepare_to_run() assert_array_equal( model._compartment_infectiousness["default"], np.array([0, 0, 0, 1, 3, 0.5, 0, 0, 0]), ) # Do pre-iteration force of infection calcs model._prepare_time_step(0, model.initial_population) # Get multipliers infectees = model.compartments[0:3] infectors = model.compartments[3:6] expected_density = 10 * 1 + 30 * 3 + 60 * 0.5 expected_frequency = expected_density / 1000 for infectee, infector in zip(infectees, infectors): assert model._get_infection_density_multiplier(infectee, infector) == expected_density for infectee, infector in zip(infectees, infectors): assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency # Stratify again, now with overwrites strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"]) strat.add_infectiousness_adjustments( "I", {"urban": adjust.Overwrite(1), "rural": adjust.Multiply(7)} ) model.stratify_with(strat) assert_array_equal( model.initial_population, np.array([45, 45, 135, 135, 270.0, 270, 5, 5, 15, 15, 30, 30, 0, 0, 0, 0, 0, 0]), ) # Do pre-run force of infection calcs. model._prepare_to_run() assert_array_equal( model._compartment_infectiousness["default"], np.array([0, 0, 0, 0, 0, 0, 1, 7, 1, 21, 1, 3.5, 0, 0, 0, 0, 0, 0]), ) # Do pre-iteration force of infection calcs model._prepare_time_step(0, model.initial_population) # Get multipliers infectees = model.compartments[0:6] infectors = model.compartments[6:12] expected_density = 5 * 1 + 5 * 7 + 15 * 1 + 15 * 21 + 30 * 1 + 30 * 3.5 expected_frequency = expected_density / 1000 for infectee, infector in zip(infectees, infectors): assert model._get_infection_density_multiplier(infectee, infector) == expected_density for infectee, infector in zip(infectees, infectors): assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency