Adding AZSS correction per subscan

sotodlib/preprocess/processes.py

@@ -839,14 +839,16 @@ def calc_and_save(self, aman, proc_aman):
             self.save(proc_aman, aman[self.azss_stats_name])
         else:
             if self.proc_aman_turnaround_info:
-                _f = attergetter(self.proc_aman_turnaround_info)
+                _f = attrgetter(self.proc_aman_turnaround_info)
                 turnaround_info = _f(proc_aman)
             else:
                 turnaround_info = None
-            azss_stats, _ = tod_ops.azss.get_azss(aman, turnaround_info=turnaround_info,
-                                                  azss_stats_name=self.azss_stats_name,
-                                                  merge_stats = False, merge_model=False,
-                                                  **self.calc_cfgs)
+            _azss = tod_ops.azss.get_azss(aman, azss_stats_name=self.azss_stats_name,
+                                  turnaround_info=turnaround_info,
+                                  merge_stats=False, merge_model=False,
+                                  subtract_in_place=self.process_cfgs["subtract"], 
+                                  **self.calc_cfgs)
+            azss_stats = _azss[0]
             self.save(proc_aman, azss_stats)
 
     def save(self, proc_aman, azss_stats):
@@ -856,22 +858,22 @@ def save(self, proc_aman, azss_stats):
             proc_aman.wrap(self.azss_stats_name, azss_stats)
 
     def process(self, aman, proc_aman):
-        subtract = self.process_cfgs.get('subtract', False)
         if self.proc_aman_turnaround_info:
-            _f = attergetter(self.proc_aman_turnaround_info)
+            _f = attrgetter(self.proc_aman_turnaround_info)
             turnaround_info = _f(proc_aman)
         else:
             turnaround_info = None
         if self.azss_stats_name in proc_aman:
