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