finish plot_wfdb

cx1111 · cx1111 · commit 761dbc010c67 · 2018-02-02T14:15:48.000-05:00
diff --git a/wfdb/io/annotation.py b/wfdb/io/annotation.py
@@ -1198,26 +1198,34 @@ def rdann(record_name, extension, sampfrom=0, sampto=None, shift_samps=False,
 
     """
 
-    return_label_elements = check_read_inputs(sampfrom, sampto, return_label_elements)
+    return_label_elements = check_read_inputs(sampfrom, sampto,
+                                              return_label_elements)
 
     # Read the file in byte pairs
     filebytes = load_byte_pairs(record_name, extension, pb_dir)
 
     # Get wfdb annotation fields from the file bytes
-    sample, label_store, subtype, chan, num, aux_note = proc_ann_bytes(filebytes, sampto)
+    (sample, label_store, subtype,
+     chan, num, aux_note) = proc_ann_bytes(filebytes, sampto)
 
     # Get the indices of annotations that hold definition information about
     # the entire annotation file, and other empty annotations to be removed.
-    potential_definition_inds, rm_inds = get_special_inds(sample, label_store, aux_note)
+    potential_definition_inds, rm_inds = get_special_inds(sample, label_store,
+                                                          aux_note)
 
     # Try to extract information describing the annotation file
-    fs, custom_labels = interpret_defintion_annotations(potential_definition_inds, aux_note)
+    (fs,
+     custom_labels) = interpret_defintion_annotations(potential_definition_inds,
+                                                      aux_note)
 
     # Remove annotations that do not store actual sample and label information
-    sample, label_store, subtype, chan, num, aux_note = rm_empty_indices(rm_inds, sample, label_store, subtype, chan, num, aux_note)
+    (sample, label_store, subtype,
+     chan, num, aux_note) = rm_empty_indices(rm_inds, sample, label_store,
+                                             subtype, chan, num, aux_note)
 
-    # Convert lists to numpy arrays
-    sample, label_store, subtype, chan, num= lists_to_arrays(sample, label_store, subtype, chan, num)
+    # Convert lists to numpy arrays dtype='int'
+    (sample, label_store, subtype,
+     chan, num) = lists_to_int_arrays(sample, label_store, subtype, chan, num)
 
     # Obtain annotation sample relative to the starting signal index
     if shift_samps and len(sample) > 0 and sampfrom:
@@ -1232,11 +1240,12 @@ def rdann(record_name, extension, sampfrom=0, sampto=None, shift_samps=False,
             pass
 
     # Create the annotation object
-    annotation = Annotation(os.path.split(record_name)[1], extension, sample=sample, label_store=label_store,
-                            subtype=subtype, chan=chan, num=num, aux_note=aux_note, fs=fs,
+    annotation = Annotation(record_name=os.path.split(record_name)[1],
+                            extension=extension, sample=sample,
+                            label_store=label_store,  subtype=subtype,
+                            chan=chan, num=num, aux_note=aux_note, fs=fs,
                             custom_labels=custom_labels)
 
-
     # Get the set of unique label definitions contained in this annotation
     if summarize_labels:
         annotation.get_contained_labels(inplace=True)
@@ -1509,9 +1518,9 @@ def rm_empty_indices(*args):
 
     return [[a[i] for i in keep_inds] for a in args[1:]]
 
-def lists_to_arrays(*args):
+def lists_to_int_arrays(*args):
     """
-    Convert lists to numpy arrays
+    Convert lists to numpy int arrays
     """
     return [np.array(a, dtype='int') for a in args]
 
diff --git a/wfdb/plot/plot.py b/wfdb/plot/plot.py
@@ -8,9 +8,6 @@
 from ..io.annotation import Annotation
 
 
-import pdb
-
-
 def plot_items(signal=None, ann_samp=None, ann_sym=None, fs=None,
                time_units='samples', sig_name=None, sig_units=None,
                ylabel=None, title=None, sig_style=[''], ann_style=['r*'],
@@ -296,7 +293,7 @@ def calc_ecg_grids(minsig, maxsig, sig_units, fs, maxt, time_units):
 
 
 def label_figure(axes, n_subplots, time_units, sig_name, sig_units, ylabel,
-                 title)
+                 title):
     "Add title, and axes labels"
     if title:
         axes[0].set_title(title)
@@ -391,8 +388,8 @@ def plot_wfdb(record=None, annotation=None, plot_sym=False,
 
     """
     (signal, ann_samp, ann_sym, fs, sig_name,
-        sig_units) = get_wfdb_plot_items(record, annotation)
-
+        sig_units) = get_wfdb_plot_items(record=record, annotation=annotation,
+                                         plot_sym=plot_sym)
 
     return plot_items(signal=signal, ann_samp=ann_samp, ann_sym=ann_sym, fs=fs,
                       time_units=time_units, sig_name=sig_name,
@@ -424,18 +421,33 @@ def get_wfdb_plot_items(record, annotation, plot_sym):
     if annotation:
         # Get channels
         all_chans = set(annotation.chan)
+
+        n_chans = max(all_chans) + 1
         
         # Just one channel. Place content in one list index.
-        if len(all_chans) == 1:
-            ann_samp = annotation.chan[0]*[None] + [annotation.sample]
-            if plot_sym:
-                ann_sym = annotation.chan[0]*[None] + [annotation.symbol]
-            else:
-                ann_sym = None
-        # Split annotations by channel
+        # if len(all_chans) == 1:
+        #     ann_samp = annotation.chan[0]*[None] + [annotation.sample]
+        #     if plot_sym:
+        #         ann_sym = annotation.chan[0]*[None] + [annotation.symbol]
+        #     else:
+        #         ann_sym = None
+        # # Split annotations by channel
+        # else:
+
+        # Indices for each channel
+        chan_inds = n_chans * [np.empty(0)]
+
+        for chan in all_chans:
+            chan_inds[chan] = np.where(annotation.chan == chan)[0]
+
+        ann_samp = [annotation.sample[ci] for ci in chan_inds]
+
+        if plot_sym:
+            ann_sym = n_chans * [None]
+            for ch in all_chans:
+                ann_sym[ch] = [annotation.symbol[ci] for ci in chan_inds[ch]]
         else:
-            for chan in all_chans:
-                pass
+            ann_sym = None
 
         # Try to get fs from annotation if not already in record
         if fs is None:
@@ -466,7 +478,7 @@ def plot_all_records(directory=os.getcwd()):
     records.sort()
 
     for record_name in records:
-        record = rdrecord(record_name)
+        record = rdrecord(os.path.join(directory, record_name))
 
-        plot_wfdb(record, title='Record - %s' % record.recordname)
+        plot_wfdb(record, title='Record - %s' % record.record_name)
         input('Press enter to continue...')
