Toast improvements (#1505)

tskisner · web-flow · commit 550bc462bee1 · 2025-12-19T13:54:29.000-08:00
* Make coadd stats robust against bad wafer obs

* Update SOFocalplane HDF5 API to conform to upstream toast changes.

* Add support for save / load of HKManager to hdf5

* Fixes and unit test for serializing HKManager to HDF5

* Add downsample method to HKManager

* Add demodulation / downsample test to HKManager test
diff --git a/sotodlib/toast/hkmanager.py b/sotodlib/toast/hkmanager.py
@@ -1,12 +1,14 @@
 # Copyright (c) 2025-2025 Simons Observatory.
 # Full license can be found in the top level "LICENSE" file.
 
+import copy
 from collections.abc import MutableMapping
 
 import numpy as np
 from scipy.interpolate import CubicSpline, PchipInterpolator
 
 from toast.utils import Logger
+from toast.io.hdf_utils import load_meta_object, save_meta_object
 
 from ..io import hkdb
 
@@ -49,8 +51,8 @@ class HKManager(MutableMapping):
 
     def __init__(
         self,
-        comm,
-        timestamps,
+        comm=None,
+        timestamps=None,
         site_root=None,
         site_db=None,
         site_fields=None,
@@ -63,8 +65,12 @@ def __init__(
         self._internal = dict()
         self._aliases = dict()
         self._stamps = timestamps
-        self._start_time = timestamps[0]
-        self._stop_time = timestamps[-1]
+        if timestamps is not None:
+            self._start_time = timestamps[0]
+            self._stop_time = timestamps[-1]
+        elif site_root is not None or plat_root is not None:
+            msg = "If loading site or platform data, timestamps are required"
+            raise RuntimeError(msg)
         rank = 0
         if comm is not None:
             rank = comm.rank
@@ -194,13 +200,118 @@ def __repr__(self):
 
     def __eq__(self, other):
         log = Logger.get()
-        if self._internal != other._internal:
-            log.verbose(f"  data {self._internal} != {other._internal}")
+        if set(self._internal.keys()) != set(other._internal.keys()):
+            log.verbose(f"  keys {self._internal} != {other._internal}")
             return False
+        for k, v in self._internal.items():
+            try:
+                if not np.allclose(v, other._internal[k]):
+                    log.verbose(f"  data array {v} != {other._internal[k]}")
+                    return False
+            except Exception:
+                if self._internal[k] != other._internal[k]:
+                    log.verbose(f"  data value {v} != {other._internal[k]}")
+                    return False
         if self._aliases != other._aliases:
             log.verbose(f"  aliases {self._aliases} != {other._aliases}")
             return False
         return True
 
     def __ne__(self, other):
         return not self.__eq__(other)
+
+    def save_hdf5(self, handle, obs, **kwargs):
+        """Save the HKManager object to an HDF5 file.
+
+        Args:
+            handle (h5py.Group):  The group to populate.
+            obs (Observation):  The parent observation.
+
+        Returns:
+            None
+
+        """
+        gcomm = obs.comm.comm_group
+        if (gcomm is None) or (gcomm.rank == 0):
+            # The rank zero process should always be writing
+            if handle is None:
+                raise RuntimeError("HDF5 group is not open on the root process")
+        if handle is not None:
+            tdata = handle.create_dataset(
+                "timestamps", self._stamps.shape, dtype=self._stamps.dtype
+            )
+            hslices = [slice(0, x) for x in self._stamps.shape]
+            hslices = tuple(hslices)
+            tdata.write_direct(self._stamps, hslices, hslices)
+            del tdata
+            save_meta_object(handle, "data", self._internal)
+            save_meta_object(handle, "aliases", self._aliases)
+
+    def load_hdf5(self, handle, obs, **kwargs):
+        """Load the HKManager object from an HDF5 file.
+
+        Args:
+            handle (h5py.Group):  The group containing noise model.
+            obs (Observation):  The parent observation.
+
+        Returns:
+            None
+
+        """
+        gcomm = obs.comm.comm_group
+        if (gcomm is None) or (gcomm.rank == 0):
+            # The rank zero process should always be reading
+            if handle is None:
+                raise RuntimeError("HDF5 group is not open on the root process")
+        stamps = None
+        data = None
+        aliases = None
+        if handle is not None:
+            tdata = handle["timestamps"]
+            stamps = np.copy(tdata[:])
+            data = load_meta_object(handle["data"])
+            aliases = load_meta_object(handle["aliases"])
+            del tdata
+        if gcomm is not None:
+            stamps = gcomm.bcast(stamps, root=0)
+            data = gcomm.bcast(data, root=0)
+            aliases = gcomm.bcast(aliases, root=0)
+
+        self._stamps = stamps
+        self._start_time = stamps[0]
+        self._stop_time = stamps[-1]
+        self._internal = data
+        self._aliases = aliases
+
+    def duplicate(self):
+        """Create a copy of the HKManager"""
+        new_hk = HKManager()
+        new_hk._stamps = np.copy(self._stamps)
+        new_hk._start_time = new_hk._stamps[0]
+        new_hk._stop_time = new_hk._stamps[-1]
+        new_hk._internal = copy.deepcopy(self._internal)
+        new_hk._aliases = copy.deepcopy(self._aliases)
+        return new_hk
+
+    def downsample(self, demod_times):
+        """Downsample the HKManager
+
+        The demodulation operator with downsample any observation
+        attributes with a `downsample()` method.  For this class
+        that operation is trivial, since it just involves replacing
+        the timestamps.
+
+        Args:
+            demod_times (array):  The new downsampled timestamps.
+
+        Returns:
+            (HKManager):  The downsampled HKManager.
+
+        """
+        new_hk = HKManager()
+        new_hk._stamps = demod_times
+        new_hk._start_time = new_hk._stamps[0]
+        new_hk._stop_time = new_hk._stamps[-1]
+        new_hk._internal = copy.deepcopy(self._internal)
+        new_hk._aliases = copy.deepcopy(self._aliases)
+        return new_hk
diff --git a/sotodlib/toast/instrument.py b/sotodlib/toast/instrument.py
@@ -10,6 +10,8 @@
 
 import astropy.units as u
 from astropy.table import QTable
+
+import toast
 from toast.instrument import Focalplane, GroundSite, Telescope
 from toast.utils import Logger, name_UID
 from toast.weather import SimWeather
@@ -366,11 +368,13 @@ def get_par_float(ddata, key, default):
             )
             # Get noise parameters.  If detector-specific entries are
             # absent, use band averages
-            net_corr = get_par_float(
-                det_data, "NET_corr", band_data["NET_corr"]
+            net_corr = get_par_float(det_data, "NET_corr", band_data["NET_corr"])
+            net = (
+                get_par_float(det_data, "NET", band_data["NET"])
+                * 1.0e-6
+                * u.K
+                * u.s**0.5
             )
-            net = get_par_float(det_data, "NET", band_data["NET"]) \
-                * 1.0e-6 * u.K * u.s**0.5
             if apply_net_corr:
                 net *= net_corr
                 net_corr = 1.0
@@ -388,8 +392,9 @@ def get_par_float(ddata, key, default):
             bandcenters.append(0.5 * (lower + upper))
             bandwidths.append(upper - lower)
 
-        meta["platescale"] = hw.data["telescopes"][meta["telescope"]]["platescale"] \
-                             * u.deg / u.mm
+        meta["platescale"] = (
+            hw.data["telescopes"][meta["telescope"]]["platescale"] * u.deg / u.mm
+        )
 
         detdata = QTable(
             [
@@ -471,6 +476,29 @@ def get_par_float(ddata, key, default):
             sample_rate=sample_rate,
         )
 
+    @classmethod
+    def _load_hdf5(
+        cls, handle, comm=None, detectors=None, file_det_sets=None, **kwargs
+    ):
+        """Load simulated focalplane from HDF5.
+
+        Although SOFocalplane inherits from toast.instrument.Focalplane, it
+        essentially is just a custom constructor and the actual content is
+        stored in the Focalplane base class.  When loading data we just
+        dispatch to the base class method.
+        """
+        return toast.instrument.Focalplane._load_hdf5(
+            handle,
+            comm=comm,
+            detectors=detectors,
+            file_det_sets=file_det_sets,
+            **kwargs,
+        )
+
+    def _save_hdf5(self, handle, comm=None, **kwargs):
+        """Save to HDF5 - dispatch to base class method."""
+        return super()._save_hdf5(handle, comm=comm, **kwargs)
+
 
 def update_creation_time(det_data, creation_time):
     """Update the readout_id column of a focalplane with a new creation time.
diff --git a/sotodlib/toast/io/load_book.py b/sotodlib/toast/io/load_book.py
@@ -1,8 +1,6 @@
 # Copyright (c) 2020-2023 Simons Observatory.
 # Full license can be found in the top level "LICENSE" file.
-"""Tools for loading Level-3 Book format data.
-
-"""
+"""Tools for loading Level-3 Book format data."""
 
 import os
 import sys
@@ -251,9 +249,9 @@ def _unpack_frames(self, buffer):
                 "el_enc"
             ]
             if "flags" in ancil:
-                self.obs.shared[self.obs_fields["shared_flags"]].data[
-                    off : off + n
-                ] = np.array(ancil["flags"]).astype(np.uint8)
+                self.obs.shared[self.obs_fields["shared_flags"]].data[off : off + n] = (
+                    np.array(ancil["flags"]).astype(np.uint8)
+                )
 
             self.obs.shared[self.obs_fields["boresight_azel"]].data[
                 off : off + n, :
@@ -265,9 +263,9 @@ def _unpack_frames(self, buffer):
                 ] = t3g.from_g3_quats(frm["qboresight_radec"])
 
             if "hwp_enc" in ancil:
-                self.obs.shared[self.obs_fields["hwp_angle"]].data[
-                    off : off + n
-                ] = np.array(ancil["hwp_enc"]).astype(np.float64)
+                self.obs.shared[self.obs_fields["hwp_angle"]].data[off : off + n] = (
+                    np.array(ancil["hwp_enc"]).astype(np.float64)
+                )
 
             if "corotation_enc" in ancil:
                 self.obs.shared[self.obs_fields["corotator_angle"]].data[
@@ -290,12 +288,12 @@ def _unpack_frames(self, buffer):
             # print(f"DBG LOAD frame {off}:{off+n} signal = {frm['signal'].data[:5,:]}")
             # print(f"LOAD frame {off}:{off+n} signal = {frm['signal'].data[detmap,:]}")
             for idet in range(len(detmap)):
-                self.obs.detdata[self.obs_fields["det_data"]][
-                    idet, off : off + n
-                ] = frm["signal"].data[detmap[idet], :]
-                self.obs.detdata[self.obs_fields["det_flags"]][
-                    idet, off : off + n
-                ] = np.array(frm["flags"].data[detmap[idet], :], dtype=np.uint8)
+                self.obs.detdata[self.obs_fields["det_data"]][idet, off : off + n] = (
+                    frm["signal"].data[detmap[idet], :]
+                )
+                self.obs.detdata[self.obs_fields["det_flags"]][idet, off : off + n] = (
+                    np.array(frm["flags"].data[detmap[idet], :], dtype=np.uint8)
+                )
             if "signal_units" in frm and frm["signal_units"] is not None:
                 # We have some unit information
                 # print(f"Frame units = {frm['signal_units']}")
@@ -589,11 +587,12 @@ def read_book(
                 # print(f"Loading focalplane {fp_file}", flush=True)
                 if os.path.isfile(fp_file):
                     # We have a file, load it
-                    focalplane = toast.instrument.Focalplane(sample_rate=sample_rate)
                     with toast.io.H5File(
                         fp_file, "r", comm=comm.comm_group, force_serial=True
                     ) as f:
-                        focalplane.load_hdf5(f.handle, comm=comm.comm_group)
+                        focalplane = toast.instrument.Focalplane.load_hdf5(
+                            f.handle, comm=comm.comm_group
+                        )
                     update_creation_time(focalplane.detector_data, creation_time)
                     use_nominal = False
                 else:
diff --git a/sotodlib/toast/scripts/so_coadd.py b/sotodlib/toast/scripts/so_coadd.py
@@ -86,6 +86,8 @@ def get_weights(args, obs):
         cl_bb = spec["cl_BB"].to_value(u.uK**2)
         high_mean = np.mean(cl_bb[ell_high])
         low_mean = np.mean(cl_bb[ell_slc])
+        if np.isnan(low_mean) or np.isnan(high_mean) or low_mean == 0:
+            continue
         noise[oindex] = low_mean
         weights[oindex] = high_mean / low_mean
         print(
diff --git a/tests/test_toast_context.py b/tests/test_toast_context.py
