ValueError with save_intermediates=True for coronagraphic calcuations in nircam

[LouisDesdoigts]

So I have encountered a bug when doing coronagraphic psf calculations with @itroitskaya using the nircam class.

The error is thrown for all pupil and image mask combinations when save_intermediates=True.

Here is a minimal example of the bug using the nircam coronagraph example:

import webbpsf
nc = webbpsf.NIRCam()
nc.filter='F430M'
nc.image_mask='MASK430R'
nc.pupil_mask='CIRCLYOT'
nc.calc_psf('coronagraphic.fits', save_intermediates=True)

> ValueError: Wavefronts can only be added if they have the same pixelscale: 0.022083140339507784 arcsec pix vs 0.02240968746126646 arcsec pix`

This is using webbpsf v1.1.1 and poppy v1.0.3.

Any help is greatly appreciated!

[LouisDesdoigts]

Also the same error is thrown for return_intermediates=True!

[mperrin]

@obi-wan76 Marcio can you take a look at this, please? Let me know if you can't figure it out, in which case I can try to take a look later.

[obi-wan76]

Without trying to ouput the intermediate planes there is no error message when creating the PSF

When running the example above I get the same error but with a different PS:
ValueError: Wavefronts can only be added if they have the same pixelscale: 0.032515086329876194 arcsec / pix vs 0.032748085873623066 arcsec / pix
my first reaction is a numerical/mathematical precision, both numbers are the same to their second significant figure. So, we could enforce a number check to that level of precision. However, I am not sure why the same example is returning different PS numbers.

[mperrin]

I strongly doubt this is just a numerical precision issue. Much more likely a real bug. The code already allows for close but not precisely identical values through use of np.isclose. And in any case a 1% difference in pixel scale is a large and significant difference at the level of precision we care about! Pretty sure this is a real bug that you will need to dig in and understand.

[mperrin]

Can confirm, yes this is a real bug. Or at least a limitation in how the code is set up.

Reminder for how coronagraphic propagations in NIRCam work (by default):

1. Pupil plane with OTE wavefront
2. Fast Fourier transform to focal plane. Apply model of coronagraph focal plane mask.
3. Inverse transform back to pupil plane. Apply coronagraph Lyot stop and NIRCam internal WFE..
4. Matrix Fourier transform to focal plane, at the NIRCam detector pixel scale.

The issue is that the FFT to the focal plane in step 2 has a pixel scale which is fixed in lambda/D per pixel, and therefore varies in arcsec per pixel. Thus for each wavelength, the pixel scale at that plane will vary slightly (linearly with wavelength), and as a result one cannot simply add the wavefront objects without resampling.

How to proceed? I think it may depend in part about what @LouisDesdoigts is trying to achieve. One could loop over the wavelengths outside of the calc_psf function, and save the intermediate wavefronts per each wavelength, and interpolate as needed to put them on a common scale for adding together. One could modify poppy to resample wavefronts when adding wavefronts with inconsistent pixel scales. There's also the semi-analytic coronagraph propagation method (i.e. Soummer et al. 2007) which allows setting the pixel scale directly in the intermediate plane and making it the same for all wavelengths. There is support for that in WebbPSF, and in fact at one point it was the default for NIRCam calculations with the round masks, but that method does not work well for the bar occulters so it's disabled by default.

Louis can you explain more about what you're trying to do? Maybe we can brainstorm a bit on which approach or workaround would be of most use for you. Can I guess you are working on trying to get NIRCam calculations running in dLux?

[LouisDesdoigts]

Thanks for this detailed and speedy response! Yes we are getting some dLux coronagraph models going at the moment.

Part of our goal is to have our models match webbpsf to machine precision as a form of validation, so a single resampled/added wavefront is less useful than the all the individual native monochromatic wavefronts anyway. If this is out-of-bound for webbpsf/poppy then we can always perform the monochromatic calculations as you said.

One question I do have though is how the chromatic calculations are performed. Because internal propagation are done using FFTs the pixel scale for each wavelength will be inconsistent to the pixel scale of the Lyot plane stop (assuming you are using a static array). Do you interpolate the wavefront to the optic or optic to the wavefront, OR are the intermediate mask planes generated dynamically based on the wavefront pixel scale at run time?

[mperrin]

The internal propagation FFTs to the image plane, where the coronagraph focal plane mask is then generated dynamically based on the wavelength-dependent pixel scale. Then FFT back to the pupil plane, returning to the same fixed wavelength-independent pixel scale as was started with. So the same Lyot plane stop can be applied for all wavelengths. Then the MFT to the detector pixel scale or a subsampled fraction thereof.

[mperrin]

Btw I would be very happy to help collaborate with you and Ben et al on this. Please feel free to start an email thread if you’d like, or we can stick with here on GitHub, however you prefer.

I think you may have heard that I’ll be visiting Sydney in September, and it’s been high on my list to plan some time then to collaborate and learn more about what you’re doing with differentiable optical modeling. But had not yet had time to reach out about that, so it’s timely that you messaged on here. I’m also interested in trying to apply this more broadly across other modes of JWST, not just coronagraphy. Have you tried any of the regular imaging modes first as a potentially simpler case than coronagraphy? There’s definitely open work to be done in tuning just the imaging mode PSFs to better match flight performances, for instance.

[LouisDesdoigts]

Ah so you're using parametric image plane occulters. I think the best way forwards for us would be to create a series of dLux layers mirroring the BandLimitedCoron classes. Having some insider knowledge on this will definitely be valuable! I've had a look at the source code, it seems like the circular occulting masks are constructed from bessel functions. I had a look at the Krist et al. 2010 SPIE paper which is referenced in the docstring but it doesn't seem to have the analytic model. Do you have any resources on these masks?

Yes I'm looking forwards to your visit! We have some papers on dLux in review and being written currently so we will have some nice things to show soon. If you want we also have some tutorials showing the cool ways to use autodiff capabilities for optics here. We have some basic phase retrieval, high-dimensional instrumental calibration, Fisher information optimisation, HMC inference etc. These should be able to give you an overview of some of the thing we are thinking about.

We have done some work with the other imaging modes, primarily focused on AMI data. I have local models that get a machine-precision to webbpsf for AMI that I have been using. You're probably aware that the interferometric observables are a factor of 10-100x worse than expected. We believe this is due to detector effect, primarily the Brighter-Fatter Effect (BFE). Our models are proving to be very useful in learning and calibrating these effects since we can directly examine our residuals and probe the optics in a way that the interferometric pipelines are not capable of. Similarly the AMI psf is a perfect test-case for these detector level effects since it's highly structured with a high dynamic range which helps decouple the BFE from things like jitter and defocus. We're also thinking about working with ramp-level data to help learn these effects directly.

We are definitely looking to develop a JWST module though, so it would be great to collaborate on calibrating models with on-sky data!

[mperrin]

Apologies for the delayed response - this slipped off my plate while juggling too many things. (Like I said, feel free to email if you end up stuck waiting on me for anything)

The analytic functions and equations for the band limited coronagraphs are in the code in webbpsf optics.py. See https://github.com/spacetelescope/webbpsf/blob/develop/webbpsf/optics.py#L799-L813
I would think that would have what you need, but let me know if not.

Yes I'm well aware of the challenges to AMI contrast from detector effects. Glad to hear you're digging into those complexities with the dLux models. Let me know if there's anything I can do to help!