Description
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.
Activity