JP-4192: Prepare output for resample (#10145)

melanieclarke · web-flow · commit d447c6c30505 · 2026-01-16T16:40:14.000Z
diff --git a/jwst/resample/resample_spec_step.py b/jwst/resample/resample_spec_step.py
@@ -53,70 +53,79 @@ def process(self, input_data):
         SlitModel or MultiSlitModel
             The resampled output, one slit per source.
         """
-        with datamodels.open(input_data) as input_new:
-            # Check if input_new is a MultiSlitModel
-            model_is_msm = isinstance(input_new, MultiSlitModel)
-
-            #  If input is a 3D rateints MultiSlitModel (unsupported) skip the step
-            if model_is_msm and len((input_new[0]).shape) == 3:
-                log.warning("Resample spec step will be skipped")
-                result = input_new.copy()
-                result.meta.cal_step.resample = "SKIPPED"
-
-                return result
-
-            # Convert ImageModel to SlitModel (needed for MIRI LRS)
-            if isinstance(input_new, ImageModel):
-                input_new = datamodels.SlitModel(input_new)
-
-            if isinstance(input_new, ModelContainer):
-                input_models = input_new
-
-                try:
-                    output = input_models.meta.asn_table.products[0].name
-                except AttributeError:
-                    # NIRSpec MOS data goes through this path, as the container
-                    # is only ModelContainer-like, and doesn't have an asn_table
-                    # attribute attached.  Output name handling gets done in
-                    # _process_multislit() via the update method
-                    # TODO: the container-like object should retain asn_table
-                    output = None
-            else:
-                input_models = ModelContainer([input_new])
-                output = input_new.meta.filename
-                self.blendheaders = False
+        output_model = self.prepare_output(input_data)
 
-            # Setup drizzle-related parameters
-            kwargs = self.get_drizpars()
-            kwargs["output"] = output
-            self.drizpars = kwargs
+        # Check if input model is a MultiSlitModel
+        model_is_msm = isinstance(output_model, MultiSlitModel)
 
-            # Call resampling
-            if isinstance(input_models[0], MultiSlitModel):
-                result = self._process_multislit(input_models)
+        #  If input is a 3D rateints MultiSlitModel (unsupported) skip the step
+        if model_is_msm and len((output_model[0]).shape) == 3:
+            log.warning("Resample spec step will be skipped")
+            output_model.meta.cal_step.resample = "SKIPPED"
+            return output_model
 
-            elif len(input_models[0].data.shape) != 2:
-                # resample can only handle 2D images, not 3D cubes, etc
-                raise TypeError(f"Input {input_models[0]} is not a 2D image.")
+        # Convert ImageModel to SlitModel (may be needed for older MIRI LRS data)
+        if isinstance(output_model, ImageModel):
+            slit_model = datamodels.SlitModel(output_model)
+            if output_model is not input_data:
+                output_model.close()
+            output_model = slit_model
 
-            else:
-                # result is a SlitModel
-                result = self._process_slit(input_models)
-
-            # Update ASNTABLE in output
-            result.meta.cal_step.resample = "COMPLETE"
-            result.meta.asn.table_name = input_models[0].meta.asn.table_name
-            result.meta.asn.pool_name = input_models[0].meta.asn.pool_name
-
-            # populate the result wavelength attribute for MultiSlitModel
-            if isinstance(result, MultiSlitModel):
-                for slit_idx, _slit in enumerate(result.slits):
-                    wl_array = get_wavelengths(result.slits[slit_idx])
-                    result.slits[slit_idx].wavelength = wl_array
-            else:
-                # populate the result wavelength attribute for SlitModel
-                wl_array = get_wavelengths(result)
-                result.wavelength = wl_array
+        if isinstance(output_model, ModelContainer):
+            input_models = output_model
+
+            try:
+                output = input_models.meta.asn_table.products[0].name
+            except AttributeError:
+                # NIRSpec MOS data goes through this path, as the container
+                # is only ModelContainer-like, and doesn't have an asn_table
+                # attribute attached.  Output name handling gets done in
+                # _process_multislit() via the update method
+                # TODO: the container-like object should retain asn_table
+                output = None
+        else:
+            input_models = ModelContainer([output_model])
+            output = output_model.meta.filename
+            self.blendheaders = False
+
+        # Setup drizzle-related parameters
+        kwargs = self.get_drizpars()
+        kwargs["output"] = output
+        self.drizpars = kwargs
+
+        # Call resampling
+        if isinstance(input_models[0], MultiSlitModel):
+            result = self._process_multislit(input_models)
+
+        elif len(input_models[0].data.shape) != 2:
+            # resample can only handle 2D images, not 3D cubes, etc
+            raise TypeError(f"Input {input_models[0]} is not a 2D image.")
+
+        else:
+            # result is a SlitModel
+            result = self._process_slit(input_models)
+
+        # Update ASNTABLE in output
+        result.meta.cal_step.resample = "COMPLETE"
+        result.meta.asn.table_name = input_models[0].meta.asn.table_name
+        result.meta.asn.pool_name = input_models[0].meta.asn.pool_name
+
+        # populate the result wavelength attribute for MultiSlitModel
+        if isinstance(result, MultiSlitModel):
+            for slit_idx, _slit in enumerate(result.slits):
+                wl_array = get_wavelengths(result.slits[slit_idx])
+                result.slits[slit_idx].wavelength = wl_array
+        else:
+            # populate the result wavelength attribute for SlitModel
+            wl_array = get_wavelengths(result)
+            result.wavelength = wl_array
+
+        # Output is a new datamodel.
+        # Clean up the input model(s) if they were opened here.
+        if output_model is not input_data:
+            del output_model
+        if input_models is not input_data:
+            del input_models
 
         return result
 
diff --git a/jwst/resample/resample_step.py b/jwst/resample/resample_step.py
@@ -1,9 +1,8 @@
 import logging
 
 from stdatamodels import filetype
-from stdatamodels.jwst import datamodels as dm
 
-from jwst.datamodels import ImageModel, ModelLibrary  # type: ignore[attr-defined]
+from jwst.datamodels import ImageModel, ModelContainer, ModelLibrary  # type: ignore[attr-defined]
 from jwst.lib.pipe_utils import match_nans_and_flags
 from jwst.resample import resample
 from jwst.resample.resample_utils import load_custom_wcs
@@ -70,19 +69,25 @@ def process(self, input_data):
         WCS parameters such as ``output_shape`` (now computed from by
         ``output_wcs.bounding_box``) and ``crpix``.
         """
-        if isinstance(input_data, str):
-            ext = filetype.check(input_data)
-            if ext in ("fits", "asdf"):
-                input_data = dm.open(input_data)
-        if isinstance(input_data, ModelLibrary):
-            input_models = input_data
-        elif isinstance(input_data, (str, dict, list)):
-            input_models = ModelLibrary(input_data, on_disk=not self.in_memory)
-        elif isinstance(input_data, ImageModel):
-            input_models = ModelLibrary([input_data], on_disk=not self.in_memory)
-            output = input_data.meta.filename
+        # Make a copy if needed for an input model.
+        # Don't open filenames if they're not already models --
+        # leave it to the ModelLibrary call below to open them.
+        input_model = self.prepare_output(input_data, open_models=False)
+
+        if isinstance(input_model, ModelLibrary):
+            # Input is already a library: leave it alone.
+            input_models = input_model
+        elif isinstance(input_model, ImageModel) or (
+            isinstance(input_model, str) and filetype.check(input_model) in ["fits", "asdf"]
+        ):
+            # Input is a single file: pass it to ModelLibrary in a list
+            input_models = ModelLibrary([input_model], on_disk=not self.in_memory)
             self.blendheaders = False
+        elif isinstance(input_model, (str, dict, list, ModelContainer)):
+            # Input is an association or list of models/files
+            input_models = ModelLibrary(input_model, on_disk=not self.in_memory)
         else:
+            # Input is not recognized
             raise TypeError(f"Input {input_data} is not a 2D image.")
 
         try:
@@ -141,6 +146,13 @@ def process(self, input_data):
             )
             result = resamp.resample_many_to_one()
 
