Clean up src/plot_image/plot_image.py.
authorW. Trevor King <wking@drexel.edu>
Tue, 9 Nov 2010 14:38:23 +0000 (09:38 -0500)
committerW. Trevor King <wking@drexel.edu>
Tue, 9 Nov 2010 14:38:23 +0000 (09:38 -0500)
src/plot_image/plot_image.py

index 641e2852d32f0ee6860f5486cbf85516b17af992..c77cfd678339e3c0a38dd5f27b2fca2aa7c5f36c 100755 (executable)
 #! /usr/bin/env python
 
-from __future__ import division
-from pylab import *
-import sys
-import Numeric
-
-
-#
-#   plot_image.py
-#
-#   Step 1: make this file executable
-#
-#           chmod +x plot_image.py
-#
-#   Step 2: pipe data in python script
-#
-#           ./gen_data | ./plot_image -s nx ny -c nc -t 'image title'
-#
-#   with optional arguments
-#          -s nx ,ny    image size [16x16]
-#          -c nc        number of contour levels [none]
-#          -t ' '       image title ['some like it hot']
-#
-#   additional:  -g  gray map  [jet map]
-#                -h  hot map
-#
-#   ref: matplotlib web site
-#                                             Michel Vallieres, 2007
-#
-
-                        # dummy function to initialize Z
-def func3(x,y):
-    return 0.005*x*y
-
-                        # defaults
-mycontours = 0
-nx = 16
-ny = 16
-mytitle = 'Some like it hot'
-mymap = cm.jet
-
-                        # parse command line arguments
-n = len( sys.argv )
-i = 1
-while i < n:
-    if sys.argv[i].find("s") == 1:
-        nx = int( sys.argv[i+1] )
-        ny = int( sys.argv[i+2] )
-        i = i + 2
-    elif sys.argv[i].find("c") == 1:
-        mycontours = int( sys.argv[i+1] )
-        i = i + 1
-    elif sys.argv[i].find("t") == 1:
-        mytitle = sys.argv[i+1]
-        i = i + 1
-    elif sys.argv[i].find("g") == 1:
-        mymap = cm.gray
-    elif sys.argv[i].find("h") == 1:
-        mymap = cm.hot
-    else:
-        print " Syntax:    script -s nx ny -c "
-    i = i + 1
+"""Generate an image from ASCII data.
 
-                       # identification
-print "Plot_image"
-print "Title:            ", mytitle
-print "Image size:       ", nx, ny
-print "# countour lines: ", mycontours
+Step 1: make this file executable::
 
+    chmod +x plot_image.py
 
-                       # set grid
-x = range( nx )
-y = range( ny )
+Step 2: pipe data into python script::
 
-X,Y = meshgrid( x, y )
+    ./gen_data | ./plot_image.py -s nx,ny -c nc -t 'image title'
 
-Z = func3( X, Y )
+Data should be one ASCII float per line,
 
-                       # read in data
-for j in y:
-   for i in x:
-      Z[j,i] = input()
+For  usage info, run::
 
-                       # min & max
-min_data = Z[0,0]
-max_data = Z[0,0]
-for i in x:
-   for j in y:
-      if Z[j,i] < min_data:
-          min_data = Z[j,i]
-      if Z[j,i] > max_data:
-          max_data = Z[j,i]
+    ./plot_image.py --help
 
-print "Data range:       ", min_data, max_data
+For other ideas, see the Matplotlib website [1]_.
 
+.. [1] http://matplotlib.sourceforge.net/
+"""
 
-                       # colored image
-im = imshow( Z, interpolation='bilinear', origin='lower',
-            cmap=mymap, extent=(1,nx-1.0,1,ny-1.0) )
+import optparse
+import sys
 
