Better handle flagged samples in trend cuts and dark templates

- In trend cuts, disregard the slope of chunks without a sufficient
  number of good samples.

- In dark templates, handle cases where many samples are flagged,
  and also cases where there are no remaining dark detectors after
  cuts.

- Other small cleanups

Author: tskisner

sotodlib/scripts/so_batch_control.py

@@ -52,8 +52,7 @@ def get_state_file(out_root, obs, name="state"):
 
 
 def get_obs_state(out_root, obs):
-    """Find the current processing state of an observation.
-    """
+    """Find the current processing state of an observation."""
     state_file = get_state_file(out_root, obs)
     if os.path.isfile(state_file):
         with open(state_file, "r") as f:
@@ -63,8 +62,7 @@ def get_obs_state(out_root, obs):
 
 
 def clear_obs_state(out_root, obs):
-    """Clear the current processing state of an observation.
-    """
+    """Clear the current processing state of an observation."""
     obs_dir = os.path.join(out_root, obs)
     if not os.path.isdir(obs_dir):
         return
@@ -74,8 +72,7 @@ def clear_obs_state(out_root, obs):
 
 
 def set_obs_state(out_root, obs, state):
-    """Set the current processing state of an observation.
-    """
+    """Set the current processing state of an observation."""
     obs_dir = os.path.join(out_root, obs)
     if not os.path.exists(obs_dir):
         os.makedirs(obs_dir)
@@ -175,6 +172,13 @@ def main():
         default=None,
         help="Set the state of this observation to 'done'",
     )
+    parser.add_argument(
+        "--set_state_failed",
+        required=False,
+        type=str,
+        default=None,
+        help="Set the state of this observation to 'failed'",
+    )
     parser.add_argument(
         "--set_state_running",
         required=False,
@@ -196,13 +200,22 @@ def main():
         default=False,
         help="Remove the running state from all observations",
     )
+    parser.add_argument(
+        "--ignore_running",
+        required=False,
+        action="store_true",
+        default=False,
+        help="Ignore stale running state when considering observations",
+    )
 
     args = parser.parse_args()
 
     if args.get_batch is not None:
         # We are getting the next batch of observations
         all_obs = load_obs_file(args.observations)
-        batch_obs = find_obs(all_obs, args.get_batch, args.out_root)
+        batch_obs = find_obs(
+            all_obs, args.get_batch, args.out_root, ignore_running=args.ignore_running
+        )
         if args.batch_list:
             batch_str = ",".join(batch_obs)
             print(f"{batch_str}", flush=True)
@@ -216,6 +229,8 @@ def main():
         clear_obs_state(args.out_root, args.clear_state)
     elif args.set_state_done is not None:
         set_obs_state(args.out_root, args.set_state_done, "done")
+    elif args.set_state_failed is not None:
+        set_obs_state(args.out_root, args.set_state_done, "failed")
     elif args.set_state_running is not None:
         set_obs_state(args.out_root, args.set_state_running, "running")
     elif args.cleanup:

sotodlib/toast/ops/dark_template.py

@@ -17,7 +17,7 @@
 import toast
 from toast.traits import trait_docs, Unicode, Int, Instance, Bool, Float, List
 from toast.ops import Operator
-from toast.utils import Logger, Environment, rate_from_times
+from toast.utils import Logger, flagged_noise_fill
 from toast.timing import function_timer, Timer
 from toast.observation import default_values as defaults
 
@@ -101,24 +101,48 @@ def __init__(self, **kwargs):
 
     @function_timer
     def _gather_dark_tod(self, ob):
