Optimization of wfdb.io.annotation.field2bytes function (#406)

Authored-by: Nicolas Beaudoin-Gagnon <[email protected]>

Author: Fegalf

tests/test_annotation.py

@@ -260,15 +260,23 @@ def test_5(self):
         ann_idx = np.array([1, 1000, 2000, 3000])
         ann_chan = np.array([3, 1, 2, 3])
         # write custom labels
-        ann_label_store = np.array([ 4, 2, 1, 3])
-        ann_custom_labels = {'label_store': [1, 2, 3, 4],
-                             'symbol': ['v','l','r','z'],
-                             'description':['pvc','lbbb','rbbb','pac']}
+        ann_label_store = np.array([4, 2, 1, 3])
+        ann_custom_labels = {
+            "label_store": [1, 2, 3, 4],
+            "symbol": ["v", "l", "r", "z"],
+            "description": ["pvc", "lbbb", "rbbb", "pac"],
+        }
         ann_custom_labels = pd.DataFrame(data=ann_custom_labels)
-        wfdb.wrann('CustomLabel', 'atr', ann_idx, chan=ann_chan,
-                   custom_labels=ann_custom_labels, label_store=ann_label_store)
-        ann = wfdb.rdann('CustomLabel', 'atr')
-        self.assertEqual(ann.symbol, ['z', 'l', 'v', 'r'])
+        wfdb.wrann(
+            "CustomLabel",
+            "atr",
+            ann_idx,
+            chan=ann_chan,
+            custom_labels=ann_custom_labels,
+            label_store=ann_label_store,
+        )
+        ann = wfdb.rdann("CustomLabel", "atr")
+        self.assertEqual(ann.symbol, ["z", "l", "v", "r"])
 
     @classmethod
     def tearDownClass(cls):
@@ -277,7 +285,7 @@ def tearDownClass(cls):
             "1003.atr",
             "12726.anI",
             "huge.qrs",
-            "CustomLabel.atr"
+            "CustomLabel.atr",
         ]
         for file in writefiles:
             if os.path.isfile(file):

tests/test_record.py

@@ -589,7 +589,10 @@ def test_to_dataframe(self):
         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)))
+        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):

wfdb/io/_signal.py

@@ -2505,4 +2505,3 @@ def _infer_sig_len(
     sig_len = int(data_size / (BYTES_PER_SAMPLE[fmt] * tsamps_per_frame))
 
     return sig_len
-

wfdb/io/annotation.py

@@ -1165,19 +1165,34 @@ def calc_core_bytes(self):
 
         data_bytes = []
 
+        # Allow use of custom labels
+        label_table = ann_label_table
+        if self.custom_labels is not None:
+            label_table = pd.concat(
+                [label_table, self.custom_labels], ignore_index=True
+            )
+
+        # Generate typecodes from annotation label table
+        typecodes = {
+            label_table.iloc[i]["symbol"]: label_table.iloc[i]["label_store"]
+            for i in range(len(label_table))
+        }
+
         # Iterate across all fields one index at a time
         for i in range(len(sampdiff)):
 
             # Process the samp (difference) and sym items
             data_bytes.append(
-                field2bytes("samptype", [sampdiff[i], self.symbol[i]], self.custom_labels)
+                field2bytes(
+                    "samptype", [sampdiff[i], self.symbol[i]], typecodes
+                )
             )
 
             # Process the extra optional fields
             for field in extra_write_fields:
                 value = getattr(compact_annotation, field)[i]
                 if value is not None:
-                    data_bytes.append(field2bytes(field, value, self.custom_labels))
+                    data_bytes.append(field2bytes(field, value, typecodes))
 
         # Flatten and convert to correct format
         data_bytes = np.array(
@@ -1600,7 +1615,7 @@ def compact_carry_field(full_field):
     return compact_field
 
 
-def field2bytes(field, value, custom_labels=None):
+def field2bytes(field, value, typecodes):
     """
     Convert an annotation field into bytes to write.
 
@@ -1610,6 +1625,8 @@ def field2bytes(field, value, custom_labels=None):
         The annotation field of the value to be converted to bytes.
     value : list
         The value to be converted to bytes.
+    typecodes : dict
+        The mapping between each annotation label an its corresponding typecode.
 
     Returns
     -------
@@ -1619,18 +1636,10 @@ def field2bytes(field, value, custom_labels=None):
     """
     data_bytes = []
 
-    # allow use of custom labels
-    label_table = ann_label_table
-    if custom_labels is not None:
-        label_table = pd.concat([label_table, custom_labels], ignore_index=True)
-
     # samp and sym bytes come together
     if field == "samptype":
         # Numerical value encoding annotation symbol
-        typecode = label_table.loc[
-            label_table["symbol"] == value[1], "label_store"
-        ].values[0]
-
+        typecode = typecodes[value[1]]
         # sample difference
         sd = value[0]
 

wfdb/io/record.py

@@ -1021,11 +1021,7 @@ def to_dataframe(self) -> pd.DataFrame:
         else:
             raise ValueError("No signal in record.")
 
-        return pd.DataFrame(
-            data=data,
-            index=index,
-            columns=self.sig_name
-        )
+        return pd.DataFrame(data=data, index=index, columns=self.sig_name)
 
 
 class MultiRecord(BaseRecord, _header.MultiHeaderMixin):