def trim_file(infile):
removebranches = ('lab[0-9]_MC12TuneV[0-9]_ProbNN', 'lab[0-9]_ProbNN',
'lab[0-9]_PID', 'lab0_DTF_vtx_DstMass',
'lab0_DTF_vtx_BothMass', 'lab0_DTF_DstMass',
'lab0_DTF_D0Mass', 'lab0_DTF_BothMass',
'lab[0-9]_hasMuon', 'lab[0-9]_isMuon',
'lab[0-9]_hasRich', 'lab[0-9]_UsedRich',
'lab[0-9]_RichAbove', 'lab[0-9]_hasCalo')
outfile = ROOT.TFile.Open(infile + '.trim', 'recreate')
infile = ROOT.TFile.Open(infile)
for k in infile.GetListOfKeys():
if 'KK' in k.GetName() or 'WIDEMASS' in k.GetName():
continue
outfile.mkdir(k.GetName())
outfile.cd(k.GetName())
for tname in infile.Get(k.GetName()).GetListOfKeys():
print 'Copy tree', k.GetName() + '/' + tname.GetName()
tcopy = copy_tree(infile.Get(k.GetName() + '/' + tname.GetName()),
removebranches=removebranches)
tcopy.Write()
outfile.cd()
outfile.Close()
infile.Close()
shutil.move(outfile.GetName(), infile.GetName())
def filter_2015_pipi():
sels = {
'Resolved': {
'': selection_R,
'_LowMass': selection_R_low,
'_HighMass': selection_R_high
},
'Merged': {}
}
for finalstate in 'pipi', : # 'Kpi' :
# Merged doesn't work currently cause the BDT expects lab6&7 to be the photons.
for pi0 in 'Resolved', : # 'Merged' :
for mag in 'Up', 'Down':
dataset = 'Data_2015_{finalstate}pi0_{pi0}_Mag{mag}_full'.format(
**locals())
print dataset
info = datalib.get_data_info(dataset)
tree = datalib.get_data(dataset)
for suff, sel in sels[pi0].items():
outputname = dataset.replace('_full', suff)
print outputname
outputdir = os.path.join(filtereddatadir, outputname)
if not os.path.exists(outputdir):
os.makedirs(outputdir)
outputfile = os.path.join(outputdir, outputname + '.root')
fout = ROOT.TFile.Open(outputfile, 'recreate')
treeout = copy_tree(tree, sel, write=True)
fout.Close()
def _make_dataloader(self):
'''Make the DataLoader for training.'''
# Load the data.
self.dataloader = TMVA.DataLoader(self.name)
# Add training variables.
for var in self.variables:
if not isinstance(var, (tuple, list)):
var = (var, )
try:
self.dataloader.AddVariable(*var)
except:
print 'Failed to call dataloader.AddVariable with args', var
raise
# Add spectator variables.
for var in self.spectators:
if not isinstance(var, (tuple, list)):
var = (var, )
try:
self.dataloader.AddSpectator(*var)
except:
print 'Failed to call dataloader.AddSpectator with args', var
raise
# Register trees.
# If we have explicit cuts for training and testing, we need to copy the TTrees first,
# applying these cuts.
if self.trainingcut:
pwd = ROOT.gROOT.CurrentDirectory()
self.tmpfile = ROOT.TFile.Open(
os.path.abspath('DataLoader_' + random_string() + '.root'),
'recreate')
self.tmpfile.cd()
signal_usedleaves, background_usedleaves = self.used_leaves()
usedleaves = {
'Signal': signal_usedleaves,
'Background': background_usedleaves
}
aliases = {
'Signal': get_aliases(self.signaltree),
'Background': get_aliases(self.backgroundtree)
}
addtreeargs = []
for name in 'Signal', 'Background':
lname = name.lower()
namecut = getattr(self, lname + 'cut')
for tname, ttype, cut in ('Training', TMVA.Types.kTraining,
self.trainingcut), (
'Testing', TMVA.Types.kTesting,
self.testingcut):
classname = self.name + '_' + name + '_' + tname + '_'
cut = AND(*filter(None, [namecut, cut]))
tree = getattr(self, lname + 'tree')
seltree, copyfriends = copy_tree(
tree,
selection=cut,
keepbranches=usedleaves[name],
rename=(
lambda name: classname + name.replace('/', '_')),
write=True,
returnfriends=True)
addtreeargs.append((seltree.GetName(), name,
getattr(self, lname + 'globalweight'),
ROOT.TCut(''), ttype))
weight = getattr(self, lname + 'weight')
if weight:
self.dataloader.SetWeightExpression(weight, name)
fname = self.tmpfile.GetName()
self.tmpfile.Close()
self.tmpfile = ROOT.TFile.Open(fname)
for args in addtreeargs:
tree = self.tmpfile.Get(args[0])
_aliases = aliases['Signal'] if 'Signal' in args else aliases[
'Background']
for name, alias in _aliases.items():
tree.SetAlias(name, alias)
self.dataloader.AddTree(tree, *args[1:])
self.dataloader.GetDataSetInfo().SetSplitOptions(
str(self.splitoptions))
if pwd:
pwd.cd()
else:
self.dataloader.AddSignalTree(self.signaltree,
self.signalglobalweight)
self.dataloader.AddBackgroundTree(self.backgroundtree,
self.backgroundglobalweight)
# Set weight expressions.
if self.signalweight:
self.dataloader.SetSignalWeightExpression(self.signalweight)
if self.backgroundweight:
self.dataloader.SetBackgroundWeightExpression(
self.backgroundweight)
# Prepare the training.
self.dataloader.PrepareTrainingAndTestTree(
ROOT.TCut(self.signalcut), ROOT.TCut(self.backgroundcut),
str(self.splitoptions))
return True
default=-1,
help='Number of entries to copy (default all).')
optParse.add_argument("--selection",
default='',
help='Selection to apply (default none).')
optParse.add_argument(
'--keepBranches',
nargs='*',
help='List of regexes to be matched to branches to be kept.')
optParse.add_argument(
'--removeBranches',
nargs='*',
help='List of regexes to be matched to branches to be removed.')
args = optParse.parse_args()
from ROOT import TTree, TFile
inFile = TFile(options.inFile)
inTree = TTree()
inFile.GetObject(options.inTree, inTree)
outFile = TFile(options.outFile, "recreate")
outTree = copy_tree(inTree,
selection=args.selection,
nentries=args.nEntries,
keepbranches=args.keepBranches,
removebranches=args.removeBranches)
outTree.Write()
outFile.Close()