def test_readme():
momentum = from_root('TMath::Sqrt(X_PX**2 + X_PY**2 + X_PZ**2)')
assert momentum.to_numexpr(
) == 'sqrt(((X_PX ** 2) + (X_PY ** 2) + (X_PZ ** 2)))'
assert momentum.to_root(
) == 'TMath::Sqrt(((X_PX ** 2) + (X_PY ** 2) + (X_PZ ** 2)))'
my_selection = from_numexpr('X_PT > 5 & (Mu_NHits > 3 | Mu_PT > 10)')
assert my_selection.to_root(
) == '(X_PT > 5) && ((Mu_NHits > 3) || (Mu_PT > 10))'
assert my_selection.to_numexpr(
) == '(X_PT > 5) & ((Mu_NHits > 3) | (Mu_PT > 10))'
my_sum = from_auto('True + False')
assert my_sum.to_root() == 'true + false'
assert my_sum.to_numexpr() == 'True + False'
my_check = from_auto('(X_THETA*TMath::DegToRad() > pi/4) && D_PE > 9.2')
assert my_check.variables == {'D_PE', 'X_THETA'}
assert my_check.named_constants == {'DEG2RAD', 'PI'}
assert my_check.unnamed_constants == {'4', '9.2'}
new_selection = (momentum > 100) and (my_check or (np.sqrt(my_sum) < 1))
def numexpr_eval(string):
return numexpr.evaluate(string,
local_dict=dict(X_THETA=1234, D_PE=678))
assert pytest.approx(
numexpr_eval(new_selection.to_numexpr()),
numexpr_eval(
'((X_THETA * 0.017453292519943295) > (3.141592653589793 / 4)) & (D_PE > 9.2)'
))
def test_unnamed_constants():
assert from_root('pi').unnamed_constants == set()
assert from_numexpr('2').unnamed_constants == {'2'}
assert from_numexpr('2e-3').unnamed_constants == {'2e-3'}
assert from_numexpr('A').unnamed_constants == set()
assert from_numexpr('A + A').unnamed_constants == set()
assert from_numexpr('A + B').unnamed_constants == set()
assert from_numexpr('A + A*A - 3e7').unnamed_constants == {'3e7'}
assert from_numexpr('arctan2(A, A)').unnamed_constants == set()
assert from_numexpr('arctan2(A, B)').unnamed_constants == set()
assert from_root('arctan2(A, pi)').unnamed_constants == set()
assert from_numexpr('arctan2(arctan2(A, B), C)').unnamed_constants == set()
for base, expect in [(UC('2'), {'2'}), (Variable('A'), set()), (NC(ConstantIDs.PI), set())]:
expr = base
for i in list(range(100)):
expr = Expression(IDs.SQRT, expr)
assert expr.unnamed_constants == expect
def test():
root_expression = from_root(root_input)
numexpr_expression = from_numexpr(numexpr_input)
assert_equal_expressions(root_expression, numexpr_expression)
assert to_numexpr(root_expression) == to_numexpr(numexpr_expression)
assert to_root(root_expression) == to_root(numexpr_expression)
assert root_expression.to_numexpr() == numexpr_expression.to_numexpr()
assert root_expression.to_root() == numexpr_expression.to_root()
def test_get_variables():
assert from_root('pi').variables == set()
assert from_numexpr('2').variables == set()
assert from_numexpr('2e-3').variables == set()
assert from_numexpr('A').variables == set(['A'])
assert from_numexpr('A + A').variables == set(['A'])
assert from_numexpr('A + B').variables == set(['A', 'B'])
assert from_numexpr('A + A*A - 3e7').variables == set(['A'])
assert from_numexpr('arctan2(A, A)').variables == set(['A'])
assert from_numexpr('arctan2(A, B)').variables == set(['A', 'B'])
assert from_root('arctan2(A, pi)').variables == set(['A'])
assert from_numexpr('arctan2(arctan2(A, B), C)').variables == set(
['A', 'B', 'C'])
for base, expect in [(UC('2'), set()), (Variable('A'), set(['A'])),
(NC(ConstantIDs.PI), set())]:
expr = base
for i in list(range(100)):
expr = Expression(IDs.SQRT, expr)
assert expr.variables == expect
def parse_args(args):
parser = argparse.ArgumentParser(
description='Convert between different types of formulae')
from_group = parser.add_mutually_exclusive_group(required=True)
from_group.add_argument('--from-root')
from_group.add_argument('--from-numexpr')
to_group = parser.add_mutually_exclusive_group(required=True)
to_group.add_argument('--to-root', action='store_true')
to_group.add_argument('--to-numexpr', action='store_true')
to_group.add_argument('--variables', action='store_true')
to_group.add_argument('--named-constants', action='store_true')
to_group.add_argument('--unnamed-constants', action='store_true')
args = parser.parse_args(args)
if args.from_root is not None:
expression = from_root(args.from_root)
elif args.from_numexpr is not None:
expression = from_numexpr(args.from_numexpr)
else:
raise NotImplementedError()
if args.to_root:
result = to_root(expression)
elif args.to_numexpr:
result = to_numexpr(expression)
elif args.variables:
result = '\n'.join(sorted(expression.variables))
elif args.named_constants:
result = '\n'.join(sorted(expression.named_constants))
elif args.unnamed_constants:
result = '\n'.join(sorted(expression.unnamed_constants))
else:
raise NotImplementedError()
return result
def get_root_from_root_file(file_name, tree_name, kwargs):
"""Load a ROOT tree into a `ROOT.RooDataSet`.
Needed keys in `kwargs` are:
+ `name`: Name of the `RooDataSet`.
+ `title`: Title of the `RooDataSet`.
Optional keys are:
+ `variables`: List of variables to load.
+ `selection`: Selection to apply.
+ `ranges`: Range to apply to some variables.
Arguments:
file_name (str): File to load.
tree_name (str): Tree to load.
kwargs (dict): Extra configuration.
Return:
ROOT.RooDataSet: ROOT file converted to RooDataSet.
Raise:
KeyError: If there are errors in `kwargs`.
ValueError: If the requested variables cannot be found in the input file.
OSError: If the ROOT file cannot be found.
"""
def get_list_of_leaves(tree):
"""Get list of leave names from a tree matching a certain regex.
Arguments:
tree (`ROOT.TTree`): Tree to extract the leaves from.
Return:
list: Leaves of the tree.
"""
object_list = tree.GetListOfLeaves()
it = object_list.MakeIterator()
output = set()
for _ in range(object_list.GetSize()):
obj = it.Next()
if obj:
output.add(obj.GetName())
return output
logger.debug("Loading ROOT file in RooDataSet format -> %s:%s", file_name,
tree_name)
if not os.path.exists(file_name):
raise OSError("Cannot find input file -> {}".format(file_name))
try:
name = kwargs['name']
title = kwargs.get('title', name)
except KeyError as error:
raise KeyError("Missing configuration key -> {}".format(error))
tfile = ROOT.TFile.Open(file_name)
tree = tfile.Get(tree_name)
if not tree:
raise KeyError(
"Cannot find tree in input file -> {}".format(tree_name))
leaves = get_list_of_leaves(tree)
variables = set(kwargs.get('variables', leaves))
# Acceptance
if 'acceptance' in kwargs:
raise NotImplementedError(
"Acceptance weights are not implemented for ROOT files")
# Check weights
try:
weight_var, weights_to_normalize, weights_not_to_normalize = _analyze_weight_config(
kwargs)
except KeyError:
raise KeyError("Badly specified weights")
if variables and weight_var:
variables = set(variables) | set(weights_to_normalize) | set(
weights_not_to_normalize)
# Crosscheck leaves
if variables - leaves:
raise ValueError("Cannot find leaves in input -> {}".format(variables -
leaves))
selection = kwargs.get('selection')
leave_set = ROOT.RooArgSet()
leave_list = []
if selection:
selection_expr = formulate.from_root(selection)
for var in selection_expr.variables.union(variables):
leave_list.append(ROOT.RooRealVar(var, var, 0.0))
leave_set.add(leave_list[-1])
name = ''.join(random.SystemRandom().choice(string.ascii_letters +
string.digits)
for _ in range(10))
temp_ds = ROOT.RooDataSet(name, name, leave_set,
ROOT.RooFit.Import(tree),
ROOT.RooFit.Cut(selection))
destruct_object(tree)
tree = temp_ds
var_set = ROOT.RooArgSet()
var_list = {}
for var in variables:
var_list[var] = ROOT.RooRealVar(var, var, 0.0)
var_set.add(var_list[var])
if kwargs.get('ranges'):
for var_name, range_val in kwargs['ranges'].items():
if var_name not in var_list:
raise KeyError(
"Range specified for a variable not included in the dataset -> {}"
.format(var_name))
try:
if isinstance(range_val, str):
min_, max_ = range_val.split()
else:
min_, max_ = range_val
except ValueError:
raise KeyError(
"Malformed range specification for {} -> {}".format(
var_name, range_val))
var_set[var_name].setMin(float(min_))
var_set[var_name].setMax(float(max_))
dataset = ROOT.RooDataSet(name, title, var_set, ROOT.RooFit.Import(tree))
if weight_var:
# Weights to normalize
to_normalize_w = ROOT.RooFormulaVar(
"{}_not_normalized".format(weight_var),
"{}_not_normalized".format(weight_var),
"*".join(weights_to_normalize),
list_to_rooarglist(var_list[weight]
for weight in weights_to_normalize))
var_set.append(to_normalize_w)
dataset.addColumn(to_normalize_w)
sum_weights = sum(
dataset.get(entry)["{}_not_normalized".format(
weight_var)].getVal() for entry in dataset.sumEntries())
normalized_w = ROOT.RooFormulaVar(
"{}_normalized".format(weight_var),
"{}_normalized".format(weight_var),
"{}_not_normalized/{}".format(weight_var, sum_weights),
ROOT.RooArgList(to_normalize_w))
var_set.append(normalized_w)
dataset.addColumn(normalized_w)
# Non-normalized weights
weights = ROOT.RooFormulaVar(
weight_var, weight_var,
"*".join(weights_not_to_normalize +
["{}_normalized".format(weight_var)]),
list_to_rooarglist(
[var_list[weight]
for weight in weights_not_to_normalize] + [normalized_w]))
var_set.append(weights)
dataset.addColumn(weights)
dataset_w = ROOT.RooDataSet(name, title, var_set,
ROOT.RooFit.Import(dataset),
ROOT.RooFit.WeightVar(weight_var))
destruct_object(dataset)
dataset = dataset_w
# ROOT Cleanup
destruct_object(tree)
tfile.Close()
destruct_object(tfile)
if selection:
for leave in leave_list:
destruct_object(leave)
for var in variables:
destruct_object(var_list[var])
# Let's return
dataset.SetName(name)
dataset.SetTitle(title)
return dataset