Extend trace model to NIRSpec and MIRI slit and slitless modes

Author: melanieclarke

jwst/adaptive_trace_model/adaptive_trace_model_step.py

@@ -95,24 +95,58 @@ def process(self, input_data):
         else:
             models = [output_model]
 
+        # Update each model in place
         for model in models:
-            if not isinstance(model, datamodels.IFUImageModel):
-                log.warning("The adaptive_trace_model step is only implemented for IFU data.")
-                log.warning("Skipping processing for datamodel type %s.", str(output_model))
+            log.info("Fitting trace model for %s", model.meta.filename)
+            if isinstance(model, datamodels.MultiSlitModel):
+                results = None
+                if self.save_intermediate_results:
+                    new_model = datamodels.MultiSlitModel()
+                    new_model.update(model, only="PRIMARY")
+                    if self.oversample == 1.0:
+                        results = [None, new_model, new_model.copy(), None, None]
+                    else:
+                        results = [
+                            None,
+                            new_model,
+                            new_model.copy(),
+                            new_model.copy(),
+                            new_model.copy(),
+                        ]
+                for slit in model.slits:
+                    log.info(f"Working on slit {slit.name}")
+                    log.debug(f"Slit is of type {type(slit)}")
+
+                    slit_results = fit_and_oversample(
+                        slit,
+                        fit_threshold=self.fit_threshold,
+                        slope_limit=self.slope_limit,
+                        oversample_factor=self.oversample,
+                        psf_optimal=self.psf_optimal,
+                        return_intermediate_models=self.save_intermediate_results,
+                    )
+                    if self.save_intermediate_results:
+                        for i, intermediate_model in enumerate(slit_results[1:]):
+                            if intermediate_model is not None:
+                                results[i + 1].slits.append(intermediate_model)
+
+            elif isinstance(model, (datamodels.SlitModel, datamodels.IFUImageModel)):
+                results = fit_and_oversample(
+                    model,
+                    fit_threshold=self.fit_threshold,
+                    slope_limit=self.slope_limit,
+                    oversample_factor=self.oversample,
+                    psf_optimal=self.psf_optimal,
+                    return_intermediate_models=self.save_intermediate_results,
+                )
+            else:
+                log.warning(
+                    "The adaptive_trace_model step is not implemented for %s.", str(output_model)
+                )
+                log.warning("Skipping processing.")
                 model.meta.cal_step.adaptive_trace_model = "SKIPPED"
                 continue
 
-            # Update the model in place
-            log.info("Fitting trace model for %s", model.meta.filename)
-            results = fit_and_oversample(
-                model,
-                fit_threshold=self.fit_threshold,
-                slope_limit=self.slope_limit,
-                oversample_factor=self.oversample,
-                psf_optimal=self.psf_optimal,
-                return_intermediate_models=self.save_intermediate_results,
-            )
-
             model.meta.cal_step.adaptive_trace_model = "COMPLETE"
             if self.save_intermediate_results:
                 _, full_spline, used_spline, linear, residual = results

jwst/adaptive_trace_model/tests/helpers.py

@@ -2,11 +2,16 @@
 from stdatamodels.jwst import datamodels
 
 from jwst.assign_wcs.assign_wcs_step import AssignWcsStep
-from jwst.assign_wcs.tests.test_miri import create_hdul
+from jwst.assign_wcs.tests.test_miri import (
+    create_datamodel_cube,
+    create_hdul,
+    create_hdul_lrs_slitless,
+)
 from jwst.assign_wcs.tests.test_nirspec import create_nirspec_fs_file, create_nirspec_ifu_file
 from jwst.extract_2d.extract_2d_step import Extract2dStep
 
 __all__ = [
+    "miri_lrs_slitless_model_with_source",
     "miri_mrs_model",
     "miri_mrs_model_with_source",
     "nirspec_ifu_model",
@@ -16,6 +21,39 @@
 ]
 
 
+def miri_lrs_slitless_model_with_source():
+    """
+    Create a mock MIRI LRS slitless model with a simple spectral source in the data array.
+
+    Returns
+    -------
+    model : `~stdatamodels.jwst.datamodels.SlitModel`
+        The LRS slitless datamodel.
+    """
+    shape = (5, 416, 72)
+    hdul = create_hdul_lrs_slitless()
+    cube_model = create_datamodel_cube(hdul, shape)
+    hdul.close()
+
+    model = datamodels.SlitModel(cube_model)
+    cube_model.close()
+
+    model.data = np.full(shape, np.nan)
+    model.err = np.zeros(shape)
+    model.dq = np.zeros(shape, dtype=np.uint32)
+    model.var_poisson = np.zeros(shape)
+    model.var_rnoise = np.zeros(shape)
+
+    ysize, xsize = shape[-2:]
+    x, y = np.meshgrid(np.arange(xsize), np.arange(ysize))
+    _, _, lam = model.meta.wcs(x, y)
+
+    region_map = (~np.isnan(lam)).astype(int)
+    _add_source(model, region_map, along_x=False)
+
+    return model
+
+
 def miri_mrs_model(detector="MIRIFUSHORT", channel="12", band="SHORT", shape=(1024, 1032)):
     """
     Create a mock MIRI MRS model.
@@ -141,9 +179,9 @@ def nirspec_ifu_model_with_source(wcs_style="coordinates"):
     return model
 
 
-def nirspec_slit_model_with_source():
+def nirspec_slit_model():
     """
