-
Notifications
You must be signed in to change notification settings - Fork 1
/
step2_adddiscriminants.py
executable file
·143 lines (125 loc) · 5.23 KB
/
step2_adddiscriminants.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
#!/usr/bin/env python
import argparse
from collections import namedtuple
class stringandlambda(namedtuple("stringandlambda", "string function")):
def __new__(cls, string):
return super(stringandlambda, cls).__new__(cls, string, eval(string))
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--submitjobs", action="store_true")
p.add_argument("--filter", type=stringandlambda, default=None)
p.add_argument("--skipxcheck", dest="doxcheck", action="store_false")
args = p.parse_args()
from array import array
from collections import OrderedDict
from helperstuff import config
from helperstuff import xrd
from helperstuff.enums import hffhypotheses, ProductionMode, productions, pythiasystematics
from helperstuff.samples import allsamples, Sample
from helperstuff.submitjob import submitjob
from helperstuff.treewrapper import TreeWrapper, TreeWrapperFactory
from helperstuff.utilities import cdtemp_slurm, cleanupscratchdir, deletemelastuff, KeepWhileOpenFile, LSB_JOBID, LSF_creating, mkdir_p
import os
import pipes
import ROOT
import sys
def adddiscriminants(*args):
sample = Sample(*args)
newfilename = sample.withdiscriminantsfile()
print newfilename
mkdir_p(os.path.dirname(newfilename))
if sample.copyfromothersample is not None:
if sample.CJLSTfile() != sample.copyfromothersample.CJLSTfile() and xrd.exists(sample.CJLSTfile()):
raise ValueError("{} exists, why not use it?".format(sample.CJLSTfile()))
return
inputfiles = []
if xrd.exists(sample.CJLSTfile()): inputfiles.append(sample.CJLSTfile())
with cdtemp_slurm(), KeepWhileOpenFile(newfilename+".tmp") as kwof:
if not kwof:
return
if os.path.exists(newfilename):
return
with LSF_creating(newfilename, ignorefailure=True, inputfiles=inputfiles) as LSF:
treewrapper = TreeWrapperFactory(sample, LSF=LSF)
if os.path.exists(newfilename):
return
failed = False
try:
newf = ROOT.TFile.Open(LSF.basename(newfilename), "recreate")
if isinstance(treewrapper, TreeWrapper):
newt = treewrapper.tree.CloneTree(0)
else:
newt = ROOT.TTree("candTree", "candTree")
discriminants = OrderedDict()
for discriminant in treewrapper.toaddtotree:
discriminants[discriminant] = array('d', [0])
newt.Branch(discriminant, discriminants[discriminant], discriminant + "/D")
for discriminant in treewrapper.toaddtotree_int:
discriminants[discriminant] = array('i', [0])
newt.Branch(discriminant, discriminants[discriminant], discriminant + "/I")
for discriminant in treewrapper.toaddtotree_float:
discriminants[discriminant] = array('f', [0])
newt.Branch(discriminant, discriminants[discriminant], discriminant + "/F")
try:
for i, entry in enumerate(treewrapper, start=1):
for discriminant in discriminants:
try:
discriminants[discriminant][0] = getattr(treewrapper, discriminant)()
except:
print "Error while calculating", discriminant
raise
newt.Fill()
if i % 50000 == 0 and LSB_JOBID():
with open(LSF.basename(newfilename)): pass #access it, hopefully preventing tmp from being deleted
with open("touch.txt", "w") as f: pass #touch another file, same idea
except:
treewrapper.Show()
raise
except:
failed = True
raise
finally:
try:
newf.Write()
newf.Close()
except:
failed = True
if failed:
try:
os.remove(newfilename)
except:
pass
def submitjobs(filter=None, doxcheck=True):
njobs = 0
for sample in allsamples(doxcheck=doxcheck):
if filter and not filter.function(sample): continue
if not os.path.exists(sample.withdiscriminantsfile()) and KeepWhileOpenFile(sample.withdiscriminantsfile()+".tmp").wouldbevalid:
njobs += 1
for i in range(njobs):
submitjobkwargs = {"jobname": str(i), "jobtime": "1-0:0:0"}
if config.host == "MARCC":
submitjobkwargs["queue"] = "shared"
submitjobkwargs["memory"] = "4000M"
submitjobkwargs["docd"] = True #since cdtemp_slurm happens in thisfile
job = ["unbuffer", os.path.join(config.repositorydir, "step2_adddiscriminants.py")]
if filter:
job += ["--filter", filter.string]
job = " ".join(pipes.quote(_) for _ in job)
submitjob(job, **submitjobkwargs)
if __name__ == '__main__':
cleanupscratchdir()
if args.submitjobs:
submitjobs(filter=args.filter, doxcheck=args.doxcheck)
else:
try:
for sample in allsamples(doxcheck=args.doxcheck):
if sample.productionmode == "ggZZ" and sample.flavor == "4tau" and not sample.copyfromothersample:
adddiscriminants(sample)
for sample in allsamples(doxcheck=args.doxcheck):
if args.filter and not args.filter.function(sample): continue
while sample.copyfromothersample:
sample = sample.copyfromothersample
adddiscriminants(sample)
finally:
if not any(KeepWhileOpenFile(sample.withdiscriminantsfile()+".tmp").wouldbevalid for sample in allsamples(doxcheck=args.doxcheck)):
deletemelastuff()