convert_root2pq_EBshower_byEvt.py

import ROOT
import numpy as np
import glob, os
import pyarrow as pa
import pyarrow.parquet as pq

import argparse
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-i', '--infile', default=['output.root'], nargs='+', type=str, help='Input root file.')
parser.add_argument('-o', '--outdir', default='parquet', type=str, help='Output pq file dir.')
parser.add_argument('-d', '--decay', default='', type=str, help='Decay name.')
parser.add_argument('-n', '--idx', default=0, type=int, help='Input root file index.')
args = parser.parse_args()

doPho = True
doPlots = False
#infile_str = args.infile[0].split('/')[-4].replace('_reg','')
infile_str = args.infile[0].split('/')[-1]
print(infile_str)

rhTreeStr = args.infile
print " >> Input file[0]:",rhTreeStr[0]
rhTree = ROOT.TChain("fevt/RHTree")
for f in rhTreeStr:
  rhTree.Add(f)
nEvts = rhTree.GetEntries()
assert nEvts > 0
print " >> nEvts:",nEvts
outStr = '%s/test.parquet'%(args.outdir)
print " >> Output file:",outStr

##### EVENT SELECTION START #####

# Event range to process
iEvtStart = 0
#iEvtEnd   = 1000
iEvtEnd   = nEvts
assert iEvtEnd <= nEvts
print " >> Processing entries: [",iEvtStart,"->",iEvtEnd,")"

nPhos = 0
nEvts = 0
nHLTs = 0
d = {} # 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

    d['idx'] = [rhTree.runId, rhTree.lumiId, rhTree.eventId]
    #d['evt_wgt'] = rhTree.evt_weight
    #print(d['idx'])

    d['hltA'] = rhTree.hltAccept
    #nHLTs += d['hltA']

    #d['m'] = list(rhTree.A_mass)
    #d['pt'] = list(rhTree.A_pT)
    d['A_recoIdx'] = list(rhTree.A_recoIdx)

    d['mHgen'] = rhTree.mHgen
    d['A_p4'] = np.transpose([rhTree.A_E, rhTree.A_pT, rhTree.A_eta, rhTree.A_phi], [1,0])
    #d['A_m'] = list(rhTree.A_mass)
    #d['A_dR'] = list(rhTree.A_DR)
    #d['A_pdgId'] = list(rhTree.A_pdgId)
    #d['A_mothPdgId'] = list(rhTree.A_mothPdgId)
    #d['A_jetM'] = list(rhTree.A_jetM)

    #d['OutPart_pdgId'] = list(rhTree.OutPart_pdgId)
    #d['A_ancestry'] = np.transpose([rhTree.A_pdgId, rhTree.A_mothPdgId, rhTree.OutPart_pdgId], [1,0])

    d['mGG']  = rhTree.m0
    #if d['mGG'] < 100. or d['mGG'] > 110.: continue
    #if d['mGG'] < 100.: continue
    #if d['mGG'] > 110. and d['mGG'] < 140.: continue
    #if d['mGG'] > 180.: continue

    d['nRecoPho'] = rhTree.nRecoPho
    #d['minDR'] = list(rhTree.minDR)

    d['pho_p4'] = np.transpose([rhTree.pho_E, rhTree.pho_pT, rhTree.pho_eta, rhTree.pho_phi], [1,0])
    d['pho_id'] = np.transpose([
            rhTree.pho_r9
            ,rhTree.pho_sieie
            ,rhTree.pho_phoIso
            ,rhTree.pho_chgIso
            ,rhTree.pho_chgIsoWrongVtx
            ,rhTree.pho_Eraw
            ,rhTree.pho_phiWidth
            ,rhTree.pho_etaWidth
            ,rhTree.pho_scEta
            ,rhTree.pho_sieip
            ,rhTree.pho_s4
        ], [1,0])
    d['pho_vars'] = np.transpose([
            rhTree.pho_r9
            ,rhTree.pho_HoE
            ,rhTree.pho_hasPxlSeed
            ,rhTree.pho_sieie
            ,rhTree.pho_phoIso
            ,rhTree.pho_trkIso
            ,rhTree.pho_chgIsoCorr
            ,rhTree.pho_neuIsoCorr
            ,rhTree.pho_phoIsoCorr
            ,rhTree.pho_bdt
        ], [1,0])
    #d['pho_Ecorr'] = list(rhTree.pho_ecalEPostCorr)

    d['iphi'] = list(rhTree.SC_iphi)
    d['ieta'] = list(rhTree.SC_ieta)

    d['X'] = np.array(rhTree.SC_energy).reshape(-1,1,32,32) # nPho, nChannel, nRows, nCols
    d['Xtz'] = np.transpose([rhTree.SC_energyT, rhTree.SC_energyZ], [1,0,2]).reshape(-1,2,32,32) # nPho, nChannel, nRows, nCols
    d['X_aod'] = np.array(rhTree.SCaod_energy).reshape(-1,1,32,32) # nPho, nChannel, nRows, nCols
    d['Xtz_aod'] = np.transpose([rhTree.SCaod_energyT, rhTree.SCaod_energyZ], [1,0,2]).reshape(-1,2,32,32) # nPho, nChannel, nRows, nCols

    #print(d['Xtz'][0].min(), d['Xtz'][0].max(), d['iphi'][0], d['ieta'][0])
    #print(d['Xtz'][1].min(), d['Xtz'][1].max(), d['iphi'][1], d['ieta'][1])
    #d['X_EB'] = np.array(rhTree.EB_energy).reshape(170,360)

    pqdata = [pa.array([d_]) if np.isscalar(d_) or type(d_) == list else pa.array([d_.tolist()]) for d_ in d.values()]
    table = pa.Table.from_arrays(pqdata, d.keys())

    if nEvts == 0:
        writer = pq.ParquetWriter(outStr, table.schema, compression='snappy')
    writer.write_table(table)

    nEvts += 1

sw.Stop()
print " >> nHLTpass:",nHLTs
print " >> nDiPhoEvents:",nEvts
print " >> nPhos:",nPhos
print " >> Real time:",sw.RealTime()/60.,"minutes"
print " >> CPU time: ",sw.CpuTime() /60.,"minutes"
print " >> ======================================"