如何访问Awkward Array中指定索引以外的元素？

问题背景 #

现有3个事件，每个事件包含1个Jet和若干Lepton，每个Lepton带有对应标记，且已记录每个Jet关联的Lepton索引：

import awkward as ak

jet_lepton_indices = ak.Array([[0, 2], [1], [2,3]])
print(f'jet_lepton_indices\n{jet_lepton_indices}\n')

lepton_flags = ak.Array([[0, 10, 20, 30], [0, 10, 20, 30], [0, 10, 20, 30, 40]])
print(f'lepton_flags\n{lepton_flags}\n')

输出：

jet_lepton_indices
[[0, 2], [1], [2, 3]]

lepton_flags
[[0, 10, 20, 30], [0, 10, 20, 30], [0, 10, 20, 30, 40]]

通过lepton_flags[jet_lepton_indices]可以获取Jet关联Lepton的标记：

[[0, 20],
 [10],
 [20, 30]]

但需要获取所有未关联Jet的Lepton标记，期望结果：

[[10, 30],
 [0, 20, 30],
 [0, 10, 40]]

直接使用lepton_flags[~jet_lepton_indices]无法得到预期结果，因为jet_lepton_indices是索引集合而非布尔掩码，不能直接取反。

解决方案 #

需要先生成每个事件中Lepton的完整索引数组，再通过求差集得到未关联的索引，最后提取对应标记：

步骤1：生成每个事件的Lepton全索引 #

用ak.local_index()生成每个子数组内的局部索引，对应每个Lepton在事件中的位置：

all_lepton_indices = ak.local_index(lepton_flags)
print(f'all_lepton_indices\n{all_lepton_indices}\n')

输出：

all_lepton_indices
[[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3, 4]]

步骤2：计算未关联的Lepton索引 #

使用ak.setdiff1d()对每个事件的全索引和关联索引求差集，得到未关联的索引：

unassociated_indices = ak.setdiff1d(all_lepton_indices, jet_lepton_indices, axis=1)
print(f'unassociated_indices\n{unassociated_indices}\n')

输出：

unassociated_indices
[[1, 3], [0, 2, 3], [0, 1, 4]]

步骤3：提取未关联Lepton的标记 #

用得到的未关联索引从lepton_flags中取值：

unassociated_flags = lepton_flags[unassociated_indices]
print(f'unassociated_flags\n{unassociated_flags}\n')

输出：

unassociated_flags
[[10, 30], [0, 20, 30], [0, 10, 40]]

完整代码 #

import awkward as ak

# 输入数据
jet_lepton_indices = ak.Array([[0, 2], [1], [2,3]])
lepton_flags = ak.Array([[0, 10, 20, 30], [0, 10, 20, 30], [0, 10, 20, 30, 40]])

# 生成全索引
all_lepton_indices = ak.local_index(lepton_flags)
# 求差集得到未关联索引
unassociated_indices = ak.setdiff1d(all_lepton_indices, jet_lepton_indices, axis=1)
# 提取未关联标记
unassociated_flags = lepton_flags[unassociated_indices]

print(f'预期结果:\n[[10, 30],\n [0, 20, 30],\n [0, 10, 40]]\n')
print(f'实际结果:\n{unassociated_flags}')

内容的提问来源于stack exchange，提问作者Matt Bellis