-    Create a mock NIRSpec FS model with a simple spectral source in the data array.
+    Create a mock NIRSpec FS model with no source in the data array.
 
     Calls assign_wcs and extract_2d.
 
@@ -157,8 +195,8 @@ def nirspec_slit_model_with_source():
     hdul.close()
 
     shape = (2048, 2048)
-    model.data = np.full(shape, np.nan)
-    model.err = np.zeros(shape)
+    model.data = np.ones(shape)
+    model.err = model.data * 0.01
     model.dq = np.zeros(shape, dtype=np.uint32)
     model.var_poisson = np.zeros(shape)
     model.var_rnoise = np.zeros(shape)
@@ -172,6 +210,23 @@ def nirspec_slit_model_with_source():
             slit.meta.bunit_data = "MJy"
         else:
             slit.meta.bunit_data = "MJy/sr"
+
+    return model
+
+
+def nirspec_slit_model_with_source():
+    """
+    Create a mock NIRSpec FS model with a simple spectral source in the data array.
+
+    Calls assign_wcs and extract_2d.
+
+    Returns
+    -------
+    model : `~stdatamodels.jwst.datamodels.MultiSlitModel`
+        The FS datamodel.
+    """
+    model = nirspec_slit_model()
+    for slit in model.slits:
         region_map = (~np.isnan(slit.wavelength)).astype(int)
         _add_source(slit, region_map)
 

jwst/adaptive_trace_model/tests/test_adaptive_trace_model_step.py

@@ -37,6 +37,20 @@ def nirspec_ifu_slice_wcs():
     model.close()
 
 
