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
相关产品推荐
相关产品推荐

