Пример #1
0
def tree_to_table(tree):
    """Convert a ROOT TTree to an astropy Table.
    
    Parameters
    ----------
    
    Returns
    -------
    table : astropy.table.Table
        
    """
    from rootpy import asrootpy
    tree = asrootpy(tree)
    array = tree.to_array()

    
    from rootpy.io import open
    from rootpy.root2array import tree_to_recarray
    print('Reading %s' % infile)
    file = open(infile)
    tree_name = 'TableAllMu_WithoutNan' # 'ParTree_Postselect'
    print('Getting %s' % tree_name)
    tree = file.get(tree_name, ignore_unsupported=True)
    print('Converting tree to recarray')
    array = tree_to_recarray(tree)
    print('Converting recarray to atpy.Table')
    table = recarray_to_table(array)

    print('Removing empty columns:')
    empty_columns = []
    for col in table.columns:
        #if table['col'].min() == table['col'].max()
        if (table['col'] == 0).all():
            empty_columns.append(col)
    print(empty_columns)
    table.remove_columns(empty_columns)
    
    for name in array.dtype.names:
        # FITS can't save these types.
        data = array[name]
        if data.dtype == np.dtype('uint64'):
            data = data.copy().astype('int64')
        table.add_column(name, data)
    return table
Пример #2
0
def tree_to_table(tree, tree_name):
    """Convert a ROOT TTree to an astropy Table.

    Parameters
    ----------
    tree : ROOT.TTree
        ROOT TTree

    Returns
    -------
    table : `~astropy.table.Table`
        ROOT tree data as an astropy table.
    """
    from rootpy import asrootpy
    from rootpy.io import open
    from rootpy.root2array import tree_to_recarray

    tree = asrootpy(tree)
    array = tree.to_array()

    file = open(infile)
    tree_name = 'TableAllMu_WithoutNan'  # 'ParTree_Postselect'
    tree = file.get(tree_name, ignore_unsupported=True)
    array = tree_to_recarray(tree)
    table = recarray_to_table(array)

    # Remove empty columns
    empty_columns = []
    for col in table.columns:
        #if table['col'].min() == table['col'].max()
        if (table['col'] == 0).all():
            empty_columns.append(col)
    print(empty_columns)
    table.remove_columns(empty_columns)

    for name in array.dtype.names:
        # FITS can't save these types.
        data = array[name]
        if data.dtype == np.dtype('uint64'):
            data = data.copy().astype('int64')
        table.add_column(name, data)
    return table
Пример #3
0
        table = None

        ROOT.TTreeCache.SetLearnEntries(1)

        time_in_read = 0
        time_in_append = 0
        processed_files = 0

        try:
            for file_chunk in chunk_files(flat_files, args.chainsize):
                t1 = time.time()
                chunk_chain = ROOT.TChain(tree)
                chunk_chain.SetCacheSize(10000000)
                for file in file_chunk:
                    chunk_chain.Add(file)
                recarray = tree_to_recarray(chunk_chain, None, False)
                t2 = time.time()
                time_in_read += t2 - t1

                t1 = time.time()
                if table is None:
                    log.info("Creating table")
                    table = hfile.createTable(args.where, h5name,
                                              recarray,
                                              "", filters=filters,
                                              expectedrows=entries)
                else:
                    table.append(recarray)
                table.flush()
                t2 = time.time()
                time_in_append += t2 - t1
Пример #4
0
    print "time without profiler overhead:"
    t1 = time.time()
    arr1 = tree_to_recarray_py(tree, branches=branches,
                               include_weight=True)
    t2 = time.time()
    print "%f seconds" % (t2 - t1)

    print '=' * 40

    print "Using compiled C extension..."
    cProfile.run('arr2 = tree_to_recarray(tree, branches=branches,'
                 'include_weight=True)')

    print "time without profiler overhead:"
    t1 = time.time()
    arr2 = tree_to_recarray(tree, branches=branches,
                            include_weight=True)
    t2 = time.time()
    print "%f seconds" % (t2 - t1)

    print '=' * 40
    print "Comparison of output:"

    print
    print "Python result:"
    print arr1
    print arr1['a_x']
    print arr1['weight']

    print
    print "C result:"
    print arr2
Пример #5
0

# define the model
class Event(TreeModel):

    x = FloatCol()
    y = FloatCol()
    z = FloatCol()
    i = IntCol()

tree = Tree("test", model=Event)

# fill the tree
for i in xrange(100000):
    tree.x = gauss(.5, 1.)
    tree.y = gauss(.3, 2.)
    tree.z = gauss(13., 42.)
    tree.i = i
    tree.fill()
tree.write()

# convert tree into a numpy record array
from rootpy.root2array import tree_to_recarray, tree_to_ndarray
array = tree_to_recarray(tree)
print array
print array.x
print array.i
print tree_to_ndarray(tree)

f.close()
Пример #6
0
        table = None

        ROOT.TTreeCache.SetLearnEntries(1)

        time_in_read = 0
        time_in_append = 0
        processed_files = 0

        try:
            for file_chunk in chunk_files(flat_files, args.chainsize):
                t1 = time.time()
                chunk_chain = ROOT.TChain(tree)
                chunk_chain.SetCacheSize(10000000)
                for file in file_chunk:
                    chunk_chain.Add(file)
                recarray = tree_to_recarray(chunk_chain, None, False)
                t2 = time.time()
                time_in_read += t2 - t1

                t1 = time.time()
                if table is None:
                    log.info("Creating table")
                    table = hfile.createTable(args.where,
                                              h5name,
                                              recarray,
                                              "",
                                              filters=filters,
                                              expectedrows=entries)
                else:
                    table.append(recarray)
                table.flush()
Пример #7
0
    print "time without profiler overhead:"
    t1 = time.time()
    arr1 = tree_to_recarray_py(tree, branches=branches, include_weight=True)
    t2 = time.time()
    print "%f seconds" % (t2 - t1)

    print '=' * 40

    print "Using compiled C extension..."
    cProfile.run('arr2 = tree_to_recarray(tree, branches=branches,'
                 'include_weight=True)')

    print "time without profiler overhead:"
    t1 = time.time()
    arr2 = tree_to_recarray(tree, branches=branches, include_weight=True)
    t2 = time.time()
    print "%f seconds" % (t2 - t1)

    print '=' * 40
    print "Comparison of output:"

    print
    print "Python result:"
    print arr1
    print arr1['a_x']
    print arr1['weight']

    print
    print "C result:"
    print arr2