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