Fix ann2rr intervals split by non-beat annotations

tests/test_hr.py

@@ -0,0 +1,119 @@
+"""Tests for RR interval extraction from annotation files."""
+
+import numpy as np
+import pytest
+
+import wfdb
+from wfdb import processing
+
+
+@pytest.fixture
+def mixed_annotations(tmp_path):
+    wfdb.wrann(
+        "mixed",
+        "atr",
+        np.array([0, 100, 150, 200, 250, 300, 400]),
+        symbol=["+", "N", "+", "V", "+", "N", '"'],
+        aux_note=["(N", "", "(AFIB", "", "(N", "", "note"],
+        fs=100,
+        write_dir=str(tmp_path),
+    )
+    return str(tmp_path / "mixed")
+
+
+@pytest.mark.parametrize(
+    "format, divisor", [(None, 1), ("s", 100), ("m", 6000), ("h", 360000)]
+)
+def test_ann2rr_ignores_non_beats(mixed_annotations, format, divisor):
+    intervals = processing.ann2rr(mixed_annotations, "atr", format=format)
+    np.testing.assert_allclose(intervals, np.array([100, 100, 100]) / divisor)
+    assert intervals.dtype.kind == ("i" if format is None else "f")
+
+
+def test_ann2rr_all_annotations(mixed_annotations):
+    intervals = processing.ann2rr(mixed_annotations, "atr", all_annotations=True)
+    np.testing.assert_array_equal(intervals, [0, 100, 50, 50, 50, 50, 100])
+
+
+@pytest.mark.parametrize("symbol", list("NLRaVFJASEj/QB?!enfr"))
+def test_ann2rr_beat_types(tmp_path, symbol):
+    # Include codes above the gaps in the standard annotation label table.
+    wfdb.wrann(
+        "beats",
+        "atr",
+        np.array([100, 200]),
+        symbol=[symbol, symbol],
+        fs=100,
+        write_dir=str(tmp_path),
+    )
+    intervals = processing.ann2rr(str(tmp_path / "beats"), "atr")
+    np.testing.assert_array_equal(intervals, [100, 100])
+
+
+@pytest.mark.parametrize("symbol", list('~|sT*D"=p^t+u[]@x()'))
+def test_ann2rr_no_beats(tmp_path, symbol):
+    wfdb.wrann(
+        "nonbeats",
+        "atr",
+        np.array([100, 200]),
+        symbol=[symbol, symbol],
+        fs=100,
+        write_dir=str(tmp_path),
+    )
+    intervals = processing.ann2rr(str(tmp_path / "nonbeats"), "atr")
+    assert intervals.shape == (0,)
+    assert intervals.dtype == np.int64
+
+
+@pytest.mark.parametrize("sample", [0, 100])
+def test_ann2rr_single_beat(tmp_path, sample):
+    wfdb.wrann(
+        "single",
+        "atr",
+        np.array([sample, sample]),
+        symbol=["+", "N"],
+        fs=100,
+        write_dir=str(tmp_path),
+    )
+    intervals = processing.ann2rr(str(tmp_path / "single"), "atr")
+    # The first interval remains measured from time zero, not the first marker.
+    np.testing.assert_array_equal(intervals, [sample])
+
+
+def test_ann2rr_custom_symbol(tmp_path):
+    wfdb.wrann(
+        "custom",
+        "atr",
+        np.array([100, 200]),
+        symbol=["Z", "Z"],
+        custom_labels=[(1, "Z", "Renamed normal beat")],
+        fs=100,
+        write_dir=str(tmp_path),
+    )
+    intervals = processing.ann2rr(str(tmp_path / "custom"), "atr")
+    np.testing.assert_array_equal(intervals, [100, 100])
+
+
+def test_ann2rr_print(mixed_annotations, capsys):
+    assert processing.ann2rr(mixed_annotations, "atr", as_array=False) is None
+    assert capsys.readouterr().out == "100\n100\n100\n"
+
+
+def test_ann2rr_empty_annotation(monkeypatch):
+    ann = wfdb.Annotation(
+        record_name="empty",
+        extension="atr",
+        fs=100,
+        sample=np.array([], dtype=np.int64),
+        label_store=np.array([], dtype=int),
+    )
+    monkeypatch.setattr(processing.hr, "rdann", lambda *args, **kwargs: ann)
+    for all_annotations in [False, True]:
+        intervals = processing.ann2rr("empty", "atr", all_annotations=all_annotations)
+        assert intervals.shape == (0,)
+
+
+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

