JP-3930: Step function for opening models (#9723)

Author: melanieclarke

changes/9723.stpipe.rst

@@ -0,0 +1 @@
+Add a ``prepare_output`` method to the JwstStep class, to support input datamodel handling and conditional copies within processing steps.

jwst/assign_mtwcs/assign_mtwcs_step.py

@@ -35,13 +35,17 @@ def process(self, input_lib):
         `~jwst.datamodels.library.ModelLibrary`
             The modified data models.
         """
-        if not isinstance(input_lib, ModelLibrary):
+        # Open the input, making a copy if necessary
+        output_lib = self.prepare_output(input_lib)
+        if not isinstance(output_lib, ModelLibrary):
             try:
-                input_lib = ModelLibrary(input_lib)
-            except Exception:
+                output_lib = ModelLibrary(output_lib)
+            except (ValueError, TypeError) as err:
                 log.warning("Input data type is not supported.")
-                record_step_status(input_lib, "assign_mtwcs", False)
-                return input_lib
+                log.debug(f"Error was: {err}")
+                record_step_status(output_lib, "assign_mtwcs", False)
+                return output_lib
 
-        result = assign_moving_target_wcs(input_lib)
+        result = assign_moving_target_wcs(output_lib)
+        record_step_status(result, "assign_mtwcs", True)
         return result

jwst/assign_mtwcs/tests/test_mtwcs.py

@@ -3,7 +3,7 @@
 from stdatamodels.jwst import datamodels
 
 from jwst.assign_mtwcs import AssignMTWcsStep
-from jwst.datamodels import ModelLibrary
+from jwst.datamodels import ModelContainer, ModelLibrary
 
 
 @pytest.mark.parametrize("errtype", ["ra", "dec", "both", "none"])
@@ -70,3 +70,47 @@ def test_mt_slitmodel(errtype):
         zero = result.borrow(0)
         assert zero.meta.wcs.output_frame.name == expected_frame
         result.shelve(zero, 0, modify=False)
+
+
+@pytest.mark.parametrize("success", [True, False])
+def test_output_is_not_input(monkeypatch, success):
+    """
+    Test that input is not modified by the step.
+
+    This is specific to the use case of calling the step on non-library
+    model input.  When the input is already a ModelLibrary, it's assumed
+    that performance is the most important thing and extra copies are
+    not desired.
+    """
+    # Mock a failure in the ModelLibrary init, to exercise the "skipped" condition
+    if not success:
+
+        def raise_error(*args, **kwargs):
+            raise ValueError("test")
+
+        monkeypatch.setattr(ModelLibrary, "__init__", raise_error)
+
+    file_path = get_pkg_data_filename("data/test_mt_asn.json", package="jwst.assign_mtwcs.tests")
+    with datamodels.open(file_path) as container:
+        result = AssignMTWcsStep.call(container)
+        if success:
+            assert isinstance(result, ModelLibrary)
+        else:
+            assert isinstance(result, ModelContainer)
+        with result:
+            for im, input_im in zip(result, container):
+                if success:
+                    assert im.meta.cal_step.assign_mtwcs == "COMPLETE"
+                else:
+                    assert im.meta.cal_step.assign_mtwcs == "SKIPPED"
+                assert im is not input_im
+                assert input_im.meta.cal_step.assign_mtwcs is None
+
+                if success:
+                    result.shelve(im, modify=False)
+
+
+def test_input_not_supported(caplog):
+    input_data = datamodels.ImageModel()
+    AssignMTWcsStep.call(input_data)
+    assert "Input data type is not supported" in caplog.text

jwst/datamodels/container.py

@@ -310,7 +310,7 @@ def from_asn(self, asn_data):
                     if attr in RECOGNIZED_MEMBER_FIELDS:
                         if attr == "tweakreg_catalog":
                             if val.strip():
-                                val = asn_dir / val
+                                val = str(asn_dir / val)
                             else:
                                 val = None
 

jwst/outlier_detection/outlier_detection_step.py

@@ -76,11 +76,17 @@ def process(self, input_data):
                         `~jwst.datamodels.library.ModelLibrary`
             The modified input data with DQ flags set for detected outliers.
         """
-        # Open the input data, making a copy as needed.
-        result_models = self._open_models(input_data)
-
         # determine the "mode" (if not set by the pipeline)
         mode = self._guess_mode(input_data)
+
+        # Open the input data, making a copy as needed.
+        if mode != "imaging":
+            result_models = self.prepare_output(input_data)
+        else:
+            # Skip loading datamodels into memory in this case - allow the
+            # ModelLibrary to handle it, later
+            result_models = self.prepare_output(input_data, open_models=False)
+
         if mode is None:
             record_step_status(result_models, "outlier_detection", False)
             return result_models
@@ -233,48 +239,3 @@ def _get_asn_id(self, input_models):
             self._make_output_path = partial(_make_output_path, asn_id=asn_id)
         log.info(f"Outlier Detection asn_id: {asn_id}")
         return
-
-    def _open_models(self, input_models):
-        """
-        Open the input data, making a copy if necessary.
-
-        If the input data is a filename or path, it is opened
-        and the open model is returned.
-
-        If it is a list of models, it is opened as a ModelContainer.
-        In this case, or if the input is a simple datamodel or a
-        ModelContainer, a deep copy of the model/container is returned,
-        in order to avoid modifying the input models.
-
-        If the input is a ModelLibrary, it is simply returned, in order
-        to avoid making unnecessary copies for performance-critical
-        use cases.
-
-        Parameters
-        ----------
-        input_models : str, list, JwstDataModel, ModelContainer, or ModelLibrary
-            Input data to open.
-
-        Returns
-        -------
-        JwstDataModel, ModelContainer, or ModelLibrary
-            The opened datamodel(s).
-        """
-        # Check whether input contains datamodels
-        make_copy = False
-        if isinstance(input_models, list):
-            is_datamodel = [isinstance(m, datamodels.JwstDataModel) for m in input_models]
-            if any(is_datamodel):
-                make_copy = True
-        elif isinstance(input_models, (datamodels.JwstDataModel, ModelContainer)):
-            make_copy = True
-
-        if not isinstance(input_models, (datamodels.JwstDataModel, ModelLibrary, ModelContainer)):
-            # Input might be a filename or path.
-            # In that case, open it.
-            input_models = datamodels.open(input_models)
-        if make_copy:
-            # For regular models, make a copy to avoid modifying the input.
-            # Leave libraries alone for memory management reasons.
-            input_models = input_models.copy()
-        return input_models

jwst/skymatch/skymatch_step.py

@@ -71,10 +71,15 @@ def process(self, input_models):
         ModelLibrary
             A library of datamodels with the skymatch step applied.
         """
-        if isinstance(input_models, ModelLibrary):
-            library = input_models
+        # Check the input for open models and make a copy if necessary
+        # to avoid modifying input data.
+        # If there are no open models already, do not open them.  Leave
+        # that to the ModelLibrary call below.
+        output_models = self.prepare_output(input_models, open_models=False)
+        if isinstance(output_models, ModelLibrary):
+            library = output_models
         else:
-            library = ModelLibrary(input_models, on_disk=not self.in_memory)
+            library = ModelLibrary(output_models, on_disk=not self.in_memory)
 
         # Method: "user". Use user-provided sky values, and bypass skymatch() altogether.
         if self.skymethod == "user":

jwst/skymatch/tests/test_skymatch.py

@@ -208,13 +208,13 @@ def test_skymatch(nircam_rate, skymethod, subtract, skystat, match_down, grouped
     )
 
     if skymethod == "match" and grouped:
-        # nothing to "match" when there is only one group:
-        assert im.meta.background.method is None
-        assert im.meta.background.subtracted is None
-
         # test that output models have original sky levels on failure:
         with result:
             for im, lev in zip(result, levels):
+                # nothing to "match" when there is only one group:
+                assert im.meta.background.method is None
+                assert im.meta.background.subtracted is None
+
                 assert abs(np.mean(im.data[dq_mask]) - lev) < 0.01
                 result.shelve(im, modify=False)
 
@@ -603,3 +603,37 @@ def test_user_skyfile(tmp_cwd, nircam_rate, subtract):
 
     with pytest.raises(ValueError):
         SkyMatchStep.call(container, skymethod="user", skylist=skyfile)
+
+
+@pytest.mark.parametrize("success", [True, False])
+def test_output_is_not_input(nircam_rate, success):
+    """
+    Test that input is not modified by the step.
+
+    This is specific to the use case of calling the step on non-library
+    model input.  When the input is already a ModelLibrary, it's assumed
+    that performance is the most important thing and extra copies are
+    not desired.
+    """
+    im1 = nircam_rate.copy()
+    im2 = im1.copy()
+    im3 = im1.copy()
+    if success:
+        # Set sequence IDs to allow skymatch to proceed.
+        # Without this, the step will be skipped.
+        im2.meta.observation.sequence_id = "2"
+        im3.meta.observation.sequence_id = "3"
+    container = ModelContainer([im1, im2, im3])
+
+    result = SkyMatchStep.call(container)
+
+    with result:
+        for im, input_im in zip(result, container):
+            if success:
+                assert im.meta.cal_step.skymatch == "COMPLETE"
+            else:
+                assert im.meta.cal_step.skymatch == "SKIPPED"
+            assert im is not input_im
+            assert input_im.meta.cal_step.skymatch is None
+
+            result.shelve(im, modify=False)

jwst/stpipe/core.py

@@ -145,6 +145,92 @@ def load_as_level3_asn(self, obj):
         update_key_value(asn, "expname", (), mod_func=self.make_input_path)
         return asn
 
+    def prepare_output(self, init, make_copy=None, open_models=True, **kwargs):
+        """
+        Open the input data as a model, making a copy if necessary.
+
+        If the input data is a filename or path, it is opened
+        and the open model is returned.
+
+        If it is a list of models, it is opened as a ModelContainer.
+        In this case, or if the input is a simple datamodel or a
+        ModelContainer, a deep copy of the model/container is returned,
+        in order to avoid modifying the input models.
+
+        If the input is a ModelLibrary, it is simply returned, in order
+        to avoid making unnecessary copies for performance-critical
+        use cases.
+
+        All copies are skipped if this step has a parent (i.e. it is
+        called as part of a pipeline).
+
+        Set make_copy explicitly to True or False to override the above
+        behavior.
+
+        Parameters
+        ----------
+        init : str, list, JwstDataModel, ModelContainer, or ModelLibrary
+            Input data to open.
+        make_copy : bool or None
+            If True, a copy of the input will always be made.
+            If False, a copy will never be made.  If None, a copy is
+            conditionally made, depending on the input and whether the
+            step is called in a standalone context.
+        open_models : bool
+            If True and the input is a filename or list of filenames,
+            then datamodels.open will be called to open the input.
+            If False, the input is returned as is.
+        **kwargs
+            Additional keyword arguments to pass to datamodels.open. Used
+            only if the input is a str or list.
+
+        Returns
+        -------
+        JwstDataModel, ModelContainer, or ModelLibrary
+            The opened datamodel(s).
+
+        Raises
+        ------
+        TypeError
+            If make_copy=True and the input is a type that cannot be copied.
+        """
+        # Check whether input contains datamodels
+        copy_needed = False
+        if isinstance(init, list):
+            is_datamodel = [isinstance(m, datamodels.JwstDataModel) for m in init]
+            if any(is_datamodel):
+                # Make the list into a ModelContainer, since it contains models
+                init = ModelContainer(init)
+                copy_needed = True
+        elif isinstance(init, (datamodels.JwstDataModel, ModelContainer)):
+            copy_needed = True
+
+        # Input might be a filename or path.
+        # In that case, open it if desired.
+        if not isinstance(init, (datamodels.JwstDataModel, ModelLibrary, ModelContainer)):
+            if open_models:
+                input_models = datamodels.open(init, **kwargs)
+            else:
+                input_models = init
+        else:
+            # Use the init model directly.
+            input_models = init
+
+        # Make a copy if needed
+        if make_copy is None:
+            make_copy = copy_needed and self.parent is None
+        if make_copy:
+            try:
+                input_models = input_models.copy()
+            except AttributeError:
+                # This should only happen if make_copy is explicitly set to
+                # True and the input is a string or a ModelLibrary.
+                raise TypeError(
+                    f"Copy is not possible for input type {type(input_models)}"
+                ) from None
+
+        return input_models
+
     def finalize_result(self, result, reference_files_used):
         """
         Update the result with the software version and reference files used.