Store 2D array into TTree use python

ROOT
,
2

Dear zhangaw325,

Here is an example of how to store a 2D array in a TTree from Python. For the sake of illustrating the different options you have, one of the dimensions is a variable that corresponds to another branch of the tree, the other dimension is defined with a constant. Numpy is used to construct the multidimensional arrays in Python.

Here is the code that stores the array:

from ROOT import TFile, TTree
import numpy as np

f = TFile('example.root', 'recreate')
t = TTree('mytree', 'example tree')

n = np.array(2, dtype=np.int32)
t.Branch('mynum', n, 'mynum/I')
x = np.array([[1., 2., 3.], [4., 5., 6.]])
t.Branch('myarray', x, 'myarray[mynum][3]/D')

nentries = 25
for i in range(nentries):
   t.Fill()

f.Write()
f.Close()

Here’s the code that reads it back into a numpy array:

import ROOT
import numpy as np

tfile = ROOT.TFile("example.root")
ttree = tfile.mytree

nentries = 25
for i in range(nentries):
  ttree.GetEntry(i)
  print(ttree.mynum)
  print(np.reshape(ttree.myarray, (2,3)))

Best,
Enric