xjtulyc/epi-modeling

Epidemic Compartmental Modeling Skill

This skill covers the mathematical and computational workflow for building, fitting, and analysing compartmental epidemic models in Python. Starting from classic SIR and SEIR formulations, it extends to age-structured models, hospitalisation compartments (SEIHR), real-time reproduction number estimation, and global sensitivity analysis via Partial Rank Correlation Coefficients (PRCC).

All numerical integration is done with scipy.integrate.solve_ivp using the Radau solver, which handles stiff ODE systems arising from wide parameter ranges common in outbreak modeling.

---

Setup

pip install numpy scipy pandas matplotlib numba

---

Core Functions

"""
epi_modeling.py
---------------
Compartmental epidemic model utilities: SIR, SEIR, SEIHR, Rt estimation,
parameter fitting, age-structured models, and sensitivity analysis.
"""

from __future__ import annotations

import warnings
from typing import Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from scipy.integrate import solve_ivp
from scipy.optimize import minimize, differential_evolution
from scipy.stats import pearsonr, spearmanr

warnings.filterwarnings("ignore", category=RuntimeWarning)


# ---------------------------------------------------------------------------
# 1. SEIR Model
# ---------------------------------------------------------------------------

def build_seir_model(
    beta: float,
    sigma: float,
    gamma: float,
    N: int,
    mu: float = 0.0,
) -> Callable:
    """
    Construct and return a SEIR ODE right-hand-side function.

    Model compartments
    ------------------
    S : Susceptible
    E : Exposed (infected but not yet infectious)
    I : Infectious
    R : Recovered / Removed

    Equations
    ---------
        dS/dt = mu*(N - S) - beta*S*I/N
        dE/dt = beta*S*I/N - (sigma + mu)*E
        dI/dt = sigma*E - (gamma + mu)*I
        dR/dt = gamma*I - mu*R

    Parameters
    ----------
    beta : float
        Transmission rate (contacts per day × probability of transmission).
    sigma : float
        Incubation rate (1 / mean_incubation_period_days).
    gamma : float
        Recovery rate (1 / mean_infectious_period_days).
    N : int
        Total population size (assumed constant when mu=0).
    mu : float, optional
        Birth/death rate (1 / life_expectancy_days). Default 0 (no demography).

    Returns
    -------
    Callable
        ODE function compatible with ``scipy.integrate.solve_ivp``.
        Signature: f(t, y) -> [dS, dE, dI, dR]

    Examples
    --------
    >>> seir = build_seir_model(beta=0.3, sigma=1/5.2, gamma=1/10, N=1_000_000)
    >>> sol = solve_ivp(seir, [0, 180], [999990, 5, 5, 0],
    ...                 t_eval=np.arange(181), method="Radau")
    """
    def seir_ode(t: float, y: np.ndarray) -> List[float]:
        S, E, I, R = y
        force_of_infection = beta * S * I / N
        dS = mu * (N - S) - force_of_infection
        dE = force_of_infection - (sigma + mu) * E
        dI = sigma * E - (gamma + mu) * I
        dR = gamma * I - mu * R
        return [dS, dE, dI, dR]

    seir_ode.__doc__ = (f"SEIR ODE: beta={beta:.4f}, sigma={sigma:.4f}, "
                        f"gamma={gamma:.4f}, N={N}, R0={beta/gamma:.2f}")
    return seir_ode


def build_sir_model(beta: float, gamma: float, N: int) -> Callable:
    """
    Construct and return an SIR ODE function.

    dS/dt = -beta*S*I/N
    dI/dt =  beta*S*I/N - gamma*I
    dR/dt =  gamma*I
    """
    def sir_ode(t, y):
        S, I, R = y
        foi = beta * S * I / N
        return [-foi, foi - gamma * I, gamma * I]

    return sir_ode


