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