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.

[1]
import numpy as np
import matplotlib.pyplot as plt
from summer import CompartmentalModel

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


def plot_compartments(model, times=[]):
    """Plot model compartment sizes over time"""
    fig, ax = plt.subplots(1, 1, figsize=(12, 6), dpi=120)
    for i in range(model.outputs.shape[1]):
        ax.plot(model.times, model.outputs.T[i])

    for t in times:
        ax.axvline(x=t, color='k', linestyle='--', alpha=0.3)

    ax.set_title("Population")
    ax.set_xlabel("Days")
    ax.set_ylabel("Compartment size")
    start, end = ax.get_xlim()
    ax.xaxis.set_ticks(np.arange(start + 1.5, end, 5))
    ax.legend([str(c) for c in model.compartments], loc='lower right')
    plt.show()

Unstratified model

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

[2]
model = build_model()
model.run()
plot_compartments(model)
../_images/examples_11-flows-between-strata_3_0.png

Stratified model

Next lets split the population into urban and rural.

[3]
from summer 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()
plot_compartments(model)
../_images/examples_11-flows-between-strata_5_0.png

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. A function flow could be used to more plausibly balance births between locations, based on their respective populations.

Stratified model with migration

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

[4]
from summer 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()
plot_compartments(model)
../_images/examples_11-flows-between-strata_8_0.png

Stratified model witha 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.

[5]
from summer 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.add_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()
plot_compartments(model)

# Plot rural/urban split
fig, ax = plt.subplots(1, 1, figsize=(12, 6), dpi=120)
ax.plot(model.times, model.derived_outputs['urban_pop'])
ax.plot(model.times, model.derived_outputs['rural_pop'])
ax.set_title("Population")
ax.set_xlabel("Days")
ax.set_ylabel("Compartment size")
start, end = ax.get_xlim()
ax.xaxis.set_ticks(np.arange(start + 1.5, end, 5))
ax.legend(['urban_pop', 'rural_pop'], loc='lower right')
plt.show()
../_images/examples_11-flows-between-strata_10_0.png
../_images/examples_11-flows-between-strata_10_1.png

Summary

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

[ ]