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, computed_values):
"""
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, computed_values):
"""
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', split_imports=True)
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)
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)
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)
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)
