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[xN,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 You can use the `--xyz` option to read an alternative data format,
  29 which is::
  30
  31     x1 y1 z[x1,y1]
  32     x2 y2 z[x1,y1]
  33     ...
  34
  35 If you use this method, the ordering of lines in the data file is
  36 irrelevant, and `Nx` and `Ny` are extracted from the data file itself.
  37 However, you still need to use a rectangular grid (i.e. for every
  38 `xi`, you need to have entries for every `yi`).
  39
  40 For usage info, run::
  41
  42     ./plot_image.py --help
  43
  44 When run with interactive output (i.e. no `-o ...` option), the
  45 interactive figure is displayed using `pylab.show`, which means that
  46 you'll have to kill `plot_image.py` using ^C or similar [1]_.
  47
  48 For other ideas, see the Matplotlib website [2]_.
  49
  50 .. [1] http://matplotlib.sourceforge.net/faq/howto_faq.html#use-show
  51 .. [2] http://matplotlib.sourceforge.net/
  52 """
  53
  54 import optparse
  55 import sys
  56
  57 import numpy
  58
  59 # Depending on your Matplotlib configuration, you may need to adjust
  60 # your backend.  Do this before importing pylab or matplotlib.backends.
  61 #import matplotlib
  62 #matplotlib.use('Agg')     # select backend that doesn't require X Windows
  63 #matplotlib.use('GTKAgg')  # select backend that supports pylab.show()
  64
  65 import pylab
  66
  67
  68 _DOC = __doc__
  69
  70
  71 def read_data_1d(stream, nx, ny):
  72     """Read in data, one entry per line.
  73
  74     >>> from StringIO import StringIO
  75     >>> s = StringIO('\\n'.join(map(str, range(10)))+'\\n')
  76     >>> X,Y,Z = read_data_1d(s, 5, 2)
  77     >>> X
  78     array([[0, 1, 2, 3, 4, 5],
  79            [0, 1, 2, 3, 4, 5],
  80            [0, 1, 2, 3, 4, 5]])
  81     >>> Y
  82     array([[0, 0, 0, 0, 0, 0],
  83            [1, 1, 1, 1, 1, 1],
  84            [2, 2, 2, 2, 2, 2]])
  85     >>> Z
  86     array([[ 0.,  1.,  2.,  3.,  4.],
  87            [ 5.,  6.,  7.,  8.,  9.]])
  88     """
  89     X,Y = pylab.meshgrid(range(nx+1), range(ny+1))
  90     Z = numpy.loadtxt(stream)
  91     assert Z.size == nx*ny, 'Z.size = %d != %d = %dx%d' % (
  92         Z.size, nx*ny, nx, ny)
  93     Z = Z.reshape([x-1 for x in X.shape])
  94     return (X,Y,Z)
  95
  96 def read_data_3d(stream):
  97     """Read in data, one `(x, y, z)` tuple per line.
  98
  99     >>> from StringIO import StringIO
 100     >>> lines = []
 101     >>> for x in range(5):
 102     ...     for y in range(2):
 103     ...         lines.append('\t'.join(map(str, [x, y, x+y*5])))
 104     >>> s = StringIO('\\n'.join(lines)+'\\n')
 105     >>> X,Y,Z = read_data_3d(s)
 106     >>> X
 107     array([[ 0.,  1.,  2.,  3.,  4.,  5.],
 108            [ 0.,  1.,  2.,  3.,  4.,  5.],
 109            [ 0.,  1.,  2.,  3.,  4.,  5.]])
 110     >>> Y
 111     array([[ 0.,  0.,  0.,  0.,  0.,  0.],
 112            [ 1.,  1.,  1.,  1.,  1.,  1.],
 113            [ 2.,  2.,  2.,  2.,  2.,  2.]])
 114     >>> Z
 115     array([[ 0.,  1.,  2.,  3.,  4.],
 116            [ 5.,  6.,  7.,  8.,  9.]])
 117     """
 118     XYZ = numpy.loadtxt(stream)
 119     assert len(XYZ.shape) == 2 and XYZ.shape[1] == 3, XYZ.shape()
 120     Xs = numpy.array(sorted(set(XYZ[:,0])))
 121     Ys = numpy.array(sorted(set(XYZ[:,1])))
 122     Z = numpy.ndarray((len(Ys), len(Xs)), dtype=float)
 123     xyz = {}  # dict of z values keyed by (x,y) tuples
 124     for i in range(XYZ.shape[0]):
 125         xyz[(XYZ[i,0], XYZ[i,1])] = XYZ[i,2]
 126     for i,x in enumerate(Xs):
 127         for j,y in enumerate(Ys):
 128             Z[j,i] = xyz[x,y]
 129     # add dummy row/column for pcolor
 130     dx = Xs[-1] - Xs[-2]
 131     dy = Ys[-1] - Ys[-2]
 132     Xs = numpy.append(Xs, Xs[-1] + dx)
 133     Ys = numpy.append(Ys, Ys[-1] + dy)
 134     X,Y = pylab.meshgrid(Xs, Ys)
 135     return (X,Y,Z)
 136
 137 def plot(X, Y, Z, full=False, title=None, contours=None, cmap=None):
 138     """Plot Z over the mesh X, Y.
 139
 140     >>> X, Y = pylab.meshgrid(range(6), range(2))
 141     >>> Z = X[:-1,:-1]**2 + Y[:-1,:-1]
 142     >>> plot(X, Y, Z)  # doctest: +ELLIPSIS
 143     <matplotlib.figure.Figure object at 0x...>
 144     """
 145     X_min = X[0,0]
 146     X_max = X[-1,-1]
 147     Y_min = Y[0,0]
 148     Y_max = Y[-1,-1]
 149
 150     fig = pylab.figure()
 151     if full:
 152         axes = pylab.axes([0, 0, 1, 1])
 153     else:
 154         axes = pylab.axes()
 155         if title:
 156             axes.set_title(title)
 157     axes.set_axis_off()
 158
 159     if contours:
 160         cset = axes.contour(X[:-1,:-1], Y[:-1,:-1], Z, contours, cmap=cmap)
 161         # [:-1,:-1] to strip dummy last row & column from X&Y.
 162         pylab.clabel(cset, inline=1, fmt='%1.1f', fontsize=10)
 163     else:
 164         # pcolor() is much slower than imshow.
 165         #plot = axes.pcolor(X, Y, Z, cmap=cmap, edgecolors='none')
 166         #axes.autoscale_view(tight=True)
 167         plot = axes.imshow(Z, aspect='auto', interpolation='bilinear',
 168                            origin='lower', cmap=cmap,
 169                            extent=(X_min, X_max, Y_min, Y_max))
 170         if not full:
 171             fig.colorbar(plot)
 172     return fig
 173
 174
 175 def test():
 176     import doctest
 177     results = doctest.testmod()
 178     return results.failed
 179
 180
 181 def main(argv=None):
 182     """Read in data and plot it.
 183
 184     >>> from tempfile import NamedTemporaryFile
 185     >>> i = NamedTemporaryFile(prefix='tmp-input', suffix='.dat')
 186     >>> i.write('\\n'.join([str(x) for x in range(10)])+'\\n')
 187     >>> i.flush()
 188     >>> o = NamedTemporaryFile(prefix='tmp-output', suffix='.png')
 189     >>> main(['-i', i.name, '-s', '5,2', '-o', o.name, '-m', 'binary'])
 190     Plot_image
 191     Title:             Some like it hot
 192     Image size:        5 2
 193     False color
 194     X range:           0 4
 195     X range:           0 1
 196     Z range:           0.0 9.0
 197     >>> img = o.read()
 198     >>> img.startswith('\\x89PNG')
 199     True
 200     >>> i.close()
 201     >>> o.close()
 202     """
 203     if argv == None:
 204         argv = sys.argv[1:]
 205
 206     usage = '%prog [options]'
 207     epilog = _DOC
 208     p = optparse.OptionParser(usage=usage, epilog=epilog)
 209     p.format_epilog = lambda formatter: epilog+'\n'
 210
 211     p.add_option('-i', '--input', dest='input', metavar='FILE',
 212                  help='If set, read data from FILE rather than stdin.')
 213     p.add_option('-o', '--output', dest='output', metavar='FILE',
 214                  help=('If set, save the figure to FILE rather than '
 215                        'displaying it immediately'))
 216     p.add_option('-s', '--size', dest='size', default='%d,%d' % (16, 16),
 217                  help='Data size (columns,rows; default: %default)')
 218     p.add_option('-3', '--xyz', dest='xyz', default=False, action='store_true',
 219                  help=('If set, read (x,y,z) tuples from the input data rather'
 220                        'then reading `z` and calculating `x` and `y` from '
 221                        '`--size`.'))
 222     p.add_option('-c', '--contours', dest='contours', type='int',
 223                  help=('Number of contour lines (if not set, draw false color '
 224                        'instead of contour lines; default: %default)'))
 225     p.add_option('-f', '--full-figure', dest='full', action='store_true',
 226                  help=('Set axes to fill the figure (i.e. no title or color '
 227                        'bar'))
 228     p.add_option('-t', '--title', dest='title', default='Some like it hot',
 229                  help='Title (%default)')
 230     p.add_option('--test', dest='test', action='store_true',
 231                  help='Run internal tests and exit.')
 232     maps=[m for m in pylab.cm.datad if not m.endswith("_r")]
 233     maps.sort()
 234     p.add_option('-m', '--color-map', dest='cmap', default='jet',
 235                  help='Select color map from %s (%%default)' % ', '.join(maps))
 236
 237     options,args = p.parse_args(argv)
 238
 239     if options.test:
 240         sys.exit(test())
 241
 242     nx,ny = [int(x) for x in options.size.split(',')]
 243     try:
 244         cmap = getattr(pylab.cm, options.cmap)
 245     except AttributeError:
 246         raise Exception('no color map named %s in %s'
 247                         % (options.cmap, ', '.join(maps)))
 248
 249     print 'Plot_image'
 250     print 'Title:            ', options.title
 251     if not options.xyz:
 252         print 'Image size:       ', nx, ny
 253     if options.contours:
 254         print '# countour lines: ', options.contours
 255     else:
 256         print 'False color'
 257     if options.input:
 258         fin = open(options.input, 'r')
 259     else:
 260         fin = sys.stdin
 261
 262     if options.xyz:
 263         X,Y,Z = read_data_3d(fin)
 264     else:
 265         X,Y,Z = read_data_1d(fin, nx, ny)
 266
 267     if options.input:
 268         fin.close()
 269
 270     Z_min = numpy.min(Z.flat)
 271     Z_max = numpy.max(Z.flat)
 272     print 'X range:          ', X[0,0], X[0,-2]
 273     print 'X range:          ', Y[0,0], Y[-2,0]
 274     print 'Z range:          ', Z_min, Z_max
 275
 276     fig = plot(X, Y, Z, full=options.full, title=options.title,
 277                contours=options.contours, cmap=cmap)
 278
 279     if options.output:
 280         fig.savefig(options.output)
 281     else:
 282         pylab.ion()
 283         pylab.show()
 284
 285
 286 if __name__ == '__main__':
 287     main()