Epidemic Modeling with Pytorch

Documentation | Paper | Web Interface

EpiLearn is a Pytorch-based machine learning tool-kit for epidemic data modeling and analysis. We provide numerour features including:

• Implementation of Epidemic Models
• Simulation of Epidemic Spreading
• Visualization of Epidemic Data
• Unified Pipeline for Epidemic Tasks

For more models, please refer to our Awesome-Epidemic-Modeling-Papers

Installation

From Source

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

conda create -n epilearn python=3.9
conda activate epilearn

python setup.py install

From Pypi

pip install epilearn

EpiLearn also requires pytorch>=1.20, torch_geometric and torch_scatter. For cpu version, we simply use pip install torch, pip install torch_geometric and pip install torch_scatter. For GPU version, please refer to Pytorch, PyG and torch_scatter.

Tutorial

We provide a complete tutorial of EpiLearn in our documentation, including pipelines, simulations and other utilities. For more examples, please refer to the example folder. For the overal framework of EpiLearn, please check our paper.

Here we also offer a quickstart of how to use the EpiLearn for forecast and detection task.

Forecast Pipeline

from epilearn.models.SpatialTemporal.STGCN import STGCN
from epilearn.data import UniversalDataset
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 = UniversalDataset()
dataset.load_toy_dataset()
# Adding Transformations
transformation = transforms.Compose({
                "features": [transforms.normalize_feat()],
                "graph": [transforms.normalize_adj()]})
dataset.transforms = transformation
# Initialize Task
task = Forecast(prototype=STGCN,
                dataset=None, 
                lookback=lookback, 
                horizon=horizon, 
                device='cpu')
# Training
result = task.train_model(dataset=dataset, 
                          loss='mse', 
                          epochs=50, 
                          batch_size=5, 
                          permute_dataset=True)
# Evaluation
evaluation = task.evaluate_model()

Detection Pipeline

from epilearn.models.Spatial.GCN import GCN
from epilearn.data import UniversalDataset
from epilearn.utils import transforms
from epilearn.tasks.detection import Detection
# initialize settings
lookback = 1 # inputs size
horizon = 2 # predicts size; also seen as number of classes
# load toy dataset
dataset = UniversalDataset()
dataset.load_toy_dataset()
# Adding Transformations
transformation = transforms.Compose({
                " features": [],
                " graph": []})
dataset.transforms = transformation
# Initialize Task
task = Detection(prototype=GCN, 
                 dataset=None, 
                 lookback=lookback, 
                 horizon=horizon, 
                 device='cpu')
# Training
result = task.train_model(dataset=dataset, 
                          loss='ce', 
                          epochs=50, 
                          batch_size=5)
# Evaluation
evaluation = task.evaluate_model()

Web Interface

Our web application is deployed online using streamlit. But it also can be initiated using:

python -m streamlit run interface/app.py to activate the interface

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!