+        # The output is a new datamodel.
+        # Clean up the input model(s) if they were opened here.
+        if input_model is not input_data:
+            del input_model
+        if input_models is not input_data:
+            del input_models
+
         return result
 
     @staticmethod
diff --git a/jwst/resample/tests/test_resample_step.py b/jwst/resample/tests/test_resample_step.py
@@ -435,7 +435,15 @@ def test_single_image_file_input(nircam_rate, tmp_cwd):
     # Create a temporary file with the input data
     im = AssignWcsStep.call(nircam_rate, sip_approx=False)
     im.meta.filename = "test_input.fits"
+
+    # Add a bad pixel in the error plane, not matched with a NaN in data.
+    # This will test if the match_nans_and_flags call at the beginning
+    # of the step modifies the input model.
+    im.err[0, 0] = np.nan
+
+    # Save a copy to disk, keep a copy of the data for testing.
     im.save("test_input.fits")
+    im_copy = im.data.copy()
 
     # Run the step on the file
     result_from_memory = ResampleStep.call(im)
@@ -449,12 +457,52 @@ def test_single_image_file_input(nircam_rate, tmp_cwd):
     # Check that input model was not modified
     assert im is not result_from_memory
     assert im.meta.cal_step.resample is None
+    assert_allclose(im.data, im_copy)
 
     result_from_file.close()
     result_from_memory.close()
     im.close()
 
 
