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

Question

我知道有类似问题的解决方案：但据我了解，它仅适用于扁平的 ROOT TTrees。我想要通用解决方案：

固定大小维度但嵌套数据，如粒子动量 (px, py, pz)，在 ROOT TTree 中表示为 vector<double>
任意尺寸维度数据

我为它申请 asjagged 的所有尝试都失败了。对于情况 (1) 是否可以避免 jaggedarray？

Answer 1

如果数据是 fixed-size 但存储为 vector<double>，那么它们将被视为不是 fixed-size。 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 向量的 header 正好是 10 个字节长。之后，其内容按顺序序列化，从头到尾，然后进入下一个 STL 向量。对于三个 8 字节浮点数，即 10 + 8 + 8 + 8 = 34 字节，大小始终相同。它始终恰好大小相同的事实对于以下内容至关重要。

A NumPy structured array 可以将 fixed-size 结构的数组表示为数据类型：

>>> 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')])

我们不喜欢看 header，所以我们可以使用普通的 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"] 得到一个 non-structured 数组）。因为我们可以用普通的 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 必须至少与您要读取的数据集一样大并读取它 returns 该数组的修剪 view .它不分配新数组——数据存储在 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 个字节 header，这不是我们希望大多数用户知道的事情.

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

How to get uproot.iterate() output identical to the root_numpy root2array() output with jagged arrays

uproot