-            if subtract:
+            if self.process_cfgs["subtract"]:
                 tod_ops.azss.subtract_azss(aman, proc_aman[self.azss_stats_name],
                                            signal = self.calc_cfgs.get('signal'),
                                            in_place=True)
         else:
             tod_ops.azss.get_azss(aman, azss_stats_name=self.azss_stats_name,
                                   turnaround_info=turnaround_info,
-                                  merge_stats = True, merge_model=False,
-                                  subtract_in_place=subtract, **self.calc_cfgs)
+                                  merge_stats=True, merge_model=False,
+                                  subtract_in_place=self.process_cfgs["subtract"], 
+                                  **self.calc_cfgs)
 
 class GlitchFill(_Preprocess):
     """Fill glitches. All process configs go to `fill_glitches`.
@@ -1518,7 +1520,7 @@ class UnionFlags(_Preprocess):
 
     Saves results for aman under the "flags.[total_flags_label]" field.
 
-     Example config block:
+     Example config block::
 
         - name : "union_flags"
           process:
@@ -1583,6 +1585,7 @@ def __init__(self, step_cfgs):
         super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
+        coeff_aman = get_qu_common_mode_coeffs(aman, Q_signal, U_signal, merge)
         self.save(proc_aman, aman)
 
     def save(self, proc_aman, aman):
@@ -1629,7 +1632,7 @@ 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):
         if self.select_cfgs is None:
             return meta
@@ -1680,24 +1683,20 @@ def __init__(self, step_cfgs):
         super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
-        if 'flags' not in proc_aman._fields:
-            from sotodlib.core import FlagManager
-            proc_aman.wrap('flags', FlagManager.for_tod(proc_aman))
-        proc_aman.flags.wrap("left_scan", aman.flags.left_scan)
-        proc_aman.flags.wrap("right_scan", aman.flags.right_scan)
-
-        subscan_stats_T = proc_aman[self.stats_name+"_T"]
-        subscan_stats_Q = proc_aman[self.stats_name+"_Q"]
-        subscan_stats_U = proc_aman[self.stats_name+"_U"]
-
-        msk_ss, msk_det = tod_ops.flags.get_noisy_subscan_flags(
-            aman, subscan_stats_T = subscan_stats_T,
-            subscan_stats_Q = subscan_stats_Q, 
-            subscan_stats_U = subscan_stats_U, **self.calc_cfgs)
+        from so3g.proj import RangesMatrix, Ranges
+        msk_ss = RangesMatrix.zeros((aman.dets.count, aman.samps.count))
+        msk_det = np.ones(aman.dets.count, dtype=bool)
+        signal = self.calc_cfgs["subscan_stats"] if isinstance(self.calc_cfgs["subscan_stats"], list) else [self.calc_cfgs["subscan_stats"]]
+        for sig in signal:
+            self.calc_cfgs["subscan_stats"] = proc_aman[self.stats_name+"_"+sig[-1]]
+            _msk_ss, _msk_det = tod_ops.flags.get_noisy_subscan_flags(
+                aman, **self.calc_cfgs)
+            msk_ss += _msk_ss
+            msk_det &= _msk_det
         ss_aman = core.AxisManager(aman.dets, aman.samps)
-        ss_aman.wrap("valid_subscan", ~msk_ss, [(0, 'dets'), (1, 'samps')])
+        ss_aman.wrap("valid_subscans", msk_ss, [(0, 'dets'), (1, 'samps')])
         det_aman = core.AxisManager(aman.dets)
-        det_aman.wrap("valid_dets", ~msk_det)
+        det_aman.wrap("valid_dets", msk_det)
         self.save(proc_aman, ss_aman, "noisy_subscan_flags")
         self.save(proc_aman, det_aman, "noisy_dets_flags")
 
@@ -1712,7 +1711,8 @@ def select(self, meta, proc_aman=None):
             return meta
         if proc_aman is None:
             proc_aman = meta.preprocess
-        meta.restrict('dets', proc_aman.dets.vals[proc_aman.noisy_dets_flags.valid_dets])
+        keep = proc_aman.noisy_dets_flags.valid_dets
+        meta.restrict('dets', proc_aman.dets.vals[keep])
         return meta
 
 _Preprocess.register(SplitFlags)

sotodlib/tod_ops/azss.py

@@ -1,11 +1,12 @@
 """Module for estimating Azimuth Synchronous Signal (azss)"""
 import numpy as np
+from operator import attrgetter
 from numpy.polynomial import legendre as L
 from scipy.optimize import curve_fit
 from scipy.interpolate import interp1d
 from sotodlib import core, tod_ops
 from sotodlib.tod_ops import bin_signal, apodize, filters
-from so3g.proj import Ranges
+from so3g.proj import Ranges, RangesMatrix
 import logging
 
 logger = logging.getLogger(__name__)
@@ -219,7 +220,7 @@ def get_azss(aman, signal='signal', az=None, frange=None, bins=100, flags=None,
         In 'interpolate', binned signal is used directly.
         In 'fit', fitting is applied to the binned signal.
         Defaults to 'interpolate'.
-    max_mode: integer, optinal
+    max_mode: integer, optional
         The number of Legendre modes to use for azss when method is 'fit'. Required when method is 'fit'.
     subtract_in_place: bool
         If True, it subtract the modeled tod from original signal. The aman.signal will be modified.
@@ -273,7 +274,9 @@ def get_azss(aman, signal='signal', az=None, frange=None, bins=100, flags=None,
 
     if flags is None:
         flags = Ranges.from_mask(np.zeros(aman.samps.count).astype(bool))
-
+    else: 
+        flags = aman.flags.reduce(flags=flags, method='union', wrap=False)
+
     if left_right:
         if turnaround_info is None:
             turnaround_info = aman.flags
@@ -287,21 +290,22 @@ def get_azss(aman, signal='signal', az=None, frange=None, bins=100, flags=None,
             left_mask = turnaround_info.left_scan
             right_mask = turnaround_info.right_scan
         else:
-            print("")
-            print("HEREEEEEEE")
             left_mask = turnaround_info.valid_left_scans
             right_mask = turnaround_info.valid_right_scans
 
         azss_left = _prepare_azss_stats(aman, signal, az, frange, bins, flags+left_mask, apodize_edges,
                                         apodize_edges_samps, apodize_flags, apodize_flags_samps,
                                         method=method, max_mode=max_mode)
         azss_left.add_axis(aman.samps)
-        azss_left.wrap('mask', left_mask, [(0, 'samps')])
+        azss_left.add_axis(aman.dets)
+        azss_left.wrap('mask', left_mask, [(0, 'dets'), (1, 'samps')])
         azss_right = _prepare_azss_stats(aman, signal, az, frange, bins, flags+right_mask, apodize_edges,
                                         apodize_edges_samps, apodize_flags, apodize_flags_samps,
                                         method=method, max_mode=max_mode)
         azss_right.add_axis(aman.samps)
-        azss_right.wrap('mask', right_mask, [(0, 'samps')])
+        azss_right.add_axis(aman.dets)
+        azss_right.wrap('mask', right_mask, [(0, 'dets'), (1, 'samps')])
+
         azss_stats = core.AxisManager(aman.dets)
         azss_stats.wrap('azss_stats_left', azss_left)
         azss_stats.wrap('azss_stats_right', azss_right)
@@ -322,14 +326,14 @@ def get_azss(aman, signal='signal', az=None, frange=None, bins=100, flags=None,
             if subtract_in_place:
                 if signal_name is None:
                     lmask = left_mask.mask()
-                    signal[:,lmask] -= model_left[:,lmask].astype(signal.dtype)
+                    signal[lmask] -= model_left[lmask].astype(signal.dtype)
                     rmask = right_mask.mask()
-                    signal[:,rmask] -= model_right[:,rmask].astype(signal.dtype)
+                    signal[rmask] -= model_right[rmask].astype(signal.dtype)
                 else:
                     lmask = left_mask.mask()
-                    aman[signal_name][:,lmask] -= model_left[:,lmask].astype(aman[signal_name].dtype)
+                    aman[signal_name][lmask] -= model_left[lmask].astype(aman[signal_name].dtype)
                     rmask = right_mask.mask()
-                    aman[signal_name][:,rmask] -= model_right[:,rmask].astype(aman[signal_name].dtype)
+                    aman[signal_name][rmask] -= model_right[rmask].astype(aman[signal_name].dtype)
         else:
             azss_stats, model, _ = get_model_sig_tod(aman, azss_stats, az)
             if merge_model:
@@ -449,7 +453,7 @@ def subtract_azss(aman, azss_stats, signal='signal', subtract_name='azss_remove'
             if azss_stats.left_right:
                 for model, azss_fld in zip([model_left, model_right], ['azss_stats_left', 'azss_stats_right']):
                     mask = azss_stats[azss_fld]['mask'].mask()
-                    aman[signal_name][:, mask] -= model[:, mask].astype(aman[signal_name].dtype)
+                    aman[signal_name][mask] -= model[mask].astype(aman[signal_name].dtype)
             else:
                 aman[signal_name] -= model_left.astype(aman[signal_name].dtype)
     else:

sotodlib/tod_ops/deproject.py

@@ -49,7 +49,7 @@ def get_qu_common_mode_coeffs(aman, Q_signal=None, U_signal=None, merge=False):
     return output_aman
 
 def subtract_qu_common_mode(aman, Q_signal=None, U_signal=None, coeff_aman=None,
-                            merge=False):
+                            merge=False, subtract=True):
     """
     Subtracts the median signal (template) from each detector scaled by the a 
     coupling coefficient per detector.

