Как быстро получить вывод uproot.iterate (), например, root_numpy root2array () - PullRequest
1 голос
/ 02 ноября 2019

array2root возвращает список кортежей с dtype, содержащим имена ветвей. Есть ли способ вернуть тот же тип формата из uproot.iterate() без дорогостоящего изменения его потом?

Вывод должен быть таким же, как и из

array = root2array(['file.root'], treename = 'tree', branches = ['pt', 'eta'])

, который выглядит как np.array([(pt0, eta0), (pt1, eta1), ... dtype=[('pt', '<f4'), ('eta', '<f4')]]

1 Ответ

0 голосов
/ 02 ноября 2019

Если у вас есть верхняя граница размера массива (т.е. вы получаете его от iterate, поэтому вы можете передать entrysteps=10000 и знать, что он никогда не будет больше 10000), тогдаВы можете предварительно выделить свой массив и передать его на выкорчевывание и заполнить выкорчеванным вместо того, чтобы создавать новые массивы. В вашем случае вы можете сделать это массивом записей:

buffer = numpy.empty(20000, dtype=[("pt", "f8"), ("eta", "f8")])
pt_buffer = buffer["pt"]
eta_buffer = buffer["eta"]

pt_buffer и eta_buffer - это представления buffer, которые, как оказалось, чередуются, но работают так же хорошо, какмассивы. (Причина, по которой я выделил 20000, а не просто 10000, будет объяснена ниже.)

Теперь скажите, что вас интересуют две ветви, по умолчанию interpretation которых является uproot.asdtype(">f8", "f8"). Запросите эти массивы с интерпретацией uproot.asarray(">f8", pt_buffer) и uproot.asarray(">f8", eta_buffer). Первый аргумент - это Numpy dtype, который будет использоваться для интерпретации необработанных данных из файла ROOT (big-endian, следовательно, ">"), а второй аргумент - это массив, в который вы собираетесь читать данные, вplace.

for arrays in tree.iterate({"pt": uproot.asarray(">f8", pt_buffer),
                            "eta": uproot.asarray(">f8", eta_buffer)},
                           outputtype=tuple, entrysteps=10000):
    start = int((arrays[0].ctypes.data - buffer.ctypes.data) / buffer.itemsize)
    stop = start + len(arrays[0])
    array_of_tuples = buffer[start:stop]
    print(array_of_tuples)

См. документацию об этой редко используемой и мало разрекламированной функции.

Даже если iterate заполняет и отправляет вам массивы в формате dictони называются arrays, они представляют собой столбцы представлений массива записей buffer («массив кортежей»). Посмотрев на исходную buffer, мы увидим нужную вам структуру.

Однако выкорчевывание фактически заполняет buffer содержимым целой корзины, начиная с начала первой соответствующей корзины и заканчивая наконец последней релевантной корзины, охватывающей каждый поддиапазон: [0, 10000), [10000, 20000), [20000, 30000) и т. д. Поэтому нужная вам часть buffer может начать несколько записей в (start != 0) и, вероятно, закончится раньше20000 (stop - start != len(buffer)). Поскольку arrays[0] - это представление первого столбца в buffer, содержащего только те записи, которые вы хотите, разница между arrays[0].ctypes.data и buffer.ctypes.data - это количество байтов в buffer, которое вы хотите. Деление на buffer.itemsize дает количество записей. Конечную позицию легче вычислить.

Предварительное распределение buffer должно быть достаточно большим, чтобы включать все записи, которые вы хотите, и любые дополнительные записи, которые идут вместе с корзиной и должны быть обрезаны. 20000 является безопасным, если корзина не превышает 10000. Для данного tree вы можете определить наибольшее количество записей в любой корзине любой ветви с помощью:

max(branch.basket_numentries(i) for branch in tree.values()
                                for i in range(branch.numbaskets))

Очевидно, что это не то, для чего были предназначены эти функции: asarray предназначался для производительности, чтобы избежать перераспределения больших массивов, таких как buffer. Однако предполагалось, что вы захотите получить данные в столбцах: arrays[0] и arrays[0], отправленные в тело цикла for. Выше мы дополнительно хотим посмотреть на данные, отформатированные как массив записей («массив кортежей»), поэтому мы на самом деле смотрим на эту «площадку сброса», известную как buffer. Чтобы сделать это разумно - избегая записей, не относящихся к этому поддиапазону - мы должны явно вырезать их, и в библиотеке не было никаких функций, чтобы выяснить, где находится этот поддиапазон. Однако эта

    start = int((arrays[0].ctypes.data - buffer.ctypes.data) / buffer.itemsize)
    stop = start + len(arrays[0])
    array_of_tuples = buffer[start:stop]

будет общей реализацией такой функции.

...