1#!/usr/local/bin/python -t
2# $Id$
3
4# Gnuplot.py -- A pipe-based interface to the gnuplot plotting program.
5
6# Copyright (C) 1998 Michael Haggerty <mhagger@blizzard.harvard.edu>.
7
8"""A pipe-based interface to the gnuplot plotting program.
9
10This program is free software; you can redistribute it and/or modify
11it under the terms of the GNU General Public License as published by
12the Free Software Foundation; either version 2 of the License, or (at
13your option) any later version.  This program is distributed in the
14hope that it will be useful, but WITHOUT ANY WARRANTY; without even
15the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
16PURPOSE.  See the GNU General Public License for more details; it is
17available at <http://www.fsf.org/copyleft/gpl.html>, or by writing to
18the Free Software Foundation, Inc., 59 Temple Place - Suite 330,
19Boston, MA 02111-1307, USA.
20
21Written by Michael Haggerty <mhagger@blizzard.harvard.edu>.  Inspired
22by and partly derived from an earlier version by Konrad Hinsen
23<hinsen@ibs.ibs.fr>.  If you find a problem or have a suggestion,
24please let me know at <mhagger@blizzard.harvard.edu>.  Other feedback
25is also welcome.
26
27For information about how to use this module, see the comments below,
28the documentation string for class Gnuplot, and the test code at the
29bottom of the file.  You can run the test code by typing
30'python Gnuplot.py'.
31
32You should import this file with 'import Gnuplot', not with
33'from Gnuplot import *'; otherwise you will have problems with
34conflicting names (specifically, the Gnuplot module name conflicts
35with the Gnuplot class name).  To obtain gnuplot itself, see
36<http://www.cs.dartmouth.edu/gnuplot_info.html>.
37
38Features:
39
40 o  Allows the creation of two or three dimensional plots from
41    python by piping commands to the 'gnuplot' program.
42 o  A gnuplot session is an instance of class 'Gnuplot', so multiple
43    sessions can be open at once:
44        'g1 = Gnuplot.Gnuplot(); g2 = Gnuplot.Gnuplot()'
45 o  The implicitly-generated gnuplot commands can be stored to a file
46    instead of executed immediately:
47        'g = Gnuplot.Gnuplot("commands.gnuplot")'
48    The file can then be run later with gnuplot's 'load' command.
49    Beware, however, if the plot commands depend on the existence of
50    temporary files, because they might be deleted before you use
51    the command file.
52 o  Can pass arbitrary commands to the gnuplot command interpreter:
53        'g("set pointsize 2")'
54 o  A Gnuplot object knows how to plot objects of type 'PlotItem'.
55    Any PlotItem can have optional `title' and/or 'with' suboptions.
56    Builtin PlotItem types:
57
58    * 'Data(array1)' -- data from a Python list or NumPy array
59                      (permits additional option 'cols')
60    * 'File("filename")' -- data from an existing data file (permits
61                      additional option 'using')
62    * 'Func("exp(4.0 * sin(x))")' -- functions (passed as a string
63                      for gnuplot to evaluate)
64    * 'GridData(m, x, y)' -- data tabulated on a grid of (x,y) values
65                      (usually to be plotted in 3-D)
66
67    See those classes for more details.
68
69 o PlotItems are implemented as objects that can be assigned to
70    variables (including their options) and plotted repeatedly ---
71    this also saves much of the overhead of plotting the same data
72    multiple times.
73 o  Communication of data between python and gnuplot is via temporary
74    files, which are deleted automatically when their associated
75    'PlotItem' is deleted.  (Communication of commands is via a pipe.)
76    The PlotItems currently in use by a Gnuplot object are stored in
77    an internal list so that they won't be deleted prematurely.
78 o  Can use 'replot' method to add datasets to an existing plot.
79 o  Can make persistent gnuplot windows by using the constructor option
80    `persist=1'.  Such windows stay around even after the gnuplot
81    program is exited.  Note that only newer version of gnuplot support
82    this option.
83 o  Plotting to a postscript file is via new 'hardcopy' method, which
84    outputs the currently-displayed plot to either a postscript
85    printer or to a postscript file.
86 o  There is a 'plot' command which is roughly compatible with the
87    command from Konrad Hinsen's old 'Gnuplot.py'.
88
89Restrictions:
90
91 -  Relies on the numpy Python extension.  This can be obtained
92    from LLNL (See ftp://ftp-icf.llnl.gov/pub/python/README.html).
93    If you're interested in gnuplot, you would probably also want
94    NumPy anyway.
95 -  Probably depends on a unix-type environment.  Anyone who wants
96    to remedy this situation should contact me.
97 -  Only a small fraction of gnuplot functionality is implemented as
98    explicit method functions.  However, you can give arbitrary
99    commands to gnuplot manually; for example:
100        'g = Gnuplot.Gnuplot()',
101        'g('set style data linespoints')',
102        'g('set pointsize 5')',
103    etc.  I might add a more organized way of setting arbitrary
104    options.
105 -  There is no provision for missing data points in array data
106    (which gnuplot would allow by specifying '?' as a data point).
107    I can't think of a clean way to implement this; maybe one could
108    use NaN for machines that support IEEE floating point.
109 -  There is no supported way to change the plotting options of
110    PlotItems after they have been created.
111
112Bugs:
113
114 -  No attempt is made to check for errors reported by gnuplot (but
115    they will appear on stderr).
116 -  All of these classes perform their resource deallocation when
117    '__del__' is called.  If you delete things explicitly, there will
118    be no problem.  If you don't, an attempt is made to delete
119    remaining objects when the interpreter is exited, but this is
120    not completely reliable, so sometimes temporary files will be
121    left around.  If anybody knows how to fix this problem, please
122    let me know.
123"""
124
125__version__ = "1.1a"
126__cvs_version__ = "CVS version $Revision: 1.1 $"
127
128import sys, os, string, tempfile, numpy
129
130
131# Set after first call of test_persist().  This will be set from None
132# to 0 or 1 upon the first call of test_persist(), then the stored
133# value will be used thereafter.  To avoid the test, type 1 or 0 on
134# the following line corresponding to whether your gnuplot is new
135# enough to understand the -persist option.
136_recognizes_persist = None
137
138# After a hardcopy is produced, we have to set the terminal type back
139# to `on screen'.  If you are using unix, then `x11' is probably
140# correct.  If not, change the following line to the terminal type you
141# use.
142_default_term = 'x11'
143
144# Gnuplot can plot to a printer by using `set output "| ..."' where
145# ... is the name of a program that sends its stdin to a printer.  On
146# my machine the appropriate program is `lpr', as set below.  On your
147# computer it may be something different (like `lp'); you can set that
148# by changing the variable below.  You can also use the following
149# variable to add options to the print command.
150_default_lpr = '| lpr'
151
152def test_persist():
153    """Test and report whether gnuplot recognizes the option '-persist'.
154
155    Test if gnuplot is new enough to know the option '-persist'.  It it
156    isn't, it will emit an error message with '-persist' in the first
157    line.
158
159    """
160
161    global _recognizes_persist
162    if _recognizes_persist is None:
163        g = os.popen('echo | gnuplot -persist 2>&1', 'r')
164        response = g.readlines()
165        g.close()
166        _recognizes_persist = ((not response)
167                               or (string.find(response[0], '-persist') == -1))
168    return _recognizes_persist
169
170
171class OptionException(Exception):
172    """raised for unrecognized option(s)"""
173    pass
174
175class DataException(Exception):
176    """raised for data in the wrong format"""
177    pass
178
179
180class PlotItem:
181    """Plotitem represents an item that can be plotted by gnuplot.
182
183    For the finest control over the output, you can create the
184    PlotItems yourself with additional keyword options, or derive new
185    classes from PlotItem.
186
187    Members:
188
189    'basecommand' -- a string holding the elementary argument that
190                     must be passed to gnuplot's `plot' command for
191                     this item; e.g., 'sin(x)' or '"filename.dat"'.
192    'options' -- a list of strings that need to be passed as options
193                 to the plot command, in the order required; e.g.,
194                 ['title "data"', 'with linespoints'].
195    'title' -- the title requested (undefined if not requested).  Note
196               that `title=None' implies the `notitle' option,
197               whereas omitting the title option implies no option
198               (the gnuplot default is then used).
199    'with' -- the string requested as a `with' option (undefined if
200              not requested)
201
202    """
203
204    def __init__(self, basecommand, **keyw):
205        self.basecommand = basecommand
206        self.options = []
207        if keyw.has_key('title'):
208            self.title = keyw['title']
209            del keyw['title']
210            if self.title is None:
211                self.options.append('notitle')
212            else:
213                self.options.append('title "' + self.title + '"')
214        if keyw.has_key('with_'):
215            self.with_ = keyw['with_']
216            del keyw['with_']
217            self.options.append('with ' + self.with_)
218        if keyw:
219            raise OptionException(keyw)
220
221    def command(self):
222        """Build the 'plot' command to be sent to gnuplot.
223
224        Build and return the 'plot' command, with options, necessary
225        to display this item.
226
227        """
228
229        if self.options:
230            return self.basecommand + ' ' + string.join(self.options)
231        else:
232            return self.basecommand
233
234    # if the plot command requires data to be put on stdin (i.e.,
235    # `plot "-"'), this method should put that data there.
236    def pipein(self, file):
237        pass
238
239
240class Func(PlotItem):
241    """Represents a mathematical expression to plot.
242
243    Func represents a mathematical expression that is to be computed by
244    gnuplot itself, as in the example
245
246        gnuplot> plot sin(x)
247
248    The argument to the constructor is a string which is a expression.
249    Example:
250
251        g.plot(Func("sin(x)", with_="line 3"))
252
253    or the shorthand example:
254
255        g.plot("sin(x)")
256
257    """
258
259    def __init__(self, funcstring, **keyw):
260        apply(PlotItem.__init__, (self, funcstring), keyw)
261
262
263class AnyFile:
264    """An AnyFile represents any kind of file to be used by gnuplot.
265
266    An AnyFile represents a file, but presumably one that holds data
267    in a format readable by gnuplot.  This class simply remembers the
268    filename; the existence and format of the file are not checked
269    whatsoever.  Note that this is not a PlotItem, though it is used by
270    the 'File' PlotItem.  Members:
271
272    'self.filename' -- the filename of the file
273
274    """
275
276    def __init__(self, filename):
277        self.filename = filename
278
279
280class TempFile(AnyFile):
281    """A TempFile is a file that is automatically deleted.
282
283    A TempFile points to a file.  The file is deleted automatically
284    when the TempFile object is deleted.
285
286    WARNING: whatever filename you pass to this constructor **WILL BE
287    DELETED** when the TempFile object is deleted, even if it was a
288    pre-existing file! This is intended to be used as a parent class of
289    TempArrayFile.
290
291    """
292
293    def __del__(self):
294        os.unlink(self.filename)
295
296
297def write_array(f, set,
298                item_sep=' ',
299                nest_prefix='', nest_suffix='\n', nest_sep=''):
300    """Write an array of arbitrary dimension to a file.
301
302    A general recursive array writer.  The last four parameters allow a
303    great deal of freedom in choosing the output format of the
304    array.  The defaults for those parameters give output that is
305    gnuplot-readable.  But using, for example, ( ',', '{', '}', ',\\n')
306    would output an array in a format that Mathematica could read.
307    item_sep should not contain '%' (or if it does, it should be
308    escaped to '%%') since item_sep is put into a format string.
309
310    """
311
312    if len(set.shape) == 1:
313        (columns,) = set.shape
314        assert columns > 0
315        fmt = string.join(['%s'] * columns, item_sep)
316        f.write(nest_prefix)
317        f.write(fmt % tuple(set.tolist()))
318        f.write(nest_suffix)
319    elif len(set.shape) == 2:
320        # This case could be done with recursion, but `unroll' for
321        # efficiency.
322        (points, columns) = set.shape
323        assert points > 0
324        assert columns > 0
325        fmt = string.join(['%s'] * columns, item_sep)
326        f.write(nest_prefix + nest_prefix)
327        f.write(fmt % tuple(set[0].tolist()))
328        f.write(nest_suffix)
329        for point in set[1:]:
330            f.write(nest_sep + nest_prefix)
331            f.write(fmt % tuple(point.tolist()))
332            f.write(nest_suffix)
333        f.write(nest_suffix)
334    else:
335        assert set.shape[0] > 0
336        f.write(nest_prefix)
337        write_array(f, set[0], item_sep, nest_prefix, nest_suffix, nest_sep)
338        for subset in set[1:]:
339            f.write(nest_sep)
340            write_array(f, subset, item_sep, nest_prefix, nest_suffix, nest_sep)
341        f.write(nest_suffix)
342
343
344class ArrayFile(AnyFile):
345    """A file to which, upon creation, an array is written.
346
347    When an ArrayFile is constructed, it creates a file and fills it
348    with the contents of a 2-d or 3-d numpy array in the format
349    expected by gnuplot.  Specifically, for 2-d, the file organization
350    is for example:
351
352        set[0,0] set[0,1] ...
353        set[1,0] set[1,1] ...
354
355    etc.  For 3-d, it is for example:
356
357        set[0,0,0] set[0,0,1] ...
358        set[0,1,0] set[0,1,1] ...
359
360        set[1,0,0] set[1,0,1] ...
361        set[1,1,0] set[1,1,1] ...
362
363    etc.
364
365    The filename can be specified, otherwise a random filename is
366    chosen.  The file is NOT deleted automatically.
367
368    """
369
370    def __init__(self, set, filename=None):
371        if not filename:
372            filename = tempfile.mktemp()
373        f = open(filename, 'w')
374        write_array(f, set)
375        f.close()
376        AnyFile.__init__(self, filename)
377
378
379class TempArrayFile(ArrayFile, TempFile):
380    """An ArrayFile that is deleted automatically."""
381
382    def __init__(self, set, filename=None):
383        ArrayFile.__init__(self, set, filename)
384
385
386class File(PlotItem):
387    """A PlotItem representing a file that contains gnuplot data.
388
389    File is a PlotItem that represents a file that should be plotted
390    by gnuplot.  The file can either be a string holding the filename
391    of an existing file, or it can be anything derived from 'AnyFile'.
392
393    """
394
395    def __init__(self, file, using=None, **keyw):
396        """Construct a File object.
397
398        '<file>' can be either a string holding the filename of an
399        existing file, or it can be an object of a class derived from
400        'AnyFile' (such as a 'TempArrayFile').  Keyword arguments
401        recognized (in addition to those recognized by 'PlotItem'):
402
403            'using=<n>' -- plot that column against line number
404            'using=<tuple>' -- plot using a:b:c:d etc.
405            'using=<string>' -- plot `using <string>' (allows gnuplot's
406                              arbitrary column arithmetic)
407
408        Note that the 'using' option is interpreted by gnuplot, so
409        columns must be numbered starting with 1.  Other keyword
410        arguments are passed along to PlotItem.  The default 'title'
411        for an AnyFile PlotItem is 'notitle'.
412
413        """
414
415        if isinstance(file, AnyFile):
416            self.file = file
417            # If no title is specified, then use `notitle' for
418            # TempFiles (to avoid using the temporary filename as the
419            # title.)
420            if isinstance(file, TempFile) and not keyw.has_key('title'):
421                keyw['title'] = None
422        elif type(file) == type(""):
423            self.file = AnyFile(file)
424        else:
425            raise OptionException
426        apply(PlotItem.__init__, (self, '"' + self.file.filename + '"'), keyw)
427        self.using = using
428        if self.using is None:
429            pass
430        elif type(self.using) == type(""):
431            self.options.insert(0, "using " + self.using)
432        elif type(self.using) == type(()):
433            self.options.insert(0,
434                                "using " +
435                                string.join(map(repr, self.using), ':'))
436        elif type(self.using) == type(1):
437            self.options.insert(0, "using " + repr(self.using))
438        else:
439            raise OptionException('using=' + repr(self.using))
440
441
442class Data(File):
443    """Allows data from memory to be plotted with Gnuplot.
444
445    Takes a numeric array from memory and outputs it to a temporary
446    file that can be plotted by gnuplot.
447
448    """
449
450    def __init__(self, *set, **keyw):
451        """Construct a Data object from a numeric array.
452
453        Create a Data object (which is a type of PlotItem) out of one
454        or more Float Python numpy arrays (or objects that can be
455        converted to a Float numpy array).  If the routine is passed
456        one array, the last index ranges over the values comprising a
457        single data point (e.g., [x, y, and sigma]) and the rest of
458        the indices select the data point.  If the routine is passed
459        more than one array, they must have identical shapes, and then
460        each data point is composed of one point from each array.
461        I.e., 'Data(x,x**2)' is a PlotItem that represents x squared
462        as a function of x.  For the output format, see the comments
463        in ArrayFile.
464
465        The array is first written to a temporary file, then that file
466        is plotted.  Keyword arguments recognized (in addition to those
467        recognized by PlotItem):
468
469            cols=<tuple> -- write only the specified columns from each
470                            data point to the file.  Since cols is
471                            used by python, the columns should be
472                            numbered in the python style (starting
473                            from 0), not the gnuplot style (starting
474                            from 1).
475
476        The data are immediately written to the temp file; no copy is
477        kept in memory.
478
479        """
480
481        if len(set) == 1:
482            # set was passed as a single structure
483            set = numpy.asarray(set, dtype=numpy.float32)
484        else:
485            # set was passed column by column (for example, Data(x,y))
486            set = numpy.asarray(set, dtype=numpy.float32)
487            dims = len(set.shape)
488            # transpose so that the last index selects x vs. y:
489            set = numpy.transpose(set, (dims-1,) + tuple(range(dims-1)))
490        if keyw.has_key('cols') and keyw['cols'] is not None:
491            set = numpy.take(set, keyw['cols'], -1)
492            del keyw['cols']
493        apply(File.__init__, (self, TempArrayFile(set)), keyw)
494
495
496class GridData(File):
497    """Holds data representing a function of two variables, for use in splot.
498
499    GridData represents a function that has been tabulated on a
500    rectangular grid.  It is a PlotItem, so GridData objects can be
501    plotted by Gnuplot.  The data are written to a file but not stored
502    in memory.
503
504    """
505
506    def __init__(self, data, xvals=None, yvals=None, **keyw):
507        """GridData constructor.
508
509        Arguments:
510
511            'data' -- a 2-d array with dimensions (numx,numy)
512            'xvals' -- a 1-d array with dimension (numx)
513            'yvals' -- a 1-d array with dimension (numy)
514
515        'data' is meant to hold the values of a function f(x,y) tabulated
516        on a grid of points, such that 'data[i,j] == f(xvals[i],
517        yvals[j])'.  These data are written to a datafile as 'x y f(x,y)'
518        triplets that can be used by gnuplot's splot command.  Thus if you
519        have three arrays in the above format and a Gnuplot instance
520        called g, you can plot your data by typing for example:
521
522            g.splot(Gnuplot.GridData(data,xvals,yvals))
523
524        If 'xvals' and/or 'yvals' are omitted, integers (starting with
525        0) are used for that coordinate.  The data are written to a
526        temporary file; no copy of the data is kept in memory.
527
528        """
529
530        data = numpy.asarray(data, dtype=numpy.float32)
531        assert len(data.shape) == 2
532        (numx, numy) = data.shape
533
534        if xvals is None:
535            xvals = numpy.arange(numx)
536        else:
537            xvals = numpy.asarray(xvals, dtype=numpy.float32)
538            assert len(xvals.shape) == 1
539            assert xvals.shape[0] == numx
540
541        if yvals is None:
542            yvals = numpy.arange(numy)
543        else:
544            yvals = numpy.asarray(yvals, dtype=numpy.float32)
545            assert len(yvals.shape) == 1
546            assert yvals.shape[0] == numy
547
548        set = numpy.transpose(
549            numpy.array(
550                (numpy.transpose(numpy.resize(xvals, (numy, numx))),
551                 numpy.resize(yvals, (numx, numy)),
552                 data)), (1,2,0))
553
554        apply(File.__init__, (self, TempArrayFile(set)), keyw)
555
556
557def grid_function(f, xvals, yvals):
558    """Compute a function on a grid.
559
560    'xvals' and 'yvals' should be 1-D arrays listing the values of x
561    and y at which f should be tabulated.  f should be a function
562    taking two floating point arguments.  The return value is a matrix
563    M where M[i,j] = f(xvals[i],yvals[j]), which can for example be
564    used in the 'GridData' constructor.
565
566    Note that f is evaluated at each pair of points using a Python loop,
567    which can be slow if the number of points is large.  If speed is an
568    issue, you are better off computing functions matrix-wise using
569    numpy's built-in ufuncs.
570
571    """
572
573    m = numpy.zeros((len(xvals), len(yvals)), dtype=numpy.float32)
574    for xi in range(len(xvals)):
575        x = xvals[xi]
576        for yi in range(len(yvals)):
577            y = yvals[yi]
578            m[xi,yi] = f(x,y)
579    return m
580
581
582class Gnuplot:
583    """gnuplot plotting object.
584
585    A Gnuplot represents a running gnuplot program and a pipe to
586    communicate with it.  It keeps a reference to each of the
587    PlotItems used in the current plot, so that they (and their
588    associated temporary files) are not deleted prematurely.  The
589    communication is one-way; gnuplot's text output just goes to
590    stdout with no attempt to check it for error messages.
591
592    Members:
593
594    'gnuplot' -- the pipe to gnuplot or a file gathering the commands
595    'itemlist' -- a list of the PlotItems that are associated with the
596                  current plot.  These are deleted whenever a new plot
597                  command is issued via the `plot' method.
598    'debug' -- if this flag is set, commands sent to gnuplot will also
599               be echoed to stderr.
600    'plotcmd' -- 'plot' or 'splot', depending on what was the last
601                 plot command.
602
603    Methods:
604
605    '__init__' -- if a filename argument is specified, the commands
606                  will be written to that file instead of being piped
607                  to gnuplot immediately.
608    'plot' -- clear the old plot and old PlotItems, then plot the
609              arguments in a fresh plot command.  Arguments can be: a
610              PlotItem, which is plotted along with its internal
611              options; a string, which is plotted as a Func; or
612              anything else, which is plotted as a Data.
613    'hardcopy' -- replot the plot to a postscript file (if filename
614                  argument is specified) or pipe it to lpr otherwise.
615                  If the option `color' is set to true, then output
616                  color postscript.
617    'replot' -- replot the old items, adding any arguments as
618                additional items as in the plot method.
619    'refresh' -- issue (or reissue) the plot command using the current
620                 PlotItems.
621    '__call__' -- pass an arbitrary string to the gnuplot process,
622                  followed by a newline.
623    'xlabel', 'ylabel', 'title' -- set attribute to be a string.
624    'interact' -- read lines from stdin and send them, one by one, to
625                  the gnuplot interpreter.  Basically you can type
626                  commands directly to the gnuplot command processor
627                  (though without command-line editing).
628    'load' -- load a file (using the gnuplot `load' command).
629    'save' -- save gnuplot commands to a file (using gnuplot `save'
630              command) If any of the PlotItems is a temporary file, it
631              will be deleted at the usual time and the save file might
632              be pretty useless :-).
633    'clear' -- clear the plot window (but not the itemlist).
634    'reset' -- reset all gnuplot settings to their defaults and clear
635               the current itemlist.
636    'set_string' -- set or unset a gnuplot option whose value is a
637                    string.
638    '_clear_queue' -- clear the current PlotItem list.
639    '_add_to_queue' -- add the specified items to the current
640                       PlotItem list.
641
642    """
643
644    def __init__(self, filename=None, persist=0, debug=0):
645        """Create a Gnuplot object.
646
647        'Gnuplot(filename=None, persist=0, debug=0)':
648
649        Create a 'Gnuplot' object.  By default, this starts a gnuplot
650        process and prepares to write commands to it.  If 'filename'
651        is specified, the commands are instead written to that file
652        (i.e., for later use using 'load').  If 'persist' is set,
653        gnuplot will be started with the '-persist' option (which
654        creates a new X11 plot window for each plot command).  (This
655        option is not available on older versions of gnuplot.)  If
656        'debug' is set, the gnuplot commands are echoed to stderr as
657        well as being send to gnuplot.
658
659        """
660
661        if filename:
662            # put gnuplot commands into a file:
663            self.gnuplot = open(filename, 'w')
664        else:
665            if persist:
666                if not test_persist():
667                    raise OptionException(
668                        '-persist does not seem to be supported '
669                        'by your version of gnuplot!')
670                self.gnuplot = os.popen('gnuplot -persist', 'w')
671            else:
672                self.gnuplot = os.popen('gnuplot', 'w')
673        self._clear_queue()
674        self.debug = debug
675        self.plotcmd = 'plot'
676
677    def __del__(self):
678        self('quit')
679        self.gnuplot.close()
680
681    def __call__(self, s):
682        """Send a command string to gnuplot.
683
684        '__call__(s)': send the string s as a command to gnuplot,
685        followed by a newline and flush.  All interaction with the
686        gnuplot process is through this method.
687
688        """
689
690        self.gnuplot.write(s + "\n")
691        self.gnuplot.flush()
692        if self.debug:
693            # also echo to stderr for user to see:
694            sys.stderr.write("gnuplot> %s\n" % (s,))
695
696    def refresh(self):
697        """Refresh the plot, using the current PlotItems.
698
699        Refresh the current plot by reissuing the gnuplot plot command
700        corresponding to the current itemlist.
701
702        """
703
704        plotcmds = []
705        for item in self.itemlist:
706            plotcmds.append(item.command())
707        self(self.plotcmd + ' ' + string.join(plotcmds, ', '))
708        for item in self.itemlist:
709            item.pipein(self.gnuplot)
710
711    def _clear_queue(self):
712        """Clear the PlotItems from the queue."""
713
714        self.itemlist = []
715
716    def _add_to_queue(self, items):
717        """Add a list of items to the itemlist, but don't plot them.
718
719        'items' is a list or tuple of items, each of which should be a
720        'PlotItem' of some kind, a string (interpreted as a function
721        string for gnuplot to evaluate), or a numpy array (or
722        something that can be converted to a numpy array).
723
724        """
725
726        for item in items:
727            if isinstance(item, PlotItem):
728                self.itemlist.append(item)
729            elif type(item) is type(""):
730                self.itemlist.append(Func(item))
731            else:
732                # assume data is an array:
733                self.itemlist.append(Data(item))
734
735    def plot(self, *items):
736        """Draw a new plot.
737
738        'plot(item, ...)': Clear the current plot and create a new 2-d
739        plot containing the specified items.  Arguments can be of the
740        following types:
741
742        'PlotItem' (e.g., 'Data', 'File', 'Func', 'GridData') -- This
743                   is the most flexible way to call plot because the
744                   PlotItems can contain suboptions.  Moreover,
745                   PlotItems can be saved to variables so that their
746                   lifetime is longer than one plot command--thus they
747                   can be replotted with minimal overhead.
748
749        'string' (i.e., "sin(x)") -- The string is interpreted as
750                 'Func(string)' (a function that is computed by
751                 gnuplot).
752
753        Anything else -- The object, which should be convertible to an
754                         array, is converted to a Data() item, and
755                         thus plotted as data.  If the conversion
756                         fails, an exception is raised.
757
758        """
759
760        # remove old files:
761        self.plotcmd = 'plot'
762        self._clear_queue()
763        self._add_to_queue(items)
764        self.refresh()
765
766    def splot(self, *items):
767        """Draw a new three-dimensional plot.
768
769        'splot(item, ...)' -- Clear the current plot and create a new
770                3-d plot containing the specified items.  Arguments can
771                be of the following types:
772        'PlotItem' (e.g., 'Data', 'File', 'Func', 'GridData') -- This
773                is the most flexible way to call plot because the
774                PlotItems can contain suboptions.  Moreover, PlotItems
775                can be saved to variables so that their lifetime is
776                longer than one plot command--thus they can be
777                replotted with minimal overhead.
778
779        'string' (i.e., "sin(x*y)") -- The string is interpreted as a
780                'Func()' (a function that is computed by gnuplot).
781
782        Anything else -- The object is converted to a Data() item, and
783                thus plotted as data.  Note that each data point
784                should normally have at least three values associated
785                with it (i.e., x, y, and z).  If the conversion fails,
786                an exception is raised.
787
788        """
789
790        # remove old files:
791        self.plotcmd = 'splot'
792        self._clear_queue()
793        self._add_to_queue(items)
794        self.refresh()
795
796    def replot(self, *items):
797        """Replot the data, possibly adding new PlotItems.
798
799        Replot the existing graph, using the items in the current
800        itemlist.  If arguments are specified, they are interpreted as
801        additional items to be plotted alongside the existing items on
802        the same graph.  See 'plot' for details.
803
804        """
805
806        self._add_to_queue(items)
807        self.refresh()
808
809    def interact(self):
810        """Allow user to type arbitrary commands to gnuplot.
811
812        Read stdin, line by line, and send each line as a command to
813        gnuplot.  End by typing C-d.
814
815        """
816
817        sys.stderr.write("Press C-d to end interactive input\n")
818        while 1:
819            sys.stderr.write("gnuplot>>> ")
820            line = sys.stdin.readline()
821            if not line:
822                break
823            if line[-1] == "\n": line = line[:-1]
824            self(line)
825
826    def clear(self):
827        """Clear the plot window (without affecting the current itemlist)."""
828
829        self('clear')
830
831    def reset(self):
832        """Reset all gnuplot settings to their defaults and clear itemlist."""
833
834        self('reset')
835        self.itemlist = []
836
837    def load(self, filename):
838        """Load a file using gnuplot's `load' command."""
839
840        self('load "%s"' % (filename,))
841
842    def save(self, filename):
843        """Save the current plot commands using gnuplot's `save' command."""
844
845        self('save "%s"' % (filename,))
846
847    def set_string(self, option, s=None):
848        """Set a string option, or if s is omitted, unset the option."""
849
850        if s is None:
851            self('set %s' % (option,))
852        else:
853            self('set %s "%s"' % (option, s))
854
855    def xlabel(self, s=None):
856        """Set the plot's xlabel."""
857
858        self.set_string('xlabel', s)
859
860    def ylabel(self, s=None):
861        """Set the plot's ylabel."""
862
863        self.set_string('ylabel', s)
864
865    def title(self, s=None):
866        """Set the plot's title."""
867
868        self.set_string('title', s)
869
870    def hardcopy(self, filename=None, eps=0, color=0, enhanced=1):
871        """Create a hardcopy of the current plot.
872
873        Create a postscript hardcopy of the current plot.  If a
874        filename is specified, save the output in that file; otherwise
875        print it immediately using lpr.  If eps is specified, generate
876        encapsulated postscript.  If color is specified, create a
877        color plot.  If enhanced is specified (the default), then
878        generate enhanced postscript.  (Some old gnuplot versions do
879        not support enhanced postscript; if this is the case set
880        enhanced=0.)  Note that this command will return immediately
881        even though it might take gnuplot a while to actually finish
882        working.
883
884        """
885
886        if filename is None:
887            filename = _default_lpr
888        setterm = ['set', 'term', 'postscript']
889        if eps: setterm.append('eps')
890        else: setterm.append(' ')
891        if enhanced: setterm.append('enhanced')
892        if color: setterm.append('color')
893        self(string.join(setterm))
894        self.set_string('output', filename)
895        self.refresh()
896        self('set term %s' % _default_term)
897        self.set_string('output')
898
899
900# The following is a command defined for compatibility with Hinson's
901# old Gnuplot.py module.  Its use is deprecated.
902
903# When the plot command is called and persist is not available, the
904# plotters will be stored here to prevent their being closed:
905_gnuplot_processes = []
906
907def plot(*items, **keyw):
908    """plot data using gnuplot through Gnuplot.
909
910    This command is roughly compatible with old Gnuplot plot command.
911    It is provided for backwards compatibility with the old functional
912    interface only.  It is recommended that you use the new
913    object-oriented Gnuplot interface, which is much more flexible.
914
915    It can only plot numpy array data.  In this routine an NxM array
916    is plotted as M-1 separate datasets, using columns 1:2, 1:3, ...,
917    1:M.
918
919    Limitations:
920
921        - If persist is not available, the temporary files are not
922          deleted until final python cleanup.
923
924    """
925
926    newitems = []
927    for item in items:
928        # assume data is an array:
929        item = numpy.asarray(item, dtype=numpy.float32)
930        dim = len(item.shape)
931        if dim == 1:
932            newitems.append(Data(item[:, numpy.newaxis], with_='lines'))
933        elif dim == 2:
934            if item.shape[1] == 1:
935                # one column; just store one item for tempfile:
936                newitems.append(Data(item, with_='lines'))
937            else:
938                # more than one column; store item for each 1:2, 1:3, etc.
939                tempf = TempArrayFile(item)
940                for col in range(1, item.shape[1]):
941                    newitems.append(File(tempf, using=(1,col+1), with_='lines'))
942        else:
943            raise DataException("Data array must be 1 or 2 dimensional")
944    items = tuple(newitems)
945    del newitems
946
947    if keyw.has_key('file'):
948        g = Gnuplot()
949        # setup plot without actually plotting (so data don't appear
950        # on the screen):
951        g._add_to_queue(items)
952        g.hardcopy(keyw['file'])
953        # process will be closed automatically
954    elif test_persist():
955        g = Gnuplot(persist=1)
956        apply(g.plot, items)
957        # process will be closed automatically
958    else:
959        g = Gnuplot()
960        apply(g.plot, items)
961        # prevent process from being deleted:
962        _gnuplot_processes.append(g)
963
964
965# Demo code
966if __name__ == '__main__':
967    from numpy import *
968    import sys
969
970    # A straightforward use of gnuplot.  The `debug=1' switch is used
971    # in these examples so that the commands that are sent to gnuplot
972    # are also output on stderr.
973    g1 = Gnuplot(debug=1)
974    g1.title('A simple example') # (optional)
975    g1('set style data linespoints') # give gnuplot an arbitrary command
976    # Plot a list of (x, y) pairs (tuples or a numpy array would
977    # also be OK):
978    g1.plot([[0.,1.1], [1.,5.8], [2.,3.3], [3.,4.2]])
979
980    # Plot one dataset from an array and one via a gnuplot function;
981    # also demonstrate the use of item-specific options:
982    g2 = Gnuplot(debug=1)
983    x = arange(10, dtype=numpy.float32)
984    y1 = x**2
985    # Notice how this plotitem is created here but used later?  This
986    # is convenient if the same dataset has to be plotted multiple
987    # times, because the data need only be written to a temporary file
988    # once.
989    d = Data(x, y1,
990             title="calculated by python",
991             with_="points pointsize 3 pointtype 3")
992    g2.title('Data can be computed by python or gnuplot')
993    g2.xlabel('x')
994    g2.ylabel('x squared')
995    # Plot a function alongside the Data PlotItem defined above:
996    g2.plot(Func("x**2", title="calculated by gnuplot"), d)
997
998    # Save what we just plotted as a color postscript file:
999    print("\n******** Generating postscript file 'gnuplot_test1.ps' ********\n")
1000    g2.hardcopy('gnuplot_test_plot.ps', color=1)
1001
1002    # Demonstrate a 3-d plot:
1003    g3 = Gnuplot(debug=1)
1004    # set up x and y values at which the function will be tabulated:
1005    x = arange(35)/2.0
1006    y = arange(30)/10.0 - 1.5
1007    # Make a 2-d array containing a function of x and y.  First create
1008    # xm and ym which contain the x and y values in a matrix form that
1009    # can be `broadcast' into a matrix of the appropriate shape:
1010    xm = x[:,newaxis]
1011    ym = y[newaxis,:]
1012    m = (sin(xm) + 0.1*xm) - ym**2
1013    g3('set style data lines')
1014    g3('set hidden')
1015    g3('set contour base')
1016    g3.xlabel('x')
1017    g3.ylabel('y')
1018    g3.splot(GridData(m,x,y))
1019
1020    # Delay so the user can see the plots:
1021    sys.stderr.write("Three plots should have appeared on your screen "
1022                     "(they may be overlapping).\n"
1023                     "Please press return to continue...\n")
1024    sys.stdin.readline()
1025
1026    # ensure processes and temporary files are cleaned up:
1027    del g1, g2, g3, d
1028
1029    # Enable the following code to test the old-style gnuplot interface
1030    if 0:
1031        # List of (x, y) pairs
1032        plot([(0.,1),(1.,5),(2.,3),(3.,4)])
1033
1034        # List of y values, file output
1035        print("\n            Generating postscript file 'gnuplot_test2.ps'\n")
1036        plot([1, 5, 3, 4], file='gnuplot_test2.ps')
1037
1038        # Two plots; each given by a 2d array
1039        x = arange(10, typecode=Float)
1040        y1 = x**2
1041        y2 = (10-x)**2
1042        plot(transpose(array([x, y1])), transpose(array([x, y2])))
1043
1044