def test():
array = ak.Array(
[[[0.0, 1.1, 2.2], []], [[3.3, 4.4]], [], [[5.5], [], [6.6, 7.7, 8.8, 9.9]]]
)
assert ak.to_list(ak.local_index(array, axis=0)) == [0, 1, 2, 3]
assert ak.to_list(ak.local_index(array, axis=1)) == [[0, 1], [0], [], [0, 1, 2]]
assert ak.to_list(ak.local_index(array, axis=2)) == [
[[0, 1, 2], []],
[[0, 1]],
[],
[[0], [], [0, 1, 2, 3]],
]
assert ak.to_list(ak.local_index(array, axis=-1)) == [
[[0, 1, 2], []],
[[0, 1]],
[],
[[0], [], [0, 1, 2, 3]],
]
assert ak.to_list(ak.local_index(array, axis=-2)) == [[0, 1], [0], [], [0, 1, 2]]
assert ak.to_list(ak.local_index(array, axis=-3)) == [0, 1, 2, 3]
assert ak.to_list(
ak.zip(
[
ak.local_index(array, axis=0),
ak.local_index(array, axis=1),
ak.local_index(array, axis=2),
]
)
) == [
[[(0, 0, 0), (0, 0, 1), (0, 0, 2)], []],
[[(1, 0, 0), (1, 0, 1)]],
[],
[[(3, 0, 0)], [], [(3, 2, 0), (3, 2, 1), (3, 2, 2), (3, 2, 3)]],
]
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 process(self, events):
events["Electron", "pdgId"] = -11 * events.Electron.charge
events["Muon", "pdgId"] = -13 * events.Muon.charge
events["leptons"] = ak.concatenate(
[events.Electron, events.Muon],
axis=1,
)
events = events[ak.num(events.leptons) >= 3]
pair = ak.argcombinations(events.leptons, 2, fields=["l1", "l2"])
pair = pair[(
events.leptons[pair.l1].pdgId == -events.leptons[pair.l2].pdgId)]
with np.errstate(invalid="ignore"):
pair = pair[ak.singletons(
ak.argmin(
abs((events.leptons[pair.l1] +
events.leptons[pair.l2]).mass - 91.2),
axis=1,
))]
events = events[ak.num(pair) > 0]
pair = pair[ak.num(pair) > 0][:, 0]
l3 = ak.local_index(events.leptons)
l3 = l3[(l3 != pair.l1) & (l3 != pair.l2)]
l3 = l3[ak.argmax(events.leptons[l3].pt, axis=1, keepdims=True)]
l3 = events.leptons[l3][:, 0]
mt = np.sqrt(2 * l3.pt * events.MET.pt *
(1 - np.cos(events.MET.delta_phi(l3))))
return (hist.Hist.new.Reg(
100, 0, 200, name="mt",
label=r"$\ell$-MET transverse mass [GeV]").Double().fill(mt))
def matchJets(self, obj, jet, deltaRCut=0.4):
combs = ak.cartesian([obj, jet], nested=True)
jet_index = ak.local_index(delta_r(
combs['0'], combs['1']))[delta_r(combs['0'], combs['1']) < 0.4]
jet_index_pad = ak.flatten(ak.fill_none(
ak.pad_none(jet_index, target=1, clip=True, axis=2), 0),
axis=2)
mask = ak.num(jet_index,
axis=2) > 0 # a mask for obj with a matched jet
mask_match = mask * 1 + ~mask * 0
mask_nomatch = mask * 0 + ~mask * 1
return jet_index_pad, mask_match, mask_nomatch
def test_localindex():
assert ak.local_index(empty, axis=0).tolist() == []
assert ak.local_index(empty, axis=1).tolist() == []
assert ak.local_index(empty, axis=2).tolist() == []