qdsfit/gracedriver.py

__author__ = 'markusro'


import tempfile,os
import numpy as np

class grace:
    def __init__(self):
        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, "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.data_counter = 0

    def __del__(self):
        # take care of tmp files:
        for f in self.tmpfiles:
            os.remove(f)

    def plot(self,  x, y, **kwds):
        tmp_fd, tmp_name = tempfile.mkstemp()
        self.tmpfiles.append(tmp_name)
        np.savetxt(tmp_name, np.array([x, y]).T)
        #tmp_fd.close()
        self.cmds.append('READ NXY "%s"\n'%tmp_name)
        self.cmds.append('S%i SYMBOL SIZE 0.7\n'%(self.data_counter))
        self.cmds.append('S%i SYMBOL COLOR 1\n'%(self.data_counter))
        #self.cmds.append('S%i SYMBOL FILL COLOR %i\n'%(self.data_counter, 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

        #self.cmds.append('S%i LINE COLOR %i\n'%(self.data_counter, self.data_counter))
        if "label" in kwds.keys():
            label = kwds["label"]
            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:
                print "Symbol not known: %s"%sym



        # if "label" in kwds.keys():
        #     label = kwds["label"]
        #     self.cmds.append('S%i LEGEND "%s"\n'%(self.data_counter, label))

        self.data_counter += 1
        pass


    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, fname):
        self.cmds.append('SAVEALL "%s"\n'%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 tmpfile
        os.system("xmgrace -batch %s -hardcopy -nosafe -printfile tmp.tmp"%tmp_name)
        os.remove("tmp.tmp")

if __name__ == "__main__":
    print "Testing Grace driver"
    np.random.seed(1337) # make it reproducible
    nums = 30
    gr = grace()
    t = np.linspace(0,1,nums)
    for i in xrange(30):
        gr.plot(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))],
                )
    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")