Metadata-Version: 2.4
Name: epimodels
Version: 1.0.2
Summary: Library of mathematical epidemic models for use in simulation studies and inference.
Author-email: Flávio Codeço Coelho <fccoelho@gmail.com>
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.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE.txt
License-File: AUTHORS.rst
Requires-Dist: diffrax>=0.7.2
Requires-Dist: ipykernel>=7.2.0
Requires-Dist: jax>=0.9.1
Requires-Dist: matplotlib>=3.10.0
Requires-Dist: numpy>=2.2.0
Requires-Dist: pandas>=3.0.1
Requires-Dist: scipy>=1.14.1
Requires-Dist: scipy-stubs~=1.17.1
Requires-Dist: sympy>=1.13.3
Provides-Extra: dataframe
Requires-Dist: pandas>=2.0.0; extra == "dataframe"
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: 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"
Requires-Dist: lab>=8.9; 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)

This library a simple interface to simulate mathematical epidemic models.


## Getting started

Simple SIR simulation

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

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

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

```bash
# 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](https://github.com/fccoelho/EpiStochModels).
