pupil_batch_resampling

"""
Script used for synchronizing pupil data across labs by resampling/downsampling to a common timescale.
Default is 30 Hz.

Outputs are new hdf5 files with all the original contents (metadata included) and the following new datasets:
- "pupil_diameter_resampled"
- "pupil_timestamps_resampled"
- "pupil_isnan_resampled"
- "pupil_norm_pos_x_resampled"
- "pupil_norm_pos_y_resampled"

Note: screen brightness features are not resampled.
"""

import numpy as np
import h5py
import glob
import scipy.io
from scipy import interpolate
import os


# inputs
root_dir = '/media/lucab/data_hdd/screen_corners'  # directory where preprocessed pupil data lives
behav_dir = '/media/lucab/data_hdd/commgame_behav_timing_files'  # directory of behavioral data
output_dir = '/media/lucab/data_hdd/CommGame_pupil_resampled'
pairNo_range = (31, 177)  # both values are inclusive!
session_names = ['BG1', 'BG2', 'BG3', 'BG4', 'BG5', 'BG6', 'BG7', 'BG8']
interp_fs = 120  # the original sampling rate of pupil size data
desired_fs = 30  # the new sampling rate to downsample to
factor = int(interp_fs/desired_fs)

blacklist = ('/prepro_pl_data_pair36_Mordor_BG1_residual.hdf5',
             '/prepro_pl_data_pair36_Gondor_BG1_residual.hdf5',
             '/prepro_pl_data_pair36_Mordor_BG2_residual.hdf5',
             '/prepro_pl_data_pair36_Gondor_BG2_residual.hdf5',
             '/prepro_pl_data_pair36_Mordor_BG3_residual.hdf5',
             '/prepro_pl_data_pair36_Gondor_BG3_residual.hdf5',
             '/prepro_pl_data_pair36_Mordor_freeConv_residual.hdf5',
             '/prepro_pl_data_pair36_Gondor_freeConv_residual.hdf5',
             '/prepro_pl_data_pair37_Mordor_freeConv_residual.hdf5',
             '/prepro_pl_data_pair59_Mordor_freeConv_residual.hdf5',
             '/prepro_pl_data_pair36_Mordor_BG1.hdf5',
             '/prepro_pl_data_pair36_Gondor_BG1.hdf5',
             '/prepro_pl_data_pair36_Mordor_BG2.hdf5',
             '/prepro_pl_data_pair36_Gondor_BG2.hdf5',
             '/prepro_pl_data_pair36_Mordor_BG3.hdf5',
             '/prepro_pl_data_pair36_Gondor_BG3.hdf5',
             '/prepro_pl_data_pair36_Mordor_freeConv.hdf5',
             '/prepro_pl_data_pair36_Gondor_freeConv.hdf5',
             '/prepro_pl_data_pair37_Mordor_freeConv.hdf5',
             '/prepro_pl_data_pair59_Mordor_freeConv.hdf5'
             # '/prepro_pl_data_pair104_Gondor_BG3_residual.hdf5',
             # '/prepro_pl_data_pair104_Mordor_BG3_residual.hdf5',
             # '/prepro_pl_data_pair121_Gondor_BG2_residual.hdf5',
             # '/prepro_pl_data_pair121_Mordor_BG2_residual.hdf5'
             )


