Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-02-14 12:48:16

0001 ######################################################################################
0002 # Makes pkl and text files comparing PU and noPU samples for training regressor and other stuff
0003 # Usage:
0004 # python3 isotrackNtupler.py -PU root://cmseos.fnal.gov//store/user/sghosh/ISOTRACK/DIPI_2021_PUpart.root -NPU root://cmseos.fnal.gov//store/user/sghosh/ISOTRACK/DIPI_2021_noPU.root -O isotk_relval 
0005 ######################################################################################
0006 
0007 
0008 
0009 import uproot
0010 import numpy as np
0011 import pandas as pd
0012 import argparse
0013 import matplotlib.pyplot as plt
0014 
0015 parser = argparse.ArgumentParser()
0016 parser.add_argument("-PU", "--filePU",help="input PU file",default="root://cmseos.fnal.gov//store/user/sghosh/ISOTRACK/DIPI_2021_PUpart.root")
0017 parser.add_argument("-NPU", "--fileNPU",help="input no PU file",default="root://cmseos.fnal.gov//store/user/sghosh/ISOTRACK/DIPI_2021_noPU.root")
0018 parser.add_argument("-O", "--opfilename",help="ouput file name",default="isotk_relval")
0019 
0020 
0021 fName1 = parser.parse_args().filePU
0022 fName2 = parser.parse_args().fileNPU
0023 foutput = parser.parse_args().opfilename
0024 
0025 
0026 # PU
0027 tree1 = uproot.open(fName1,xrootdsource=dict(chunkbytes=1024**3, limitbytes=1024**3))['HcalIsoTrkAnalyzer/CalibTree']
0028 
0029 #no PU
0030 tree2 = uproot.open(fName2,xrootdsource=dict(chunkbytes=1024**3, limitbytes=1024**3))['HcalIsoTrkAnalyzer/CalibTree']
0031 #tree2.keys()
0032 print ("loaded files")
0033 branchespu = ['t_Run','t_Event','t_nVtx','t_ieta','t_iphi','t_p','t_pt','t_gentrackP','t_eMipDR','t_eHcal','t_eHcal10','t_eHcal30','t_hmaxNearP','t_emaxNearP','t_hAnnular','t_eAnnular','t_rhoh','t_selectTk','t_qltyFlag']
0034 branchesnpu = ['t_Event','t_ieta','t_iphi','t_eHcal']
0035 #dictn = tree.arrays(branches=branches,entrystart=0, entrystop=300)
0036 dictpu = tree1.arrays(branches=branchespu)
0037 dictnpu = tree2.arrays(branches=branchesnpu)
0038 dfspu = pd.DataFrame.from_dict(dictpu)
0039 dfspu.columns=branchespu
0040 dfsnpu = pd.DataFrame.from_dict(dictnpu)
0041 dfsnpu.columns=branchesnpu
0042 print ("loaded dicts and dfs")
0043 print ("PU sample size:",dfspu.shape[0])
0044 print ("noPU sample size:",dfsnpu.shape[0])
0045 merged = pd.merge(dfspu, dfsnpu , on=['t_Event','t_ieta','t_iphi'])
0046 print ("selected common events before cut:",merged.shape[0])
0047 
0048 #cuts = (merged['t_selectTk'])&(merged['t_qltyFlag'])&(merged['t_hmaxNearP']<10)&(merged['t_eMipDR_y']<1)
0049 keepvars =  ['t_nVtx','t_ieta','t_eHcal10','t_eHcal30','t_delta','t_hmaxNearP','t_emaxNearP','t_hAnnular','t_eAnnular','t_rhoh','t_pt','t_eHcal_x','t_eHcal_y','t_p','t_eMipDR']
0050 
0051 
0052 
0053 #########################all ietas
0054 cuts1 = (merged['t_selectTk'])&(merged['t_qltyFlag'])&(merged['t_hmaxNearP']<20)&(merged['t_eMipDR']<1)&(abs(merged['t_p'] - 50)<10)&(merged['t_eHcal_x']>10)
0055 merged1=merged.loc[cuts1]
0056 merged1 = merged1.reset_index(drop=True)
0057 print ("selected events after cut for all ietas:",merged1.shape[0])
0058 merged1['t_delta']=merged1['t_eHcal30']-merged1['t_eHcal10']
0059 final_df_all = merged1[keepvars]
0060 final_df_all.to_pickle(foutput+"_all.pkl")
0061 final_df_all.to_csv(foutput+"_all.txt")
0062 #########################split ieta < 16
0063 
0064 cuts2 = (merged['t_selectTk'])&(merged['t_qltyFlag'])&(merged['t_hmaxNearP']<20)&(merged['t_eMipDR']<1)&(abs(merged['t_ieta'])<16)&(abs(merged['t_p'] - 50)<10)&(merged['t_eHcal_x']>10)
0065 merged2=merged.loc[cuts2]
0066 merged2 = merged2.reset_index(drop=True)
0067 print ("selected events after cut for ieta < 16:",merged2.shape[0])
0068 merged2['t_delta']=merged2['t_eHcal30']-merged2['t_eHcal10']
0069 final_df_low = merged2[keepvars]
0070 final_df_low.to_pickle(foutput+"_lo.pkl")
0071 final_df_low.to_csv(foutput+"_lo.txt")
0072 
0073 #########################split ieta > 15
0074 
0075 cuts3 = (merged['t_selectTk'])&(merged['t_qltyFlag'])&(merged['t_hmaxNearP']<20)&(merged['t_eMipDR']<1)&(abs(merged['t_ieta'])>15)&(abs(merged['t_p'] - 50)<10)&(merged['t_eHcal_x']>10)
0076 merged3=merged.loc[cuts3]
0077 merged3 = merged3.reset_index(drop=True)
0078 print ("selected events after cut for ieta > 15:",merged3.shape[0])
0079 merged3['t_delta']=merged3['t_eHcal30']-merged3['t_eHcal10']
0080 final_df_hi = merged3[keepvars]
0081 final_df_hi.to_pickle(foutput+"_hi.pkl")
0082 final_df_hi.to_csv(foutput+"_hi.txt")
0083