Tree models and object collectionsΒΆ

This example demonstrates how to define a tree model and collections of objects associated to sets of tree branches.

Python source code: model.py

print(__doc__)
from rootpy.tree import Tree, TreeModel, FloatCol, IntCol
from rootpy.io import root_open
from rootpy.vector import LorentzVector
from rootpy import stl
from random import gauss, randint


f = root_open("test.root", "recreate")

# define the model
class Point(TreeModel):
    x = FloatCol()
    y = FloatCol()
    z = FloatCol()

class Event(Point.prefix('a_'), Point.prefix('b_')):
    # a_x, a_y, a_z and b_x, b_y, b_z are implicitly included here
    # define vector branches
    col_x = stl.vector("float")
    col_y = stl.vector("float")
    col_z = stl.vector("float")
    col_n = IntCol()
    # a TLorentzVector
    p = LorentzVector
    i = IntCol()

tree = Tree("test", model=Event)

# fill the tree
for i in range(10):
    tree.a_x = gauss(.5, 1.)
    tree.a_y = gauss(.3, 2.)
    tree.a_z = gauss(13., 42.)

    tree.b_x = gauss(.5, 1.)
    tree.b_y = gauss(.3, 2.)
    tree.b_z = gauss(13., 42.)

    n = randint(1, 10)
    for j in range(n):
        tree.col_x.push_back(gauss(.5, 1.))
        tree.col_y.push_back(gauss(.3, 2.))
        tree.col_z.push_back(gauss(13., 42.))
    tree.col_n = n

    tree.p.SetPtEtaPhiM(gauss(.5, 1.),
                        gauss(.5, 1.),
                        gauss(.5, 1.),
                        gauss(.5, 1.))

    tree.i = i
    tree.fill(reset=True)
tree.write()

f.close()
f = root_open("test.root")

tree = f.test

# define objects by prefix:
tree.define_object(name='a', prefix='a_')
tree.define_object(name='b', prefix='b_')

# define a mixin class to add functionality to a tree object
class Particle(object):
    def who_is_your_daddy(self):
        print("You are!")

# define collections of objects by prefix
tree.define_collection(name='particles',
                       prefix='col_',
                       size='col_n',
                       mix=Particle)

# loop over "events" in tree
for event in tree:
    print("a.x: {0:f}".format(event.a.x))
    print("b.y: {0:f}".format(event.b.y))
    # loop over "particles" in current event
    for p in event.particles:
        print("p.x: {0:f}".format(p.x))
        p.who_is_your_daddy()
    print(event.p.Eta())

f.close()