{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "058c9794",
   "metadata": {},
   "source": [
    "NOTE: This code must be run with a fork of the allensdk: https://github.com/arjunsridhar12345/AllenSDK as well as access to the NWBs for the four-day experiments, which are not part of the main data release."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57a1ac7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import glob\n",
    "from matplotlib import pyplot as plt\n",
    "from analysis_utils import exponential_convolve\n",
    "from scipy.stats import kstest, ranksums\n",
    "from statsmodels.stats.multitest import multipletests\n",
    "import vbn_utils\n",
    "import vbn_4day_utils\n",
    "import decoding_utils as du\n",
    "%matplotlib inline\n",
    "from notebook_utils import fwhm, fraction_above_half_max, get_peak_time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5faafe8a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55cc89b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "high_res = True\n",
    "if high_res:\n",
    "    plt.rcParams['figure.dpi'] = 300\n",
    "    plt.rcParams['savefig.dpi'] = 300\n",
    "    plt.rcParams['font.size'] = 12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35954b61",
   "metadata": {},
   "source": [
    "## Data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f6a74f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "probe_info = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/VBN_four_day_experiment_nwbs/probes.csv\")\n",
    "session_info = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/VBN_four_day_experiment_nwbs/sessions.csv\")\n",
    "all_units = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/VBN_four_day_experiment_nwbs/units.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79f75bdf",
   "metadata": {},
   "source": [
    "## Annotate sessions with useful metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b913976",
   "metadata": {},
   "outputs": [],
   "source": [
    "for session, sess_info in session_info.iterrows():\n",
    "\n",
    "    dirname = sess_info['session_path']\n",
    "    exp_id = sess_info['exp_id']\n",
    "    nwb_path = glob.glob(os.path.join(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/VBN_four_day_experiment_nwbs\", \n",
    "                                      str(exp_id),\n",
    "                                      '*.nwb'))[0]\n",
    "    session_info.loc[session, 'nwb_path'] = nwb_path\n",
    "\n",
    "    area_calls = glob.glob(os.path.join(dirname, '*areaClassifications.csv'))[0]\n",
    "    area_calls = pd.read_csv(area_calls)\n",
    "\n",
    "    for ip, proberow in area_calls.iterrows():\n",
    "        probe = proberow['Probe']\n",
    "\n",
    "        probe_info_index = probe_info[(probe_info['probe_name']==probe) & (probe_info['session_id']==sess_info['exp_id'])].index\n",
    "        if len(probe_info_index) == 0:\n",
    "            continue\n",
    "\n",
    "        probe_info_index = probe_info_index[0]\n",
    "        probe_area = proberow['Area']\n",
    "\n",
    "        probe_info.loc[probe_info_index, 'area'] = probe_area"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1e17f6d",
   "metadata": {},
   "source": [
    "## Unit quality filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a18a5c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "good_session_info = session_info[~session_info['exp_id'].isin([1381868515,1382052219,1382241611,1382465116])]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a69cc75",
   "metadata": {},
   "source": [
    "## Overview PSTH by image set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1166286",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_ctx = all_units['in_cortex']\n",
    "quality = all_units['quality'] #this label includes standard quality metrics for this dataset\n",
    "no_seizure = ~all_units['session_id'].isin([1381868515,1382052219,1382241611,1382465116])\n",
    "in_vis_area = all_units['area'].isin(['VISp', 'VISl', 'VISrl', 'VISam', 'VISpm', 'VISal'])\n",
    "rs = all_units['waveform_duration'] > 0.4\n",
    "unit_ids = all_units[quality & in_ctx & no_seizure & in_vis_area & rs]['unit_id'].values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb0d7ae",
   "metadata": {},
   "source": [
    "## Nonchange flash PSTHs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c355df49",
   "metadata": {},
   "outputs": [],
   "source": [
    "stim_filter = [ '~is_change',\n",
    "                '~omitted', \n",
    "                '~previous_omitted',\n",
    "                'flashes_since_change>5',\n",
    "                'flashes_since_last_lick>1',]\n",
    "bin_size = 0.005\n",
    "win_before = 1\n",
    "win_after = 1\n",
    "gpsth, bins = vbn_4day_utils.unit_averaged_psth_from_nwb(good_session_info[good_session_info['image_set']=='G']['exp_id'].values, unit_ids, stim_filter + ['~image_name.isin([\"im083_r\", \"im111_r\"])'], cut=False, win_before=win_before, win_after=win_after, bin_size=bin_size)\n",
    "h1psth, bins = vbn_4day_utils.unit_averaged_psth_from_nwb(good_session_info[good_session_info['stage']=='EPHYS_day_2']['exp_id'].values, unit_ids, stim_filter + ['~image_name.isin([\"im083_r\", \"im111_r\"])'], cut=False, win_before=win_before, win_after=win_after, bin_size=bin_size)\n",
    "h2psth, bins = vbn_4day_utils.unit_averaged_psth_from_nwb(good_session_info[good_session_info['stage']=='EPHYS_day_3']['exp_id'].values, unit_ids, stim_filter + ['~image_name.isin([\"im083_r\", \"im111_r\"])'], cut=False, win_before=win_before, win_after=win_after, bin_size=bin_size)\n",
    "kpsth, bins = vbn_4day_utils.unit_averaged_psth_from_nwb(good_session_info[good_session_info['stage']=='EPHYS_day_4']['exp_id'].values, unit_ids, stim_filter + ['~image_name.isin([\"im024_r\", \"im034_r\"])'], cut=False, win_before=win_before, win_after=win_after, bin_size=bin_size)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2bc21de",
   "metadata": {},
   "source": [
    "## Response properties: peak time, amplitude, FWHM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74cf258b",
   "metadata": {},
   "outputs": [],
   "source": [
    "win_before = 1\n",
    "win_after = 1\n",
    "bin_size = 0.005\n",
    "bins = np.arange(-win_before, win_after, bin_size)\n",
    "base_end = int(win_before/bin_size)\n",
    "stim_end = int((win_before + 0.25)/bin_size)\n",
    "resp_slice = slice(base_end + int(0.010/bin_size), stim_end + int(0.050/bin_size))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "fig_peak_time, ax_peak_time = plt.subplots(constrained_layout=True)\n",
    "fig_peak_amp, ax_peak_amp = plt.subplots()\n",
    "fig_fwhm, ax_fwhm = plt.subplots()\n",
    "peak_times = []\n",
    "peak_amps = []\n",
    "fwhms = []\n",
    "for ipsth, (psth, color, imlabel) in enumerate(zip([gpsth, h1psth, h2psth, kpsth], ('b', 'r', 'purple', '#FBAF41'), ('G', 'H1', 'H2', 'K'))):\n",
    "    goodpsth = [g for g in psth if len(g) > 0]\n",
    "    grand_psth = np.concatenate(goodpsth, axis=0)\n",
    "    base_end = int(win_before/bin_size)\n",
    "    base_start = base_end - int(0.1/bin_size)  # 0.1s baseline\n",
    "    base_sub = grand_psth - grand_psth[:,base_start:base_end].mean(axis=1)[:,np.newaxis]\n",
    "\n",
    "    mean = base_sub.mean(axis=0)\n",
    "    sem = base_sub.std(axis=0)/np.sqrt(base_sub.shape[0])\n",
    "    ax.plot(bins, mean, label=f'{imlabel}: n={base_sub.shape[0]}', color=color)\n",
    "    ax.fill_between(bins, mean-sem, mean+sem, alpha=0.5, color=color, lw=0)\n",
    "\n",
    "    time_to_peak = np.array([get_peak_time(trace, slice(resp_slice.start, int(resp_slice.start + 0.150/bin_size)), bin_size) for trace in base_sub]) - win_before\n",
    "    ax_peak_time.boxplot(time_to_peak, positions=[ipsth], widths=0.5, showfliers=False, whis=[10,90], notch=True, patch_artist=True,\n",
    "        boxprops=dict(facecolor=color, color=color), whiskerprops=dict(color=color),\n",
    "        capprops=dict(color=color), medianprops=dict(color='w', linewidth=2))\n",
    "    ax_peak_time.set_ylabel('peak time (s)')\n",
    "    ax_peak_time.set_xticks(np.arange(4))\n",
    "    ax_peak_time.set_xticklabels(['G', 'H1', 'H2', 'K'])\n",
    "    peak_times.append(time_to_peak)\n",
    "\n",
    "    peak_amp = base_sub[:,base_end:base_end+int(0.1/bin_size)].mean(axis=1)\n",
    "    peak_amp = peak_amp[~np.isnan(peak_amp)]\n",
    "    ax_peak_amp.boxplot(peak_amp, positions=[ipsth], widths=0.5, showfliers=False, whis=[10,90], notch=True, patch_artist=True,\n",
    "        boxprops=dict(facecolor=color, color=color), whiskerprops=dict(color=color),\n",
    "        capprops=dict(color=color), medianprops=dict(color='w', linewidth=2))\n",
    "    ax_peak_amp.set_ylabel('peak amplitude (Hz)')\n",
    "    ax_peak_amp.set_xticks(np.arange(4))\n",
    "    ax_peak_amp.set_xticklabels(['G', 'H1', 'H2', 'K'])\n",
    "    peak_amps.append(peak_amp)\n",
    "\n",
    "    full_width_half_max = np.array([fraction_above_half_max(bin_size, trace) for trace in base_sub[:,base_end:base_end+int(0.3/bin_size)]])\n",
    "    ax_fwhm.boxplot(full_width_half_max[~np.isnan(full_width_half_max)], positions=[ipsth], widths=0.5, showfliers=False, whis=[10,90], notch=True, patch_artist=True,\n",
    "        boxprops=dict(facecolor=color, color=color), whiskerprops=dict(color=color),\n",
    "        capprops=dict(color=color), medianprops=dict(color='w', linewidth=2))\n",
    "    ax_fwhm.set_ylabel('fwhm (s)')\n",
    "    fwhms.append(full_width_half_max[~np.isnan(full_width_half_max)])\n",
    "    ax_fwhm.set_xticks(np.arange(4))\n",
    "    ax_fwhm.set_xticklabels(['G', 'H1', 'H2', 'K'])\n",
    "\n",
    "ax.set_xlim(0, 0.3)\n",
    "ax.legend(frameon=False)\n",
    "ax.set_xlabel('Time from stimulus onset (s)')\n",
    "ax.set_ylabel('Firing rate (Hz)')\n",
    "vbn_utils.formatFigure(fig, ax)\n",
    "vbn_utils.formatFigure(fig_peak_time, ax_peak_time)\n",
    "vbn_utils.formatFigure(fig_peak_amp, ax_peak_amp)\n",
    "vbn_utils.formatFigure(fig_fwhm, ax_fwhm)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22f316e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import kstest, ranksums\n",
    "from statsmodels.stats.multitest import multipletests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c9eb31e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for vals, ax in zip([peak_times, peak_amps, fwhms], [ax_peak_time, ax_peak_amp, ax_fwhm]):\n",
    "    pvalues = []\n",
    "    comparisons = []\n",
    "    for ip1, p1 in enumerate(vals):\n",
    "        for ip2, p2 in enumerate(vals[ip1+1:]):\n",
    "            print(f'{ip1} vs {ip2+ip1+1}: {ranksums(p1, p2)}')\n",
    "            pvalues.append(ranksums(p1, p2).pvalue)\n",
    "            comparisons.append((ip1, ip2+ip1+1))\n",
    "\n",
    "    corrected_pvalues = multipletests(pvalues, method='fdr_bh')[1]\n",
    "    ax_max = ax.get_ylim()[1]\n",
    "    y_increment = ax_max/10\n",
    "    for i, (start, end) in enumerate(comparisons):\n",
    "        if corrected_pvalues[i]<0.05:\n",
    "            y = ax_max + y_increment*i\n",
    "            ax.plot([start, start, end, end], [y, y + y_increment/2, y + y_increment/2, y], lw=1.5, c='k',)\n",
    "            ax.text((start + end) * .5, y+ y_increment/2, f\"p = {corrected_pvalues[i]:.1e}\", ha='center', va='bottom', color='k', fontdict={'size': 14})\n",
    "    plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vbn_manuscript",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}