Simplify for loop when calculating sigma

hpc4cmb · Oct 11, 2024 · 517f00b · 517f00b
1 parent 30df347
commit 517f00b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/toast/ops/simple_jumpcorrect.py b/src/toast/ops/simple_jumpcorrect.py
@@ -179,7 +179,7 @@ def _find_peaks(self, toi, flag, flag_out, lim=3.0, tol=1e4, sigma_in=None):
             npeak = np.ma.sum(np.abs(mytoi) > sigma * lim)
 
         # Only one jump per iteration
-        if npeak > 0:
+        while npeak > 0:
             imax = np.argmax(np.abs(mytoi))
             amplitude = mytoi[imax]
             significance = np.abs(amplitude) / sigma
@@ -192,6 +192,7 @@ def _find_peaks(self, toi, flag, flag_out, lim=3.0, tol=1e4, sigma_in=None):
             flag_out[istart:istop] = True
             if sigma_in is None:
                 sigma = self._get_sigma(mytoi, flag_out, tol)
+                break
 
             # Excessive flagging is a sign of false detection
             if significance > 5 or (
@@ -208,10 +209,9 @@ def _get_sigma(self, toi, flag, tol):
 
         sigmas = []
         nn = len(toi)
-        for start in range(tol, nn - tol, 2 * tol):
+        # Ignore tol samples at the edge
+        for start in range(tol, nn - 3*tol + 1, 2*tol):
             stop = start + 2 * tol
-            if stop > nn - tol:
-                break
             ind = slice(start, stop)
             x = toi[ind][full_flag[ind] == 0]
             if len(x) != 0: