如何基于索引创建新的Awkward Array

问题描述 #

需要基于粒子索引为每个喷流生成对应的粒子四动量数组，现有粒子和喷流数据如下，用for循环可实现但希望用Awkward Array做高效处理；同时想了解嵌套数组场景（如每个事件包含多个喷流）的处理方式。

示例代码：

import awkward as ak
import vector

particles = ak.Array({
    "mass": ak.Array([1, 2, 3, 4, 5]),
    "px": ak.Array([1, 2, 3, 4, 5]),
    "py": ak.Array([1, 2, 3, 4, 5]),
    "pz": ak.Array([1, 2, 3, 4, 5]),
})

particles_p4 = vector.awk(ak.zip({
    "mass": particles["mass"],
    "x": particles["px"],
    "y": particles["py"],
    "z": particles["pz"]
}))

jet = ak.Array({
    "constituents": ak.Array([[1, 2, 4, 5], [3]]),
    "energy": ak.Array([1.2, 3.4])
})

期望输出：

<Array [[{x: 1, y: 1, z: 1, ... z: 5, tau: 5}], [{x:3, ...}]] type='2 * var * {"x": int64, "y"...'>

解决方案 #

基础场景：单级喷流索引 #

Awkward Array原生支持向量化索引操作，直接通过索引数组提取对应元素，无需循环：

注意：示例中jet.constituents是1-based索引，需转换为0-based才能匹配particles_p4的索引规则

# 转换为0-based索引并提取粒子四动量
jet_constituents_p4 = particles_p4[jet.constituents - 1]

# 查看结果
print(jet_constituents_p4)

输出结果与预期一致：

<Array [[{x:1, y:1, z:1, mass:1}, {x:2,...}, {x:4,...}, {x:5,...}], [{x:3,...}]] type='2 * var * {"x": int64, "y": int64, "z": int64, "mass": int64, ...}'>

嵌套场景：多事件多喷流 #

对于每个事件包含多个喷流的嵌套结构，Awkward Array同样能直接处理，自动适配嵌套层级：

# 模拟嵌套数据：2个事件，每个事件含2个喷流
events = ak.Array({
    "jets": ak.Array([
        [
            {"constituents": [0, 1], "energy": 1.2},
            {"constituents": [3, 4], "energy": 2.3}
        ],
        [
            {"constituents": [2], "energy": 3.4},
            {"constituents": [0, 2, 4], "energy": 4.5}
        ]
    ])
})

# 直接提取嵌套结构下的粒子四动量
event_jet_constituents_p4 = particles_p4[events.jets.constituents]

print(event_jet_constituents_p4)

输出为对应嵌套层级的数组：

<Array [[[{x:1,...}, {x:2,...}], [{x:4,...}, {x:5,...}]], [[{x:3,...}], [{x:1,...}, {x:3,...}, {x:5,...}]]] type='2 * var * var * {"x": int64, "y": int64, "z": int64, "mass": int64, ...}'>

核心优势 #

Awkward Array的索引操作是底层向量化实现，性能远优于Python循环，且天然适配任意嵌套结构，完美匹配粒子物理数据的典型格式。

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