Added a new convenience class to pytree module.

The PyTree class is a minial extension of ROOT's
TTree, which handles a lot of the boilerplate of
creating and writing to branches.

See the output_ntup.py program for example usage.

example/output_ntup.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+
+'''
+ Example program demonstrating the use of the
+ pytree module to easily create output ntuples,
+ with practically zero ROOT-boilerplate crap!
+'''
+
+from analysis_utils.pytree import PyTree
+import ROOT as r
+import numpy as np
+
+if __name__ == "__main__":
+	outfile = r.TFile('output_ntuple.root', 'recreate')
+
+	ntup = PyTree('example', 'example')
+
+	for i in xrange(5000):
+		# call reset on the PyTree; this will set all the 
+		# sclars to 0, and clear out any vector-valued
+		# branches.
+		ntup.reset()
+
+		# calculate some silly variables at random
+		uniform1 = np.random.rand()
+		uniform50 = np.random.rand()*50.
+		gauss = np.random.normal(0,1)
+
+		rlength = int(np.random.gamma(1, 5))
+		rvector_gamma2 = np.random.gamma(2, 25e3, rlength)
+		rvector_gamma3 = np.random.gamma(3, 15e3, rlength)
+
+		# now write them out. note that the branches
+		# are created on the fly as needed.
+		ntup.write_branch(i, 'event_number', int)
+		ntup.write_branch(uniform1, 'uniform1', float)
+		ntup.write_branch(uniform50, 'uniform50', int)
+		ntup.write_branch(gauss, 'gauss', float)
+		ntup.write_branch(rvector_gamma2, 'rvector_gamma2', [float])
+		ntup.write_branch(rvector_gamma3, 'rvector_gamma3', [float])
+		ntup.write_branch(rlength, 'rvector_n', int)
+
+		# For illustration, write one branch out later
+		# than the others. This branch will be
+		# backfilled with 0's for the first 20 events.
+		# Then it should be identical to `event_number`
+		if i>20:
+			ntup.write_branch(i, 'late_branch', int)
+
+		# call fill to commit the current row to file
+		ntup.Fill()
+
+	outfile.Write()
+

pytree.py

@@ -46,34 +46,96 @@ def setup_tree(name, treedef):
 
 	return t
 
+class PyTree(r.TTree):
+	def __init__(self, *args, **kwargs):
+		r.TTree.__init__(self, *args)
+
+		self.__branch_cache = {}
+
+	def reset(self):
+		for v, btype in self.__branch_cache.values():
+			# set the values to an appropriate
+			# state depending on thier type
+			if type(btype) == list:
+				# the pointer will be a vector type.
+				# call clear on it.
+				v.clear()
+			else:
+				# the pointer is to a numpy.array with
+				# a scalar value. set it to zero.
+				v[0] = 0
+
+	def write_branch(self, value, bname, btype):
+		try:
+			v, _ = self.__branch_cache[bname]
+		except KeyError:
+			# oops haven't made this branch yet.
+			v = bind_and_backfill(self, bname, btype)
+			self.__branch_cache[bname] = (v, btype)
+
+		# now we have a pointer to the branch memory.
+		# write the value depending on the type
+		if type(btype) == list:
+			# some vector type
+			if btype[0] == list:
+				# TODO!
+				raise TypeError('Unsupported type: %s' % btype)
+			# okay, just a normal vector.
+			# copy the input array with push_back
+			map(v.push_back, value)
+		else:
+			# some sclar type
+			v[0] = value
+
 typenames_long = { float: 'double', int: 'int' }
 typenames_short = { float: 'D', int: 'I' }
 
+def bind_and_backfill(t, bname, btype):
+	v = bind_any(t, bname, btype)
+
+	b = t.GetBranch(bname)
+	if t.GetEntries() > 0:
+		# tree has already started filling! backfill the new branch
+		# with blank values
+		for i in xrange(t.GetEntries()):
+			b.Fill()
+
+	return v
+
 def bind_any(t, bname, btype):
 	if not type(btype) == list:
-		bind_scalar(t, bname, btype)
+		v = bind_scalar(t, bname, btype)
 	elif len(btype) and type(btype[0]) == list:
 		if len(btype[0]):
-			bind_vector2(t, bname, btype[0][0])
+			v = bind_vector2(t, bname, btype[0][0])
 		else:
-			bind_vector2(t, bname, float)
+			v = bind_vector2(t, bname, float)
 	elif len(btype):
-		bind_vector(t, bname, btype[0])
+		v = bind_vector(t, bname, btype[0])
 	else:
-		bind_vector(t, bname, float)
-
+		v = bind_vector(t, bname, float)
+
+	return v
+
 
 def bind_scalar(t, bname, btype):
 	v = np.array([0], dtype=btype)
-	t.Branch(bname, v, '%s/%s' % (bname, typenames_short[btype]))
+	b = t.Branch(bname, v, '%s/%s' % (bname, typenames_short[btype]))
 	setattr(t, bname, v)
 
+	return v
+
+
 def bind_vector(t, bname, btype):
 	v = r.vector(typenames_long[btype])()
-	t.Branch(bname, 'vector<%s>'%typenames_long[btype], v)
+	b = t.Branch(bname, 'vector<%s>'%typenames_long[btype], v)
 	setattr(t, bname, v)
 
+	return v
+
 def bind_vector2(t, bname, btype):
 	v = r.vector('vector<%s>'%typenames_long[btype])()
-	t.Branch(bname, 'vector<vector<%s> >'%typenames_long[btype], v)
+	b = t.Branch(bname, 'vector<vector<%s> >'%typenames_long[btype], v)
 	setattr(t, bname, v)
+
+	return v