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_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)