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 test_full_age_strat():
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
age_strat = AgeStratification(name="age",
strata=["0", "5", "10"],
compartments=["S", "I", "R"])
model.stratify_with(age_strat)
def test_partial_age_strat_fails():
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
age_strat = AgeStratification(name="age",
strata=["0", "5", "10"],
compartments=["I", "R"])
with pytest.raises(AssertionError):
model.stratify_with(age_strat)
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_stratify_exit_flows():
"""
Ensure exit flows are stratified correctly.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_death_flow("d_S", 3, "S")
model.add_death_flow("d_I", 5, "I")
model.add_death_flow("d_R", 7, "R")
expected_flows = [
DeathFlow("d_S", C("S"), 3),
DeathFlow("d_I", C("I"), 5),
DeathFlow("d_R", C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
# Apply partial stratification
strat = Stratification("location", ["urban", "rural"], ["S", "I"])
model.stratify_with(strat)
expected_flows = [
DeathFlow("d_S", C("S", {"location": "urban"}), 3),
DeathFlow("d_S", C("S", {"location": "rural"}), 3),
DeathFlow("d_I", C("I", {"location": "urban"}), 5),
DeathFlow("d_I", C("I", {"location": "rural"}), 5),
DeathFlow("d_R", C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
# Apply partial stratification with flow adjustments
strat = Stratification("age", ["young", "old"], ["I", "R"])
strat.add_flow_adjustments("d_I", {"young": Multiply(0.5), "old": Multiply(2)})
strat.add_flow_adjustments("d_R", {"young": Multiply(0.5), "old": Multiply(2)})
model.stratify_with(strat)
expected_flows = [
DeathFlow("d_S", C("S", {"location": "urban"}), 3),
DeathFlow("d_S", C("S", {"location": "rural"}), 3),
DeathFlow("d_I", C("I", {"location": "urban", "age": "young"}), 5, [Multiply(0.5)]),
DeathFlow("d_I", C("I", {"location": "urban", "age": "old"}), 5, [Multiply(2)]),
DeathFlow("d_I", C("I", {"location": "rural", "age": "young"}), 5, [Multiply(0.5)]),
DeathFlow("d_I", C("I", {"location": "rural", "age": "old"}), 5, [Multiply(2)]),
DeathFlow("d_R", C("R", {"age": "young"}), 7, [Multiply(0.5)]),
DeathFlow("d_R", C("R", {"age": "old"}), 7, [Multiply(2)]),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
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_add_entry_flows_post_stratification():
"""
Ensure we can add flows after a model is stratified.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert len(model._flows) == 0
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
model.stratify_with(strat)
with pytest.raises(AssertionError):
model.add_importation_flow("imports", 10, "S", expected_flow_count=1)
assert len(model._flows) == 0
model.add_importation_flow("imports", 10, "S", expected_flow_count=2)
assert len(model._flows) == 2
expected_flows = [
ImportFlow("imports", C("S", {"location": "urban"}), 10, []),
ImportFlow("imports", C("S", {"location": "rural"}), 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_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_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 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_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_add_exit_flows_post_stratification():
"""
Ensure user can add exit flows post stratification.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert len(model._flows) == 0
# Apply partial stratification
strat = Stratification("location", ["urban", "rural"], ["S", "I"])
model.stratify_with(strat)
assert len(model._flows) == 0
model.add_death_flow("d_S", 3, "S")
model.add_death_flow("d_I", 5, "I")
model.add_death_flow("d_R", 7, "R")
expected_flows = [
DeathFlow("d_S", C("S", {"location": "urban"}), 3),
DeathFlow("d_S", C("S", {"location": "rural"}), 3),
DeathFlow("d_I", C("I", {"location": "urban"}), 5),
DeathFlow("d_I", C("I", {"location": "rural"}), 5),
DeathFlow("d_R", C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_add_mixing_matrix_fails():
"""
Ensure validation works when trying to add a mixing matrix.
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
strat = Stratification(name="agegroup",
strata=["child", "adult"],
compartments=["S", "R"])
mixing_matrix = np.array([[2, 3], [5, 7]])
strat.set_mixing_matrix(mixing_matrix)
# Expect this to fail because it's not a full stratification (no I compartment).
with pytest.raises(AssertionError):
model.stratify_with(strat)
def test_set_times__with_failure(start_time, end_time, time_step):
with pytest.raises(AssertionError):
CompartmentalModel(
times=(start_time, end_time),
compartments=["S", "I", "R"],
infectious_compartments=["I"],
timestep=time_step,
)
def test_no_mixing_matrix__with_previous_strat():
"""
Test that we are using the default 'null-op' mixing matrix when
we have a no user-specified mixing matrix and a stratification has already been applied
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
# Apply first stratification with a mixing matrix.
strat = Stratification(name="agegroup",
strata=["child", "adult"],
compartments=["S", "I", "R"])
first_strat_matrix = np.array([[2, 3], [5, 7]])
strat.set_mixing_matrix(first_strat_matrix)
model.stratify_with(strat)
# We should get the default mixing matrix
actual_mixing = model._get_mixing_matrix(0)
assert_array_equal(actual_mixing, first_strat_matrix)
# Static matrices shouldn't change over time
actual_mixing = model._get_mixing_matrix(123)
assert_array_equal(actual_mixing, first_strat_matrix)
# Agegroup mixing categories have been added.
assert model._mixing_categories == [{
"agegroup": "child"
}, {
"agegroup": "adult"
}]
# Apply second stratification with no mixing matrix.
strat = Stratification(name="location",
strata=["work", "home"],
compartments=["S", "I", "R"])
model.stratify_with(strat)
# We should get the same results as before.
actual_mixing = model._get_mixing_matrix(0)
assert_array_equal(actual_mixing, first_strat_matrix)
actual_mixing = model._get_mixing_matrix(123)
assert_array_equal(actual_mixing, first_strat_matrix)
assert model._mixing_categories == [{
"agegroup": "child"
}, {
"agegroup": "adult"
}]
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 _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_stratify_entry_flows__with_no_explicit_adjustments():
"""
Ensure entry flows are stratified correctly when no adjustments are requested.
Expect flow to be conserved, split evenly over the new strata.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_importation_flow("imports", 10, "S")
expected_flows = [ImportFlow("imports", C("S"), 10)]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
model.stratify_with(strat)
expected_flows = [
ImportFlow("imports", C("S", {"location": "urban"}), 10, [Multiply(0.5)]),
ImportFlow("imports", C("S", {"location": "rural"}), 10, [Multiply(0.5)]),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
strat = Stratification("age", ["young", "old"], ["S", "I", "R"])
model.stratify_with(strat)
expected_flows = [
ImportFlow(
"imports",
C("S", {"location": "urban", "age": "young"}),
10,
[Multiply(0.5), Multiply(0.5)],
),
ImportFlow(
"imports",
C("S", {"location": "urban", "age": "old"}),
10,
[Multiply(0.5), Multiply(0.5)],
),
ImportFlow(
"imports",
C("S", {"location": "rural", "age": "young"}),
10,
[Multiply(0.5), Multiply(0.5)],
),
ImportFlow(
"imports",
C("S", {"location": "rural", "age": "old"}),
10,
[Multiply(0.5), Multiply(0.5)],
),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_create_model():
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
assert_array_equal(model.times, np.array([0, 1, 2, 3, 4, 5]))
assert model.compartments == [Compartment("S"), Compartment("I"), Compartment("R")]
assert model._infectious_compartments == [Compartment("I")]
assert_array_equal(model.initial_population, np.array([0, 0, 0]))
# Times out of order
with pytest.raises(AssertionError):
CompartmentalModel(
times=[5, 0], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
# Infectious compartment not a compartment
with pytest.raises(AssertionError):
CompartmentalModel(
times=[-1, 5], compartments=["S", "I", "R"], infectious_compartments=["E"]
)
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_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_single_static_mixing_matrix():
"""
Test that we are using the correct mixing matrix when
we have a single static mixing matrix
"""
model = CompartmentalModel(times=[0, 5],
compartments=["S", "I", "R"],
infectious_compartments=["I"])
# Apply first stratification with a mixing matrix.
strat = Stratification(name="agegroup",
strata=["child", "adult"],
compartments=["S", "I", "R"])
mixing_matrix = np.array([[2, 3], [5, 7]])
strat.set_mixing_matrix(mixing_matrix)
model.stratify_with(strat)
# We should get the default mixing matrix
actual_mixing = model._get_mixing_matrix(0)
assert_array_equal(actual_mixing, mixing_matrix)
# Static matrices shouldn't change over time
actual_mixing = model._get_mixing_matrix(123)
assert_array_equal(actual_mixing, mixing_matrix)
# Agegroup mixing categories have been added.
assert model._mixing_categories == [{
"agegroup": "child"
}, {
"agegroup": "adult"
}]
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 test_stratify_transition_flows__with_dest_only_stratified__with_adjustments_and_strains():
"""
Ensure transition flows are stratified correctly when only the flow destination is stratified.
Expect adjustments to override the automatic person-conserving adjustment when using a strain strat.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_sojourn_flow("recovery", 7, "I", "R")
expected_flows = [
SojournFlow("recovery", C("I"), C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
# Apply stratification
strat = StrainStratification("location", ["urban", "rural"], ["R"])
strat.add_flow_adjustments("recovery", {"urban": Overwrite(0.7), "rural": Overwrite(0.1)})
model.stratify_with(strat)
expected_flows = [
SojournFlow("recovery", C("I"), C("R", {"location": "urban"}), 7, [Overwrite(0.7)]),
SojournFlow("recovery", C("I"), C("R", {"location": "rural"}), 7, [Overwrite(0.1)]),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_stratify_transition_flows__with_dest_only_stratified():
"""
Ensure transition flows are stratified correctly when only the flow destination is stratified.
Expect an person-conserving adjustment of 1/N to be applied to each flow - N being the number of new strata.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_sojourn_flow("recovery", 7, "I", "R")
expected_flows = [
SojournFlow("recovery", C("I"), C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
# Apply stratification
strat = Stratification("location", ["urban", "rural"], ["R"])
model.stratify_with(strat)
expected_flows = [
SojournFlow("recovery", C("I"), C("R", {"location": "urban"}), 7, [Multiply(0.5)]),
SojournFlow("recovery", C("I"), C("R", {"location": "rural"}), 7, [Multiply(0.5)]),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def test_stratify_transition_flows__with_source_only_stratified():
"""
Ensure transition flows are stratified correctly when only the flow source is stratified.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_sojourn_flow("recovery", 7, "I", "R")
expected_flows = [
SojournFlow("recovery", C("I"), C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
# Apply stratification
strat = Stratification("location", ["urban", "rural"], ["S", "I"])
model.stratify_with(strat)
expected_flows = [
SojournFlow("recovery", C("I", {"location": "urban"}), C("R"), 7),
SojournFlow("recovery", C("I", {"location": "rural"}), C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
def apply_post_cluster_strat_hacks(params: Parameters,
model: CompartmentalModel):
metro_clusters = [
region.replace("-", "_") for region in Region.VICTORIA_METRO
]
regional_clusters = [
region.replace("-", "_") for region in Region.VICTORIA_RURAL
]
vic = params.victorian_clusters
country = params.country
# A bit of a hack - to get rid of all the infectious populations from the regional clusters
for i_comp, comp in enumerate(model.compartments):
if any([
comp.has_stratum("cluster", cluster)
for cluster in regional_clusters
]) and not comp.has_name(Compartment.SUSCEPTIBLE):
model.initial_population[i_comp] = 0.0
elif any([
comp.has_stratum("cluster", cluster)
for cluster in metro_clusters
]) and not comp.has_name(Compartment.SUSCEPTIBLE):
model.initial_population[i_comp] *= 9.0 / 4.0
"""
Hack in a custom (144x144) mixing matrix where each region is adjusted individually
based on its time variant mobility data.
"""
# Get the inter-cluster mixing matrix
intercluster_mixing_matrix = create_assortative_matrix(
vic.intercluster_mixing, CLUSTER_STRATA)
# Replace regional Victoria maximum effect calibration parameters with the metro values for consistency
for microdist_process in ["face_coverings", "behaviour"]:
vic.regional.mobility.microdistancing[
f"{microdist_process}_adjuster"].parameters.effect = vic.metro.mobility.microdistancing[
f"{microdist_process}_adjuster"].parameters.effect
# Get new mixing matrix
static_mixing_matrix = inputs.get_country_mixing_matrix(
"all_locations", country.iso3)
get_mixing_matrix = build_victorian_mixing_matrix_func(
static_mixing_matrix,
vic.metro.mobility,
vic.regional.mobility,
country,
intercluster_mixing_matrix,
)
setattr(model, "_get_mixing_matrix", MethodType(get_mixing_matrix, model))
def test_set_times__with_success(start_time, end_time, time_step):
model = CompartmentalModel(
times=(start_time, end_time),
compartments=["S", "I", "R"],
infectious_compartments=["I"],
timestep=time_step,
)
assert all([start_time <= t <= end_time for t in model.times])
assert_array_equal(model.times, np.sort(model.times))
assert_array_equal(model.times, np.unique(model.times))
assert start_time in model.times # Start time should always be included in evaluation times.
assert end_time in model.times # End Time should always be included in evaluation times.
# Check time step size has been applied correctly for each time step
for i_time in range(2, len(model.times) - 1):
assert abs(model.times[i_time] - model.times[i_time - 1] - time_step) < 1e-6
def test_stratify_transition_flows__with_source_and_dest_stratified():
"""
Ensure transition flows are stratified correctly when both the flow source and dest are stratified.
"""
model = CompartmentalModel(
times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
)
model.add_infection_frequency_flow("infection", 0.03, "S", "I")
model.add_sojourn_flow("recovery", 7, "I", "R")
expected_flows = [
InfectionFrequencyFlow(
"infection", C("S"), C("I"), 0.03, model._get_infection_frequency_multiplier
),
SojournFlow("recovery", C("I"), C("R"), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
# Apply stratification
strat = Stratification("location", ["urban", "rural"], ["S", "I", "R"])
model.stratify_with(strat)
expected_flows = [
InfectionFrequencyFlow(
"infection",
C("S", {"location": "urban"}),
C("I", {"location": "urban"}),
0.03,
model._get_infection_frequency_multiplier,
),
InfectionFrequencyFlow(
"infection",
C("S", {"location": "rural"}),
C("I", {"location": "rural"}),
0.03,
model._get_infection_frequency_multiplier,
),
SojournFlow("recovery", C("I", {"location": "urban"}), C("R", {"location": "urban"}), 7),
SojournFlow("recovery", C("I", {"location": "rural"}), C("R", {"location": "rural"}), 7),
]
assert len(expected_flows) == len(model._flows)
assert all([a._is_equal(e) for e, a in zip(expected_flows, model._flows)])
