diff --git a/tests/test_processing.py b/tests/test_processing.py index 0f371660..2bfe7377 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -4,7 +4,7 @@ from wfdb import processing -class test_processing: +class TestProcessing: """ Test processing functions """ @@ -62,6 +62,13 @@ def test_find_peaks_empty(self): assert hp.shape == (0,) assert sp.shape == (0,) + def test_find_local_peaks(self): + # the sample `radius` positions to the right is part of the window + x = np.array([0, 1, 2, 3, 2, 1, 0], dtype=float) + assert np.array_equal(processing.find_local_peaks(x, 1), [3]) + x = np.array([0, 1, 3, 2, 10, 5, 1, 0], dtype=float) + assert np.array_equal(processing.find_local_peaks(x, 2), [4]) + def test_gqrs(self): record = wfdb.rdrecord( "sample-data/100", @@ -178,7 +185,7 @@ def test_correct_peaks(self): assert np.array_equal(yz, expected_peaks) -class test_qrs: +class TestQrs: """ Testing QRS detectors """ diff --git a/wfdb/processing/peaks.py b/wfdb/processing/peaks.py index 47d7ce48..75a05327 100644 --- a/wfdb/processing/peaks.py +++ b/wfdb/processing/peaks.py @@ -82,22 +82,8 @@ def find_local_peaks(sig, radius): peak_inds = [] i = 0 - while i < radius + 1: - if sig[i] == max(sig[: i + radius]): - peak_inds.append(i) - i += radius - else: - i += 1 - - while i < len(sig): - if sig[i] == max(sig[i - radius : i + radius]): - peak_inds.append(i) - i += radius - else: - i += 1 - while i < len(sig): - if sig[i] == max(sig[i - radius :]): + if sig[i] == max(sig[max(0, i - radius) : i + radius + 1]): peak_inds.append(i) i += radius else: