在 DataFrames 中保存 TLorentzVector 信息以供将来分析

Question

我想知道将 .root 文件中的 TLorentzVector 信息转换为 Pandas DataFrame 的推荐协议。到目前为止，我的策略是为我关心的每个粒子保存 pT、eta 和 phi 信息。然后我编写自己的函数（基于 TLorentzVector 定义）来计算我可能偶尔需要的任何其他量，例如 DeltaR、mT 等。

然后我想知道我是否可以只将 TLorentzVector 保存到我的 DataFrame 并使用 uproot 来即时获取 pT、eta、phi 等数量，使用类似这样的东西（当我运行在我刚刚从 .root 文件转换的 DataFrame 上：

for row in df.index:
    print(df.at[row,"leptons_p4_0"].pt)

不过，我很快意识到，Pandas 本身并不能理解 TLorentzVector 是什么，所以当我稍后使用 pd.read_csv 重新加载文件时，这不起作用。

那么我的问题是，其他人如何建议我将 TLorentzVector 信息保存在稍后将在 pandas 中打开的 DataFrame 中，而不是连根拔起？似乎我的选择是要么为每个粒子保存 (pT, eta, phi) 列然后编写我自己的函数，要么保存 TLorentzVector 组件 (E, px, py, pz) 并使用 uproot_methods 来每次我重新加载 DataFrame 时，将这些组件转换回 TLorentzVector。或者，希望有另一个我还没有遇到的更简单的解决方案！

非常感谢任何建议。

Answer 1

由于 Pandas 没有任何处理洛伦兹向量的工具，用它们的分量（pT、eta、phi、质量）来表达它们并编写自己的函数来转换它们是唯一的方法去，特别是如果你想保存到 CSV 和从 CSV 保存。

也就是说，可以创建在 Pandas 内保留其 "Lorentziness" 的洛伦兹矢量对象，但存在限制。您可以将结构化数据创建为 Awkward Arrays:

>>> import awkward1 as ak
>>> import pandas as pd
>>> import numpy as np
>>> class Lorentz:
...     @property
...     def p(self):
...         return self.pt * np.cosh(self.eta)
... 
>>> class LorentzRecord(Lorentz, ak.Record): pass
... 
>>> class LorentzArray(Lorentz, ak.Array): pass
... 
>>> ak.behavior["lorentz"] = LorentzRecord
>>> ak.behavior["*", "lorentz"] = LorentzArray
>>> array = ak.Array([{"pt": 1.1, "eta": 2.2},
...                   {"pt": 3.3, "eta": 4.4},
...                   {"pt": 5.5, "eta": -2.2}],
...                  with_name="lorentz")
>>> array
<LorentzArray [{pt: 1.1, eta: 2.2}, ... eta: -2.2}] type='3 * lorentz["pt": floa...'>

上面定义了一个 array 字段 pt 和 eta 的记录，并为单记录和记录数组视图提供了一个新的属性 p，由pt和eta派生而来。

>>> # Each record has a pt, eta, and p.
>>> array[0].pt
1.1
>>> array[0].eta
2.2
>>> array[0].p
5.024699161788051
>>> # The whole array has a pt, eta, and p (columns).
>>> array.pt
<Array [1.1, 3.3, 5.5] type='3 * float64'>
>>> array.eta
<Array [2.2, 4.4, -2.2] type='3 * float64'>
>>> array.p
<Array [5.02, 134, 25.1] type='3 * float64'>

您可以将一组洛伦兹记录放入 Pandas DataFrame:

>>> df = pd.DataFrame({"column": array})
>>> df
                 column
0   {pt: 1.1, eta: 2.2}
1   {pt: 3.3, eta: 4.4}
2  {pt: 5.5, eta: -2.2}

并用它做同样的事情：

>>> df.column.values.pt
<Array [1.1, 3.3, 5.5] type='3 * float64'>
>>> df.column.values.eta
<Array [2.2, 4.4, -2.2] type='3 * float64'>
>>> df.column.values.p
<Array [5.02, 134, 25.1] type='3 * float64'>

但那是因为我们将 Awkward Array 拉回以应用这些操作。

>>> df.column.values
<LorentzArray [{pt: 1.1, eta: 2.2}, ... eta: -2.2}] type='3 * lorentz["pt": floa...'>

应用于 DataFrame 的任何 NumPy 函数，例如取反（隐式调用 np.negative），无需解包即可传递给 Awkward Array。

>>> -df
                  column
0  {pt: -1.1, eta: -2.2}
1  {pt: -3.3, eta: -4.4}
2   {pt: -5.5, eta: 2.2}

但目前是错误的操作：它不应该否定pt。可以进一步重载 that:

>>> def negative_Lorentz(x):
...     return ak.zip({"pt": x.pt, "eta": -x.eta})
... 
>>> ak.behavior[np.negative, "lorentz"] = negative_Lorentz
>>> -df
                 column
0  {pt: 1.1, eta: -2.2}
1  {pt: 3.3, eta: -4.4}
2   {pt: 5.5, eta: 2.2}

我们仍在为洛伦兹数组构建一套函数，但现在它们以 Pandas 运行的一次数组模式工作。有一个名为 [=30= 的项目] 为 2D、3D 和洛伦兹向量定义所有这些函数，但它处于开发的早期阶段。

回到保存的问题——以上所有内容对您没有帮助，因为 Pandas "saves" 这些数据通过打印出来：

>>> df.to_csv("whatever.csv")

写

,column
0,"{pt: 1.1, eta: 2.2}"
1,"{pt: 3.3, eta: 4.4}"
2,"{pt: 5.5, eta: -2.2}"

这不是可以回读的东西。我们可以试试,

>>> df2 = pd.read_csv("whatever.csv")
>>> df2
   Unnamed: 0                column
0           0   {pt: 1.1, eta: 2.2}
1           1   {pt: 3.3, eta: 4.4}
2           2  {pt: 5.5, eta: -2.2}
>>> df2.column.values
array(['{pt: 1.1, eta: 2.2}', '{pt: 3.3, eta: 4.4}',
       '{pt: 5.5, eta: -2.2}'], dtype=object)

到目前为止，它看起来不错，但并不好：

>>> df2.column.values
array(['{pt: 1.1, eta: 2.2}', '{pt: 3.3, eta: 4.4}',
       '{pt: 5.5, eta: -2.2}'], dtype=object)

它们是 字符串 。它们不再是可计算的。所以如果你想保存到文件，把它分解成组件。

也许所有这些都可以整合到一个可用的系统中，但有些方面，比如保存这些数组 "Lorentizness" 完好无损，还没有准备好。

在 DataFrames 中保存 TLorentzVector 信息以供将来分析

Saving TLorentzVector info in DataFrames for future analysis

uproot