def build_seihr_model(
    beta: float,
    sigma: float,
    gamma: float,
    gamma_h: float,
    hosp_rate: float,
    N: int,
) -> Callable:
    """
    SEIR extended with Hospitalisation (H) and fatality-adjusted Removed (R).

    Additional compartments
    -----------------------
    H : Hospitalised (a fraction of infectious individuals)
    R : Recovered (from community) + hospital discharges

    Parameters
    ----------
    gamma_h : float
        Hospital discharge rate (1 / mean_hospital_stay_days).
    hosp_rate : float
        Fraction of infectious individuals requiring hospitalisation (0–1).
    """
    def seihr_ode(t, y):
        S, E, I, H, R = y
        foi = beta * S * I / N
        dS = -foi
        dE = foi - sigma * E
        dI = sigma * E - gamma * I
        dH = hosp_rate * gamma * I - gamma_h * H
        dR = (1 - hosp_rate) * gamma * I + gamma_h * H
        return [dS, dE, dI, dH, dR]

    return seihr_ode


# ---------------------------------------------------------------------------
# 2. Parameter Fitting
# ---------------------------------------------------------------------------

def fit_model_to_data(
    incidence: np.ndarray,
    dates: Union[np.ndarray, pd.DatetimeIndex],
    model_func: Callable,
    N: int,
    param_bounds: Dict[str, Tuple[float, float]],
    compartment_idx: int = 2,
    initial_exposed: int = 10,
    method: str = "differential_evolution",
) -> Dict:
    """
    Fit a compartmental model to observed incidence data using least-squares
    optimisation.

    Parameters
    ----------
    incidence : np.ndarray
        Daily or weekly new case counts (length T).
    dates : array-like
        Date array corresponding to incidence (length T).
    model_func : Callable
        A function that accepts (beta, ...) and returns a scipy solve_ivp ODE
        callable. Must be consistent with ``param_bounds`` keys.
    N : int
        Population size.
    param_bounds : dict
        Dictionary of {parameter_name: (lower_bound, upper_bound)}.
        Keys determine order of optimisation variables.
    compartment_idx : int
        Index of the ODE state vector corresponding to the observed count
        (e.g. 2 for I in SEIR with states [S,E,I,R]).
    initial_exposed : int
        Assumed initial exposed individuals (E0).
    method : str
        Optimisation method: 'differential_evolution' (global, robust) or
        'Nelder-Mead' (fast, local).

    Returns
    -------
    dict with keys:
        - 'params'     : fitted parameter dict
        - 'solution'   : solve_ivp solution object
        - 'fitted_incidence': np.ndarray of modelled new cases
        - 'rmse'       : root mean squared error
        - 'R0'         : basic reproduction number (beta/gamma if present)
    """
    T = len(incidence)
    t_eval = np.arange(T)

    # Initial conditions: one infectious seed
    I0 = 1
    S0 = N - I0 - initial_exposed
    param_names = list(param_bounds.keys())
    bounds_list = [param_bounds[k] for k in param_names]

    def _residuals(params_vec):
        kwargs = dict(zip(param_names, params_vec))
        kwargs["N"] = N
        try:
            ode = model_func(**kwargs)
        except TypeError:
            return 1e12

        y0_len = ode.__code__.co_varcount - 2  # rough estimate; override below
        # Determine y0 from compartment count via a probe call
        n_states = 4 if "seir" in getattr(ode, "__doc__", "").lower() else 3
        if "seihr" in getattr(model_func, "__name__", "").lower():
            n_states = 5

        y0_map = {
            3: [S0,             I0, 0],
            4: [S0, initial_exposed, I0, 0],
            5: [S0, initial_exposed, I0, 0, 0],
        }
        y0 = y0_map.get(n_states, [S0, initial_exposed, I0, 0])

        try:
            sol = solve_ivp(
                ode, [0, T - 1], y0,
                t_eval=t_eval, method="Radau",
                max_step=1.0, dense_output=False,
            )
            if not sol.success:
                return 1e12
            # Convert cumulative to incidence via daily difference
            cumulative = sol.y[compartment_idx - 1] + sol.y[compartment_idx]
            modelled = np.maximum(np.diff(cumulative, prepend=cumulative[0]), 0)
            return float(np.sqrt(np.mean((modelled - incidence) ** 2)))
        except Exception:
            return 1e12

    if method == "differential_evolution":
        result = differential_evolution(
            _residuals, bounds_list,
            maxiter=1000, tol=1e-6, seed=42,
            workers=1, polish=True,
        )
    else:
        x0 = [(lo + hi) / 2 for lo, hi in bounds_list]
        result = minimize(_residuals, x0, method="Nelder-Mead",
                          options={"maxiter": 5000, "xatol": 1e-6})

    fitted_params = dict(zip(param_names, result.x))
    fitted_params["N"] = N
    ode_fitted = model_func(**fitted_params)

    n_states_f = 4
    y0_final = [S0, initial_exposed, I0, 0]
    sol_final = solve_ivp(
        ode_fitted, [0, T - 1], y0_final,
        t_eval=t_eval, method="Radau", max_step=1.0,
    )
    cumulative_f = sol_final.y[compartment_idx - 1] + sol_final.y[compartment_idx]
    fitted_inc = np.maximum(np.diff(cumulative_f, prepend=cumulative_f[0]), 0)

    R0 = None
    if "beta" in fitted_params and "gamma" in fitted_params:
        R0 = fitted_params["beta"] / fitted_params["gamma"]

    return {
        "params": {k: v for k, v in fitted_params.items() if k != "N"},
        "solution": sol_final,
        "fitted_incidence": fitted_inc,
        "rmse": result.fun,
        "R0": R0,
    }


