drsaversp.py

#!/bin/python3
import numpy as np
import ROOT as rt
import os, subprocess, sys
import datetime
from tqdm import tqdm, trange
from particle import PDGID
import scipy.sparse

def index2img(phiindex,thetaindex,energy,from_width,from_pix_per_tower,**kwargs):
    to_width=kwargs.get('to_width',from_width)
    to_pix_per_tower=kwargs.get('to_pix_per_tower',from_pix_per_tower)
    if(from_width<to_width):
        print("to_width must be smaller then from_width")
        to_width=from_width
    if(from_pix_per_tower<to_pix_per_tower):
        print("to_pix_per_tower must be smaller then from_pix_per_tower")
        to_pix_per_tower=from_pix_per_tower
    from_height=(from_width*2+1)*from_pix_per_tower
    to_height=(to_width*2+1)*to_pix_per_tower
    outside=(from_width-to_width)*from_pix_per_tower
    _img=np.zeros((to_height,to_height))
    for i in range(len(phiindex)):
        if(phiindex[i]<outside or phiindex[i]>from_height-outside):continue
        if(thetaindex[i]<outside or thetaindex[i]>from_height-outside):continue
        phi=int(to_height*(phiindex[i]-outside)/(from_height-2*outside))
        theta=int(to_height*(thetaindex[i]-outside)/(from_height-2*outside))
        _img[theta,phi]+=energy[i]
    return _img
def asCartesian(r,theta,phi):
    x = r * np.sin( theta ) * np.cos( phi )
    y = r * np.sin( theta ) * np.sin( phi )
    z = r * np.cos( theta )
    return np.array([x,y,z])
def asSpherical(x,y,z):
    r       =  np.sqrt(x*x + y*y + z*z)
    theta   =  np.arccos(z/r)
    phi     =  np.arctan2(y,x)
    return np.array([r,theta,phi])
def project(x,y,z,notheta,nophi,ttheta,tphi,tntheta,tnphi):
    pull=200
    r=np.zeros(len(x))
    theta=np.zeros(len(x))
    phi=np.zeros(len(x))
    for i in range(len(x)):
        if(abs(notheta[i]>10)):continue
        tower=np.where((tntheta==notheta[i]) & (tnphi==nophi[i]))[0][0]
        towerxyz=np.array(asCartesian(pull,ttheta[tower],tphi[tower]))
        x[i],y[i],z[i]=np.array([x[i],y[i],z[i]])-towerxyz
        r[i],theta[i],phi[i]=asSpherical(x[i],y[i],z[i])
    return x,y,z,theta,phi

batch_size=5
now=datetime.datetime.now()
fin=""
_cb=0
version="0.0.2.3"
#_fb=56*3
#main_dir1="/pad/yulee/dream/repo/buf23/origin_buf/shelf2"
#main_dir2="/hdfs/user/yulee/DRsim/shelf"
#is_append=True
main_dir1="/pad/yulee/dream/repo/0023_img/tools/pack"
main_dir2="/hdfs/user/yulee/DRsim/simdata/pack"
main_dir=[main_dir1,main_dir2]
#energys=["20","70","50","100"]
energys=["10-100"]
#energys=["20-100"]
#saver_list=['gamma','e-','pi0','pi+','kaon0','kaon+','neutron','proton']
#saver_list=['kaon0','kaon+','neutron','proton']
#saver_list=['gamma','pi0','pi+']
saver_list=['e-','gamma']
#saver_list=['gg','uu','dd','ss','cc']
skip_name=["gg_20GeV","uu_20GeV"]
#saver_list=['gg','uu']
_fb=162
is_append=True
is_force=False
name_file=[]
name_cls=[]
table={}
particles={
  "e-":{"root":"{}/e-_{}GeV_{}.root","isjet":0},
  "gamma":{"root":"{}/gamma_{}GeV_{}.root","isjet":0},
  "pi+":{"root":"{}/pi+_{}GeV_{}.root","isjet":0},
  "pi0":{"root":"{}/pi0_{}GeV_{}.root","isjet":0},
  "kaon+":{"root":"{}/kaon+_{}GeV_{}.root","isjet":0},
  "kaon0":{"root":"{}/kaon0L_{}GeV_{}.root","isjet":0},
  "neutron":{"root":"{}/neutron_{}GeV_{}.root","isjet":0},
  "proton":{"root":"{}/proton_{}GeV_{}.root","isjet":0},
  "gg":{"root":"{}/gg_{}GeV_{}.root","isjet":1},
  "dd":{"root":"{}/dd_{}GeV_{}.root","isjet":1},
  "uu":{"root":"{}/uu_{}GeV_{}.root","isjet":1},
  "ss":{"root":"{}/ss_{}GeV_{}.root","isjet":1},
  "cc":{"root":"{}/cc_{}GeV_{}.root","isjet":1},
}
for i in range(len(saver_list)):
  _name=saver_list[i]
  name_file.append(particles[_name]["root"])
  name_cls.append("{}".format(_name))
