Mss3 tweaks

sotodlib/toast/ops/mlmapmaker.py

@@ -488,6 +488,21 @@ def _wrap_obs(self, ob, dets, passinfo):
 
         return axobs
 
+    @function_timer
+    def _add_obs(self, mapmaker, name, axobs, nmat, signal_estimate):
+        """ Add data to the mapmaker
+
+        Singled out for more granular timing
+        """
+        mapmaker.add_obs(
+            name,
+            axobs,
+            deslope=self.deslope,
+            noise_model=nmat,
+            signal_estimate=signal_estimate,
+        )
+        return
+
     @function_timer
     def _init_mapmaker(
             self, mapmaker, signal_map, mapmaker_prev, eval_prev, comm, gcomm, prefix,
@@ -749,13 +764,8 @@ def _exec(self, data, detectors=None, **kwargs):
                 else:
                     signal_estimate = None
 
-                mapmaker.add_obs(
-                    ob.name,
-                    axobs,
-                    deslope=self.deslope,
-                    noise_model=nmat,
-                    signal_estimate=signal_estimate,
-                )
+                self._add_obs(mapmaker, ob.name, axobs, nmat, signal_estimate)
+
                 del axobs
                 del signal_estimate
 

sotodlib/toast/ops/splits.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024 Simons Observatory.
+# Copyright (c) 2024-2025 Simons Observatory.
 # Full license can be found in the top level "LICENSE" file.
 
 import os
@@ -12,7 +12,7 @@
 from toast.timing import function_timer
 from toast.intervals import IntervalList
 from toast.observation import default_values as defaults
-from toast.ops import Operator
+from toast.ops import Operator, FlagIntervals
 from toast import qarray as qa
 
 
@@ -442,6 +442,53 @@ class Splits(Operator):
 
     splits = List([], help="The list of named splits to apply")
 
+    splits_as_flags = Bool(
+        False,
+        help="If True, splits are applied through flagging rather than views. "
+        "This can be more efficient if there are only a few samples in each "
+        "interval.",
+    )
+
+    skip_existing = Bool(
+        False,
+        help="If the mapmaker reports that all requested output files exist "
+        "for a split, skip the mapmaking step."
+    )
+
+    shared_flags = Unicode(
+        defaults.shared_flags,
+        allow_none=True,
+        help="Observation shared key for telescope flags to use when "
+        "splits_as_flags is True",
+    )
+
+    shared_flag_mask = Int(
+        128,
+        help="Bit mask value for flagging when splits_as_flags is True",
+    )
+
+    write_hits = Bool(True, help="If True, write the hits map")
+
+    write_cov = Bool(
+        False, help="If True, write the white noise covariance matrices."
+    )
+
+    write_invcov = Bool(
+        True,
+        help="If True, write the inverse white noise covariance matrices.",
+    )
+
+    write_rcond = Bool(
+        False, help="If True, write the reciprocal condition numbers."
+    )
+
+    write_map = Bool(True, help="If True, write the filtered/destriped map")
+
+    write_noiseweighted_map = Bool(
+        False,
+        help="If True, write the noise-weighted map (for fast co-add)",
+    )
+
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         return
@@ -523,6 +570,7 @@ def _exec(self, data, detectors=None, **kwargs):
         else:
             map_binner = self.mapmaker.binning
         pointing_view_save = map_binner.pixel_pointing.view
+        shared_flag_mask_save = map_binner.shared_flag_mask
 
         # If the pixel distribution does yet exit, we create it once prior to
         # doing any splits.
@@ -542,36 +590,70 @@ def _exec(self, data, detectors=None, **kwargs):
             if hasattr(self.mapmaker, trt):
                 setattr(self.mapmaker, trt, True)
 
+        # Optionally make copies of the shared flags
+        if self.splits_as_flags:
+            save_shared_flags = {}
+            for ob in data.obs:
+                if ob.comm_col_rank == 0:
+                    save_shared_flags[ob.name] = ob.shared[self.shared_flags].data
+                else:
+                    save_shared_flags[ob.name] = None
+        else:
+            save_shared_flags = None
+
         # Loop over splits
         for split_name, spl in self._split_obj.items():
-            log.info_rank(f"Running Split '{split_name}'", comm=data.comm.comm_world)
             # Set mapmaker name based on split and the name of this
             # Splits operator.
             mname = f"{self.name}_{split_name}"
             self.mapmaker.name = mname
+            log.info_rank(
+                f"Running Split '{split_name}'", comm=data.comm.comm_world
+            )
 
             # Apply this split
             for ob in data.obs:
                 spl.create_split(ob)
 
-            # Set mapmaking tools to use the current split interval list
-            map_binner.pixel_pointing.view = spl.split_intervals
-            if not map_binner.full_pointing:
-                # We are not using full pointing and so we clear the
-                # residual pointing for this split
-                toast.ops.Delete(
-                    detdata=[
-                        map_binner.pixel_pointing.pixels,
-                        map_binner.stokes_weights.weights,
-                        map_binner.pixel_pointing.detector_pointing.quats,
-                    ],
-                ).apply(data)
-
-            # Run mapmaking
-            self.mapmaker.apply(data)
-
-            # Write
-            self.write_splits(data, split_name=split_name)
+            if self.skip_existing and self.splits_exist(data, split_name=split_name):
+                log.info_rank(
+                    f"All outputs for split '{split_name}' exist, skipping...",
+                    comm=data.comm.comm_world,
+                )
+            else:
+                if self.splits_as_flags:
+                    # Split is applied through flagging, not as a view
+                    # Flag samples outside the intervals by prefixing '~'
+                    # to the view name
+                    FlagIntervals(
+                        shared_flags=self.shared_flags,
+                        shared_flag_bytes=1,
+                        view_mask=[
+                            (f"~{spl.split_intervals}", np.uint8(self.shared_flag_mask))
+                        ],
+                        reset=True,
+                    ).apply(data)
+                    map_binner.shared_flag_mask |= self.shared_flag_mask
+                else:
+                    # Set mapmaking tools to use the current split interval list
+                    map_binner.pixel_pointing.view = spl.split_intervals
+
+                if not map_binner.full_pointing:
+                    # We are not using full pointing and so we clear the
+                    # residual pointing for this split
+                    toast.ops.Delete(
+                        detdata=[
+                            map_binner.pixel_pointing.pixels,
+                            map_binner.stokes_weights.weights,
+                            map_binner.pixel_pointing.detector_pointing.quats,
+                        ],
+                    ).apply(data)
+
+                # Run mapmaking
+                self.mapmaker.apply(data)
+
+                # Write
+                self.write_splits(data, split_name=split_name)
 
             # Remove split
             for ob in data.obs:
@@ -581,13 +663,78 @@ def _exec(self, data, detectors=None, **kwargs):
         for k, v in mapmaker_save_traits.items():
             setattr(self.mapmaker, k, v)
         map_binner.pixel_pointing.view = pointing_view_save
+        map_binner.pixel_pointing.shared_flag_mask = shared_flag_mask_save
+
+        # Optionally restore shared flags
+        if save_shared_flags is not None:
+            for ob in data.obs:
+                ob.shared[self.shared_flags].set(
+                    save_shared_flags[ob.name], offset=(0,), fromrank=0
+                )
+
+        return
+
+    def splits_exist(self, data, split_name=None):
+        """Write out all split products."""
+        if not hasattr(self, "_split_obj"):
+            msg = "No splits have been created yet, cannot check existence"
+            raise RuntimeError(msg)
+
+        if hasattr(self.mapmaker, "map_binning"):
+            pixel_pointing = self.mapmaker.map_binning.pixel_pointing
+        else:
+            pixel_pointing = self.mapmaker.binning.pixel_pointing
+
+        if split_name is None:
+            to_write = dict(self._split_obj)
+        else:
+            to_write = {split_name: self._split_obj[split_name]}
+
+        if self.mapmaker.write_hdf5:
+            fname_suffix = "h5"
+        else:
+            fname_suffix = "fits"
+
+        all_exist = True
+        for spname, spl in to_write.items():
+            mname = f"{self.name}_{split_name}"
+            for prod, binner_key, write in [
+                    ("hits", None, self.write_hits),
+                    ("cov", None, self.write_cov),
+                    ("invcov", None, self.write_invcov),
+                    ("rcond", None, self.write_rcond),
+                    ("map", "binned", self.write_map),
+                    ("noiseweighted_map", "noiseweighted", self.write_noiseweighted_map),
+            ]:
+                if not write:
+                    continue
+                if binner_key is not None:
+                    # get the product name from BinMap
+                    mkey = getattr(self.mapmaker.binning, binner_key)
+                else:
+                    # hits and covariance are not made by BinMap.
+                    # Try synthesizing the product name
+                    mkey = f"{mname}_{prod}"
+                fname = os.path.join(
+                    self.output_dir,
+                    f"{self.name}_{spname}_{prod}.{fname_suffix}"
+                )
+                if not os.path.isfile(fname):
+                    all_exist = False
+                    break
+        return all_exist
 
     def write_splits(self, data, split_name=None):
         """Write out all split products."""
         if not hasattr(self, "_split_obj"):
             msg = "No splits have been created yet, cannot write"
             raise RuntimeError(msg)
 
+        if hasattr(self.mapmaker, "map_binning"):
+            pixel_pointing = self.mapmaker.map_binning.pixel_pointing
+        else:
+            pixel_pointing = self.mapmaker.binning.pixel_pointing
+
         if split_name is None:
             to_write = dict(self._split_obj)
         else:
@@ -600,10 +747,30 @@ def write_splits(self, data, split_name=None):
 
         for spname, spl in to_write.items():
             mname = f"{self.name}_{split_name}"
-            for prod in ["hits", "map", "invcov", "noiseweighted_map"]:
-                mkey = f"{mname}_{prod}"
+            for prod, binner_key, write in [
+                    ("hits", None, self.write_hits),
+                    ("cov", None, self.write_cov),
+                    ("invcov", None, self.write_invcov),
+                    ("rcond", None, self.write_rcond),
+                    ("map", "binned", self.write_map),
+                    ("noiseweighted_map", "noiseweighted", self.write_noiseweighted_map),
+            ]:
+                if not write:
+                    continue
+                if binner_key is not None:
+                    # get the product name from BinMap
+                    mkey = getattr(self.mapmaker.binning, binner_key)
+                else:
+                    # hits and covariance are not made by BinMap.
+                    # Try synthesizing the product name
+                    mkey = f"{mname}_{prod}"
+                if mkey not in data:
+                    msg = f"'{mkey}' not found in data. "
+                    msg += f"Available keys are {data.keys()}"
+                    raise RuntimeError(msg)
                 fname = os.path.join(
-                    self.output_dir, f"{self.name}_{spname}_{prod}.{fname_suffix}"
+                    self.output_dir,
+                    f"{self.name}_{spname}_{prod}.{fname_suffix}"
                 )
                 data[mkey].write(
                     fname,

sotodlib/toast/workflows/pointing.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023-2024 Simons Observatory.
+# Copyright (c) 2023-2025 Simons Observatory.
 # Full license can be found in the top level "LICENSE" file.
 
 import numpy as np
@@ -167,11 +167,11 @@ def select_pointing(job, otherargs, runargs, data):
     if hasattr(job_ops, "binner"):
         job_ops.binner.pixel_pointing = job.pixels_solve
         job_ops.binner.stokes_weights = job.weights_solve
-        job_ops.binner.full_pointing = otherargs.full_pointing
+        job_ops.binner.full_pointing |= otherargs.full_pointing
     if hasattr(job_ops, "binner_final"):
         job_ops.binner_final.pixel_pointing = job.pixels_final
         job_ops.binner_final.stokes_weights = job.weights_final
-        job_ops.binner_final.full_pointing = otherargs.full_pointing
+        job_ops.binner_final.full_pointing |= otherargs.full_pointing
         # If we are not using a different binner for our final binning, use the
         # same one as the solve.
         if not job_ops.binner_final.enabled:

sotodlib/toast/workflows/proc_characterize.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023-2023 Simons Observatory.
+# Copyright (c) 2023-2025 Simons Observatory.
 # Full license can be found in the top level "LICENSE" file.
 """Timestream processing filters.
 """
@@ -116,7 +116,7 @@ def hn_map(job, otherargs, runargs, data):
         job_ops.h_n.pixel_pointing = job.pixels_final
         job_ops.h_n.pixel_dist = job_ops.binner_final.pixel_dist
         job_ops.h_n.output_dir = otherargs.out_dir
-        job_ops.h_n.save_pointing = otherargs.full_pointing
+        job_ops.h_n.save_pointing |= otherargs.full_pointing
         job_ops.h_n.apply(data)
 
 
@@ -154,7 +154,7 @@ def cadence_map(job, otherargs, runargs, data):
         job_ops.cadence_map.pixel_pointing = job.pixels_final
         job_ops.cadence_map.pixel_dist = job_ops.binner_final.pixel_dist
         job_ops.cadence_map.output_dir = otherargs.out_dir
-        job_ops.cadence_map.save_pointing = otherargs.full_pointing
+        job_ops.cadence_map.save_pointing |= otherargs.full_pointing
         job_ops.cadence_map.apply(data)
 
 
@@ -192,5 +192,5 @@ def crosslinking_map(job, otherargs, runargs, data):
         job_ops.crosslinking.pixel_pointing = job.pixels_final
         job_ops.crosslinking.pixel_dist = job_ops.binner_final.pixel_dist
         job_ops.crosslinking.output_dir = otherargs.out_dir
-        job_ops.crosslinking.save_pointing = otherargs.full_pointing
+        job_ops.crosslinking.save_pointing |= otherargs.full_pointing
         job_ops.crosslinking.apply(data)

sotodlib/toast/workflows/proc_demodulation.py

@@ -91,7 +91,9 @@ def demodulate(job, otherargs, runargs, data):
         comm=data.comm.comm_world,
         timer=timer,
     )
-    demod_weights = toast.ops.StokesWeightsDemod()
+    demod_weights = toast.ops.StokesWeightsDemod(
+        mode=job_ops.demodulate.mode
+    )
     job_ops.weights_radec = demod_weights
     if hasattr(job_ops, "binner"):
         job_ops.binner.stokes_weights = demod_weights

sotodlib/toast/workflows/proc_flagging.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023-2024 Simons Observatory.
+# Copyright (c) 2023-2025 Simons Observatory.
 # Full license can be found in the top level "LICENSE" file.
 """Timestream flagging operations.
 """
@@ -343,5 +343,5 @@ def processing_mask(job, otherargs, runargs, data):
             # We are using the same pointing matrix as the mapmaking
             job_ops.processing_mask.pixel_dist = job_ops.binner.pixel_dist
             job_ops.processing_mask.pixel_pointing = job.pixels_solve
-        job_ops.processing_mask.save_pointing = otherargs.full_pointing
+        job_ops.processing_mask.save_pointing |= otherargs.full_pointing
         job_ops.processing_mask.apply(data)