如何使uproot.iterate()输出与带锯齿数组的root_numpy root2array()输出相同

Question

我知道有一个类似问题的解决方案：numpy root2array()输出快，但据我所知，它只适用于平面根TTrees。我希望有一个广义的解决方案：

1. 固定尺寸的但是嵌套的数据，如粒子动量(px，py，pz)，它们在根TTree中表示为vector<double>。
2. 任意尺寸数据

我申请asjagged的所有尝试都失败了。是否有可能避免jaggedarray的案例(1)？

Answer 1

如果数据是固定大小的，但存储为vector<double>，那么它们就会被当作不是固定大小的。链接总是将它们读取为锯齿数组，因此在另一个问题中描述的asarray方法是不可用的。

也就是说，如果您比您的文件的元数据有更多的知识，并且愿意尝试一种不受支持的攻击，您可以强制您的vector<double>解释总是有三个元素。下面是一个示例--首先，我们需要一个合适的根文件：

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();

现在，用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)

但是，每次读取它时，它都会分配新的数组，并执行一些操作来查找边框，您知道这些边框是正常的。(连根都不知道。)

STL向量的头恰好有10字节长。在此之后，将其内容按顺序序列化，从第一个到最后，然后再转到下一个STL向量。对于三个8字节浮点数，即10 +8+8+8=34个字节，大小总是一样的.事实是，它总是相同的大小是至关重要的，以下。

NumPy结构阵列可以将固定大小的结构数组表示为dtype：

>>> array = branch.array(uproot.asdtype([("header", "S10"), ("x", ">f8"), ("y", ">f8"), ("z", ">f8")]))
>>> array
array([(b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 101., 102., 103.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 201., 202., 203.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 301., 302., 303.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 401., 402., 403.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 501., 502., 503.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 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'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 101., 102., 103.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 201., 202., 203.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 301., 302., 303.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 401., 402., 403.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 501., 502., 503.),
       (b'@\x00\x00\x1e\x00\t\x00\x00\x00\x03', 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字节的头，这不是我们期望大多数用户知道的事情。