Stratification introduction¶
So far we’ve looked at how to create a compartmental model, add flows, request derived outputs and use different solvers. Now we’ll look into stratifying a model using the Stratification class.
Consider the following scenario, you are modelling some infectious disease that has different susceptibility and infection mortality across different age groups. We can use stratifications to model this.
In this example we’ll cover:
First, let’s import the summer library and create a new CompartmentalModel with S, I, R compartments plus infection, death and recovery flows.
[1]
from summer2 import CompartmentalModel
import pandas as pd
pd.options.plotting.backend = "plotly"
def build_model():
"""Returns a model for the stratification examples"""
model = CompartmentalModel(
times=[0, 20],
compartments=["S", "I", "R"],
infectious_compartments=["I"],
timestep=0.1,
)
# Add people to the model
model.set_initial_population(distribution={"S": 990, "I": 10})
# Susceptible people can get infected.
model.add_infection_frequency_flow(name="infection", contact_rate=2, source="S", dest="I")
# Infectious people take 3 years, on average, to recover.
model.add_transition_flow(name="recovery", fractional_rate=1/3, source="I", dest="R")
# Add an infection-specific death flow to the I compartment.
model.add_death_flow(name="infection_death", death_rate=0.05, source="I")
return model
No stratification¶
With no stratification, this is just a regular SIR model: there are 3 compartments where susceptible people get infected/infectious, some of them die, and some of them recover.
[2]
# Build and run model with no stratifications
model = build_model()
model.run()
model.get_outputs_df().plot()
Minimal stratification¶
Next, let’s try a simple stratification where we split the population into ‘young’ (say, 0 to 39 years old) and ‘old’ (aged 40 and above). Notice the following changes to the model outputs:
There are now 6 compartments instead of 3: each original compartment has been split into an “old” and “young” compartment, with the original population evenly divided between them (by default)
The model dynamics haven’t changed otherwise: we will get the same results as before if we add the old and young compartments back together. This is because there is homogeneous mixing between strata and no demographic processes, etc.
[3]
from summer2 import Stratification
# Create a stratification named 'age', applying to all compartments, which
# splits the population into 'young' and 'old'.
strata = ["young", "old"]
strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
# Build and run model with the stratification we just defined
model = build_model()
model.stratify_with(strat)
model.run()
model.get_outputs_df().plot()
Importation flows and stratification¶
We build a model with the same simple stratifaction as above; 2 age compartments, “young” and “old”, but then add an importation flow to the model. Note the following important details:
In addition to the existing (transition) infections, there are new importatation infections for both young and old, each at half the total rate specified. This is because split_imports is set to True, and therefore evenly divides its total amongst the target compartments. The increase in infections compared to the previous run is consistent with this.
The importation flow is added to the model directly, but only after the Stratification has been applied. This is because split_imports uses the model state at the time it is called in order to determine its splitting values.
[4]
from summer2 import Stratification
# Create a stratification named 'age', applying to all compartments, which
# splits the population into 'young' and 'old'.
strata = ["young", "old"]
strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
# Build and run model with the stratification we just defined
model = build_model()
# Stratify the model first
model.stratify_with(strat)
# Now the following call is aware of the changes made by the Stratification
model.add_importation_flow("infection_imports", 10, "I", split_imports=True)
model.run()
model.get_outputs_df().plot()
Population distribution¶
We may not always wish to split the population evenly between strata. For example, we might know that 60% of the population is ‘young’ while 40% is ‘old’. Notice that
The stratified compartments are now split according to a 60:40 ratio into young and old respectively
The overall model dynamics still haven’t changed otherwise
[5]
from summer2 import Stratification
# Create a stratification named 'age', applying to all compartments, which
# splits the population into 'young' and 'old'.
strata = ["young", "old"]
strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
# Set a population distribution
strat.set_population_split({"young": 0.6, "old": 0.4})
# Build and run model with the stratification we just defined
model = build_model()
model.stratify_with(strat)
model.run()
model.get_outputs_df().plot()
Flow adjustments¶
As noted so far, we’ve been successful in subdividing the population, but haven’t actually changed our model dynamics, which is kind of boring. Next let’s look at how we can adjust the flow rates based on strata. Let’s assume two new facts about our disease:
young people are twice as susceptible to infection
old people are three times as likely to die from the infectious disease, while younger people are half as likely as under the original parameters we requested
younger people take twice as long to recover
These inter-strata differences can be modelled using flow adjustments. Now we’re seeing some genuinely new model dynamics. Note how there are fewer recovered ‘old’ people at the end of the model run, because of their higher death rate.
[6]
from summer2 import Stratification, Multiply
# Create a stratification named 'age', applying to all compartments, which
# splits the population into 'young' and 'old'.
strata = ["young", "old"]
strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
# Set a population distribution
strat.set_population_split({"young": 0.6, "old": 0.4})
# Add an adjustment to the 'infection' flow
strat.set_flow_adjustments("infection", {
"old": None, # No adjustment for old people, use baseline requested value
"young": Multiply(2), # Young people are twice as susceptible to infection
})
# Add an adjustment to the 'infection_death' flow
strat.set_flow_adjustments("infection_death", {
"old": Multiply(3), # Older people die at three times the rate requested under the original parameters
"young": Multiply(0.5), # Younger people die at half the rate requested under the original parameters
})
# Add an adjustment to the 'recovery' flow
strat.set_flow_adjustments("recovery", {
"old": None, # No adjustment for old people, use baseline
"young": Multiply(0.5), # Young people take twice as long to recover
})
# Build and run model with the stratification we just defined
model = build_model()
model.stratify_with(strat)
model.run()
model.get_outputs_df().plot()
Infectiousness adjustments¶
In addition to adjusting flow rates for each strata, you can also set an infectiousness level for a given strata. This affects how likely an infectious person in that stratum is to infect someone else. For example we could consider the following:
young people are 1.2 times as infectious, because they’re not wearing face masks as much
young people are twice as susceptible to the disease, because some of them have immature immune systems
[7]
from summer2 import Stratification, Multiply, Overwrite
# Create a stratification named 'age', applying to all compartments, which
# splits the population into 'young' and 'old'.
strata = ["young", "old"]
strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
# Set a population distribution
strat.set_population_split({"young": 0.5, "old": 0.5})
# Add an adjustment to the 'infection' flow
strat.set_flow_adjustments("infection", {
"old": None, # No adjustment for old people, use baseline
"young": Multiply(2), # Young people twice as susceptible
})
# Add an adjustment to infectiousness levels for young people in the 'I' compartment
strat.add_infectiousness_adjustments("I", {
"old": None, # No adjustment for old people, use baseline
"young": Multiply(1.2), # Young people 1.2x more infectious
})
# Build and run model with the stratification we just defined
model = build_model()
model.stratify_with(strat)
model.run()
model.get_outputs_df().plot()
Partial stratifications¶
So far we’ve been stratifying all compartments, but Summer allows only some of the compartments to be stratified. For example, we can stratify only the infectious compartment to model three different levels of disease severity: asymptomatic, mild and severe.
When you do a partial stratification, flow rates into that stratified compartment will automatically be adjusted with an even split to conserve the behaviour by default, e.g. a flow rate of 3 from a source will be evenly split into (1, 1, 1) across the three destinations. This behaviour can be manually overriden with a flow adjustment.
[8]
from summer2 import Stratification, Multiply, Overwrite
# Create a stratification named 'severity', applying to the infectious, which
# splits that compartment into 'asymptomatic', 'mild' and 'severe'.
strata = ["asymptomatic", "mild", "severe"]
strat = Stratification(name="severity", strata=strata, compartments=['I'])
# Set a population distribution - everyone starts out asymptomatic.
strat.set_population_split({"asymptomatic": 1.0, "mild": 0, "severe": 0})
# Add an adjustment to the 'infection' flow, overriding default split.
strat.set_flow_adjustments("infection", {
"asymptomatic": Multiply(0.3), # 30% of incident cases are asymptomatic
"mild": Multiply(0.5), # 50% of incident cases are mild
"severe": Multiply(0.2), # 20% of incident cases are severe
})
# Add an adjustment to the 'infection_death' flow
strat.set_flow_adjustments("infection_death", {
"asymptomatic": Multiply(0.5),
"mild": None,
"severe": Multiply(1.5),
})
strat.add_infectiousness_adjustments("I", {
"asymptomatic": Multiply(0.5),
"mild": None,
"severe": Multiply(1.5),
})
# Build and run model with the stratification we just defined
model = build_model()
model.stratify_with(strat)
model.run()
model.get_outputs_df().plot()
Multiple stratifications¶
A model can have multiple stratifications applied in series. For example, we can add an ‘age’ stratification, followed by a ‘severity’ one.
[9]
from summer2 import Stratification, Multiply, Overwrite
### Age stratification
# Create a stratification named 'age', applying to all compartments,
# which splits the population into 'young' and 'old'.
strata = ["young", "old"]
age_strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
age_strat.set_population_split({"young": 0.6, "old": 0.4})
# Add an adjustment to the 'infection' flow
age_strat.set_flow_adjustments("infection", {
"old": None, # No adjustment for old people, use unstratified parameter value
"young": Multiply(2), # Young people are twice as susceptible
})
# Add an adjustment to infectiousness levels for young people the 'I' compartment
age_strat.add_infectiousness_adjustments("I", {
"old": None, # No adjustment for old people, use unstratified parameter value
"young": Multiply(1.2), # Young people are 5x more infectious
})
### Disease severity stratification
# Create a stratification named 'severity', applying to the infectious compartment, which
# splits that compartment into 'asymptomatic', 'mild' and 'severe'.
strata = ["asymptomatic", "mild", "severe"]
severity_strat = Stratification(name="severity", strata=strata, compartments=["I"])
severity_strat.set_population_split({"asymptomatic": 1.0, "mild": 0, "severe": 0})
# Add an adjustment to the 'infection' flow (overriding the default split of one third to each stratum)
severity_strat.set_flow_adjustments("infection", {
"asymptomatic": Multiply(0.3), # 30% of cases are asympt.
"mild": Multiply(0.5), # 50% of cases are mild.
"severe": Multiply(0.2), # 20% of cases are severse.
})
# Add an adjustment to the 'infection_death' flow
severity_strat.set_flow_adjustments("infection_death", {
"asymptomatic": Multiply(0.5),
"mild": None,
"severe": Multiply(1.5),
})
severity_strat.add_infectiousness_adjustments("I", {
"asymptomatic": Multiply(0.5),
"mild": None,
"severe": Multiply(1.5),
})
# Build and run model with the stratifications we just defined
model = build_model()
# Apply age, then severity stratifications
model.stratify_with(age_strat)
model.stratify_with(severity_strat)
model.run()
model.get_outputs_df().plot()
Multiple interdependent stratifications¶
In the previous example we assumed that the age and severity stratifications were independent. For example, we assumed that the proportion of infected people who have a disease severity of asymptomatic, mild and severe is the same for both young and old people. Perhaps, for a given disease, this is not true! it’s easy to imagine an infection for which younger people tend towards being more asymptomatic, and older people tend towards having a more severe infection.
This interdependency between stratifications can be modelled using Summer, where a flow adjustment for a stratification can selectively refer to strata used for previous stratifications. You can refer to the API reference for set_flow_adjustments for more details.
To clarify, let’s consider the example described above:
[10]
from summer2 import Stratification, Multiply, Overwrite
### Age stratification
# Create a stratification named 'age', applying to all compartments,
# which splits the population into 'young' and 'old'.
strata = ["young", "old"]
age_strat = Stratification(name="age", strata=strata, compartments=["S", "I", "R"])
age_strat.set_population_split({"young": 0.6, "old": 0.4})
### Disease severity stratification (depends on the age stratification)
# Create a stratification named 'severity', applying to the infectious, which
# splits that compartment into 'asymptomatic', 'mild' and 'severe'.
strata = ["asymptomatic", "mild", "severe"]
severity_strat = Stratification(name="severity", strata=strata, compartments=["I"])
severity_strat.set_population_split({"asymptomatic": 1.0, "mild": 0, "severe": 0})
# Add an adjustment to the 'infection' flow for young people
# where younger people tend towards asymptomatic infection
young_infection_adjustments = {
"asymptomatic": Multiply(0.5), # 50% of cases are asympt.
"mild": Multiply(0.4), # 40% of cases are mild.
"severe": Multiply(0.1), # 10% of cases are severe.
}
severity_strat.set_flow_adjustments(
"infection",
young_infection_adjustments,
source_strata={'age': 'young'} # Only apply this adjustment to flows of young people
)
# Add an adjustment to the 'infection' flow for old people
# where older people tend towards severe infection
old_infection_adjustments = {
"asymptomatic": Multiply(0.1), # 10% of cases are asympt.
"mild": Multiply(0.4), # 40% of cases are mild.
"severe": Multiply(0.5), # 50% of cases are severe.
}
severity_strat.set_flow_adjustments(
"infection",
old_infection_adjustments,
source_strata={'age': 'old'} # Only apply this adjustment to flows of old people
)
# Add an adjustment to the 'infection_death' flow (for all age groups)
severity_strat.set_flow_adjustments("infection_death", {
"asymptomatic": Multiply(0.5),
"mild": None,
"severe": Multiply(1.5),
})
# Adjust infectiousness levels (for all age groups)
severity_strat.add_infectiousness_adjustments("I", {
"asymptomatic": Multiply(0.5),
"mild": None,
"severe": Multiply(1.5),
})
# Build and run model with the stratifications we just defined
model = build_model()
# Apply age, then severity stratifications
model.stratify_with(age_strat)
model.stratify_with(severity_strat)
model.run()
model.get_outputs_df().plot()
[ ]