Epidemic Modeling with Python

Documentation | Paper

EpiLearn is a Python machine learning toolkit for epidemic data modeling and analysis. We provide numerous features including:

• Implementation of Epidemic Models — 65 models spanning mechanistic, statistical, deep temporal, spatiotemporal graph, and time-series foundation models
• Unified Pipeline for Epidemic Tasks — forecasting, nowcasting, scenario modeling and source detection behind one interface
• Rolling-Window Evaluation with Uncertainty — walk-forward folds and conformal prediction intervals for every task
• Config-Driven Benchmark — compare any set of models from a single YAML file
• Simulation of Epidemic Spreading
• Visualization of Epidemic Data

For more machine models in epidemic modeling, feel free to check out our curated paper list Awesome-Epidemic-Modeling-Papers.

Encounter Any Issues?

If you experience any issues, please don’t hesitate to open a GitHub Issue. We will do our best to address it within three business days. You are also warmly invited to join our User Slack Channel for more efficient communication. Alternatively, reaching out to us via email is also perfectly fine!

Installation

From Source

git clone https://github.com/Emory-Melody/EpiLearn.git
cd EpiLearn

conda create -n epilearn python=3.10
conda activate epilearn

pip install .

This installs everything EpiLearn needs, including torch and torch_geometric. A source checkout is also what gives you configs/, datasets/, examples/ and tests/ — the wheel contains only the epilearn package.

From Pypi

pip install epilearn

Note: PyPI currently serves 0.0.19. The snippets below use the 0.1.0 API, so until the 0.1.0 upload lands please install from source, or straight from git: pip install git+https://github.com/Emory-Melody/EpiLearn.git

Upgrading from 0.0.x? UniversalDataset is now Dataset (the old name still works) and train_model now takes explicit splits — see MIGRATION.md and CHANGELOG.md.

For a CUDA build of PyTorch, install it first following pytorch.org, then pip install . — pip will keep the build you already have.

Tutorial

A full tutorial lives in our documentation: quickstart, pipelines, simulations, utilities and the benchmark. For the overall framework of EpiLearn, please check our paper.

Runnable code lives in two places: the examples folder (notebooks and scripts, one per topic) and the tests folder (short per-model and per-task demo scripts — python tests/forecast.py, tests/nowcast.py, tests/scenario.py, …). Run every demo at once with python tests/run_all_demos.py.

Below we also offer a quick start on how to use EpiLearn for forecast, detection and nowcasting tasks.

Forecast Pipeline

from epilearn.models.SpatialTemporal.STGCN import STGCN
from epilearn.data import Dataset
from epilearn.utils import transforms
from epilearn.tasks.forecast import Forecast
# initialize settings
lookback = 12 # inputs size
horizon = 3 # predicts size
# load toy dataset
dataset = Dataset()
dataset.load_toy_dataset()
# Adding Transformations
transformation = transforms.Compose({
                "features": [transforms.normalize_feat()],
                "target": [transforms.normalize_target()],
                "graph": [transforms.normalize_adj()]})
dataset.transforms = transformation
# Initialize Task
task = Forecast(prototype=STGCN,
                dataset=None, 
                lookback=lookback, 
                horizon=horizon, 
                device='cpu')
# Training: rolling-window evaluation, with conformal intervals per fold
result = task.rolling_train(dataset=dataset,
                            train_size=400,
                            val_size=50,
                            test_size=50,
                            train_loss='mse',
                            epochs=10,
                            batch_size=5)
# Evaluation: rolling_train already scored every fold
print(result['aggregate_metrics'])

This trains STGCN over two rolling folds and takes well under a minute on CPU.

Detection Pipeline

from epilearn.models.Spatial.GCN import GCN
from epilearn.data import Dataset
from epilearn.utils import transforms
from epilearn.tasks.detection import Detection
# initialize settings
lookback = 1 # inputs size
horizon = 2 # number of classes
# load toy dataset
dataset = Dataset()
dataset.load_toy_dataset()
dataset.y = (dataset.y > dataset.y.median(dim=1, keepdim=True).values).long() # per-node class labels
# Adding Transformations
transformation = transforms.Compose({
                "features": [transforms.normalize_feat()],
                "graph": [transforms.normalize_adj()]})
dataset.set_transforms(transformation, apply_now=True)
# Build sliding-window splits (pass adj so the graph reaches the model)
def make_split(start, end):
    return dataset.generate_dataset(X=dataset.x[start:end], Y=dataset.y[start:end],
                                    adj=dataset.graph,
                                    lookback_window_size=lookback, horizon_size=1)
train_split, val_split, test_split = make_split(0, 300), make_split(300, 400), make_split(400, 539)
# Initialize Task
task = Detection(prototype=GCN, 
                 dataset=dataset, 
                 lookback=lookback, 
                 horizon=horizon, 
                 device='cpu')
# Training
result = task.train_model(train_split=train_split,
                          val_split=val_split,
                          test_split=test_split,
                          train_loss='ce',
                          val_loss='ce',
                          epochs=50, 
                          batch_size=5)
# Evaluation
evaluation = task.evaluate_model(dataset=test_split)

Nowcasting Pipeline

Nowcasting corrects for reporting delay: recent counts are still incomplete, and the task learns how much each day will be revised upward.

import numpy as np
from epilearn.models.Temporal import GRUModel
from epilearn.tasks.nowcast import NowcastTask
# initialize settings
lookback = 14 # inputs size
horizon = 7 # nowcast the last 7 days, whose reports are still incomplete
# build a reporting triangle: day t's cases trickle in over delays 1..9
rng = np.random.default_rng(0)
delays = np.arange(1, 10)
final = np.round(100 + 60 * np.sin(np.arange(260) / 18) + rng.normal(0, 4, 260))
share = np.diff(1 - np.exp(-np.r_[0, delays] / 2.5)); share /= share.sum()
triangle = np.cumsum([rng.multinomial(int(c), share) for c in final], axis=1)
# Initialize Task
task = NowcastTask(prototype=GRUModel,
                   lookback=lookback,
                   horizon=horizon,
                   min_delay=1,
                   max_delay=9,
                   device='cpu')
dataset = task.create_dataset(triangle, final, delays=delays)
# Training
result = task.rolling_train(dataset,
                            train_size=140,
                            val_size=40,
                            test_size=40,
                            epochs=60,
                            batch_size=32,
                            lr=1e-2)
# Evaluation: compare against simply trusting the latest report
print("nowcast MAE      :", result['aggregate_metrics']['mae_mean'])
print("latest-report MAE:", task.compute_naive_baseline(dataset)['naive_mae'])

Benchmark

To compare many models under one rolling-window protocol, describe the run in a YAML file instead of writing a script:

python -m epilearn.benchmark --config configs/quick_test_config.yaml

Per-model metrics, conformal coverage, Optuna trials and raw predictions are written to benchmark_results/<task>/models_<timestamp>/. See benchmark.md for the config schema and the full model list.

Citing

If you find this work useful, please cite: EpiLearn: A Python Library for Machine Learning in Epidemic Modeling

@article{liu2024epilearn,
title={EpiLearn: A Python Library for Machine Learning in Epidemic Modeling},
author={Liu, Zewen and Li, Yunxiao and Wei, Mingyang and Wan, Guancheng and Lau, Max SY and Jin, Wei},
journal={arXiv e-prints},
pages={arXiv--2406},
year={2024}
}

Acknowledgement

Some algorithms are adopted from the papers' implmentation and the original links can be easily found on top of each file. We also appreciate the datasets from various sources, which will be highlighted in the dataset file.

Thanks to their great work and contributions!