def strip_n_save (Min, Max, cuts, directory, saving_directory, extra_variables, particle, bins = False, tree = None): if not (bins): filename = "{0}2pKpiTuple.root".format(particle) alldata = TChain("tuple_{0}2pKpi/DecayTree".format(particle)) extra_dir = "" for job in range(Min, Max) : if os.path.exists("{0}/{1}{2}/{3}".format(directory,job,extra_dir,filename)): alldata.Add("{0}/{1}{2}/{3}".format(directory,job,extra_dir,filename)) #Check if there are any issues with the data if (alldata.GetEntries() == 0): print("Error: entries = 0 for range " + str(Min) + "-" + str(Max)) return if (alldata.GetEntries() == -1): print("Error: entries = -1 for range " + str(Min) + "-" + str(Max)) return alldata = setBranch_funct(alldata, extra_variables) extra_string = particle + "_cluster_{0}-{1}.root".format(Min, Max) else: if not (tree == None): alldata = tree extra_string = "" wfile = TFile.Open(saving_directory + extra_string, "RECREATE") subtree = alldata.CopyTree( cuts ) wfile.cd() subtree.Write() wfile.Close()
def skim(sample): print sample fileList = os.listdir(inDir+"/"+sample) if len(fileList)==0: return print fileList chain = TChain("tree") events_hist = TH1F('Events', 'Events', 1,0,1) for f in fileList: chain.Add(inDir+"/"+sample+"/"+f) print chain.GetNtrees(), chain.GetEntries(), inDir+"/"+sample+"/"+f fl = TFile(inDir+"/"+sample+"/"+f, "READ") events_hist.Add(fl.Get("Events")) fl.Close() #return # save to new file outFile = TFile(outDir+"/"+sample+".root", "RECREATE") outFile.cd() # create a new skimmed tree tree = chain.CopyTree(cutstring)#GetTree().CloneTree(-1, cutstring) tree.Write("", TObject.kOverwrite) events_hist.Write() outFile.Close() chain.Reset()
def readTreeFromFile(path, dileptonCombination, preselection="nJets >= 2", SingleLepton=False, use532=False): """ helper function path: path to .root file containing simulated events dileptonCombination: EE, EMu, or MuMu for electron-electron, electron-muon, or muon-muon events returns: tree containing events for on sample and dileptonCombination """ from ROOT import TChain chain = TChain() # chain.Add("%s/ETH2AachenNtuples/%sDileptonTree"%(path, dileptonCombination)) if SingleLepton: chain.Add("%s/cutsV22DileptonFinalTriggerTrees/%sDileptonTree" % (path, dileptonCombination)) elif use532: chain.Add("%s/cutsV22DileptonFinalTrees/%sDileptonTree" % (path, dileptonCombination)) else: chain.Add("%s/cutsV23DileptonFinalTrees/%sDileptonTree" % (path, dileptonCombination)) result = chain.CopyTree(preselection) # result = chain return result return chain
def save_after(path_in, path_out, cms, region): try: chain = TChain('save') chain.Add(path_in) except: logging.error(path_in + ' is invalid!') sys.exit() cut = '' if region == 'STDDmiss_signal': signal_low = 1.86965 - width(cms) / 2. signal_up = 1.86965 + width(cms) / 2. window_low = 1.86965 - window(cms) / 2. window_up = 1.86965 + window(cms) / 2. cut_base = '(m_rawm_D > ' + str(signal_low) + ' && m_rawm_D < ' + str( signal_up) + ' && m_chi2_kf < ' + str(chi2_kf(cms)) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + ' && m_rm_Dpipi < ' + str(window_up) + ')' cut_KS = '!(m_m_pipi > 0.491036 && m_m_pipi < 0.503471 && m_ctau_svf < 0.5)' cut_Dst = '(m_m_Dpi0 < 2.0082 || m_m_Dpi0 > 2.01269)' cut = cut_base + ' && ' + cut_window + ' && ' + cut_KS + ' && ' + cut_Dst if region == 'STDDmiss_sidebandlow': signal_low = 1.86965 - width(cms) / 2. signal_up = 1.86965 + width(cms) / 2. window_low = 1.86965 - window(cms) / 2. window_up = 1.86965 + window(cms) / 2. windowlow_up = window_low - (window_up - window_low) windowlow_low = windowlow_up - (window_up - window_low) cut_base = '(m_rawm_D > ' + str(signal_low) + ' && m_rawm_D < ' + str( signal_up) + ' && m_chi2_kf < ' + str(chi2_kf(cms)) + ')' cut_window = '(m_rm_Dpipi > ' + str( windowlow_low) + '&& m_rm_Dpipi < ' + str(windowlow_up) + ')' cut_KS = '!(m_m_pipi > 0.491036 && m_m_pipi < 0.503471 && m_ctau_svf < 0.5)' cut_Dst = '(m_m_Dpi0 < 2.0082 || m_m_Dpi0 > 2.01269)' cut = cut_base + ' && ' + cut_window + ' && ' + cut_KS + ' && ' + cut_Dst if region == 'STDDmiss_sidebandup': signal_low = 1.86965 - width(cms) / 2. signal_up = 1.86965 + width(cms) / 2. window_low = 1.86965 - window(cms) / 2. window_up = 1.86965 + window(cms) / 2. windowup_low = window_up + (window_up - window_low) windowup_up = windowup_low + (window_up - window_low) cut_base = '(m_rawm_D > ' + str(signal_low) + ' && m_rawm_D < ' + str( signal_up) + ' && m_chi2_kf < ' + str(chi2_kf(cms)) + ')' cut_window = '(m_rm_Dpipi > ' + str( windowup_low) + '&& m_rm_Dpipi < ' + str(windowup_up) + ')' cut_KS = '!(m_m_pipi > 0.491036 && m_m_pipi < 0.503471 && m_ctau_svf < 0.5)' cut_Dst = '(m_m_Dpi0 < 2.0082 || m_m_Dpi0 > 2.01269)' cut = cut_base + ' && ' + cut_window + ' && ' + cut_KS + ' && ' + cut_Dst t = chain.CopyTree(cut) t.SaveAs(path_out)
def save(file_in, file_out): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' cut = '(m_rm_pipi > 3.79) && (m_rm_D > 2.4 && m_rm_D < 2.45)' t = chain.CopyTree(cut) t.SaveAs(file_out)
def skim_data(wk): """ run this automatically w/ process() apply energy thresholds to create reduced size "skim" ROOT files for faster plotting --> output files: "./data/skim/week_{}.root" """ wk = "4" out_name = "./data/skim/week_{}.root".format(wk) # load metadata runs = run_info["runs"][wk] chans = run_info["goodres_chans"][wk] thresh = run_info["thresholds"][wk] # build a ROOT TCut string print("Skimming data from week", wk, "run list:", runs) cut = "(" for chan in chans: cut += "(Channel=={} && Energy>{}) || ".format(chan, thresh[int(chan)]) cut = cut[:-4] cut += ")" print(" Using cut:", cut) # load data files = [] for r in runs: # handle runs where we have multiple subfiles (_x.root) tmp = glob.glob("./data/run_{}/FILTERED/compassF_run_{}*.root".format( r, r)) files.extend(tmp) runtime = 0 ch = TChain("Data_F") for f in sorted(files): runtime += get_runtime(f) ch.Add(f) print(" Data loaded. Found", ch.GetEntries(), "entries.") print(" Runtime (hrs): {:.2f}".format(runtime)) # create a TTree from the TCut and write to the output file print("Writing to output file:\n ", out_name) out_file = TFile(out_name, "RECREATE") out_tree = TTree() out_tree = ch.CopyTree(cut) out_tree.Write() out_file.Close() # check it f2 = TFile(out_name) tt2 = f2.Get("Data_F") print("Done. Skim file has", tt2.GetEntries(), "entries.")
def save(file_in, file_out): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() m_D1_2420 = 2.4221 cut = 'abs(m_rm_D - ' + str( m_D1_2420) + ') > 0.01 && abs(m_rm_Dmiss - ' + str( m_D1_2420) + ') > 0.01' t = chain.CopyTree(cut) t.SaveAs(file_out)
def apply(file_in, file_out): try: chain = TChain('track') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' cut_dang = '(abs(m_pipi_dang)<0.98&&(flag==1||(flag==-1&&(abs(m_pimlp_dang)<0.98&&abs(m_piplm_dang)<0.98)))) && ' cut_chi2 = '(chi2_pipi_ll<60) && ' cut_m_jpsi = '(m_jpsi>3.08&&m_jpsi<3.12)' cut = cut_dang + cut_chi2 + cut_m_jpsi t = chain.CopyTree(cut) t.SaveAs(file_out)
def skim_data(weekend): rdata = [[1, 3, 4, 5], [9, 13], [27], [33], [34]] runList = rdata[weekend] thresh = {} thresh[1] = {0: 80, 1: 60, 2: 70, 3: 80} thresh[1] = {0: 80, 1: 60, 2: 70, 3: 80} thresh[3] = {0: 80, 1: 60, 2: 70, 3: 80} thresh[4] = {0: 80, 1: 60, 2: 70, 3: 80} thresh[5] = {0: 80, 1: 60, 2: 70, 3: 80} thresh[9] = {0: 80, 1: 70, 2: 100, 3: 90} thresh[13] = {0: 80, 1: 70, 2: 100, 3: 90} thresh[27] = {0: 80, 1: 75, 2: 70, 3: 90} thresh[33] = {0: 80, 1: 230, 2: 300, 3: 430} thresh[34] = {0: 80, 1: 230, 2: 300, 3: 430} # print(thresh) # for key in thresh: # print(key, thresh[key]) fileDir, skimDir = "./data", "./skim" for run in runList: filelist = glob.glob("{}/run_{}/FILTERED/compassF_run_{}".format( fileDir, run, run) + "*.root") print(filelist) for f in sorted(filelist): fcut = "(Channel==0 && Energy > {}) || ".format(thresh[run][0]) fcut += "(Channel==1 && Energy > {}) || ".format(thresh[run][1]) fcut += "(Channel==2 && Energy > {}) || ".format(thresh[run][2]) fcut += "(Channel==3 && Energy > {})".format(thresh[run][3]) ch = TChain("Data_F") ch.Add(f) print("Found %d entries" % (ch.GetEntries())) outName = "{}/run_{}.root".format(skimDir, run) outFile = TFile(outName, "RECREATE") outTree = TTree() outTree = ch.CopyTree(fcut) outTree.Write() outFile.Close() f2 = TFile(outName) tt2 = f2.Get("Data_F") print(tt2.GetEntries())
def readTreeFromFile(path, dileptonCombination, selection="OS"): """ helper function path: path to .root file containing simulated events dileptonCombination: EE, EMu, or MuMu for electron-electron, electron-muon, or muon-muon events returns: tree containing events for on sample and dileptonCombination """ from ROOT import TChain chain = TChain() # chain.Add("%s/ETH2AachenNtuples/%sDileptonTree"%(path, dileptonCombination)) chain.Add("%s/cutsV22DileptonFinalTrees/%sDileptonTree" % (path, dileptonCombination)) chain.SetName("%sDileptonTree" % dileptonCombination) result = chain.CopyTree("nJets >= 2") #~ result = chain return result
def save(file_in, file_out, ecms): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' width_low = 1.86965 - width(ecms) / 2. width_up = 1.86965 + width(ecms) / 2. cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_Dst = '(rm_D > 2.03)' cut = cut_base + ' && ' + cut_Dst t = chain.CopyTree(cut) t.SaveAs(file_out)
def save_before(file_in, file_out, ecms, region): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' width_low = 1.86965 - width(ecms) / 2. width_up = 1.86965 + width(ecms) / 2. window_low = 1.86965 - window(ecms) / 2. window_up = 1.86965 + window(ecms) / 2. width_side_low_up = width_low - (width_up - width_low) width_side_low_low = width_side_low_up - (width_up - width_low) width_side_up_low = width_up + (width_up - width_low) width_side_up_up = width_side_up_low + (width_up - width_low) window_side_low_up = window_low - (window_up - window_low) window_side_low_low = window_side_low_up - (window_up - window_low) window_side_up_low = window_up + (window_up - window_low) window_side_up_up = window_side_up_low + (window_up - window_low) if region == 'raw_signal': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut = cut_base if region == 'raw_sidebandlow': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + '&& m_rawm_D < ' + str( width_side_low_up) + ')' cut = cut_base if region == 'raw_sidebandup': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + '&& m_rawm_D < ' + str(width_side_up_up) + ')' cut = cut_base if region == 'STDDmiss_signal': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + ' && m_rm_Dpipi < ' + str(window_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side1_low': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_low_low) + '&& m_rm_Dpipi < ' + str( window_side_low_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side1_up': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_up_low) + '&& m_rm_Dpipi < ' + str( window_side_up_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side2_low': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + ' && m_rawm_D < ' + str( width_side_low_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + '&& m_rm_Dpipi < ' + str(window_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side2_up': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + ' && m_rawm_D < ' + str( width_side_up_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + '&& m_rm_Dpipi < ' + str(window_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side3_low': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + ' && m_rawm_D < ' + str( width_side_low_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_low_low) + '&& m_rm_Dpipi < ' + str( window_side_low_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side3_up': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + ' && m_rawm_D < ' + str( width_side_low_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_up_low) + '&& m_rm_Dpipi < ' + str( window_side_up_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side4_low': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + ' && m_rawm_D < ' + str( width_side_up_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_low_low) + '&& m_rm_Dpipi < ' + str( window_side_low_up) + ')' cut = cut_base + ' && ' + cut_window if region == 'STDDmiss_side4_up': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + ' && m_rawm_D < ' + str( width_side_up_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_up_low) + '&& m_rm_Dpipi < ' + str( window_side_up_up) + ')' cut = cut_base + ' && ' + cut_window t = chain.CopyTree(cut) t.SaveAs(file_out)
def save_after(file_in, file_out, ecms, region): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' width_low = 1.86965 - width(ecms) / 2. width_up = 1.86965 + width(ecms) / 2. window_low = 1.86965 - window(ecms) / 2. window_up = 1.86965 + window(ecms) / 2. width_side_low_up = width_low - (width_up - width_low) width_side_low_low = width_side_low_up - (width_up - width_low) width_side_up_low = width_up + (width_up - width_low) width_side_up_up = width_side_up_low + (width_up - width_low) window_side_low_up = window_low - (window_up - window_low) window_side_low_low = window_side_low_up - (window_up - window_low) window_side_up_low = window_up + (window_up - window_low) window_side_up_up = window_side_up_low + (window_up - window_low) if region == 'STDDmiss_signal': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + ' && m_rm_Dpipi < ' + str(window_up) + ')' if region == 'STDDmiss_side1_low': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_low_low) + '&& m_rm_Dpipi < ' + str( window_side_low_up) + ')' if region == 'STDDmiss_side1_up': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_up_low) + '&& m_rm_Dpipi < ' + str( window_side_up_up) + ')' if region == 'STDDmiss_side2_low': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + ' && m_rawm_D < ' + str( width_side_low_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + '&& m_rm_Dpipi < ' + str(window_up) + ')' if region == 'STDDmiss_side2_up': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + ' && m_rawm_D < ' + str( width_side_up_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_low) + '&& m_rm_Dpipi < ' + str(window_up) + ')' if region == 'STDDmiss_side3_low': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + ' && m_rawm_D < ' + str( width_side_low_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_low_low) + '&& m_rm_Dpipi < ' + str( window_side_low_up) + ')' if region == 'STDDmiss_side3_up': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + ' && m_rawm_D < ' + str( width_side_low_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_up_low) + '&& m_rm_Dpipi < ' + str( window_side_up_up) + ')' if region == 'STDDmiss_side4_low': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + ' && m_rawm_D < ' + str( width_side_up_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_low_low) + '&& m_rm_Dpipi < ' + str( window_side_low_up) + ')' if region == 'STDDmiss_side4_up': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + ' && m_rawm_D < ' + str( width_side_up_up) + ')' cut_window = '(m_rm_Dpipi > ' + str( window_side_up_low) + '&& m_rm_Dpipi < ' + str( window_side_up_up) + ')' cut_DDPI = '(abs(m_Kpipipi1-1.86483) > 0.01 && abs(m_Kpipipi2-1.86483) > 0.01)' cut_num = '(m_n_p == 0 && m_n_pbar == 0)' cut_KS = '(!((m_m_pipi > 0.491036 && m_m_pipi < 0.503471) && abs(m_L_svf/m_Lerr_svf) > 2))' cut_Vxy = '(m_Vxy_Dtrks[0] <= 0.55 && m_Vxy_Dtrks[1] <= 0.55 && m_Vxy_Dtrks[2] <= 0.55 && m_Vxy_pip <= 0.55 && m_Vxy_pim <= 0.55)' cut_Vz = '(m_Vz_Dtrks[0] <= 3 && m_Vz_Dtrks[1] <= 3 && m_Vz_Dtrks[2] <= 3 && m_Vz_pip <= 3 && m_Vz_pim <= 3)' if not (region == 'raw_signal' or region == 'raw_sidebandlow' or region == 'raw_sidebandup' or region == 'rm_Dpipi_signal'): cut = cut_base + ' && ' + cut_window + ' && ' + cut_num + ' && ' + cut_KS + ' && ' + cut_Vxy + ' && ' + cut_Vz + ' && ' + cut_DDPI if region == 'rm_Dpipi_signal': cut_window = '(m_rm_Dpipi > ' + str( window_low) + ' && m_rm_Dpipi < ' + str(window_up) + ')' cut = cut_window + ' && ' + cut_num + ' && ' + cut_KS + ' && ' + cut_Vxy + ' && ' + cut_Vz + ' && ' + cut_DDPI if region == 'raw_signal': cut_base = '(m_rawm_D > ' + str(width_low) + ' && m_rawm_D < ' + str( width_up) + ')' cut = cut_base + ' && ' + cut_num + ' && ' + cut_KS + ' && ' + cut_Vxy + ' && ' + cut_Vz + ' && ' + cut_DDPI if region == 'raw_sidebandlow': cut_base = '(m_rawm_D > ' + str( width_side_low_low) + '&& m_rawm_D < ' + str( width_side_low_up) + ')' cut = cut_base + ' && ' + cut_num + ' && ' + cut_KS + ' && ' + cut_Vxy + ' && ' + cut_Vz + ' && ' + cut_DDPI if region == 'raw_sidebandup': cut_base = '(m_rawm_D > ' + str( width_side_up_low) + '&& m_rawm_D < ' + str(width_side_up_up) + ')' cut = cut_base + ' && ' + cut_num + ' && ' + cut_KS + ' && ' + cut_Vxy + ' && ' + cut_Vz + ' && ' + cut_DDPI t = chain.CopyTree(cut) t.SaveAs(file_out)
def skim_tree(fname_patts, branches_to_keep, treename="t", fname_out="skim.root", cut_str="", flip_branches=False): ch = TChain(treename) for patt in fname_patts: ch.Add(patt) nevents = ch.GetEntries() branches_to_keep = [b for b in branches_to_keep if b] # remove empty strings # https://root-forum.cern.ch/t/pyroot-crashes-when-in-arguments/25379/3 -- figured out solution to below lines # # cut_str = "Sum$(abs(genps_id_mother)==24 && genps_isLastCopy && (abs(genps_id)==11 || abs(genps_id)==13))==2" # 2 leps # # cut_str = "abs(Sum$(genps_id*(abs(genps_id_mother)==24 && genps_isLastCopy && (abs(genps_id)==11 || abs(genps_id)==13))))>20" # SS # # apparently when using a $ sign in the cut_str, root hijacks the arguments to the script and crashes. Sigh. # cut_str = cut_str.replace("SSS","$") if len(cut_str) > 0: print ">>> [!] You specified a cut string of: %s" % cut_str print ">>> [!] Make sure that you are opting to keep all branches used in that cut string." filenames = [f.GetTitle() for f in ch.GetListOfFiles()] f1 = TFile(filenames[0]) tree = f1.Get(treename) tree.SetMakeClass(1); branches = [b.GetName() for b in tree.GetListOfBranches()] # see if the dummy user specified any branches to keep that aren't in the chain # and subtract them out to avoid segfaulttttt if not flip_branches: branches_not_in_chain = set(branches_to_keep)-set(branches) if len(branches_not_in_chain) > 0 and len(branches_to_keep) > 0: print ">>> [!] You dummy! I am going to neglect these branches which are not even in the TTree: %s" % ",".join(list(branches_not_in_chain)) branches_to_keep = list(set(branches_to_keep)-branches_not_in_chain) if len(branches_to_keep) == 0 and not flip_branches: if len(cut_str) == 0: print ">>> [!] You dummy! You want me to skim 0 branches without any cut? That's pointless." return else: print ">>> [!] You specified 0 branches to keep, but you gave me a cut string, so keeping ALL branches." branches_to_keep = branches[:] else: # whitelist the ones to copy # or reverse if we have flip_branches if not flip_branches: ch.SetBranchStatus("*",0) for bname in branches_to_keep: ch.SetBranchStatus(bname,1) else: ch.SetBranchStatus("*",1) for bname in branches_to_keep: ch.SetBranchStatus(bname,0) # need this to actually copy over any 4vectors. WTF. # https://root.cern.ch/phpBB3/viewtopic.php?t=10725 ch.SetBranchStatus("fCoordinates*",1); # actually do the skim and save the file t0 = time.time() new_file = TFile(fname_out,"RECREATE") # copy over all the histograms too - note that this only takes the first file (TODO is to actually add multiples, but this is not a use case for me right now) for key in f1.GetListOfKeys(): if key.ReadObj().InheritsFrom(TTree.Class()): continue name = key.GetName() f1.Get(name).Write() print ">>> Started skimming tree %s with %i events: %i --> %i branches" % (treename, nevents, len(branches), len(branches_to_keep)) ch_new = ch.CopyTree(cut_str) print ">>> Finished skim in %.2f seconds" % (time.time()-t0) ch_new.GetCurrentFile().Write() ch_new.GetCurrentFile().Close() # wow the user with incredible reduction stats size_before = get_filesizes(filenames) size_after = get_filesizes([fname_out]) print ">>> Size reduction: %s --> %s (factor of %.1f)" % (readable_size(size_before),readable_size(size_after),1.0*size_before/size_after) print ">>> Your output file is %s." % fname_out
#!/usr/bin/env python import sys import array, math from ROOT import gROOT, gStyle, TFile, TTree, TChain, TMVA, TCut, TCanvas, gDirectory, TH1, TGraph, gPad, TF1, THStack, TLegend, TH2D, TH1D, TGraphErrors, TSpline3 import getopt import utils gROOT.SetBatch(True) gStyle.SetOptStat(11) gStyle.SetOptFit(1) events = TChain("save") events.Add("occupancy_rs67_db01.root") events.Add("occupancy_rs67_db03.root") #infile = TFile("test.root") #events = infile.Get("save") outfilename="cut_data_occupancy" outfile = TFile(outfilename+".root","RECREATE") cutevents = events.CopyTree("NIM3") cutevents.Write() outfile.Write() outfile.Close()
cut = trkExtrpCutL1L2 if opt == '-b': cut = trkExtrpCutL2L2 if opt == '-d': cut = customCut if opt == '-e': ebeam = float(arg) if opt == '-c': clusterT = float(arg) if opt == '-z': targetZ = float(arg) if opt == '-h': print_usage() sys.exit(0) print remainder[0] #treeFile = TFile(sys.argv[1],"RECREATE") #tree = TTree("ntuple","data from text tuple "+sys.argv[2]) print remainder[1:] chain = TChain("ntuple") for i in remainder[1:]: chain.Add(i, 0) print chain.GetEntries() outFile = TFile(remainder[0], "RECREATE") events = chain.CopyTree(cut.format(ebeam, clusterT, targetZ)) print events.GetEntries() #outFile = TFile(remainder[0],"RECREATE") #events.Write() events.Write("ntuple", TTree.kOverwrite) gDirectory.ls()
gStyle.SetOptStat(11) gStyle.SetOptFit(1) dimuoncuts = utils.dimuoncuts() trackcuts = utils.trackcuts() datacuts = " && ".join([dimuoncuts, trackcuts]) print datacuts events = TChain("save") events.Add("mc_drellyan_C_M026_S002_messy_v2.root") outfilename = "cut_data_messy" outfile = TFile(outfilename + ".root", "RECREATE") cutevents = events.CopyTree(" && ".join([datacuts])) cutevents.Write() outfile.Write() outfile.Close() events = TChain("save") events.Add("mc_drellyan_C_M026_S002_clean_v2.root") outfilename = "cut_data_clean" outfile = TFile(outfilename + ".root", "RECREATE") cutevents = events.CopyTree(" && ".join([datacuts])) cutevents.Write() outfile.Write() outfile.Close()
def save(file_in, file_out, sample, mode): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' D0_W = width('D0') Dm_W = width('Dm') D0_signal_low = 1.86483 - D0_W/2. D0_signal_up = 1.86483 + D0_W/2. D0_sidebandlow_up = D0_signal_low - D0_W D0_sidebandlow_low = D0_signal_low - 2. * D0_W D0_sidebandup_low = D0_signal_up + D0_W D0_sidebandup_up = D0_signal_up + 2. * D0_W Dm_signal_low = 1.86965 - Dm_W/2. Dm_signal_up = 1.86965 + Dm_W/2. Dm_sidebandlow_up = Dm_signal_low - Dm_W Dm_sidebandlow_low = Dm_signal_low - 2. * Dm_W Dm_sidebandup_low = Dm_signal_up + Dm_W Dm_sidebandup_up = Dm_signal_up + 2. * Dm_W if mode == 'after_mass': if sample == 'signal': cut_D0Dm = '(m_m_D0 > ' + str(D0_signal_low) + ' && m_m_D0 < ' + str(D0_signal_up) + ' && m_m_Dm > ' + str(Dm_signal_low) + ' && m_m_Dm < ' + str(Dm_signal_up) + ')' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side1_low': cut_D0Dm = '(m_m_D0 > ' + str(D0_signal_low) + ' && m_m_D0 < ' + str(D0_signal_up) + ' && (m_m_Dm > ' + str(Dm_sidebandlow_low) + ' && m_m_Dm < ' + str(Dm_sidebandlow_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side1_up': cut_D0Dm = '(m_m_D0 > ' + str(D0_signal_low) + ' && m_m_D0 < ' + str(D0_signal_up) + ' && (m_m_Dm > ' + str(Dm_sidebandup_low) + ' && m_m_Dm < ' + str(Dm_sidebandup_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side2_low': cut_D0Dm = '(m_m_Dm > ' + str(Dm_signal_low) + ' && m_m_Dm < ' + str(Dm_signal_up) + ' && (m_m_D0 > ' + str(D0_sidebandlow_low) + ' && m_m_D0 < ' + str(D0_sidebandlow_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side2_up': cut_D0Dm = '(m_m_Dm > ' + str(Dm_signal_low) + ' && m_m_Dm < ' + str(Dm_signal_up) + ' && (m_m_D0 > ' + str(D0_sidebandup_low) + ' && m_m_D0 < ' + str(D0_sidebandup_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side3_low': cut_D0Dm = '(m_m_D0 > ' + str(D0_sidebandlow_low) + ' && m_m_D0 < ' + str(D0_sidebandlow_up) + ' && (m_m_Dm > ' + str(Dm_sidebandlow_low) + ' && m_m_Dm < ' + str(Dm_sidebandlow_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side3_up': cut_D0Dm = '(m_m_D0 > ' + str(D0_sidebandlow_low) + ' && m_m_D0 < ' + str(D0_sidebandlow_up) + ' && (m_m_Dm > ' + str(Dm_sidebandup_low) + ' && m_m_Dm < ' + str(Dm_sidebandup_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side4_low': cut_D0Dm = '(m_m_D0 > ' + str(D0_sidebandup_low) + ' && m_m_D0 < ' + str(D0_sidebandup_up) + ' && (m_m_Dm > ' + str(Dm_sidebandlow_low) + ' && m_m_Dm < ' + str(Dm_sidebandlow_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if sample == 'side4_up': cut_D0Dm = '(m_m_D0 > ' + str(D0_sidebandup_low) + ' && m_m_D0 < ' + str(D0_sidebandup_up) + ' && (m_m_Dm > ' + str(Dm_sidebandup_low) + ' && m_m_Dm < ' + str(Dm_sidebandup_up) + '))' cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_D0Dm + ' && ' + cut_DstD + ' && ' + cut_chi2 if mode == 'before_mass': cut_DstD = '(m_m_D0pip - m_m_D0 + 1.86483 > 2.030)' cut_chi2 = '(m_chi2_4c < 60)' cut = cut_DstD + ' && ' + cut_chi2 t = chain.CopyTree(cut) t.SaveAs(file_out)
gStyle.SetOptFit(1) events = TChain("save") events.Add("dimuonmix_rs67_db01.root") events.Add("dimuonmix_rs67_db03.root") #infile = TFile("test.root") #events = infile.Get("save") outfilename = "cut_data_mix" outfile = TFile(outfilename + ".root", "RECREATE") #MATRIX1 && abs(dx)<.25 && abs(dy-1.6)<.22 && dx*dx+(dy-1.6)*(dy-1.6)<.06 && dz>-280 && dz<-5 && abs(dpx)<1.8 && abs(dpy)<2 && dpx*dpx+dpy*dpy<5 && dpz>38 && dpz<116 && abs(trackSeparation)<270 && chisq_dimuon<18 && pchisq_target<15 && ppz1>9 && ppz1<75 && pnumHits>13 && pxT*pxT+(pyT-1.6)*(pyT-1.6)<320 && pxD*pxD+(pyD-1.6)*(pyD-1.6)<1100 && pxD*pxD+(pyD-1.6)*(pyD-1.6)>16 && pchisq_target<1.5*pchisq_upstream && pchisq_target<1.5*pchisq_dump && pz0<-5 && pz0>-320 && pchisq/(pnumHits-5)<12 && py1/py3<1 && nchisq_target<15 && npz1>9 && npz1<75 && nnumHits>13 && nxT*nxT+(nyT-1.6)*(nyT-1.6)<320 && nxD*nxD+(nyD-1.6)*(nyD-1.6)<1100 && nxD*nxD+(nyD-1.6)*(nyD-1.6)>16 && nchisq_target<1.5*nchisq_upstream && nchisq_target<1.5*nchisq_dump && nz0<-5 && nz0>-320 && nchisq/(nnumHits-5)<12 && ny1/ny3<1 #qiecut = " && ".join([i+"<5000" for i in ["RFm{0:02d}".format(j) for j in range(8,0,-1)]+["RF00"]+["RFp{0:02d}".format(j) for j in range(1,9)]]) qualitycut = utils.qualitycuts(True) print qualitycut dimuoncuts = utils.dimuoncuts() trackcuts = utils.trackcuts() datacuts = " && ".join(["pMATRIX1 && nMATRIX1", dimuoncuts, trackcuts]) print datacuts cutevents = events.CopyTree(" && ".join([qualitycut, datacuts])) cutevents.Write() outfile.Write() outfile.Close()
def save_before(file_in, file_out, D_sample, D_type): try: chain = TChain('save') chain.Add(file_in) except: logging.error(file_in + ' is invalid!') sys.exit() cut = '' if D_sample == 'Dplus' and D_type == 'raw': rawm_D_low = 1.86965 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86965 + width_rawm_D(D_sample) / 2. cut_Dplus = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' cut_chi2 = '(m_chi2_kf < 20)' cut_Dst0 = '(m_m_Dpi0 > 2.02)' cut = cut_Dplus + ' && ' + cut_chi2 + ' && ' + cut_Dst0 if D_sample == 'D0' and D_type == 'raw': rawm_D_low = 1.86483 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86483 + width_rawm_D(D_sample) / 2. cut_D0 = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' cut_chi2 = '(m_chi2_kf < 20)' cut = cut_D0 + ' && ' + cut_chi2 if D_sample == 'Dplus' and D_type == 'D': rawm_D_low = 1.86965 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86965 + width_rawm_D(D_sample) / 2. cut_Dplus = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' rm_Dpi_low = 1.86483 - width_rm_Dpi(D_sample, D_type) / 2. rm_Dpi_up = 1.86483 + width_rm_Dpi(D_sample, D_type) / 2. cut_D0 = '(m_rm_Dpi > ' + str(rm_Dpi_low) + ' && m_rm_Dpi < ' + str( rm_Dpi_up) + ')' cut_Dst0 = '(m_m_Dpi0 > 2.02)' cut_range = '(m_rm_pi_MCTruth>3.8&&m_rm_pi_MCTruth<4.4&&m_rm_D_MCTruth>2.1&&m_rm_D_MCTruth<2.8&&m_m_Dpi_MCTruth>2.1&&m_m_Dpi_MCTruth<2.8)' cut = cut_Dplus + ' && ' + cut_D0 + ' && ' + cut_Dst0 + ' && ' + cut_range if D_sample == 'Dplus' and D_type == 'Dst': rawm_D_low = 1.86965 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86965 + width_rawm_D(D_sample) / 2. cut_Dplus = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' rm_Dpi_low = 2.00685 - width_rm_Dpi(D_sample, D_type) / 2. rm_Dpi_up = 2.00685 + width_rm_Dpi(D_sample, D_type) / 2. cut_D0st = '(m_rm_Dpi > ' + str(rm_Dpi_low) + ' && m_rm_Dpi < ' + str( rm_Dpi_up) + ')' cut = cut_Dplus + ' && ' + cut_D0st if D_sample == 'D0' and D_type == 'D': rawm_D_low = 1.86483 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86483 + width_rawm_D(D_sample) / 2. cut_D0 = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' rm_Dpi_low = 1.86965 - width_rm_Dpi(D_sample, D_type) / 2. rm_Dpi_up = 1.86965 + width_rm_Dpi(D_sample, D_type) / 2. cut_Dplus = '(m_rm_Dpi > ' + str(rm_Dpi_low) + ' && m_rm_Dpi < ' + str( rm_Dpi_up) + ')' cut = cut_D0 + ' && ' + cut_Dplus if D_sample == 'D0' and D_type == 'Dst': rawm_D_low = 1.86483 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86483 + width_rawm_D(D_sample) / 2. cut_D0 = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' cut_Dplusst = '(m_rm_Dpi > 1.996 && m_rm_Dpi < 2.028)' cut = cut_D0 + ' && ' + cut_Dplusst if D_sample == 'Dplus' and D_type == 'missing': rawm_D_low = 1.86965 - width_rawm_D(D_sample) / 2. rawm_D_up = 1.86965 + width_rawm_D(D_sample) / 2. cut_Dplus = '(m_rawm_D > ' + str(rawm_D_low) + ' && m_rawm_D < ' + str( rawm_D_up) + ')' cut_DstD = '(m_rm_D > 2.03)' cut_Dst0 = '(m_m_Dpi0 < 2.0 || m_m_Dpi0 > 2.02)' cut = cut_Dplus + ' && ' + cut_DstD + ' && ' + cut_Dst0 t = chain.CopyTree(cut) t.SaveAs(file_out)
def ApplyChannelCuts(dsNum, cutType, dType): """ ./lat3.py -cut [dsNum] [cutType] This runs over whole datasets. cutTypes: fs, rn, wf, fs+rn, fs+wf, rn+wf, fs+rn+wf dTypes: bkg, cal """ # load the database calDB = db.TinyDB('../calDB.json') pars = db.Query() # setup a loop over modules and dataset ranges gROOT.ProcessLine("gErrorIgnoreLevel = 3001;") cInfo = ds.CalInfo() nMods = [1] if dsNum == 4: nMods = [2] if dsNum == 5: nMods = [1,2] # Dummy string for file writing -- adds nothing to the directories if background dString = "" if dType == "cal": dString = "cal" for modNum in nMods: # Changed so range of idx are set here to take advantage of module number nRanges = [] if dType == "bkg": nRanges = [0, ds.dsMap[dsNum]] # if dsNum==5: nRanges[0] = 80 # exclude DS-5A elif dType == "cal": nRanges = [0, len(cInfo.master['ds%d_m%d'%(dsNum, modNum)])-1] else: print "cal or bkg not set, returning" return 0 # Loop over bkgIdx, even though for calibration runs this will represent calIdx for bkgIdx in range(nRanges[0], nRanges[1]+1): # load the chains and find the right calIdx's. skimTree = TChain("skimTree") # build the file list fRegex = "" fList = [] if dType == "bkg": fRegex = "/global/homes/w/wisecg/project/bg-lat/latSkimDS%d_%d_*.root" % (dsNum, bkgIdx) fList = glob.glob(fRegex) skimTree.Add(fRegex) elif dType == "cal": calList = cInfo.GetCalList("ds%d_m%d" % (dsNum, modNum), bkgIdx, runLimit=10) for i in calList: fList += glob.glob("/global/homes/w/wisecg/project/cal-lat/latSkimDS%d_run%d_*.root"%(dsNum,i)) skimTree.Add("/global/homes/w/wisecg/project/cal-lat/latSkimDS%d_run%d_*.root" % (dsNum, i)) file0 = fList[0] print "DS-%d subset %d, Mod-%d. N_files: %d" % (dsNum, bkgIdx, modNum, len(fList)) # Print some basic info about files f = TFile(file0) firstRun, lastRun, calIdxLo, calIdxHi = 0,0,0,0 theCut = f.Get("theCut").GetTitle() if dType == "bkg": skimTree.GetEntry(0) firstRun = skimTree.run skimTree.GetEntry(skimTree.GetEntries()-1) lastRun = skimTree.run calIdxLo = cInfo.GetCalIdx("ds%d_m%d" % (dsNum, modNum), firstRun) calIdxHi = cInfo.GetCalIdx("ds%d_m%d" % (dsNum, modNum), lastRun) elif dType == "cal": # All the idx are the same for calibration! calIdxLo = calIdxHi = bkgIdx firstRun, lastRun = calList[0], calList[-1] print " Entries %d firstRun %d lastRun %d calIdxLo %d calIdxHi %d" % (skimTree.GetEntries(),firstRun,lastRun,calIdxLo,calIdxHi) # build the channel list (remove 692 and 1232 from DS5 for now.) chList = ds.GetGoodChanList(dsNum) if dsNum==5 and modNum==1: chList = [ch for ch in chList if ch < 1000 and ch!=692] if dsNum==5 and modNum==2: chList = [ch for ch in chList if ch > 1000 and ch!=1232] # -- create a dict of cuts for each channel, covering each calIdx. -- cutDict = {} for calIdx in range(calIdxLo, calIdxHi+1): runCovMin = cInfo.master["ds%d_m%d" % (dsNum, modNum)][calIdx][1] runCovMax = cInfo.master["ds%d_m%d" % (dsNum, modNum)][calIdx][2] runCut = "run>=%d && run<=%d" % (runCovMin, runCovMax) fsD = ds.getDBRecord("fitSlo_ds%d_idx%d_m%d_Peak" % (dsNum, calIdx, modNum), False, calDB, pars) rnSD = ds.getDBRecord("riseNoise_ds%d_idx%d_m%d_SoftPlus" % (dsNum, calIdx, modNum), False, calDB, pars) rnCD = ds.getDBRecord("riseNoise_ds%d_idx%d_m%d_Continuum" % (dsNum, calIdx, modNum), False, calDB, pars) wfD = ds.getDBRecord("wfstd_ds%d_idx%d_mod%d" % (dsNum, calIdx, modNum), False, calDB, pars) for ch in chList: fsCut, rnCut, wfCut, chanCut = None, None, None, None # print ch,":",fsD[ch][2] # fitSlo: check the 90% value is positive if fsD[ch][2] > 0: fsCut = "fitSlo<%.2f" % fsD[ch][2] # riseNoise: check the softplus curvature is positive if rnSD[ch][3] > 0: rnCut = "riseNoise<(%.3f+%.5f*TMath::Log(1+TMath::Exp((trapENFCal-(%.3f))/%.3f)))" % (max(rnSD[ch][0],rnCD[ch][4]), rnSD[ch][1], rnSD[ch][2], rnSD[ch][3]) # wfStd: check if ralph says this is ok to use if wfD!=0 and ch in wfD.keys() and wfD[ch][0]==u'y': wfCut = "abs(wfstd - sqrt((%.4e + %.4e*trapENFCal + %.4e*trapENFCal**2 + %.2e*pow(trapENFCal,3) + %.2e*pow(trapENFCal,4))**2 + %.4f)) < (%.2f+%.3f*trapENFCal)" % (wfD[ch][3], wfD[ch][4], wfD[ch][5], wfD[ch][6], wfD[ch][7], wfD[ch][8], wfD[ch][9], wfD[ch][10]) # set the combination channel cut if cutType == "fs" and fsCut!=None: chanCut = "(%s && %s)" % (runCut, fsCut) if cutType == "rn" and rnCut!=None: chanCut = "(%s && %s)" % (runCut, rnCut) if cutType == "wf" and wfCut!=None: chanCut = "(%s && %s)" % (runCut, wfCut) if cutType == "fs+rn" and fsCut!=None and rnCut!=None: chanCut = "(%s && %s && %s)" % (runCut, fsCut, rnCut) if cutType == "fs+wf" and fsCut!=None and wfCut!=None: chanCut = "(%s && %s && %s)" % (runCut, fsCut, wfCut) if cutType == "rn+wf" and rnCut!=None and wfCut!=None: chanCut = "(%s && %s && %s)" % (runCut, rnCut, wfCut) if cutType == "fs+rn+wf" and fsCut!=None and rnCut!=None and wfCut!=None: chanCut = "(%s && %s && %s && %s)" % (runCut, fsCut, rnCut, wfCut) # create dict entry for this channel or append to existing, taking care of parentheses and OR's. if ch in cutDict.keys() and chanCut!=None: cutDict[ch] += " || %s" % chanCut elif ch not in cutDict.keys() and chanCut!=None: cutDict[ch] = "(%s" % chanCut # close the parens for each channel entry for key in cutDict: cutDict[key] += ")" # -- finally, loop over each channel we have an entry for, get its cut, and create an output file. -- for ch in cutDict: # TODO: threshold cut (or at least save the value for each bkgIdx) chanCut = theCut + "&& gain==0 && channel==%d" % ch if cutType == "fs": outFile = "~/project/cuts/%sfs/%sfitSlo-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& fitSlo>0 && %s" % cutDict[ch] if cutType == "rn": outFile = "~/project/cuts/%srn/%sriseNoise-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& %s" % cutDict[ch] if cutType == "wf": outFile = "~/project/cuts/%swf/%swfstd-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& %s" % cutDict[ch] if cutType == "fs+rn": outFile = "~/project/cuts/%sfs_rn/%sfs_rn-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& fitSlo>0 && %s" % cutDict[ch] if cutType == "fs+wf": outFile = "~/project/cuts/%sfs_wf/%sfs_wf-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& fitSlo>0 && %s" % cutDict[ch] if cutType == "rn+wf": outFile = "~/project/cuts/%srn_wf/%srn_wf-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& %s" % cutDict[ch] if cutType == "fs+rn+wf": outFile = "~/project/cuts/%sfs_rn_wf/%sfs_rn_wf-DS%d-%d-ch%d.root" % (dString, dString, dsNum, bkgIdx, ch) chanCut += "&& fitSlo>0 && %s" % cutDict[ch] print " Writing to:",outFile print " Cut used:",chanCut,"\n" outFile = TFile(outFile,"RECREATE") outTree = TTree() outTree = skimTree.CopyTree(chanCut) outTree.Write() cutUsed = TNamed("chanCut",chanCut) cutUsed.Write() print "Wrote",outTree.GetEntries(),"entries." outFile.Close()
#!/usr/bin/env python import sys, os import getopt import ROOT from ROOT import gROOT, TFile, TTree, TChain, TMVA, TCut print sys.argv[1] #tree = TTree("ntuple","data from text tuple "+sys.argv[2]) print sys.argv[2:] chain = TChain("ntuple") for i in sys.argv[2:]: chain.Add(i) #chain.Merge(sys.argv[1]) #goodEvents = chain.CopyTree("") events = chain.CopyTree("uncM>0.03&&uncM<0.04") #goodEvents = chain.CopyTree("uncM>0.03&&uncM<0.04&&eleP<1.05*0.85&&posP<1.05*0.85&&elePY*posPY<0&&eleP>1.05*0.05&&posP>1.05*0.05&&nPos==1&&uncP<1.05*1.25&&abs(eleClT-posClT)<2&&eleHasL1&&posHasL1&&max(eleMatchChisq,posMatchChisq)<10&&uncP>1.05*0.8&&uncVZ<15") #badEvents = chain.CopyTree("uncM>0.03&&uncM<0.04&&eleP<1.05*0.85&&posP<1.05*0.85&&elePY*posPY<0&&eleP>1.05*0.05&&posP>1.05*0.05&&nPos==1&&uncP<1.05*1.25&&abs(eleClT-posClT)<2&&eleHasL1&&posHasL1&&max(eleMatchChisq,posMatchChisq)<10&&uncP>1.05*0.8&&uncVZ>15") TMVA.Tools.Instance() outFile = TFile(sys.argv[1], "RECREATE") #factory = TMVA.Factory("TMVAClassification", outFile,":".join(["!V","!Silent","Color","DrawProgressBar","Transformations=I;D;P;G,D","AnalysisType=Classification"])) factory = TMVA.Factory( "TMVAClassification", outFile, ":".join([ "!V", "!Silent", "Color", "DrawProgressBar", "Transformations=I;G,D", "AnalysisType=Classification" ])) factory.AddVariable("bscChisq", "F") #factory.AddVariable("min(minL1Iso,20.0)","F") factory.AddVariable("1/minL1Iso", "F") #factory.AddVariable("abs(uncPY/uncP)","F") #factory.AddVariable("abs(uncPX/uncP)","F")
def skim_tree(fname_patts, branches_to_keep, treename="t", fname_out="skim.root", cut_str=""): # This stuff is super necessary or else we all die from ROOT import TChain, TFile, gSystem, gROOT, TTree gSystem.Load("libFWCoreFWLite.so") gSystem.Load("libDataFormatsFWLite.so") gROOT.ProcessLine("FWLiteEnabler::enable()") ch = TChain(treename) for patt in fname_patts: ch.Add(patt) nevents = ch.GetEntries() branches_to_keep = [b for b in branches_to_keep if b] # remove empty strings if len(cut_str) > 0: print ">>> [!] You specified a cut string of: %s" % cut_str print ">>> [!] Make sure that you are opting to keep all branches used in that cut string." filenames = [f.GetTitle() for f in ch.GetListOfFiles()] f1 = TFile(filenames[0]) tree = f1.Get(treename) tree.SetMakeClass(1) branches = [b.GetName() for b in tree.GetListOfBranches()] # see if the dummy user specified any branches to keep that aren't in the chain # and subtract them out to avoid segfaulttttt branches_not_in_chain = set(branches_to_keep) - set(branches) if len(branches_not_in_chain) > 0 and len(branches_to_keep) > 0: print ">>> [!] You dummy! I am going to neglect these branches which are not even in the TTree: %s" % ",".join( list(branches_not_in_chain)) branches_to_keep = list(set(branches_to_keep) - branches_not_in_chain) if len(branches_to_keep) == 0: if len(cut_str) == 0: print ">>> [!] You dummy! You want me to skim 0 branches without any cut? That's pointless." return else: print ">>> [!] You specified 0 branches to keep, but you gave me a cut string, so keeping ALL branches." branches_to_keep = branches[:] else: # whitelist the ones to copy ch.SetBranchStatus("*", 0) for bname in branches_to_keep: ch.SetBranchStatus(bname, 1) # need this to actually copy over any 4vectors. WTF. # https://root.cern.ch/phpBB3/viewtopic.php?t=10725 ch.SetBranchStatus("fCoordinates*", 1) # actually do the skim and save the file t0 = time.time() new_file = TFile(fname_out, "RECREATE") # copy over all the histograms too - note that this only takes the first file (TODO is to actually add multiples, but this is not a use case for me right now) for key in f1.GetListOfKeys(): if key.ReadObj().InheritsFrom(TTree.Class()): continue name = key.GetName() print name f1.Get(name).Write() print ">>> Started skimming tree %s with %i events: %i --> %i branches" % ( treename, nevents, len(branches), len(branches_to_keep)) ch_new = ch.CopyTree(cut_str) print ">>> Finished skim in %.2f seconds" % (time.time() - t0) ch_new.GetCurrentFile().Write() ch_new.GetCurrentFile().Close() # wow the user with incredible reduction stats size_before = get_filesizes(filenames) size_after = get_filesizes([fname_out]) print ">>> Size reduction: %s --> %s (factor of %.1f)" % (readable_size( size_before), readable_size(size_after), size_before / size_after) print ">>> Your output file is %s." % fname_out
nBins, minTheta, maxTheta) posL12tscatter = TH2F("posL12tscatter", "posL12tscatter", nBins, minVZ, maxVZ, nBins, minTheta, maxTheta) posL12bscatter = TH2F("posL12bscatter", "posL12bscatter", nBins, minVZ, maxVZ, nBins, minTheta, maxTheta) eleposL1tscatter = TH2F("eleposL1tscatter", "eleposL1tscatter", nBins, minVZ, maxVZ, nBins, minTheta, maxTheta) eleposL1bscatter = TH2F("eleposL1bscatter", "eleposL1bscatter", nBins, minVZ, maxVZ, nBins, minTheta, maxTheta) eleposL12tscatter = TH2F("eleposL12tscatter", "eleposL12tscatter", nBins, minVZ, maxVZ, nBins, minTheta, maxTheta) eleposL12bscatter = TH2F("eleposL12bscatter", "eleposL12bscatter", nBins, minVZ, maxVZ, nBins, minTheta, maxTheta) cutevents = events.CopyTree(cut) neventscut = cutevents.GetEntries() uncVZ = array('d', [0]) eleL1tthetaY = array('d', [0]) eleL2tthetaY = array('d', [0]) eleL3tthetaY = array('d', [0]) eleL4tthetaY = array('d', [0]) eleL1bthetaY = array('d', [0]) eleL2bthetaY = array('d', [0]) eleL3bthetaY = array('d', [0]) eleL4bthetaY = array('d', [0]) eleL1tInthetaY = array('d', [0]) eleL2tInthetaY = array('d', [0])
if opt == '-b': uncTargProjYSig = float(arg) if opt == '-e': ebeam = float(arg) if opt == '-c': clusterT = float(arg) if opt == '-z': targetZ = float(arg) if opt == '-h': print_usage() sys.exit(0) print remainder[0] #treeFile = TFile(sys.argv[1],"RECREATE") #tree = TTree("ntuple","data from text tuple "+sys.argv[2]) print remainder[1:] chain = TChain("ntuple") for i in remainder[1:]: chain.Add(i, 0) print chain.GetEntries() outFile = TFile(remainder[0], "RECREATE") #events = chain.CopyTree(cut.format(ebeam,clusterT,targetZ)) events = chain.CopyTree( cut.format(ebeam, clusterT, targetZ, uncVX, uncVXSig, uncVY, uncVYSig, uncTargProjX, uncTargProjXSig, uncTargProjY, uncTargProjYSig)) print events.GetEntries() #outFile = TFile(remainder[0],"RECREATE") #events.Write() events.Write("ntuple", TTree.kOverwrite) gDirectory.ls()
elif cutType == "vertexing_noL1": baseCut = "uncP>0.8*{0}&&uncVZ<80&&!eleHasL1&&!posHasL1".format(ebeam) goodCut = baseCut + "&&abs(uncVZ)*uncM<0.1" badCut = baseCut + "&&(uncVZ+5)*uncM>0.5" elif cutType == "vertexing_eleL1": baseCut = "uncP>0.8*{0}&&uncVZ<80&&eleHasL1&&!posHasL1".format(ebeam) goodCut = baseCut + "&&abs(uncVZ)*uncM<0.1" badCut = baseCut + "&&(uncVZ+5)*uncM>0.5" elif cutType == "vertexing_posL1": baseCut = "uncP>0.8*{0}&&uncVZ<80&&!eleHasL1&&posHasL1".format(ebeam) goodCut = baseCut + "&&abs(uncVZ)*uncM<0.1" badCut = baseCut + "&&(uncVZ+5)*uncM>0.5" else: print "invalid cut type" sys.exit(-1) chain = TChain("ntuple") for i in remainder[2:]: chain.Add(i) #maincut = "uncP>0.8*1.056&&bscChisq<5&&minIso>0.5&&max(eleTrkChisq,posTrkChisq)<20" goodOutFile = TFile(remainder[0], "RECREATE") goodEvents = chain.CopyTree(goodCut) goodOutFile.Write() goodOutFile.Close() badOutFile = TFile(remainder[1], "RECREATE") badEvents = chain.CopyTree(badCut) badOutFile.Write() badOutFile.Close()
else: run = 1 XicDataCuts = "{0} && {1}".format(getDataCuts(run), "lcplus_MM > 2375") LcDataCuts = "{0} && {1}".format(getDataCuts(run), "lcplus_MM < 2375") tree = TChain("tuple_Lc2pKpi/DecayTree") PATH = "/dcache/bfys/jdevries/ntuples/LcAnalysis/ganga/" + job for subdir, dirs, files in os.walk(PATH): for file in files: if "-histos" not in file: #print subdir tree.Add(os.path.join(subdir,file)) if(run == 2): XicTree = tree.CopyTree(XicMCCuts) XicTree = XicTree.CopyTree(XicDataCuts) LcTree = tree.CopyTree(LcMCCuts) LcTree = LcTree.CopyTree(LcDataCuts) print("The nb of entries for job " + job + " is (Xic: " + str(XicTree.GetEntries()) + ";Lc: " + str(LcTree.GetEntries())+ ")") if(GRAPHS == True): XicMasshist = ROOT.TH1F("XicMasshist","Histogram of Xic mass, job: "+ job ,nbin,XicMassRange[0],XicMassRange[1]) LcMasshist = ROOT.TH1F("LcMasshist","Histogram of Lc mass, job: "+ job ,nbin,LcMassRange[0],LcMassRange[1]) LcTree.Draw("lcplus_MM>>LcMasshist") c1.SaveAs(job + "_Lc_MC.pdf") XicTree.Draw("lcplus_MM>>XicMasshist") c1.SaveAs(job + "_Xic_MC.pdf")