"""
History-dependent choice strategy, depending on previous contrast
See also https://elifesciences.org/articles/49834
@author: Anne Urai
10 May 2020
"""
# import wrappers etc
from ibl_pipeline import behavior, subject, reference
import matplotlib.pyplot as plt
from paper_behavior_functions import (seaborn_style, figpath, query_sessions_around_criterion,
group_colors, institution_map, FIGURE_HEIGHT, FIGURE_WIDTH,
dj2pandas, fit_psychfunc)
import seaborn as sns
import pandas as pd
import numpy as np
import os
from ibl_pipeline.utils import psychofit as psy
# INITIALIZE A FEW THINGS
seaborn_style()
figpath = figpath()
pal = group_colors()
institution_map, col_names = institution_map()
# ================================= #
# GRAB ALL DATA FROM DATAJOINT
# 3 days before and 3 days after starting biasedChoiceWorld
# ================================= #
use_sessions, use_days = query_sessions_around_criterion(criterion='biased',
days_from_criterion=[
2, 3],
as_dataframe=False)
# restrict by list of dicts with uuids for these sessions
b = (use_sessions * subject.Subject * subject.SubjectLab * reference.Lab
* behavior.TrialSet.Trial)
# # temporary, to speed up computations
# b = (use_sessions * subject.Subject \
# * (subject.SubjectLab & 'lab_name = "churchlandlab"') \
# * reference.Lab
# * behavior.TrialSet.Trial)
# reduce the size of the fetch
b2 = b.proj('institution_short', 'subject_nickname', 'task_protocol',
'trial_stim_contrast_left', 'trial_stim_contrast_right', 'trial_response_choice',
'task_protocol', 'trial_stim_prob_left', 'trial_feedback_type')
bdat = b2.fetch(order_by='institution_short, subject_nickname, session_start_time, trial_id',
format='frame').reset_index()
behav = dj2pandas(bdat)
behav['institution_code'] = behav.institution_short.map(institution_map)
# split the two types of task protocols (remove the pybpod version number)
behav['task'] = behav['task_protocol'].str[14:20]
# also code for the future choice (for correction)
behav['next_choice'] = behav.choice.shift(-1)
behav.loc[behav.next_choice == 0, 'next_choice'] = np.nan
behav['next_outcome'] = behav.trial_feedback_type.shift(-1)
behav.loc[behav.next_outcome == 0, 'next_outcome'] = np.nan
behav['next_contrast'] = np.abs(behav.signed_contrast.shift(-1))
behav['next_signed_contrast'] = behav['signed_contrast'].shift(-1)
behav['previous_signed_contrast'] = behav.signed_contrast.shift(1)
# behav['next_choice_name'] = behav['next_choice'].map(
# {-1: 'left', 1: 'right'})
# behav['next_outcome_name'] = behav['next_outcome'].map(
# {-1: 'pre_error', 1: 'pre_correct'})
# # easy to use names for groupby
# behav['previous_name'] = behav.previous_outcome_name + \
# ', ' + behav.previous_choice_name
# behav['next_name'] = behav.next_outcome_name + \
# ', ' + behav.next_choice_name
# remove weird contrast levels that have very few trials
allowed_contrasts = [-100., -25., -12.5, -6.25, 0, 6.25, 12.5, 25., 100.]
behav = behav[behav['signed_contrast'].isin(allowed_contrasts)]
behav = behav[behav['previous_signed_contrast'].isin(allowed_contrasts)]
behav = behav[behav['next_signed_contrast'].isin(allowed_contrasts)]
print(behav.groupby(['lab_name'])['signed_contrast', 'previous_signed_contrast', 'next_signed_contrast'].nunique())
# choose: take only those trials where the objective probability is 0.5???
# behav = behav.loc[behav.probabilityLeft == 50, :]
# ======================================== #
# HEATMAP, WITHOUT PSYCHFUNC FITTING
# see Lak et al. eLife 2020, figure 1f
# ======================================== #
update_training = pd.pivot_table(behav[(behav.task == 'traini') & (behav.previous_outcome == 1)].
groupby(['signed_contrast',
'previous_signed_contrast',
'subject_nickname'])[
'choice'].mean().reset_index(),
values='choice',
index=['signed_contrast'],
columns=['previous_signed_contrast'],
aggfunc='mean')
future_training = pd.pivot_table(behav[(behav.task == 'traini') &
(behav.previous_outcome == 1)].groupby(['signed_contrast',
'next_signed_contrast',
'subject_nickname'])[
'choice'].mean().reset_index(),
values='choice',
index=['signed_contrast'],
columns=['next_signed_contrast'],
aggfunc='mean')
update_biased = pd.pivot_table(behav[(behav.task == 'biased') & (behav.previous_outcome == 1)].
groupby(['signed_contrast',
'previous_signed_contrast',
'subject_nickname'])[
'choice'].mean().reset_index(),
values='choice',
index=['signed_contrast'],
columns=['previous_signed_contrast'],
aggfunc='mean')
future_biased = pd.pivot_table(behav[(behav.task == 'biased') & (behav.previous_outcome == 1)].
groupby(['signed_contrast',
'next_signed_contrast',
'subject_nickname'])[
'choice'].mean().reset_index(),
values='choice',
index=['signed_contrast'],
columns=['next_signed_contrast'],
aggfunc='mean')
# subtract the average psychfunc, so that current stimulus does not dominate!
behav['dummy'] = 1
avg_psychfunc = pd.pivot_table(behav.groupby(['subject_nickname',
'signed_contrast',
'dummy'])['choice'].mean().reset_index(),
values='choice', index=['signed_contrast'],
columns=['dummy'], aggfunc='mean')
plt.close('all')
fig, ax = plt.subplots(1, 2, figsize=[8, 3.5])
kwargs = {'linewidths':0, 'cmap':"PuOr", 'square': True,
'cbar_kws':{'label': 'Updating (%)', 'shrink': 0.8,
'ticks': [-0.2, 0, 0.2]},
'vmin':-0.3, 'vmax':0.3}
sns.heatmap(update_training.sub(future_training.values, axis='rows'),
ax=ax[0], cbar=False, **kwargs)
ax[0].set(xlabel='Previous signed contrast (%)',
ylabel='Current signed contrast (%)',
title='Training task')
sns.heatmap(update_biased.sub(future_biased.values),
ax=ax[1], cbar=True, **kwargs)
ax[1].set(xlabel='Previous signed contrast (%)', ylabel='',
yticklabels=[], title='Biased task')
fig.tight_layout()
fig.savefig(os.path.join(figpath, "figure5d_history_prevcontrast_heatmap.pdf"))
# ============================================== #
# DEPENDENCE ON PREVIOUS CONTRAST - SUMMARY
# ============================================== #
def pars2choicefract(group):
group['choicefract'] = psy.erf_psycho_2gammas([group.bias.item(),
group.threshold.item(),
group.lapselow.item(),
group.lapsehigh.item()], 0) * 100
return group
print('fitting psychometric functions, NOW ALSO BASED ON FUTURE CONTRAST...')
# for name, gr in behav.groupby(['subject_nickname', 'task', 'next_choice', 'next_outcome', 'next_contrast']):
# print(gr.reset_index()['subject_nickname'].unique())
# fit_psychfunc(gr.reset_index())
pars = behav.groupby(['subject_nickname', 'task', 'lab_name',
'next_choice', 'next_outcome',
'next_contrast']).apply(fit_psychfunc).reset_index()
# convert to choice fraction
pars2 = pars.groupby(['subject_nickname', 'task', 'lab_name',
'next_choice', 'next_outcome', 'next_contrast']).apply(
pars2choicefract).reset_index()
# now compute the dependence on previous choice
future_shift = pd.pivot_table(pars2, values='choicefract',
index=['task', 'lab_name', 'subject_nickname',
'next_outcome', 'next_contrast'],
columns='next_choice').reset_index()
future_shift['future_shift'] = future_shift[1.] - future_shift[-1.]
# rename, so that the columns can be matched with history shift
future_shift = future_shift.rename(columns={'next_outcome': 'previous_outcome',
'next_contrast': 'previous_contrast'})
print('fitting psychometric functions, NOW ALSO BASED ON PREVIOUS CONTRAST...')
pars = behav.groupby(['subject_nickname', 'task', 'lab_name',
'previous_choice', 'previous_outcome',
'previous_contrast']).apply(fit_psychfunc).reset_index()
# convert to choice fraction
pars2 = pars.groupby(['subject_nickname', 'task', 'lab_name',
'previous_choice', 'previous_outcome',
'previous_contrast']).apply(pars2choicefract).reset_index()
# now compute the dependence on previous choice
history_shift = pd.pivot_table(pars2, values='choicefract',
index=['task', 'lab_name', 'subject_nickname',
'previous_outcome', 'previous_contrast'],
columns='previous_choice').reset_index()
history_shift['history_shift'] = history_shift[1.] - history_shift[-1.]
# ================================= #
# merge and subtract the future shift from each history shift
# ================================= #
pars5 = pd.merge(history_shift, future_shift,
on=['subject_nickname', 'previous_outcome',
'previous_contrast', 'task', 'lab_name'])
pars5['history_shift_corrected'] = pars5['history_shift'] - pars5['future_shift']
history_shift = pars5.copy()
history_shift.previous_contrast.replace([100], [40], inplace=True)
# ================================= #
# PLOT PREVIOUS CONTRAST-DEPENDENCE
# ================================= #
plt.close('all')
fig, axes = plt.subplots(1, 2, figsize=(FIGURE_WIDTH/2.2, FIGURE_HEIGHT), sharex=True, sharey=True)
for task, taskname, ax in zip(['traini', 'biased'], ['Basic task', 'Full task'], axes):
# # thin labels, per lab
# sns.lineplot(data=history_shift[(history_shift.task == task)].groupby(['lab_name',
# 'previous_contrast',
# 'previous_outcome'
# ]).mean().reset_index(),
# x='previous_contrast', y='history_shift_corrected',
# hue='previous_outcome', ax=ax, legend=False,
# ci=None, marker=None, hue_order=[-1., 1.],
# palette=sns.color_palette(["firebrick", "forestgreen"]),
# units='lab_name', estimator=None, zorder=0,
# linewidth=1, alpha=0.2)
# thick, across labs
sns.lineplot(data=history_shift[(history_shift.task == task)],
x='previous_contrast', y='history_shift_corrected',
hue='previous_outcome', ax=ax, legend=False, estimator=np.mean,
err_style='bars', marker='o', hue_order=[-1., 1.],
palette=sns.color_palette(["firebrick", "forestgreen"]),
zorder=100, ci=95)
ax.axhline(color='grey', linestyle=':', zorder=-100)
ax.set(ylabel='Choice updating (%)',
xlabel='Previous contrast (%)',
xticks=[0, 6, 12.5, 25, 40],
xticklabels=['0', '6', '12.5', '25', '100'],
title=taskname,
ylim=[-20, 20])
sns.despine(trim=True)
fig.tight_layout()
fig.savefig(os.path.join(figpath, "figure5c_history_prevcontrast.pdf"))
fig.savefig(os.path.join(figpath, "figure5c_history_prevcontrast.png"), dpi=600)
plt.close("all")
