コード例 #1
0
ファイル: loaders.py プロジェクト: apuignav/analysis-tools
def get_pandas_from_root_file(file_name, tree_name, kwargs):
    """Load a pandas DataFrame from a ROOT file.

    Optional keys in `kwargs` are:
        + `variables`: List of variables to load.
        + `selection`: Selection to apply.

    Arguments:
        file_name (str): File to load.
        tree_name (str): Tree to load.
        kwargs (dict): Extra configuration.

    Return:
        pandas.DataFrame: ROOT file converted to pandas.

    """
    logger.debug("Loading ROOT file in pandas format -> %s:%s", file_name,
                 tree_name)
    if not os.path.exists(file_name):
        raise OSError("Cannot find input file -> {}".format(file_name))
    selection = kwargs.get('selection')
    variables = kwargs.get('variables', [])
    if selection:
        selection_expr = formulate.from_numexpr(selection)
        full_variables = variables + list(selection_expr.variables)
        output_data = read_root(file_name, tree_name,
                                columns=full_variables).query(selection)
        if variables:
            output_data = output_data[variables]
    else:
        output_data = read_root(file_name, tree_name, columns=variables)
    return output_data
コード例 #2
0
ファイル: test_backends.py プロジェクト: scikit-hep/formulate
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)'
        ))
コード例 #3
0
ファイル: test_backends.py プロジェクト: scikit-hep/formulate
 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()
コード例 #4
0
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
コード例 #5
0
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
コード例 #6
0
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
コード例 #7
0
ファイル: test_parser.py プロジェクト: scikit-hep/formulate
def test_too_many_function_arguments():
    with pytest.raises(ParsingException):
        from_numexpr('sqrt(2, 3)')
コード例 #8
0
ファイル: test_parser.py プロジェクト: scikit-hep/formulate
def test_invalid_arg_parse():
    with pytest.raises(ValueError):
        from_numexpr(Expression(IDs.SQRT, UC('2')))
コード例 #9
0
ファイル: test_parser.py プロジェクト: scikit-hep/formulate
def test_parse_invalid_expression():
    with pytest.raises(ParsingException):
        from_numexpr('saadasd()&+|()')