-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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.
- Loading branch information
Showing
2 changed files
with
125 additions
and
9 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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() | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters