<jupyter_start><jupyter_text>Figure 4This notebook recreates the figure panels included in Figure 4 of Lee et al. 2021.  Description of the DataThe data used in this notebook comes from the experiments described in Lee et al. 2021. Specifically, we have the behavioral and activity of a trained deep RL agent performing a evidence accumulation task from Engelhard et al. 2019.   Preparing the Data Importing required code packages and modules<jupyter_code>import pickle
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
import pandas as pd
import sys
from scipy.io import loadmat, savemat
import utils.cnnlstm_analysis_utils as utils
import seaborn as sns
from scipy import stats
from matplotlib import gridspec
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.lines import Line2D
import os<jupyter_output>Missing colon in file '/Users/sasha/.matplotlib/matplotlibrc', line 1 ('~/.matplotlib/matplotlibrc')<jupyter_text>downloading data<jupyter_code>load_prefix = '../../data/logs/VA_maze/'

with open(load_prefix + '5000t_mosttrain_nova_db.p', 'rb') as f:
     [actions_, rewards_, feats_, terms_, vs_, tow_counts_, episode_lengths, _]  = pickle.load(f)
f.close()

with open(load_prefix + '1000t_obses_nova_db.p', 'rb') as f:
     obses_  = pickle.load(f)[0]
f.close()

trial_info = loadmat(load_prefix + 'trialinfo_nova_db.mat')['trials']
raw_ypos_ = [x[:,1] for x in trial_info['position'][0]]

# observations from an empty run thru of the maze
obses_empty  = loadmat(load_prefix + 'emptymaze_runthru.mat')['obses']
ypos_obses  = loadmat(load_prefix + 'emptymaze_runthru.mat')['ypos']<jupyter_output><empty_output><jupyter_text>Plotting Parameters<jupyter_code># PLOTTING PARAMS
matplotlib.rcParams.update({'font.size': 15})
matplotlib.rcParams.update({'font.family': 'Arial'})

FONT_BG = 25
import matplotlib as mpl
mpl.rcParams['pdf.fonttype'] = 42 # allow text of pdf to be edited in illustrator

mpl.rcParams["axes.spines.right"] = False
mpl.rcParams["axes.spines.top"] = False<jupyter_output><empty_output><jupyter_text>Organizing DataPulling out the specific data that we will use for figure panels<jupyter_code>CUEP_LIM = 140

autocorr_start = 2
autocorr_end = -35

CUEP_LIM = 140
ypos_cuep = np.squeeze(np.dstack([ypos_t[:CUEP_LIM] for ypos_t in raw_ypos_])[:,:,0])

obses_empty_flat = obses_empty[:, :,:140].T.reshape((140, 68 * 120))<jupyter_output><empty_output><jupyter_text>Calculate Vector RPEs<jupyter_code>if os.path.exists(load_prefix + 'no_va/pes_nova_db.p'):

    with open(load_prefix + 'no_va/pes_nova_db.p', 'rb') as f:
         pes  = pickle.load(f)
    f.close()
else: 
    weights = utils.get_params_from_zip(load_prefix + 'rl_model_20800000_steps')
    w_val = np.squeeze(weights['model/vf/w:0'])

    feats = np.vstack(feats_)
    rewards = np.hstack(rewards_)
    terms = np.hstack(terms_)
    start = np.roll(terms,1)
    nsteps = len(terms)
    nfeatures = feats_[0][0].shape[0]
    gamma = 0.99
    # compute per-feature PEs
    pes = np.zeros((nsteps, nfeatures))
    for i in range(0,nsteps-1):
        if (terms[i]): # there is a one-off error-- the SECOND index of the start of the trial accurately measures the start of the trial
            pes[i,:] = rewards[i] / nfeatures -  w_val * feats[i,:]
        else:
            pes[i,:] = rewards[i] / nfeatures +  w_val * (-feats[i,:] + gamma * feats[i+1,:])

    pickle.dump(pes, open(load_prefix + "pes_nova_db.p", "wb") )


    
# pes split by 5000 trials
pes_ = utils.split_by_ep_len(pes, np.hstack((episode_lengths)))

pes_cuep = np.dstack([pes_i[:CUEP_LIM,:] for pes_i in pes_])
# helper functions (to put in a different file)

def plot_frame(frame_i, vmin=None,vmax=None):
    if vmin is None:
        vmin = np.min(frame_i)
        vmax = np.max(frame_i)
    plt.imshow(obses_trial[:,:,frame_i], cmap = 'viridis', vmin=vmin, vmax=vmax)
    plt.xticks([])
    plt.yticks([])
    
"""
autocorr calculates mean-corrected cos angle across time of a matrix as a measure of similarity
        each timestep's vector is mean-corrected, then normalized to length 1, then dot producted to itself
inputs> mat: matrix with ROWS of observations across time.(so timesteps x shape)
"""
def covar_mat(mat):

    mat_meancorr = np.vstack([mat_i - np.mean(mat_i) for mat_i in mat])
    mat_norm = np.sqrt(np.diag(mat_meancorr.dot(mat_meancorr.T)))
    return (mat_meancorr/mat_norm[:, None]).dot((mat_meancorr/mat_norm[:, None]).T)<jupyter_output><empty_output><jupyter_text>Figure 4A: Similarity of Video Frames<jupyter_code>TEXTURE_REPEAT_IDX = 67

obses_trial = obses_empty[:,:,:140]
vmin = np.min(np.hstack([obses_trial[:,:,120], obses_trial[:,:,0]]))
vmax = np.max(np.hstack([obses_trial[:,:,120], obses_trial[:,:,0]]))

plot_frame(0, vmin, vmax)
plt.title('Frame at Position = 0 cm');
plot_frame(TEXTURE_REPEAT_IDX, vmin, vmax)
plt.title('Frame at Position = ' + str(round(ypos_cuep[TEXTURE_REPEAT_IDX])) + 'cm');

obses_ac = covar_mat(obses_empty_flat)
obses_ac_lim = obses_ac[autocorr_start:autocorr_end, autocorr_start:autocorr_end]

fig, ax = plt.subplots(1,2, figsize =(4.5, 3.5), gridspec_kw=dict(width_ratios=[1,0.03]))
im = ax[0].imshow(obses_ac_lim, aspect = 'auto',
                extent = [ypos_cuep[autocorr_start], ypos_cuep[autocorr_end], ypos_cuep[autocorr_start], ypos_cuep[autocorr_end]],
                cmap = 'viridis',  interpolation = 'none')

ax[0].set_yticks([20, 40, 60])
ax[0].set_xticks([20, 40, 60])
im.set_clim(0.5,1) # min is at 0.544, clim to ensure colorbar has 0.5 tick
cbar = fig.colorbar(im, cax = ax[1], ticks = [-1, 0, 0.5, 1])
cbar.ax.yaxis.set_label_position('left')
cbar.set_label('Similarity (Cosine θ)')


ax[0].spines['right'].set_visible(True)
ax[0].spines['top'].set_visible(True)

ax[0].set_xlabel('Position (cm)');
ax[0].set_ylabel('Position (cm)');<jupyter_output><empty_output><jupyter_text>Figure 2C: Similarity of Vector RPEs<jupyter_code>pes_ac = covar_mat(np.mean(pes_cuep, -1))
pes_ac_lim = pes_ac[autocorr_start:autocorr_end, autocorr_start:autocorr_end]

fig, ax = plt.subplots(1,2, figsize =(4.5, 3.5), gridspec_kw=dict(width_ratios=[1,0.03]))
im = ax[0].imshow(pes_ac_lim, aspect = 'auto',
                extent = [ypos_cuep[autocorr_start], ypos_cuep[autocorr_end], ypos_cuep[autocorr_start], ypos_cuep[autocorr_end]],
                cmap = 'magma',  interpolation = 'none')
ax[0].set_xticks([20, 40, 60])
im.set_clim(-1,1)
ax[0].spines['right'].set_visible(True)
ax[0].spines['top'].set_visible(True)
cbar = fig.colorbar(im, cax = ax[1], ticks = [-1, 0, 1])
cbar.ax.yaxis.set_label_position('left')
cbar.set_label('Similarity (Cosine θ)')

ax[0].set_xlabel('Position (cm)');
ax[0].set_ylabel('Position (cm)');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14196521480600566, 'fraction_numerical': 0.025122640107031365, 'mean_word_length': 5.133090246125797, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '271641', 'license': 'MIT', 'n_tokens_mistral': 2783, 'n_tokens_neox': 2611, 'n_words': 554, 'path': 'figures/Figure 4.ipynb', 'repo_name': 'ndawlab/vectorRPE', 'size': '6727'}