Added in reindex_axis and handling for tdets (#1506)

mjrand · web-flow · commit df8f1c493c94 · 2025-12-22T12:54:30.000-05:00
* Added in reindex_axis and handling for tdets

Added in reindex_axis function in the AxisManager class which allows one to reindex all datasets that are assigned to a specified axis in an AxisManager. This is useful for dealing with tdets in get_obs.

Also added in handling for tdets in get_obs. The handling merges the dets+tdets axis, inserts the tdets signal into the dets signal dataset, and reindexes every dataset in the obs_aman that is assigned to the dets axis (det_info... etc.). These axes will have nans inserted where the tdets exist in the new merged dets axis.

* updated order of axis existence check

Moved the axis existence check to happen before any aman copying takes place (if in_place=False). This will prevent an aman copy from being made right before a ValueError which would cause a memory leak in certain circumstances.

* Fixed tdet_id cleanup

Fixed the portion of get_obs that added tdet_ids into the dets axis. Previously it would only operate on the highest level aman's dets axis so the sub-amans dets axes wouldn't match.

It's been rewritten to recursively (scary) delve through all lower-level amans to fix all dets axes.

* Fixed default reindexing behavior and load_book

Removed the default behavior of removing the tdets axis and data. This data will now never be removed.

Added in `reindex_dets` arg to `get_obs`. `special_dets=True` will now load load the special dets while `reindex_dets=True` will reindex them into the `dets` axis. This will copy the special dets data into the dets signal, bands, channels, and readout_ids but will not remove the tdets data (in case anyone just wants to look at fixed tones).

Used Kyohei's code for the special_dets `load_book` fix.

* Added check for tones and det_info

Added check for reindexing in `get_obs` to ensure tones and det_info data exists before attempting to reindex. Reindexing is logged and skipped if either are nonexistant

* Changed logger behavior for missing tones

Changed the logger behavior from `error` to `info` if get_obs is called with `special_channels=True` and `reindex_dets=True` when there are no tones to load.

Minor changes for if statements that retain the same behavior.

* Fixed np.full bug

Previous commit introduced a small bug on a `np.full` call. Fixed this issue.
diff --git a/sotodlib/core/axisman.py b/sotodlib/core/axisman.py
@@ -809,6 +809,144 @@ def restrict_axes(self, axes, in_place=True):
                 dest._fields[k] = v[sslice]
         return dest
 
+    def reindex_axis(self, axis, indexes, in_place=True):
+        """
+        Reindexes all data that is assigned to a specified axis
+        with a new list/array of indexes.
+        This is particularly useful if the number of detectors
+        between the meta and obs data don't match.
+        This function will recursively delve through all
+        AxisManagers in aman and will reindex every
+        data array that is found assigned to an axis
+        matching the specified axis.
+
+        Args:
+            axis (str): The name of the axis in the aman to reindex.
+            indexes (int array): an array of ints with length
+                equal to the length of the new array
+                and values equal to the idxs of the
+                values in the data to be reindexed.
+                Indexes that should be left as nan in
+                the new array should be set to -1 or nan.
+
+            For example:
+                data = [1,3,5], indexes = [0, -1, 2, 1]
+                would result in new_data = [1, nan, 5, 3]
+            
+            in_place (bool): If in_place == True, the intersection is
+            applied to self.  Otherwise, a new object is returned,
+            with data copied out.
+        """
+        # Check if axis even exists first
+        if axis not in self._axes.keys():
+            raise ValueError(f"Axis doesn't exist in aman! \
+                             Can't re-index along {axis}")
+
+        if in_place:
+            aman = self
+        else:
+            aman = self.copy(axes_only=True)
+            aman._assignments.update(self._assignments)
+
+        # Loop through ever assignment and reindex along
+        # each that is tied to the axis in question
+        new_axes = {}
+        reindexed_vs = {}
+        assignments = list(aman._assignments.keys())
+        for assignment in assignments:
+            axes = aman._assignments[assignment]
+            # If this assignment isn't connected to our axis
+            # we can skip it.
+            if axis not in axes:
+                continue
+
+            v = aman[assignment]
+
+            if isinstance(v, AxisManager):
+                # If we hit an axis manager,
+                # recursively reindex it as well. Scary!
+                new_v = v.reindex_axis(axis, indexes)
+
+            else:
+                # By this point we have a non AxisManager
+                # assignment assigned to only our axis.
+                # Build new array with the correct indexes.
+                shape = [len(indexes)]
+                if isinstance(v, np.ndarray):
+                    for s in np.shape(v)[1:]:
+                        shape.append(s)
+
+                new_v = np.empty(shape, dtype=v.dtype)
+                if isinstance(v.dtype, float):
+                    # Fill any float arrays with nans
+                    # Non float arrays may have weird
+                    # behavior for newly added indexes. 
+                    # Oh well.
+                    new_v *= np.nan  
+
+                for i, index in enumerate(indexes):
+                    if np.isnan(index) or not (0 <= index < len(v)):
+                        continue
+
+                    new_v[i] = v[int(index)]
+
+            reindexed_vs[assignment] = new_v
+            new_axes[assignment] = np.array(axes)
+
+            # Destroy the old assignment
+            aman.move(name=assignment, new_name=None)
+
+        old_axis = aman._axes[axis]
+
+        # Recreate the axis
+        if isinstance(old_axis, IndexAxis):
+            # Build a new axis that has a length equal to the indexes arg.
+            new_axis = IndexAxis(name=axis, count=len(indexes))
+
+        if isinstance(old_axis, LabelAxis):
+            # A LabelAxis dtype may vary by length,
+            # we'll insert empty values for the newly added idxs.
+            # This will produce empty strings
+            # ('') for det_ids, readout_ids, etc.
+            # It may produce strange behavior
+            # for non string like objects. Be careful!
+            vals = np.empty(len(indexes), dtype=old_axis.vals.dtype)
+            for i, index in enumerate(indexes):
+                if np.isnan(index) or not (0 <= int(index) < len(old_axis.vals)):
+                    continue
+                vals[i] = old_axis.vals[int(index)]
+
+            new_axis = LabelAxis(name=axis, vals=vals)
+
+        if isinstance(old_axis, OffsetAxis):
+            new_axis = OffsetAxis(count=len(indexes),
+                                  offset=old_axis.offset,
+                                  origin_tag=old_axis.origin_tag)
+
+        # We're done with this old axis now, destroy it.
+        del aman._axes[axis]
+        # Add in the reindexed axis.
+        aman.add_axis(new_axis)
+
+        # Now we'll go through all the reindexed data and wrap it back in.
+        for assignment, axes in new_axes.items():
+            # Build the axis map for wrapping the data.
+            ax_map = []
+            for i, ax in enumerate(axes):
+                # Axis map looks like a list of numbered tuples.
+                ax_map.append((i, ax))
+
+            vs = reindexed_vs[assignment]
+            # Need to wrap aman's with no axismap
+            if isinstance(vs, AxisManager):
+                aman.wrap(name=assignment, data=vs)
+
+            else:  # Everything else needs an axismap
+                aman.wrap(name=assignment, data=vs, axis_map=ax_map)
+
+        # Everything is now reindexed and rewrapped. Done!
+        return aman  # Return for rewrapping if recursively called.
+
     @staticmethod
     def _broadcast_selector(sslice):
         """sslice is a list of selectors, which will typically be slice(), or
diff --git a/sotodlib/core/context.py b/sotodlib/core/context.py
@@ -174,6 +174,7 @@ def get_obs(self,
                 no_signal=None,
                 no_headers=None,
                 special_channels=None,
+                reindex_dets=None,
                 loader_type=None,
     ):
         """Load TOD and supporting metadata for some observation.
@@ -217,10 +218,20 @@ def get_obs(self,
             information that tags along with the signal.
           special_channels (bool): If True, load "special" readout
             channels that are normally skipped (e.g. fixed tones).
+          reindex_dets (bool) If True, reindexes the axismanager
+            "dets" axis to include any special dets that may have
+            been loaded. All special det signals, bands, channels, 
+            and readout_ids will be inserted along the dets axis
+            and respective data arrays. Does not destroy special 
+            dets axes like the "tdets" axis. WARNING: ~Doubles
+            run time and instantaenous memory usage as the signal
+            dataset has to effectively be copied! Memory returned
+            after reindexing is complete.
           loader_type (str): Name of the registered TOD loader
             function to use (this will override whatever is specified
             in context.yaml).
 
+
         Notes:
           It is acceptable to pass the ``obs_id`` argument by position
           (first), but all other arguments should be passed by
@@ -329,6 +340,97 @@ def get_obs(self,
                         _det_info.move(k, None)
                 meta.det_info.merge(_det_info)
             aman.merge(meta)
+
+        # Deal with special channels, if they exist.
+        # We will merge the dets and tdets axes
+        # And merge the tdets signal into the dets signal, band, and channels.
+        # nans will be inserted into every other dataset assigned to the dets axis.
+        if special_channels and reindex_dets:
+            # Check for special channels (tones).
+            if 'tones' not in aman:
+                logger.info('"tones" not found in aman, no special channels to reindex!')
+
+            # Check there is det_info to reindex.
+            elif 'det_info' not in aman:
+                logger.error('"det_info" not found in aman, no dets to reindex!')
+
+            # Both tones and det_info exist.
+            else:
+                # Grab all band and channel info for dets + tdets
+                det_bands = aman.det_info.smurf.band
+                det_channels = aman.det_info.smurf.channel
+                tdet_bands = aman.tones.band
+                tdet_channels = aman.tones.channel
+    
+                # Create a sorted array of dets + tdets
+                special_band_ch = [(b, c) for b, c in zip(tdet_bands, tdet_channels)]
+                normal_band_ch = [(b, c) for b, c in zip(det_bands, det_channels)]
+                band_ch = np.array(sorted(normal_band_ch + special_band_ch))
+    
+                # Grab the det idxs from the det band + channels
+                det_indexes = np.full(len(band_ch), np.nan)
+                for i, (b, c) in enumerate(band_ch):
+                    w = np.where((det_bands == b) & (det_channels == c))[0]
+                    if len(w) == 0:
+                        continue
+    
+                    det_indexes[i] = w[0]
+    
+                # Grab the tdet idxs from the tdet band + channels
+                tdet_indexes = np.full(len(band_ch), np.nan)
+                for i, (b, c) in enumerate(band_ch):
+                    w = np.where((tdet_bands == b) & (tdet_channels == c))[0]
+                    if len(w) == 0:
+                        continue
+    
+                    tdet_indexes[i] = w[0]
+    
+                # Use the det idxs to reindex all datasets assigned to the dets axis
+                # This will set the len to dets + tdets
+                # And will insert nans where the tdet channels exist
+                # in the sorted band_ch array
+                aman.reindex_axis(axis='dets', indexes=det_indexes)
+    
+                # Finally use the tdet idxs to fill in the tdet data
+                # For the signal, band, and channels
+                for i, tidx in enumerate(tdet_indexes):
+                    if np.isnan(tidx):
+                        continue
+    
+                    aman.signal[i] = aman.tones.signal[int(tidx)]
+                    aman.det_info.smurf.channel[i] = aman.tones.channel[int(tidx)]
+                    aman.det_info.smurf.band[i] = aman.tones.band[int(tidx)]
+    
+                def add_tdet_ids(aman, tdet_indexes, tdet_ids):
+                    """
+                        Small Function for recursively delving through
+                        all amans in the arg aman to add tdet ids into
+                        the dets axes.
+                    """
+                    # Check all assignments for any AxisManagers
+                    for assignment, axes in aman._assignments.items():
+                        if isinstance(aman[assignment], AxisManager) and "dets" in axes:
+                            # If they have a Dets axis, go fix it first.
+                            add_tdet_ids(aman[assignment], tdet_indexes, tdet_ids)
+    
+                    # Skip current aman if it doesn't have a dets axis.
+                    # This is the recursion base case
+                    # Assuming someone didn't make an infinitely deep aman
+                    # Or circular amans....
+                    if "dets" not in aman._axes.keys():
+                        return
+    
+                    # If this aman has a dets axis, add in tdet ids.
+                    for i, tidx in enumerate(tdet_indexes):
+                        if np.isnan(tidx):
+                            continue
+                        aman.dets.vals[i] = tdet_ids[int(tidx)]
+    
+                # Add in tdet ids to all amans.
+                add_tdet_ids(aman, tdet_indexes, aman.tdets.vals)
+    
+                # obs_aman tdet info is merged!
+
         return aman
 
     def get_meta(self,
diff --git a/sotodlib/io/load_book.py b/sotodlib/io/load_book.py
@@ -525,19 +525,20 @@ def _concat_filesets(results, ancil=None, timestamps=None,
                 # Should look like this: sch_NONE_2_326
                 b, c = map(int, k.split('_')[2:])
                 tone_info.append((v['stream_id'], v['stream_id'] + f'_{b}_{c}', b, c))
-        ts, tk, tb, tc = map(np.array, zip(*tone_info))
-        tman = core.AxisManager(core.LabelAxis('tdets', tk),
-                                aman.samps)
-        aman.wrap('tones', tman)
-        aman.tones.wrap('stream_id', ts, axis_map=[(0, 'tdets')])
-        aman.tones.wrap('band', tb, axis_map=[(0, 'tdets')])
-        aman.tones.wrap('channel', tc, axis_map=[(0, 'tdets')])
-        aman.tones.wrap_new('signal', shape=('tdets', 'samps'), dtype='float32')
-        dets_ofs = 0
-        for v in results.values():
-            d = v['tones'].data
-            aman.tones['signal'][dets_ofs:dets_ofs + len(d)] = d
-            dets_ofs += len(d)
+        if tone_info:
+            ts, tk, tb, tc = map(np.array, zip(*tone_info))
+            tman = core.AxisManager(core.LabelAxis('tdets', tk),
+                                    aman.samps)
+            aman.wrap('tones', tman)
+            aman.tones.wrap('stream_id', ts, axis_map=[(0, 'tdets')])
+            aman.tones.wrap('band', tb, axis_map=[(0, 'tdets')])
+            aman.tones.wrap('channel', tc, axis_map=[(0, 'tdets')])
+            aman.tones.wrap_new('signal', shape=('tdets', 'samps'), dtype='float32')
+            dets_ofs = 0
+            for v in results.values():
+                d = v['tones'].data
+                aman.tones['signal'][dets_ofs:dets_ofs + len(d)] = d
+                dets_ofs += len(d)
 
     # In sims, or if no_headers, the primary block may be unpopulated.
     if any([(v['primary'] is not None and v['primary'].data is not None)
