From e2c8b1049738cdbc67cc2774d87689996a913ddb Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Tue, 12 Nov 2024 19:03:22 +0000
Subject: [PATCH 1/9] Add subscan operations to preprocess

Add basic statistics, PSDs, and glitches on subscans.
---
 sotodlib/preprocess/processes.py |  81 +++++++++++++++++-
 sotodlib/tod_ops/fft_ops.py      |  28 ++++++-
 sotodlib/tod_ops/flags.py        | 139 ++++++++++++++++++++++++++++++-
 3 files changed, 241 insertions(+), 7 deletions(-)

diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index d9df285b5..54287f581 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -172,6 +172,7 @@ class GlitchDetection(_FracFlaggedMixIn, _Preprocess):
           buffer: 10
           hp_fc: 1
           n_sig: 10
+          subscan: False
         save: True
         plot:
             plot_ds_factor: 50
@@ -340,6 +341,7 @@ def plot(self, aman, proc_aman, filename):
                              plot_ds_factor=self.plot_cfgs.get("plot_ds_factor", 50), filename=filename.replace('{name}', f'{ufm}_jump_signal_diff'))
             plot_flag_stats(aman, proc_aman[name], flag_type='jumps', filename=filename.replace('{name}', f'{ufm}_jumps_stats'))
 
+
 class PSDCalc(_Preprocess):
     """ Calculate the PSD of the data and add it to the AxisManager under the
     "psd" field.
@@ -349,6 +351,7 @@ class PSDCalc(_Preprocess):
       - "name : "psd"
         "signal: "signal" # optional
         "wrap": "psd" # optional
+        "subscan": False # optional
         "process":
           "psd_cfgs": # optional, kwargs to scipy.welch
             "nperseg": 1024
@@ -363,13 +366,19 @@ class PSDCalc(_Preprocess):
     def __init__(self, step_cfgs):
         self.signal = step_cfgs.get('signal', 'signal')
         self.wrap = step_cfgs.get('wrap', 'psd')
+        self.subscan = step_cfgs.get('subscan', False)
 
         super().__init__(step_cfgs)
         
 
     def process(self, aman, proc_aman):
-        freqs, Pxx = tod_ops.fft_ops.calc_psd(aman, signal=aman[self.signal],
-                                              **self.process_cfgs)
+        if not self.subscan:
+            calc_psd = tod_ops.fft_ops.calc_psd
+        else:
+            calc_psd = tod_ops.fft_ops.calc_psd_subscan
+
+        freqs, Pxx = calc_psd(aman, signal=aman[self.signal],
+                              **self.process_cfgs)
         fft_aman = core.AxisManager(
             aman.dets, 
             core.OffsetAxis("nusamps",len(freqs))
@@ -385,6 +394,74 @@ def save(self, proc_aman, fft_aman):
         if not(self.save_cfgs is None):
             proc_aman.wrap(self.wrap, fft_aman)
 
+
+class IdentifySubscans(_Preprocess):
+    """ Get subscan info and wrap it into aman.subscans.
+
+    Example config block::
+
+      - "name : "subscans"
+        "calc": True
+        "save": True
+
+    """
+
+    name = "subscans"
+
+    def __init__(self, step_cfgs):
+        super().__init__(step_cfgs)
+
+    def process(self, aman, proc_aman):
+        tod_ops.flags.get_subscans(aman, merge=True)
+
+    def calc_and_save(self, aman, proc_aman):
+        self.save(proc_aman, aman.subscans)
+
+    def save(self, proc_aman, subscan_aman):
+        if not(self.save_cfgs is None):
+            proc_aman.wrap('subscans', subscan_aman)
+
+class TODStats(_Preprocess):
+    """ Get basic statistics from a TOD or its power spectrum.
+
+    Example config block:
+
+      - "name : "stats"
+        "signal: "signal" # optional
+        "wrap": "stats" # optional
+        "calc":
+          "stat_names": ["median", "std"]
+          "split_subscans": False # optional
+          "mask": # optional, for cutting a power spectrum in frequency
+            "psd_aman": "psd"
+            "low_f": 1
+            "high_f": 10
+        "save": True
+
+    """
+    name = "stats"
+    def __init__(self, step_cfgs):
+        self.signal = step_cfgs.get('signal', 'signal')
+        self.wrap = step_cfgs.get('wrap', 'stats')
+
+        super().__init__(step_cfgs)
+
+    def calc_and_save(self, aman, proc_aman):
+        if self.calc_cfgs.get('mask') is not None:
+            mask_dict = self.calc_cfgs.get('mask')
+            freqs = aman[mask_dict['psd_aman']]['freqs']
+            low_f, high_f = mask_dict['low_f'], mask_dict['high_f']
+            fmask = np.all([freqs >= low_f, freqs <= high_f], axis=0)
+            self.calc_cfgs['mask'] = fmask
+
+        stats_aman = tod_ops.flags.get_stats(aman, aman[self.signal], **self.calc_cfgs)
+        self.save(proc_aman, stats_aman)
+
+    def save(self, proc_aman, stats_aman):
+        if not(self.save_cfgs is None):
+            proc_aman.wrap(self.wrap, stats_aman)
+
+
 class Noise(_Preprocess):
     """Estimate the white noise levels in the data. Assumes the PSD has been
     wrapped into the AxisManager. All calculation configs goes to `calc_wn`. 
diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py
index 4f173e8e6..dc90cd221 100644
--- a/sotodlib/tod_ops/fft_ops.py
+++ b/sotodlib/tod_ops/fft_ops.py
@@ -10,7 +10,7 @@
 from sotodlib import core
 
 from . import detrend_tod
-
+from .flags import get_subscan_signal
 
 def _get_num_threads():
     # Guess how many threads we should be using in FFT ops...
@@ -281,6 +281,32 @@ def calc_psd(
         aman.wrap("Pxx", Pxx, [(0,"dets"),(1,"nusamps")])
     return freqs, Pxx
 
+def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=None, **kwargs):
+    if signal is None:
+        signal = aman.signal
+
+    fs = 1 / np.nanmedian(np.diff(aman.timestamps))
+    if "nperseg" not in kwargs:
+        if freq_spacing is not None:
+            nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing))))
+        else:
+            duration_samps = np.around((aman.subscans.stop_time - aman.subscans.start_time) * fs)
+            nperseg = int(2 ** (np.around(np.log2(np.median(duration_samps) / 4))))
+        kwargs["nperseg"] = nperseg
+
+    Pxx = []
+    for iss in range(aman.subscans.subscans.count):
+        signal_ss = get_subscan_signal(aman, signal, iss)
+        freqs, pxx_sub = welch(signal_ss, fs, **kwargs)
+        Pxx.append(pxx_sub)
+    Pxx = np.array(Pxx)
+    Pxx = Pxx.transpose(1, 2, 0) # Dets, nusamps, subscans
+    if merge:
+        wrap = "Pxx_ss" if wrap is None else wrap
+        aman.merge( core.AxisManager(core.OffsetAxis("nusamps_ss", len(freqs))))
+        aman.wrap("freqs_ss", freqs, [(0,"nusamps_ss")])
+        aman.wrap(wrap, Pxx, [(0,aman.dets),(1,"nusamps_ss"), (2, "subscans")])
+    return freqs, Pxx
 
 def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10):
     """
diff --git a/sotodlib/tod_ops/flags.py b/sotodlib/tod_ops/flags.py
index 4f0088eb7..39c32c622 100644
--- a/sotodlib/tod_ops/flags.py
+++ b/sotodlib/tod_ops/flags.py
@@ -319,7 +319,8 @@ def get_glitch_flags(aman,
                      overwrite=False,
                      name="glitches",
                      full_output=False,
-                     edge_guard=2000):
+                     edge_guard=2000,
+                     subscan=False):
     """
     Find glitches with fourier filtering. Translation from moby2 as starting point
 
