drsaver.py

#!/usr/bin/python2.7
import os
import sys
import numpy as np
import random
import argparse
from driter import *
parser=argparse.ArgumentParser()
now=datetime.datetime.now()
parser.add_argument("--pt",type=int,default=200,help='end ratio')
parser.add_argument("--num",type=int,default=1,help='end ratio')
args=parser.parse_args()
data_path=["/pad/yulee/geant4/tester/analysis/fast/uJet50GeV_fastsim_{}.root","/pad/yulee/geant4/tester/analysis/fast/gJet50GeV_fastsim_{}.root"]
tree_name="event"
num_file=10
l=args.num
chain=[]
for i in range(len(data_path)):
  chain.append(rt.TChain(tree_name))
  for j in range(num_file):
    chain[i].Add(data_path[i].format(j+10*l))
#data=DataGenerator(chain,data_form="pixel",target=4)
data=DataGenerator(chain,data_form="point",num_point=2048,target=4)
print(sum(data.total_len))
x,y=data._DataGenerator__get_data(sum(data.total_len))
label,cmult,nmult,chad_mult,nhad_mult,photon_mult,width,major_axis,minor_axis,ptd,E_Gen,mass_Gen=np.transpose(y)
#np.savez_compressed("npzs/drpixel{}vars".format(l),x=x,label=label,cmult=cmult,nmult=nmult,chad_mult=chad_mult,nhad_mult=nhad_mult,photon_mult=photon_mult,width=width,major_axis=major_axis,minor_axis=minor_axis,ptd=ptd,E_Gen=E_Gen,mass_Gen=mass_Gen)
np.savez_compressed("npzs/drpoint{}vars".format(l),x=x,label=label,cmult=cmult,nmult=nmult,chad_mult=chad_mult,nhad_mult=nhad_mult,photon_mult=photon_mult,width=width,major_axis=major_axis,minor_axis=minor_axis,ptd=ptd,E_Gen=E_Gen,mass_Gen=mass_Gen)
del x
del y
del data
del chain
print("save done")
print(datetime.datetime.now()-now,l)