epimodels 1.0.0

Library of mathematical epidemic models for use in simulation studies and inference.

Epimodels

This library a simple interface to simulate mathematical epidemic models.

Getting started

Simple SIR simulation

from epimodels.continuous.models import SIR

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

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

Backend	Class	Methods	Best For
scipy	ScipySolver	RK45, RK23, DOP853, Radau, BDF, LSODA	General use, CPU-bound
diffrax	DiffraxSolver	Tsit5, Dopri5, Dopri8, Euler, Heun, Midpoint, Ralston	GPU acceleration, batch simulations

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)

Performance Benchmarks

Benchmarks run on SIR model with N=1,000,000, t=[0,365], β=0.4, γ=0.1.

Scipy Methods

Method	Time (ms)	Accuracy*	Stiff Handling	Notes
LSODA	2.4	Good	Excellent	Auto stiffness detection
RK23	6.5	Good	Poor	Fastest explicit method
DOP853	4.9	Excellent	Poor	Highest accuracy
RK45	48.3	Good	Poor	Default, robust
Radau	23.5	Excellent	Excellent	Implicit, for stiff systems
BDF	31.5	Good	Excellent	Implicit multi-step

*Accuracy measured as deviation from DOP853 reference solution.

Diffrax Methods (JAX)

Method	CPU Time	GPU Time*	Notes
Tsit5	~2x faster than scipy	10-50x faster	Recommended default
Dopri5	Similar to Tsit5	10-50x faster	Classic Dormand-Prince
Dopri8	Slower	5-20x faster	High accuracy

*GPU speedup observed on batch simulations (100+ concurrent models)

When to Use Each Solver

Scenario	Recommended Solver	Reason
General use	ScipySolver('LSODA')	Fast, handles stiffness automatically
High accuracy needed	ScipySolver('DOP853')	8th order method
Stiff systems	ScipySolver('BDF') or ScipySolver('Radau')	Implicit methods
Batch simulations	DiffraxSolver('Tsit5')	GPU parallelization
Parameter sweeps	DiffraxSolver	JAX JIT compilation
Quick prototyping	Default (RK45)	Robust and reliable

Installing Diffrax

For GPU acceleration, install the JAX backend:

# CPU only
pip install diffrax jax

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

# GPU (CUDA 11)
pip install diffrax "jax[cuda11]"

Related libraries

For stochastic epidemic models check this.