An SLIAR model for the 1957 influenza pandemic implemented using diffrax.

diffrax_SLIAR_1957_simulation.md

Consider the following compartmental SLIAR model (with births or deaths) for influenza during the 1957 pandemic:

$$ \begin{aligned} S' &= -\beta S (I + \delta A)\\ L' &= \beta S (I + \delta A) - \kappa L\\ I' &= p \kappa L - \alpha I \\ A' &= (1 - p) \kappa L - \eta A\\ R' &= \alpha I + \eta A \end{aligned} $$

where $\beta$ is the transmission rate, $\delta$ the reduced infectivity factor in asymptomatics, $\kappa$ the rate of latent depletion, $p$ the probability of becoming symptomatic, $\alpha$ the recovery rate for the infectious, and $\eta$ the recovery rate for symptomatics.

NOTE: The author encountered this model in the following resource,

@book{brauer2019mathematical,
  title={Mathematical models in epidemiology},
  author={Brauer, Fred and Castillo-Chavez, Carlos and Feng, Zhilan and others},
  volume={32},
  year={2019},
  publisher={Springer}
}

which referenced this resource in the discussion of the above model's parameter values:

@article{longini2004containing,
  title={Containing pandemic influenza with antiviral agents},
  author={Longini Jr, Ira M and Halloran, M Elizabeth and Nizam, Azhar and Yang, Yang},
  journal={American journal of epidemiology},
  volume={159},
  number={7},
  pages={623--633},
  year={2004},
  publisher={Oxford University Press}
}

The implemented model:

"""
Using diffrax to implement a basic compartmental
SLIAR (Susceptible, Latent, Infected, Asymptomatic,
Recovered) for the 1957 pandemic. 

Notes
-----
The parameters for this model were mentioned in [1], 
which referenced [2]. 

References
----------
..  [1] Brauer, Fred, Carlos Castillo-Chavez, 
    and Zhilan Feng. Mathematical models in epidemiology. 
    Vol. 32. New York: Springer, 2019.
..  [2] Longini, I.M., M.E. Halloran, A. Nizam, & Y. Yang (2004) 
    Containing pandemic influenza with antiviral agents, Am. J. 
    Epidem. 159: 623–633.
"""

import diffrax
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt


@jax.jit
def ODE(t, y, params):  # numpydoc ignore=GL08
    S, L, I, A, R = y
    beta, delta, kappa, p, alpha, eta = params
    dS = -(beta * S * (I + (delta * A)))
    dL = (beta * S * (I + (delta * A))) - (kappa * L)
    dI = (p * kappa * L) - (alpha * I)
    dA = ((1 - p) * kappa * L) - (eta * A)
    dR = (alpha * I) + (eta * A)
    return jnp.array([dS, dL, dI, dA, dR])


# probability symptomatic v. asymptomatic
p = 2 / 3

# asymptomatic and symptomatic recovery depletion rate
eta = alpha = 1 / 4.1

# reduced infectivity factor in asymptomatics
delta = 0.5

# latent depletion rate
kappa = 1 / 1.9

# population size
N = 2000

# transmission rate (contact rate x tran
beta = 0.398 / N

# initial compartment sizes and state
I0 = 12
S0 = N - I0
L0 = R0 = A0 = 0
init_state = jnp.array([S0, L0, I0, A0, R0])

# parameters
params = (beta, delta, kappa, p, alpha, eta)

# epidemic evolution
ts = list(range(1, 100, 1))
dt0 = (ts[-1] - ts[0]) / len(ts)

# retrieving solution
solution = diffrax.diffeqsolve(
    diffrax.ODETerm(ODE),
    solver=diffrax.Tsit5(),
    t0=ts[0],
    t1=ts[-1],
    dt0=dt0,
    args=params,
    y0=init_state,
    saveat=diffrax.SaveAt(ts=ts),
)