import re import PyQt4.QtGui __author__ = 'markusro' import tempfile,os import numpy as np class GracePlot(object): def __init__(self, fname): self.fname = fname self.ls_map = {"None":0, "-":1, ":":2, "--":3, "-.":4 } # Symbols: 0:None 1:Circle 2:Square 3:Diamond 4:Triangle 5:up 6:left 7:down 8:right, 9:PLus 10:X 11:Star self.sym_map = {None:0, "None":0, "o":1, "s":2, "d":3, "D":4, "2":5, "v":6, "3":6, "4":7,"1":8, "+":8,"x":9,"*":10 } self.tmpfiles = [] self.cmds = [] self.color_map = {} # color map will pbe prepended to the final agr file self.data_counter = 0 self.color_counter = 8 # keep first 8 colors already defined? def __del__(self): # take care of tmp files: for f in self.tmpfiles: os.remove(f) def plot(self, x, y, **kwds): # save data to temporary file tmp_fd, tmp_name = tempfile.mkstemp() self.tmpfiles.append(tmp_name) np.savetxt(tmp_name, np.array([x, y]).T) # read data from temporary file self.cmds.append('READ NXY "%s"\n'%tmp_name) self.cmds.append('S%i SYMBOL SIZE 0.5\n'%(self.data_counter)) self.cmds.append('S%i SYMBOL FILL PATTERN 1\n'%(self.data_counter)) self.cmds.append('S%i SYMBOL 1\n'%(self.data_counter)) # No line if "label" in kwds.keys(): # TODO: implement at least greek symbols and lower upper case (_ and ^)? # with translate method? label = unicode(kwds["label"]).encode('ascii', 'ignore') self.cmds.append('S%i LEGEND "%s"\n'%(self.data_counter, label)) if "ls" in kwds.keys(): ls = kwds["ls"] if ls in self.ls_map.keys(): self.cmds.append('S%i LINE LINESTYLE %i\n'%(self.data_counter, self.ls_map[ls])) # Line if "sym" in kwds.keys(): sym = kwds["sym"] if sym in self.sym_map.keys(): self.cmds.append('S%i SYMBOL %i\n'%(self.data_counter, self.sym_map[sym])) if sym in ['+', 'x', '*']: self.cmds.append('S%i SYMBOL COLOR %i\n'%(self.data_counter, self.data_counter)) else: self.cmds.append('S%i SYMBOL COLOR %i\n' % (self.data_counter, 1)) # vlack symbol outline else: print "Symbol not known: %s"%sym if "color" in kwds.keys(): color = str(PyQt4.QtGui.QColor(kwds["color"]).getRgb()[:3]) if color in self.color_map.keys(): color_id = self.color_map[color][0] else: self.color_map[color] = [ self.color_counter, "@MAP COLOR %i TO %s, \"qds%i\" \n"% (self.color_counter, color, self.color_counter) ] color_id = self.color_counter self.color_counter += 1 self.cmds.append('S%i SYMBOL FILL COLOR %i\n' % (self.data_counter, color_id)) self.cmds.append('S%i LINE COLOR %i\n' % (self.data_counter, color_id)) self.data_counter += 1 def loglog(self, x,y, **kwds): self.cmds.append('YAXES SCALE LOGARITHMIC\n') self.cmds.append("YAXIS TICKLABEL FORMAT POWER\n") self.cmds.append("YAXIS TICKLABEL PREC 0\n") self.cmds.append('XAXES SCALE LOGARITHMIC\n') self.cmds.append("XAXIS TICKLABEL FORMAT POWER\n") self.cmds.append("XAXIS TICKLABEL PREC 0\n") self.plot(x, y, **kwds) def semilogx(self, x, y, **kwds): self.cmds.append('XAXES SCALE LOGARITHMIC\n') self.cmds.append("xAXIS TICKLABEL FORMAT POWER\n") self.cmds.append("xAXIS TICKLABEL PREC 0\n") self.plot(x, y, **kwds) def semilogy(self, x, y, **kwds): self.cmds.append('YAXES SCALE LOGARITHMIC\n') self.cmds.append("YAXIS TICKLABEL FORMAT POWER\n") self.cmds.append("YAXIS TICKLABEL PREC 0\n") self.plot(x, y, **kwds) def xlabel(self, label ): self.cmds.append('XAXIS LABEL "%s"\n'%label) pass def ylabel(self, label ): self.cmds.append('YAXIS LABEL "%s"\n'%label) pass def legend(self, on=True): if on: self.cmds.append('LEGEND ON\n') else: self.cmds.append('LEGEND OFF\n') def save(self): self.cmds.append('AUTOSCALE\n') self.cmds.append('SAVEALL "%s"\n'%self.fname) # write cmds to tmpfile tmp_fd, tmp_name = tempfile.mkstemp() self.tmpfiles.append(tmp_name) tmp_file = open(tmp_name, 'w') tmp_file.writelines(self.cmds) tmp_file.close() # excecute temporary xmgrace file to create final agr os.system("xmgrace -batch %s -hardcopy -nosafe -printfile tmp.tmp" % tmp_name) os.remove("tmp.tmp") # prepend color map to the new file with file(self.fname, 'r') as original_agr: data = original_agr.readlines() # get the last "@map color ..." line last_color_lineno = 0 for i,line in enumerate(data): if line.lower().startswith("@map color"): last_color_lineno = i+1 with file(self.fname, 'w') as new_agr: new_agr.writelines(data[:last_color_lineno]) for color in self.color_map: new_agr.write(self.color_map[color][1]) new_agr.writelines(data[last_color_lineno:]) if __name__ == "__main__": print "Testing Grace driver" np.random.seed(1337) # make it reproducible nums = 30 gr = GracePlot() t = np.linspace(0,1,nums) for i in xrange(30): gr.loglog(t, i + np.sin(2*np.pi * 3 * t + i*0.33 ) + 0.3*np.random.random(nums), label="label %i"%i, ls=gr.ls_map.keys()[i%(len(gr.ls_map))], sym=gr.sym_map.keys()[i%(len(gr.sym_map))], color="red" ) gr.xlabel(r"xlabel / \xm\sl\N\0") gr.ylabel(r"ylabel / \xt\sS\N\0") gr.save("test.agr") print "created test.agr" os.system("xmgrace test.agr") #print "deleting test.agr" #os.remove("test.agr")