Integrating External Infrastructure

[LouisDesdoigts]

Ideally ∂Lux is able to become the industry standard package for astrophysical optical modelling, and there is currently already a massive infrastructure for this primarily maintained by STSci. With this in mind I think it is very important that we are able to integrate ∂Lux with these packages as seamlessly as possible. Both to minimise the difficulty for other researchers to integrate ∂Lux into their work, and to build out as much functionality as possible without actually building it!

We are lucky that the most advanced space telescope ever has just been launched that will provide us with an unprecedented view of our universe. Its complicated mess though (they only use 5 different coordinates systems 🤷‍♂️, enjoy this , more info here).

Good thing is they have already built all the code to do everything for us primarily in pysiaf!. This basically integrates with a centralised and maintained data-base containing all the info one could need to model a telescope, from pupil and filter wheel information, to focal plane distortions to map pixel-coordinates to on-sky coordinates. I've gone over most of the source code and it seems that everything that is non-differentiable could be done in the ∂Lux OpticalSystem.init() methods, with the run-time code segmented into beautiful scripts that we can most-likely replace import numpy as np with import jax.numpy as np and then whack-a-mole a few DeviceArray.at[].set() errors and wash our hand of it.

I think this would be massively beneficial and we could even just make a pysaif_def property of the OpticalSystem that gets all the information at init() and populates the relevant parameters/layers to be called at run-time. There's obviously some nuance in how that would work, and I'm only familiar with pysaif at a surface level but I suspect this would be an awesome feature to help bring ∂Lux to the forefront of optical modelling, rather than a specialised library.

Pysaif relies on a few other packages too like pysynphot for stellar spectrums, I think we could take similar approach with this package - pre calc parameters in init(), spit out layers/jax arrays to populate OpticalSystem object.

Here is an example of a package that somewhat integrates jax into the astropy infrastructure. This could serve as a good reference point moving forward.

Also if someone is feeling masochistic I have a hair-brained idea to write a jax-based astropy units module that will make everyone life easier (except whoever writes it).

[LouisDesdoigts]

So I've done a lot more investigating into how to integrate ∂Lux into the STSci/JWST infrastructure, there's a few moving pieces so here's a short overview. There are two semi-separate areas that deal with JWST things, but the parts are kinda of all interconnected.

So when dealing with real JWST data, the first step is to clean the data using the JWST pipieline (IR detectors are kinda of garbage turns out). It uses the header information from the data file to talk to both the CRDS which contains all of their data about the telescope which is updated and maintained over time, and with MIRAGE which models JWST images using a massive library of psfs generated by webbpsf.

There is also the optical modelling side, built from webbpsf and poppy. SO webbpsf is the high-level driver, which talks to the CRDS in order to get up-to-date and accurate data files, containing their measured field dependent phase aberrations, aperture files, filter info etc etc. webbpsf is designed to loosely function in the same mechanical way as the actual telescope (ie you can't apply two filters on the same wheel, or add a chronographic occulter to an optical system without one). webbpsf is designed to spit out fits files with matching header info to the data that would be spit out of the JWST in the specific configuration, pointing, readout mode etc. This makes it the perfect interface for us with the real JWST data. The package itself doesn't perform any optical calcuations, instead creating a poppy instrument object and modelling that data from its internally stored poppy optical system.

I've played around with the webbpsf and JWST packages, processed some mock data and read through lots of the source code and docs, I think we have a really neat way to integrate into the system without actually constructing much more code directly. I think we can actually just create a fork of webbpsf (found here) and actually work directly with the webbpsf instrument objects. In this fork we should really only need to change one line in the JWSTInstrument.calc_psf() call that takes the poppy object that is created and spits out an equivalent ∂Lux object. This should actually be relatively simple as most of the data is already formatted how we need it. This may necessitate also creating a fork of poppy in order to reveal the actual operations performed on the wavefront, but I think for starting out we should be able to directly take the data loaded into the poppy.osys.planes list, or modify the get_optical_system() function in webbpsf to create a ∂Lux object instead.

I think this method has a lot of benefits - It massively down scopes the amount of code and structure we need to build (they have already made a wheel, we don't need to make a new one), it directly hooks ∂Lux objects into the full JWST infrastructure with up-to-date files from the CRDS, we can directly go from a real data file to a ∂Lux object (so we can do sick autodiff stuff from there), and should only require relatively minimal amounts of code modifications to webbpsf package (which we can keep directly updated with them as a fork of the actual library).

Anyway, I think this is by far the best way to go forward!

[Jordan-Dennis]

Sounds lit. I'm currently dev'ing the hexikes. Should I continue with this or is there a new priority?

[benjaminpope]

I think we'll need hexikes whatever we do, and it shouldn't be a huge problem to build these - I reckon go ahead

[benjaminpope]

Yep - I think this is the way forward (as we have been hoping since Louis' Queensland trip!). We should programmatically extract exactly the arrays and filters used by WebbPSF and run them through dLux - these have been extensively tested at STScI by serious professionals and we want to be able to match and use their work.

[LouisDesdoigts]

Yeah we can let them get us 90% of the way there, and then we go the last 10% with autodiff (and then presumably, eventually have our data feed back into the CDRS)

[Jordan-Dennis]

Probably a waste of time but I have written code to generate the primary aperture ourselves 😮‍💨. Not sure if we will even want to use it long term because as @benjaminpope says the WebbPSF code has be tested very seriously and this has not. The code for this is on the branch JamesWebb and I plan to merge it down for no other reason than that it implements functionality that will probably be useful elsewhere.