Merge branch 'master' into jp-3664-np2p0

CHANGES.rst

@@ -26,12 +26,16 @@ assign_wcs
 - Moved `update_s_region_imaging`, `update_s_region_keyword`, and `wcs_from_footprints`
   into stcal. [#8624]
 
+- Add helper functions to copy only the necessary parts of the WCS so that
+  these parts can be used within loops, avoiding copying the full WCS within
+  a loop [#8793]
+
 associations
 ------------
 
 - Restored slit name to level 3 product names for NIRSpec BOTS and background
   fixed slit targets. [#8699]
-  
+
 - Update warning message about use of paths in associations. [#8752]
 
 - Remove ``MultilineLogger`` and no longer set it as the default logger. [#8781]
@@ -70,12 +74,14 @@ cube_build
 
 - Removed direct setting of the ``self.skip`` attribute from within the step
   itself. [#8600]
-  
+
 - Fixed a bug when ``cube_build`` was called from the ``mrs_imatch`` step. [#8728]
 
 - Ensure that NaNs and DO_NOT_USE flags match up in all input data before
   building a cube. [#8557]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops. [#8793]
+
 datamodels
 ----------
 
@@ -102,6 +108,8 @@ flat_field
 - Ensure that NaNs and DO_NOT_USE flags match up in all science, error,
   variance, and DQ extensions for all modes. [#8557]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops [#8793]
+
 general
 -------
 
@@ -127,6 +135,8 @@ master_background
 - Either of ``"background"`` or ``"bkg"`` in slit name now defines the slit
   as a background slit, instead of ``"bkg"`` only. [#8600]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops [#8793]
+
 model_blender
 -------------
 
@@ -137,6 +147,11 @@ mrs_imatch
 
 - Added a deprecation warning and set the default to skip=True for the step. [#8728]
 
+msaflagopen
+-----------
+
+- Replaced deep copies of NIRSpec WCS objects within most loops. [#8793]
+
 nsclean
 -------
 
@@ -149,6 +164,8 @@ nsclean
   can still be called from the ``calwebb_spec2`` pipeline on NIRSpec rate
   data, but it is now deprecated. [#8669]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops. [#8793]
+
 outlier_detection
 -----------------
 
@@ -174,12 +191,16 @@ pathloss
 - Ensure that NaNs and DO_NOT_USE flags match up in all output science, error,
   variance, and DQ extensions. [#8557]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops [#8793]
+
 photom
 ------
 
 - Ensure that NaNs and DO_NOT_USE flags match up in all output science, error,
   variance, and DQ extensions. [#8557]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops. [#8793]
+
 pipeline
 --------
 
@@ -188,10 +209,17 @@ pipeline
   in memory or on disk. [#8683]
 
 - Updated ``calwebb_spec2`` to run ``nsclean`` on NIRSpec imprint and background 
-  association members. [#8786]
+  association members. [#8786, #8809]
 
 - Updated `calwebb_spec3` to not save the `pixel_replacement` output by default.[#8765]
 
+- Replaced deep copies of NIRSpec WCS objects within most loops. [#8793]
+
+pixel_replace
+-------------
+
+- Replaced deep copies of NIRSpec WCS objects within most loops. [#8793]
+
 ramp_fitting
 ------------
 
@@ -522,6 +550,7 @@ master_background
   wavelength range instead of NaN to avoid NaN-ing out entire
   sets of science data when backgrounds are missing. [#8597]
 
+
 master_background_mos
 ---------------------
 
@@ -598,6 +627,7 @@ photom
 - Added a hook to bypass the ``photom`` step when the ``extract_1d`` step
   was bypassed for non-TSO NIRISS SOSS exposures. [#8575]
 
+
 pipeline
 --------
 

jwst/assign_wcs/nirspec.py

@@ -6,6 +6,7 @@
 """
 import logging
 import numpy as np
+import copy
 
 from astropy.modeling import models
 from astropy.modeling.models import Mapping, Identity, Const1D, Scale, Tabular1D
@@ -1696,6 +1697,131 @@ def gwa_to_ymsa(msa2gwa_model, lam_cen=None, slit=None, slit_y_range=None):
     return tab
 
 
+def _get_transforms(input_model, slitnames, return_slits=False):
+
+    """
+    Return a WCS object with necessary transforms for all slits.
+
+    This function enables the JWST pipeline to avoid excessive deep
+    copying of WCS objects in later steps. It is used internally in
+    the pipeline only and should not be used if any of the WCSs is
+    modified.
+
+    Parameters
+    ----------
+    input_model : `~jwst.datamodels.JwstDataModel`
+        A data model with a WCS object for the all open slitlets in
+        an observation.
+    slitnames : list of int or str
+        Slit.name of all open slits.
+    return_slits : bool, optional
+        Return the open slits
+
+    Returns
+    -------
+    wcsobj : `~gwcs.wcs.WCS`
+        WCS object deep copied from input_model.meta.wcs
+    sca2gwa : `~astropy.modeling.core.Model`
+        Transform from ``sca`` to ``gwa``
+    gwa2slit : list of `~astropy.modeling.core.Model`
+        Transform from ``gwa`` to ``slit`` for each input slit
+    slit2slicer : list of `~astropy.modeling.core.Model`
+        Transform from ``slit_frame`` to ``slicer`` for each input slit
+    open_slits : list of `~stdatamodels.jwst.transforms.models.Slit`
+        open slits from wcs.get_transform('gwa', 'slit_frame').slits
+        Only returned if return_slits is True
+    """
+
+    wcs = copy.deepcopy(input_model.meta.wcs)
+
+    sca2gwa = copy.deepcopy(wcs.pipeline[1].transform[1:])
+    wcs.set_transform('sca', 'gwa', sca2gwa)
+
+    gwa2slit = [copy.deepcopy(wcs.pipeline[2].transform.get_model(slit_name))
+              for slit_name in slitnames]
+
+    slit2slicer = [copy.deepcopy(wcs.pipeline[3].transform.get_model(slit_name))
+              for slit_name in slitnames]
+
+    if return_slits:
+        g2s = wcs.get_transform('gwa', 'slit_frame')
+        open_slits = g2s.slits
+        return wcs, sca2gwa, gwa2slit, slit2slicer, copy.deepcopy(open_slits)
+    else:
+        return wcs, sca2gwa, gwa2slit, slit2slicer
+
+
+def _nrs_wcs_set_input_lite(input_model, input_wcs, slit_name, transforms,
+                           wavelength_range=None, open_slits=None,
+                           slit_y_low=None, slit_y_high=None):
+
+    """
+    Return a WCS object for a specific slit, slice or shutter
+
+    The lite version of the routine is distinguished from the legacy
+    routine because it does not make a deep copy of the input WCS object.
+
+    Parameters
+    ----------
+    input_model : `~jwst.datamodels.JwstDataModel`
+        A WCS object for the all open slitlets in an observation.
+    input_wcs : `~gwcs.wcs.WCS`
+        A WCS object for the all open slitlets in an observation.  This
+        will be modified and returned.
+    slit_name : int or str
+        Slit.name of an open slit.
+    transforms : list of `~astropy.modeling.core.Model`
+        Model transforms output from ``_get_transforms``
+    wavelength_range: list
+        Wavelength range for the combination of filter and grating.  Optional.
+    open_slits : list of slits
+        List of open slits.  Optional.
+
+    Returns
+    -------
+    wcsobj : `~gwcs.wcs.WCS`
+        WCS object for this slit.
+    """
+
+
+    def _get_y_range(input_model, open_slits):
+        if open_slits is None:
+            log_message = 'nrs_wcs_set_input_lite must be called with open_slits if not in ifu mode'
+            log.critical(log_message)
+            raise RuntimeError(log_message)
+        # Need the open slits to get the slit ymin,ymax
+        slit = [s for s in open_slits if s.name == slit_name][0]
+        return slit.ymin, slit.ymax
+
+    if wavelength_range is None:
+        _, wavelength_range = spectral_order_wrange_from_model(input_model)
+
+    slit_wcs = copy.copy(input_wcs)
+
+    slit_wcs.set_transform('sca', 'gwa', transforms[0])
+    slit_wcs.set_transform('gwa', 'slit_frame', transforms[1])
+
+    is_nirspec_ifu = is_nrs_ifu_lamp(input_model) or input_model.meta.exposure.type.lower() == 'nrs_ifu'
+
+    if is_nirspec_ifu:
+        slit_wcs.set_transform('slit_frame', 'slicer', transforms[2] & Identity(1))
+    else:
+        slit_wcs.set_transform('slit_frame', 'msa_frame', transforms[2] & Identity(1))
+
+    transform = slit_wcs.get_transform('detector', 'slit_frame')
+
+    if is_nirspec_ifu:
+        bb = compute_bounding_box(transform, wavelength_range)
+    else:
+        if slit_y_low is None or slit_y_high is None:
+            slit_y_low, slit_y_high = _get_y_range(input_model, open_slits)
+        bb = compute_bounding_box(transform, wavelength_range,
+                                  slit_ymin=slit_y_low, slit_ymax=slit_y_high)
+
+    slit_wcs.bounding_box = bb
+    return slit_wcs
+
+
 def _nrs_wcs_set_input(input_model, slit_name):
     """
     Returns a WCS object for a specific slit, slice or shutter.
@@ -1713,7 +1839,7 @@ def _nrs_wcs_set_input(input_model, slit_name):
     wcsobj : `~gwcs.wcs.WCS`
         WCS object for this slit.
     """
-    import copy
+
     wcsobj = input_model.meta.wcs
 
     slit_wcs = copy.deepcopy(wcsobj)