ODE solvers

Summer’s compartmental models produce results by solving a system of ordinary differential equations (ODE) defined by the inter-compartmental flows. You can see an explicit example of how this kind of thing is calculated here.

This document will give you an overview of the ODE solvers available when running Summer’s compartmental models. To help demonstrate the strengths and weaknesses of the available solvers, we will look at an example SIR model where the model:

  • Has three compartments (S, I, R)

  • Runs from 0 to 20 days

  • Has three flows

    • Infection (S -> I) with a constant contact rate

    • Recovery (I -> R) with a period of higher recovery rates from days 5 to 10

    • Import of infected people (into I), occuring as four transient 1-day events during 10, 14, 16 and 19

The purpose of this example model is to show you how the different solvers handle transient events. The model and ploting code is defined below:

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

RECOVERY_TIMES = [5, 10]
IMPORT_TIMES = [10.9, 14.3, 16.7, 19.4]
TIMES = set(RECOVERY_TIMES + IMPORT_TIMES)

def build_model():
    """Returns a new SIR 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})
    # Add a infection flow
    model.add_infection_frequency_flow("infection", contact_rate=1, source="S", dest="I")

    # Add a recovery flow.
    def recovery_rate(time):
        """
        Returns the recovery rate for a given time.
        People recover faster after day 5 due to a magic drug
        """
        start, end = RECOVERY_TIMES
        if time < start or time > end:
            return 0.1
        else:
            return 0.6

    model.add_transition_flow("recovery", recovery_rate, "I", "R")

    # Add an import flow.
    def get_infected_imports(time):
        """
        Returns the number of infected people imported at a given timestep.
        Import 100 people per day during each import event.
        """
        if any([t < time < t + 1 for  t in IMPORT_TIMES]):
            return 200
        else:
            return 0


    model.add_importation_flow('infected_imports', get_infected_imports, 'I')

    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("SIR Model Outputs")
    ax.set_xlabel("Days")
    ax.set_ylabel("Compartment size")
    start, end = ax.get_xlim()
    ax.xaxis.set_ticks(np.arange(start + 1, end, 5))
    ax.legend(["S", "I", "R"])
    plt.show()

Default ODE solver (solve_ivp)

By default, when you call model.run(), Summer will use SciPy’s solve_ivp function (docs). This method defaults to using an adaptive Runge-Kutta solver, which tends to run the fastest out of the currently available deterministic options. There are two potential problems with using this solver:

  • #1 The number of iterations required (and rumtime) can depend on the model dynamics. Changing the model can affect the runtime in unpredictable ways, because the step size can become very small.

  • #2 The adaptive step size of this solver can completely miss transient events. For example, in the below plot, note how the day 14 import event was not picked up by the solver, because the system wasn’t evaluated at the times necessary for this to have an effect.

[2]
model = build_model()
# Run model with solve_ivp, with default arguments.
model.run()
# Also equivalent:
#     model.run('solve_ivp')
#     model.run(solver='solve_ivp')
plot_compartments(model, times=TIMES)
../_images/examples_4-ode-solvers_3_0.png

Default ODE solver (solve_ivp) with additional arguments

You can pass extra arguments to the solver function to adjust its behavior. For example, we can force the solve_ivp function to use a maximum step size of 0.1 so that it does not miss any transient import events (docs). By default, Summer uses a maximum step size of 1.0 for the default ODE solver (solve_ivp). The downside of using a smaller maximum step size is that the solver will need to evaluate the model at more timesteps and will run more slowly.

[3]
model = build_model()
# Run model with solve_ivp, with a custom argument.
model.run(max_step=0.1)
plot_compartments(model, times=TIMES)
../_images/examples_4-ode-solvers_5_0.png

Fixed-step Runge-Kutta 4 solver

The model can also be evaluated with a hand-rolled RK4 solver which runs with a fixed step size. It tends to be slower (sometimes much slower) than the using SciPy’s solve_ivp, but it is guaranteed to evaluate every time step. It can be useful for debugging, because it does not use an adaptive step size.

[4]
model = build_model()
model.run("rk4", step_size=0.1)
plot_compartments(model, times=TIMES)
../_images/examples_4-ode-solvers_7_0.png

Of course, it is still important to choose an appropriate step size. If your step size is too large, you will miss some events.

[5]
model = build_model()
model.run("rk4", step_size=2)
plot_compartments(model, times=TIMES)
../_images/examples_4-ode-solvers_9_0.png

Summary

That’s it for now. Now you know how to use different ODE solvers available in Summer. Please be aware they each solver has limitations that should be interrogated.

A detailed API reference of the solvers available can be found here.

[ ]