src/plot_image/plot_image.py

   1 #! /usr/bin/env python
   2
   3 """Generate an image from ASCII data.
   4
   5 Step 1: make this file executable::
   6
   7     chmod +x plot_image.py
   8
   9 Step 2: pipe data into python script::
  10
  11     ./gen_data | ./plot_image.py -s nx,ny -c nc -t 'image title'
  12
  13 Data should be one ASCII float per line, in the following order::
  14
  15     z[x1,y1]
  16     z[x2,y1]
  17     ...
  18     z[nN,y1]
  19     z[x1,y2]
  20     ...
  21     z[xN,y2]
  22     ...
  23     z[xN,yN]
  24
  25 where `x` increases from `x1` to `xN` and `y` increases from `y1`
  26 through `yN`.
  27
  28 For  usage info, run::
  29
  30     ./plot_image.py --help
  31
  32 When run with interactive output (i.e. no `-o ...` option), the
  33 interactive figure is displayed using `pylab.show`, which means that
  34 you'll have to kill `plot_image.py` using ^C or similar [1]_.
  35
  36 For other ideas, see the Matplotlib website [2]_.
  37
  38 .. [1] http://matplotlib.sourceforge.net/faq/howto_faq.html#use-show
  39 .. [2] http://matplotlib.sourceforge.net/
  40 """
  41
  42 import optparse
  43 import sys
  44
  45 import numpy
  46
  47 # Depending on your Matplotlib configuration, you may need to adjust
  48 # your backend.  Do this before importing pylab or matplotlib.backends.
  49 #import matplotlib
  50 #matplotlib.use('Agg')     # select backend that doesn't require X Windows
  51 #matplotlib.use('GTKAgg')  # select backend that supports pylab.show()
  52
  53 import pylab
  54
  55
  56 _DOC = __doc__
  57
  58
  59 def read_data(stream, nx, ny):
  60     """Read in data, one entry per line.
  61
  62     >>> from StringIO import StringIO
  63     >>> s = StringIO('\\n'.join([str(x) for x in range(10)])+'\\n')
  64     >>> X,Y,Z = read_data(s, 5, 2)
  65     >>> X
  66     array([[0, 1, 2, 3, 4, 5],
  67            [0, 1, 2, 3, 4, 5],
  68            [0, 1, 2, 3, 4, 5]])
  69     >>> Y
  70     array([[0, 0, 0, 0, 0, 0],
  71            [1, 1, 1, 1, 1, 1],
  72            [2, 2, 2, 2, 2, 2]])
  73     >>> Z
  74     array([[ 0.,  1.,  2.,  3.,  4.],
  75            [ 5.,  6.,  7.,  8.,  9.]])
  76     """
  77     X,Y = pylab.meshgrid(range(nx+1), range(ny+1))
  78     Z = numpy.loadtxt(stream)
  79     assert Z.size == nx*ny, 'Z.size = %d != %d = %dx%d' % (
  80         Z.size, nx*ny, nx, ny)
  81     Z = Z.reshape([x-1 for x in X.shape])
  82     return (X,Y,Z)
  83
  84
  85 def plot(X, Y, Z, full=False, title=None, contours=None, cmap=None):
  86     """Plot Z over the mesh X, Y.
  87
  88     >>> X, Y = pylab.meshgrid(range(6), range(2))
  89     >>> Z = X[:-1,:-1]**2 + Y[:-1,:-1]
  90     >>> plot(X, Y, Z)  # doctest: +ELLIPSIS
  91     <matplotlib.figure.Figure object at 0x...>
  92     """
  93     X_min = X[0,0]
  94     X_max = X[-1,-1]
  95     Y_min = Y[0,0]
  96     Y_max = Y[-1,-1]
  97
  98     fig = pylab.figure()
  99     if full:
 100         axes = pylab.axes([0, 0, 1, 1])
 101     else:
 102         axes = pylab.axes()
 103         if title:
 104             axes.set_title(title)
 105     axes.set_axis_off()
 106
 107     if contours:
 108         cset = axes.contour(X[:-1,:-1], Y[:-1,:-1], Z, contours, cmap=cmap)
 109         # [:-1,:-1] to strip dummy last row & column from X&Y.
 110         pylab.clabel(cset, inline=1, fmt='%1.1f', fontsize=10)
 111     else:
 112         # pcolor() is much slower than imshow.
 113         #plot = axes.pcolor(X, Y, Z, cmap=cmap, edgecolors='none')
 114         #axes.autoscale_view(tight=True)
 115         plot = axes.imshow(Z, aspect='auto', interpolation='bilinear',
 116                            origin='lower', cmap=cmap,
 117                            extent=(X_min, X_max, Y_min, Y_max))
 118         if not full:
 119             fig.colorbar(plot)
 120     return fig
 121
 122
 123 def test():
 124     import doctest
 125     results = doctest.testmod()
 126     return results.failed
 127
 128
 129 def main(argv=None):
 130     """Read in data and plot it.
 131
 132     >>> from tempfile import NamedTemporaryFile
 133     >>> i = NamedTemporaryFile(prefix='tmp-input', suffix='.dat')
 134     >>> i.write('\\n'.join([str(x) for x in range(10)])+'\\n')
 135     >>> i.flush()
 136     >>> o = NamedTemporaryFile(prefix='tmp-output', suffix='.png')
 137     >>> main(['-i', i.name, '-s', '5,2', '-o', o.name, '-m', 'binary'])
 138     Plot_image
 139     Title:             Some like it hot
 140     Image size:        5 2
 141     False color
 142     Data range:        0.0 9.0
 143     >>> img = o.read()
 144     >>> img.startswith('\\x89PNG')
 145     True
 146     >>> i.close()
 147     >>> o.close()
 148     """
 149     if argv == None:
 150         argv = sys.argv[1:]
 151
 152     usage = '%prog [options]'
 153     epilog = _DOC
 154     p = optparse.OptionParser(usage=usage, epilog=epilog)
 155     p.format_epilog = lambda formatter: epilog+'\n'
 156
 157     p.add_option('-i', '--input', dest='input', metavar='FILE',
 158                  help='If set, read data from FILE rather than stdin.')
 159     p.add_option('-o', '--output', dest='output', metavar='FILE',
 160                  help=('If set, save the figure to FILE rather than '
 161                        'displaying it immediately'))
 162     p.add_option('-s', '--size', dest='size', default='%d,%d' % (16, 16),
 163                  help='Data size (columns,rows; default: %default)')
 164     p.add_option('-c', '--contours', dest='contours', type='int',
 165                  help=('Number of contour lines (if not set, draw false color '
 166                        'instead of contour lines; default: %default)'))
 167     p.add_option('-f', '--full-figure', dest='full', action='store_true',
 168                  help=('Set axes to fill the figure (i.e. no title or color '
 169                        'bar'))
 170     p.add_option('-t', '--title', dest='title', default='Some like it hot',
 171                  help='Title (%default)')
 172     p.add_option('--test', dest='test', action='store_true',
 173                  help='Run internal tests and exit.')
 174     maps=[m for m in pylab.cm.datad if not m.endswith("_r")]
 175     maps.sort()
 176     p.add_option('-m', '--color-map', dest='cmap', default='jet',
 177                  help='Select color map from %s (%%default)' % ', '.join(maps))
 178
 179     options,args = p.parse_args(argv)
 180
 181     if options.test:
 182         sys.exit(test())
 183
 184     nx,ny = [int(x) for x in options.size.split(',')]
 185     try:
 186         cmap = getattr(pylab.cm, options.cmap)
 187     except AttributeError:
 188         raise Exception('no color map named %s in %s'
 189                         % (options.cmap, ', '.join(maps)))
 190
 191     print 'Plot_image'
 192     print 'Title:            ', options.title
 193     print 'Image size:       ', nx, ny
 194     if options.contours:
 195         print '# countour lines: ', options.contours
 196     else:
 197         print 'False color'
 198     if options.input:
 199         fin = open(options.input, 'r')
 200     else:
 201         fin = sys.stdin
 202
 203     X,Y,Z = read_data(fin, nx, ny)
 204
 205     if options.input:
 206         fin.close()
 207
 208     Z_min = numpy.min(Z.flat)
 209     Z_max = numpy.max(Z.flat)
 210     print 'Data range:       ', Z_min, Z_max
 211
 212     fig = plot(X, Y, Z, full=options.full, title=options.title,
 213                contours=options.contours, cmap=cmap)
 214
 215     if options.output:
 216         fig.savefig(options.output)
 217     else:
 218         pylab.ion()
 219         pylab.show()
 220
 221
 222 if __name__ == '__main__':
 223     main()