如何为awkward数组内Vector对象统一设置固定质量值并重算能量

我的问题围绕scikit-hep中的Vector模块展开，现有一个存储vector对象的awkward数组，需要将数组内所有矢量的质量设置为统一公共值。
对于单个vector对象，可以通过如下代码实现需求：

x = vector.obj(pt=2, eta=1.5, phi=1, energy=10)
y = x.from_rhophietatau(rho=x.rho, eta=x.eta, phi=x.phi, tau=20)

print(f"{x.mass:6.3f}  {x.pt}  {x.eta}  {x.phi}  {x.energy:6.2f}")
print(f"{y.mass:6.3f}  {y.pt}  {y.eta}  {y.phi}  {y.energy:6.2f}")

运行输出为：

8.824  2  1.5  1   10.00
20.000  2  1.5  1   20.55

但不清楚如何对存储矢量的awkward数组执行同类批量操作。
测试使用的公开数据可通过如下命令获取：

curl http://opendata.cern.ch/record/12361/files/SMHiggsToZZTo4L.root --output SMHiggsToZZTo4L.root

数据加载与处理代码如下：

import numpy as np
import matplotlib.pylab as plt

import uproot
import awkward as ak

import vector
vector.register_awkward()


infile = uproot.open("/tmp/SMHiggsToZZTo4L.root")

muon_branch_arrays = infile["Events"].arrays(filter_name="Muon_*")
electron_branch_arrays = infile["Events"].arrays(filter_name="Electron_*")

muons = ak.zip({
    "pt": muon_branch_arrays["Muon_pt"],
    "phi": muon_branch_arrays["Muon_phi"],
    "eta": muon_branch_arrays["Muon_eta"],
    "mass": muon_branch_arrays["Muon_mass"],
    "charge": muon_branch_arrays["Muon_charge"],
}, with_name="Momentum4D")

quads = ak.combinations(muons, 4)
mu1, mu2, mu3, mu4 = ak.unzip(quads)

p4 = mu1 + mu2 + mu3 + mu4

得到的p4对象类型为<class 'vector._backends.awkward_.MomentumArray4D'>，需要将p4内所有四动量对象的质量统一设置为固定值（例如125），后续将使用调整后的p4将各个muX对象boost到p4的质心系计算相关相对角度，需要确认该批量设置固定质量需求的实现方法。

完全可以实现。Vector对awkward后端的四动量数组做了接口对齐，数组对象和单个矢量对象的方法名、参数规则完全一致，传入的标量参数会自动按numpy/awkward的广播规则作用到所有元素，不需要手动写循环或者拆包重组数组。
你单个对象使用的from_rhophietatau方法可以直接在数组上调用，保持空间三动量分量不变，替换质量为目标值即可：

# 目标固定质量值
target_mass = 125
# 批量生成质量固定的新四动量数组，空间动量部分和原p4完全一致
p4_fixed_mass = p4.from_rhophietatau(
    rho=p4.rho,
    eta=p4.eta,
    phi=p4.phi,
    tau=target_mass
)

可以用如下代码验证设置结果：

# 检查所有元素质量是否符合目标值（浮点比较留数值误差余量）
print(ak.all(abs(p4_fixed_mass.mass - target_mass) < 1e-12))
# 输出True即为设置成功

补充说明 #

• 如果场景需要保持笛卡尔坐标系下的px/py/pz分量不变，替换为目标质量，可以改用from_xyzt方法，传入原数组的px/py/pz，再传入对应目标质量计算得到的能量值即可。做质心系boost的场景下，上面给出的圆柱坐标写法最简便。
• 后续的洛伦兹boost操作也支持批量调用，不需要逐事例循环，例如将mu1变换到修正后p4的质心系，直接写mu1_cm = mu1.boostCM_of(p4_fixed_mass)即可，返回的数组结构和原mu1完全一致。

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