sotodlib/tod_ops/flags.py

@@ -546,30 +546,6 @@ def get_dark_dets(aman, merge=True, overwrite=True, dark_flags_name='darks'):
     """
     Identify and flag dark detectors in the given aman object.
 
-    Parameters:
-    ----------
-    aman : AxisManager
-        The tod.
-    merge : bool, optional
-        If True, merge the dark detector flags into the aman.flags. Default is True.
-    overwrite : bool, optional
-        If True, overwrite existing flags with the same name. Default is True.
-    dark_flags_name : str, optional
-        The name to use for the dark detector flags in aman.flags. Default is 'darks'.
-
-    Returns:
-    -------
-    mskdarks: RangesMatrix
-        A matrix of ranges indicating the dark detectors.
-
-    Raises:
-    -------
-    ValueError
-        If merge is True and dark_flags_name already exists in aman.flags and overwrite is False.
-    """
-    """
-    Identify and flag dark detectors in the given aman object.
-
     Parameters:
     ----------
     aman : AxisManager
@@ -741,7 +717,7 @@ def get_inv_var_flags(aman, signal_name='signal', nsigma=5,
 
     return mskinvar
 
-def get_subscans(aman, merge=True, include_turnarounds=False, overwrite=True):
+def get_subscans(aman, flags="flags", merge=True, include_turnarounds=False, overwrite=True):
     """
     Returns an axis manager with information about subscans.
     This includes direction and a ranges matrix (subscans samps)
@@ -767,6 +743,7 @@ def get_subscans(aman, merge=True, include_turnarounds=False, overwrite=True):
     """
     if not include_turnarounds:
         ss_ind = (~aman.flags.turnarounds).ranges() # sliceable indices (first inclusive, last exclusive) for subscans
+        #ss_ind = (~aman[flags]["turnarounds"]).ranges()
     else:
         left = aman.flags.left_scan.ranges()
         right = aman.flags.right_scan.ranges()
@@ -1022,13 +999,18 @@ def noise_fit_flags(aman, low_wn, high_wn, high_fk):
     else:
         return None
 
-def get_noisy_subscan_flags(aman, subscan_stats_T, subscan_stats_Q, subscan_stats_U,
+def get_noisy_subscan_flags(aman, subscan_stats, flags="flags",
                          nstd_lim=None, ptp_lim=None, kurt_lim=None, 
-                         skew_lim=None, merge=False, overwrite=False, 
-                         name="bad_subscan", cut_bad=False):
+                         skew_lim=None, noisy_detector_lim=0.9,
+                         merge=False, overwrite=False, 
+                         name="noisy_subscan"):
     """
     Identify and flag bad subscans based on various statistical thresholds.