+def test_list_model_input(nircam_rate, tmp_cwd):
+    """Ensure step can be run on a list of models without modifying them."""
+    # Create a temporary file with the input data
+    im = AssignWcsStep.call(nircam_rate, sip_approx=False)
+    im.meta.filename = "test_input_1.fits"
+
+    # Add a bad pixel in the error plane, not matched with a NaN in data.
+    # This will test if the match_nans_and_flags call at the beginning
+    # of the step modifies the input model.
+    im.err[0, 0] = np.nan
+
+    # Keep a copy of the data for testing
+    im_copy = im.data.copy()
+
+    # Make a list of input models to run
+    im2 = im.copy()
+    im2.meta.filename = "test_input_2.fits"
+    im_list = [im, im2]
+
+    # Run the step on the file
+    result = ResampleStep.call(im_list)
+
+    # Check that the output is as expected
+    assert isinstance(result, ImageModel)
+    assert result.meta.cal_step.resample == "COMPLETE"
+
+    # Check that input models were not modified
+    assert result is not im
+    assert result is not im2
+    assert im.meta.cal_step.resample is None
+    assert im2.meta.cal_step.resample is None
+    assert_allclose(im.data, im_copy)
+    assert_allclose(im2.data, im_copy)
+
+    result.close()
+    im.close()
+    im2.close()
+
+
 @pytest.mark.parametrize("ratio", [0.5, 0.7, 1.0])
 def test_pixel_scale_ratio_imaging(nircam_rate, ratio):
     im = AssignWcsStep.call(nircam_rate, sip_approx=False)
@@ -1616,15 +1664,33 @@ def test_nirspec_lamp_pixscale(nirspec_lamp, tmp_path):
     result4.close()
 
 
-def test_spec_output_is_not_input(nirspec_cal):
-    im = ResampleSpecStep.call(nirspec_cal)
+@pytest.mark.parametrize("input_list", [True, False])
+def test_spec_input_not_modified(nirspec_cal, input_list):
+    # Add a bad pixel in the error plane of one slit, not matched
+    # with a NaN in data.
+    # This will test if the match_nans_and_flags call at the beginning
+    # of the step modifies the input model.
+    nirspec_cal.slits[0].err[15, 100] = np.nan
+    data_copy = nirspec_cal.slits[0].data.copy()
+
+    if input_list:
+        input_models = [nirspec_cal, nirspec_cal.copy()]
+    else:
+        input_models = nirspec_cal
+
+    im = ResampleSpecStep.call(input_models)
 
     # Step is complete
     assert im.meta.cal_step.resample == "COMPLETE"
 
     # Input is not modified
     assert im is not nirspec_cal
     assert nirspec_cal.meta.cal_step.resample is None
+    if input_list:
+        for model in input_models:
+            assert_allclose(model.slits[0].data, data_copy)
+    else:
+        assert_allclose(input_models.slits[0].data, data_copy)
 
 
 def test_spec_skip_cube():
