Update for PSF datamodel changes

Author: melanieclarke

jwst/extract_1d/psf_profile.py

@@ -18,40 +18,43 @@
 NOD_PAIR_PATTERN = ["ALONG-SLIT-NOD", "2-POINT-NOD"]
 
 
-def open_psf(psf_refname, exp_type):
+def open_psf(psf_refname, slit_name):
     """
     Open the PSF reference file.
 
     Parameters
     ----------
     psf_refname : str
         The name of the psf reference file.
-    exp_type : str
-        The exposure type of the data.
+    slit_name : str or None
+        The slit name for the data.
 
     Returns
     -------
     psf_model : SpecPsfModel
         Returns the EPSF model.
     """
-    if exp_type == "MIR_LRS-FIXEDSLIT":
-        # The information we read in from PSF file is:
-        # center_col: psf_model.meta.psf.center_col
-        # super sample factor: psf_model.meta.psf.subpix)
-        # psf : psf_model.data (2d)
-        # wavelength of PSF planes: psf_model.wave
+    try:
         psf_model = SpecPsfModel(psf_refname)
+    except (ValueError, AttributeError):
+        raise NotImplementedError(
+            f"PSF file {psf_refname} could not be read as SpecPsfModel."
+        ) from None
 
-    else:
-        # 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 = SpecPsfModel(psf_refname)
-        except (ValueError, AttributeError):
-            raise NotImplementedError(
-                f"PSF file for EXP_TYPE {exp_type} could not be read as SpecPsfModel."
-            ) from None
-    return psf_model
+    # Get the right PSF aperture
+    slit_name = str(slit_name).upper()
+    psf_aperture = None
+    for aperture in psf_model.apertures:
+        aper_name = str(aperture.name).upper()
+        if aper_name in ["NONE", "ANY"] or aper_name == slit_name:
+            psf_aperture = aperture
+            break
+    psf_model.close()
+
+    if psf_aperture is None:
+        raise ValueError(f"No matching aperture found in SpecPsfModel for slit={slit_name}")
+
+    return psf_aperture
 
 
 def _normalize_profile(profile, dispaxis):
@@ -253,9 +256,10 @@ def psf_profile(
         For PSF profiles, this is always set to the upper edge of the bounding box,
         since the full array may have non-zero weight.
     """
-    # Read in reference files
+    # Read in reference file
     exp_type = input_model.meta.exposure.type
-    psf_model = open_psf(psf_ref_name, exp_type)
+    slit_name = getattr(input_model, "name", None)
+    psf_model = open_psf(psf_ref_name, slit_name)
 
     # Get the data cutout
     data_shape = input_model.data.shape[-2:]
@@ -343,14 +347,14 @@ def psf_profile(
 
     # 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_subpix = psf_model.subpix
     psf_location = trace - bbox[0][0]
     if dispaxis == HORIZONTAL:
-        psf_shift = psf_model.meta.psf.center_row - (psf_location * psf_subpix)
+        psf_shift = psf_model.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)
+        psf_shift = psf_model.center_col - (psf_location * psf_subpix)
         xidx = _x * psf_subpix + psf_shift[:, None]
         yidx = wave_idx
 

jwst/extract_1d/tests/conftest.py

@@ -425,11 +425,12 @@ def psf_reference():
         The mock model.
     """
     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
+    psf_model.apertures.append({})
+    psf_model.apertures[0].data = np.ones((50, 50), dtype=float)
+    psf_model.apertures[0].wave = np.linspace(0, 10, 50)
+    psf_model.apertures[0].subpix = 1.0
+    psf_model.apertures[0].center_col = 25
+    psf_model.apertures[0].center_row = 25
     yield psf_model
     psf_model.close()
 
@@ -460,13 +461,14 @@ def psf_reference_with_source():
         The mock model.
     """
     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
+    psf_model.apertures.append({})
+    psf_model.apertures[0].data = np.full((50, 50), 1e-6)
+    psf_model.apertures[0].data[:, 24:27] += 1.0
+
+    psf_model.apertures[0].wave = np.linspace(0, 10, 50)
+    psf_model.apertures[0].subpix = 1.0
+    psf_model.apertures[0].center_col = 25
+    psf_model.apertures[0].center_row = 25
     yield psf_model
     psf_model.close()
 

jwst/extract_1d/tests/test_psf_profile.py

@@ -5,20 +5,28 @@
 from jwst.extract_1d import psf_profile as pp
 
 
-@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 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, SpecPsfModel)
+@pytest.mark.parametrize("slit_name", [None, "ANY", "UNKNOWN"])
+def test_open_psf(psf_reference_file, slit_name):
+    # for any slit name passed in, the first aperture will be matched,
+    # since there's only one available and it has no specific slit name
+    aperture = pp.open_psf(psf_reference_file, slit_name)
+    assert aperture.data.size != 0
+    assert aperture.wave.size != 0
 
 
 def test_open_psf_fail():
     with pytest.raises(NotImplementedError, match="could not be read"):
         pp.open_psf("bad_file", "UNKNOWN")
 
 
+def test_open_psf_no_aperture_found(psf_reference_file):
+    # If no PSF aperture matches the slit, an error is raised
+    with SpecPsfModel(psf_reference_file) as psf_model:
+        psf_model.apertures[0].name = "GOOD_SLIT"
+        with pytest.raises(ValueError, match="No matching aperture"):
+            pp.open_psf(psf_model, "BAD_SLIT")
+
+
 @pytest.mark.parametrize("dispaxis", [1, 2])
 def test_normalize_profile(nod_profile, dispaxis):
     profile = 2 * nod_profile
@@ -42,11 +50,12 @@ def test_normalize_profile_with_nans(nod_profile, dispaxis):
 
 @pytest.mark.parametrize("dispaxis", [1, 2])
 def test_make_cutout_profile_default(psf_reference, dispaxis):
-    data_shape = psf_reference.data.shape
+    psf_aper = psf_reference.apertures[0]
+    data_shape = psf_aper.data.shape
     yidx, xidx = np.mgrid[: data_shape[0], : data_shape[1]]
 
-    psf_subpix = psf_reference.meta.psf.subpix
-    profiles = pp._make_cutout_profile(xidx, yidx, psf_subpix, psf_reference.data, dispaxis)
+    psf_subpix = psf_aper.subpix
+    profiles = pp._make_cutout_profile(xidx, yidx, psf_subpix, psf_aper.data, dispaxis)
     assert len(profiles) == 1
     assert profiles[0].shape == data_shape
 
@@ -60,12 +69,13 @@ def test_make_cutout_profile_default(psf_reference, dispaxis):
 @pytest.mark.parametrize("dispaxis", [1, 2])
 @pytest.mark.parametrize("extra_shift", [1, 2])
 def test_make_cutout_profile_shift_down(psf_reference, dispaxis, extra_shift):
-    data_shape = psf_reference.data.shape
+    psf_aper = psf_reference.apertures[0]
+    data_shape = psf_aper.data.shape
     yidx, xidx = np.mgrid[: data_shape[0], : data_shape[1]]
-    psf_subpix = psf_reference.meta.psf.subpix
+    psf_subpix = psf_aper.subpix
 
     profiles = pp._make_cutout_profile(
-        xidx, yidx, psf_subpix, psf_reference.data, dispaxis, extra_shift=extra_shift
+        xidx, yidx, psf_subpix, psf_aper.data, dispaxis, extra_shift=extra_shift
     )
     assert len(profiles) == 1
     assert profiles[0].shape == data_shape
@@ -84,12 +94,13 @@ def test_make_cutout_profile_shift_down(psf_reference, dispaxis, extra_shift):
 @pytest.mark.parametrize("dispaxis", [1, 2])
 @pytest.mark.parametrize("extra_shift", [-1, -2])
 def test_make_cutout_profile_shift_up(psf_reference, dispaxis, extra_shift):
-    data_shape = psf_reference.data.shape
+    psf_aper = psf_reference.apertures[0]
+    data_shape = psf_aper.data.shape
     yidx, xidx = np.mgrid[: data_shape[0], : data_shape[1]]
-    psf_subpix = psf_reference.meta.psf.subpix
+    psf_subpix = psf_aper.subpix
 
     profiles = pp._make_cutout_profile(
-        xidx, yidx, psf_subpix, psf_reference.data, dispaxis, extra_shift=extra_shift
+        xidx, yidx, psf_subpix, psf_aper.data, dispaxis, extra_shift=extra_shift
     )
     assert len(profiles) == 1
     assert profiles[0].shape == data_shape
@@ -107,13 +118,14 @@ def test_make_cutout_profile_shift_up(psf_reference, dispaxis, extra_shift):
 
 @pytest.mark.parametrize("dispaxis", [1, 2])
 def test_make_cutout_profile_with_nod(psf_reference, dispaxis):
-    data_shape = psf_reference.data.shape
+    psf_aper = psf_reference.apertures[0]
+    data_shape = psf_aper.data.shape
     yidx, xidx = np.mgrid[: data_shape[0], : data_shape[1]]
-    psf_subpix = psf_reference.meta.psf.subpix
+    psf_subpix = psf_aper.subpix
 
     offset = 2
     profiles = pp._make_cutout_profile(
-        xidx, yidx, psf_subpix, psf_reference.data, dispaxis, nod_offset=offset
+        xidx, yidx, psf_subpix, psf_aper.data, dispaxis, nod_offset=offset
     )
     assert len(profiles) == 2
     source, nod = profiles
@@ -138,9 +150,10 @@ def test_make_cutout_profile_with_nod(psf_reference, dispaxis):
 
 @pytest.mark.parametrize("dispaxis", [1, 2])
 def test_profile_residual(psf_reference, dispaxis):
+    psf_aper = psf_reference.apertures[0]
     data_shape = (50, 50)
     yidx, xidx = np.mgrid[: data_shape[0], : data_shape[1]]
-    psf_subpix = psf_reference.meta.psf.subpix
+    psf_subpix = psf_aper.subpix
 
     # Set data to all ones, so residual should be zero
     # when background is not fit
@@ -149,16 +162,17 @@ def test_profile_residual(psf_reference, dispaxis):
 
     param = [0, None]
     residual = pp._profile_residual(
-        param, data, var, xidx, yidx, psf_subpix, psf_reference.data, dispaxis, fit_bkg=False
+        param, data, var, xidx, yidx, psf_subpix, psf_aper.data, dispaxis, fit_bkg=False
     )
     assert np.isclose(residual, 0.0)
 
 
 @pytest.mark.parametrize("dispaxis", [1, 2])
 def test_profile_residual_with_bkg(psf_reference, dispaxis):
+    psf_aper = psf_reference.apertures[0]
     data_shape = (50, 50)
     yidx, xidx = np.mgrid[: data_shape[0], : data_shape[1]]
-    psf_subpix = psf_reference.meta.psf.subpix
+    psf_subpix = psf_aper.subpix
 
     # Set data to all ones, so it is all background - residual
     # should be all of the data
@@ -167,7 +181,7 @@ def test_profile_residual_with_bkg(psf_reference, dispaxis):
 
     param = [0, None]
     residual = pp._profile_residual(
-        param, data, var, xidx, yidx, psf_subpix, psf_reference.data, dispaxis, fit_bkg=True
+        param, data, var, xidx, yidx, psf_subpix, psf_aper.data, dispaxis, fit_bkg=True
     )
     assert np.isclose(residual, np.sum(data**2 / var))
 

pyproject.toml

@@ -30,7 +30,8 @@ dependencies = [
     "scikit-image>=0.21.0",
     "scipy>=1.14.1",
     # "stdatamodels>=5.0.1,<6",
-    "stdatamodels @ git+https://github.com/spacetelescope/stdatamodels.git@main",
+    # "stdatamodels @ git+https://github.com/spacetelescope/stdatamodels.git@main",
+    "stdatamodels @ git+https://github.com/melanieclarke/stdatamodels.git@jp-4273",
     "stcal>=1.18,<2",
     "stpipe>=0.12,<2",
     "synphot>=1.3",