{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d53e772d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n",
    "from matplotlib import pyplot as plt\n",
    "import scipy.stats\n",
    "from scipy.stats import wilcoxon\n",
    "import matplotlib.gridspec as gridspec\n",
    "from matplotlib.patches import Rectangle, Polygon\n",
    "from matplotlib.collections import PolyCollection\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "from functools import partial\n",
    "import warnings\n",
    "import vbn_utils\n",
    "import decoding_utils as du\n",
    "from analysis_utils import exponential_convolve\n",
    "import ccf_utils\n",
    "from vbn_utils import cumulative_hist, formatFigure, mean_sem_plot, make_iterable, get_unit_ids, bootstrap_ci\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f904ae3",
   "metadata": {},
   "source": [
    "## Data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f111098",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Paths to all of the useful supplemental tables and tensors\n",
    "active_tensor_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/vbnAllUnitSpikeTensor.hdf5\"\n",
    "passive_tensor_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/vbnAllUnitSpikeTensor_passive.hdf5\"\n",
    "\n",
    "stim_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_stim_table_no_filter.csv\"\n",
    "unit_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_units_with_responsiveness.csv\"\n",
    "\n",
    "sessions_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_sessions_table.csv\" #\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/vbn_s3_cache/visual-behavior-neuropixels-0.5.0/project_metadata/ecephys_sessions.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c18394",
   "metadata": {},
   "outputs": [],
   "source": [
    "units = pd.read_csv(unit_table_file)\n",
    "units['cortical_layer'] = units['cortical_layer'].replace('3-Feb','2/3') # necessary since 2/3 sometimes gets incorrectly reformatted as a date\n",
    "\n",
    "stim_table = pd.read_csv(stim_table_file)\n",
    "stim_table = stim_table.drop(columns='Unnamed: 0')\n",
    "\n",
    "active_tensor = h5py.File(active_tensor_file)\n",
    "\n",
    "sessions_table = pd.read_csv(sessions_table_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5027840a",
   "metadata": {},
   "outputs": [],
   "source": [
    "structure_tree = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/ccf_structure_tree_2017.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85efc869",
   "metadata": {},
   "outputs": [],
   "source": [
    "stim_table['is_shared'] = stim_table['image_name'].isin(['im083_r', 'im111_r'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17c5700e",
   "metadata": {},
   "source": [
    "## Time from last lick"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c270bea5",
   "metadata": {},
   "source": [
    "### Cluster-wise responses as function of time from last lick (running matched)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06078842",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.dpi'] = 300\n",
    "plt.rcParams['font.size'] = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b37d1fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "plt.rcParams['font.size'] = 15\n",
    "\n",
    "fig_agg, ax_agg = plt.subplots()\n",
    "\n",
    "cluster_colors = {\n",
    "    1: 'navy',  # Blue\n",
    "    2: 'mediumblue',  # Lighter blue\n",
    "    3: 'royalblue',  # Even lighter blue\n",
    "    4: 'cornflowerblue',  # Lighter blue\n",
    "    5: 'slateblue',  # Lightest blue\n",
    "    6: 'red',  # Red\n",
    "    7: 'coral',  # Lighter red\n",
    "    8: 'gold'   # Lightest red\n",
    "}\n",
    "\n",
    "cluster_names = {\n",
    "\t1: 'On Transient',\n",
    "\t2: 'On Sustained',\n",
    "\t3: 'Off Transient',\n",
    "\t4: 'Off Sustained',\n",
    "\t5: 'Image suppressed',\n",
    "\t6: 'Lick anticipation',\n",
    "\t7: 'Licking',\n",
    "\t8: 'Running'\n",
    "}\n",
    "\n",
    "time = np.arange(-750,750)\n",
    "\n",
    "all_cluster_means = {}\n",
    "for cluster in range(6,8):\n",
    "\tfig, ax = plt.subplots()\n",
    "\tfig.suptitle(cluster_names[cluster])\n",
    "\tunit_filter = du.get_units_in_cluster(units, cluster, clustering='new') & du.apply_unit_quality_filter(units) & \\\n",
    "\t\t\tdu.getUnitsInRegion(units, 'SCMRN')\n",
    "\t\t\t\n",
    "\tunit_ids = units.loc[unit_filter]['unit_id'].values\n",
    "\tif len(unit_ids)==0:\n",
    "\t\tcontinue\n",
    "\tsession_list = units.loc[unit_filter]['ecephys_session_id'].unique()\n",
    "\n",
    "\tmax_flashes = 7\n",
    "\n",
    "\tstim_filter_base = ['engaged', '~is_change', '~omitted', '~previous_omitted', '~grace_period_after_hit']# 'flashes_since_change>5']\n",
    "\tcond_filters = [\n",
    "\t\t\t['engaged', 'is_change', 'hit'],\n",
    "\t\t\t['lickbout_for_flash_during_response_window'] + stim_filter_base,\n",
    "\t\t\t['engaged', 'is_change', 'miss']\n",
    "\t\t\t] + [[f'flashes_since_last_lick=={flashes_since_lick}', '~lickbout_for_flash_during_response_window'] + \n",
    "\t\t\t\t\tstim_filter_base for flashes_since_lick in range(2, max_flashes+1)]\n",
    "\n",
    "\tunit_data, stim_indices, unitIds, returned_session_ids = vbn_utils.unit_averaged_psth_col_matched(active_tensor_file, stim_table, session_list, unit_ids, cond_filters, 'baseline_running', \n",
    "\t\t\t\t\t\t\tbaseline_length=750, resp_window_length=750, num_iterations=1)\n",
    "\t\n",
    "\tpassive_unit_data, passive_stim_indices, passive_unitIds, passive_returned_session_ids = vbn_utils.unit_averaged_psth_col_matched(passive_tensor_file, stim_table, session_list, unit_ids, cond_filters, 'baseline_running', \n",
    "\t\t\t\t\t\t\tbaseline_length=750, resp_window_length=750, num_iterations=1)\n",
    "\t\n",
    "\tunit_data = [u for u in unit_data if u.size>0]\n",
    "\tunit_data = np.concatenate(unit_data, axis=1)\n",
    "\tpassive_unit_data = [u for u in passive_unit_data if u.size>0]\n",
    "\tpassive_unit_data = np.concatenate(passive_unit_data, axis=1)\n",
    "\tcluster_means = []\n",
    "\tfor icond, cond in enumerate(unit_data):\n",
    "\t\tif icond==0:\n",
    "\t\t\tlabel = 'change lick'\n",
    "\t\t\talpha = 1\n",
    "\t\t\tcolor = 'g'\n",
    "\t\t\tls = 'solid'\n",
    "\t\t\n",
    "\t\telif icond==1: \n",
    "\t\t\tlabel = 'non change lick'\n",
    "\t\t\tcolor = 'g'\n",
    "\t\t\talpha = 1\n",
    "\t\t\tls = 'dotted'\n",
    "\t\t\n",
    "\t\telif icond== 2:\n",
    "\t\t\tlabel = 'miss'\n",
    "\t\t\talpha = 1\n",
    "\t\t\tcolor = 'goldenrod'\n",
    "\t\t\tls = 'solid'\n",
    "\n",
    "\t\telif icond>2:\n",
    "\t\t\tlabel = f'{icond-1} flashes since lick'\n",
    "\t\t\talpha = (icond-1)/max_flashes\n",
    "\t\t\tcolor = 'k'\n",
    "\t\t\tls = 'solid'\n",
    "\t\t\n",
    "\t\t\n",
    "\t\t\n",
    "\t\tcond = np.array([exponential_convolve(c, 7, symmetrical=True) for c in cond])\n",
    "\t\tax.plot(time, np.nanmean(cond, axis=0)*1000, color=color, alpha=alpha, ls=ls,)# label=icond)\n",
    "\t\tcluster_means.append(np.nanmean(cond[:, 550:750]*1000, axis=1))\n",
    "    \n",
    "\tcluster_means = np.array(cluster_means)\n",
    "\tcluster_means = np.roll(cluster_means, shift=-3, axis=0)\n",
    "\tall_cluster_means[cluster] = cluster_means\n",
    "\tcluster_means = cluster_means - cluster_means[0][None,:]\n",
    "\tcluster_mean = np.nanmean(cluster_means, axis=1)\n",
    "\tcluster_sem = np.nanstd(cluster_means, axis=1)/cluster_means.shape[1]**0.5\n",
    "\t\n",
    "\tax_agg.plot(np.arange(2,max_flashes+1), cluster_mean[:-3], 'o', color=cluster_colors[cluster], label=cluster_names[cluster])\n",
    "\tax_agg.errorbar(np.arange(2,max_flashes+1), cluster_mean[:-3], yerr=cluster_sem[:-3], color=cluster_colors[cluster])\n",
    "\tax_agg.plot(np.arange(max_flashes+1,max_flashes+4), cluster_mean[-3:], 'o', color=cluster_colors[cluster],)# label=cluster)\n",
    "\tax_agg.errorbar(np.arange(max_flashes+1,max_flashes+4), cluster_mean[-3:], yerr=cluster_sem[-3:], color=cluster_colors[cluster])\n",
    "\t\n",
    "\tax.set_xlim(-300, 500)\n",
    "\tax.set_xlabel('Time from image onset (ms)')\n",
    "\tax.set_ylabel('Firing rate (Hz)')\n",
    "\n",
    "\tpassive_hit = np.array([exponential_convolve(c, 7, symmetrical=True) for c in passive_unit_data[0]])\n",
    "\tax.plot(time, np.nanmean(passive_hit, axis=0)*1000, color='g', alpha=0.5, ls='-.', label='passive hit')\n",
    "\tax.legend(['hit', 'false alarm', 'miss'] + list(np.arange(2,max_flashes+1)) + ['passive hit',], frameon=False, loc='upper left', bbox_to_anchor=(0, 1.05))\n",
    "\n",
    "\n",
    "\tvbn_utils.formatFigure(fig, ax)\n",
    "\n",
    "ax_agg.set_xticks(np.arange(2,max_flashes+4))\n",
    "ax_agg.set_xticklabels(list(np.arange(2,max_flashes+1)) + ['hit', 'false alarm', 'miss'])\n",
    "for il, label in enumerate(ax_agg.get_xticklabels()):\n",
    "\tif il > max_flashes-2:\n",
    "\t\tlabel.set_rotation(30)\n",
    "\n",
    "ax_agg.legend(frameon=False)\n",
    "ax_agg.set_xlabel('Image presentations since last lick')\n",
    "ax_agg.set_ylabel('Baseline firing rate (Hz)')\n",
    "vbn_utils.formatFigure(fig_agg, ax_agg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c9b3cda",
   "metadata": {},
   "source": [
    "## Analysis of behavior since last lick"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95af3d23",
   "metadata": {},
   "outputs": [],
   "source": [
    "stims = stim_table[stim_table['no_abnorm']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c281b80",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "table = stims[stims['engaged']]\n",
    "ax.plot(table.pivot_table(index='flashes_since_last_lick', values='is_change', aggfunc='mean'), 'ko-',)\n",
    "ax.set_xlim(0, 15.5)\n",
    "ax.set_xlabel('Image presentations since last lick')\n",
    "ax.set_ylabel('change probability')\n",
    "ax.set_ylim(-0.01,0.35)\n",
    "\n",
    "ax2 = ax.twinx()\n",
    "table = stims[(~stims['is_change'])&(~stims['omitted'])&(stims['engaged'])&(~stims['previous_omitted'])]\n",
    "ax2.plot(table.pivot_table(index='flashes_since_last_lick', \n",
    "            values='lickbout_for_flash_during_response_window', aggfunc='mean'), 'ro-',)\n",
    "ax2.set_xlim(0, 15.5)\n",
    "ax2.set_ylabel('false alarm rate', color='r')\n",
    "ax2.set_ylim(-0.01,0.35)\n",
    "ax2.tick_params(axis='y', colors='r')\n",
    "\n",
    "[a.spines['top'].set_visible(False) for a in [ax, ax2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c52c179a",
   "metadata": {},
   "outputs": [],
   "source": [
    "stims['rt_zscore'] = stims.groupby('session_id')['reaction_time'].transform(lambda x: (x-np.nanmean(x))/np.nanstd(x))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "hit_table = stims[stims['is_change']&stims['engaged']&(stims['flashes_since_last_lick']>1)&(stims['flashes_since_last_lick']<16)]\n",
    "vals = hit_table.pivot_table(index='flashes_since_last_lick', values='rt_zscore', aggfunc=('mean', 'sem'))\n",
    "vals.plot(y='mean', yerr='sem', kind='line', marker='o', color='k', ax=ax, label='hits')\n",
    "\n",
    "\n",
    "fa_table = stims[(stims['engaged'])&(~stims['is_change'])&(~stims['omitted'])&(~stims['previous_omitted'])&(~stims['grace_period_after_hit']) &\\\n",
    "    (stims['flashes_since_change']>1)&(stims['flashes_since_last_lick']>1)&(stims['flashes_since_last_lick']<16)]\n",
    "vals = fa_table.pivot_table(index='flashes_since_last_lick', values='rt_zscore', aggfunc=('mean', 'sem'))\n",
    "vals.plot(y='mean', yerr='sem', kind='line', marker='o', color='r', ax=ax, label='false alarms')\n",
    "ax.set_xlim(0,15.5)\n",
    "ax.set_ylim(-0.55,0.5)\n",
    "ax.set_xlabel('Image presentations since last lick')\n",
    "ax.set_ylabel('z-scored response time')\n",
    "\n",
    "plt.legend()\n",
    "vbn_utils.formatFigure(fig, ax)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "hit_table = stims[stims['is_change']&stims['engaged']&(stims['flashes_since_last_lick']>1)&(stims['flashes_since_last_lick']<16)]\n",
    "vals = hit_table.pivot_table(index='flashes_since_last_lick', values='reaction_time', aggfunc=('mean', 'sem'))\n",
    "vals.plot(y='mean', yerr='sem', kind='line', marker='o', color='k', ax=ax, label='hits')\n",
    "\n",
    "\n",
    "fa_table = stims[(stims['engaged'])&(~stims['is_change'])&(~stims['omitted'])&(~stims['previous_omitted'])&(~stims['grace_period_after_hit']) &\\\n",
    "    (stims['flashes_since_change']>1)&(stims['flashes_since_last_lick']>1)&(stims['flashes_since_last_lick']<16)]\n",
    "vals = fa_table.pivot_table(index='flashes_since_last_lick', values='reaction_time', aggfunc=('mean', 'sem'))\n",
    "vals.plot(y='mean', yerr='sem', kind='line', marker='o', color='r', ax=ax, label='false alarms')\n",
    "ax.set_xlim(0,15.5)\n",
    "ax.set_xlabel('Image presentations since last lick')\n",
    "ax.set_ylabel('Response time (s)')\n",
    "ax.set_ylim(0.35, 0.5)\n",
    "plt.legend()\n",
    "vbn_utils.formatFigure(fig, ax)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "185f1e43",
   "metadata": {},
   "source": [
    "## Responses by reaction time quintile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c56337e5",
   "metadata": {},
   "source": [
    "### Running matched"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf04482e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "fig_rt, ax_rt = plt.subplots()\n",
    "for q in range(5):\n",
    "\n",
    "    rts = stim_table[stim_table['rt_quintiles']==q]['reaction_time'].values\n",
    "    ax_rt.boxplot(rts, notch=True, positions=[q,], showfliers=False)\n",
    "ax_rt.set_xticks(np.arange(5))\n",
    "ax_rt.set_xticklabels(np.arange(1,6))\n",
    "ax_rt.set_ylabel('Response time (s)')\n",
    "ax_rt.set_xlabel('Response time quintile')\n",
    "vbn_utils.formatFigure(fig_rt, ax_rt)\n",
    "\n",
    "fig_agg, ax_agg = plt.subplots()\n",
    "\n",
    "cluster_colors = {\n",
    "    1: 'navy',  # Blue\n",
    "    2: 'mediumblue',  # Lighter blue\n",
    "    3: 'royalblue',  # Even lighter blue\n",
    "    4: 'cornflowerblue',  # Lighter blue\n",
    "    5: 'slateblue',  # Lightest blue\n",
    "    6: 'red',  # Red\n",
    "    7: 'coral',  # Lighter red\n",
    "    8: 'gold'   # Lightest red\n",
    "}\n",
    "\n",
    "cluster_names = {\n",
    "\t1: 'On Transient',\n",
    "\t2: 'On Sustained',\n",
    "\t3: 'Off Transient',\n",
    "\t4: 'Off Sustained',\n",
    "\t5: 'Image suppressed',\n",
    "\t6: 'Lick anticipation',\n",
    "\t7: 'Licking',\n",
    "\t8: 'Running'\n",
    "}\n",
    "\n",
    "time = np.arange(-750, 750)\n",
    "for cluster in range(6,8):\n",
    "\tfig, ax = plt.subplots()\n",
    "\tfig.suptitle(cluster_names[cluster])\n",
    "\t# fig.set_size_inches(12,6)\n",
    "\tunit_filter = du.get_units_in_cluster(units, cluster, clustering='new') & du.apply_unit_quality_filter(units) & \\\n",
    "\t\t\tdu.getUnitsInRegion(units, 'SCMRN')\n",
    "\t\t\t\n",
    "\tunit_ids = units.loc[unit_filter]['unit_id'].values\n",
    "\tif len(unit_ids)==0:\n",
    "\t\tcontinue\n",
    "\tsession_list = units.loc[unit_filter]['ecephys_session_id'].unique()\n",
    "\n",
    "\tmax_flashes = 7\n",
    "\n",
    "\tstim_filter_base = ['engaged', 'is_change', '~omitted', '~previous_omitted',]# 'flashes_since_change>5']\n",
    "\tcond_filters = [\n",
    "\t\t\t['engaged', 'is_change', 'hit'],\n",
    "\t\t\t['engaged', '~is_change', '~omitted', '~previous_omitted',\n",
    "                '~lickbout_for_flash_during_response_window', 'flashes_since_last_lick>1', 'flashes_since_change>5'],\n",
    "\t\t\t] + [[f'rt_quintiles=={rt_quintile}', 'lickbout_for_flash_during_response_window'] + \n",
    "\t\t\t\t\tstim_filter_base for rt_quintile in range(5)]\n",
    "\n",
    "\tunit_data, stim_indices, unitIds, returned_session_ids = vbn_utils.unit_averaged_psth_col_matched(active_tensor_file, \n",
    "\t\t\t\t\t\t\tstim_table, session_list, unit_ids, cond_filters, 'baseline_running', \n",
    "\t\t\t\t\t\t\tbaseline_length=750, resp_window_length=750, num_iterations=5)\n",
    "\t\n",
    "\tunit_data = [u for u in unit_data if u.size>0]\n",
    "\tunit_data = np.concatenate(unit_data, axis=1)\n",
    "\tcluster_means = []\n",
    "\tfor icond, cond in enumerate(unit_data):\n",
    "\t\tif icond==0:\n",
    "\t\t\tlabel = 'hit'\n",
    "\t\t\talpha = 0\n",
    "\t\t\tcolor = 'g'\n",
    "\t\t\tls = 'solid'\n",
    "\t\t\n",
    "\t\telif icond==1: \n",
    "\t\t\tlabel = 'correct reject'\n",
    "\t\t\tcolor = 'g'\n",
    "\t\t\talpha = 1\n",
    "\t\t\tls = 'dotted'\n",
    "\t\t\n",
    "\t\telif icond>=2:\n",
    "\t\t\tlabel = f'{icond-2}'\n",
    "\t\t\talpha = 1.2-(icond-1)/5\n",
    "\t\t\tcolor = 'k'\n",
    "\t\t\tls = 'solid'\n",
    "\t\t\n",
    "\t\tcond = np.array([exponential_convolve(c, 7, symmetrical=True) for c in cond])\n",
    "\t\tif icond>0:\n",
    "\t\t\tax.plot(time, np.nanmean(cond, axis=0)*1000, color=color, alpha=alpha, ls=ls, label=label)\n",
    "\t\tcluster_means.append(np.nanmean(cond[:, 550:750]*1000, axis=1))\n",
    "    \n",
    "\tax.set_xlabel('Time from flash (ms)')\n",
    "\n",
    "\tcluster_means = np.array(cluster_means)\n",
    "\tcluster_means = np.roll(cluster_means, shift=-2, axis=0)\n",
    "\tcluster_means = cluster_means - cluster_means[0][None,:]\n",
    "\n",
    "\tcluster_means_norm = cluster_means/np.nanmax(np.abs(cluster_means), axis=0)[None, :]\n",
    "\n",
    "\tcluster_mean = np.nanmean(cluster_means, axis=1)\n",
    "\tcluster_sem = np.nanstd(cluster_means, axis=1)/cluster_means.shape[1]**0.5\n",
    "\t\n",
    "\tax_agg.plot(np.arange(5), cluster_mean[:-2], 'o', color=cluster_colors[cluster], label=cluster_names[cluster])\n",
    "\tax_agg.errorbar(np.arange(5), cluster_mean[:-2], yerr=cluster_sem[:-2], color=cluster_colors[cluster])\n",
    "\tax_agg.plot(np.arange(5,7), cluster_mean[-2:], 'o', color=cluster_colors[cluster],)# label=cluster)\n",
    "\tax_agg.errorbar(np.arange(5,7), cluster_mean[-2:], yerr=cluster_sem[-2:], color=cluster_colors[cluster])\n",
    "\t\n",
    "\tax.legend(['correct reject',] + list(np.arange(5)+1), frameon=False, loc='upper left', bbox_to_anchor=(0, 1.05))\n",
    "\tax.set_xlim(-400, 750)\n",
    "\tax.set_xlabel('Time from image onset (ms)')\n",
    "\tax.set_ylabel('Firing rate (Hz)')\n",
    "\tvbn_utils.formatFigure(fig, ax)\n",
    "\n",
    "\n",
    "ax_agg.set_xticks(np.arange(7))\n",
    "ax_agg.set_xticklabels(list(np.arange(5)+1) + ['hit', 'cr'])\n",
    "ax_agg.legend(frameon=False, loc='upper left', bbox_to_anchor=(0, 0.2))\n",
    "ax_agg.set_xlabel('Response time quintile')\n",
    "ax_agg.set_ylabel('Baseline firing rate (Hz)')\n",
    "vbn_utils.formatFigure(fig_agg, ax_agg)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vbn_manuscript",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}