def main(args):
"""Main"""
treename = get_treename(args.inputfile)
branchname = args.branchname
while branch_already_present(args.inputfile, treename, branchname):
logging.warning(
'\'{}\' is already present in tree \'{}\' in file {}'.format(
branchname, treename, args.inputfile))
branchname = raw_input('Please enter a different branch name: ')
map_file = r.TFile.Open(args.corrmapfile)
corr_map = map_file.Get(args.name)
if not corr_map:
logging.fatal('Cannot find acceptance map \'{}\' in file {}'.format(
args.name, args.corrmapfile))
sys.exit(1)
acc_prov = AcceptanceCorrectionProvider(corr_map)
variables = ['{costh,phi}_PX']
if acc_prov.dim == 3:
variables.append(args.variable)
data = get_dataframe(args.inputfile, treename, columns=variables)
if acc_prov.dim == 2:
corr_weights = acc_prov.eval(data.costh_PX, data.phi_PX)
else:
corr_weights = acc_prov.eval(data.costh_PX, data.phi_PX,
_get_var(data, args.variable))
add_branch(corr_weights, branchname, args.inputfile, treename)
def test_return_tree_name(self):
"""Test if the returned treename is the expected one for one TTree"""
one_tree_files = ['one_tree.root', 'one_tree_plus_others.root']
for filen in one_tree_files:
full_path = os.path.join(self.test_data_dir, filen)
# NOTE: 'tree1' is hardcoded here as it is in the creation script
self.assertEqual(get_treename(full_path), 'tree1')
def main(args):
"""Main"""
# In order to not have to load the whole file first determine the names of
# the branches that are necessary
var_names = [('costh', 'phi')]
if args.genlevel:
logging.info('Also adding generator level folding')
var_names.append(('gen_costh', 'gen_phi'))
frames = args.frames.split(',')
load_variables = ['_'.join(p) for p in product(flatten(var_names), frames)]
for infile in args.inputfiles:
logging.info('Processing file {}'.format(infile))
if not args.treename:
treename = get_treename(infile)
else:
treename = args.treename
df = get_dataframe(infile, treename, columns=load_variables)
for var_pair in var_names:
for frame in frames:
costh_f, phi_f = get_folded_frame(df, frame, *var_pair)
add_branch(costh_f, '_'.join([var_pair[0], frame, 'fold']),
infile, treename)
add_branch(phi_f, '_'.join([var_pair[1], frame, 'fold']),
infile, treename)
def test_return_none_no_trees(self, mock_logger):
"""Test if None is returned on no found TTrees"""
no_tree_files = ['no_tree.root', 'no_tree_plus_others.root']
for filen in no_tree_files:
full_path = os.path.join(self.test_data_dir, filen)
exp_warning = 'Found no TTrees in {}'.format(full_path)
self.assertTrue(get_treename(full_path) is None)
mock_logger.warning.assert_called_with(exp_warning)
def test_return_none_mult_trees(self, mock_logger):
"""Test if None is returned on multiple TTrees"""
mult_tree_files = [
'multiple_trees.root', 'multiple_trees_plus_others.root'
]
for filen in mult_tree_files:
full_path = os.path.join(self.test_data_dir, filen)
exp_warning = 'Found more than one TTrees in {}: {}'\
.format(full_path, ['tree1', 'tree2'])
self.assertTrue(get_treename(full_path) is None)
mock_logger.warning.assert_called_with(exp_warning)
def get_load_vars(infile, trigger):
"""
Get the list of variables to load
"""
triggers = [trigger + '*']
rfile = r.TFile.Open(infile)
tree = rfile.Get(get_treename(infile))
# Updating a branch is not possible.
if check_branch_available(tree, 'trigger', True):
logging.warn('The branch \'trigger\' is already in the TTree of the '
'input file. In place updating of the Branch is not '
'possible!')
return triggers
def main(args):
"""Main"""
muon_effs = get_effs(args.efficiencies, args.muoneffs, MuonEfficiencies)
phot_effs = get_effs(args.efficiencies, args.photoneffs,
PhotonEfficiencies)
treename = get_treename(
args.datafile) # do this here and fail early in case
# only read in the necessary columns and add newly created branches to file
columns = [n + '{Pt,Eta}' for n in MUON_NAMES + [PHOTON_NAME]]
data = get_dataframe(args.datafile, treename, columns=columns)
if phot_effs is not None:
photon_effs = np.array(calc_effs(data, phot_effs, PHOTON_NAME))
add_branch(photon_effs, '_'.join(['gamma_sm', args.name]),
args.datafile, treename)
if muon_effs is not None:
for muon in MUON_NAMES:
mu_effs = np.array(calc_effs(data, muon_effs, muon))
lep = 'lepP' if 'P' in muon else 'lepN' # for consistency
add_branch(mu_effs, '_'.join([lep, 'sm', args.name]),
args.datafile, treename)
def main(args):
"""Main"""
in_data = get_dataframe(args.inputfile, columns=get_load_vars(args.inputfile, args.trigger))
triggers = get_trigger_bit_event(in_data, args.trigger)
add_branch(triggers, 'trigger', args.inputfile, get_treename(args.inputfile))