cds2.epidemiology

Compartmental epidemic models: SIR/SEIR simulation and analysis.

Compartmental epidemic models: SIR/SEIR trajectories and threshold quantities.

SIRResult dataclass

Daily SIR trajectory produced by :func:simulate_sir.

Source code in src/cds2/epidemiology.py

@dataclass(frozen=True)
class SIRResult:
    """Daily SIR trajectory produced by :func:`simulate_sir`."""

    times: FloatArray
    susceptible: FloatArray
    infected: FloatArray
    recovered: FloatArray
    population: float
    beta: float
    gamma: float

    @property
    def r0(self) -> float:
        """Basic reproduction number ``beta / gamma``."""
        return (
            (0.0 if self.beta == 0.0 else math.inf) if self.gamma == 0.0 else self.beta / self.gamma
        )

    @property
    def peak_day(self) -> float:
        """Day on which the infected curve attains its maximum."""
        return float(self.times[int(np.argmax(self.infected))])

    @property
    def attack_rate(self) -> float:
        """Fraction of the population ultimately recovered."""
        return float(self.recovered[-1] / self.population)

r0 property

Basic reproduction number beta / gamma.

peak_day property

Day on which the infected curve attains its maximum.

attack_rate property

Fraction of the population ultimately recovered.

SEIRResult dataclass

Daily SEIR trajectory produced by :func:simulate_seir.

Source code in src/cds2/epidemiology.py

@dataclass(frozen=True)
class SEIRResult:
    """Daily SEIR trajectory produced by :func:`simulate_seir`."""

    times: FloatArray
    susceptible: FloatArray
    exposed: FloatArray
    infected: FloatArray
    recovered: FloatArray
    population: float
    beta: float
    sigma: float
    gamma: float

    @property
    def r0(self) -> float:
        """Basic reproduction number ``beta / gamma``."""
        return (
            (0.0 if self.beta == 0.0 else math.inf) if self.gamma == 0.0 else self.beta / self.gamma
        )

    @property
    def peak_day(self) -> float:
        """Day on which the infected curve attains its maximum."""
        return float(self.times[int(np.argmax(self.infected))])

    @property
    def attack_rate(self) -> float:
        """Fraction of the population ultimately recovered."""
        return float(self.recovered[-1] / self.population)

r0 property

Basic reproduction number beta / gamma.

peak_day property

Day on which the infected curve attains its maximum.

attack_rate property

Fraction of the population ultimately recovered.

simulate_sir(population, beta, gamma, days, i0=1.0, steps_per_day=10)

Integrate the SIR model with RK4 and sample the state once per day.

The system is dS/dt = -beta*S*I/N, dI/dt = beta*S*I/N - gamma*I and dR/dt = gamma*I; times runs from 0 to days inclusive.

Source code in src/cds2/epidemiology.py

def simulate_sir(
    population: float,
    beta: float,
    gamma: float,
    days: int,
    i0: float = 1.0,
    steps_per_day: int = 10,
) -> SIRResult:
    """Integrate the SIR model with RK4 and sample the state once per day.

    The system is ``dS/dt = -beta*S*I/N``, ``dI/dt = beta*S*I/N - gamma*I``
    and ``dR/dt = gamma*I``; ``times`` runs from ``0`` to ``days`` inclusive.
    """
    _validate_common(population, days, steps_per_day, beta, gamma)
    if not np.isfinite(i0) or not 0.0 <= i0 <= population:
        msg = "initial infections must lie within the population"
        raise ValueError(msg)

    def deriv(y: FloatArray) -> FloatArray:
        transmission = beta * y[0] * y[1] / population
        return np.array(
            [-transmission, transmission - gamma * y[1], gamma * y[1]],
            dtype=np.float64,
        )

    times = np.linspace(0.0, float(days), days + 1)
    susceptible = np.empty(days + 1, dtype=np.float64)
    infected = np.empty(days + 1, dtype=np.float64)
    recovered = np.empty(days + 1, dtype=np.float64)
    state: FloatArray = np.array([population - i0, i0, 0.0], dtype=np.float64)
    susceptible[0] = state[0]
    infected[0] = state[1]
    recovered[0] = state[2]

    for day in range(1, days + 1):
        state = _rk4_integrate(state, deriv, 1.0 / steps_per_day, steps_per_day)
        susceptible[day] = state[0]
        infected[day] = state[1]
        recovered[day] = state[2]

    return SIRResult(
        times=times,
        susceptible=susceptible,
        infected=infected,
        recovered=recovered,
        population=population,
        beta=beta,
        gamma=gamma,
    )

simulate_seir(population, beta, sigma, gamma, days, i0=1.0, e0=0.0, steps_per_day=10)

Integrate the SEIR model with RK4 and sample the state once per day.

The system is dS/dt = -beta*S*I/N, dE/dt = beta*S*I/N - sigma*E, dI/dt = sigma*E - gamma*I and dR/dt = gamma*I; times runs from 0 to days inclusive.

Source code in src/cds2/epidemiology.py

