PriorCVAE in JAX

This repository is based on the following two papers:

Semenova, Elizaveta, et al. "PriorVAE: encoding spatial priors with variational autoencoders for small-area estimation." Journal of the Royal Society Interface 19.191 (2022): 20220094. Original code is avilable here.
Semenova, Elizaveta, Max Cairney-Leeming, and Seth Flaxman. "PriorCVAE: scalable MCMC parameter inference with Bayesian deep generative modelling." arXiv preprint arXiv:2304.04307 (2023). Original code is avilable here.

Environment

We recommend setting up a conda environment.

conda create -n prior_cvae -c conda-forge python==3.10.1
conda activate prior_cvae

Within the virtual environment, install the dependencies by running

pip install -r requirements.txt

Note: The code has been tested with Python 3.10.1. There is a known issue with Python 3.10.0 related to loading a saved model because of the bug which is resolved in Python 3.10.1.

Install the package

python setup.py install

To install in the develop mode:

python setup.py develop

Examples

Example notebooks can be found in the examples\ directory. Remember to install the priorCVAE package before running the notebooks.

Sample command:

cd examples/
jupyter notebook GP-PriorCVAE.ipynb

Note: For experiments it is recommended to use float64 precision to avoid numerical instability:

import jax.config as config
config.update("jax_enable_x64", True)

To runs tests

First install the test-requirements by running the following command from within the conda environment:

pip install -r requirements-test.txt

Then, run the following command:

pytest -v tests/

Projects using PriorVAE or PriorCVAE

Project	Description	Publication	Uses current library
aggVAE	"Deep learning and MCMC with aggVAE for shifting administrative boundaries: mapping malaria prevalence in Kenya", Elizaveta Semenova, Swapnil Mishra, Samir Bhatt, Seth Flaxman, Juliette Unwin	arvix Accepted to the "Epistemic Uncertainty in Artificial Intelligence" workshop of the "Uncertainty in Artificial Intelligence (UAI 2023)" conference.	no