def test_stratify_compartments_partial(): """Ensure compartment names and values get partially stratified properly""" model = StratifiedModel( times=[0, 1, 2, 3, 4, 5], compartment_names=["sus", "inf"], initial_conditions={"inf": 200}, parameters={}, requested_flows=[], starting_population=1000, infectious_compartments=["inf"], birth_approach=BirthApproach.NO_BIRTH, entry_compartment="sus", ) assert model.compartment_names == ["sus", "inf"] vals = np.array([800, 200]) assert_array_equal(model.compartment_values, vals) model.stratify( "location", strata_request=["home", "work", "other"], compartments_to_stratify=["sus"], requested_proportions={"home": 0.6}, ) assert model.compartment_names == [ "susXlocation_home", "susXlocation_work", "susXlocation_other", "inf", ] vals = np.array([480, 160, 160, 200]) assert_array_equal(model.compartment_values, vals) assert model.all_stratifications == {"location": ["home", "work", "other"]} assert model.full_stratification_list == []
def test_no_mixing_matrix(): """ Test that we are using the default 'null-op' mixing matrix when we have a no user-specified mixing matrix """ model = StratifiedModel(**MODEL_KWARGS) default_matrix = np.array([[2, 3], [5, 7]]) model.stratify( stratification_name="agegroup", strata_request=["child", "adult"], compartments_to_stratify=["S", "I", "R"], mixing_matrix=default_matrix, ) # We should get the default mixing matrix actual_mixing = model.get_mixing_matrix(0) assert_array_equal(actual_mixing, default_matrix) # Static matrices shouldn't change over time actual_mixing = model.get_mixing_matrix(123) assert_array_equal(actual_mixing, default_matrix) model.stratify( stratification_name="location", strata_request=["work", "home"], # These kids have jobs. compartments_to_stratify=["S", "I", "R"], mixing_matrix=None, ) # We should get the default mixing matrix actual_mixing = model.get_mixing_matrix(0) assert_array_equal(actual_mixing, default_matrix) # Static matrices shouldn't change over time actual_mixing = model.get_mixing_matrix(123) assert_array_equal(actual_mixing, default_matrix)
def test_multiple_dynamic_mixing_matrices(): """ Test that we are using the correct mixing matrix when we have multiple dynamic mixing matrices """ model = StratifiedModel(**MODEL_KWARGS) agegroup_matrix = lambda t: t * np.array([[2, 3], [5, 7]]) model.stratify( stratification_name="agegroup", strata_request=["child", "adult"], compartments_to_stratify=["S", "I", "R"], mixing_matrix=agegroup_matrix, ) location_matrix = lambda t: t * np.array([[11, 13], [17, 19]]) model.stratify( stratification_name="location", strata_request=["work", "home"], # These kids have jobs. compartments_to_stratify=["S", "I", "R"], mixing_matrix=location_matrix, ) expected_mixing = np.array([ [2 * 11, 2 * 13, 3 * 11, 3 * 13], [2 * 17, 2 * 19, 3 * 17, 3 * 19], [5 * 11, 5 * 13, 7 * 11, 7 * 13], [5 * 17, 5 * 19, 7 * 17, 7 * 19], ]) # We should get the Kronecker product of the two matrices actual_mixing = model.get_mixing_matrix(1) assert_array_equal(actual_mixing, expected_mixing) # Double check that we calculated the Kronecker product correctly kron_mixing = np.kron(agegroup_matrix(1), location_matrix(1)) assert_array_equal(expected_mixing, kron_mixing) # Dynamic matrices should change over time actual_mixing = model.get_mixing_matrix(5) assert_array_equal(actual_mixing, 25 * expected_mixing)
def _get_model(): pop = 1000 params = { "contact_rate": 0.1, "recovery_rate": 0.3, "infect_death": 0.5, "crude_birth_rate": 0.7, "universal_death_rate": 0.11, } flows = [ { "type": Flow.INFECTION_FREQUENCY, "origin": "S", "to": "I", "parameter": "contact_rate" }, { "type": Flow.STANDARD, "to": "R", "origin": "I", "parameter": "recovery_rate" }, { "type": Flow.DEATH, "origin": "I", "parameter": "infect_death" }, ] model = StratifiedModel( times=np.array([5.0, 6.0, 7.0, 8.0, 9.0, 10.0]), compartment_names=["S", "I", "R"], initial_conditions={"S": pop}, parameters=params, requested_flows=flows, starting_population=pop, infectious_compartments=["I"], birth_approach=BirthApproach.ADD_CRUDE, entry_compartment="S", ) # Add basic age stratification model.stratify("agegroup", strata_request=[0, 5, 15, 60], compartments_to_stratify=["S", "I"]) # Pretend we ran the model model.prepare_to_run() model.outputs = np.array([ [250, 250, 250, 250, 0, 0, 0, 0, 0], [150, 150, 150, 150, 100, 100, 100, 100, 0], [100, 100, 100, 100, 100, 100, 100, 100, 200], [50, 50, 50, 50, 150, 150, 150, 150, 200], [50, 50, 50, 50, 100, 100, 100, 100, 400], [0, 0, 0, 0, 150, 150, 150, 150, 400], ]) return model
def test_add_extra_flows__before_strat_then_stratify(): """ Ensure that adding an extra flow without any stratifications creates a new flow. Then check that stratifying the model creates the new stratified flows correctly. """ # Create the model model = StratifiedModel( times=np.array([1.0, 2.0, 3.0, 4.0, 5.0]), compartment_names=["S", "I", "R"], initial_conditions={"S": 100}, parameters={}, requested_flows=[], starting_population=1000, infectious_compartments=["I"], entry_compartment="S", birth_approach=BirthApproach.NO_BIRTH, ) # Add a new flow. model.add_extra_flow( flow={ "type": Flow.INFECTION_FREQUENCY, "origin": "S", "to": "I", "parameter": "contact_rate", }, source_strata={}, dest_strata={}, expected_flow_count=1, ) # Ensure the flow was added correctly. assert len(model.flows) == 1 flow = model.flows[0] assert type(flow) is InfectionFrequencyFlow assert flow.source == "S" assert flow.dest == "I" assert flow.adjustments == [] assert flow.param_name == "contact_rate" # Stratify the model. model.stratify("agegroup", strata_request=["young", "old"], compartments_to_stratify=["S", "I", "R"]) # Ensure the new flow was stratified correctly into two flows.. assert len(model.flows) == 2 flow_1 = model.flows[0] assert type(flow_1) is InfectionFrequencyFlow assert flow_1.source == "SXagegroup_young" assert flow_1.dest == "IXagegroup_young" assert flow_1.adjustments == [] assert flow_1.param_name == "contact_rate" flow_2 = model.flows[1] assert type(flow_2) is InfectionFrequencyFlow assert flow_2.source == "SXagegroup_old" assert flow_2.dest == "IXagegroup_old" assert flow_2.adjustments == [] assert flow_2.param_name == "contact_rate"
def test_single_dynamic_mixing_matrix(): """ Test that we are using the correct mixing matrix when we have a single dynamic mixing matrix """ model = StratifiedModel(**MODEL_KWARGS) agegroup_matrix = lambda t: t * np.array([[2, 3], [5, 7]]) model.stratify( stratification_name="agegroup", strata_request=["child", "adult"], compartments_to_stratify=["S", "I", "R"], mixing_matrix=agegroup_matrix, ) # We should get the age mixing matrix actual_mixing = model.get_mixing_matrix(1) assert_array_equal(actual_mixing, agegroup_matrix(1)) # Dynamic matrices should change over time actual_mixing = model.get_mixing_matrix(123) assert_array_equal(actual_mixing, agegroup_matrix(123))
def test_strat_model__with_age_and_starting_proportion__expect_ageing(): """ Ensure that a module with age stratification and starting proporptions produces ageing flows, and the correct output. """ pop = 1000 model = StratifiedModel( times=_get_integration_times(2000, 2005, 1), compartment_names=["S", "I"], initial_conditions={"S": pop}, parameters={}, requested_flows=[], starting_population=pop, infectious_compartments=["I"], birth_approach=BirthApproach.NO_BIRTH, entry_compartment="S", ) # Add basic age stratification model.stratify( "age", strata_request=[0, 5, 15, 60], compartments_to_stratify=["S", "I"], comp_split_props={ "0": 0.8, "5": 0.1, "15": 0.1 }, ) # Run the model for 5 years. model.run_model(integration_type=IntegrationType.ODE_INT) # Expect everyone to generally get older, but no one should die or get sick. # Expect initial distribution of ages to be set according to "requested_proportions". expected_arr = np.array([ [800.0, 100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0], [655.0, 228.3, 114.4, 2.4, 0.0, 0.0, 0.0, 0.0], [536.3, 319.3, 139.3, 5.2, 0.0, 0.0, 0.0, 0.0], [439.1, 381.3, 171.1, 8.6, 0.0, 0.0, 0.0, 0.0], [359.5, 420.5, 207.2, 12.8, 0.0, 0.0, 0.0, 0.0], [294.4, 442.4, 245.4, 17.8, 0.0, 0.0, 0.0, 0.0], ]) assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_strat_model__with_locations__expect_no_change(): """ Ensure that a module with location stratification populates locations correctly. """ pop = 1000 model = StratifiedModel( times=_get_integration_times(2000, 2005, 1), compartment_names=["S", "I"], initial_conditions={"S": pop}, parameters={}, requested_flows=[], starting_population=pop, infectious_compartments=["I"], birth_approach=BirthApproach.NO_BIRTH, entry_compartment="S", ) # Add basic location stratification model.stratify( "location", strata_request=["rural", "urban", "prison"], compartments_to_stratify=["S", "I"], comp_split_props={ "rural": 0.44, "urban": 0.55, "prison": 0.01 }, ) # Run the model for 5 years. model.run_model(integration_type=IntegrationType.ODE_INT) # Expect everyone to start in their locations, then nothing should change, expected_arr = np.array([ [440.0, 550.0, 10.0, 0.0, 0.0, 0.0], [440.0, 550.0, 10.0, 0.0, 0.0, 0.0], [440.0, 550.0, 10.0, 0.0, 0.0, 0.0], [440.0, 550.0, 10.0, 0.0, 0.0, 0.0], [440.0, 550.0, 10.0, 0.0, 0.0, 0.0], [440.0, 550.0, 10.0, 0.0, 0.0, 0.0], ]) assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_strat_model__with_age__expect_ageing(): """ Ensure that a module with age stratification produces ageing flows, and the correct output. """ pop = 1000 model = StratifiedModel( times=_get_integration_times(2000, 2005, 1), compartment_names=["S", "I"], initial_conditions={"S": pop}, parameters={}, requested_flows=[], starting_population=pop, infectious_compartments=["I"], birth_approach=BirthApproach.NO_BIRTH, entry_compartment="S", ) # Add basic age stratification model.stratify("age", strata_request=[0, 5, 15, 60], compartments_to_stratify=["S", "I"]) # Run the model for 5 years. model.run_model(integration_type=IntegrationType.ODE_INT) # Expect everyone to generally get older, but no one should die or get sick expected_arr = np.array([ [250.0, 250.0, 250.0, 250.0, 0.0, 0.0, 0.0, 0.0], [204.8, 269.1, 270.3, 255.8, 0.0, 0.0, 0.0, 0.0], [167.7, 278.8, 291.5, 262.0, 0.0, 0.0, 0.0, 0.0], [137.3, 281.2, 312.8, 268.7, 0.0, 0.0, 0.0, 0.0], [112.5, 277.9, 333.7, 275.9, 0.0, 0.0, 0.0, 0.0], [92.1, 270.8, 353.5, 283.6, 0.0, 0.0, 0.0, 0.0], ]) assert_allclose(model.outputs, expected_arr, atol=0.1, verbose=True)
def test_add_extra_flows__after_single_strat_with_cross_cut(): """ Ensure that adding an extra flow with an existing stratification creates a new flow. This flow cuts across strata. """ # Create the model model = StratifiedModel( times=np.array([1.0, 2.0, 3.0, 4.0, 5.0]), compartment_names=["S", "I", "R"], initial_conditions={"S": 100}, parameters={}, requested_flows=[], starting_population=1000, infectious_compartments=["I"], entry_compartment="S", birth_approach=BirthApproach.NO_BIRTH, ) # Ensure there are no flows yet. assert len(model.flows) == 0 # Stratify the model. model.stratify("agegroup", strata_request=["young", "old"], compartments_to_stratify=["S", "I", "R"]) # Add a new flow. model.add_extra_flow( flow={ "type": Flow.INFECTION_FREQUENCY, "origin": "S", "to": "I", "parameter": "contact_rate", }, source_strata={"agegroup": "young"}, dest_strata={"agegroup": "old"}, expected_flow_count=1, ) # Ensure the new flow was added correctly. assert len(model.flows) == 1 flow = model.flows[0] assert type(flow) is InfectionFrequencyFlow assert flow.source == "SXagegroup_young" assert flow.dest == "IXagegroup_old" assert flow.adjustments == [] assert flow.param_name == "contact_rate" # Stratify the model again. model.stratify("location", strata_request=["urban", "rural"], compartments_to_stratify=["S", "I", "R"]) # Ensure the new flow was stratified correctly into two flows. assert len(model.flows) == 2 flow_1 = model.flows[0] assert type(flow_1) is InfectionFrequencyFlow assert flow_1.source == "SXagegroup_youngXlocation_urban" assert flow_1.dest == "IXagegroup_oldXlocation_urban" assert flow_1.adjustments == [] assert flow_1.param_name == "contact_rate" flow_2 = model.flows[1] assert type(flow_2) is InfectionFrequencyFlow assert flow_2.source == "SXagegroup_youngXlocation_rural" assert flow_2.dest == "IXagegroup_oldXlocation_rural" assert flow_2.adjustments == [] assert flow_2.param_name == "contact_rate"
def test_stratify_flows_full__with_adjustment_requests(): """Ensure flows get stratified properly in full strat""" adjustment_requests = { "contact_rate": { "home": "home_contact_rate", "work": 0.5, }, "recovery_rate": { "home": 1, "work": 2, }, "infect_death": { "home": 1, "work": 2, }, } requested_flows = [ { "type": Flow.INFECTION_FREQUENCY, "parameter": "contact_rate", "origin": "sus", "to": "inf", }, { "type": Flow.STANDARD, "parameter": "recovery_rate", "origin": "inf", "to": "sus", }, { "type": Flow.DEATH, "parameter": "infect_death", "origin": "inf", }, ] parameters = { "contact_rate": 1000, "recovery_rate": "recovery_rate", "infect_death": 10, } home_contact_rate_func = lambda t: t recovery_rate_func = lambda t: 2 * t model = StratifiedModel( times=[0, 1, 2, 3, 4, 5], compartment_names=["sus", "inf"], initial_conditions={"inf": 200}, parameters=parameters, requested_flows=requested_flows, starting_population=1000, infectious_compartments=["inf"], birth_approach=BirthApproach.NO_BIRTH, entry_compartment="sus", ) model.time_variants["recovery_rate"] = recovery_rate_func model.time_variants["home_contact_rate"] = home_contact_rate_func assert model.compartment_names == ["sus", "inf"] assert model.flows == requested_flows assert model.parameters == parameters assert model.time_variants == { "recovery_rate": recovery_rate_func, "home_contact_rate": home_contact_rate_func, } assert model.overwrite_parameters == [] assert model.adaptation_functions == {} vals = np.array([800, 200]) assert_array_equal(model.compartment_values, vals) model.stratify( "location", strata_request=["home", "work"], compartments_to_stratify=[], requested_proportions={"home": 0.6}, adjustment_requests=adjustment_requests, ) assert model.flows == [ { "origin": "susXlocation_home", "parameter": "contact_rateXlocation_home", "to": "infXlocation_home", "type": Flow.INFECTION_FREQUENCY, }, { "origin": "susXlocation_work", "parameter": "contact_rateXlocation_work", "to": "infXlocation_work", "type": Flow.INFECTION_FREQUENCY, }, { "origin": "infXlocation_home", "parameter": "recovery_rateXlocation_home", "to": "susXlocation_home", "type": Flow.STANDARD, }, { "origin": "infXlocation_work", "parameter": "recovery_rateXlocation_work", "to": "susXlocation_work", "type": Flow.STANDARD, }, { "origin": "infXlocation_home", "parameter": "infect_deathXlocation_home", "type": Flow.DEATH, }, { "origin": "infXlocation_work", "parameter": "infect_deathXlocation_work", "type": Flow.DEATH, }, ] assert model.time_variants == { "recovery_rate": recovery_rate_func, "home_contact_rate": home_contact_rate_func, } assert model.parameters["contact_rate"] == 1000 assert model.parameters["infect_death"] == 10 assert model.parameters["recovery_rate"] == "recovery_rate" assert model.parameters[ "contact_rateXlocation_home"] == "home_contact_rate" assert model.parameters["contact_rateXlocation_work"] == 0.5 assert model.parameters["recovery_rateXlocation_home"] == 1 assert model.parameters["recovery_rateXlocation_work"] == 2 assert model.parameters["infect_deathXlocation_home"] == 1 assert model.parameters["infect_deathXlocation_work"] == 2
def build_model(params: dict) -> StratifiedModel: """ Build the master function to run a tuberculosis model """ validate_params(params) # perform validation of parameter format check_param_values(params) # perform validation of some parameter values # Define model times. time = params["time"] integration_times = get_model_times_from_inputs(round(time["start"]), time["end"], time["step"], time["critical_ranges"]) # Define model compartments. compartments = [ Compartment.SUSCEPTIBLE, Compartment.EARLY_LATENT, Compartment.LATE_LATENT, Compartment.INFECTIOUS, Compartment.ON_TREATMENT, Compartment.RECOVERED, ] infectious_comps = [ Compartment.INFECTIOUS, Compartment.ON_TREATMENT, ] # Define initial conditions - 1 infectious person. init_conditions = { Compartment.INFECTIOUS: 1, } # prepare infectiousness adjustment for individuals on treatment treatment_infectiousness_adjustment = [{ "comp_name": Compartment.ON_TREATMENT, "comp_strata": {}, "value": params["on_treatment_infect_multiplier"], }] # Define inter-compartmental flows. flows = deepcopy(preprocess.flows.DEFAULT_FLOWS) # is ACF implemented? implement_acf = len(params["time_variant_acf"]) > 0 if implement_acf: flows.append(preprocess.flows.ACF_FLOW) # is ltbi screening implemented? implement_ltbi_screening = len(params["time_variant_ltbi_screening"]) > 0 if implement_ltbi_screening: flows += preprocess.flows.get_preventive_treatment_flows( params["pt_destination_compartment"]) # Set some parameter values or parameters that require pre-processing ( params, treatment_recovery_func, treatment_death_func, relapse_func, detection_rate_func, acf_detection_rate_func, preventive_treatment_func, contact_rate_functions, ) = preprocess.flows.process_unstratified_parameter_values( params, implement_acf, implement_ltbi_screening) # Create the model. tb_model = StratifiedModel( times=integration_times, compartment_names=compartments, initial_conditions=init_conditions, parameters=params, requested_flows=flows, infectious_compartments=infectious_comps, birth_approach=BirthApproach.ADD_CRUDE, entry_compartment=Compartment.SUSCEPTIBLE, starting_population=int(params["start_population_size"]), ) # register acf_detection_func if acf_detection_rate_func is not None: tb_model.time_variants["acf_detection_rate"] = acf_detection_rate_func # register preventive_treatment_func if preventive_treatment_func is not None: tb_model.time_variants[ "preventive_treatment_rate"] = preventive_treatment_func # register time-variant contact-rate functions: for param_name, func in contact_rate_functions.items(): tb_model.time_variants[param_name] = func # Apply infectiousness adjustment for individuals on treatment tb_model.individual_infectiousness_adjustments = treatment_infectiousness_adjustment # apply age stratification if "age" in params["stratify_by"]: stratify_by_age(tb_model, params, compartments) else: # set time-variant functions for treatment death and relapse rates tb_model.time_variants[ "treatment_recovery_rate"] = treatment_recovery_func tb_model.time_variants["treatment_death_rate"] = treatment_death_func tb_model.time_variants["relapse_rate"] = relapse_func # Load time-variant birth rates birth_rates, years = inputs.get_crude_birth_rate(params["iso3"]) birth_rates = [b / 1000.0 for b in birth_rates ] # birth rates are provided / 1000 population tb_model.time_variants["crude_birth_rate"] = scale_up_function( years, birth_rates, smoothness=0.2, method=5) tb_model.parameters["crude_birth_rate"] = "crude_birth_rate" # apply user-defined stratifications user_defined_stratifications = [ s for s in list(params["user_defined_stratifications"].keys()) if s in params["stratify_by"] ] for stratification in user_defined_stratifications: assert "_" not in stratification, "Stratification name should not include '_'" stratification_details = params["user_defined_stratifications"][ stratification] apply_user_defined_stratification( tb_model, compartments, stratification, stratification_details, implement_acf, implement_ltbi_screening, ) if "organ" in params["stratify_by"]: stratify_by_organ(tb_model, params) else: tb_model.time_variants["detection_rate"] = detection_rate_func # Register derived output functions, which are calculations based on the model's compartment values or flows. # These are calculated after the model is run. outputs.get_all_derived_output_functions(params["calculated_outputs"], params["outputs_stratification"], tb_model) return tb_model
def test_strains__with_two_symmetric_strains(): """ Adding two strains with the same properties should yield the same infection dynamics and outputs as having no strains at all. We expect the force of infection for each strain to be 1/2 of the unstratified model, but the stratification process will not apply the usual conservation fraction to the fan out flows. """ params = { "contact_rate": 0.2, "recovery_rate": 0.1, } flows = ( { "type": Flow.INFECTION_FREQUENCY, "parameter": "contact_rate", "origin": "S", "to": "I", }, { "type": Flow.STANDARD, "parameter": "recovery_rate", "origin": "I", "to": "R", }, ) # Create an unstratified model kwargs = {**MODEL_KWARGS, "parameters": params, "requested_flows": flows} model = StratifiedModel(**kwargs) # Do pre-run force of infection calcs. model.prepare_to_run() model.prepare_time_step(0, model.compartment_values) # Check infectiousness multipliers susceptible = model.compartment_names[0] infectious = model.compartment_names[1] assert model.get_infection_density_multiplier(susceptible, infectious) == 100.0 assert model.get_infection_frequency_multiplier(susceptible, infectious) == 0.1 model.run_model() # Create a stratified model where the two strains are symmetric strain_model = StratifiedModel(**kwargs) strain_model.stratify( stratification_name="strain", strata_request=["a", "b"], compartments_to_stratify=["I"], # Use defaults - equally split compartments, flows, etc. comp_split_props={}, flow_adjustments={}, infectiousness_adjustments={}, mixing_matrix=None, ) strain_model.run_model() merged_outputs = np.zeros_like(model.outputs) merged_outputs[:, 0] = strain_model.outputs[:, 0] merged_outputs[:, 1] = strain_model.outputs[:, 1] + strain_model.outputs[:, 2] merged_outputs[:, 2] = strain_model.outputs[:, 3] assert_allclose(merged_outputs, model.outputs, atol=0.01, rtol=0.01, verbose=True)
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. """ contact_rate = 0.2 params = { "contact_rate": contact_rate, } flows = ( { "type": Flow.INFECTION_FREQUENCY, "parameter": "contact_rate", "origin": "S", "to": "I", }, ) kwargs = {**MODEL_KWARGS, "parameters": params, "requested_flows": flows} model = StratifiedModel(**kwargs) model.stratify( stratification_name="agegroup", strata_request=["child", "adult"], compartments_to_stratify=["S", "I", "R"], comp_split_props={ "child": 0.6, # 600 people "adult": 0.4, # 400 people }, # Higher mixing among adults or children, # than between adults or children. mixing_matrix=np.array([[1.5, 0.5], [0.5, 1.5]]), ) model.stratify( stratification_name="strain", strata_request=["a", "b", "c"], compartments_to_stratify=["I"], comp_split_props={ "a": 0.7, # 70 people "b": 0.2, # 20 people "c": 0.1, # 10 people }, infectiousness_adjustments={ "a": 0.5, # 0.5x as infectious "b": 3, # 3x as infectious "c": 2, # 2x as infectious }, mixing_matrix=None, ) # 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.compartment_values) 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.compartment_names[0] sus_adult = model.compartment_names[1] inf_child_a = model.compartment_names[2] inf_child_b = model.compartment_names[3] inf_child_c = model.compartment_names[4] inf_adult_a = model.compartment_names[5] inf_adult_b = model.compartment_names[6] inf_adult_c = model.compartment_names[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_flow_rates(model.compartment_values, 0) assert_allclose(expected_flow_rates, flow_rates, verbose=True)
def test_strain__with_flow_adjustments(): """ Test infectious multiplier and flow rate calculations for 3 strains which have different flow adjustments. These flow adjustments would correspond to some physical process that we're modelling, and they should be effectively the same as applying infectiousness multipliers. """ contact_rate = 0.2 params = { "contact_rate": contact_rate, } flows = ( { "type": Flow.INFECTION_FREQUENCY, "parameter": "contact_rate", "origin": "S", "to": "I", }, ) kwargs = {**MODEL_KWARGS, "parameters": params, "requested_flows": flows} model = StratifiedModel(**kwargs) model.stratify( stratification_name="strain", strata_request=["a", "b", "c"], compartments_to_stratify=["I"], comp_split_props={ "a": 0.7, # 70 people "b": 0.2, # 20 people "c": 0.1, # 10 people }, flow_adjustments={ "contact_rate": { "a": 0.5, # 0.5x as infectious "b": 3, # 3x as infectious "c": 2, # 2x as infectious } }, mixing_matrix=None, ) # Do pre-run force of infection calcs. model.prepare_to_run() assert_array_equal(model.compartment_infectiousness["a"], np.array([0, 1, 0, 0, 0])) assert_array_equal(model.compartment_infectiousness["b"], np.array([0, 0, 1, 0, 0])) assert_array_equal(model.compartment_infectiousness["c"], np.array([0, 0, 0, 1, 0])) assert model.category_lookup == {0: 0, 1: 0, 2: 0, 3: 0, 4: 0} assert_array_equal(model.category_matrix, np.array([[1, 1, 1, 1, 1]])) # Do pre-iteration force of infection calcs model.prepare_time_step(0, model.compartment_values) assert_array_equal(model.category_populations, np.array([[1000]])) assert_array_equal(model.infection_density["a"], np.array([[70]])) assert_array_equal(model.infection_density["b"], np.array([[20]])) assert_array_equal(model.infection_density["c"], np.array([[10]])) assert_array_equal(model.infection_frequency["a"], np.array([[70 / 1000]])) assert_array_equal(model.infection_frequency["b"], np.array([[20 / 1000]])) assert_array_equal(model.infection_frequency["c"], np.array([[10 / 1000]])) # Get multipliers susceptible = model.compartment_names[0] infectious_a = model.compartment_names[1] infectious_b = model.compartment_names[2] infectious_c = model.compartment_names[3] assert model.get_infection_density_multiplier(susceptible, infectious_a) == 70 assert model.get_infection_density_multiplier(susceptible, infectious_b) == 20 assert model.get_infection_density_multiplier(susceptible, infectious_c) == 10 assert model.get_infection_frequency_multiplier(susceptible, infectious_a) == 70 / 1000 assert model.get_infection_frequency_multiplier(susceptible, infectious_b) == 20 / 1000 assert model.get_infection_frequency_multiplier(susceptible, infectious_c) == 10 / 1000 # Get infection flow rates flow_rates = model.get_flow_rates(model.compartment_values, 0) sus_pop = 900 flow_to_a = sus_pop * contact_rate * (70 * 0.5 / 1000) flow_to_b = sus_pop * contact_rate * (20 * 3 / 1000) flow_to_c = sus_pop * contact_rate * (10 * 2 / 1000) expected_flow_rates = np.array( [-flow_to_a - flow_to_b - flow_to_c, flow_to_a, flow_to_b, flow_to_c, 0.0] ) assert_allclose(expected_flow_rates, flow_rates, verbose=True)