def test_method_file_cd():
file1 = TemporaryFile()
foo = Foo()
foo.SetDirectory(file1)
file2 = TemporaryFile()
foo.write()
def create_tree():
f = TemporaryFile()
tree = Tree("tree", model=create_model())
# fill the tree
for i in xrange(1000):
assert_equal(tree.a_vect, LorentzVector(0, 0, 0, 0))
random_vect = LorentzVector(gauss(.5, 1.), gauss(.5, 1.),
gauss(.5, 1.), gauss(.5, 1.))
tree.a_vect.copy_from(random_vect)
assert_equal(tree.a_vect, random_vect)
tree.a_x = gauss(.5, 1.)
tree.a_y = gauss(.3, 2.)
tree.a_z = gauss(13., 42.)
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(gauss(.5, 1.), gauss(.5, 1.), gauss(.5, 1.),
gauss(.5, 1.))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(.3, 2.))
tree.i = i
assert_equal(tree.b_n, tree.b_vect.size())
assert_equal(tree.b_n, tree.b_x.size())
assert_equal(tree.b_n, tree.b_y.size())
tree.fill(reset=True)
tree.write()
# TFile.Close the file but keep the underlying
# tempfile file descriptor open
ROOT.TFile.Close(f)
FILES.append(f)
FILE_PATHS.append(f.GetName())
def test_file_assoc():
with TemporaryFile() as f1:
t = Tree()
with TemporaryFile() as f2:
pass
#f1.cd() <== this should not be needed!
# the tree should "remember" what file it was created in
t.Write()
def test_file():
f = TemporaryFile()
assert_raises(DoesNotExist, f.Get, 'blah')
hist = Hist(1, 0, 1, name='test')
hist.Write()
hist2 = f.test
assert hist2.__class__ == hist.__class__
os.unlink(f.GetName())
def test_context():
with MemFile() as a:
assert_equal(ROOT.gDirectory, a)
with MemFile() as b:
d = Directory('test')
with d:
assert_equal(ROOT.gDirectory, d)
assert_equal(ROOT.gDirectory, b)
assert_equal(ROOT.gDirectory, a)
# test out of order
f1 = MemFile()
f2 = MemFile()
with f1:
assert_equal(ROOT.gDirectory, f1)
assert_equal(ROOT.gDirectory, f2)
f1.Close()
f2.Close()
d = Directory('test')
d.cd()
# test without with statement
f1 = MemFile()
f2 = TemporaryFile()
assert_equal(ROOT.gDirectory, f2)
f2.Close()
assert_equal(ROOT.gDirectory, f1)
f1.Close()
def test_file_item():
with TemporaryFile() as f:
h = Hist(1, 0, 1, name='test')
f['myhist'] = h
f.myhist
assert_equals(f['myhist'].name, 'test')
def to_root(self, dst, optimize=False):
"""Write the selected events and count histograms to a ROOT file.
Parameters
----------
dst : path
The path of the output ROOT file.
optimize : bool, optional
If True, optimize the storage of events in the output ROOT file.
The default is False.
"""
with root_open(dst, 'w') as outfile:
for hist in self._count_histograms:
hist.Write()
if optimize:
with TemporaryFile() as tmp:
tree_tmp = self.CopyTree('')
tmp.Write()
tmp.Flush()
outfile.cd()
tree_opt = tree_tmp.CloneTree(-1,
'fast SortBasketsByEntry')
tree_opt.OptimizeBaskets()
tree_opt.Write()
else:
tree_new = self.CopyTree(
'') if self.eventlist else self.CloneTree(-1, 'fast')
tree_new.Write()
def test_ntuple():
with TemporaryFile():
ntuple = Ntuple(('a', 'b', 'c'), name='test')
for i in range(100):
ntuple.Fill(gauss(.3, 2.), gauss(0, 1.), gauss(-1., 5))
ntuple.Write()
def test_histfactory():
# create some Samples
data = Data('data')
data.hist = get_random_hist()
a = Sample('QCD')
b = Sample('QCD')
for sample in (a, b):
sample.hist = get_random_hist()
# include some histosysts
for sysname in ('x', 'y', 'z'):
histosys = HistoSys(sysname)
histosys.high = get_random_hist()
histosys.low = get_random_hist()
sample.AddHistoSys(histosys)
# include some normfactors
for normname in ('x', 'y', 'z'):
norm = NormFactor(normname)
norm.value = 1
norm.high = 2
norm.low = 0
norm.const = False
sample.AddNormFactor(norm)
# samples must be compatible here
c = a + b
c = sum([a, b])
# create Channels
channel_a = Channel('VBF')
channel_a.data = data
channel_a.AddSample(a)
channel_b = Channel('VBF')
channel_b.data = data
channel_b.AddSample(b)
combined_channel = channel_a + channel_b
combined_channel = sum([channel_a, channel_b])
# create a Measurement
meas = Measurement('MyAnalysis')
meas.AddChannel(channel_a)
# create the workspace containing the model
workspace = make_workspace(meas, silence=True)
with TemporaryFile():
workspace.Write()
assert_true(channel_a.GetSample(a.name) is not None)
channel_a.RemoveSample(a.name)
assert_true(channel_a.GetSample(a.name) is None)
assert_true(meas.GetChannel(channel_a.name) is not None)
meas.RemoveChannel(channel_a.name)
assert_true(meas.GetChannel(channel_a.name) is None)
def test_file_contains():
with TemporaryFile() as f:
assert_equal('some/thing' in f, False)
rdir = f.mkdir('some')
thing = Hist(10, 0, 1, name='thing')
rdir.thing = thing
assert_equal('some/thing' in f, True)
assert_equal('thing' in rdir, True)
def test_tempfile():
with TemporaryFile() as f:
assert_equal(os.path.isfile(f.GetName()), True)
assert_raises(DoesNotExist, f.Get, 'blah')
hist = Hist(1, 0, 1, name='test')
hist.Write()
hist2 = f.test
assert_equal(hist2.__class__, hist.__class__)
assert_equal(os.path.isfile(f.GetName()), False)
def setup_class(cls):
class ObjectA(TreeModel):
# A simple tree object
x = FloatCol()
y = FloatCol()
z = FloatCol()
class ObjectB(TreeModel):
# A tree object collection
x = stl.vector('int')
y = stl.vector('float')
vect = stl.vector('TLorentzVector')
# collection size
n = IntCol()
class Event(ObjectA.prefix('a_') + ObjectB.prefix('b_')):
i = IntCol()
for i in range(5):
f = TemporaryFile()
tree = Tree("tree", model=Event)
# fill the tree
for i in xrange(10000):
tree.a_x = gauss(.5, 1.)
tree.a_y = gauss(.3, 2.)
tree.a_z = gauss(13., 42.)
tree.b_vect.clear()
tree.b_x.clear()
tree.b_y.clear()
tree.b_n = randint(1, 5)
for j in xrange(tree.b_n):
vect = LorentzVector(gauss(.5, 1.), gauss(.5, 1.),
gauss(.5, 1.), gauss(.5, 1.))
tree.b_vect.push_back(vect)
tree.b_x.push_back(randint(1, 10))
tree.b_y.push_back(gauss(.3, 2.))
tree.i = i
tree.fill()
tree.write()
# TFile.Close the file but keep the underlying
# tempfile file descriptor open
ROOT.TFile.Close(f)
cls.files.append(f)
cls.file_paths.append(f.GetName())
def test_file_attr():
with TemporaryFile() as f:
h = Hist(1, 0, 1, name='test')
f.myhist = h
f.Get('myhist')
assert_equal(f.myhist.name, 'test')
f.something = 123
f.mkdir('hello')
f.hello.something = h
assert_equal(f['hello/something'].name, 'test')
def test_correlated_values():
try:
import uncertainties
except ImportError:
raise SkipTest("uncertainties package is not installed")
from rootpy.stats.correlated_values import correlated_values
# construct pdf and toy data following example at
# http://root.cern.ch/drupal/content/roofit
# --- Observable ---
mes = RooRealVar("mes", "m_{ES} (GeV)", 5.20, 5.30)
# --- Parameters ---
sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)
# --- Build Gaussian PDF ---
signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)
# --- Build Argus background PDF ---
argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
background = RooArgusBG("background", "Argus PDF",
mes, RooFit.RooConst(5.291), argpar)
# --- Construct signal+background PDF ---
nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
model = RooAddPdf("model", "g+a",
RooArgList(signal,background),
RooArgList(nsig,nbkg))
# --- Generate a toyMC sample from composite PDF ---
data = model.generate(RooArgSet(mes), 2000)
# --- Perform extended ML fit of composite PDF to toy data ---
fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1))
nsig, nbkg = correlated_values(["nsig", "nbkg"], fitresult)
# Arbitrary math expression according to what the `uncertainties`
# package supports, automatically computes correct error propagation
sum_value = nsig + nbkg
value, error = sum_value.nominal_value, sum_value.std_dev
workspace = Workspace(name='workspace')
# import the data
assert_false(workspace(data))
with TemporaryFile():
workspace.Write()
def test_pickler():
hlist = list()
for i in range(10):
hlist.append(Hist(10, 0, 10))
with TemporaryFile() as tmpfile:
dump(hlist, tmpfile)
hlist_out = load(tmpfile)
assert_equal([h.name for h in hlist_out], [h.name for h in hlist])
hdict = dict()
for i in range(100):
hist = Hist(10, 0, 1, type=random.choice('CSIFD'))
hdict[hist.name] = hist
with TemporaryFile() as tmpfile:
rdir = tmpfile.mkdir('pickle')
dump(hdict, rdir)
hdict_out = load(rdir)
assert_equal(len(hdict_out), 100)
for name, hist in hdict_out.items():
assert_equal(name, hist.name)
assert_equal(hist.TYPE, hdict[hist.name].TYPE)
def test_plottable():
# construct pdf and toy data following example at
# http://root.cern.ch/drupal/content/roofit
# Observable
mes = RooRealVar("mes", "m_{ES} (GeV)", 5.20, 5.30)
# Parameters
sigmean = RooRealVar("sigmean", "B^{#pm} mass", 5.28, 5.20, 5.30)
sigwidth = RooRealVar("sigwidth", "B^{#pm} width", 0.0027, 0.001, 1.)
# Build Gaussian PDF
signal = RooGaussian("signal", "signal PDF", mes, sigmean, sigwidth)
# Build Argus background PDF
argpar = RooRealVar("argpar", "argus shape parameter", -20.0, -100., -1.)
background = RooArgusBG("background", "Argus PDF",
mes, RooFit.RooConst(5.291), argpar)
# Construct signal+background PDF
nsig = RooRealVar("nsig", "#signal events", 200, 0., 10000)
nbkg = RooRealVar("nbkg", "#background events", 800, 0., 10000)
model = RooAddPdf("model", "g+a",
RooArgList(signal, background),
RooArgList(nsig, nbkg))
# Generate a toyMC sample from composite PDF
data = model.generate(RooArgSet(mes), 2000)
# Perform extended ML fit of composite PDF to toy data
fitresult = model.fitTo(data, RooFit.Save(), RooFit.PrintLevel(-1))
# Plot toy data and composite PDF overlaid
mesframe = asrootpy(mes.frame())
data.plotOn(mesframe)
model.plotOn(mesframe)
for obj in mesframe.objects:
assert_true(obj)
for curve in mesframe.curves:
assert_true(curve)
for hist in mesframe.data_hists:
assert_true(hist)
assert_true(mesframe.plotvar)
with TemporaryFile():
mesframe.Write()
def split_train_test(src, dst, subset=None, selection='', branches=None):
"""Split a sample's events into training and testing sets based on the
parity of an event. Odd numbered events are reserved for training,
while even numbered events are reserved for testing.
Parameters
----------
src : path
The path to the input sample.
dst : path
The path to the output sample containing the train and test trees.
subset : string, optional
A selection applied to the sample's tree to choose only a
subset of events to split into training and teseting events.
The default is None to consider all events.
selection : string, optional
A selection applied to the sample's tree to define training
and testing events. The default is None for no selection.
branches : list of strings, optional
The list of branches to be kept in the output sample.
The default is None to keep all branches.
"""
if subset:
if selection:
selection = '({0})&&({1})'.format(selection, subset)
else:
selection = subset
with root_open(src) as infile, TemporaryFile():
tree = infile.Get('tree')
train_temp = asrootpy(tree.CopyTree(
'{}&&(evt%2==1)'.format(selection)))
train_temp.SetName('train')
test_temp = asrootpy(tree.CopyTree('{}&&(evt%2==0)'.format(selection)))
test_temp.SetName('test')
with root_open(dst, 'w') as outfile:
if branches:
train_temp.activate(branches, exclusive=True)
test_temp.activate(branches, exclusive=True)
train = train_temp.CloneTree()
test = test_temp.CloneTree()
outfile.Write()
def test_tempfile():
with TemporaryFile():
Hist(1, 0, 1, name='test').write()
import os
import atexit
# # pytables imports
# import tables
# rootpy imports
from rootpy.io import root_open, TemporaryFile
# local imports
from .. import NTUPLE_PATH, DEFAULT_STUDENT
from . import log
log = log[__name__]
FILES = {}
TEMPFILE = TemporaryFile()
PILEUP_FILES = {}
def get_file(ntuple_path=NTUPLE_PATH,
file_name=None,
student=DEFAULT_STUDENT,
hdf=False,
suffix='_train',
force_reopen=False,
**kwargs):
if file_name is None:
ext = '.h5' if hdf else '.root'
filename = student + suffix + ext
import platform
import matplotlib.pyplot as plt
import os
with open('hardware.pkl', 'r') as pkl:
info = pickle.load(pkl)
# construct system hardware information string
hardware = '{cpu}\nStorage: {hdd}\nROOT-{root}\nPython-{python}\nNumPy-{numpy}'.format(
cpu=info['CPU'],
hdd=info['HDD'],
root=rootpy.ROOT_VERSION,
python=platform.python_version(),
numpy=np.__version__)
rfile = TemporaryFile()
def randomword(length):
return ''.join(random.choice(string.lowercase) for i in range(length))
def make_tree(entries, branches=1, dtype=np.double):
dtype = np.dtype([(randomword(20), dtype) for idx in range(branches)])
array = np.zeros(entries, dtype=dtype)
return array2tree(array, name=uuid.uuid4().hex)
# warm up
print("warming up... ", end="")
for i in range(30):
def test_method_file_check_good():
foo = Foo()
with TemporaryFile():
foo.something(42)
def test_require_file_good():
with TemporaryFile():
t = Tree()
