Active development?

[benjaminpope]

Hi Ian,

Love this repo - seems great! I was just wondering if this is still under active development - our team were hoping to build a small package for regularized image reconstruction from the JWST AMI sparse aperture masking instrument, and I was going to suggest we build this in Jax, but was curious about adapting MPoL too which is much more mature.

What do you think is the best way forward?

All the best,

Ben

[iancze]

Hi Ben,

Thanks so much for reaching out, it's great to hear from you! We are still very much actively developing this package, despite what the recent commit history might suggest :-p I think the reasons for the drop in public-facing commit frequency are

• we've finally reached a fairly stable codebase (at least as concerns application to ALMA observations)
• we (and by that I mean @briannazawadzki ) have been focusing on writing up a paper describing our results. We plan to submit Zawadzki, Czekala, Loomis et al. to AAS journals in the near future.

You can view a list of planned development items here. I think highest on our list is the "loose visibility" connector using TorchKbNufft (WIP), followed by exploration of regularizers for the spectral dimension. There are a couple ongoing ALMA programs for which we are using MPoL as part of the RML/imaging analysis, too.

Your idea to utilize/adapt MPoL for JWST AMI sounds very exciting! I'm definitely interested in learning more and, if possible, adapting MPoL to work for these type of data. As I've been learning more about imaging, RML, and developing this package, I have always maintained a hope that this could become a general purpose RML package for all types of interferometric data, since, as far as I can tell, the basic math is quite universal. I think the modular nature of these neural network packages (in our case, PyTorch, but also applies to JAX) lends itself very well to construction of flexible, instrument-specific workflows. Thus far, I've lacked a knowledgeable collaborator to help extend to SAM, so hopefully this is an opportunity for growth :)

About ~1.5 yrs ago I tried to run a case study of PyTorch v. JAX to see whether it made sense to switch backends. At the time, I concluded that PyTorch was more mature/documented (and easier to understand, for me) and the modular NN design made a lot of sense for our applications. I don't have much recent experience with JAX, so those conclusions might need to be updated in 2022. However, all of the functionality that led me to choose PyTorch in the first place is still there and maintained so I think the functionality of MPoL itself is on solid ground. I was also comforted by the fact that (at least by "feel") the fraction of MPoL code that is PyTorch-specific feels relatively small and contained. This has helped me sleep at night in case the code needs to be updated in a few years if/when there are radical changes in the autodiff landscape. In your opinion, what would be the benefits of using JAX over PyTorch here? I see that @dfm has created an awesome set of JAX bindings to finufft, so this could be a candidate for replacement for TorchKbNufft (though we might need to wait for GPU support, since that's critical for some of our image-cube applications).

If you wouldn't mind sketching out your requirements for JWST/AMI, I could comment on whether the functionality exists in MPoL and/or how difficult it would be to add it. I think it would be really great to use the JWST/AMI application to help drive MPoL towards a more general purpose imaging package, and am happy to support changing the codebase to make it possible.

Best,
Ian

[benjaminpope]

Hi Ian,

Awesome, love to hear it. Very briefly re Jax, the truth is the main reason I want to do it is that our differentiable optics simulation code dLux is in Jax, and more generally a lot of packages are moving that way in the near term. There is an excellent package equinox which abstracts away a lot of the more restrictive syntax in Jax and gives a nice object-oriented framework, and PyMC now has a Jax NUTS backend. It has overall excellent performance on GPUs (though sadly not Apple silicon yet) and is pretty friendly. My general feeling is that Jax is going to eat other frameworks over the next few years, though I certainly wouldn't bet my house on it and PyTorch is of course a powerful and mature system itself.

For AMI, we have a reasonably stable ecosystem of packages (eg our team's package AMICAL, the STScI package ImPlaneIA, and the closed-source SAMPip) for extracting calibrated squared visibilities and closure phases from the JWST NIRISS mask (and other aperture masking instruments). Unlike in the radio, we do not have any complex visibilities or receivers with uncorrupted phase information. AMICAL and CANDID then do pretty well on fitting parametric models to data. But we aren't settled on what to do for nonparametric image reconstruction, and personally I want to differentiably simulate the entire thing end-to-end and include the 'calibration' in the model fitting, which is partly what dLux is designed to do.

There are already pretty good nonparametric image reconstruction codes (SQUEEZE, Mira, and the as yet unpublished generative NN libraries by our collaborators Dori Blakely and Joel Sanchez-Bermudez), but more can't hurt, especially when at very high resolution, high SNR, but poor Fourier coverage, every result is likely to be controversial unless we can nail down the whole pipeline. We have some GTO AMI data coming up and we want to nail it first go, no controversies. So my thought had been to build a Jax differentiable deconvolution code to progressively include MaxEnt, TV, smoothness, etc regularizers, and connect this to dLux. But if your MPoL has this functionality, or you're likely to build something like this in the near future, it could be worth collaborating!

All the best,

Ben

[iancze]

Thanks for your thoughts on JAX, that makes a lot of sense. It also sounds like you have (and are growing) a wonderful ecosystem of tools for optics simulation, data reduction, and modeling.

We do already have (all?) of the regularizers you mentioned implemented in MPoL: https://mpol-dev.github.io/MPoL/api.html#module-mpol.losses and we're actively thinking of more. If there's any you'd like to see us implement, please flag an issue or discussion and we'll take a closer look about feasibility. Long term, we do have plans to think about how to incorporate regularization/prior info using NN (e.g., Deep Image Prior), so we're curious to hear about your and your collaborator's progress on this front.

I think that we would only need some straightforward (hopefully minor) development required to fit your squared visibilities and closure phases:

1. Finish the WIP PR on NUFFT integration
2. Write a new "Connector" that takes the complex-valued visibilities output from a NUFFT layer and calculates squared vis and closure phases for relevant baselines
3. Write a new loss (likelihood) function to compare these model quantities to your data products
4. (?)
5. Profit

[benjaminpope]

Re NN regularization -

• Dori Blakely has some awesome (unpublished) work with a variety of NN architectures, mapping both the pixels of the AMI interferogram and also the extracted closure phases and visibilities to images;
• Joel Sanchez Bermudez has a GAN architecture but I am not sure of the details, but I think it is progressing well;
• and Laurence Perreault Levasseur & Alexandre Adam have been using a Recurrent Inference Machine.

Happy to work with you on using this with AMI data - I'm not a big PyTorch user but it should be ok. It might be worth getting in touch again after stage 1 about stage 2 ; stage 3 is trivial, and JWST data in stage 4 should lead to substantial, though entirely non-monetary, profit.