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