import glob,math,sys,subprocess,copy,os,commands,ROOT pandapath=os.environ.get('PANDAPATH') sys.path.append(pandapath+'/macro/tpc/FOPI/mberger') sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1') import argparse #from functions import set_palette class Elektron: def __init__(self,x,y,z,time): self.pos=ROOT.TVector3(x,y,z) # in m self.time=time # in s self.v=0 #in m/s self.crpos=ROOT.TVector3(x,y,z) # in m self.crtime=time # in s self.vVec=ROOT.TVector3(0,0,0) # in m/s def setVelocity(self,v): self.v=v def setCreation(self,x,y,z,t): self.crpos.SetXYZ(x,y,z) self.crtime=t self.vVec=self.pos-self.crpos self.vVec*=1./(self.time-t) # divide by time def getV(self): return self.vVec class ElektronColl: def __init__(self,z): self.elektrons=[] self.zPos=z*1e-6 #mum->m self.zCut=0. self.zCutIsSet=False self.numAval=-1 self.functions={} self.histos={} nbins=500 self.nbins=nbins self.histos['xy']=ROOT.TH2D('hxy'+str(self.zPos), 'XY (cm) of electrons Z='+str(self.zPos), nbins,-0.001,0.001,nbins,-0.001,0.001) self.histos['v']=ROOT.TH1D('hv'+str(self.zPos), 'Velocity (m/s) of electrons Z='+str(self.zPos), nbins,0,0) self.histos['x']=ROOT.TH1D('hx'+str(self.zPos),'X pos (m) Z='+str(self.zPos), nbins,-0.001,0.001) self.histos['y']=ROOT.TH1D('hy'+str(self.zPos),'Y pos (m) Z='+str(self.zPos), nbins,-0.001,0.001) self.histos['t']=ROOT.TH1D('ht'+str(self.zPos),'time (s) Z='+str(self.zPos), nbins,0,0) self.histos['xcr']=ROOT.TH1D('hxc'+str(self.zPos),'Xcr pos (m) Z='+str(self.zPos), nbins,-0.001,0.001) self.histos['ycr']=ROOT.TH1D('hyc'+str(self.zPos),'Ycr pos (m) Z='+str(self.zPos), nbins,-0.001,0.001) self.histos['zcr']=ROOT.TH1D('hzc'+str(self.zPos),'Zcr pos (m) Z='+str(self.zPos), nbins,0,0) self.histos['tcr']=ROOT.TH1D('htc'+str(self.zPos),'time cr (s) Z='+str(self.zPos), nbins,0,0) self.histos['vx']=ROOT.TH1D('hvx'+str(self.zPos),'Vx (m/s) Z='+str(self.zPos), nbins,0,0) self.histos['vy']=ROOT.TH1D('hvy'+str(self.zPos),'Vy (m/s) Z='+str(self.zPos), nbins,0,0) self.histos['vz']=ROOT.TH1D('hvz'+str(self.zPos),'Vz (m/s) Z='+str(self.zPos), nbins,0,0) self.histos['vxVsx']=ROOT.TH1D('hvxVsx'+str(self.zPos),'Vx as function of x Z='+str(self.zPos), nbins,-0.001,0.001) self.histos['vyVsy']=ROOT.TH1D('hvyVsy'+str(self.zPos),'Vy as function of y Z='+str(self.zPos), nbins,-0.001,0.001) def norm(self): self.histos['vxVsx'].Divide(self.histos['x']) self.histos['vyVsy'].Divide(self.histos['y']) self.histos['vzVsz'].Divide(self.histos['z']) def setZcut(self,cut): self.zCut=cut self.zCutIsSet=True self.histos['z']=ROOT.TH1D('hz'+str(self.zPos),'Z pos (m) Z='+str(self.zPos), self.nbins,cut-0.001,cut+0.001) self.histos['vzVsz']=ROOT.TH1D('hvzVsz'+str(self.zPos),'Vz as function of z Z='+str(self.zPos), self.nbins,cut-0.001,cut+0.001) def setNumAval(self,num): self.numAval=num def readFromRootFile(self,filename): inf=ROOT.TFile(filename,'read') tree=inf.Get('ava') print tree.GetEntries() print tree.e_x2 avalcounter=-1 for a in tree: avalcounter+=1 if self.numAval!=-1 and self.numAvalself.zCut: continue self.addElec(a.e_x2[i],a.e_y2[i],a.e_z2[i],a.e_t2[i]) self.setElecCr(-1,a.e_x1[i],a.e_y1[i],a.e_z1[i],a.e_t1[i]) def readFromComsolFile(self,filename): inf=open(filename) for line in inf: if line[0]=='\n' or line[0]=='%': continue words=line.split() print line self.addElec(float(words[5]),float(words[6]),float(words[7]),float(words[8])) self.setElecCr(-1,float(words[1]),float(words[2]),float(words[3]),float(words[4])) def writeToFile(self): #make sure a file is open for h in self.histos: self.histos[h].Write() def addElec(self,x,y,z,t): x/=100 #cm->m y/=100 z/=100 t/=1e9 #ns->s self.elektrons.append(Elektron(x,y,z,t)) v=(self.zPos)/(t) self.elektrons[-1].setVelocity(v) self.histos['xy'].Fill(x,y) self.histos['x'].Fill(x) self.histos['y'].Fill(y) self.histos['y'].Fill(z) self.histos['v'].Fill(v) self.histos['t'].Fill(t) def setElecCr(self,i,x,y,z,t): x/=100 #cm->m y/=100 z/=100 t/=1e9 #ns->s self.elektrons[i].setCreation(x,y,z,t) self.histos['xcr'].Fill(x) self.histos['ycr'].Fill(y) self.histos['zcr'].Fill(z) self.histos['tcr'].Fill(t) self.histos['vx'].Fill(self.elektrons[i].getV().X()) self.histos['vy'].Fill(self.elektrons[i].getV().Y()) self.histos['vz'].Fill(self.elektrons[i].getV().Z()) self.histos['vxVsx'].Fill(self.elektrons[i].pos.X(),self.elektrons[i].getV().X()) self.histos['vyVsy'].Fill(self.elektrons[i].pos.Y(),self.elektrons[i].getV().Y()) self.histos['vzVsz'].Fill(self.elektrons[i].pos.Z(),self.elektrons[i].getV().Z()) def getDistr(self,dis,val): if self.functions.get(dis,None)==None: self.functions[dis]=ROOT.TF1('gaus'+dis,'gaus',-0.1,0.1) self.functions[dis].SetLineColor(2) if self.histos.get(dis,None)==None: return -99999 tmpcanv=ROOT.TCanvas() self.histos[dis].Fit(self.functions[dis],'QO') del tmpcanv return self.functions[dis].GetParameter(val) parser=argparse.ArgumentParser(description='read in takus grfield files and plot stuff') parser.add_argument('infilelist',help='the input files with electron information',type=str) parser.add_argument('--numAval',help='number of avalanches to analyze',type=int, default=5) parser.add_argument('--rfile',help='the root file to write everything into',type=str,default='./avalanches.root') args=parser.parse_args() ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")') #set_palette() zcuts=[-0.425,-0.465,-0.515,-0.565] DriftStates=[] ifile=-1 canv=[] for line in open(args.infilelist): ifile+=1 if line[0]=='#' or line[0]=='\n': continue words=line.split(';') DriftStates.append(ElektronColl(float(words[1]))) DriftStates[-1].setNumAval(args.numAval) DriftStates[-1].setZcut(zcuts[ifile]) if words[0].find('.root')!=-1: DriftStates[-1].readFromRootFile(words[0]) elif words[0].find('.comsol')!=-1: DriftStates[-1].readFromComsolFile(words[0]) else: print 'unknown file format' exit() print 'pos:',words[1] print 'mean time:',DriftStates[-1].getDistr('t',1) print 'sigma x:',DriftStates[-1].getDistr('x',2) print 'sigma y:',DriftStates[-1].getDistr('y',2) DriftStates[-1].norm() canv.append(ROOT.TCanvas(str(words[1]),str(words[1]),1500,1000)) canv[-1].Divide(4,3) canv[-1].cd(1) DriftStates[-1].histos['xy'].Draw('colz') canv[-1].cd(2) DriftStates[-1].histos['v'].Draw() canv[-1].cd(3) DriftStates[-1].histos['x'].Draw() canv[-1].cd(4) DriftStates[-1].histos['y'].Draw() canv[-1].cd(5) #DriftStates[-1].histos['t'].Draw() DriftStates[-1].histos['vxVsx'].Draw() canv[-1].cd(6) #DriftStates[-1].histos['xcr'].Draw() DriftStates[-1].histos['vyVsy'].Draw() canv[-1].cd(7) #DriftStates[-1].histos['ycr'].Draw() DriftStates[-1].histos['vzVsz'].Draw() canv[-1].cd(8) DriftStates[-1].histos['zcr'].Draw() canv[-1].cd(9) DriftStates[-1].histos['tcr'].Draw() canv[-1].cd(10) DriftStates[-1].histos['vx'].Draw() canv[-1].cd(11) DriftStates[-1].histos['vy'].Draw() canv[-1].cd(12) DriftStates[-1].histos['vz'].Draw() canv[-1].Update() graphs=[] cgraph=ROOT.TCanvas('graph','graph',1000,1000) if len(DriftStates)>1: deltat=[] deltax=[] deltay=[] deltaz=[] for i in range(len(DriftStates)-1): print i,len(DriftStates) deltat.append(abs(DriftStates[i].getDistr('t',1)-DriftStates[i+1].getDistr('t',1))) deltax.append(abs(DriftStates[i].getDistr('x',2)-DriftStates[i+1].getDistr('x',2))) deltay.append(abs(DriftStates[i].getDistr('y',2)-DriftStates[i+1].getDistr('y',2))) deltaz.append(abs(DriftStates[i].zPos-DriftStates[i+1].zPos)) print 'dt=',deltat[-1],'dx=',deltax[-1],'dy=',deltay[-1],'vx=',deltax[-1]/deltat[-1],'vy=',deltay[-1]/deltat[-1] velx=ROOT.TGraph() velx.SetTitle('velocity X') velx.SetName('velocity_X') velx.SetMarkerStyle(20) vely=ROOT.TGraph() vely.SetTitle('velocity Y') vely.SetName('velocity_Y') vely.SetMarkerStyle(20) for i in range(len(DriftStates)-1): velx.SetPoint(i,DriftStates[i+1].zPos,deltax[i]/deltat[i]) vely.SetPoint(i,DriftStates[i+1].zPos,deltay[i]/deltat[i]) graphs.append(velx) graphs.append(vely) velx2=ROOT.TGraphErrors() velx2.SetTitle('velocity X v2') velx2.SetName('velocity_X_v2') velx2.SetMarkerStyle(20) vely2=ROOT.TGraphErrors() vely2.SetTitle('velocity Y v2') vely2.SetName('velocity_Y_v2') vely2.SetMarkerStyle(20) velz2=ROOT.TGraphErrors() velz2.SetTitle('velocity Z v2') velz2.SetName('velocity_Z_v2') velz2.SetMarkerStyle(20) graphs.append(velx2) graphs.append(vely2) graphs.append(velz2) posx=ROOT.TGraph() posx.SetTitle('position spread X') posx.SetName('position_spread_X') posx.SetMarkerStyle(20) posy=ROOT.TGraph() posy.SetTitle('position spread Y') posy.SetName('position_spread_Y') posy.SetMarkerStyle(20) graphs.append(posx) graphs.append(posy) print 'generating velocity graph:' for i in range(len(DriftStates)): print 'vx=',DriftStates[i].getDistr('vx',1),', sigma=',DriftStates[i].getDistr('vx',2) velx2.SetPoint(i,DriftStates[i].zPos,DriftStates[i].getDistr('vx',1)) velx2.SetPointError(i,0,DriftStates[i].getDistr('vx',2)) print 'vy=',DriftStates[i].getDistr('vy',1),', sigma=',DriftStates[i].getDistr('vy',2) vely2.SetPoint(i,DriftStates[i].zPos,DriftStates[i].getDistr('vy',1)) vely2.SetPointError(i,0,DriftStates[i].getDistr('vy',2)) print 'vz=',DriftStates[i].getDistr('vz',1),', sigma=',DriftStates[i].getDistr('vz',2) velz2.SetPoint(i,DriftStates[i].zPos,DriftStates[i].getDistr('vz',1)) velz2.SetPointError(i,0,DriftStates[i].getDistr('vz',2)) posx.SetPoint(i,DriftStates[i].zPos,DriftStates[i].getDistr('x',2)) posy.SetPoint(i,DriftStates[i].zPos,DriftStates[i].getDistr('y',2)) cgraph.Divide(3,3) cgraph.cd(1) velx.Draw('AP') cgraph.cd(2) vely.Draw('AP') cgraph.cd(4) velx2.Draw('AP') cgraph.cd(5) vely2.Draw('AP') cgraph.cd(6) velz2.Draw('AP') cgraph.cd(7) posx.Draw('AP') cgraph.cd(8) posy.Draw('AP') cgraph.Update() rfile=ROOT.TFile(args.rfile,'recreate') for D in DriftStates: D.writeToFile() for c in canv: c.Write() for g in graphs: g.Write() cgraph.Write() rfile.Close() u='w' while u!='q': u=raw_input('done?')