def test_stratify_compartments(strata, proportions, to_stratify,
expected_names, expected_values):
"""
Ensure that `stratify_compartments` splits up compartment names and values correctly.
"""
model = StratifiedModel(**_get_model_kwargs())
model.stratify_compartments("test", strata, proportions, to_stratify)
assert model.compartment_names == expected_names
assert model.compartment_values == expected_values
def test_find_transition_indices_to_implement(back_one, include_change,
all_stratifications, flows,
expected_idxs):
"""
Ensure `find_transition_indices_to_implement` returns the correct list of indices.
"""
model = StratifiedModel(**_get_model_kwargs())
cols = [
"type", "parameter", "origin", "to", "implement", "strain",
"force_index"
]
model.transition_flows = pd.DataFrame(flows, columns=cols).astype(object)
model.all_stratifications = all_stratifications
actual_idxs = model.find_transition_indices_to_implement(
back_one, include_change)
assert expected_idxs == actual_idxs
def test_set_ageing_rates(age_strata, expected_flows, expected_ageing):
"""
Ensure that `set_ageing_rates` adds ageing flows to the transition flows dataframe
"""
model = StratifiedModel(**_get_model_kwargs())
cols = [
"type", "parameter", "origin", "to", "implement", "strain",
"force_index"
]
# Ensure there are no initial flows
initial_df = pd.DataFrame([], columns=cols).astype(object)
assert_frame_equal(initial_df, model.transition_flows)
# Set ageing rates
model.set_ageing_rates(age_strata)
# Check ageing flows are set
expected_df = pd.DataFrame(expected_flows, columns=cols).astype(object)
assert_frame_equal(expected_df, model.transition_flows)
# Check ageing params are set
for k, v in expected_ageing.items():
assert model.parameters[k] == v
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_types=[Compartment.SUSCEPTIBLE, Compartment.INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
)
# Run the model for 5 years.
model.run_model()
# Expect everyone to generally get older, but no one should die or get sick
expected_output = [
[250.0, 250.0, 250.0, 250.0, 0.0, 0.0, 0.0, 0.0],
[205.0, 269.0, 270.0, 256.0, 0.0, 0.0, 0.0, 0.0],
[168.0, 279.0, 291.0, 262.0, 0.0, 0.0, 0.0, 0.0],
[137.0, 281.0, 313.0, 269.0, 0.0, 0.0, 0.0, 0.0],
[112.0, 278.0, 334.0, 276.0, 0.0, 0.0, 0.0, 0.0],
[92.0, 271.0, 354.0, 284.0, 0.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_stratify_transition_flows(
flows,
custom_func,
params,
adjustment,
comps,
expected_flows,
expected_custom_func,
expected_params,
):
"""
Ensure that `stratify_compartments` splits up transition flows correctly.
"""
model = StratifiedModel(**_get_model_kwargs())
cols = [
"type", "parameter", "origin", "to", "implement", "strain",
"force_index"
]
model.transition_flows = pd.DataFrame(flows, columns=cols).astype(object)
model.parameters = params
strat_name = "test"
strata_names = ["foo", "bar"]
model.customised_flow_functions = custom_func
model.all_stratifications = {strat_name: strata_names}
model.stratify_compartments(strat_name, strata_names, {
"foo": 0.5,
"bar": 0.5
}, comps)
model.stratify_transition_flows(strat_name, strata_names, adjustment,
comps)
# Check flows df stratified
expected_flows_df = pd.DataFrame(expected_flows,
columns=cols).astype(object)
assert_frame_equal(expected_flows_df, model.transition_flows)
# Check custom flow func is updated
assert model.customised_flow_functions == expected_custom_func
# Check params are stratified
for k, v in expected_params.items():
assert model.parameters[k] == v
def test_strat_model__with_locations_and_mixing__expect_varied_transmission():
"""
Ensure that location-based mixing works.
Expect urbanites to be highly infectious t other locations, but not the reverse.
"""
pop = 1000
model = StratifiedModel(
times=_get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.INFECTIOUS],
initial_conditions={Compartment.INFECTIOUS: 100},
parameters={"contact_rate": 3},
requested_flows=[{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.INFECTIOUS,
}],
starting_population=pop,
)
# Add basic location stratification
model.stratify(
Stratification.LOCATION,
strata_request=["rural", "urban", "prison"],
compartment_types_to_stratify=[],
requested_proportions={},
mixing_matrix=np.array([
# Rural people catch disease from urbanites
[0.0, 1.0, 0.0],
# Urbanites cannot catch disease from anyone
[0.0, 0.0, 0.0],
# Prisoners catch disease from urbanites
[0.0, 1.0, 0.0],
]),
)
# Run the model for 5 years.
model.run_model()
# Expect everyone to generally get older, but no one should die or get sick
expected_output = [
[300.0, 300.0, 300.0, 33.0, 33.0, 33.0],
[271.0, 300.0, 271.0, 62.0, 33.0, 62.0],
[246.0, 300.0, 246.0, 88.0, 33.0, 88.0],
[222.0, 300.0, 222.0, 111.0, 33.0, 111.0],
[201.0, 300.0, 201.0, 132.0, 33.0, 132.0],
[182.0, 300.0, 182.0, 151.0, 33.0, 151.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
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_types=[Compartment.SUSCEPTIBLE, Compartment.INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={
"0": 0.8,
"5": 0.1,
"15": 0.1
},
)
# Run the model for 5 years.
model.run_model()
# 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_output = [
[800.0, 100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[655.0, 228.0, 114.0, 2.0, 0.0, 0.0, 0.0, 0.0],
[536.0, 319.0, 139.0, 5.0, 0.0, 0.0, 0.0, 0.0],
[439.0, 381.0, 171.0, 9.0, 0.0, 0.0, 0.0, 0.0],
[360.0, 421.0, 207.0, 13.0, 0.0, 0.0, 0.0, 0.0],
[294.0, 442.0, 245.0, 18.0, 0.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
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_types=[Compartment.SUSCEPTIBLE, Compartment.INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic location stratification
model.stratify(
Stratification.LOCATION,
strata_request=["rural", "urban", "prison"],
compartment_types_to_stratify=[],
requested_proportions={
"rural": 0.44,
"urban": 0.55,
"prison": 0.01
},
)
# Run the model for 5 years.
model.run_model()
# Expect everyone to start in their locations, then nothing should change,
expected_output = [
[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],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()