RatCAVE_Round2_Backup/Code/WallAnalysis/scripts/merge_data.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_labshare/Projects/VRExperiments_Round2/Data/motion_tracking/processed'
fnames = glob(loadpath + '/*/VRWall*.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 = ['RAT', 'PAPER_LOG_CODE', 'Capture Start Time', 'VR_WALL_X_OFFSET']
    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]
                
                phase_frames = f['events/phaseStartFrameNum'][:]
            except KeyError:
                continue
            sess_dict['id'] = idx
            print(phase_frames)
            for phase, frame in enumerate(phase_frames):
                sess_dict[f'PhaseFrame{phase}'] = frame
        
        sessions.append(sess_dict)
    sessions = pd.DataFrame(sessions)


    save_fname = '../data/VRWall_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')