Flows between strata

In the previous introduction to flows and introduction to stratification we saw a workflow where:

• An unstratified model is defined

• Flows are added to that model

• The model is then stratified, splitting the flows between new strata

This approach works fine for many workflows, but in some cases, we want to define flows that move people between strata. For example, we might want to model people migrating from a rural location to an urban location over time.

This example will show you how to implement flows between strata. Let’s start with a baseline model, stratified by location.

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=[1990, 2020],
        compartments=["pop"],
        infectious_compartments=[],
        timestep=0.1,
    )
    model.set_initial_population(distribution={"pop": 20e6})
    model.add_crude_birth_flow("birth", 0.02, "pop")
    model.add_death_flow("death", 0.01, "pop")
    return model

Unstratified model

In our example model, there is only one compartment with a birth and death rate.

model = build_model()
model.run()
model.get_outputs_df().plot()

Stratified model

Next lets split the population into urban and rural.

from summer2 import Stratification

model = build_model()

strat = Stratification('location', ['urban', 'rural'], ['pop'])
strat.set_population_split({'rural': 0.7, 'urban': 0.3})
model.stratify_with(strat)

model.run()
model.get_outputs_df().plot()

Note that, by default, 50% of the births, which are based on total population, are born into the urban/rural stratum respectively. This isn’t physically realistic but we’ll ignore it for simplicity’s sake.

Stratified model with migration

Now we can add a transition flow where 2% of the rural population migrates to the urban compartment per year.

from summer2 import Stratification

model = build_model()

# Add an urban/rural stratification with an inter-location migration flow.
strat = Stratification('location', ['urban', 'rural'], ['pop'])
strat.set_population_split({'rural': 0.7, 'urban': 0.3})
model.stratify_with(strat)
model.add_transition_flow(
    'migration',
    fractional_rate=0.02,
    source='pop',
    dest='pop',
    source_strata={'location': 'rural'},
    dest_strata={'location': 'urban'},
    # Expected flow count can be used as a sanity check,
    # to assert that the expected number of flows was added.
    expected_flow_count=1
)

model.run()
model.get_outputs_df().plot()

Stratified model with age-based migration

We can take this example one step further with the observation that:

• people aged 0-19 are unlikely to migrate

• people aged 20-39 are likely to migrate

• people aged 40+ are less likely to migrate

We can use an age stratification to model the age strata and ageing flows.

from summer2 import Stratification, AgeStratification, Overwrite

model = build_model()

# Add an urban/rural stratification with an inter-location migration flow.
strat = Stratification('location', ['urban', 'rural'], ['pop'])
strat.set_population_split({'rural': 0.7, 'urban': 0.3})
model.stratify_with(strat)
model.add_transition_flow(
    'migration',
    fractional_rate=0,  # To be overwritten
    source='pop',
    dest='pop',
    source_strata={'location': 'rural'},
    dest_strata={'location': 'urban'},
    # Expected flow count can be used as a sanity check,
    # to assert that the expected number of flows was added.
    expected_flow_count=1
)

# Set age-specific migration rates.
age_strat = AgeStratification('age', [0, 20, 40], ['pop'])
age_strat.set_population_split({'0': 0.2, '20': 0.4, '40': 0.4})
age_strat.set_flow_adjustments("migration", {
    '0': Overwrite(0),  # No migration
    '20': Overwrite(0.05),  # 5% of 20-39 year olds per year
    '40': Overwrite(0.01),  # 1% of 40+ year olds per year
})
model.stratify_with(age_strat)

# Track urban and rural populations
model.request_output_for_compartments(
    'urban_pop',
    compartments=["pop"],
    strata={"location": "urban"}
)
model.request_output_for_compartments(
    'rural_pop',
    compartments=["pop"],
    strata={"location": "rural"}
)

model.run()
model.get_outputs_df().plot()

model.get_derived_outputs_df().plot()

Summary

Now you know how to add flows between strata after a model has been stratified.