@@ -350,6 +351,8 @@ def get_glitch_flags(aman,
     edge_guard : int
         Number of samples at the beginning and end of the tod to exclude from
         the returned glitch RangesMatrix. Defaults to 2000 samples (10 sec).
+    subscan : bool
+        If True, compute the glitch threshold on a per-subscan basis. Includes turnarounds.
 
     Returns
     -------
@@ -370,9 +373,18 @@ def get_glitch_flags(aman,
         ds = int(fvec.shape[1]/20000)
     else: 
         ds = 1
-    iqr_range = 0.741 * stats.iqr(fvec[:,::ds], axis=1)
-    # get flags
-    msk = fvec > iqr_range[:, None] * n_sig
+
+    if subscan:
+        # We include turnarounds
+        subscan_indices = np.concatenate([aman.flags.left_scan.ranges(), (~aman.flags.left_scan).ranges()])
+    else:
+        subscan_indices = np.array([[0, fvec.shape[1]]])
+
+    msk = np.zeros_like(fvec, dtype='bool')
+    for ss in subscan_indices:
+        iqr_range = 0.741 * stats.iqr(fvec[:,ss[0]:ss[1]:ds], axis=1)
+        # get flags
+        msk[:,ss[0]:ss[1]] = fvec[:,ss[0]:ss[1]] > iqr_range[:, None] * n_sig
     msk[:,:edge_guard] = False
     msk[:,-edge_guard:] = False
     flag = RangesMatrix([Ranges.from_bitmask(m) for m in msk])
@@ -672,3 +684,122 @@ def get_inv_var_flags(aman, signal_name='signal', nsigma=5,
             aman.flags.wrap(inv_var_flag_name, mskinvar, [(0, 'dets'), (1, 'samps')])
 
     return mskinvar
+
+def get_subscans(aman, merge=True):
+    """
+    Returns an axis manager with information about subscans.
+    This includes direction, start time, stop time, and a ranges matrix (subscans samps)
+    True inside each subscan. Subscans are defined excluding turnarounds.
+    """
+    ss_ind = (~aman.flags.turnarounds).ranges() # sliceable indices (first inclusive, last exclusive) for subscans
+    start_inds, end_inds = ss_ind.T
+    n_subscan = ss_ind.shape[0]
+    tt = aman.timestamps
+    subscan_aman = core.AxisManager(aman.samps, core.IndexAxis("subscans", n_subscan))
+
+    is_left = aman.flags.left_scan.mask()[start_inds]
+    subscan_aman.wrap('direction', np.array(['left' if is_left[ii] else 'right' for ii in range(n_subscan)]), [(0, 'subscans')])
+
+    subscan_aman.wrap('start_time', tt[start_inds], [(0, 'subscans')])
+    stop_time = aman.timestamps[ss_ind[:-1,1]]
+    last_time = [2*tt[-1] - tt[-2] if end_inds[-1] == tt.size else tt[end_inds[-1]]]
+    subscan_aman.wrap('stop_time', np.concatenate([stop_time, last_time]), [(0, 'subscans')])
+    rm = RangesMatrix([Ranges.from_array(np.atleast_2d(ss), tt.size) for ss in ss_ind])
+    subscan_aman.wrap('subscan_flags', rm, [(0, 'subscans'), (1, 'samps')]) # True in the subscan
+    if merge:
+        aman.wrap('subscans', subscan_aman)
+    return subscan_aman
+
+def get_subscan_signal(aman, arr, isub=None):
+    """
+    Split an array into subscans.
+
+    Parameters
+    ----------
+    aman : AxisManager
+        Input AxisManager.
+    arr : Array
+        Input array.
+    isub : int
+        (Optional). Index of the desired subscan. May also be a list of indices.
+        If None, all are used.
+    """
+    if np.isscalar(isub):
+        return apply_rng(arr, aman.subscans.subscan_flags[isub])
+    else:
+        if isub is None:
+            isub = range(len(aman.subscans.subscan_flags))
+        return [apply_rng(arr, aman.subscans.subscan_flags[ii]) for ii in isub]
+
+def apply_rng(arr, rng):
+    """
+    Helper function to get signals on subcsans.
+
+    Parameters
+    ----------
+    arr : Array
+        Array containing the signal. Should have one axis of len (samps)
+    rng : Ranges
+        Ranges object of len (samps) selecting the desired range
+    """
+    slc = slice(*rng.ranges()[0])
+    isamps = np.where(np.array(arr.shape) == rng.count)[0][0]
+    ndslice = tuple((slice(None) if ii != isamps else slc for ii in range(arr.ndim)))
+    return arr[ndslice]
+
+def wrap_info(aman, info_aman_name, info, info_names, merge=True):
+    """Wrap info into an aman. Assumed to be (dets,) or (dets, subscans)"""
+    if info[0].ndim == 1:
+        info_aman = core.AxisManager(aman.dets)
+        axmap = [(0, 'dets')]
+    elif info[0].ndim == 2:
+        info_aman = core.AxisManager(aman.dets, aman.subscans.subscans)
+        axmap = [(0, 'dets'), (1, 'subscans')]
+
+    for ii in range(len(info_names)):
+        info_aman.wrap(info_names[ii], info[ii], axmap)
+    if merge:
+        if info_aman_name in aman.keys():
+            aman[info_aman_name].merge(info_aman)
+        else:
+            aman.wrap(info_aman_name, info_aman)
+    return info_aman
+
+def get_stats(aman, signal, stat_names, split_subscans=False, mask=None, name="stats", merge=False):
+    """
+    Calculate basic statistics on a TOD or power spectrum.
+
+    Parameters
+    ----------
+    aman : AxisManager
+        Input AxisManager.
+    signal : Array
+        Input signal. Statistics will be computed over *axis 1*.
+    stat_names : list
+        List of strings identifying which statistics to run.
+    split_subscans : bool
+        (Optional). If True statistics will be computed on subscans. Assumes aman.subscans exists already.
+    mask : Array
+        (Optional). Mask to apply before computation. 1d array for advanced indexing, or a slice object.
+    name : str
+        Name of axis manager to add to aman if merge is True.
+    """
+    stat_names = np.atleast_1d(stat_names)
+    fn_dict = {'mean': np.mean, 'median': np.median, 'ptp': np.ptp, 'std': np.std,
+                     'kurtosis': stats.kurtosis, 'skew': stats.skew}
+
+    if split_subscans:
+        if mask is not None:
+            raise ValueError("Cannot mask samples and split subscans")
+        stats_arr = []
+        for iss in range(aman.subscans.subscans.count):
+            data = get_subscan_signal(aman, signal, iss)
+            stats_arr.append([fn_dict[name](data, axis=1) for name in stat_names]) # Samps axis assumed to be 1
+        stats_arr = np.array(stats_arr).transpose(1, 2, 0) # stat, dets, subscan
+    else:
+        if mask is None:
+            mask = slice(None)
+        stats_arr = np.array([fn_dict[name](data[:, mask], axis=1) for name in stat_names]) # Samps axis assumed to be 1
+
+    info_aman = wrap_info(aman, name, stats_arr, stat_names, merge)
+    return info_aman

From 67b59b075e993c1cf56230585227821509d8e2cc Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Wed, 20 Nov 2024 15:42:50 +0000
Subject: [PATCH 2/9] subscan preprocess fixes and upgrades

---
 sotodlib/preprocess/processes.py |  37 +++-------
 sotodlib/tod_ops/fft_ops.py      | 114 +++++++++++++++++++++++++++----
 sotodlib/tod_ops/flags.py        |  52 +++++++++++---
 3 files changed, 151 insertions(+), 52 deletions(-)

diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index 54287f581..6da06fd9d 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -395,32 +395,6 @@ def save(self, proc_aman, fft_aman):
             proc_aman.wrap(self.wrap, fft_aman)
 
 
-class IdentifySubscans(_Preprocess):
-    """ Get subscan info and wrap it into aman.subscans.
-
-    Example config block::
-
-      - "name : "subscans"
-        "calc": True
-        "save": True
-
-    """
-
-    name = "subscans"
-
-    def __init__(self, step_cfgs):
-        super().__init__(step_cfgs)
-
-    def process(self, aman, proc_aman):
-        tod_ops.flags.get_subscans(aman, merge=True)
-
-    def calc_and_save(self, aman, proc_aman):
-        self.save(proc_aman, aman.subscans)
-
-    def save(self, proc_aman, subscan_aman):
-        if not(self.save_cfgs is None):
-            proc_aman.wrap('subscans', subscan_aman)
-
 class TODStats(_Preprocess):
     """ Get basic statistics from a TOD or its power spectrum.
 
@@ -868,13 +842,20 @@ def calc_and_save(self, aman, proc_aman):
             calc_aman = core.AxisManager(aman.dets, aman.samps)
             calc_aman.wrap('turnarounds', ta, [(0, 'dets'), (1, 'samps')])
 
-        self.save(proc_aman, calc_aman)
+        if ('merge_subscans' not in self.calc_cfgs) or (self.calc_cfgs['merge_subscans']):
+            subscan_aman = aman.subscans
+        else:
+            subscan_aman = None
+
+        self.save(proc_aman, calc_aman, subscan_aman)
 
-    def save(self, proc_aman, turn_aman):
+    def save(self, proc_aman, turn_aman, subscan_aman):
         if self.save_cfgs is None:
             return
         if self.save_cfgs:
             proc_aman.wrap("turnaround_flags", turn_aman)
+            if subscan_aman is not None:
+                proc_aman.wrap("subscans", subscan_aman)
 
     def process(self, aman, proc_aman):
         tod_ops.flags.get_turnaround_flags(aman, **self.process_cfgs)
diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py
index dc90cd221..4d4b46be4 100644
--- a/sotodlib/tod_ops/fft_ops.py
+++ b/sotodlib/tod_ops/fft_ops.py
@@ -10,7 +10,6 @@
 from sotodlib import core
 
 from . import detrend_tod
-from .flags import get_subscan_signal
 
 def _get_num_threads():
     # Guess how many threads we should be using in FFT ops...
@@ -282,6 +281,13 @@ def calc_psd(
     return freqs, Pxx
 
 def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=None, **kwargs):
+    """
+    Calculate the power spectrum density of subscans using signal.welch().
+    Data defaults to aman.signal. aman.timestamps is used for times.
+    aman.subscans is used to identify subscans.
+    See calc_psd for arguments.
+    """
+    from .flags import get_subscan_signal
     if signal is None:
         signal = aman.signal
 
@@ -290,15 +296,20 @@ def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=Non
         if freq_spacing is not None:
             nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing))))
         else:
-            duration_samps = np.around((aman.subscans.stop_time - aman.subscans.start_time) * fs)
+            duration_samps = np.sum(aman.subscans.subscan_flags.mask(), axis=1)
+            duration_samps = duration_samps[duration_samps > 0]
             nperseg = int(2 ** (np.around(np.log2(np.median(duration_samps) / 4))))
         kwargs["nperseg"] = nperseg
 
     Pxx = []
     for iss in range(aman.subscans.subscans.count):
         signal_ss = get_subscan_signal(aman, signal, iss)
-        freqs, pxx_sub = welch(signal_ss, fs, **kwargs)
-        Pxx.append(pxx_sub)
+        axis = -1 if "axis" not in kwargs else kwargs["axis"]
+        if signal_ss.shape[axis] >= kwargs["nperseg"]:
+            freqs, pxx_sub = welch(signal_ss, fs, **kwargs)
+            Pxx.append(pxx_sub)
+        else:
+            Pxx.append(np.full((signal.shape[0], kwargs["nperseg"]//2+1), np.nan)) # Add nans if subscan is too short
     Pxx = np.array(Pxx)
     Pxx = Pxx.transpose(1, 2, 0) # Dets, nusamps, subscans
     if merge:
@@ -379,6 +390,8 @@ def fit_noise_model(
     f_max=100,
     merge_name="noise_fit_stats",
     merge_psd=True,
+    freq_spacing=None,
+    approx_fit=False,
 ):
     """
     Fits noise model with white and 1/f noise to the PSD of signal.
@@ -418,6 +431,10 @@ def fit_noise_model(
         If ``merge_fit`` is True then addes into axis manager with merge_name.
     merge_psd : bool
         If ``merg_psd`` is True then adds fres and Pxx to the axis manager.
+    freq_spacing : float
+        The approximate desired frequency spacing of the PSD. Passed to calc_psd.
+    approx_fit : bool
+        Get a rough fit instead of minimizing loglike.
     Returns
     -------
     noise_fit_stats : AxisManager
@@ -431,13 +448,14 @@ def fit_noise_model(
     if f is None or pxx is None:
         if psdargs is None:
             f, pxx = calc_psd(
-                aman, signal=signal, timestamps=aman.timestamps, merge=merge_psd
+                aman, signal=signal, timestamps=aman.timestamps, freq_spacing=freq_spacing, merge=merge_psd
             )
         else:
             f, pxx = calc_psd(
                 aman,
                 signal=signal,
                 timestamps=aman.timestamps,
+                freq_spacing=freq_spacing,
                 merge=merge_psd,
                 **psdargs,
             )
@@ -454,17 +472,21 @@ def fit_noise_model(
         pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1)
         fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest))
         p0 = [f[fidx], wnest, -pfit[0]]
-        bounds = [(0, None), (sys.float_info.min, None), (None, None)]
-        res = minimize(neglnlike, p0, args=(f, p), bounds=bounds, method="Nelder-Mead")
-        try:
-            Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True)
-            hessian_ndt, _ = Hfun(res["x"])
-            # Inverse of the hessian is an estimator of the covariance matrix
-            # sqrt of the diagonals gives you the standard errors.
-            covout[i] = np.linalg.inv(hessian_ndt)
-        except np.linalg.LinAlgError:
+        if approx_fit:
             covout[i] = np.full((3, 3), np.nan)
-        fitout[i] = res.x
+            fitout[i] = p0
+        else:
+            bounds = [(0, None), (sys.float_info.min, None), (None, None)]
+            res = minimize(neglnlike, p0, args=(f, p), bounds=bounds, method="Nelder-Mead")
+            try:
+                Hfun = ndt.Hessian(lambda params: neglnlike(params, f, p), full_output=True)
+                hessian_ndt, _ = Hfun(res["x"])
+                # Inverse of the hessian is an estimator of the covariance matrix
+                # sqrt of the diagonals gives you the standard errors.
+                covout[i] = np.linalg.inv(hessian_ndt)
+            except np.linalg.LinAlgError:
+                covout[i] = np.full((3, 3), np.nan)
+            fitout[i] = res.x
 
     noise_model_coeffs = ["fknee", "white_noise", "alpha"]
     noise_fit_stats = core.AxisManager(
@@ -485,6 +507,68 @@ def fit_noise_model(
     return noise_fit_stats
 
 
+def fit_noise_model_subscan(
+    aman,
+    signal=None,
+    f=None,
+    pxx=None,
+    psdargs={},
+    fwhite=(10, 100),
+    lowf=1,
+    merge_fit=False,
+    f_max=100,
+    merge_name="noise_fit_stats_ss",
+    merge_psd=True,
+    freq_spacing=None,
+    approx_fit=False,
+):
+    """
+    Fits noise model with white and 1/f noise to the PSD of signal subscans.
+    Args are as for fit_noise_model.
+    """
+    fitout = np.empty((aman.dets.count, 3, aman.subscans.subscans.count))
+    covout = np.empty((aman.dets.count, 3, 3, aman.subscans.subscans.count))
+
+    if signal is None:
+        signal = aman.signal
+
+    if f is None or pxx is None:
+        f, pxx = calc_psd_subscan(
+            aman,
+            signal=signal,
+            freq_spacing=freq_spacing,
+            merge=merge_psd,
+            **psdargs,
+            )
+
+    for isub in range(aman.subscans.subscans.count):
+        if np.all(np.isnan(pxx[...,isub])): # Subscan has been fully cut
+            fitout[..., isub] = np.full((aman.dets.count, 3), np.nan)
+            covout[..., isub] = np.full((aman.dets.count, 3, 3), np.nan)
+        else:
+            noise_model = fit_noise_model(aman, f=f, pxx=pxx[...,isub], fwhite=fwhite, lowf=lowf,
+                                          merge_fit=False, f_max=f_max, merge_psd=False, approx_fit=approx_fit)
+
+            fitout[..., isub] = noise_model.fit
+            covout[..., isub] = noise_model.cov
+
+    noise_fit_stats = core.AxisManager(
+        aman.dets,
+        noise_model.noise_model_coeffs,
+        aman.subscans.subscans
+        )
+    noise_fit_stats.wrap("fit", fitout, [(0, "dets"), (1, "noise_model_coeffs"), (2, "subscans")])
+    noise_fit_stats.wrap(
+        "cov",
+        covout,
+        [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs"), (3, "subscans")],
+    )
+
+    if merge_fit:
+        aman.wrap(merge_name, noise_fit_stats)
+    return noise_fit_stats
+
+
 def build_hpf_params_dict(
     filter_name,
     noise_fit=None,
diff --git a/sotodlib/tod_ops/flags.py b/sotodlib/tod_ops/flags.py
index 39c32c622..a0511978f 100644
--- a/sotodlib/tod_ops/flags.py
+++ b/sotodlib/tod_ops/flags.py
@@ -138,7 +138,7 @@ def get_det_bias_flags(aman, detcal=None, rfrac_range=(0.1, 0.7),
 def get_turnaround_flags(aman, az=None, method='scanspeed', name='turnarounds',
                          merge=True, merge_lr=True, overwrite=True, 
                          t_buffer=2., kernel_size=400, peak_threshold=0.1, rel_distance_peaks=0.3,
-                         truncate=False, qlim=1):
+                         truncate=False, qlim=1, merge_subscans=True, turnarounds_in_subscan=False):
     """
     Compute turnaround flags for a dataset.
 
@@ -172,6 +172,10 @@ def get_turnaround_flags(aman, az=None, method='scanspeed', name='turnarounds',
         (Optional). Truncate unstable scan segments if True in ``scanspeed`` method.
     qlim : float
         (Optional). Azimuth threshold percentile for ``az`` method turnaround detection.
+    merge_subscans : bool
+        (Optional). Also merge an AxisManager with subscan information.
+    turnarounds_in_subscan : bool
+        (Optional). Turnarounds are included as part of a subscan.
 
     Returns
     -------
@@ -299,6 +303,10 @@ def get_turnaround_flags(aman, az=None, method='scanspeed', name='turnarounds',
             aman.flags[name] = ta_flag
         else:
             aman.flags.wrap(name, ta_flag)   
+
+    if merge_subscans:
+        get_subscans(aman, merge=True, include_turnarounds=turnarounds_in_subscan)
+
     if method == 'az':
         ta_exp = RangesMatrix([ta_flag for i in range(aman.dets.count)])
         return ta_exp
@@ -685,13 +693,22 @@ def get_inv_var_flags(aman, signal_name='signal', nsigma=5,
 
     return mskinvar
 
-def get_subscans(aman, merge=True):
+def get_subscans(aman, merge=True, include_turnarounds=False):
     """
     Returns an axis manager with information about subscans.
     This includes direction, start time, stop time, and a ranges matrix (subscans samps)
     True inside each subscan. Subscans are defined excluding turnarounds.
     """
-    ss_ind = (~aman.flags.turnarounds).ranges() # sliceable indices (first inclusive, last exclusive) for subscans
+    if not include_turnarounds:
+        ss_ind = (~aman.flags.turnarounds).ranges() # sliceable indices (first inclusive, last exclusive) for subscans
+    else:
+        left = aman.flags.left_scan.ranges()
+        right = aman.flags.right_scan.ranges()
+        start_left = 0 if (left[0,0] < right[0,0]) else 1
+        ss_ind = np.empty((left.shape[0] + right.shape[0], 2), dtype=left.dtype)
+        ss_ind[start_left::2] = left
+        ss_ind[(start_left-1)%2::2] = right
+
     start_inds, end_inds = ss_ind.T
     n_subscan = ss_ind.shape[0]
     tt = aman.timestamps
@@ -710,7 +727,7 @@ def get_subscans(aman, merge=True):
         aman.wrap('subscans', subscan_aman)
     return subscan_aman
 
-def get_subscan_signal(aman, arr, isub=None):
+def get_subscan_signal(aman, arr, isub=None, trim=False):
     """
     Split an array into subscans.
 
@@ -724,12 +741,21 @@ def get_subscan_signal(aman, arr, isub=None):
         (Optional). Index of the desired subscan. May also be a list of indices.
         If None, all are used.
     """
+    if isinstance(arr, str):
+        arr = aman[arr]
     if np.isscalar(isub):
-        return apply_rng(arr, aman.subscans.subscan_flags[isub])
+        out = apply_rng(arr, aman.subscans.subscan_flags[isub])
+        if trim and out.size == 0:
+            out = None
     else:
         if isub is None:
             isub = range(len(aman.subscans.subscan_flags))
-        return [apply_rng(arr, aman.subscans.subscan_flags[ii]) for ii in isub]
+        out = [apply_rng(arr, aman.subscans.subscan_flags[ii]) for ii in isub]
+        if trim:
+            out = [x for x in out if x.size > 0]
+
+    return out
+
 
 def apply_rng(arr, rng):
     """
@@ -742,7 +768,10 @@ def apply_rng(arr, rng):
     rng : Ranges
         Ranges object of len (samps) selecting the desired range
     """
-    slc = slice(*rng.ranges()[0])
+    if rng.ranges().size == 0:
+        slc = slice(0,0) #  Return an empty array if rng is empty
+    else:
+        slc = slice(*np.squeeze(rng.ranges()))
     isamps = np.where(np.array(arr.shape) == rng.count)[0][0]
     ndslice = tuple((slice(None) if ii != isamps else slc for ii in range(arr.ndim)))
     return arr[ndslice]
@@ -788,18 +817,23 @@ def get_stats(aman, signal, stat_names, split_subscans=False, mask=None, name="s
     fn_dict = {'mean': np.mean, 'median': np.median, 'ptp': np.ptp, 'std': np.std,
                      'kurtosis': stats.kurtosis, 'skew': stats.skew}
 
+    if isinstance(signal, str):
+        signal = aman[signal]
     if split_subscans:
         if mask is not None:
             raise ValueError("Cannot mask samples and split subscans")
         stats_arr = []
         for iss in range(aman.subscans.subscans.count):
             data = get_subscan_signal(aman, signal, iss)
-            stats_arr.append([fn_dict[name](data, axis=1) for name in stat_names]) # Samps axis assumed to be 1
+            if data.size > 0:
+                stats_arr.append([fn_dict[name](data, axis=1) for name in stat_names]) # Samps axis assumed to be 1
+            else:
+                stats_arr.append(np.full((len(stat_names), signal.shape[0]), np.nan)) # Add nans if subscan has been entirely cut
         stats_arr = np.array(stats_arr).transpose(1, 2, 0) # stat, dets, subscan
     else:
         if mask is None:
             mask = slice(None)
-        stats_arr = np.array([fn_dict[name](data[:, mask], axis=1) for name in stat_names]) # Samps axis assumed to be 1
+        stats_arr = np.array([fn_dict[name](signal[:, mask], axis=1) for name in stat_names]) # Samps axis assumed to be 1
 
     info_aman = wrap_info(aman, name, stats_arr, stat_names, merge)
     return info_aman

From 2e9b247898932753f8a592280ae8dc3364df99b6 Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Thu, 21 Nov 2024 07:31:13 -0800
Subject: [PATCH 3/9] subscan preprocess fixes

---
 sotodlib/preprocess/pcore.py     | 15 +++++--
 sotodlib/preprocess/processes.py | 77 ++++++++++++++++++++------------
 sotodlib/tod_ops/fft_ops.py      | 12 ++---
 sotodlib/tod_ops/flags.py        | 17 ++++---
 4 files changed, 76 insertions(+), 45 deletions(-)

diff --git a/sotodlib/preprocess/pcore.py b/sotodlib/preprocess/pcore.py
index e6bf5a42e..e0df0ed6a 100644
--- a/sotodlib/preprocess/pcore.py
+++ b/sotodlib/preprocess/pcore.py
@@ -270,7 +270,7 @@ def _expand(new, full, wrap_valid=True):
                 continue
             out.wrap_new( k, new._assignments[k], cls=_zeros_cls(v))
             oidx=[]; nidx=[]
-            for a in new._assignments[k]:
+            for ii, a in enumerate(new._assignments[k]):
                 if a == 'dets':
                     oidx.append(fs_dets)
                     nidx.append(ns_dets)
@@ -278,8 +278,17 @@ def _expand(new, full, wrap_valid=True):
                     oidx.append(fs_samps)
                     nidx.append(ns_samps)
                 else:
-                    oidx.append(slice(None))
-                    nidx.append(slice(None))
+                    if (ii == 0) and isinstance(out[k], RangesMatrix): # Treat like dets
+                        if a in full._axes:
+                            _, fs, ns = full[a].intersection(new[a], return_slices=True)
+                        else:
+                            fs = range(new[a].count)
+                            ns = range(new[a].count)
+                        oidx.append(fs)
+                        nidx.append(ns)
+                    else: # Treat like samps
+                        oidx.append(slice(None))
+                        nidx.append(slice(None))
             oidx = tuple(oidx)
             nidx = tuple(nidx)
             if isinstance(out[k], RangesMatrix):
diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index 6da06fd9d..f7b7a7713 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -372,19 +372,22 @@ def __init__(self, step_cfgs):
         
 
     def process(self, aman, proc_aman):
-        if not self.subscan:
-            calc_psd = tod_ops.fft_ops.calc_psd
-        else:
+        if self.subscan:
             calc_psd = tod_ops.fft_ops.calc_psd_subscan
+        else:
+            calc_psd = tod_ops.fft_ops.calc_psd
 
         freqs, Pxx = calc_psd(aman, signal=aman[self.signal],
                               **self.process_cfgs)
-        fft_aman = core.AxisManager(
-            aman.dets, 
-            core.OffsetAxis("nusamps",len(freqs))
-        )
+        axis_list = [aman.dets, core.OffsetAxis("nusamps", len(freqs))]
+        pxx_axis_map = [(0, "dets"), (1, "nusamps")]
+        if self.subscan:
+            axis_list.append(aman.subscan_info.subscans)
+            pxx_axis_map.append((2, "subscans"))
+
+        fft_aman = core.AxisManager(*axis_list)
         fft_aman.wrap("freqs", freqs, [(0,"nusamps")])
-        fft_aman.wrap("Pxx", Pxx, [(0,"dets"), (1,"nusamps")])
+        fft_aman.wrap("Pxx", Pxx, pxx_axis_map)
         aman.wrap(self.wrap, fft_aman)
 
     def calc_and_save(self, aman, proc_aman):
@@ -400,17 +403,17 @@ class TODStats(_Preprocess):
 
     Example config block:
 
-      - "name : "stats"
-        "signal: "signal" # optional
-        "wrap": "stats" # optional
-        "calc":
-          "stat_names": ["median", "std"]
-          "split_subscans": False # optional
-          "mask": # optional, for cutting a power spectrum in frequency
-            "psd_aman": "psd"
-            "low_f": 1
-            "high_f": 10
-        "save": True
+      - name : "stats"
+        signal: "signal" # optional
+        wrap: "stats" # optional
+        calc:
+          stat_names: ["median", "std"]
+          split_subscans: False # optional
+          mask: # optional, for cutting a power spectrum in frequency
+            freqs: "psd.freqs"
+            low_f: 1
+            high_f: 10
+        save: True
 
     """
     name = "stats"
@@ -421,14 +424,18 @@ def __init__(self, step_cfgs):
         super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
+        def get_sub(aman, name): # Helper fn to access nested aman entries eg aman.psd.Pxx
+            cmd = "aman"+"".join([str([xx]) for xx in name.split(".")])
+            return eval(cmd)
+
         if self.calc_cfgs.get('mask') is not None:
             mask_dict = self.calc_cfgs.get('mask')
-            freqs = aman[mask_dict['psd_aman']]['freqs']
+            freqs = get_sub(aman, mask_dict['freqs'])
             low_f, high_f = mask_dict['low_f'], mask_dict['high_f']
             fmask = np.all([freqs >= low_f, freqs <= high_f], axis=0)
             self.calc_cfgs['mask'] = fmask
 
-        stats_aman = tod_ops.flags.get_stats(aman, aman[self.signal], **self.calc_cfgs)
+        stats_aman = tod_ops.flags.get_stats(aman, get_sub(aman, self.signal), **self.calc_cfgs)
         self.save(proc_aman, stats_aman)
 
     def save(self, proc_aman, stats_aman):
@@ -464,6 +471,7 @@ class Noise(_Preprocess):
     def __init__(self, step_cfgs):
         self.psd = step_cfgs.get('psd', 'psd')
         self.fit = step_cfgs.get('fit', False)
+        self.subscan = step_cfgs.get('subscan', False)
 
         super().__init__(step_cfgs)
 
@@ -476,7 +484,11 @@ def calc_and_save(self, aman, proc_aman):
             self.calc_cfgs = {}
         
         if self.fit:
-            calc_aman = tod_ops.fft_ops.fit_noise_model(aman, pxx=psd.Pxx, 
+            if not self.subscan:
+                fit_noise_model = tod_ops.fft_ops.fit_noise_model
+            else:
+                fit_noise_model = tod_ops.fft_ops.fit_noise_model_subscan
+            calc_aman = fit_noise_model(aman, pxx=psd.Pxx,
                                                         f=psd.freqs, 
                                                         merge_fit=True,
                                                         **self.calc_cfgs)
@@ -484,8 +496,12 @@ def calc_and_save(self, aman, proc_aman):
             wn = tod_ops.fft_ops.calc_wn(aman, pxx=psd.Pxx,
                                          freqs=psd.freqs,
                                          **self.calc_cfgs)
-            calc_aman = core.AxisManager(aman.dets)
-            calc_aman.wrap("white_noise", wn, [(0,"dets")])
+            if not self.subscan:
+                calc_aman = core.AxisManager(aman.dets)
+                calc_aman.wrap("white_noise", wn, [(0,"dets")])
+            else:
+                calc_aman = core.AxisManager(aman.dets, aman.subscan_info.subscans)
+                calc_aman.wrap("white_noise", wn, [(0,"dets"), (1,"subscans")])
 
         self.save(proc_aman, calc_aman)
     
@@ -513,10 +529,12 @@ def select(self, meta, proc_aman=None):
         self.select_cfgs['name'] = self.select_cfgs.get('name','noise')
 
         if self.fit:
-            keep = proc_aman[self.select_cfgs['name']].fit[:,1] <= self.select_cfgs["max_noise"]
+            wn = proc_aman[self.select_cfgs['name']].fit[:,1]
         else:
-            keep = proc_aman[self.select_cfgs['name']].white_noise <= self.select_cfgs["max_noise"]
-
+            wn = proc_aman[self.select_cfgs['name']].white_noise
+        if self.subscan:
+            wn = np.nanmean(wn, axis=-1) # Mean over subscans
+        keep = wn <= np.float64(self.select_cfgs["max_noise"])
         meta.restrict("dets", meta.dets.vals[keep])
         return meta
     
@@ -843,7 +861,7 @@ def calc_and_save(self, aman, proc_aman):
             calc_aman.wrap('turnarounds', ta, [(0, 'dets'), (1, 'samps')])
 
         if ('merge_subscans' not in self.calc_cfgs) or (self.calc_cfgs['merge_subscans']):
-            subscan_aman = aman.subscans
+            subscan_aman = aman.subscan_info
         else:
             subscan_aman = None
 
@@ -855,7 +873,7 @@ def save(self, proc_aman, turn_aman, subscan_aman):
         if self.save_cfgs:
             proc_aman.wrap("turnaround_flags", turn_aman)
             if subscan_aman is not None:
-                proc_aman.wrap("subscans", subscan_aman)
+                proc_aman.wrap("subscan_info", subscan_aman)
 
     def process(self, aman, proc_aman):
         tod_ops.flags.get_turnaround_flags(aman, **self.process_cfgs)
@@ -1385,3 +1403,4 @@ def process(self, aman, proc_aman):
 _Preprocess.register(DarkDets)
 _Preprocess.register(SourceFlags)
 _Preprocess.register(HWPAngleModel)
+_Preprocess.register(TODStats)
diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py
index 4d4b46be4..947465663 100644
--- a/sotodlib/tod_ops/fft_ops.py
+++ b/sotodlib/tod_ops/fft_ops.py
@@ -296,13 +296,13 @@ def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=Non
         if freq_spacing is not None:
             nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing))))
         else:
-            duration_samps = np.sum(aman.subscans.subscan_flags.mask(), axis=1)
+            duration_samps = np.sum(aman.subscan_info.subscan_flags.mask(), axis=1)
             duration_samps = duration_samps[duration_samps > 0]
             nperseg = int(2 ** (np.around(np.log2(np.median(duration_samps) / 4))))
         kwargs["nperseg"] = nperseg
 
     Pxx = []
-    for iss in range(aman.subscans.subscans.count):
+    for iss in range(aman.subscan_info.subscans.count):
         signal_ss = get_subscan_signal(aman, signal, iss)
         axis = -1 if "axis" not in kwargs else kwargs["axis"]
         if signal_ss.shape[axis] >= kwargs["nperseg"]:
@@ -526,8 +526,8 @@ def fit_noise_model_subscan(
     Fits noise model with white and 1/f noise to the PSD of signal subscans.
     Args are as for fit_noise_model.
     """
-    fitout = np.empty((aman.dets.count, 3, aman.subscans.subscans.count))
-    covout = np.empty((aman.dets.count, 3, 3, aman.subscans.subscans.count))
+    fitout = np.empty((aman.dets.count, 3, aman.subscan_info.subscans.count))
+    covout = np.empty((aman.dets.count, 3, 3, aman.subscan_info.subscans.count))
 
     if signal is None:
         signal = aman.signal
@@ -541,7 +541,7 @@ def fit_noise_model_subscan(
             **psdargs,
             )
 
-    for isub in range(aman.subscans.subscans.count):
+    for isub in range(aman.subscan_info.subscans.count):
         if np.all(np.isnan(pxx[...,isub])): # Subscan has been fully cut
             fitout[..., isub] = np.full((aman.dets.count, 3), np.nan)
             covout[..., isub] = np.full((aman.dets.count, 3, 3), np.nan)
@@ -555,7 +555,7 @@ def fit_noise_model_subscan(
     noise_fit_stats = core.AxisManager(
         aman.dets,
         noise_model.noise_model_coeffs,
-        aman.subscans.subscans
+        aman.subscan_info.subscans
         )
     noise_fit_stats.wrap("fit", fitout, [(0, "dets"), (1, "noise_model_coeffs"), (2, "subscans")])
     noise_fit_stats.wrap(
diff --git a/sotodlib/tod_ops/flags.py b/sotodlib/tod_ops/flags.py
index a0511978f..537244c4c 100644
--- a/sotodlib/tod_ops/flags.py
+++ b/sotodlib/tod_ops/flags.py
@@ -693,7 +693,7 @@ def get_inv_var_flags(aman, signal_name='signal', nsigma=5,
 
     return mskinvar
 
-def get_subscans(aman, merge=True, include_turnarounds=False):
+def get_subscans(aman, merge=True, include_turnarounds=False, overwrite=True):
     """
     Returns an axis manager with information about subscans.
     This includes direction, start time, stop time, and a ranges matrix (subscans samps)
@@ -724,7 +724,10 @@ def get_subscans(aman, merge=True, include_turnarounds=False):
     rm = RangesMatrix([Ranges.from_array(np.atleast_2d(ss), tt.size) for ss in ss_ind])
     subscan_aman.wrap('subscan_flags', rm, [(0, 'subscans'), (1, 'samps')]) # True in the subscan
     if merge:
-        aman.wrap('subscans', subscan_aman)
+        name = 'subscan_info'
+        if overwrite and name in aman:
+            aman.move(name, None)
+        aman.wrap(name, subscan_aman)
     return subscan_aman
 
 def get_subscan_signal(aman, arr, isub=None, trim=False):
@@ -744,13 +747,13 @@ def get_subscan_signal(aman, arr, isub=None, trim=False):
     if isinstance(arr, str):
         arr = aman[arr]
     if np.isscalar(isub):
-        out = apply_rng(arr, aman.subscans.subscan_flags[isub])
+        out = apply_rng(arr, aman.subscan_info.subscan_flags[isub])
         if trim and out.size == 0:
             out = None
     else:
         if isub is None:
-            isub = range(len(aman.subscans.subscan_flags))
-        out = [apply_rng(arr, aman.subscans.subscan_flags[ii]) for ii in isub]
+            isub = range(len(aman.subscan_info.subscan_flags))
+        out = [apply_rng(arr, aman.subscan_info.subscan_flags[ii]) for ii in isub]
         if trim:
             out = [x for x in out if x.size > 0]
 
@@ -782,7 +785,7 @@ def wrap_info(aman, info_aman_name, info, info_names, merge=True):
         info_aman = core.AxisManager(aman.dets)
         axmap = [(0, 'dets')]
     elif info[0].ndim == 2:
-        info_aman = core.AxisManager(aman.dets, aman.subscans.subscans)
+        info_aman = core.AxisManager(aman.dets, aman.subscan_info.subscans)
         axmap = [(0, 'dets'), (1, 'subscans')]
 
     for ii in range(len(info_names)):
@@ -823,7 +826,7 @@ def get_stats(aman, signal, stat_names, split_subscans=False, mask=None, name="s
         if mask is not None:
             raise ValueError("Cannot mask samples and split subscans")
         stats_arr = []
-        for iss in range(aman.subscans.subscans.count):
+        for iss in range(aman.subscan_info.subscans.count):
             data = get_subscan_signal(aman, signal, iss)
             if data.size > 0:
                 stats_arr.append([fn_dict[name](data, axis=1) for name in stat_names]) # Samps axis assumed to be 1

From 84158e2eeb805671e2e7546da375e9f079467f06 Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Thu, 21 Nov 2024 10:58:53 -0800
Subject: [PATCH 4/9] Replace custom fn with attrgetter

---
 sotodlib/preprocess/processes.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index f7b7a7713..68a30aec7 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -424,18 +424,16 @@ def __init__(self, step_cfgs):
         super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
-        def get_sub(aman, name): # Helper fn to access nested aman entries eg aman.psd.Pxx
-            cmd = "aman"+"".join([str([xx]) for xx in name.split(".")])
-            return eval(cmd)
-
         if self.calc_cfgs.get('mask') is not None:
             mask_dict = self.calc_cfgs.get('mask')
-            freqs = get_sub(aman, mask_dict['freqs'])
+            _f = attrgetter(mask_dict['freqs'])
+            freqs = _f(aman)
             low_f, high_f = mask_dict['low_f'], mask_dict['high_f']
             fmask = np.all([freqs >= low_f, freqs <= high_f], axis=0)
             self.calc_cfgs['mask'] = fmask
 
-        stats_aman = tod_ops.flags.get_stats(aman, get_sub(aman, self.signal), **self.calc_cfgs)
+        _f = attrgetter(self.signal)
+        stats_aman = tod_ops.flags.get_stats(aman, _f(aman), **self.calc_cfgs)
         self.save(proc_aman, stats_aman)
 
     def save(self, proc_aman, stats_aman):

From 217279392a7d08d876427cff0ccd81056663d36e Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Thu, 28 Nov 2024 14:11:43 +0000
Subject: [PATCH 5/9] Address Max's comments

---
 sotodlib/preprocess/pcore.py     |   2 +
 sotodlib/preprocess/processes.py |  28 ++---
 sotodlib/tod_ops/fft_ops.py      | 195 +++++++++++++------------------
 sotodlib/tod_ops/flags.py        | 103 ++++++++++++----
 4 files changed, 176 insertions(+), 152 deletions(-)

diff --git a/sotodlib/preprocess/pcore.py b/sotodlib/preprocess/pcore.py
index e0df0ed6a..b01fb50e4 100644
--- a/sotodlib/preprocess/pcore.py
+++ b/sotodlib/preprocess/pcore.py
@@ -279,6 +279,8 @@ def _expand(new, full, wrap_valid=True):
                     nidx.append(ns_samps)
                 else:
                     if (ii == 0) and isinstance(out[k], RangesMatrix): # Treat like dets
+                        # _ranges_matrix_match expects oidx[0] and nidx[0] to be list(inds), not slice.
+                        # Unknown axes treated as dets if first entry, else like samps. Added to support (subscans, samps) RangesMatrix.
                         if a in full._axes:
                             _, fs, ns = full[a].intersection(new[a], return_slices=True)
                         else:
diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index 2c5f190bf..ec9b513f3 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -388,34 +388,34 @@ def save(self, proc_aman, fft_aman):
             proc_aman.wrap(self.wrap, fft_aman)
 
 
-class TODStats(_Preprocess):
+class GetStats(_Preprocess):
     """ Get basic statistics from a TOD or its power spectrum.
 
     Example config block:
 
-      - name : "stats"
+      - name : "get_stats"
         signal: "signal" # optional
-        wrap: "stats" # optional
+        wrap: "tod_stats" # optional
         calc:
           stat_names: ["median", "std"]
           split_subscans: False # optional
-          mask: # optional, for cutting a power spectrum in frequency
+          psd_mask: # optional, for cutting a power spectrum in frequency
             freqs: "psd.freqs"
             low_f: 1
             high_f: 10
         save: True
 
     """
-    name = "stats"
+    name = "tod_stats"
     def __init__(self, step_cfgs):
         self.signal = step_cfgs.get('signal', 'signal')
-        self.wrap = step_cfgs.get('wrap', 'stats')
+        self.wrap = step_cfgs.get('wrap', 'tod_stats')
 
         super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
-        if self.calc_cfgs.get('mask') is not None:
-            mask_dict = self.calc_cfgs.get('mask')
+        if self.calc_cfgs.get('psd_mask') is not None:
+            mask_dict = self.calc_cfgs.get('psd_mask')
             _f = attrgetter(mask_dict['freqs'])
             freqs = _f(aman)
             low_f, high_f = mask_dict['low_f'], mask_dict['high_f']
@@ -844,19 +844,15 @@ def calc_and_save(self, aman, proc_aman):
             calc_aman.wrap('turnarounds', ta, [(0, 'dets'), (1, 'samps')])
 
         if ('merge_subscans' not in self.calc_cfgs) or (self.calc_cfgs['merge_subscans']):
-            subscan_aman = aman.subscan_info
-        else:
-            subscan_aman = None
+            calc_aman.wrap('subscan_info', aman.subscan_info)
 
-        self.save(proc_aman, calc_aman, subscan_aman)
+        self.save(proc_aman, calc_aman)
 
-    def save(self, proc_aman, turn_aman, subscan_aman):
+    def save(self, proc_aman, turn_aman):
         if self.save_cfgs is None:
             return
         if self.save_cfgs:
             proc_aman.wrap("turnaround_flags", turn_aman)
-            if subscan_aman is not None:
-                proc_aman.wrap("subscan_info", subscan_aman)
 
     def process(self, aman, proc_aman):
         tod_ops.flags.get_turnaround_flags(aman, **self.process_cfgs)
@@ -1448,4 +1444,4 @@ def save(self, proc_aman, split_flg_aman):
 _Preprocess.register(DarkDets)
 _Preprocess.register(SourceFlags)
 _Preprocess.register(HWPAngleModel)
-_Preprocess.register(TODStats)
+_Preprocess.register(GetStats)
diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py
index 947465663..06b1fd308 100644
--- a/sotodlib/tod_ops/fft_ops.py
+++ b/sotodlib/tod_ops/fft_ops.py
@@ -212,7 +212,8 @@ def calc_psd(
     prefer='center',
     freq_spacing=None,
     merge=False, 
-    overwrite=True, 
+    overwrite=True,
+    subscan=False,
     **kwargs
 ):
     """Calculates the power spectrum density of an input signal using signal.welch().
@@ -233,6 +234,7 @@ def calc_psd(
             If an nperseg is explicitly passed then that will be used.
         merge (bool): if True merge results into axismanager.
         overwrite (bool): if true will overwrite f, pxx axes.
+        subscan (bool): if True, compute psd on subscans.
         **kwargs: keyword args to be passed to signal.welch().
 
     Returns:
@@ -241,34 +243,40 @@ def calc_psd(
     """
     if signal is None:
         signal = aman.signal
-    if timestamps is None:
-        timestamps = aman.timestamps
-
-    n_samps = signal.shape[-1]
-    if n_samps <= max_samples:
-        start = 0
-        stop = n_samps
+    if subscan:
+        freqs, Pxx = _calc_psd_subscan(aman, signal=signal, freq_spacing=freq_spacing, **kwargs)
+        axis_map_pxx = [(0, "dets"), (1, "nusamps"), (2, "subscans")]
     else:
-        offset = n_samps - max_samples
-        if prefer == "left":
-            offset = 0
-        elif prefer == "center":
-            offset //= 2
-        elif prefer == "right":
-            pass
-        else:
-            raise ValueError(f"Invalid choise prefer='{prefer}'")
-        start = offset
-        stop = offset + max_samples
-    fs = 1 / np.nanmedian(np.diff(timestamps[start:stop]))
-    if "nperseg" not in kwargs:
-        if freq_spacing is not None:
-            nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing))))
+        if timestamps is None:
+            timestamps = aman.timestamps
+
+        n_samps = signal.shape[-1]
+        if n_samps <= max_samples:
+            start = 0
+            stop = n_samps
         else:
-            nperseg = int(2 ** (np.around(np.log2((stop - start) / 50.0))))
-        kwargs["nperseg"] = nperseg
+            offset = n_samps - max_samples
+            if prefer == "left":
+                offset = 0
+            elif prefer == "center":
+                offset //= 2
+            elif prefer == "right":
+                pass
+            else:
+                raise ValueError(f"Invalid choice prefer='{prefer}'")
+            start = offset
+            stop = offset + max_samples
+        fs = 1 / np.nanmedian(np.diff(timestamps[start:stop]))
+        if "nperseg" not in kwargs:
+            if freq_spacing is not None:
+                nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing))))
+            else:
+                nperseg = int(2 ** (np.around(np.log2((stop - start) / 50.0))))
+            kwargs["nperseg"] = nperseg
+
+        freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs)
+        axis_map_pxx = [(0, aman.dets), (1, "nusamps")]
 
