170 lines
6.1 KiB
Python
170 lines
6.1 KiB
Python
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") |