def regression_pokes_aligned_warped_simple(experiment, all_sessions):
C_warped = []
cpd_warped = []
# Finding correlation coefficients for task 1
for s, session in enumerate(experiment):
all_neurons_all_spikes_raster_plot_task = all_sessions[s]
if all_neurons_all_spikes_raster_plot_task.shape[1] > 0:
all_neurons_all_spikes_raster_plot_task = all_neurons_all_spikes_raster_plot_task[
1::2, :, :]
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1,\
different_outcome_task_2, different_outcome_task_3 = re.predictors_include_previous_trial(session)
n_trials, n_neurons, n_timepoints = all_neurons_all_spikes_raster_plot_task.shape
predictor_A = predictor_A_Task_1 + predictor_A_Task_2 + predictor_A_Task_3
predictor_A = predictor_A[:all_neurons_all_spikes_raster_plot_task.
shape[0]]
reward = reward[:all_neurons_all_spikes_raster_plot_task.shape[0]]
# interaction = reward*predictor_A
ones = np.ones(len(predictor_A))
predictors = OrderedDict([('A', predictor_A), ('reward', reward),
('ones', ones)])
#('interaction',interaction)])
X = np.vstack(
predictors.values()).T[:len(predictor_A), :].astype(float)
ones = np.ones(X.shape[0]).reshape(X.shape[0], 1)
X_check_rank = np.hstack([X, ones])
X_check_rank = X
print(X_check_rank.shape[1])
rank = matrix_rank(X_check_rank)
print(rank)
n_predictors = X.shape[1]
y = all_neurons_all_spikes_raster_plot_task.reshape(
[all_neurons_all_spikes_raster_plot_task.shape[0],
-1]) # Activity matrix [n_trials, n_neurons*n_timepoints]
ols = LinearRegression(copy_X=True, fit_intercept=False)
ols.fit(X, y)
C_warped.append(
ols.coef_.reshape(n_neurons, n_timepoints,
n_predictors)) # Predictor loadings
cpd_warped.append(
_CPD(X, y).reshape(n_neurons, n_timepoints, n_predictors))
C_warped = np.concatenate(C_warped, 0)
cpd_warped = np.nanmean(np.concatenate(cpd_warped, 0),
axis=0) # Population CPD is mean over neurons.
def simulate_Q4(session, params):
# Unpack trial events.
choice_a_ind = []
choice_b_ind = []
# Unpack trial events.
choices, outcomes,task = (session.trial_data['choices'], session.trial_data['outcomes'], session.trial_data['task'])
#forced_trials = session.trial_data['forced_trial']
#non_forced_array = np.where(forced_trials == 0)[0]
#task_non_forced = task[non_forced_array]
#choices = choices[non_forced_array]
#outcomes = outcomes[non_forced_array]
task_1 = np.where(task == 1)[0]
task_2 = np.where(task == 2)[0]
predictor_A_Task_1_forced, predictor_A_Task_2_forced, predictor_A_Task_3_forced,\
predictor_B_Task_1_forced, predictor_B_Task_2_forced, predictor_B_Task_3_forced, reward_forced,\
predictor_a_good_task_1_forced, predictor_a_good_task_2_forced, predictor_a_good_task_3_forced = ft.predictors_forced(session)
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1, different_outcome_task_2, different_outcome_task_3, switch = re.predictors_include_previous_trial(session)
choices, outcomes = (session.trial_data['choices'], session.trial_data['outcomes'])
index_non_forced = np.where(session.trial_data['forced_trial'] == 0)[0]
index_forced = np.where(session.trial_data['forced_trial'] == 1)[0]
non_forced_choices = predictor_A_Task_1 + predictor_A_Task_2 + predictor_A_Task_3
forced_choices = predictor_A_Task_1_forced + predictor_A_Task_2_forced + predictor_A_Task_3_forced
choices_forced_unforced = np.zeros(len(choices))
choices_forced_unforced[index_forced] = forced_choices[:len(index_forced)]
choices_forced_unforced[index_non_forced] = non_forced_choices
choices_forced_unforced = np.asarray(choices_forced_unforced, dtype = int)
ind_task_2 = len(task_1)
ind_task_3 = len(task_1)+len(task_2)
n_trials = choices.shape[0]
# Unpack parameters.
alpha, iTemp, h, k = params
#Variables.
Q_td = np.zeros([n_trials + 1, 2]) # Model free action values.
Q_chosen = [0]
for i, (c, o) in enumerate(zip(choices_forced_unforced, outcomes)): # loop over trials.
if i == ind_task_2 or i == ind_task_3:
nc = 1 - c # Not chosen action.
# update action values simple RW
Q_td[i+1, c] = 0
Q_td[i+1,nc] = 0
Q_chosen.append(Q_td[i+1, c])
else:
nc = 1 - c # Not chosen action
# update action values simple RW
Q_td[i+1, c] = (1-alpha)*Q_td[i, c]+alpha*o
Q_td[i+1, nc] = (1-alpha*abs(h))*Q_td[i, nc]+alpha*h*o
Q_chosen.append(Q_td[i+1, c])
# Evaluate choice probabilities
for c,choice in enumerate(choices_forced_unforced):
if choice == 0:
if choices_forced_unforced[c] == choices_forced_unforced[c-1]:
choice_a_ind.append(c-1)
elif choice == 1:
if choices_forced_unforced[c] == choices_forced_unforced[c-1]:
choice_b_ind.append(c-1)
Q_td[choice_a_ind,0] *= k
Q_td[choice_b_ind,1] *= k
choice_probs = array_softmax(Q_td, iTemp)
return choice_probs, Q_td, Q_chosen
def session_likelihood(self, session, params, return_Qs = False):
# Unpack trial events.
choice_a_ind = []
choice_b_ind = []
choices, outcomes = (session.trial_data['choices'], session.trial_data['outcomes'])
predictor_A_Task_1_forced, predictor_A_Task_2_forced, predictor_A_Task_3_forced,\
predictor_B_Task_1_forced, predictor_B_Task_2_forced, predictor_B_Task_3_forced, reward_forced,\
predictor_a_good_task_1_forced, predictor_a_good_task_2_forced, predictor_a_good_task_3_forced = ft.predictors_forced(session)
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1, different_outcome_task_2, different_outcome_task_3, switch = re.predictors_include_previous_trial(session)
index_non_forced = np.where(session.trial_data['forced_trial'] == 0)[0]
index_forced = np.where(session.trial_data['forced_trial'] == 1)[0]
non_forced_choices = predictor_A_Task_1 + predictor_A_Task_2 + predictor_A_Task_3
forced_choices = predictor_A_Task_1_forced + predictor_A_Task_2_forced + predictor_A_Task_3_forced
choices_forced_unforced = np.zeros(len(choices))
choices_forced_unforced[index_forced] = forced_choices[:len(index_forced)]
choices_forced_unforced[index_non_forced] = non_forced_choices
n_trials = choices.shape[0]
# Unpack parameters.
alpha, iTemp, k = params
# Variables.
Q_td = np.zeros([n_trials + 1, 2]) # Action values.
for i, (c, o) in enumerate(zip(choices_forced_unforced, outcomes)): # loop over trials.
nc = 1 - c # Not chosen action.
# update action values simple RW
Q_td[i+1, c] = Q_td[i, c] + alpha*(o - Q_td[i, c])
Q_td[i+1,nc] = 1-Q_td[i+1, c]
# Evaluate choice probabilities and likelihood.
for c,choice in enumerate(choices_forced_unforced):
if choice == 0:
if choices_forced_unforced[c] == choices_forced_unforced[c-1]:
choice_a_ind.append(c-1)
elif choice == 1:
if choices_forced_unforced[c] == choices_forced_unforced[c-1]:
choice_b_ind.append(c-1)
Q_td[choice_a_ind,0] *= k
Q_td[choice_b_ind,1] *= k
choice_probs = array_softmax(Q_td, iTemp)
trial_log_likelihood = protected_log(choice_probs[np.arange(n_trials), choices_forced_unforced])
session_log_likelihood = np.sum(trial_log_likelihood)
if return_Qs:
return Q_td
else:
return session_log_likelihood
def tim_create_mat(experiment, title):
all_sessions_list = []
firing_rates = []
for s, session in enumerate(experiment):
index_non_forced = np.where(session.trial_data['forced_trial'] == 0)[0]
index_forced = np.where(session.trial_data['forced_trial'] == 1)[0]
firing_rate_non_forced = session.aligned_rates[index_non_forced]
firing_rate_forced = session.aligned_rates[index_forced]
choices = session.trial_data['choices']
trials, neurons, time = firing_rate_non_forced.shape
firing_rate = np.zeros((len(choices), neurons, time))
index_non_forced = np.where(session.trial_data['forced_trial'] == 0)[0]
index_forced = np.where(session.trial_data['forced_trial'] == 1)[0]
task = session.trial_data['task']
forced_trials = session.trial_data['forced_trial']
block = session.trial_data['block']
non_forced_array = np.where(forced_trials == 0)[0]
non_forced_choices = choices[non_forced_array]
# Getting out task indicies and choices
task = session.trial_data['task']
forced_trials = session.trial_data['forced_trial']
non_forced_array = np.where(forced_trials == 0)[0]
task_non_forced = task[non_forced_array]
task_1 = np.where(task == 1)[0]
task_2 = np.where(task == 2)[0]
task_3 = np.where(task == 3)[0]
task_2_non_forced = np.where(task_non_forced == 2)[0]
task_3_non_forced = np.where(task_non_forced == 3)[0]
forced_trials = session.trial_data['forced_trial']
outcomes = session.trial_data['outcomes']
predictor_A_Task_1_forced, predictor_A_Task_2_forced, predictor_A_Task_3_forced,\
predictor_B_Task_1_forced, predictor_B_Task_2_forced, predictor_B_Task_3_forced, reward_forced,\
predictor_a_good_task_1_forced, predictor_a_good_task_2_forced, predictor_a_good_task_3_forced = ft.predictors_forced(session)
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1, different_outcome_task_2, different_outcome_task_3, switch = re.predictors_include_previous_trial(session)
non_forced_choices = predictor_A_Task_1 + predictor_A_Task_2 + predictor_A_Task_3
forced_choices = predictor_A_Task_1_forced + predictor_A_Task_2_forced + predictor_A_Task_3_forced
choices_forced_unforced = np.zeros(len(choices))
choices_forced_unforced[
index_forced] = forced_choices[:len(index_forced)]
choices_forced_unforced[index_non_forced] = non_forced_choices
state = np.zeros(len(choices))
forced_state = predictor_a_good_task_1_forced + predictor_a_good_task_2_forced + predictor_a_good_task_3_forced
non_forced_state = np.zeros(len(non_forced_array))
non_forced_state[predictor_a_good_task_1] = 1
non_forced_state[predictor_a_good_task_2 + task_2_non_forced[0]] = 1
non_forced_state[predictor_a_good_task_3 + task_3_non_forced[0]] = 1
state[index_forced] = forced_state[:len(index_forced)]
state[index_non_forced] = non_forced_state
choices_forced_unforced[
index_forced] = forced_choices[:len(index_forced)]
choices_forced_unforced[index_non_forced] = non_forced_choices
ones = np.ones(len(choices))
firing_rate[index_forced] = firing_rate_forced[:len(index_forced)]
firing_rate[
index_non_forced] = firing_rate_non_forced[:len(index_non_forced)]
# =============================================================================
# Extracting identity of pokes in each task
# =============================================================================
poke_A = 'poke_' + str(session.trial_data['poke_A'][0])
poke_A_task_2 = 'poke_' + str(session.trial_data['poke_A'][task_2[0]])
poke_A_task_3 = 'poke_' + str(session.trial_data['poke_A'][task_3[0]])
poke_B = 'poke_' + str(session.trial_data['poke_B'][0])
poke_B_task_2 = 'poke_' + str(session.trial_data['poke_B'][task_2[0]])
poke_B_task_3 = 'poke_' + str(session.trial_data['poke_B'][task_3[0]])
configuration = session.trial_data['configuration_i']
i_pokes = np.unique(configuration)
#print('These are I pokes')
i_poke_task_1 = configuration[0]
i_poke_task_2 = configuration[task_2[0]]
i_poke_task_3 = configuration[task_3[0]]
#print(i_poke_task_1, i_poke_task_2, i_poke_task_3)
poke_A1_A2_A3, poke_A1_B2_B3, poke_A1_B2_A3, poke_A1_A2_B3, poke_B1_B2_B3, poke_B1_A2_A3, poke_B1_A2_B3, poke_B1_B2_A3 = ep.poke_A_B_make_consistent(
session)
if poke_A1_A2_A3 == True:
constant_poke_a = poke_A
poke_b_1 = poke_B
poke_b_2 = poke_B_task_2
poke_b_3 = poke_B_task_3
if poke_A1_B2_B3 == True:
constant_poke_a = poke_A
poke_b_1 = poke_B
poke_b_2 = poke_A_task_2
poke_b_3 = poke_A_task_3
if poke_A1_B2_A3 == True:
constant_poke_a = poke_A
poke_b_1 = poke_B
poke_b_2 = poke_A_task_2
poke_b_3 = poke_B_task_3
if poke_A1_A2_B3 == True:
constant_poke_a = poke_A
poke_b_1 = poke_B
poke_b_2 = poke_B_task_2
poke_b_3 = poke_A_task_3
if poke_B1_B2_B3 == True:
constant_poke_a = poke_B
poke_b_1 = poke_A
poke_b_2 = poke_A_task_2
poke_b_3 = poke_A_task_3
if poke_B1_A2_A3 == True:
constant_poke_a = poke_B
poke_b_1 = poke_A
poke_b_2 = poke_B_task_2
poke_b_3 = poke_B_task_3
if poke_B1_A2_B3 == True:
constant_poke_a = poke_B
poke_b_1 = poke_A
poke_b_2 = poke_B_task_2
poke_b_3 = poke_A_task_3
if poke_B1_B2_A3 == True:
constant_poke_a = poke_B
poke_b_1 = poke_A
poke_b_2 = poke_A_task_2
poke_b_3 = poke_B_task_3
a_pokes = np.zeros(len(choices))
a_pokes[:] = constant_poke_a[-1]
b_pokes = np.zeros(len(choices))
b_pokes[:task_1[-1] + 1] = poke_b_1[-1]
b_pokes[task_1[-1] + 1:task_2[-1] + 1] = poke_b_2[-1]
b_pokes[task_2[-1] + 1:] = poke_b_3[-1]
i_pokes = np.zeros(len(choices))
i_pokes[:task_1[-1] + 1] = i_poke_task_1
i_pokes[task_1[-1] + 1:task_2[-1] + 1] = i_poke_task_2
i_pokes[task_2[-1] + 1:] = i_poke_task_3
# chosen_Q1 = experiment_sim_Q1[s][:len(choices)]
# chosen_Q4 = experiment_sim_Q4[s][:len(choices)]
# Q1_value_a = experiment_sim_Q1_value_a[s][:len(choices)]
# Q1_value_b = experiment_sim_Q1_value_b[s][:len(choices)]
# Q4_value_a = experiment_sim_Q4_values[s][:len(choices)]
# Q1_prediction_error_chosen = experiment_sim_Q1_prediction_error_chosen[s][:len(choices)]
predictors_all = OrderedDict([
('latent_state', state),
('choice', choices_forced_unforced),
('reward', outcomes),
('forced_trials', forced_trials),
('block', block),
('task', task),
('A', a_pokes),
('B', b_pokes),
('Initiation', i_pokes),
# ('Chosen_Simple_RW',chosen_Q1),
# ('Chosen_Cross_learning_RW', chosen_Q4),
# ('Value_A_RW', Q1_value_a),
# ('Value_B_RW', Q1_value_b),
# ('Value_A_Cross_learning', Q4_value_a),
# ('Prediction_Error_Q',Q1_prediction_error_chosen),
('ones', ones)
])
X = np.vstack(
predictors_all.values()).T[:len(choices), :].astype(float)
# Save all sessions
all_sessions_list.append(X)
firing_rates.append(firing_rate)
scipy.io.savemat('/Users/veronikasamborska/Desktop/' + title + '.mat', {
'Data': firing_rates,
'DM': all_sessions_list
})
data = {'Data': firing_rates, 'DM': all_sessions_list}
return data
def regression_Q_values_choices(experiment):
C_1_choice = []
cpd_1_choice = []
C_2_choice = []
cpd_2_choice= []
C_3_choice = []
cpd_3_choice = []
# Finding correlation coefficients for task 1
for s,session in enumerate(experiment):
aligned_spikes= session.aligned_rates[:]
if aligned_spikes.shape[1] > 0: # sessions with neurons?
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1,\
different_outcome_task_2, different_outcome_task_3 = re.predictors_include_previous_trial(session)
n_trials, n_neurons, n_timepoints = aligned_spikes.shape
# Task indicies
task = session.trial_data['task']
forced_trials = session.trial_data['forced_trial']
non_forced_array = np.where(forced_trials == 0)[0]
task_non_forced = task[non_forced_array]
task_1 = np.where(task_non_forced == 1)[0]
task_2 = np.where(task_non_forced == 2)[0]
# (1) Better correlated with the *previous* choice, or the *current choice*
predictor_A_Task_1 = predictor_A_Task_1[1:len(task_1)]
previous_trial_task_1 = previous_trial_task_1[1:len(task_1)]
aligned_spikes_task_1 = aligned_spikes[1:len(task_1)]
ones = np.ones(len(predictor_A_Task_1))
# Task 1
predictors = OrderedDict([('Current Choice', predictor_A_Task_1),
('Previous Choice', previous_trial_task_1),
('ones', ones)])
X = np.vstack(predictors.values()).T[:len(predictor_A_Task_1),:].astype(float)
n_predictors = X.shape[1]
y = aligned_spikes_task_1.reshape([len(aligned_spikes_task_1),-1]) # Activity matrix [n_trials, n_neurons*n_timepoints]
ols = LinearRegression(copy_X = True,fit_intercept= True)
ols.fit(X,y)
C_1_choice.append(ols.coef_.reshape(n_neurons,n_timepoints, n_predictors)) # Predictor loadings
cpd_1_choice.append(re._CPD(X,y).reshape(n_neurons,n_timepoints, n_predictors))
# Task 2
# (1) Better correlated with the *previous* choice, or the *current choice*
predictor_A_Task_2 = predictor_A_Task_2[len(task_1)+1:len(task_1)+len(task_2)]
previous_trial_task_2 = previous_trial_task_2[len(task_1)+1:len(task_1)+len(task_2)]
aligned_spikes_task_2 = aligned_spikes[len(task_1)+1:len(task_1)+len(task_2)]
ones = np.ones(len(predictor_A_Task_2))
predictors = OrderedDict([('Current Choice', predictor_A_Task_2),
('Previous CHoice', previous_trial_task_2),
('ones', ones)])
X = np.vstack(predictors.values()).T[:len(predictor_A_Task_2),:].astype(float)
n_predictors = X.shape[1]
y = aligned_spikes_task_2.reshape([len(aligned_spikes_task_2),-1]) # Activity matrix [n_trials, n_neurons*n_timepoints]
ols = LinearRegression(copy_X = True,fit_intercept= True)
ols.fit(X,y)
C_2_choice.append(ols.coef_.reshape(n_neurons,n_timepoints, n_predictors)) # Predictor loadings
cpd_2_choice.append(re._CPD(X,y).reshape(n_neurons,n_timepoints, n_predictors))
# Task 3
# (1) Better correlated with the *previous* choice, or the *current choice*
predictor_A_Task_3 = predictor_A_Task_3[len(task_1)+len(task_2)+1:]
previous_trial_task_3 = previous_trial_task_3[len(task_1)+len(task_2)+1:]
aligned_spikes_task_3 = aligned_spikes[len(task_1)+len(task_2)+1:]
aligned_spikes_task_3 = aligned_spikes_task_3[:len(predictor_A_Task_3)]
ones = np.ones(len(predictor_A_Task_3))
predictors = OrderedDict([('Current Choice', predictor_A_Task_3),
('Previous Choice', previous_trial_task_3),
('ones', ones)])
X = np.vstack(predictors.values()).T[:len(predictor_A_Task_3),:].astype(float)
n_predictors = X.shape[1]
y = aligned_spikes_task_3.reshape([len(aligned_spikes_task_3),-1]) # Activity matrix [n_trials, n_neurons*n_timepoints]
ols = LinearRegression(copy_X = True,fit_intercept= True)
ols.fit(X,y)
C_3_choice.append(ols.coef_.reshape(n_neurons,n_timepoints, n_predictors)) # Predictor loadings
cpd_3_choice.append(re._CPD(X,y).reshape(n_neurons,n_timepoints, n_predictors))
cpd_1_choice = np.nanmean(np.concatenate(cpd_1_choice,0), axis = 0) # Population CPD is mean over neurons.
cpd_2_choice = np.nanmean(np.concatenate(cpd_2_choice,0), axis = 0) # Population CPD is mean over neurons.
cpd_3_choice = np.nanmean(np.concatenate(cpd_3_choice,0), axis = 0) # Population CPD is mean over neurons.
return C_1_choice, C_2_choice, C_3_choice, cpd_1_choice, cpd_2_choice, cpd_3_choice, predictors
def simulate_bayes(session, params):
# Unpack trial events.
choices, outcomes = (session.trial_data['choices'],
session.trial_data['outcomes'])
forced_trials = session.trial_data['forced_trial']
non_forced_array = np.where(forced_trials == 0)[0]
choices = choices[non_forced_array]
outcomes = outcomes[non_forced_array]
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1,\
different_outcome_task_2, different_outcome_task_3, switch = re.predictors_include_previous_trial(session)
stay = np.asarray(switch)
n_trials = choices.shape[0]
#Unpack parameters.
alpha, iTemp, sigma = params
reward_prob_correct = 0.75
reward_prob_incorrect = 0.25
no_reward_prob_correct = 0.25
no_reward_prob_incorrect = 0.75
# Variables.
Posterior_correct_incorrect = np.zeros(
[n_trials, 2]) # Array for correct and incorrect posteriors
Posterior_correct_incorrect[
0, :] = 0.5 # First choices start with 0.5 priors for both correct and incorrect choice
Prior_correct_incorrect = np.zeros(
[n_trials, 2]) # Array for correct and incorrect priors
# Switch array --> ones when choice is the same; reverse such that 1s mean there was a switch between current and last trial
switch_choice = stay - 1
switch_choice = switch_choice * (-1)
outcomes = outcomes[1:]
for i, (o, c) in enumerate(zip(outcomes,
switch_choice)): # loop over trials.
if c == 1:
# Prior that the decision is correct if a switch occured;
# Reversal P * Posterior past choice was correct + (1-Reversal P)* Posterior past choice was incorrect
Prior_correct_incorrect[i, 0] = (
(sigma) * Posterior_correct_incorrect[i, 0]) + (
(1 - sigma) * Posterior_correct_incorrect[i, 1])
Prior_correct_incorrect[i, 1] = 1 - Prior_correct_incorrect[i, 0]
# Prior that the decision is correct if no switch occured;
# Choice is correct with probability (1-Reversal P) & given past choice being correct + (Reversal P)* Posterior past choice was incorrect
elif c == 0:
Prior_correct_incorrect[i, 0] = (
(1 - sigma) * Posterior_correct_incorrect[i, 0]) + (
(sigma) * Posterior_correct_incorrect[i, 1])
Prior_correct_incorrect[i, 1] = 1 - Prior_correct_incorrect[i, 0]
if o == 1:
# Decision is correct based on the probability of getting a reward weighted by the prior of the trial being correct
#/Probability of getting a reward in correct prior; probability of getting no reward given the incorrect prior
Posterior_correct_incorrect[
i + 1,
0] = reward_prob_correct * Prior_correct_incorrect[i, 0] / (
(reward_prob_correct * Prior_correct_incorrect[i, 0]) +
(reward_prob_incorrect * Prior_correct_incorrect[i, 1]))
Posterior_correct_incorrect[
i + 1, 1] = 1 - Posterior_correct_incorrect[i + 1, 0]
elif o == 0:
# Decision is correct based on the probability of getting no reward weighted by the prior of the trial being correct
#/Probability of getting a reward in an correct prior; probability of getting no reward given the incorrect prior
Posterior_correct_incorrect[
i + 1,
0] = no_reward_prob_correct * Prior_correct_incorrect[i, 0] / (
((no_reward_prob_correct) * Prior_correct_incorrect[i, 0])
+ ((no_reward_prob_incorrect) *
Prior_correct_incorrect[i, 1]))
Posterior_correct_incorrect[
i + 1, 1] = 1 - Posterior_correct_incorrect[i + 1, 0]
Posterior_correct_incorrect = Posterior_correct_incorrect[:-1]
Prior_correct_incorrect = Prior_correct_incorrect[:-1]
return Posterior_correct_incorrect, Prior_correct_incorrect
def session_likelihood(self, session, params):
# Unpack trial events.
choices, outcomes = (session.trial_data['choices'],
session.trial_data['outcomes'])
forced_trials = session.trial_data['forced_trial']
non_forced_array = np.where(forced_trials == 0)[0]
choices = choices[non_forced_array]
outcomes = outcomes[non_forced_array]
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1, different_outcome_task_2, different_outcome_task_3, switch = re.predictors_include_previous_trial(session)
stay = np.asarray(switch)
n_trials = choices.shape[0]
#Unpack parameters.
alpha, iTemp, sigma = params
reward_prob_correct = 0.75
reward_prob_incorrect = 0.25
no_reward_prob_correct = 0.25
no_reward_prob_incorrect = 0.75
#iTemp = 2.46759374
#alpha = 0.35039014
# Variables.
Posterior_correct_incorrect = np.zeros(
[n_trials, 2]) # Array for correct and incorrect posteriors
Posterior_correct_incorrect[
0, :] = 0.5 # First choices start with 0.5 priors for both correct and incorrect choice
Prior_correct_incorrect = np.zeros(
[n_trials, 2]) # Array for correct and incorrect priors
# Switch array --> ones when choice is the same; reverse such that 1s mean there was a switch between current and last trial
switch_choice = stay - 1
switch_choice = switch_choice * (-1)
outcomes = outcomes[1:]
for i, (o, c) in enumerate(zip(outcomes,
switch_choice)): # loop over trials.
if c == 1:
# Prior that the decision is correct if a switch occured;
# Reversal P * Posterior past choice was correct + (1-Reversal P)* Posterior past choice was incorrect
Prior_correct_incorrect[i, 0] = (
(sigma) * Posterior_correct_incorrect[i, 0]) + (
(1 - sigma) * Posterior_correct_incorrect[i, 1])
Prior_correct_incorrect[i,
1] = 1 - Prior_correct_incorrect[i, 0]
# Prior that the decision is correct if no switch occured;
# Choice is correct with probability (1-Reversal P) & given past choice being correct + (Reversal P)* Posterior past choice was incorrect
elif c == 0:
Prior_correct_incorrect[i, 0] = (
(1 - sigma) * Posterior_correct_incorrect[i, 0]) + (
(sigma) * Posterior_correct_incorrect[i, 1])
Prior_correct_incorrect[i,
1] = 1 - Prior_correct_incorrect[i, 0]
if o == 1:
# Decision is correct based on the probability of getting a reward weighted by the prior of the trial being correct
#/Probability of getting a reward in correct prior; probability of getting no reward given the incorrect prior
Posterior_correct_incorrect[
i + 1,
0] = reward_prob_correct * Prior_correct_incorrect[i, 0] / (
(reward_prob_correct * Prior_correct_incorrect[i, 0]) +
(reward_prob_incorrect *
Prior_correct_incorrect[i, 1]))
Posterior_correct_incorrect[
i + 1, 1] = 1 - Posterior_correct_incorrect[i + 1, 0]
elif o == 0:
# Decision is correct based on the probability of getting no reward weighted by the prior of the trial being correct
#/Probability of getting a reward in an correct prior; probability of getting no reward given the incorrect prior
Posterior_correct_incorrect[
i + 1,
0] = no_reward_prob_correct * Prior_correct_incorrect[
i, 0] / (((no_reward_prob_correct) *
Prior_correct_incorrect[i, 0]) +
((no_reward_prob_incorrect) *
Prior_correct_incorrect[i, 1]))
Posterior_correct_incorrect[
i + 1, 1] = 1 - Posterior_correct_incorrect[i + 1, 0]
# Evaluate choice probabilities and likelihood.
Posterior_correct_incorrect = Posterior_correct_incorrect[:-1, :]
choice_probs = array_sigmoid(iTemp, Posterior_correct_incorrect, alpha)
trial_log_likelihood_switch = switch_choice * protected_log(
choice_probs[:, 0])
trial_log_likelihood_stay = stay * protected_log(choice_probs[:, 1])
session_log_likelihood = (
np.sum(trial_log_likelihood_switch) / np.sum(switch_choice)) + (
np.sum(trial_log_likelihood_stay) / sum(stay))
## Checking everything is working okay
# plt.figure(3)
# plt.plot(outcomes, "v", color = 'red', alpha = 0.7, markersize=3)
# plt.plot(switch_choice,"x", color = 'green', alpha = 0.7, markersize=3)
# plt.plot(Posterior_correct_incorrect[:,1], color = 'grey', alpha = 0.7)
# plt.plot(trial_log_likelihood_switch, color = 'red')
# plt.plot(trial_log_likelihood_stay, color = 'blue')
# plt.plot(choice_probs[:,0],'black', alpha = 0.4)
# plt.plot(Posterior_correct_incorrect[:,1], color = 'grey', alpha = 0.7)
# plt.title(str(session_log_likelihood))
return session_log_likelihood
def average_firing_rates(experiment, all_sessions):
all_sessions_Ar = []
all_sessions_Br = []
all_sessions_Anr = []
all_sessions_Bnr = []
for s, session in enumerate(experiment):
all_neurons_all_spikes_raster_plot_task = all_sessions[s]
if all_neurons_all_spikes_raster_plot_task.shape[1] > 0:
all_neurons_all_spikes_raster_plot_task = np.asarray(
all_neurons_all_spikes_raster_plot_task)
all_neurons_all_spikes_raster_plot_task = all_neurons_all_spikes_raster_plot_task[:,
1::
2, :]
predictor_A_Task_1, predictor_A_Task_2, predictor_A_Task_3,\
predictor_B_Task_1, predictor_B_Task_2, predictor_B_Task_3, reward,\
predictor_a_good_task_1,predictor_a_good_task_2, predictor_a_good_task_3,\
reward_previous,previous_trial_task_1,previous_trial_task_2,previous_trial_task_3,\
same_outcome_task_1, same_outcome_task_2, same_outcome_task_3,different_outcome_task_1,\
different_outcome_task_2, different_outcome_task_3 = re.predictors_include_previous_trial(session)
n_neurons, n_trials, n_timepoints = all_neurons_all_spikes_raster_plot_task.shape
all_neurons_all_spikes_raster_plot_task = all_neurons_all_spikes_raster_plot_task[:, :len(
predictor_A_Task_1), :]
predictor_A = predictor_A_Task_1 + predictor_A_Task_2 + predictor_A_Task_3
predictor_A = predictor_A[:all_neurons_all_spikes_raster_plot_task.
shape[1]]
reward = reward[:all_neurons_all_spikes_raster_plot_task.shape[1]]
#Indexing A rewarded, B rewarded firing rates in 3 tasks
aligned_rates_A_reward = all_neurons_all_spikes_raster_plot_task[:,
np
.
where((
predictor_A
==
1
) & (
reward
==
1
)), :]
aligned_rates_A_Nreward = all_neurons_all_spikes_raster_plot_task[:,
np
.
where((
predictor_A
==
1
) & (
reward
==
0
)), :]
aligned_rates_B_reward = all_neurons_all_spikes_raster_plot_task[:,
np
.
where((
predictor_A
==
0
) & (
reward
==
1
)), :]
aligned_rates_B_Nreward = all_neurons_all_spikes_raster_plot_task[:,
np
.
where((
predictor_A
==
0
) & (
reward
==
0
)), :]
mean_aligned_rates_A_reward = np.mean(aligned_rates_A_reward,
axis=2)
mean_aligned_rates_A_Nreward = np.mean(aligned_rates_A_Nreward,
axis=2)
mean_aligned_rates_B_reward = np.mean(aligned_rates_B_reward,
axis=2)
mean_aligned_rates_B_Nreward = np.mean(aligned_rates_B_Nreward,
axis=2)
mean_aligned_rates_A_reward = np.mean(mean_aligned_rates_A_reward,
axis=0)
mean_aligned_rates_A_Nreward = np.mean(
mean_aligned_rates_A_Nreward, axis=0)
mean_aligned_rates_B_reward = np.mean(mean_aligned_rates_B_reward,
axis=0)
mean_aligned_rates_B_Nreward = np.mean(
mean_aligned_rates_B_Nreward, axis=0)
all_sessions_Ar.append(mean_aligned_rates_A_reward[0, :])
all_sessions_Br.append(mean_aligned_rates_B_reward[0, :])
all_sessions_Anr.append(mean_aligned_rates_A_Nreward[0, :])
all_sessions_Bnr.append(mean_aligned_rates_B_Nreward[0, :])
all_sessions_Ar = np.asarray(all_sessions_Ar)
all_sessions_Br = np.asarray(all_sessions_Br)
all_sessions_Anr = np.asarray(all_sessions_Anr)
all_sessions_Bnr = np.asarray(all_sessions_Bnr)
average_Ar = np.mean(all_sessions_Ar, axis=0)
average_Br = np.mean(all_sessions_Br, axis=0)
average_Anr = np.mean(all_sessions_Anr, axis=0)
average_Bnr = np.mean(all_sessions_Bnr, axis=0)
plt.plot(average_Ar, label='A reward')
plt.plot(average_Br, label='B reward')
plt.plot(average_Anr, label='A No reward')
plt.plot(average_Bnr, label='B No reward')
plt.legend()
