Rework detector selection for preprocess loading (#1116)

* rework det selection for preprocess loading

* add select=False when running pipelines in load functions

* add in_place for select

* remove unneeded import

* fix missing in_places

* base in_place should be true

* fix missing in_place

* update invvar flag select

* fix type

---------

Co-authored-by: Michael McCrackan <[email protected]>
Co-authored-by: Michael McCrackan <[email protected]>
Co-authored-by: Michael McCrackan <[email protected]>

Author: 4 people

sotodlib/preprocess/pcore.py

@@ -98,8 +98,8 @@ def save(self, proc_aman, *args):
         if self.save_cfgs is None:
             return
         raise NotImplementedError
-        
-    def select(self, meta, proc_aman=None):
+
+    def select(self, meta, proc_aman=None, in_place=True):
         """ This function runs any desired data selection of the preprocessing
         pipeline results. Assumes the pipeline has already been run and that the
         resulting proc_aman is now saved under the ``preprocess`` key in the
@@ -111,12 +111,18 @@ def select(self, meta, proc_aman=None):
 
         Arguments
         ---------
-        meta : AxisManager 
+        meta : AxisManager
             Metadata related to the specific observation
+        proc_aman : AxisManager
+            Optional. Any information generated by previous elements in the
+            preprocessing pipeline.
+        in_place : bool
+            Optional.  Apply selection and return restricted axis manager if
+            True, else return the flag array.
 
         Returns
         -------
-        meta : AxisManager 
+        meta : AxisManager
             Metadata where non-selected detectors have been removed
         """
 

sotodlib/preprocess/preprocess_util.py

@@ -319,8 +319,13 @@ def load_preprocess_det_select(obs_id, configs, context=None,
     pipe = Pipeline(configs["process_pipe"], logger=logger)
 
     meta = context.get_meta(obs_id, dets=dets, meta=meta)
-    logger.info(f"Cutting on the last process: {pipe[-1].name}")
-    pipe[-1].select(meta)
+    logger.info("Restricting detectors on all processes")
+    keep_all = np.ones(meta.dets.count,dtype=bool)
+    for process in pipe[:]:
+        keep = process.select(meta, in_place=False)
+        if isinstance(keep, np.ndarray):
+            keep_all &= keep
+    meta.restrict("dets", meta.dets.vals[keep_all])
     return meta
 
 
@@ -367,7 +372,7 @@ def load_and_preprocess(obs_id, configs, context=None, dets=None, meta=None,
     else:
         pipe = Pipeline(configs["process_pipe"], logger=logger)
         aman = context.get_obs(meta, no_signal=no_signal)
-        pipe.run(aman, aman.preprocess)
+        pipe.run(aman, aman.preprocess, select=False)
         return aman
 
 
@@ -436,21 +441,29 @@ def multilayer_load_and_preprocess(obs_id, configs_init, configs_proc,
 
         if check_cfg_match(aman_cfgs_ref, meta_proc.preprocess['pcfg_ref'],
                            logger=logger):
+            pipe_proc = Pipeline(configs_proc["process_pipe"], logger=logger)
 
+            logger.info("Restricting detectors on all proc pipeline processes")
+            keep_all = np.ones(meta_proc.dets.count, dtype=bool)
+            for process in pipe_proc[:]:
+                keep = process.select(meta_proc, in_place=False)
+                if isinstance(keep, np.ndarray):
+                    keep_all &= keep
+            meta_proc.restrict("dets", meta_proc.dets.vals[keep_all])
             meta_init.restrict('dets', meta_proc.dets.vals)
+
             aman = context_init.get_obs(meta_init, no_signal=no_signal)
             logger.info("Running initial pipeline")
-            pipe_init.run(aman, aman.preprocess)
+            pipe_init.run(aman, aman.preprocess, select=False)
             if init_only:
                 return aman
 
-            pipe_proc = Pipeline(configs_proc["process_pipe"], logger=logger)
             logger.info("Running dependent pipeline")
             proc_aman = context_proc.get_meta(obs_id, meta=aman)
 
             aman.preprocess.merge(proc_aman.preprocess)
 
-            pipe_proc.run(aman, aman.preprocess)
+            pipe_proc.run(aman, aman.preprocess, select=False)
 
             return aman
         else:

sotodlib/preprocess/processes.py

@@ -68,14 +68,17 @@ def save(self, proc_aman, dbc_aman):
         if self.save_cfgs:
             proc_aman.wrap("det_bias_flags", dbc_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
-        keep = ~proc_aman.det_bias_flags.det_bias_flags
-        meta.restrict("dets", meta.dets.vals[has_all_cut(keep)])
-        return meta
+        keep = ~has_all_cut(proc_aman.det_bias_flags.det_bias_flags)
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
     def plot(self, aman, proc_aman, filename):
         if self.plot_cfgs is None:
@@ -132,7 +135,7 @@ def save(self, proc_aman, trend_aman):
         if self.save_cfgs:
             proc_aman.wrap("trends", trend_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
@@ -144,8 +147,11 @@ def select(self, meta, proc_aman=None):
         else:
             raise ValueError(f"Entry '{self.select_cfgs['kind']}' not"
                                 "understood. Expect 'any' or 'all'")
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
     def plot(self, aman, proc_aman, filename):
         if self.plot_cfgs is None:
@@ -210,7 +216,7 @@ def save(self, proc_aman, glitch_aman):
         if self.save_cfgs:
             proc_aman.wrap(self.glitch_name, glitch_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
@@ -220,8 +226,11 @@ def select(self, meta, proc_aman=None):
         )
         n_cut = count_cuts(flag)
         keep = n_cut <= self.select_cfgs["max_n_glitch"]
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
     def plot(self, aman, proc_aman, filename):
         if self.plot_cfgs is None:
@@ -325,7 +334,7 @@ def save(self, proc_aman, jump_aman):
             name = self.save_cfgs.get('jumps_name', 'jumps')
             proc_aman.wrap(name, jump_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
@@ -334,8 +343,10 @@ def select(self, meta, proc_aman=None):
 
         n_cut = count_cuts(proc_aman[name].jump_flag)
         keep = n_cut <= self.select_cfgs["max_n_jumps"]
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        return keep
 
     def plot(self, aman, proc_aman, filename):
         if self.plot_cfgs is None:
@@ -619,7 +630,7 @@ def save(self, proc_aman, noise):
         else:
             proc_aman.wrap(self.save_cfgs['wrap_name'], noise)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
 
@@ -648,8 +659,11 @@ def select(self, meta, proc_aman=None):
             keep &= (wn >= np.float64(self.select_cfgs["min_noise"]))
         if fk is not None and "max_fknee" in self.select_cfgs.keys():
             keep &= (fk <= np.float64(self.select_cfgs["max_fknee"]))
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
 class Calibrate(_Preprocess):
     """Calibrate the timestreams based on some provided information.
@@ -1244,14 +1258,17 @@ def save(self, proc_aman, dark_aman):
         if self.save_cfgs:
             proc_aman.wrap("darks", dark_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
         keep = ~has_all_cut(proc_aman.darks.darks)
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
 class SourceFlags(_Preprocess):
     """Calculate the source flags in the data.
@@ -1315,7 +1332,7 @@ def save(self, proc_aman, source_aman):
         if self.save_cfgs:
             proc_aman.wrap("source_flags", source_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
@@ -1330,12 +1347,20 @@ def select(self, meta, proc_aman=None):
             if len(source_list) == 0:
                 raise ValueError("No tags match source list")
 
+        keep_all = np.ones(meta.dets.count, dtype=bool)
+
         for source in source_list:
             if source in source_flags._fields:
                 keep = ~has_all_cut(source_flags[source])
-                meta.restrict("dets", meta.dets.vals[keep])
-                source_flags.restrict("dets", source_flags.dets.vals[keep])
-        return meta
+                if in_place:
+                    meta.restrict("dets", meta.dets.vals[keep])
+                    source_flags.restrict("dets", source_flags.dets.vals[keep])
+                else:
+                    keep_all &= keep
+        if in_place:
+            return meta
+        else:
+            return keep_all
 
 class HWPAngleModel(_Preprocess):
     """Apply hwp angle model to the TOD.
@@ -1545,14 +1570,17 @@ def save(self, proc_aman, pca_aman):
         if self.save_cfgs:
             proc_aman.wrap(self.run_name, pca_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
         keep = ~proc_aman[self.run_name]['pca_det_mask']
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
     def plot(self, aman, proc_aman, filename):
         if self.plot_cfgs is None:
@@ -1701,14 +1729,17 @@ def save(self, proc_aman, ptp_aman):
         if self.save_cfgs:
             proc_aman.wrap("ptp_flags", ptp_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
         keep = ~has_all_cut(proc_aman.ptp_flags.ptp_flags)
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        else:
+            return keep
 
 class InvVarFlags(_Preprocess):
     """Find detectors with too high inverse variance.
@@ -1741,14 +1772,16 @@ def save(self, proc_aman, inv_var_aman):
         if self.save_cfgs:
             proc_aman.wrap("inv_var_flags", inv_var_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
         keep = ~has_all_cut(proc_aman.inv_var_flags.inv_var_flags)
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        return keep
 
 class EstimateT2P(_Preprocess):
     """Estimate T to P leakage coefficients.
@@ -2011,15 +2044,17 @@ def save(self, proc_aman, fp_aman):
             return
         if self.save_cfgs:
             proc_aman.wrap("fp_flags", fp_aman)
-
-    def select(self, meta, proc_aman=None):
+    
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
         keep = ~has_all_cut(proc_aman.fp_flags.fp_nans)
-        meta.restrict("dets", meta.dets.vals[keep])
-        return meta
+        if in_place:
+            meta.restrict("dets", meta.dets.vals[keep])
+            return meta
+        return keep
 
 class PointingModel(_Preprocess):
     """Apply pointing model to the TOD.
@@ -2086,15 +2121,16 @@ def save(self, proc_aman, calc_aman, name):
         if self.save_cfgs:
             proc_aman.wrap(name, calc_aman)
 
-    def select(self, meta, proc_aman=None):
+    def select(self, meta, proc_aman=None, in_place=True):
         if self.select_cfgs is None:
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
-        if hasattr(proc_aman.noisy_dets_flags, "valid_dets"):
-            keep = proc_aman.noisy_dets_flags.valid_dets
+        keep = proc_aman.noisy_dets_flags.valid_dets
+        if in_place:
             meta.restrict('dets', proc_aman.dets.vals[keep])
-        return meta
+            return meta
+        return keep
 
 class CorrectIIRParams(_Preprocess):
     """Correct missing iir_params by default values.