def test_groupby():
array = ak.Array(
[
{"x": 1, "y": 1.1},
{"x": 2, "y": 2.2},
{"x": 1, "y": 1.1},
{"x": 3, "y": 3.3},
{"x": 1, "y": 1.1},
{"x": 2, "y": 2.2},
]
)
sorted = array[ak.argsort(array.x)]
assert sorted.x.tolist() == [1, 1, 1, 2, 2, 3]
assert ak.run_lengths(sorted.x).tolist() == [3, 2, 1]
assert ak.unflatten(sorted, ak.run_lengths(sorted.x)).tolist() == [
[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}],
[{"x": 2, "y": 2.2}, {"x": 2, "y": 2.2}],
[{"x": 3, "y": 3.3}],
]
array = ak.Array(
[
[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 1, "y": 1.1}],
[{"x": 3, "y": 3.3}, {"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
]
)
sorted = array[ak.argsort(array.x)]
assert sorted.x.tolist() == [[1, 1, 2], [1, 2, 3]]
assert ak.run_lengths(sorted.x).tolist() == [[2, 1], [1, 1, 1]]
counts = ak.flatten(ak.run_lengths(sorted.x), axis=None)
assert ak.unflatten(sorted, counts, axis=-1).tolist() == [
[[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}]],
[[{"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}], [{"x": 3, "y": 3.3}]],
]
def get_taus(self, apply_selection=True):
events = self.get_events()
taus_dict = {"e": events.tau_e, "pt": events.tau_pt, "eta": events.tau_eta, "phi": events.tau_phi,
"looseIsoAbs": events.tau_looseIsoAbs, "looseIsoRel": events.tau_looseIsoRel,
"mediumIsoAbs": events.tau_mediumIsoAbs, "mediumIsoRel": events.tau_mediumIsoRel,
"tightIsoAbs": events.tau_tightIsoAbs, "tightIsoRel": events.tau_tightIsoRel,
"gen_e": events.gen_tau_e, "gen_pt": events.gen_tau_pt, "gen_eta": events.gen_tau_eta,
"gen_phi": events.gen_tau_phi,
"lepton_gen_match": events.lepton_gen_match, "deepTau_VSjet": events.deepTau_VSjet}
if self.is_old:
taus = ak.zip(taus_dict)
index = ak.argsort(taus.pt, ascending=False)
taus = taus[index]
tau_1, tau_2 = ak.unzip(ak.combinations(taus, 2, axis=1))
else:
taus_dict["vz"] = events.tau_vz
taus = ak.zip(taus_dict)
index = ak.argsort(taus.pt, ascending=False)
taus = taus[index]
tau_1, tau_2 = ak.unzip(ak.combinations(taus, 2, axis=1))
if apply_selection:
L1taus = ak.zip({"e": events.L1tau_e, "pt": events.L1tau_pt, "eta": events.L1tau_eta,
"phi": events.L1tau_phi})
# apply L1seed correction in case Pt28 and Pt30 seeds are considered
L1taus, taus = L1seed_correction(L1taus, taus)
# match taus with L1 taus
taus = L1THLTTauMatching(L1taus, taus)
tau_1, tau_2 = HLTJetPairDzMatchFilter(taus)
# Return all possible pairs of tau which pass preselection
return tau_1, tau_2
def test_NumpyArray():
array = ak.layout.NumpyArray(np.array([3.3, 2.2, 1.1, 5.5, 4.4]))
assert ak.to_list(ak.argsort(array, axis=0, ascending=True,
stable=False)) == [
2,
1,
0,
4,
3,
]
assert ak.to_list(ak.argsort(array, axis=0, ascending=False,
stable=False)) == [
3,
4,
0,
1,
2,
]
assert ak.to_list(ak.sort(array, axis=0, ascending=True,
stable=False)) == [
1.1,
2.2,
3.3,
4.4,
5.5,
]
assert ak.to_list(ak.sort(array, axis=0, ascending=False,
stable=False)) == [
5.5,
4.4,
3.3,
2.2,
1.1,
]
array2 = ak.layout.NumpyArray(np.array([[3.3, 2.2, 4.4], [1.1, 5.5, 3.3]]))
assert ak.to_list(ak.sort(array2, axis=1, ascending=True,
stable=False)) == ak.to_list(
np.sort(array2, axis=1))
assert ak.to_list(ak.sort(array2, axis=0, ascending=True,
stable=False)) == ak.to_list(
np.sort(array2, axis=0))
assert ak.to_list(ak.argsort(array2, axis=1, ascending=True,
stable=False)) == ak.to_list(
np.argsort(array2, 1))
assert ak.to_list(ak.argsort(array2, axis=0, ascending=True,
stable=False)) == ak.to_list(
np.argsort(array2, 0))
with pytest.raises(ValueError) as err:
ak.sort(array2, axis=2, ascending=True, stable=False)
assert str(err.value).startswith(
"axis=2 exceeds the depth of the nested list structure (which is 2)")
def test_EmptyArray():
array = ak.layout.EmptyArray()
assert ak.to_list(ak.sort(array)) == []
assert ak.to_list(ak.argsort(array)) == []
assert str(ak.type(ak.sort(array))) == "0 * float64"
assert str(ak.type(ak.argsort(array))) == "0 * int64"
array2 = ak.Array([[], [], []])
assert ak.to_list(ak.argsort(array2)) == [[], [], []]
assert str(ak.type(ak.argsort(array2))) == "3 * var * int64"
def test_argsort():
array = ak.Array([1, 2, None, 3, 0, None])
assert ak.argsort(array).tolist() == [4, 0, 1, 3, 2, 5]
assert array[ak.argsort(array)].tolist() == [
0,
1,
2,
3,
None,
None,
]
def test_3d():
array = ak.layout.NumpyArray(
np.array(
[
# axis 2: 0 1 2 3 4 # axis 1:
[
[1.1, 2.2, 3.3, 4.4, 5.5], # 0
[6.6, 7.7, 8.8, 9.9, 10.10], # 1
[11.11, 12.12, 13.13, 14.14, 15.15],
], # 2
[
[-1.1, -2.2, -3.3, -4.4, -5.5], # 3
[-6.6, -7.7, -8.8, -9.9, -10.1], # 4
[-11.11, -12.12, -13.13, -14.14, -15.15],
],
]
)
) # 5
assert ak.to_list(
ak.argsort(array, axis=2, ascending=True, stable=False)
) == ak.to_list(np.argsort(array, 2))
assert ak.to_list(
ak.sort(array, axis=2, ascending=True, stable=False)
) == ak.to_list(np.sort(array, 2))
assert ak.to_list(
ak.argsort(array, axis=1, ascending=True, stable=False)
) == ak.to_list(np.argsort(array, 1))
assert ak.to_list(
ak.sort(array, axis=1, ascending=True, stable=False)
) == ak.to_list(np.sort(array, 1))
assert ak.to_list(ak.sort(array, axis=1, ascending=False, stable=False)) == [
[
[11.11, 12.12, 13.13, 14.14, 15.15],
[6.6, 7.7, 8.8, 9.9, 10.1],
[1.1, 2.2, 3.3, 4.4, 5.5],
],
[
[-1.1, -2.2, -3.3, -4.4, -5.5],
[-6.6, -7.7, -8.8, -9.9, -10.1],
[-11.11, -12.12, -13.13, -14.14, -15.15],
],
]
assert ak.to_list(
ak.sort(array, axis=0, ascending=True, stable=False)
) == ak.to_list(np.sort(array, 0))
assert ak.to_list(
ak.argsort(array, axis=0, ascending=True, stable=False)
) == ak.to_list(np.argsort(array, 0))
def test_ByteMaskedArray():
content = ak.from_iter(
[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]],
highlevel=False)
mask = ak.layout.Index8(np.array([0, 0, 1, 1, 0], dtype=np.int8))
array = ak.layout.ByteMaskedArray(mask, content, valid_when=False)
assert ak.to_list(ak.argsort(array, axis=0, ascending=True,
stable=False)) == [
[0, 0, 0],
[],
[2, 2, 2, 2],
None,
None,
]
assert ak.to_list(ak.sort(array, axis=0, ascending=True,
stable=False)) == [
[0.0, 1.1, 2.2],
[],
[6.6, 7.7, 8.8, 9.9],
None,
None,
]
assert ak.to_list(ak.sort(array, axis=0, ascending=False,
stable=False)) == [
[6.6, 7.7, 8.8],
[],
[0.0, 1.1, 2.2, 9.9],
None,
None,
]
assert ak.to_list(ak.argsort(array, axis=1, ascending=True,
stable=False)) == [
[0, 1, 2],
[],
None,
None,
[0, 1, 2, 3],
]
assert ak.to_list(array.sort(1, False, False)) == [
[2.2, 1.1, 0.0],
[],
None,
None,
[9.9, 8.8, 7.7, 6.6],
]
def run_deltar_matching(store,
target,
drname='deltaR',
radius=0.4,
unique=False,
sort=False):
"""
Running a delta R matching of some object collection "store" of dimension NxS
with some target collection "target" of dimension NxT, The return object will
have dimension NxSxT' where objects in the T' contain all "target" objects
within the delta R radius. The delta R between the store and target object will
be stored in the field `deltaR`. If the unique flag is turned on, then objects
in the target collection will only be associated to the closest object. If the
sort flag is turned on, then the target collection will be sorted according to
the computed `deltaR`.
"""
_, target = ak.unzip(ak.cartesian([store.eta, target], nested=True))
target[drname] = delta_r(store, target)
if unique: # Additional filtering
t_index = ak.argmin(target[drname], axis=-2)
s_index = ak.local_index(store.eta, axis=-1)
_, t_index = ak.unzip(ak.cartesian([s_index, t_index], nested=True))
target = target[s_index == t_index]
# Cutting on the computed delta R
target = target[target[drname] < radius]
# Sorting according to the computed delta R
if sort:
idx = ak.argsort(target[drname], axis=-1)
target = target[idx]
return target
def test_sort_bytestrings():
array = ak.from_iter(
[b"one", b"two", b"three", b"two", b"two", b"one", b"three"], highlevel=False
)
assert ak.to_list(array) == [
b"one",
b"two",
b"three",
b"two",
b"two",
b"one",
b"three",
]
assert ak.to_list(ak.sort(array, axis=0, ascending=True, stable=False)) == [
b"one",
b"one",
b"three",
b"three",
b"two",
b"two",
b"two",
]
assert ak.to_list(ak.argsort(array, axis=0, ascending=True, stable=True)) == [
0,
5,
2,
6,
1,
3,
4,
]
def sort_constituents_by_var(taus, c_type, var, ascending=True):
"""
Sort constituents in array `taus[c_type]`` inplace according to the values of `var`.
This transformation is needed for `fill_feature_tensor()` to resolve the case of multiple cell entries,
where the highest pt candidate needs to be filled into the grid cell.
"""
idx = ak.argsort(taus[c_type][var], ascending=ascending)
taus[c_type] = taus[c_type][idx]
def test_sort_zero_length_arrays():
array = ak.layout.IndexedArray64(
ak.layout.Index64([]), ak.layout.NumpyArray([1, 2, 3])
)
assert ak.to_list(array) == []
assert ak.to_list(ak.sort(array)) == []
assert ak.to_list(ak.argsort(array)) == []
content0 = ak.from_iter([[1.1, 2.2, 3.3], [], [4.4, 5.5]], highlevel=False)
content1 = ak.from_iter(["one", "two", "three", "four", "five"], highlevel=False)
tags = ak.layout.Index8([])
index = ak.layout.Index32([])
array = ak.layout.UnionArray8_32(tags, index, [content0, content1])
assert ak.to_list(array) == []
assert ak.to_list(ak.sort(array)) == []
content = ak.from_iter(
[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]], highlevel=False
)
mask = ak.layout.Index8([])
array = ak.layout.ByteMaskedArray(mask, content, valid_when=False)
assert ak.to_list(array) == []
assert ak.to_list(ak.sort(array)) == []
assert ak.to_list(ak.argsort(array)) == []
array = ak.layout.NumpyArray([])
assert ak.to_list(array) == []
assert ak.to_list(ak.sort(array)) == []
assert ak.to_list(ak.argsort(array)) == []
content = ak.layout.NumpyArray(
np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
)
starts1 = ak.layout.Index64([])
stops1 = ak.layout.Index64([])
offsets1 = ak.layout.Index64(np.array([0]))
array = ak.layout.ListArray64(starts1, stops1, content)
assert ak.to_list(array) == []
assert ak.to_list(ak.sort(array)) == []
assert ak.to_list(ak.argsort(array)) == []
array = ak.layout.ListOffsetArray64(offsets1, content)
assert ak.to_list(array) == []
assert ak.to_list(ak.sort(array)) == []
assert ak.to_list(ak.argsort(array)) == []
def get_gen_taus(self):
events = self.get_gen_events()
gen_taus = ak.zip({"gen_e": events.gen_tau_e, "gen_pt": events.gen_tau_pt, "gen_eta": events.gen_tau_eta,
"gen_phi": events.gen_tau_phi,
"lepton_gen_match": events.lepton_gen_match})
index = ak.argsort(gen_taus.gen_pt, ascending=False)
gen_taus = gen_taus[index]
gen_tau_1, gen_tau_2 = ak.unzip(ak.combinations(gen_taus, 2, axis=1))
return gen_tau_1, gen_tau_2
def test_empty_slice():
# muon = ak.Array([[{"pt": 1.0}], []], with_name="muon")
electron = ak.Array([[], [{"pt": 1.0}]], with_name="electron")
electron = electron[electron.pt > 5]
id = ak.argsort(electron.pt, axis=1)
assert ak.to_list(electron[id]) == [[], []]
def test_keep_None_in_place_test():
array = ak.Array([[3, 2, 1], [], None, [4, 5]])
assert ak.to_list(ak.argsort(array, axis=1)) == [
[2, 1, 0],
[],
None,
[0, 1],
]
assert ak.to_list(ak.sort(array, axis=1)) == [
[1, 2, 3],
[],
None,
[4, 5],
]
assert ak.to_list(array[ak.argsort(array, axis=1)]) == ak.to_list(
ak.sort(array, axis=1))
def test_second_issue():
a = ak.layout.NumpyArray(np.arange(122))
idx = ak.layout.Index64([0, 2, 4, 6, 8, 10, 12])
a = ak.layout.ListOffsetArray64(idx, a)
idx = ak.layout.Index64([0, -1, 1, 2, -1, 3, 4, 5])
a = ak.layout.IndexedOptionArray64(idx, a)
a = ak.Array(a)
assert ak.is_valid(a)
assert ak.is_valid(ak.argsort(a))
assert a[ak.argsort(a)].tolist() == [
[0, 1],
None,
[2, 3],
[4, 5],
None,
[6, 7],
[8, 9],
[10, 11],
]
def sorted_by_E(data):
if isinstance(data, ak.Array):
try:
tempE = data.E
except AttributeError:
raise AttributeError(
"Needs either correct coordinates or embedded vector backend"
) from None
tmpsort = ak.argsort(tempE, axis=-1)
return data[tmpsort]
else:
return fastjet._swig.sorted_by_E(data)
def cluster(inter, classify_by_energy=False):
"""
Function which clusters the found clusters together.
To cluster events a weighted mean is computed for time and position.
The individual interactions are weighted by their energy.
The energy of clustered interaction is given by the total sum.
Events can be classified either by the first interaction in time in the
cluster or by the highest energy deposition.
Args:
inter (awkward.Array): Array containing at least the following
fields: x,y,z,t,ed,cluster_ids, type, parenttype, creaproc,
edproc.
Kwargs:
classify_by_energy (bool): If true events are classified
according to the properties of the highest energy deposit
within the cluster. If false cluster is classified according
to first interaction.
Returns:
awkward.Array: Clustered events with nest conform
classification.
"""
if len(inter) == 0:
result_cluster_dtype = [('x', 'float64'), ('y', 'float64'),
('z', 'float64'), ('t', 'float64'),
('ed', 'float64'), ('nestid', 'int64'),
('A', 'int64'), ('Z', 'int64')]
return ak.from_numpy(np.empty(0, dtype=result_cluster_dtype))
# Sort interactions by cluster_ids to simplify looping
inds = ak.argsort(inter['cluster_ids'])
inter = inter[inds]
# TODO: Better way to do this with awkward?
x = inter['x']
y = inter['y']
z = inter['z']
ed = inter['ed']
time = inter['t']
ci = inter['cluster_ids']
types = inter['type']
parenttype = inter['parenttype']
creaproc = inter['creaproc']
edproc = inter['edproc']
# Init result and cluster:
res = ak.ArrayBuilder()
_cluster(x, y, z, ed, time, ci, types, parenttype, creaproc, edproc,
classify_by_energy, res)
return res.snapshot()
def test_argsort_2d():
array = ak.Array([[1, 2, None, 3, 0, None], [1, 2, None, 3, 0, None]])
assert ak.argsort(array).tolist() == [[4, 0, 1, 3, 2, 5], [4, 0, 1, 3, 2, 5]]
assert array[ak.argsort(array)].tolist() == [
[
0,
1,
2,
3,
None,
None,
],
[
0,
1,
2,
3,
None,
None,
],
]
def test_sort():
assert ak.sort(ak.from_iter([[2 + 4j, 1 + 5j], [],
[3 + 3j]])).tolist() == [
[1 + 5j, 2 + 4j],
[],
[3 + 3j],
]
assert ak.argsort(ak.from_iter([[2 + 4j, 1 + 5j], [],
[3 + 3j]])).tolist() == [
[1, 0],
[],
[0],
]
def test_nan():
array = ak.Array([1, 2, np.nan, 3, 0, np.nan])
assert ak.argsort(array).tolist() == [
2,
5,
4,
0,
1,
3,
]
# Note, `nan` comparison with `nan` returns False
assert str(ak.sort(array).tolist()) == "[nan, nan, 0.0, 1.0, 2.0, 3.0]"
def test_argsort():
array = ak.Array(["one", "two", "three", "four", "five", "six", "seven", "eight"])
assert ak.argsort(array, axis=-1).tolist() == [7, 4, 3, 0, 6, 5, 2, 1]
array = ak.Array(
[["twotwo", "two", "three"], ["four", "five"], [], ["six", "seven", "eight"]]
)
assert ak.argsort(array, axis=-1).tolist() == [[2, 1, 0], [1, 0], [], [2, 1, 0]]
array = ak.Array(
[
[["twotwo", "two"], ["three"]],
[["four", "five"]],
[],
[["six"], ["seven", "eight"]],
]
)
assert ak.argsort(array, axis=-1).tolist() == [
[[1, 0], [0]],
[[1, 0]],
[],
[[0], [1, 0]],
]
def test():
array = ak.Array([[0, 1, 2, 3], [3, 3, 3, 2, 1]])
is_valid = array != 3
assert array.mask[is_valid].tolist() == [[0, 1, 2, None], [None, None, None, 2, 1]]
assert ak.sort(array.mask[is_valid]).tolist() == [
[0, 1, 2, None],
[1, 2, None, None, None],
]
assert ak.argsort(array.mask[is_valid]).tolist() == [
[0, 1, 2, 3],
[4, 3, 0, 1, 2],
]
def test_argsort():
array = ak.Array([[1.1, 2.2], [3.3, 3.1]])
assert ak.argsort(array).tolist() == [[0, 1], [1, 0]]
assert str(ak.type(ak.argsort(array))) == "2 * var * int64"
empty_array = ak.Array([[], []])
assert ak.argsort(empty_array).tolist() == [[], []]
assert str(ak.type(ak.argsort(empty_array))) == "2 * var * int64"
select_array = array[array > 5]
assert select_array.tolist() == [[], []]
assert str(ak.type(select_array)) == "2 * var * float64"
assert ak.argsort(select_array).tolist() == [[], []]
assert str(ak.type(ak.argsort(select_array))) == "2 * var * int64"
assert ak.argsort(array[array > 5]).tolist() == [[], []]
assert str(ak.type(ak.argsort(array[array > 5]))) == "2 * var * int64"
def process(self, events):
output = self.accumulator.identity()
# we can use a very loose preselection to filter the events. nothing is done with this presel, though
presel = ak.num(events.Jet) > 0
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
#output['totalEvents']['all'] += len(events)
#output['skimmedEvents']['all'] += len(ev)
if self.year == 2018:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL
elif self.year == 2017:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL
elif self.year == 2016:
triggers = ev.HLT.Ele23_Ele12_CaloIdL_TrackIdL_IsoVL_DZ
if self.year == 2018:
lumimask = LumiMask(
'processors/Cert_314472-325175_13TeV_Legacy2018_Collisions18_JSON.txt'
)
## Electrons
electron = Collections(ev, "Electron", "tight").get()
electron = electron[(electron.pt > 25) & (np.abs(electron.eta) < 2.4)]
loose_electron = Collections(ev, "Electron", "veto").get()
loose_electron = loose_electron[(loose_electron.pt > 25)
& (np.abs(loose_electron.eta) < 2.4)]
SSelectron = (ak.sum(electron.charge, axis=1) != 0) & (ak.num(electron)
== 2)
OSelectron = (ak.sum(electron.charge, axis=1) == 0) & (ak.num(electron)
== 2)
dielectron = choose(electron, 2)
dielectron_mass = (dielectron['0'] + dielectron['1']).mass
dielectron_pt = (dielectron['0'] + dielectron['1']).pt
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[(leading_electron_idx)]
leading_electron = leading_electron[(leading_electron.pt > 30)]
trailing_electron_idx = ak.singletons(ak.argmin(electron.pt, axis=1))
trailing_electron = electron[trailing_electron_idx]
##Muons
loose_muon = Collections(ev, "Muon", "veto").get()
loose_muon = loose_muon[(loose_muon.pt > 20)
& (np.abs(loose_muon.eta) < 2.4)]
#jets
jet = getJets(ev, minPt=40, maxEta=2.4, pt_var='pt', UL=False)
jet = jet[ak.argsort(
jet.pt, ascending=False
)] # need to sort wrt smeared and recorrected jet pt
jet = jet[~match(jet, loose_muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
#selections
filters = getFilters(ev, year=self.year, dataset=dataset)
mask = lumimask(ev.run, ev.luminosityBlock)
ss = (SSelectron)
os = (OSelectron)
mass = (ak.min(np.abs(dielectron_mass - 91.2), axis=1) < 15)
lead_electron = (ak.min(leading_electron.pt, axis=1) > 30)
jet1 = (ak.num(jet) >= 1)
jet2 = (ak.num(jet) >= 2)
num_loose = ((ak.num(loose_electron) == 2) & (ak.num(loose_muon) == 0))
selection = PackedSelection()
selection.add('filter', (filters))
selection.add('mask', (mask))
selection.add('ss', ss)
selection.add('os', os)
selection.add('mass', mass)
selection.add('leading', lead_electron)
selection.add('triggers', triggers)
selection.add('one jet', jet1)
selection.add('two jets', jet2)
selection.add('num_loose', num_loose)
bl_reqs = ['filter'] + ['mass'] + ['mask'] + ['triggers'] + [
'leading'
] + ['num_loose']
#bl_reqs = ['filter'] + ['mass'] + ['triggers'] + ['leading'] + ['num_loose']
bl_reqs_d = {sel: True for sel in bl_reqs}
baseline = selection.require(**bl_reqs_d)
s_reqs = bl_reqs + ['ss']
s_reqs_d = {sel: True for sel in s_reqs}
ss_sel = selection.require(**s_reqs_d)
o_reqs = bl_reqs + ['os']
o_reqs_d = {sel: True for sel in o_reqs}
os_sel = selection.require(**o_reqs_d)
j1s_reqs = s_reqs + ['one jet']
j1s_reqs_d = {sel: True for sel in j1s_reqs}
j1ss_sel = selection.require(**j1s_reqs_d)
j1o_reqs = o_reqs + ['one jet']
j1o_reqs_d = {sel: True for sel in j1o_reqs}
j1os_sel = selection.require(**j1o_reqs_d)
j2s_reqs = s_reqs + ['two jets']
j2s_reqs_d = {sel: True for sel in j2s_reqs}
j2ss_sel = selection.require(**j2s_reqs_d)
j2o_reqs = o_reqs + ['two jets']
j2o_reqs_d = {sel: True for sel in j2o_reqs}
j2os_sel = selection.require(**j2o_reqs_d)
output["N_jet"].fill(
dataset=dataset,
multiplicity=ak.num(jet)[os_sel],
)
return output
def select_normal(genparts, w_decay_momid):
columns = ["pt", "eta", "phi", "mass", "pdgId", "charge", "decaytype"]
# get tops, defined as last copy
gen_tops = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == 6)]
gen_tops['charge'] = ak.ones_like(gen_tops.pt) * (2. / 3.)
gen_tbars = genparts[(genparts.hasFlags(['isLastCopy']))
& (genparts.pdgId == -6)]
gen_tbars['charge'] = ak.ones_like(gen_tbars.pt) * (-2. / 3.)
# get direct top decay products (children are "isHardProcess" and "isFirstCopy")
gen_bs = ak.flatten(gen_tops.children[(gen_tops.children.pdgId == 5)],
axis=2)
gen_bs['charge'] = ak.ones_like(gen_bs.pt) * (-1. / 3.)
gen_bbars = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -5)], axis=2)
gen_bbars['charge'] = ak.ones_like(gen_bbars.pt) * (1. / 3.)
gen_wplus = ak.flatten(
gen_tops.children[(gen_tops.children.pdgId == w_decay_momid)], axis=2)
gen_wplus['charge'] = ak.ones_like(gen_wplus.pt)
gen_wminus = ak.flatten(
gen_tbars.children[(gen_tbars.children.pdgId == -1 * w_decay_momid)],
axis=2)
gen_wminus['charge'] = ak.ones_like(gen_wminus.pt) * (-1.)
# get w decay products
# get last copy of W bosons
last_gen_wplus = ak.flatten(gen_tops.distinctChildren[
(gen_tops.distinctChildren.pdgId == w_decay_momid)
& (gen_tops.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
last_gen_wminus = ak.flatten(gen_tbars.distinctChildren[
(gen_tbars.distinctChildren.pdgId == -1 * w_decay_momid)
& (gen_tbars.distinctChildren.hasFlags(['isLastCopy']))],
axis=2)
gen_wplus_decaytype = np.zeros(len(gen_wplus))
gen_wminus_decaytype = np.zeros(len(gen_wminus))
# up/down partons from last W+
gen_wpartons_up_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (2. / 3.)
gen_wplus_decaytype[ak.num(gen_wpartons_up_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWplus) > 0)) * 2
gen_wpartons_dw_fromWplus = ak.flatten(
last_gen_wplus.children[(np.mod(last_gen_wplus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wplus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWplus['charge'] = np.sign(
gen_wpartons_up_fromWplus.pdgId) * (-1. / 3.)
# up/down partons from last W-
gen_wpartons_up_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 0)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_up_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (2. / 3.)
gen_wminus_decaytype[ak.num(gen_wpartons_up_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_wpartons_up_fromWminus) > 0)) * 2
gen_wpartons_dw_fromWminus = ak.flatten(last_gen_wminus.children[
(np.mod(last_gen_wminus.children.pdgId, 2) == 1)
& (np.abs(last_gen_wminus.children.pdgId) < 6)],
axis=2)
gen_wpartons_dw_fromWminus['charge'] = np.sign(
gen_wpartons_up_fromWminus.pdgId) * (-1. / 3.)
# charged leps from last W+
gen_charged_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 11) |
(np.abs(last_gen_wplus.children.pdgId) == 13) |
(np.abs(last_gen_wplus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWplus['charge'] = ak.ones_like(
gen_charged_leps_fromWplus.pdgId)
gen_wplus_decaytype[ak.num(gen_charged_leps_fromWplus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWplus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wplus['decaytype'] = ak.unflatten(gen_wplus_decaytype,
ak.num(gen_wplus))
gen_tops['decaytype'] = gen_wplus['decaytype'] # set decaytype for tops
gen_bs['decaytype'] = gen_wplus['decaytype'] # set decaytype for bs
# neutral leps from last W+
gen_neutral_leps_fromWplus = ak.flatten(
last_gen_wplus.children[(np.abs(last_gen_wplus.children.pdgId) == 12) |
(np.abs(last_gen_wplus.children.pdgId) == 14) |
(np.abs(last_gen_wplus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWplus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWplus.pt)
# charged leps from last W-
gen_charged_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 11) |
(np.abs(last_gen_wminus.children.pdgId) == 13) |
(np.abs(last_gen_wminus.children.pdgId) == 15)],
axis=2)
gen_charged_leps_fromWminus['charge'] = ak.ones_like(
gen_charged_leps_fromWminus.pdgId) * (-1.)
gen_wminus_decaytype[ak.num(gen_charged_leps_fromWminus) > 0] = np.ones(
ak.sum(ak.num(gen_charged_leps_fromWminus) > 0))
# add decaytype (0 is INVALID, 1 is LEPTONIC, 2 is HADRONIC)
gen_wminus['decaytype'] = ak.unflatten(gen_wminus_decaytype,
ak.num(gen_wminus))
gen_tbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for tbars
gen_bbars['decaytype'] = gen_wminus['decaytype'] # set decaytype for bbars
# neutral leps from last W-
gen_neutral_leps_fromWminus = ak.flatten(last_gen_wminus.children[
(np.abs(last_gen_wminus.children.pdgId) == 12) |
(np.abs(last_gen_wminus.children.pdgId) == 14) |
(np.abs(last_gen_wminus.children.pdgId) == 16)],
axis=2)
gen_neutral_leps_fromWminus['charge'] = ak.zeros_like(
gen_neutral_leps_fromWminus.pt)
gen_wpartons_up = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_wpartons_dw = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_charged_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
gen_neutral_leps = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
if column == 'decaytype': continue
gen_wpartons_up[column] = ak.flatten(
ak.concatenate([
gen_wpartons_up_fromWplus[column],
gen_wpartons_up_fromWminus[column]
],
axis=1)) # (up-type partons from W+, W-)
gen_wpartons_dw[column] = ak.flatten(
ak.concatenate([
gen_wpartons_dw_fromWplus[column],
gen_wpartons_dw_fromWminus[column]
],
axis=1)) # (dw-type partons from W+, W-)
gen_charged_leps[column] = ak.flatten(
ak.concatenate([
gen_charged_leps_fromWplus[column],
gen_charged_leps_fromWminus[column]
],
axis=1)) # (charged leps from W+, W-)
gen_neutral_leps[column] = ak.flatten(
ak.concatenate([
gen_neutral_leps_fromWplus[column],
gen_neutral_leps_fromWminus[column]
],
axis=1)) # (neutral leps from W+, W-)
gen_wpartons_up = ak.unflatten(
gen_wpartons_up,
ak.num(gen_wpartons_up_fromWplus) + ak.num(gen_wpartons_up_fromWminus))
gen_wpartons_dw = ak.unflatten(
gen_wpartons_dw,
ak.num(gen_wpartons_dw_fromWplus) + ak.num(gen_wpartons_dw_fromWminus))
gen_charged_leps = ak.unflatten(
gen_charged_leps,
ak.num(gen_charged_leps_fromWplus) +
ak.num(gen_charged_leps_fromWminus))
gen_neutral_leps = ak.unflatten(
gen_neutral_leps,
ak.num(gen_neutral_leps_fromWplus) +
ak.num(gen_neutral_leps_fromWminus))
gen_taus = gen_charged_leps[np.abs(gen_charged_leps.pdgId) == 15]
gen_nu_taus = gen_neutral_leps[np.abs(gen_neutral_leps.pdgId) == 16]
# fully hadronic evts
had_evts = (ak.num(gen_charged_leps) == 0) & (
ak.num(gen_neutral_leps)
== 0) & (ak.num(gen_wpartons_up) == 2) & (ak.num(gen_wpartons_dw) == 2)
#set_trace()
# get direct tau decay products from hard processes (subset of gen_taus events above)
tau_decay_prods = genparts[genparts.hasFlags(
['isDirectHardProcessTauDecayProduct'])]
# only need decays to leptons (e/mu, tau nu) for event classification
tau_TO_tau_nu = tau_decay_prods[np.abs(tau_decay_prods.pdgId) == 16]
tau_TO_charged_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 11)
|
(np.abs(tau_decay_prods.pdgId) == 13)]
tau_TO_neutral_lep = tau_decay_prods[(np.abs(tau_decay_prods.pdgId) == 12)
|
(np.abs(tau_decay_prods.pdgId) == 14)]
# set decaytype for gen taus
charged_lep_decaytype_array = ak.to_numpy(
ak.flatten(ak.zeros_like(gen_charged_leps['pt'])))
# semilep evts
semilep_evts = (ak.num(gen_charged_leps) == 1) & (
ak.num(gen_neutral_leps)
== 1) & (ak.num(gen_wpartons_up) == 1) & (ak.num(gen_wpartons_dw) == 1)
tau_jets_evts = semilep_evts & (ak.num(gen_taus) == 1)
sl_evts_mask = np.repeat(ak.to_numpy(semilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
semilep_decaytype_array = np.zeros(ak.to_numpy(semilep_evts).sum(),
dtype=int)
# tau -> l
semilep_tau_leptonic_decay = (tau_jets_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
semilep_tau_hadronic_decay = (tau_jets_evts) & (
~semilep_tau_leptonic_decay)
semilep_decaytype_array[
semilep_tau_leptonic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_leptonic_decay).sum(), dtype=int)
semilep_decaytype_array[
semilep_tau_hadronic_decay[semilep_evts]] = np.ones(
ak.to_numpy(semilep_tau_hadronic_decay).sum(), dtype=int) * 2
# set charged_lep_decatype_array for semileptonic events
charged_lep_decaytype_array[sl_evts_mask] = semilep_decaytype_array
# dilep evts
dilep_evts = (ak.num(gen_charged_leps) == 2) & (
ak.num(gen_neutral_leps)
== 2) & (ak.num(gen_wpartons_up) == 0) & (ak.num(gen_wpartons_dw) == 0)
lep_lep_evts = dilep_evts & (ak.num(gen_taus) == 0)
lep_tau_evts = dilep_evts & (ak.num(gen_taus) == 1)
tau_tau_evts = dilep_evts & (ak.num(gen_taus) == 2)
dl_evts_mask = np.repeat(ak.to_numpy(dilep_evts),
ak.to_numpy(ak.num(gen_charged_leps)))
dilep_decaytype_array = np.zeros((ak.to_numpy(dilep_evts).sum(), 2),
dtype=int)
# tau + tau
# tau + tau -> ll
dilep_TauTau_ll_decay = (tau_tau_evts) & (
ak.num(tau_TO_charged_lep) == 2) & (
ak.num(tau_TO_neutral_lep) == 2) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_ll_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_ll_decay).sum(), 2), dtype=int)
# tau + tau -> hh
dilep_TauTau_hh_decay = (tau_tau_evts) & (
(ak.num(tau_TO_charged_lep) + ak.num(tau_TO_neutral_lep))
== 0) & (ak.num(tau_TO_tau_nu) == 2)
dilep_decaytype_array[dilep_TauTau_hh_decay[dilep_evts]] = np.ones(
(ak.to_numpy(dilep_TauTau_hh_decay).sum(), 2), dtype=int) * 2
# tau + tau -> lh
dilep_TauTau_lh_decay = (
tau_tau_evts) & ~(dilep_TauTau_ll_decay | dilep_TauTau_hh_decay)
# set index corresponding to leptonically decaying tau to 1, default array is set to 2
dl_TauTau_to_lh_decaytype_array = np.ones(
ak.to_numpy(dilep_TauTau_lh_decay).sum() * 2, dtype=int) * 2
lep_tau_mask = (np.repeat(
ak.to_numpy(
ak.flatten(
tau_TO_charged_lep[dilep_TauTau_lh_decay].parent.pdgId)),
2) == ak.flatten(gen_charged_leps[dilep_TauTau_lh_decay].pdgId))
dl_TauTau_to_lh_decaytype_array[lep_tau_mask] = np.ones(lep_tau_mask.sum(),
dtype=int)
dilep_decaytype_array[dilep_TauTau_lh_decay[
dilep_evts]] = dl_TauTau_to_lh_decaytype_array.reshape(
ak.to_numpy(dilep_TauTau_lh_decay).sum(), 2)
# lep + tau
# tau -> l
dilep_LepTau_l_decay = (lep_tau_evts) & (
ak.num(tau_TO_charged_lep) == 1) & (
ak.num(tau_TO_neutral_lep) == 1) & (ak.num(tau_TO_tau_nu) == 1)
# set index corresponding to tau to 1
dl_LepTau_to_Lep_decaytype_array = np.zeros(
ak.to_numpy(dilep_LepTau_l_decay).sum() * 2, dtype=int)
dl_LepTau_to_Lep_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_l_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_l_decay), dtype=int)
dilep_decaytype_array[dilep_LepTau_l_decay[
dilep_evts]] = dl_LepTau_to_Lep_decaytype_array.reshape(
ak.sum(dilep_LepTau_l_decay), 2)
# tau -> h
dilep_LepTau_h_decay = (lep_tau_evts) & ~(dilep_LepTau_l_decay)
# set index corresponding to tau to 2
dl_LepTau_to_Had_decaytype_array = np.zeros(ak.sum(dilep_LepTau_h_decay) *
2,
dtype=int)
dl_LepTau_to_Had_decaytype_array[ak.flatten(
np.abs(gen_charged_leps[dilep_LepTau_h_decay].pdgId) == 15)] = np.ones(
ak.sum(dilep_LepTau_h_decay), dtype=int) * 2
dilep_decaytype_array[dilep_LepTau_h_decay[
dilep_evts]] = dl_LepTau_to_Had_decaytype_array.reshape(
ak.sum(dilep_LepTau_h_decay), 2)
# set charged_lep_decatype_array for dileptonic events
charged_lep_decaytype_array[dl_evts_mask] = dilep_decaytype_array.flatten()
# set charged lepton decaytype (defined for taus only, e/mu are 0) (1 is LEPTONIC, 2 is HADRONIC)
gen_charged_leps['decaytype'] = ak.unflatten(charged_lep_decaytype_array,
ak.num(gen_charged_leps))
# make awkward arrays of (top decay prods, tbar decay prods)
Gen_Top_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_B_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_W_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
Gen_Wparton_Pairs = ak.Array({}, with_name="PtEtaPhiMLorentzVector")
for column in columns:
Gen_Top_Pairs[column] = ak.flatten(
ak.concatenate([gen_tops[column], gen_tbars[column]],
axis=1)) # (top, tbar)
Gen_B_Pairs[column] = ak.flatten(
ak.concatenate([gen_bs[column], gen_bbars[column]],
axis=1)) # (b, bbar)
Gen_W_Pairs[column] = ak.flatten(
ak.concatenate([gen_wplus[column], gen_wminus[column]],
axis=1)) # (W+, W-)
if column is not "decaytype":
Gen_Wparton_Pairs[column] = ak.flatten(
ak.concatenate(
[
ak.pad_none(gen_wpartons_up[column], 1, axis=1),
ak.pad_none(gen_wpartons_dw[column], 1, axis=1)
],
axis=1)) # (up-type wpartons, down-type wpartons)
Gen_Top_Pairs = ak.unflatten(Gen_Top_Pairs,
ak.num(gen_tops) + ak.num(gen_tbars))
Gen_B_Pairs = ak.unflatten(Gen_B_Pairs, ak.num(gen_bs) + ak.num(gen_bbars))
Gen_W_Pairs = ak.unflatten(Gen_W_Pairs,
ak.num(gen_wplus) + ak.num(gen_wminus))
Gen_Wparton_Pairs = ak.unflatten(
Gen_Wparton_Pairs,
ak.num(ak.pad_none(gen_wpartons_up, 1, axis=1)) +
ak.num(ak.pad_none(gen_wpartons_dw, 1, axis=1)))
Gen_Wparton_Pairs = Gen_Wparton_Pairs[ak.argsort(
Gen_Wparton_Pairs["pt"], ascending=False)] # sort by pt
Gen_TTbar = ak.Array(
{
"pt": (gen_tops + gen_tbars).pt,
"eta": (gen_tops + gen_tbars).eta,
"phi": (gen_tops + gen_tbars).phi,
"mass": (gen_tops + gen_tbars).mass,
"decaytype": gen_tops["decaytype"] + gen_tbars[
"decaytype"], # 0 is for INVALID, 2 for DILEP, 3 for SEMILEP, 4 for HADRONIC
},
with_name="PtEtaPhiMLorentzVector")
## make "table" of gen objects for certain decays
## DILEP
DILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[dilep_evts] if ak.any(dilep_evts) else ak.unflatten(
Gen_TTbar[dilep_evts], ak.values_astype(dilep_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts] if ak.any(dilep_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][dilep_evts],
ak.values_astype(dilep_evts, int)),
"First_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"First_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # charged lepton always made leading
"Second_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts,
int)), # neutral lepton always made subleading
"Up_plus":
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second plus
"Down_plus":
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge > 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first plus
"Up_minus":
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_neutral_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as second minus
"Down_minus":
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts]
if ak.any(dilep_evts) else ak.unflatten(
gen_charged_leps[gen_charged_leps.charge < 0][dilep_evts],
ak.values_astype(dilep_evts, int)), # same as first minus
})
## HAD
HAD_evts = ak.zip({
"TTbar":
Gen_TTbar[had_evts] if ak.any(had_evts) else ak.unflatten(
Gen_TTbar[had_evts], ak.values_astype(had_evts, int)),
"Top":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Tbar":
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_Top_Pairs[np.sign(Gen_Top_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"B":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == -1][had_evts],
ak.values_astype(had_evts, int)),
"Bbar":
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts]
if ak.any(had_evts) else ak.unflatten(
Gen_B_Pairs[np.sign(Gen_B_Pairs.charge) == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wplus":
Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts] if ak.any(had_evts) else
ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == 1][had_evts],
ak.values_astype(had_evts, int)),
"Wminus":
Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts] if ak.any(had_evts)
else ak.unflatten(Gen_W_Pairs[Gen_W_Pairs.charge == -1][had_evts],
ak.values_astype(had_evts, int)),
"First_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading positively-charged parton
"Second_plus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge > 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading positively-charged parton
"First_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
0]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 0], ak.values_astype(
had_evts, int)), # leading negatively-charged parton
"Second_minus":
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0][:,
1]
if ak.any(had_evts) else ak.unflatten(
Gen_Wparton_Pairs[had_evts][Gen_Wparton_Pairs[had_evts].charge < 0]
[:, 1], ak.values_astype(
had_evts, int)), # subleading negatively-charged parton
"Up_plus":
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged up-type parton
"Down_plus":
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge > 0][had_evts],
ak.values_astype(had_evts,
int)), # positively-charged down-type parton
"Up_minus":
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_up[gen_wpartons_up.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged up-type parton
"Down_minus":
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts]
if ak.any(had_evts) else ak.unflatten(
gen_wpartons_dw[gen_wpartons_dw.charge < 0][had_evts],
ak.values_astype(had_evts,
int)), # negatively-charged down-type parton
})
## SEMILEP
SEMILEP_evts = ak.zip({
"TTbar":
Gen_TTbar[semilep_evts] if ak.any(semilep_evts) else ak.unflatten(
Gen_TTbar[semilep_evts], ak.values_astype(semilep_evts, int)),
"THad":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"TLep":
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_Top_Pairs[Gen_Top_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BHad":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"BLep":
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_B_Pairs[Gen_B_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WHad":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 2][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WLep":
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts]
if ak.any(semilep_evts) else ak.unflatten(
Gen_W_Pairs[Gen_W_Pairs.decaytype == 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Lepton":
gen_charged_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_charged_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Nu":
gen_neutral_leps[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_neutral_leps[semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJa":
Gen_Wparton_Pairs[:, 0][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 0][semilep_evts],
ak.values_astype(semilep_evts, int)),
"WJb":
Gen_Wparton_Pairs[:, 1][semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(Gen_Wparton_Pairs[:, 1][semilep_evts],
ak.values_astype(semilep_evts, int)),
"Up_Had":
gen_wpartons_up[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_up[semilep_evts],
ak.values_astype(semilep_evts, int)),
"Down_Had":
gen_wpartons_dw[semilep_evts] if ak.any(semilep_evts) else
ak.unflatten(gen_wpartons_dw[semilep_evts],
ak.values_astype(semilep_evts, int)),
})
# make dictionary to return
GenObjects = dict({
"SL": SEMILEP_evts,
"DL": DILEP_evts,
"Had": HAD_evts,
})
return GenObjects
def process(self, events):
output = self.accumulator.identity()
output['total']['all'] += len(events)
# use a very loose preselection to filter the events
presel = ak.num(events.Jet) > 2
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
## Muons
muon = Collections(ev, "Muon", "vetoTTH").get()
tightmuon = Collections(ev, "Muon", "tightTTH").get()
dimuon = choose(muon, 2)
SSmuon = ak.any((dimuon['0'].charge * dimuon['1'].charge) > 0, axis=1)
leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1))
leading_muon = muon[leading_muon_idx]
## Electrons
electron = Collections(ev, "Electron", "vetoTTH").get()
tightelectron = Collections(ev, "Electron", "tightTTH").get()
dielectron = choose(electron, 2)
SSelectron = ak.any(
(dielectron['0'].charge * dielectron['1'].charge) > 0, axis=1)
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[leading_electron_idx]
## Merge electrons and muons - this should work better now in ak1
dilepton = cross(muon, electron)
SSlepton = ak.any((dilepton['0'].charge * dilepton['1'].charge) > 0,
axis=1)
lepton = ak.concatenate([muon, electron], axis=1)
leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1))
leading_lepton = lepton[leading_lepton_idx]
trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1))
trailing_lepton = lepton[trailing_lepton_idx]
dilepton_mass = (leading_lepton + trailing_lepton).mass
dilepton_pt = (leading_lepton + trailing_lepton).pt
dilepton_dR = delta_r(leading_lepton, trailing_lepton)
mt_lep_met = mt(lepton.pt, lepton.phi, ev.MET.pt, ev.MET.phi)
min_mt_lep_met = ak.min(mt_lep_met, axis=1)
## Jets
jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom')
jet = jet[ak.argsort(
jet.pt_nom, ascending=False
)] # need to sort wrt smeared and recorrected jet pt
jet = jet[~match(jet, muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
central = jet[(abs(jet.eta) < 2.4)]
btag = getBTagsDeepFlavB(
jet, year=self.year) # should study working point for DeepJet
light = getBTagsDeepFlavB(jet, year=self.year, invert=True)
fwd = getFwdJet(light)
fwd_noPU = getFwdJet(light, puId=False)
tau = getTaus(ev)
track = getIsoTracks(ev)
## forward jets
j_fwd = fwd[ak.singletons(ak.argmax(
fwd.p, axis=1))] # highest momentum spectator
high_score_btag = central[ak.argsort(central.btagDeepFlavB)][:, :2]
bl = cross(lepton, high_score_btag)
bl_dR = delta_r(bl['0'], bl['1'])
min_bl_dR = ak.min(bl_dR, axis=1)
jf = cross(j_fwd, jet)
mjf = (jf['0'] + jf['1']).mass
j_fwd2 = jf[ak.singletons(
ak.argmax(mjf, axis=1)
)]['1'] # this is the jet that forms the largest invariant mass with j_fwd
delta_eta = ak.fill_none(
ak.pad_none(abs(j_fwd2.eta - j_fwd.eta), 1, clip=True), 0)
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
## other variables
ht = ak.sum(jet.pt, axis=1)
st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt,
axis=1)
## event selectors
filters = getFilters(ev, year=self.year, dataset=dataset)
dilep = ((ak.num(tightelectron) + ak.num(tightmuon)) == 2)
lep0pt = ((ak.num(electron[(electron.pt > 25)]) +
ak.num(muon[(muon.pt > 25)])) > 0)
lep1pt = ((ak.num(electron[(electron.pt > 20)]) +
ak.num(muon[(muon.pt > 20)])) > 1)
lepveto = ((ak.num(electron) + ak.num(muon)) == 2)
selection = PackedSelection()
selection.add('lepveto', lepveto)
selection.add('dilep', dilep)
selection.add('filter', (filters))
selection.add('p_T(lep0)>25', lep0pt)
selection.add('p_T(lep1)>20', lep1pt)
selection.add('SS', (SSlepton | SSelectron | SSmuon))
selection.add('N_jet>3', (ak.num(jet) >= 4))
selection.add('N_central>2', (ak.num(central) >= 3))
selection.add('N_btag>0', (ak.num(btag) >= 1))
selection.add('N_fwd>0', (ak.num(fwd) >= 1))
#ss_reqs = ['lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS']
ss_reqs = [
'lepveto', 'dilep', 'filter', 'p_T(lep0)>25', 'p_T(lep1)>20', 'SS'
]
#bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0', 'N_fwd>0']
bl_reqs = ss_reqs + ['N_jet>3', 'N_central>2', 'N_btag>0']
ss_reqs_d = {sel: True for sel in ss_reqs}
ss_selection = selection.require(**ss_reqs_d)
bl_reqs_d = {sel: True for sel in bl_reqs}
BL = selection.require(**bl_reqs_d)
weight = Weights(len(ev))
if not dataset == 'MuonEG':
# lumi weight
weight.add("weight", ev.weight)
# PU weight - not in the babies...
weight.add("PU",
ev.puWeight,
weightUp=ev.puWeightUp,
weightDown=ev.puWeightDown,
shift=False)
# b-tag SFs
weight.add("btag", self.btagSF.Method1a(btag, light))
# lepton SFs
weight.add("lepton", self.leptonSF.get(electron, muon))
#cutflow = Cutflow(output, ev, weight=weight)
#cutflow_reqs_d = {}
#for req in bl_reqs:
# cutflow_reqs_d.update({req: True})
# cutflow.addRow( req, selection.require(**cutflow_reqs_d) )
labels = {
'topW_v3': 0,
'TTW': 1,
'TTZ': 2,
'TTH': 3,
'ttbar': 4,
'ttbar1l_MG': 4
}
if dataset in labels:
label_mult = labels[dataset]
else:
label_mult = 5
label = np.ones(len(ev[BL])) * label_mult
output["n_lep"] += processor.column_accumulator(
ak.to_numpy((ak.num(electron) + ak.num(muon))[BL]))
output["n_lep_tight"] += processor.column_accumulator(
ak.to_numpy((ak.num(tightelectron) + ak.num(tightmuon))[BL]))
output["lead_lep_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].pt, axis=1)))
output["lead_lep_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].eta, axis=1)))
output["lead_lep_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].phi, axis=1)))
output["lead_lep_charge"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(leading_lepton[BL].charge, axis=1)))
output["sublead_lep_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].pt, axis=1)))
output["sublead_lep_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].eta, axis=1)))
output["sublead_lep_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].phi, axis=1)))
output["sublead_lep_charge"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(trailing_lepton[BL].charge, axis=1)))
output["lead_jet_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].pt, axis=1)))
output["lead_jet_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].eta, axis=1)))
output["lead_jet_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 0:1][BL].phi, axis=1)))
output["sublead_jet_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].pt, axis=1)))
output["sublead_jet_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].eta, axis=1)))
output["sublead_jet_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(jet[:, 1:2][BL].phi, axis=1)))
output["lead_btag_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].pt, axis=1)))
output["lead_btag_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].eta, axis=1)))
output["lead_btag_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 0:1][BL].phi, axis=1)))
output["sublead_btag_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].pt, axis=1)))
output["sublead_btag_eta"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].eta, axis=1)))
output["sublead_btag_phi"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(high_score_btag[:, 1:2][BL].phi, axis=1)))
output["fwd_jet_p"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(ak.pad_none(j_fwd[BL].p, 1, clip=True),
0),
axis=1)))
output["fwd_jet_pt"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].pt, 1, clip=True), 0),
axis=1)))
output["fwd_jet_eta"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].eta, 1, clip=True), 0),
axis=1)))
output["fwd_jet_phi"] += processor.column_accumulator(
ak.to_numpy(
ak.flatten(ak.fill_none(
ak.pad_none(j_fwd[BL].phi, 1, clip=True), 0),
axis=1)))
output["mjj_max"] += processor.column_accumulator(
ak.to_numpy(ak.fill_none(ak.max(mjf[BL], axis=1), 0)))
output["delta_eta_jj"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(delta_eta[BL], axis=1)))
output["met"] += processor.column_accumulator(ak.to_numpy(met_pt[BL]))
output["ht"] += processor.column_accumulator(ak.to_numpy(ht[BL]))
output["st"] += processor.column_accumulator(ak.to_numpy(st[BL]))
output["n_jet"] += processor.column_accumulator(
ak.to_numpy(ak.num(jet[BL])))
output["n_btag"] += processor.column_accumulator(
ak.to_numpy(ak.num(btag[BL])))
output["n_fwd"] += processor.column_accumulator(
ak.to_numpy(ak.num(fwd[BL])))
output["n_central"] += processor.column_accumulator(
ak.to_numpy(ak.num(central[BL])))
output["n_tau"] += processor.column_accumulator(
ak.to_numpy(ak.num(tau[BL])))
output["n_track"] += processor.column_accumulator(
ak.to_numpy(ak.num(track[BL])))
output["dilepton_pt"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(dilepton_pt[BL], axis=1)))
output["dilepton_mass"] += processor.column_accumulator(
ak.to_numpy(ak.flatten(dilepton_mass[BL], axis=1)))
output["min_bl_dR"] += processor.column_accumulator(
ak.to_numpy(min_bl_dR[BL]))
output["min_mt_lep_met"] += processor.column_accumulator(
ak.to_numpy(min_mt_lep_met[BL]))
output["label"] += processor.column_accumulator(label)
output["weight"] += processor.column_accumulator(weight.weight()[BL])
output["presel"]["all"] += len(ev[ss_selection])
output["sel"]["all"] += len(ev[BL])
return output
def sort_by_pt(ele, pho, jet):
ele = ele[ak.argsort(ele.pt, ascending=False, axis=1)]
pho = pho[ak.argsort(pho.pt, ascending=False, axis=1)]
jet = jet[ak.argsort(jet.pt, ascending=False, axis=1)]
return ele, pho, jet
def get(self, ele, mu, variation='central'):
multiplier = {'central': 0, 'up': 1, 'down': -1}
from Tools.helpers import pad_and_flatten
from Tools.objects import cross, choose
import numpy as np
# get a lepton collection
lep = ak.concatenate([mu, ele], axis=1)
# now sort them
lep = lep[ak.argsort(lep.pt, ascending=False)]
l0_is_ele = (abs(pad_and_flatten(lep[:, 0:1].pdgId)) == 11)
l1_is_ele = (abs(pad_and_flatten(lep[:, 1:2].pdgId)) == 11)
if self.year == 2016:
ee_sf = self.evaluator["ee_2016"](pad_and_flatten(ele[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
ee_sf = ee_sf + multiplier[variation] * self.evaluator[
"ee_2016_error"](pad_and_flatten(ele[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
emu_sf = self.evaluator["emu_2016"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(ele[:,
1:2].pt))
emu_sf = emu_sf + multiplier[variation] * self.evaluator[
"emu_2016_error"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
mumu_sf = self.evaluator["mumu_2016"](pad_and_flatten(mu[:,
0:1].pt),
pad_and_flatten(mu[:,
1:2].pt))
mumu_sf = mumu_sf + multiplier[variation] * self.evaluator[
"mumu_2016_error"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(mu[:, 1:2].pt))
elif self.year == 2017:
ee_sf = self.evaluator["ee_2017"](pad_and_flatten(ele[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
ee_sf = ee_sf + multiplier[variation] * self.evaluator[
"ee_2017_error"](pad_and_flatten(ele[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
emu_sf = self.evaluator["emu_2017"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(ele[:,
1:2].pt))
emu_sf = emu_sf + multiplier[variation] * self.evaluator[
"emu_2017_error"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
mumu_sf = self.evaluator["mumu_2017"](pad_and_flatten(mu[:,
0:1].pt),
pad_and_flatten(mu[:,
1:2].pt))
mumu_sf = mumu_sf + multiplier[variation] * self.evaluator[
"mumu_2017_error"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(mu[:, 1:2].pt))
elif self.year == 2018:
ee_sf = self.evaluator["ee_2018"](pad_and_flatten(ele[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
ee_sf = ee_sf + multiplier[variation] * self.evaluator[
"ee_2018_error"](pad_and_flatten(ele[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
emu_sf = self.evaluator["emu_2018"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(ele[:,
1:2].pt))
emu_sf = emu_sf + multiplier[variation] * self.evaluator[
"emu_2018_error"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(ele[:, 1:2].pt))
mumu_sf = self.evaluator["mumu_2018"](pad_and_flatten(mu[:,
0:1].pt),
pad_and_flatten(mu[:,
1:2].pt))
mumu_sf = mumu_sf + multiplier[variation] * self.evaluator[
"mumu_2018_error"](pad_and_flatten(mu[:, 0:1].pt),
pad_and_flatten(mu[:, 1:2].pt))
sf = (ee_sf * (l0_is_ele & l1_is_ele)) + (
emu_sf * (l0_is_ele ^ l1_is_ele)) + (mumu_sf *
(~l0_is_ele & ~l1_is_ele))
#ee_mm_em = ak.concatenate([choose(ele), choose(mu), cross(mu, ele)], axis=1)
#sf = (pad_and_flatten(sf*(ee_mm_em.pt/ee_mm_em.pt)))
#sf = (sf/sf)*sf
#sf = (np.where(np.isnan(sf), 1, sf))
return sf
def process(self, events):
output = self.accumulator.identity()
# use a very loose preselection to filter the events
presel = ak.num(events.Jet) > 2
ev = events[presel]
dataset = ev.metadata['dataset']
# load the config - probably not needed anymore
cfg = loadConfig()
output['totalEvents']['all'] += len(events)
output['skimmedEvents']['all'] += len(ev)
## Generated leptons
'''gen_lep = ev.GenL
leading_gen_lep = gen_lep[ak.singletons(ak.argmax(gen_lep.pt, axis=1))]
trailing_gen_lep = gen_lep[ak.singletons(ak.argmin(gen_lep.pt, axis=1))]'''
## Muons
muon = Collections(ev, "Muon", "tightTTH").get()
vetomuon = Collections(ev, "Muon", "vetoTTH").get()
leading_muon_idx = ak.singletons(ak.argmax(muon.pt, axis=1))
leading_muon = muon[leading_muon_idx]
## Electrons
electron = Collections(ev, "Electron", "tightTTH").get()
vetoelectron = Collections(ev, "Electron", "vetoTTH").get()
leading_electron_idx = ak.singletons(ak.argmax(electron.pt, axis=1))
leading_electron = electron[leading_electron_idx]
## Merge electrons and muons - this should work better now in ak1
dilepton = cross(muon, electron)
dimuon = choose(muon, 2)
OS_dimuon = dimuon[(dimuon['0'].charge * dimuon['1'].charge < 0)]
dielectron = choose(electron, 2)
OS_dielectron = dielectron[(
dielectron['0'].charge * dielectron['1'].charge < 0)]
OS_dimuon_bestZmumu = OS_dimuon[ak.singletons(
ak.argmin(abs(OS_dimuon.mass - 91.2), axis=1))]
OS_dielectron_bestZee = OS_dielectron[ak.singletons(
ak.argmin(abs(OS_dielectron.mass - 91.2), axis=1))]
OS_dilepton_mass = ak.fill_none(
ak.pad_none(ak.concatenate(
[OS_dimuon_bestZmumu.mass, OS_dielectron_bestZee.mass],
axis=1),
1,
clip=True), -1)
lepton = ak.concatenate([muon, electron], axis=1)
leading_lepton_idx = ak.singletons(ak.argmax(lepton.pt, axis=1))
leading_lepton = lepton[leading_lepton_idx]
trailing_lepton_idx = ak.singletons(ak.argmin(lepton.pt, axis=1))
trailing_lepton = lepton[trailing_lepton_idx]
## Jets
jet = getJets(ev, minPt=25, maxEta=4.7, pt_var='pt_nom')
jet = jet[ak.argsort(
jet.pt_nom, ascending=False
)] # need to sort wrt smeared and recorrected jet pt
jet = jet[~match(jet, muon,
deltaRCut=0.4)] # remove jets that overlap with muons
jet = jet[~match(
jet, electron,
deltaRCut=0.4)] # remove jets that overlap with electrons
central = jet[(abs(jet.eta) < 2.4)]
btag = getBTagsDeepFlavB(
jet, year=self.year) # should study working point for DeepJet
light = getBTagsDeepFlavB(jet, year=self.year, invert=True)
fwd = getFwdJet(light)
fwd_noPU = getFwdJet(light, puId=False)
## forward jets
j_fwd = fwd[ak.singletons(ak.argmax(
fwd.p, axis=1))] # highest momentum spectator
jf = cross(j_fwd, jet)
mjf = (jf['0'] + jf['1']).mass
# j_fwd2 = jf[ak.singletons(ak.argmax(mjf, axis=1))]['1'] # this is the jet that forms the largest invariant mass with j_fwd
# delta_eta = abs(j_fwd2.eta - j_fwd.eta)
## MET -> can switch to puppi MET
met_pt = ev.MET.pt
met_phi = ev.MET.phi
## other variables
ht = ak.sum(jet.pt, axis=1)
st = met_pt + ht + ak.sum(muon.pt, axis=1) + ak.sum(electron.pt,
axis=1)
# define the weight
weight = Weights(len(ev))
if not re.search(re.compile('MuonEG|DoubleMuon|DoubleEG|EGamma'),
dataset):
# lumi weight
weight.add("weight", ev.weight * cfg['lumi'][self.year])
# PU weight - not in the babies...
weight.add("PU",
ev.puWeight,
weightUp=ev.puWeightUp,
weightDown=ev.puWeightDown,
shift=False)
# b-tag SFs
weight.add("btag", self.btagSF.Method1a(btag, light))
# lepton SFs
# weight.add("lepton", self.leptonSF.get(electron, muon))
cutflow = Cutflow(output, ev, weight=weight)
sel = Selection(
dataset=dataset,
events=ev,
year=self.year,
ele=electron,
ele_veto=vetoelectron,
mu=muon,
mu_veto=vetomuon,
jet_all=jet,
jet_central=central,
jet_btag=btag,
jet_fwd=fwd,
met=ev.MET,
)
BL = sel.trilep_baseline(cutflow=cutflow)
# first, make a few super inclusive plots
output['ST'].fill(dataset=dataset,
ht=st[BL],
weight=weight.weight()[BL])
output['PV_npvs'].fill(dataset=dataset,
multiplicity=ev.PV[BL].npvs,
weight=weight.weight()[BL])
output['PV_npvsGood'].fill(dataset=dataset,
multiplicity=ev.PV[BL].npvsGood,
weight=weight.weight()[BL])
output['N_jet'].fill(dataset=dataset,
multiplicity=ak.num(jet)[BL],
weight=weight.weight()[BL])
output['N_b'].fill(dataset=dataset,
multiplicity=ak.num(btag)[BL],
weight=weight.weight()[BL])
output['N_central'].fill(dataset=dataset,
multiplicity=ak.num(central)[BL],
weight=weight.weight()[BL])
output['N_ele'].fill(dataset=dataset,
multiplicity=ak.num(vetoelectron)[BL],
weight=weight.weight()[BL])
output['N_mu'].fill(dataset=dataset,
multiplicity=ak.num(vetomuon)[BL],
weight=weight.weight()[BL])
output['N_fwd'].fill(dataset=dataset,
multiplicity=ak.num(fwd)[BL],
weight=weight.weight()[BL])
'''output['nLepFromTop'].fill(dataset=dataset, multiplicity=ev[BL].nLepFromTop, weight=weight.weight()[BL])
output['nLepFromTau'].fill(dataset=dataset, multiplicity=ev.nLepFromTau[BL], weight=weight.weight()[BL])
output['nLepFromZ'].fill(dataset=dataset, multiplicity=ev.nLepFromZ[BL], weight=weight.weight()[BL])
output['nLepFromW'].fill(dataset=dataset, multiplicity=ev.nLepFromW[BL], weight=weight.weight()[BL])
output['nGenTau'].fill(dataset=dataset, multiplicity=ev.nGenTau[BL], weight=weight.weight()[BL])
output['nGenL'].fill(dataset=dataset, multiplicity=ak.num(ev.GenL[BL], axis=1), weight=weight.weight()[BL])'''
# make a plot of the dilepton mass, but without applying the cut on the dilepton mass itself (N-1 plot)
output['dilep_mass'].fill(
dataset=dataset,
mass=ak.flatten(
OS_dilepton_mass[sel.trilep_baseline(omit=['offZ'])]),
weight=weight.weight()[sel.trilep_baseline(omit=['offZ'])])
output['MET'].fill(dataset=dataset,
pt=ev.MET[BL].pt,
phi=ev.MET[BL].phi,
weight=weight.weight()[BL])
'''output['lead_gen_lep'].fill(
dataset = dataset,
pt = ak.to_numpy(ak.flatten(leading_gen_lep[BL].pt)),
eta = ak.to_numpy(ak.flatten(leading_gen_lep[BL].eta)),
phi = ak.to_numpy(ak.flatten(leading_gen_lep[BL].phi)),
weight = weight.weight()[BL]
)
output['trail_gen_lep'].fill(
dataset = dataset,
pt = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].pt)),
eta = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].eta)),
phi = ak.to_numpy(ak.flatten(trailing_gen_lep[BL].phi)),
weight = weight.weight()[BL]
)'''
output['lead_lep'].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(leading_lepton[BL].pt)),
eta=ak.to_numpy(ak.flatten(leading_lepton[BL].eta)),
phi=ak.to_numpy(ak.flatten(leading_lepton[BL].phi)),
weight=weight.weight()[BL])
output['trail_lep'].fill(
dataset=dataset,
pt=ak.to_numpy(ak.flatten(trailing_lepton[BL].pt)),
eta=ak.to_numpy(ak.flatten(trailing_lepton[BL].eta)),
phi=ak.to_numpy(ak.flatten(trailing_lepton[BL].phi)),
weight=weight.weight()[BL])
output['j1'].fill(dataset=dataset,
pt=ak.flatten(jet.pt_nom[:, 0:1][BL]),
eta=ak.flatten(jet.eta[:, 0:1][BL]),
phi=ak.flatten(jet.phi[:, 0:1][BL]),
weight=weight.weight()[BL])
output['j2'].fill(dataset=dataset,
pt=ak.flatten(jet[:, 1:2][BL].pt_nom),
eta=ak.flatten(jet[:, 1:2][BL].eta),
phi=ak.flatten(jet[:, 1:2][BL].phi),
weight=weight.weight()[BL])
#output['j3'].fill(
# dataset = dataset,
# pt = ak.flatten(jet[:, 2:3][BL].pt_nom),
# eta = ak.flatten(jet[:, 2:3][BL].eta),
# phi = ak.flatten(jet[:, 2:3][BL].phi),
# weight = weight.weight()[BL]
#)
output['fwd_jet'].fill(dataset=dataset,
pt=ak.flatten(j_fwd[BL].pt),
eta=ak.flatten(j_fwd[BL].eta),
phi=ak.flatten(j_fwd[BL].phi),
weight=weight.weight()[BL])
output['high_p_fwd_p'].fill(dataset=dataset,
p=ak.flatten(j_fwd[BL].p),
weight=weight.weight()[BL])
vetolepton = ak.concatenate([vetomuon, vetoelectron], axis=1)
trilep = choose3(vetolepton, 3)
trilep_m = trilep.mass
output['m3l'].fill(dataset=dataset,
mass=ak.flatten(trilep_m[BL]),
weight=weight.weight()[BL])
return output
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
