{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c8acb36",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.stats\n",
    "from scipy.stats import binned_statistic_2d, binned_statistic\n",
    "import decoding_utils as du\n",
    "import vbn_utils as vbn\n",
    "import ccf_utils\n",
    "from matplotlib.colors import LinearSegmentedColormap, to_rgba\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68069bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_utils import scatter_ccf, binned_stat_ccf, weighted_gaussian_filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "758404b4",
   "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\n",
    "    plt.rcParams['pdf.fonttype'] = 42\n",
    "    \n",
    "    plt.rcParams['figure.facecolor'] = 'white'\n",
    "    plt.rcParams['axes.facecolor'] = 'white'\n",
    "    plt.rcParams['savefig.facecolor'] = 'white'  # affects clipboard copy too\n",
    "    plt.rcParams['savefig.transparent'] = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce8e5968",
   "metadata": {},
   "source": [
    "## Data loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236ecb9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Paths to all of the useful supplemental tables and tensors\n",
    "unit_table_file = \"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/master_units_with_responsiveness.csv\"\n",
    "units = pd.read_csv(unit_table_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00a59f91",
   "metadata": {},
   "source": [
    "## Load GLM dropout scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e933657f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dropouts = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/np-behavior/vbn_data_release/supplemental_tables/GLM_dropout_with_unit_info_active_passive.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25278960",
   "metadata": {},
   "source": [
    "## CCF coordinate setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3580c453",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nrrd\n",
    "import ccf_utils as ccf\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from matplotlib.colors import to_rgba\n",
    "\n",
    "annotation_volume = nrrd.read(\"/Volumes/programs/mindscope/workgroups/np-behavior/annotation_25.nrrd\")[0]\n",
    "structure_tree = pd.read_csv(\"/Volumes/programs/mindscope/workgroups/dynamicrouting/dynamic_gating_insertions/ccf_structure_tree_2017.csv\")\n",
    "\n",
    "def get_area_mask(area, annotation_volume, structure_tree):\n",
    "\n",
    "    if len(structure_tree[structure_tree['acronym']==area]) == 0:\n",
    "        print(f'could not find area {area} in structure tree')\n",
    "        return\n",
    "    \n",
    "    area_id = structure_tree[structure_tree['acronym']==area]['id'].values[0]\n",
    "    if not area_id in annotation_volume:\n",
    "        area_ccf = structure_tree[structure_tree['parent_structure_id']==area_id]\n",
    "        area_mask = np.sum([annotation_volume==id for id in area_ccf['id'].values], axis=0)\n",
    "    else:\n",
    "        area_mask = annotation_volume==area_id\n",
    "    \n",
    "    return area_mask\n",
    "\n",
    "def get_area_projection(area, annotation_volume, structure_tree, plane, spacing=25, padding=0):\n",
    "\n",
    "    axis_to_average = {'coronal': 0, 'sagittal': 2, 'horizontal': 1}\n",
    "\n",
    "    if isinstance(area, list):\n",
    "        area_mask = np.full_like(annotation_volume, False)\n",
    "        for a in area:\n",
    "            amask = get_area_mask(a, annotation_volume, structure_tree)\n",
    "            area_mask = np.logical_or(area_mask, amask)\n",
    "        #area_mask = np.sum([get_area_mask(a, annotation_volume, structure_tree) for a in area], axis=0)\n",
    "    else:\n",
    "        area_mask = get_area_mask(area, annotation_volume, structure_tree)\n",
    "    \n",
    "    area_projection = np.nanmax(area_mask, axis=axis_to_average[plane])\n",
    "    if plane in ['coronal', 'horizontal']:\n",
    "        area_projection = area_projection[:, :228]\n",
    "    \n",
    "    rows, cols = np.where(area_projection)\n",
    "    min_row, max_row = np.min(rows)-padding, np.max(rows)+padding\n",
    "    min_col, max_col = np.min(cols)-padding, np.max(cols)+padding\n",
    "    bounding_box = np.array([min_row, max_row, min_col, max_col])\n",
    "\n",
    "    scale_factor = 25//spacing\n",
    "    projection_box = np.repeat(np.repeat(area_projection[min_row:max_row+1, min_col:max_col+1], scale_factor, axis=0), scale_factor, axis=1)\n",
    "    if plane in ['coronal', 'horizontal']:\n",
    "        projection_box = np.flipud(projection_box)\n",
    "        bounding_box = bounding_box[[2, 3, 0, 1]]\n",
    "    else:\n",
    "        projection_box = np.flipud(projection_box.T)\n",
    "        \n",
    "    return projection_box, bounding_box*scale_factor\n",
    "\n",
    "\n",
    "def get_border(array, coord):\n",
    "\n",
    "    if len(coord)==3:\n",
    "        x, y, z = coord\n",
    "        array_box = np.logical_not(array[x-1:x+2, y-1:y+2, z-1:z+2])\n",
    "    elif len(coord)==2:\n",
    "        x, y = coord\n",
    "        array_box = np.logical_not(array[x-1:x+2, y-1:y+2])\n",
    "\n",
    "    return any(array_box.flatten())\n",
    "\n",
    "\n",
    "def get_area_boundary(area, annotation_volume, structure_tree, plane, spacing=25):\n",
    "    area_projection, bounding_box = get_area_projection(area, annotation_volume, structure_tree, plane, spacing, padding=2)\n",
    "    \n",
    "    bounding_box = np.array(bounding_box)*25\n",
    "\n",
    "    coords = np.where(area_projection)\n",
    "    area_edges = np.zeros(area_projection.shape)\n",
    "    for coord in tuple(zip(*coords)):\n",
    "        area_edges[coord] = get_border(area_projection, coord)\n",
    "\n",
    "    return area_edges, bounding_box\n",
    "\n",
    "\n",
    "def draw_area_boundary(area, plane, area_edges, bounding_box, ax=None, cmap=None):\n",
    "\n",
    "    ax_lims = None\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots()\n",
    "    else:\n",
    "        ax_lims = ax.get_xlim(), ax.get_ylim()\n",
    "\n",
    "    if cmap is None:\n",
    "        color = ccf.get_area_color(vbn.make_iterable(area)[0], structure_tree)\n",
    "        colors = [(1, 1, 1, 0), to_rgba(color)]  # RGBA for transparent (0) and red (1)\n",
    "        n_bins = 10  # Discretize the interpolation into bins\n",
    "        area_cmap = LinearSegmentedColormap.from_list('area_map', colors, N=n_bins)\n",
    "    else:\n",
    "        area_cmap = cmap\n",
    "    \n",
    "    if plane == 'coronal':\n",
    "        \n",
    "        im = ax.imshow(area_edges, extent=[bounding_box[0], bounding_box[1], bounding_box[2], bounding_box[3]], \n",
    "                       cmap=area_cmap, aspect='auto')\n",
    "        #ax.invert_xaxis()\n",
    "        ax.set_xlabel('left->right')\n",
    "        ax.set_ylabel('ventral->dorsal')\n",
    "\n",
    "    elif plane == 'sagittal':\n",
    "        im = ax.imshow(area_edges, extent=[bounding_box[0], bounding_box[1], bounding_box[2], bounding_box[3]], \n",
    "                       cmap=area_cmap, aspect='auto')\n",
    "        ax.set_xlabel('anterior->posterior')\n",
    "        ax.set_ylabel('ventral->dorsal')\n",
    "    \n",
    "    elif plane == 'horizontal':\n",
    "        im = ax.imshow(area_edges, extent=[bounding_box[0], bounding_box[1], bounding_box[2], bounding_box[3]], \n",
    "                       cmap=area_cmap, aspect='auto')\n",
    "        #ax.invert_xaxis()\n",
    "        ax.set_xlabel('left->right')\n",
    "        ax.set_ylabel('posterior->anterior')\n",
    "\n",
    "    if ax_lims:\n",
    "        ax.set_xlim(ax_lims[0])\n",
    "        ax.set_ylim(ax_lims[1])\n",
    "    \n",
    "    return ax\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b4a14bc",
   "metadata": {},
   "source": [
    "## Annotate dropout data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b800630b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dropouts['has_ccf'] = dropouts['left_right_ccf_coordinate']>0\n",
    "dropouts['cortical_layer'] = dropouts['cortical_layer'].replace('3-Feb','2/3') # 2/3 sometimes gets incorrectly reformatted as a date\n",
    "dropouts['cortical_layer'] = dropouts['cortical_layer'].replace('1','2/3')\n",
    "dropouts['cortical_layer'] = dropouts['cortical_layer'].replace('6a','6')\n",
    "dropouts['cortical_layer'] = dropouts['cortical_layer'].replace('6b','6')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a91baef0",
   "metadata": {},
   "source": [
    "## CCF spatial distribution of lick vs. image dropout scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77b3054c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['font.size'] = 18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d61db7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_area_ccf_cluster_distribution(area, metric, plane='sagittal', binsize=100, clim='auto'):\n",
    "    fig, ax = plt.subplots(2,2, figsize=(10,10))\n",
    "    fig.suptitle(f'{area}, {metric}')\n",
    "    projection = get_area_projection(area, annotation_volume, structure_tree, plane=plane, spacing=25)\n",
    "    ax[0][0].imshow(projection[0], extent=np.array((projection[1][0], projection[1][1], projection[1][2], projection[1][3]))*25, aspect='auto', cmap='gray', alpha=0.1)    \n",
    "    \n",
    "    area_neurons = vbn.get_unit_ids(dropouts[(dropouts['has_ccf'])&(dropouts['firing_rate']>=0.5)], area)\n",
    "    area_dropouts = dropouts[dropouts['unit_id'].isin(area_neurons)]\n",
    "    print(f'{area} mean: {area_dropouts[metric].mean()}')\n",
    "    if clim=='auto':\n",
    "        scatter_ccf(area_dropouts, plane=plane, c=metric, size=15, ax=ax[0][0],)# clim=[0,-0.05] )\n",
    "    else:\n",
    "        scatter_ccf(area_dropouts, plane=plane, c=metric, size=15, ax=ax[0][0], clim=clim)\n",
    "    \n",
    "    ax[0][1].imshow(projection[0], extent=np.array((projection[1][0], projection[1][1], projection[1][2], projection[1][3]))*25, aspect='auto', cmap='gray', alpha=0.1)    \n",
    "    values, _, _ = binned_stat_ccf(area_dropouts, binsize=binsize, plane=plane, c=metric, statistic=np.nanmean, ax=ax[0][1],)\n",
    "    counts, x_edge, y_edge = binned_stat_ccf(area_dropouts, binsize=binsize, plane=plane, c=metric, statistic='count', ax=ax[1][0])\n",
    "    \n",
    "    ax[1][1].imshow(projection[0], extent=np.array((projection[1][0], projection[1][1], projection[1][2], projection[1][3]))*25, aspect='auto', cmap='gray', alpha=0.1)    \n",
    "    smoothed = weighted_gaussian_filter(values, counts, 1).T\n",
    "    im = ax[1][1].imshow(-smoothed, extent=np.array((x_edge[0], x_edge[-1], y_edge[-1], y_edge[0])), aspect='auto', cmap='viridis_r',)\n",
    "    for area in vbn.make_iterable(area):\n",
    "        area_boundary, bounding_box = get_area_boundary(area, annotation_volume, structure_tree, plane, spacing=25)\n",
    "        for a in ax.flatten():\n",
    "            draw_area_boundary(area, plane, area_boundary, bounding_box, ax=a)\n",
    "            a.set_aspect('equal')\n",
    "\n",
    "    divider = make_axes_locatable(ax[1][1])\n",
    "    cbar_ax = divider.append_axes(\"right\", size=\"5%\", pad=0.1)  # Adjust size and padding\n",
    "    colorbar = fig.colorbar(im, cax=cbar_ax)\n",
    "\n",
    "    ax[0][0].set_title('individual units')\n",
    "    ax[0][1].set_title('binned mean')\n",
    "    ax[1][0].set_title('binned count')\n",
    "    ax[1][1].set_title('binned mean smoothed')\n",
    "\n",
    "    plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc87c24e",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['font.size'] = 18\n",
    "plot_area_ccf_cluster_distribution(['SCm', 'MRN'], \n",
    "                                   \"('absolute_change_from_full', 'licks')\", \n",
    "                                   plane='sagittal', binsize=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59b1c801",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_area_ccf_cluster_distribution(['SCm', 'MRN'], \n",
    "                                   \"('absolute_change_from_full', 'all-images')\", \n",
    "                                   plane='sagittal', binsize=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff60205a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_area_ccf_cluster_distribution(['VISam', 'VISpm'], \n",
    "                                   \"('absolute_change_from_full', 'licks')\", \n",
    "                                   plane='sagittal', binsize=50 ,clim=[0, 0.05])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e653b000",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "for ia, area in enumerate(['VISp', 'VISl', 'VISal', 'VISpm', 'VISam', 'VISrl']):\n",
    "    lick_vals = dropouts[dropouts['structure_acronym']==area][\"('absolute_change_from_full', 'licks')\"].values\n",
    "    \n",
    "    iteration_proportions = []\n",
    "    for iteration in range(1000):\n",
    "        iter_vals = np.random.choice(lick_vals, size=len(lick_vals), replace=True)\n",
    "        proportion_high_lick_units = np.sum(iter_vals<-0.01)/len(iter_vals)\n",
    "        iteration_proportions.append(proportion_high_lick_units)\n",
    "    \n",
    "    \n",
    "    ax.boxplot(iteration_proportions, positions=[ia], widths=0.5, patch_artist=True, \n",
    "               boxprops=dict(facecolor='gray'), showfliers=False, whis=[10, 90], medianprops=dict(color='black'), notch=True)\n",
    "\n",
    "ax.set_xticks(np.arange(len(['VISp', 'VISl', 'VISal', 'VISpm', 'VISam', 'VISrl'])))\n",
    "ax.set_xticklabels(['VISp', 'VISl', 'VISal', 'VISpm', 'VISam', 'VISrl'])\n",
    "ax.set_ylabel('Proportion of units \\n with lick dropout > 0.01')\n",
    "vbn.formatFigure(fig, ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49a1c5ca",
   "metadata": {},
   "source": [
    "## Overall GLM model performance across regions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3bb4c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "quality = du.apply_unit_quality_filter(units, no_abnorm=True)\n",
    "quality_units = units[quality]\n",
    "\n",
    "quality_units = quality_units[quality_units['brain_division'] != 'not in list']\n",
    "quality_units.loc[quality_units['structure_acronym'].isin(['SCig', 'SCiw']).astype(bool), 'structure_acronym'] = 'SCm'\n",
    "quality_units.loc[quality_units['structure_acronym'].isin(['MGv', 'MGd', 'MGm']).astype(bool), 'structure_acronym'] = 'MG'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e54bbc2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "areas_meet_threshold = quality_units['structure_acronym'].value_counts() > 100\n",
    "areas_meet_threshold = areas_meet_threshold[areas_meet_threshold].index\n",
    "areas_meet_threshold"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "glm-variance-explained-by-area",
   "metadata": {},
   "source": [
    "## GLM variance explained by area"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "glm-variance-boxplot",
   "metadata": {},
   "outputs": [],
   "source": [
    "dropouts_areas_100 = dropouts[dropouts['structure_acronym'].isin(areas_meet_threshold)]\n",
    "sorted_areas = dropouts_areas_100.groupby('structure_acronym').median()[\"('variance_explained', 'Full')\"].sort_values(ascending=True).index.values\n",
    "\n",
    "plt.rcParams['font.size'] = 16\n",
    "fig, ax = plt.subplots(figsize=(12, 6))\n",
    "for ia, area in enumerate(sorted_areas):\n",
    "    area_data = dropouts_areas_100[dropouts_areas_100['structure_acronym'] == area]\n",
    "    area_color = ccf.get_area_color(area, structure_tree)\n",
    "    ax.boxplot(area_data[\"('variance_explained_full', 'Full')\"].values, positions=[ia], widths=0.5, patch_artist=True, boxprops=dict(facecolor=area_color), medianprops=dict(color='k', linewidth=1),\n",
    "                showfliers=False, whis=[10, 90], notch=True)\n",
    "vbn.formatFigure(fig, ax, yLabel='Variance Explained (cross-validated)')\n",
    "ax.set_xticks(np.arange(len(sorted_areas)))\n",
    "ax.set_xticklabels(sorted_areas, rotation=90)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3051e641",
   "metadata": {},
   "outputs": [],
   "source": [
    "dropouts_areas_edited = dropouts.merge(quality_units[['unit_id', 'structure_acronym']], left_on='unit_id', right_on='unit_id', suffixes=('', '_edited'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a83b3630",
   "metadata": {},
   "outputs": [],
   "source": [
    "areas_meet_threshold = quality_units['structure_acronym'].value_counts() > 1000\n",
    "areas_meet_threshold = areas_meet_threshold[areas_meet_threshold].index\n",
    "areas_meet_threshold_units = vbn.get_unit_ids(quality_units, areas_meet_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "302a8adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2,1)\n",
    "ax = axes[0]\n",
    "ax2 = axes[1]\n",
    "fig.set_size_inches(12,5)\n",
    "area_dropout_pivot = dropouts_areas_edited[dropouts_areas_edited['structure_acronym_edited'].isin(areas_meet_threshold)].pivot_table(index='structure_acronym_edited', \n",
    "                                    values=[\"('absolute_change_from_full', 'all-images')\",\"('absolute_change_from_full', 'licks')\", '(\\'variance_explained\\', \\'Full\\')'], aggfunc='mean')\n",
    "\n",
    "areas_ordered = area_dropout_pivot.sort_values(by=\"('absolute_change_from_full', 'all-images')\").index.values\n",
    "\n",
    "for ia, area in enumerate(areas_ordered):\n",
    "    vis_vals = -dropouts_areas_edited[dropouts_areas_edited['structure_acronym_edited']==area][\"('absolute_change_from_full', 'all-images')\"].values\n",
    "    lick_vals = -dropouts_areas_edited[dropouts_areas_edited['structure_acronym_edited']==area][\"('absolute_change_from_full', 'licks')\"].values\n",
    "\n",
    "    area_color = ccf_utils.get_area_color(area, structure_tree)\n",
    "    ax.boxplot(vis_vals, positions=[ia], widths=0.5, patch_artist=True, showfliers=False, \n",
    "               whis=[10, 90], boxprops=dict(facecolor=area_color), medianprops=dict(color='k', linewidth=1), notch=True)\n",
    "    ax2.boxplot(lick_vals, positions=[ia], widths=0.5, patch_artist=True, showfliers=False, \n",
    "                whis=[10, 90],  boxprops=dict(facecolor=area_color), medianprops=dict(color='k', linewidth=1), notch=True)\n",
    "\n",
    "ax.set_title('Image dropout scores')\n",
    "ax2.set_title('Lick dropout scores')\n",
    "\n",
    "ax.xaxis.set_visible(False)\n",
    "    \n",
    "ax2.set_xticks(np.arange(len(areas_ordered)))\n",
    "ax2.set_xticklabels(areas_ordered, rotation=90)\n",
    "\n",
    "[vbn.formatFigure(fig, a,) for a in [ax, ax2]]\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for ia, area in enumerate(areas_ordered):\n",
    "    if not 'VIS' in area:\n",
    "        continue\n",
    "    lick_vals = dropouts_areas_edited[dropouts_areas_edited['structure_acronym_edited']==area][\"('absolute_change_from_full', 'licks')\"].values\n",
    "    print(f'{area} {np.sum(lick_vals<-0.01)/len(lick_vals)} {np.median(lick_vals)}')\n",
    "    area_color = ccf_utils.get_area_color(area, structure_tree)\n",
    "    ax.boxplot(lick_vals, positions=[ia], widths=0.5, patch_artist=True, boxprops=dict(facecolor=area_color), showfliers=False, whis=[5, 95])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d534c4be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mixed encoding in SCRMN stronger than in VIS (looking at lick dropout scores for sensory cluster neurons)\n",
    "fig, ax = plt.subplots()\n",
    "area_drops = []\n",
    "for ia, area in enumerate(['VISall', 'SCMRN']):\n",
    "    unit_ids = vbn.get_unit_ids(units, areas=area, clusters='sensory', clustering='new')\n",
    "    drops = dropouts_areas_edited[dropouts_areas_edited['unit_id'].isin(unit_ids)][\"('absolute_change_from_full', 'licks')\"].values\n",
    "    area_drops.append(drops)\n",
    "    ax.boxplot(-drops, positions=[ia], widths=0.5, patch_artist=True, boxprops=dict(facecolor='gray'), \n",
    "               showfliers=False, whis=[10, 90], medianprops=dict(color='k', linewidth=1), notch=True)\n",
    "\n",
    "pval = scipy.stats.ranksums(area_drops[0], area_drops[1])\n",
    "fig.suptitle(f'Lick dropout for sensory cluster; pval: {pval.pvalue:.2e}')\n",
    "\n",
    "print(pval)\n",
    "ax.set_xticks(np.arange(len(['VISall', 'SCm/MRN'])))\n",
    "ax.set_xticklabels(['VISall', 'SCm/MRN'])\n",
    "ax.set_ylabel('Lick dropout score')\n",
    "vbn.formatFigure(fig, ax,) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2939b835",
   "metadata": {},
   "outputs": [],
   "source": [
    "import vbn_utils\n",
    "for region in ['VISall', 'SCMRN']:\n",
    "\n",
    "    good_units = vbn_utils.get_unit_ids(dropouts, region)\n",
    "    region_dropouts = dropouts.set_index('unit_id').loc[good_units]\n",
    "\n",
    "    # print(region, region_dropouts[\"('absolute_change_from_full', 'licks')\"].mean())\n",
    "    print(region, 'images', np.sum(region_dropouts[\"('absolute_change_from_full', 'all-images')\"].values<-0.01)/len(region_dropouts), len(region_dropouts))\n",
    "    print(region, 'licks', np.sum(region_dropouts[\"('absolute_change_from_full', 'licks')\"].values<-0.01)/len(region_dropouts))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2134ca3f",
   "metadata": {},
   "source": [
    "## Running modulation for cortical interneurons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574c9894",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import nwb_session_utils as nwb\n",
    "from analysis_utils import exponential_convolve, makePSTH_numba\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8646303",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lagged_pearson_correlation(x, y, maxlag = 20):\n",
    "    \"\"\"\n",
    "    Calculates the Pearson correlation coefficient for all possible lags \n",
    "    using numpy.correlate for cross-correlation.\n",
    "    \"\"\"\n",
    "    # 1. Detrend the signals by removing the mean\n",
    "    x_dm = x - np.mean(x)\n",
    "    y_dm = y - np.mean(y)\n",
    "\n",
    "    # 2. Compute the cross-correlation of the detrended signals\n",
    "    cross_corr = np.correlate(x_dm, y_dm, mode='full')\n",
    "\n",
    "    # 3. Calculate normalization factors for each lag\n",
    "    n = len(x) if len(x) < len(y) else len(y)\n",
    "    length_diff = abs(len(x) - len(y))\n",
    "    \n",
    "    # The number of overlapping points for each lag\n",
    "    overlap_lengths = np.arange(1, n + 1)\n",
    "    overlap_lengths = np.concatenate([overlap_lengths, [n]*length_diff, overlap_lengths[:-1][::-1]])\n",
    "    lags = np.arange(-len(x) + 1, len(y))\n",
    "\n",
    "    # 4. Normalize the cross-correlation by the product of standard deviations \n",
    "    # and the number of overlapping points\n",
    "    std_dev_product = np.std(x) * np.std(y)\n",
    "    # Handle potential division by zero if std dev is 0\n",
    "    if std_dev_product == 0:\n",
    "        return np.zeros_like(cross_corr)[np.abs(lags) <= maxlag], lags[np.abs(lags) <= maxlag]\n",
    "\n",
    "    # Normalization factor at each lag\n",
    "    normalization = std_dev_product * overlap_lengths\n",
    "    pearson_r_lags = cross_corr / normalization\n",
    "\n",
    "    if maxlag is not None:\n",
    "        pearson_r_lags = pearson_r_lags[np.abs(lags) <= maxlag]\n",
    "        lags = lags[np.abs(lags) <= maxlag]\n",
    "    \n",
    "    return pearson_r_lags, lags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67a2fbbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "session_list = units[units['no_anomalies']]['ecephys_session_id'].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c72f0a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "quality_units = du.apply_unit_quality_filter(units)\n",
    "quality_units = units.loc[quality_units]['unit_id']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5008fd66",
   "metadata": {},
   "outputs": [],
   "source": [
    "stimulus_block = 3 # use spontaneous activity\n",
    "binsize = 0.1\n",
    "unit_ccg_dict = {}\n",
    "unit_activity_dict = {}\n",
    "session_running_dict = {}\n",
    "\n",
    "for isess, session in enumerate(session_list):\n",
    "    print(f'processing session {isess+1} of {len(session_list)}: {session}')\n",
    "    nwbfile = nwb.get_session(session)\n",
    "    stims = nwbfile.stimulus_presentations\n",
    "    block_start = stims[stims['stimulus_block']==stimulus_block]['start_time'].iloc[0]\n",
    "    block_end = stims[stims['stimulus_block']==stimulus_block]['end_time'].iloc[-1]\n",
    "\n",
    "    running = nwbfile.running_speed\n",
    "    session_units = nwbfile.get_units()\n",
    "    session_unit_ids = [uid for uid in session_units.index.values if uid in quality_units.values]\n",
    "    if len(session_unit_ids)==0:\n",
    "        continue\n",
    "\n",
    "    runvals, runtimes = nwb.resample_df_to_times(running, 'timestamps', 'speed', np.arange(block_start, block_end+binsize, binsize)*1000)\n",
    "    session_running_dict[session] = (runvals, runtimes)\n",
    "\n",
    "    for unit_id in session_unit_ids:\n",
    "        \n",
    "        spike_times = nwbfile.spike_times[unit_id]\n",
    "        block3_firing_rate, time_bins = makePSTH_numba(spike_times, [block_start,], block_end - block_start, binsize)\n",
    "\n",
    "        minlength = min(len(block3_firing_rate), len(runvals))\n",
    "\n",
    "        unit_ccg, unit_lags = lagged_pearson_correlation(block3_firing_rate[:minlength], runvals[:minlength], maxlag=50)\n",
    "        unit_ccg_dict[unit_id] = (unit_ccg, unit_lags)\n",
    "        unit_activity_dict[unit_id] = block3_firing_rate[:minlength]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b30c2a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "binsize = 0.1\n",
    "cell_types = ['FS', 'SST', 'VIP']\n",
    "cell_type_color_dict = {'RS':'k', 'FS':'red', 'SST':'dodgerblue', 'VIP':'orchid'}\n",
    "for cell_type in cell_types:\n",
    "\n",
    "    if 'RS' in cell_type and len(cell_type)>2:\n",
    "        layer = cell_type.split('RS')[-1][1:]\n",
    "        cell_type = 'RS'\n",
    "    else:\n",
    "        layer = 'all'\n",
    "\n",
    "    unit_filter = du.getUnitsInRegion(units, 'VISall', cell_type=cell_type, layer = layer) & \\\n",
    "                        du.apply_unit_quality_filter(units)\n",
    "    unit_ids = units.loc[unit_filter]['unit_id'].values\n",
    "\n",
    "    cell_type_ccgs = []\n",
    "    for unit_id in unit_ids:\n",
    "        # if unit_id in unit_ccg_dict:\n",
    "        unit_ccg, unit_lags = unit_ccg_dict[unit_id]\n",
    "        cell_type_ccgs.append(unit_ccg)\n",
    "    \n",
    "    cell_type_ccgs = np.array(cell_type_ccgs)\n",
    "\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.set_title(cell_type + ' ' + layer)\n",
    "    vbn_utils.mean_sem_plot(np.array(cell_type_ccgs), ax, unit_lags*binsize, color=cell_type_color_dict[cell_type])\n",
    "    ax.set_xlabel('lag (s)')\n",
    "    ax.set_ylabel('Correlation')\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.hist(cell_type_ccgs[:, unit_lags==0], bins=np.arange(-1, 1, 0.05))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2844fe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "binsize = 0.1\n",
    "cell_types = ['RSl2/3', 'RSl4', 'RSl5', 'RSl6', 'FS', 'SST', 'VIP']\n",
    "cell_types = ['FS', 'SST', 'VIP']\n",
    "cell_type_color_dict = {'RS':'k', 'FS':'red', 'SST':'dodgerblue', 'VIP':'orchid'}\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "zero_lags = []\n",
    "for cell_type in cell_types:\n",
    "\n",
    "    if 'RS' in cell_type and len(cell_type)>2:\n",
    "        layer = cell_type.split('RS')[-1][1:]\n",
    "        cell_type = 'RS'\n",
    "    else:\n",
    "        layer = 'all'\n",
    "\n",
    "    unit_filter = du.getUnitsInRegion(units, 'VISp', cell_type=cell_type, layer = layer) & \\\n",
    "                        du.apply_unit_quality_filter(units) \n",
    "    unit_ids = units.loc[unit_filter]['unit_id'].values\n",
    "\n",
    "    cell_type_ccgs = []\n",
    "    for unit_id in unit_ids:\n",
    "        # if unit_id in unit_ccg_dict:\n",
    "        unit_ccg, unit_lags = unit_ccg_dict[unit_id]\n",
    "        cell_type_ccgs.append(unit_ccg)\n",
    "    \n",
    "    cell_type_ccgs = np.array(cell_type_ccgs)\n",
    "    vbn_utils.mean_sem_plot(np.array(cell_type_ccgs), ax, unit_lags*binsize, color=cell_type_color_dict[cell_type])\n",
    "    zero_lags.append(cell_type_ccgs[:, unit_lags==0])\n",
    "\n",
    "ax.set_xlabel('lag (s)')\n",
    "ax.set_ylabel('Correlation')\n",
    "vbn_utils.formatFigure(fig, ax)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7147323b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.bar(cell_types, [np.nanmean(zl) for zl in zero_lags], yerr=np.array([np.nanstd(zl)/np.sqrt(sum(~np.isnan(zl))) for zl in zero_lags]).squeeze(), color=['red', 'dodgerblue', 'orchid'], alpha=0.7,)\n",
    "ax.set_ylabel('Zero lag correlation')\n",
    "\n",
    "vbn_utils.formatFigure(fig, ax)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vbn_manuscript",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}