Revision edc453189104a1f76f4b2ab230cd86f2140e3f63 authored by Anne Urai on 08 April 2021, 13:13 UTC, committed by Anne Urai on 08 April 2021, 13:13 UTC
1 parent 22583a6
Raw File
suppfig_history_bycontrast.py
"""
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")
back to top