如何像uproot.iterate()一样快速输出root_numpy root2array()

Question

array2root返回包含分支名称的dtype的元组列表。有没有一种方法可以从uproot.iterate()中返回相同类型的格式，而不需要在事后对其进行昂贵的整形？

输出应与

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)，那么您可以预先分配数组并将其传递给根，并将其连根填充，而不是创建新数组。在您的例子中，您可以使它成为一个记录数组：

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，它将用于解释根文件中的原始数据(大端，因此是">")，第二个参数是要将数据读入的数组。

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正在用名为arrays的dict填充和发送数组，但它们是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这样的大数组重新分配。但是，假设您需要列中的数据:发送到for循环主体的arrays[0]和arrays[0]。在上面，我们还希望查看格式化为记录数组(“元组数组”)的数据，因此我们实际上是在查看这个称为buffer的“倾倒场”。为了明智地做到这一点--避免与这个子范围无关的条目--我们必须显式地删除它们，而且库中没有任何函数来确定子范围的位置。然而，这

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

将是这一职能的普遍执行。