diff --git a/changelog.d/ann2rr-time-range.rst b/changelog.d/ann2rr-time-range.rst new file mode 100644 index 00000000..56f98277 --- /dev/null +++ b/changelog.d/ann2rr-time-range.rst @@ -0,0 +1,2 @@ +**Select intervals by time in the record in ann2rr** + The ``start_time`` and ``stop_time`` arguments of ``ann2rr`` now select intervals by when they end in the record, as the WFDB ``ann2rr`` program's ``-f`` and ``-t`` options do, and the first interval is measured from ``start_time``. They used to filter intervals by their length instead. A ``start_time`` after ``stop_time`` now raises a ``ValueError``. diff --git a/src/wfdb/processing/hr.py b/src/wfdb/processing/hr.py index aabce0b7..4621346d 100644 --- a/src/wfdb/processing/hr.py +++ b/src/wfdb/processing/hr.py @@ -198,9 +198,12 @@ def ann2rr( directory from which to find the required annotation file. eg. For record '100' in 'http://physionet.org/content/mitdb': pn_dir='mitdb'. start_time : float, optional - The time to start the intervals in seconds. + The time in the record, in seconds, from which to start. Only + intervals ending at or after this time are returned, and the first + interval is measured from this time. stop_time : float, optional - The time to stop the intervals in seconds. + The time in the record, in seconds, at which to stop. Only + intervals ending at or before this time are returned. format : str, optional Print intervals in the specified format. By default, intervals are printed in units of sample intervals. Other formats include @@ -218,6 +221,11 @@ def ann2rr( The intervals in the requested units, or None when printing. An empty array is returned if no annotations are selected. + Raises + ------ + ValueError + If `start_time` is after `stop_time`. + Examples -------- >>> from wfdb import processing @@ -229,6 +237,9 @@ def ann2rr( 257 """ + if start_time is not None and stop_time is not None and start_time > stop_time: + raise ValueError("start_time must not be after stop_time") + if (pn_dir is not None) and ("." not in pn_dir): dir_list = pn_dir.split("/") pn_dir = posixpath.join( @@ -243,15 +254,21 @@ def ann2rr( if not all_annotations: samples = samples[np.isin(ann.label_store, np.flatnonzero(is_qrs))] + # Select annotations by their time in the record, as WFDB's ann2rr -f + # and -t do, and measure the first interval from the start time. Halves + # round up, as in the WFDB library. + start = 0 + if start_time is not None: + start = int(start_time * ann.fs + 0.5) + samples = samples[samples >= start] + if stop_time is not None: + samples = samples[samples <= int(stop_time * ann.fs + 0.5)] + rr_interval = calc_rr(samples, fs=ann.fs) if samples.size: - rr_interval = np.insert(rr_interval, 0, samples[0]) + rr_interval = np.insert(rr_interval, 0, samples[0] - start) time_interval = rr_interval / ann.fs - if start_time is not None: - time_interval = time_interval[(time_interval > start_time).astype(bool)] - if stop_time is not None: - time_interval = time_interval[(time_interval < stop_time).astype(bool)] # Already given in seconds (format == 's') if format == "s": diff --git a/tests/test_hr.py b/tests/test_hr.py index b0f3b37b..98c6f4c2 100644 --- a/tests/test_hr.py +++ b/tests/test_hr.py @@ -117,3 +117,29 @@ def test_ann2rr_sample_record(): intervals = processing.ann2rr("sample-data/100", "atr") np.testing.assert_array_equal(intervals[:3], [77, 293, 292]) assert len(intervals) == 2273 + + +@pytest.mark.parametrize( + "start_time, stop_time, expected", + [ + (None, None, [100, 100, 100]), + (1.5, None, [50, 100]), + (None, 2.5, [100, 100]), + (1.5, 3.0, [50, 100]), + (2.0, 3.0, [0, 100]), + (3.5, None, []), + # Half a sample rounds up, as in WFDB. + (0.125, None, [87, 100, 100]), + ], +) +def test_ann2rr_time_range(mixed_annotations, start_time, stop_time, expected): + # Beats are at 1, 2 and 3 seconds into the record. + intervals = processing.ann2rr( + mixed_annotations, "atr", start_time=start_time, stop_time=stop_time + ) + np.testing.assert_array_equal(intervals, expected) + + +def test_ann2rr_start_after_stop(mixed_annotations): + with pytest.raises(ValueError, match="start_time"): + processing.ann2rr(mixed_annotations, "atr", start_time=2, stop_time=1)