analyzing_units_stability_across_blocks.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load unit stability information\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "from scipy.stats import ttest_rel, wilcoxon\n",
    "from scipy.io import loadmat\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.gridspec import GridSpec\n",
    "import seaborn as sns\n",
    "\n",
    "%matplotlib widget\n",
    "\n",
    "\n",
    "main_dir = '/home/kouroshmaboudi/Documents/Learned_tuning_Python/Datasets/'\n",
    "rr = os.listdir(main_dir)\n",
    "\n",
    "current_sessions = [x for x in range(5)] # Grosmark sessions\n",
    "# current_sessions = [x for x in range(6, 11)] # Giri sessions\n",
    "# current_sessions = [x for x in range(11, 17) if x not in [12, 13]] # Hiro's sessions\n",
    "# current_sessions.append(5)\n",
    "\n",
    "number_of_sessions = len(current_sessions)\n",
    "\n",
    "num_units = np.empty((number_of_sessions,), dtype=int)\n",
    "each_unit_session_number = np.empty((number_of_sessions,), dtype=object) # to keep track of session identity for each unit when we pool together and sort the units from diferent sessions\n",
    "\n",
    "spatial_tunings_maze = np.empty((number_of_sessions,), dtype=object) # spatial tunings on MAZE\n",
    "\n",
    "spikes_all_sessions = np.empty((number_of_sessions,), dtype=object)\n",
    "\n",
    "for session_idx in range(number_of_sessions):\n",
    "\n",
    "    session_number = current_sessions[session_idx]\n",
    "    session_name = rr[session_number]\n",
    "\n",
    "    print(session_name + \" \" + str(session_idx+1) + \"/\" + str(number_of_sessions))\n",
    "\n",
    "    session_dataset_path = os.path.join(main_dir, session_name)\n",
    "\n",
    "\n",
    "    # Load epochs information\n",
    "    filename = f'{session_name}.fileInfo_for_python.mat'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    mat_file = loadmat(file_path)\n",
    "    session_info = mat_file[\"fileInfo\"]\n",
    "\n",
    "    epochs = session_info[\"behavior\"][0][0][0][0][\"time\"]\n",
    "    epoch_durations = epochs[:, 1] - epochs[:, 0]\n",
    "\n",
    "\n",
    "    # Load spike data\n",
    "    filename = f'{session_name}.spikes_for_python.mat'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    mat_file = loadmat(file_path)\n",
    "    spikes_pyr = mat_file[\"spikes_pyr\"]\n",
    "\n",
    "    spikes_pyr = loadmat(os.path.join(session_dataset_path, session_name + '.spikes_for_python.mat'))['spikes_pyr']\n",
    "    \n",
    "    if session_number in [6, 7]: # RatN and RatS\n",
    "        num_units_total = spikes_pyr[\"spatialTuning_smoothed\"].shape[0] # for RatN only\n",
    "    else:\n",
    "        num_units_total = spikes_pyr[\"spatialTuning_smoothed\"][0].shape[0]\n",
    "\n",
    "    num_pos_bins = spikes_pyr[\"spatialTuning_smoothed\"][0][0]['uni'][0][0].size\n",
    "\n",
    "\n",
    "    # Load unit stability information\n",
    "    filename = f'{session_name}.cluster_quality_by_block'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    mat_file = loadmat(file_path)\n",
    "    cluster_quality_by_block = mat_file['cluster_quality_by_block'][0]\n",
    "\n",
    "\n",
    "    spikes = [] # spike data and place field info of each unit\n",
    "    running_directions = {'LR', 'RL', 'uni'}\n",
    "\n",
    "    iter = 0\n",
    "    for unit in range(num_units_total):\n",
    "\n",
    "        unit_spikes = dict()\n",
    "\n",
    "        if session_number in [9, 10]: # for Rat V sessions\n",
    "            unit_spikes['spike_times'] = spikes_pyr['time'][0][unit]\n",
    "            unit_spikes['shank_id']    = spikes_pyr['id'][0][unit][0][1] \n",
    "            unit_spikes['cluster_id']  = spikes_pyr['id'][0][unit][0][0]\n",
    "\n",
    "        elif session_number in [6, 7]: # for RatN and RatS\n",
    "            unit_spikes['spike_times'] = spikes_pyr['time'][unit][0] \n",
    "            unit_spikes['shank_id']    = spikes_pyr['id'][unit][0][0][0]-1\n",
    "            unit_spikes['cluster_id']  = spikes_pyr['id'][unit][0][0][1]\n",
    "\n",
    "        elif session_number == 8: # RatU  \n",
    "            unit_spikes['spike_times'] = spikes_pyr['time'][0][unit]\n",
    "            unit_spikes['shank_id']    = spikes_pyr['id'][0][unit][0][0] # shank indices already starts at zero\n",
    "            unit_spikes['cluster_id']  = spikes_pyr['id'][0][unit][0][1]\n",
    "        else:\n",
    "            unit_spikes['spike_times'] = spikes_pyr['time'][0][unit]\n",
    "            unit_spikes['shank_id']    = spikes_pyr['id'][0][unit][0][0]-1 # need to go one down for the other datasets\n",
    "            unit_spikes['cluster_id']  = spikes_pyr['id'][0][unit][0][1]\n",
    "\n",
    "\n",
    "        # Extract the cluster quality information by block for the current unit  \n",
    "        curr_unit_idx = np.where(cluster_quality_by_block['cluster_ids'][unit_spikes['shank_id']] == unit_spikes['cluster_id'])[0]\n",
    "        \n",
    "        spike_amplitude_by_block = cluster_quality_by_block['spike_amplitude_by_block'][unit_spikes['shank_id']][curr_unit_idx]\n",
    "        spike_amplitude_by_block = np.nan_to_num(spike_amplitude_by_block, nan=0)\n",
    "        spike_amplitude_by_block_percent = spike_amplitude_by_block/cluster_quality_by_block['session_mean_spike_amplitude'][unit_spikes['shank_id']][curr_unit_idx]# as a percentage of session mean\n",
    "        unit_spikes['spike_amplitude_by_block'] = spike_amplitude_by_block_percent\n",
    "\n",
    "        firing_rate_by_block = cluster_quality_by_block['firing_rate_by_block'][unit_spikes['shank_id']][curr_unit_idx]\n",
    "        firing_rate_by_block = np.nan_to_num(firing_rate_by_block, nan=0)\n",
    "        # firing_rate_by_block_percent = firing_rate_by_block/cluster_quality_by_block['session_mean_firing_rate'][unit_spikes['shank_id']][curr_unit_idx]\n",
    "        if firing_rate_by_block.shape[1] == 2: \n",
    "            sleep_firing_rate = (firing_rate_by_block[0][0]*epoch_durations[0] + firing_rate_by_block[0][1]*epoch_durations[2])/np.sum(epoch_durations[[0,2]])\n",
    "        elif firing_rate_by_block.shape[1] == 3:\n",
    "            sleep_firing_rate = (firing_rate_by_block[0][0]*epoch_durations[0] + firing_rate_by_block[0][1]*4*3600 + firing_rate_by_block[0][2]*(epoch_durations[2]-4*3600))/np.sum(epoch_durations[[0,2]])\n",
    "\n",
    "        if sleep_firing_rate > 0:\n",
    "            firing_rate_by_block_percent = firing_rate_by_block/sleep_firing_rate\n",
    "        else:\n",
    "            firing_rate_by_block_percent = np.zeros((len(firing_rate_by_block),))\n",
    "\n",
    "        unit_spikes['firing_rate_by_block'] = firing_rate_by_block_percent                \n",
    "\n",
    "        isolation_distance_by_block = cluster_quality_by_block['isolation_distance_by_block'][unit_spikes['shank_id']][curr_unit_idx]\n",
    "        isolation_distance_by_block = np.nan_to_num(isolation_distance_by_block, nan=0)\n",
    "        unit_spikes['isolation_distance_by_block'] = isolation_distance_by_block\n",
    "\n",
    "\n",
    "        unit_spikes['pre_post_unit_stability'] = (\n",
    "            (unit_spikes['spike_amplitude_by_block'][:2] > 0.67) & \n",
    "            (unit_spikes['firing_rate_by_block'][:2] > 0.33) & \n",
    "            (unit_spikes['isolation_distance_by_block'][:2] > 15)\n",
    "        ).all()\n",
    "        \n",
    "        # unit_spikes['unit_stability_latePOST'] = (\n",
    "        #     (unit_spikes['spike_amplitude_by_block'][0][2] > 0.67) &\n",
    "        #     (unit_spikes['firing_rate_by_block'][0][2] > 0.33) &\n",
    "        #     (unit_spikes['isolation_distance_by_block'][0][2] > 15)\n",
    "        #     ).all()\n",
    "\n",
    "        unit_spikes['place_fields_maze']  = {}\n",
    "        unit_spikes['peak_pos_bins_maze'] = {}\n",
    "        # unit_spikes['peak_firing_rate'] = {}\n",
    "\n",
    "        for direction in running_directions:\n",
    "            try:\n",
    "                if session_number in [6, 7]:    \n",
    "                    unit_spikes['place_fields_maze'][direction] = spikes_pyr[\"spatialTuning_smoothed\"][unit][0][direction][0][0].reshape(num_pos_bins) \n",
    "                    unit_spikes['peak_pos_bins_maze'][direction] = spikes_pyr['peakPosBin'][unit][0][direction][0][0][0][0]\n",
    "                else:\n",
    "                    unit_spikes['place_fields_maze'][direction] = spikes_pyr[\"spatialTuning_smoothed\"][0][unit][direction][0][0].reshape(num_pos_bins) \n",
    "                    unit_spikes['peak_pos_bins_maze'][direction] = spikes_pyr['peakPosBin'][0][unit][direction][0][0][0][0]\n",
    "\n",
    "            except ValueError:\n",
    "                if iter == 0:\n",
    "                    print(\"This session has only one running direction\")\n",
    "                iter += 1\n",
    "\n",
    "        spikes.append(unit_spikes) \n",
    "\n",
    "\n",
    "    # interpolate the spatial tunings to be consistent with Figure2\n",
    "    interp_pos_bins = np.linspace(0, num_pos_bins, 200)\n",
    "    num_pos_bins_interp = len(interp_pos_bins)\n",
    "\n",
    "    spatial_tunings_maze[session_idx] = np.zeros((num_units_total, num_pos_bins_interp))\n",
    "    for unit in range(len(spikes)):\n",
    "        spatial_tuning_maze_curr_unit = spikes[unit]['place_fields_maze']['uni']\n",
    "        spatial_tuning_maze_curr_unit = np.interp(interp_pos_bins, np.arange(1, num_pos_bins+1), spatial_tuning_maze_curr_unit)\n",
    "        spatial_tunings_maze[session_idx][unit] = spatial_tuning_maze_curr_unit\n",
    "\n",
    "    active_units = np.where(np.nanmax(spatial_tunings_maze[session_idx], axis=1) > 1)[0]\n",
    "    num_units[session_idx] = len(active_units)\n",
    "    \n",
    "\n",
    "    spikes_all_sessions[session_idx] = []\n",
    "    for unit in active_units: \n",
    "        spikes_all_sessions[session_idx].append(spikes[unit])\n",
    "    \n",
    "    each_unit_session_number[session_idx] = np.full((num_units[session_idx],), session_number)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the same for each individual session\n",
    "\n",
    "def get_pval_statement(pvalue):\n",
    "    if pvalue < 0.001:\n",
    "        pvalue_statement = 'P<0.001'\n",
    "    else:\n",
    "        pvalue_statement = f'P={pvalue:.3f}'\n",
    "    return pvalue_statement\n",
    "\n",
    "# Define a function to plot violin plots\n",
    "def plot_violin(ax, data, color):\n",
    "    alpha = 0.8\n",
    "\n",
    "    Q1 = np.percentile(data, 25, axis=0)\n",
    "    Q3 = np.percentile(data, 75, axis=0)\n",
    "    IQR = Q3 - Q1\n",
    "    lower_bound = Q1 - 1.5 * IQR\n",
    "    upper_bound = Q3 + 1.5 * IQR\n",
    "\n",
    "    # Filter out the outliers\n",
    "    data = data[~((data < lower_bound) | (data > upper_bound)).any(axis=1)]\n",
    "\n",
    "    sns.violinplot(data=data, ax=ax, inner='quartiles', linewidth=0, palette=color)\n",
    "\n",
    "    # for i, violin in enumerate(ax.collections[::2]):\n",
    "    #     violin.set_facecolor(color[i])\n",
    "\n",
    "    for violin, curr_alpha in zip(ax.collections[::1], [alpha] * 3):\n",
    "        violin.set_alpha(curr_alpha)\n",
    "    for l in ax.lines:\n",
    "        l.set_linestyle('-')\n",
    "        l.set_linewidth(0.75)\n",
    "        l.set_color('white')\n",
    "        l.set_alpha(1)\n",
    "    for l in ax.lines[1::3]:\n",
    "        l.set_linestyle('-')\n",
    "        l.set_linewidth(1.5)\n",
    "        l.set_color('white')\n",
    "        l.set_alpha(1)\n",
    "\n",
    "    ax.set_xlim([-0.5, 2.5])\n",
    "    if data.shape[1] == 2:\n",
    "        ax.set_xticklabels(['PRE', 'POST'], rotation=45, ha='center')\n",
    "    elif data.shape[1] == 3:\n",
    "        ax.set_xticklabels(['PRE', 'POST', 'latePOST'], rotation=45, ha='center')\n",
    "\n",
    "    for i , label in enumerate(ax.get_xticklabels()):\n",
    "        label.set_color(colors[i])\n",
    "\n",
    "    ax.tick_params(axis='both', which='major', labelsize=7, length=2, pad=0.5)\n",
    "\n",
    "    ax.grid(axis='y', color='gray', linewidth=1)\n",
    "    for axis in ['left', 'bottom']:\n",
    "        ax.spines[axis].set_linewidth(1.5)\n",
    "\n",
    "    # Add p-values corresponding to the statistical test of difference in mean between the blocks\n",
    "    data = np.nan_to_num(data, nan=0)\n",
    "    significance_bar_height = np.max(data)*1.05\n",
    "    p_value = np.empty((num_blocks, num_blocks), dtype=float)\n",
    "    for i in range(num_blocks):\n",
    "        for j in range(i+1, num_blocks):\n",
    "            p_value[i,j] = wilcoxon(data[:, i],\n",
    "                                data[:, j]).pvalue\n",
    "            ax.plot([i,j], [significance_bar_height, significance_bar_height], lw = 1, color = 'black')\n",
    "            ax.text((i+j)/2, significance_bar_height*1.01, get_pval_statement(p_value[i, j]), ha = 'center', va = \"bottom\", fontsize=6)\n",
    "            significance_bar_height = significance_bar_height + significance_bar_height*0.1\n",
    "\n",
    "colors = [\n",
    "    '#005CE9', # PRE\n",
    "    '#DD335D',  # POST\n",
    "    '#DC9A5D' # late POST\n",
    "    ] \n",
    "\n",
    "sns.set_style('whitegrid')\n",
    "sns.set_context('paper')\n",
    "\n",
    "custom_params = {\"axes.spines.right\": False, \"axes.spines.top\": False}\n",
    "sns.set_theme(style=\"ticks\", rc=custom_params)\n",
    "\n",
    "plt.rcParams['axes.linewidth'] = 1.5\n",
    "plt.rcParams['pdf.fonttype'] = 42\n",
    "plt.rcParams['ps.fonttype'] = 42\n",
    "\n",
    "\n",
    "plotheight = 800\n",
    "plotwidth = 300\n",
    "font_size = 6\n",
    "\n",
    "fig = plt.figure()\n",
    "fig.set_size_inches([plotwidth/72, plotheight/72])\n",
    "gs = GridSpec(nrows=number_of_sessions+1, ncols=3, figure=fig, height_ratios=[1]*(number_of_sessions+1))\n",
    "\n",
    "for session_idx in range(number_of_sessions):\n",
    "\n",
    "    session_number = current_sessions[session_idx]\n",
    "    session_name = rr[session_number]\n",
    "\n",
    "    num_blocks = spikes_all_sessions[0][0]['spike_amplitude_by_block'].shape[1]\n",
    "    num_units_current_session = num_units[session_idx]\n",
    "\n",
    "    spike_amplitude_by_block = np.empty((num_units_current_session, num_blocks), dtype=float)\n",
    "    firing_rate_by_block = np.empty((num_units_current_session, num_blocks), dtype=float)\n",
    "    isolation_distance_by_block = np.empty((num_units_current_session, num_blocks), dtype=float)\n",
    "\n",
    "    for unit in range(num_units_current_session):\n",
    "        spike_amplitude_by_block[unit, :] = spikes_all_sessions[session_idx][unit]['spike_amplitude_by_block']\n",
    "        firing_rate_by_block[unit, :] = spikes_all_sessions[session_idx][unit]['firing_rate_by_block']\n",
    "        isolation_distance_by_block[unit, :] = spikes_all_sessions[session_idx][unit]['isolation_distance_by_block']\n",
    "        \n",
    "\n",
    "    # Plot the distribution of spike amplitude for each block/epoch \n",
    "    ax = fig.add_subplot(gs[session_idx, 0])\n",
    "    plot_violin(ax, spike_amplitude_by_block*100, colors)\n",
    "    ax.set_ylabel('Spike amplitude (%)', fontsize=8)\n",
    "    ax.set_title(session_name, fontsize= 9, fontweight='normal', loc = 'left')\n",
    "    old_ylim = ax.get_ylim()\n",
    "    ax.set_ylim([0, old_ylim[1]])\n",
    "    # axs[0].set_ylim([0, 300])\n",
    "\n",
    "    # Plot the distribution of isolation distance for each block/epoch \n",
    "    ax = fig.add_subplot(gs[session_idx, 1])\n",
    "    plot_violin(ax, isolation_distance_by_block, colors)\n",
    "    ax.set_ylabel('Iso. distance (A.U.)', fontsize=8)\n",
    "    old_ylim = ax.get_ylim()\n",
    "    ax.set_ylim([0, old_ylim[1]])\n",
    "    # axs[1].set_ylim([0, 200])\n",
    "\n",
    "    # Plot the distribution of firing rate for each block/epoch \n",
    "    ax = fig.add_subplot(gs[session_idx, 2])\n",
    "    plot_violin(ax, firing_rate_by_block*100, colors)\n",
    "    ax.set_ylabel('Firing rate (%)', fontsize=8)\n",
    "    # axs[2].set_ylim([0, 400])\n",
    "\n",
    "\n",
    "# Concatenate the data from all sessions\n",
    "num_unit_all_sessions = np.sum(num_units)\n",
    "\n",
    "spike_amplitude_by_block_session_concat = np.empty((num_unit_all_sessions, num_blocks), dtype=float)\n",
    "firing_rate_by_block_session_concat = np.empty((num_unit_all_sessions, num_blocks), dtype=float)\n",
    "isolation_distance_by_block_session_concat = np.empty((num_unit_all_sessions, num_blocks), dtype=float)\n",
    "\n",
    "index = 0;\n",
    "for session_idx in range(number_of_sessions):\n",
    "    for unit in range(num_units[session_idx]):\n",
    "\n",
    "        spike_amplitude_by_block_session_concat[index, :] = spikes_all_sessions[session_idx][unit]['spike_amplitude_by_block']\n",
    "        firing_rate_by_block_session_concat[index, :] = spikes_all_sessions[session_idx][unit]['firing_rate_by_block']\n",
    "        isolation_distance_by_block_session_concat[index, :] = spikes_all_sessions[session_idx][unit]['isolation_distance_by_block']\n",
    "\n",
    "        index += 1\n",
    "\n",
    "\n",
    " # Plot the distribution of spike amplitude for each block/epoch \n",
    "ax = fig.add_subplot(gs[number_of_sessions, 0])\n",
    "plot_violin(ax, spike_amplitude_by_block_session_concat*100, colors)\n",
    "ax.set_ylabel('Spike amplitude (%)', fontsize=8)\n",
    "ax.set_title('Pooled', fontsize= 8, fontweight='bold', loc = 'left')\n",
    "old_ylim = ax.get_ylim()\n",
    "ax.set_ylim([0, old_ylim[1]])\n",
    "# axs[0].set_ylim([0, 300])\n",
    "\n",
    "# Plot the distribution of isolation distance for each block/epoch \n",
    "ax = fig.add_subplot(gs[number_of_sessions, 1])\n",
    "plot_violin(ax, isolation_distance_by_block_session_concat, colors)\n",
    "ax.set_ylabel('Iso. distance (A.U.)', fontsize=8)\n",
    "old_ylim = ax.get_ylim()\n",
    "ax.set_ylim([0, old_ylim[1]])\n",
    "# axs[1].set_ylim([0, 200])\n",
    "\n",
    "# Plot the distribution of firing rate for each block/epoch \n",
    "ax = fig.add_subplot(gs[number_of_sessions, 2])\n",
    "plot_violin(ax, firing_rate_by_block_session_concat*100, colors)\n",
    "ax.set_ylabel('Firing rate (%)', fontsize=8)\n",
    "# axs[2].set_ylim([0, 400])\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.subplots_adjust(wspace=0.5)\n",
    "sns.despine()\n",
    "\n",
    "filename = 'unit_stability_Grosmark_indiv_sessions.pdf'\n",
    "file_path = os.path.join(main_dir, filename)\n",
    "plt.savefig(file_path, format='pdf', dpi=300)\n",
    "\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Experimenting with printing on an exisiting PDF. Please ignore the code below as it is not yet functional "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PyPDF2\n",
    "\n",
    "\n",
    "\n",
    "def get_pval_statement(pvalue):\n",
    "    if pvalue < 0.001:\n",
    "        pvalue_statement = 'P<0.001'\n",
    "    else:\n",
    "        pvalue_statement = f'P={pvalue:.3f}'\n",
    "    return pvalue_statement\n",
    "\n",
    "colors = [\n",
    "    '#005CE9', # PRE\n",
    "    '#DD335D',  # POST\n",
    "    '#DC9A5D' # late POST\n",
    "    ] \n",
    "\n",
    "pdf_file_name = 'figure_test.pdf'\n",
    "\n",
    "with open(pdf_file_name, \"rb\") as file:\n",
    "    pdf_reader = PyPDF2.PdfReader(file)\n",
    "\n",
    "    page = pdf_reader.pages[0]\n",
    "\n",
    "\n",
    "\n",
    "    sns.set_style('whitegrid')\n",
    "    sns.set_context('paper')\n",
    "\n",
    "    custom_params = {\"axes.spines.right\": False, \"axes.spines.top\": False}\n",
    "    sns.set_theme(style=\"ticks\", rc=custom_params)\n",
    "\n",
    "    plt.rcParams['axes.linewidth'] = 1.5\n",
    "    plt.rcParams['pdf.fonttype'] = 42\n",
    "    plt.rcParams['ps.fonttype'] = 42\n",
    "\n",
    "    fig, axs = plt.subplots(nrows=1, ncols=3, sharey=False)\n",
    "    fig.set_size_inches([plotwidth/72, plotheight/72])\n",
    "\n",
    "\n",
    "    # Define a function to plot violin plots\n",
    "    def plot_violin(ax, data, color):\n",
    "        alpha = 0.8\n",
    "\n",
    "        Q1 = np.percentile(data, 25, axis=0)\n",
    "        Q3 = np.percentile(data, 75, axis=0)\n",
    "        IQR = Q3 - Q1\n",
    "        lower_bound = Q1 - 1.5 * IQR\n",
    "        upper_bound = Q3 + 1.5 * IQR\n",
    "\n",
    "        # Filter out the outliers\n",
    "        data = data[~((data < lower_bound) | (data > upper_bound)).any(axis=1)]\n",
    "\n",
    "        sns.violinplot(data=data, ax=ax, inner='quartiles', linewidth=0, palette=color)\n",
    "\n",
    "        # for i, violin in enumerate(ax.collections[::2]):\n",
    "        #     violin.set_facecolor(color[i])\n",
    "\n",
    "        for violin, curr_alpha in zip(ax.collections[::1], [alpha] * 3):\n",
    "            violin.set_alpha(curr_alpha)\n",
    "        for l in ax.lines:\n",
    "            l.set_linestyle('-')\n",
    "            l.set_linewidth(0.75)\n",
    "            l.set_color('white')\n",
    "            l.set_alpha(1)\n",
    "        for l in ax.lines[1::3]:\n",
    "            l.set_linestyle('-')\n",
    "            l.set_linewidth(1.5)\n",
    "            l.set_color('white')\n",
    "            l.set_alpha(1)\n",
    "\n",
    "        ax.set_xlim([-0.5, 2.5])\n",
    "        if data.shape[1] == 2:\n",
    "            ax.set_xticklabels(['PRE', 'POST'], rotation=45, ha='center')\n",
    "        elif data.shape[1] == 3:\n",
    "            ax.set_xticklabels(['PRE', 'POST', 'latePOST'], rotation=45, ha='center')\n",
    "\n",
    "        for i , label in enumerate(ax.get_xticklabels()):\n",
    "            label.set_color(colors[i])\n",
    "\n",
    "        ax.tick_params(axis='both', which='major', labelsize=7, length=2, pad=0.5)\n",
    "\n",
    "        ax.grid(axis='y', color='gray', linewidth=1)\n",
    "        for axis in ['left', 'bottom']:\n",
    "            ax.spines[axis].set_linewidth(1.5)\n",
    "\n",
    "        # Add p-values corresponding to the statistical test of difference in mean between the blocks\n",
    "        data = np.nan_to_num(data, nan=0)\n",
    "        significance_bar_height = np.max(data)*1.05\n",
    "        p_value = np.empty((num_blocks, num_blocks), dtype=float)\n",
    "        for i in range(num_blocks):\n",
    "            for j in range(i+1, num_blocks):\n",
    "                p_value[i,j] = wilcoxon(data[:, i],\n",
    "                                    data[:, j]).pvalue\n",
    "                ax.plot([i,j], [significance_bar_height, significance_bar_height], lw = 1, color = 'black')\n",
    "                ax.text((i+j)/2, significance_bar_height*1.01, get_pval_statement(p_value[i, j]), ha = 'center', va = \"bottom\", fontsize=6)\n",
    "                significance_bar_height = significance_bar_height + significance_bar_height*0.1\n",
    "        \n",
    "\n",
    "    # Plot the distribution of spike amplitude for each block/epoch \n",
    "    plot_violin(axs[0], spike_amplitude_by_block_session_concat*100, colors)\n",
    "    axs[0].set_ylabel('Spike amplitude (%)', fontsize=8)\n",
    "    # axs[0].set_ylim([0, 300])\n",
    "\n",
    "    # Plot the distribution of isolation distance for each block/epoch \n",
    "    plot_violin(axs[1], isolation_distance_by_block_session_concat, colors)\n",
    "    axs[1].set_ylabel('Iso. distance (A.U.)', fontsize=8)\n",
    "    # axs[1].set_ylim([0, 200])\n",
    "\n",
    "    # Plot the distribution of firing rate for each block/epoch \n",
    "    plot_violin(axs[2], firing_rate_by_block_session_concat*100, colors)\n",
    "    axs[2].set_ylabel('Firing rate (%)', fontsize=8)\n",
    "    # axs[2].set_ylim([0, 400])\n",
    "\n",
    "\n",
    "    sns.despine()\n",
    "\n",
    "    fig.tight_layout(rect=[0,0,1,1])\n",
    "    fig.canvas.draw()\n",
    "\n",
    "    pdf_figure = \"figure.pdf\"\n",
    "    \n",
    "    fig.savefig(pdf_figure, format=\"pdf\")\n",
    "\n",
    "    existing_pdf_stream = page.get_object()\n",
    "    figure_pdf = PyPDF2.PdfReader(open(pdf_figure, 'rb')).pages[0]\n",
    "    existing_pdf_stream.merge_page(figure_pdf)\n",
    "\n",
    "    with open(pdf_file_name, \"wb\") as output_file:\n",
    "        pdf_writer = PyPDF2.PdfWriter()\n",
    "        pdf_writer.add_page(page)\n",
    "        for page_num in range(1,pdf_reader.getNumPages()):\n",
    "            pdf_writer.addPage(pdf_reader.getPage(page_num))\n",
    "        pdf_writer.write(output_file)\n",
    "\n",
    "    plt.close(fig)\n",
    "    os.remove(pdf_figure)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}