Получить данные в Pandas - PullRequest
Новая
       10

Получить данные в Pandas

3 голосов
/ 07 февраля 2020

Я использую pandas и root для чтения данных из файла. root, и я получаю таблицу, подобную следующей:

enter image description here

Вышеупомянутая таблица составлена ​​со следующим кодом: 

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.

Если бы кто-то мог дать мне любой совет, совет и / или предложение, я бы быть очень благодарным.

Ответы [ 2 ]

2 голосов
/ 07 февраля 2020

Я сделал несколько фиктивных данных, так как вы не предоставили свои в любом легком формате. Я думаю, это то, что вы ищете. 

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)

возвращает: 

    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
1 голос
/ 27 февраля 2020

Для этого вы можете использовать интерфейс 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()

Первый будет быстрее, если впоследствии вам больше не понадобятся столбцы, иначе последний может быть лучше.

...