def test_epi_model__with_death_rate__expect_pop_decrease(ModelClass):
"""
Ensure that a model with two compartments and only death rate dynamics results in fewer people.
"""
# Set up a model with 100 people, all susceptible, no transmission possible.
# Add some dying at ~2 people / 100 / year.
pop = 100
model = ModelClass(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={"universal_death_rate": 2e-2},
requested_flows=[],
birth_approach=BirthApproach.NO_BIRTH,
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that we have more people in the population
expected_output = [
[100.0, 0.0],
[98.0, 0.0],
[96.0, 0.0],
[94.0, 0.0],
[92.0, 0.0],
[91.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_epi_model__with_infection_density__expect_all_infected(ModelClass):
"""
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.
pop = 100
model = ModelClass(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.EARLY_INFECTIOUS: 1},
parameters={"contact_rate": 0.03},
requested_flows=[
{
"type": Flow.INFECTION_DENSITY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.EARLY_INFECTIOUS,
}
],
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that everyone gets infected eventually.
expected_output = [
[[99.0, 1.0], [83.0, 17.0], [20.0, 80.0], [1.0, 99.0], [0.0, 100.0], [0.0, 100.0],]
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_strat_model__with_locations__expect_no_change():
"""
Ensure that a module with location stratification populates locations correctly.
"""
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic location stratification
model.stratify(
Stratification.LOCATION,
strata_request=["rural", "urban", "prison"],
compartment_types_to_stratify=[],
requested_proportions={"rural": 0.44, "urban": 0.55, "prison": 0.01},
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to start in their locations, then nothing should change,
expected_output = [
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
[440.0, 550.0, 10.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_epi_model__with_no_infected__expect_no_change(ModelClass):
"""
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.
pop = 100
model = ModelClass(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={"contact_rate": 10},
requested_flows=[
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.EARLY_INFECTIOUS,
}
],
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that no one has moved from sucsceptible to infections at any point in time
expected_output = [
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_epi_model__with_static_dynamics__expect_no_change(ModelClass):
"""
Ensure that a model with two compartments and no internal dynamics results in no change.
"""
# Set up a model with 100 people, all susceptible, no transmission possible.
pop = 100
model = ModelClass(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that no one has moved from sucsceptible to infections at any point in time
expected_output = [
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
[100.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_strat_model__with_age__expect_ageing():
"""
Ensure that a module with age stratification produces ageing flows,
and the correct output.
"""
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to generally get older, but no one should die or get sick
expected_output = [
[250.0, 250.0, 250.0, 250.0, 0.0, 0.0, 0.0, 0.0],
[205.0, 269.0, 270.0, 256.0, 0.0, 0.0, 0.0, 0.0],
[168.0, 279.0, 291.0, 262.0, 0.0, 0.0, 0.0, 0.0],
[137.0, 281.0, 313.0, 269.0, 0.0, 0.0, 0.0, 0.0],
[112.0, 278.0, 334.0, 276.0, 0.0, 0.0, 0.0, 0.0],
[92.0, 271.0, 354.0, 284.0, 0.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def _build_model(*args, **kwargs):
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
)
return model
def test_epi_model__with_recovery_rate__expect_all_recover(ModelClass):
"""
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.
# Add recovery dynamics.
pop = 100
model = ModelClass(
times=get_integration_times(2000, 2005, 1),
compartment_types=[
Compartment.SUSCEPTIBLE,
Compartment.EARLY_INFECTIOUS,
Compartment.RECOVERED,
],
initial_conditions={Compartment.EARLY_INFECTIOUS: pop},
parameters={"recovery": 1},
requested_flows=[
{
"type": Flow.STANDARD,
"parameter": "recovery",
"origin": Compartment.EARLY_INFECTIOUS,
"to": Compartment.RECOVERED,
}
],
birth_approach=BirthApproach.NO_BIRTH,
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that almost everyone recovers
expected_output = [
[0.0, 100.0, 0.0],
[0.0, 37.0, 63.0],
[0.0, 14.0, 86.0],
[0.0, 5.0, 95.0],
[0.0, 2.0, 98.0],
[0.0, 1.0, 99.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_epi_model__with_infect_death_rate__expect_infected_pop_decrease(ModelClass):
"""
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 susceptible, no transmission possible.
# Add some dying at ~2 people / 100 / year.
pop = 100
model = ModelClass(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.EARLY_INFECTIOUS: 50},
parameters={"infect_death": 2e-2},
requested_flows=[
{
"type": Flow.COMPARTMENT_DEATH,
"parameter": "infect_death",
"origin": Compartment.EARLY_INFECTIOUS,
}
],
birth_approach=BirthApproach.NO_BIRTH,
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that we have more people in the population
expected_output = [
[50.0, 50.0],
[50.0, 49.0],
[50.0, 48.0],
[50.0, 47.0],
[50.0, 46.0],
[50.0, 45.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_strat_model__with_age_and_starting_proportion__expect_ageing():
"""
Ensure that a module with age stratification and starting proporptions
produces ageing flows, and the correct output.
"""
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={},
requested_flows=[],
starting_population=pop,
)
# Add basic age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={"0": 0.8, "5": 0.1, "15": 0.1},
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to generally get older, but no one should die or get sick.
# Expect initial distribution of ages to be set according to "requested_proportions".
expected_output = [
[800.0, 100.0, 100.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[655.0, 228.0, 114.0, 2.0, 0.0, 0.0, 0.0, 0.0],
[536.0, 319.0, 139.0, 5.0, 0.0, 0.0, 0.0, 0.0],
[439.0, 381.0, 171.0, 9.0, 0.0, 0.0, 0.0, 0.0],
[360.0, 421.0, 207.0, 13.0, 0.0, 0.0, 0.0, 0.0],
[294.0, 442.0, 245.0, 18.0, 0.0, 0.0, 0.0, 0.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_strat_model__with_age_and_infectiousness__expect_age_based_infectiousness():
"""
Ensure that a module with age stratification produces ageing flows,
and the correct output. Ensure that age-speific mortality is used.
"""
pop = 1000
model = StratifiedModel(
times=get_integration_times(2000, 2005, 1),
compartment_types=[Compartment.SUSCEPTIBLE, Compartment.EARLY_INFECTIOUS],
initial_conditions={Compartment.SUSCEPTIBLE: pop},
parameters={"contact_rate": None},
requested_flows=[
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.EARLY_INFECTIOUS,
}
],
starting_population=pop,
)
# Add age stratification
model.stratify(
Stratification.AGE,
strata_request=[0, 5, 15, 60],
compartment_types_to_stratify=[],
requested_proportions={},
adjustment_requests={"contact_rate": {"0": 0.0, "5": 3.0, "15": 0.0, "60": 0.0}},
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect everyone to generally get older, but no one should die or get sick
expected_output = []
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()
def test_epi_model__with_complex_dynamics__expect_correct_outputs(ModelClass):
"""
Ensure that a EpiModel with the "full suite" of TB dynamics produces correct results:
- 5 compartments
- birth rate + universal death rate
- standard inter-compartment flows
FIXME: Change parameter values to be vaugely realistic
"""
# Set up a model with 1000 people, 100 intially infected
pop = 1000
times = get_integration_times(2000, 2005, 1)
compartments = [
Compartment.SUSCEPTIBLE,
Compartment.EARLY_INFECTIOUS,
Compartment.LATE_LATENT,
Compartment.EARLY_LATENT,
Compartment.RECOVERED,
]
initial_pop = {Compartment.EARLY_INFECTIOUS: 100}
params = {
# Global birth / death params
"crude_birth_rate": 2e-2, # ~ 2 babies / 100 / year
"universal_death_rate": 2e-2, # ~ 2 deaths / 100 / year
# Compartment flow params
"infect_death": 0.4,
"recovery": 0.2,
"contact_rate": 14,
"contact_rate_recovered": 14,
"contact_rate_late_latent": 3,
"stabilisation": 3,
"early_progression": 2,
"late_progression": 1,
"case_detection": 1,
}
flows = [
# Susceptible becoming latent
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate",
"origin": Compartment.SUSCEPTIBLE,
"to": Compartment.EARLY_LATENT,
},
# Recovered becoming latent again.
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate_recovered",
"origin": Compartment.RECOVERED,
"to": Compartment.EARLY_LATENT,
},
# Late latent going back to early latent
{
"type": Flow.INFECTION_FREQUENCY,
"parameter": "contact_rate_late_latent",
"origin": Compartment.LATE_LATENT,
"to": Compartment.EARLY_LATENT,
},
# Early progression from latent to infectious
{
"type": Flow.STANDARD,
"parameter": "early_progression",
"origin": Compartment.EARLY_LATENT,
"to": Compartment.EARLY_INFECTIOUS,
},
# Transition from early to late latent
{
"type": Flow.STANDARD,
"parameter": "stabilisation",
"origin": Compartment.EARLY_LATENT,
"to": Compartment.LATE_LATENT,
},
# Late latent becoming infectious.
{
"type": Flow.STANDARD,
"parameter": "late_progression",
"origin": Compartment.LATE_LATENT,
"to": Compartment.EARLY_INFECTIOUS,
},
# Infected people dying.
{
"type": Flow.COMPARTMENT_DEATH,
"parameter": "infect_death",
"origin": Compartment.EARLY_INFECTIOUS,
},
# Infected people recovering naturally.
{
"type": Flow.STANDARD,
"parameter": "recovery",
"origin": Compartment.EARLY_INFECTIOUS,
"to": Compartment.RECOVERED,
},
# Infectious people recovering via manual intervention
{
"type": Flow.STANDARD,
"parameter": "case_detection",
"origin": Compartment.EARLY_INFECTIOUS,
"to": Compartment.RECOVERED,
},
]
model = ModelClass(
times,
compartments,
initial_pop,
params,
flows,
birth_approach=BirthApproach.ADD_CRUDE,
starting_population=pop,
)
# Run the model for 5 years.
model.run_model(integration_type=IntegrationType.ODE_INT)
# Expect that the results are consistent, nothing crazy happens.
expected_output = [
[900.0, 100.0, 0.0, 0.0, 0.0],
[66.0, 307.0, 274.0, 204.0, 75.0],
[3.0, 345.0, 221.0, 151.0, 69.0],
[2.0, 297.0, 176.0, 127.0, 58.0],
[2.0, 249.0, 146.0, 106.0, 48.0],
[1.0, 208.0, 121.0, 89.0, 40.0],
]
actual_output = np.round(model.outputs)
assert (actual_output == np.array(expected_output)).all()