Add to_dataframe() to Record. (#380)

Author: Ivorforce

demo.ipynb

@@ -4,6 +4,7 @@
 import unittest
 
 import numpy as np
+import pandas as pd
 
 import wfdb
 
@@ -581,6 +582,16 @@ def test_4d(self):
 
         assert np.array_equal(sig_round, sig_target)
 
+    def test_to_dataframe(self):
+        record = wfdb.rdrecord("sample-data/test01_00s")
+        df = record.to_dataframe()
+
+        self.assertEqual(record.sig_name, list(df.columns))
+        self.assertEqual(len(df), record.sig_len)
+        self.assertEqual(df.index[0], pd.Timedelta(0))
+        self.assertEqual(df.index[-1], pd.Timedelta(seconds=1 / record.fs * (record.sig_len - 1)))
+        assert np.array_equal(record.p_signal, df.values)
+
     def test_header_with_non_utf8(self):
         """
         Ignores non-utf8 characters in the header part.

tests/test_record.py

@@ -986,6 +986,46 @@ def _arrange_fields(self, channels, sampfrom, smooth_frames):
         # Adjust date and time if necessary
         self._adjust_datetime(sampfrom=sampfrom)
 
+    def to_dataframe(self) -> pd.DataFrame:
+        """
+        Create a dataframe containing the data from this record.
+
+
+        Returns
+        -------
+        A dataframe, with sig_name in the columns. The index is a DatetimeIndex
+        if both base_date and base_time were set, otherwise a TimedeltaIndex.
+        """
+        if self.base_datetime is not None:
+            index = pd.date_range(
+                start=self.base_datetime,
+                periods=self.sig_len,
+                freq=pd.Timedelta(seconds=1 / self.fs),
+            )
+        else:
+            index = pd.timedelta_range(
+                start=pd.Timedelta(0),
+                periods=self.sig_len,
+                freq=pd.Timedelta(seconds=1 / self.fs),
+            )
+
+        if self.p_signal is not None:
+            data = self.p_signal
+        elif self.d_signal is not None:
+            data = self.d_signal
+        elif self.e_p_signal is not None:
+            data = np.array(self.e_p_signal).T
+        elif self.e_d_signal is not None:
+            data = np.array(self.e_d_signal).T
+        else:
+            raise ValueError("No signal in record.")
+
+        return pd.DataFrame(
+            data=data,
+            index=index,
+            columns=self.sig_name
+        )
+
 
 class MultiRecord(BaseRecord, _header.MultiHeaderMixin):
     """