diff --git a/Phys/B2KShh/job/selection.py b/Phys/B2KShh/job/selection.py new file mode 100644 index 0000000..0ba90bf --- /dev/null +++ b/Phys/B2KShh/job/selection.py @@ -0,0 +1,59 @@ +import sys +import os +import ROOT as r + +source = str(sys.argv[1]) +a = str(sys.argv[2]) +sel_names = [c for c in a.split(",")] + +source_path = '/disk/users/elena/B2KShh/ntuples/jobs/' +target_path = '/disk/users/elena/B2KShh/ntuples/sel/' +sel_strings = { + 'L0' : '(L0HadronDecision_TOS || L0Global_TIS)', + 'HLT1' : 'Hlt1TrackAllL0Decision_TOS', +} + + +def getall(d, basepath="/"): + "Generator function to recurse into a ROOT file/dir and yield (path, obj) pairs" + keylist = d.GetListOfKeys() + for key in keylist: + kname = key.GetName() + if key.IsFolder() and d.Get(kname).ClassName() == 'TDirectoryFile' : + for i in getall(d.Get(kname), basepath+kname+"/"): + yield i + else: + yield basepath+kname, d.Get(kname) + + +infile = r.TFile(source_path + source, 'read') +trees = [] + +for k, o in getall(infile): + if 'TTree' in o.ClassName(): + trees.append(o) + +recotrees = [t for t in trees if 'MCTruth' not in t.GetName()] +truthtrees = [t for t in trees if 'MCTruth' in t.GetName()] +if len(truthtrees) == 1: + truth_tree = truthtrees[0] +else: + truth_tree = None + +os.system('mkdir -p %s' % target_path) +if not os.path.isfile(target_path + source): + os.system('cp %s %s' % (source_path + source, target_path + source)) + +outfilename = source.split('.')[0] + '_' + '_'.join(sel_names) + '.root' +outfile = r.TFile(outfilename, 'recreate') +clones = [] +if truth_tree: + print 'Copying MCTruth tree...' + truth_tree.AutoSave() +for t in recotrees: + print 'Copying %s tree...' % t.GetName() + clones.append(t.CopyTree(' && '.join([sel_strings[sel_name] for sel_name in sel_names]))) +for c in clones: + c.AutoSave() +outfile.Close() +infile.Close() \ No newline at end of file