# 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)
```

## 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)
```

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

## 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()
```