Example: Epidemic SIR with policy switching

The SIR model is a simple epidemic model that divides a population into susceptible, infected, and recovered compartments. In this example we will use it to illustrate a hybrid model where the transmission rate changes when a public-health policy switches on or off.

Let

\[ \mathbf{y} = \begin{bmatrix} S \\ I \\ R \end{bmatrix} \]

where \(S\) is the susceptible population, \(I\) is the infected population, and \(R\) is the recovered population. For a population of size \(P\), recovery rate \(\gamma\), and transmission rate \(\beta\), the standard SIR equations are

\[ \begin{align} \frac{dS}{dt} &= -\beta \frac{SI}{P}, \\ \frac{dI}{dt} &= \beta \frac{SI}{P} - \gamma I, \\ \frac{dR}{dt} &= \gamma I. \end{align} \]

To model policy switching, we will make \(\beta\) depend on a discrete model index \(N\). The open-policy mode has a larger transmission rate, and the lockdown mode has a smaller transmission rate:

\[ \beta_N = \begin{cases} \beta_\text{open}, & N = 0, \\ \beta_\text{lockdown}, & N = 1. \end{cases} \]

The policy uses hysteresis so that lockdown is not switched on and off at the same threshold. Starting in the open-policy mode, lockdown starts when the infected population rises above a high threshold \(I_\text{high}\). Once lockdown is active, it remains active until the infected population falls below a lower threshold \(I_\text{low}\):

\[ N^+ = \begin{cases} 1, & N^- = 0 \text{ and } I = I_\text{high}, \\ 0, & N^- = 1 \text{ and } I = I_\text{low}. \end{cases} \]

This can be encoded with a single mode-dependent stopping condition,

\[ \mathbf{s}(\mathbf{y}) = I - \mathbf{h}, \]

where

\[ h = \begin{bmatrix} I_\text{low} \\ I_\text{high} \end{bmatrix} \]

In DiffSL we write the two transmission rates and two thresholds as vectors indexed by N. We implement the changes in model index \(N\) by setting the model index to whichever root index is true for \(\mathbf{s}\). As for the dosing model, we use a Solution struct and solve_soln methods to stop through the different segments of the hybrid model, after each segment using the root index found (root_idx) to set the value for \(N\).

use diffsol::{
    CraneliftJitModule, MatrixCommon, OdeBuilder, OdeEquations, OdeSolverMethod,
    OdeSolverStopReason, Solution,
};
use plotly::{
    common::{AxisSide, Line, LineShape, Mode},
    layout::{Axis, Layout},
    Plot, Scatter,
};
use std::{fs, path::PathBuf};

type M = diffsol::NalgebraMat<f64>;
type CG = CraneliftJitModule;
type LS = diffsol::NalgebraLU<f64>;

fn main() {
    let model = "
        population { 1000.0 }
        gamma { 0.1 }
        beta_i {
            0.3,
            0.08,
        }
        threshold_i {
            20.0,
            100.0,
        }
        u_i {
            S = 999.0,
            I = 1.0,
            R = 0.0,
        }
        F_i {
            -beta_i[N] * S * I / population,
            beta_i[N] * S * I / population - gamma * I,
            gamma * I,
        }
        stop_i {
            I - threshold_i,
        }
    ";
    let mut problem = OdeBuilder::<M>::new()
        .build_from_diffsl::<CG>(model)
        .unwrap();

    let final_time = 300.0;
    let mut soln = Solution::new(final_time);
    let mut state = problem.bdf_state::<LS>().unwrap();
    let mut switch_times = vec![(0.0, 0)];

    while !soln.is_complete() {
        state = problem
            .bdf_solver::<LS>(state)
            .unwrap()
            .solve_soln(&mut soln)
            .unwrap()
            .into_state();
        if let Some(OdeSolverStopReason::RootFound(t_root, root_idx)) = soln.stop_reason {
            switch_times.push((t_root, root_idx));
            problem.eqn_mut().set_model_index(root_idx);
        }
    }

    let time = soln.ts.clone();
    let susceptible: Vec<_> = soln.ys.inner().row(0).into_iter().copied().collect();
    let infected: Vec<_> = soln.ys.inner().row(1).into_iter().copied().collect();
    let recovered: Vec<_> = soln.ys.inner().row(2).into_iter().copied().collect();
    let lockdown: Vec<_> = time
        .iter()
        .map(|t| {
            switch_times
                .iter()
                .rev()
                .find(|(switch_time, _)| switch_time <= t)
                .map(|(_, index)| *index as f64)
                .unwrap_or(0.0)
        })
        .collect();

    let mut plot = Plot::new();
    plot.add_trace(
        Scatter::new(time.clone(), susceptible)
            .mode(Mode::Lines)
            .name("S"),
    );
    plot.add_trace(
        Scatter::new(time.clone(), infected)
            .mode(Mode::Lines)
            .name("I"),
    );
    plot.add_trace(
        Scatter::new(time.clone(), recovered)
            .mode(Mode::Lines)
            .name("R"),
    );
    plot.add_trace(
        Scatter::new(time, lockdown)
            .mode(Mode::Lines)
            .name("lockdown")
            .line(Line::new().shape(LineShape::Hv))
            .y_axis("y2"),
    );
    plot.set_layout(
        Layout::new()
            .x_axis(Axis::new().title("t [days]"))
            .y_axis(Axis::new().title("people"))
            .y_axis2(
                Axis::new()
                    .title("lockdown")
                    .range(vec![-0.05, 1.05])
                    .overlaying("y")
                    .side(AxisSide::Right)
                    .position(1.0),
            ),
    );

    let output_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("../../book/src/primer/images/sir-policy-switching.html");
    let plot_html = plot.to_inline_html(Some("sir-policy-switching"));
    fs::write(output_path, plot_html).expect("Unable to write file");
}