This repository contains the official implementation of the paper "Deep Generative Spatiotemporal Engression for Probabilistic Forecasting of Epidemics" by Rajdeep Pathak and Tanujit Chakraborty.
We introduce Deep Spatiotemporal Engression, a novel method for generating accurate and reliable probabilistic forecasts, specifically designed for low-frequency spatiotemporal datasets. Acting as distributional lenses, these methods generate out-of-sample probabilistic forecasts by sampling from trained models.
- Lightweight Deep Generative Architecture: Replaces heavy conventional models while maintaining high accuracy.
- Endogenous Uncertainty Quantification: Uncertainty is driven by a pre-additive noise component during model construction.
We propose three spatiotemporal engression models:
- Graph Convolutional Engression Network (GCEN)
- Spatio-Temporal Engression Network (STEN)
- Multivariate Engression Network (MVEN)
GCEN Architecture Figure: Architecture of the Graph Convolutional Engression Network (GCEN).
You can install the package directly via PyPI or clone the repository to install it from the source.
pip install stengression
If you want to modify the code or run the latest development version, you can clone the repository:
git clone https://github.com/PyCoder913/stengression.git
cd stengression
pip install -r requirements.txtIf you use this code, models, or find our work helpful in your research, please consider citing our paper:
@article{pathak2026deep, title={Deep Generative Spatiotemporal Engression for Probabilistic Forecasting of Epidemics}, author={Pathak, Rajdeep and Chakraborty, Tanujit}, journal={Transactions on Machine Learning Research}, year={2026}, url={https://openreview.net/pdf?id=7AfAztCd5A} }