This repository contains code for Bayesian estimation of compartmental models for COVID-19 using numpyro and jax.

We are experimenting with different Bayesian compartmental models. The basic ingredients are:

• classical compartmental models from epidemiology
• prior distributions on parameters
• models for time-varying dynamics
• models for partial/noisy observations of confirmed cases and deaths
• Bayesian inference using numpyro

This Bayesian SEIRD model is current as of April 27, 2020. Model code here.

The team is based at the College of Information and Computer Sciences and School of Public Health & Health Sciences at UMass. The model contributors are:

• Dan Sheldon
• Casey Gibson
• Nick Reich

Dr. Reich directs the CDC Influenza Forecasting Center of Excellence at UMass.

Our code depends on recent development versions of numpyro and jax. If you don't have these packages and don't care about tracking their versions, our installation routine will pull and install them as unversioned packages:

git clone https://github.com/dsheldon/covid
cd covid
pip install -e .

If you need to manually install jax and numpyro, here are rough instructions. More details can be found at the project sites.

Install jaxlib (more info here)

pip install --upgrade jaxlib

Install jax (more into here)

git clone https://github.com/google/jax
cd jax
pip install -e .

Install numpyro (more details here)

git clone https://github.com/pyro-ppl/numpyro.git
cd numpyro
pip install -e .[dev]