Updates for PSF datamodel changes

jwst/extract_1d/psf_profile.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 from scipy import ndimage, optimize
-from stdatamodels.jwst.datamodels import MiriLrsPsfModel
+from stdatamodels.jwst.datamodels import SpecPsfModel
 
 from jwst.extract_1d.extract1d import extract1d
 from jwst.extract_1d.source_location import (
@@ -31,7 +31,7 @@ def open_psf(psf_refname, exp_type):
 
     Returns
     -------
-    psf_model : MiriLrsPsfModel
+    psf_model : SpecPsfModel
         Currently only works on MIRI LRS-FIXEDSLIT exposures.
         Returns the EPSF model.
 
@@ -42,16 +42,16 @@ def open_psf(psf_refname, exp_type):
         # super sample factor: psf_model.meta.psf.subpix)
         # psf : psf_model.data (2d)
         # wavelength of PSF planes: psf_model.wave
-        psf_model = MiriLrsPsfModel(psf_refname)
+        psf_model = SpecPsfModel(psf_refname)
 
     else:
-        # So far, only MIRI LRS has a PSF datamodel defined. Try to use it
-        # for the input model.
+        # So far, only MIRI LRS has a PSF datamodel defined. For any other
+        # exposure type, try to use the model MIRI LRS uses to open the input model
         try:
-            psf_model = MiriLrsPsfModel(psf_refname)
+            psf_model = SpecPsfModel(psf_refname)
         except (ValueError, AttributeError):
             raise NotImplementedError(f'PSF file for EXP_TYPE {exp_type} '
-                                      f'could not be read as MiriLrsPsfModel.') from None
+                                      f'could not be read as SpecPsfModel.') from None
     return psf_model 
 
 
@@ -262,11 +262,6 @@ def psf_profile(input_model, trace, wl_array, psf_ref_name,
     else:
         trace = trace[y0:y1]
 
-    # Scale the trace location to the subsampled psf
-    psf_subpix = psf_model.meta.psf.subpix
-    psf_location = trace - bbox[0][0]
-    psf_shift = psf_model.meta.psf.center_col - (psf_location * psf_subpix)
-
     # Check if we need to add a negative nod pair trace
     nod_offset = None
     if model_nod_pair:
@@ -294,11 +289,16 @@ def psf_profile(input_model, trace, wl_array, psf_ref_name,
     cutout_shape = cutout.shape
     _y, _x = np.mgrid[:cutout_shape[0], :cutout_shape[1]]
 
-    # Add the wavelength and spatial shifts to the coordinates to map to
+    # Scale the trace location to the subsampled psf and
+    # add the wavelength and spatial shifts to the coordinates to map to
+    psf_subpix = psf_model.meta.psf.subpix
+    psf_location = trace - bbox[0][0]
     if dispaxis == HORIZONTAL:
+        psf_shift = psf_model.meta.psf.center_row - (psf_location * psf_subpix)
         xidx = wave_idx
         yidx = _y * psf_subpix + psf_shift
     else:
+        psf_shift = psf_model.meta.psf.center_col - (psf_location * psf_subpix)
         xidx = _x * psf_subpix + psf_shift[:, None]
         yidx = wave_idx
 

jwst/extract_1d/tests/conftest.py

@@ -483,11 +483,12 @@ def nirspec_fs_apcorr_file(tmp_path, nirspec_fs_apcorr):
 
 @pytest.fixture()
 def psf_reference():
-    psf_model = dm.MiriLrsPsfModel()
+    psf_model = dm.SpecPsfModel()
     psf_model.data = np.ones((50, 50), dtype=float)
     psf_model.wave = np.linspace(0, 10, 50)
     psf_model.meta.psf.subpix = 1.0
     psf_model.meta.psf.center_col = 25
+    psf_model.meta.psf.center_row = 25
     yield psf_model
     psf_model.close()
 
@@ -501,13 +502,14 @@ def psf_reference_file(tmp_path, psf_reference):
 
 @pytest.fixture()
 def psf_reference_with_source():
-    psf_model = dm.MiriLrsPsfModel()
+    psf_model = dm.SpecPsfModel()
     psf_model.data = np.full((50, 50), 1e-6)
     psf_model.data[:, 24:27] += 1.0
 
     psf_model.wave = np.linspace(0, 10, 50)
     psf_model.meta.psf.subpix = 1.0
     psf_model.meta.psf.center_col = 25
+    psf_model.meta.psf.center_row = 25
     yield psf_model
     psf_model.close()
 

jwst/extract_1d/tests/test_psf_profile.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 import pytest
-from stdatamodels.jwst.datamodels import MiriLrsPsfModel
+from stdatamodels.jwst.datamodels import SpecPsfModel
 
 from jwst.extract_1d import psf_profile as pp
 from jwst.tests.helpers import LogWatcher
@@ -22,10 +22,10 @@ def log_watcher(monkeypatch):
 @pytest.mark.parametrize('exp_type', ['MIR_LRS-FIXEDSLIT', 'NRS_FIXEDSLIT', 'UNKNOWN'])
 def test_open_psf(psf_reference_file, exp_type):
     # for any exptype, a model that can be read
-    # as MiriLrsPsfModel will be, since it's the only
+    # as SpecPsfModel will be, since it's the only
     # one implemented so far
     with pp.open_psf(psf_reference_file, exp_type=exp_type) as model:
-        assert isinstance(model, MiriLrsPsfModel)
+        assert isinstance(model, SpecPsfModel)
 
 
 def test_open_psf_fail():