检索 Pandas 中的数据

Question

我正在使用 pandas 和 uproot 从 .root 文件中读取数据，我得到一个 table 如下所示：

前面提到的table是用下面的代码制作的：

fname = 'ZZ4lAnalysis_VBFH.root' 
key = 'ZZTree/candTree'
ttree = uproot.open(fname)[key]
branches = ['Z1Flav', 'Z2Flav', 'nCleanedJetsPt30', 'LepPt', 'LepLepId'] 
df = ttree.pandas.df(branches, flatten=False)

我需要在 LepPt 中找到最大值，一旦找到最大值，我还需要检索该最大值的 LepLepId。我没问题找到最大值：

Pt_l1 = [max(i) for i in df.LepPt]

通过这种方式，我得到了一个包含所有最大值的数组。但是，我必须根据 LepLepId 将这些值分开。所以我需要一个具有最大 LepPt 和 |LepLepId|=11 的数组，以及一个具有最大 LepPt 和 |LepLepId|=13 的数组。

如果有人能给我任何提示、建议and/or建议，我将不胜感激。

Answer 1

我制作了一些模拟数据，因为您没有以任何简单的格式提供您的数据。我想这就是您要找的。

import pandas as pd

df = pd.DataFrame.from_records(
    [   [[1,2,3], [4,5,6]],
        [[4,6,5], [7,8,9]]
    ],
    columns=['LepPt', 'LepLepld']
)

df['max_LepPt'] = [max(i) for i in df.LepPt]

def f(row):
    # get index position within list
    pos = row['LepPt'].index(row['max_LepPt']).tolist()
    return row['LepLepld'][pos]

df['same_index_LepLepld'] = df.apply(lambda x: f(x), axis=1)

returns:

    LepPt       LepLepld    max_LepPt   same_index_LepLepld
0   [1, 2, 3]   [4, 5, 6]   3           6
1   [4, 6, 5]   [7, 8, 9]   6           8

Answer 2

您可以为此使用 awkward.JaggedArray 接口（uproot 的依赖项之一），它允许您拥有不规则大小的数组。

为此，您需要稍微更改加载数据的方式，但它允许您使用与普通 numpy 数组相同的方法，即 argmax:

fname = 'ZZ4lAnalysis_VBFH.root' 
key = 'ZZTree/candTree'
ttree = uproot.open(fname)[key]
# branches = ['Z1Flav', 'Z2Flav', 'nCleanedJetsPt30', 'LepPt', 'LepLepId']
branches = ['LepPt', 'LepLepId']   # to save memory, only load what you need

# df = ttree.pandas.df(branches, flatten=False)
a = ttree.arrays(branches)    # use awkward array interface

max_pt_idx = a[b'LepPt'].argmax()
max_pt_lepton_id = a[b'LepLepld'][max_pt_idx].flatten()

这只是一个普通的 numpy 数组，如果需要，您可以将其分配给 pandas 数据框的列。它应该具有正确的维度和顺序。它也应该比使用内置 Python 函数更快。

请注意，键是字节串，而不是普通字符串，如果有没有轻子的事件，您将不得不采取一些额外的步骤（在这种情况下，flatten 将忽略那些空事件，破坏对齐）。

或者，您也可以在之后转换列：

import awkward

df = ttree.pandas.df(branches, flatten=False)

max_pt_idx = awkward.fromiter(df["LepPt"]).argmax()
lepton_id = awkward.fromiter(df["LepLepld"])
df["max_pt_lepton_id"] = lepton_id[max_pt_idx].flatten()

如果以后不再需要这些列，前者会更快，否则后者可能更好。

检索 Pandas 中的数据

Retrieve data in Pandas

python

physics

pandas

uproot