add ecg grids

cx1111 · cx1111 · commit 32265d681014 · 2017-05-03T15:04:48.000-04:00
diff --git a/devtests.ipynb b/devtests.ipynb
diff --git a/wfdb/_headers.py b/wfdb/_headers.py
@@ -582,4 +582,41 @@ def __init__(self, allowedtypes, delimiter, dependency, write_req, read_def, wri
 
 
 
+# For storing WFDB Signal definitions.
+
+# SignalType class with all its parameters
+class SignalType():
+    def __init__(self, description, measurement=None, default_display=None, signalnames=None):
+        self.description = description
+        self.unitscale = unitscale
+        # Tuple pair (a, b). The plot displays range a, of unit b.
+        self.default_display = default_display
+        self.signalnames = signalnames
+
+unitscale = {
+    'Voltage': ['pV', 'nV', 'uV', 'mV', 'V', 'kV'],
+    'Temperature': ['C'],
+    'Pressure': ['mmHg'],
+}
+
+# All signal types
+signaltypes = {
+    'BP': SignalType('Blood Pressure', 'Pressure'),
+    'CO2': SignalType('Carbon Dioxide'),
+    'CO': SignalType('Carbon Monoxide'),
+    'ECG': SignalType('Electrocardiogram'),
+    'EEG': SignalType('Electroencephalogram'),
+    'EMG': SignalType('Electromyograph'),
+    'EOG': SignalType('Electrooculograph'),
+    'HR': SignalType('Heart Rate'),
+    'MMG': SignalType('Magnetomyograph'),
+    'O2': SignalType('Oxygen'),
+    'PLETH': SignalType('Plethysmograph'),
+    'RESP': SignalType('Respiration'),
+    'SCG': SignalType('Seismocardiogram'),
+    'STAT': SignalType('Status'), # small integers indicating status
+    'ST': SignalType('ECG ST Segment'),
+    'TEMP': SignalType('Temperature'),
+}
+
 
diff --git a/wfdb/plots.py b/wfdb/plots.py
@@ -1,12 +1,13 @@
 import numpy as np
 import matplotlib.pyplot as plt
+import math
 from . import records
 from . import _headers
 from . import annotations
 
 # Plot a WFDB Record's signals
 # Optionally, overlay annotation locations
-def plotrec(record=None, title = None, annotation = None, annch = [0], timeunits='samples', returnfig = False): 
+def plotrec(record=None, title = None, annotation = None, annch = [0], timeunits='samples', returnfig = False, ecggrids=False): 
     """ Subplot and label each channel of a WFDB Record.
     Optionally, subplot annotation locations over selected channels.
     
@@ -70,12 +71,65 @@ def plotrec(record=None, title = None, annotation = None, annch = [0], timeunits
             unitlabel='NU'
         plt.ylabel(chanlabel+"/"+unitlabel)
         
+        # Show standard ecg grids if specified.
+        if ecggrids:
+            major_ticks_x, minor_ticks_x, major_ticks_y, minor_ticks_y = calc_ecg_grids(
+                record.p_signals[:,ch], record.units[ch], record.fs, t, timeunits)                                       
+            ax.set_xticks(major_ticks_x)
+            ax.set_xticks(minor_ticks_x, minor=True)
+            ax.set_yticks(major_ticks_y)
+            ax.set_yticks(minor_ticks_y, minor=True)
+            ax.grid(which='both')
+
     plt.show(fig)
     
     # Return the figure if requested
     if returnfig:
         return fig
 
+# Calculate tick intervals for ecg grids
+def calc_ecg_grids(signal, units, fs, t, timeunits):
+
+    # 5mm 0.2s major grids, 0.04s minor grids
+    # 0.5mV major grids, 0.125 minor grids 
+    # 10 mm is equal to 1mV in voltage.
+    
+    # Get the grid interval of the x axis
+    if timeunits == 'samples':
+        majorx = 0.2*fs
+        minorx = 0.04*fs
+    elif timeunits == 'seconds':
+        majorx = 0.2
+        minorx = 0.04
+    elif timeunits == 'minutes':
+        majorx = 0.2/60
+        minorx = 0.04/60
+    elif timeunits == 'hours':
+        majorx = 0.2/3600
+        minorx = 0.04/3600
+
+    # Get the grid interval of the y axis
+    if units.lower()=='uv':
+        majory = 500
+        minory = 125
+    elif units.lower()=='mv':
+        majory = 0.5
+        minory = 0.125
+    elif units.lower()=='v':
+        majory = 0.0005
+        minory = 0.000125
+    else:
+        raise ValueError('Signal units must be uV, mV, or V to plot the ECG grid.')
+
+
+    major_ticks_x = np.arange(0, upround(max(t), majorx), majorx)
+    minor_ticks_x = np.arange(0, upround(max(t), majorx), minorx)
+
+    major_ticks_y = np.arange(downround(min(signal), majory), upround(max(signal), majory), majory)
+    minor_ticks_y = np.arange(downround(min(signal), majory), upround(max(signal), majory), minory)
+
+    return (major_ticks_x, minor_ticks_x, major_ticks_y, minor_ticks_y)
+
 # Check the validity of items used to make the plot
 # Return the x axis time values to plot for the record (and annotation if any)
 def checkplotitems(record, title, annotation, annch, timeunits):
@@ -233,3 +287,11 @@ def checkannplotitems(annotation, title, timeunits):
         raise TypeError("The 'title' field must be a string")
     
     return plotvals
+
+# Round down to nearest <base>
+def downround(x, base):
+    return base * round(float(x)/base)
+
+# Round up to nearest <base>
+def upround(x, base):
+    return base * math.ceil(float(x)/base)
