diff --git a/Phys/B2KShh/job/selection.py b/Phys/B2KShh/job/selection.py index a8e927b..69bcf26 100644 --- a/Phys/B2KShh/job/selection.py +++ b/Phys/B2KShh/job/selection.py @@ -40,7 +40,7 @@ def truth_match_string(source, hypo): - string = '' + string = '(' mother = 'MOTHER' b_id = truth_match['Lb'] if 'Lc' in source: @@ -50,29 +50,34 @@ for h in ['h1', 'h2', 'KS']: string += ' && %s_TRUEID == %s' % (h, truth_match[hypos[hypo][h]]) string += ' && %s_%s_TRUEID == %s' % (h, mother, b_id) + string += ')' return string -sel_strings = { - 'L0' : '(L0HadronDecision_TOS || L0Global_TIS)', - 'HLT1' : 'Hlt1TrackAllL0Decision_TOS', - '2011' : '(runNumber <= %s)' % end2011, - 'HLT2-2011' : '((Hlt2Topo2BodyBBDTDecision_TOS || Hlt2Topo3BodyBBDTDecision_TOS || Hlt2Topo4BodyBBDTDecision_TOS) || (Hlt2Topo2BodySimpleDecision_TOS || Hlt2Topo3BodySimpleDecision_TOS || Hlt2Topo4BodySimpleDecision_TOS))', - '2012pre-june' : '(runNumber > %s && runNumber <= %s)' % (end2011, endjune), - '2012post-june': '(runNumber > %s)' % endjune, - 'HLT2-2012' : '(Hlt2Topo2BodyBBDTDecision_TOS || Hlt2Topo3BodyBBDTDecision_TOS || Hlt2Topo4BodyBBDTDecision_TOS)', -} - -sel_strings['Trigger'] = ' && '.join([sel_strings['L0'], sel_strings['HLT1'], sel_strings['HLT2-2011']]) -sel_strings['Trigger'] += ' && ' -if 'MC_2011' in source: - sel_strings['Trigger'] += sel_strings['HLT2-2011'] -elif 'MC_2012' in source: - sel_strings['Trigger'] += sel_strings['HLT2-2012'] -elif not 'MC' in source: - pass -else: - raise Exception('DANGER: problem with file name parsing and trigger conditions!') +def make_sel_trings(source, recotree): + sel_strings = { + 'L0' : '(L0HadronDecision_TOS || L0Global_TIS)', + 'HLT1' : 'Hlt1TrackAllL0Decision_TOS', + '2011' : '(runNumber <= %s)' % end2011, + 'HLT2-2011' : '((Hlt2Topo2BodyBBDTDecision_TOS || Hlt2Topo3BodyBBDTDecision_TOS || Hlt2Topo4BodyBBDTDecision_TOS) || (Hlt2Topo2BodySimpleDecision_TOS || Hlt2Topo3BodySimpleDecision_TOS || Hlt2Topo4BodySimpleDecision_TOS))', + '2012pre-june' : '(runNumber > %s && runNumber <= %s)' % (end2011, endjune), + '2012post-june': '(runNumber > %s)' % endjune, + 'HLT2-2012' : '(Hlt2Topo2BodyBBDTDecision_TOS || Hlt2Topo3BodyBBDTDecision_TOS || Hlt2Topo4BodyBBDTDecision_TOS)', + } + + sel_strings['Trigger'] = ' && '.join([sel_strings['L0'], sel_strings['HLT1'], sel_strings['HLT2-2011']]) + sel_strings['Trigger'] += ' && ' + if 'MC_2011' in source: + sel_strings['Trigger'] += sel_strings['HLT2-2011'] + elif 'MC_2012' in source: + sel_strings['Trigger'] += sel_strings['HLT2-2012'] + elif not 'MC' in source: + pass + else: + raise Exception('DANGER: problem with file name parsing and trigger conditions!') + + sel_strings['TruthMatch'] = truth_match_string(source, recotree) + return sel_strings def getall(d, basepath="/"): @@ -104,7 +109,23 @@ print "\tSource file not in target directory, I'll bring it along..." os.system('cp %s %s' % (source_path + source, target_path + source)) -infile = r.TFile(target_path + source, 'read') + +def find_input(target_path, source, sel_names): + """ + If some selections were already performed, + use the selected sample as input! + """ + infile_name = target_path + source + available_files = os.listdir(target_path) + for i in range(len(sel_names)): + sel_subset = sel_names[:i+1] + for f in available_files: + if f == source.split('.')[0] + '_' + '_'.join(sel_subset) + '.root': + infile_name = target_path + f + return infile_name + + +infile = r.TFile(find_input(target_path, source, sel_names), 'read') trees = [] for k, o in getall(infile): @@ -127,6 +148,7 @@ print '\tMaking directory %s ...' % t.location outfile.mkdir(t.location, t.location) outfile.cd(t.location) + sel_strings = make_sel_strings(source, t.GetName()) t.cptree = t.tree.CopyTree(' && '.join([sel_strings[sel_name] for sel_name in sel_names])) n_new, n_old = t.cptree.GetEntriesFast(), t.tree.GetEntriesFast() print '\t Copying {0}/{1} tree: {2:.1%} out of {3} events selected...'.format(t.location, t.tree.GetName(), n_new / n_old, n_old)