如何快速获得 uproot.iterate() 类似 root_numpy root2array() 输出的输出

Question

array2root returns 一个包含分支名称的 dtype 的元组列表。有没有一种方法可以从 uproot.iterate() return 获得相同类型的格式，而无需在事后进行昂贵的重塑？

输出应与

相同

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

这就像 np.array([(pt0, eta0), (pt1, eta1), ... dtype=[('pt', '<f4'), ('eta', '<f4')]]

Answer 1

如果您对数组的大小有一个上限（即您从 iterate 获取它，那么您可以传递 entrysteps=10000 并且知道它永远不会大于10000), 然后你可以预分配你的数组并将它传递给 uproot 并让 uproot 填充它而不是创建新数组。在您的情况下，您可以将其设为记录数组：

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，因此是 ">"），第二个参数是您要将数据读入的数组，就地。

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)

请参阅 the documentation 了解此很少使用且未广泛宣传的功能。

即使 iterate 在名为 arrays 的字典中填充并发送数组给您，它们也是 buffer 记录数组 ("array of tuples") 的列视图.通过查看原始 buffer，我们看到了您想要的结构。

然而，uproot 实际上用整个篮子的内容填充 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 之间的差异是 [=] 中的字节数19=] 你想要的。除以 buffer.itemsize 得出条目数。结束位置更容易计算。

buffer 的预分配必须足够大，以包含您想要的所有条目以及篮子中附带的任何其他条目，需要被切断。如果没有篮子大于 10000，20000 是安全的。对于给定的 tree，您可以确定任何分支的任何篮子中的最大条目数：

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

显然，这不是这些函数的设计目的：asarray 是为了提高性能，以避免重新分配像 buffer 这样的大数组。然而，假设您希望列中的数据：arrays[0] 和 arrays[0] 发送到 for 循环的主体。在上面，我们另外想查看格式化为记录数组 ("array of tuples") 的数据，所以我们实际上正在查看这个 "dumping ground" 称为buffer。为了明智地做到这一点——避免与该子范围无关的条目——我们必须明确地将它们删除，并且库中没有任何函数可以确定该子范围的位置。然而，这

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

将是此类功能的一般实现。

如何快速获得 uproot.iterate() 类似 root_numpy root2array() 输出的输出

How to get uproot.iterate() output like the root_numpy root2array() output fast

uproot