Fix correct_peaks for peaks near the start of the signal

docs/changes.rst

@@ -38,6 +38,12 @@ Unreleased
 **Use Ruff for formatting and linting**
   The ``dev`` extra now includes ``ruff`` instead of ``black`` and ``pylint``, and the code is formatted with a line length of 88 (the default).
 
+**Fix search window in correct_peaks**
+  ``correct_peaks`` now searches ``search_radius`` samples on both sides of each peak, including the last sample of the signal. It used to stop one sample short on the right, so a peak exactly ``search_radius`` samples after the original index, or on the last sample, could not be found.
+
+**Fix correct_peaks for peaks near the start of the signal**
+  ``correct_peaks`` no longer moves peaks within ``search_radius`` samples of the start of the signal to the wrong place, often before the start (for example ``-6`` instead of ``0``). The offset for these peaks was subtracted instead of added.
+
 **Fix resample_sig for signals starting with NaN**
   ``resample_sig`` (and the functions that use it, such as ``resample_singlechan`` and ``resample_multichan``) returned only NaN values if the signal started with NaN. Leading NaN values now take the first valid value before resampling, as trailing ones already took the last.
 

src/wfdb/processing/peaks.py

@@ -225,7 +225,7 @@ def shift_peaks(sig, peak_inds, search_radius, peak_up):
     for i in range(n_peaks):
         ind = peak_inds[i]
         local_sig = sig[
-            max(0, ind - search_radius) : min(ind + search_radius, sig_len - 1)
+            max(0, ind - search_radius) : min(ind + search_radius + 1, sig_len)
         ]
 
         if peak_up:
@@ -238,7 +238,8 @@ def shift_peaks(sig, peak_inds, search_radius, peak_up):
         ind = peak_inds[i]
         if ind >= search_radius:
             break
-        shift_inds[i] -= search_radius - ind
+        # The window starts at 0 rather than ind - search_radius
+        shift_inds[i] += search_radius - ind
 
     shifted_peak_inds = peak_inds + shift_inds - search_radius
 

tests/test_processing.py

@@ -195,6 +195,32 @@ def test_correct_peaks(self):
 
         assert np.array_equal(yz, expected_peaks)
 
+    def test_correct_peaks_window_edges(self):
+        """Peaks at the edge of the search radius or the signal are found."""
+        sig = np.zeros(30)
+        sig[13] = 5
+        sig[29] = 5
+
+        # Peak exactly search_radius samples after the guess
+        assert processing.correct_peaks(
+            sig, np.array([10]), search_radius=3, smooth_window_size=9, peak_dir="up"
+        ).tolist() == [13]
+        # Peak on the last sample of the signal
+        assert processing.correct_peaks(
+            sig, np.array([27]), search_radius=5, smooth_window_size=9, peak_dir="up"
+        ).tolist() == [29]
+
+    @pytest.mark.parametrize("peak", [0, 1, 4, 6])
+    @pytest.mark.parametrize("peak_dir", ["up", "down"])
+    def test_correct_peaks_near_start(self, peak, peak_dir):
+        """Peaks closer to the start than search_radius aren't shifted below 0."""
+        sig = np.zeros(30)
+        sig[peak] = 5 if peak_dir == "up" else -5
+
+        assert processing.correct_peaks(
+            sig, np.array([2]), search_radius=5, smooth_window_size=9, peak_dir=peak_dir
+        ).tolist() == [peak]
+
 
 class TestComputeHr:
     """Tests for compute_hr."""