Bayesian Inference (posterior exploration) with MPoL and parametric models

[iancze]

Is your feature request related to a problem or opportunity? Please describe.
The MPoL framework allows the rapid calculation of gradients w.r.t. model parameters. If the sky-plane model is specified parametrically (i.e., as a sky-plane Gaussian, or Gaussian ring, etc...) and is reasonably low-dimensional (< 100) then the posterior distribution of those parameters can be explored using techniques like Variational Inference or Hamiltonian Monte Carlo, conveniently provided by the Pyro package.

This would make MPoL useable for more than just non-parametric imaging.

Describe the solution you'd like
In theory, using Pyro with MPoL should already be possible (no modifications to MPoL needed).

We would like to test whether this is possible by making a tutorial showing how to

1. Generate a fake dataset for testing (using existing baselines... + noise) and inspect with dirty image
2. Write a parametric model as input to mpol.images.ImageCube
3. Optimize the parametric model using Gradient Descent
4. Explore the parameter posteriors using Pyro and HMC or VI sampling. We already did a little bit of this in PyMC3, but we'd like this to work with MPoL and Pyro.

We should start with a simple Gaussian blob model.

But for an application to real data, a useful (and popular) dataset would be the DSHARP AS 209 continuum. We could make a direct comparison to the posteriors calculated by Guzmán+18 using their ringed Gaussian model with 27 parameters, sampled by emcee.

[iancze]

Some updates on this issue.

• PR Introduce geometry utilities #121 will introduce the geometric utilities that are needed for disk projections/deprojections
• We have the AS 209 continuum visibilities processed with mpoldatasets
• I am making progress understanding Pyro and how we can relatively easily wrap a torch.nn.Module with PyroModule and introduce stochastic parameters.
• I am making progress understanding optimization, SVI, and HMC in the Pyro ecosystem. It seems to work well on small examples though we haven't tried the full machinery on the AS 209 dataset yet.
• I'm not sure if there's anything we actually need to change to the MPoL codebase itself to address this issue, since most of the model-writing and Pyro-interfacing will be specific to the problem at hand.
• I think this is a stepping stone to solving Parametric modeling for solar system bodies #105, since I think in that use case we'll want to sample (or at least optimize) the parameters of the model during the fitting process, at least as the first part of a workflow. The final RML loop might not use Pyro at all, though.
• However, I think a nice deliverable on this specific issue would be a tutorial showing how to use Pyro with MPoL on the AS 209 dataset itself, either continuously integrated or as a "large" tutorial like the HD 143006 ones.
• It would be cool to make a GIF iterating through posterior predictive distributions of the images and 1D profiles, to show the uncertainty on the fit parameters.

[iancze]

Some updates on this issue:

Beyond the already-merged #121, most of the additional code for this is outside of any changes to the MPoL codebase itself. So I think this makes sense just to demonstrate these capabilities as a long-form tutorial.

To do that, I need to

• create new branch for this
• create Jupyter notebook + jupytext link
• ascertain device link (cpu / GPU)
• upload rescaled AS209.asdf dataset to Zenodo
• iterate development on a GPU-enabled device so it goes more quickly

Overview of tutorial contents

• overview of MPoL as FFT link, possibly include mermaid.js graph showing where parametric model fits in
• summary of DSHARP dataset, dirty images, review of Guzman et al. 2018 approach
• describe Pyro PPL and highlight a few examples using static code of PyroSample
• build and heavily comment the disk model, including MPoL geometry routines, deterministic statements
• run SVI inference loop on GPU
• analyze samples
• explore MultiNormal fits to see if posterior changes
• run HMC loop on GPU and analyze samples
• show scatter in 1D profile as draws or movie

We made a movie showing iterative posterior predictions of the images but we found the actual changes were very small, and thus it didn't make for a great visualization.