# FUNCTIONS
def data_collect(root_dir, pairNo_range, session_names):

    pairNos = range(pairNo_range[0], pairNo_range[1] + 1)
    labNames = ["Mordor", "Gondor"]

    mordor_files = []
    mordor_mat_files = []
    mordor_ids_pairNo = []
    mordor_ids_session = []

    for pairNo in pairNos:
            for session in session_names:

                mordor_file = glob.glob(f'{root_dir}/**/prepro_pl_data_pair{pairNo}_Mordor_{session}_residual.hdf5',
                                        recursive=True)
                mordor_mat_file = glob.glob(f'{behav_dir}/**/pair{pairNo}_Mordor_behav/pair{pairNo}_Mordor_{session}_**imes.mat',
                              recursive=True)
                if len(mordor_file) == 0:
                    mordor_file = ['Invalid']
                    mordor_mat_file = ['Invalid']
                mordor_files.append(mordor_file[0])
                mordor_mat_files.append(mordor_mat_file[0])
                mordor_ids_pairNo.append(pairNo)
                mordor_ids_session.append(session)

    gondor_files = []
    gondor_mat_files = []
    gondor_ids_pairNo = []
    gondor_ids_session = []

    for pairNo in pairNos:
            for session in session_names:
                gondor_file = glob.glob(f'{root_dir}/**/prepro_pl_data_pair{pairNo}_Gondor_{session}_residual.hdf5',
                                        recursive=True)
                gondor_mat_file = glob.glob(f'{behav_dir}/**/pair{pairNo}_Gondor_behav/pair{pairNo}_Gondor_{session}_**imes.mat',
                              recursive=True)
                if len(gondor_file) == 0:
                    gondor_file = ['Invalid']
                    gondor_mat_file = ['Invalid']
                gondor_files.append(gondor_file[0])
                gondor_mat_files.append(gondor_mat_file[0])
                gondor_ids_pairNo.append(pairNo)
                gondor_ids_session.append(session)


    # check for asymmetry in input data
    if not len(mordor_files) == len(gondor_files):
        print('\n!!!WARNING!!! '
              '\nFound an asymmetry in the number of recordings across labs!!!'
              '\n!!!WARNING!!!\n'
              '\nNumber of files for Mordor: {}'
              '\nNUmber of files for Gondor: {}' .format(len(mordor_files), len(gondor_files)))

    print(mordor_files, gondor_files)

    paired_files = list(zip(mordor_files, gondor_files))
    paired_mat_files = list(zip(mordor_mat_files, gondor_mat_files))
    paired_ids_pairNo = np.column_stack((mordor_ids_pairNo, gondor_ids_pairNo))
    paired_ids_session = np.column_stack((mordor_ids_session, gondor_ids_session))

    return paired_files, paired_mat_files, paired_ids_pairNo, paired_ids_session


def h5_reader(filepath):
    if filepath == 'Invalid':
        print("\n Not a valid path, all data will be NaN!")
        pupil_timestamps = np.nan
        pupil_diameter = np.nan
        pupil_norm_pos_x = np.nan
        pupil_norm_pos_y = np.nan
        pupil_isnan = np.nan
    else:
        print('\nReading {} ...' .format(filepath))
        f = h5py.File(filepath, 'r')
        f.visititems(print)
        for key in f.keys():
            pair = f[key]
            for lab in pair.keys():
                labName = pair[lab]
                for sess in labName.keys():
                    session = labName[sess]
                    for key in session.attrs.keys():
                        print('{} : {}' .format(key, session.attrs[key]))
                    pupil_timestamps = np.array(session['pupil_timestamps'])
                    pupil_diameter = np.array(session['pupil_diameter'])
                    pupil_norm_pos_x = np.array(session['pupil_norm_pos_x'])
                    pupil_norm_pos_y = np.array(session['pupil_norm_pos_y'])
                    pupil_isnan = np.array(session['pupil_isnan'])
                    #pupil_brightness = session['pupil_brightness']
    #f.close()
    return pupil_timestamps, pupil_diameter, pupil_norm_pos_x, pupil_norm_pos_y, pupil_isnan, pair


def get_timing_info(mat_filepath, delay=3):
    video_times = scipy.io.loadmat(mat_filepath)
    startTime = float(video_times["sharedStartTime"].flatten())
    stopTime = float(video_times["stopCaptureTime"].flatten())
    # trim data some more since we want only the actual task
    startTime = startTime + delay
    stopTime = stopTime - delay

    return startTime, stopTime


