def test_entry_flow_stratify__with_no_flow_adjustments(): strat = Stratification( name="location", strata=["1", "2"], compartments=["infect", "recovery"], comp_split_props={}, flow_adjustments={}, ) flow = EntryFlow( dest=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert len(new_flows) == 2 assert new_flows[0].param_name == "foo" assert new_flows[0].param_func == _get_param_value assert new_flows[0].adjustments == [(FlowAdjustment.MULTIPLY, 0.5)] assert new_flows[0].dest == Compartment("infect", strat_names=["location"], strat_values={"location": "1"}) assert new_flows[1].param_name == "foo" assert new_flows[1].param_func == _get_param_value assert new_flows[1].adjustments == [(FlowAdjustment.MULTIPLY, 0.5)] assert new_flows[1].dest == Compartment("infect", strat_names=["location"], strat_values={"location": "2"})
def test_exit_flow_stratify__with_no_flow_adjustments(): strat = Stratification( name="age", strata=["1", "2"], compartments=["infect", "recovery"], comp_split_props={}, flow_adjustments={}, ) flow = ExitFlow( source=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert len(new_flows) == 2 assert new_flows[0].param_name == "foo" assert new_flows[0].param_func == _get_param_value assert new_flows[0].adjustments == [] assert new_flows[0].source == Compartment("infect", strat_names=["age"], strat_values={"age": "1"}) assert new_flows[1].param_name == "foo" assert new_flows[1].param_func == _get_param_value assert new_flows[1].adjustments == [] assert new_flows[1].source == Compartment("infect", strat_names=["age"], strat_values={"age": "2"})
def test_stratification_with_basic_setup(): """ Ensure we can create a simple stratification. """ strat = Stratification( name="foo", strata=[1, "2", "bar", "baz"], compartments=["sus", "inf", "rec"], comp_split_props={}, flow_adjustments={}, ) assert strat.name == "foo" assert strat.compartments == [ Compartment("sus"), Compartment("inf"), Compartment("rec") ] assert strat.strata == ["1", "2", "bar", "baz"] assert strat.comp_split_props == { "1": 0.25, "2": 0.25, "bar": 0.25, "baz": 0.25 } assert strat.flow_adjustments == {}
def test_get_infection_density_multiplier(): model = EpiModel(**MODEL_KWARGS) model.prepare_to_run() model.prepare_time_step(0, model.compartment_values) c_src = Compartment("S") c_dst = Compartment("I") multiplier = model.get_infection_density_multiplier(c_src, c_dst) assert model.population_infectious == 10 assert model.population_total == 1000 assert multiplier == 10
def test_serialize(): c = Compartment( "infected", strat_names=["location", "age"], strat_values={ "location": "hawaii", "age": "15" }, ) assert c.serialize() == "infectedXlocation_hawaiiXage_15"
def test_standard_flow_get_net_flow(): flow = StandardFlow( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, adjustments=[], ) flow.source.idx = 1 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 2 * 3 * 7
def test_entry_flow_stratify_with_ageing(): strat = Stratification( name="age", strata=["0", "1", "2"], compartments=["infect", "recovery"], comp_split_props={}, flow_adjustments={}, ) # Expect these to be ignored in favour of birth specific adjustments. strat.flow_adjustments = { "foo": [{ "strata": {}, "adjustments": { "1": (FlowAdjustment.MULTIPLY, 0.1) } }] } flow = EntryFlow( dest=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[(FlowAdjustment.OVERWRITE, 0.2)], ) new_flows = flow.stratify(strat) assert len(new_flows) == 3 assert new_flows[0].param_name == "foo" assert new_flows[0].param_func == _get_param_value assert new_flows[0].adjustments == [ (FlowAdjustment.OVERWRITE, 0.2), (FlowAdjustment.MULTIPLY, 1), ] assert new_flows[0].dest == Compartment("infect", strat_names=["age"], strat_values={"age": "0"}) assert new_flows[1].param_name == "foo" assert new_flows[1].param_func == _get_param_value assert new_flows[1].adjustments == [ (FlowAdjustment.OVERWRITE, 0.2), (FlowAdjustment.OVERWRITE, 0), ] assert new_flows[1].dest == Compartment("infect", strat_names=["age"], strat_values={"age": "1"}) assert new_flows[2].param_name == "foo" assert new_flows[2].param_func == _get_param_value assert new_flows[2].adjustments == [ (FlowAdjustment.OVERWRITE, 0.2), (FlowAdjustment.OVERWRITE, 0), ] assert new_flows[2].dest == Compartment("infect", strat_names=["age"], strat_values={"age": "2"})
def test_standard_flow_get_net_flow_with_adjust(): flow = StandardFlow( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, adjustments=[(FlowAdjustment.MULTIPLY, 13)], ) flow.source.idx = 2 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 2 * 5 * 7 * 13
def test_infect_density_get_net_flow_with_adjust(FlowClass): flow = FlowClass( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, find_infectious_multiplier=find_infectious_multiplier, adjustments=[(FlowAdjustment.MULTIPLY, 13)], ) flow.source.idx = 2 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 2 * 5 * 7 * 13 * 23
def test_infect_get_net_flow(FlowClass): flow = FlowClass( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, find_infectious_multiplier=find_infectious_multiplier, adjustments=[], ) flow.source.idx = 1 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 2 * 3 * 7 * 23
def test_stratify(): c = Compartment("infected") assert c._name == "infected" assert c._strat_names == tuple() assert c._strat_values == {} assert c.serialize() == "infected" c_age = c.stratify("age", "15") assert c_age._name == "infected" assert c_age._strat_names == ("age", ) assert c_age._strat_values == {"age": "15"} assert c_age.serialize() == "infectedXage_15" c_loc = c_age.stratify("location", "work") assert c_loc._name == "infected" assert c_loc._strat_names == ("age", "location") assert c_loc._strat_values == {"age": "15", "location": "work"} assert c_loc.serialize() == "infectedXage_15Xlocation_work"
def _add_flow(self, requested_flow): flow = None f_source = requested_flow.get("origin") f_source = Compartment.deserialize(f_source) if f_source else None f_dest = requested_flow.get("to") f_dest = Compartment.deserialize(f_dest) if f_dest else None f_param = requested_flow["parameter"] f_type = requested_flow["type"] if f_type == Flow.STANDARD: flow = StandardFlow( source=f_source, dest=f_dest, param_name=f_param, param_func=self.get_parameter_value, ) elif f_type == Flow.INFECTION_FREQUENCY: flow = InfectionFrequencyFlow( source=f_source, dest=f_dest, param_name=f_param, param_func=self.get_parameter_value, find_infectious_multiplier=self.get_infection_frequency_multiplier, ) elif f_type == Flow.INFECTION_DENSITY: flow = InfectionDensityFlow( source=f_source, dest=f_dest, param_name=f_param, param_func=self.get_parameter_value, find_infectious_multiplier=self.get_infection_density_multiplier, ) elif f_type == Flow.DEATH: flow = InfectionDeathFlow( source=f_source, param_name=f_param, param_func=self.get_parameter_value ) elif f_type == Flow.IMPORT: flow = ImportFlow( dest=self.entry_compartment, param_name=f_param, param_func=self.get_parameter_value, ) if flow: self.flows.append(flow) return flow
def test_transition_flow_stratify_with_no_matching_compartments(): strat = Stratification( name="age", strata=["1", "2", "3"], compartments=["recovery"], # Flow compartments not included comp_split_props={}, flow_adjustments={}, ) flow = TransitionFlow( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert new_flows == [flow]
def test_get_stratified_compartment_values__with_extisting_strat(): """ Stratify compartments for the second time, expect that compartments are are split according to proportions and old compartments are removed. """ comp_values = np.array( [250.0, 500.0, 250.0, 25.0, 50.0, 25.0, 0.0, 0.0, 0.0]) comp_names = [ Compartment("S", strat_names=["age"], strat_values={"age": "0"}), Compartment("S", strat_names=["age"], strat_values={"age": "10"}), Compartment("S", strat_names=["age"], strat_values={"age": "20"}), Compartment("I", strat_names=["age"], strat_values={"age": "0"}), Compartment("I", strat_names=["age"], strat_values={"age": "10"}), Compartment("I", strat_names=["age"], strat_values={"age": "20"}), Compartment("R", strat_names=["age"], strat_values={"age": "0"}), Compartment("R", strat_names=["age"], strat_values={"age": "10"}), Compartment("R", strat_names=["age"], strat_values={"age": "20"}), ] strat = Stratification( name="location", strata=["rural", "urban"], compartments=["S", "I", "R"], comp_split_props={ "rural": 0.1, "urban": 0.9 }, flow_adjustments={}, ) new_comp_values = get_stratified_compartment_values( strat, comp_names, comp_values) expected_arr = np.array([ 25, 225.0, 50.0, 450.0, 25.0, 225.0, 2.5, 22.5, 5.0, 45.0, 2.5, 22.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ]) assert_array_equal(expected_arr, new_comp_values)
def test_replace_deaths_birth_flow_get_net_flow(): flow = ReplacementBirthFlow( dest=Compartment("susceptible"), get_total_deaths=_get_total_deaths, adjustments=[], ) vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 23
def test_get_stratified_compartment_values__with_subset_stratified(): strat = Stratification( name="age", strata=["0", "10", "20"], compartments=["S"], comp_split_props={ "0": 0.25, "10": 0.5, "20": 0.25 }, flow_adjustments={}, ) comp_names = [Compartment("S"), Compartment("I"), Compartment("R")] comp_values = np.array([1000.0, 100.0, 0.0]) new_comp_values = get_stratified_compartment_values( strat, comp_names, comp_values) expected_arr = np.array([250.0, 500.0, 250.0, 100.0, 0.0]) assert_array_equal(expected_arr, new_comp_values)
def test_replace_deaths_birth_flow_get_net_flow_with_adjust(): flow = ReplacementBirthFlow( dest=Compartment("susceptible"), get_total_deaths=_get_total_deaths, adjustments=[(FlowAdjustment.MULTIPLY, 13)], ) flow.dest.idx = 2 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 23 * 13
def test_crude_birth_flow_get_net_flow(): flow = CrudeBirthFlow( dest=Compartment("susceptible"), param_name="crude_birth_rate", param_func=_get_param_value, adjustments=[], ) vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 0.1 * 7 * (1 + 3 + 5)
def test_setup__with_double_strat(): c = Compartment( "infected", strat_names=["location", "age"], strat_values={ "location": "hawaii", "age": "15" }, ) assert str(c) == "infectedXlocation_hawaiiXage_15" assert hash(c) == hash("infectedXlocation_hawaiiXage_15")
def test_death_flow_get_net_flow(FlowClass): flow = FlowClass( source=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[], ) flow.source.idx = 1 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 2 * 3 * 7
def test_death_flow_get_net_flow_with_adjust(FlowClass): flow = FlowClass( source=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[(FlowAdjustment.MULTIPLY, 13)], ) flow.source.idx = 2 vals = np.array([1, 3, 5]) net_flow = flow.get_net_flow(vals, 7) assert net_flow == 2 * 5 * 7 * 13
def test_deserialzie(): s = "infectedXlocation_hawaiiXage_15" c = Compartment.deserialize(s) assert c == Compartment( "infected", strat_names=["location", "age"], strat_values={ "location": "hawaii", "age": "15" }, ) assert c != Compartment( "infected", strat_names=["age", "location"], strat_values={ "location": "hawaii", "age": "15" }, ) assert str(c) == "infectedXlocation_hawaiiXage_15" assert hash(c) == hash("infectedXlocation_hawaiiXage_15")
def test_exit_flow_stratify_with_flow_adjustments(): strat = Stratification( name="age", strata=["1", "2"], compartments=["infect", "recovery"], comp_split_props={}, flow_adjustments={}, ) strat.flow_adjustments = { "foo": [{ "strata": {}, "adjustments": { "1": (FlowAdjustment.MULTIPLY, 0.1) } }] } flow = ExitFlow( source=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[(FlowAdjustment.OVERWRITE, 0.2)], ) new_flows = flow.stratify(strat) assert len(new_flows) == 2 assert new_flows[0].param_name == "foo" assert new_flows[0].param_func == _get_param_value assert new_flows[0].adjustments == [ (FlowAdjustment.OVERWRITE, 0.2), (FlowAdjustment.MULTIPLY, 0.1), ] assert new_flows[0].source == Compartment("infect", strat_names=["age"], strat_values={"age": "1"}) assert new_flows[1].param_name == "foo" assert new_flows[1].param_func == _get_param_value assert new_flows[1].adjustments == [(FlowAdjustment.OVERWRITE, 0.2)] assert new_flows[1].source == Compartment("infect", strat_names=["age"], strat_values={"age": "2"})
def test_entry_flow_stratify__when_not_applicable(): strat = Stratification( name="location", strata=["1", "2", "3"], compartments=["recovery"], # infect compartment not included comp_split_props={}, flow_adjustments={}, ) flow = EntryFlow( dest=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert new_flows == [flow]
def test_exit_flow_stratify__when_no_compartment_match(): strat = Stratification( name="age", strata=["1", "2", "3"], compartments=["recovery"], # infect compartment not included comp_split_props={}, flow_adjustments={}, ) flow = ExitFlow( source=Compartment("infect"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert new_flows == [flow]
def __init__( self, name: str, strata: List[str], compartments: List[str], comp_split_props: Dict[str, float], flow_adjustments: Dict[str, Dict[str, float]], ): self.name = name self.strata = list(map(str, strata)) self.comp_split_props = get_all_proportions(self.strata, comp_split_props) self.flow_adjustments = parse_flow_adjustments(flow_adjustments) self.compartments = [ Compartment(c) if type(c) is str else c for c in compartments ]
def test_get_stratified_compartment_names__with_no_extisting_strat_and_subset_only( ): strat = Stratification( name="age", strata=["0", "10", "20"], compartments=["S"], comp_split_props={}, flow_adjustments={}, ) comp_names = [Compartment("S"), Compartment("I"), Compartment("R")] strat_comp_names = get_stratified_compartment_names(strat, comp_names) assert strat_comp_names == [ Compartment("S", strat_names=["age"], strat_values={"age": "0"}), Compartment("S", strat_names=["age"], strat_values={"age": "10"}), Compartment("S", strat_names=["age"], strat_values={"age": "20"}), Compartment("I"), Compartment("R"), ]
def __init__( self, times: List[float], compartment_names: List[str], initial_conditions: Dict[str, float], parameters: Dict[str, float], requested_flows: List[dict], infectious_compartments: List[str], birth_approach: str, entry_compartment: str, starting_population: int, ): super().__init__( times=times, compartment_names=compartment_names, initial_conditions=initial_conditions, parameters=parameters, requested_flows=requested_flows, infectious_compartments=infectious_compartments, birth_approach=birth_approach, entry_compartment=entry_compartment, starting_population=starting_population, ) # Keeps track of Stratifications that have been applied. self.stratifications = [] # Keeps track of original, pre-stratified compartment names. self.original_compartment_names = [Compartment.deserialize(n) for n in compartment_names] # A list of the different sub-categories ('strains') of the diease that we are modelling. self.disease_strains = [self.DEFAULT_DISEASE_STRAIN] # Strata-based multipliers for compartmental infectiousness levels. self.infectiousness_levels = {} # Compartment-based overwrite values for compartmental infectiousness levels. self.individual_infectiousness_adjustments = [] # Mixing matrices a list of NxN arrays used to calculate force of infection. self.mixing_matrices = [] # A list of dicts that has the strata required to match a row in the mixing matrix. self.mixing_categories = [{}]
def test_transition_flow_stratify_dest_but_not_source__with_flow_adjustments(): """ Ensure flow is adjusted to account for fan out. """ strat = Stratification( name="age", strata=["1", "2"], compartments=["happy", "recovery"], comp_split_props={}, flow_adjustments={}, ) strat.flow_adjustments = { "foo": [{ "strata": {}, "adjustments": { "1": (FlowAdjustment.MULTIPLY, 0.1) } }] } flow = TransitionFlow( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert len(new_flows) == 2 assert new_flows[0].param_name == "foo" assert new_flows[0].param_func == _get_param_value assert new_flows[0].adjustments == [(FlowAdjustment.MULTIPLY, 0.1)] assert new_flows[0].source == Compartment("infect", strat_names=[], strat_values={}) assert new_flows[0].dest == Compartment("happy", strat_names=["age"], strat_values={"age": "1"}) assert new_flows[1].param_name == "foo" assert new_flows[1].param_func == _get_param_value assert new_flows[1].adjustments == [(FlowAdjustment.MULTIPLY, 0.5)] assert new_flows[1].source == Compartment("infect", strat_names=[], strat_values={}) assert new_flows[1].dest == Compartment("happy", strat_names=["age"], strat_values={"age": "2"})
def test_transition_flow_stratify_source_and_dest(): strat = Stratification( name="age", strata=["1", "2"], compartments=["infect", "happy", "recovery"], comp_split_props={}, flow_adjustments={}, ) flow = TransitionFlow( source=Compartment("infect"), dest=Compartment("happy"), param_name="foo", param_func=_get_param_value, adjustments=[], ) new_flows = flow.stratify(strat) assert len(new_flows) == 2 assert new_flows[0].param_name == "foo" assert new_flows[0].param_func == _get_param_value assert new_flows[0].adjustments == [] assert new_flows[0].source == Compartment("infect", strat_names=["age"], strat_values={"age": "1"}) assert new_flows[0].dest == Compartment("happy", strat_names=["age"], strat_values={"age": "1"}) assert new_flows[1].param_name == "foo" assert new_flows[1].param_func == _get_param_value assert new_flows[1].adjustments == [] assert new_flows[1].source == Compartment("infect", strat_names=["age"], strat_values={"age": "2"}) assert new_flows[1].dest == Compartment("happy", strat_names=["age"], strat_values={"age": "2"})