+@pytest.fixture(scope="module")
+def nirspec_slit_model():
+    model = helpers.nirspec_slit_model()
+    yield model
+    model.close()
+
+
+@pytest.fixture(scope="module")
+def nirspec_slit_model_with_source():
+    model = helpers.nirspec_slit_model_with_source()
+    yield model
+    model.close()
+
+
 @pytest.fixture()
 def asn_input(tmp_path, miri_mrs_model):
     """
@@ -74,6 +88,13 @@ def asn_input(tmp_path, miri_mrs_model):
     return asn
 
 
+@pytest.fixture(scope="module")
+def miri_lrs_slitless_model_with_source():
+    model = helpers.miri_lrs_slitless_model_with_source()
+    yield model
+    model.close()
+
+
 @pytest.mark.parametrize("dataset", ["miri_mrs_model", "nirspec_ifu_model_with_source"])
 def test_adaptive_trace_model_step_success(request, dataset):
     model = request.getfixturevalue(dataset)
@@ -239,7 +260,7 @@ def test_adaptive_trace_model_step_with_container(miri_mrs_model):
 def test_adaptive_trace_model_unsupported_model(caplog):
     model = datamodels.ImageModel()
     result = AdaptiveTraceModelStep.call(model)
-    assert "only implemented for IFU" in caplog.text
+    assert "not implemented for <ImageModel>" in caplog.text
 
     assert result is not model
     assert result.meta.cal_step.adaptive_trace_model == "SKIPPED"
@@ -290,9 +311,11 @@ def test_save_container_asn_id_missing(tmp_path, asn_input):
         assert output.exists()
 
 
+@pytest.mark.parametrize("dataset", ["miri_mrs_model", "nirspec_slit_model"])
 @pytest.mark.parametrize("oversample", [1.0, 2.0])
-def test_adaptive_trace_model_step_save_intermediate(tmp_path, miri_mrs_model, oversample):
-    model = miri_mrs_model
+def test_adaptive_trace_model_step_save_intermediate(tmp_path, request, dataset, oversample):
+    model = request.getfixturevalue(dataset).copy()
+    model.meta.filename = "test_input_cal.fits"
     AdaptiveTraceModelStep.call(
         model,
         oversample=oversample,
@@ -304,20 +327,98 @@ def test_adaptive_trace_model_step_save_intermediate(tmp_path, miri_mrs_model, o
 
     # Check for expected files
     expected = [
-        "test12SHORT_atm.fits",
-        "test12SHORT_spline_full.fits",
-        "test12SHORT_spline_used.fits",
+        "test_input_atm.fits",
+        "test_input_spline_full.fits",
+        "test_input_spline_used.fits",
     ]
+    # For a model with no source, the spline models are expected to be empty
     expect_empty = [False, True, True]
     if oversample > 1:
         # Extra files expected if oversampling is done
-        expected.extend(["test12SHORT_linear_interp.fits", "test12SHORT_spline_residual.fits"])
+        expected.extend(["test_input_linear_interp.fits", "test_input_spline_residual.fits"])
         expect_empty.extend([False, True])
     for filename, is_empty in zip(expected, expect_empty):
         assert (tmp_path / filename).exists()
-        with datamodels.open(str(tmp_path / filename)) as model:
-            assert isinstance(model, datamodels.IFUImageModel)
+        with datamodels.open(str(tmp_path / filename)) as new_model:
+            assert isinstance(new_model, type(model))
+            if isinstance(new_model, datamodels.MultiSlitModel):
+                data = new_model.slits[0].data
+            else:
+                data = new_model.data
             if is_empty:
-                assert np.all(np.isnan(model.data))
+                assert np.all(np.isnan(data))
             else:
-                assert not np.all(np.isnan(model.data))
+                assert not np.all(np.isnan(data))
+
+
+def test_adaptive_trace_model_step_tso(miri_lrs_slitless_model_with_source):
+    model = miri_lrs_slitless_model_with_source
+    result = AdaptiveTraceModelStep.call(model, oversample=1)
+    assert result.meta.cal_step.adaptive_trace_model == "COMPLETE"
+
+    # data is unchanged with oversample=1
+    np.testing.assert_equal(result.data, model.data)
+
+    # trace_model is attached, contains non-NaN trace from the median image
+    assert result.trace_model.shape == result.data.shape[-2:]
+    indx = ~np.isnan(result.data[0]) & ~np.isnan(result.trace_model)
+    assert np.all(np.isnan(result.trace_model[~indx]))
+    assert np.all(~np.isnan(result.trace_model[indx]))
+
+    # fit trace is a reasonable model of the slit but not perfect
+    atol = 0.25 * np.nanmax(model.data)
+    np.testing.assert_allclose(result.data[0, indx], result.trace_model[indx], atol=atol)
+
+    result.close()
+
+
+def test_adaptive_trace_model_step_tso_oversample(miri_lrs_slitless_model_with_source):
+    model = miri_lrs_slitless_model_with_source
+    with pytest.raises(ValueError, match="Oversampling is not supported for TSO data"):
+        AdaptiveTraceModelStep.call(model, oversample=2)
+
+
+def test_adaptive_trace_model_step_oversample_slit(nirspec_slit_model_with_source):
+    model = nirspec_slit_model_with_source
+    result = AdaptiveTraceModelStep.call(model, oversample=2, slope_limit=0.05, fit_threshold=0.0)
+    assert result.meta.cal_step.adaptive_trace_model == "COMPLETE"
+    assert isinstance(result, datamodels.MultiSlitModel)
+
+    # data is twice the size of the input along the x axis
+    extnames = ["data", "dq", "err", "var_poisson", "var_rnoise", "var_flat"]
+    input_models = model.slits
+    output_models = result.slits
+    for input_model, output_model in zip(input_models, output_models, strict=True):
+        for extname in extnames:
+            # check for extension presence
+            if not input_model.hasattr(extname):
+                assert not output_model.hasattr(extname)
+                continue
+
+            # check that shape is expected
+            input_ext = getattr(input_model, extname)
+            output_ext = getattr(output_model, extname)
+            assert output_ext.shape == (input_ext.shape[0] * 2, input_ext.shape[1])
+
+        # trace_model is attached, contains non-NaN trace for the one bright slit only,
+        # and only over the core of the source trace
+        assert output_model.trace_model.shape == output_model.data.shape
+
+        indx = ~np.isnan(output_model.wavelength)
+        assert np.all(np.isnan(output_model.trace_model[~indx]))
+        assert np.sum(~np.isnan(output_model.trace_model[indx])) > 0.15 * np.sum(indx)
+
+        # fit trace is a reasonable model of the slice but not identical -
+        # the slice is mostly linearly interpolated
+        valid = indx & ~np.isnan(output_model.data) & ~np.isnan(output_model.trace_model)
+        atol = 0.25 * np.nanmax(input_model.data)
+        if output_model.meta.bunit_data == "MJy":
+            # for flux density units, we need to account for flux conservation,
+            # due to pixel size change
+            factor = 2.0
+        else:
+            factor = 1.0
+        np.testing.assert_allclose(
+            output_model.data[valid] * factor, output_model.trace_model[valid], atol=atol
+        )
+    result.close()