Merge branch 'main' into update/pre-commit

Author: tapastro

changes/9570.pipeline.rst

@@ -0,0 +1 @@
+Use ``read_metadata`` where possible in calwebb_image3; result is a few percent runtime improvement if in_memory=False

jwst/assign_mtwcs/moving_target_wcs.py

@@ -51,19 +51,12 @@ def assign_moving_target_wcs(input_models):
     mt_valid = True
     with input_models:
         for i in ind:
-            model = input_models.borrow(i)
-            if model.meta.wcsinfo.mt_ra is None or model.meta.wcsinfo.mt_dec is None:
-                mt_valid = False
-            else:
-                mt_ra[i] = model.meta.wcsinfo.mt_ra
-                mt_dec[i] = model.meta.wcsinfo.mt_dec
-
-            if mt_valid and isinstance(model, datamodels.MultiSlitModel):
-                for slit in model.slits:
-                    if slit.meta.wcsinfo.mt_ra is None or slit.meta.wcsinfo.mt_dec is None:
-                        mt_valid = False
-
-            input_models.shelve(model, i, modify=False)
+            meta = input_models.read_metadata(i, flatten=False)
+            mt_valid = _is_mt_meta_valid(meta)
+            if not mt_valid:
+                break
+            mt_ra[i] = meta["meta"]["wcsinfo"]["mt_ra"]
+            mt_dec[i] = meta["meta"]["wcsinfo"]["mt_dec"]
 
     # Compute the mean MT RA/Dec over all exposures
     if not mt_valid:
@@ -148,3 +141,37 @@ def add_mt_frame(wcs, ra_average, dec_average, mt_ra, mt_dec):
     pipeline.append((mt, None))
     new_wcs = WCS(pipeline)
     return new_wcs
+
+
+def _is_mt_meta_valid(meta):
+    """
+    Check if the metadata contains valid moving target RA/DEC.
+
+    Checks both the top-level wcsinfo as well as the wcsinfo for all
+    the slits in a MultiSlitModel.
+
+    Parameters
+    ----------
+    meta : dict
+        Nested metadata dictionary from a data model, as output from `read_metadata`
+        with `flatten=False`.
+
+    Returns
+    -------
+    bool
+        True if valid, False otherwise.
+    """
+    # check all the slits
+    for slit in meta.get("slits", []):
+        if (
+            slit["meta"]["wcsinfo"].get("mt_ra", None) is None
+            or slit["meta"]["wcsinfo"].get("mt_dec", None) is None
+        ):
+            return False
+    # check the top-level wcsinfo
+    if (
+        meta["meta"]["wcsinfo"].get("mt_ra", None) is None
+        or meta["meta"]["wcsinfo"].get("mt_dec", None) is None
+    ):
+        return False
+    return True

jwst/datamodels/library.py

@@ -1,8 +1,13 @@
 import warnings
 from pathlib import Path
+import numpy as np
+from datetime import datetime
+from astropy.time import Time
+from asdf.tags.core import NDArrayType
+
 from stdatamodels.jwst.datamodels.util import open as datamodels_open
 from stdatamodels.jwst.datamodels import read_metadata
-from stpipe.library import AbstractModelLibrary, NoGroupID
+from stpipe.library import AbstractModelLibrary, NoGroupID, BorrowError
 
 from jwst.associations import AssociationNotValidError, load_asn
 from jwst.datamodels.utils import attrs_to_group_id
@@ -218,25 +223,161 @@ def get_crds_parameters(self):
                 )
                 idx = 0
 
-            # find model in _loaded_models, temp_filenames, or asn_dir
-            if self._on_disk:
-                if idx in self._temp_filenames:
-                    # if model has been modified, find its temp filename
-                    filename = self._temp_filenames[idx]
-                else:
-                    # otherwise, find the filename in the asn_dir
-                    member = self._members[idx]
-                    filename = Path(self._asn_dir) / member["expname"]
+            return self.read_metadata(idx)
+
+    def read_metadata(self, idx, flatten=True):
+        """
+        Read metadata for a model at the given index.
+
+        Parameters
+        ----------
+        idx : int
+            Index of the model in the library.
+
+        Returns
+        -------
+        dict
+            The metadata dictionary for the model.
+        """
+        # if model was already borrowed, calling code should just read the metadata
+        # from the model itself, so we raise an error here
+        if idx in self._ledger:
+            raise BorrowError("Attempt to read metadata from model that is already borrowed")
+
+        # find model in _loaded_models, temp_filenames, or asn_dir
+        if self._on_disk:
+            if idx in self._temp_filenames:
+                # if model has been modified, find its temp filename
+                filename = self._temp_filenames[idx]
             else:
