Test RooDataSet handling of TTree branch aliases.

workspace_aliases.py

"""Test RooDataSet handling of TTree branch aliases.

The TTree::SetAlias method allows you to create "branch aliases", names that
point to formulas of other branches.
It seems RooFit will not load alias branches when importing a TTree in to a
RooDatSet, filling the variable with zeros.

This script demonstrates that, firstly by fitting the non-aliased branch 'x',
which works fine, and then fitting the alias to x, 'y', which fails because the
values are all zero.
"""
from __future__ import print_function
from array import array

import ROOT

NENTRIES = int(10e3)
MEAN = 0
SIGMA = 1

# Create file with tree
f = ROOT.TFile('workspace_aliases.root', 'recreate')
t = ROOT.TTree('WorkspaceAliases', 'WorkspaceAliases')

# Add branches
x = array('d', [0])
t.Branch('x', x, 'x/d')

# Fill tree and write it
rand = ROOT.TRandom3()
for i in range(NENTRIES):
    x[0] = rand.Gaus()
    t.Fill()
print('Tree has {0} entries after filling'.format(t.GetEntries()))
f.Write()
f.Close()

# Reopen the tree
f = ROOT.TFile('workspace_aliases.root')
t = f.Get('WorkspaceAliases')
print('Tree has {0} entries after loading'.format(t.GetEntries()))
# Create an alias to a branch
t.SetAlias('y', 'x')

# Create RooWorkspace
w = ROOT.RooWorkspace('w')
x = w.factory('x[-5, 5]')
y = w.factory('y[-5, 5]')
d = ROOT.RooDataSet('d', 'd', t, ROOT.RooArgSet(x, y))
print('RDS has {0} entries after loading TTree'.format(d.sumEntries()))

# Fit the normal branch
pdf_x = w.factory('RooGaussian::f(x, mean[0, -2, 2], sigma[1, 0.3, 2])')
pdf_x.fitTo(d)
frame = x.frame()
d.plotOn(frame)
pdf_x.plotOn(frame)
c = ROOT.TCanvas('c_x', 'c_x', 400, 400)
frame.Draw()
c.SaveAs('workspace_aliases_x.pdf')

# Fit the alias branch
pdf_y = w.factory('RooGaussian::g(y, mean, sigma)')
pdf_y.fitTo(d)
frame = y.frame()
d.plotOn(frame)
pdf_y.plotOn(frame)
c = ROOT.TCanvas('c_y', 'c_y', 400, 400)
frame.Draw()
c.SaveAs('workspace_aliases_y.pdf')

workspace_aliases_x.png

workspace_aliases_y.png