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]
from summer2 import CompartmentalModel

import pandas as pd
pd.options.plotting.backend = "plotly"
[2]
# 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()

# Plot the outputs
model.get_outputs_df().plot()

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 summer2 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.

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

# View a description of the model compartments
model.compartments
[3]:
[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.

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

# View the initial population

# This will be performed automatically when a model is run,
# but must be run manually to inspect the value interactively
from summer2.population import calculate_initial_population

calculate_initial_population(model)
[4]:
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

[5]
# Susceptible people can get infected.
model.add_infection_frequency_flow(name="infection", contact_rate=1.0, 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
[5]:
[<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.

[6]
model.run()

View the model outputs

The model’s results are available as a Pandas DataFrame, via the get_outputs_df() method This is available after the model has been run. Let’s have a look at what’s inside:

[7]
model.get_outputs_df()
[7]:
S I R
0.0 990.000000 10.000000 0.000000
0.1 988.979868 10.624938 0.343647
0.2 987.897131 11.287804 0.708753
0.3 986.748142 11.990752 1.096614
0.4 985.529061 12.736043 1.508605
... ... ... ...
19.6 82.061414 20.386139 780.480388
19.7 81.874865 19.802453 781.150158
19.8 81.694045 19.235097 781.800746
19.9 81.518771 18.683640 782.432686
20.0 81.348868 18.147658 783.046500

201 rows × 3 columns

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.

[8]
model.get_outputs_df().plot(title = "SIR Model Outputs")

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. Euler’s method is easy to reason about, but inappropriate for most problems; summer2 uses an adaptive Runge-Kutta method as its default solver.

[9]
import numpy as np

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.
pd.DataFrame(outputs,columns=["S","I","R"]).plot()