RatCAVE_Round2_Backup/Code/CliffAnalysis/scripts/merge_sessions.py

"""
Takes all processed_data .h5 session files and puts just the Rat data (position and orientation) into a single
file, along with the metadata.  Makes for easier analysis.
"""

import os
from glob import glob
import pandas as pd
import h5py


loadpath = '/home/nickdg/theta_storage/data/VR_Experiments_Round_2/processed_data_by_experiment/VRCliffExp'
fnames = glob(loadpath + '/*/*.h5', recursive=True)

def merge_dataframes(fnames):
    pos, ori = [], []
    for idx, fname in enumerate(fnames):
        try:
            for key, dfs in (('Position', pos), ('Orientation', ori)):
                df = pd.read_hdf(fname, '/preprocessed/Rigid Body/Rat/'+key).dropna()
                df['Session_id'] = idx
                dfs.append(df)
        except KeyError:
            print(f'Skipped {fname}')

    pos = pd.concat(pos)
    ori = pd.concat(ori)
    return pos, ori
        

if __name__ == '__main__':
    names = ['CLIFF_TYPE', 'CLIFF_SIDE', 'RAT', 'PAPER_LOG_CODE', 'Capture Start Time']
    sessions = []
    for idx, fname in enumerate(fnames):
        sess_dict = {}
        with h5py.File(fname, 'r') as f:
            try:
                for name in names:
                    new_name = name.replace('_', ' ').title().replace(' ', '')
                    sess_dict[new_name] = f.attrs[name]
            except KeyError:
                continue
            sess_dict['id'] = idx

        board_height = pd.read_hdf(fname, '/preprocessed/Rigid Body/Board/Position').dropna().Y
        sess_dict['BoardHeight'] = board_height.median()
        
        sessions.append(sess_dict)
    sessions = pd.DataFrame(sessions)


    save_fname = '../data/VRCliff_AllData.h5'
    sessions.to_hdf(save_fname, '/Session')
    pos, ori = merge_dataframes(fnames)
    pos.to_hdf(save_fname, '/Position')
    ori.to_hdf(save_fname, '/Orientation')