如何获得与锯齿状数组的root_numpy root2array()输出相同的root.iterate()输出



我知道类似的问题有一个解决方案:如何快速获得root_numpy root2array((输出那样的root.iterate((输出但据我所知,它只适用于平坦的ROOT树。我想有一个通用的解决方案:

  1. 固定大小的维度但嵌套数据,如粒子动量(px,py,pz(,这些数据在ROOT TTree中表示为vector<double>
  2. 任意大小的维度数据

我为其应用asjagged的所有尝试都失败了。对于情况(1(,是否可以避免jaggedarray

如果数据是固定大小的,但存储为vector<double>,则将其视为非固定大小。Uproot总是将它们读取为锯齿状数组,因此另一个问题中描述的asarray方法不可用。

也就是说,如果你比文件的元数据了解更多,并且愿意尝试不受支持的破解,你可以强制对vector<double>的解释始终包含三个元素。这里有一个例子——首先,我们需要一个合适的ROOT文件:

TFile *f = new TFile("tmp.root", "recreate");
TTree *t = new TTree("t", "t");
std::vector<double> x;
t->Branch("x", &x);
x = {101, 102, 103};
t->Fill();
x = {201, 202, 203};
t->Fill();
x = {301, 302, 303};
t->Fill();
x = {401, 402, 403};
t->Fill();
x = {501, 502, 503};
t->Fill();
x = {601, 602, 603};
t->Fill();
t->Write();
f->Close();

现在,在Uproot中读取此内容的正常方法是使用JaggedArray:

>>> import uproot
>>> branch = uproot.open("tmp.root")["t"]["x"]
>>> branch.array()
<JaggedArray [[101.0 102.0 103.0] [201.0 202.0 203.0] [301.0 302.0 303.0] [401.0 402.0 403.0] [501.0 502.0 503.0] [601.0 602.0 603.0]] at 0x7f325ab4eb90>
>>> branch.interpretation
asjagged(asdtype('>f8'), 10)

但这会在每次读取时分配新的数组,并进行一些操作以找到边界,您知道边界是规则的。(Uproot并不知道这一点。(

STL向量的标头恰好有10个字节长。之后,它的内容按顺序进行序列化,从第一个到最后一个,然后再转到下一个STL向量。对于三个8字节的浮点数字,即10+8+8=34字节,大小始终相同。事实上,它总是碰巧是相同的大小,这对以下内容至关重要。

NumPy结构化数组可以将固定大小结构的数组表示为dtype:

>>> array = branch.array(uproot.asdtype([("header", "S10"), ("x", ">f8"), ("y", ">f8"), ("z", ">f8")]))
>>> array
array([(b'@x00x00x1ex00tx00x00x00x03', 101., 102., 103.),
(b'@x00x00x1ex00tx00x00x00x03', 201., 202., 203.),
(b'@x00x00x1ex00tx00x00x00x03', 301., 302., 303.),
(b'@x00x00x1ex00tx00x00x00x03', 401., 402., 403.),
(b'@x00x00x1ex00tx00x00x00x03', 501., 502., 503.),
(b'@x00x00x1ex00tx00x00x00x03', 601., 602., 603.)],
dtype=[('header', 'S10'), ('x', '<f8'), ('y', '<f8'), ('z', '<f8')])

我们不喜欢看标题,所以我们可以使用普通的NumPy切片将其剥离:

>>> array[["x", "y", "z"]]
array([(101., 102., 103.), (201., 202., 203.), (301., 302., 303.),
(401., 402., 403.), (501., 502., 503.), (601., 602., 603.)],
dtype={'names':['x','y','z'], 'formats':['<f8','<f8','<f8'], 'offsets':[10,18,26], 'itemsize':34})

(或者只是请求array["x"]以获得非结构化数组(。因为我们可以用普通的uproot.asdtype来实现,所以我们可以用预分配的uproot.asarray来实现。

>>> import numpy as np
>>> big = np.empty(1000, dtype=[("header", "S10"), ("x", ">f8"), ("y", ">f8"), ("z", ">f8")])
>>> branch.array(uproot.asarray(big.dtype, big))
array([(b'@x00x00x1ex00tx00x00x00x03', 101., 102., 103.),
(b'@x00x00x1ex00tx00x00x00x03', 201., 202., 203.),
(b'@x00x00x1ex00tx00x00x00x03', 301., 302., 303.),
(b'@x00x00x1ex00tx00x00x00x03', 401., 402., 403.),
(b'@x00x00x1ex00tx00x00x00x03', 501., 502., 503.),
(b'@x00x00x1ex00tx00x00x00x03', 601., 602., 603.)],
dtype=[('header', 'S10'), ('x', '>f8'), ('y', '>f8'), ('z', '>f8')])

上面,big必须至少与您想要读取的数据集一样大,并且读取它会返回该数组的修剪后的视图。它不分配新的数组——数据存储在big:中

>>> big[["x", "y", "z"]][:10]
array([( 1.01000000e+002,  1.02000000e+002,  1.03000000e+002),
( 2.01000000e+002,  2.02000000e+002,  2.03000000e+002),
( 3.01000000e+002,  3.02000000e+002,  3.03000000e+002),
( 4.01000000e+002,  4.02000000e+002,  4.03000000e+002),
( 5.01000000e+002,  5.02000000e+002,  5.03000000e+002),
( 6.01000000e+002,  6.02000000e+002,  6.03000000e+002),
( 1.22164945e-309,  5.26335088e-310,  1.22167067e-309),
(-5.70498984e+158,  5.97958175e+158, -5.97958175e+158),
(-4.92033505e+032, -4.92033505e+032, -4.92033505e+032),
( 3.77957352e+011,  3.77957221e+011,  3.77957320e+011)],
dtype={'names':['x','y','z'], 'formats':['>f8','>f8','>f8'], 'offsets':[10,18,26], 'itemsize':34})

读取的6个事件之外的所有内容都是未初始化的垃圾,所以我建议使用branch.array函数的修剪输出;这只是为了表明big正在被填充——您没有得到一个新的数组。

根据您的问题(2(,如果数据不规则(每个条目的字节数(,那么您就无法做到这一点。此外,请记住,上面的技术并没有得到官方支持:你必须知道你的数据是规则的,你必须知道STL向量有一个10字节的头,这不是我们希望大多数用户知道的。

最新更新