Convenience wrappers for loading a ROOT file into a pandas DataFrame

root_pandas.py

def get_matching_variables(fname, tree, patterns):
    from fnmatch import fnmatch
    from root_numpy import list_branches

    branches = list_branches(fname, tree)

    selected = []

    for p in patterns:
        for b in branches:
            if fnmatch(b, p) and not b in selected:
                selected.append(b)
    return selected


def load_root(fname, tree=None, patterns=None, *kargs, **kwargs):
    """
    Loads a root file into a pandas DataFrame.
    Further *kargs and *kwargs are passed to root_numpy's root2array.

    >>> df = load_root('test.root', 'MyTree', patterns=['x_*', 'y_*'], selection='x_1 > 100')

    If the root file contains a branch called index, it will become the DataFrame's index.
    """
    from pandas import DataFrame
    from root_numpy import root2array, list_trees

    if tree == None:
        branches = list_trees(fname)
        if len(branches) == 1:
            tree = branches[0]
        else:
            raise ValueError('More than one tree found in {}'.format(fname))

    if not patterns:
        all_vars = None
    else:
        # index is always loaded if it exists
        patterns.append('index')
        all_vars = get_matching_variables(fname, tree, patterns)

    arr = root2array(fname, tree, all_vars, *kargs, **kwargs)
    if 'index' in arr.dtype.names:
        df = DataFrame.from_records(arr, index='index')
    else:
        df = DataFrame.from_records(arr)
    return df


def save_root(df, fname, tree_name, *kargs, **kwargs):
    """
    Saves a pandas DataFrame as a root file.
    Further *kargs and *kwargs are passed to root_numpy's array2root.

    >>> df = DataFrame({'x': [1,2,3], 'y': [4,5,6]})
    >>> save_root('test.root', 'MyTree', df)
    
    The DataFrame index will be saved as an 'index' branch.
    """
    from root_numpy import array2root
    arr = df.to_records()
    array2root(arr, fname, tree_name, 'recreate', *kargs, **kwargs)