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.

Features

• 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

Installation

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

Getting Started

Simulation

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())

Parameter Fitting

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())

Symbolic Analysis

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}")

Available Models

Continuous (ODE)

Discrete (Difference Equations)

Solvers

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

Available Solvers

Usage Examples

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)

When to Use Each Solver

Installing Diffrax

For GPU acceleration, install the JAX backend:

# CPU only
pip install diffrax jax

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

Model Fitting

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

Loss Functions

Optimizers

Related Libraries

For stochastic epidemic models check EpiStochModels.

Model Registry

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.

Intervention Scenarios

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")

Uncertainty Ensembles

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")

Bayesian Inference

# ... 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())

Rt Estimation

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

Stochastic Differential Equations

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")

Network Models

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")

Saving and Loading Models

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

Documentation

Full documentation is available at epimodels.readthedocs.io.

License

MIT License - see LICENSE.txt for details.