我决定给出一个完整的工作示例(跟进上面的 cmets),但发现有很多选择要做。你只想复制输入的 TTree——你不想做出选择——所以你真的想要一个高级的“复制整个 TTree”功能,但这样的功能并不存在。 (这将是对 Uproot 的一个很好的补充,或者是一个使用 Uproot 来完成hadd 类型工作的新模块。如果有人感兴趣的话,这是一个很好的项目!)
我从 this file 开始,可以通过多种方式获得:
file_path = "root://eospublic.cern.ch//eos/opendata/cms/derived-data/AOD2NanoAODOutreachTool/Run2012BC_DoubleMuParked_Muons.root"
file_path = "http://opendata.cern.ch/record/12341/files/Run2012BC_DoubleMuParked_Muons.root"
file_path = "/tmp/Run2012BC_DoubleMuParked_Muons.root"
它足够大,应该分块复制,而不是一次全部复制。第一个块设置类型,因此可以通过将新分支名称分配给数组来执行,但后续块必须调用WritableFile.extend,因为你不想替换新的 TTree,你想添加它。这些都没有明确处理类型;类型是从数组中提取的。
这是第一次尝试,使用 "100 MB" 作为块大小。 (这将是输出 TTree 中跨 TBranches 的 TBasket 大小的总和。我们在这里所做的不仅仅是复制;它是将数据重新分区为新的块大小。)
with uproot.recreate("/tmp/output.root") as output_file:
first_chunk = True
with uproot.open(file_path) as input_file:
input_ttree = input_file["Events"]
for arrays_chunk in input_ttree.iterate(step_size="100 MB"):
if first_chunk:
output_file["Events"] = arrays_chunk
first_chunk = False
else:
output_file["Events"].extend(arrays_chunk)
但是,它失败了,因为赋值和 extend 期望数组字典,而不是单个数组。
所以我们可以请求 TTree.iterate 给我们一个 Awkward Arrays 的字典,每个 TBranch 一个,而不是一个代表所有 TBranches 的 Awkward Array。看起来像这样:
with uproot.recreate("/tmp/output.root") as output_file:
first_chunk = True
with uproot.open(file_path) as input_file:
input_ttree = input_file["Events"]
for dict_of_arrays in input_ttree.iterate(step_size="100 MB", how=dict):
if first_chunk:
output_file["Events"] = dict_of_arrays
first_chunk = False
else:
output_file["Events"].extend(dict_of_arrays)
它复制文件,但是原始文件有 TBranches 之类的
name | typename | interpretation
---------------------+--------------------------+-------------------------------
nMuon | uint32_t | AsDtype('>u4')
Muon_pt | float[] | AsJagged(AsDtype('>f4'))
Muon_eta | float[] | AsJagged(AsDtype('>f4'))
Muon_phi | float[] | AsJagged(AsDtype('>f4'))
Muon_mass | float[] | AsJagged(AsDtype('>f4'))
Muon_charge | int32_t[] | AsJagged(AsDtype('>i4'))
新文件有 TBranches 之类的
name | typename | interpretation
---------------------+--------------------------+-------------------------------
nMuon | uint32_t | AsDtype('>u4')
nMuon_pt | int32_t | AsDtype('>i4')
Muon_pt | float[] | AsJagged(AsDtype('>f4'))
nMuon_eta | int32_t | AsDtype('>i4')
Muon_eta | float[] | AsJagged(AsDtype('>f4'))
nMuon_phi | int32_t | AsDtype('>i4')
Muon_phi | float[] | AsJagged(AsDtype('>f4'))
nMuon_mass | int32_t | AsDtype('>i4')
Muon_mass | float[] | AsJagged(AsDtype('>f4'))
nMuon_charge | int32_t | AsDtype('>i4')
Muon_charge | int32_t[] | AsJagged(AsDtype('>i4'))
发生的事情是 Uproot 不知道每个尴尬数组的每个条目都有相同数量的项目(一个事件中 pt 值的数量与一个事件中 eta 值的数量相同) .如果 T 分支不是全部都是 μ 介子,而是有些是 μ 介子,有些是电子或喷流,那就不是真的了。
这些 nMuon_pt、nMuon_eta 等 TBranches 被包含在内的原因是因为 ROOT 需要它们。 Muon_pt、Muon_eta 等。TBranches 在 ROOT 中被读取为可变长度的 C++ 数组,C++ 用户需要知道预分配一个数组有多大,以及在哪个数组条目之后内容是未初始化的垃圾。 Python 中不需要这些(Awkward Array 会阻止用户看到未初始化的垃圾)。
所以你可以忽略它们。但是如果你真的需要/想要摆脱它们,这里有一个方法:构建你想要写的数组。现在我们正在处理类型,我们将使用WritableDirectory.mktree 并明确指定类型。由于每次写入都是extend,我们将不必再跟踪我们是否正在写入first_chunk 或后续块。
对于Muon_pt、Muon_eta等TBranch共享一个计数器TBranch,nMuons,你希望Muon字段是一个可变长度的μ子对象列表数组,pt,eta等领域。该类型可以从字符串构造:
import awkward as ak
muons_type = ak.types.from_datashape("""var * {
pt: float32,
eta: float32,
phi: float32,
mass: float32,
charge: int32
}""", highlevel=False)
给定 chunk 类型为 var * float32 的分隔数组,您可以创建类型为 var * {pt: float32, eta: float32, ...} 和 ak.zip 的单个数组。
muons = ak.zip({
"pt": chunk["Muon_pt"],
"eta": chunk["Muon_eta"],
"phi": chunk["Muon_phi"],
"mass": chunk["Muon_mass"],
"charge": chunk["Muon_charge"],
})
(打印muons.type 返回类型字符串。)这是您可能用于数据分析的表单。假设用户会将数据作为读取和写入之间的对象进行分析,而不是从一个文件中读取并在不进行任何修改的情况下写入另一个文件。
这是一个读写器,使用muons_type:
with uproot.recreate("/tmp/output.root") as output_file:
output_ttree = output_file.mktree("Events", {"Muon": muons_type})
with uproot.open(file_path) as input_file:
input_ttree = input_file["Events"]
for chunk in input_ttree.iterate(step_size="100 MB"):
muons = ak.zip({
"pt": chunk["Muon_pt"],
"eta": chunk["Muon_eta"],
"phi": chunk["Muon_phi"],
"mass": chunk["Muon_mass"],
"charge": chunk["Muon_charge"],
})
output_ttree.extend({"Muon": muons})
或者您可以通过再次跟踪 first_chunk 来显式构造 muons_type 来完成它:
with uproot.recreate("/tmp/output.root") as output_file:
first_chunk = True
with uproot.open(file_path) as input_file:
input_ttree = input_file["Events"]
for chunk in input_ttree.iterate(step_size="100 MB"):
muons = ak.zip({
"pt": chunk["Muon_pt"],
"eta": chunk["Muon_eta"],
"phi": chunk["Muon_phi"],
"mass": chunk["Muon_mass"],
"charge": chunk["Muon_charge"],
})
if first_chunk:
output_file["Events"] = {"Muon": muons}
first_chunk = False
else:
output_file["Events"].extend({"Muon": muons})
诚然,它很复杂(因为我展示了很多备选方案,各有优缺点),但那是因为不加修改地复制 TTree 并不是 TTree 编写函数的可预见用例。由于这是一个重要的用例,因此隐藏这些细节的专门功能将是一个受欢迎的补充。