Skip to content

Commit 7a6feff

Browse files
mafrahmmafrahm
committed
add remastered sl selection
1 parent 8f3a12d commit 7a6feff

File tree

1 file changed

+279
-0
lines changed

1 file changed

+279
-0
lines changed

‎hbw/selection/sl_remastered.py

+279
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
# coding: utf-8
2+
3+
"""
4+
Selection modules for HH -> bbWW(qqlnu).
5+
"""
6+
7+
from collections import defaultdict
8+
from typing import Tuple
9+
10+
from columnflow.util import maybe_import
11+
from columnflow.selection import Selector, SelectionResult, selector
12+
13+
from hbw.selection.common import (
14+
jet_selection, lepton_definition,
15+
masked_sorted_indices, sl_boosted_jet_selection, vbf_jet_selection,
16+
pre_selection, post_selection,
17+
)
18+
from hbw.production.weights import event_weights_to_normalize
19+
from hbw.selection.cutflow_features import cutflow_features
20+
21+
np = maybe_import("numpy")
22+
ak = maybe_import("awkward")
23+
24+
25+
@selector(
26+
uses={lepton_definition, "Electron.charge", "Muon.charge"},
27+
produces={lepton_definition},
28+
e_pt=None, mu_pt=None, trigger=None,
29+
)
30+
def sl_lepton_selection(
31+
self: Selector,
32+
events: ak.Array,
33+
stats: defaultdict,
34+
**kwargs,
35+
) -> Tuple[ak.Array, SelectionResult]:
36+
37+
events, lepton_results = self[lepton_definition](events, stats, **kwargs)
38+
39+
# tau veto
40+
lepton_results.steps["VetoTau"] = events.cutflow.n_veto_tau == 0
41+
42+
# number of electrons
43+
lepton_results.steps["nRecoElectron1"] = ak.num(events.Electron) >= 1
44+
lepton_results.steps["nLooseElectron1"] = events.cutflow.n_loose_electron >= 1
45+
lepton_results.steps["nFakeableElectron1"] = events.cutflow.n_fakeable_electron >= 1
46+
lepton_results.steps["nTightElectron1"] = events.cutflow.n_tight_electron >= 1
47+
48+
# number of muons
49+
lepton_results.steps["nRecoMuon1"] = ak.num(events.Muon) >= 1
50+
lepton_results.steps["nLooseMuon1"] = events.cutflow.n_loose_muon >= 1
51+
lepton_results.steps["nFakeableMuon1"] = events.cutflow.n_fakeable_muon >= 1
52+
lepton_results.steps["nTightMuon1"] = events.cutflow.n_tight_muon >= 1
53+
54+
# select events
55+
mu_mask_fakeable = lepton_results.x.mu_mask_fakeable
56+
e_mask_fakeable = lepton_results.x.e_mask_fakeable
57+
58+
# NOTE: leading lepton pt could be reduced to trigger threshold + 1
59+
leading_mu_mask = (mu_mask_fakeable) & (events.Muon.cone_pt > self.config_inst.x.mu_pt)
60+
leading_e_mask = (e_mask_fakeable) & (events.Electron.cone_pt > self.config_inst.x.e_pt)
61+
62+
# NOTE: we might need pt > 15 for lepton SFs. Needs to be checked in Run 3.
63+
subleading_mu_mask = (mu_mask_fakeable) & (events.Muon.cone_pt > 15)
64+
subleading_e_mask = (e_mask_fakeable) & (events.Electron.cone_pt > 15)
65+
66+
# For further analysis after Reduction, we consider all tight leptons with pt > 15 GeV
67+
lepton_results.objects["Electron"]["Electron"] = masked_sorted_indices(subleading_e_mask, events.Electron.pt)
68+
lepton_results.objects["Muon"]["Muon"] = masked_sorted_indices(subleading_mu_mask, events.Muon.pt)
69+
electron = events.Electron[subleading_e_mask]
70+
muon = events.Muon[subleading_mu_mask]
71+
72+
# Create a temporary lepton collection
73+
lepton = ak.concatenate(
74+
[
75+
electron * 1,
76+
muon * 1,
77+
],
78+
axis=1,
79+
)
80+
lepton = lepton_results.aux["lepton"] = lepton[ak.argsort(lepton.pt, axis=-1, ascending=False)]
81+
82+
lepton_results.steps["DileptonVeto"] = ak.num(lepton, axis=1) <= 1
83+
84+
lepton_results.steps["Lep_e"] = e_mask = (
85+
(ak.sum(leading_e_mask, axis=1) == 1) &
86+
(ak.sum(subleading_mu_mask, axis=1) == 0)
87+
)
88+
lepton_results.steps["Lep_mu"] = mu_mask = (
89+
(ak.sum(leading_mu_mask, axis=1) == 1) &
90+
(ak.sum(subleading_e_mask, axis=1) == 0)
91+
)
92+
93+
lepton_results.steps["Lepton"] = (e_mask | mu_mask)
94+
95+
lepton_results.steps["Fake"] = (
96+
lepton_results.steps.Lepton &
97+
(ak.sum(electron.is_tight, axis=1) + ak.sum(muon.is_tight, axis=1) == 0)
98+
)
99+
lepton_results.steps["SR"] = (
100+
lepton_results.steps.Lepton &
101+
(ak.sum(electron.is_tight, axis=1) + ak.sum(muon.is_tight, axis=1) == 1)
102+
)
103+
104+
for channel, trigger_columns in self.config_inst.x.trigger.items():
105+
# apply the "or" of all triggers of this channel
106+
trigger_mask = ak.any([events.HLT[trigger_column] for trigger_column in trigger_columns], axis=0)
107+
lepton_results.steps[f"Trigger_{channel}"] = trigger_mask
108+
109+
# ensure that Lepton channel is in agreement with trigger
110+
lepton_results.steps[f"TriggerAndLep_{channel}"] = (
111+
lepton_results.steps[f"Trigger_{channel}"] & lepton_results.steps[f"Lep_{channel}"]
112+
)
113+
114+
# combine results of each individual channel
115+
lepton_results.steps["Trigger"] = ak.any([
116+
lepton_results.steps[f"Trigger_{channel}"]
117+
for channel in self.config_inst.x.trigger.keys()
118+
], axis=0)
119+
120+
lepton_results.steps["TriggerAndLep"] = ak.any([
121+
lepton_results.steps[f"TriggerAndLep_{channel}"]
122+
for channel in self.config_inst.x.trigger.keys()
123+
], axis=0)
124+
125+
return events, lepton_results
126+
127+
128+
@sl_lepton_selection.init
129+
# @call_once_on_instance()
130+
def sl_lepton_selection_init(self: Selector) -> None:
131+
# update selector steps labels
132+
self.config_inst.x.selector_step_labels = self.config_inst.x("selector_step_labels", {})
133+
self.config_inst.x.selector_step_labels.update({
134+
"DileptonVeto": r"$N_{lepton} \leq 1$",
135+
"Lepton": r"$N_{lepton} = 1$",
136+
"Lep_e": r"$N_{e} = 1$ and $N_{\mu} = 0$",
137+
"Lep_mu": r"$N_{\mu} = 1$ and $N_{e} = 0$",
138+
"Fake": r"$N_{lepton}^{tight} = 0$",
139+
"SR": r"$N_{lepton}^{tight} = 1$",
140+
"TriggerAndLep": "Trigger matches Lepton Channel",
141+
"VetoTau": r"$N_{\tau}^{veto} = 0$",
142+
})
143+
144+
year = self.config_inst.campaign.x.year
145+
146+
# setup lepton pt and trigger requirements in the config
147+
# when the lepton selector does not define the values, resort to defaults
148+
# NOTE: this is not doing what I was intending: this allows me to share the selector info
149+
# with other tasks, but I want other selectors to be able to change these attributes...
150+
if year == 2016:
151+
self.config_inst.x.mu_pt = self.mu_pt or 25
152+
self.config_inst.x.e_pt = self.e_pt or 27
153+
self.config_inst.x.trigger = self.trigger or {
154+
"e": ["Ele27_WPTight_Gsf"],
155+
"mu": ["IsoMu24"],
156+
}
157+
elif year == 2017:
158+
self.config_inst.x.mu_pt = self.mu_pt or 28
159+
self.config_inst.x.e_pt = self.e_pt or 36
160+
self.config_inst.x.trigger = self.trigger or {
161+
"e": ["Ele35_WPTight_Gsf"],
162+
"mu": ["IsoMu27"],
163+
}
164+
elif year == 2018:
165+
self.config_inst.x.mu_pt = self.mu_pt or 25
166+
self.config_inst.x.e_pt = self.e_pt or 33
167+
self.config_inst.x.trigger = self.trigger or {
168+
"e": ["Ele32_WPTight_Gsf"],
169+
"mu": ["IsoMu24"],
170+
}
171+
elif year == 2022:
172+
self.config_inst.x.mu_pt = self.mu_pt or 25
173+
self.config_inst.x.e_pt = self.e_pt or 31
174+
self.config_inst.x.trigger = self.trigger or {
175+
"e": ["Ele30_WPTight_Gsf"],
176+
"mu": ["IsoMu24"],
177+
}
178+
else:
179+
raise Exception(f"Single lepton trigger not implemented for year {year}")
180+
181+
# add all required trigger to the uses
182+
for trigger_columns in self.config_inst.x.trigger.values():
183+
for column in trigger_columns:
184+
self.uses.add(f"HLT.{column}")
185+
186+
187+
@selector(
188+
uses={
189+
pre_selection, post_selection,
190+
vbf_jet_selection, sl_boosted_jet_selection,
191+
jet_selection, sl_lepton_selection,
192+
},
193+
produces={
194+
pre_selection, post_selection,
195+
vbf_jet_selection, sl_boosted_jet_selection,
196+
jet_selection, sl_lepton_selection,
197+
},
198+
exposed=True,
199+
)
200+
def sl(
201+
self: Selector,
202+
events: ak.Array,
203+
stats: defaultdict,
204+
**kwargs,
205+
) -> Tuple[ak.Array, SelectionResult]:
206+
# prepare events
207+
events, results = self[pre_selection](events, stats, **kwargs)
208+
209+
# lepton selection
210+
events, lepton_results = self[sl_lepton_selection](events, stats, **kwargs)
211+
results += lepton_results
212+
213+
# jet selection
214+
events, jet_results = self[jet_selection](events, lepton_results, stats, **kwargs)
215+
results += jet_results
216+
217+
# boosted selection
218+
events, boosted_results = self[sl_boosted_jet_selection](events, lepton_results, jet_results, stats, **kwargs)
219+
results += boosted_results
220+
221+
# vbf-jet selection
222+
events, vbf_jet_results = self[vbf_jet_selection](events, results, stats, **kwargs)
223+
results += vbf_jet_results
224+
225+
results.steps["Resolved"] = (results.steps.nJet3 & results.steps.nBjet1)
226+
227+
results.steps["ResolvedOrBoosted"] = (
228+
(results.steps.nJet3 & results.steps.nBjet1) | results.steps.HbbJet
229+
)
230+
231+
# combined event selection after all steps except b-jet selection
232+
results.steps["all_but_bjet"] = (
233+
results.steps.cleanup &
234+
(results.steps.nJet3 | results.steps.HbbJet_no_bjet) &
235+
results.steps.ll_lowmass_veto &
236+
results.steps.ll_zmass_veto &
237+
results.steps.DileptonVeto &
238+
results.steps.Lepton &
239+
results.steps.VetoTau &
240+
results.steps.Trigger &
241+
results.steps.TriggerAndLep
242+
)
243+
244+
# combined event selection after all steps
245+
# NOTE: we only apply the b-tagging step when no AK8 Jet is present; if some event with AK8 jet
246+
# gets categorized into the resolved category, we might need to cut again on the number of b-jets
247+
results.event = (
248+
results.steps.all_but_bjet &
249+
((results.steps.nJet3 & results.steps.nBjet1) | results.steps.HbbJet)
250+
)
251+
results.steps["all"] = results.event
252+
253+
# build categories
254+
events, results = self[post_selection](events, results, stats, **kwargs)
255+
256+
return events, results
257+
258+
259+
@sl.init
260+
def sl_init(self: Selector) -> None:
261+
# define mapping from selector step to labels used in cutflow plots
262+
self.config_inst.x.selector_step_labels = self.config_inst.x("selector_step_labels", {})
263+
self.config_inst.x.selector_step_labels.update({
264+
"Resolved": r"$N_{jets}^{AK4} \geq 3$ and $N_{jets}^{BTag} \geq 1$",
265+
"ResolvedOrBoosted": (
266+
r"($N_{jets}^{AK4} \geq 3$ and $N_{jets}^{BTag} \geq 1$) "
267+
r"or $N_{H \rightarrow bb}^{AK8} \geq 1$"
268+
),
269+
})
270+
271+
if self.config_inst.x("do_cutflow_features", False):
272+
self.uses.add(cutflow_features)
273+
self.produces.add(cutflow_features)
274+
275+
if not getattr(self, "dataset_inst", None) or self.dataset_inst.is_data:
276+
return
277+
278+
self.uses.add(event_weights_to_normalize)
279+
self.produces.add(event_weights_to_normalize)

0 commit comments

Comments
 (0)