Putting the mapmaker in the same split convention as the preprocessing (

#1060) * Changing the convention of focal plane splits, also fixing a bug in preprocess * Fixing docstring in obs_ops.splits.get_split_flags * Update preproc split method to match demod mapmaker. * Remove import of get_flags * Looks for glitch_flags in obs.flags * Fixing the cuts flags so that dets/samps we want to keep are False * Fix sample restriction issue * make preproc splits optional. * Missed noise split. --------- Co-authored-by: Max Silva-Feaver <[email protected]>
simonsobs · Dec 13, 2024 · 60fb595 · 60fb595
1 parent e334326
commit 60fb595
Show file tree

Hide file tree

Showing 4 changed files with 89 additions and 70 deletions.
diff --git a/sotodlib/mapmaking/demod_mapmaker.py b/sotodlib/mapmaking/demod_mapmaker.py
@@ -13,7 +13,7 @@
 from .. import core
 from .. import coords
 from .utilities import recentering_to_quat_lonlat, evaluate_recentering, MultiZipper, unarr, safe_invert_div
-from .utilities import import_optional, get_flags
+from .utilities import import_optional
 from .noise_model import NmatWhite
 
 hp = import_optional('healpy')
@@ -291,10 +291,10 @@ def add_obs(self, id, obs, nmat, Nd, pmap=None, split_labels=None):
                 else: rot = None
                 if self.Nsplits == 1:
                     # this is the case with no splits
-                    flagnames = ['glitch_flags']
+                    cuts = obs.flags.glitch_flags
                 else:
-                    flagnames = ['glitch_flags', split_labels[n_split]]
-                cuts = get_flags(obs, flagnames)
+                    # remember that the dets or samples you want to keep should be false, hence we negate
+                    cuts = obs.flags.glitch_flags + ~obs.preprocess.split_flags.cuts[split_labels[n_split]]
                 if self.pix_scheme == "rectpix":
                     threads='domdir'
                     geom = self.rhs.geometry

diff --git a/sotodlib/mapmaking/utilities.py b/sotodlib/mapmaking/utilities.py
@@ -527,29 +527,6 @@ def downsample_obs(obs, down):
     # doesn't matter anyway
     return res
 
-def get_flags(obs, flagnames):
-    """Parse detector-set splits"""
-    cuts_out = None
-    if flagnames is None:
-        return so3g.proj.RangesMatrix.zeros(obs.shape)
-    det_splits = ['det_left','det_right','det_in','det_out','det_upper','det_lower']
-    for flagname in flagnames:
-        if flagname in det_splits:
-            cuts = obs.det_flags[flagname]
-        elif flagname == 'scan_left':
-            cuts = obs.flags.left_scan
-        elif flagname == 'scan_right':
-            cuts = obs.flags.right_scan
-        else:
-            cuts = getattr(obs.flags, flagname) # obs.flags.flagname
-
-        ## Add to the output matrix
-        if cuts_out is None:
-            cuts_out = cuts
-        else:
-            cuts_out += cuts
-    return cuts_out
-
 def import_optional(module_name):
     try:
         module = importlib.import_module(module_name)

diff --git a/sotodlib/obs_ops/splits.py b/sotodlib/obs_ops/splits.py
@@ -64,8 +64,10 @@ def det_splits_relative(aman, det_left_right=False, det_upper_lower=False, det_i
 
 def get_split_flags(aman, proc_aman=None, split_cfg=None):
     '''
-    Function returning flags used for null splits consumed by the mapmaking and bundling codes.                             Fields labeled ``field_name_flag`` contain boolean masks and ``_avg`` are the mean
-    of the numerical based split flags to be used for observation level splits.
+    Function returning flags used for null splits consumed by the mapmaking
+    and bundling codes. Fields labeled ``field_name_flag`` contain boolean
+    masks and ``_avg`` are the mean of the numerical based split flags to 
+    be used for observation level splits.
 
     Arguments
     ---------
@@ -80,55 +82,96 @@ def get_split_flags(aman, proc_aman=None, split_cfg=None):
     -------
     split_aman: AxisManager
         Axis manager containing splitting flags.
+        ``cuts`` field is a FlagManager containing the detector and subscan based splits used in the mapmaker.
+        ``<split_name>_threshold`` fields contain the threshold used for the split.
+        Other fields conatain info for obs-level splits.
     '''
-    if proc_aman is None:
-        try:
-            proc_aman = aman.preprocess
-        except:
-            raise ValueError('proc_aman is None and no preprocess field in aman provide valid preprocess metadata')
-
-    if (not 't2p' in proc_aman) | (not 'hwpss_stats' in proc_aman):
-        raise ValueError('t2p or hwpss_stats not in proc_aman must run after those steps in the pipeline.')
-
     # Set default set of splits
-    default_cfg = {'high_gain': 0.115, 'high_noise': 3.5e-5, 'high_tau': 1.5e-3,
-                   'det_A': 'A', 'pol_angle': 35, 'det_top': 'B', 'high_leakage': 1e-3,
-                   'high_2f': 1.5e-3, 'right_focal_plane': 0, 'top_focal_plane': 0,
-                   'central_pixels': 0.071 }
+    default_cfg = {'high_gain': 0.115, 'high_tau': 1.5e-3,
+                   'det_A': 'A', 'pol_angle': 35, 'det_top': 'B', 'right_focal_plane': 0,
+                   'top_focal_plane': 0, 'central_pixels': 0.071 }
     if split_cfg is None:
         split_cfg = default_cfg
 
     split_aman = AxisManager(aman.dets)
+    fm = FlagManager.for_tod(aman)
     # If provided split config doesn't include all of the splits in default
     for k in default_cfg.keys():
         if not k in split_cfg:
             split_cfg[k] = default_cfg[k]
-        split_aman.wrap(f'{k}_threshold', split_cfg[k])
 
-    split_aman.wrap('high_gain_flag', aman.det_cal.phase_to_pW > split_cfg['high_gain'],
-                    [(0, 'dets')])
+    # Gain split
+    fm.wrap_dets('high_gain', aman.det_cal.phase_to_pW > split_cfg['high_gain'])
+    fm.wrap_dets('low_gain', aman.det_cal.phase_to_pW <= split_cfg['high_gain'])
     split_aman.wrap('gain_avg', np.nanmean(aman.det_cal.phase_to_pW))
-    split_aman.wrap('high_noise_flag', proc_aman.noiseQ_fit.fit[:,1] > split_cfg['high_noise'],
-                    [(0, 'dets')])
-    split_aman.wrap('noise_avg', np.nanmean(proc_aman.noiseQ_fit.fit[:,1]))
-    split_aman.wrap('high_tau_flag', aman.det_cal.tau_eff > split_cfg['high_tau'],
-                    [(0, 'dets')])
+    # Time constant split
+    fm.wrap_dets('high_tau', aman.det_cal.tau_eff > split_cfg['high_tau'])
+    fm.wrap_dets('low_tau', aman.det_cal.tau_eff <= split_cfg['high_tau'])
     split_aman.wrap('tau_avg', np.nanmean(aman.det_cal.tau_eff))
-    split_aman.wrap('det_A_flag', aman.det_info.wafer.pol <= split_cfg['det_A'],
-                    [(0, 'dets')])
-    split_aman.wrap('pol_angle_flag', aman.det_info.wafer.angle > split_cfg['pol_angle'],
-                    [(0, 'dets')])
-    split_aman.wrap('det_top_flag', aman.det_info.wafer.crossover > split_cfg['det_top'],
-                    [(0, 'dets')])
-    split_aman.wrap('high_leakage_flag', np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2) > split_cfg['high_leakage'],
-                    [(0, 'dets')])
-    split_aman.wrap('leakage_avg', np.nanmean(np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2)),
-                    [(0, 'dets')])
-    a2 = aman.det_cal.phase_to_pW*np.sqrt(proc_aman.hwpss_stats.coeffs[:,2]**2 + proc_aman.hwpss_stats.coeffs[:,3]**2)
-    split_aman.wrap('high_2f_flag', a2 > split_cfg['high_2f'], [(0, 'dets')])
-    split_aman.wrap('2f_avg', np.nanmean(a2), [(0, 'dets')])
-    split_aman.wrap('right_focal_plane_flag', aman.focal_plane.xi > split_cfg['right_focal_plane'], [(0, 'dets')])
-    split_aman.wrap('top_focal_plane_flag', aman.focal_plane.eta > split_cfg['top_focal_plane'], [(0, 'dets')])
-    split_aman.wrap('central_pixels_flag', np.sqrt(aman.focal_plane.xi**2 + aman.focal_plane.eta**2) < split_cfg['central_pixels'],
+    # detAB split
+    fm.wrap_dets('det_A', aman.det_info.wafer.pol <= split_cfg['det_A'])
+    fm.wrap_dets('det_B', aman.det_info.wafer.pol > split_cfg['det_A'])
+    # def pol split
+    fm.wrap_dets('high_pol_angle', aman.det_info.wafer.angle > split_cfg['pol_angle'])
+    fm.wrap_dets('low_pol_angle', aman.det_info.wafer.angle <= split_cfg['pol_angle'])
+    # det top/bottom split
+    fm.wrap_dets('det_type_top', aman.det_info.wafer.crossover > split_cfg['det_top'])
+    fm.wrap_dets('det_type_bottom', aman.det_info.wafer.crossover <= split_cfg['det_top'])
+    # Right/left focal plane split
+    fm.wrap_dets('det_right', aman.focal_plane.xi > split_cfg['right_focal_plane'])
+    fm.wrap_dets('det_left', aman.focal_plane.xi <= split_cfg['right_focal_plane'])
+    # Top/bottom focal plane split
+    fm.wrap_dets('det_upper', aman.focal_plane.eta > split_cfg['top_focal_plane'])
+    fm.wrap_dets('det_lower', aman.focal_plane.eta <= split_cfg['top_focal_plane'])
+    # Inner/outter pixel split
+    r =  np.sqrt(aman.focal_plane.xi**2 + aman.focal_plane.eta**2)
+    fm.wrap_dets('det_in', r < split_cfg['central_pixels'])
+    fm.wrap_dets('det_out', r >= split_cfg['central_pixels'])
+
+    # Preproc dependent splits
+    if proc_aman is None:
+        try:
+            proc_aman = aman.preprocess
+        except:
+            print('Preprocess information not present, cannot generate preprocess dependent splits.')
+            for k in split_cfg.keys():
+                split_aman.wrap(f'{k}_threshold', split_cfg[k])
+            split_aman.wrap('cuts', fm)
+            return split_aman
+
+    # This one is a bit funky to be forcing it to be noiseQ_fit, and units matter!
+    if 'noiseQ_fit' in proc_aman:
+        # Noise split
+        if not 'high_noise' in split_cfg:
+            split_cfg['high_noise'] = 3.5e-5
+        fm.wrap_dets('high_noise', proc_aman.noiseQ_fit.fit[:,1] > split_cfg['high_noise'])
+        fm.wrap_dets('low_noise', proc_aman.noiseQ_fit.fit[:,1] <= split_cfg['high_noise'])
+        split_aman.wrap('noise_avg', np.nanmean(proc_aman.noiseQ_fit.fit[:,1]))
+
+    if 't2p' in proc_aman:
+        # T2P Leakage split
+        if not 'high_leakage' in split_cfg:
+            split_cfg['high_leakage'] = 1e-3
+        fm.wrap_dets('high_leakage', np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2) > split_cfg['high_leakage'])
+        fm.wrap_dets('low_leakage', np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2) <= split_cfg['high_leakage'])
+        split_aman.wrap('leakage_avg', np.nanmean(np.sqrt(proc_aman.t2p.lamQ**2 + proc_aman.t2p.lamU**2)),
                     [(0, 'dets')])
