add coron_include_pre_lyot_plane option for extra output plane

webbpsf/tests/test_nircam.py

@@ -592,3 +592,18 @@ def test_coron_shift(offset_npix_x=4, offset_npix_y=-3, plot=False):
     assert np.isclose(cutout_1.sum(), cutout_2.sum()), "PSF cutout sums should be consistent"
 
     assert np.allclose(cutout_1, cutout_2), "PSF cutouts should be consistent"
+
+def test_coron_extra_lyot_plane():
+    # Test adding the optional output of the WFE prior to the Lyot stop plane
+    nrc = webbpsf_core.NIRCam()
+    nrc.pupil_mask = 'MASKLWB'
+    nrc.image_mask = 'MASKLWB'
+    nrc.filter='F460M'
+
+    psf, planes = nrc.calc_psf(nlambda=1, return_intermediates=True, display=True)
+
+    nrc.options['coron_include_pre_lyot_plane'] = True
+    psf2, planes2 = nrc.calc_psf(nlambda=1, return_intermediates=True, display=True)
+
+    assert len(planes2) == len(planes)+1, "There should be an added plane for coron_include_pre_lyot_plane"
+    assert np.allclose(psf[0].data, psf2[0].data), "The PSF output should be the same either way"

webbpsf/webbpsf_core.py

@@ -1932,6 +1932,10 @@ def make_fqpm_wrapper(name, wavelength):
         shift_x, shift_y = self._get_pupil_shift()
         rotation = self.options.get('pupil_rotation', None)
 
+        if self.options.get('coron_include_pre_lyot_plane', False) and self.pupil_mask.startswith('MASK'):
+            optsys.add_pupil(poppy.ScalarTransmission(name='Pre Lyot Stop'))
+            optsys.planes[3].wavefront_display_hint = 'intensity'
+
         if self.pupil_mask == 'MASKFQPM':
             optsys.add_pupil(transmission=self._datapath + "/optics/MIRI_FQPMLyotStop.fits.gz",
                              name=self.pupil_mask,
@@ -2476,6 +2480,13 @@ def _addAdditionalOptics(self, optsys, oversample=2):
             optsys.add_pupil(transmission=self._WebbPSF_basepath + "/tricontagon_oversized_4pct.fits.gz",
                              name='filter stop', shift_x=shift_x, shift_y=shift_y, rotation=rotation)
 
+        if self.options.get('coron_include_pre_lyot_plane', False) and self.pupil_mask.startswith('MASK'):
+            optsys.add_pupil(poppy.ScalarTransmission(name='Pre Lyot Stop'), index=3)  # this is before the above plane, but do the insertion here
+                                                             # because of all the hard-coded index=3 above
+
+            optsys.planes[3].wavefront_display_hint = 'intensity'
+
+
         return (optsys, trySAM, SAM_box_size)
 
     def _get_fits_header(self, hdulist, options):