add_attributes.py

Go to the documentation of this file.

from __future__ import print_function
import h5py
import argparse
import numpy as np

"""
Turns out it is useful for hdf5 groups to be self-describing for the 
PandAna framework. This is done by adding an HDF5 attribute to each group
that specifies the columns within that group that should be used as indices 
when constructing a pandas DataFrame. 

HDF5Maker was updated to produces these attributes at the time of file creation
For those files that were created before HDF5Maker was updated,
this script adds the attributes by hand

author: Derek.Doyle@colostate.edu
date: Feb 25, 2021
"""

# Here's the bit that makes this script NOvA-specific. Other
# experiments will have different indexing schemas 
# All groups will be searched for datasets matching these names.
# When groups contain these datasets, we will add these to the index_cols 
# attribute. Additional index columns will be determined by looking for
# datasets ending with '_idx'
required_indices = ['run', 'subrun', 'cycle', 'evt']
other_indices = ['subevt']

def determine_index_cols(h5group):
    index_cols = []
    dsets = h5group.keys()
    for required in required_indices:
        if required in dsets: index_cols.append(required)
    for other in other_indices:
        if other in dsets: index_cols.append(other)

    for dset in h5group.keys():
        if dset.endswith('_idx'): index_cols.append(dset)
    return index_cols

def encode_index_cols(index_cols):
    return [','.join(index_cols)]


# nova group naming scheme provides a full "path" to the data
# as arranged in standard record where nested objects
# are separated by a '.'
# The parentage can then be determined by removing 
# name after the last '.'
def parent_table(group_name):
    if '.' not in group_name: 
        return ['/']
    else: 
        return ['.'.join(group_name.split('.')[:-1])]


# function to write all attributes
def write_attributes(h5group, verbose=False, dry_run=True):
    if verbose:
        print(h5group.name)
    # index_cols
    index_cols = determine_index_cols(g)
    attr_index_cols = encode_index_cols(index_cols)
    if verbose:
        try:
            existing_index_cols = g.attrs['index_cols']
        except KeyError:
            existing_index_cols = 'None'
            
        finally:
            print('\t|index_cols: {} ---> {}'.format(existing_index_cols, attr_index_cols))

    if not dry_run: h5group.attrs['index_cols'  ] = attr_index_cols

    # parent_table
    attr_parent_table = parent_table(h5group.name)
    if verbose:
        try:
            existing_parent_table = g.attrs['parent_table']
        except KeyError:
            existing_parent_table = 'None'
        finally:
            print('\t|parent_table: {} ---> {}'.format(existing_parent_table, attr_parent_table))

    if not dry_run: h5group.attrs['parent_table'] = attr_parent_table



# progress bar for funzies    
def insert(astring, substr, pos, overwrite=True):
    if overwrite:
        return astring[:pos] + substr + astring[pos + len(substr):]
    else:
        return astring[:pos] + substr + astring[pos:]

class progbar:
    invert_open  = '\033[F\033[7m['
    invert_close = '\033[0m'
    def __init__(self, message, nfields=50):
        print(message, '....')
        self._nfields = nfields

    def __call__(self, prog):
        percent = prog
        percent_str = '{:.1f}%'.format(percent*100)
        current = int(percent * self._nfields)
        bar = progbar.invert_open + ' ' * (self._nfields - 1) + ']'
        bar = insert(bar, percent_str         , pos=len(progbar.invert_open)        , overwrite=True)
        bar = insert(bar, progbar.invert_close, pos=len(progbar.invert_open)+current, overwrite=False)
        print(bar)
        if current == self._nfields: print()

parser = argparse.ArgumentParser('Add index columns to the input h5 file(s)')
parser.add_argument('files', nargs='*', help='Input file(s). Wildcard allowed')
parser.add_argument('--dry_run', '-d', action='store_true', help='Don\'t actually add the attributes yet')
parser.add_argument('--verbose', '-v', action='store_true', help='Print out old and new attributes')

args = parser.parse_args()

files = args.files

if not args.verbose: prog = progbar('Adding attributes')
for i, fname in enumerate(files):
    with h5py.File(fname, 'r+') as h5file:
        for group_name in h5file.keys():
            if group_name != 'MetaData':
                g = h5file.get(group_name)

                write_attributes(g, 
                                 verbose=args.verbose, 
                                 dry_run=args.dry_run)

    if not args.verbose: prog(i / len(files))