################################################################################################################

# MAIN PART #

# collect pupil and behavioral data
paired_files, paired_mat_files, paired_ids_pairNo, paired_ids_session = data_collect(root_dir,
                                                                                     pairNo_range,
                                                                                     session_names)
# loop through the file lists (pupil hdf5 and timing .mat files)
for paired_file, paired_mat_file, pairNo, session in zip(paired_files,
                                                         paired_mat_files,
                                                         paired_ids_pairNo,
                                                         paired_ids_session):

    # check for missing files
    if (paired_file[0] == 'Invalid' or paired_file[1] == 'Invalid'
            or paired_mat_file[0] == 'Invalid' or paired_mat_file[1] == 'Invalid'):
        print("Will skip this pair, missing / corrupt files!")
        continue
    if paired_file[0].endswith(blacklist) or paired_file[1].endswith(blacklist):
        print('Blacklist file!')
        continue

    else:

        # load hdf5 files for mordor and for gondor
        (mordor_timestamps, mordor_diameter, mordor_norm_pos_x,
         mordor_norm_pos_y, mordor_isnan, mordor_h5group) = h5_reader(paired_file[0])
        (gondor_timestamps, gondor_diameter, gondor_norm_pos_x,
         gondor_norm_pos_y, gondor_isnan, gondor_h5group) = h5_reader(paired_file[1])

        # load timing info for later trimming
        startTimeM, stopTimeM = get_timing_info(paired_mat_file[0])
        startTimeG, stopTimeG = get_timing_info(paired_mat_file[1])

        # decide which one to use to trim both recordings
        if startTimeM < startTimeG:
            shared_startTime = startTimeG
        else:
            shared_startTime = startTimeM
        if stopTimeM < stopTimeG:
            shared_stopTime = stopTimeM
        else:
            shared_stopTime = stopTimeG

        # exception for when the pupil recording abruptly ends and the timestamps are not quite right
        if mordor_timestamps[-1] < shared_stopTime:
            shared_stopTime = mordor_timestamps[-1]
        if gondor_timestamps[-1] < shared_stopTime:
            shared_stopTime = gondor_timestamps[-1]

        # cut data between common start and common stop
        idx_startM = np.where(mordor_timestamps > shared_startTime)
        idx_startG = np.where(gondor_timestamps > shared_startTime)
        idx_stopM = np.where(mordor_timestamps >= shared_stopTime)
        idx_stopG = np.where(gondor_timestamps >= shared_stopTime)

        mordor_timestamps_trimmed = mordor_timestamps[idx_startM[0][0]:idx_stopM[0][0]]
        gondor_timestamps_trimmed = gondor_timestamps[idx_startG[0][0]:idx_stopG[0][0]]
        mordor_diameter_trimmed = mordor_diameter[idx_startM[0][0]:idx_stopM[0][0]]
        gondor_diameter_trimmed = gondor_diameter[idx_startG[0][0]:idx_stopG[0][0]]
        mordor_isnan_trimmed = mordor_isnan[idx_startM[0][0]:idx_stopM[0][0]]
        gondor_isnan_trimmed = gondor_isnan[idx_startG[0][0]:idx_stopG[0][0]]
        mordor_norm_pos_x_trimmed = mordor_norm_pos_x[idx_startM[0][0]:idx_stopM[0][0]]
        gondor_norm_pos_x_trimmed = gondor_norm_pos_x[idx_startG[0][0]:idx_stopG[0][0]]
        mordor_norm_pos_y_trimmed = mordor_norm_pos_y[idx_startM[0][0]:idx_stopM[0][0]]
        gondor_norm_pos_y_trimmed = gondor_norm_pos_y[idx_startG[0][0]:idx_stopG[0][0]]

        # total elapsed time should be really close between labs since we already trimmed it
        total_timeM = mordor_timestamps_trimmed[-1] - mordor_timestamps_trimmed[0]
        real_fsM = len(mordor_timestamps_trimmed) / total_timeM
        total_timeG = gondor_timestamps_trimmed[-1] - gondor_timestamps_trimmed[0]
        real_fsG = len(gondor_timestamps_trimmed) / total_timeG
        # still, we need to make sure to use a common time window
        if total_timeM > total_timeG:
            total_time_max = total_timeG
        else:
            total_time_max = total_timeM


        # ACTUAL RESAMPLING PART #

        # Interpolate both mordor and gondor recordings to common timestamps
        # in order to smooth out sampling rate (120 Hz)

        print("\nResampling...")

        # first, define common timestamps (we use this later to resample every datatype)
        interp_times = np.arange(shared_startTime,
                                 shared_startTime + total_time_max,
                                 1 / interp_fs)

        fM = interpolate.interp1d(mordor_timestamps_trimmed, mordor_diameter_trimmed,
                                  bounds_error=False, fill_value=np.nan)
        fG = interpolate.interp1d(gondor_timestamps_trimmed, gondor_diameter_trimmed,
                                  bounds_error=False, fill_value=np.nan)
        diameter_interpM = fM(interp_times)
        diameter_interpG = fG(interp_times)

        # replace possible NaN value(s) on edges caused by interpolation
        if np.sum(np.isnan(diameter_interpM)) > 0:
            valid_indices = np.where(~np.isnan(diameter_interpM))[0]
            diameter_interpM[0:valid_indices[0]] = diameter_interpM[valid_indices[0]]
            diameter_interpM[valid_indices[-1]:] = diameter_interpM[valid_indices[-1]]
        if np.sum(np.isnan(diameter_interpG)) > 0:
            valid_indices = np.where(~np.isnan(diameter_interpG))[0]
            diameter_interpG[0:valid_indices[0]] = diameter_interpG[valid_indices[0]]
            diameter_interpG[valid_indices[-1]:] = diameter_interpG[valid_indices[-1]]

        # now we can just take every n-th element from both to downsample the signal
        s = slice(0, None, factor)
        mordor_diameter_rs = diameter_interpM[s]
        gondor_diameter_rs = diameter_interpG[s]
        mordor_timestamps_rs = interp_times[s]
        gondor_timestamps_rs = interp_times[s]

        # resample isnan vector
        fM = interpolate.interp1d(mordor_timestamps_trimmed, mordor_isnan_trimmed,
                                  bounds_error=False, fill_value=np.nan, kind="nearest")
        fG = interpolate.interp1d(gondor_timestamps_trimmed, gondor_isnan_trimmed,
                                  bounds_error=False, fill_value=np.nan, kind="nearest")
        isnan_interpM = fM(interp_times)
        isnan_interpG = fG(interp_times)

        # replace possible NaN value(s) on edges caused by interpolation
        if np.sum(np.isnan(isnan_interpM)) > 0:
            valid_indices = np.where(~np.isnan(isnan_interpM))[0]
            isnan_interpM[0:valid_indices[0]] = isnan_interpM[valid_indices[0]]
            isnan_interpM[valid_indices[-1]:] = isnan_interpM[valid_indices[-1]]
        if np.sum(np.isnan(isnan_interpG)) > 0:
            valid_indices = np.where(~np.isnan(isnan_interpG))[0]
            isnan_interpG[0:valid_indices[0]] = isnan_interpG[valid_indices[0]]
            isnan_interpG[valid_indices[-1]:] = isnan_interpG[valid_indices[-1]]

        # again, take every n-th element to downsample
        mordor_isnan_rs = isnan_interpM[s].astype(bool)
        gondor_isnan_rs = isnan_interpG[s].astype(bool)

        # interpolate gaze position data
        # for this datatype we don't need to cut extra NaNs caused by interpolation, since we keep them anyways

        # x position
        fM = interpolate.interp1d(mordor_timestamps_trimmed, mordor_norm_pos_x_trimmed,
                                  bounds_error=False, fill_value=np.nan, kind="nearest")
        fG = interpolate.interp1d(gondor_timestamps_trimmed, gondor_norm_pos_x_trimmed,
                                  bounds_error=False, fill_value=np.nan, kind="nearest")
        norm_pos_x_interpM = fM(interp_times)
        norm_pos_x_interpG = fG(interp_times)
        # y position
        fM = interpolate.interp1d(mordor_timestamps_trimmed, mordor_norm_pos_y_trimmed,
                                  bounds_error=False, fill_value=np.nan, kind="nearest")
        fG = interpolate.interp1d(gondor_timestamps_trimmed, gondor_norm_pos_y_trimmed,
                                  bounds_error=False, fill_value=np.nan, kind="nearest")
        norm_pos_y_interpM = fM(interp_times)
        norm_pos_y_interpG = fG(interp_times)

        # downsample
        mordor_norm_pos_x_rs = norm_pos_x_interpM[s]
        gondor_norm_pos_x_rs = norm_pos_x_interpG[s]
        mordor_norm_pos_y_rs = norm_pos_y_interpM[s]
        gondor_norm_pos_y_rs = norm_pos_y_interpG[s]


        # SAVE EVERYTHING IN HDF5 #

        print("\nSaving the resampled data...")

        # define paths, filenames
        mordor_new_filename = ("prepro_pl_data_pair{}_Mordor_{}_residual_resampled.hdf5"
                               .format(pairNo[0], session[0]))
        gondor_new_filename = ("prepro_pl_data_pair{}_Gondor_{}_residual_resampled.hdf5"
                               .format(pairNo[1], session[1]))

        mordor_out_filepath = os.path.join(output_dir, mordor_new_filename)
        gondor_out_filepath = os.path.join(output_dir, gondor_new_filename)

        # copy original hdf5 file, then create new datasets
        f_srcM = h5py.File(paired_file[0])
        with h5py.File(mordor_out_filepath, 'a') as f_destM:
            for key in f_srcM.keys():
                f_srcM.copy(key, f_destM)

                for key in f_destM.keys():
                    pair = f_destM[key]
                    for lab in pair.keys():
                        labName = pair[lab]
                        for sess in labName.keys():
                            session = labName[sess]
                            session.create_dataset("pupil_diameter_resampled", data=mordor_diameter_rs)
                            session.create_dataset("pupil_timestamps_resampled", data=mordor_timestamps_rs)
                            session.create_dataset("pupil_isnan_resampled", data=mordor_isnan_rs)
                            session.create_dataset("pupil_norm_pos_x_resampled", data=mordor_norm_pos_x_rs)
                            session.create_dataset("pupil_norm_pos_y_resampled", data=mordor_norm_pos_y_rs)

                # f_dest.visititems(print)
                # for key in session.attrs.keys():
                #     print('{} : {}'.format(key, session.attrs[key]))

        f_srcG = h5py.File(paired_file[1])
        with h5py.File(gondor_out_filepath, 'a') as f_destG:
            for key in f_srcG.keys():
                f_srcG.copy(key, f_destG)

                for key in f_destG.keys():
                    pair = f_destG[key]
                    for lab in pair.keys():
                        labName = pair[lab]
                        for sess in labName.keys():
                            session = labName[sess]
                            session.create_dataset("pupil_diameter_resampled", data=gondor_diameter_rs)
                            session.create_dataset("pupil_timestamps_resampled", data=gondor_timestamps_rs)
                            session.create_dataset("pupil_isnan_resampled", data=gondor_isnan_rs)
                            session.create_dataset("pupil_norm_pos_x_resampled", data=gondor_norm_pos_x_rs)
                            session.create_dataset("pupil_norm_pos_y_resampled", data=gondor_norm_pos_y_rs)