You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

Uproot创建TTree时如何为多个数组分支复用同一计数器分支

问题说明

目标是输出与输入结构完全一致的TTree,输入TTree结构如下:

ttree.show(datatypes.keys())

name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
Weight               | float                    | AsDtype('>f4')
E_Beam               | float                    | AsDtype('>f4')
Px_Beam              | float                    | AsDtype('>f4')
Py_Beam              | float                    | AsDtype('>f4')
Pz_Beam              | float                    | AsDtype('>f4')
NumFinalState        | int32_t                  | AsDtype('>i4')
E_FinalState         | float[]                  | AsJagged(AsDtype('>f4'))
Px_FinalState        | float[]                  | AsJagged(AsDtype('>f4'))
Py_FinalState        | float[]                  | AsJagged(AsDtype('>f4'))
Pz_FinalState        | float[]                  | AsJagged(AsDtype('>f4'))

其中NumFinalState分支存储所有*_FinalState数组分支的元素个数,该对应关系在当前场景下固定成立。如果直接按默认方式创建树,uproot会为每个可变长度数组自动生成独立的n*开头的计数器分支,产生冗余:

outfile = uproot.recreate("myData_OUT.root")
datatypes = {"Weight": "float32", "E_Beam": "float32", "Px_Beam": "float32", "Py_Beam": "float32", "Pz_Beam": "float32", "NumFinalState": "int32", "E_FinalState": "var * float32", "Px_FinalState": "var * float32", "Py_FinalState": "var * float32", "Pz_FinalState": "var * float32"}
outfile.mktree("kin", datatypes)
outfile["kin"].show()

生成的冗余结构如下:

name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
Weight               | float                    | AsDtype('>f4')
E_Beam               | float                    | AsDtype('>f4')
Px_Beam              | float                    | AsDtype('>f4')
Py_Beam              | float                    | AsDtype('>f4')
Pz_Beam              | float                    | AsDtype('>f4')
NumFinalState        | int32_t                  | AsDtype('>i4')
nE_FinalState        | int32_t                  | AsDtype('>i4')
E_FinalState         | float[]                  | AsJagged(AsDtype('>f4'))
nPx_FinalState       | int32_t                  | AsDtype('>i4')
Px_FinalState        | float[]                  | AsJagged(AsDtype('>f4'))
nPy_FinalState       | int32_t                  | AsDtype('>i4')
Py_FinalState        | float[]                  | AsJagged(AsDtype('>f4'))
nPz_FinalState       | int32_t                  | AsDtype('>i4')
Pz_FinalState        | float[]                  | AsJagged(AsDtype('>f4'))

尝试通过mktree的counter_name参数将所有数组分支的计数器指定为NumFinalState时会触发报错,测试代码如下:

outfile = uproot.recreate("myData_OUT.root")
datatypes = {"Weight": "float32", "E_Beam": "float32", "Px_Beam": "float32", "Py_Beam": "float32", "Pz_Beam": "float32", "NumFinalState": "int32", "E_FinalState": "var * float32", "Px_FinalState": "var * float32", "Py_FinalState": "var * float32", "Pz_FinalState": "var * float32"}
def counter_name(in_str: str) -> str:
    if "FinalState" in in_str:
        return "NumFinalState"
    return f"n{in_str}"
outfile.mktree("kin", datatypes, counter_name=counter_name)

报错信息:

---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
/var/folders/td/j379rd296477649qvl1k8n180000gn/T/ipykernel_24226/1447106219.py in <module>
      5         return "NumFinalState"
      6     return f"n{in_str}"
----> 7 outfile.mktree("kin", datatypes, counter_name=counter_name)

/opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/uproot/writing/writable.py in mktree(self, name, branch_types, title, counter_name, field_name, initial_basket_capacity, resize_factor)
   1268             path,
   1269             directory._file,
-> 1270             directory._cascading.add_tree(
   1271                 directory._file.sink,
   1272                 treename,

/opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/uproot/writing/_cascade.py in add_tree(self, sink, name, title, branch_types, counter_name, field_name, initial_basket_capacity, resize_factor)
   1796             resize_factor,
   1797         )
-> 1798         tree.write_anew(sink)
   1799         return tree
   1800 

/opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/uproot/writing/_cascadetree.py in write_anew(self, sink)
   1114                 # reference to fLeafCount
   1115                 out.append(
-> 1116                     uproot.deserialization._read_object_any_format1.pack(
   1117                         datum["counter"]["tleaf_reference_number"]
   1118                     )

error: required argument is not an integer

报错原因很简单:用字符串简写定义分支类型配合counter_name回调时,uproot每处理一个可变长度数组,都会尝试独立创建对应名称的计数器分支,就算你返回的计数器名完全一样,它也不会复用已经建好的分支对象,最后就会因为拿不到计数器的TLeaf引用编号抛出参数类型错误。

解决方案

别用字符串简写定义可变长度分支类型,直接导入uproot内置的类型对象,显式让所有*_FinalState数组分支绑定同一个计数器分支实例就行,参考代码如下:

import uproot
from uproot.types import Float32, Int32, AsJagged

outfile = uproot.recreate("myData_OUT.root")
# 先实例化公共计数器分支,所有数组共享这个实例
shared_counter = Int32()

datatypes = {
    "Weight": Float32(),
    "E_Beam": Float32(),
    "Px_Beam": Float32(),
    "Py_Beam": Float32(),
    "Pz_Beam": Float32(),
    "NumFinalState": shared_counter,
    "E_FinalState": AsJagged(Float32(), counter=shared_counter),
    "Px_FinalState": AsJagged(Float32(), counter=shared_counter),
    "Py_FinalState": AsJagged(Float32(), counter=shared_counter),
    "Pz_FinalState": AsJagged(Float32(), counter=shared_counter),
}

outfile.mktree("kin", datatypes)
# 验证输出结构
outfile["kin"].show()

运行后生成的TTree结构和输入完全一致,不会生成冗余的n*计数器分支,文件紧凑性和原生ROOT生成的同结构文件一致。填充数据时只需要正常传入NumFinalState的值,以及对应长度的四个FinalState数组即可,uproot会自动校验数组长度和计数器值的匹配性,行为和原生ROOT TTree完全一致。


内容的提问来源于stack exchange,提问作者Nathaniel D. Hoffman

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.08.29 21:21:32