计算射流不变质量时 uproot4/vector 中的 TLorentz 矢量特征

Question

我想总结喷流中成分的所有 4 动量。在 uproot3（+ uproot3-methods）中，有创建 TLorentzVectorArray 的功能，只需执行 .sum()

所以这很好用：

import uproot3
import akward0 as ak

input_file = uproot3.open(input_path)
tree = input_file['Jets']
pt = tree.array('Constituent_pt')
phi = tree.array('Constituent_phi')
eta = tree.array('Constituent_eta')
energy = tree.array('Constituent_energy')
mass = tree.array('Constituent_mass')
p4 = uproot3_methods.TLorentzVectorArray.from_ptetaphie(pt, eta, phi, energy)
jet_p4_u3 = p4.sum()
jet_pt_u3 = jet_p4.pt
jet_eta_u3 = jet_p4.eta
jet_phi_u3 = jet_p4.phi
jet_energy_u3 = jet_p4.energy

但是，由于 uproot3 已被弃用，因此请按照 TLorentz vector in Uproot 4 seems to be the vector 包的方式进行。我尝试的是以下内容。

import uproot
import awkward
import vector

input_file = uproot.open(input_path)
tree = input_file['Jets']
pt = tree.arrays()['Constituent_pt']
phi = tree.arrays()['Constituent_phi']
eta = tree.arrays()['Constituent_eta']
energy = tree.arrays()['Constituent_energy']
mass = tree.arrays()['Constituent_mass']
p4 = vector.awk({"pt": pt, "phi": phi, "eta": eta, "energy": energy})

现在的问题是这个功能 p4.sum() 似乎不存在。 vector discussion #117 中显示了我发现的另一种可能性。所以，现在我在导入 vector.register_awkward() 之后添加到最后 jet_p4_u4 = ak.Array(p4, with_name="Momentum4D"),

import uproot
import awkward
import vector
vector.register_awkward()

input_file = uproot.open(input_path)
tree = input_file['Jets']
pt = tree.arrays()['Constituent_pt']
phi = tree.arrays()['Constituent_phi']
eta = tree.arrays()['Constituent_eta']
energy = tree.arrays()['Constituent_energy']
mass = tree.arrays()['Constituent_mass']
p4 = ak.Array({"pt": pt, "phi": phi, "eta": eta, "energy": energy})
jet_p4_u4 = ak.Array(p4, with_name="Momentum4D")

问题仍然存在，如何对 4-momenta 求和？当做 ak.sum(jet_p4_u4, axis=-1) 时，只有 pt 和 energy 似乎有正确的值，然而 eta 和 phi 与 uproot3.

的结果完全不同

更新：似乎由于```ak.sum`` 函数无法以所需的方式将角度相加，因此将求和部分替换为求和 x、y、z 和能量像这样构建向量可以解决问题。但是，我相信一定有比这更好的方法。所以当前工作版本：

import uproot
import awkward
import vector

input_file = uproot.open(input_path)
tree = input_file['Jets']
pt = tree.arrays()['Constituent_pt']
phi = tree.arrays()['Constituent_phi']
eta = tree.arrays()['Constituent_eta']
energy = tree.arrays()['Constituent_energy']
mass = tree.arrays()['Constituent_mass']
p4 = vector.awk({"pt": pt, "phi": phi, "eta": eta, "energy": energy})
p4_lz = vector.awk({"x": p4.x, "y": p4.y, "z": p4.z, "t": energy})
lz_sum = ak.sum(p4_lz, axis=-1)
jet_p4 = vector.awk({
    "x": lz_sum.x,
    "y": lz_sum.y,
    "z": lz_sum.z,
    "t": lz_sum.t
})
jet_energy = jet_p4.t
jet_mass = jet_p4.tau
jet_phi = jet_p4.phi
jet_pt = jet_p4.rho

Answer 1

对于对平面洛伦兹向量数组和锯齿状洛伦兹向量数组同样有效的解决方案，试试这个：

import uproot
import awkward as ak
import vector
vector.register_awkward()   # any record named "Momentum4D" will be Lorentz

with uproot.open(input_path) as input_file:
    tree = input_file["Jets"]
    arrays = tree.arrays(filter_name="Constituent_*")
    p4 = ak.zip({
        "pt": arrays.Constituent_pt,
        "phi": arrays.Constituent_phi,
        "eta": arrays.Constituent_eta,
        "energy": arrays.Constituent_energy,
    }, with_name="Momentum4D")
    jet_p4 = ak.zip({
        "px": ak.sum(p4.px, axis=-1),
        "py": ak.sum(p4.py, axis=-1),
        "pz": ak.sum(p4.pz, axis=-1),
        "energy": ak.sum(p4.energy, axis=-1)
    }, with_name="Momentum4D")

注意 uproot.TTree.arrays 函数，如果没有参数，将读取 TTree 中的所有 TBranches。在您的函数中，您读取了所有数据四次，每次都从已读取的数据中选择一个不同的列，然后将其余数据丢弃。

此外，我不喜欢 vector.awk 函数，因为它可以构造类型数组：

N * Momentum4D[px: var * float64, py: var * float64, pz: var * float64, E: var * float64]

（换句话说，每个“px”值是一个 浮点列表 ），而不是你想要的：

N * var * Momentum4D[px: float64, py: float64, pz: float64, E: float64]

ak.zip 组合列表，以便每个洛伦兹向量的“px”只是一个数字，但您可以嵌套洛伦兹向量列表。如果你有参差不齐的数组，这只会有所不同，但我会指出来，这样就不会有人落入这个陷阱。

with_name="Momentum4D" 参数用该名称标记记录，并使用 vector.register_awkward() 注册 Lorentz-vector 行为会给出所有此类记录的洛伦兹向量方法。在这种情况下，我们使用它，以便根据 pt、phi、eta、energy 定义的 p4 具有属性 px, py, pz——换句话说，按需进行坐标变换。

没有一种洛伦兹向量求和方法可以对锯齿状数组中的每一项求和（uproot-methods 一个是仅适用于锯齿状洛伦兹向量数组的 hack，不适用于其他结构，例如 jagged-jagged, 等), 所以在笛卡尔坐标系中对 ak.sum 的分量求和.

计算射流不变质量时 uproot4/vector 中的 TLorentz 矢量特征

TLorentz vector features in uproot4/vector when calculating invariant mass of a jet

vector

uproot

awkward-array