A Python library for modeling infectious disease dynamics, based on Keeling & Rohani (2008) Modeling Infectious Diseases in Humans and Animals.
kr_epi provides deterministic ODE models, stochastic simulation engines, seasonal forcing functions, and analysis tools for studying epidemic dynamics in closed and open populations.
pip install git+https://github.com/kimon-ang/kr_epi.gitFor development:
git clone https://github.com/kimon-ang/kr_epi.git
cd kr_epi
pip install -e .- Python >= 3.10
- NumPy
- SciPy
- Matplotlib
- Pandas
from kr_epi.models.ode import SIR
model = SIR(beta=0.5, gamma=0.1)
print(f"R0 = {model.R0()}") # R0 = 5.0
t, y = model.integrate(
{"S": 0.99, "I": 0.01, "R": 0.0},
t_span=(0, 100),
)Deterministic ODE models — SI, SIS, SIR, SIRS, SEIR, SEIRS for closed populations, plus SIR/SIS/SEIR/SEIRS/MSIR variants with demography (births, deaths, vaccination) and disease-induced mortality.
Stochastic simulation — Gillespie's Direct Method (SSA) for exact simulation and adaptive Tau-leaping for approximate simulation of large populations. Both operate on a shared reaction system representation.
Seasonal forcing — Sinusoidal forcing, step-function term-time forcing, and smooth term-time forcing for modeling seasonal variation in transmission rates.
Analysis tools — Equilibrium calculations (endemic equilibria, final size relations, vaccination thresholds), Jacobian-based stability analysis, power spectral density via Welch's method, and local sensitivity analysis.
Validation framework — Automated checks for R0 threshold behavior, conservation laws, equilibrium stability, final size relations, and numerical accuracy.
Parameter sweeps — Grid-based parameter sweeps with metric collection for both deterministic and stochastic models.
from kr_epi.models.ode_counts import SIRDemographyCounts
model = SIRDemographyCounts(
beta=0.5, gamma=1/13, v=1/(70*365), mu=1/(70*365)
)
print(f"R0 = {model.R0():.2f}")
eq = model.endemic_equilibrium()
print(f"Endemic I* = {eq['Y']:.4f}")
t, y = model.integrate(
{"X": 1e6 - 10, "Y": 10, "Z": 0},
t_span=(0, 365 * 50),
)from kr_epi.models.ode import SIR
from kr_epi.forcing.sinusoid import BetaSinusoid
forcing = BetaSinusoid(beta0=0.5, amp=0.2, period=365)
model = SIR(beta=0.5, gamma=0.1)
t, y = model.integrate(
{"S": 0.99, "I": 0.01, "R": 0.0},
t_span=(0, 365 * 10),
beta_fn=forcing,
)from kr_epi.models.reactions import sir_demography_counts_reactions
from kr_epi.stochastic import Direct
from kr_epi.sweeps.runners import run_ensemble
import numpy as np
system = sir_demography_counts_reactions(
beta=0.5, gamma=1/13, v=1/(70*365), mu=1/(70*365)
)
engine = Direct(system, seed=42)
t, y = engine.run(
x0={"X": 990, "Y": 10, "Z": 0},
t_max=365,
record_times=np.arange(0, 365, 0.5),
)
# Or run an ensemble
ens = run_ensemble(
engine,
x0={"X": 990, "Y": 10, "Z": 0},
t_max=365,
n_runs=50,
record_times=np.arange(0, 365, 0.5),
)from kr_epi.models.ode import SIR
from kr_epi.validation import validate_model
model = SIR(beta=2.0, gamma=1.0)
results = validate_model(model)
# Runs: R0 threshold, conservation, equilibrium stability,
# final size relation, numerical accuracyfrom kr_epi.analysis.spectra import psd_welch, dominant_peaks
# Assuming `infected_ts` is a time series of infected counts
f, Pxx = psd_welch(infected_ts, fs=1.0) # daily sampling
freqs, powers, periods = dominant_peaks(f, Pxx)
for freq, period in zip(freqs, periods):
print(f"Period: {period:.0f} days")kr_epi/
├── models/
│ ├── _helpers.py # Shared helpers (Mixing type, _foi, _incidence)
│ ├── base.py # ODEBase class
│ ├── ode.py # Closed-population models (SI, SIS, SIR, SIRS, SEIR, SEIRS)
│ ├── ode_counts.py # Demography models (births, deaths, vaccination)
│ ├── ode_counts_mortality.py # Disease-induced mortality
│ └── reactions.py # Reaction system for stochastic engines
├── stochastic/
│ ├── direct.py # Gillespie's Direct Method
│ └── tau_leap.py # Adaptive Tau-leaping
├── forcing/
│ ├── sinusoid.py # Sinusoidal seasonal forcing
│ ├── term_time.py # Step-function term-time forcing
│ └── term_time_smooth.py # Smooth term-time forcing
├── analysis/
│ ├── equilibria.py # Equilibrium and final size calculations
│ ├── stability.py # Jacobian and eigenvalue stability
│ ├── spectra.py # Power spectral density analysis
│ └── sensitivity.py # Local sensitivity analysis
├── sweeps/
│ ├── runners.py # Ensemble runner and summary statistics
│ └── parameter_sweep.py # Grid-based parameter sweeps
├── validation/
│ └── model_validator.py # Automated model validation suite
├── plotting.py # Plotting utilities
└── utils/
└── params.py # Parameter conversion helpers
Keeling, M. J., & Rohani, P. (2008). Modeling Infectious Diseases in Humans and Animals. Princeton University Press.
MIT License. See LICENSE for details.
Kimon Anagnostopoulos