Add --plot option to slow_bend.py, for displaying strip-charts.

[blog.git] / posts / slow_bend / slow_bend.py

diff --git a/posts/slow_bend/slow_bend.py b/posts/slow_bend/slow_bend.py
index e7231f0d747518ecbda3f3c75dd0d48b16bfc078..63cd2e05dcdf19a92bf5f5cc5bd16f786b3fe5c2 100755 (executable)
--- a/posts/slow_bend/slow_bend.py
+++ b/posts/slow_bend/slow_bend.py
@@ -34,6 +34,12 @@ import os.path as _os_path
 import sys as _sys
 import time as _time
 
+try:
+    import numpy as _numpy
+    from matplotlib import pyplot as _pyplot
+except ImportError, _matplotlib_import_error:
+    _pyplot = None
+
 import pycomedi.device as _pycomedi_device
 import pycomedi.subdevice as _pycomedi_subdevice
 import pycomedi.channel as _pycomedi_channel
@@ -54,13 +60,15 @@ class Monitor (object):
     """
 
     def __init__(self, channels=[0], temperature=False, dt=4,
-                 comedi_device='/dev/comedi0', data_stream=None, logger=None):
+                 comedi_device='/dev/comedi0', data_stream=None, logger=None,
+                 plotter=None):
         self.channel_indexes = channels
         self.with_temperature = temperature
         self.dt = dt
         self.comedi_device = comedi_device
         self.data_stream = data_stream
         self._logger = logger
+        self._plotter = plotter
         self.temperature = None
 
     def log(self, level=_logging.DEBUG, msg=None):
@@ -112,7 +120,7 @@ class Monitor (object):
         self.log(msg='teardown slow bend monitor')
         self._teardown_channels()
         self._teardown_temperature()
- 
+
     def _teardown_channels(self):
         self.log(msg='teardown input channels')
         if hasattr(self, 'device') and self.device is not None:
@@ -124,7 +132,7 @@ class Monitor (object):
         if self.temperature is not None:
             self.temperature.cleanup()
             self.temperature = None
-       
+
     def _run(self):
         self.log(msg='running slow bend monitor')
         self.start_time = _time.time()
@@ -174,8 +182,10 @@ class Monitor (object):
             fields.extend(['{} ({})'.format(name, u) for u in ['bit', unit]])
         headline = '#{}'.format('\t'.join(fields))
         self.data_stream.write(headline + '\n')
-        self.log(_logging.INFO, headline)
         self.data_stream.flush()
+        self.log(_logging.INFO, headline)
+        if self._plotter:
+            self._plotter.add_header(fields)
 
     def _save_reading(self, time, bitvals, physicals):
         self.log(msg='save measurement')
@@ -184,8 +194,10 @@ class Monitor (object):
             '\t'.join('{:d}\t{:g}'.format(b, p)
                       for b,p in zip(bitvals, physicals)))
         self.data_stream.write(dataline + '\n')
-        self.log(_logging.INFO, dataline)
         self.data_stream.flush()
+        self.log(_logging.INFO, dataline)
+        if self._plotter:
+            self._plotter.add_points(time, bitvals, physicals)
 
     def _wait_until(self, target_time):
         self.log(msg='sleep until {}'.format(target_time))
@@ -194,6 +206,39 @@ class Monitor (object):
             _time.sleep(dt)
 
 
+class Plotter (object):
+    """Matplotlib-based strip-chart
+    """
+    def __init__(self, count=100):
+        self.count = count
+        if _pyplot is None:
+            raise _matplotlib_import_error
+        _pyplot.ion()
+
+    def add_header(self, fields):
+        self.figure = _pyplot.figure()
+        self.channels = len(fields)/2  # integer division
+        self.axes = []
+        self.lines = []
+        for i in range(self.channels):
+            self.axes.append(self.figure.add_subplot(self.channels, 1, i+1))
+            self.axes[i].set_title(fields[2*i+1])
+            self.lines.append(self.axes[i].plot(
+                    range(self.count),
+                    _numpy.zeros((self.count,), dtype=_numpy.float), 'r.'))
+        #_pyplot.show(block=False)  # block is an experimental kwarg
+        _pyplot.show()
+
+    def add_points(self, time, bit_values, physical_values):
+        for i in range(self.channels):
+            phys_values = self.lines[i][0].get_ydata()
+            phys_values = _numpy.roll(phys_values, -1)
+            phys_values[-1] = physical_values[i]
+            self.lines[i][0].set_ydata(phys_values)
+            self.axes[i].relim()
+            self.axes[i].autoscale(axis='y')
+        self.figure.canvas.draw()
+
 def _get_data_stream(data_dir=None, logger=None):
     if data_dir is None:
         return _sys.stdout
@@ -240,6 +285,10 @@ if __name__ == '__main__':
         '-T', '--temperature', dest='temperature',
         default=False, action='store_const', const=True,
         help='Also record the temperature')
+    parser.add_argument(
+        '-p', '--plot', dest='plot',
+        default=False, action='store_const', const=True,
+        help='Display recorded physical values on a strip-chart')
     parser.add_argument(
         '-v', '--verbose', dest='verbose',
         default=0, action='count',
@@ -255,12 +304,16 @@ if __name__ == '__main__':
         logger = _get_logger(level=level)
     else:
         logger = None
+    if args.plot:
+        plotter = Plotter()
+    else:
+        plotter = None
 
     data_stream = _get_data_stream(data_dir=args.data_dir, logger=logger)
     try:
         m = Monitor(
             channels=args.channels, temperature=args.temperature, dt=args.dt,
-            data_stream=data_stream, logger=logger)
+            data_stream=data_stream, logger=logger, plotter=plotter)
         m.run()
     finally:
         if data_stream != _sys.stdout: