Viral Spread

This notebook builds a susceptible-infected-recovered (SIR) model as an AutoReduce System and simulates the full ODE model. The example is intentionally lightweight so it can be used as a documentation smoke test.

SIR model

The model is

\[\begin{split}\begin{aligned} \dot{S} &= -\beta\frac{SI}{N}, \\ \dot{I} &= \beta\frac{SI}{N} - \gamma I, \\ \dot{R} &= \gamma I. \end{aligned}\end{split}\]

The total population \(N = S + I + R\) is conserved by the dynamics.

import numpy as np
from sympy import Symbol

from autoreduce import System, solve_ode

S = Symbol("S")
I = Symbol("I")
R = Symbol("R")
beta = Symbol("beta")
gamma = Symbol("gamma")
N = Symbol("N")

x = [S, I, R]
f = [
    -beta * S * I / N,
    beta * S * I / N - gamma * I,
    gamma * I,
]

population = 1000.0
system = System(
    x,
    f,
    params=[beta, gamma, N],
    params_values=[0.3, 0.08, population],
    x_init=[990.0, 10.0, 0.0],
    C=np.eye(3),
)

Simulate the full model

The ODE solver returns one row per time point and one column per state.

timepoints = np.linspace(0.0, 160.0, 41)
solution = solve_ode(system, timepoints)

solution.shape

Check population conservation

The total population should remain constant up to numerical integration error.

total_population = solution.sum(axis=1)
max_conservation_error = np.max(np.abs(total_population - population))

final_state = dict(zip([str(state) for state in system.x], solution[-1]))
max_conservation_error, final_state