-        """ Gather dark tod to root process
+        """Gather dark tod to root process
 
         """
         comm = ob.comm.comm_group
+        rank = ob.comm.group_rank
         dark_tod = []
         fp = ob.telescope.focalplane
         for det in ob.select_local_detectors(flagmask=self.det_mask):
             if fp[det]["det_info:wafer:type"] != "DARK":
                 continue
-            tod = ob.detdata[self.det_data][det]
+            tod = ob.detdata[self.det_data][det].copy()
+            if self.shared_flags is not None:
+                flags = ob.shared[self.shared_flags].data & self.shared_flag_mask
+            else:
+                flags = np.zeros(len(tod), dtype=np.uint8)
+            if self.det_flags is not None:
+                flags |= ob.detdata[self.det_flags] & self.det_flag_mask
             # Check for typical pathologies
             if np.any(np.isnan(tod)) or np.std(tod) == 0:
                 continue
-            # Subtract the mean
-            dark_tod.append(tod - np.mean(tod))
+            if np.count_nonzero(flags) < 0.5 * len(tod):
+                continue
+            # Subtract the mean and fill flagged samples with noise.  For the gap
+            # filling, base the buffer size on the lowpass kernel size.
+            good = flags == 0
+            avg = np.mean(tod[good])
+            tod -= avg
+            flagged_noise_fill(tod, flags, self.naverage // 2, poly_order=5)
+            dark_tod.append(tod)
+        fail = False
         if comm is not None:
-            dark_tod = comm.gather(dark_tod)
-        return dark_tod
+            dark_tod = comm.gather(dark_tod, root=0)
+        if rank == 0:
+            total_dark = np.sum([len(x) for x in dark_tod])
+            if total_dark == 0:
+                fail = True
+        if comm is not None:
+            fail = comm.bcast(fail, root=0)
+        if fail:
+            return None
+        else:
+            return dark_tod
 
     @function_timer
     def _derive_templates(self, dark_tod):
@@ -163,6 +187,7 @@ def _load_dark_templates(self, ob):
         """
         log = Logger.get()
         comm = ob.comm.comm_group
+        rank = ob.comm.group_rank
 
         # See if the templates were already cached
 
@@ -171,17 +196,25 @@ def _load_dark_templates(self, ob):
                 self.cache_dir,
                 f"dark_templates.{ob.name}.pck",
             )
-            if os.path.isfile(fname_cache):
-                if comm is None or comm.rank == 0:
+            exists = False
+            if rank == 0:
+                if os.path.isfile(fname_cache):
+                    exists = True
+            if comm is not None:
+                exists = comm.bcast(exists, root=0)
+            if exists:
+                if rank == 0:
                     with open(fname_cache, "rb") as f:
                         templates = pickle.load(f)
                     log.debug(f"Loaded dark templates from {fname_cache}")
                 else:
                     templates = None
                 self._share_templates(ob, templates)
                 return ob.shared[self.key]
-
-        return None
+            else:
+                return None
+        else:
+            return None
 
     @function_timer
     def _share_templates(self, ob, templates):
@@ -196,8 +229,11 @@ def _share_templates(self, ob, templates):
             ntemplate, nsample = templates.shape
 
         if comm is not None:
-            ntemplate = comm.bcast(ntemplate)
-            nsample = comm.bcast(nsample)
+            ntemplate = comm.bcast(ntemplate, root=0)
+            nsample = comm.bcast(nsample, root=0)
+
+        if ntemplate is None:
+            return
 
         ob.shared.create_column(
             self.key,
@@ -209,7 +245,6 @@ def _share_templates(self, ob, templates):
         else:
             temp_trans = None
         ob.shared[self.key].set(temp_trans, offset=(0, 0), fromrank=0)
-        return
 
     @function_timer
     def _save_dark_templates(self, ob, templates):
@@ -220,48 +255,68 @@ def _save_dark_templates(self, ob, templates):
 
         log = Logger.get()
         comm = ob.comm.comm_group
+        rank = ob.comm.group_rank
 
         if self.cache_dir is None:
             return
 
-        if comm is None or comm.rank == 0:
-            fname_cache = os.path.join(
-                self.cache_dir,
-                f"dark_templates.{ob.name}.pck",
-            )
-            with open(fname_cache, "wb") as f:
-                pickle.dump(templates.data.T, f)
-            log.debug(f"Wrote dark templates to {fname_cache}")
-
-        return
+        if rank == 0:
+            if templates is not None:
+                fname_cache = os.path.join(
+                    self.cache_dir,
+                    f"dark_templates.{ob.name}.pck",
+                )
+                with open(fname_cache, "wb") as f:
+                    pickle.dump(templates.data.T, f)
+                log.debug(f"Wrote dark templates to {fname_cache}")
+        if comm is not None:
+            comm.barrier()
 
     @function_timer
     def _get_dark_templates(self, ob):
-        """ Construct low-passed dark bolometer templates
+        """ Construct low-passed dark bolometer templates.
+
+        If no good dark detector data is found, the return value is None.
         """
-        log = Logger.get()
         comm = ob.comm.comm_group
+        rank = ob.comm.group_rank
 
         # gather all dark TOD to the root process
 
         dark_tod = self._gather_dark_tod(ob)
 
         # Construct dark templates through PCA and lowpass
 
-        if comm is None or comm.rank == 0:
-            templates = self._derive_templates(dark_tod)
+        if rank == 0:
+            if dark_tod is None:
+                fail = True
+            else:
+                fail = False
+                templates = self._derive_templates(dark_tod)
         else:
+            fail = False
             templates = None
+        if comm is not None:
+            fail = comm.bcast(fail, root=0)
 
-        self._share_templates(ob, templates)
-
-        return ob.shared[self.key]
+        if fail:
+            return None
+        else:
+            self._share_templates(ob, templates)
+            return ob.shared[self.key]
 
     @function_timer
     def _project_dark_templates(self, ob, templates):
         """ Project the provided templates out of every optical detector
 
         """
+        log = Logger.get()
+        if templates is None:
+            # No templates available
+            msg = f"No dark templates available for {ob.name}"
+            log.info_rank(msg, comm=ob.comm.comm_group)
+            return
+        
         if self.shared_flags is not None:
             common_flags = ob.shared[self.shared_flags].data & self.shared_flag_mask
         else:
@@ -303,8 +358,6 @@ def _project_dark_templates(self, ob, templates):
                 coeff = np.dot(cov, proj)
                 tod[ind] -= np.dot(coeff, templates_ind.T)
 
-        return
-
     @function_timer
     def _exec(self, data, detectors=None, **kwargs):
         log = Logger.get()
@@ -330,8 +383,6 @@ def _exec(self, data, detectors=None, **kwargs):
                 self._save_dark_templates(ob, templates)
             self._project_dark_templates(ob, templates)
 
-        return
-
     @function_timer
     def _finalize(self, data, **kwargs):
         pass

sotodlib/toast/ops/load_context_utils.py

@@ -637,14 +637,14 @@ def _process_flag(mask, inbuf, detdata, det_begin, det_end, invert):
     # temporary buffers of ranges flags.
     del send_data
 
-    # Every process checks its local data for NaN values.  If any are found, a warning
+    # Every process checks its local data for NaN values.  If any are found, a message
     # is printed and the detector is cut.
     dflags = dict()
     for det in obs.local_detectors:
         nnan = np.count_nonzero(np.isnan(obs.detdata[field][det]))
         if nnan > 0:
             msg = f"{obs.name}:{det} has {nnan} NaN values.  Cutting."
-            log.warning(msg)
+            log.debug(msg)
             dflags[det] = defaults.det_mask_invalid
     obs.update_local_detector_flags(dflags)
 

sotodlib/toast/ops/trend_cut.py

@@ -164,15 +164,21 @@ def _compute_trend(self, obs, reltime, slices, det):
         slopes = list()
         for slc in slices:
             good = flags[slc] == 0
-            t = reltime[slc][good]
-            s = signal[slc][good]
-
-            t_mean = np.mean(t)
-            t_var = np.var(t, mean=t_mean)
-            s_mean = np.mean(s)
-            ts_mean = np.mean(t * s)
-
-            slopes.append((ts_mean - t_mean * s_mean) / t_var)
+            if np.count_nonzero(good) < 0.1 * len(flags[slc]):
+                # Not enough good data
+                slopes.append(0.0)
+            else:
+                t = reltime[slc][good]
+                s = signal[slc][good]
+
+                t_mean = np.mean(t)
+                t_var = np.var(t, mean=t_mean)
+                s_mean = np.mean(s)
+                ts_mean = np.mean(t * s)
+                if t_var == 0:
+                    slopes.append(0.0)
+                else:
+                    slopes.append((ts_mean - t_mean * s_mean) / t_var)
         return np.array(slopes)
 
     def _finalize(self, data, **kwargs):