-    freqs, Pxx = welch(signal[:, start:stop], fs, **kwargs)
     if merge:
         aman.merge( core.AxisManager(core.OffsetAxis("nusamps", len(freqs))))
         if overwrite:
@@ -277,14 +285,14 @@ def calc_psd(
             if "Pxx" in aman._fields:
                 aman.move("Pxx", None)
         aman.wrap("freqs", freqs, [(0,"nusamps")])
-        aman.wrap("Pxx", Pxx, [(0,"dets"),(1,"nusamps")])
+        aman.wrap("Pxx", Pxx, axis_map_pxx)
     return freqs, Pxx
 
-def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=None, **kwargs):
+def _calc_psd_subscan(aman, signal=None, freq_spacing=None, **kwargs):
     """
     Calculate the power spectrum density of subscans using signal.welch().
     Data defaults to aman.signal. aman.timestamps is used for times.
-    aman.subscans is used to identify subscans.
+    aman.subscan_info is used to identify subscans.
     See calc_psd for arguments.
     """
     from .flags import get_subscan_signal
@@ -296,7 +304,7 @@ def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=Non
         if freq_spacing is not None:
             nperseg = int(2 ** (np.around(np.log2(fs / freq_spacing))))
         else:
-            duration_samps = np.sum(aman.subscan_info.subscan_flags.mask(), axis=1)
+            duration_samps = np.asarray([np.ptp(x.ranges()) if x.ranges().size > 0 else 0 for x in aman.subscan_info.subscan_flags])
             duration_samps = duration_samps[duration_samps > 0]
             nperseg = int(2 ** (np.around(np.log2(np.median(duration_samps) / 4))))
         kwargs["nperseg"] = nperseg
