Probabilistic modeling of tabular data with normalizing flows.
If your data consists of continuous variables that can be put into a Pandas DataFrame, pzflow can model the joint probability distribution of your data set.
The Flow class makes building and training a normalizing flow simple.
It also allows you to easily sample from the normalizing flow (e.g. for forward modeling or data augmentation), and to calculate posteriors over any of your variables.
See this Jupyter notebook for an introduction. See this notebook for a more complicated reshift example.
You can install pzflow from PyPI with pip:
pip install pzflow
If you want to run pzflow on a GPU with CUDA, you need to follow the GPU-enabled installation instructions for jaxlib here.
You may also need to add the following to your .bashrc:
# cuda setup
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export PATH=$PATH:/usr/local/cuda/bin
If you have the GPU enabled version of jax installed, but would like to run on a CPU, add the following to the top of your scripts/notebooks:
import jax
# Global flag to set a specific platform, must be used at startup.
jax.config.update('jax_platform_name', 'cpu')
Note that if you run jax on GPU in multiple Jupyter notebooks simultaneously, you may get RuntimeError: cuSolver internal error. Read more here and here.
We are preparing a paper on pzflow. If you are using this package in your research, please check back here for a citation before publication.
To work on pzflow, after forking and cloning the repo:
- Create a virtual environment with Python
E.g., with condaconda create -n pzflow - Activate the environment.
E.g.,conda activate pzflow - Install pzflow in edit mode with the
devflag
I.e., in the root directory,pip install -e .[dev]
pzflow was originally designed for forward modeling of photometric redshifts as a part of the Creation Module of the DESC RAIL project. The idea to use normalizing flows for photometric redshifts originated with Bryce Kalmbach. The earliest version of the normalizing flow in RAIL was based on a notebook by Francois Lanusse and included contributions from Alex Malz.
The functional jax structure of the bijectors was originally based on jax-flows by Chris Waites. The implementation of the Neural Spline Coupling is largely based on the Tensorflow implementation, with some inspiration from nflows.
Neural Spline Flows are based on the following papers:
NICE: Non-linear Independent Components Estimation
Laurent Dinh, David Krueger, Yoshua Bengio
arXiv:1410.8516
Density estimation using Real NVP
Laurent Dinh, Jascha Sohl-Dickstein, Samy Bengio
arXiv:1605.08803
Neural Spline Flows
Conor Durkan, Artur Bekasov, Iain Murray, George Papamakarios
arXiv:1906.04032