Merge pull request MIT-LCP#226 from MIT-LCP/edit_dtype_225

Lucas-Mc · web-flow · commit ca29df65f679 · 2020-06-02T12:11:41.000-04:00
Adds datatype parameter in rdrecord/rdsamp MIT-LCP#224 MIT-LCP#225
diff --git a/wfdb/io/_signal.py b/wfdb/io/_signal.py
@@ -856,9 +856,10 @@ def smooth_frames(self, sigtype='physical'):
 
 #------------------- Reading Signals -------------------#
 
+
 def _rd_segment(file_name, dir_name, pn_dir, fmt, n_sig, sig_len, byte_offset,
                 samps_per_frame, skew, sampfrom, sampto, channels,
-                smooth_frames, ignore_skew):
+                smooth_frames, ignore_skew, return_res=64):
     """
     Read the digital samples from a single segment record's associated
     dat file(s).
@@ -896,6 +897,11 @@ def _rd_segment(file_name, dir_name, pn_dir, fmt, n_sig, sig_len, byte_offset,
         Specifies whether to apply the skew to align the signals in the
         output variable (False), or to ignore the skew field and load in
         all values contained in the dat files unaligned (True).
+    return_res : int, optional
+        The numpy array dtype of the returned signals. Options are: 64,
+        32, 16, and 8, where the value represents the numpy int or float
+        dtype. Note that the value cannot be 8 when physical is True
+        since there is no float8 format.
 
     Returns
     -------
@@ -930,6 +936,10 @@ def _rd_segment(file_name, dir_name, pn_dir, fmt, n_sig, sig_len, byte_offset,
     if ignore_skew:
         skew = [0]*n_sig
 
+    # Change format if requested
+    if return_res != 64:
+        fmt = len(fmt) * [str(return_res)]
+
     # Get the set of dat files, and the
     # channels that belong to each file.
     file_name, datchannel = describe_list_indices(file_name)
diff --git a/wfdb/io/record.py b/wfdb/io/record.py
@@ -1517,7 +1517,8 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
                                                   record.samps_per_frame,
                                                   record.skew, sampfrom, sampto,
                                                   channels, smooth_frames,
-                                                  ignore_skew)
+                                                  ignore_skew,
+                                                  return_res=return_res)
 
             # Arrange/edit the object fields to reflect user channel
             # and/or signal range input
@@ -1538,7 +1539,8 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
                                                     record.samps_per_frame,
                                                     record.skew, sampfrom,
                                                     sampto, channels,
-                                                    smooth_frames, ignore_skew)
+                                                    smooth_frames, ignore_skew,
+                                                    return_res=return_res)
 
             # Arrange/edit the object fields to reflect user channel
             # and/or signal range input
@@ -1588,7 +1590,8 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
                 record.segments[seg_num] = rdrecord(
                     os.path.join(dir_name, record.seg_name[seg_num]),
                     sampfrom=seg_ranges[i][0], sampto=seg_ranges[i][1],
-                    channels=seg_channels[i], physical=physical, pn_dir=pn_dir)
+                    channels=seg_channels[i], physical=physical, pn_dir=pn_dir,
+                    return_res=return_res)
 
         # Arrange the fields of the layout specification segment, and
         # the overall object, to reflect user input.
@@ -1609,7 +1612,7 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
 
 
 def rdsamp(record_name, sampfrom=0, sampto=None, channels=None, pn_dir=None,
-           channel_names=None, warn_empty=False):
+           channel_names=None, warn_empty=False, return_res=64):
     """
     Read a WFDB record, and return the physical signals and a few important
     descriptor fields.
@@ -1641,6 +1644,11 @@ def rdsamp(record_name, sampfrom=0, sampto=None, channels=None, pn_dir=None,
         Whether to display a warning if the specified channel indices
         or names are not contained in the record, and no signal is
         returned.
+    return_res : int, optional
+        The numpy array dtype of the returned signals. Options are: 64,
+        32, 16, and 8, where the value represents the numpy int or float
+        dtype. Note that the value cannot be 8 when physical is True
+        since there is no float8 format.
 
     Returns
     -------
@@ -1682,8 +1690,8 @@ def rdsamp(record_name, sampfrom=0, sampto=None, channels=None, pn_dir=None,
 
     record = rdrecord(record_name=record_name, sampfrom=sampfrom,
                       sampto=sampto, channels=channels, physical=True,
-                      pn_dir=pn_dir, m2s=True, channel_names=channel_names,
-                      warn_empty=warn_empty)
+                      pn_dir=pn_dir, m2s=True, return_res=return_res,
+                      channel_names=channel_names, warn_empty=warn_empty)
 
     signals = record.p_signal
     fields = {}