#for nf,nc in zip(name_file,name_cls):
for energy in energys:
  for i in range(len(saver_list)):
    _name=saver_list[i]
    nf=particles[_name]["root"]
    nc="{}".format(_name)
    isjet=particles[_name]["isjet"]
    if("{}_{}GeV".format(_name,energy) in skip_name):continue
    for batch in range(int(100/batch_size)):
        buf_now=datetime.datetime.now()
        event=rt.TChain("event")
        filecount=0
        nf_list=[]
        entry_cut=[]
        batch_begin=batch*batch_size+0
        batch_end=batch*batch_size+batch_size
        for i in range(batch_begin,batch_end):
        #for i in range(0,26):
          for md in main_dir:
              nf_buf=nf.format(md,energy,i)
              if os.path.isfile(nf_buf):
                has_prt=0
                event_buf=rt.TChain("event")
                event_buf.Add(nf_buf)
                for leaf in event_buf.GetListOfLeaves():
                  if(leaf.GetName()=="prt_pid"):
                      has_prt+=17
                  if(leaf.GetName()=="point_phi_idx"):
                      has_prt+=13
                if(has_prt!=30):
                  print("###",nf_buf,event.GetEntries(),"no ptr")
                else:
                  print("###",nf_buf,event.GetEntries())
                  event.Add(nf_buf)
                  nf_list.append(nf_buf)
                  entry_cut.append(event.GetEntries())
                  filecount+=1
        if(filecount==0):continue
        images=[]
        image_gen=[]
        points=[]
        points_proj=[]
        ptcs=[]
        prts=[]
        pt_Gen=[]
        ptd_Gen=[]
        E_Gen=[]
        width_Gen=[]
        SC_ratio=[]
        CS_ratio=[]
        dr_ecorr=[]
        npz_name="npzs/{}_{}GeV_{}-{}".format(nc,energy,batch_begin,batch_end)
        print(npz_name,energy,"entries",event.GetEntries())
        if(is_append):
          #box=np.load("/hdfs/ml/dualreadout/{}/{}.npz".format(version,npz_name.split("/")[-1]),allow_pickle=True)
          #box=np.load("/hdfs/user/yulee/DRsim/npzs/{}/{}.npz".format(version,npz_name.split("/")[-1]),allow_pickle=True)
          if(os.path.isfile("{}.npz".format(npz_name))):
            box=np.load("{}.npz".format(npz_name),allow_pickle=True)
            E_Gen_box=box['E_Gen']
            dr_ecorr_box=box['dr_ecorr']
            print("loaded",len(E_Gen_box),len(dr_ecorr_box))
          else:
            is_append=False

        has_prt=0
        for leaf in event.GetListOfLeaves():
            if(leaf.GetName()=="prt_pid"):
                has_prt=1
        if(has_prt==0):
            print("no prt")
            break
        tq = trange(int(event.GetEntries()))
        for i in tq:
            tq.set_description("{} ".format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
            if(i>200000):break
            #if(i>100):i=i-100
            event.GetEntry(i)
            if(event.E_S==0 or event.E_C==0):
              continue
            if(len(event.fiber_phi)==0):
              continue
            if(event.E_Gen in E_Gen):
              if(event.E_DRcorr in dr_ecorr):
                #if(event.E_C/event.E_S in CS_ratio):
                continue
            if(is_append):
              if(event.E_Gen in E_Gen_box):
                if(event.E_DRcorr in dr_ecorr_box):
                  continue
            pt_Gen.append(event.pt_Gen)
            ptd_Gen.append(event.ptd)
            E_Gen.append(event.E_Gen)
            width_Gen.append(event.width_Gen)
            SC_ratio.append(event.E_S/event.E_C)
            CS_ratio.append(event.E_C/event.E_S)
            dr_ecorr.append(event.E_DRcorr)
            if(isjet==1):
              _fb=162
              prt=np.array([
                  list(event.prt_phi),
                  list(event.prt_theta),
                  list(event.prt_E),
                  list(event.prt_pid),
                  list(event.prt_px),
                  list(event.prt_py),
                  list(event.prt_pz),
                  ]).transpose()
              G4ptc=np.array([
                  list(event.ptc_phi),
                  list(event.ptc_theta),
                  list(event.ptc_p),
                  list(event.ptc_E),
                  list(event.ptc_pid),
                  list(event.ptc_px),
                  list(event.ptc_py),
                  list(event.ptc_pz),
                  list(event.ptc_vx),
                  list(event.ptc_vy),
                  list(event.ptc_vz),
                  ]).transpose()#phi,theta,p,E,px,py,pz,ch,nh,cl,nl
            if(isjet==0):
              _fb=168
              prt=np.array([
                  list(event.prt_phi),
                  list(event.prt_theta),
                  list(event.prt_E),
                  list(event.prt_pid),
                  list(event.prt_px),
                  list(event.prt_py),
                  list(event.prt_pz),
                  ]).transpose()
              if(len(event.ptc_phi)==len(event.ptc_vx)):
                G4ptc=np.array([
                  list(event.ptc_phi),
                  list(event.ptc_theta),
                  list(event.ptc_p),
                  list(event.ptc_E),
                  list(event.ptc_pid),
                  list(event.ptc_px),
                  list(event.ptc_py),
                  list(event.ptc_pz),
                  list(event.ptc_vx),
                  list(event.ptc_vy),
                  list(event.ptc_vz),
                  ]).transpose()#phi,theta,p,E,px,py,pz,ch,nh,cl,nl
              else:
                G4ptc=np.array([
                  list(event.ptc_phi),
                  list(event.ptc_theta),
                  list(event.ptc_p),
                  list(event.ptc_E),
                  list(event.ptc_pid),
                  list(event.ptc_px),
                  list(event.ptc_py),
                  list(event.ptc_pz),
                  list(event.ptc_vx)[-1:],
                  list(event.ptc_vy)[-1:],
                  list(event.ptc_vz)[-1:],#old process didn't cleared ptc_v reprocess please
                  ]).transpose()#phi,theta,p,E,px,py,pz,ch,nh,cl,nl
                  #np.array(list(event.ptc_px)),
                  #np.array(list(event.ptc_py)),
                  #np.array(list(event.ptc_pz)),
            #old image/====
            #images.append([np.array(list(event.image_ecor_s)[:_fb*_fb]).reshape((_fb,_fb))[_cb:_fb-_cb,_cb:_fb-_cb],
            #    np.array(list(event.image_ecor_c)[:_fb*_fb]).reshape((_fb,_fb))[_cb:_fb-_cb,_cb:_fb-_cb],
            #    np.array(list(event.image_n_s)[:_fb*_fb]).reshape((_fb,_fb))[_cb:_fb-_cb,_cb:_fb-_cb],
            #    np.array(list(event.image_n_c)[:_fb*_fb]).reshape((_fb,_fb))[_cb:_fb-_cb,_cb:_fb-_cb]])#s c ns nc
            #====/old image
            image_gen.append([np.array(list(event.image_gen)[:_fb*_fb]).reshape((_fb,_fb))[_cb:_fb-_cb,_cb:_fb-_cb]])
            #index point====/
            pix_height=(40*2+1)*56
            pix_cenphi=pix_height/2#+centher_theta_gen
            pix_centheta=pix_height/2#+centher_phi_gen
            phis_list=[]
            thetas_list=[]
            point_rt=np.array([event.point_phi_idx,event.point_theta_idx,event.point_ecor_s,event.point_ecor_c,event.point_depth]).transpose()
            #cut=0
            #for j in range(len(point_rt)):
            #    if(max(point_rt[j][0],point_rt[j][1])==0):
            #        cut=j
            #        break
            img_es=index2img(point_rt[:,0],point_rt[:,1],point_rt[:,2],from_width=40,from_pix_per_tower=56,to_width=1,to_pix_per_tower=56)
            img_ec=index2img(point_rt[:,0],point_rt[:,1],point_rt[:,3],from_width=40,from_pix_per_tower=56,to_width=1,to_pix_per_tower=56)
            images.append([img_es,img_ec])
            #for i in range(len(point_rt)): # index to polar angle
            #    if(point_rt[i][0]==0 and point_rt[i][1]==0):
            #        point_rt[i][0]=0
            #        point_rt[i][1]=0
            #    else:
            #        point_rt[i][0]=((point_rt[i][0]-pix_cenphi)/pix_height)
            #        point_rt[i][1]=((point_rt[i][1]-pix_centheta)/pix_height)
            point_rt=np.array(sorted(point_rt, key=lambda pnt:max(pnt[2],pnt[3]),reverse=True))
            points.append(point_rt)
            #====/index point

          
            x,y,z,notheta,nophi=np.array([list(event.fiber_x),list(event.fiber_y),list(event.fiber_z),list(event.tower_no_eta),list(event.tower_no_phi)])
            ttheta,tphi,tntheta,tnphi=np.array([list(event.tower_theta),list(event.tower_phi),list(event.tower_notheta),list(event.tower_nophi)])
            X,Y,Z,Theta,Phi=project(x,y,z,notheta,nophi,ttheta,tphi,tntheta,tnphi)
            point_proj=np.array([np.array(Phi),
                np.array(Theta),
                np.array(list(event.fiber_depth))/1000.,
                np.array(list(event.fiber_ecor_s)),
                np.array(list(event.fiber_ecor_c)),
                np.array(X),
                np.array(Y),
                np.array(Z),
                ]).transpose()#phi theta depth s c x y z

            #print("point",i,len(list(event.fiber_phi)),event.E_Gen)
            point_proj=sorted(point_proj, key=lambda pnt:max(pnt[3],pnt[4]),reverse=True)
            #print(len(point),len(gen))
            #print([len(i) for i in gen])

            if(len(point_proj)<2048):
                point_proj=np.concatenate([point_proj,np.zeros((2048-len(point_proj),8))],axis=0)
            if(len(point_proj)>2048):
                point_proj=point_proj[:2048]
            points_proj.append(point_proj)
            #====/old point
            if(len(G4ptc)<64):
                G4ptc=np.concatenate([G4ptc,np.zeros((64-len(G4ptc),11))],axis=0)
            G4ptc=sorted(G4ptc, key=lambda ptc:ptc[3],reverse=True)
            if(len(G4ptc)>64):
                G4ptc=G4ptc[:64]
            #for j in range(len(G4ptc)): #For What?
            #  is_l=-1
            #  if(PDGID(G4ptc[j][4]).is_hadron==0):is_l=0
            #  if(PDGID(G4ptc[j][4]).is_lepton==0):is_l=1
            #  if(is_l!=-1):G4ptc[j][5+2*is_l+int(PDGID(G4ptc[j][4]).charge==0)]=1

            ptcs.append(G4ptc)
            prts.append(prt)

        #print("depth",np.mean(list(event.fiber_depth)))
        images=np.array(images,dtype='float32')
        image_gen=np.array(image_gen,dtype='float32')
        #print([len(i) for i in point])
        points=np.array(points,dtype='float32')
        points_proj=np.array(points_proj,dtype='float32')
        ptcs=np.array(ptcs,dtype='float32')
        #prts=np.array(prts,dtype='int32')
        if(len(images)==0 and len(points)==0 and len(points_proj)==0):
          if(is_force):
            info=box['info'].item()
            images=box['image'].item().toarray().reshape(info['image_shape'])
            image_gen=box['image_gen'].item().toarray().reshape(info['image_gen_shape'])
            points=box['point'].item().toarray().reshape(info['point_shape'])
            points_proj=box['point_proj'].item().toarray().reshape(info['point_proj_shape'])
            ptcs=box['ptc']
            #prts=box['prt']
            pt_Gen=box['pt_Gen']
            ptd_Gen=box['ptd_Gen']
            E_Gen=box['E_Gen']
            width_Gen=box['width_Gen']
            SC_ratio=box['SC_ratio']
            CS_ratio=box['CS_ratio']
            dr_ecorr=box['dr_ecorr']
          else:
            continue
        elif(is_append):
          info=box['info'].item()
          images=np.concatenate([box['image'].item().toarray().reshape(info['image_shape']),images])
          image_gen=np.concatenate([box['image_gen'].reshape(info['image_gen_shape']),image_gen])
          points=np.concatenate([box['point'].reshape(info['point_shape']),points])
          print(box['point_proj'].shape,info['point_proj_shape'])
          points_proj=np.concatenate([box['point_proj'].reshape(info['point_proj_shape']),points_proj])
          ptcs=np.concatenate([box['ptc'],ptcs])
          #prts=np.concatenate([box['prt'],prts])
          pt_Gen=np.concatenate([box['pt_Gen'],pt_Gen])
          ptd_Gen=np.concatenate([box['ptd_Gen'],ptd_Gen])
          E_Gen=np.concatenate([box['E_Gen'],E_Gen])
          width_Gen=np.concatenate([box['width_Gen'],width_Gen])
          SC_ratio=np.concatenate([box['SC_ratio'],SC_ratio])
          CS_ratio=np.concatenate([box['CS_ratio'],CS_ratio])
          dr_ecorr=np.concatenate([box['dr_ecorr'],dr_ecorr])
          print("loaded")

        fin+="{} concat shape image {}, point {}, point_proj {}\n".format(npz_name,str(images.shape),str(points.shape),str(points_proj.shape))
        info={"name":npz_name,"image_shape":images.shape,"image_gen_shape":image_gen.shape,"point_shape":points.shape,"point_proj_shape":points_proj.shape,"ptcs_shape":ptcs.shape,"energy":energy,"version":version,"_fb":_fb,"_cb":_cb,"flist":nf_list,"entry_cut":entry_cut}
        #info["last_event_mean_depth"]=np.mean(list(event.fiber_depth))
        info["date"]=datetime.datetime.now()
        info["time_processing"]=datetime.datetime.now()-buf_now
        print(npz_name,"time costs ",datetime.datetime.now()-buf_now)
        #sp_points=scipy.sparse.csr_matrix(points.reshape(len(points),-1))
        sp_images=scipy.sparse.csr_matrix(images.reshape(len(images),-1))
        #sp_image_gen=scipy.sparse.csr_matrix(image_gen.reshape(len(image_gen),-1))
        table[npz_name]=[info['image_shape'],info['point_shape'],info['point_proj_shape']]
        try:
          np.savez_compressed(npz_name,image=sp_images,image_gen=image_gen,point=points,point_proj=points_proj,ptc=ptcs,prt=prts,pt_Gen=pt_Gen,ptd_Gen=ptd_Gen,dr_ecorr=dr_ecorr,E_Gen=E_Gen,width_Gen=width_Gen,SC_ratio=SC_ratio,CS_ratio=CS_ratio,info=info)
          #np.savez_compressed(npz_name+"test_imagesparse",image=sp_images)
          #np.savez_compressed(npz_name+"test_image_",image=images)
          #np.savez_compressed(npz_name+"test_imagegensparse_",image_gen=sp_image_gen)
          #np.savez_compressed(npz_name+"test_imagegen_",image_gen=image_gen)
          #np.savez_compressed(npz_name+"test_pointsparse_",point=sp_points)
          #np.savez_compressed(npz_name+"test_point_",point=points)
          #np.savez_compressed(npz_name+"test_prt_",ptc=ptcs,prt=prts)
          #np.savez_compressed(npz_name+"test_else_",pt_Gen=pt_Gen,ptd_Gen=ptd_Gen,dr_ecorr=dr_ecorr,E_Gen=E_Gen,width_Gen=width_Gen,SC_ratio=SC_ratio,CS_ratio=CS_ratio,info=info)
        except:
          print("Cut")
          np.savez_compressed(npz_name,image=sp_images[:11000],point=sp_points[:11000],ptc=ptcs[:11000],prt=prts[:11000],pt_Gen=pt_Gen[:11000],ptd_Gen=ptd_Gen[:11000],dr_ecorr=dr_ecorr[:11000],E_Gen=E_Gen[:11000],width_Gen=width_Gen[:11000],SC_ratio=SC_ratio[:11000],CS_ratio=CS_ratio[:11000],info=info)
        del points
        del images,image_gen

        #del event,images,points
        #del event, sp_images,sp_image_gen, sp_points
        del event, sp_images
        if(not is_append):
          #subprocess.call(["hdfs","dfs","-moveFromLocal","-f","{}.npz".format(npz_name),"/user/yulee/DRsim/npzs/{}/".format(version)])
          print("/hdfs/user/yulee/DRsim/npzs/{}/{}.npz".format(version,npz_name))
        else:
          print("/pad/yulee/keras/npzs/{}.npz".format(version,npz_name))

print("time costs ",datetime.datetime.now()-now)
f=open("npzs/table.txt","a")
f.write(str(table)+"\n")
f.close()
print(fin)
print(datetime.datetime.now())