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): print "\tSource file not in target directory, I'll bring it along..." os.system('cp %s %s' % (source_path + source, target_path + source)) outfilename = source.split('.')[0] + '_' + '_'.join(sel_names) + '.root' print '\tCreating file %s ...' % outfilename outfile = r.TFile(target_path + outfilename, 'recreate') clones = [] if truth_tree: print '\tCopying MCTruth tree...' truth_tree.AutoSave() for t in recotrees: print '\tCopying %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() print '\tAll done.' print