Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Contextmanagers for ROOT.TTree
import ROOT
import contextlib2
class TTree:
"""
A contextmanager for TTree, which loads from a file directly.
Can disable all branches except the needed branches (default: all enabled).
"""
def __init__(self, tree_name, file_name, active_branches=None):
"""
>>> import ROOT
>>> from utils import TTree
>>> with TTree("DYTuple/DecayTree", "../ntuple/dy_tuple_12_mu_reduced.root", ["Z0_MM"]) as tree:
... tree.Draw("Z0_MM")
"""
self.file_name = file_name
self.tree_name = tree_name
self.active_branches = active_branches
def __enter__(self):
self.file = ROOT.TFile(self.file_name)
tree = self.file.Get(self.tree_name)
if self.active_branches is not None:
tree.SetBranchStatus("*", False)
for branch in self.active_branches:
tree.SetBranchStatus(branch, True)
return tree
def __exit__(self, *args):
self.file.Close()
class TTreeStack(object):
"""
A stack of TTrees, can be used to create multiple trees at once:
Example:
>>> with TTreeStack("tree", "file1.root", "file2.root") as (tree1, tree2):
... tree1.Print()
... tree2.Print()
"""
def __init__(self, tree_name, *file_names, **kwargs):
self.tree_name = tree_name
self.file_names = file_names
self.active_branches = kwargs.get("active_branches", None)
def __enter__(self):
with contextlib2.ExitStack() as stack:
self.trees = [stack.enter_context(
TTree(self.tree_name, fname, active_branches=self.branches)) for fname in self.file_names]
self.close = stack.pop_all().close
return self.trees
def __exit__(self, exception_type, exception_value, traceback):
self.close()
@swang373

This comment has been minimized.

Copy link

@swang373 swang373 commented Mar 17, 2017

Do you have any thoughts on how best to work with the separate trees as one? I guess you could create a temporary ROOT file with the trees merged using TTree::Merge() and return that, but that sounds unwieldy. Even less feasible is hadd-ing all the ROOT files together just to work with one tree. A TChain sounds like the natural solution, but then we wouldn't need the context manager... I ask because of situations where one would have to pass a single TTree, e.g. TMVA::Factory::AddSignalTree()

@graipher

This comment has been minimized.

Copy link
Owner Author

@graipher graipher commented Mar 19, 2017

@swang373 In that case I think a TChain is preferable, because it is already designed to behave like a TTree but pulling from multiple trees in the background. The above TTreeStack is good if you need to open all these trees but want to do different stuff with them afterwards (I need to extract templates for a fit from each tree, doing some reweighting on them first).

@graipher

This comment has been minimized.

Copy link
Owner Author

@graipher graipher commented Mar 19, 2017

Also note that the TTreeStack is just a shorter way to write:
with TTree("tree", "file1.root") as tree1, TTree("tree", "file2.root") as tree2, ...:

@graipher

This comment has been minimized.

Copy link
Owner Author

@graipher graipher commented Mar 19, 2017

I also wrote a contextmanager for a TChain, but haven't really needed it so far (see https://gist.github.com/graipher/f4f35d792a97c65c6c458c65c3cc9295).

@swang373

This comment has been minimized.

Copy link

@swang373 swang373 commented Mar 29, 2017

@graipher That use case for distinct operations on the separate trees makes sense, I was stuck thinking of operating on them uniformly. Thanks for the examples!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.