Epimodels

Sponsored by

AI-powered epidemiological intelligence

Epimodels is a Python library for simulating and fitting mathematical epidemic models. It provides deterministic models in both continuous (ODE-based) and discrete (difference equation) time, along with a comprehensive parameter inference framework, symbolic analysis tools, and multiple ODE solver backends.

• 27 model classes across continuous, discrete and stochastic (CTMC) families (SIR, SIS, SIRS, SEIR, SEQIAHR, multi-strain, vector-borne, and more)
• Model registry -- get_model("SIR", family="continuous"), string-based lookup across all families
• Model fitting -- parameter estimation from observed data with 7 loss functions and 4 optimizers, plus a one-call model.fit(...) sugar API
• Bayesian inference -- DE-MCMC posterior sampling with normal/Poisson/negative-binomial observation models
• Uncertainty ensembles -- run many simulations with sampled parameters and get quantile bands
• Intervention scenarios -- time-bounded parameter changes (lockdowns, vaccination) with scenario comparison
• Rt estimation -- Cori/EpiEstim-style time-varying reproduction number from incidence data
• SDE models -- stochastic differential equation versions of any continuous model (JAX backend)
• Network models -- event-driven SIR/SIS on networkx graphs, adjacency dicts or matrices
• Model serialization -- save/load model specs (and traces) to JSON/YAML
• Symbolic analysis -- R0 computation, equilibrium finding, stability analysis, sensitivity analysis
• Multiple solvers -- scipy (CPU) and diffrax/JAX (GPU) backends with a unified interface
• Phase space tools -- time delay embedding, mutual information, phase portraits
• Mermaid diagrams -- auto-generated compartment flow diagrams for every model

pip install epimodels

Optional extras:

pip install epimodels[plot]       # matplotlib plotting
pip install epimodels[dataframe]  # pandas DataFrame support
pip install epimodels[jax]        # diffrax/JAX GPU solvers

from epimodels.continuous.models import SIR

model = SIR()
model([1000, 1, 0], [0, 50], 1001, {'beta': 2, 'gamma': .1})
model.plot_traces()
print(f"R0 = {model.R0}")
print(model.summary())

from epimodels.continuous.models import SIR
from epimodels.fitting import fit_model, Dataset

model = SIR()
dataset = Dataset()
dataset.add_series("I", times=[0, 1, 2, 3, 5, 7, 10], values=[1, 3, 8, 20, 50, 80, 60])

result = fit_model(
    model,
    dataset,
    params={"beta": (0.1, 5.0), "gamma": (0.01, 1.0)},
    initial_conditions=[1000, 1, 0],
    time_range=[0, 10],
)
print(result.best_params)
print(result.fitted_model.summary())

from epimodels.validation import SymbolicModel

sym = SymbolicModel()
sym.add_parameter("beta", positive=True, real=True)
sym.add_parameter("gamma", positive=True, real=True)
sym.add_variable("S", positive=True)
sym.add_variable("I", positive=True)
sym.add_variable("R", positive=True)
sym.set_total_population("N")

sym.define_ode("S", "-beta*S*I/N")
sym.define_ode("I", "beta*S*I/N - gamma*I")
sym.define_ode("R", "gamma*I")

R0 = sym.compute_R0_next_generation()
print(f"R0 = {R0}")

Epimodels supports multiple ODE solvers through a unified interface. You can choose between scipy (CPU-only) and diffrax (JAX-accelerated with GPU support) backends.

from epimodels.continuous import SIR
from epimodels.solvers import ScipySolver, DiffraxSolver

# Default scipy solver (RK45)
model = SIR()
model([999, 1, 0], [0, 100], 1000, {'beta': 0.3, 'gamma': 0.1})

# Explicit scipy solver with specific method
solver = ScipySolver(method='LSODA')
model = SIR()
model([999, 1, 0], [0, 100], 1000, {'beta': 0.3, 'gamma': 0.1}, solver=solver)

# JAX-accelerated solver (requires: pip install diffrax jax)
solver = DiffraxSolver(solver='Tsit5', rtol=1e-6, atol=1e-9)
model = SIR()
model([999, 1, 0], [0, 100], 1000, {'beta': 0.3, 'gamma': 0.1}, solver=solver)

For GPU acceleration, install the JAX backend:

# CPU only
pip install diffrax jax

# GPU (CUDA 12)
pip install diffrax "jax[cuda12]"

The epimodels.fitting module provides parameter estimation from observed epidemiological data.

For stochastic epidemic models check EpiStochModels.

Look up models by name across the continuous, discrete and stochastic families:

from epimodels import get_model, list_models

SIR = get_model("SIR", family="continuous")
model = SIR()
print(list_models())

Custom models can be registered with the @register_model decorator from epimodels.registry.

from epimodels.continuous import SIR
from epimodels.interventions import Intervention, Scenario, ScenarioComparison

model = SIR()
base = Scenario("baseline", model, params={"beta": 2.0, "gamma": 0.5},
                initial_conditions=[999, 1, 0], trange=[0, 100], totpop=1000)
lockdown = Scenario("lockdown", model, params={"beta": 2.0, "gamma": 0.5},
                    initial_conditions=[999, 1, 0], trange=[0, 100], totpop=1000,
                    interventions=[Intervention("beta", start=10, end=40, factor=0.4)])

cmp = ScenarioComparison(base, [lockdown]).run()
print(cmp.peak("I"))
cmp.plot("I")

from epimodels.continuous import SIR
from epimodels.ensembles import simulate_ensemble

model = SIR()
rng = np.random.default_rng(0)
ensemble = simulate_ensemble(
    model, n_sims=200,
    param_sampler=lambda: {"beta": rng.uniform(1.5, 2.5), "gamma": 0.5},
    initial_conditions=[999, 1, 0], trange=[0, 100], totpop=1000,
)
ensemble.quantiles([0.025, 0.5, 0.975])
ensemble.plot_band("I")

# ... or simply, using the sugar API:
result = model.fit(
    {"I": observed_I}, times=times,
    params_to_fit={"beta": (0.1, 5.0), "gamma": (0.01, 1.0)},
    total_population=10000,
    method="bayes", likelihood="poisson",
)
print(result.summary())

Model-free, real-time reproduction number estimation from incidence data (Cori et al. 2013, the EpiEstim method):

from epimodels.rt import estimate_rt

result = estimate_rt(incidence, window=7, si_mean=4.0, si_sd=2.0)
print(result.rt_mean)   # posterior mean Rt per window
result.plot()           # Rt with 95% credible band

Add demographic (square-root) noise to any continuous model via diffrax/JAX (pip install epimodels[jax]):

from epimodels.sde import SDEModel

sde = SDEModel(SIR())
sde([999, 1, 0], [0, 100], 1000, {"beta": 2.0, "gamma": 0.5},
    n_sims=50, seed=42)
sde.get_quantiles(0.95)
sde.plot_traces("I")

Event-driven SIR/SIS on contact networks (pip install epimodels[network]). Accepts networkx graphs, adjacency dicts or adjacency matrices:

import networkx as nx
from epimodels.network import NetworkSIR

G = nx.barabasi_albert_graph(1000, 3, seed=0)
model = NetworkSIR(G)
model(5, [0, 50], {"beta": 0.3, "gamma": 0.1}, n_sims=20, seed=0)
print(model.final_size().mean())   # mean attack rate
model.plot_traces("I")

from epimodels.io import save_model, load_model

model.simulate([1000, 1, 0], [0, 50], 1001, {"beta": 2, "gamma": 0.1})
save_model(model, "sir_run.json", include_traces=True)
clone = load_model("sir_run.json")   # registry-based reconstruction

Full documentation is available at epimodels.readthedocs.io.

The models and algorithms implemented in this library are described in the following works (full bibliography in the docs/references.rst documentation page):

• SIR / SIS / SEIR / SIRS models — Kermack & McKendrick (1927), doi:10.1098/rspa.1927.0118; Anderson & May (1991), Infectious Diseases of Humans (Oxford University Press); Compartmental models in epidemiology (Wikipedia)
• Discrete-time models — Finkenstädt & Grenfell (2000), doi:10.1111/1467-9876.00187
• Multi-strain dengue models (Dengue4Strain, SIR2Strain) — Bartley, Donnelly & Garnett (2002), doi:10.1016/S0035-9203(02)90371-8; Aguiar et al. (2011), doi:10.1016/j.jtbi.2011.08.043
• Climate-forced vector-borne models (SIRSEI, SEIRS_SEI) — Parham & Michael (2010), doi:10.1289/ehp.0901256
• SEQIAHR COVID-19 model — Nana-Kyere et al. (2022), doi:10.1155/2022/9491847
• Heterogeneous susceptibility SIR (NeipelHeterogeneousSIR) — Neipel et al. (2020), doi:10.1371/journal.pone.0239678
• Ebola model with ring vaccination (EbolaSEIHFRV) — Legrand et al. (2007), doi:10.1017/S0950268806007217; Henao-Restrepo et al. (2017), doi:10.1016/S0140-6736(16)32621-6
• Stochastic CTMC models / Gillespie SSA — Gillespie (1976), doi:10.1016/0021-9991(76)90041-3; Gillespie (1977), doi:10.1021/j100540a008; Allen (2017), doi:10.1016/j.idm.2017.03.001
• R0 (next-generation matrix) and sensitivity analysis — Diekmann, Heesterbeek & Metz (1990), doi:10.1007/BF00178324; van den Driessche & Watmough (2002), doi:10.1016/S0025-5564(02)00108-6; Chitnis et al. (2008), doi:10.1007/s11538-008-9299-0
• Parameter fitting — King et al. (2015), doi:10.1098/rspb.2015.0347; Huber (1964), doi:10.1214/aoms/1177703732; Byrd et al. (1995) (L-BFGS-B), doi:10.1137/0916069; Kingma & Ba (2015) (Adam); Raue et al. (2009) (profile likelihood), doi:10.1093/bioinformatics/btp358
• Phase space analysis — Takens (1981), doi:10.1007/BFb0091924; Fraser & Swinney (1986), doi:10.1103/PhysRevA.33.1134; Cao (1997), doi:10.1016/S0167-2789(97)00118-8
• ODE solvers — Dormand & Prince (1980) (RK45), doi:10.1016/0771-050X(80)90013-3; Tsitouras (2011) (Tsit5), doi:10.1016/j.apnum.2010.09.005; Hairer, Nørsett & Wanner (1993, 1996); Petzold (1983) (LSODA), doi:10.1137/0904010

MIT License - see LICENSE.txt for details.