+    aman : AxisManager 
         The tod.
+    subscan_stats : dict
+        Dictionary containing statistical metrics for subscans. Keys should 
+        include 'std', 'ptp', 'kurtosis', and 'skew'.
     nstd_lim : float [optional]
         Threshold for standard deviation.
     ptp_lim : float [optional]
@@ -1048,46 +1030,49 @@ def get_noisy_subscan_flags(aman, subscan_stats_T, subscan_stats_Q, subscan_stat
     -------
     noisy_subscan_flags : RangesMatrix
         RangesMatrix of bad subscans.
+    noisy_detector_flags : ndarray
+        Array indicating detectors that are too noisy for more than 50% of the subscan duration.
+
     """
+
     if 'flags' not in aman:
         overwrite = False
         merge = False
     if overwrite and name in aman.flags:
         aman.flags.move(name, None)
 
-    median_Qstd = np.median(subscan_stats_Q["std"], axis=1)[:, np.newaxis]
-    median_Ustd = np.median(subscan_stats_U["std"], axis=1)[:, np.newaxis]
-
-    noisy_subscan_indicator = np.zeros_like(subscan_stats_Q["std"], dtype=bool)
+    median_std = np.median(subscan_stats["std"], axis=1)[:, np.newaxis]
+    noisy_subscan_indicator = np.zeros_like(subscan_stats["std"], dtype=bool)
 
-    if ptp_lim is not None:
-        noisy_subscan_indicator |= (subscan_stats_T['ptp'] > ptp_lim)
+    if ptp_lim is not None and 'ptp' in subscan_stats:
+        noisy_subscan_indicator |= (subscan_stats['ptp'] > ptp_lim)
 
-    if nstd_lim is not None:
-        noisy_subscan_indicator |= (subscan_stats_Q['std'] > median_Qstd * nstd_lim)
-        noisy_subscan_indicator |= (subscan_stats_U['std'] > median_Ustd * nstd_lim)
+    if nstd_lim is not None and 'std' in subscan_stats:
+        noisy_subscan_indicator |= (subscan_stats['std'] > median_std * nstd_lim)
 
-    if kurt_lim is not None:
-        noisy_subscan_indicator |= (np.abs(subscan_stats_Q['kurtosis']) > kurt_lim)
-        noisy_subscan_indicator |= (np.abs(subscan_stats_U['kurtosis']) > kurt_lim)
+    if kurt_lim is not None and 'kurtosis' in subscan_stats:
+        noisy_subscan_indicator |= (np.abs(subscan_stats['kurtosis']) > kurt_lim)
 
-    if skew_lim is not None:
-        noisy_subscan_indicator |= (np.abs(subscan_stats_Q['skew']) > skew_lim)
-        noisy_subscan_indicator |= (np.abs(subscan_stats_U['skew']) > skew_lim)
-
-    ss_info = aman.subscan_info if 'subscan_info' in aman else get_subscans(aman, merge=False)
-    ssf = ss_info.subscan_flags
-    zeros=RangesMatrix.zeros((1, aman.samps.count)) # Needed when no subcsans are True
+    if skew_lim is not None and 'skew' in subscan_stats:
+        noisy_subscan_indicator |= (np.abs(subscan_stats['skew']) > skew_lim)
+
+    if "subscan_info" in aman:
+        ssf = aman.subscan_info.subscan_flags
+    elif "subscans" in aman: 
+        # If reading from proc_aman
+        ssf = aman.subscans
+    else:
+        get_subscans(aman, flags=flags, merge=False)
 
+    zeros=RangesMatrix.zeros((1, aman.samps.count)) # Needed when no subscans are True
     # For each det, select flagged subscans and merge their Ranges
     noisy_subscan_flags = RangesMatrix([np.sum(RangesMatrix.concatenate([ssf[noisy_subscan_indicator[idet]], zeros]), axis=0) for idet in range(aman.dets.count)])
-
-    # Detectors which are too noisy for > 50% of the the subscan duration
-    noisy_detector_flags = np.array([np.sum(np.diff(rng.ranges(), axis=1)) / aman.samps.count for rng in noisy_subscan_flags.ranges]) > 0.5
-
+    # Detectors which are too noisy for > 50% of the subscan duration
+    noisy_detector_flags = np.array([np.sum(np.diff(rng.ranges(), axis=1)) / aman.samps.count for rng in noisy_subscan_flags.ranges]) > noisy_detector_lim
+
     if merge:
         if name in aman.flags and not overwrite:
             raise ValueError(f"Flag name {name} already exists in aman.flags")
         aman.flags.wrap(name, noisy_subscan_flags)
 
-    return noisy_subscan_flags, noisy_detector_flags
+    return noisy_subscan_flags, ~noisy_detector_flags