def test_0166_IndexedOptionArray(): array = ak.Array([[2, 3, 5], None, [], [7, 11], None, [13], None, [17, 19]]) assert ak.to_list(ak.prod( array, axis=-1)) == [30, None, 1, 77, None, 13, None, 323] array = ak.Array([[[2, 3], [5]], None, [], [[7], [11]], None, [[13]], None, [[17, 19]]]) assert ak.to_list(ak.prod(array, axis=-1)) == [ [6, 5], None, [], [7, 11], None, [13], None, [323], ] array = ak.Array([[[2, 3], None, [5]], [], [[7], [11]], [[13]], [None, [17], [19]]]) ak.to_list(ak.prod(array, axis=-1)) == [ [6, None, 5], [], [7, 11], [13], [None, 17, 19], ] array = ak.Array([[6, None, 5], [], [7, 11], [13], [None, 17, 19]]) assert ak.to_list(ak.prod(array, axis=-1)) == [30, 1, 77, 13, 323]
def combine(eff, sf): # tagged SF = SF*eff / eff = SF tagged_sf = awkward.prod(sf[passbtag], axis=-1) # untagged SF = (1 - SF*eff) / (1 - eff) untagged_sf = awkward.prod(((1 - sf * eff) / (1 - eff))[~passbtag], axis=-1) return awkward.fill_none(tagged_sf * untagged_sf, 1.) # TODO: move None guard to coffea
def test(): assert ak.to_list(ak.prod(ak.Array([[[2, 3, 5]], [[7], [11]], [[]]]), axis=-1)) == [ [30], [7, 11], [1], ] assert ak.to_list(ak.prod(ak.Array([[[2, 3, 5]], [[7], [11]], []]), axis=-1)) == [ [30], [7, 11], [], ]
def flip_weight(self, electron): #f_1 = self.evaluator['el'](electron.pt[:,0:1], abs(electron.eta[:,0:1])) #f_2 = self.evaluator['el'](electron.pt[:,1:2], abs(electron.eta[:,1:2])) # For custom measurements f_1 = yahist_2D_lookup(self.ratio, electron.pt[:, 0:1], abs(electron.eta[:, 0:1])) f_2 = yahist_2D_lookup(self.ratio, electron.pt[:, 1:2], abs(electron.eta[:, 1:2])) # I'm using ak.prod and ak.sum to replace empty arrays by 1 and 0, respectively weight = ak.sum(f_1 / (1 - f_1), axis=1) * ak.prod( 1 - f_2 / (1 - f_2), axis=1) + ak.sum( f_2 / (1 - f_2), axis=1) * ak.prod(1 - f_1 / (1 - f_1), axis=1) return weight
def test_reducers(): # axis=None reducers are implemented in NumPy. assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == 6 + 6j assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == -12 + 12j # axis != None reducers are implemented in libawkward; this should be ReducerSum. assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ 3 + 3j, 0 + 0j, 3 + 3j, ] # And this is in ReducerProd. assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ 0 + 4j, 1 + 0j, 3 + 3j, ] # ReducerCount, ReducerCountNonzero, ReducerAny, and ReducerAll work. assert ak.count(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [2, 0, 1] assert ak.count_nonzero(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [2, 0, 1] assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ True, False, True, ] assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]), axis=1).tolist() == [ True, True, True, ] assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]), axis=1).tolist() == [True, False, True] assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]), axis=1).tolist() == [False, True, True]
def test_other_axis_values(): ak_array = ak.Array([ [[2, 3, 5], [], [], [7, 11, 13]], [[17, 19, 23], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-1).tolist() == [ [2 * 3 * 5, 1, 1, 7 * 11 * 13], [17 * 19 * 23, 1, 1, 29 * 31 * 37], ] assert ak.prod(ak_array, axis=-2).tolist() == [ [2 * 7, 3 * 11, 5 * 13], [17 * 29, 19 * 31, 23 * 37], ] assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19, 5 * 23], [], [], [7 * 29, 11 * 31, 13 * 37], ]
def flip_ratio(self, lepton1, lepton2): """takes a dilepton event and weights it based on the odds that one of the leptons has a charge flip""" flip1 = yahist_2D_lookup(self.ratio, lepton1.pt, abs(lepton1.eta)) flip2 = yahist_2D_lookup(self.ratio, lepton2.pt, abs(lepton2.eta)) flip_rate1 = (ak.prod(flip1, axis=1) * ak.prod(1 / (1 - flip1), axis=1) * ak.prod(1 - flip2 / (1 - flip2), axis=1)) + ( ak.prod(flip2, axis=1) * ak.prod(1 / (1 - flip2), axis=1) * ak.prod(1 - flip1 / (1 - flip1), axis=1)) return flip_rate1
def test_prod(): array = ak.Array( np.array(["2020-07-27T10:41:11", "2019-01-01", "2020-01-01"], "datetime64[s]") ) with pytest.raises(ValueError): ak.prod(array, axis=-1)
def process(self, events): np.random.seed( 10 ) # sets seed so values from random distributions are reproducible (JER corrections) output = self.accumulator.identity() self.sample_name = events.metadata['dataset'] ## make event weights # data or MC distinction made internally mu_evt_weights = MCWeights.get_event_weights( events, year=args.year, corrections=self.corrections) el_evt_weights = MCWeights.get_event_weights( events, year=args.year, corrections=self.corrections) ## initialize selections and regions selection = PackedSelection() regions = { 'Muon': { 'Tight': { 'btagPass': { '3Jets': { 'lep_and_filter_pass', 'passing_jets', 'jets_3', 'tight_MU', 'DeepCSV_pass' }, '4PJets': { 'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'tight_MU', 'DeepCSV_pass' }, }, #'btagFail' : { # '3Jets' : {'lep_and_filter_pass', 'passing_jets', 'jets_3' , 'tight_MU', 'DeepCSV_fail'}, # '4PJets' : {'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'tight_MU', 'DeepCSV_fail'}, #}, }, #'Loose' : { # 'btagPass' : { # '3Jets' : {'lep_and_filter_pass', 'passing_jets', 'jets_3' , 'loose_MU', 'DeepCSV_pass'}, # '4PJets' : {'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'loose_MU', 'DeepCSV_pass'}, # }, # 'btagFail' : { # '3Jets' : {'lep_and_filter_pass', 'passing_jets', 'jets_3' , 'loose_MU', 'DeepCSV_fail'}, # '4PJets' : {'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'loose_MU', 'DeepCSV_fail'}, # }, #}, }, 'Electron': { 'Tight': { 'btagPass': { '3Jets': { 'lep_and_filter_pass', 'passing_jets', 'jets_3', 'tight_EL', 'DeepCSV_pass' }, '4PJets': { 'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'tight_EL', 'DeepCSV_pass' }, }, #'btagFail' : { # '3Jets' : {'lep_and_filter_pass', 'passing_jets', 'jets_3' , 'tight_EL', 'DeepCSV_fail'}, # '4PJets' : {'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'tight_EL', 'DeepCSV_fail'}, #}, }, #'Loose' : { # 'btagPass' : { # '3Jets' : {'lep_and_filter_pass', 'passing_jets', 'jets_3' , 'loose_EL', 'DeepCSV_pass'}, # '4PJets' : {'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'loose_EL', 'DeepCSV_pass'}, # }, # 'btagFail' : { # '3Jets' : {'lep_and_filter_pass', 'passing_jets', 'jets_3' , 'loose_EL', 'DeepCSV_fail'}, # '4PJets' : {'lep_and_filter_pass', 'passing_jets', 'jets_4p', 'loose_EL', 'DeepCSV_fail'}, # }, #}, }, } # get all passing leptons lep_and_filter_pass = objsel.select_leptons(events, year=args.year, cutflow=output['cutflow']) selection.add('lep_and_filter_pass', lep_and_filter_pass) # jet selection passing_jets = objsel.jets_selection(events, year=args.year, corrections=self.corrections, cutflow=output['cutflow']) output['cutflow'][ 'nEvts passing jet and lepton obj selection'] += ak.sum( passing_jets & lep_and_filter_pass) selection.add('passing_jets', passing_jets) selection.add('jets_3', ak.num(events['SelectedJets']) == 3) selection.add('jets_4p', ak.num(events['SelectedJets']) > 3) # only for getting btag weights selection.add('DeepCSV_pass', ak.sum(events['SelectedJets'][btag_wps[0]], axis=1) >= 2) # sort jets by btag value events['SelectedJets'] = events['SelectedJets'][ak.argsort( events['SelectedJets']['btagDeepB'], ascending=False )] if btaggers[0] == 'DeepCSV' else events['SelectedJets'][ak.argsort( events['SelectedJets']['btagDeepFlavB'], ascending=False)] # btag fail sideband deepcsv_sorted = events['SelectedJets'][ak.argsort( events['SelectedJets']['btagDeepB'], ascending=False)]['btagDeepB'] valid_counts_inds = ak.where(ak.num(events['SelectedJets']) > 1)[0] deepcsv_fail = np.zeros(len(events)).astype(bool) deepcsv_fail[valid_counts_inds] = ( deepcsv_sorted[valid_counts_inds][:, 0] < btag_values[args.year]['btagDeepB']['DeepCSV' + wps_to_use[0]]) & ( deepcsv_sorted[valid_counts_inds][:, 1] < btag_values[args.year]['btagDeepB']['DeepCSV' + wps_to_use[0]]) selection.add( 'DeepCSV_fail', deepcsv_fail ) # highest and second highest DeepCSV values don't pass tight and loose WPs #set_trace() if isData_: isSE_Data = self.sample_name.startswith('data_SingleElectron') isSM_Data = self.sample_name.startswith('data_SingleMuon') runs = events.run lumis = events.luminosityBlock Golden_Json_LumiMask = lumi_tools.LumiMask(lumiMask_path) LumiMask = Golden_Json_LumiMask.__call__( runs, lumis) ## returns array of valid events selection.add('lumimask', LumiMask) ## object selection and add different selections if isSM_Data: del regions['Electron'] ## muons selection.add( 'tight_MU', ak.sum(events['Muon']['TIGHTMU'], axis=1) == 1) # one muon passing TIGHT criteria selection.add( 'loose_MU', ak.sum(events['Muon']['LOOSEMU'], axis=1) == 1) # one muon passing LOOSE criteria if isSE_Data: del regions['Muon'] ## electrons selection.add( 'tight_EL', ak.sum(events['Electron']['TIGHTEL'], axis=1) == 1) # one electron passing TIGHT criteria selection.add( 'loose_EL', ak.sum(events['Electron']['LOOSEEL'], axis=1) == 1) # one electron passing LOOSE criteria for lepton in regions.keys(): for leptype in regions[lepton].keys(): for btagregion in regions[lepton][leptype].keys(): for jmult in regions[lepton][leptype][btagregion].keys( ): regions[lepton][leptype][btagregion][jmult].update( {'lumimask'}) else: ## add different selections ## muons selection.add( 'tight_MU', ak.sum(events['Muon']['TIGHTMU'], axis=1) == 1) # one muon passing TIGHT criteria selection.add( 'loose_MU', ak.sum(events['Muon']['LOOSEMU'], axis=1) == 1) # one muon passing LOOSE criteria ## electrons selection.add( 'tight_EL', ak.sum(events['Electron']['TIGHTEL'], axis=1) == 1) # one electron passing TIGHT criteria selection.add( 'loose_EL', ak.sum(events['Electron']['LOOSEEL'], axis=1) == 1) # one electron passing LOOSE criteria ### apply lepton SFs to MC (only applicable to tight leptons) if 'LeptonSF' in corrections.keys(): tight_mu_cut = selection.require( tight_MU=True ) # find events passing muon object selection with one tight muon tight_muons = events['Muon'][tight_mu_cut][( events['Muon'][tight_mu_cut]['TIGHTMU'] == True)] muSFs_dict = MCWeights.get_lepton_sf( year=args.year, lepton='Muons', corrections=self.corrections['LeptonSF'], pt=ak.flatten(tight_muons['pt']), eta=ak.flatten(tight_muons['eta'])) mu_reco_cen = np.ones(len(events)) mu_reco_err = np.zeros(len(events)) mu_trig_cen = np.ones(len(events)) mu_trig_err = np.zeros(len(events)) mu_reco_cen[tight_mu_cut] = muSFs_dict['RECO_CEN'] mu_reco_err[tight_mu_cut] = muSFs_dict['RECO_ERR'] mu_trig_cen[tight_mu_cut] = muSFs_dict['TRIG_CEN'] mu_trig_err[tight_mu_cut] = muSFs_dict['TRIG_ERR'] mu_evt_weights.add('RECO', mu_reco_cen, mu_reco_err, mu_reco_err, shift=True) mu_evt_weights.add('TRIG', mu_trig_cen, mu_trig_err, mu_trig_err, shift=True) tight_el_cut = selection.require( tight_EL=True ) # find events passing electron object selection with one tight electron tight_electrons = events['Electron'][tight_el_cut][( events['Electron'][tight_el_cut]['TIGHTEL'] == True)] elSFs_dict = MCWeights.get_lepton_sf( year=args.year, lepton='Electrons', corrections=self.corrections['LeptonSF'], pt=ak.flatten(tight_electrons['pt']), eta=ak.flatten(tight_electrons['etaSC'])) el_reco_cen = np.ones(len(events)) el_reco_err = np.zeros(len(events)) el_trig_cen = np.ones(len(events)) el_trig_err = np.zeros(len(events)) el_reco_cen[tight_el_cut] = elSFs_dict['RECO_CEN'] el_reco_err[tight_el_cut] = elSFs_dict['RECO_ERR'] el_trig_cen[tight_el_cut] = elSFs_dict['TRIG_CEN'] el_trig_err[tight_el_cut] = elSFs_dict['TRIG_ERR'] el_evt_weights.add('RECO', el_reco_cen, el_reco_err, el_reco_err, shift=True) el_evt_weights.add('TRIG', el_trig_cen, el_trig_err, el_trig_err, shift=True) ## apply btagging SFs to MC if corrections['BTagSF'] == True: deepcsv_cen = np.ones(len(events)) threeJets_cut = selection.require(lep_and_filter_pass=True, passing_jets=True, jets_3=True) deepcsv_3j_wts = self.corrections['BTag_Constructors'][ 'DeepCSV']['3Jets'].get_scale_factor( jets=events['SelectedJets'][threeJets_cut], passing_cut='DeepCSV' + wps_to_use[0]) deepcsv_cen[threeJets_cut] = ak.prod(deepcsv_3j_wts['central'], axis=1) fourplusJets_cut = selection.require(lep_and_filter_pass=True, passing_jets=True, jets_4p=True) deepcsv_4pj_wts = self.corrections['BTag_Constructors'][ 'DeepCSV']['4PJets'].get_scale_factor( jets=events['SelectedJets'][fourplusJets_cut], passing_cut='DeepCSV' + wps_to_use[0]) deepcsv_cen[fourplusJets_cut] = ak.prod( deepcsv_4pj_wts['central'], axis=1) if isTTbar_: # find gen level particles for ttbar system genpsel.select(events, mode='NORMAL') selection.add('semilep', ak.num(events['SL']) > 0) if isTTSL_: for lepton in regions.keys(): for leptype in regions[lepton].keys(): for btagregion in regions[lepton][leptype].keys(): for jmult in regions[lepton][leptype][ btagregion].keys(): regions[lepton][leptype][btagregion][ jmult].update({'semilep'}) ## fill hists for each region for lepton in regions.keys(): evt_weights = mu_evt_weights if lepton == 'Muon' else el_evt_weights for leptype in regions[lepton].keys(): for btagregion in regions[lepton][leptype].keys(): for jmult in regions[lepton][leptype][btagregion].keys(): cut = selection.all( *regions[lepton][leptype][btagregion][jmult]) #set_trace() #if args.debug: print(lepton, leptype, btagregion, jmult) if cut.sum() > 0: # get leptons ltype = 'MU' if lepton == 'Muon' else 'EL' if 'loose_or_tight_%s' % ltype in regions[lepton][ leptype][btagregion][jmult]: leptons = events[lepton][cut][( (events[lepton][cut]['TIGHT%s' % ltype] == True) | (events[lepton][cut]['LOOSE%s' % ltype] == True))] elif 'tight_%s' % ltype in regions[lepton][ leptype][btagregion][jmult]: leptons = events[lepton][cut][( events[lepton][cut]['TIGHT%s' % ltype] == True)] elif 'loose_%s' % ltype in regions[lepton][ leptype][btagregion][jmult]: leptons = events[lepton][cut][( events[lepton][cut]['LOOSE%s' % ltype] == True)] else: raise ValueError( "Not sure what lepton type to choose for event" ) # get jets and MET jets, met = events['SelectedJets'][cut], events[ 'SelectedMET'][cut] # find best permutations best_perms = ttpermutator.find_best_permutations( jets=jets, leptons=leptons, MET=met, btagWP=btag_wps[0], btag_req=False if btagregion == 'btagFail' else True) valid_perms = ak.num(best_perms['TTbar'].pt) > 0 bp_status = np.zeros( cut.size, dtype=int ) # 0 == '' (no gen matching), 1 == 'right', 2 == 'matchable', 3 == 'unmatchable', 4 == 'sl_tau', 5 == 'noslep' # get matched permutation (semilep ttbar only) if isTTbar_: semilep_evts = selection.require(semilep=True) bp_status[~semilep_evts] = 5 if semilep_evts.sum() > 0: # find matched permutations mp = ttmatcher.best_match( gen_hyp=events['SL'][cut], jets=jets, leptons=leptons, met=met) perm_cat_array = compare_matched_best_perms( mp, best_perms, njets=jmult) bp_status[cut] = perm_cat_array sl_tau_evts = ak.where( ak.fill_none( ak.pad_none( np.abs(events['SL'] ['Lepton'].pdgId) == 15, 1), False) == True)[0] bp_status[sl_tau_evts] = 4 ## create MT regions MT = make_vars.MT(leptons, met) MTHigh = ak.flatten(MT[valid_perms] >= MTcut) wts = evt_weights.weight( )[cut][valid_perms][MTHigh] if isData_ else ( evt_weights.weight() * deepcsv_cen)[cut][valid_perms][MTHigh] #set_trace() # fill hists if isInt_: # fill hists for positive weights pos_evts = np.where(wts > 0) self.sample_name = '%s_pos' % events.metadata[ 'dataset'] output = self.fill_selection_hists( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=bp_status[cut][valid_perms] [MTHigh][pos_evts], perm=best_perms[valid_perms][MTHigh] [pos_evts], MTvals=MT[valid_perms][MTHigh][pos_evts], evt_wts=wts[pos_evts]) output = self.fill_jet_hists( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=bp_status[cut][valid_perms] [MTHigh][pos_evts], obj=jets[valid_perms][MTHigh][pos_evts], evt_wts=wts[pos_evts]) # fill hists for negative weights neg_evts = np.where(wts < 0) self.sample_name = '%s_neg' % events.metadata[ 'dataset'] output = self.fill_selection_hists( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=bp_status[cut][valid_perms] [MTHigh][neg_evts], perm=best_perms[valid_perms][MTHigh] [neg_evts], MTvals=MT[valid_perms][MTHigh][neg_evts], evt_wts=wts[neg_evts]) output = self.fill_jet_hists( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=bp_status[cut][valid_perms] [MTHigh][neg_evts], obj=jets[valid_perms][MTHigh][neg_evts], evt_wts=wts[neg_evts]) else: output = self.fill_selection_hists( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=bp_status[cut][valid_perms] [MTHigh], perm=best_perms[valid_perms][MTHigh], MTvals=MT[valid_perms][MTHigh], evt_wts=wts) output = self.fill_jet_hists( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=bp_status[cut][valid_perms] [MTHigh], obj=jets[valid_perms][MTHigh], evt_wts=wts) return output
def test(): ak_array = ak.Array([ [[2, 3], [], [], [7, 11, 13]], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [-1], [], [7, 11, 13]], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [-1], [], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [], [], [7, 11, 13]], [[17, 19], [-1], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [-1], [], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [], [-1], [7, 11, 13]], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [-1], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [], [], [7, 11, 13]], [[17, 19], [], [-1], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [-1], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [-1], [], [7, 11, 13]], [[17, 19], [39], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [-39], [], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [], [-1], [7, 11, 13]], [[17, 19], [], [39], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [-39], [7 * 29, 11 * 31, 13 * 37], ] ak_array = ak.Array([ [[2, 3], [], []], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], []], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3]], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], [-1], []], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [-1], [], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], [], []], [[17, 19], [-1], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [-1], [], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], [], [-1]], [[17, 19], [], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [-1], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], [], []], [[17, 19], [], [-1], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [-1], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], [-1], []], [[17, 19], [39], [], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [-39], [], [29, 31, 37], ] ak_array = ak.Array([ [[2, 3], [], [-1]], [[17, 19], [], [39], [29, 31, 37]], ]) assert ak.prod(ak_array, axis=-3).tolist() == [ [2 * 17, 3 * 19], [], [-39], [29, 31, 37], ]
def process(self, events): # Dataset parameters dataset = events.metadata['dataset'] histAxisName = self._samples[dataset]['histAxisName'] year = self._samples[dataset]['year'] xsec = self._samples[dataset]['xsec'] sow = self._samples[dataset]['nSumOfWeights'] isData = self._samples[dataset]['isData'] datasets = [ 'SingleMuon', 'SingleElectron', 'EGamma', 'MuonEG', 'DoubleMuon', 'DoubleElectron' ] for d in datasets: if d in dataset: dataset = dataset.split('_')[0] # Initialize objects met = events.MET e = events.Electron mu = events.Muon tau = events.Tau j = events.Jet # Muon selection mu['isPres'] = isPresMuon(mu.dxy, mu.dz, mu.sip3d, mu.looseId) mu['isTight'] = isTightMuon(mu.pt, mu.eta, mu.dxy, mu.dz, mu.pfRelIso03_all, mu.sip3d, mu.mvaTTH, mu.mediumPromptId, mu.tightCharge, mu.looseId, minpt=10) mu['isGood'] = mu['isPres'] & mu['isTight'] leading_mu = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] leading_mu = leading_mu[leading_mu.isGood] mu = mu[mu.isGood] mu_pres = mu[mu.isPres] # Electron selection e['isPres'] = isPresElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.lostHits, minpt=15) e['isTight'] = isTightElec(e.pt, e.eta, e.dxy, e.dz, e.miniPFRelIso_all, e.sip3d, e.mvaTTH, e.mvaFall17V2Iso, e.lostHits, e.convVeto, e.tightCharge, e.sieie, e.hoe, e.eInvMinusPInv, minpt=15) e['isClean'] = isClean(e, mu, drmin=0.05) e['isGood'] = e['isPres'] & e['isTight'] & e['isClean'] leading_e = e[ak.argmax(e.pt, axis=-1, keepdims=True)] leading_e = leading_e[leading_e.isGood] e = e[e.isGood] e_pres = e[e.isPres & e.isClean] # Tau selection tau['isPres'] = isPresTau(tau.pt, tau.eta, tau.dxy, tau.dz, tau.leadTkPtOverTauPt, tau.idAntiMu, tau.idAntiEle, tau.rawIso, tau.idDecayModeNewDMs, minpt=20) tau['isClean'] = isClean(tau, e_pres, drmin=0.4) & isClean( tau, mu_pres, drmin=0.4) tau['isGood'] = tau['isPres'] # & tau['isClean'], for the moment tau = tau[tau.isGood] nElec = ak.num(e) nMuon = ak.num(mu) nTau = ak.num(tau) twoLeps = (nElec + nMuon) == 2 threeLeps = (nElec + nMuon) == 3 twoElec = (nElec == 2) twoMuon = (nMuon == 2) e0 = e[ak.argmax(e.pt, axis=-1, keepdims=True)] m0 = mu[ak.argmax(mu.pt, axis=-1, keepdims=True)] # Attach the lepton SFs to the electron and muons collections AttachElectronSF(e, year=year) AttachMuonSF(mu, year=year) # Create a lepton (muon+electron) collection and calculate a per event lepton SF leps = ak.concatenate([e, mu], axis=-1) events['lepSF_nom'] = ak.prod(leps.sf_nom, axis=-1) events['lepSF_hi'] = ak.prod(leps.sf_hi, axis=-1) events['lepSF_lo'] = ak.prod(leps.sf_lo, axis=-1) # Jet selection jetptname = 'pt_nom' if hasattr(j, 'pt_nom') else 'pt' ### Jet energy corrections if not isData: j["pt_raw"] = (1 - j.rawFactor) * j.pt j["mass_raw"] = (1 - j.rawFactor) * j.mass j["pt_gen"] = ak.values_astype(ak.fill_none(j.matched_gen.pt, 0), np.float32) j["rho"] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, j.pt)[0] events_cache = events.caches[0] corrected_jets = jet_factory.build(j, lazy_cache=events_cache) #print('jet pt: ',j.pt) #print('cor pt: ',corrected_jets.pt) #print('jes up: ',corrected_jets.JES_jes.up.pt) #print('jes down: ',corrected_jets.JES_jes.down.pt) #print(ak.fields(corrected_jets)) ''' # SYSTEMATICS jets = corrected_jets if(self.jetSyst == 'JERUp'): jets = corrected_jets.JER.up elif(self.jetSyst == 'JERDown'): jets = corrected_jets.JER.down elif(self.jetSyst == 'JESUp'): jets = corrected_jets.JES_jes.up elif(self.jetSyst == 'JESDown'): jets = corrected_jets.JES_jes.down ''' j['isGood'] = isTightJet(getattr(j, jetptname), j.eta, j.jetId, j.neHEF, j.neEmEF, j.chHEF, j.chEmEF, j.nConstituents) #j['isgood'] = isGoodJet(j.pt, j.eta, j.jetId) #j['isclean'] = isClean(j, e, mu) j['isClean'] = isClean(j, e, drmin=0.4) & isClean( j, mu, drmin=0.4) # & isClean(j, tau, drmin=0.4) goodJets = j[(j.isClean) & (j.isGood)] njets = ak.num(goodJets) ht = ak.sum(goodJets.pt, axis=-1) j0 = goodJets[ak.argmax(goodJets.pt, axis=-1, keepdims=True)] #nbtags = ak.num(goodJets[goodJets.btagDeepFlavB > 0.2770]) # Loose DeepJet WP if year == 2017: btagwpl = 0.0532 #WP loose else: btagwpl = 0.0490 #WP loose isBtagJetsLoose = (goodJets.btagDeepB > btagwpl) isNotBtagJetsLoose = np.invert(isBtagJetsLoose) nbtagsl = ak.num(goodJets[isBtagJetsLoose]) # Medium DeepJet WP if year == 2017: btagwpm = 0.3040 #WP medium else: btagwpm = 0.2783 #WP medium isBtagJetsMedium = (goodJets.btagDeepB > btagwpm) isNotBtagJetsMedium = np.invert(isBtagJetsMedium) nbtagsm = ak.num(goodJets[isBtagJetsMedium]) # Btag SF following 1a) in https://twiki.cern.ch/twiki/bin/viewauth/CMS/BTagSFMethods btagSF = np.ones_like(ht) btagSFUp = np.ones_like(ht) btagSFDo = np.ones_like(ht) if not isData: pt = goodJets.pt abseta = np.abs(goodJets.eta) flav = goodJets.hadronFlavour bJetSF = GetBTagSF(abseta, pt, flav) bJetSFUp = GetBTagSF(abseta, pt, flav, sys=1) bJetSFDo = GetBTagSF(abseta, pt, flav, sys=-1) bJetEff = GetBtagEff(abseta, pt, flav, year) bJetEff_data = bJetEff * bJetSF bJetEff_dataUp = bJetEff * bJetSFUp bJetEff_dataDo = bJetEff * bJetSFDo pMC = ak.prod(bJetEff[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff[isNotBtagJetsMedium]), axis=-1) pData = ak.prod(bJetEff_data[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_data[isNotBtagJetsMedium]), axis=-1) pDataUp = ak.prod( bJetEff_dataUp[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_dataUp[isNotBtagJetsMedium]), axis=-1) pDataDo = ak.prod( bJetEff_dataDo[isBtagJetsMedium], axis=-1) * ak.prod( (1 - bJetEff_dataDo[isNotBtagJetsMedium]), axis=-1) pMC = ak.where(pMC == 0, 1, pMC) # removeing zeroes from denominator... btagSF = pData / pMC btagSFUp = pDataUp / pMC btagSFDo = pDataUp / pMC ################################################################## ### 2 same-sign leptons ################################################################## # emu singe = e[(nElec == 1) & (nMuon == 1) & (e.pt > -1)] singm = mu[(nElec == 1) & (nMuon == 1) & (mu.pt > -1)] em = ak.cartesian({"e": singe, "m": singm}) emSSmask = (em.e.charge * em.m.charge > 0) emSS = em[emSSmask] nemSS = len(ak.flatten(emSS)) # ee and mumu # pt>-1 to preserve jagged dimensions ee = e[(nElec == 2) & (nMuon == 0) & (e.pt > -1)] mm = mu[(nElec == 0) & (nMuon == 2) & (mu.pt > -1)] sumcharge = ak.sum(e.charge, axis=-1) + ak.sum(mu.charge, axis=-1) eepairs = ak.combinations(ee, 2, fields=["e0", "e1"]) eeSSmask = (eepairs.e0.charge * eepairs.e1.charge > 0) eeonZmask = (np.abs((eepairs.e0 + eepairs.e1).mass - 91.2) < 10) eeoffZmask = (eeonZmask == 0) mmpairs = ak.combinations(mm, 2, fields=["m0", "m1"]) mmSSmask = (mmpairs.m0.charge * mmpairs.m1.charge > 0) mmonZmask = (np.abs((mmpairs.m0 + mmpairs.m1).mass - 91.2) < 10) mmoffZmask = (mmonZmask == 0) eeSSonZ = eepairs[eeSSmask & eeonZmask] eeSSoffZ = eepairs[eeSSmask & eeoffZmask] mmSSonZ = mmpairs[mmSSmask & mmonZmask] mmSSoffZ = mmpairs[mmSSmask & mmoffZmask] neeSS = len(ak.flatten(eeSSonZ)) + len(ak.flatten(eeSSoffZ)) nmmSS = len(ak.flatten(mmSSonZ)) + len(ak.flatten(mmSSoffZ)) print('Same-sign events [ee, emu, mumu] = [%i, %i, %i]' % (neeSS, nemSS, nmmSS)) # Cuts eeSSmask = (ak.num(eeSSmask[eeSSmask]) > 0) mmSSmask = (ak.num(mmSSmask[mmSSmask]) > 0) eeonZmask = (ak.num(eeonZmask[eeonZmask]) > 0) eeoffZmask = (ak.num(eeoffZmask[eeoffZmask]) > 0) mmonZmask = (ak.num(mmonZmask[mmonZmask]) > 0) mmoffZmask = (ak.num(mmoffZmask[mmoffZmask]) > 0) emSSmask = (ak.num(emSSmask[emSSmask]) > 0) ################################################################## ### 3 leptons ################################################################## # eem muon_eem = mu[(nElec == 2) & (nMuon == 1) & (mu.pt > -1)] elec_eem = e[(nElec == 2) & (nMuon == 1) & (e.pt > -1)] ee_eem = ak.combinations(elec_eem, 2, fields=["e0", "e1"]) ee_eemZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) < 10) ee_eemOffZmask = (ee_eem.e0.charge * ee_eem.e1.charge < 1) & (np.abs( (ee_eem.e0 + ee_eem.e1).mass - 91.2) > 10) ee_eemZmask = (ak.num(ee_eemZmask[ee_eemZmask]) > 0) ee_eemOffZmask = (ak.num(ee_eemOffZmask[ee_eemOffZmask]) > 0) eepair_eem = (ee_eem.e0 + ee_eem.e1) trilep_eem = eepair_eem + muon_eem #ak.cartesian({"e0":ee_eem.e0,"e1":ee_eem.e1, "m":muon_eem}) # mme muon_mme = mu[(nElec == 1) & (nMuon == 2) & (mu.pt > -1)] elec_mme = e[(nElec == 1) & (nMuon == 2) & (e.pt > -1)] mm_mme = ak.combinations(muon_mme, 2, fields=["m0", "m1"]) mm_mmeZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) < 10) mm_mmeOffZmask = (mm_mme.m0.charge * mm_mme.m1.charge < 1) & (np.abs( (mm_mme.m0 + mm_mme.m1).mass - 91.2) > 10) mm_mmeZmask = (ak.num(mm_mmeZmask[mm_mmeZmask]) > 0) mm_mmeOffZmask = (ak.num(mm_mmeOffZmask[mm_mmeOffZmask]) > 0) mmpair_mme = (mm_mme.m0 + mm_mme.m1) trilep_mme = mmpair_mme + elec_mme mZ_mme = mmpair_mme.mass mZ_eem = eepair_eem.mass m3l_eem = trilep_eem.mass m3l_mme = trilep_mme.mass # eee and mmm eee = e[(nElec == 3) & (nMuon == 0) & (e.pt > -1)] mmm = mu[(nElec == 0) & (nMuon == 3) & (mu.pt > -1)] eee_leps = ak.combinations(eee, 3, fields=["e0", "e1", "e2"]) mmm_leps = ak.combinations(mmm, 3, fields=["m0", "m1", "m2"]) ee_pairs = ak.combinations(eee, 2, fields=["e0", "e1"]) mm_pairs = ak.combinations(mmm, 2, fields=["m0", "m1"]) ee_pairs_index = ak.argcombinations(eee, 2, fields=["e0", "e1"]) mm_pairs_index = ak.argcombinations(mmm, 2, fields=["m0", "m1"]) mmSFOS_pairs = mm_pairs[ (np.abs(mm_pairs.m0.pdgId) == np.abs(mm_pairs.m1.pdgId)) & (mm_pairs.m0.charge != mm_pairs.m1.charge)] offZmask_mm = ak.all( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) > 10., axis=1, keepdims=True) & (ak.num(mmSFOS_pairs) > 0) onZmask_mm = ak.any( np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2) < 10., axis=1, keepdims=True) eeSFOS_pairs = ee_pairs[ (np.abs(ee_pairs.e0.pdgId) == np.abs(ee_pairs.e1.pdgId)) & (ee_pairs.e0.charge != ee_pairs.e1.charge)] offZmask_ee = ak.all( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) > 10, axis=1, keepdims=True) & (ak.num(eeSFOS_pairs) > 0) onZmask_ee = ak.any( np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2) < 10, axis=1, keepdims=True) # Create masks **for event selection** eeeOnZmask = (ak.num(onZmask_ee[onZmask_ee]) > 0) eeeOffZmask = (ak.num(offZmask_ee[offZmask_ee]) > 0) mmmOnZmask = (ak.num(onZmask_mm[onZmask_mm]) > 0) mmmOffZmask = (ak.num(offZmask_mm[offZmask_mm]) > 0) # Now we need to create masks for the leptons in order to select leptons from the Z boson candidate (in onZ categories) ZeeMask = ak.argmin(np.abs((eeSFOS_pairs.e0 + eeSFOS_pairs.e1).mass - 91.2), axis=1, keepdims=True) ZmmMask = ak.argmin(np.abs((mmSFOS_pairs.m0 + mmSFOS_pairs.m1).mass - 91.2), axis=1, keepdims=True) Zee = eeSFOS_pairs[ZeeMask] Zmm = mmSFOS_pairs[ZmmMask] eZ0 = Zee.e0[ak.num(eeSFOS_pairs) > 0] eZ1 = Zee.e1[ak.num(eeSFOS_pairs) > 0] eZ = eZ0 + eZ1 mZ0 = Zmm.m0[ak.num(mmSFOS_pairs) > 0] mZ1 = Zmm.m1[ak.num(mmSFOS_pairs) > 0] mZ = mZ0 + mZ1 mZ_eee = eZ.mass mZ_mmm = mZ.mass # And for the W boson ZmmIndices = mm_pairs_index[ZmmMask] ZeeIndices = ee_pairs_index[ZeeMask] eW = eee[~ZeeIndices.e0 | ~ZeeIndices.e1] mW = mmm[~ZmmIndices.m0 | ~ZmmIndices.m1] triElec = eee_leps.e0 + eee_leps.e1 + eee_leps.e2 triMuon = mmm_leps.m0 + mmm_leps.m1 + mmm_leps.m2 m3l_eee = triElec.mass m3l_mmm = triMuon.mass ################################################################## ### >=4 leptons ################################################################## # 4lep cat is4lmask = ((nElec + nMuon) >= 4) muon_4l = mu[(is4lmask) & (mu.pt > -1)] elec_4l = e[(is4lmask) & (e.pt > -1)] # selecting 4 leading leptons leptons = ak.concatenate([e, mu], axis=-1) leptons_sorted = leptons[ak.argsort(leptons.pt, axis=-1, ascending=False)] lep4l = leptons_sorted[:, 0:4] e4l = lep4l[abs(lep4l.pdgId) == 11] mu4l = lep4l[abs(lep4l.pdgId) == 13] nElec4l = ak.num(e4l) nMuon4l = ak.num(mu4l) # Triggers trig_eeSS = passTrigger(events, 'ee', isData, dataset) trig_mmSS = passTrigger(events, 'mm', isData, dataset) trig_emSS = passTrigger(events, 'em', isData, dataset) trig_eee = passTrigger(events, 'eee', isData, dataset) trig_mmm = passTrigger(events, 'mmm', isData, dataset) trig_eem = passTrigger(events, 'eem', isData, dataset) trig_mme = passTrigger(events, 'mme', isData, dataset) trig_4l = triggerFor4l(events, nMuon, nElec, isData, dataset) # MET filters # Weights genw = np.ones_like(events['event']) if ( isData or len(self._wc_names_lst) > 0) else events['genWeight'] ### We need weights for: normalization, lepSF, triggerSF, pileup, btagSF... weights = {} for r in [ 'all', 'ee', 'mm', 'em', 'eee', 'mmm', 'eem', 'mme', 'eeee', 'eeem', 'eemm', 'mmme', 'mmmm' ]: # weights[r] = coffea.analysis_tools.Weights(len(events)) weights[r] = coffea.analysis_tools.Weights(len(events), storeIndividual=True) if len(self._wc_names_lst) > 0: sow = np.ones_like( sow ) # Not valid in nanoAOD for EFT samples, MUST use SumOfEFTweights at analysis level weights[r].add('norm', genw if isData else (xsec / sow) * genw) weights[r].add('btagSF', btagSF, btagSFUp, btagSFDo) weights[r].add('lepSF', events.lepSF_nom, events.lepSF_hi, events.lepSF_lo) # Extract the EFT quadratic coefficients and optionally use them to calculate the coefficients on the w**2 quartic function # eft_coeffs is never Jagged so convert immediately to numpy for ease of use. eft_coeffs = ak.to_numpy(events['EFTfitCoefficients']) if hasattr( events, "EFTfitCoefficients") else None if eft_coeffs is not None: # Check to see if the ordering of WCs for this sample matches what want if self._samples[dataset]['WCnames'] != self._wc_names_lst: eft_coeffs = efth.remap_coeffs( self._samples[dataset]['WCnames'], self._wc_names_lst, eft_coeffs) eft_w2_coeffs = efth.calc_w2_coeffs(eft_coeffs, self._dtype) if ( self._do_errors and eft_coeffs is not None) else None # Selections and cuts selections = PackedSelection() #(dtype='uint64') channels2LSS = ['eeSSonZ', 'eeSSoffZ', 'mmSSonZ', 'mmSSoffZ', 'emSS'] selections.add('eeSSonZ', (eeonZmask) & (eeSSmask) & (trig_eeSS)) selections.add('eeSSoffZ', (eeoffZmask) & (eeSSmask) & (trig_eeSS)) selections.add('mmSSonZ', (mmonZmask) & (mmSSmask) & (trig_mmSS)) selections.add('mmSSoffZ', (mmoffZmask) & (mmSSmask) & (trig_mmSS)) selections.add('emSS', (emSSmask) & (trig_emSS)) channels3L = ['eemSSonZ', 'eemSSoffZ', 'mmeSSonZ', 'mmeSSoffZ'] selections.add('eemSSonZ', (ee_eemZmask) & (trig_eem)) selections.add('eemSSoffZ', (ee_eemOffZmask) & (trig_eem)) selections.add('mmeSSonZ', (mm_mmeZmask) & (trig_mme)) selections.add('mmeSSoffZ', (mm_mmeOffZmask) & (trig_mme)) channels3L += ['eeeSSonZ', 'eeeSSoffZ', 'mmmSSonZ', 'mmmSSoffZ'] selections.add('eeeSSonZ', (eeeOnZmask) & (trig_eee)) selections.add('eeeSSoffZ', (eeeOffZmask) & (trig_eee)) selections.add('mmmSSonZ', (mmmOnZmask) & (trig_mmm)) selections.add('mmmSSoffZ', (mmmOffZmask) & (trig_mmm)) channels4L = ['eeee', 'eeem', 'eemm', 'mmme', 'mmmm'] selections.add('eeee', ((nElec4l == 4) & (nMuon4l == 0)) & (trig_4l)) selections.add('eeem', ((nElec4l == 3) & (nMuon4l == 1)) & (trig_4l)) selections.add('eemm', ((nElec4l == 2) & (nMuon4l == 2)) & (trig_4l)) selections.add('mmme', ((nElec4l == 1) & (nMuon4l == 3)) & (trig_4l)) selections.add('mmmm', ((nElec4l == 0) & (nMuon4l == 4)) & (trig_4l)) selections.add('ch+', (sumcharge > 0)) selections.add('ch-', (sumcharge < 0)) selections.add('ch0', (sumcharge == 0)) levels = ['base', '1+bm2+bl', '1bm', '2+bm'] selections.add('base', (nElec + nMuon >= 2)) selections.add('1+bm2+bl', (nElec + nMuon >= 2) & ((nbtagsm >= 1) & (nbtagsl >= 2))) selections.add('1bm', (nElec + nMuon >= 2) & (nbtagsm == 1)) selections.add('2+bm', (nElec + nMuon >= 2) & (nbtagsm >= 2)) # Variables invMass_eeSSonZ = (eeSSonZ.e0 + eeSSonZ.e1).mass invMass_eeSSoffZ = (eeSSoffZ.e0 + eeSSoffZ.e1).mass invMass_mmSSonZ = (mmSSonZ.m0 + mmSSonZ.m1).mass invMass_mmSSoffZ = (mmSSoffZ.m0 + mmSSoffZ.m1).mass invMass_emSS = (emSS.e + emSS.m).mass varnames = {} varnames['met'] = met.pt varnames['ht'] = ht varnames['njets'] = njets varnames['invmass'] = { 'eeSSonZ': invMass_eeSSonZ, 'eeSSoffZ': invMass_eeSSoffZ, 'mmSSonZ': invMass_mmSSonZ, 'mmSSoffZ': invMass_mmSSoffZ, 'emSS': invMass_emSS, 'eemSSonZ': mZ_eem, 'eemSSoffZ': mZ_eem, 'mmeSSonZ': mZ_mme, 'mmeSSoffZ': mZ_mme, 'eeeSSonZ': mZ_eee, 'eeeSSoffZ': mZ_eee, 'mmmSSonZ': mZ_mmm, 'mmmSSoffZ': mZ_mmm, } varnames['m3l'] = { 'eemSSonZ': m3l_eem, 'eemSSoffZ': m3l_eem, 'mmeSSonZ': m3l_mme, 'mmeSSoffZ': m3l_mme, 'eeeSSonZ': m3l_eee, 'eeeSSoffZ': m3l_eee, 'mmmSSonZ': m3l_mmm, 'mmmSSoffZ': m3l_mmm, } varnames['e0pt'] = e0.pt varnames['e0eta'] = e0.eta varnames['m0pt'] = m0.pt varnames['m0eta'] = m0.eta varnames['j0pt'] = j0.pt varnames['j0eta'] = j0.eta varnames['counts'] = np.ones_like(events['event']) # systematics systList = [] if isData == False: systList = ['nominal'] if self._do_systematics: systList = systList + [ 'lepSFUp', 'lepSFDown', 'btagSFUp', 'btagSFDown' ] else: systList = ['noweight'] # fill Histos hout = self.accumulator.identity() normweights = weights['all'].weight().flatten( ) # Why does it not complain about .flatten() here? sowweights = np.ones_like(normweights) if len( self._wc_names_lst) > 0 else normweights hout['SumOfEFTweights'].fill(sample=histAxisName, SumOfEFTweights=varnames['counts'], weight=sowweights, eft_coeff=eft_coeffs, eft_err_coeff=eft_w2_coeffs) for syst in systList: for var, v in varnames.items(): for ch in channels2LSS + channels3L + channels4L: for sumcharge in ['ch+', 'ch-', 'ch0']: for lev in levels: #find the event weight to be used when filling the histograms weightSyst = syst #in the case of 'nominal', or the jet energy systematics, no weight systematic variation is used (weightSyst=None) if syst in [ 'nominal', 'JERUp', 'JERDown', 'JESUp', 'JESDown' ]: weightSyst = None # no weight systematic for these variations if syst == 'noweight': weight = np.ones(len(events)) # for data else: # call weights.weight() with the name of the systematic to be varied if ch in channels3L: ch_w = ch[:3] elif ch in channels2LSS: ch_w = ch[:2] else: ch_w = ch weight = weights['all'].weight( weightSyst ) if isData else weights[ch_w].weight( weightSyst) cuts = [ch] + [lev] + [sumcharge] cut = selections.all(*cuts) weights_flat = weight[cut].flatten( ) # Why does it not complain about .flatten() here? weights_ones = np.ones_like(weights_flat, dtype=np.int) eft_coeffs_cut = eft_coeffs[ cut] if eft_coeffs is not None else None eft_w2_coeffs_cut = eft_w2_coeffs[ cut] if eft_w2_coeffs is not None else None # filling histos if var == 'invmass': if ((ch in [ 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]) or (ch in channels4L)): continue else: values = ak.flatten(v[ch][cut]) hout['invmass'].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, invmass=values, weight=weights_flat, systematic=syst) elif var == 'm3l': if ((ch in channels2LSS) or (ch in [ 'eeeSSoffZ', 'mmmSSoffZ', 'eeeSSonZ', 'mmmSSonZ' ]) or (ch in channels4L)): continue values = ak.flatten(v[ch][cut]) hout['m3l'].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, m3l=values, weight=weights_flat, systematic=syst) else: values = v[cut] # These all look identical, do we need if/else here? if var == 'ht': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, ht=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'met': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, met=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'njets': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, njets=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'nbtags': hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, nbtags=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'counts': hout[var].fill(counts=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_ones, systematic=syst) elif var == 'j0eta': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, j0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'e0pt': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmmm' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, e0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst ) # Crashing here, not sure why. Related to values? elif var == 'm0pt': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eeee' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, m0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'e0eta': if ch in [ 'mmSSonZ', 'mmSSoffZ', 'mmmSSoffZ', 'mmmSSonZ', 'mmmm' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, e0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'm0eta': if ch in [ 'eeSSonZ', 'eeSSoffZ', 'eeeSSoffZ', 'eeeSSonZ', 'eeee' ]: continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, m0eta=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) elif var == 'j0pt': if lev == 'base': continue values = ak.flatten(values) #values=np.asarray(values) hout[var].fill( eft_coeff=eft_coeffs_cut, eft_err_coeff=eft_w2_coeffs_cut, j0pt=values, sample=histAxisName, channel=ch, cut=lev, sumcharge=sumcharge, weight=weights_flat, systematic=syst) return hout
def GetLeptonSF(pt1, eta1, flavor1, pt2, eta2, flavor2, pt3=None, eta3=None, flavor3=None, pt4=None, eta4=None, flavor4=None, year=2018, sys=0): if sys == 0: if flavor1 == 'm': SF1 = ak.prod( SFevaluator['MuonLooseSF_%i' % year](np.abs(eta1), pt1) * SFevaluator['MuonTightSF_%i' % year](np.abs(eta1), pt1), axis=-1) elif flavor1 == 'e': SF1 = ak.prod( SFevaluator['ElecRecoSF_%i' % year](eta1, pt1) * SFevaluator['ElecLooseSF_%i' % year](np.abs(eta1), pt1) * SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta1), pt1) * SFevaluator['ElecTightSF_%i' % year](np.abs(eta1), pt1), axis=-1) else: print(flavor1, ' is not a valid flavor. Valid flavors: "m" or "e"') if flavor2 == 'm': SF2 = ak.prod( SFevaluator['MuonLooseSF_%i' % year](np.abs(eta2), pt2) * SFevaluator['MuonTightSF_%i' % year](np.abs(eta2), pt2), axis=-1) elif flavor2 == 'e': SF2 = ak.prod( SFevaluator['ElecRecoSF_%i' % year](eta2, pt2) * SFevaluator['ElecLooseSF_%i' % year](np.abs(eta2), pt2) * SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta2), pt2) * SFevaluator['ElecTightSF_%i' % year](np.abs(eta2), pt2), axis=-1) else: print(flavor2, ' is not a valid flavor. Valid flavors: "m" or "e"') if flavor3 == None: return (np.multiply(SF1, SF2)) elif flavor3 == 'm': SF3 = ak.prod( SFevaluator['MuonLooseSF_%i' % year](np.abs(eta3), pt3) * SFevaluator['MuonTightSF_%i' % year](np.abs(eta3), pt3), axis=-1) elif flavor3 == 'e': SF3 = ak.prod( SFevaluator['ElecRecoSF_%i' % year](eta3, pt3) * SFevaluator['ElecLooseSF_%i' % year](np.abs(eta3), pt3) * SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta3), pt3) * SFevaluator['ElecTightSF_%i' % year](np.abs(eta3), pt3), axis=-1) else: print(flavor3, ' is not a valid flavor. Valid flavors: "m" , "e" or None') if flavor3 != None and flavor4 == None: return (np.multiply(SF3, np.multiply(SF1, SF2))) if flavor4 == 'm': SF4 = ak.prod( SFevaluator['MuonLooseSF_%i' % year](np.abs(eta4), pt4) * SFevaluator['MuonTightSF_%i' % year](np.abs(eta4), pt4), axis=-1) elif flavor4 == 'e': SF4 = ak.prod( SFevaluator['ElecRecoSF_%i' % year](eta4, pt4) * SFevaluator['ElecLooseSF_%i' % year](np.abs(eta4), pt4) * SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta4), pt4) * SFevaluator['ElecTightSF_%i' % year](np.abs(eta4), pt4), axis=-1) else: print(flavor4, ' is not a valid flavor. Valid flavors: "m" , "e" or None') if flavor4 != None: return (np.multiply(SF4, np.multiply(SF3, np.multiply(SF1, SF2)))) elif sys == 1: if flavor1 == 'm': SF1 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta1), pt1) + SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta1), pt1)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta1), pt1) + SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta1), pt1)), axis=-1) elif flavor1 == 'e': SF1 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta1, pt1) + SFevaluator['ElecRecoSF_%i_er' % year](eta1, pt1)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta1), pt1) + SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta1), pt1)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta1), pt1) + SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta1), pt1)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta1), pt1) + SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta1), pt1)), axis=-1) else: print(flavor1, ' is not a valid flavor. Valid flavors: "m" or "e"') if flavor2 == 'm': SF2 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta2), pt2) + SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta2), pt2)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta2), pt2) + SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta2), pt2)), axis=-1) elif flavor2 == 'e': SF2 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta2, pt2) + SFevaluator['ElecRecoSF_%i_er' % year](eta2, pt2)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta2), pt2) + SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta2), pt2)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta2), pt2) + SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta2), pt2)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta2), pt2) + SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta2), pt2)), axis=-1) else: print(flavor2, ' is not a valid flavor. Valid flavors: "m" or "e"') if flavor3 == None: return (np.multiply(SF1, SF2)) if flavor3 == 'm': SF3 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta3), pt3) + SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta3), pt3)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta3), pt3) + SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta3), pt3)), axis=-1) elif flavor3 == 'e': SF3 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta3, pt3) + SFevaluator['ElecRecoSF_%i_er' % year](eta3, pt3)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta3), pt3) + SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta3), pt3)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta3), pt3) + SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta3), pt3)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta3), pt3) + SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta3), pt3)), axis=-1) else: print(flavor3, ' is not a valid flavor. Valid flavors: "m" , "e" or None') if flavor3 != None and flavor4 == None: return (np.multiply(SF3, np.multiply(SF1, SF2))) if flavor4 == 'm': SF4 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta4), pt4) + SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta4), pt4)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta4), pt4) + SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta4), pt4)), axis=-1) elif flavor4 == 'e': SF4 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta4, pt4) + SFevaluator['ElecRecoSF_%i_er' % year](eta4, pt4)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta4), pt4) + SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta4), pt4)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta4), pt4) + SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta4), pt4)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta4), pt4) + SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta4), pt4)), axis=-1) else: print(flavor4, ' is not a valid flavor. Valid flavors: "m" , "e" or None') if flavor4 != None: return (np.multiply(SF4, np.multiply(SF3, np.multiply(SF1, SF2)))) elif sys == -1: if flavor1 == 'm': SF1 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta1), pt1) - SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta1), pt1)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta1), pt1) - SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta1), pt1)), axis=-1) elif flavor1 == 'e': SF1 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta1, pt1) - SFevaluator['ElecRecoSF_%i_er' % year](eta1, pt1)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta1), pt1) - SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta1), pt1)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta1), pt1) - SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta1), pt1)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta1), pt1) - SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta1), pt1)), axis=-1) else: print(flavor1, ' is not a valid flavor. Valid flavors: "m" or "e"') if flavor2 == 'm': SF2 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta2), pt2) - SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta2), pt2)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta2), pt2) - SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta2), pt2)), axis=-1) elif flavor2 == 'e': SF2 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta2, pt2) - SFevaluator['ElecRecoSF_%i_er' % year](eta2, pt2)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta2), pt2) - SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta2), pt2)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta2), pt2) - SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta2), pt2)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta2), pt2) - SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta2), pt2)), axis=-1) else: print(flavor2, ' is not a valid flavor. Valid flavors: "m" or "e"') if flavor3 == None: return (np.multiply(SF1, SF2)) if flavor3 == 'm': SF3 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta3), pt3) - SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta3), pt3)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta3), pt3) - SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta3), pt3)), axis=-1) elif flavor3 == 'e': SF3 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta3, pt3) - SFevaluator['ElecRecoSF_%i_er' % year](eta3, pt3)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta3), pt3) - SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta3), pt3)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta3), pt3) - SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta3), pt3)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta3), pt3) - SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta3), pt3)), axis=-1) else: print(flavor3, ' is not a valid flavor. Valid flavors: "m" , "e" or None') if flavor3 != None and flavor4 == None: return (np.multiply(SF3, np.multiply(SF1, SF2))) if flavor4 == 'm': SF4 = ak.prod( (SFevaluator['MuonLooseSF_%i' % year](np.abs(eta4), pt4) - SFevaluator['MuonLooseSF_%i_er' % year](np.abs(eta4), pt4)) * (SFevaluator['MuonTightSF_%i' % year](np.abs(eta4), pt4) - SFevaluator['MuonTightSF_%i_er' % year](np.abs(eta4), pt4)), axis=-1) elif flavor4 == 'e': SF4 = ak.prod( (SFevaluator['ElecRecoSF_%i' % year](eta4, pt4) - SFevaluator['ElecRecoSF_%i_er' % year](eta4, pt4)) * (SFevaluator['ElecLooseSF_%i' % year](np.abs(eta4), pt4) - SFevaluator['ElecLooseSF_%i_er' % year](np.abs(eta4), pt4)) * (SFevaluator['ElecLoosettHSF_%i' % year](np.abs(eta4), pt4) - SFevaluator['ElecLoosettHSF_%i_er' % year](np.abs(eta4), pt4)) * (SFevaluator['ElecTightSF_%i' % year](np.abs(eta4), pt4) - SFevaluator['ElecTightSF_%i_er' % year](np.abs(eta4), pt4)), axis=-1) else: print(flavor4, ' is not a valid flavor. Valid flavors: "m" , "e" or None') if flavor4 != None: return (np.multiply(SF4, np.multiply(SF3, np.multiply(SF1, SF2))))
def process(self, events): np.random.seed( 10 ) # sets seed so values from random distributions are reproducible (JER corrections) output = self.accumulator.identity() self.sample_name = events.metadata['dataset'] ## make event weights # data or MC distinction made internally mu_evt_weights = MCWeights.get_event_weights( events, year=args.year, corrections=self.corrections) el_evt_weights = MCWeights.get_event_weights( events, year=args.year, corrections=self.corrections) ## initialize selections and regions selection = PackedSelection() regions = { 'Muon': { 'Tight': { 'btagPass': { '3Jets': { 'objselection', 'jets_3', 'tight_MU', 'DeepCSV_pass', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'tight_MU', 'DeepCSV_pass', 'semilep' }, }, 'btagFail': { '3Jets': { 'objselection', 'jets_3', 'tight_MU', 'DeepCSV_fail', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'tight_MU', 'DeepCSV_fail', 'semilep' }, }, }, 'Loose': { 'btagPass': { '3Jets': { 'objselection', 'jets_3', 'loose_MU', 'DeepCSV_pass', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'loose_MU', 'DeepCSV_pass', 'semilep' }, }, 'btagFail': { '3Jets': { 'objselection', 'jets_3', 'loose_MU', 'DeepCSV_fail', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'loose_MU', 'DeepCSV_fail', 'semilep' }, }, }, }, 'Electron': { 'Tight': { 'btagPass': { '3Jets': { 'objselection', 'jets_3', 'tight_EL', 'DeepCSV_pass', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'tight_EL', 'DeepCSV_pass', 'semilep' }, }, 'btagFail': { '3Jets': { 'objselection', 'jets_3', 'tight_EL', 'DeepCSV_fail', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'tight_EL', 'DeepCSV_fail', 'semilep' }, }, }, 'Loose': { 'btagPass': { '3Jets': { 'objselection', 'jets_3', 'loose_EL', 'DeepCSV_pass', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'loose_EL', 'DeepCSV_pass', 'semilep' }, }, 'btagFail': { '3Jets': { 'objselection', 'jets_3', 'loose_EL', 'DeepCSV_fail', 'semilep' }, '4PJets': { 'objselection', 'jets_4p', 'loose_EL', 'DeepCSV_fail', 'semilep' }, }, }, }, } ## object selection objsel_evts = objsel.select(events, year=args.year, corrections=self.corrections, cutflow=output['cutflow']) output['cutflow'][ 'nEvts passing jet and lepton obj selection'] += ak.sum( objsel_evts) selection.add('objselection', objsel_evts) selection.add('jets_3', ak.num(events['Jet']) == 3) selection.add('jets_4p', ak.num(events['Jet']) > 3) selection.add( 'DeepCSV_pass', ak.sum(events['Jet']['DeepCSV' + wps_to_use[0]], axis=1) >= 2) # sort jets by btag value events['Jet'] = events['Jet'][ak.argsort( events['Jet']['btagDeepB'], ascending=False)] if btagger == 'DeepCSV' else events['Jet'][ ak.argsort(events['Jet']['btagDeepFlavB'], ascending=False)] # btag fail sideband deepcsv_sorted = events['Jet'][ak.argsort( events['Jet']['btagDeepB'], ascending=False)]['btagDeepB'] valid_counts_inds = ak.where(ak.num(events['Jet']) > 1)[0] deepcsv_fail = np.zeros(len(events)).astype(bool) deepcsv_fail[valid_counts_inds] = ( deepcsv_sorted[valid_counts_inds][:, 0] < btag_values[args.year]['btagDeepB']['DeepCSV' + wps_to_use[0]]) & ( deepcsv_sorted[valid_counts_inds][:, 1] < btag_values[args.year]['btagDeepB']['DeepCSV' + wps_to_use[0]]) selection.add( 'DeepCSV_fail', deepcsv_fail ) # highest and second highest DeepCSV values don't pass tight and loose WPs ## add different selections ## muons selection.add('tight_MU', ak.sum(events['Muon']['TIGHTMU'], axis=1) == 1) # one muon passing TIGHT criteria selection.add('loose_MU', ak.sum(events['Muon']['LOOSEMU'], axis=1) == 1) # one muon passing LOOSE criteria ## electrons selection.add( 'tight_EL', ak.sum(events['Electron']['TIGHTEL'], axis=1) == 1) # one electron passing TIGHT criteria selection.add( 'loose_EL', ak.sum(events['Electron']['LOOSEEL'], axis=1) == 1) # one electron passing LOOSE criteria ### apply lepton SFs to MC (only applicable to tight leptons) if 'LeptonSF' in corrections.keys(): tight_mu_cut = selection.require( objselection=True, tight_MU=True ) # find events passing muon object selection with one tight muon tight_muons = events['Muon'][tight_mu_cut][( events['Muon'][tight_mu_cut]['TIGHTMU'] == True)] muSFs_dict = MCWeights.get_lepton_sf( year=args.year, lepton='Muons', corrections=self.corrections['LeptonSF'], pt=ak.flatten(tight_muons['pt']), eta=ak.flatten(tight_muons['eta'])) mu_reco_cen = np.ones(len(events)) mu_reco_err = np.zeros(len(events)) mu_trig_cen = np.ones(len(events)) mu_trig_err = np.zeros(len(events)) mu_reco_cen[tight_mu_cut] = muSFs_dict['RECO_CEN'] mu_reco_err[tight_mu_cut] = muSFs_dict['RECO_ERR'] mu_trig_cen[tight_mu_cut] = muSFs_dict['TRIG_CEN'] mu_trig_err[tight_mu_cut] = muSFs_dict['TRIG_ERR'] mu_evt_weights.add('RECO', mu_reco_cen, mu_reco_err, mu_reco_err, shift=True) mu_evt_weights.add('TRIG', mu_trig_cen, mu_trig_err, mu_trig_err, shift=True) tight_el_cut = selection.require( objselection=True, tight_EL=True ) # find events passing electron object selection with one tight electron tight_electrons = events['Electron'][tight_el_cut][( events['Electron'][tight_el_cut]['TIGHTEL'] == True)] elSFs_dict = MCWeights.get_lepton_sf( year=args.year, lepton='Electrons', corrections=self.corrections['LeptonSF'], pt=ak.flatten(tight_electrons['pt']), eta=ak.flatten(tight_electrons['etaSC'])) el_reco_cen = np.ones(len(events)) el_reco_err = np.zeros(len(events)) el_trig_cen = np.ones(len(events)) el_trig_err = np.zeros(len(events)) el_reco_cen[tight_el_cut] = elSFs_dict['RECO_CEN'] el_reco_err[tight_el_cut] = elSFs_dict['RECO_ERR'] el_trig_cen[tight_el_cut] = elSFs_dict['TRIG_CEN'] el_trig_err[tight_el_cut] = elSFs_dict['TRIG_ERR'] el_evt_weights.add('RECO', el_reco_cen, el_reco_err, el_reco_err, shift=True) el_evt_weights.add('TRIG', el_trig_cen, el_trig_err, el_trig_err, shift=True) if corrections['BTagSF'] == True: deepcsv_cen = np.ones(len(events)) threeJets_cut = selection.require(objselection=True, jets_3=True) deepcsv_3j_wts = self.corrections['BTag_Constructors']['DeepCSV'][ '3Jets'].get_scale_factor(jets=events['Jet'][threeJets_cut], passing_cut='DeepCSV' + wps_to_use[0]) deepcsv_cen[threeJets_cut] = ak.prod(deepcsv_3j_wts['central'], axis=1) fourplusJets_cut = selection.require(objselection=True, jets_4p=True) deepcsv_4pj_wts = self.corrections['BTag_Constructors']['DeepCSV'][ '4PJets'].get_scale_factor( jets=events['Jet'][fourplusJets_cut], passing_cut='DeepCSV' + wps_to_use[0]) deepcsv_cen[fourplusJets_cut] = ak.prod(deepcsv_4pj_wts['central'], axis=1) # make dict of btag weights btag_weights = { 'DeepCSV_CEN': deepcsv_cen, } # find gen level particles for ttbar system genpsel.select(events, mode='NORMAL') selection.add('semilep', ak.num(events['SL']) > 0) ## fill hists for each region for lepton in regions.keys(): evt_weights = mu_evt_weights if lepton == 'Muon' else el_evt_weights for leptype in regions[lepton].keys(): for btagregion in regions[lepton][leptype].keys(): for jmult in regions[lepton][leptype][btagregion].keys(): cut = selection.all( *regions[lepton][leptype][btagregion][jmult]) if args.debug: print(lepton, leptype, btagregion, jmult) if cut.sum() > 0: ltype = 'MU' if lepton == 'Muon' else 'EL' if 'loose_or_tight_%s' % ltype in regions[lepton][ leptype][btagregion][jmult]: lep_mask = ( (events[lepton][cut]['TIGHT%s' % ltype] == True) | (events[lepton][cut]['LOOSE%s' % ltype] == True)) elif 'tight_%s' % ltype in regions[lepton][ leptype][btagregion][jmult]: lep_mask = (events[lepton][cut]['TIGHT%s' % ltype] == True) elif 'loose_%s' % ltype in regions[lepton][ leptype][btagregion][jmult]: lep_mask = (events[lepton][cut]['LOOSE%s' % ltype] == True) else: raise ValueError( "Not sure what lepton type to choose for event" ) # find matched perm mp = ttmatcher.best_match( gen_hyp=events['SL'][cut], jets=events['Jet'][cut], leptons=events[lepton][cut][lep_mask], met=events['MET'][cut]) ## create MT regions MT = make_vars.MT(events[lepton][cut][lep_mask], events['MET'][cut]) MTHigh = ak.flatten(MT >= MTcut) wts = (evt_weights.weight() * btag_weights['%s_CEN' % btaggers[0]])[cut][MTHigh] mp_status = np.zeros( cut.sum(), dtype=int ) # 0 == '' (no gen matching), 1 == 'right', 2 == 'matchable', 3 == 'unmatchable', 4 == 'sl_tau', 5 == 'noslep' if jmult == '3Jets': valid_evts = (ak.num(mp['TTbar'].pt) > 0) & ( ak.flatten(mp['unique_matches'] >= 3)) # merged events merged_evts = valid_evts & ak.flatten( mp['Merged_Event']) correct_merged = merged_evts & ak.flatten( mp['Merged_BHadWJa'] | mp['Merged_BHadWJb'] | mp['Merged_WJets']) wrong_merged = merged_evts & ~(ak.flatten( mp['Merged_BHadWJa'] | mp['Merged_BHadWJb'] | mp['Merged_WJets'])) # lost events lost_evts = valid_evts & ak.flatten( mp['Lost_Event']) correct_lost = lost_evts & ak.flatten( (mp['Lost_WJa'] | mp['Lost_WJb'])) wrong_lost = lost_evts & ~(ak.flatten( mp['Lost_WJa'] | mp['Lost_WJb'])) # event categorization # unmatchable events unmatchable_evts = (~valid_evts | wrong_merged | wrong_lost) # matchable events matchable_evts = ( correct_lost | correct_merged ) # matched perm is correct event type irregardless of object matching else: valid_evts = (ak.num(mp['TTbar'].pt) > 0) & ( ak.flatten(mp['unique_matches'] == 4)) # unmatchable events unmatchable_evts = ~valid_evts # matchable events matchable_evts = valid_evts mp_status[matchable_evts] = 2 mp_status[unmatchable_evts] = 3 output = self.make_categories( acc=output, jetmult=jmult, leptype=lepton, lepcat=leptype, btagregion=btagregion, permarray=mp_status[MTHigh], genttbar=events['SL'][cut][MTHigh], matched_perm=mp[MTHigh], evt_weights=wts) return output
def test_highlevel(): array = ak.Array( [[[2, 3, 5], [], [7, 11], [13]], [], [[17, 19], [23]]], check_valid=True ) assert ak.count(array) == 9 assert ak.to_list(ak.count(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]] assert ak.to_list(ak.count(array, axis=2)) == [[3, 0, 2, 1], [], [2, 1]] assert ak.to_list(ak.count(array, axis=-1, keepdims=True)) == [ [[3], [0], [2], [1]], [], [[2], [1]], ] assert ak.to_list(ak.count(array, axis=-2)) == [[3, 2, 1], [], [2, 1]] assert ak.to_list(ak.count(array, axis=1)) == [[3, 2, 1], [], [2, 1]] assert ak.to_list(ak.count(array, axis=-2, keepdims=True)) == [ [[3, 2, 1]], [[]], [[2, 1]], ] assert ak.count_nonzero(array) == 9 assert ak.to_list(ak.count_nonzero(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]] assert ak.to_list(ak.count_nonzero(array, axis=-2)) == [[3, 2, 1], [], [2, 1]] assert ak.sum(array) == 2 + 3 + 5 + 7 + 11 + 13 + 17 + 19 + 23 assert ak.to_list(ak.sum(array, axis=-1)) == [ [2 + 3 + 5, 0, 7 + 11, 13], [], [17 + 19, 23], ] assert ak.to_list(ak.sum(array, axis=-2)) == [ [2 + 7 + 13, 3 + 11, 5], [], [17 + 23, 19], ] assert ak.prod(array) == 2 * 3 * 5 * 7 * 11 * 13 * 17 * 19 * 23 assert ak.to_list(ak.prod(array, axis=-1)) == [ [2 * 3 * 5, 1, 7 * 11, 13], [], [17 * 19, 23], ] assert ak.to_list(ak.prod(array, axis=-2)) == [ [2 * 7 * 13, 3 * 11, 5], [], [17 * 23, 19], ] assert ak.min(array) == 2 assert ak.to_list(ak.min(array, axis=-1)) == [[2, None, 7, 13], [], [17, 23]] assert ak.to_list(ak.min(array, axis=-2)) == [[2, 3, 5], [], [17, 19]] assert ak.max(array) == 23 assert ak.to_list(ak.max(array, axis=-1)) == [[5, None, 11, 13], [], [19, 23]] assert ak.to_list(ak.max(array, axis=-2)) == [[13, 11, 5], [], [23, 19]] array = ak.Array( [ [[True, False, True], [], [False, False], [True]], [], [[False, True], [True]], ], check_valid=True, ) assert ak.any(array) == True assert ak.to_list(ak.any(array, axis=-1)) == [ [True, False, False, True], [], [True, True], ] assert ak.to_list(ak.any(array, axis=-2)) == [[True, False, True], [], [True, True]] assert ak.all(array) == False assert ak.to_list(ak.all(array, axis=-1)) == [ [False, True, False, True], [], [False, True], ] assert ak.to_list(ak.all(array, axis=-2)) == [ [False, False, True], [], [False, True], ]
def test_0166_ByteMaskedArray(): content = ak.from_iter( [[2, 3, 5], [999], [], [7, 11], [], [13], [123, 999], [17, 19]], highlevel=False) mask = ak.layout.Index8(np.array([0, 1, 0, 0, 1, 0, 1, 0], dtype=np.int8)) array = ak.Array(ak.layout.ByteMaskedArray(mask, content, valid_when=False)) assert ak.to_list(array) == [ [2, 3, 5], None, [], [7, 11], None, [13], None, [17, 19], ] assert ak.to_list(ak.prod( array, axis=-1)) == [30, None, 1, 77, None, 13, None, 323] content = ak.from_iter( [ [[2, 3], [5]], [[999]], [], [[7], [11]], [], [[13]], [[123], [999]], [[17, 19]], ], highlevel=False, ) mask = ak.layout.Index8(np.array([0, 1, 0, 0, 1, 0, 1, 0], dtype=np.int8)) array = ak.Array(ak.layout.ByteMaskedArray(mask, content, valid_when=False)) assert ak.to_list(array) == [ [[2, 3], [5]], None, [], [[7], [11]], None, [[13]], None, [[17, 19]], ] assert ak.to_list(ak.prod(array, axis=-1)) == [ [6, 5], None, [], [7, 11], None, [13], None, [323], ] content = ak.from_iter( [[2, 3], [999], [5], [7], [11], [13], [], [17], [19]], highlevel=False) mask = ak.layout.Index8( np.array([0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8)) bytemasked = ak.layout.ByteMaskedArray(mask, content, valid_when=False) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 9], dtype=np.int64)) array = ak.Array(ak.layout.ListOffsetArray64(offsets, bytemasked)) array = ak.Array([[[2, 3], None, [5]], [], [[7], [11]], [[13]], [None, [17], [19]]]) assert ak.to_list(ak.prod(array, axis=-1)) == [ [6, None, 5], [], [7, 11], [13], [None, 17, 19], ] content = ak.from_iter([6, None, 5, 7, 11, 13, None, 17, 19], highlevel=False) mask = ak.layout.Index8( np.array([0, 1, 0, 0, 0, 0, 1, 0, 0], dtype=np.int8)) bytemasked = ak.layout.ByteMaskedArray(mask, content, valid_when=False) offsets = ak.layout.Index64(np.array([0, 3, 3, 5, 6, 9], dtype=np.int64)) array = ak.Array(ak.layout.ListOffsetArray64(offsets, bytemasked)) assert ak.to_list(array) == [[6, None, 5], [], [7, 11], [13], [None, 17, 19]] assert ak.to_list(ak.prod(array, axis=-1)) == [30, 1, 77, 13, 323]
def process(self, events): np.random.seed( 10 ) # sets seed so values from random distributions are reproducible (JER corrections) output = self.accumulator.identity() self.sample_name = events.metadata['dataset'] ## make event weights # data or MC distinction made internally mu_evt_weights = MCWeights.get_event_weights( events, year=args.year, corrections=self.corrections) el_evt_weights = MCWeights.get_event_weights( events, year=args.year, corrections=self.corrections) ## initialize selections and regions selection = PackedSelection() regions = { 'Muon': { 'lep_and_filter_pass', 'passing_jets', 'jets_3', 'tight_MU', 'btag_pass', 'semilep' }, 'Electron': { 'lep_and_filter_pass', 'passing_jets', 'jets_3', 'tight_EL', 'btag_pass', 'semilep' }, } # get all passing leptons lep_and_filter_pass = objsel.select_leptons(events, year=args.year) selection.add('lep_and_filter_pass', lep_and_filter_pass ) # add passing leptons requirement to all systematics ## add different selections ## muons tight_mu_sel = ak.sum(events['Muon']['TIGHTMU'], axis=1) == 1 selection.add('tight_MU', tight_mu_sel) # one muon passing TIGHT criteria ## electrons tight_el_sel = ak.sum(events['Electron']['TIGHTEL'], axis=1) == 1 selection.add('tight_EL', tight_el_sel) # one electron passing TIGHT criteria ### apply lepton SFs to MC (only applicable to tight leptons) if 'LeptonSF' in self.corrections.keys(): tight_muons = events['Muon'][tight_mu_sel][( events['Muon'][tight_mu_sel]['TIGHTMU'] == True)] muSFs_dict = MCWeights.get_lepton_sf( year=args.year, lepton='Muons', corrections=self.corrections['LeptonSF'], pt=ak.flatten(tight_muons['pt']), eta=ak.flatten(tight_muons['eta'])) mu_reco_cen = np.ones(len(events)) mu_reco_err = np.zeros(len(events)) mu_trig_cen = np.ones(len(events)) mu_trig_err = np.zeros(len(events)) mu_reco_cen[tight_mu_sel] = muSFs_dict['RECO_CEN'] mu_reco_err[tight_mu_sel] = muSFs_dict['RECO_ERR'] mu_trig_cen[tight_mu_sel] = muSFs_dict['TRIG_CEN'] mu_trig_err[tight_mu_sel] = muSFs_dict['TRIG_ERR'] mu_evt_weights.add('Lep_RECO', mu_reco_cen, mu_reco_err, mu_reco_err, shift=True) mu_evt_weights.add('Lep_TRIG', mu_trig_cen, mu_trig_err, mu_trig_err, shift=True) tight_electrons = events['Electron'][tight_el_sel][( events['Electron'][tight_el_sel]['TIGHTEL'] == True)] elSFs_dict = MCWeights.get_lepton_sf( year=args.year, lepton='Electrons', corrections=self.corrections['LeptonSF'], pt=ak.flatten(tight_electrons['pt']), eta=ak.flatten(tight_electrons['etaSC'])) el_reco_cen = np.ones(len(events)) el_reco_err = np.zeros(len(events)) el_trig_cen = np.ones(len(events)) el_trig_err = np.zeros(len(events)) el_reco_cen[tight_el_sel] = elSFs_dict['RECO_CEN'] el_reco_err[tight_el_sel] = elSFs_dict['RECO_ERR'] el_trig_cen[tight_el_sel] = elSFs_dict['TRIG_CEN'] el_trig_err[tight_el_sel] = elSFs_dict['TRIG_ERR'] el_evt_weights.add('Lep_RECO', el_reco_cen, el_reco_err, el_reco_err, shift=True) el_evt_weights.add('Lep_TRIG', el_trig_cen, el_trig_err, el_trig_err, shift=True) ## build corrected jets and MET events['Jet'], events['MET'] = IDJet.process_jets( events, args.year, self.corrections['JetCor']) # jet selection passing_jets = objsel.jets_selection(events, year=args.year) selection.add('passing_jets', passing_jets) selection.add('jets_3', ak.num(events['SelectedJets']) == 3) selection.add('btag_pass', ak.sum(events['SelectedJets'][btag_wps[0]], axis=1) >= 2) events['SelectedJets'] = events['SelectedJets'][ak.argsort( events['SelectedJets']['btagDeepB'], ascending=False )] if btaggers[0] == 'DeepCSV' else events['SelectedJets'][ak.argsort( events['SelectedJets']['btagDeepFlavB'], ascending=False)] if self.corrections['BTagSF'] == True: #set_trace() deepcsv_cen = np.ones(len(events)) threeJets_cut = selection.require(lep_and_filter_pass=True, passing_jets=True, jets_3=True) deepcsv_3j_wts = self.corrections['BTag_Constructors']['DeepCSV'][ '3Jets'].get_scale_factor( jets=events['SelectedJets'][threeJets_cut], passing_cut='DeepCSV' + wps_to_use[0]) deepcsv_cen[threeJets_cut] = ak.prod(deepcsv_3j_wts['central'], axis=1) # make dict of btag weights btag_weights = { 'DeepCSV_CEN': deepcsv_cen, } # find gen level particles for ttbar system genpsel.select(events, mode='NORMAL') selection.add('semilep', ak.num(events['SL']) > 0) if 'NNLO_Rewt' in self.corrections.keys(): nnlo_wts = MCWeights.get_nnlo_weights( self.corrections['NNLO_Rewt'], events) mu_evt_weights.add( '%s_reweighting' % self.corrections['NNLO_Rewt']['Var'], nnlo_wts) el_evt_weights.add( '%s_reweighting' % self.corrections['NNLO_Rewt']['Var'], nnlo_wts) ## fill hists for each region for lepton in regions.keys(): evt_weights = mu_evt_weights if lepton == 'Muon' else el_evt_weights cut = selection.all(*regions[lepton]) #set_trace() if cut.sum() > 0: leptype = 'MU' if lepton == 'Muon' else 'EL' if 'loose_or_tight_%s' % leptype in regions[lepton]: leptons = events[lepton][cut][( (events[lepton][cut]['TIGHT%s' % leptype] == True) | (events[lepton][cut]['LOOSE%s' % leptype] == True))] elif 'tight_%s' % leptype in regions[lepton]: leptons = events[lepton][cut][( events[lepton][cut]['TIGHT%s' % leptype] == True)] elif 'loose_%s' % leptype in regions[lepton]: leptons = events[lepton][cut][( events[lepton][cut]['LOOSE%s' % leptype] == True)] else: raise ValueError( "Not sure what lepton type to choose for event") # get jets and MET jets, met = events['SelectedJets'][cut], events['SelectedMET'][ cut] # find matched permutations mp = ttmatcher.best_match(gen_hyp=events['SL'][cut], jets=jets, leptons=leptons, met=met) # find best permutations best_perms = ttpermutator.find_best_permutations( jets=jets, leptons=leptons, MET=met, btagWP=btag_wps[0]) valid_perms = ak.num(best_perms['TTbar'].pt) > 0 # compare matched per to best perm bp_status = np.zeros( cut.size, dtype=int ) # 0 == '' (no gen matching), 1 == 'right', 2 == 'matchable', 3 == 'unmatchable', 4 == 'sl_tau', 5 == 'noslep' perm_cat_array = compare_matched_best_perms(mp, best_perms, njets='3Jets') bp_status[cut] = perm_cat_array if ak.any(ak.num(events['SL']['Lepton'].pdgId) != 1): raise ValueError( "Number of leptons is incorrect for classifying tau+jets events" ) sl_tau_evts = ak.where( np.abs(events['SL']['Lepton'].pdgId) == 15)[0] bp_status[sl_tau_evts] = 4 ## create MT regions MT = make_vars.MT(leptons, met) MTHigh = ak.flatten(MT[valid_perms] >= MTcut) wts = (evt_weights.weight() * btag_weights['%s_CEN' % btaggers[0]])[cut][valid_perms][MTHigh] output = self.make_3j_categories( acc=output, leptype=lepton, permarray=bp_status[cut][valid_perms][MTHigh], genttbar=events['SL'][cut][valid_perms][MTHigh], bp=best_perms[valid_perms][MTHigh], evt_wts=wts) return output