pipeline/analysis/loading.py

import sys, os
sys.path.append(os.path.join(os.getcwd(), '..'))
sys.path.append(os.path.join(os.getcwd(), '..', '..'))

from imports import *
from target import build_tgt_matrix


def load_session_data(session):
    all_areas = ['A1', 'PPC', 'HPC']

    animal      = session.split('_')[0]
    sessionpath = os.path.join(source, animal, session)
    aeps_file   = os.path.join(sessionpath, 'AEPs.h5')
    h5_file     = os.path.join(sessionpath, session + '.h5')
    report_path = os.path.join(report, 'PSTH', session)
    if not os.path.exists(report_path):
        os.makedirs(report_path)
        
    # load timeline and configuration
    with h5py.File(h5_file, 'r') as f:
        tl = np.array(f['processed']['timeline'])  # time, X, Y, speed, etc.
        trials = np.array(f['processed']['trial_idxs'])  # t_start_idx, t_end_idx, x_tgt, y_tgt, r_tgt, result
        cfg = json.loads(f['processed'].attrs['parameters'])

    # load AEPs
    aeps = {}
    with h5py.File(aeps_file, 'r') as f:
        for area in all_areas:
            if not area in f:
                continue
            aeps[area] = np.array(f[area]['aeps'])
        aeps_events = np.array(f['aeps_events'])

    areas = list(aeps.keys())
    # TODO find better way. Remove outliers
    if 'A1' in areas:
        aeps['A1'][aeps['A1'] >  5000]   =   5000
        aeps['A1'][aeps['A1'] < -5000]   =  -5000
    if 'PPC' in areas:
        aeps['PPC'][aeps['PPC'] >  1500] =   1500
        aeps['PPC'][aeps['PPC'] < -1500] =  -1500
    if 'HPC' in areas:
        aeps['HPC'][aeps['HPC'] >  1500] =   1500
        aeps['HPC'][aeps['HPC'] < -1500] =  -1500
    aeps[areas[0]].shape
    
    # load LFP
    lfp = {}
    with h5py.File(aeps_file, 'r') as f:
        for area in areas:
            if 'LFP' in f[area]:
                lfp[area] = np.array(f[area]['LFP'])
    
    # load AEP metrics
    AEP_metrics_lims = dict([(area, {}) for area in areas])
    AEP_metrics_raw  = dict([(area, {}) for area in areas])
    AEP_metrics_norm = dict([(area, {}) for area in areas])
    with h5py.File(aeps_file, 'r') as f:
        for area in areas:
            grp = f[area]

            for metric_name in grp['raw']:
                AEP_metrics_raw[area][metric_name]  = np.array(grp['raw'][metric_name])
                AEP_metrics_norm[area][metric_name] = np.array(grp['norm'][metric_name])
                AEP_metrics_lims[area][metric_name] = [int(x) for x in grp['raw'][metric_name].attrs['limits'].split(',')]
                
    # build target matrix
    tgt_matrix = build_tgt_matrix(tl, trials, aeps_events)

    # get available units
    unit_names = []
    with h5py.File(h5_file, 'r') as f:
        unit_names = [x for x in f['units']]

    # read single units
    single_units = {}
    spike_times = {}

    with h5py.File(h5_file, 'r') as f:
        for unit_name in unit_names:
            spike_times[unit_name] = np.array(f['units'][unit_name][H5NAMES.spike_times['name']])
            single_units[unit_name] = np.array(f['units'][unit_name][H5NAMES.inst_rate['name']])
            #single_units[unit_name] = instantaneous_rate(unit_times, tl[:, 0], k_width=50)
            
    return {
        'tl': tl,
        'trials': trials,
        'cfg': cfg,
        'areas': areas,
        'aeps': aeps,
        'aeps_events': aeps_events,
        'lfp': lfp,
        'AEP_metrics_lims': AEP_metrics_lims,
        'AEP_metrics_raw': AEP_metrics_raw,
        'AEP_metrics_norm': AEP_metrics_norm,
        'tgt_matrix': tgt_matrix,
        'single_units': single_units,
        'spike_times': spike_times,
        'unit_names': unit_names,
        'animal': animal,
        'aeps_file': aeps_file,
        'h5_file': h5_file,
        'report_path': report_path
    }