Repository: simonsobs/sotodlib

Issue:
Jump Improvements
We want to address 3 main issues with the current jump finding:

1. Speed it up
2. Reduce false positives
3. Reduce false negatives

This is to track progress on this. Current PR open addressing this is PR #807

Discussion:
@skhrg: Sub task that is also related to #808 , but to reduce false negatives and positives we need to know what is a real jump somehow...

I suppose we manually look at the standard TODs for testing to check this, but I think we should also standardize a set of diagnostic plots or something to help with this. I have seen versions of these plots from myself and others, I remember thinking @msilvafe's were easier to tell at a glance if something wasn't working, but regardless a standard plotting function to glance at what it getting flagged/not flagged would be nice.@kwolz: Here is an example of the type of diagnostic plot I used to check the jump fixing algorithm. Left is the whole TOD, right is a zoom in to the jump-flagged region (yes, this requires the jump finder to have found a jump before). The function (tentatively documented) is [this one](https://github.com/simonsobs/pwg-scripts/blob/b1835a85ca19b39adbe785f28591be000cf65938/flp/sat-map-processing/kwolz/scripts/jump_validation.py#L170). ![Image](https://github.com/user-attachments/assets/4ca1a73d-3b88-49c3-8e69-f4ef5af9964d).
@msilvafe: Should be fixed with #1081 and testing from Kevin

Solution patch:
diff --git a/sotodlib/core/metadata/loader.py b/sotodlib/core/metadata/loader.py
index d0aa0967..a344f582 100644
--- a/sotodlib/core/metadata/loader.py
+++ b/sotodlib/core/metadata/loader.py
@@ -637,7 +637,9 @@ def merge_det_info(det_info, new_info, multi=True, on_missing='trim'):
                 f'in the det_id, maybe add {k} to the match_keys?')
 
     if len(det_info) != len(i1) and on_missing != 'trim':
-        raise IncompleteMetadataError('{len(det_info)} -> {len(i1)})')
+        raise IncompleteMetadataError(
+            f"det_info had n_dets={len(det_info)}; After merge it has n_dets={len(i1)}"
+        )
 
     logger.debug(f' ... updating det_info (row count '
                  f'{len(det_info)} -> {len(i1)})')
diff --git a/sotodlib/io/bookbinder.py b/sotodlib/io/bookbinder.py
index a1bdd1e6..c0c0a97e 100644
--- a/sotodlib/io/bookbinder.py
+++ b/sotodlib/io/bookbinder.py
@@ -635,6 +635,10 @@ class SmurfStreamProcessor:
         # Sample idx within each in-frame for every input sample
         in_offset_idxs = np.arange(len(self.times)) - offsets[self.frame_idxs]
 
+        # Mask to use to separate out non-resonator signal from resonator signal
+        tracked_chans = np.array(['NONE' not in r for r in self.readout_ids], dtype=bool)
+        readout_id_arr = np.array(self.readout_ids)
+
         # Handle file writers
         writer = None
         cur_file_idx = None
@@ -751,7 +755,8 @@ class SmurfStreamProcessor:
             oframe['flag_smurfgaps'] = core.G3VectorBool(~filled)
 
             ts = core.G3VectorTime(ts * core.G3Units.s)
-            oframe['signal'] = so3g.G3SuperTimestream(self.readout_ids, ts, data)
+            oframe['signal'] = so3g.G3SuperTimestream(readout_id_arr[tracked_chans], ts, data[tracked_chans, :])
+            oframe['untracked'] = so3g.G3SuperTimestream(readout_id_arr[~tracked_chans], ts, data[~tracked_chans, :])
             oframe['primary'] = so3g.G3SuperTimestream(self.primary_names, ts, primary)
             oframe['tes_biases'] = so3g.G3SuperTimestream(self.bias_names, ts, biases)
             oframe['stream_id'] = self.stream_id
diff --git a/sotodlib/io/hkdb.py b/sotodlib/io/hkdb.py
index 1592b474..7f8f9f73 100644
--- a/sotodlib/io/hkdb.py
+++ b/sotodlib/io/hkdb.py
@@ -488,7 +488,8 @@ def load_hk(load_spec: Union[LoadSpec, dict], show_pb=False):
 
     nframes = np.sum([len(offsets) for offsets in file_spec.values()])
     pb = tqdm(total=nframes, disable=(not show_pb))
-    for path, offsets in file_spec.items():
+    for path in sorted(list(file_spec.keys())):
+        offsets = file_spec[path]
         reader = so3g.G3IndexedReader(path)
         for offset in sorted(offsets):
             reader.Seek(offset)
