jammy_flows

This package implements (conditional) PDFs with Joint Autoregressive Manifold (MY) normalizing-flows. It grew out of work for the paper Unifying supervised learning and VAEs - coverage, systematics and goodness-of-fit in normalizing-flow based neural network models for astro-particle reconstructions [arXiv:2008.05825] and includes the paper's described methodology for coverage calculation of PDFs on tensor products of manifolds. For Euclidean manifolds, it includes an updated implementation of the offical implementation of Gaussianization flows [arXiv:2003.01941], where now the inverse is differentiable (adding Newton iterations to the bisection) and made more stable using better approximations of the inverse Gaussian CDF. Several other state-of-the art flows are implemented sometimes using slight modifications or extensions.

The package has a simple syntax that lets the user define a PDF and get going with a single line of code that should just work. To define a 10-d PDF, with 4 Euclidean dimensions, followed by a 2-sphere, followed again by 4 Euclidean dimensions, one could for example write

import jammy_flows

pdf=jammy_flows.pdf("e4+s2+e4", "gggg+n+gggg")

The first argument describes the manifold structure, the second argument the flow layers for a particular manifold. Here "g" and "n" stand for particular normalizing flow layers that are pre-implemented (see Features below). The Euclidean parts in this example use 4 "g" layers each.

Have a look at the script that generates the above animation.

Documentation

The docs can be found here.

Also check out the example notebook.

Features

General

Autoregressive conditional structure is taken care of behind the scenes and connects manifolds
Coverage is straightforward. Everything (including spherical, interval and simplex flows) is based on a Gaussian base distribution (arXiv:2008.0582).
Bisection & Newton iterations for differentiable inverse (used for certain non-analytic inverse flow functions)
amortizable MLPs that can use low-rank approximations
amortizable PDFs - the total PDF can be the output of another neural network
unit tests that make sure backwards / and forward flow passes of all implemented flow-layers agree
include log-lambda as an additional flow parameter to define parametrized Poisson-Processes
easily extendible: define new Euclidean / spherical flow layers by subclassing Euclidean or spherical base classes

Euclidean flows:

Generic affine flow (Multivariate normal distribution) ("t")
Gaussianization flow arXiv:2003.01941 ("g")
Hybrid of nonlinear scalings and rotations ("Polynomial Stretch flow") ("p")

Spherical flows:

S1:

Moebius transformations (described in arXiv:2002.02428) ("m")
Circular rational-quadratic splines (described in arXiv:2002.02428) ("o")

S2:

Autorregressive flow for N-Spheres (arXiv:2002.02428) ("n")
Exponential map flow (arXiv:0906.0874/arXiv:2002.02428) ("v")
Neural Manifold Ordinary Differential Equations arXiv:2006.10254 ("c")

Interval Flows:

"Neural Spline Flows" (Rational-quadratic splines) arXiv:1906.04032 ("r")

Simplex Flows:

Autoregressive simplex flow arXiv:2008.05456 ("w")

For a description of all flows and abbreviations, have a look in the docs here.

Requirements

pytorch (>=1.7)
numpy (>=1.18.5)
scipy (>=1.5.4)
matplotlib (>=3.3.3)
torchdiffeq (>=0.2.1)

The package has been built and tested with these versions, but might work just fine with older ones.

Installation

specific version:

pip install git+https://github.com/thoglu/jammy_flows.git@*tag*

e.g.

pip install git+https://github.com/thoglu/jammy_flows.git@1.0.0

to install release 1.0.0.

master:

pip install git+https://github.com/thoglu/jammy_flows.git

Contributions

If you want to implement your own layer or have bug / feature suggestions, just file an issue.