utilities.py

# Utility functions for the aCCE model scripts.

import os

import numpy as np
import pandas as pd

import model_core


OFF = 0
LEFT = 1
RIGHT = 2
THUMB = 1
INDEX = 2
UPRIGHT = 1
UPSIDE_DOWN = -1


def is_congruent(trial):
    if trial[0] == UPRIGHT:
        if trial[1] == THUMB:
            return trial[2] == LEFT
        return trial[2] == RIGHT
    if trial[1] == THUMB:
        return trial[2] == RIGHT
    return trial[2] == LEFT


def get_rt_df_(RTs):
    is_congr = np.array([is_congruent(trial) for trial in RTs[:, 1:]])
    congruent_upright_rts = RTs[:, 0][is_congr & (RTs[:, 1] == UPRIGHT)]
    incongruent_upright_rts = RTs[:, 0][~is_congr & (RTs[:, 1] == UPRIGHT)]
    congruent_upside_down_rts = RTs[:, 0][is_congr & (RTs[:, 1] == UPSIDE_DOWN)]
    incongruent_upside_down_rts = RTs[:, 0][~is_congr & (RTs[:, 1] == UPSIDE_DOWN)]
    return pd.DataFrame({
        "Condition": (
            ["Congruent Upright"] * len(congruent_upright_rts)
            + ["Incongruent Upright"] * len(incongruent_upright_rts)
            + ["Congruent Upside-Down"] * len(congruent_upside_down_rts)
            + ["Incongruent Upside-Down"] * len(incongruent_upside_down_rts)
        ),
        "RT": np.concatenate([
            congruent_upright_rts,
            incongruent_upright_rts,
            congruent_upside_down_rts,
            incongruent_upside_down_rts,
        ]),
    })


def get_means_datasets(trial_datasets):
    datasets = np.zeros((8, trial_datasets.shape[1], trial_datasets.shape[2]))
    num_runs = trial_datasets.shape[0] // 8

    for i in range(2):
        for j in range(2):
            for k in range(2):
                tr = [i * 4 + j * 2 + k + r * 8 for r in range(num_runs)]
                trial_data_ = trial_datasets[tr, :, :]
                trial_data = np.mean(trial_data_, axis=0)

                response_starts = [
                    np.where(trial_data_[trial_idx, :, 10] != 0)[0]
                    for trial_idx in range(len(trial_data_))
                ]
                response_starts = [starts for starts in response_starts if starts.size != 0]
                response_starts = [starts[0] for starts in response_starts]
                resp = np.mean(response_starts)

                trial_data[:, 10] = 0
                try:
                    trial_data[int(np.round(resp)), 10] = 1
                except:
                    trial_data[0, 10] = 1

                datasets[i * 4 + j * 2 + k, :, :] = trial_data

    return datasets


def get_means_entropies(entropies):
    datasets = np.zeros((8, entropies.shape[1], entropies.shape[2], entropies.shape[3]))
    num_runs = entropies.shape[0] // 8

    for i in range(2):
        for j in range(2):
            for k in range(2):
                tr = [i * 4 + j * 2 + k + r * 8 for r in range(num_runs)]
                datasets[i * 4 + j * 2 + k, :, :, :] = np.mean(entropies[tr, :, :, :], axis=0)

    return datasets


def get_means_further_trials(trial_datasets):
    trial_data = np.mean(trial_datasets, axis=0)
    response_starts = [
        np.where(trial_datasets[trial_idx, :, 10] != 0)[0]
        for trial_idx in range(len(trial_datasets))
    ]
    response_starts = [starts for starts in response_starts if starts.size != 0]
    response_starts = [starts[0] for starts in response_starts]
    resp = np.mean(response_starts)

    trial_data[:, 10] = 0
    try:
        trial_data[int(np.round(resp)), 10] = 1
    except:
        trial_data[0, 10] = 1
    return trial_data


def get_file_name(test, dlr_curr, prs_curr, div_curr, r_curr, decay_par_curr, b_curr,
                  bins_curr, l_r_curr, sfn_curr, noise_curr,
                  timesteps_reach_shorter_curr, timesteps_reach_longer_curr,
                  timestep_bottle_appearance_curr, timestep_stim_curr, len_reps_curr,
                  timesteps_transport_curr):
    test_curr = test
    test_curr += f"dlr{dlr_curr}_"
    test_curr += (f"prs{prs_curr}_" if prs_curr != "noprs" else "noprs_")
    test_curr += f"div_{div_curr}_r_{r_curr}_decay{decay_par_curr}_b{b_curr}_{bins_curr}bins_t_"
    test_curr += f"lr_{str(l_r_curr)}_sfn_{str(sfn_curr)}_"

    if not isinstance(noise_curr, list):
        test_curr += f"nb_{str(noise_curr)}_"
    else:
        test_curr += f"nb_{str(noise_curr[0])}_"
        if noise_curr[1] is not None:
            test_curr += f"r_{str(noise_curr[1])}_"
        if noise_curr[2] is not None:
            test_curr += f"uni_{str(noise_curr[2])}_"

    if timesteps_reach_shorter_curr != 100 or timesteps_reach_longer_curr != 120:
        test_curr += f"t_reach_{timesteps_reach_shorter_curr}_{timesteps_reach_longer_curr}_"
    if timesteps_transport_curr != 100:
        test_curr += f"t_tr_{timesteps_transport_curr}_"
    test_curr += f"t_bottle_{timestep_bottle_appearance_curr}_stim_{timestep_stim_curr}"
    if len_reps_curr != 100:
        test_curr += f"_reps_{len_reps_curr}"
    return test_curr