# ---------------------------------------------------------------------------
# 3. Rt Estimation (sliding-window exponential growth)
# ---------------------------------------------------------------------------

def estimate_rt(
    incidence: np.ndarray,
    si_mean: float,
    si_sd: float,
    window: int = 7,
    quantiles: Tuple[float, float] = (0.025, 0.975),
) -> pd.DataFrame:
    """
    Estimate the time-varying reproduction number Rt using the Wallinga-Lipsitch
    exponential growth method with a discretised serial interval distribution.

    Parameters
    ----------
    incidence : np.ndarray
        Daily new case counts (length T). Should be smoothed (7-day rolling average).
    si_mean : float
        Mean serial interval in days.
    si_sd : float
        Standard deviation of serial interval in days.
    window : int
        Rolling window size in days for growth rate estimation.
    quantiles : tuple
        Lower and upper quantile for the uncertainty interval (default 95% CI).

    Returns
    -------
    pd.DataFrame with columns: ['Rt', 'Rt_lower', 'Rt_upper', 'r', 'valid']

    Notes
    -----
    The relationship Rt = 1 / M(-r) where M is the moment generating function of
    the serial interval distribution is used (Wallinga & Lipsitch 2007).
    For a gamma-distributed SI with mean mu and sd sigma:
        Rt = (1 + r * sigma^2 / mu)^(mu^2 / sigma^2)
    """
    from scipy.stats import gamma as gamma_dist

    T = len(incidence)
    # Fit gamma distribution to serial interval
    cv2 = (si_sd / si_mean) ** 2
    k = 1.0 / cv2          # shape
    theta = si_mean * cv2  # scale

    rt_vals   = np.full(T, np.nan)
    rt_lower  = np.full(T, np.nan)
    rt_upper  = np.full(T, np.nan)
    r_vals    = np.full(T, np.nan)

    inc_smooth = pd.Series(incidence).rolling(window, center=True,
                                               min_periods=3).mean().values

    for t in range(window, T - window):
        segment = inc_smooth[t - window:t + window + 1]
        if np.any(segment <= 0) or np.any(np.isnan(segment)):
            continue
        # Log-linear regression to estimate instantaneous growth rate r
        log_inc = np.log(segment + 1)
        time_idx = np.arange(len(segment)) - window
        slope, intercept = np.polyfit(time_idx, log_inc, 1)
        r = slope

        # Wallinga-Lipsitch Rt
        rt = (1 + r * theta) ** k
        # Approximate CI via bootstrap over segment
        boot_rt = []
        rng = np.random.default_rng(t)
        for _ in range(200):
            noise = rng.normal(0, 0.05 * np.abs(r) + 1e-6, len(segment))
            log_b = log_inc + noise
            b_slope, _ = np.polyfit(time_idx, log_b, 1)
            boot_rt.append((1 + b_slope * theta) ** k)
        rt_lower[t] = max(0, np.quantile(boot_rt, quantiles[0]))
        rt_upper[t] = np.quantile(boot_rt, quantiles[1])
        rt_vals[t]  = rt
        r_vals[t]   = r

    df = pd.DataFrame({
        "Rt": rt_vals,
        "Rt_lower": rt_lower,
        "Rt_upper": rt_upper,
        "r": r_vals,
        "valid": ~np.isnan(rt_vals),
    })
    return df


