Merge pull request MIT-LCP#384 from MIT-LCP/time-conversion

tompollard · web-flow · commit 84cbefbc6334 · 2022-06-03T13:35:46.000-04:00
Convenience methods for time conversion
diff --git a/docs/io.rst b/docs/io.rst
@@ -12,10 +12,12 @@ WFDB Records
     :members: rdrecord, rdheader, rdsamp, wrsamp
 
 .. autoclass:: wfdb.io.Record
-    :members: wrsamp, adc, dac
+    :members: get_frame_number, get_elapsed_time, get_absolute_time,
+              wrsamp, adc, dac
 
 .. autoclass:: wfdb.io.MultiRecord
-    :members: multi_to_single
+    :members: get_frame_number, get_elapsed_time, get_absolute_time,
+              multi_to_single
 
 
 WFDB Anotations
diff --git a/docs/wfdb.rst b/docs/wfdb.rst
@@ -12,10 +12,12 @@ WFDB Records
     :members: rdrecord, rdheader, rdsamp, wrsamp
 
 .. autoclass:: wfdb.Record
-    :members: wrsamp, adc, dac
+    :members: get_frame_number, get_elapsed_time, get_absolute_time,
+              wrsamp, adc, dac
 
 .. autoclass:: wfdb.MultiRecord
-    :members: multi_to_single
+    :members: get_frame_number, get_elapsed_time, get_absolute_time,
+              multi_to_single
 
 
 WFDB Anotations
diff --git a/tests/test_record.py b/tests/test_record.py
@@ -1,3 +1,4 @@
+import datetime
 import os
 import shutil
 import unittest
@@ -817,6 +818,68 @@ def test_multi_variable_d(self):
         assert record.__eq__(record_named)
 
 
+class TestTimeConversion(unittest.TestCase):
+    """
+    Test cases for time conversion
+    """
+
+    def test_single(self):
+        """
+        Time conversion for a single-segment record
+
+        This checks the get_frame_number, get_elapsed_time, and
+        get_absolute_time methods for a Record object.  The example record
+        has no base date defined, so attempting to convert to/from absolute
+        time should raise an exception.
+
+        """
+        header = wfdb.rdheader("sample-data/test01_00s")
+
+        # these time values should be equivalent
+        n = 123 * header.fs
+        t = datetime.timedelta(seconds=123)
+        self.assertEqual(header.get_frame_number(n), n)
+        self.assertEqual(header.get_frame_number(t), n)
+        self.assertEqual(header.get_elapsed_time(n), t)
+        self.assertEqual(header.get_elapsed_time(t), t)
+
+        # record test01_00s has no base date, so absolute time conversions
+        # should fail
+        self.assertIsNone(header.base_date)
+        d = datetime.datetime(2001, 1, 1, 12, 0, 0)
+        self.assertRaises(ValueError, header.get_frame_number, d)
+        self.assertRaises(ValueError, header.get_absolute_time, n)
+        self.assertRaises(ValueError, header.get_absolute_time, t)
+
+    def test_multisegment_with_date(self):
+        """
+        Time conversion for a multi-segment record with base date
+
+        This checks the get_frame_number, get_elapsed_time, and
+        get_absolute_time methods for a MultiRecord object.  The example
+        record has a base date, so we can convert timestamps between all
+        three of the supported representations.
+
+        """
+        header = wfdb.rdheader(
+            "sample-data/multi-segment/p000878/p000878-2137-10-26-16-57"
+        )
+
+        # these time values should be equivalent
+        n = 123 * header.fs
+        t = datetime.timedelta(seconds=123)
+        d = t + header.base_datetime
+        self.assertEqual(header.get_frame_number(n), n)
+        self.assertEqual(header.get_frame_number(t), n)
+        self.assertEqual(header.get_frame_number(d), n)
+        self.assertEqual(header.get_elapsed_time(n), t)
+        self.assertEqual(header.get_elapsed_time(t), t)
+        self.assertEqual(header.get_elapsed_time(d), t)
+        self.assertEqual(header.get_absolute_time(n), d)
+        self.assertEqual(header.get_absolute_time(t), d)
+        self.assertEqual(header.get_absolute_time(d), d)
+
+
 class TestSignal(unittest.TestCase):
     """
     For lower level signal tests
diff --git a/wfdb/io/record.py b/wfdb/io/record.py
@@ -249,6 +249,105 @@ def base_datetime(self, value):
         else:
             raise TypeError(f"invalid base_datetime value: {value!r}")
 
+    def get_frame_number(self, time_value):
+        """
+        Convert a time value to a frame number.
+
+        A time value may be specified as:
+        - An integer or floating-point number, representing the number of
+          WFDB frames elapsed from the start of the record.
+        - A `datetime.timedelta` object, representing elapsed time from the
+          start of the record.
+        - A `datetime.datetime` object, representing an absolute date and
+          time (if the record starting time is known.)
+
+        Note that this function may return a value that is less than zero
+        or greater than the actual length of the record.
+
+        Parameters
+        ----------
+        time_value : number or timedelta or datetime
+            A time value.
+
+        Returns
+        -------
+        frame_number : float
+            Frame number (possibly a fractional frame number).
+
+        """
+        if hasattr(time_value, "__float__"):
+            return float(time_value)
+
+        if isinstance(time_value, datetime.datetime):
+            if not self.base_datetime:
+                raise ValueError(
+                    "base_datetime is unknown; cannot convert absolute "
+                    "date/time to a frame number"
+                )
+            time_value -= self.base_datetime
+
+        if isinstance(time_value, datetime.timedelta):
+            return time_value.total_seconds() * self.fs
+
+        raise TypeError(f"invalid time value: {time_value!r}")
+
+    def get_elapsed_time(self, time_value):
+        """
+        Convert a time value to an elapsed time in seconds.
+
+        A time value may be specified as:
+        - An integer or floating-point number, representing the number of
+          WFDB frames elapsed from the start of the record.
+        - A `datetime.timedelta` object, representing elapsed time from the
+          start of the record.
+        - A `datetime.datetime` object, representing an absolute date and
+          time (if the record starting time is known.)
+
+        Parameters
+        ----------
+        time_value : number or timedelta or datetime
+            A time value.
+
+        Returns
+        -------
+        elapsed_time : timedelta
+            Elapsed time from the start of the record.
+
+        """
+        time_value = self.get_frame_number(time_value)
+        return datetime.timedelta(seconds=time_value / self.fs)
+
+    def get_absolute_time(self, time_value):
+        """
+        Convert a time value to an absolute date and time.
+
+        A time value may be specified as:
+        - An integer or floating-point number, representing the number of
+          WFDB frames elapsed from the start of the record.
+        - A `datetime.timedelta` object, representing elapsed time from the
+          start of the record.
+        - A `datetime.datetime` object, representing an absolute date and
+          time (if the record starting time is known.)
+
+        Parameters
+        ----------
+        time_value : number or timedelta or datetime
+            A time value.
+
+        Returns
+        -------
+        absolute_time : datetime
+            Absolute date and time.
+
+        """
+        time_value = self.get_elapsed_time(time_value)
+        if not self.base_datetime:
+            raise ValueError(
+                "base_datetime is unknown; cannot convert frame number "
+                "to an absolute date/time"
+            )
+        return time_value + self.base_datetime
+
     def check_field(self, field, required_channels="all"):
         """
         Check whether a single field is valid in its basic form. Does