wfdb/processing/hr.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from wfdb.io import download
-from wfdb.io.annotation import rdann, wrann
+from wfdb.io.annotation import is_qrs, rdann, wrann
 
 
 def compute_hr(sig_len, qrs_inds, fs):
@@ -161,10 +161,15 @@ def ann2rr(
     stop_time=None,
     format=None,
     as_array=True,
+    all_annotations=False,
 ):
     """
     Obtain RR interval series from ECG annotation files.
 
+    By default, only QRS (beat) annotations are used. The first interval
+    is measured from the beginning of the record to the first selected
+    annotation.
+
     Parameters
     ----------
     record_name : str
@@ -186,20 +191,27 @@ def ann2rr(
         printed in units of sample intervals. Other formats include
         's' (seconds), 'm' (minutes), 'h' (hours). Set to 'None' for samples.
     as_array : bool, optional
-        If True, return an an 'ndarray', else print the output.
+        If True, return an 'ndarray', else print the output.
+    all_annotations : bool, optional
+        If True, include all annotation types, rather than only QRS
+        annotations. This retains the behavior of earlier versions and
+        corresponds to the WFDB ann2rr application's -A option.
 
     Returns
     -------
-    N/A
+    out_interval : ndarray or None
+        The intervals in the requested units, or None when printing.
+        An empty array is returned if no annotations are selected.
 
     Examples
     --------
-    >>> wfdb.ann2rr('sample-data/100', 'atr', as_array=False)
-    >>> 18
-    >>> 59
-    >>> ...
-    >>> 250
-    >>> 257
+    >>> from wfdb import processing
+    >>> processing.ann2rr('sample-data/100', 'atr', as_array=False)
+    77
+    293
+    ...
+    250
+    257
 
     """
     if (pn_dir is not None) and ("." not in pn_dir):
@@ -208,10 +220,17 @@ def ann2rr(
             dir_list[0], download.get_version(dir_list[0]), *dir_list[1:]
         )
 
-    ann = rdann(record_name, extension, pn_dir=pn_dir)
+    ann = rdann(
+        record_name, extension, pn_dir=pn_dir, return_label_elements=["label_store"]
+    )
+
+    samples = ann.sample
+    if not all_annotations:
+        samples = samples[np.isin(ann.label_store, np.flatnonzero(is_qrs))]
 
-    rr_interval = calc_rr(ann.sample, fs=ann.fs)
-    rr_interval = np.insert(rr_interval, 0, ann.sample[0])
+    rr_interval = calc_rr(samples, fs=ann.fs)
+    if samples.size:
+        rr_interval = np.insert(rr_interval, 0, samples[0])
 
     time_interval = rr_interval / ann.fs
     if start_time is not None:
@@ -250,7 +269,8 @@ def rr2ann(rr_array, record_name, extension, fs=250, as_time=False):
         the shape of the input array is '(n_annot,2)', then treat the second
         column as the annotation mnemonic ('N', 'V', etc.). If a second column
         is not specified, then the default annotation will the '"' which
-        specifies a comment.
+        specifies a comment. Use ``ann2rr(..., all_annotations=True)`` to
+        include these comment annotations when extracting intervals.
     record_name : str
         The record name of the WFDB annotation file. ie. for file '100.atr',
         record_name='100'.