# ---------------------------------------------------------------------------
# 4. Epidemic Curve Plot
# ---------------------------------------------------------------------------

def plot_epidemic_curve(
    dates: Union[pd.DatetimeIndex, np.ndarray],
    cases: np.ndarray,
    model_output: Optional[np.ndarray] = None,
    rt_df: Optional[pd.DataFrame] = None,
    title: str = "Epidemic Curve",
    output_path: Optional[str] = None,
) -> plt.Figure:
    """
    Plot observed incidence bars with optional model fit line and Rt panel.
    """
    n_panels = 2 if rt_df is not None else 1
    fig, axes = plt.subplots(n_panels, 1, figsize=(14, 5 * n_panels),
                             sharex=True)
    if n_panels == 1:
        axes = [axes]

    ax = axes[0]
    ax.bar(dates, cases, color="steelblue", alpha=0.5, label="Observed cases", width=1)
    if model_output is not None:
        ax.plot(dates, model_output, color="firebrick", linewidth=2,
                label="Model fit")
    ax.set_ylabel("Daily new cases")
    ax.set_title(title, fontsize=13)
    ax.legend()
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    ax.xaxis.set_major_locator(mdates.MonthLocator())

    if rt_df is not None:
        ax2 = axes[1]
        valid = rt_df["valid"]
        ax2.plot(np.array(dates)[valid], rt_df.loc[valid, "Rt"],
                 color="darkorange", linewidth=2, label="Rt")
        ax2.fill_between(
            np.array(dates)[valid],
            rt_df.loc[valid, "Rt_lower"],
            rt_df.loc[valid, "Rt_upper"],
            color="darkorange", alpha=0.2, label="95% CI",
        )
        ax2.axhline(1.0, color="black", linewidth=0.8, linestyle="--")
        ax2.set_ylabel("Rt")
        ax2.set_ylim(0, max(4.0, rt_df["Rt_upper"].quantile(0.99) * 1.1))
        ax2.legend()

    fig.autofmt_xdate()
    fig.tight_layout()

    if output_path:
        fig.savefig(output_path, dpi=150, bbox_inches="tight")
        print(f"Saved: {output_path}")

    return fig


# ---------------------------------------------------------------------------
# 5. Age-Structured SEIR
# ---------------------------------------------------------------------------

def age_structured_model(
    contact_matrix: np.ndarray,
    beta: float,
    sigma: float,
    gamma: float,
    N_by_age: np.ndarray,
) -> Callable:
    """
    Build an age-structured SEIR ODE system.

    Parameters
    ----------
    contact_matrix : np.ndarray
        (n_age x n_age) POLYMOD-style contact matrix. Entry C[i,j] is the
        average number of contacts individuals in age group i make with
        age group j per day.
    beta : float
        Per-contact transmission probability.
    sigma : float
        Incubation rate (scalar, shared across age groups).
    gamma : float
        Recovery rate (scalar, shared across age groups).
    N_by_age : np.ndarray
        Population size in each age group (length n_age).

    Returns
    -------
    Callable
        ODE function f(t, y) with y flattened as [S0,S1,...,E0,E1,...,I0,I1,...,R0,R1,...].
    """
    n = len(N_by_age)
    N_total = N_by_age.sum()

    def ode(t, y):
        y = np.reshape(y, (4, n))
        S, E, I, R = y

        # Force of infection for each age group
        foi = beta * (contact_matrix @ (I / N_by_age))

        dS = -foi * S
        dE =  foi * S - sigma * E
        dI =  sigma * E - gamma * I
        dR =  gamma * I
        return np.concatenate([dS, dE, dI, dR])

    return ode

---

Example 1: Fit SEIR to COVID-19 Data and Estimate Rt

"""
example_covid_seir_rt.py
-------------------------
Fit a SEIR model to a COVID-19-like incidence curve and estimate Rt
using the sliding-window exponential growth method.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from epi_modeling import (
    build_seir_model,
    fit_model_to_data,
    estimate_rt,
    plot_epidemic_curve,
)

# ---- 1. Generate synthetic "observed" COVID-like incidence ----
# In production, replace this with real data from:
#   WHO: https://covid19.who.int/data
#   CDC: https://data.cdc.gov
rng = np.random.default_rng(7)
N_pop = 5_000_000

# True parameters (to be recovered by fitting)
BETA_TRUE  = 0.28
SIGMA_TRUE = 1 / 5.1    # mean incubation 5.1 days
GAMMA_TRUE = 1 / 10.0   # mean infectious period 10 days

from scipy.integrate import solve_ivp
true_ode = build_seir_model(BETA_TRUE, SIGMA_TRUE, GAMMA_TRUE, N_pop)
T_days = 120
sol_true = solve_ivp(
    true_ode, [0, T_days - 1],
    [N_pop - 20, 10, 10, 0],
    t_eval=np.arange(T_days), method="Radau",
)
S_true, E_true, I_true, R_true = sol_true.y
new_infections_true = np.maximum(np.diff(SIGMA_TRUE * E_true, prepend=0), 0)
observed = rng.poisson(np.clip(new_infections_true, 0, None)).astype(float)

start_date = pd.Timestamp("2023-01-01")
dates = pd.date_range(start_date, periods=T_days, freq="D")

# ---- 2. Fit SEIR model ----
print("Fitting SEIR model (this may take ~30-60 seconds) ...")
fit_result = fit_model_to_data(
    incidence=observed,
    dates=dates,
    model_func=build_seir_model,
    N=N_pop,
    param_bounds={
        "beta":  (0.05, 1.0),
        "sigma": (1/14, 1/2),
        "gamma": (1/21, 1/4),
    },
    compartment_idx=2,
    initial_exposed=10,
    method="differential_evolution",
)

fp = fit_result["params"]
print(f"\n--- Fitted Parameters ---")
print(f"  beta  : {fp['beta']:.4f}  (true: {BETA_TRUE:.4f})")
print(f"  sigma : {fp['sigma']:.4f}  (true: {SIGMA_TRUE:.4f})")
print(f"  gamma : {fp['gamma']:.4f}  (true: {GAMMA_TRUE:.4f})")
print(f"  R0    : {fit_result['R0']:.2f}  (true: {BETA_TRUE/GAMMA_TRUE:.2f})")
print(f"  RMSE  : {fit_result['rmse']:.1f} cases/day")

# ---- 3. Estimate Rt ----
rt_df = estimate_rt(
    incidence=observed,
    si_mean=7.5,    # mean serial interval (days)
    si_sd=3.4,
    window=7,
)
rt_df.index = range(len(rt_df))

# ---- 4. Plot epidemic curve + Rt ----
fig = plot_epidemic_curve(
    dates=dates,
    cases=observed,
    model_output=fit_result["fitted_incidence"],
    rt_df=rt_df,
    title="COVID-19-like Outbreak: SEIR Fit and Rt Estimation",
    output_path="seir_covid_fit.png",
)
plt.show()

# ---- 5. Peak incidence statistics ----
peak_day = np.argmax(fit_result["fitted_incidence"])
print(f"\nPeak incidence (model): {fit_result['fitted_incidence'][peak_day]:.0f} cases "
      f"on {dates[peak_day].strftime('%Y-%m-%d')} (day {peak_day})")

valid_rt = rt_df[rt_df["valid"]]
rt_above_1 = (valid_rt["Rt"] > 1).sum()
print(f"Days with Rt > 1: {rt_above_1} / {len(valid_rt)}")

---

Example 2: Compare SIR vs SEIR vs SEIHR

"""
example_model_comparison.py
----------------------------
Run SIR, SEIR, and SEIHR models with identical R0 and compare
epidemic dynamics: peak timing, peak magnitude, hospitalisation burden.
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from epi_modeling import build_sir_model, build_seir_model, build_seihr_model

N_pop    = 1_000_000
R0       = 2.5
gamma    = 1 / 10.0      # recovery rate
beta     = R0 * gamma    # transmission rate
sigma    = 1 / 5.2       # incubation rate (SEIR / SEIHR only)
gamma_h  = 1 / 7.0       # hospital discharge rate (SEIHR only)
hosp_rate = 0.03         # 3% of infectious cases hospitalised (SEIHR only)

T_days = 200
t_span = [0, T_days - 1]
t_eval = np.arange(T_days)

# ---- SIR ----
sir_ode  = build_sir_model(beta, gamma, N_pop)
sol_sir  = solve_ivp(sir_ode,  t_span, [N_pop - 1, 1, 0],      t_eval=t_eval, method="Radau")

# ---- SEIR ----
seir_ode = build_seir_model(beta, sigma, gamma, N_pop)
sol_seir = solve_ivp(seir_ode, t_span, [N_pop - 11, 10, 1, 0], t_eval=t_eval, method="Radau")

# ---- SEIHR ----
seihr_ode = build_seihr_model(beta, sigma, gamma, gamma_h, hosp_rate, N_pop)
sol_seihr = solve_ivp(seihr_ode, t_span, [N_pop - 11, 10, 1, 0, 0], t_eval=t_eval, method="Radau")

# ---- Plot ----
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# Infectious curves
ax = axes[0, 0]
ax.plot(t_eval, sol_sir.y[1]  / N_pop * 100, label="SIR  (I)",   color="steelblue")
ax.plot(t_eval, sol_seir.y[2] / N_pop * 100, label="SEIR (I)",   color="darkorange")
ax.plot(t_eval, sol_seihr.y[2]/ N_pop * 100, label="SEIHR (I)",  color="firebrick")
ax.set_xlabel("Day"); ax.set_ylabel("% of population")
ax.set_title("Infectious Prevalence")
ax.legend()

# Hospitalised (SEIHR only)
ax = axes[0, 1]
ax.plot(t_eval, sol_seihr.y[3] / N_pop * 100, color="purple", linewidth=2)
ax.set_xlabel("Day"); ax.set_ylabel("% of population")
ax.set_title("Hospitalised Prevalence (SEIHR)")
ax.fill_between(t_eval, 0, sol_seihr.y[3] / N_pop * 100, color="purple", alpha=0.2)

# Susceptible depletion
ax = axes[1, 0]
ax.plot(t_eval, sol_sir.y[0]  / N_pop * 100, label="SIR",   color="steelblue")
ax.plot(t_eval, sol_seir.y[0] / N_pop * 100, label="SEIR",  color="darkorange")
ax.plot(t_eval, sol_seihr.y[0]/ N_pop * 100, label="SEIHR", color="firebrick")
ax.set_xlabel("Day"); ax.set_ylabel("% of population")
ax.set_title("Susceptible Depletion")
ax.legend()

