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_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_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_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_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_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)