@@ -312,14 +320,9 @@ def calc_psd_subscan(aman, signal=None, freq_spacing=None, merge=False, wrap=Non
             Pxx.append(np.full((signal.shape[0], kwargs["nperseg"]//2+1), np.nan)) # Add nans if subscan is too short
     Pxx = np.array(Pxx)
     Pxx = Pxx.transpose(1, 2, 0) # Dets, nusamps, subscans
-    if merge:
-        wrap = "Pxx_ss" if wrap is None else wrap
-        aman.merge( core.AxisManager(core.OffsetAxis("nusamps_ss", len(freqs))))
-        aman.wrap("freqs_ss", freqs, [(0,"nusamps_ss")])
-        aman.wrap(wrap, Pxx, [(0,aman.dets),(1,"nusamps_ss"), (2, "subscans")])
     return freqs, Pxx
 
-def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10):
+def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10, lowf_fk=1):
     """
     Function that calculates the white noise level as a median PSD value between
     two frequencies. Defaults to calculation of white noise between 5 and 10Hz.
@@ -359,6 +362,7 @@ def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10):
         wn2 = np.median(pxx[:, fmsk], axis=1)
 
     wn = np.sqrt(wn2)
+
     return wn
 
 
@@ -383,7 +387,7 @@ def fit_noise_model(
     signal=None,
     f=None,
     pxx=None,
-    psdargs=None,
+    psdargs={},
     fwhite=(10, 100),
     lowf=1,
     merge_fit=False,
@@ -391,7 +395,7 @@ def fit_noise_model(
     merge_name="noise_fit_stats",
     merge_psd=True,
     freq_spacing=None,
-    approx_fit=False,
+    subscan=False
 ):
     """
     Fits noise model with white and 1/f noise to the PSD of signal.
@@ -433,8 +437,6 @@ def fit_noise_model(
         If ``merg_psd`` is True then adds fres and Pxx to the axis manager.
     freq_spacing : float
         The approximate desired frequency spacing of the PSD. Passed to calc_psd.
-    approx_fit : bool
-        Get a rough fit instead of minimizing loglike.
     Returns
     -------
     noise_fit_stats : AxisManager
@@ -446,36 +448,35 @@ def fit_noise_model(
         signal = aman.signal
 
     if f is None or pxx is None:
-        if psdargs is None:
-            f, pxx = calc_psd(
-                aman, signal=signal, timestamps=aman.timestamps, freq_spacing=freq_spacing, merge=merge_psd
-            )
-        else:
-            f, pxx = calc_psd(
-                aman,
-                signal=signal,
-                timestamps=aman.timestamps,
-                freq_spacing=freq_spacing,
-                merge=merge_psd,
-                **psdargs,
-            )
-    eix = np.argmin(np.abs(f - f_max))
-    f = f[1:eix]
-    pxx = pxx[:, 1:eix]
-
-    fitout = np.zeros((aman.dets.count, 3))
-    # This is equal to np.sqrt(np.diag(cov)) when doing curve_fit
-    covout = np.zeros((aman.dets.count, 3, 3))
-    for i in range(aman.dets.count):
-        p = pxx[i]
-        wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])
-        pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1)
-        fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest))
-        p0 = [f[fidx], wnest, -pfit[0]]
-        if approx_fit:
-            covout[i] = np.full((3, 3), np.nan)
-            fitout[i] = p0
-        else:
+        f, pxx = calc_psd(
+            aman,
+            signal=signal,
+            timestamps=aman.timestamps,
+            freq_spacing=freq_spacing,
+            merge=merge_psd,
+            subscan=subscan,
+            **psdargs,
+        )
+    if subscan:
+       fitout, covout = _fit_noise_model_subscan(aman, signal,  f, pxx, psdargs=psdargs,
+                                                  fwhite=fwhite, lowf=lowf, f_max=f_max,
+                                                  freq_spacing=freq_spacing)
+       axis_map_fit = [(0, "dets"), (1, "noise_model_coeffs"), (2, aman.subscans)]
+       axis_map_cov = [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs"), (3, aman.subscans)]
+    else:
+        eix = np.argmin(np.abs(f - f_max))
+        f = f[1:eix]
+        pxx = pxx[:, 1:eix]
+
+        fitout = np.zeros((aman.dets.count, 3))
+        # This is equal to np.sqrt(np.diag(cov)) when doing curve_fit
+        covout = np.zeros((aman.dets.count, 3, 3))
+        for i in range(aman.dets.count):
+            p = pxx[i]
+            wnest = np.median(p[((f > fwhite[0]) & (f < fwhite[1]))])
+            pfit = np.polyfit(np.log10(f[f < lowf]), np.log10(p[f < lowf]), 1)
+            fidx = np.argmin(np.abs(10 ** np.polyval(pfit, np.log10(f)) - wnest))
+            p0 = [f[fidx], wnest, -pfit[0]]
             bounds = [(0, None), (sys.float_info.min, None), (None, None)]
             res = minimize(neglnlike, p0, args=(f, p), bounds=bounds, method="Nelder-Mead")
             try:
@@ -487,6 +488,8 @@ def fit_noise_model(
             except np.linalg.LinAlgError:
                 covout[i] = np.full((3, 3), np.nan)
             fitout[i] = res.x
+        axis_map_fit = [(0, "dets"), (1, "noise_model_coeffs")]
+        axis_map_cov = [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs")]
 
     noise_model_coeffs = ["fknee", "white_noise", "alpha"]
     noise_fit_stats = core.AxisManager(
@@ -495,32 +498,24 @@ def fit_noise_model(
             name="noise_model_coeffs", vals=np.array(noise_model_coeffs, dtype="<U8")
         ),
     )
-    noise_fit_stats.wrap("fit", fitout, [(0, "dets"), (1, "noise_model_coeffs")])
-    noise_fit_stats.wrap(
-        "cov",
-        covout,
-        [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs")],
-    )
+    noise_fit_stats.wrap("fit", fitout, axis_map_fit)
+    noise_fit_stats.wrap("cov", covout, axis_map_cov)
 
     if merge_fit:
         aman.wrap(merge_name, noise_fit_stats)
     return noise_fit_stats
 
 
-def fit_noise_model_subscan(
+def _fit_noise_model_subscan(
     aman,
-    signal=None,
-    f=None,
-    pxx=None,
+    signal,
+    f,
+    pxx,
     psdargs={},
     fwhite=(10, 100),
     lowf=1,
-    merge_fit=False,
     f_max=100,
-    merge_name="noise_fit_stats_ss",
-    merge_psd=True,
     freq_spacing=None,
-    approx_fit=False,
 ):
     """
     Fits noise model with white and 1/f noise to the PSD of signal subscans.
@@ -529,44 +524,18 @@ def fit_noise_model_subscan(
     fitout = np.empty((aman.dets.count, 3, aman.subscan_info.subscans.count))
     covout = np.empty((aman.dets.count, 3, 3, aman.subscan_info.subscans.count))
 
-    if signal is None:
-        signal = aman.signal
-
-    if f is None or pxx is None:
-        f, pxx = calc_psd_subscan(
-            aman,
-            signal=signal,
-            freq_spacing=freq_spacing,
-            merge=merge_psd,
-            **psdargs,
-            )
-
     for isub in range(aman.subscan_info.subscans.count):
         if np.all(np.isnan(pxx[...,isub])): # Subscan has been fully cut
             fitout[..., isub] = np.full((aman.dets.count, 3), np.nan)
             covout[..., isub] = np.full((aman.dets.count, 3, 3), np.nan)
         else:
-            noise_model = fit_noise_model(aman, f=f, pxx=pxx[...,isub], fwhite=fwhite, lowf=lowf,
-                                          merge_fit=False, f_max=f_max, merge_psd=False, approx_fit=approx_fit)
+            noise_model = fit_noise_model(aman, f=f, pxx=pxx[...,isub], fwhite=fwhite, lowf=lowf, merge_fit=False,
+                                          f_max=f_max, merge_psd=False, subscan=False)
 
             fitout[..., isub] = noise_model.fit
             covout[..., isub] = noise_model.cov
 
-    noise_fit_stats = core.AxisManager(
-        aman.dets,
-        noise_model.noise_model_coeffs,
-        aman.subscan_info.subscans
-        )
-    noise_fit_stats.wrap("fit", fitout, [(0, "dets"), (1, "noise_model_coeffs"), (2, "subscans")])
-    noise_fit_stats.wrap(
-        "cov",
-        covout,
-        [(0, "dets"), (1, "noise_model_coeffs"), (2, "noise_model_coeffs"), (3, "subscans")],
-    )
-
-    if merge_fit:
-        aman.wrap(merge_name, noise_fit_stats)
-    return noise_fit_stats
+    return fitout, covout
 
 
 def build_hpf_params_dict(
diff --git a/sotodlib/tod_ops/flags.py b/sotodlib/tod_ops/flags.py
index 537244c4c..b30cb8328 100644
--- a/sotodlib/tod_ops/flags.py
+++ b/sotodlib/tod_ops/flags.py
@@ -696,8 +696,26 @@ def get_inv_var_flags(aman, signal_name='signal', nsigma=5,
 def get_subscans(aman, merge=True, include_turnarounds=False, overwrite=True):
     """
     Returns an axis manager with information about subscans.
-    This includes direction, start time, stop time, and a ranges matrix (subscans samps)
-    True inside each subscan. Subscans are defined excluding turnarounds.
+    This includes direction and a ranges matrix (subscans samps)
+    True inside each subscan.
+
+    Parameters
+    ----------
+    aman : AxisManager
+        Input AxisManager.
+    merge : bool
+        Merge into aman as 'subscan_info'
+    include_turnarounds : bool
+        Include turnarounds in the subscan ranges
+    overwrite : bool
+        If true, write over subscan_info.
+
+    Returns
+    -------
+    subscan_aman : AxisManager
+        AxisManager containing information about the subscans.
+        "direction" is a (subscans,) array of strings 'left' or 'right'
+        "subscan_flags" is a (subscans, samps) RangesMatrix; True inside the subscan.
     """
     if not include_turnarounds:
         ss_ind = (~aman.flags.turnarounds).ranges() # sliceable indices (first inclusive, last exclusive) for subscans
@@ -717,10 +735,6 @@ def get_subscans(aman, merge=True, include_turnarounds=False, overwrite=True):
     is_left = aman.flags.left_scan.mask()[start_inds]
     subscan_aman.wrap('direction', np.array(['left' if is_left[ii] else 'right' for ii in range(n_subscan)]), [(0, 'subscans')])
 
-    subscan_aman.wrap('start_time', tt[start_inds], [(0, 'subscans')])
-    stop_time = aman.timestamps[ss_ind[:-1,1]]
-    last_time = [2*tt[-1] - tt[-2] if end_inds[-1] == tt.size else tt[end_inds[-1]]]
-    subscan_aman.wrap('stop_time', np.concatenate([stop_time, last_time]), [(0, 'subscans')])
     rm = RangesMatrix([Ranges.from_array(np.atleast_2d(ss), tt.size) for ss in ss_ind])
     subscan_aman.wrap('subscan_flags', rm, [(0, 'subscans'), (1, 'samps')]) # True in the subscan
     if merge:
@@ -741,8 +755,16 @@ def get_subscan_signal(aman, arr, isub=None, trim=False):
     arr : Array
         Input array.
     isub : int
-        (Optional). Index of the desired subscan. May also be a list of indices.
+        Index of the desired subscan. May also be a list of indices.
         If None, all are used.
+    trim : bool
+        Do not include size-zero arrays from empty subscans in the output.
+
+    Returns
+    -------
+    out : list
+        If isub is a scalar, return an Array of arr cut on the samps axis to the given subscan.
+        If isub is a list or None, return a list of such Arrays.
     """
     if isinstance(arr, str):
         arr = aman[arr]
@@ -762,29 +784,64 @@ def get_subscan_signal(aman, arr, isub=None, trim=False):
 
 def apply_rng(arr, rng):
     """
-    Helper function to get signals on subcsans.
+    Apply a Ranges object to an array. rng should be True on the samples you want to keep.
 
     Parameters
     ----------
     arr : Array
-        Array containing the signal. Should have one axis of len (samps)
+        Array containing the signal. Should have one axis of len (samps).
     rng : Ranges
-        Ranges object of len (samps) selecting the desired range
+        Ranges object of len (samps) selecting the desired range.
     """
     if rng.ranges().size == 0:
-        slc = slice(0,0) #  Return an empty array if rng is empty
+        slices = [slice(0,0)] #  Return an empty array if rng is empty
     else:
-        slc = slice(*np.squeeze(rng.ranges()))
-    isamps = np.where(np.array(arr.shape) == rng.count)[0][0]
-    ndslice = tuple((slice(None) if ii != isamps else slc for ii in range(arr.ndim)))
-    return arr[ndslice]
-
-def wrap_info(aman, info_aman_name, info, info_names, merge=True):
-    """Wrap info into an aman. Assumed to be (dets,) or (dets, subscans)"""
-    if info[0].ndim == 1:
+        slices = [slice(*irng) for irng in rng.ranges()]
+
+    # Identify the samps axis
+    isamps = np.where(np.array(arr.shape) == rng.count)[0]
+    if isamps.size != 1:
+        # Check for axis mismatch between arr and rng, or multiple axes with the same size
+        raise RuntimeError("Could not identify axis matching Ranges")
+    # Apply ranges
+    out = []
+    for slc in slices:
+        ndslice = tuple((slice(None) if ii != isamps[0] else slc for ii in range(arr.ndim)))
+        out.append(arr[ndslice])
+    return np.concatenate(out, axis=isamps[0])
+
+def wrap_stats(aman, info_aman_name, info, info_names, merge=True):
+    """
+    Wrap multiple stats into a new aman, checking for subscan information. Stats can be (dets,) or (dets, subscans).
+
+    Parameters
+    ----------
+    aman : AxisManager
+        Input AxisManager.
+    info_aman_name : str
+        Name for info_aman when wrapped into input.
+    info : Array
+        (stats, dets,) or (stats, dets, subscans) containing the information you want to wrap.
+    info_names : list
+        List of str names for each entry in the new aman.
+    merge : bool
+        If True merge info_aman into aman.
+
+    Returns
+    -------
+    info_aman : AxisManager
+        (dets,) or (dets, subscans) aman with a field for each item in info_names.
+    """
+    info_names = np.atleast_1d(info_names)
+    info = np.atleast_2d(info)
+    if info.shape == (len(info_names), aman.dets.count): # (stats, dets)
+        if len(info_names) == aman.dets.count and aman.dets.count == aman.subscan_info.subscans.count:
+            raise RuntimeError("Cannot infer axis mapping") # Catch corner case
         info_aman = core.AxisManager(aman.dets)
         axmap = [(0, 'dets')]
-    elif info[0].ndim == 2:
+
+    else:
+        info = np.atleast_3d(info) # (stats, dets, subscans)
         info_aman = core.AxisManager(aman.dets, aman.subscan_info.subscans)
         axmap = [(0, 'dets'), (1, 'subscans')]
 
@@ -810,9 +867,9 @@ def get_stats(aman, signal, stat_names, split_subscans=False, mask=None, name="s
     stat_names : list
         List of strings identifying which statistics to run.
     split_subscans : bool
-        (Optional). If True statistics will be computed on subscans. Assumes aman.subscans exists already.
+        If True statistics will be computed on subscans. Assumes aman.subscan_info exists already.
     mask : Array
-        (Optional). Mask to apply before computation. 1d array for advanced indexing, or a slice object.
+        Mask to apply before computation. 1d array for advanced indexing (keep True), or a slice object.
     name : str
         Name of axis manager to add to aman if merge is True.
     """
@@ -838,5 +895,5 @@ def get_stats(aman, signal, stat_names, split_subscans=False, mask=None, name="s
             mask = slice(None)
         stats_arr = np.array([fn_dict[name](signal[:, mask], axis=1) for name in stat_names]) # Samps axis assumed to be 1
 
-    info_aman = wrap_info(aman, name, stats_arr, stat_names, merge)
+    info_aman = wrap_stats(aman, name, stats_arr, stat_names, merge)
     return info_aman

From 1e9df8cb11dfb7641e0a5ea1c1b3dd7a197285b8 Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Thu, 28 Nov 2024 14:57:23 +0000
Subject: [PATCH 6/9] Populate psd.Pxx_ss and psd.Pxx together

---
 sotodlib/preprocess/processes.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index ec9b513f3..2db5b225c 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -371,13 +371,13 @@ def calc_and_save(self, aman, proc_aman):
         freqs, Pxx = tod_ops.fft_ops.calc_psd(aman, signal=aman[self.signal],
                                               **self.calc_cfgs)
 
-        axis_list = [aman.dets, core.OffsetAxis("nusamps", len(freqs))]
+        fft_aman = core.AxisManager(aman.dets,
+                                    core.OffsetAxis("nusamps", len(freqs)))
         pxx_axis_map = [(0, "dets"), (1, "nusamps")]
         if self.calc_cfgs.get('subscan', False):
-            axis_list.append(aman.subscan_info.subscans)
-            pxx_axis_map.append((2, "subscans"))
+            fft_aman.wrap("Pxx_ss", Pxx, pxx_axis_map+[(2, aman.subscans)])
+            Pxx = np.mean(Pxx, axis=-1) # Mean of subscans
 
-        fft_aman = core.AxisManager(*axis_list)
         fft_aman.wrap("freqs", freqs, [(0,"nusamps")])
         fft_aman.wrap("Pxx", Pxx, pxx_axis_map)
 

From 91d55c690763b66a355a8e97a56bafb5b0c6f868 Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Thu, 28 Nov 2024 07:15:25 -0800
Subject: [PATCH 7/9] Add tod and psd plotting

---
 sotodlib/preprocess/preprocess_plot.py | 34 ++++++++++++++++++++++++++
 sotodlib/preprocess/processes.py       | 26 ++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/sotodlib/preprocess/preprocess_plot.py b/sotodlib/preprocess/preprocess_plot.py
index 83f3325d0..542e283ec 100644
--- a/sotodlib/preprocess/preprocess_plot.py
+++ b/sotodlib/preprocess/preprocess_plot.py
@@ -390,6 +390,40 @@ def plot_trending_flags(aman, trend_aman, filename='./trending_flags.png'):
     os.makedirs(head_tail[0], exist_ok=True)
     plt.savefig(filename)
 
+def plot_signal(aman, signal_name="signal", x_name="timestamps", plot_ds_factor=50, plot_ds_factor_dets=None, xlim=None, alpha=0.2, yscale='linear', y_unit=None, filename="./signal.png"):
+    from operator import attrgetter
+    if plot_ds_factor_dets is None:
+        plot_ds_factor_dets = plot_ds_factor
+    yy = attrgetter(signal_name)(aman)[::plot_ds_factor_dets, 1::plot_ds_factor].copy() # (dets, samps); (dets, nusamps); (dets, nusamps, subscans)
+    xx = attrgetter(x_name)(aman)[1::plot_ds_factor].copy() # (samps); (nusamps)
+    if x_name == "timestamps":
+        xx -= xx[0]
+    if yy.ndim > 2: # Flatten subscan axis into dets
+        yy = yy.swapaxes(1,2).reshape(-1, yy.shape[1])
+
+    if xlim is not None:
+        xinds = np.logical_and(xx >= xlim[0], xx <= xlim[1])
+        xx = xx[xinds]
+        yy = yy[:,xinds]
+
+    fig, ax = plt.subplots(1, 1, figsize=(6.4, 4.8))
+    ax.plot(xx, yy.T, color='k', alpha=0.2)
+    ax.set_yscale(yscale)
+    if "freqs" in x_name:
+        ax.set_xlabel("freq [Hz]")
+    else:
+        ax.set_xlabel(f"{x_name} [s]")
+    y_unit = "" if y_unit is None else f" [{y_unit}]"
+    ax.set_ylabel(f"{signal_name.replace('.Pxx', '')}{y_unit}")
+    plt.suptitle(f"{aman.obs_info.obs_id}, dT = {np.ptp(aman.timestamps)/60:.1f} min")
+    plt.tight_layout()
+    head_tail = os.path.split(filename)
+    os.makedirs(head_tail[0], exist_ok=True)
+    plt.savefig(filename)
+
+def plot_psd(aman, signal_name="psd.Pxx", x_name="psd.freqs", plot_ds_factor=4, plot_ds_factor_dets=20, xlim=None, alpha=0.2, yscale='log', y_unit=None, filename="./psd.png"):
+    return plot_signal(aman, signal_name, x_name, plot_ds_factor, plot_ds_factor_dets, xlim, alpha, yscale, y_unit, filename)
+
 def plot_signal_diff(aman, flag_aman, flag_type="glitches", flag_threshold=10, plot_ds_factor=50, filename="./glitch_signal_diff.png"):
     """
     Function for plotting the difference in signal before and after cuts from either glitches or jumps.
diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index 2db5b225c..a788ffec5 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -386,6 +386,19 @@ def calc_and_save(self, aman, proc_aman):
     def save(self, proc_aman, fft_aman):
         if not(self.save_cfgs is None):
             proc_aman.wrap(self.wrap, fft_aman)
+    def plot(self, aman, proc_aman, filename):
+        if self.plot_cfgs is None:
+            return
+        if self.plot_cfgs:
+            from .preprocess_plot import plot_psd
+
+            filename = filename.replace('{ctime}', f'{str(aman.timestamps[0])[:5]}')
+            filename = filename.replace('{obsid}', aman.obs_info.obs_id)
+            det = aman.dets.vals[0]
+            ufm = det.split('_')[2]
+            filename = filename.replace('{name}', f'{ufm}_{self.wrap}')
+
+            plot_psd(aman, signal_name=f"{self.wrap}.Pxx", x_name=f"{self.wrap}.freqs", filename=filename, **self.plot_cfgs)
 
 
 class GetStats(_Preprocess):
@@ -430,6 +443,19 @@ def save(self, proc_aman, stats_aman):
         if not(self.save_cfgs is None):
             proc_aman.wrap(self.wrap, stats_aman)
 
+    def plot(self, aman, proc_aman, filename):
+        if self.plot_cfgs is None:
+            return
+        if self.plot_cfgs:
+            from .preprocess_plot import plot_signal
+
+            filename = filename.replace('{ctime}', f'{str(aman.timestamps[0])[:5]}')
+            filename = filename.replace('{obsid}', aman.obs_info.obs_id)
+            det = aman.dets.vals[0]
+            ufm = det.split('_')[2]
+            filename = filename.replace('{name}', f'{ufm}_{self.signal}')
+
+            plot_signal(aman, signal_name=self.signal, x_name="timestamps", filename=filename, **self.plot_cfgs)
 
 class Noise(_Preprocess):
     """Estimate the white noise levels in the data. Assumes the PSD has been

From 449866449ce8002fea68d2913ce4116f5670e541 Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Fri, 29 Nov 2024 03:31:41 -0800
Subject: [PATCH 8/9] subscan preproc fixes

---
 sotodlib/preprocess/pcore.py           |  2 ++
 sotodlib/preprocess/preprocess_plot.py | 14 ++++++---
 sotodlib/preprocess/processes.py       | 42 ++++++++++++++++++--------
 sotodlib/tod_ops/fft_ops.py            |  2 +-
 4 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/sotodlib/preprocess/pcore.py b/sotodlib/preprocess/pcore.py
index b01fb50e4..0f99f9037 100644
--- a/sotodlib/preprocess/pcore.py
+++ b/sotodlib/preprocess/pcore.py
@@ -5,6 +5,7 @@
 from .. import core
 from so3g.proj import Ranges, RangesMatrix
 from scipy.sparse import csr_array
+from matplotlib import pyplot as plt
 
 class _Preprocess(object):
     """The base class for Preprocessing modules which defines the required
@@ -467,6 +468,7 @@ def run(self, aman, proc_aman=None, select=True, sim=False, update_plot=False):
                 update_full_aman( proc_aman, full, self.wrap_valid)
             if update_plot:
                 process.plot(aman, proc_aman, filename=os.path.join(self.plot_dir, '{ctime}/{obsid}', f'{step+1}_{{name}}.png'))
+            plt.close()
             if select:
                 process.select(aman, proc_aman)
                 proc_aman.restrict('dets', aman.dets.vals)
diff --git a/sotodlib/preprocess/preprocess_plot.py b/sotodlib/preprocess/preprocess_plot.py
index 542e283ec..3985cfc81 100644
--- a/sotodlib/preprocess/preprocess_plot.py
+++ b/sotodlib/preprocess/preprocess_plot.py
@@ -390,12 +390,16 @@ def plot_trending_flags(aman, trend_aman, filename='./trending_flags.png'):
     os.makedirs(head_tail[0], exist_ok=True)
     plt.savefig(filename)
 
-def plot_signal(aman, signal_name="signal", x_name="timestamps", plot_ds_factor=50, plot_ds_factor_dets=None, xlim=None, alpha=0.2, yscale='linear', y_unit=None, filename="./signal.png"):
+def plot_signal(aman, signal=None, xx=None, signal_name="signal", x_name="timestamps", plot_ds_factor=50, plot_ds_factor_dets=None, xlim=None, alpha=0.2, yscale='linear', y_unit=None, filename="./signal.png"):
     from operator import attrgetter
     if plot_ds_factor_dets is None:
         plot_ds_factor_dets = plot_ds_factor
-    yy = attrgetter(signal_name)(aman)[::plot_ds_factor_dets, 1::plot_ds_factor].copy() # (dets, samps); (dets, nusamps); (dets, nusamps, subscans)
-    xx = attrgetter(x_name)(aman)[1::plot_ds_factor].copy() # (samps); (nusamps)
+    if signal is None:
+        signal = attrgetter(signal_name)(aman)
+    if xx is None:
+        xx = attrgetter(x_name)(aman)
+    yy = signal[::plot_ds_factor_dets, 1::plot_ds_factor].copy() # (dets, samps); (dets, nusamps); (dets, nusamps, subscans)
+    xx = xx[1::plot_ds_factor].copy() # (samps); (nusamps)
     if x_name == "timestamps":
         xx -= xx[0]
     if yy.ndim > 2: # Flatten subscan axis into dets
@@ -421,8 +425,8 @@ def plot_signal(aman, signal_name="signal", x_name="timestamps", plot_ds_factor=
     os.makedirs(head_tail[0], exist_ok=True)
     plt.savefig(filename)
 
-def plot_psd(aman, signal_name="psd.Pxx", x_name="psd.freqs", plot_ds_factor=4, plot_ds_factor_dets=20, xlim=None, alpha=0.2, yscale='log', y_unit=None, filename="./psd.png"):
-    return plot_signal(aman, signal_name, x_name, plot_ds_factor, plot_ds_factor_dets, xlim, alpha, yscale, y_unit, filename)
+def plot_psd(aman, signal=None, xx=None, signal_name="psd.Pxx", x_name="psd.freqs", plot_ds_factor=4, plot_ds_factor_dets=20, xlim=None, alpha=0.2, yscale='log', y_unit=None, filename="./psd.png"):
+    return plot_signal(aman, signal, xx, signal_name, x_name, plot_ds_factor, plot_ds_factor_dets, xlim, alpha, yscale, y_unit, filename)
 
 def plot_signal_diff(aman, flag_aman, flag_type="glitches", flag_threshold=10, plot_ds_factor=50, filename="./glitch_signal_diff.png"):
     """
diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index a788ffec5..5b77ab094 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -376,7 +376,7 @@ def calc_and_save(self, aman, proc_aman):
         pxx_axis_map = [(0, "dets"), (1, "nusamps")]
         if self.calc_cfgs.get('subscan', False):
             fft_aman.wrap("Pxx_ss", Pxx, pxx_axis_map+[(2, aman.subscans)])
-            Pxx = np.mean(Pxx, axis=-1) # Mean of subscans
+            Pxx = np.nanmean(Pxx, axis=-1) # Mean of subscans
 
         fft_aman.wrap("freqs", freqs, [(0,"nusamps")])
         fft_aman.wrap("Pxx", Pxx, pxx_axis_map)
@@ -398,7 +398,8 @@ def plot(self, aman, proc_aman, filename):
             ufm = det.split('_')[2]
             filename = filename.replace('{name}', f'{ufm}_{self.wrap}')
 
-            plot_psd(aman, signal_name=f"{self.wrap}.Pxx", x_name=f"{self.wrap}.freqs", filename=filename, **self.plot_cfgs)
+            plot_psd(aman, signal=attrgetter(f"{self.wrap}.Pxx")(proc_aman),
+                     xx=attrgetter(f"{self.wrap}.freqs")(proc_aman), filename=filename, **self.plot_cfgs)
 
 
 class GetStats(_Preprocess):
@@ -406,7 +407,7 @@ class GetStats(_Preprocess):
 
     Example config block:
 
-      - name : "get_stats"
+      - name : "tod_stats"
         signal: "signal" # optional
         wrap: "tod_stats" # optional
         calc:
@@ -430,13 +431,21 @@ def calc_and_save(self, aman, proc_aman):
         if self.calc_cfgs.get('psd_mask') is not None:
             mask_dict = self.calc_cfgs.get('psd_mask')
             _f = attrgetter(mask_dict['freqs'])
-            freqs = _f(aman)
+            try:
+                freqs = _f(aman)
+            except KeyError:
+                freqs = _f(proc_aman)
             low_f, high_f = mask_dict['low_f'], mask_dict['high_f']
             fmask = np.all([freqs >= low_f, freqs <= high_f], axis=0)
             self.calc_cfgs['mask'] = fmask
+            del self.calc_cfgs['psd_mask']
 
         _f = attrgetter(self.signal)
-        stats_aman = tod_ops.flags.get_stats(aman, _f(aman), **self.calc_cfgs)
+        try:
+            signal = _f(aman)
+        except KeyError:
+            signal = _f(proc_aman)
+        stats_aman = tod_ops.flags.get_stats(aman, signal, **self.calc_cfgs)
         self.save(proc_aman, stats_aman)
 
     def save(self, proc_aman, stats_aman):
@@ -485,6 +494,7 @@ class Noise(_Preprocess):
     def __init__(self, step_cfgs):
         self.psd = step_cfgs.get('psd', 'psd')
         self.fit = step_cfgs.get('fit', False)
+        self.subscan = step_cfgs.get('subscan', False)
 
         super().__init__(step_cfgs)
 
@@ -492,20 +502,23 @@ def calc_and_save(self, aman, proc_aman):
         if self.psd not in proc_aman:
             raise ValueError("PSD is not saved in Preprocessing AxisManager")
         psd = proc_aman[self.psd]
-        
+        pxx = psd.Pxx_ss if self.subscan else psd.Pxx
+
         if self.calc_cfgs is None:
             self.calc_cfgs = {}
-        
+
         if self.fit:
-            calc_aman = tod_ops.fft_ops.fit_noise_model(aman, pxx=psd.Pxx,
+            if self.calc_cfgs.get('subscan') is None:
+                self.calc_cfgs['subscan'] = self.subscan
+            calc_aman = tod_ops.fft_ops.fit_noise_model(aman, pxx=pxx,
                                                         f=psd.freqs, 
                                                         merge_fit=True,
                                                         **self.calc_cfgs)
         else:
-            wn = tod_ops.fft_ops.calc_wn(aman, pxx=psd.Pxx,
+            wn = tod_ops.fft_ops.calc_wn(aman, pxx=pxx,
                                          freqs=psd.freqs,
                                          **self.calc_cfgs)
-            if not self.calc_cfgs.get('subscan', False):
+            if not self.subscan:
                 calc_aman = core.AxisManager(aman.dets)
                 calc_aman.wrap("white_noise", wn, [(0,"dets")])
             else:
@@ -1169,9 +1182,9 @@ class PCARelCal(_Preprocess):
                 yfac: 1.5
                 calc_good_medianw: True
             lpf:
-                type: "low_pass_sine2"
+                type: "sine2"
                 cutoff: 1
-                width: 0.1
+                trans_width: 0.1
             trim_samps: 2000
         save: True
         plot:
@@ -1188,6 +1201,7 @@ def __init__(self, step_cfgs):
         super().__init__(step_cfgs)
 
     def calc_and_save(self, aman, proc_aman):
+        self.plot_signal = self.signal
         if self.calc_cfgs.get("lpf") is not None:
             filt = tod_ops.filters.get_lpf(self.calc_cfgs.get("lpf"))
             filt_tod = tod_ops.fourier_filter(aman, filt, signal_name='signal')
@@ -1203,6 +1217,8 @@ def calc_and_save(self, aman, proc_aman):
                                              proc_aman.samps.offset + proc_aman.samps.count - trim))
                 filt_aman.restrict('samps', (filt_aman.samps.offset + trim,
                                              filt_aman.samps.offset + filt_aman.samps.count - trim))
+            if self.plot_cfgs:
+                self.plot_signal = filt_aman[self.signal]
 
         bands = np.unique(aman.det_info.wafer.bandpass)
         bands = bands[bands != 'NC']
@@ -1270,7 +1286,7 @@ def plot(self, aman, proc_aman, filename):
             for band in bands:
                 pca_aman = aman.restrict('dets', aman.dets.vals[proc_aman[self.run_name][f'{band}_idx']], in_place=False)
                 band_aman = proc_aman[self.run_name].restrict('dets', aman.dets.vals[proc_aman[self.run_name][f'{band}_idx']], in_place=False)
-                plot_pcabounds(pca_aman, band_aman, filename=filename.replace('{name}', f'{ufm}_{band}_pca'), signal=self.signal, band=band, plot_ds_factor=self.plot_cfgs.get('plot_ds_factor', 20))
+                plot_pcabounds(pca_aman, band_aman, filename=filename.replace('{name}', f'{ufm}_{band}_pca'), signal=self.plot_signal, band=band, plot_ds_factor=self.plot_cfgs.get('plot_ds_factor', 20))
 
 
 class PTPFlags(_Preprocess):
diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py
index 06b1fd308..fb9b727af 100644
--- a/sotodlib/tod_ops/fft_ops.py
+++ b/sotodlib/tod_ops/fft_ops.py
@@ -322,7 +322,7 @@ def _calc_psd_subscan(aman, signal=None, freq_spacing=None, **kwargs):
     Pxx = Pxx.transpose(1, 2, 0) # Dets, nusamps, subscans
     return freqs, Pxx
 
-def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10, lowf_fk=1):
+def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10):
     """
     Function that calculates the white noise level as a median PSD value between
     two frequencies. Defaults to calculation of white noise between 5 and 10Hz.

From 26f8e8d8d1ef8f56af49dd38c6e9e0ef62ed14d2 Mon Sep 17 00:00:00 2001
From: Erik Rosenberg <errosenberg47@gmail.com>
Date: Mon, 2 Dec 2024 07:59:50 -0800
Subject: [PATCH 9/9] Small docs changes

---
 docs/preprocess.rst              | 1 +
 sotodlib/preprocess/processes.py | 3 +++
 sotodlib/tod_ops/fft_ops.py      | 4 +++-
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/docs/preprocess.rst b/docs/preprocess.rst
index 6ce35a95a..1845208f5 100644
--- a/docs/preprocess.rst
+++ b/docs/preprocess.rst
@@ -252,6 +252,7 @@ Flagging and Products
 .. autoclass:: sotodlib.preprocess.processes.FlagTurnarounds
 .. autoclass:: sotodlib.preprocess.processes.DarkDets
 .. autoclass:: sotodlib.preprocess.processes.SourceFlags
+.. autoclass:: sotodlib.preprocess.processes.GetStats
 
 HWP Related
 :::::::::::
diff --git a/sotodlib/preprocess/processes.py b/sotodlib/preprocess/processes.py
index 5b77ab094..fc284d9d2 100644
--- a/sotodlib/preprocess/processes.py
+++ b/sotodlib/preprocess/processes.py
@@ -355,6 +355,7 @@ class PSDCalc(_Preprocess):
           "psd_cfgs": # optional, kwargs to scipy.welch
             "nperseg": 1024
           "wrap_name": "psd" # optional
+          "subscan": False
         "save": True
 
     .. autofunction:: sotodlib.tod_ops.fft_ops.calc_psd
@@ -477,6 +478,8 @@ class Noise(_Preprocess):
     Example config block::
 
      - name: "noise"
+       fit: False
+       subscan: False
        calc:
          low_f: 5
          high_f: 10
diff --git a/sotodlib/tod_ops/fft_ops.py b/sotodlib/tod_ops/fft_ops.py
index fb9b727af..9b64f43fd 100644
--- a/sotodlib/tod_ops/fft_ops.py
+++ b/sotodlib/tod_ops/fft_ops.py
@@ -11,6 +11,7 @@
 
 from . import detrend_tod
 
+
 def _get_num_threads():
     # Guess how many threads we should be using in FFT ops...
     return so3g.useful_info().get("omp_num_threads", 4)
@@ -362,7 +363,6 @@ def calc_wn(aman, pxx=None, freqs=None, low_f=5, high_f=10):
         wn2 = np.median(pxx[:, fmsk], axis=1)
 
     wn = np.sqrt(wn2)
-
     return wn
 
 
@@ -437,6 +437,8 @@ def fit_noise_model(
         If ``merg_psd`` is True then adds fres and Pxx to the axis manager.
     freq_spacing : float
         The approximate desired frequency spacing of the PSD. Passed to calc_psd.
+    subscan : bool
+        If True, fit noise on subscans.
     Returns
     -------
     noise_fit_stats : AxisManager