# Summary bar chart
ax = axes[1, 1]
models  = ["SIR", "SEIR", "SEIHR"]
peak_I  = [
    sol_sir.y[1].max() / N_pop * 100,
    sol_seir.y[2].max() / N_pop * 100,
    sol_seihr.y[2].max() / N_pop * 100,
]
peak_day = [
    sol_sir.y[1].argmax(),
    sol_seir.y[2].argmax(),
    sol_seihr.y[2].argmax(),
]
final_R = [
    sol_sir.y[2][-1] / N_pop * 100,
    sol_seir.y[3][-1] / N_pop * 100,
    sol_seihr.y[4][-1] / N_pop * 100,
]
x = np.arange(len(models))
width = 0.3
ax.bar(x - width, peak_I,  width, label="Peak I (%)", color="steelblue")
ax.bar(x,         final_R, width, label="Final R (%)", color="darkorange")
ax.bar(x + width, peak_day, width, label="Peak day",   color="firebrick")
ax.set_xticks(x); ax.set_xticklabels(models)
ax.set_title("Model Comparison Summary")
ax.legend()

fig.suptitle(f"Epidemic Model Comparison — R0 = {R0}", fontsize=14)
fig.tight_layout()
fig.savefig("model_comparison.png", dpi=150, bbox_inches="tight")
print("Saved: model_comparison.png")
plt.show()

# ---- Print table ----
print("\n{:<8} {:>12} {:>12} {:>12}".format("Model", "Peak I (%)", "Peak day", "Final R (%)"))
print("-" * 46)
for name, pi, pd_, fr in zip(models, peak_I, peak_day, final_R):
    print(f"{name:<8} {pi:>12.2f} {pd_:>12} {fr:>12.2f}")

---

PRCC Sensitivity Analysis

