Basic model introduction

This page introduces the processes for building and running a simple compartmental disease model with Summer. In the following example, we will create an SIR compartmental model for a general, unspecified emerging infectious disease spreading through a fully susceptible population. In this model there will be:

  • three compartments: susceptible (S), infected (I) and recovered (R)

  • a starting population of 1000 people, with 10 of them infected (and infectious)

  • an evaluation timespan from day zero to 20 in 0.1 day steps

  • inter-compartmental flows for infection, deaths and recovery

First, let’s look at a complete example of this model in action, and then examine the details of each step. This is the complete example model that we will be working with:

[1]
import numpy as np
import matplotlib.pyplot as plt

from summer import CompartmentalModel

# Define the model
model = CompartmentalModel(
    times=[0, 20],
    compartments=["S", "I", "R"],
    infectious_compartments=["I"],
    timestep=0.1,
)
model.set_initial_population(distribution={"S": 990, "I": 10})
model.add_infection_frequency_flow(name="infection", contact_rate=1, source="S", dest="I")
model.add_transition_flow(name="recovery", fractional_rate=1/3, source="I", dest="R")
model.add_death_flow(name="infection_death", death_rate=0.05, source="I")

# Run the model
model.run()

# Visualize the results.
subplot = {"title": "SIR Model Outputs", "xlabel": "Days", "ylabel": "Compartment size"}
fig, ax = plt.subplots(1, 1, figsize=(12, 6), dpi=120, subplot_kw=subplot)
ax.legend(["S", "I", "R"])
for compartment_idx in range(model.outputs.shape[1]):
    ax.plot(model.times, model.outputs.T[compartment_idx])

start, end = ax.get_xlim()
ax.xaxis.set_ticks(np.arange(start + 1, end, 5))
plt.show()
../_images/examples_1-basic-model_1_0.png

Now let’s inspect each step of the example in more details. To start, here’s how to create a new model: let’s import the summer library and create a new CompartmentalModel object. You can see that our model has an attribute called compartments, which contains a description of each modelled compartment.

[2]
from summer import CompartmentalModel

model = CompartmentalModel(
    times=[0, 20],
    compartments=["S", "I", "R"],
    infectious_compartments=["I"],
    timestep=0.1,
)

# View a description of the model compartments
model.compartments
[2]:
[S, I, R]

Adding a population

Initially the model compartments are all empty. Let’s add:

  • 990 people to the susceptible (S) compartment, plus

  • 10 in the infectious (I) compartment.

[3]
# Add people to the model
model.set_initial_population(distribution={"S": 990, "I": 10})

# View the initial population
model.initial_population
[3]:
array([990.,  10.,   0.])

Adding inter-compartmental flows

Now, let’s add some flows for people to transition between the compartments. These flows will define the dynamics of our infection. We will add:

  • an infection flow from S to I (using frequency-dependent transmission)

  • a recovery flow from I to R

  • an infection death flow, that depletes people from the I compartment

[4]
# Susceptible people can get infected.
model.add_infection_frequency_flow(name="infection", contact_rate=1, source="S", dest="I")

# Infectious people take 3 days, on average, to recover.
# If the model was run at this stage of construction,
# then the basic reproduction number (R0) of this infection would be 3.
model.add_transition_flow(name="recovery", fractional_rate=1/3, source="I", dest="R")

# Add an infection-specific death flow to the I compartment.
# This now slightly reduces the actual sojourn time in the I compartment
# from the original request of 3 days, and so slightly reduces R0 as well.
model.add_death_flow(name="infection_death", death_rate=0.05, source="I")

# Inspect the new flows, which we just added to the model.
model._flows
[4]:
[<InfectionFrequencyFlow 'infection' from S to I>,
 <TransitionFlow 'recovery' from I to R>,
 <DeathFlow 'infection_death' from I>]

Running the model

Now we can calculate the outputs for the model over the requested time period. The model calculates the compartment sizes by solving a system of differential equations (defined by the flows we just added) over the requested time period.

[5]
model.run()

Plot the outputs

You can get a better idea of what is going on inside the model by visualizing how the compartment sizes change over time.

[7]
import numpy as np
import matplotlib.pyplot as plt

subplot = {"title": "SIR Model Outputs", "xlabel": "Days", "ylabel": "Compartment size"}
fig, ax = plt.subplots(1, 1, figsize=(12, 6), dpi=120, subplot_kw=subplot)
ax.legend(["S", "I", "R"])
for compartment_idx in range(model.outputs.shape[1]):
    ax.plot(model.times, model.outputs.T[compartment_idx])

start, end = ax.get_xlim()
ax.xaxis.set_ticks(np.arange(start + 1, end, 5))
plt.show()
../_images/examples_1-basic-model_13_0.png

Summary

That’s it for now, now you know how to:

  • Create a model

  • Add a population

  • Add flows

  • Run the model

  • Access and visualize the outputs

A detailed API reference for the CompartmentalModel class can be found here

Bonus: how the model works inside

This section presents a code snippet that shows an approximation of what is happening inside the model we just built and ran.

In the example code below we use the Euler method to solve an ordinary differential equation (ODE) which is defined by the model’s flows. We don’t actually use Euler in Summer though, you can read more about the actual ODE solvers available to evaluate models here.

[8]
import numpy as np
import matplotlib.pyplot as plt

TIMESTEP = 0.1
START_TIME = 0
END_TIME = 20

# Get times
time_period = END_TIME - START_TIME + 1
num_steps = time_period / TIMESTEP
times = np.linspace(START_TIME, END_TIME, num=int(num_steps))

# Define initial conditions
initial_conditions = np.array([990.0, 10.0, 0.0])  # S, I, R

# Define outputs
outputs = np.zeros((int(num_steps), 3))
outputs[0] = initial_conditions

# Model parameters
contact_rate = 1.0
sojourn_time = 3.0
death_rate = 0.05

# Calculate outputs for each timestep
for t_idx, t in enumerate(times):
    if t_idx == 0:
        continue

    flow_rates = np.zeros(3)
    compartment_sizes = outputs[t_idx - 1 ]

    # Susceptible people can get infected (frequency-dependent).
    num_sus = compartment_sizes[0]
    num_inf = compartment_sizes[1]
    num_pop = compartment_sizes.sum()
    force_of_infection = contact_rate * num_inf / num_pop
    infection_flow_rate = force_of_infection * num_sus
    flow_rates[0] -= infection_flow_rate
    flow_rates[1] += infection_flow_rate

    # Infectious take some time to recover.
    num_inf = compartment_sizes[1]
    recovery_flow_rate = num_inf / sojourn_time
    flow_rates[1] -= recovery_flow_rate
    flow_rates[2] += recovery_flow_rate

    # Add an infection-specific death flow to the I compartment.
    num_inf = compartment_sizes[1]
    recovery_flow_rate = num_inf * death_rate
    flow_rates[1] -= recovery_flow_rate

    # Calculate compartment sizes at next timestep given flowrates.
    outputs[t_idx] = compartment_sizes + flow_rates * TIMESTEP

# Plot the results as a function of time for S, I, R respectively.
fig, ax = plt.subplots(1, 1, figsize=(12, 6), dpi=120)

# Add each compartment to the plot.
for i in range(outputs.shape[1]):
    ax.plot(times, outputs.T[i])

ax.set_title("SIR Model Outputs")
ax.set_xlabel("Days")
ax.set_ylabel("Compartment size")
ax.legend(["S", "I", "R"])
start, end = ax.get_xlim()
ax.xaxis.set_ticks(np.arange(start + 1, end, 5))
plt.show()
../_images/examples_1-basic-model_16_0.png
[ ]