书写树数,篮数和压缩(连根拔起)



我试图优化树在pyroot中编写的方式,并遇到了uproot。最后,我的应用程序应该将事件(由数组组成)写入不断传入的树。

第一种方法是经典的方法:

event= [1.,2.,3.]
f = ROOT.TFile("my_tree.root", "RECREATE")
tree = ROOT.TTree("tree", "An Example Tree")
pt = array.array('f', [0.]*3)

tree.Branch("pt", pt, "pt[3]/F")
#for loop to simulate incoming events
for _ in range(10000):
for i, element in enumerate(event):
pt[i] = element
tree.Fill()
tree.Print()
tree.Write("", ROOT.TObject.kOverwrite);
f.Close()

这给出了以下树和执行时间:

树对策我的代码看起来是这样的:

np_array = np.array([[1,2,3]])
ak_array = ak.from_numpy(np_array)
with uproot.recreate("testing.root", compression=None) as fout:
fout.mktree("tree", {"branch": ak_array.type})

for _ in range(10000):

fout["tree"].extend({"branch": ak_array})

给出如下树:

树特征所以uproot方法需要更长的时间,文件大小更大,每个事件都有一个单独的篮子。我尝试了不同的压缩设置,但没有改变任何东西。有什么优化方法吗?这是uproot的一个合理的用途吗?与第一种方法相比,写树的过程能加快吗?

extend方法假定在每次调用时写入一个新的TBasket。(请参阅文档,特别是橙色警告框。这样做的目的是为了控制TBasket的大小。如果你调用它10000次来写一个值(值[1, 2, 3]),这是一个最大的效率低下的使用。

从根本上说,您是以逐条目的方式来考虑这个问题的,而不是按列来考虑,这是在Python中通常进行科学处理的方式。相反,您需要做的是在内存中收集一个大型数据集,并将其写入一个块中的文件。如果你最终要寻址的数据比你电脑的内存还大,你会用"足够大"来寻址。块,可能在几百兆字节或千兆字节的量级。

例如,从你的例子开始,

import time
import uproot
import numpy as np
import awkward as ak
np_array = np.array([[1, 2, 3]])
ak_array = ak.from_numpy(np_array)
starttime = time.time()
with uproot.recreate("bad.root") as fout:
fout.mktree("tree", {"branch": ak_array.type})
for _ in range(10000):
fout["tree"].extend({"branch": ak_array})
print("Total time:", time.time() - starttime)

总时间(在我的计算机上)是1.9秒,tree的特征是可怕的:

******************************************************************************
*Tree    :tree      :                                                        *
*Entries :    10000 : Total =         1170660 bytes  File  Size =    2970640 *
*        :          : Tree compression factor =   1.00                       *
******************************************************************************
*Br    0 :branch    : branch[3]/L                                            *
*Entries :    10000 : Total  Size=    1170323 bytes  File Size  =     970000 *
*Baskets :    10000 : Basket Size=      32000 bytes  Compression=   1.00     *
*............................................................................*

相反,我们希望数据在单个数组中(或一些循环产生~GB规模的数组):

np_array = np.array([[1, 2, 3]] * 10000)

(这不一定是您得到np_array的方式,因为* 10000生成了一个大型的中间Python列表。可以这么说,您以某种方式获得了数据。

现在,我们通过对extend的单个调用来执行写入操作,这将生成单个TBasket:

np_array = np.array([[1, 2, 3]] * 10000)
ak_array = ak.from_numpy(np_array)
starttime = time.time()
with uproot.recreate("good.root") as fout:
fout.mktree("tree", {"branch": ak_array.type})
fout["tree"].extend({"branch": ak_array})
print("Total time:", time.time() - starttime)

总时间(在我的计算机上)为0.0020秒,tree特性要好得多:

******************************************************************************
*Tree    :tree      :                                                        *
*Entries :    10000 : Total =          240913 bytes  File  Size =       3069 *
*        :          : Tree compression factor = 107.70                       *
******************************************************************************
*Br    0 :branch    : branch[3]/L                                            *
*Entries :    10000 : Total  Size=     240576 bytes  File Size  =       2229 *
*Baskets :        1 : Basket Size=      32000 bytes  Compression= 107.70     *
*............................................................................*

因此,写入速度几乎提高了1000倍,压缩速度提高了100倍。(在前面的例子中,每个TBasket有一个条目,没有压缩,因为任何压缩的数据都会比原始数据大!)

相比之下,如果我们使用PyROOT逐项写入,

import time
import array
import ROOT
data = [1, 2, 3]
holder = array.array("q", [0]*3)
file = ROOT.TFile("pyroot.root", "RECREATE")
tree = ROOT.TTree("tree", "An Example Tree")
tree.Branch("branch", holder, "branch[3]/L")
starttime = time.time()
for _ in range(10000):
for i, x in enumerate(data):
holder[i] = x
tree.Fill()
tree.Write("", ROOT.TObject.kOverwrite)
file.Close()
print("Total time:", time.time() - starttime)

总时间(在我的计算机上)是0.062秒,tree特性很好:

******************************************************************************
*Tree    :tree      : An Example Tree                                        *
*Entries :    10000 : Total =          241446 bytes  File  Size =       3521 *
*        :          : Tree compression factor =  78.01                       *
******************************************************************************
*Br    0 :branch    : branch[3]/L                                            *
*Entries :    10000 : Total  Size=     241087 bytes  File Size  =       3084 *
*Baskets :        8 : Basket Size=      32000 bytes  Compression=  78.01     *
*............................................................................*

因此,PyROOT在这里要慢30倍,但压缩效果几乎一样好。ROOT决定创建8个tbasket,可以通过AutoFlush参数进行配置。

但是请记住,这是对技术的比较,而不是库。如果你用RDataFrame包装一个NumPy数组并写入它,那么你可以跳过Python for循环所涉及的所有开销,并获得列处理的优势。

但是列式处理只有在处理大数据时才有用。就像压缩一样,如果你将它应用于非常小的数据集(或多次应用于非常小的数据集),那么它可能会带来伤害,而不是帮助。

最新更新