"""
sensitivity_analysis.py
-----------------------
Partial Rank Correlation Coefficient (PRCC) global sensitivity analysis
for epidemic model outputs (e.g., peak prevalence, final attack rate).
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import spearmanr, rankdata
from scipy.integrate import solve_ivp
from epi_modeling import build_seir_model


def prcc_sensitivity(
    param_distributions: dict,
    model_func: callable,
    output_func: callable,
    N: int,
    n_samples: int = 1000,
    seed: int = 42,
) -> pd.DataFrame:
    """
    Compute PRCC indices for each parameter against a scalar model output.

    Parameters
    ----------
    param_distributions : dict
        {param_name: (low, high)} uniform distribution bounds.
    model_func : callable
        Function(params_dict, N) -> solve_ivp solution.
    output_func : callable
        Function(solution) -> scalar output of interest.
    N : int
        Population size.
    n_samples : int
        LHS sample size.
    seed : int

    Returns
    -------
    pd.DataFrame with columns: param, prcc, p_value.
    """
    rng = np.random.default_rng(seed)
    param_names = list(param_distributions.keys())
    k = len(param_names)

    # Latin Hypercube Sampling
    lhs = np.zeros((n_samples, k))
    for j, name in enumerate(param_names):
        lo, hi = param_distributions[name]
        perm = rng.permutation(n_samples)
        u = (perm + rng.uniform(size=n_samples)) / n_samples
        lhs[:, j] = lo + u * (hi - lo)

    # Evaluate model
    outputs = np.full(n_samples, np.nan)
    for i in range(n_samples):
        params = dict(zip(param_names, lhs[i]))
        try:
            sol = model_func(params, N)
            outputs[i] = output_func(sol)
        except Exception:
            pass

    valid = ~np.isnan(outputs)
    lhs_v = lhs[valid]
    out_v = outputs[valid]

    # Rank-transform
    ranked = np.column_stack([rankdata(lhs_v[:, j]) for j in range(k)] +
                             [rankdata(out_v)])

    # Partial correlations via residuals
    prcc_vals, pvals = [], []
    for j in range(k):
        other = [jj for jj in range(k) if jj != j]
        X = ranked[:, j]
        Y = ranked[:, k]
        Z = ranked[:, other]

        # Regress X on Z
        coef_x = np.linalg.lstsq(np.column_stack([np.ones(len(X)), Z]), X, rcond=None)[0]
        res_x = X - np.column_stack([np.ones(len(X)), Z]) @ coef_x

        # Regress Y on Z
        coef_y = np.linalg.lstsq(np.column_stack([np.ones(len(Y)), Z]), Y, rcond=None)[0]
        res_y = Y - np.column_stack([np.ones(len(Y)), Z]) @ coef_y

        r, p = spearmanr(res_x, res_y)
        prcc_vals.append(r)
        pvals.append(p)

    df = pd.DataFrame({
        "param": param_names,
        "prcc": prcc_vals,
        "p_value": pvals,
    }).sort_values("prcc", key=abs, ascending=False)
    return df


if __name__ == "__main__":
    N = 1_000_000

    def run_seir(params, N):
        ode = build_seir_model(params["beta"], params["sigma"], params["gamma"], N)
        return solve_ivp(ode, [0, 200], [N - 11, 10, 1, 0],
                         t_eval=np.arange(201), method="Radau")

    def peak_prevalence(sol):
        return sol.y[2].max() / N * 100

    results = prcc_sensitivity(
        param_distributions={
            "beta":  (0.1, 0.6),
            "sigma": (1/14, 1/3),
            "gamma": (1/21, 1/3),
        },
        model_func=run_seir,
        output_func=peak_prevalence,
        N=N,
        n_samples=500,
    )

    print("\nPRCC Sensitivity Analysis — Peak Infectious Prevalence")
    print(results.to_string(index=False))

    fig, ax = plt.subplots(figsize=(7, 4))
    colors = ["firebrick" if v > 0 else "steelblue" for v in results["prcc"]]
    ax.barh(results["param"], results["prcc"], color=colors)
    ax.axvline(0, color="black", linewidth=0.8)
    ax.set_xlabel("PRCC")
    ax.set_title("Sensitivity Analysis: Peak Infectious Prevalence")
    fig.tight_layout()
    fig.savefig("prcc_sensitivity.png", dpi=150)
    print("Saved: prcc_sensitivity.png")
    plt.show()

---

Tips and Best Practices

• Solver choice: Use method="Radau" for stiff systems (large populations, wide R0 ranges). LSODA is also stiff-capable. RK45 is fast but may fail for stiff problems.
• Initial conditions: Seed with a small number of exposed individuals (E0 > 0) to allow the latent period to seed infectious cases naturally; avoids discontinuities.
• Identifiability: SEIR with only I(t) observed is structurally identifiable only when sigma is fixed from literature. Fitting all three parameters simultaneously risks converging to incorrect minima.
• Data smoothing: Apply a 7-day rolling average to raw case counts before Rt estimation to suppress weekly reporting artefacts.
• Serial interval vs generation time: The serial interval (time between symptom onsets) is observable; the generation time (time between infections) is not. For Rt estimation, use the serial interval.
• Age-structured models: Use POLYMOD contact matrices (Mossong et al., 2008) or country-specific matrices from the socialmixr R package (exportable to CSV).

---

References

• Kermack & McKendrick (1927). A contribution to the mathematical theory of epidemics. Proc. R. Soc. A, 115, 700–721.
• Wallinga & Lipsitch (2007). How generation intervals shape the relationship between growth rates and reproductive numbers. Proc. R. Soc. B, 274, 599–604.
• Cori et al. (2013). A new framework and software to estimate time-varying reproduction numbers during epidemics. Am. J. Epidemiol., 178, 1505–1512.
• Marino et al. (2008). A methodology for performing global uncertainty and sensitivity analysis in systems biology. J. Theor. Biol., 254, 178–196.