Merge pull request #424 from MIT-LCP/wrsamp-checksum-spf

Set correct checksums and samps_per_frame in Record.wrsamp

Author: tompollard

tests/test_record.py

@@ -633,6 +633,22 @@ def test_4d(self):
 
         assert np.array_equal(sig_round, sig_target)
 
+    def test_write_smoothed(self):
+        """
+        Test writing a record after reading with smooth_frames
+        """
+        record = wfdb.rdrecord(
+            "sample-data/drive02",
+            physical=False,
+            smooth_frames=True,
+        )
+        record.wrsamp(write_dir=self.temp_path)
+        record2 = wfdb.rdrecord(
+            os.path.join(self.temp_path, "drive02"),
+            physical=False,
+        )
+        np.testing.assert_array_equal(record.d_signal, record2.d_signal)
+
     def test_to_dataframe(self):
         record = wfdb.rdrecord("sample-data/test01_00s")
         df = record.to_dataframe()

wfdb/io/_header.py

@@ -278,7 +278,7 @@ def set_defaults(self):
         for f in sfields:
             self.set_default(f)
 
-    def wrheader(self, write_dir=""):
+    def wrheader(self, write_dir="", expanded=True):
         """
         Write a WFDB header file. The signals are not used. Before
         writing:
@@ -290,6 +290,10 @@ def wrheader(self, write_dir=""):
         ----------
         write_dir : str, optional
             The output directory in which the header is written.
+        expanded : bool, optional
+            Whether the header file should include `samps_per_frame` (this
+            should only be true if the signal files are written using
+            `expanded=True`).
 
         Returns
         -------
@@ -305,6 +309,8 @@ def wrheader(self, write_dir=""):
         # sig_write_fields is a dictionary of
         # {field_name:required_channels}
         rec_write_fields, sig_write_fields = self.get_write_fields()
+        if not expanded:
+            sig_write_fields.pop("samps_per_frame", None)
 
         # Check the validity of individual fields used to write the header
         # Record specification fields (and comments)

wfdb/io/_signal.py

@@ -142,7 +142,7 @@ def wr_dats(self, expanded, write_dir):
         # Get all the fields used to write the header
         # Assuming this method was called through wrsamp,
         # these will have already been checked in wrheader()
-        write_fields = self.get_write_fields()
+        _, _ = self.get_write_fields()
 
         if expanded:
             # Using list of arrays e_d_signal
@@ -152,8 +152,10 @@ def wr_dats(self, expanded, write_dir):
             self.check_field("d_signal")
 
         # Check the cohesion of the d_signal field against the other
-        # fields used to write the header
-        self.check_sig_cohesion(write_fields, expanded)
+        # fields used to write the header.  (Note that for historical
+        # reasons, this doesn't actually check any of the optional
+        # header fields.)
+        self.check_sig_cohesion([], expanded)
 
         # Write each of the specified dat files
         self.wr_dat_files(expanded=expanded, write_dir=write_dir)
@@ -192,10 +194,8 @@ def check_sig_cohesion(self, write_fields, expanded):
             for ch in range(self.n_sig):
                 if len(self.e_d_signal[ch]) != spf[ch] * self.sig_len:
                     raise ValueError(
-                        "Length of channel "
-                        + str(ch)
-                        + "does not match samps_per_frame["
-                        + str(ch + "]*sig_len")
+                        f"Length of channel {ch} does not match "
+                        f"samps_per_frame[{ch}]*sig_len"
                     )
 
             # For each channel (if any), make sure the digital format has no values out of bounds

wfdb/io/record.py

@@ -922,9 +922,19 @@ def wrsamp(self, expanded=False, write_dir=""):
         N/A
 
         """
+        # Update the checksum field (except for channels that did not have
+        # a checksum to begin with, or where the checksum was already
+        # valid.)
+        if self.checksum is not None:
+            checksums = self.calc_checksum(expanded=expanded)
+            for ch, old_val in enumerate(self.checksum):
+                if old_val is None or (checksums[ch] - old_val) % 65536 == 0:
+                    checksums[ch] = old_val
+            self.checksum = checksums
+
         # Perform field validity and cohesion checks, and write the
         # header file.
-        self.wrheader(write_dir=write_dir)
+        self.wrheader(write_dir=write_dir, expanded=expanded)
         if self.n_sig > 0:
             # Perform signal validity and cohesion checks, and write the
             # associated dat files.