diff --git a/changelog.d/wrann-float-samples.rst b/changelog.d/wrann-float-samples.rst new file mode 100644 index 00000000..109a40dc --- /dev/null +++ b/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``. diff --git a/src/wfdb/io/annotation.py b/src/wfdb/io/annotation.py index a6b94df8..64f80df3 100644 --- a/src/wfdb/io/annotation.py +++ b/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. diff --git a/tests/test_annotation.py b/tests/test_annotation.py index ef0f55d7..f0b88354 100644 --- a/tests/test_annotation.py +++ b/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")