Metadata-Version: 2.4
Name: epimodels
Version: 1.4.0
Summary: Library of mathematical epidemic models for use in simulation studies and inference.
Author-email: Flávio Codeço Coelho <fccoelho@gmail.com>
License-Expression: GPL-3.0-or-later
Project-URL: Homepage, https://github.com/fccoelho/epimodels
Project-URL: Documentation, https://epimodels.readthedocs.io
Classifier: Intended Audience :: Developers
Classifier: Natural Language :: English
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE.txt
License-File: AUTHORS.rst
Requires-Dist: numpy>=2.2.0
Requires-Dist: scipy>=1.14.1
Requires-Dist: sympy>=1.13.3
Provides-Extra: plot
Requires-Dist: matplotlib>=3.10.0; extra == "plot"
Provides-Extra: dataframe
Requires-Dist: pandas>=2.0.0; extra == "dataframe"
Provides-Extra: jax
Requires-Dist: jax>=0.9.1; extra == "jax"
Requires-Dist: diffrax>=0.7.2; extra == "jax"
Provides-Extra: network
Requires-Dist: networkx>=3.0; extra == "network"
Provides-Extra: yaml
Requires-Dist: pyyaml>=6.0; extra == "yaml"
Provides-Extra: dev
Requires-Dist: cython>=3.0.11; extra == "dev"
Requires-Dist: marimo>=0.13.10; extra == "dev"
Requires-Dist: mypy>=1.13.0; extra == "dev"
Requires-Dist: nbsphinx>=0.9.6; extra == "dev"
Requires-Dist: notebook>=7.4.2; extra == "dev"
Requires-Dist: pytest>=8.3.4; extra == "dev"
Requires-Dist: pytest-cov>=6.0.0; extra == "dev"
Requires-Dist: ruff>=0.8.3; extra == "dev"
Requires-Dist: scipy-stubs>=1.17.1; extra == "dev"
Requires-Dist: sphinx>=8.1.3; extra == "dev"
Requires-Dist: sphinxcontrib-mermaid>=1.0.0; extra == "dev"
Requires-Dist: ipython>=9.11.0; extra == "dev"
Requires-Dist: jupyter>=1.1.1; extra == "dev"
Requires-Dist: jupyterlab>=4.5.5; extra == "dev"
Requires-Dist: mermaid-python>=0.1; extra == "dev"
Dynamic: license-file


# Epimodels

