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()