diff --git a/webbpsf/tests/test_nircam.py b/webbpsf/tests/test_nircam.py index c56f1a44..7639537e 100644 --- a/webbpsf/tests/test_nircam.py +++ b/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" diff --git a/webbpsf/webbpsf_core.py b/webbpsf/webbpsf_core.py index 00563f35..812385ff 100644 --- a/webbpsf/webbpsf_core.py +++ b/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):