+    if 'hwpss_stats' in proc_aman:
+        # High 2f amplitude split
+        if not 'high_2f' in split_cfg:
+            split_cfg['high_2f'] = 1.5e-3
+        a2 = aman.det_cal.phase_to_pW*np.sqrt(proc_aman.hwpss_stats.coeffs[:,2]**2 + proc_aman.hwpss_stats.coeffs[:,3]**2)
+        fm.wrap_dets('high_2f', a2 > split_cfg['high_2f'])
+        fm.wrap_dets('low_2f', a2 <= split_cfg['high_2f'])
+        split_aman.wrap('2f_avg', np.nanmean(a2), [(0, 'dets')])
+    # Left/right subscans
+    if 'turnaround_flags' in proc_aman:
+        fm.wrap('scan_left', proc_aman.turnaround_flags.left_scan)
+        fm.wrap('scan_right', proc_aman.turnaround_flags.right_scan)
+
+    for k in split_cfg.keys():
+        split_aman.wrap(f'{k}_threshold', split_cfg[k])
+
+    split_aman.wrap('cuts', fm)
+
     return split_aman
diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
@@ -1214,8 +1214,6 @@ def calc_and_save(self, aman, proc_aman):
 
             if self.calc_cfgs.get("trim_samps") is not None:
                 trim = self.calc_cfgs["trim_samps"]