-                       # contour lines
-if mycontours > 0:
-    dcont = ( max_data - min_data ) / ( mycontours - 1 )
-    cset = contour( Z, arange(min_data,max_data,dcont),
-               origin='lower',
-               linewidths=2,
-               extent=(0,nx-1,0,ny-1)
-               )
+import numpy
+
+# Depending on your Matplotlib configuration, you may need to adjust
+# your backend.  Do this before importing pylab or matplotlib.backends.
+#import matplotlib
+#matplotlib.use('Agg')     # select backend that doesn't require X Windows
+#matplotlib.use('GTKAgg')  # select backend that supports pylab.show()
+
+import pylab
+
+
+_DOC = __doc__
+
+
+def read_data(stream, nx, ny):
+    """Read in data, one entry per line.
+
+    >>> from StringIO import StringIO
+    >>> s = StringIO('\\n'.join([str(x) for x in range(10)])+'\\n')
+    >>> X,Y,Z = read_data(s, 5, 2)
+    >>> X
+    array([[0, 1, 2, 3, 4, 5],
+           [0, 1, 2, 3, 4, 5],
+           [0, 1, 2, 3, 4, 5]])
+    >>> Y
+    array([[0, 0, 0, 0, 0, 0],
+           [1, 1, 1, 1, 1, 1],
+           [2, 2, 2, 2, 2, 2]])
+    >>> Z
+    array([[ 0.,  1.,  2.,  3.,  4.],
+           [ 5.,  6.,  7.,  8.,  9.]])
+    """
+    X,Y = pylab.meshgrid(range(nx+1), range(ny+1))
+    Z = numpy.loadtxt(stream)
+    Z = Z.reshape([x-1 for x in X.shape])
+    return (X,Y,Z)
+
+
+def plot(X, Y, Z, full=False, title=None, contours=None, cmap=None):
+    """Plot Z over the mesh X, Y.
+
+    >>> X, Y = pylab.meshgrid(range(6), range(2))
+    >>> Z = X[:-1,:-1]**2 + Y[:-1,:-1]
+    >>> plot(X, Y, Z)  # doctest: +ELLIPSIS
+    <matplotlib.figure.Figure object at 0x...>
+    """
+    X_min = X[0,0]
+    X_max = X[-1,-1]
+    Y_min = Y[0,0]
+    Y_max = Y[-1,-1]
+
+    fig = pylab.figure()
+    if full:
+        axes = pylab.axes([0, 0, 1, 1])
+    else:
+        axes = pylab.axes()
+        if title:
+            axes.set_title(title)
+    axes.set_axis_off()
+
+    if contours:
+        cset = axes.contour(X[:-1,:-1], Y[:-1,:-1], Z, contours, cmap=cmap)
+        # [:-1,:-1] to strip dummy last row & column from X&Y.
+        pylab.clabel(cset, inline=1, fmt='%1.1f', fontsize=10)
+    else:
+        # pcolor() is much slower than imshow.
+        #plot = axes.pcolor(X, Y, Z, cmap=cmap, edgecolors='none')
+        #axes.autoscale_view(tight=True)
+        plot = axes.imshow(Z, aspect='auto', interpolation='bilinear',
+                           origin='lower', cmap=cmap,
+                           extent=(X_min, X_max, Y_min, Y_max))
+        if not full:
+            fig.colorbar(plot)
+    return fig
+
+
+def test():
+    import doctest
+    results = doctest.testmod()
+    return results.failed
+
+
+def main(argv=None):
+    """Read in data and plot it.
+
+    >>> from tempfile import NamedTemporaryFile
+    >>> i = NamedTemporaryFile(prefix='tmp-input', suffix='.dat')
+    >>> i.write('\\n'.join([str(x) for x in range(10)])+'\\n')
+    >>> i.flush()
+    >>> o = NamedTemporaryFile(prefix='tmp-output', suffix='.png')
+    >>> main(['-i', i.name, '-s', '5,2', '-o', o.name, '-m', 'binary'])
+    Plot_image
+    Title:             Some like it hot
+    Image size:        5 2
+    False color
+    Data range:        0.0 9.0
+    >>> img = o.read()
+    >>> img.startswith('\\x89PNG')
+    True
+    >>> i.close()
+    >>> o.close()
+    """
+    if argv == None:
+        argv = sys.argv[1:]
+
+    usage = '%prog [options]'
+    epilog = _DOC
+    p = optparse.OptionParser(usage=usage, epilog=epilog)
+    p.format_epilog = lambda formatter: epilog+'\n'
+
+    p.add_option('-i', '--input', dest='input', metavar='FILE',
+                 help='If set, read data from FILE rather than stdin.')
+    p.add_option('-o', '--output', dest='output', metavar='FILE',
+                 help=('If set, save the figure to FILE rather than '
+                       'displaying it immediately'))
+    p.add_option('-s', '--size', dest='size', default='%d,%d' % (16, 16),
+                 help='Data size (columns,rows; default: %default)')
+    p.add_option('-c', '--contours', dest='contours', type='int',
+                 help=('Number of contour lines (if not set, draw false color '
+                       'instead of contour lines; default: %default)'))
+    p.add_option('-f', '--full-figure', dest='full', action='store_true',
+                 help=('Set axes to fill the figure (i.e. no title or color '
+                       'bar'))
+    p.add_option('-t', '--title', dest='title', default='Some like it hot',
+                 help='Title (%default)')
+    p.add_option('--test', dest='test', action='store_true',
+                 help='Run internal tests and exit.')
+    maps=[m for m in pylab.cm.datad if not m.endswith("_r")]
+    maps.sort()
+    p.add_option('-m', '--color-map', dest='cmap', default='jet',
+                 help='Select color map from %s (%%default)' % ', '.join(maps))
+
+    options,args = p.parse_args(argv)
+
+    if options.test:
+        sys.exit(test())
+
+    nx,ny = [int(x) for x in options.size.split(',')]
+    try:
+        cmap = getattr(pylab.cm, options.cmap)
+    except AttributeError:
+        raise Exception('no color map named %s in %s'
+                        % (options.cmap, ', '.join(maps)))
+
+    print 'Plot_image'
+    print 'Title:            ', options.title
+    print 'Image size:       ', nx, ny
+    if options.contours:
+        print '# countour lines: ', options.contours
+    else:
+        print 'False color'
+    if options.input:
+        fin = open(options.input, 'r')
+    else:
+        fin = sys.stdin
+
+    X,Y,Z = read_data(fin, nx, ny)
+
+    if options.input:
+        fin.close()
 
-    clabel( cset, inline=1, fmt='%1.1f', fontsize=10 )
+    Z_min = numpy.min(Z.flat)
+    Z_max = numpy.max(Z.flat)
+    print 'Data range:       ', Z_min, Z_max
 
+    fig = plot(X, Y, Z, full=options.full, title=options.title,
+               contours=options.contours, cmap=cmap)
+
+    if options.output:
+        fig.savefig(options.output)
+    else:
+        pylab.ion()
+        pylab.show()
 
-                       # render picture
-axis('off')
 
-colorbar()
-title( mytitle )
-show()
+if __name__ == '__main__':
+    main()