""" Psychometric function and choice shifts in the biased task @author: Anne Urai 15 January 2020 """ import pandas as pd import numpy as np import os from os.path import join import matplotlib.pyplot as plt import seaborn as sns import statsmodels.api as sm from statsmodels.formula.api import ols from paper_behavior_functions import (seaborn_style, figpath, group_colors, institution_map, query_sessions_around_criterion, EXAMPLE_MOUSE, FIGURE_HEIGHT, FIGURE_WIDTH, QUERY, load_csv, dj2pandas, plot_psychometric, fit_psychfunc, plot_chronometric, break_xaxis) # import wrappers etc from ibl_pipeline import reference, subject, behavior from ibl_pipeline.utils import psychofit as psy # Initialize seaborn_style() figpath = figpath() pal = group_colors() institution_map, col_names = institution_map() col_names = col_names[:-1] # colors for biased blocks cmap = sns.color_palette([[0.8984375, 0.37890625, 0.00390625], [0.3, 0.3, 0.3], [0.3671875, 0.234375, 0.59765625]]) sns.set_palette(cmap) # ================================= # # GET DATA FROM TRAINED ANIMALS # ================================= # if QUERY is True: # query sessions use_sessions, _ = query_sessions_around_criterion(criterion='ephys', days_from_criterion=[2, 0], force_cutoff=True) use_sessions = use_sessions & 'task_protocol LIKE "%biased%"' # only get biased sessions # restrict by list of dicts with uuids for these sessions b = (use_sessions * subject.Subject * subject.SubjectLab * reference.Lab * behavior.TrialSet.Trial) # reduce the size of the fetch b2 = b.proj('institution_short', 'subject_nickname', 'task_protocol', 'session_uuid', 'trial_stim_contrast_left', 'trial_stim_contrast_right', 'trial_response_choice', 'task_protocol', 'trial_stim_prob_left', 'trial_feedback_type', 'trial_response_time', 'trial_stim_on_time') # construct pandas dataframe 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) else: behav = load_csv('Fig4.csv') # how many mice are there for each lab? N = behav.groupby(['institution_code'])['subject_nickname'].nunique().to_dict() behav['n_mice'] = behav.institution_code.map(N) behav['institution_name'] = behav.institution_code + '\n' + behav.n_mice.apply(str) + ' mice' # ================================= # # PSYCHOMETRIC FUNCTIONS # FOR OUR EXAMPLE ANIMAL # ================================= # fig = sns.FacetGrid(behav[behav['subject_nickname'] == EXAMPLE_MOUSE], hue="probabilityLeft", palette=cmap, sharex=True, sharey=True, height=FIGURE_HEIGHT, aspect=(FIGURE_WIDTH/4)/FIGURE_HEIGHT) fig.map(plot_psychometric, "signed_contrast", "choice_right", "session_uuid") fig.set_axis_labels('\u0394 Contrast (%)', 'Rightward choices (%)') fig.ax.annotate('20:80', xy=(-5, 0.6), xytext=(-25, 0.8), color=cmap[0], fontsize=7) fig.ax.annotate('80:20', xy=(5, 0.4), xytext=(13, 0.18), color=cmap[2], fontsize=7) fig.despine(trim=True) fig.axes[0][0].set_title('Example mouse', fontweight='bold', color='k') fig.savefig(os.path.join(figpath, "figure4b_psychfuncs_biased_example.pdf")) fig.savefig(os.path.join( figpath, "figure4b_psychfuncs_biased_example.png"), dpi=600) plt.close('all') # ================================= # # PSYCHOMETRIC FUNCTIONS # one for all labs combined # ================================= # fig = sns.FacetGrid(behav, hue="probabilityLeft", palette=cmap, sharex=True, sharey=True, height=FIGURE_HEIGHT, aspect=(FIGURE_WIDTH/4)/FIGURE_HEIGHT) fig.map(plot_psychometric, "signed_contrast", "choice_right", "subject_nickname") fig.set_axis_labels('\u0394 Contrast (%)', '') fig.ax.annotate('20:80', xy=(-5, 0.6), xytext=(-25, 0.8), color=cmap[0], fontsize=7) fig.ax.annotate('80:20', xy=(5, 0.4), xytext=(13, 0.18), color=cmap[2], fontsize=7) fig.despine(trim=True) fig.axes[0][0].set_title('All mice: n = %d' % behav.subject_nickname.nunique(), fontweight='bold', color='k') fig.axes[0][0].set(yticklabels=[]) fig.savefig(os.path.join(figpath, "figure4b_psychfuncs_biased.pdf")) fig.savefig(os.path.join( figpath, "figure4b_psychfuncs_biased.png"), dpi=600) plt.close('all') # ================================================================== # # DIFFERENCE BETWEEN TWO PSYCHOMETRIC FUNCTIONS # FOR EACH ANIMAL + for each lab (in 'lab color') # ================================================================== # print('fitting psychometric functions...') pars = behav.groupby(['institution_code', 'subject_nickname', 'probabilityLeft']).apply(fit_psychfunc).reset_index() # now read these out at the presented levels of signed contrast behav2 = pd.DataFrame([]) xvec = behav.signed_contrast.unique() for index, group in pars.groupby(['institution_code', 'subject_nickname', 'probabilityLeft']): # expand yvec = psy.erf_psycho_2gammas([group.bias.item(), group.threshold.item(), group.lapselow.item(), group.lapsehigh.item()], xvec) group2 = group.loc[group.index.repeat( len(yvec))].reset_index(drop=True).copy() group2['signed_contrast'] = xvec group2['choice'] = 100 * yvec # add this behav2 = behav2.append(group2) # now subtract these to compute a bias shift behav3 = pd.pivot_table(behav2, values='choice', index=['institution_code', 'subject_nickname', 'signed_contrast'], columns=['probabilityLeft']).reset_index() behav3['biasshift'] = behav3[20] - behav3[80] # %% PLOT # plot one curve for each animal, one panel per lab plt.close('all') fig = sns.FacetGrid(behav3, col="institution_code", col_wrap=7, col_order=col_names, sharex=True, sharey=True, hue="subject_nickname", height=FIGURE_HEIGHT, aspect=(FIGURE_WIDTH/7)/FIGURE_HEIGHT) fig.map(plot_chronometric, "signed_contrast", "biasshift", "subject_nickname", color='gray', alpha=0.7) # overlay the example mouse tmpdat = behav3[behav3['subject_nickname'].str.contains(EXAMPLE_MOUSE)] plot_chronometric(tmpdat.signed_contrast, tmpdat.biasshift, tmpdat.subject_nickname, color='black', ax=fig.axes[0], legend=False) fig.set_titles("{col_name}") fig.despine(trim=True) ymin = fig.axes[0].get_ylim()[0]-0.2 fig.map(break_xaxis, y=ymin) # add lab means on top for axidx, ax in enumerate(fig.axes.flat): tmp_behav = behav3.loc[behav3.institution_code == behav3.institution_code.unique()[axidx], :] plot_chronometric(tmp_behav.signed_contrast, tmp_behav.biasshift, tmp_behav.institution_code, ax=ax, legend=False, color=pal[axidx], linewidth=2) ax.set_title(sorted(behav.institution_name.unique())[axidx], color=pal[axidx]) fig.set_axis_labels('\u0394 Contrast (%)', '\u0394 Rightward choices (%)') plt.tight_layout(w_pad=0) fig.savefig(os.path.join(figpath, "figure4e_biasshift.pdf")) fig.savefig(os.path.join(figpath, "figure4e_biasshift.png"), dpi=300) plt.close('all') # %% PLOT fig, ax1 = plt.subplots(1, 1, figsize=(FIGURE_WIDTH/4, FIGURE_HEIGHT)) for i, inst in enumerate(behav.institution_code.unique()): tmp_behav = behav3[behav3['institution_code'].str.contains(inst)] plot_chronometric(tmp_behav.signed_contrast, tmp_behav.biasshift, tmp_behav.subject_nickname, ax=ax1, legend=False, color=pal[i]) # ax1.set_title('All labs', color='k', fontweight='bold') ax1.set(xlabel='\u0394 Contrast (%)', ylabel='\u0394 Rightward choices (%)', yticks=[0, 10, 20, 30, 40]) sns.despine(trim=True) plt.tight_layout() ymin = ax1.get_ylim()[0]-0.15 break_xaxis(y=ymin) fig.savefig(os.path.join(figpath, "figure4d_biasshift_all_labs.pdf")) fig.savefig(os.path.join(figpath, "figure4d_biasshift_all_labs.png"), dpi=300) # ================================================================== # # Plot behavioral metrics per lab # ================================================================== # bias = behav3.loc[behav3.signed_contrast == 0, :] # stats on bias shift between laboratories: sm_lm = ols('biasshift ~ C(institution_code)', data=bias).fit() table = sm.stats.anova_lm(sm_lm, typ=2) # Type 2 ANOVA DataFrame print(table) # Add all mice to dataframe seperately for plotting bias_all = bias.copy() print('average bias shift across all mice: ') print(bias_all['biasshift'].mean()) bias_all['institution_code'] = 'All' bias_all = bias.append(bias_all) # Set color palette use_palette = [[0.6, 0.6, 0.6]] * len(np.unique(bias['institution_code'])) use_palette = use_palette + [[1, 1, 0.2]] sns.set_palette(use_palette) # plot f, ax1 = plt.subplots(1, 1, figsize=(3, 3.5)) sns.set_palette(use_palette) sns.boxplot(y='biasshift', x='institution_code', data=bias_all, ax=ax1) ax1.set(ylabel='\u0394 Rightward choices (%)\n at 0% contrast', ylim=[0, 51], xlabel='') [tick.set_color(pal[i]) for i, tick in enumerate(ax1.get_xticklabels()[:-1])] plt.setp(ax1.xaxis.get_majorticklabels(), rotation=40) plt.tight_layout(pad=2) seaborn_style() # plt.savefig(os.path.join(figpath, 'figure4e_bias_per_lab.pdf')) # plt.savefig(os.path.join(figpath, 'figure4e_bias_per_lab.png'), dpi=300)