1# Pizza.py toolkit, www.cs.sandia.gov/~sjplimp/pizza.html
2# Steve Plimpton, sjplimp@sandia.gov, Sandia National Laboratories
3#
4# Copyright (2005) Sandia Corporation.  Under the terms of Contract
5# DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
6# certain rights in this software.  This software is distributed under
7# the GNU General Public License.
8
9# gnu tool
10
11oneline = "Create plots via GnuPlot plotting program"
12
13docstr = """
14g = gnu()		       start up GnuPlot
15g.stop()		       shut down GnuPlot process
16
17g.plot(a)                      plot vector A against linear index
18g.plot(a,b)	 	       plot B against A
19g.plot(a,b,c,d,...)	       plot B against A, D against C, etc
20g.mplot(M,N,S,"file",a,b,...)  multiple plots saved to file0000.eps, etc
21
22  each plot argument can be a tuple, list, or Numeric/NumPy vector
23  mplot loops over range(M,N,S) and create one plot per iteration
24    last args are same as list of vectors for plot(), e.g. 1, 2, 4 vectors
25    each plot is made from a portion of the vectors, depending on loop index i
26      Ith plot is of b[0:i] vs a[0:i], etc
27    series of plots saved as file0000.eps, file0001.eps, etc
28    if use xrange(),yrange() then plot axes will be same for all plots
29
30g("plot 'file.dat' using 2:3 with lines")      execute string in GnuPlot
31
32g.enter()	   	                enter GnuPlot shell
33gnuplot> plot sin(x) with lines         type commands directly to GnuPlot
34gnuplot> exit, quit	      	        exit GnuPlot shell
35
36g.export("data",range(100),a,...)       create file with columns of numbers
37
38  all vectors must be of equal length
39  could plot from file with GnuPlot command: plot 'data' using 1:2 with lines
40
41g.select(N)  	               figure N becomes the current plot
42
43  subsequent commands apply to this plot
44
45g.hide(N)  	               delete window for figure N
46g.save("file")	               save current plot as file.eps
47
48Set attributes for current plot:
49
50g.erase()                      reset all attributes to default values
51g.aspect(1.3)                  aspect ratio
52g.xtitle("Time")               x axis text
53g.ytitle("Energy")             y axis text
54g.title("My Plot")             title text
55g.title("title","x","y")       title, x axis, y axis text
56g.xrange(xmin,xmax)            x axis range
57g.xrange()                     default x axis range
58g.yrange(ymin,ymax)            y axis range
59g.yrange()                     default y axis range
60g.xlog()                       toggle x axis between linear and log
61g.ylog()                       toggle y axis between linear and log
62g.label(x,y,"text")            place label at x,y coords
63g.curve(N,'r')                 set color of curve N
64
65  colors: 'k' = black, 'r' = red, 'g' = green, 'b' = blue
66          'm' = magenta, 'c' = cyan, 'y' = yellow
67"""
68
69# History
70#   8/05, Matt Jones (BYU): original version
71#   9/05, Steve Plimpton: added mplot() method
72
73# ToDo list
74#   allow choice of JPG or PNG or GIF when saving ?
75#     can this be done from GnuPlot or have to do via ImageMagick convert ?
76#     way to trim EPS plot that is created ?
77#   hide does not work on Mac aqua
78#   select does not pop window to front on Mac aqua
79
80# Variables
81#   current = index of current figure (1-N)
82#   figures = list of figure objects with each plot's attributes
83#             so they aren't lost between replots
84
85# Imports and external programs
86
87import types, os
88
89try: from DEFAULTS import PIZZA_GNUPLOT
90except: PIZZA_GNUPLOT = "gnuplot"
91try: from DEFAULTS import PIZZA_GNUTERM
92except: PIZZA_GNUTERM = "x11"
93
94# Class definition
95
96class gnu:
97
98  # --------------------------------------------------------------------
99
100  def __init__(self):
101    self.GNUPLOT = os.popen(PIZZA_GNUPLOT,'w')
102    self.file = "tmp.gnu"
103    self.figures = []
104    self.select(1)
105
106  # --------------------------------------------------------------------
107
108  def stop(self):
109    self.__call__("quit")
110    del self.GNUPLOT
111
112  # --------------------------------------------------------------------
113
114  def __call__(self,command):
115    self.GNUPLOT.write(command + '\n')
116    self.GNUPLOT.flush()
117
118  # --------------------------------------------------------------------
119
120  def enter(self):
121    while 1:
122      command = raw_input("gnuplot> ")
123      if command == "quit" or command == "exit": return
124      self.__call__(command)
125
126  # --------------------------------------------------------------------
127  # write plot vectors to files and plot them
128
129  def plot(self,*vectors):
130    if len(vectors) == 1:
131      file = self.file + ".%d.1" % self.current
132      linear = range(len(vectors[0]))
133      self.export(file,linear,vectors[0])
134      self.figures[self.current-1].ncurves = 1
135    else:
136      if len(vectors) % 2: raise StandardError,"vectors must come in pairs"
137      for i in range(0,len(vectors),2):
138        file = self.file + ".%d.%d" % (self.current,i/2+1)
139        self.export(file,vectors[i],vectors[i+1])
140      self.figures[self.current-1].ncurves = len(vectors)/2
141    self.draw()
142
143  # --------------------------------------------------------------------
144  # create multiple plots from growing vectors, save to numbered files
145  # don't plot empty vector, create a [0] instead
146
147  def mplot(self,start,stop,skip,file,*vectors):
148    n = 0
149    for i in range(start,stop,skip):
150      partial_vecs = []
151      for vec in vectors:
152        if i: partial_vecs.append(vec[:i])
153        else: partial_vecs.append([0])
154      self.plot(*partial_vecs)
155
156      if n < 10:     newfile = file + "000" + str(n)
157      elif n < 100:  newfile = file + "00" + str(n)
158      elif n < 1000: newfile = file + "0" + str(n)
159      else:          newfile = file + str(n)
160
161      self.save(newfile)
162      n += 1
163
164  # --------------------------------------------------------------------
165  # write list of equal-length vectors to filename
166
167  def export(self,filename,*vectors):
168    n = len(vectors[0])
169    for vector in vectors:
170      if len(vector) != n: raise StandardError,"vectors must be same length"
171    f = open(filename,'w')
172    nvec = len(vectors)
173    for i in xrange(n):
174      for j in xrange(nvec):
175        print >>f,vectors[j][i],
176      print >>f
177    f.close()
178
179  # --------------------------------------------------------------------
180  # select plot N as current plot
181
182  def select(self,n):
183    self.current = n
184    if len(self.figures) < n:
185      for i in range(n - len(self.figures)):
186        self.figures.append(figure())
187    cmd = "set term " + PIZZA_GNUTERM + ' ' + str(n)
188    self.__call__(cmd)
189    if self.figures[n-1].ncurves: self.draw()
190
191  # --------------------------------------------------------------------
192  # delete window for plot N
193
194  def hide(self,n):
195    cmd = "set term %s close %d" % (PIZZA_GNUTERM,n)
196    self.__call__(cmd)
197
198  # --------------------------------------------------------------------
199  # save plot to file.eps
200  # final re-select will reset terminal
201  # do not continue until plot file is written out
202  #   else script could go forward and change data file
203  #   use tmp.done as semaphore to indicate plot is finished
204
205  def save(self,file):
206    self.__call__("set terminal postscript enhanced solid lw 2 color portrait")
207    cmd = "set output '%s.eps'" % file
208    self.__call__(cmd)
209    if os.path.exists("tmp.done"): os.remove("tmp.done")
210    self.draw()
211    self.__call__("!touch tmp.done")
212    while not os.path.exists("tmp.done"): continue
213    self.__call__("set output")
214    self.select(self.current)
215
216  # --------------------------------------------------------------------
217  # restore default attributes by creating a new fig object
218
219  def erase(self):
220    fig = figure()
221    fig.ncurves = self.figures[self.current-1].ncurves
222    self.figures[self.current-1] = fig
223    self.draw()
224
225  # --------------------------------------------------------------------
226
227  def aspect(self,value):
228    self.figures[self.current-1].aspect = value
229    self.draw()
230
231  # --------------------------------------------------------------------
232
233  def xrange(self,*values):
234    if len(values) == 0:
235      self.figures[self.current-1].xlimit = 0
236    else:
237      self.figures[self.current-1].xlimit = (values[0],values[1])
238    self.draw()
239
240  # --------------------------------------------------------------------
241
242  def yrange(self,*values):
243    if len(values) == 0:
244      self.figures[self.current-1].ylimit = 0
245    else:
246      self.figures[self.current-1].ylimit = (values[0],values[1])
247    self.draw()
248
249  # --------------------------------------------------------------------
250
251  def label(self,x,y,text):
252    self.figures[self.current-1].labels.append((x,y,text))
253    self.figures[self.current-1].nlabels += 1
254    self.draw()
255
256  # --------------------------------------------------------------------
257
258  def nolabels(self):
259    self.figures[self.current-1].nlabel = 0
260    self.figures[self.current-1].labels = []
261    self.draw()
262
263  # --------------------------------------------------------------------
264
265  def title(self,*strings):
266    if len(strings) == 1:
267      self.figures[self.current-1].title = strings[0]
268    else:
269      self.figures[self.current-1].title = strings[0]
270      self.figures[self.current-1].xtitle = strings[1]
271      self.figures[self.current-1].ytitle = strings[2]
272    self.draw()
273
274  # --------------------------------------------------------------------
275
276  def xtitle(self,label):
277    self.figures[self.current-1].xtitle = label
278    self.draw()
279
280  # --------------------------------------------------------------------
281
282  def ytitle(self,label):
283    self.figures[self.current-1].ytitle = label
284    self.draw()
285
286  # --------------------------------------------------------------------
287
288  def xlog(self):
289    if self.figures[self.current-1].xlog:
290      self.figures[self.current-1].xlog = 0
291    else:
292      self.figures[self.current-1].xlog = 1
293    self.draw()
294
295  # --------------------------------------------------------------------
296
297  def ylog(self):
298    if self.figures[self.current-1].ylog:
299      self.figures[self.current-1].ylog = 0
300    else:
301      self.figures[self.current-1].ylog = 1
302    self.draw()
303
304  # --------------------------------------------------------------------
305
306  def curve(self,num,color):
307    fig = self.figures[self.current-1]
308    while len(fig.colors) < num: fig.colors.append(0)
309    fig.colors[num-1] = colormap[color]
310    self.draw()
311
312  # --------------------------------------------------------------------
313  # draw a plot with all its settings
314  # just return if no files of vectors defined yet
315
316  def draw(self):
317    fig = self.figures[self.current-1]
318    if not fig.ncurves: return
319
320    cmd = 'set size ratio ' + str(1.0/float(fig.aspect))
321    self.__call__(cmd)
322
323    cmd = 'set title ' + '"' + fig.title + '"'
324    self.__call__(cmd)
325    cmd = 'set xlabel ' + '"' + fig.xtitle + '"'
326    self.__call__(cmd)
327    cmd = 'set ylabel ' + '"' + fig.ytitle + '"'
328    self.__call__(cmd)
329
330    if fig.xlog: self.__call__("set logscale x")
331    else: self.__call__("unset logscale x")
332    if fig.ylog: self.__call__("set logscale y")
333    else: self.__call__("unset logscale y")
334    if fig.xlimit:
335      cmd = 'set xr [' + str(fig.xlimit[0]) + ':' + str(fig.xlimit[1]) + ']'
336      self.__call__(cmd)
337    else: self.__call__("set xr [*:*]")
338    if fig.ylimit:
339      cmd = 'set yr [' + str(fig.ylimit[0]) + ':' + str(fig.ylimit[1]) + ']'
340      self.__call__(cmd)
341    else: self.__call__("set yr [*:*]")
342
343    self.__call__("set nolabel")
344    for i in range(fig.nlabels):
345      x = fig.labels[i][0]
346      y = fig.labels[i][1]
347      text = fig.labels[i][2]
348      cmd = 'set label ' + '\"' + text + '\" at ' + str(x) + ',' + str(y)
349      self.__call__(cmd)
350
351    self.__call__("set key off")
352    cmd = 'plot '
353    for i in range(fig.ncurves):
354      file = self.file + ".%d.%d" % (self.current,i+1)
355      if len(fig.colors) > i and fig.colors[i]:
356        cmd += "'" + file + "' using 1:2 with line %d, " % fig.colors[i]
357      else:
358        cmd += "'" + file + "' using 1:2 with lines, "
359    self.__call__(cmd[:-2])
360
361# --------------------------------------------------------------------
362# class to store settings for a single plot
363
364class figure:
365
366  def __init__(self):
367    self.ncurves = 0
368    self.colors  = []
369    self.title   = ""
370    self.xtitle  = ""
371    self.ytitle  = ""
372    self.aspect  = 1.3
373    self.xlimit  = 0
374    self.ylimit  = 0
375    self.xlog    = 0
376    self.ylog    = 0
377    self.nlabels = 0
378    self.labels  = []
379
380# --------------------------------------------------------------------
381# line color settings
382
383colormap = {'k':-1, 'r':1, 'g':2, 'b':3, 'm':4, 'c':5, 'y':7}
384