def test_stochastic_exit_flows(pop, rtol, deathrate):
"""
Check that death flows produce outputs that tend towards mean as pop increases.
"""
model = CompartmentalModel(
times=[0, 10], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
s_pop = 0.80 * pop
i_pop = 0.20 * pop
model.set_initial_population(distribution={"S": s_pop, "I": i_pop})
model.add_universal_death_flows("deaths", deathrate)
model.run_stochastic(RANDOM_SEED)
# No change to recovered compartments
assert_array_equal(model.outputs[:, 2], 0)
# Calculate births using mean birth rate
mean_s = np.zeros_like(model.times)
mean_i = np.zeros_like(model.times)
mean_s[0] = s_pop
mean_i[0] = i_pop
for i in range(1, len(model.times)):
mean_s[i] = mean_s[i - 1] - deathrate * mean_s[i - 1]
mean_i[i] = mean_i[i - 1] - deathrate * mean_i[i - 1]
# All S and I compartment sizes are are within the error range
# of the mean of the multinomial dist that determines exit.
assert_allclose(model.outputs[:, 0], mean_s, rtol=rtol)
assert_allclose(model.outputs[:, 1], mean_i, rtol=rtol)
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.
"""
model = CompartmentalModel(times=[0, 5], compartments=["S", "I"], infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 1000, "I": 0})
strat = AgeStratification("age", [0, 5, 15, 60], ["S", "I"])
strat.set_population_split({"0": 0.8, "5": 0.1, "15": 0.1, "60": 0})
model.stratify_with(strat)
# Run the model for 5 years.
model.run()
# 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.2, 319.4, 139.2, 5.2, 0.0, 0.0, 0.0, 0.0],
[439.0, 381.3, 171.1, 8.6, 0.0, 0.0, 0.0, 0.0],
[359.5, 420.6, 207.1, 12.8, 0.0, 0.0, 0.0, 0.0],
[294.3, 442.5, 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_stochastic_transition_flows(pop, rtol, recovery_rate):
"""
Check that transition flows produce outputs that tend towards mean as pop increases.
"""
model = CompartmentalModel(
times=[0, 10], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
s_pop = 0.10 * pop
i_pop = 0.90 * pop
model.set_initial_population(distribution={"S": s_pop, "I": i_pop})
model.add_transition_flow("recovery", recovery_rate, "I", "R")
model.run_stochastic(RANDOM_SEED)
# No change to susceptible compartments
assert_array_equal(model.outputs[:, 0], s_pop)
# Calculate recoveries using mean recovery rate
mean_i = np.zeros_like(model.times)
mean_r = np.zeros_like(model.times)
mean_i[0] = i_pop
for i in range(1, len(model.times)):
recovered = mean_i[i - 1] * recovery_rate
mean_i[i] = mean_i[i - 1] - recovered
mean_r[i] = mean_r[i - 1] + recovered
# All I and R compartment sizes are are within the error range
# of the mean of the multinomial dist that determines transition.
assert_allclose(model.outputs[:, 1], mean_i, rtol=rtol)
assert_allclose(model.outputs[:, 2], mean_r, rtol=rtol)
def test_strat_model__with_age__expect_ageing():
"""
Ensure that a module with age stratification produces ageing flows,
and the correct output.
"""
model = CompartmentalModel(times=[0, 5], compartments=["S", "I"], infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 1000, "I": 0})
strat = AgeStratification("age", [0, 5, 15, 60], ["S", "I"])
model.stratify_with(strat)
# Run the model for 5 years.
model.run()
# 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.7, 269.3, 270.3, 255.8, 0.0, 0.0, 0.0, 0.0],
[167.6, 278.9, 291.5, 262.0, 0.0, 0.0, 0.0, 0.0],
[137.2, 281.2, 312.8, 268.7, 0.0, 0.0, 0.0, 0.0],
[112.3, 278.1, 333.7, 275.9, 0.0, 0.0, 0.0, 0.0],
[92.0, 270.9, 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_model__with_recovery_rate__expect_all_recover():
"""
Ensure that a model with three compartments and only recovery dynamics
results in (almost) everybody recovering.
"""
# Set up a model with 100 people, all infectious.
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"I": 100})
# Add recovery dynamics.
model.add_transition_flow("recovery", 1, "I", "R")
model.run()
# Expect that almost everyone recovers
expected_outputs = np.array(
[
[0.00, 100.00, 0.00], # Initial conditions
[0.00, 36.79, 63.21],
[0.00, 13.53, 86.47],
[0.00, 4.98, 95.02],
[0.00, 1.83, 98.17],
[0.00, 0.67, 99.33],
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
def build_model(params: dict) -> CompartmentalModel:
"""
Build the master function to run a simple SIR model
"""
# Define model compartments.
compartments = ["S", "I", "R"]
time = params["time"]
model = CompartmentalModel(
times=[time["start"], time["end"]],
compartments=compartments,
infectious_compartments=["I"],
timestep=time["step"],
)
model.set_initial_population({
"I": 1,
"S": 999999,
})
# Add flows
model.add_infection_frequency_flow("infection", params["contact_rate"],
"S", "I")
model.add_fractional_flow("recovery", params["recovery_rate"], "I", "R")
# Request derived outputs
model.request_output_for_compartments("prevalence_susceptible",
compartments=["S"])
model.request_output_for_compartments("prevalence_infectious",
compartments=["I"])
return model
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_stochastic_recovery_exitinction(recovery_rate, contact_rate):
"""
A smokey test to make sure the disease goes extinct sometimes,
because the infectious person recovers before they can infect someone else.
Calculations similar to test_stochastic_death_exitinction
"""
pr_recovery = 1 - np.exp(-recovery_rate)
pr_infected = 1 - np.exp(-contact_rate / 1000)
pr_noone_infected = binom.pmf(0, 1000, pr_infected)
pr_extinction = pr_recovery * pr_noone_infected
expected_extinctions = _find_num_successes(pr_extinction, TRIALS,
ERROR_RATE)
count_extinctions = 0
for _ in range(TRIALS):
model = CompartmentalModel(
times=[0, 1],
compartments=["S", "I", "R"],
infectious_compartments=["I"],
)
model.set_initial_population(distribution={"S": 999, "I": 1})
model.add_transition_flow("recovery", recovery_rate, "I", "R")
model.add_infection_frequency_flow("infection", contact_rate, "S", "I")
model.run_stochastic()
is_extinct = model.outputs[1, 1] == 0
if is_extinct:
count_extinctions += 1
assert count_extinctions >= expected_extinctions
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_stochastic_death_exitinction(death_rate, contact_rate):
"""
A smokey test to make sure the disease goes extinct around the right amount,
because the infectious person dies before they can infect someone else.
See here for how this stuff is calculated
https://autumn-files.s3-ap-southeast-2.amazonaws.com/Switching_to_stochastic_mode.pdf
Consider the following flow rates:
- 0.5 infected deaths timestep
- 2 people infected per timestep
- infection frequency force of infection of 1 inf / 1000 pop
- sus pop of 999
- contact rate of 2
- flow rate of 2 * 999 / 1000 = 1.998 ~= 2
Based on stochastic model (per person)
- P(infect_death) ~=40%(1 - e^(-0.5/1))
- P(infected) ~= 0.2% (1 - e^(-2/1000))
Using a binomial calculator, we get
- ~86% chance of 1 or more people getting infected
- ~14% chance of noone getting infected
Death and infection are independent processes within the model.
So then we expect a ~6% chance of exctinction (infected person dies, no one infected) (40% * 14%)
Given this there is a > 0.999999 chance that we see at least 25
disease exctinctions in 1000 runs (using binomial calculation)
"""
pr_death = 1 - np.exp(-death_rate)
pr_infected = 1 - np.exp(-contact_rate / 1000)
pr_noone_infected = binom.pmf(0, 1000, pr_infected)
pr_extinction = pr_death * pr_noone_infected
expected_extinctions = _find_num_successes(pr_extinction, TRIALS,
ERROR_RATE)
count_extinctions = 0
for _ in range(TRIALS):
model = CompartmentalModel(
times=[0, 1],
compartments=["S", "I", "R"],
infectious_compartments=["I"],
)
model.set_initial_population(distribution={"S": 999, "I": 1})
model.add_death_flow("infect_death", death_rate, "I")
model.add_infection_frequency_flow("infection", contact_rate, "S", "I")
model.run_stochastic()
is_extinct = model.outputs[1, 1] == 0
if is_extinct:
count_extinctions += 1
assert count_extinctions >= expected_extinctions
def test_model__with_two_symmetric_stratifications():
"""
Adding two strata with the same properties should yield the exact same infection dynamics and outputs as having no strata at all.
This does not test strains directly, but if this doesn't work then further testing is pointless.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
model.add_infection_frequency_flow("infection", 0.2, "S", "I")
model.add_sojourn_flow("recovery", 10, "I", "R")
# Do pre-run force of infection calcs.
model._prepare_to_run()
model._prepare_time_step(0, model.initial_population)
# Check infectiousness multipliers
susceptible = model.compartments[0]
infectious = model.compartments[1]
assert model._get_infection_density_multiplier(susceptible,
infectious) == 100.0
assert model._get_infection_frequency_multiplier(susceptible,
infectious) == 0.1
model.run()
# Create a stratified model where the two non-strain strata are symmetric
stratified_model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
stratified_model.set_initial_population(distribution={"S": 900, "I": 100})
stratified_model.add_infection_frequency_flow("infection", 0.2, "S", "I")
stratified_model.add_sojourn_flow("recovery", 10, "I", "R")
strat = Stratification("clinical", ["home", "hospital"], ["I"])
stratified_model.stratify_with(strat)
stratified_model.run()
# Ensure stratified model has the same results as the unstratified model.
merged_outputs = np.zeros_like(model.outputs)
merged_outputs[:, 0] = stratified_model.outputs[:, 0]
merged_outputs[:,
1] = stratified_model.outputs[:,
1] + stratified_model.outputs[:,
2]
merged_outputs[:, 2] = stratified_model.outputs[:, 3]
assert_allclose(merged_outputs,
model.outputs,
atol=0.01,
rtol=0.01,
verbose=True)
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.
"""
# Create an unstratified model
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
model.add_infection_frequency_flow("infection", 0.2, "S", "I")
model.add_sojourn_flow("recovery", 10, "I", "R")
# Do pre-run force of infection calcs.
model._prepare_to_run()
model._prepare_time_step(0, model.initial_population)
# Check infectiousness multipliers
susceptible = model.compartments[0]
infectious = model.compartments[1]
assert model._get_infection_density_multiplier(susceptible,
infectious) == 100.0
assert model._get_infection_frequency_multiplier(susceptible,
infectious) == 0.1
model.run()
# Create a stratified model where the two strain strata are symmetric
strain_model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
strain_model.set_initial_population(distribution={"S": 900, "I": 100})
strain_model.add_infection_frequency_flow("infection", 0.2, "S", "I")
strain_model.add_sojourn_flow("recovery", 10, "I", "R")
strat = StrainStratification("strain", ["a", "b"], ["I"])
strain_model.stratify_with(strat)
strain_model.run()
# Ensure stratified model has the same results as the unstratified 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_model__with_birth_and_death_rate_replace_deaths__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})
model.add_replacement_birth_flow("births", "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_stratify__single_with_pop_split__validate_compartments():
"""
Ensure stratifying a model correctly adjusts the model compartments.
Also the population split should be applied.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population({"S": 900, "I": 90, "R": 10})
# Compartments exist
assert model.compartments == [
Compartment("S"), Compartment("I"),
Compartment("R")
]
# Each compartment knows its index
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
# Compartments have the correct population
assert_array_equal(model.initial_population, np.array([900, 90, 10]))
# Stratify the model
strat = Stratification(name="age",
strata=["child", "adult"],
compartments=["S", "I", "R"])
strat.set_population_split({"child": 0.8, "adult": 0.2})
model.stratify_with(strat)
assert model._stratifications == [strat]
# Ensure compartments are stratified correctly
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
assert model.compartments == [
Compartment("S", {"age": "child"}),
Compartment("S", {"age": "adult"}),
Compartment("I", {"age": "child"}),
Compartment("I", {"age": "adult"}),
Compartment("R", {"age": "child"}),
Compartment("R", {"age": "adult"}),
]
expected_pop_arr = np.array([720, 180, 72, 18, 8, 2])
assert_array_equal(model.initial_population, expected_pop_arr)
def test_model__with_static_dynamics__expect_no_change():
"""
Ensure that a model with two compartments and no internal dynamics results in no change.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 990, "I": 10})
model.run()
# Expect that no one has moved from sucsceptible to infections at any point in time
expected_outputs = np.array(
[
[990, 10, 0], # Initial conditions
[990, 10, 0],
[990, 10, 0],
[990, 10, 0],
[990, 10, 0],
[990, 10, 0],
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
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_no_infected__expect_no_change():
"""
Ensure that if no one has the disease, then no one gets the disease in the future.
"""
# Set up a model with 100 people, all susceptible, transmission highly likely, but no one is infected.
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 100, "I": 0})
model.add_infection_frequency_flow("infection", 10, "S", "I")
model.run()
# Expect that no one has moved from sucsceptible to infections at any point in time
expected_outputs = np.array(
[
[100.0, 0.0, 0], # Initial conditions
[100.0, 0.0, 0],
[100.0, 0.0, 0],
[100.0, 0.0, 0],
[100.0, 0.0, 0],
[100.0, 0.0, 0],
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
def test_model__with_death_rate__expect_pop_decrease():
"""
Ensure that a model with two compartments and only death rate dynamics results in fewer people.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 100, "I": 100})
# Add some dying at ~2 people / 100 / year.
model.add_universal_death_flows("deaths", 0.02)
model.run()
# Expect that we have fewer people in the population per year
expected_outputs = np.array(
[
[100.0, 100.0, 0], # Initial conditions
[98.0, 98.0, 0],
[96.1, 96.1, 0],
[94.2, 94.2, 0],
[92.3, 92.3, 0],
[90.5, 90.5, 0],
]
)
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_infection_density__expect_all_infected():
"""
Ensure that a model with two compartments and one-way internal dynamics results in all infected.
"""
# Set up a model with 100 people, all susceptible execept 1 infected.
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 99, "I": 1})
model.add_infection_density_flow("infection", 0.03, "S", "I")
model.run()
# Expect that everyone gets infected eventually.
expected_outputs = np.array(
[
[99.00, 1.00, 0], # Initial conditions
[83.13, 16.87, 0],
[19.70, 80.30, 0],
[1.21, 98.79, 0],
[0.06, 99.94, 0],
[0.00, 100.00, 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 test_model__with_infect_death_rate__expect_infected_pop_decrease():
"""
Ensure that a model with two compartments and only infected death rate dynamics
results in fewer infected people, but no change to susceptible pop.
"""
# Set up a model with 100 people, all infectious.
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 50, "I": 50})
# Add some dying at ~2 people / 100 / year.
model.add_death_flow("infect_death", 0.02, "I")
model.run()
expected_outputs = np.array(
[
[50.00, 50.00, 0], # Initial conditions
[50.00, 49.01, 0],
[50.00, 48.04, 0],
[50.00, 47.09, 0],
[50.00, 46.16, 0],
[50.00, 45.24, 0],
]
)
assert_allclose(model.outputs, expected_outputs, atol=0.1, verbose=True)
def test_set_initial_population():
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert_array_equal(model.initial_population, np.array([0, 0, 0]))
model.set_initial_population({"S": 100})
assert_array_equal(model.initial_population, np.array([100, 0, 0]))
model.set_initial_population({"I": 100})
assert_array_equal(model.initial_population, np.array([0, 100, 0]))
model.set_initial_population({"R": 1, "S": 50, "I": 99})
assert_array_equal(model.initial_population, np.array([50, 99, 1]))
def _get_test_model(timestep=1, times=[0, 150]):
comps = ["S", "EE", "LE", "EA", "LA", "R"]
infectious_comps = ["LE", "EA", "LA"]
model = CompartmentalModel(
times=times,
compartments=comps,
infectious_compartments=infectious_comps,
timestep=timestep,
)
model.set_initial_population({"S": int(20e6), "LA": 100})
# Add flows
model.add_infection_frequency_flow(name="infection", contact_rate=0.03, source="S", dest="EE")
model.add_sojourn_flow(name="infect_onset", sojourn_time=7, source="EE", dest="LE")
model.add_sojourn_flow(name="incidence", sojourn_time=7, source="LE", dest="EA")
model.add_sojourn_flow(name="progress", sojourn_time=7, source="EA", dest="LA")
model.add_sojourn_flow(name="recovery", sojourn_time=7, source="LA", dest="R")
model.add_death_flow(name="infect_death", death_rate=0.005, source="LA")
model.add_transition_flow(name="warning_immunity", fractional_rate=0.01, source="R", dest="S")
# Stratify by age
age_strat = Stratification("age", AGE_STRATA, comps)
age_strat.set_population_split(AGE_SPLIT_PROPORTIONS)
age_strat.set_mixing_matrix(AGE_MIXING_MATRIX)
age_strat.add_flow_adjustments(
"infection", {s: Multiply(v) for s, v in AGE_SUSCEPTIBILITY.items()}
)
model.stratify_with(age_strat)
# Stratify by clinical status
clinical_strat = Stratification("clinical", CLINICAL_STRATA, infectious_comps)
clinical_strat.add_infectiousness_adjustments("LE", {**ADJ_BASE, "non_sympt": Overwrite(0.25)})
clinical_strat.add_infectiousness_adjustments("EA", {**ADJ_BASE, "non_sympt": Overwrite(0.25)})
clinical_strat.add_infectiousness_adjustments(
"LA",
{
**ADJ_BASE,
"non_sympt": Overwrite(0.25),
"sympt_isolate": Overwrite(0.2),
"hospital": Overwrite(0.2),
"icu": Overwrite(0.2),
},
)
clinical_strat.add_flow_adjustments(
"infect_onset",
{
"non_sympt": Multiply(0.26),
"icu": Multiply(0.01),
"hospital": Multiply(0.04),
"sympt_public": Multiply(0.66),
"sympt_isolate": Multiply(0.03),
},
)
model.stratify_with(clinical_strat)
# Request derived outputs.
model.request_output_for_flow(name="incidence", flow_name="incidence")
model.request_output_for_flow(name="progress", flow_name="progress")
for age in AGE_STRATA:
for clinical in NOTIFICATION_STRATA:
model.request_output_for_flow(
name=f"progressXage_{age}Xclinical_{clinical}",
flow_name="progress",
dest_strata={"age": age, "clinical": clinical},
)
hospital_sources = []
icu_sources = []
for age in AGE_STRATA:
icu_sources.append(f"progressXage_{age}Xclinical_icu")
hospital_sources += [
f"progressXage_{age}Xclinical_icu",
f"progressXage_{age}Xclinical_hospital",
]
model.request_aggregate_output(
name="new_hospital_admissions",
sources=hospital_sources,
)
model.request_aggregate_output(name="new_icu_admissions", sources=icu_sources)
# Get notifications, which may included people detected in-country as they progress, or imported cases which are detected.
notification_sources = [
f"progressXage_{a}Xclinical_{c}" for a in AGE_STRATA for c in NOTIFICATION_STRATA
]
model.request_aggregate_output(name="notifications", sources=notification_sources)
# Infection deaths.
model.request_output_for_flow(name="infection_deaths", flow_name="infect_death")
model.request_cumulative_output(name="accum_deaths", source="infection_deaths")
# Track hospital occupancy.
# We count all ICU and hospital late active compartments and a proportion of early active ICU cases.
model.request_output_for_compartments(
"_late_active_hospital",
compartments=["LA"],
strata={"clinical": "hospital"},
save_results=False,
)
model.request_output_for_compartments(
"icu_occupancy",
compartments=["LA"],
strata={"clinical": "icu"},
)
model.request_output_for_compartments(
"_early_active_icu",
compartments=["EA"],
strata={"clinical": "icu"},
save_results=False,
)
proportion_icu_patients_in_hospital = 0.25
model.request_function_output(
name="_early_active_icu_proportion",
func=lambda patients: patients * proportion_icu_patients_in_hospital,
sources=["_early_active_icu"],
save_results=False,
)
model.request_aggregate_output(
name="hospital_occupancy",
sources=[
"_late_active_hospital",
"icu_occupancy",
"_early_active_icu_proportion",
],
)
# Proportion seropositive
model.request_output_for_compartments(
name="_total_population", compartments=comps, save_results=False
)
model.request_output_for_compartments(name="_recovered", compartments=["R"], save_results=False)
model.request_function_output(
name="proportion_seropositive",
sources=["_recovered", "_total_population"],
func=lambda recovered, total: recovered / total,
)
return model
def test_strat_infectiousness__with_adjustments():
"""
Ensure multiply infectiousness adjustment is applied.
"""
# Create a model
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.set_initial_population(distribution={"S": 900, "I": 100})
strat = Stratification("age", ["baby", "child", "adult"], ["S", "I", "R"])
strat.set_population_split({"baby": 0.1, "child": 0.3, "adult": 0.6})
strat.add_infectiousness_adjustments(
"I", {"child": adjust.Multiply(3), "adult": adjust.Multiply(0.5), "baby": None}
)
model.stratify_with(strat)
assert_array_equal(
model.initial_population,
np.array([90, 270, 540, 10, 30, 60, 0, 0, 0]),
)
# Do pre-run force of infection calcs.
model._prepare_to_run()
assert_array_equal(
model._compartment_infectiousness["default"],
np.array([0, 0, 0, 1, 3, 0.5, 0, 0, 0]),
)
# Do pre-iteration force of infection calcs
model._prepare_time_step(0, model.initial_population)
# Get multipliers
infectees = model.compartments[0:3]
infectors = model.compartments[3:6]
expected_density = 10 * 1 + 30 * 3 + 60 * 0.5
expected_frequency = expected_density / 1000
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_density_multiplier(infectee, infector) == expected_density
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency
# Stratify again, now with overwrites
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
strat.add_infectiousness_adjustments(
"I", {"urban": adjust.Overwrite(1), "rural": adjust.Multiply(7)}
)
model.stratify_with(strat)
assert_array_equal(
model.initial_population,
np.array([45, 45, 135, 135, 270.0, 270, 5, 5, 15, 15, 30, 30, 0, 0, 0, 0, 0, 0]),
)
# Do pre-run force of infection calcs.
model._prepare_to_run()
assert_array_equal(
model._compartment_infectiousness["default"],
np.array([0, 0, 0, 0, 0, 0, 1, 7, 1, 21, 1, 3.5, 0, 0, 0, 0, 0, 0]),
)
# Do pre-iteration force of infection calcs
model._prepare_time_step(0, model.initial_population)
# Get multipliers
infectees = model.compartments[0:6]
infectors = model.compartments[6:12]
expected_density = 5 * 1 + 5 * 7 + 15 * 1 + 15 * 21 + 30 * 1 + 30 * 3.5
expected_frequency = expected_density / 1000
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_density_multiplier(infectee, infector) == expected_density
for infectee, infector in zip(infectees, infectors):
assert model._get_infection_frequency_multiplier(infectee, infector) == expected_frequency
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.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
contact_rate = 0.2
model.add_infection_frequency_flow("infection", contact_rate, "S", "I")
age_strat = Stratification("age", ["child", "adult"], ["S", "I", "R"])
age_strat.set_population_split({
"child": 0.6, # 600 people
"adult": 0.4, # 400 people
})
# Higher mixing among adults or children,
# than between adults or children.
age_strat.set_mixing_matrix(np.array([[1.5, 0.5], [0.5, 1.5]]))
model.stratify_with(age_strat)
strain_strat = StrainStratification("strain", ["a", "b", "c"], ["I"])
strain_strat.set_population_split({
"a": 0.7, # 70 people
"b": 0.2, # 20 people
"c": 0.1, # 10 people
})
strain_strat.add_infectiousness_adjustments(
"I",
{
"a": adjust.Multiply(0.5), # 0.5x as susceptible
"b": adjust.Multiply(3), # 3x as susceptible
"c": adjust.Multiply(2), # 2x as susceptible
},
)
model.stratify_with(strain_strat)
# 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.initial_population)
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.compartments[0]
sus_adult = model.compartments[1]
inf_child_a = model.compartments[2]
inf_child_b = model.compartments[3]
inf_child_c = model.compartments[4]
inf_adult_a = model.compartments[5]
inf_adult_b = model.compartments[6]
inf_adult_c = model.compartments[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_compartment_rates(model.initial_population, 0)
assert_allclose(expected_flow_rates, flow_rates, verbose=True)
def test_stratify__double_with_split_and_partial__validate_compartments():
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population({"S": 900, "I": 90, "R": 10})
# Compartments exist
assert model.compartments == [
Compartment("S"), Compartment("I"),
Compartment("R")
]
# Each compartment knows its index
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
# Compartments have the correct population
assert_array_equal(model.initial_population, np.array([900, 90, 10]))
# Stratify the model
age_strat = Stratification(name="age",
strata=["child", "adult"],
compartments=["S", "R"])
age_strat.set_population_split({"child": 0.8, "adult": 0.2})
model.stratify_with(age_strat)
assert model._stratifications == [age_strat]
# Ensure compartments are stratified correctly
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
assert model.compartments == [
Compartment("S", {"age": "child"}),
Compartment("S", {"age": "adult"}),
Compartment("I"),
Compartment("R", {"age": "child"}),
Compartment("R", {"age": "adult"}),
]
expected_pop_arr = np.array([720, 180, 90, 8, 2])
assert_array_equal(model.initial_population, expected_pop_arr)
# Stratify the model again!
loc_strat = Stratification(name="location",
strata=["urban", "rural", "alpine"],
compartments=["S", "I"])
loc_strat.set_population_split({"urban": 0.7, "rural": 0.2, "alpine": 0.1})
model.stratify_with(loc_strat)
assert model._stratifications == [age_strat, loc_strat]
# Ensure compartments are stratified correctly
assert [c.idx for c in model.compartments
] == list(range(len(model.compartments)))
assert model.compartments == [
Compartment("S", {
"age": "child",
"location": "urban"
}),
Compartment("S", {
"age": "child",
"location": "rural"
}),
Compartment("S", {
"age": "child",
"location": "alpine"
}),
Compartment("S", {
"age": "adult",
"location": "urban"
}),
Compartment("S", {
"age": "adult",
"location": "rural"
}),
Compartment("S", {
"age": "adult",
"location": "alpine"
}),
Compartment("I", {"location": "urban"}),
Compartment("I", {"location": "rural"}),
Compartment("I", {"location": "alpine"}),
Compartment("R", {"age": "child"}),
Compartment("R", {"age": "adult"}),
]
expected_pop_arr = np.array([504, 144, 72, 126, 36, 18, 63, 18, 9, 8, 2])
assert_allclose(model.initial_population,
expected_pop_arr,
atol=1e-9,
rtol=0)
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_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.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
model.set_initial_population(distribution={"S": 900, "I": 100})
contact_rate = 0.2
model.add_infection_frequency_flow("infection", contact_rate, "S", "I")
strat = StrainStratification("strain", ["a", "b", "c"], ["I"])
strat.set_population_split({
"a": 0.7, # 70 people
"b": 0.2, # 20 people
"c": 0.1, # 10 people
})
strat.add_flow_adjustments(
"infection",
{
"a": adjust.Multiply(0.5), # 0.5x as susceptible
"b": adjust.Multiply(3), # 3x as susceptible
"c": adjust.Multiply(2), # 2x as susceptible
},
)
model.stratify_with(strat)
# 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.initial_population)
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.compartments[0]
infectious_a = model.compartments[1]
infectious_b = model.compartments[2]
infectious_c = model.compartments[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_compartment_rates(model.initial_population, 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)
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)