Accept sample numbers stored as floats in wrann

changelog.d/wrann-float-samples.rst

@@ -0,0 +1,2 @@
+**Accept sample numbers stored as floats in wrann**
+  ``wrann`` now accepts a ``sample`` array of floats if all the values are whole numbers (as from a detector that returns floats), and writes them as integers. Other floats, such as ``20.5`` or NaN, raise a ``ValueError``. Float arrays used to fail with ``TypeError: ufunc 'bitwise_and' not supported for the input types``.

src/wfdb/io/annotation.py

@@ -565,6 +565,30 @@ def check_field(self, field):
                         )
 
         elif field == "sample":
+            # Sample numbers may be stored as floats (for example, the
+            # output of a detector), as long as they are whole numbers.
+            # NaN and infinity are rejected before rounding, so that no
+            # floating-point exception is raised for them.
+            if np.issubdtype(self.sample.dtype, np.floating):
+                if not np.all(np.isfinite(self.sample)) or np.any(
+                    self.sample != np.round(self.sample)
+                ):
+                    raise ValueError(
+                        "The 'sample' field must only contain whole numbers"
+                    )
+            # Object arrays may hold any type of number. Check that
+            # converting them to integers does not change any value.
+            elif self.sample.dtype == object:
+                try:
+                    int_sample = self.sample.astype(np.int64)
+                except (TypeError, ValueError, OverflowError):
+                    raise ValueError(
+                        "The 'sample' field must only contain whole numbers"
+                    ) from None
+                if np.any(int_sample != self.sample):
+                    raise ValueError(
+                        "The 'sample' field must only contain whole numbers"
+                    )
             if len(self.sample) == 1:
                 sampdiffs = np.array([self.sample[0]])
             elif len(self.sample) > 1:
@@ -1156,11 +1180,13 @@ def calc_core_bytes(self):
             All of the bytes to be written to the annotation file.
 
         """
-        # The difference sample to write
-        if len(self.sample) == 1:
-            sampdiff = np.array([self.sample[0]])
+        # The difference sample to write. Sample numbers stored as floats
+        # have been checked to be whole numbers.
+        sample = self.sample.astype(np.int64)
+        if len(sample) == 1:
+            sampdiff = np.array([sample[0]])
         else:
-            sampdiff = np.concatenate(([self.sample[0]], np.diff(self.sample)))
+            sampdiff = np.concatenate(([sample[0]], np.diff(sample)))
 
         # Create a copy of the annotation object with a
         # compact version of fields to write
@@ -1734,7 +1760,8 @@ def wrann(
         The string annotation file extension.
     sample : ndarray
         A numpy array containing the annotation locations in samples relative to
-        the beginning of the record.
+        the beginning of the record. Floats are accepted if they are whole
+        numbers.
     symbol : list, numpy array, optional
         The symbols used to display the annotation labels. List or numpy array.
         If this field is present, `label_store` must not be present.

tests/test_annotation.py

@@ -1,6 +1,8 @@
 import os
 import pathlib
 import re
+from decimal import Decimal
+from fractions import Fraction
 
 import numpy as np
 import pandas as pd
@@ -360,6 +362,60 @@ def test_contained_labels_from_symbols():
     assert labels["n_occurrences"].sum() == len(ann.sample)
 
 
+def test_wrann_float_samples(temp_path):
+    """Sample numbers stored as whole-number floats are written as integers"""
+    sample = np.array([100, 2500, 70000])
+    symbol = ["N", "V", "N"]
+    wfdb.wrann("int", "atr", sample, symbol=symbol, fs=360, write_dir=temp_path)
+
+    float_sample = sample.astype("float64")
+    wfdb.wrann("float", "atr", float_sample, symbol=symbol, fs=360, write_dir=temp_path)
+    assert float_sample.dtype == np.float64
+
+    with (
+        open(os.path.join(temp_path, "int.atr"), "rb") as f1,
+        open(os.path.join(temp_path, "float.atr"), "rb") as f2,
+    ):
+        assert f1.read() == f2.read()
+
+
+@pytest.mark.parametrize("sample", [[10.0, 20.5], [10.0, np.nan], [10.0, np.inf]])
+def test_wrann_non_whole_samples(temp_path, sample):
+    with pytest.raises(ValueError, match="whole numbers"):
+        wfdb.wrann("x", "atr", np.array(sample), symbol=["N", "N"], write_dir=temp_path)
+
+
+@pytest.mark.parametrize(
+    "sample", [[10, 20], [10.0, 20.0], [np.float32(10), Fraction(20)]]
+)
+def test_wrann_object_samples(temp_path, sample):
+    """Whole numbers in an object array are written as integers"""
+    wfdb.wrann(
+        "x",
+        "atr",
+        np.array(sample, dtype=object),
+        symbol=["N", "N"],
+        write_dir=temp_path,
+    )
+    ann = wfdb.rdann(os.path.join(temp_path, "x"), "atr")
+    assert ann.sample.tolist() == [10, 20]
+
+
+@pytest.mark.parametrize(
+    "sample",
+    [[10.0, 20.5], [10, Fraction(41, 2)], [10, Decimal("20.5")], [10.0, np.nan]],
+)
+def test_wrann_non_whole_object_samples(temp_path, sample):
+    with pytest.raises(ValueError, match="whole numbers"):
+        wfdb.wrann(
+            "x",
+            "atr",
+            np.array(sample, dtype=object),
+            symbol=["N", "N"],
+            write_dir=temp_path,
+        )
+
+
 def test_rdann_path():
     """rdann accepts a pathlib.Path as the record name"""
     expected = wfdb.rdann("sample-data/100", "atr")