-                aman.restrict('samps', (aman.samps.offset + trim,
-                                        aman.samps.offset + aman.samps.count - trim))
                 proc_aman.restrict('samps', (proc_aman.samps.offset + trim,
                                              proc_aman.samps.offset + proc_aman.samps.count - trim))
                 filt_aman.restrict('samps', (filt_aman.samps.offset + trim,
@@ -1225,7 +1223,8 @@ def calc_and_save(self, aman, proc_aman):
 
         bands = np.unique(aman.det_info.wafer.bandpass)
         bands = bands[bands != 'NC']
-        rc_aman = core.AxisManager(aman.dets, aman.samps)
+        # align samps w/ proc_aman to include samps restriction when loading back from db.
+        rc_aman = core.AxisManager(proc_aman.dets, proc_aman.samps)
         pca_det_mask = np.full(aman.dets.count, False, dtype=bool)
         relcal = np.zeros(aman.dets.count)
         pca_weight0 = np.zeros(aman.dets.count)
@@ -1450,7 +1449,7 @@ class SplitFlags(_Preprocess):
     name = "split_flags"
 
     def calc_and_save(self, aman, proc_aman):
-        split_flg_aman = obs_ops.flags.get_split_flags(aman, proc_aman, split_cfg=self.calc_cfgs)
+        split_flg_aman = obs_ops.splits.get_split_flags(aman, proc_aman, split_cfg=self.calc_cfgs)
 
         self.save(proc_aman, split_flg_aman)