-
Notifications
You must be signed in to change notification settings - Fork 3
/
MeritSkimmerWithFriends.py
108 lines (94 loc) · 5.27 KB
/
MeritSkimmerWithFriends.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
#!/usr/bin/env python2
import sys
import os
import os.path as path
import subprocess
import time
import click
import ROOT
from ROOT import gROOT, gDirectory, gPad, gSystem, gStyle, kTRUE, kFALSE, TEntryList, TEventList, TTree, TChain, TFile, TList
ROOT.gROOT.SetBatch()
from logging import getLogger,StreamHandler,DEBUG,INFO,WARNING,ERROR,CRITICAL
##### Logger #####
logger = getLogger(__name__)
handler = StreamHandler()
loglevel = 'INFO'
handler.setLevel(loglevel)
logger.setLevel(loglevel)
logger.addHandler(handler)
@click.command()
@click.argument('pathin', type=str)
@click.option('--cut', type=str, default='U01_V200909_62IRFBK_020RAWE20woRW_09_BDTG200D05>=0.0 && EvtEventId%5==0 && CalOnly==1 && Cal1MomZDir>=0.2 && log10(WP8CalOnlyEnergy)>=4.35 && log10(WP8CalOnlyEnergy)<=5.75 && Cal1MomZCrossSide840>=0 && (GltGemSummary&0x20)==0') #"U01_V200909_62IRFBK_020RAWE20woRW_09_BDTG200D05>=0.0 && CalOnly==1 && Cal1MomZDir>=0.2 && log10(WP8CalOnlyEnergy)>=4.35 && log10(WP8CalOnlyEnergy)<=5.75 && McSourceId!=7000 && Cal1MomZCrossSide840>=0 && (GltGemSummary&0x20)==0")
@click.option('--branches', type=str, default='/u/gl/mtakahas/work/data/lists/branchlist_forEval_2019Oct.yaml')
@click.option('--suffix', type=str, default='Skimmed')
@click.option('--friend', '-f', multiple=True, default=[])
@click.option('--skipskim', is_flag=True)
@click.option('--outdir', type=str, default='.')
@click.option('--loglevel', type=click.Choice(['DEBUG', 'INFO', 'WARNING', 'CRITICAL']), default='INFO')
@click.option('--logout', type=str, default='./MeritSkimmerWithFriends.log')
def main(pathin, cut, branches, suffix, friend, skipskim, outdir, loglevel, logout):
##### Logger #####
handler.setLevel(loglevel)
logger.setLevel(loglevel)
logger.addHandler(handler)
filein = ROOT.TFile(pathin, "READ")
treein = filein.Get("MeritTuple")
logger.info('Input tree has {0} events.'.format(treein.GetEntries()))
pathskimmed = '/'.join([outdir, path.basename(pathin).replace(".root", "_{0}.root".format(suffix))])
if skipskim==False:
# mscmd = ['bsub', '-o{0}'.format(logout), '-JskimFr', '-W1000', '-We100', 'python', '/u/gl/mtakahas/eventSelectForTMVA42/scripts/MeritSkimmer.py', '--aliases=pass8', '--selection="{0}"'.format(cut), '--extra_trees=jobinfo', pathin, '--output={0}'.format(pathskimmed), '--branches={0}'.format(branches)]
mscmd = ['python', '/u/gl/mtakahas/eventSelectForTMVA42/scripts/MeritSkimmer.py', '--aliases=pass8', '--selection="{0}"'.format(cut), '--extra_trees=jobinfo', pathin, '--output={0}'.format(pathskimmed), '--branches={0}'.format(branches)]
logger.debug(mscmd)
proc_skim = subprocess.Popen(mscmd)
#time.sleep(1000*60)
treein.Draw(">>enlist", cut) #, "entrylist")
enlist = gDirectory.Get("enlist")
logger.info("Event list: {0}".format(enlist.GetN()))
dictfriends = {}
dict_friendsuffix = {'WP8CalOnlyBEP_E_BDT': 'WP8CalOnlyBEPCaseE_myBDT',
'UpDown':'U01_V200909_62IRFBK_020RAWE20woRW_09',
'UpDownZDir':'U01_V200909_62IRFBK_020RAWE20woRW_09ZDir_MinNode2_D06_MinNode00005'}
if len(friend)==0:
logger.warning('List of friends in {0} is referred.'.format(treein.GetName()))
rlistfriends = treein.GetListOfFriends()
for ifr in range(rlistfriends.GetEntries()):
logger.debug(rlistfriends.At(ifr).GetName())
dictfriends[rlistfriends.At(ifr).GetName()] = rlistfriends.At(ifr).GetTitle()
else:
for fr in friend:
logger.info(fr)
if fr in dict_friendsuffix.keys():
dictfriends[fr] = pathin.replace(".root", "_{0}.root".format(dict_friendsuffix[fr]))
else:
dictfriends[fr] = pathin.replace(".root", "_{0}.root".format(fr))
logger.debug(dictfriends)
for namefriend, pathfriendin in dictfriends.items():
logger.info(pathfriendin)
filefriendin = ROOT.TFile(pathfriendin, "READ")
treefriendin = filefriendin.Get("MeritTuple")
logger.info('Friend tree has {0} events.'.format(treefriendin.GetEntries()))
treefriendin.SetEventList(enlist)
#treefriendin.SetEntryList(enlist)
if namefriend in dict_friendsuffix.keys():
friendsuffix = dict_friendsuffix[namefriend]
else:
friendsuffix = namefriend
pathfriendskimmed = '/'.join([outdir, path.basename(pathskimmed).replace(suffix, '_'.join([suffix,friendsuffix]))]) #pathskimmed.replace(".root", "_{0}.root".format(namefriend))
logger.debug(pathfriendskimmed)
filefriendskimmed = ROOT.TFile(pathfriendskimmed, "RECREATE")
filefriendskimmed.cd()
treefriendskimmed = treefriendin.CopyTree("")
logger.info('Skimmed friend tree has {0} events.'.format(treefriendskimmed.GetEntries()))
treefriendskimmed.Write()
#treeskimmed.AddFriend(treefriendskimmed, namefriend)
#fileskimmed.cd()
#treeskimmed.Write()
#fileskimmed = ROOT.TFile(pathskimmed, "UPDATE")
#treeskimmed = fileskimmed.Get("MeritTuple")
#logger.info('Skimmed tree has {0} events.'.format(treeskimmed.GetEntries()))
#treeskimmed.GetListOfFriends().RemoveAll()
if skipskim==False:
logger.info('Waiting the skimming process...')
proc_skim.wait()
if __name__ == '__main__':
main()