improve flux conservation and testing for MRS IFU sims

stpsf/detectors.py

@@ -621,6 +621,8 @@ def apply_miri_ifu_broadening(hdulist, options, slice_width=0.196):
             offset_cruciform = constants.INSTRUMENT_IFU_BROADENING_PARAMETERS["MIRI"]["offset_cruciform"] * oversample_factor
             out = _miri_mrs_empirical_cruciform(psf_model=out, amp=amplitude_cruciform,
                                                 fwhm=fwhm_cruciform, x_0=offset_cruciform)
+    # enforce strict flux conservation
+    out *= hdulist[ext].data.sum() / out.sum()
 
     hdulist[ext].data = out
 

stpsf/stpsf_core.py

@@ -1278,10 +1278,11 @@ def _calc_psf_format_output(self, result, options):
                 )  # apply detector charge transfer model
             elif self.name == 'MIRI':
                 # Apply distortion effects to MIRI psf: Distortion and MIRI Scattering
-                _log.debug('MIRI: Adding optical distortion and Si:As detector internal scattering')
                 if self.mode != 'IFU':
+                    _log.debug('MIRI imager: Adding optical distortion and Si:As detector internal scattering')
                     if self._detector_geom_info.aperture.AperType != 'SLIT':
                         psf_siaf = distortion.apply_distortion(result)  # apply siaf distortion
+                        _log.debug('MIRI: Applied optical distortion based on SIAF parameters')
                     else:
                         # slit type aperture, specifically LRS SLIT, does not have distortion polynomials
                         # therefore omit apply_distortion if a SLIT aperture is selected.
@@ -1291,20 +1292,22 @@ def _calc_psf_format_output(self, result, options):
                         psf_siaf_rot, options
                     )  # apply detector charge transfer model
                 else:
+                    _log.debug('MIRI MRS: Adding IFU PSF broadening effects.')
                     # there is not yet any distortion calibration for the IFU, and
                     # we don't want to apply charge diffusion directly here
                     psf_distorted = detectors.apply_miri_ifu_broadening(result, options, slice_width=self._ifu_slice_width)
             elif self.name == 'NIRSpec':
                 # Apply distortion effects to NIRSpec psf: Distortion only
                 # (because applying detector effects would only make sense after simulating spectral dispersion)
-                _log.debug('NIRSpec: Adding optical distortion')
                 if self.mode != 'IFU':
+                    _log.debug('NIRSpec: Adding optical distortion and detector charge transfer')
                     psf_siaf = distortion.apply_distortion(result)  # apply siaf distortion model
                     psf_distorted = detectors.apply_detector_charge_diffusion(
                         psf_siaf, options
                     )  # apply detector charge transfer model
 
                 else:
+                    _log.debug('NIRSpec IFU: Adding IFU PSF broadening')
                     # there is not yet any distortion calibration for the IFU.
                     psf_distorted = detectors.apply_nirspec_ifu_broadening(result, options)
 

stpsf/tests/test_miri.py

@@ -243,7 +243,10 @@ def test_miri_ifu_broadening():
     fwhm_detdist = stpsf.measure_fwhm(psf, ext='DET_DIST')
     assert fwhm_overdist > fwhm_oversamp, "IFU broadening model should increase the FWHM for the distorted extensions"
 
-    assert np.isclose(psf['DET_SAMP'].data.sum(), psf['DET_DIST'].data.sum(), rtol=5e-3), "IFU broadening should not change total flux much"
+    # test flux conservation. THis is close but not exact, due to the optical distortion part which is distinct from but
+    # happens at same point in the calculation as the IFU broadening effect.
+    assert np.isclose(psf['DET_SAMP'].data.sum(), psf['DET_DIST'].data.sum(), rtol=2e-4), "IFU broadening should not change total flux much"
+
 
     # Now test that we can also optionally turn off that effect
     miri.options['ifu_broadening'] = None
@@ -253,3 +256,8 @@ def test_miri_ifu_broadening():
     fwhm_overdist = stpsf.measure_fwhm(psf_nb, ext='OVERDIST')
     # The PSF will still be a little broader in this case due to the IPC model, but not by a lot..
     assert fwhm_oversamp < fwhm_overdist <= 1.1 * fwhm_oversamp, "IFU broadening model should be disabled for this test case"
+
+    # test flux conservation with and without the IFU broadening. This should be a more exact match
+    assert np.isclose(psf_nb['DET_DIST'].data.sum(), psf['DET_DIST'].data.sum() ), "IFU broadening should not change total flux much"
+
+