Add --xyz option to plot_image.py for reading in (x,y,z) tuples.

[parallel_computing.git] / src / plot_image / plot_image.py

diff --git a/src/plot_image/plot_image.py b/src/plot_image/plot_image.py
index 4854eeadfd9aef2ce52a095ce18f36d1fb323fe9..a62cbc322c5233aacf45b4020bdc8021c27097fa 100755 (executable)
--- a/src/plot_image/plot_image.py
+++ b/src/plot_image/plot_image.py
@@ -1,4 +1,4 @@
-#! /usr/bin/env python
+#!/usr/bin/env python
 
 """Generate an image from ASCII data.
 
@@ -15,7 +15,7 @@ Data should be one ASCII float per line, in the following order::
     z[x1,y1]
     z[x2,y1]
     ...
-    z[nN,y1]
+    z[xN,y1]
     z[x1,y2]
     ...
     z[xN,y2]
@@ -25,7 +25,19 @@ Data should be one ASCII float per line, in the following order::
 where `x` increases from `x1` to `xN` and `y` increases from `y1`
 through `yN`.
 
-For  usage info, run::
+You can use the `--xyz` option to read an alternative data format,
+which is::
+
+    x1 y1 z[x1,y1]
+    x2 y2 z[x1,y1]
+    ...
+
+If you use this method, the ordering of lines in the data file is
+irrelevant, and `Nx` and `Ny` are extracted from the data file itself.
+However, you still need to use a rectangular grid (i.e. for every
+`xi`, you need to have entries for every `yi`).
+
+For usage info, run::
 
     ./plot_image.py --help
 
@@ -56,12 +68,12 @@ import pylab
 _DOC = __doc__
 
 
-def read_data(stream, nx, ny):
+def read_data_1d(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)
+    >>> s = StringIO('\\n'.join(map(str, range(10)))+'\\n')
+    >>> X,Y,Z = read_data_1d(s, 5, 2)
     >>> X
     array([[0, 1, 2, 3, 4, 5],
            [0, 1, 2, 3, 4, 5],
@@ -81,6 +93,46 @@ def read_data(stream, nx, ny):
     Z = Z.reshape([x-1 for x in X.shape])    
     return (X,Y,Z)
 
+def read_data_3d(stream):
+    """Read in data, one `(x, y, z)` tuple per line.
+
+    >>> from StringIO import StringIO
+    >>> lines = []
+    >>> for x in range(5):
+    ...     for y in range(2):
+    ...         lines.append('\t'.join(map(str, [x, y, x+y*5])))
+    >>> s = StringIO('\\n'.join(lines)+'\\n')
+    >>> X,Y,Z = read_data_3d(s)
+    >>> 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.]])
+    """
+    XYZ = numpy.loadtxt(stream)
+    assert len(XYZ.shape) == 2 and XYZ.shape[1] == 3, XYZ.shape()
+    Xs = numpy.array(sorted(set(XYZ[:,0])))
+    Ys = numpy.array(sorted(set(XYZ[:,1])))
+    Z = numpy.ndarray((len(Ys), len(Xs)), dtype=float)
+    xyz = {}  # dict of z values keyed by (x,y) tuples
+    for i in range(XYZ.shape[0]):
+        xyz[(XYZ[i,0], XYZ[i,1])] = XYZ[i,2]
+    for i,x in enumerate(Xs):
+        for j,y in enumerate(Ys):
+            Z[j,i] = xyz[x,y]
+    # add dummy row/column for pcolor
+    dx = Xs[-1] - Xs[-2]
+    dy = Ys[-1] - Ys[-2]
+    Xs = numpy.append(Xs, Xs[-1] + dx)
+    Ys = numpy.append(Ys, Ys[-1] + dy)
+    X,Y = pylab.meshgrid(Xs, Ys)
+    return (X,Y,Z)
 
 def plot(X, Y, Z, full=False, title=None, contours=None, cmap=None):
     """Plot Z over the mesh X, Y.
@@ -139,7 +191,9 @@ def main(argv=None):
     Title:             Some like it hot
     Image size:        5 2
     False color
-    Data range:        0.0 9.0
+    X range:           0 4
+    X range:           0 1
+    Z range:           0.0 9.0
     >>> img = o.read()
     >>> img.startswith('\\x89PNG')
     True
@@ -161,6 +215,10 @@ def main(argv=None):
                        '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('-3', '--xyz', dest='xyz', default=False, action='store_true',
+                 help=('If set, read (x,y,z) tuples from the input data rather'
+                       'then reading `z` and calculating `x` and `y` from '
+                       '`--size`.'))
     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)'))
@@ -190,7 +248,8 @@ def main(argv=None):
 
     print 'Plot_image'
     print 'Title:            ', options.title
-    print 'Image size:       ', nx, ny
+    if not options.xyz:
+        print 'Image size:       ', nx, ny
     if options.contours:
         print '# countour lines: ', options.contours
     else:
@@ -200,14 +259,19 @@ def main(argv=None):
     else:
         fin = sys.stdin
 
-    X,Y,Z = read_data(fin, nx, ny)
+    if options.xyz:
+        X,Y,Z = read_data_3d(fin)
+    else:
+        X,Y,Z = read_data_1d(fin, nx, ny)
 
     if options.input:
         fin.close()
 
     Z_min = numpy.min(Z.flat)
     Z_max = numpy.max(Z.flat)
-    print 'Data range:       ', Z_min, Z_max
+    print 'X range:          ', X[0,0], X[0,-2]
+    print 'X range:          ', Y[0,0], Y[-2,0]
+    print 'Z range:          ', Z_min, Z_max
 
     fig = plot(X, Y, Z, full=options.full, title=options.title,
                contours=options.contours, cmap=cmap)