Normalising Flows as NDEs for production-grade plots

[htjb]

Had a brief discussion with @williamjameshandley offline about using normalizing flows (NFs) as Neural Density Estimators (NDEs) instead of KDEs for production-grade anesthetic plots. I've been doing some work recently that shows that NFs typically out-perform KDEs when used to estimate the KL divergence/BMD of target distributions. This suggests that they can be used to better represent the underlying samples in a distribution for production grade plots.

We would like to add a kind=nde option to the plotting functionality in anesthetic and integrate in margarine for NF training. A simple example is shown below.

import numpy as np
from margarine.maf import MAF
import matplotlib.pyplot as plt

samples = np.random.multivariate_normal([0, 0], 
                                        [[1, 0.], [2, 0.1]], size=5000)
f = MAF(samples)
f.train(1000, early_stop=True)

x = np.linspace(samples[:, 0].min(), samples[:, 0].max(), 100).astype(np.float32)
y = np.linspace(samples[:, 1].min(), samples[:, 1].max(), 100).astype(np.float32)
xv, yv = np.meshgrid(x, y, indexing='ij')

fig, axes = plt.subplots(1, 1)
    
lp = f.log_prob(np.array([xv.flatten(), yv.flatten()]).T).numpy()
z = np.exp(lp - lp.max()).reshape(xv.shape)

plt.scatter(samples[:, 0], samples[:, 1], s=1, c='k', alpha=0.5)
axes.contourf(xv, yv, z, cmap='Blues', levels=[0.68, 0.95, 1.00], alpha=0.8)

plt.tight_layout()
plt.savefig('kind=nde.png', dpi=300)
plt.show()

Which produces the following plot

For multi-modal distributions we can take advantage of the clustering built into margarine. The flows take seconds to minutes to train depending on number of samples and dimensionality.

[williamjameshandley]

Hi @htjb,

in principle this is a pretty easy addition. The only thing that needs to be got right is the computation of the level sets for the contours (for which the example in anesthetic.plot.kde_contour_plot_2d shows the right way to do it with iso_probability_contours.

To plumb this in, you would need to create something akin to (i.e. in large part copy-paste)

• anesthetic.plot.kde_contour_plot_2d
• anesthetic.plot.kde_plot_1d
• anesthetic.plotting._matplotlib.hist.Kde1dPlot
• anesthetic.plotting._matplotlib.hist.Kde2dPlot
• anesthetic.plotting._core.PlotAccessor.kde_1d
• anesthetic.plotting._core.PlotAccessor.kde_2d

and adjust:

• anesthetic.plotting._matplotlib.init.PLOT_CLASSES
• anesthetic.samples.Samples.doc

It's this messy in order to give us pandas-like plotting functionality (e.g. samples.x0.plot.nde_1d())

In general a `grep -ri kde anesthetic tests' will show you most of what needs to be adjusted.

[htjb]

Nice, this sounds good! I remember playing with iso_probability_contours for a previous PR I think. I will give this a go and put a PR together in the coming week(s).