def simulate_seir(
    population: float,
    beta: float,
    sigma: float,
    gamma: float,
    days: int,
    i0: float = 1.0,
    e0: float = 0.0,
    steps_per_day: int = 10,
) -> SEIRResult:
    """Integrate the SEIR model with RK4 and sample the state once per day.

    The system is ``dS/dt = -beta*S*I/N``, ``dE/dt = beta*S*I/N - sigma*E``,
    ``dI/dt = sigma*E - gamma*I`` and ``dR/dt = gamma*I``; ``times`` runs from
    ``0`` to ``days`` inclusive.
    """
    _validate_common(population, days, steps_per_day, beta, gamma)
    if not np.isfinite(sigma) or sigma <= 0.0:
        msg = "sigma must be positive"
        raise ValueError(msg)
    if not np.isfinite(i0) or not 0.0 <= i0 <= population:
        msg = "initial infections must lie within the population"
        raise ValueError(msg)
    if not np.isfinite(e0) or e0 < 0.0:
        msg = "initial exposures must be non-negative"
        raise ValueError(msg)
    if e0 + i0 > population:
        msg = "initial exposed and infected must fit within the population"
        raise ValueError(msg)

    def deriv(y: FloatArray) -> FloatArray:
        transmission = beta * y[0] * y[2] / population
        onset = sigma * y[1]
        return np.array(
            [
                -transmission,
                transmission - onset,
                onset - gamma * y[2],
                gamma * y[2],
            ],
            dtype=np.float64,
        )

    times = np.linspace(0.0, float(days), days + 1)
    susceptible = np.empty(days + 1, dtype=np.float64)
    exposed = np.empty(days + 1, dtype=np.float64)
    infected = np.empty(days + 1, dtype=np.float64)
    recovered = np.empty(days + 1, dtype=np.float64)
    state: FloatArray = np.array([population - i0 - e0, e0, i0, 0.0], dtype=np.float64)
    susceptible[0] = state[0]
    exposed[0] = state[1]
    infected[0] = state[2]
    recovered[0] = state[3]

    for day in range(1, days + 1):
        state = _rk4_integrate(state, deriv, 1.0 / steps_per_day, steps_per_day)
        susceptible[day] = state[0]
        exposed[day] = state[1]
        infected[day] = state[2]
        recovered[day] = state[3]

    return SEIRResult(
        times=times,
        susceptible=susceptible,
        exposed=exposed,
        infected=infected,
        recovered=recovered,
        population=population,
        beta=beta,
        sigma=sigma,
        gamma=gamma,
    )

herd_immunity_threshold(r0)

Critical immune fraction; zero when transmission is already subcritical.

Source code in src/cds2/epidemiology.py

def herd_immunity_threshold(r0: float) -> float:
    """Critical immune fraction; zero when transmission is already subcritical."""
    if not np.isfinite(r0) or r0 < 0.0:
        raise ValueError("r0 must be non-negative and finite")
    if r0 <= 1.0:
        return 0.0
    return float(1.0 - 1.0 / r0)

effective_reproduction(r0_value, susceptible_fraction)

Effective reproduction number r0 * susceptible_fraction.

Source code in src/cds2/epidemiology.py

def effective_reproduction(r0_value: float, susceptible_fraction: float) -> float:
    """Effective reproduction number ``r0 * susceptible_fraction``."""
    if not np.isfinite(r0_value) or r0_value < 0.0:
        raise ValueError("r0_value must be non-negative and finite")
    if not np.isfinite(susceptible_fraction) or not 0.0 <= susceptible_fraction <= 1.0:
        raise ValueError("susceptible_fraction must be in [0, 1]")
    return float(r0_value * susceptible_fraction)

final_size_iteration(r0, tol=1e-10, max_iter=200)

Final epidemic size solving z = 1 - exp(-r0*z) for the nonzero root.

Source code in src/cds2/epidemiology.py

def final_size_iteration(r0: float, tol: float = 1e-10, max_iter: int = 200) -> float:
    """Final epidemic size solving ``z = 1 - exp(-r0*z)`` for the nonzero root."""
    if not np.isfinite(r0) or r0 < 0.0:
        raise ValueError("r0 must be non-negative and finite")
    if not np.isfinite(tol) or tol <= 0.0:
        raise ValueError("tol must be positive and finite")
    if not isinstance(max_iter, (int, np.integer)) or isinstance(max_iter, bool) or max_iter < 1:
        raise ValueError("max_iter must be a positive integer")
    if r0 <= 1.0:
        return 0.0

    def objective(z: float) -> float:
        return z - (1.0 - math.exp(-r0 * z))

    epsilon = min(1e-8, 0.1 * (r0 - 1.0) / r0)
    try:
        return float(
            __import__("scipy").optimize.brentq(
                objective,
                epsilon,
                1.0 - np.finfo(float).eps,
                xtol=tol,
                rtol=max(4.0 * np.finfo(float).eps, tol),
                maxiter=max_iter,
            )
        )
    except RuntimeError as exc:
        raise RuntimeError("final-size iteration did not converge") from exc