def main():
# Arguments:
assert len(sys.argv) > 1
f_in = sys.argv[1]
with open("cuts.yaml") as f:
known_cuts = yaml.load(f)
cut_key = "pre"
if len(sys.argv) > 2:
if sys.argv[2] in known_cuts:
cut_key = sys.argv[2]
tcut = TCut(cut_key, known_cuts[cut_key])
cut = tcut.GetTitle()
print "[..] Applying the following cut to {}:\n{}\n".format(f_in, cut)
f_out = f_in.replace(".root", "_{}.root".format(cut_key))
# Apply the cut:
rf = root.rfile(f_in)
tts = rf.get_ttrees()
tf_out = TFile(f_out, "RECREATE")
for tt_name, tt in tts.items():
print "[..] Cutting {}.".format(tt.GetName())
tt_out = tt.CopyTree(cut)
tf_out.WriteTObject(tt_out)
tf_out.Close()
return True
def get_inner_efficiency(roc=1, q=.2):
i_pl = 2 if roc == 1 else 1
cut = TCut(f'n_clusters[0] == 1 & n_clusters[3] == 1 & n_clusters[{i_pl}] == 1') # select only events with exactly one cluster in the other three planes
n_tracks = t.GetEntries(cut.GetTitle())
info(f'Found {n_tracks} tracks')
cut += f'n_clusters[{roc}] > 0'
xfits, chi_x, yfits, chi_y = fit_tracks(roc, cut)
z_ = get_z_positions()
(x, y), n = z.get_tree_vec(get_track_vars(roc=roc), cut=cut), z.get_tree_vec(f'n_clusters[{roc}]', cut, dtype='i2')
chi_cut = (chi_x < quantile(chi_x, q)) & (chi_y < quantile(chi_y, q))
dx, dy = array([polyval(xfits.repeat(n, axis=1), z_[roc]) - x, polyval(yfits.repeat(n, axis=1), z_[roc]) - y]) * 10000 # to um
rcut = get_res_cut(dx, n, chi_cut, 2 * z.Plane.PX * 1000) & get_res_cut(dy, n, chi_cut, 2 * z.Plane.PY * 1000)
# n_good = round(mean([get_res_cut(dx, n, chi_cut), get_res_cut(dy, n, chi_cut)]))
n_max = count_nonzero(chi_cut) # TODO fix wrong number of events
n_good = rcut[rcut].size
eff = calc_eff(n_good, n_max)
info('Efficiency for plane {}: ({:1.1f}+{:1.1f}-{:1.1f})%'.format(roc, *eff))
return eff
def makecuts(self, Tree=None, cutstring="", extraCuts="", name=None):
''' This Function is to apply selection on specific tree'''
from ROOT import TCut
#print(name)
CUTtext = open(self.outdir + "/cuts.txt", "w+")
CUTtext.write(cutstring + extraCuts + '\n')
cutstring = TCut(cutstring + extraCuts)
#print(cutstring)
cut = cutstring.GetTitle()
# when doing the cutflow, make a temp root file to avoid using a lot of memory
if name:
fname = self.outdir + "/" + name + ".root"
#check if the file is already existing to avoid reproducing it
if os.path.exists(fname):
#open it and try to check if the tree is stored inside
tt_out = self.makeChain(filesList=[fname], Tname="t")
if tt_out.LoadTree(0) >= 0:
return tt_out
else:
print(
name,
"the chain is not saved properly, will reproduce it")
tempFile = ROOT.TFile(self.outdir + "/" + name + ".root",
"recreate")
#print(Tree.LoadTree(0))
tt_out = Tree.CopyTree(cut)
# need to use Write() here otherwise it will not be overwritten
tt_out.Write()
tempFile.Close()
# Stupied to write/Close the file to be able to load the chain correctly
tt_out = self.makeChain(filesList=[fname], Tname="t")
return tt_out
# if file is not existing, poduce it
else:
tempFile = ROOT.TFile(self.outdir + "/" + name + ".root",
"recreate")
# the normal workflow
if Tree == None:
print('no Tree founded please cheack')
pass
else:
tt_out = Tree.CopyTree(cut)
return tt_out
pass
def main(o, args):
# Import TMVA classes from ROOT
from ROOT import TMVA, TFile, TCut
print o
# Output file
if o.restrict:
o.label += "_%s" % o.restrict
outputFile = TFile(o.outfile % {"label": o.label}, 'RECREATE')
atype = "Classification"
if hasattr(o, "type"):
atype = str(o.type)
factory = TMVA.Factory(
"TMVAClassification", outputFile,
"!V:!Silent:!Color:!DrawProgressBar:Transformations=I:AnalysisType=%s"
% atype)
# Set verbosity
factory.SetVerbose(o.verbose)
TMVA.Config.Instance().GetIONames().fWeightFileDir = o.weightsdir
# variables
if type(o.variables) == str:
o.variables = [
v.lstrip().rstrip() for v in o.variables.split(":") if v != ""
]
allvars = ""
for v in o.variables:
factory.AddVariable(str(v))
if allvars != "": allvars += ":"
allvars += v.split(":=")[0].lstrip(" ").rstrip(" ")
print "variables %s" % allvars
print o.spectators
for s in o.spectators:
if not s in o.variables:
factory.AddSpectator(str(s))
# categories and sub categories
categories = []
subcategories = []
if hasattr(o, "subcategories") and len(o.subcategories) > 0:
subcategories = o.subcategories[0]
for sc in o.subcategories[1:]:
subcategories = map(
lambda x: (TCut(x[0][0]) * TCut(x[1][0]), "%s_%s" %
(x[0][1], x[1][1])),
itertools.product(subcategories, sc))
for cut, name, vrs in o.categories:
myvars = allvars
if vrs != "":
for v in vrs.split(":"):
myvars = myvars.replace(v, "").replace("::", ":")
myvars = myvars.rstrip(":")
vrs = str(myvars)
print vrs
if len(subcategories) > 0:
for subcut, subname in subcategories:
if subname == "":
subname = subname.replace(" ", "").replace(
">", "_gt_").replace("<", "_lt_").replace(
"=", "_eq_").replace("&", "_and_")
fullname = "%s_%s" % (name, subname)
categories.append(
(TCut(cut) * TCut(subcut), str(fullname), vrs))
else:
categories.append((TCut(cut), str(name), vrs))
if o.restrict:
categories = filter(lambda x: x[1] == o.restrict, categories)
# load tree
selection = TCut(o.selection)
if len(categories) == 1:
selection *= categories[0][0]
allvars = categories[0][2]
categories = []
o.categories = []
aliases = {}
for a in o.aliases:
name, vdef = [str(t.lstrip(" ").rstrip(" ")) for t in a.split(":=", 1)]
aliases[name] = vdef
if len(o.categories) > 0:
allcats = "(%s)" % ")||(".join([c[0] for c in o.categories])
print allcats
selection = selection * TCut(allcats)
tmp = TFile.Open("/tmp/musella/tmp%s.root" % o.label, "recreate")
for evclass, info in o.classes.iteritems():
samples = info["samples"]
for samp in samples:
if len(samp) == 4:
name, weight, cut, ttype = samp
friend = None
else:
name, weight, cut, ttype, friend = samp
tcut = TCut(cut) * selection
fullTree = mkChain(getListOfFiles(o.indir, o.files), name, aliases)
selTree = fullTree.CopyTree(tcut.GetTitle())
if friend:
if type(friend) != list:
friends = [friend]
else:
friends = friend
for ifr, friend in enumerate(friends):
print friend
fname, tname = friend.split("::")
friend = mkChain(getListOfFiles(o.indir, fname), tname)
friend.AddFriend(fullTree)
selFriend = friend.CopyTree(tcut.GetTitle())
if len(friends) == 1:
frname = "fr"
else:
frname = "fr%d" % ifr
selTree.AddFriend(selFriend, frname)
setAliases(selTree, aliases)
factory.AddTree(selTree, str(evclass), float(weight), TCut(""),
int(ttype))
# weights
if "weight" in info:
weight = info["weight"]
factory.AddSpectator(str("%s_wei := %s" % (evclass, weight)))
factory.SetWeightExpression(str(weight), str(evclass))
else:
factory.SetWeightExpression("1.", str(evclass))
# "SplitMode=Random" means that the input events are randomly shuffled before
# splitting them into training and test samples
factory.PrepareTrainingAndTestTree(
TCut(""), "SplitMode=Random:NormMode=NumEvents:!V")
# --------------------------------------------------------------------------------------------------
defaultSettings = {
"BDT":
"!H:!V:!CreateMVAPdfs:BoostType=Grad:UseBaggedGrad"
":GradBaggingFraction=0.6:SeparationType=GiniIndex:nCuts=20:NNodesMax=5"
":Shrinkage=0.3:NTrees=1000",
"Cuts":
"!H:!V:FitMethod=MC:EffSel:SampleSize=200000:VarProp=FSmart"
}
if "FisherD" in o.methods:
mname = "FisherD%s" % o.label
fcats = factory.BookMethod(TMVA.Types.kCategory, mname)
for cut, name, vars in categories:
print "booking sub-category classifier : %s %s %s" % (cut, name,
vars)
fcats.AddMethod(cut, vars, TMVA.Types.kFisher,
"%s_%s" % (mname, name),
"!H:!V:Fisher:!CreateMVAPdfs:VarTransform=D")
if "Fisher" in o.methods or "LD" in o.methods:
mname = str("Fisher%s" % o.label)
fcats = factory.BookMethod(TMVA.Types.kCategory, mname)
for cut, name, vars in categories:
print "booking sub-category classifier : %s %s %s" % (cut, name,
vars)
fcats.AddMethod(TCut(cut), vars, TMVA.Types.kFisher,
"%s_%s" % (mname, name),
"!H:!V:Fisher:!CreateMVAPdfs")
if "Likelihood" in o.methods:
mname = "Likelihood%s" % o.label
fcats = factory.BookMethod(TMVA.Types.kCategory, mname)
for cut, name, vars in categories:
print "booking sub-category classifier : %s %s %s" % (cut, name,
vars)
fcats.AddMethod(
cut, vars, TMVA.Types.kLikelihood, "%s_%s" % (mname, name),
"!H:!V:!CreateMVAPdfs:!TransformOutput:PDFInterpol=KDE:KDEtype=Gauss:KDEiter=Adaptive:KDEFineFactor=0.3:KDEborder=None:NAvEvtPerBin=150"
)
if "LikelihoodD" in o.methods:
mname = "LikelihoodD%s" % o.label
fcats = factory.BookMethod(TMVA.Types.kCategory, mname)
for cut, name, vars in categories:
print "booking sub-category classifier : %s %s %s" % (cut, name,
vars)
fcats.AddMethod(
cut, vars, TMVA.Types.kLikelihood, "%s_%s" % (mname, name),
"!H:!V:!CreateMVAPdfs:!TransformOutput:VarTransform=D:PDFInterpol=KDE:KDEtype=Gauss:KDEiter=Adaptive:KDEFineFactor=0.3:KDEborder=None:NAvEvtPerBin=150"
)
if "BDT" in o.methods:
mname = str("BDT%s" % o.label)
settings = defaultSettings["BDT"]
if hasattr(o, "settings") and "BDT" in o.settings:
settings = str(o.settings["BDT"])
if len(categories) == 0:
cats = factory.BookMethod(TMVA.Types.kBDT, mname, settings)
else:
cats = factory.BookMethod(TMVA.Types.kCategory, mname)
for cut, name, vars in categories:
print "booking sub-category classifier : %s %s %s" % (
cut, name, vars)
cats.AddMethod(cut, vars, TMVA.Types.kBDT,
"%s_%s" % (mname, name), settings)
if "Cuts" in o.methods:
mname = str("Cuts%s" % o.label)
settings = defaultSettings["Cuts"]
if hasattr(o, "settings") and "Cuts" in o.settings:
settings = str(o.settings["Cuts"])
if len(categories) == 0:
cats = factory.BookMethod(TMVA.Types.kCuts, mname, settings)
else:
cats = factory.BookMethod(TMVA.Types.kCategory, mname)
for cut, name, vars in categories:
print "booking sub-category classifier : %s %s %s" % (
cut, name, vars)
cats.AddMethod(cut, vars, TMVA.Types.kCuts,
"%s_%s" % (mname, name), settings)
# ---- Now you can tell the factory to train, test, and evaluate the MVAs.
if o.optimize:
print "Optimizing?"
factory.OptimizeAllMethods()
factory.TrainAllMethods()
factory.TestAllMethods()
factory.EvaluateAllMethods()
# Save the output.
outputFile.Close()
from ROOT import TROOT, TCanvas, TH1F, TChain, TLine, TCut
# load files
chain = TChain('D2KSPiPiPi_DD/DecayTree')
chain.Add('/eos/lhcb/user/m/mamartin/D2Kshhh/Test/MD/*.root')
print 'Number of entries: ', chain.GetEntries()
#prints in the command line the number of entries
SW1 = TCut("D_M>1830")
SW2 = TCut("D_M<1900")
SW = TCut(SW1.GetTitle() + '&&' + SW2.GetTitle())
SW.Print()
#D_M
h0_D_M = TH1F("h0_D_M", "D_M", 100, 1600, 2200)
#from 0 to 600000, 100 bins
#chain.Project('h0_D_M', 'D_M', SW.GetTitle())
chain.Project(
'h0_D_M', 'D_M',
"((D_M>1600)&&(D_M<1700))||((D_M>1700)&&(D_M>1800))||((D_M>1800)&&(D_M>1900))"
)
#this fills the histogram h1 with variable D_M
c0_D_M = TCanvas('c0_D_M', 'c0_D_M', 200, 10, 700, 500)
h0_D_M.Draw()
c0_D_M.Print("D_M0.pdf")
# cuts
from ROOT import TCut
cut1 = TCut('a<10')
cut2 = TCut('b%2==0')
cut = TCut('(%s)&&(%s)' % (cut1.GetTitle(), cut2.GetTitle()))
print cut.GetTitle()