-                if idx in self._loaded_models:
-                    # if this model is in memory, retrieve parameters from it directly
-                    model = self._loaded_models[idx]
-                    return model.get_crds_parameters()
-                else:
-                    # otherwise, find the filename in the asn_dir
-                    member = self._members[idx]
-                    filename = Path(self._asn_dir) / member["expname"]
+                # otherwise, find the filename in the asn_dir
+                member = self._members[idx]
+                filename = Path(self._asn_dir) / member["expname"]
+        else:
+            if idx in self._loaded_models:
+                # if this model is in memory, retrieve parameters from it directly
+                model = self._loaded_models[idx]
+                return _read_meta_from_open_model(model, flatten=flatten)
+            else:
+                # otherwise, find the filename in the asn_dir
+                member = self._members[idx]
+                filename = Path(self._asn_dir) / member["expname"]
 
-            meta = read_metadata(filename, flatten=True)
+        meta = read_metadata(filename, flatten=flatten)
+        meta = self._assign_member_to_meta(meta, self._members[idx], flatten)
+        return meta
 
+    def _assign_member_to_meta(self, meta, member, flatten):
+        """
+        Update meta dict with asn-related attributes, similar to _assign_member_to_model.
+
+        Parameters
+        ----------
+        meta : dict
+            The metadata dictionary to update.
+        member : dict
+            The member dictionary containing association attributes.
+        flatten : bool
+            If True, the metadata will be flattened to a single level.
+            If False, the metadata will be nested.
+
+        Returns
+        -------
+        dict
+            The updated metadata dictionary with association attributes.
+        """
+        if flatten:
+            meta["meta.asn.exptype"] = member["exptype"]
+            for attr in ("group_id", "tweakreg_catalog"):
+                if attr in member:
+                    meta["meta." + attr] = member[attr]
+            if "table_name" in self.asn.keys():
+                meta["meta.asn.table_name"] = self.asn["table_name"]
+            if "asn_pool" in self.asn.keys():
+                meta["meta.asn.pool_name"] = self.asn["asn_pool"]
+        else:
+            meta["meta"].setdefault("asn", {})
+            meta["meta"]["asn"]["exptype"] = member["exptype"]
+            for attr in ("group_id", "tweakreg_catalog"):
+                if attr in member:
+                    meta["meta"][attr] = member[attr]
+            if "table_name" in self.asn.keys():
+                meta["meta"]["asn"]["table_name"] = self.asn["table_name"]
+            if "asn_pool" in self.asn.keys():
+                meta["meta"]["asn"]["pool_name"] = self.asn["asn_pool"]
         return meta
+
+
+def _read_meta_from_open_model(model, flatten):
+    """
+    Read metadata from an open model.
+
+    Parameters
+    ----------
+    model : DataModel
+        The model from which to read metadata.
+    flatten : bool
+        If True, the metadata will be flattened to a single level.
+        If False, the metadata will be nested.
+
+    Returns
+    -------
+    dict
+        The metadata dictionary for the model.
+    """
+    if flatten:
+        return model.get_crds_parameters()
+
+    def convert_val(val):
+        if isinstance(val, datetime):
+            return val.isoformat()
+        elif isinstance(val, Time):
+            return str(val)
+        return val
+
+    def recurse(tree):
+        """
+        Walk the tree and add metadata to a new dictionary.
+
+        Do all conversions and exclusions model.get_crds_parameters does,
+        but without flattening.
+
+        Parameters
+        ----------
+        tree : dict
+            An asdf-like nested dictionary structure, e.g. model._instance.
+
+        Returns
+        -------
+        dict
+            A new nested dictionary containing the metadata.
+        """
+        new_tree = {}
+        for key, val in tree.items():
+            if key == "wcs":
+                continue
+            val = convert_val(val)
+            if isinstance(val, dict):
+                new_tree[key] = recurse(val)
+            elif isinstance(val, list):
+                new_list = []
+                for item in val:
+                    if isinstance(item, dict):
+                        new_list.append(recurse(item))
+                    elif isinstance(item, list):
+                        # handle nested lists
+                        new_list.append(recurse({"_idx": item})["_idx"])  # placeholder key
+                    elif isinstance(item, (np.ndarray, NDArrayType)):
+                        continue
+                    else:
+                        new_list.append(convert_val(item))
+                new_tree[key] = new_list
+            elif isinstance(val, (np.ndarray, NDArrayType)):
+                continue
+            elif isinstance(val, (str, int, float, complex, bool)):
+                new_tree[key] = val
+            else:
+                continue  # skip unsupported types
+        return new_tree
+
+    return recurse(model._instance)  # noqa: SLF001