[![PyPI version](https://badge.fury.io/py/epimodels.svg)](https://badge.fury.io/py/epimodels)
[![Python](https://img.shields.io/pypi/pyversions/epimodels.svg)](https://pypi.org/project/epimodels/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Documentation Status](https://readthedocs.org/projects/epimodels/badge/?version=latest)](https://epimodels.readthedocs.io/en/latest/?badge=latest)
[![GitHub Actions Workflow Status](https://github.com/fccoelho/epimodels/actions/workflows/python-package.yml/badge.svg)](https://github.com/fccoelho/epimodels/actions/workflows/python-package.yml)
[![GitHub stars](https://img.shields.io/github/stars/fccoelho/epimodels.svg?style=social)](https://github.com/fccoelho/epimodels/stargazers)
[![GitHub forks](https://img.shields.io/github/forks/fccoelho/epimodels.svg?style=social)](https://github.com/fccoelho/epimodels/network/members)
[![GitHub issues](https://img.shields.io/github/issues/fccoelho/epimodels.svg)](https://github.com/fccoelho/epimodels/issues)
[![GitHub last commit](https://img.shields.io/github/last-commit/fccoelho/epimodels.svg)](https://github.com/fccoelho/epimodels/commits/master)
[![Code style: ruff](https://img.shields.io/badge/code%20style-ruff-orange.svg)](https://docs.astral.sh/ruff/)
[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://github.com/pre-commit/pre-commit)
[![EpidBot Ready](https://img.shields.io/badge/EpidBot-Ready-brightgreen?style=flat-square&logo=robot&logoColor=white)](https://epidbot.kwar-ai.com.br)

<p align="center">
  <strong>Sponsored by</strong><br>
  <a href="https://kwar-ai.com.br" target="_blank">
    <img src="docs/images/kwar-ai-logo.jpg" alt="Kwar-AI" width="200">
  </a><br>
  <em>AI-powered epidemiological intelligence</em>
</p>

---

**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

```bash
pip install epimodels
```

Optional extras:

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

## Getting Started

### Simulation

```python
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

```python
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

```python
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)

| Model | Compartments | Key Features |
|-------|-------------|--------------|
| `SIR` | S, I, R | Classic susceptible-infectious-removed |
| `SIS` | S, I | No immunity, reinfection |
| `SIRS` | S, I, R | Waning immunity |
| `SEIR` | S, E, I, R | Latent period |
| `SEQIAHR` | S, E, I, A, H, R, C, D | COVID-like with quarantine, hospitalization |
| `Dengue4Strain` | 49 compartments | 4-strain dengue with cross-immunity |
| `SIRSEI` | 7 compartments | Malaria vector-host with climate forcing |
| `SIRSEIData` | 7 compartments | Malaria with real climate data |
| `SEIRS_SEI` | 7 compartments | Vector-borne with deforestation/fire effects |
| `SIR2Strain` | 10 compartments | Two-strain SIR with cross-immunity |
| `SIR1D` | S, I | 1D reduced SIR (beta/gamma tracking) |
| `SISLogistic` | S, I | SIS with logistic population growth |
| `SIRSNonAutonomous` | S, I, R | Time-dependent parameters (callables) |
| `NeipelHeterogeneousSIR` | I, tau | Heterogeneous susceptibility |

### Discrete (Difference Equations)

| Model | Compartments | Key Features |
|-------|-------------|--------------|
| `SIR` | S, I, R | Classic discrete SIR |
| `SIS` | S, I | No immunity |
| `SIRS` | S, I, R | Waning immunity |
| `SEIR` | S, E, I, R | Latent period |
| `SEIS` | S, E, I | Exposed, no immunity |
| `SIpRpS` | S, I, R | Partial immunity |
| `SIpR` | S, I, R | Secondary infections from recovered |
| `SEIpRpS` | S, E, I, R | Exposed + partial immunity |
| `SEIpR` | S, E, I, R | Exposed + secondary infections from R |
| `Influenza` | 20 compartments | Age-structured (4 groups) |
| `SEQIAHR` | S, E, I, A, H, R, C, D | COVID-like discrete version |

## 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

```python
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

| 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:

```bash
# 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

| Loss Function | Best For |
|---------------|----------|
| `SumOfSquaredErrors` | General purpose |
| `WeightedSSE` | Variable importance weighting |
| `PoissonLikelihood` | Count data |
| `NegativeBinomialLikelihood` | Overdispersed count data |
| `NormalLikelihood` | Continuous data with noise |
| `HuberLoss` | Robust to outliers |
| `CustomLoss` | User-defined objectives |

### Optimizers

| Optimizer | Methods | Notes |
|-----------|---------|-------|
| `ScipyOptimizer` | L-BFGS-B, BFGS, Nelder-Mead, Powell, CG, differential_evolution | CPU, most methods |
| `JAXOptimizer` | Adam, SGD, RMSprop | GPU-accelerated |
| `NevergradOptimizer` | Derivative-free | No gradients needed |
| `MultiStartOptimizer` | Multi-start wrapper | Avoids local minima |

## Related Libraries

For stochastic epidemic models check [EpiStochModels](https://github.com/fccoelho/EpiStochModels).

### Model Registry

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

```python
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

```python
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

```python
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

```python
# ... 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):

```python
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]`):

```python
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:

```python
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

```python
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](https://epimodels.readthedocs.io).

## License

MIT License - see [LICENSE.txt](LICENSE.txt) for details.
