-
Notifications
You must be signed in to change notification settings - Fork 18
/
Copy pathconvert_root2pq_jet.py
148 lines (118 loc) · 5.56 KB
/
convert_root2pq_jet.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import pyarrow.parquet as pq
import pyarrow as pa # pip install pyarrow==0.7.1
import ROOT
import numpy as np
import glob, os
from skimage.measure import block_reduce # pip install scikit-image
from numpy.lib.stride_tricks import as_strided
import argparse
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-i', '--infile', default='output_qqgg.root', type=str, help='Input root file.')
parser.add_argument('-o', '--outdir', default='.', type=str, help='Output pq file dir.')
parser.add_argument('-d', '--decay', default='test', type=str, help='Decay name.')
parser.add_argument('-n', '--idx', default=0, type=int, help='Input root file index.')
args = parser.parse_args()
def upsample_array(x, b0, b1):
r, c = x.shape # number of rows/columns
rs, cs = x.strides # row/column strides
x = as_strided(x, (r, b0, c, b1), (rs, 0, cs, 0)) # view as a larger 4D array
return x.reshape(r*b0, c*b1)/(b0*b1) # create new 2D array with same total occupancy
def resample_EE(imgECAL, factor=2):
# EE-
imgEEm = imgECAL[:140-85] # EE- in the first 55 rows
imgEEm = np.pad(imgEEm, ((1,0),(0,0)), 'constant', constant_values=0) # for even downsampling, zero pad 55 -> 56
imgEEm_dn = block_reduce(imgEEm, block_size=(factor, factor), func=np.sum) # downsample by summing over [factor, factor] window
imgEEm_dn_up = upsample_array(imgEEm_dn, factor, factor)/(factor*factor) # upsample will use same values so need to correct scale by factor**2
imgECAL[:140-85] = imgEEm_dn_up[1:] ## replace the old EE- rows
# EE+
imgEEp = imgECAL[140+85:] # EE+ in the last 55 rows
imgEEp = np.pad(imgEEp, ((0,1),(0,0)), 'constant', constant_values=0) # for even downsampling, zero pad 55 -> 56
imgEEp_dn = block_reduce(imgEEp, block_size=(factor, factor), func=np.sum) # downsample by summing over [factor, factor] window
imgEEp_dn_up = upsample_array(imgEEp_dn, factor, factor)/(factor*factor) # upsample will use same values so need to correct scale by factor*factor
imgECAL[140+85:] = imgEEp_dn_up[:-1] # replace the old EE+ rows
return imgECAL
def crop_jet(imgECAL, iphi, ieta, jet_shape=125):
# NOTE: jet_shape here should correspond to the one used in RHAnalyzer
off = jet_shape//2
iphi = int(iphi*5 + 2) # 5 EB xtals per HB tower
ieta = int(ieta*5 + 2) # 5 EB xtals per HB tower
# Wrap-around on left side
if iphi < off:
diff = off-iphi
img_crop = np.concatenate((imgECAL[:,ieta-off:ieta+off+1,-diff:],
imgECAL[:,ieta-off:ieta+off+1,:iphi+off+1]), axis=-1)
# Wrap-around on right side
elif 360-iphi < off:
diff = off - (360-iphi)
img_crop = np.concatenate((imgECAL[:,ieta-off:ieta+off+1,iphi-off:],
imgECAL[:,ieta-off:ieta+off+1,:diff+1]), axis=-1)
# Nominal case
else:
img_crop = imgECAL[:,ieta-off:ieta+off+1,iphi-off:iphi+off+1]
return img_crop
rhTreeStr = args.infile
rhTree = ROOT.TChain("fevt/RHTree")
rhTree.Add(rhTreeStr)
nEvts = rhTree.GetEntries()
assert nEvts > 0
print " >> Input file:",rhTreeStr
print " >> nEvts:",nEvts
outStr = '%s/%s.parquet.%d'%(args.outdir, args.decay, args.idx)
print " >> Output file:",outStr
##### MAIN #####
# Event range to process
iEvtStart = 0
#iEvtEnd = 10
iEvtEnd = nEvts
assert iEvtEnd <= nEvts
print " >> Processing entries: [",iEvtStart,"->",iEvtEnd,")"
nJets = 0
data = {} # Arrays to be written to parquet should be saved to data dict
sw = ROOT.TStopwatch()
sw.Start()
for iEvt in range(iEvtStart,iEvtEnd):
# Initialize event
rhTree.GetEntry(iEvt)
if iEvt % 10000 == 0:
print " .. Processing entry",iEvt
ECAL_energy = np.array(rhTree.ECAL_energy).reshape(280,360)
ECAL_energy = resample_EE(ECAL_energy)
HBHE_energy = np.array(rhTree.HBHE_energy).reshape(56,72)
HBHE_energy = upsample_array(HBHE_energy, 5, 5) # (280, 360)
TracksAtECAL_pt = np.array(rhTree.ECAL_tracksPt).reshape(280,360)
data['X_CMSII'] = np.stack([TracksAtECAL_pt, ECAL_energy, HBHE_energy], axis=0) # (3, 280, 360)
# Jet attributes
ys = rhTree.jetIsQuark
pts = rhTree.jetPt
iphis = rhTree.jetSeed_iphi
ietas = rhTree.jetSeed_ieta
pdgIds = rhTree.jetPdgIds
njets = len(ys)
for i in range(njets):
data['y'] = ys[i]
data['pt'] = pts[i]
data['iphi'] = iphis[i]
data['ieta'] = ietas[i]
data['pdgId'] = pdgIds[i]
data['X_jet'] = crop_jet(data['X_CMSII'], data['iphi'], data['ieta']) # (3, 125, 125)
# Create pyarrow.Table
pqdata = [pa.array([d]) if (np.isscalar(d) or type(d) == list) else pa.array([d.tolist()]) for d in data.values()]
table = pa.Table.from_arrays(pqdata, data.keys())
if nJets == 0:
writer = pq.ParquetWriter(outStr, table.schema, compression='snappy')
writer.write_table(table)
nJets += 1
writer.close()
print " >> nJets:",nJets
print " >> Real time:",sw.RealTime()/60.,"minutes"
print " >> CPU time: ",sw.CpuTime() /60.,"minutes"
print "========================================================"
# Verify output file
pqIn = pq.ParquetFile(outStr)
print(pqIn.metadata)
print(pqIn.schema)
X = pqIn.read_row_group(0, columns=['y','pt','iphi','ieta','pdgId']).to_pydict()
print(X)
#X = pqIn.read_row_group(0, columns=['X_jet.list.item.list.item.list.item']).to_pydict()['X_jet'] # read row-by-row
#X = pqIn.read(['X_jet.list.item.list.item.list.item', 'y']).to_pydict()['X_jet'] # read entire column(s)
#X = np.float32(X)