def test_stochastic_entry_flows(pop, rtol, birthrate):
"""
Check that entry flow produces outputs that tend towards mean as pop increases.
"""
model = CompartmentalModel(
times=[0, 10], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
s_pop = 0.99 * pop
i_pop = 0.01 * pop
model.set_initial_population(distribution={"S": s_pop, "I": i_pop})
model.add_crude_birth_flow("births", birthrate, "S")
model.run_stochastic(RANDOM_SEED)
# No change to infected or recovered compartments
assert_array_equal(model.outputs[:, 1], i_pop)
assert_array_equal(model.outputs[:, 2], 0)
# Calculate births using mean birth rate
mean_s = np.zeros_like(model.times)
mean_s[0] = s_pop
for i in range(1, len(model.times)):
mean_s[i] = mean_s[i - 1] + birthrate * (i_pop + mean_s[i - 1])
# All S compartment sizes are are within the error range
# of the mean of the poisson dist that determines entry.
assert_allclose(model.outputs[:, 0], mean_s, rtol=rtol)
def test_stratify__age__validate_ageing_flows_added():
"""
Ensure, when using an age stratification, that ageing flows are automatically added
and that birth flows are all sent to age 0.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert len(model._flows) == 0
model.add_crude_birth_flow("births", 0.02, "S")
assert len(model._flows) == 1
strat = AgeStratification("age", ["0", "5", "15"], ["S", "I", "R"])
model.stratify_with(strat)
# Expect ageing flows amongst age group and a birth flow that only goes to age 0.
expected_flows = [
CrudeBirthFlow("births", C("S", {"age": "0"}), 0.02),
SojournFlow("ageing_SXage_0_to_SXage_5", C("S", {"age": "0"}), C("S", {"age": "5"}), 5),
SojournFlow("ageing_IXage_0_to_IXage_5", C("I", {"age": "0"}), C("I", {"age": "5"}), 5),
SojournFlow("ageing_RXage_0_to_RXage_5", C("R", {"age": "0"}), C("R", {"age": "5"}), 5),
SojournFlow("ageing_SXage_5_to_SXage_15", C("S", {"age": "5"}), C("S", {"age": "15"}), 10),
SojournFlow("ageing_IXage_5_to_IXage_15", C("I", {"age": "5"}), C("I", {"age": "15"}), 10),
SojournFlow("ageing_RXage_5_to_RXage_15", C("R", {"age": "5"}), C("R", {"age": "15"}), 10),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def _build_model(self, params):
model = CompartmentalModel(
times=[params["time"]["start"], 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"],
)
model.set_initial_population(distribution={"S": 1000, "I": 1000})
model.add_crude_birth_flow("birth", params["birth_rate"], "S")
model.add_fractional_flow("recovery", params["recovery_rate"], "I",
"R")
return model
def test_model__with_higher_birth_than_and_death_rate__expect_pop_increase():
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 100, "I": 100})
# Add some babies at ~10 babies / 100 / year.
model.add_crude_birth_flow("births", 0.1, "S")
# Add some dying at ~2 people / 100 / year.
model.add_universal_death_flows("deaths", 0.02)
model.run()
expected_outputs = np.array(
[
[100.0, 100.0, 0], # Initial conditions
[118.6, 98.0, 0], # Tweaked ~0.1
[138.6, 96.1, 0], # Tweaked ~0.4
[160.1, 94.2, 0], # Tweaked ~0.9
[183.1, 92.3, 0], # Tweaked ~1.7
[207.9, 90.5, 0], # Tweaked ~2.7
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
def test_model__with_birth_and_death_rate__expect_pop_static_overall():
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 100, "I": 100})
# Add some babies at ~2 babies / 100 / year.
model.add_crude_birth_flow("births", 0.02, "S")
# Add some dying at ~2 people / 100 / year.
model.add_universal_death_flows("deaths", 0.02)
model.run()
expected_outputs = np.array(
[
[100.0, 100.0, 0], # Initial conditions
[102.0, 98.0, 0],
[104.0, 96.0, 0],
[105.8, 94.2, 0], # Tweaked.
[107.7, 92.3, 0], # Tweaked.
[109.5, 90.5, 0], # Tweaked.
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
def test_model__with_birth_rate__expect_pop_increase():
"""
Ensure that a model with two compartments and only birth rate dynamics results in more people.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 100, "I": 100})
# Add some babies at ~2 babies / 100 / year.
model.add_crude_birth_flow("births", 0.02, "S")
model.run()
# Expect that we have more people in the population per year
expected_outputs = np.array(
[
[100.0, 100, 0], # Initial conditions
[104.0, 100, 0],
[108.2, 100, 0],
[112.4, 100, 0],
[116.7, 100, 0],
[121.0, 100, 0],
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
def test_model__with_complex_dynamics__expect_correct_outputs():
"""
Ensure that a model with the "full suite" of TB dynamics produces correct results:
- 5 compartments
- birth rate + universal death rate
- standard inter-compartment flows
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "EL", "LL", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 900, "I": 100})
model.add_crude_birth_flow("births", 0.02, "S")
model.add_universal_death_flows("universal_deaths", 0.02)
model.add_infection_frequency_flow("infection", 14, "S", "EL")
model.add_infection_frequency_flow("reinfection", 14, "R", "EL")
model.add_infection_frequency_flow("regression", 3, "LL", "EL")
model.add_transition_flow("early_progression", 2, "EL", "I")
model.add_transition_flow("stabilisation", 3, "EL", "LL")
model.add_transition_flow("early_progression", 1, "LL", "I")
model.add_death_flow("infect_death", 0.4, "I")
model.add_transition_flow("recovery", 0.2, "I", "R")
model.add_transition_flow("case_detection", 1, "I", "R")
model.run()
# Expect that the results are consistent, nothing crazy happens.
# These results were not independently calculated, so this is more of an "acceptance test".
expected_outputs = np.array(
[
[900.0, 0.0, 0.0, 100.0, 0.0],
[66.1, 203.8, 274.2, 307.2, 75.3],
[2.9, 150.9, 220.5, 345.3, 69.4],
[2.2, 127.3, 175.6, 297.0, 58.1],
[1.8, 106.4, 145.6, 248.8, 48.5],
[1.5, 88.8, 121.5, 207.8, 40.5],
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.2, verbose=True)
def build_model(params: dict) -> CompartmentalModel:
time = params["time"]
model = CompartmentalModel(
times=[time["start"], time["end"]],
compartments=COMPARTMENTS,
infectious_compartments=INFECTIOUS_COMPS,
timestep=time["step"],
)
# Add initial population
init_pop = {
Compartment.EARLY_LATENT: params["initial_early_latent_population"],
Compartment.LATE_LATENT: params["initial_late_latent_population"],
Compartment.INFECTIOUS: params["initial_infectious_population"],
Compartment.DETECTED: params["initial_detected_population"],
Compartment.ON_TREATMENT: params["initial_on_treatment_population"],
Compartment.RECOVERED: 0,
}
sum_init_pop = sum(init_pop.values())
init_pop[Compartment.
SUSCEPTIBLE] = params["start_population_size"] - sum_init_pop
model.set_initial_population(init_pop)
# Add inter-compartmental flows
params = _get_derived_params(params)
# Entry flows
model.add_crude_birth_flow(
"birth",
params["crude_birth_rate"],
Compartment.SUSCEPTIBLE,
)
# Infection flows.
model.add_infection_frequency_flow(
"infection",
params["contact_rate"],
Compartment.SUSCEPTIBLE,
Compartment.EARLY_LATENT,
)
model.add_infection_frequency_flow(
"infection_from_latent",
params["contact_rate_from_latent"],
Compartment.LATE_LATENT,
Compartment.EARLY_LATENT,
)
model.add_infection_frequency_flow(
"infection_from_recovered",
params["contact_rate_from_recovered"],
Compartment.RECOVERED,
Compartment.EARLY_LATENT,
)
# Transition flows.
model.add_fractional_flow(
"treatment_early",
params["preventive_treatment_rate"],
Compartment.EARLY_LATENT,
Compartment.RECOVERED,
)
model.add_fractional_flow(
"treatment_late",
params["preventive_treatment_rate"],
Compartment.LATE_LATENT,
Compartment.RECOVERED,
)
model.add_fractional_flow(
"stabilisation",
params["stabilisation_rate"],
Compartment.EARLY_LATENT,
Compartment.LATE_LATENT,
)
model.add_fractional_flow(
"early_activation",
params["early_activation_rate"],
Compartment.EARLY_LATENT,
Compartment.INFECTIOUS,
)
model.add_fractional_flow(
"late_activation",
params["late_activation_rate"],
Compartment.LATE_LATENT,
Compartment.INFECTIOUS,
)
# Post-active-disease flows
model.add_fractional_flow(
"detection",
params["detection_rate"],
Compartment.INFECTIOUS,
Compartment.DETECTED,
)
model.add_fractional_flow(
"treatment_commencement",
params["treatment_commencement_rate"],
Compartment.DETECTED,
Compartment.ON_TREATMENT,
)
model.add_fractional_flow(
"missed_to_active",
params["missed_to_active_rate"],
Compartment.DETECTED,
Compartment.INFECTIOUS,
)
model.add_fractional_flow(
"self_recovery_infectious",
params["self_recovery_rate"],
Compartment.INFECTIOUS,
Compartment.LATE_LATENT,
)
model.add_fractional_flow(
"self_recovery_detected",
params["self_recovery_rate"],
Compartment.DETECTED,
Compartment.LATE_LATENT,
)
model.add_fractional_flow(
"treatment_recovery",
params["treatment_recovery_rate"],
Compartment.ON_TREATMENT,
Compartment.RECOVERED,
)
model.add_fractional_flow(
"treatment_default",
params["treatment_default_rate"],
Compartment.ON_TREATMENT,
Compartment.INFECTIOUS,
)
model.add_fractional_flow(
"failure_retreatment",
params["failure_retreatment_rate"],
Compartment.ON_TREATMENT,
Compartment.DETECTED,
)
model.add_fractional_flow(
"spontaneous_recovery",
params["spontaneous_recovery_rate"],
Compartment.ON_TREATMENT,
Compartment.LATE_LATENT,
)
# Death flows
# Universal death rate to be overriden by a multiply in age stratification.
uni_death_flow_names = model.add_universal_death_flows("universal_death",
death_rate=1)
model.add_death_flow(
"infectious_death",
params["infect_death_rate"],
Compartment.INFECTIOUS,
)
model.add_death_flow(
"detected_death",
params["infect_death_rate"],
Compartment.DETECTED,
)
model.add_death_flow(
"treatment_death",
params["treatment_death_rate"],
Compartment.ON_TREATMENT,
)
# Apply age-stratification
age_strat = _build_age_strat(params, uni_death_flow_names)
model.stratify_with(age_strat)
# Add vaccination stratification.
vac_strat = _build_vac_strat(params)
model.stratify_with(vac_strat)
# Apply organ stratification
organ_strat = _build_organ_strat(params)
model.stratify_with(organ_strat)
# Apply strain stratification
strain_strat = _build_strain_strat(params)
model.stratify_with(strain_strat)
# Add amplification flow
model.add_fractional_flow(
name="amplification",
fractional_rate=params["amplification_rate"],
source=Compartment.ON_TREATMENT,
dest=Compartment.ON_TREATMENT,
source_strata={"strain": "ds"},
dest_strata={"strain": "mdr"},
expected_flow_count=9,
)
# Add cross-strain reinfection flows
model.add_infection_frequency_flow(
name="reinfection_ds_to_mdr",
contact_rate=params["reinfection_rate"],
source=Compartment.EARLY_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "ds"},
dest_strata={"strain": "mdr"},
expected_flow_count=3,
)
model.add_infection_frequency_flow(
name="reinfection_mdr_to_ds",
contact_rate=params["reinfection_rate"],
source=Compartment.EARLY_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "mdr"},
dest_strata={"strain": "ds"},
expected_flow_count=3,
)
model.add_infection_frequency_flow(
name="reinfection_late_ds_to_mdr",
contact_rate=params["reinfection_rate"],
source=Compartment.LATE_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "ds"},
dest_strata={"strain": "mdr"},
expected_flow_count=3,
)
model.add_infection_frequency_flow(
name="reinfection_late_mdr_to_ds",
contact_rate=params["reinfection_rate"],
source=Compartment.LATE_LATENT,
dest=Compartment.EARLY_LATENT,
source_strata={"strain": "mdr"},
dest_strata={"strain": "ds"},
expected_flow_count=3,
)
# Apply classification stratification
class_strat = _build_class_strat(params)
model.stratify_with(class_strat)
# Apply retention stratification
retention_strat = _build_retention_strat(params)
model.stratify_with(retention_strat)
# Register derived output functions, which are calculations based on the model's compartment values or flows.
# These are calculated after the model is run.
model.request_output_for_flow("notifications", flow_name="detection")
model.request_output_for_flow("early_activation",
flow_name="early_activation")
model.request_output_for_flow("late_activation",
flow_name="late_activation")
model.request_output_for_flow("infectious_deaths",
flow_name="infectious_death")
model.request_output_for_flow("detected_deaths",
flow_name="detected_death")
model.request_output_for_flow("treatment_deaths",
flow_name="treatment_death")
model.request_output_for_flow("progression_early",
flow_name="early_activation")
model.request_output_for_flow("progression_late",
flow_name="late_activation")
model.request_aggregate_output("progression",
["progression_early", "progression_late"])
model.request_output_for_compartments("population_size", COMPARTMENTS)
model.request_aggregate_output(
"_incidence",
sources=["early_activation", "late_activation"],
save_results=False)
model.request_function_output("incidence",
sources=["_incidence", "population_size"],
func=lambda i, p: 1e5 * i / p)
model.request_aggregate_output(
"disease_deaths",
sources=["infectious_deaths", "detected_deaths", "treatment_deaths"])
cum_start_time = params["cumulative_output_start_time"]
model.request_cumulative_output("cumulative_diseased",
source="_incidence",
start_time=cum_start_time)
model.request_cumulative_output("cumulative_deaths",
source="disease_deaths",
start_time=cum_start_time)
model.request_output_for_compartments("_count_infectious",
INFECTIOUS_COMPS,
save_results=False)
model.request_function_output(
"prevalence_infectious",
sources=["_count_infectious", "population_size"],
func=lambda c, p: 1e5 * c / p,
)
model.request_output_for_compartments(
"_count_latent", [Compartment.EARLY_LATENT, Compartment.LATE_LATENT],
save_results=False)
model.request_function_output(
"percentage_latent",
sources=["_count_latent", "population_size"],
func=lambda c, p: 100 * c / p,
)
return model
def test_stratify__age__validate_ageing_flows_added_second():
"""
Ensure that age stratification works when applied after a previous stratification.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert len(model._flows) == 0
model.add_crude_birth_flow("births", 0.02, "S")
assert len(model._flows) == 1
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
model.stratify_with(strat)
strat = AgeStratification("age", ["0", "5", "15"], ["S", "I", "R"])
model.stratify_with(strat)
# Expect ageing flows amongst age group and a birth flow that only goes to age 0.
expected_flows = [
CrudeBirthFlow("births", C("S", {"location": "urban", "age": "0"}), 0.02, [Multiply(0.5)]),
CrudeBirthFlow("births", C("S", {"location": "rural", "age": "0"}), 0.02, [Multiply(0.5)]),
SojournFlow(
"ageing_SXlocation_urbanXage_0_to_SXlocation_urbanXage_5",
C("S", {"location": "urban", "age": "0"}),
C("S", {"location": "urban", "age": "5"}),
5,
),
SojournFlow(
"ageing_SXlocation_ruralXage_0_to_SXlocation_ruralXage_5",
C("S", {"location": "rural", "age": "0"}),
C("S", {"location": "rural", "age": "5"}),
5,
),
SojournFlow(
"ageing_IXlocation_urbanXage_0_to_IXlocation_urbanXage_5",
C("I", {"location": "urban", "age": "0"}),
C("I", {"location": "urban", "age": "5"}),
5,
),
SojournFlow(
"ageing_IXlocation_ruralXage_0_to_IXlocation_ruralXage_5",
C("I", {"location": "rural", "age": "0"}),
C("I", {"location": "rural", "age": "5"}),
5,
),
SojournFlow(
"ageing_RXlocation_urbanXage_0_to_RXlocation_urbanXage_5",
C("R", {"location": "urban", "age": "0"}),
C("R", {"location": "urban", "age": "5"}),
5,
),
SojournFlow(
"ageing_RXlocation_ruralXage_0_to_RXlocation_ruralXage_5",
C("R", {"location": "rural", "age": "0"}),
C("R", {"location": "rural", "age": "5"}),
5,
),
SojournFlow(
"ageing_SXlocation_urbanXage_5_to_SXlocation_urbanXage_15",
C("S", {"location": "urban", "age": "5"}),
C("S", {"location": "urban", "age": "15"}),
10,
),
SojournFlow(
"ageing_SXlocation_ruralXage_5_to_SXlocation_ruralXage_15",
C("S", {"location": "rural", "age": "5"}),
C("S", {"location": "rural", "age": "15"}),
10,
),
SojournFlow(
"ageing_IXlocation_urbanXage_5_to_IXlocation_urbanXage_15",
C("I", {"location": "urban", "age": "5"}),
C("I", {"location": "urban", "age": "15"}),
10,
),
SojournFlow(
"ageing_IXlocation_ruralXage_5_to_IXlocation_ruralXage_15",
C("I", {"location": "rural", "age": "5"}),
C("I", {"location": "rural", "age": "15"}),
10,
),
SojournFlow(
"ageing_RXlocation_urbanXage_5_to_RXlocation_urbanXage_15",
C("R", {"location": "urban", "age": "5"}),
C("R", {"location": "urban", "age": "15"}),
10,
),
SojournFlow(
"ageing_RXlocation_ruralXage_5_to_RXlocation_ruralXage_15",
C("R", {"location": "rural", "age": "5"}),
C("R", {"location": "rural", "age": "15"}),
10,
),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_solve_stochastic(monkeypatch):
"""
Test that _solve_stochastic glue code works.
Don't test the actual flow rate calculations or stochastic sampling bits.
"""
model = CompartmentalModel(
times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"],
)
# Add some people to the model, expect initial conditions of [990, 10, 0]
model.set_initial_population(distribution={"S": 990, "I": 10})
# Add flows - the parameters add here will be overidden by `mock_get_rates`
# but the flow directions will be used.
model.add_crude_birth_flow("birth", 8, "S")
model.add_infection_frequency_flow("infection", 6, "S", "I")
model.add_death_flow("infect_death", 3, "I")
model.add_transition_flow("recovery", 2, "I", "R")
# Mock out flow rate calculation - tested elsewhere and tricky to predict.
def mock_get_rates(comp_vals, time):
# Return the flow rates that will be used to solve the model
return None, np.array([float(f.param) for f in model._flows])
monkeypatch.setattr(model, "_get_rates", mock_get_rates)
# Mock out stochastic flow sampling - tested elsewhere.
def mock_sample_entry_flows(seed, entry_flow_rates, timestep):
assert not seed
assert 0 < timestep <= 5
expected_flow_rates = np.array([8, 0, 0])
assert_array_equal(entry_flow_rates, expected_flow_rates)
return np.array([8, 0, 0])
def mock_sample_transistion_flows(seed, flow_rates, flow_map, comp_vals, timestep):
assert not seed
assert 0 < timestep <= 5
# Flows get re-arranged by setup process
expected_flow_map = np.array([[0, 1, -1], [2, 0, 1], [3, 1, 2]])
assert_array_equal(flow_map, expected_flow_map)
expected_flow_rates = np.array(
[[0.0, 3.0, 0.0], [0.0, 0.0, 0.0], [6.0, 0.0, 0.0], [0.0, 2.0, 0.0]]
)
assert_array_equal(flow_rates, expected_flow_rates)
return np.array([-6, 1, 2])
monkeypatch.setattr(stochastic, "sample_entry_flows", mock_sample_entry_flows)
monkeypatch.setattr(stochastic, "sample_transistion_flows", mock_sample_transistion_flows)
model.run_stochastic()
expected_outputs = np.array(
[
[990, 10, 0],
[992, 11, 2],
[994, 12, 4],
[996, 13, 6],
[998, 14, 8],
[1000, 15, 10],
]
)
assert_array_equal(model.outputs, expected_outputs)