Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Normalising Flows as NDEs for production-grade plots #334

Open
htjb opened this issue Sep 6, 2023 · 2 comments · May be fixed by #353
Open

Normalising Flows as NDEs for production-grade plots #334

htjb opened this issue Sep 6, 2023 · 2 comments · May be fixed by #353

Comments

@htjb
Copy link
Collaborator

htjb commented Sep 6, 2023

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

kind=nde

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
Copy link
Collaborator

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
Copy link
Collaborator Author

htjb commented Sep 6, 2023

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).

@htjb htjb linked a pull request Nov 8, 2023 that will close this issue
Draft
6 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

NDE implementation handley-lab/anesthetic
2 participants
@williamjameshandley @htjb