diff --git a/sotodlib/tod_ops/jumps.py b/sotodlib/tod_ops/jumps.py
index 0b270735..77e6d1a7 100644
--- a/sotodlib/tod_ops/jumps.py
+++ b/sotodlib/tod_ops/jumps.py
@@ -1,18 +1,13 @@
 from typing import Literal, Optional, Tuple, Union, cast, overload
+from warnings import warn
 
 import numpy as np
 import scipy.ndimage as simg
 import scipy.stats as ss
+import so3g
 from numpy.typing import NDArray
-from pixell.utils import block_expand, block_reduce, moveaxis
 from scipy.sparse import csr_array
 from skimage.restoration import denoise_tv_chambolle
-from so3g import (
-    matched_jumps,
-    matched_jumps64,
-    find_quantized_jumps,
-    find_quantized_jumps64,
-)
 from so3g.proj import Ranges, RangesMatrix
 from sotodlib.core import AxisManager
 
@@ -56,7 +51,7 @@ def std_est(
         x = x.reshape(tuple(shape))
         x = np.moveaxis(x, -2, 0)
         diff = np.diff(x[::ds], axis=-1)
-        diff = moveaxis(diff, 0, -2)
+        diff = np.moveaxis(diff, 0, -2)
         diff = diff.reshape(shape[:-1])
         diff = np.moveaxis(diff, -1, axis)
     else:
@@ -108,9 +103,9 @@ def _jumpfinder(
     if len(x.shape) > 2:
         raise ValueError("x may not have more than 2 dimensions")
     if dtype == "float32":
-        matched_filt = matched_jumps
+        matched_filt = so3g.matched_jumps
     elif dtype == "float64":
-        matched_filt = matched_jumps64
+        matched_filt = so3g.matched_jumps64
     else:
         raise TypeError("x must be float32 or float64")
 
@@ -124,8 +119,8 @@ def _jumpfinder(
 
     # Flag with a matched filter
     win_size += win_size % 2  # Odd win size adds a wierd phasing issue
-    _x = np.ascontiguousarray(x[msk])
-    _jumps = np.ascontiguousarray(np.empty_like(_x), "int32")
+    _x = np.asarray(x[msk], dtype=x.dtype, order="C")
+    _jumps = np.empty_like(_x, dtype="int32", order="C")
     if isinstance(min_size, np.ndarray):
         _min_size = min_size[msk].astype(_x.dtype)
     elif min_size is None:
@@ -182,33 +177,45 @@ def jumpfix_subtract_heights(
         x_fixed: x with jumps removed.
                  If inplace is True this is just a reference to x.
     """
-
-    def _fix(jump_ranges, heights, x_fixed):
-        for j, jump_range in enumerate(jump_ranges):
-            for start, end in jump_range.ranges():
-                _heights = heights[j, start:end]
-                height = _heights[np.argmax(np.abs(_heights))]
-                x_fixed[j, int((start + end) / 2) :] -= height
-
-    x_fixed = x
-    if not inplace:
-        x_fixed = x.copy()
     orig_shape = x.shape
-    x_fixed = np.atleast_2d(x_fixed)
+    force_copy = not inplace
+    if inplace and not x.flags['C_CONTIGUOUS']:
+        warn("Requested in place jump fixing but signal is not contigious, a copy will be made as a buffer")
+        force_copy = True
+    x_use = np.asarray(np.atleast_2d(x), order="C")  # type: ignore
+    x_fixed = x_use
+    # When we switch to numpy 2.0 asarray can take copy=(False+force_copy if inplace else (True if force_copy else None))
+    if force_copy and np.may_share_memory(x_use, x_fixed):
+        x_fixed = x_fixed.copy()
+
     if isinstance(jumps, np.ndarray):
-        jumps = RangesMatrix.from_mask(np.atleast_2d(jumps))
+        jumps_rm = RangesMatrix.from_mask(np.atleast_2d(jumps))
     elif isinstance(jumps, Ranges):
-        jumps = RangesMatrix.from_mask(np.atleast_2d(jumps.mask()))
-    if not isinstance(jumps, RangesMatrix):
+        jumps_rm = RangesMatrix.from_mask(np.atleast_2d(jumps.mask()))
+    elif isinstance(jumps, RangesMatrix):
+        jumps_rm = jumps
+    else:
         raise TypeError("jumps not RangesMatrix or convertable to RangesMatrix")
 
     if heights is None:
-        heights = estimate_heights(x_fixed, jumps.mask(), **kwargs)
+        heights = estimate_heights(x_fixed, jumps_rm.mask(), **kwargs)
     elif isinstance(heights, csr_array):
         heights = heights.toarray()
     heights = cast(NDArray[np.floating], heights)
+    heights = np.asarray(heights, dtype=x.dtype, order="C")
 
-    _fix(jumps.ranges, heights, x_fixed)
+    dtype = x.dtype.name
+    if len(x.shape) > 2:
+        raise ValueError("x may not have more than 2 dimensions")
+    if dtype == "float32":
+        fix = so3g.subtract_jump_heights
+    elif dtype == "float64":
+        fix = so3g.subtract_jump_heights64
+    else:
+        raise TypeError("x must be float32 or float64")
+    fix(x_use, x_fixed, heights, jumps_rm)
+    if force_copy and inplace:
+        x[:] = x_fixed[:]
 
     return x_fixed.reshape(orig_shape)
 
@@ -254,6 +261,7 @@ def estimate_heights(
     twopi: bool = False,
     make_step: bool = False,
     diff_buffed: Optional[NDArray[np.floating]] = None,
+    clean: float = 80,
 ) -> NDArray[np.floating]:
     """
     Simple jump estimation routine.
@@ -273,6 +281,9 @@ def estimate_heights(
         diff_buffed: Difference between signal and a signal shifted by win_size.
                      If None will be computed.
 
+        clean: Make each jump range a constant height set by this percentile.
+               If this is 0 then it is ignored.
+
     Returns:
 
         heights: Array of jump heights.
@@ -290,6 +301,17 @@ def estimate_heights(
     heights = np.zeros_like(jumps, dtype=float)
     heights[jumps] = diff_buffed[jumps]
 
+    if clean > 0:
+        ranges = RangesMatrix.from_mask(jumps)
+        for i, det in enumerate(ranges.ranges):
+            for r in det.ranges():
+                sign = (
+                    np.sign(np.mean(np.sign(heights[i, r[0] : r[1]]))) - 1
+                ) / 2  # 0 pos, -1 neg
+                heights[i, r[0] : r[1]] = np.percentile(
+                    heights[i, r[0] : r[1]], abs(clean + 100 * sign)
+                )
+
     return heights
 
 
@@ -319,30 +341,29 @@ def _filter(
 @overload
 def twopi_jumps(
     aman,
-    signal=...,
-    win_size=...,
-    nsigma=...,
-    atol=...,
-    max_tol=...,
-    fix: Literal[True] = True,
-    inplace=...,
-    merge=...,
-    overwrite=...,
-    name=...,
-    ds=...,
+    signal,
+    win_size,
+    nsigma,
+    atol,
+    max_tol,
+    fix: Literal[True],
+    inplace,
+    merge,
+    overwrite,
+    name,
+    ds,
     **filter_pars,
-) -> Tuple[RangesMatrix, csr_array, NDArray[np.floating]]:
-    ...
+) -> Tuple[RangesMatrix, csr_array, NDArray[np.floating]]: ...
 
 
 @overload
 def twopi_jumps(
     aman,
-    signal=...,
-    win_size=...,
-    nsigma=...,
-    atol=...,
-    max_tol=...,
+    signal,
+    win_size,
+    nsigma,
+    atol,
+    max_tol,
     fix: Literal[False] = False,
     inplace=...,
     merge=...,
@@ -350,8 +371,7 @@ def twopi_jumps(
     name=...,
     ds=...,
     **filter_pars,
-) -> Tuple[RangesMatrix, csr_array]:
-    ...
+) -> Tuple[RangesMatrix, csr_array]: ...
 
 
 def twopi_jumps(
@@ -419,7 +439,7 @@ def twopi_jumps(
 
         fixed: signal with jump fixed. Only returned if fix is set.
     """
-    if signal is None:
+    if signal is None and isinstance(aman.signal, np.ndarray):
         signal = aman.signal
     if not isinstance(signal, np.ndarray):
         raise TypeError("Signal is not an array")
@@ -438,13 +458,13 @@ def twopi_jumps(
     if len(atol) != len(signal):
         raise ValueError(f"Non-scalar atol provided with length {len(atol)}")
 
-    _signal = np.ascontiguousarray(_signal)
-    heights = np.empty_like(_signal)
-    atol = np.ascontiguousarray(atol, dtype=_signal.dtype)
+    _signal = np.asarray(_signal, dtype=_signal.dtype, order="C")
+    heights = np.empty_like(_signal, order="C")
+    atol = np.asarray(atol, dtype=_signal.dtype, order="C")
     if _signal.dtype.name == "float32":
-        find_quantized_jumps(_signal, heights, atol, win_size, 2 * np.pi)
+        so3g.find_quantized_jumps(_signal, heights, atol, win_size, 2 * np.pi)
     elif _signal.dtype.name == "float64":
-        find_quantized_jumps64(_signal, heights, atol, win_size, 2 * np.pi)
+        so3g.find_quantized_jumps64(_signal, heights, atol, win_size, 2 * np.pi)
     else:
         raise TypeError("signal must be float32 or float64")
 
@@ -467,35 +487,35 @@ def twopi_jumps(
 @overload
 def slow_jumps(
     aman,
-    signal=...,
-    win_size=...,
-    thresh=...,
-    abs_thresh=...,
-    fix: Literal[True] = True,
-    inplace=...,
-    merge=...,
-    overwrite=...,
-    name=...,
+    signal,
+    win_size,
+    thresh,
+    abs_thresh,
+    fix: Literal[True],
+    inplace,
+    merge,
+    overwrite,
+    name,
+    clean,
     **filter_pars,
-) -> Tuple[RangesMatrix, csr_array, NDArray[np.floating]]:
-    ...
+) -> Tuple[RangesMatrix, csr_array, NDArray[np.floating]]: ...
 
 
 @overload
 def slow_jumps(
     aman,
-    signal=...,
-    win_size=...,
-    thresh=...,
-    abs_thresh=...,
+    signal,
+    win_size,
+    thresh,
+    abs_thresh,
     fix: Literal[False] = False,
     inplace=...,
     merge=...,
     overwrite=...,
     name=...,
+    clean=...,
     **filter_pars,
-) -> Tuple[RangesMatrix, csr_array]:
-    ...
+) -> Tuple[RangesMatrix, csr_array]: ...
 
 
 def slow_jumps(
@@ -509,6 +529,7 @@ def slow_jumps(
     merge: bool = True,
     overwrite: bool = False,
     name: str = "jumps_slow",
+    clean: float = 80,
     **filter_pars,
 ) -> Union[
     Tuple[RangesMatrix, csr_array], Tuple[RangesMatrix, csr_array, NDArray[np.floating]]
@@ -544,6 +565,8 @@ def slow_jumps(
 
         name: String used to populate field in flagmanager if merge is True.
 
+        clean: Cleaning value to pass to ``estimate_heights``. See that function for details.
+
         **filter_pars: Parameters to pass to _filter
 
     Returns:
@@ -556,24 +579,34 @@ def slow_jumps(
 
         fixed: signal with jump fixed. Only returned if fix is set.
     """
-    if signal is None:
+    if signal is None and isinstance(aman.signal, np.ndarray):
         signal = aman.signal
     if not isinstance(signal, np.ndarray):
         raise TypeError("Signal is not an array")
 
     _signal = _filter(signal, **filter_pars)
+    _signal = np.asarray(np.atleast_2d(_signal), dtype=_signal.dtype, order="C")
 
     # Block ptp
-    bptp = block_reduce(_signal, win_size, op=np.ptp, inclusive=True)
+    dtype = _signal.dtype.name
+    if len(_signal.shape) > 2:
+        raise ValueError("signal may not have more than 2 dimensions")
+    if dtype == "float32":
+        get_ptp = so3g.block_minmax
+    elif dtype == "float64":
+        get_ptp = so3g.block_minmax64
+    else:
+        raise TypeError("signal must be float32 or float64")
+    bptp = np.zeros_like(_signal, order="C")
+    get_ptp(_signal, bptp, win_size, 2, 0)
 
     if not abs_thresh:
         thresh = float(np.quantile(bptp.ravel(), thresh))
-    bptp = block_expand(bptp, win_size, _signal.shape[-1], inclusive=True)
     jumps = bptp > thresh
 
     jump_ranges = RangesMatrix.from_mask(jumps).buffer(int(win_size / 2))
     heights = estimate_heights(
-        _signal, jump_ranges.mask(), win_size=win_size, make_step=True
+        _signal, jump_ranges.mask(), win_size=win_size, make_step=True, clean=clean
     )
 
     if merge:
@@ -589,39 +622,39 @@ def slow_jumps(
 @overload
 def find_jumps(
     aman,
-    signal=...,
-    min_sigma=...,
-    min_size=...,
-    win_size=...,
-    exact=...,
+    signal,
+    min_sigma,
+    min_size,
+    win_size,
+    exact,
     fix: Literal[False] = False,
     inplace=...,
     merge=...,
     overwrite=...,
     name=...,
     ds=...,
+    clean=...,
     **filter_pars,
-) -> Tuple[RangesMatrix, csr_array]:
-    ...
+) -> Tuple[RangesMatrix, csr_array]: ...
 
 
 @overload
 def find_jumps(
     aman,
-    signal=...,
-    min_sigma=...,
-    min_size=...,
-    win_size=...,
-    exact=...,
-    fix: Literal[True] = True,
-    inplace=...,
-    merge=...,
-    overwrite=...,
-    name=...,
-    ds=...,
+    signal,
+    min_sigma,
+    min_size,
+    win_size,
+    exact,
+    fix: Literal[True],
+    inplace,
+    merge,
+    overwrite,
+    name,
+    ds,
+    clean,
     **filter_pars,
-) -> Tuple[RangesMatrix, csr_array, NDArray[np.floating]]:
-    ...
+) -> Tuple[RangesMatrix, csr_array, NDArray[np.floating]]: ...
 
 
 def find_jumps(
@@ -637,6 +670,7 @@ def find_jumps(
     overwrite: bool = False,
     name: str = "jumps",
     ds: int = 10,
+    clean: float = 80,
     **filter_pars,
 ) -> Union[
     Tuple[RangesMatrix, csr_array], Tuple[RangesMatrix, csr_array, NDArray[np.floating]]
@@ -680,6 +714,8 @@ def find_jumps(
 
         ds: Downsample factor used when computing noise level, the actual factor used is `ds*win_size`.
 
+        clean: Cleaning value to pass to ``estimate_heights``. See that function for details.
+
         **filter_pars: Parameters to pass to _filter
 
     Returns:
@@ -693,7 +729,7 @@ def find_jumps(
 
         fixed: signal with jump fixed. Only returned if fix is set.
     """
-    if signal is None:
+    if signal is None and isinstance(aman.signal, np.ndarray):
         signal = aman.signal
     if not isinstance(signal, np.ndarray):
         raise TypeError("Signal is not an array")
@@ -724,7 +760,7 @@ def find_jumps(
 
     jump_ranges = RangesMatrix.from_mask(jumps).buffer(int(win_size / 2))
     jumps = jump_ranges.mask()
-    heights = estimate_heights(signal, jumps, win_size=win_size)
+    heights = estimate_heights(signal, jumps, win_size=win_size, clean=clean)
 
     if merge:
         _merge(aman, jump_ranges, name, overwrite)


Test patch:
diff --git a/tests/test_tod_ops.py b/tests/test_tod_ops.py
index 831d215a..0db0fa6a 100644
--- a/tests/test_tod_ops.py
+++ b/tests/test_tod_ops.py
@@ -284,6 +284,7 @@ class JumpfindTest(unittest.TestCase):
         """Test that jumpfinder finds jumps in white noise."""
         np.random.seed(0)
         tod = get_tod('white')
+        ptp_orig = np.ptp(tod.signal)
         sig_jumps = tod.signal[0]
         jump_locs = np.array([200, 400, 700])
         sig_jumps[jump_locs[0]:] += 10
@@ -293,8 +294,8 @@ class JumpfindTest(unittest.TestCase):
         tod.wrap('sig_jumps', sig_jumps, [(0, 'samps')])
 
         # Find jumps without filtering
-        jumps_nf, _ = tod_ops.jumps.find_jumps(tod, signal=tod.sig_jumps, min_size=5, win_size=23)
-        jumps_nf = jumps_nf.ranges().flatten()
+        jumps, _, fixed = tod_ops.jumps.find_jumps(tod, signal=tod.sig_jumps, min_size=5, win_size=23, fix=True)
+        jumps_nf = jumps.ranges().flatten()
         
         # Remove double counted jumps and round to remove uncertainty
         jumps_nf = np.unique(np.round(jumps_nf, -2))
@@ -310,6 +311,10 @@ class JumpfindTest(unittest.TestCase):
         heights = heights[heights.nonzero()].ravel()
         self.assertTrue(np.all(np.abs(np.array([10, -13, -8]) - np.round(heights)) < 3))
 
+        # Check fixing
+        ptp_fix = np.ptp(fixed.ravel()[~jumps.buffer(10).mask().ravel()])
+        self.assertTrue(ptp_fix < 1.1*ptp_orig)
+
 
 class FFTTest(unittest.TestCase):
     def test_psd(self):