def run_rts_over_time(testfolder, folder="Exp-Data/", alpha_lr=model_core.alpha_lr,
                      threshold=model_core.threshold, t_to_run=None):
    res_folder = folder + "RTs_sequence/"
    os.makedirs(res_folder, exist_ok=True)

    res_file_name_base = res_folder + testfolder + "_RTs_seq"
    res_file_name_compr = res_file_name_base + "_c.csv"
    res_file_name_uncompr = res_file_name_base + "_all.csv"

    if os.path.exists(res_file_name_uncompr):
        RTs_seq = np.loadtxt(res_file_name_uncompr, delimiter="\t")
        print("Loaded existing RTs_seq data (uncompressed)")
    elif os.path.exists(res_file_name_compr):
        RTs_seq = np.loadtxt(res_file_name_compr, delimiter="\t")
        print("Loaded existing RTs_seq data (compressed)")
    else:
        if t_to_run is None:
            t_stim_to_run = range(100)
        else:
            t_stim_to_run = range(t_to_run)
        RTs_seq = np.zeros((len(t_stim_to_run), 17))

        filename_uncompressed = folder + testfolder + "_all.npz"
        try:
            all_data = np.load(filename_uncompressed, allow_pickle=True)
            compressed = False
        except:
            all_data = np.load(folder + testfolder + ".npz", allow_pickle=True)
            compressed = True

        trial_datasets = all_data["res_data"]
        if len(trial_datasets) == 8:
            compressed = False

        RTs_seq_all = np.zeros((len(t_stim_to_run), len(trial_datasets)))
        orig_rt_data = all_data["rt_data"]
        compressed_data = len(trial_datasets) < len(orig_rt_data)
        if compressed_data:
            orig_rt_data = orig_rt_data[0:8, :]

        is_congr = [is_congruent(trial) for trial in orig_rt_data[:, 1:]]
        is_upright = orig_rt_data[:, 1] == model_core.UPRIGHT
        is_congr_idxes = [i for i in range(len(is_congr)) if is_congr[i]]
        is_incongr_idxes = [i for i in range(len(is_congr)) if not is_congr[i]]
        is_upright_idxes = [i for i in range(len(is_upright)) if is_upright[i]]
        is_updown_idxes = [i for i in range(len(is_upright)) if not is_upright[i]]
        is_upr_congr_idxes = [
            i for i in range(len(is_upright)) if is_upright[i] and is_congr[i]
        ]
        is_upr_incongr_idxes = [
            i for i in range(len(is_upright)) if is_upright[i] and not is_congr[i]
        ]
        is_updown_congr_idxes = [
            i for i in range(len(is_upright)) if not is_upright[i] and is_congr[i]
        ]
        is_updown_incongr_idxes = [
            i for i in range(len(is_upright)) if not is_upright[i] and not is_congr[i]
        ]

        for timestep_stim in t_stim_to_run:
            for trial_idx in range(len(trial_datasets)):
                trial_data = trial_datasets[trial_idx, :, :]
                F_all = trial_data[:, -5:]
                t_bottle = int(trial_data[0, 3])
                done = False
                t = 0
                M_grasp = model_core.M_default
                P_S = np.array([0.5, 0.5])
                P_R = np.array([0.5, 0.5])
                distractor = np.array([0, 0])
                stimulation = model_core.to_one_hot(model_core.OFF)
                t_stim_orig_data = int(trial_data[0, 4])
                stimulation_value = trial_data[t_stim_orig_data, 1]
                distractor_value = trial_data[t_stim_orig_data, 2]

                while not done:
                    if t == t_bottle:
                        bottle = trial_data[t_bottle, 0]
                        M_grasp = (
                            model_core.M_upsidedown
                            if bottle == model_core.UPSIDE_DOWN
                            else model_core.M_upright
                        )
                    if t == timestep_stim:
                        stimulation = model_core.to_one_hot(stimulation_value)
                        distractor = model_core.to_one_hot(distractor_value)

                    V_in = 0
                    M = [
                        model_core.M_idle,
                        model_core.M_reach,
                        M_grasp,
                        model_core.M_transport,
                        model_core.M_putdown,
                    ]
                    F = F_all[t]
                    for i in range(len(M)):
                        V_in = V_in + F[i] * M[i] @ distractor
                    T_in = stimulation
                    MP = model_core.normalize(model_core.W_vis * V_in + model_core.W_tac * T_in)
                    G_S = np.array(np.mean(MP, axis=0))
                    P_S = P_S + alpha_lr * G_S
                    P_S = P_S / sum(P_S)
                    G_R = P_S
                    P_R = P_R + alpha_lr * G_R
                    P_R = P_R / sum(P_R)
                    if np.max(P_R) >= threshold:
                        done = True
                        RTs_seq_all[timestep_stim, trial_idx] = t
                    t = t + 1

            RTs_cur = RTs_seq_all[timestep_stim]
            mean_rt_upright = np.mean(RTs_cur[is_upright_idxes])
            mean_rt_upsidedown = np.mean(RTs_cur[is_updown_idxes])
            mean_rt_congruent = np.mean(RTs_cur[is_congr_idxes])
            mean_rt_incongruent = np.mean(RTs_cur[is_incongr_idxes])

            mean_rt_congruent_upright = np.mean(RTs_cur[is_upr_congr_idxes])
            mean_rt_congruent_upsidedown = np.mean(RTs_cur[is_updown_congr_idxes])
            mean_rt_incongruent_upright = np.mean(RTs_cur[is_upr_incongr_idxes])
            mean_rt_incongruent_upsidedown = np.mean(RTs_cur[is_updown_incongr_idxes])

            std_congruent = np.std(RTs_cur[is_congr_idxes])
            std_incongruent = np.std(RTs_cur[is_incongr_idxes])
            std_upright = np.std(RTs_cur[is_upright_idxes])
            std_upsidedown = np.std(RTs_cur[is_updown_congr_idxes])
            std_congruent_upright = np.std(RTs_cur[is_upr_congr_idxes])
            std_congruent_upsidedown = np.std(RTs_cur[is_updown_congr_idxes])
            std_incongruent_upright = np.std(RTs_cur[is_upr_incongr_idxes])
            std_incongruent_upsidedown = np.std(RTs_cur[is_updown_incongr_idxes])

            RTs_seq[timestep_stim, 0] = mean_rt_incongruent - mean_rt_congruent
            RTs_seq[timestep_stim, 1] = mean_rt_congruent
            RTs_seq[timestep_stim, 2] = mean_rt_incongruent
            RTs_seq[timestep_stim, 3] = mean_rt_upright
            RTs_seq[timestep_stim, 4] = mean_rt_upsidedown
            RTs_seq[timestep_stim, 5] = mean_rt_congruent_upright
            RTs_seq[timestep_stim, 6] = mean_rt_congruent_upsidedown
            RTs_seq[timestep_stim, 7] = mean_rt_incongruent_upright
            RTs_seq[timestep_stim, 8] = mean_rt_incongruent_upsidedown
            RTs_seq[timestep_stim, 9] = std_congruent
            RTs_seq[timestep_stim, 10] = std_incongruent
            RTs_seq[timestep_stim, 11] = std_upright
            RTs_seq[timestep_stim, 12] = std_upsidedown
            RTs_seq[timestep_stim, 13] = std_congruent_upright
            RTs_seq[timestep_stim, 14] = std_congruent_upsidedown
            RTs_seq[timestep_stim, 15] = std_incongruent_upright
            RTs_seq[timestep_stim, 16] = std_incongruent_upsidedown

            print("Finished timestep_stim = ", timestep_stim, "(acce = ", RTs_seq[timestep_stim, 0], ")")

        if compressed:
            np.savetxt(res_file_name_compr, RTs_seq, delimiter="\t")
        else:
            np.savetxt(res_file_name_uncompr, RTs_seq, delimiter="\t")

    return RTs_seq