learned_tuning_mathcing_firing_bins.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Achilles_10252013\n",
      "Achilles_11012013\n",
      "Buddy_06272013\n",
      "Cicero_09172014\n",
      "Gatsby_08282013\n",
      "RatN_Day2_2019-10-11_03-58-54\n",
      "RatS-Day2-2020-11-27_10-22-29\n",
      "RatU_Day2NSD_2021-07-24_08-16-38\n",
      "RatV_Day1NSD_2021-10-02_08-10-23\n",
      "RatV_Day3NSD_2021-10-07_08-10-12\n",
      "Roy-maze1\n",
      "Ted-maze1\n",
      "Ted-maze2\n",
      "Ted-maze3\n",
      "Kevin-maze1\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import scipy.io\n",
    "from scipy.stats import ranksums\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import h5py\n",
    "from learned_tuning.learned_tuning import calculate_place_field_fidelity_of_learned_tuning, calculate_learned_tuning_matched_participation\n",
    "\n",
    "# Define functions\n",
    "def p_score(arr):\n",
    "    percentile_scores = []\n",
    "    n = len(arr)\n",
    "    for i in range(n):\n",
    "        current_element = arr[i]\n",
    "        count = np.sum(current_element > arr)  # Count elements greater than current_element\n",
    "        percentile = (count / n) * 100  # Calculate the percentile score\n",
    "        percentile_scores.append(percentile)\n",
    "\n",
    "    return percentile_scores\n",
    "\n",
    "normalize_to_max = lambda x:(x/np.nanmax(x))\n",
    "\n",
    "\n",
    "data_dir = r'/home/kouroshmaboudi/Documents/Learned_tuning_Python/Datasets'\n",
    "state_detection_dir =  r'/home/kouroshmaboudi/Documents/NCMLproject/StateDetectionResults/'\n",
    "\n",
    "sessions = [d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d))]\n",
    "\n",
    "included_session_IDX = [x for x in range(17) if x not in (12, 13, 5)]\n",
    "included_session_IDX.append(5)\n",
    "# included_session_IDX = [0]\n",
    "\n",
    "sessions = [sessions[i] for i in included_session_IDX]\n",
    "\n",
    "session_name_mapping = {'Achilles_10252013': 'Rat A linear',\n",
    "                        'Achilles_11012013': 'Rat A circular',\n",
    "                        'Buddy_06272013': 'Rat B',\n",
    "                        'Cicero_09172014': 'Rat C',\n",
    "                        'Gatsby_08282013':'Rat G',\n",
    "                        'RatN_Day2_2019-10-11_03-58-54': 'Rat N',\n",
    "                        'RatS-Day2-2020-11-27_10-22-29': 'Rat S',\n",
    "                        'RatU_Day2NSD_2021-07-24_08-16-38': 'Rat U',\n",
    "                        'RatV_Day1NSD_2021-10-02_08-10-23': 'Rat V linear',\n",
    "                        'RatV_Day3NSD_2021-10-07_08-10-12': 'Rat V semicircular',\n",
    "                        'Roy-maze1': 'Rat R',\n",
    "                        'Ted-maze1': 'Rat T linear',\n",
    "                        'Ted-maze2': 'Rat T L shape',\n",
    "                        'Ted-maze3': 'Rat T U shape',\n",
    "                        'Kevin-maze1': 'Rat K'}\n",
    "\n",
    "number_of_sessions = len(sessions)\n",
    "\n",
    "total_duration_per_brain_state = np.empty((number_of_sessions,), dtype=object)\n",
    "epoch_durations = np.empty((number_of_sessions,), dtype=object)\n",
    "\n",
    "session_names = []\n",
    "\n",
    "for session_idx, session_name in enumerate(sessions):\n",
    "\n",
    "    print(session_name)\n",
    "    session_names.append(session_name)\n",
    "\n",
    "    session_dataset_path = os.path.join(data_dir, session_name)\n",
    "    session_number = included_session_IDX[session_idx]\n",
    "\n",
    "\n",
    "    session_state_detection_results_path = os.path.join(state_detection_dir, session_name)\n",
    "\n",
    "    #--------------------------------------------------------------------------------------------\n",
    "    # Load epochs information\n",
    "\n",
    "    filename = f'{session_name}.fileInfo_for_python.mat'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    mat_file = scipy.io.loadmat(file_path)\n",
    "    session_info = mat_file[\"fileInfo\"]\n",
    "\n",
    "    epochs = session_info[\"behavior\"][0][0][0][0][\"time\"]\n",
    "    epoch_durations[session_idx] = epochs[:, 1] - epochs[:, 0]\n",
    "\n",
    "\n",
    "    #---------------------------------------------------------------------------------------------\n",
    "    # Brain state detection results\n",
    "\n",
    "    if 0<=session_number<=4 or 6<=session_number<=10: # Grosmark and Giri datasets\n",
    "        filename = f'{session_name}.brainStateDetection_HMMtheta_EMG_SWS_SchmidtTrigger.mat'\n",
    "        file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "        mat_file = scipy.io.loadmat(file_path)\n",
    "        brainStates_bouts_label = mat_file['brainState']['bouts'][0][0][:, :-1]\n",
    "        bouts_start_end = brainStates_bouts_label[:, :-1]\n",
    "        bout_durations = bouts_start_end[:, 1] - bouts_start_end[:, 0] \n",
    "        bout_labels = brainStates_bouts_label[:, -1].astype(int)\n",
    "\n",
    "        brainStates_names = []\n",
    "        for i in range(4):\n",
    "            brainStates_names.append(mat_file['brainState']['names'][0][0][i][0][0])\n",
    "\n",
    "    else: # Miyawaki dataset\n",
    "        filename = f'{session_name}.fileInfo_for_python.mat'\n",
    "        file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "        mat_file = scipy.io.loadmat(file_path)\n",
    "        brainStates_bouts_label = mat_file['fileInfo']['brainStates'].item()\n",
    "\n",
    "        bouts_start_end = brainStates_bouts_label[:, :-1]\n",
    "        bout_durations = bouts_start_end[:, 1] - bouts_start_end[:, 0] \n",
    "\n",
    "        bout_labels = brainStates_bouts_label[:, -1].astype(int)\n",
    "\n",
    "        # swapping the 3s and 4s, because the 3s in Hiro's brain_state_df are QWAKE while in the other datasets they are active wake(WAKE)\n",
    "        bout_labels_temp = bout_labels.copy()\n",
    "        bout_labels_temp[bout_labels == 3] = 4\n",
    "        bout_labels_temp[bout_labels == 4] = 3\n",
    "        bout_labels = bout_labels_temp[:]\n",
    "        del bout_labels_temp\n",
    "\n",
    "    brainStates_names = ['NREM', 'REM', 'WAKE', 'QWAKE'] # , 'Undetermined'\n",
    "\n",
    "\n",
    "    # Temporary for calculating the percentage of each brain state during different epochs\n",
    "    epoch_names = ['pre', 'post']\n",
    "    \n",
    "    total_duration_per_brain_state[session_idx] = dict({\n",
    "        'pre': dict.fromkeys(brainStates_names),\n",
    "        'post': dict.fromkeys(brainStates_names)\n",
    "        })\n",
    "\n",
    "    \n",
    "    for epoch in epoch_names:\n",
    "        if epoch == 'pre':\n",
    "            epoch_idx = 0\n",
    "            epoch_duration = epochs[0,1] - epochs[0,0]\n",
    "        elif epoch == 'post':\n",
    "            epoch_idx = 2\n",
    "            epoch_duration = 4*60*60 # limiting to the first 4 hours of post sleep\n",
    "\n",
    "        if_inside_epoch = (bouts_start_end[:, 0] > epochs[epoch_idx,0]) & (bouts_start_end[:, 0] < (epochs[epoch_idx,0]+epoch_duration))\n",
    "\n",
    "        for brain_state_index, brain_state in enumerate(brainStates_names):\n",
    "            \n",
    "            current_brain_state_bout_indexes = if_inside_epoch & (bout_labels == (brain_state_index+1))\n",
    "            current_brain_state_bout_durations = bout_durations[current_brain_state_bout_indexes]\n",
    "\n",
    "            total_duration_per_brain_state[session_idx][epoch][brain_state] = f'{np.sum(current_brain_state_bout_durations):.1f}'\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([array([18079.5   ,  2067.5   , 14714.1032]),\n",
       "       array([19330.   ,  2674.1  , 14625.112]),\n",
       "       array([10717.9,  2328.1,  8008.7]),\n",
       "       array([16000.7,  3077.8, 14481.7]),\n",
       "       array([16761.1,  2692.6, 15894.2]),\n",
       "       array([10270,  3265, 39718], dtype=uint16),\n",
       "       array([ 2358,  3671, 30866], dtype=uint16),\n",
       "       array([ 9544,  3310, 32460], dtype=uint16),\n",
       "       array([10145,  3124, 32450,  3163], dtype=uint16),\n",
       "       array([10135,  3217, 32172,  2421], dtype=uint16),\n",
       "       array([10799.957333, 10828.8992  , 10800.008533]),\n",
       "       array([10799.968, 11147.632, 10800.016]),\n",
       "       array([10799.968, 10815.856, 10800.016]),\n",
       "       array([10799.968, 10909.776, 10800.016]),\n",
       "       array([10799.984   , 10822.044188, 10800.      ])], dtype=object)"
      ]
     },
     "execution_count": 86,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epoch_durations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataframe(total_duration_per_brain_state, session_names, session_name_mapping):\n",
    "\n",
    "    sesseion_names = []\n",
    "    epoch_names = []\n",
    "    NREM_durations = []\n",
    "    REM_durations = []\n",
    "    AWAKE_durations = []\n",
    "    QWAKE_durations = []\n",
    "\n",
    "    for session_idx in range(len(total_duration_per_brain_state)):\n",
    "        for epoch_idx, epoch in enumerate(['pre', 'post']):\n",
    "            current_session_name = session_name_mapping[session_names[session_idx]]\n",
    "\n",
    "            if epoch_idx == 0:\n",
    "                sesseion_names.append(current_session_name)\n",
    "            else:\n",
    "                sesseion_names.append('')\n",
    "\n",
    "            epoch_names.append(epoch)\n",
    "\n",
    "            NREM_durations.append(total_duration_per_brain_state[session_idx][epoch]['NREM'])\n",
    "            REM_durations.append(total_duration_per_brain_state[session_idx][epoch]['REM'])\n",
    "            AWAKE_durations.append(total_duration_per_brain_state[session_idx][epoch]['WAKE'])\n",
    "            QWAKE_durations.append(total_duration_per_brain_state[session_idx][epoch]['QWAKE'])\n",
    "\n",
    "    df = pd.DataFrame({\n",
    "        \"Session\": sesseion_names,\n",
    "        \"epoch\": epoch_names,\n",
    "        \"SWS\": NREM_durations,\n",
    "        \"REM\": REM_durations,\n",
    "        \"Active wake\": AWAKE_durations,\n",
    "        \"Quiet wake\": QWAKE_durations   \n",
    "    })\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "total_duration_per_brain_state = list(total_duration_per_brain_state[5:10]) + list(total_duration_per_brain_state[10:]) + list(total_duration_per_brain_state[:5])\n",
    "session_names = session_names[5:10] + session_names[10:] + session_names[:5]\n",
    "\n",
    "\n",
    "duration_per_brain_state_df = get_dataframe(total_duration_per_brain_state, session_names, session_name_mapping)\n",
    "\n",
    "\n",
    "from docx import Document\n",
    "\n",
    "doc = Document()\n",
    "\n",
    "doc.add_heading('Sleep_state_durations', level = 1)\n",
    "\n",
    "table = doc.add_table(rows = 1, cols = len(duration_per_brain_state_df.columns))\n",
    "table.autofit = True\n",
    "\n",
    "for col_num, col_name in enumerate(duration_per_brain_state_df.columns):\n",
    "    table.cell(0, col_num).text = col_name\n",
    "\n",
    "for row_num in range(len(duration_per_brain_state_df)):\n",
    "    table.add_row()\n",
    "    for col_num, value in enumerate(duration_per_brain_state_df.iloc[row_num]):\n",
    "        table.cell(row_num+1, col_num).text = str(value)\n",
    "\n",
    "doc.save(os.path.join(data_dir, 'Sleep_state_durations.docx'))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Session</th>\n",
       "      <th>epoch</th>\n",
       "      <th>SWS</th>\n",
       "      <th>REM</th>\n",
       "      <th>Active wake</th>\n",
       "      <th>Quiet wake</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Rat N</td>\n",
       "      <td>pre</td>\n",
       "      <td>307.2</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3237.1</td>\n",
       "      <td>6714.6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>7947.5</td>\n",
       "      <td>2583.6</td>\n",
       "      <td>2289.5</td>\n",
       "      <td>1563.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Rat S</td>\n",
       "      <td>pre</td>\n",
       "      <td>486.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>965.0</td>\n",
       "      <td>906.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>9012.8</td>\n",
       "      <td>1293.7</td>\n",
       "      <td>2481.4</td>\n",
       "      <td>1681.6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Rat U</td>\n",
       "      <td>pre</td>\n",
       "      <td>1251.9</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4292.8</td>\n",
       "      <td>3987.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>9429.0</td>\n",
       "      <td>1626.7</td>\n",
       "      <td>2461.7</td>\n",
       "      <td>909.7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>Rat V linear</td>\n",
       "      <td>pre</td>\n",
       "      <td>789.9</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4188.6</td>\n",
       "      <td>5157.7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>8059.3</td>\n",
       "      <td>1566.5</td>\n",
       "      <td>2705.4</td>\n",
       "      <td>2068.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>Rat V semicircular</td>\n",
       "      <td>pre</td>\n",
       "      <td>1923.7</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3104.8</td>\n",
       "      <td>5105.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>7563.9</td>\n",
       "      <td>1971.0</td>\n",
       "      <td>2497.9</td>\n",
       "      <td>2348.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>Rat R</td>\n",
       "      <td>pre</td>\n",
       "      <td>5925.7</td>\n",
       "      <td>519.7</td>\n",
       "      <td>513.4</td>\n",
       "      <td>3271.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>5923.8</td>\n",
       "      <td>353.0</td>\n",
       "      <td>1536.5</td>\n",
       "      <td>2977.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>Rat T linear</td>\n",
       "      <td>pre</td>\n",
       "      <td>7277.8</td>\n",
       "      <td>899.4</td>\n",
       "      <td>511.4</td>\n",
       "      <td>2022.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>5612.0</td>\n",
       "      <td>386.0</td>\n",
       "      <td>1006.2</td>\n",
       "      <td>3474.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>Rat T L shape</td>\n",
       "      <td>pre</td>\n",
       "      <td>6356.7</td>\n",
       "      <td>983.9</td>\n",
       "      <td>768.7</td>\n",
       "      <td>2360.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>4205.3</td>\n",
       "      <td>851.6</td>\n",
       "      <td>1616.9</td>\n",
       "      <td>4125.9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>Rat T U shape</td>\n",
       "      <td>pre</td>\n",
       "      <td>6202.9</td>\n",
       "      <td>1316.7</td>\n",
       "      <td>494.7</td>\n",
       "      <td>2785.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>4376.7</td>\n",
       "      <td>383.8</td>\n",
       "      <td>1686.0</td>\n",
       "      <td>4332.7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>Rat K</td>\n",
       "      <td>pre</td>\n",
       "      <td>8526.7</td>\n",
       "      <td>1870.0</td>\n",
       "      <td>204.0</td>\n",
       "      <td>190.7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>3352.4</td>\n",
       "      <td>361.0</td>\n",
       "      <td>2400.5</td>\n",
       "      <td>4677.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>Rat A linear</td>\n",
       "      <td>pre</td>\n",
       "      <td>9832.2</td>\n",
       "      <td>1799.8</td>\n",
       "      <td>3717.7</td>\n",
       "      <td>2706.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>4883.3</td>\n",
       "      <td>838.7</td>\n",
       "      <td>3979.7</td>\n",
       "      <td>4558.2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>Rat A circular</td>\n",
       "      <td>pre</td>\n",
       "      <td>10449.7</td>\n",
       "      <td>1968.5</td>\n",
       "      <td>5004.1</td>\n",
       "      <td>1907.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>5996.3</td>\n",
       "      <td>549.7</td>\n",
       "      <td>5465.7</td>\n",
       "      <td>2595.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>Rat B</td>\n",
       "      <td>pre</td>\n",
       "      <td>5891.7</td>\n",
       "      <td>950.5</td>\n",
       "      <td>3297.1</td>\n",
       "      <td>573.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>2013.2</td>\n",
       "      <td>355.7</td>\n",
       "      <td>4692.0</td>\n",
       "      <td>946.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>Rat C</td>\n",
       "      <td>pre</td>\n",
       "      <td>5032.8</td>\n",
       "      <td>1679.6</td>\n",
       "      <td>4207.8</td>\n",
       "      <td>5079.9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>5126.8</td>\n",
       "      <td>1291.1</td>\n",
       "      <td>4867.7</td>\n",
       "      <td>3115.8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>Rat G</td>\n",
       "      <td>pre</td>\n",
       "      <td>9022.5</td>\n",
       "      <td>1504.7</td>\n",
       "      <td>3903.7</td>\n",
       "      <td>2328.4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td></td>\n",
       "      <td>post</td>\n",
       "      <td>6854.7</td>\n",
       "      <td>412.9</td>\n",
       "      <td>5666.0</td>\n",
       "      <td>1475.8</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               Session epoch      SWS     REM Active wake Quiet wake\n",
       "0                Rat N   pre    307.2     0.0      3237.1     6714.6\n",
       "1                       post   7947.5  2583.6      2289.5     1563.4\n",
       "2                Rat S   pre    486.0     0.0       965.0      906.2\n",
       "3                       post   9012.8  1293.7      2481.4     1681.6\n",
       "4                Rat U   pre   1251.9     0.0      4292.8     3987.3\n",
       "5                       post   9429.0  1626.7      2461.7      909.7\n",
       "6         Rat V linear   pre    789.9     0.0      4188.6     5157.7\n",
       "7                       post   8059.3  1566.5      2705.4     2068.1\n",
       "8   Rat V semicircular   pre   1923.7     0.0      3104.8     5105.5\n",
       "9                       post   7563.9  1971.0      2497.9     2348.4\n",
       "10               Rat R   pre   5925.7   519.7       513.4     3271.5\n",
       "11                      post   5923.8   353.0      1536.5     2977.5\n",
       "12        Rat T linear   pre   7277.8   899.4       511.4     2022.2\n",
       "13                      post   5612.0   386.0      1006.2     3474.5\n",
       "14       Rat T L shape   pre   6356.7   983.9       768.7     2360.2\n",
       "15                      post   4205.3   851.6      1616.9     4125.9\n",
       "16       Rat T U shape   pre   6202.9  1316.7       494.7     2785.4\n",
       "17                      post   4376.7   383.8      1686.0     4332.7\n",
       "18               Rat K   pre   8526.7  1870.0       204.0      190.7\n",
       "19                      post   3352.4   361.0      2400.5     4677.5\n",
       "20        Rat A linear   pre   9832.2  1799.8      3717.7     2706.2\n",
       "21                      post   4883.3   838.7      3979.7     4558.2\n",
       "22      Rat A circular   pre  10449.7  1968.5      5004.1     1907.3\n",
       "23                      post   5996.3   549.7      5465.7     2595.0\n",
       "24               Rat B   pre   5891.7   950.5      3297.1      573.4\n",
       "25                      post   2013.2   355.7      4692.0      946.0\n",
       "26               Rat C   pre   5032.8  1679.6      4207.8     5079.9\n",
       "27                      post   5126.8  1291.1      4867.7     3115.8\n",
       "28               Rat G   pre   9022.5  1504.7      3903.7     2328.4\n",
       "29                      post   6854.7   412.9      5666.0     1475.8"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "duration_per_brain_state_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "    too_short_state_bout_index = np.where(bout_durations < 6)\n",
    "    bout_labels[too_short_state_bout_index] = 5\n",
    "\n",
    "\n",
    "\n",
    "    unique_state_labels = np.unique(bout_labels)\n",
    "    num_bouts = bouts_start_end.shape[0]\n",
    "\n",
    "    time_bins_centers = np.arange(0, bouts_start_end[-1,1]+1, 1)+0.5    \n",
    "    time_bins_brain_state = np.full(time_bins_centers.shape, np.nan)\n",
    "\n",
    "    for bout_idx, bout_timing in enumerate(bouts_start_end):\n",
    "        inside_bout_idx = np.logical_and(bout_timing[0] < time_bins_centers,  time_bins_centers < bout_timing[1]) \n",
    "        time_bins_brain_state[inside_bout_idx] = bout_labels[bout_idx]\n",
    "\n",
    "\n",
    "    chunk_size_in_hours = 1/60 # one minutes\n",
    "    chunk_size = int(chunk_size_in_hours*3600) # in seconds\n",
    "    num_chunks = len(time_bins_brain_state) // chunk_size\n",
    "    chunk_start_times = [i*chunk_size_in_hours for i in range(num_chunks)] # centers of the time chunks\n",
    "    \n",
    "    reshaped_time_bins_brain_state = time_bins_brain_state[:num_chunks*chunk_size].reshape(num_chunks, chunk_size)\n",
    "\n",
    "    bins = np.arange(1, unique_state_labels.shape[0]+2)\n",
    "    counts = np.apply_along_axis(lambda x:np.histogram(x, bins = bins)[0] , axis = 1, arr=reshaped_time_bins_brain_state)\n",
    "\n",
    "\n",
    "    # Normalize the counts to get percentages\n",
    "    percentages = (counts / counts.sum(axis=1, keepdims=True))*100\n",
    "\n",
    "    # Create a DataFrame with percentages for each category and each chunk\n",
    "    brain_state_df = pd.DataFrame(percentages, columns=brainStates_names)\n",
    "\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 = scipy.io.loadmat(file_path)\n",
    "    spikes_pyr = mat_file[\"spikes_pyr\"]\n",
    "\n",
    "\n",
    "    # Load unit stability information\n",
    "\n",
    "    filename = f'{session_name}.cluster_quality_by_block'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    mat_file = scipy.io.loadmat(file_path)\n",
    "    cluster_quality_by_block = mat_file['cluster_quality_by_block'][0]\n",
    "\n",
    "\n",
    "\n",
    "    #### Extracting all place fields from the imported .mat file\n",
    "    spatial_tuning_smoothed = spikes_pyr[\"spatialTuning_smoothed\"]\n",
    "\n",
    "\n",
    "    if session_number in [6,7]:\n",
    "        num_units = spatial_tuning_smoothed.shape[0] # for RatN and RatS\n",
    "    else:\n",
    "        num_units = spatial_tuning_smoothed[0].shape[0]\n",
    "        \n",
    "    num_pos_bins = spatial_tuning_smoothed[0][0]['uni'][0][0].size\n",
    "\n",
    "    # print(num_units, num_pos_bins)\n",
    "\n",
    "    spikes = []; # spike data and place field info of each unit\n",
    "\n",
    "    # attributes = list(spikes_pyr.dtype.names) % if we want to work on all variable in the imported .mat data structure\n",
    "    running_directions = {'LR', 'RL', 'uni'}\n",
    "    other_attributes   = {'spike_times', 'shank_id','cluster_id'}\n",
    "\n",
    "    iter = 0\n",
    "    for unit in range(num_units):\n",
    "        \n",
    "        # Create dictionaries for each unit and store the matrices\n",
    "        \n",
    "        unit_spikes = dict()\n",
    "        \n",
    "        unit_spikes['place_fields']  = {}\n",
    "        unit_spikes['peak_pos_bins'] = {}\n",
    "\n",
    "        \n",
    "        for direction in running_directions:\n",
    "            try:\n",
    "                if session_number in [6,7]:\n",
    "                    unit_spikes['place_fields'][direction] = spatial_tuning_smoothed[unit][0][direction][0][0].reshape(num_pos_bins) \n",
    "                    unit_spikes['peak_pos_bins'][direction] = spikes_pyr['peakPosBin'][unit][0][direction][0][0][0][0]\n",
    "                else:\n",
    "                    unit_spikes['place_fields'][direction] = spatial_tuning_smoothed[0][unit][direction][0][0].reshape(num_pos_bins) \n",
    "                    unit_spikes['peak_pos_bins'][direction] = spikes_pyr['peakPosBin'][0][unit][direction][0][0][0][0]\n",
    "\n",
    "\n",
    "            except ValueError:\n",
    "                if iter == 0:\n",
    "                    print(\"This session has only one running direction\")\n",
    "                iter += 1\n",
    "\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['shank_id']    += 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]\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['shank_id']    += 1\n",
    "            unit_spikes['cluster_id']  = spikes_pyr['id'][0][unit][0][1]\n",
    "        else: # Grosmark, Hiro, and all other sessions\n",
    "            unit_spikes['spike_times'] = spikes_pyr['time'][0][unit]\n",
    "            unit_spikes['shank_id']    = spikes_pyr['id'][0][unit][0][0] # need to go one down for the other datasets\n",
    "            unit_spikes['cluster_id']  = spikes_pyr['id'][0][unit][0][1]\n",
    "\n",
    "\n",
    "\n",
    "        spikes.append(unit_spikes) \n",
    "\n",
    "\n",
    "\n",
    "    # place fields by pooling spikes across both running directions\n",
    "    place_fields_uni = []\n",
    "    for unit in range(num_units):\n",
    "        place_fields_uni.append(spikes[unit]['place_fields']['uni'])\n",
    "    place_fields_uni = np.array(place_fields_uni)\n",
    "\n",
    "    place_fields_uni[place_fields_uni == 0] = 1e-4\n",
    "    active_units = np.where(np.nanmax(place_fields_uni, axis=1) > 1)[0]\n",
    "\n",
    "\n",
    "\n",
    "    # --------------------------------------------------------------------------------------------\n",
    "    # Load cluster quality data (L-ratios)\n",
    "    \n",
    "    filename = f'{session_name}.clusterQuality.mat'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    mat_file = scipy.io.loadmat(file_path)\n",
    "\n",
    "    # Access data structure\n",
    "    num_shanks = len(mat_file[\"clusterQuality\"][\"Lratio\"][0])\n",
    "\n",
    "    L_ratios = list()\n",
    "    for shank in range(num_shanks):    \n",
    "        curr_shank_L_ratios = dict()\n",
    "        curr_shank_L_ratios[\"L_ratios\"] = mat_file[\"clusterQuality\"][\"Lratio\"][0][shank]\n",
    "        curr_shank_L_ratios[\"cluster_ids\"] = mat_file[\"clusterQuality\"][\"clus\"][0][shank]\n",
    "        \n",
    "        L_ratios.append(curr_shank_L_ratios)\n",
    "\n",
    "\n",
    "\n",
    "    #-------------------------------------------------------------------------------------------\n",
    "    # Loading the population burst evenst with all their corresponding measured lfp features\n",
    "\n",
    "    overwrite = False # in case we need to read the .mat file again, if there was a change\n",
    "\n",
    "    filename = f'{session_name}.PBEs.pkl'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    if os.path.exists(file_path) and overwrite == False:\n",
    "\n",
    "        # PBEs = np.load(file_path, allow_pickle=True)\n",
    "        PBEs = pd.read_pickle(file_path)\n",
    "    else: # if it doesn't exist then read it from the .mat file\n",
    "\n",
    "        filename = f'{session_name}.PBEInfo_replayScores_with_spindle_and_deltaPowers.mat'\n",
    "        file_path = os.path.join(session_dataset_path, filename)\n",
    "        f = h5py.File(file_path, \"r\")\n",
    "        PBEInfo = f['PBEInfo_replayScores']    \n",
    "       \n",
    "\n",
    "        # Store the population burst events in a pandas DataFrame\n",
    "        num_PBEs = PBEInfo[\"fr_1msbin\"].shape[0]\n",
    "        attributes = list(PBEInfo.keys())\n",
    "\n",
    "        PBEs = pd.DataFrame(columns=[attr for attr in attributes[1:] if attr not in ['posteriorProbMat', 'postMat_nonNorm']])\n",
    "\n",
    "        # Loop over the PBEs and add each one as a row to the DataFrame\n",
    "        num_dots = int(num_PBEs * (10/100))\n",
    "        count = 0\n",
    "\n",
    "        for pbe in range(num_PBEs): #  \n",
    "  \n",
    "            for attr in PBEs.columns:\n",
    "                ref = PBEInfo[attr][pbe][0]\n",
    "                obj = f[ref]\n",
    "\n",
    "                if attr in ['epoch', 'brainState']: # convert the ascii code to string\n",
    "                    arr = np.array(obj).flatten()\n",
    "                    epoch = \"\".join(chr(code) for code in arr)\n",
    "                    PBEs.at[pbe, attr] = epoch\n",
    "                elif attr in ['fr_1msbin', 'fr_20msbin', 'posteriorProbMat', 'postMat_nonNorm']: # no need to flatten\n",
    "                    arr = np.array(obj)\n",
    "                    PBEs.at[pbe, attr] = arr\n",
    "                else: \n",
    "                    arr = np.array(obj).flatten()\n",
    "                    PBEs.at[pbe, attr] = arr\n",
    "\n",
    "        if (pbe+1) % num_dots == 1:\n",
    "            count += 1\n",
    "            message = \"Importing PBEs\" + \".\" * count\n",
    "            print(message, end=\"\\r\")\n",
    "\n",
    "        print(\"All PBEs were imported\") \n",
    "    \n",
    "        filename = f'{session_name}.PBEs.pkl'\n",
    "        file_path = os.path.join(session_dataset_path, filename)\n",
    "        PBEs.to_pickle(file_path)\n",
    "\n",
    "    num_PBEs = PBEs.shape[0]\n",
    "\n",
    "    # add a brain state label depending on the time bin of the PBE\n",
    "    for pbe in range(num_PBEs):\n",
    "        idx = np.where(np.logical_and(time_bins_centers-0.5 <= PBEs.at[pbe, 'peakT'],  time_bins_centers+0.5 > PBEs.at[pbe, 'peakT']))[0]\n",
    "        if len(idx) == 1:\n",
    "                PBEs.at[pbe, 'brain_state'] = time_bins_brain_state[idx]\n",
    "        else:\n",
    "            PBEs.at[pbe, 'brain_state'] = ''\n",
    "\n",
    "\n",
    "\n",
    "    # # Adding SWA that Bapun calculated - only for RatU\n",
    "    # filename = f'{session_name}.SWA_activity_for_Kourosh.npy'\n",
    "    # file_path = os.path.join(session_dataset_path, filename)\n",
    "    # swa_arr = np.load(file_path, allow_pickle=True).item()\n",
    "\n",
    "    # swa = swa_arr['swa']\n",
    "    # swa_t_start = swa_arr['t_start']\n",
    "    # # swa_window_length = swa_arr['window']\n",
    "    # # swa_windows_overlap = swa_arr['overlap']\n",
    "    # swa_time_pnts = np.arange(swa_t_start, swa.shape[0]+.1, 1)\n",
    "    \n",
    "    # PBEs['swa_bapun'] = np.interp(PBEs['peakT'], swa_time_pnts, swa)\n",
    "\n",
    "\n",
    "    # # Adding delta and spindle that I calculated using the multitaper method using a channel that showed the highest ripple power\n",
    "    # filename = f'{session_name}.delta_spindle_multitaper.mat'\n",
    "    # file_path = os.path.join(session_dataset_path, filename)\n",
    "\n",
    "    # mat_file = scipy.io.loadmat(file_path)\n",
    "    # tPnts = mat_file['t'].flatten()\n",
    "    # spindle_power_multitaper = mat_file['spindle_power'].flatten()\n",
    "    # delta_power_multitaper = mat_file['delta_power'].flatten()\n",
    "\n",
    "\n",
    "    # PBEs['delta_power_multitaper'] = np.interp(PBEs['peakT'], tPnts, delta_power_multitaper)\n",
    "    # PBEs['spindle_power_multitaper'] = np.interp(PBEs['peakT'], tPnts, spindle_power_multitaper)\n",
    "\n",
    "\n",
    "    # Adding delta and spindle that I calculated using spectrogram and without whitening\n",
    "\n",
    "    # ## using the channel that showed the highest bipolarity in slow wave slope - using MATLAB Buzcode\n",
    "    # filename = f'{session_name}.delta_spindle_spectrogram_tukey_prior_zscoring_best_sw_slope_channel.mat'\n",
    "    # file_path = os.path.join(session_state_detection_results_path, filename)\n",
    "\n",
    "    # mat_file = scipy.io.loadmat(file_path)\n",
    "    # tPnts = mat_file['t'].flatten()\n",
    "    # delta_power = mat_file['delta_power'].flatten()\n",
    "    # spindle_power = mat_file['spindle_power'].flatten()\n",
    "\n",
    "    # PBEs['delta_sw_slope_chan'] = np.interp(PBEs['peakT'], tPnts, delta_power)\n",
    "    # PBEs['spindle_sw_slope_chan'] = np.interp(PBEs['peakT'], tPnts, spindle_power)\n",
    "\n",
    "\n",
    "    # Using the channel with the highest ripple power\n",
    "    filename = f'{session_name}.delta_spindle_spectrogram_tukey_prior_zscoring_best_ripple_channel_no_overlap_windows.mat'\n",
    "    file_path = os.path.join(session_state_detection_results_path, filename)\n",
    "\n",
    "    mat_file = scipy.io.loadmat(file_path)\n",
    "\n",
    "    tPnts_delta = mat_file['t_delta'].flatten()\n",
    "    delta_power = mat_file['delta_power'].flatten()\n",
    "    PBEs['delta_ripple_chan'] = np.interp(PBEs['peakT'], tPnts_delta, delta_power)\n",
    "\n",
    "\n",
    "    tPnts_spindle = mat_file['t_spindle'].flatten()\n",
    "    spindle_power = mat_file['spindle_power'].flatten()\n",
    "    PBEs['spindle_ripple_chan'] = np.interp(PBEs['peakT'], tPnts_spindle, spindle_power)\n",
    "\n",
    "\n",
    "    # low versus high spindle within each 2s window used to measure delta\n",
    "    spindle_p_score_in_each_delta_window = np.full((len(tPnts_spindle), ), np.nan)\n",
    "    for delta_t in tPnts_delta:\n",
    "        # current 500 ms time bins within the 2s window\n",
    "        idx =  (tPnts_spindle > (delta_t-1)) & (tPnts_spindle < (delta_t+1))\n",
    "        spindle_p_score_in_each_delta_window[idx] = p_score(spindle_power[idx])\n",
    "\n",
    "    PBEs['spindle_prc_in_delta_window'] = np.interp(PBEs['peakT'], tPnts_spindle, spindle_p_score_in_each_delta_window)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    # ## using the channel with highest ripple power\n",
    "    # filename = f'{session_name}.delta_spindle_spectrogram_tukey_prior_zscoring_best_ripple_channel.mat'\n",
    "    # file_path = os.path.join(session_state_detection_results_path, filename)\n",
    "\n",
    "    # mat_file = scipy.io.loadmat(file_path)\n",
    "    # tPnts_delta = mat_file['t'].flatten()\n",
    "    # delta_power = mat_file['delta_power'].flatten()\n",
    "    # # spindle_power = mat_file['spindle_power'].flatten()\n",
    "\n",
    "    # PBEs['delta_ripple_chan'] = np.interp(PBEs['peakT'], tPnts_delta, delta_power)\n",
    "    # # PBEs['spindle_ripple_chan'] = np.interp(PBEs['peakT'], tPnts, spindle_power)\n",
    "\n",
    "    # # Spindle calculated using 500 ms windows with 250 ms overlaps\n",
    "    # filename = f'{session_name}.spindle_spectrogram_tukey_prior_zscoring_best_ripple_channel.mat'\n",
    "    # file_path = os.path.join(session_state_detection_results_path, filename)\n",
    "\n",
    "    # mat_file = scipy.io.loadmat(file_path)\n",
    "    # tPnts_spindles = mat_file['t'].flatten()\n",
    "    # spindle_power = mat_file['spindle_power'].flatten()\n",
    "\n",
    "    # PBEs['spindle_ripple_chan'] = np.interp(PBEs['peakT'], tPnts_spindles, spindle_power)\n",
    "\n",
    "\n",
    "    #------------------------------------------------------------------------------------------------------\n",
    "    # Learned tunings calculated based on NREM PBEs with values of certain lfp features like ripple power, slow wave amplitude, and spindle power within a low or high range (along the median)\n",
    "\n",
    "    num_PBEs = PBEs.shape[0]\n",
    "    brain_state_first_letter = np.empty((num_PBEs,), dtype='str')\n",
    "\n",
    "    epoch_names = ['pre', 'post']\n",
    "    # lfp_features = ['swa_bapun']\n",
    "    # lfp_features = ['SWA', 'spindleP', 'peakRippleA', 'peakMUA', 'nFiringUnits', 'duration']\n",
    "    # lfp_features = ['delta_sw_slope_chan', 'spindle_sw_slope_chan', 'delta_ripple_chan', 'spindle_ripple_chan'] # 'delta_power_multitaper', 'spindle_power_multitaper', \n",
    "    lfp_features = ['delta_ripple_chan', 'spindle_ripple_chan', 'spindle_prc_in_delta_window']\n",
    "    # lfp_feature_combinations = ['spindle_ripple_chan_high_delta', 'spindle_ripple_chan_low_delta', 'delta_ripple_chan_high_delta', 'delta_ripple_chan_low_delta']\n",
    "    # lfp_features.extend(lfp_feature_combinations)\n",
    "    ripple_categories = ['NREM', 'QW']\n",
    "    lfp_features.extend(ripple_categories)\n",
    "\n",
    "\n",
    "    lfp_feature_threshold = {epoch:{lfp_feature:None for lfp_feature in lfp_features} for epoch in epoch_names} \n",
    "    epoch_PBEs_lfp_feature_value = {epoch:{lfp_feature:None for lfp_feature in lfp_features} for epoch in epoch_names} # The distribution of a specific lfp feature within each epoch\n",
    "\n",
    "    PBEs_indices_low_vs_high_lfp_feature = {epoch: {lfp_feature: {'low':None, 'high': None} for lfp_feature in lfp_features} for epoch in epoch_names}\n",
    "    num_PBEs_low_vs_high_lfp_feature = {epoch: {lfp_feature: {'low':None, 'high': None} for lfp_feature in lfp_features} for epoch in epoch_names}\n",
    "    learned_tunings_low_vs_high_lfp_feature_matched_particpation = {epoch: {lfp_feature: {'low':None, 'high': None} for lfp_feature in lfp_features} for epoch in epoch_names}\n",
    "\n",
    "\n",
    "\n",
    "    time_bin_duration = 0.02\n",
    "    num_PF_shuffles = 10000\n",
    "\n",
    "    for epoch in epoch_names:\n",
    "        if epoch == 'pre':\n",
    "            epoch_idx = 0\n",
    "            epoch_duration = epochs[0,1] - epochs[0,0]\n",
    "        elif epoch == 'post':\n",
    "            epoch_idx = 2\n",
    "            epoch_duration = 4*60*60 # limiting to the first 4 hours of post sleep\n",
    "\n",
    "        if_inside_epoch = PBEs['peakT'].between(epochs[epoch_idx,0], epochs[epoch_idx,0]+epoch_duration)\n",
    "        if_PBEs_is_NREM_and_within_epoch = if_inside_epoch & (PBEs['brain_state'] == 1) # within the epoch and only includes NREM PBEs\n",
    "\n",
    "        for lfp_feature_index, lfp_feature in enumerate(lfp_features[:3]):\n",
    "            \n",
    "            # calculate median value of each lfp feature across all NREM PBEs ocurring within the epoch\n",
    "            epoch_PBEs_lfp_feature_value[epoch][lfp_feature] = PBEs[lfp_feature][if_PBEs_is_NREM_and_within_epoch]\n",
    "\n",
    "            if lfp_feature_index == 2:\n",
    "                lfp_feature_threshold[epoch][lfp_feature] = 50\n",
    "            else:\n",
    "                lfp_feature_threshold[epoch][lfp_feature] = np.nanmedian(epoch_PBEs_lfp_feature_value[epoch][lfp_feature])\n",
    "\n",
    "            # Indices of PBEs ocurring during each epoch and during high or low values of lfp-feature\n",
    "            low_indices = np.where(if_PBEs_is_NREM_and_within_epoch & (PBEs[lfp_feature] < lfp_feature_threshold[epoch][lfp_feature]))[0]\n",
    "            high_indices = np.where(if_PBEs_is_NREM_and_within_epoch & (PBEs[lfp_feature] >= lfp_feature_threshold[epoch][lfp_feature]))[0]\n",
    "\n",
    "            PBEs_indices_low_vs_high_lfp_feature[epoch][lfp_feature]['low'] = low_indices\n",
    "            num_PBEs_low_vs_high_lfp_feature[epoch][lfp_feature]['low'] = low_indices.shape[0]\n",
    "\n",
    "            PBEs_indices_low_vs_high_lfp_feature[epoch][lfp_feature]['high'] = high_indices\n",
    "            num_PBEs_low_vs_high_lfp_feature[epoch][lfp_feature]['high'] = high_indices.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "    # Maybe PRE vs POST for NREM and QW can be calculated here too?\n",
    "\n",
    "    # Set PRE to low, POST to high\n",
    "    # Set lfp_feature to NREM, and QW\n",
    "\n",
    "    # set epoch to something like a placeholder like 'placeholder'\n",
    "\n",
    "    for epoch in epoch_names:\n",
    "        \n",
    "        ## PRE\n",
    "        epoch_idx = 0\n",
    "        epoch_duration = epochs[0,1] - epochs[0,0]\n",
    "        if_inside_epoch = PBEs['peakT'].between(epochs[epoch_idx,0], epochs[epoch_idx,0]+epoch_duration)\n",
    "\n",
    "        # NREM\n",
    "        if_PBEs_is_NREM_and_within_epoch = if_inside_epoch & (PBEs['brain_state'] == 1)\n",
    "        low_indices =  np.where(if_PBEs_is_NREM_and_within_epoch)[0]\n",
    "        \n",
    "        PBEs_indices_low_vs_high_lfp_feature[epoch]['NREM']['low'] = low_indices\n",
    "        num_PBEs_low_vs_high_lfp_feature[epoch]['NREM']['low'] = low_indices.shape[0]\n",
    "\n",
    "        # QW\n",
    "        if_PBEs_is_QW_and_within_epoch = if_inside_epoch & (PBEs['brain_state'] == 4)\n",
    "        low_indices =  np.where(if_PBEs_is_QW_and_within_epoch)[0]\n",
    "        \n",
    "        PBEs_indices_low_vs_high_lfp_feature[epoch]['QW']['low'] = low_indices\n",
    "        num_PBEs_low_vs_high_lfp_feature[epoch]['QW']['low'] = low_indices.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "        ## POST\n",
    "        epoch_idx = 2\n",
    "        epoch_duration = 4*60*60\n",
    "        if_inside_epoch = PBEs['peakT'].between(epochs[epoch_idx,0], epochs[epoch_idx,0]+epoch_duration)\n",
    "\n",
    "        # NREM\n",
    "        if_PBEs_is_NREM_and_within_epoch = if_inside_epoch & (PBEs['brain_state'] == 1)\n",
    "        high_indices =  np.where(if_PBEs_is_NREM_and_within_epoch)[0]\n",
    "        \n",
    "        PBEs_indices_low_vs_high_lfp_feature[epoch]['NREM']['high'] = high_indices\n",
    "        num_PBEs_low_vs_high_lfp_feature[epoch]['NREM']['high'] = high_indices.shape[0]\n",
    "\n",
    "\n",
    "        # QW\n",
    "        if_PBEs_is_QW_and_within_epoch = if_inside_epoch & (PBEs['brain_state'] == 4)\n",
    "        high_indices =  np.where(if_PBEs_is_QW_and_within_epoch)[0]\n",
    "        \n",
    "        PBEs_indices_low_vs_high_lfp_feature[epoch]['QW']['high'] = high_indices\n",
    "        num_PBEs_low_vs_high_lfp_feature[epoch]['QW']['high'] = high_indices.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        # # Indices of PBEs ocurring during high or low delta power and high or low spindle power\n",
    "        # for lfp_feature_combination in lfp_feature_combinations[:2]:\n",
    "\n",
    "        #     if lfp_feature_combination == 'spindle_ripple_chan_high_delta':\n",
    "        #         current_delta_level_indices = if_PBEs_is_NREM_and_within_epoch & (PBEs['delta_ripple_chan'] > lfp_feature_threshold[epoch]['delta_ripple_chan'])\n",
    "        #     elif lfp_feature_combination == 'spindle_ripple_chan_low_delta':\n",
    "        #         current_delta_level_indices = if_PBEs_is_NREM_and_within_epoch & (PBEs['delta_ripple_chan'] < lfp_feature_threshold[epoch]['delta_ripple_chan'])\n",
    "\n",
    "        #     epoch_PBEs_lfp_feature_value[epoch][lfp_feature_combination] = PBEs['spindle_ripple_chan'][current_delta_level_indices] \n",
    "        #     lfp_feature_threshold[epoch][lfp_feature_combination] = np.nanmedian(epoch_PBEs_lfp_feature_value[epoch][lfp_feature_combination])         \n",
    "\n",
    "        #     low_indices = np.where(current_delta_level_indices & (PBEs['spindle_ripple_chan'] < lfp_feature_threshold[epoch][lfp_feature_combination]))[0]\n",
    "        #     high_indices = np.where(current_delta_level_indices & (PBEs['spindle_ripple_chan'] > lfp_feature_threshold[epoch][lfp_feature_combination]))[0]\n",
    "\n",
    "        #     PBEs_indices_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['low'] = low_indices\n",
    "        #     num_PBEs_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['low'] = low_indices.shape[0]\n",
    "\n",
    "        #     PBEs_indices_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['high'] = high_indices\n",
    "        #     num_PBEs_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['high'] = high_indices.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "        # # Further splitting based on delta within each low or high delta categories\n",
    "        # for lfp_feature_combination in lfp_feature_combinations[2:]:\n",
    "\n",
    "        #     if lfp_feature_combination == 'delta_ripple_chan_high_delta':\n",
    "        #         current_delta_level_indices = if_PBEs_is_NREM_and_within_epoch & (PBEs['delta_ripple_chan'] > lfp_feature_threshold[epoch]['delta_ripple_chan'])\n",
    "        #     elif lfp_feature_combination == 'delta_ripple_chan_low_delta':\n",
    "        #         current_delta_level_indices = if_PBEs_is_NREM_and_within_epoch & (PBEs['delta_ripple_chan'] < lfp_feature_threshold[epoch]['delta_ripple_chan'])\n",
    "\n",
    "        #     epoch_PBEs_lfp_feature_value[epoch][lfp_feature_combination] = PBEs['delta_ripple_chan'][current_delta_level_indices] \n",
    "        #     lfp_feature_threshold[epoch][lfp_feature_combination] = np.nanmedian(epoch_PBEs_lfp_feature_value[epoch][lfp_feature_combination])         \n",
    "\n",
    "        #     low_indices = np.where(current_delta_level_indices & (PBEs['delta_ripple_chan'] < lfp_feature_threshold[epoch][lfp_feature_combination]))[0]\n",
    "        #     high_indices = np.where(current_delta_level_indices & (PBEs['delta_ripple_chan'] > lfp_feature_threshold[epoch][lfp_feature_combination]))[0]\n",
    "\n",
    "        #     PBEs_indices_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['low'] = low_indices\n",
    "        #     num_PBEs_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['low'] = low_indices.shape[0]\n",
    "\n",
    "        #     PBEs_indices_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['high'] = high_indices\n",
    "        #     num_PBEs_low_vs_high_lfp_feature[epoch][lfp_feature_combination]['high'] = high_indices.shape[0]\n",
    "\n",
    "\n",
    "\n",
    "    # calculate learned tunings within each epoch *using all offline PBEs*, regardless of their corresponding brain state   \n",
    "    learned_tunings_low_vs_high_lfp_feature_matched_particpation = calculate_learned_tuning_matched_participation(PBEs, spikes, PBEs_indices_low_vs_high_lfp_feature, L_ratios, time_bin_duration)\n",
    "\n",
    "    filename = f'{session_name}.learned_tunings_low_versus_high_lfp_feature_matched_particpation_NREM_delta_spindle_spectrogram_500ms_windows_for_spindles_brain_state_corrected.npy'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "    np.save(file_path, learned_tunings_low_vs_high_lfp_feature_matched_particpation)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "    # # Plot the distributions of LT fidelities for LTs calculated separately based on low or high level of lfp features values\n",
    "\n",
    "    # for current_lfp_feature in lfp_features:\n",
    "\n",
    "    #     # Plot the distributions of learned tunings for low and high SWA within each sleep epoch\n",
    "    #     colors = sns.color_palette(\"husl\", 2) # Set color palette\n",
    "    #     sns.set_style('whitegrid') # Set style and context\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",
    "    #     fig, axes = plt.subplots(1,2, figsize = (10, 4))\n",
    "\n",
    "    #     for i, epoch in enumerate(epoch_names):\n",
    "            \n",
    "    #         learned_tunings_pearson_corr_NREM_low_lfp_feature_value, _, median_LT_PF_pearson_corr_NREM_low_lfp_feature_value = calculate_place_field_fidelity_of_learned_tuning(learned_tunings_low_vs_high_lfp_feature_matched_particpation[epoch][current_lfp_feature]['low'][active_units], place_fields_uni[active_units], num_PF_shuffles)\n",
    "    #         learned_tunings_pearson_corr_NREM_high_lfp_feature_value, _, median_LT_PF_pearson_corr_NREM_high_lfp_feature_value = calculate_place_field_fidelity_of_learned_tuning(learned_tunings_low_vs_high_lfp_feature_matched_particpation[epoch][current_lfp_feature]['high'][active_units], place_fields_uni[active_units], num_PF_shuffles)\n",
    "\n",
    "    #         # Plot the distributions using ECDFs with overlaid ticks\n",
    "    #         sns.ecdfplot(learned_tunings_pearson_corr_NREM_low_lfp_feature_value, ax = axes[i], label=f'{epoch} Low {current_lfp_feature}', color=colors[0], linewidth = 2)\n",
    "    #         sns.ecdfplot(learned_tunings_pearson_corr_NREM_high_lfp_feature_value, ax = axes[i], label=f'{epoch} high {current_lfp_feature}', color=colors[1], linewidth = 2)\n",
    "\n",
    "    #         axes[i].set_xlim([-1,1])\n",
    "    #         axes[i].set_xlabel(f'{epoch} LT fidelity', fontsize=10)\n",
    "    #         axes[i].set_ylabel('Proportion of units', fontsize=10)\n",
    "    #         axes[i].tick_params(labelsize=8)\n",
    "    #         axes[i].legend(fontsize=8)\n",
    "\n",
    "\n",
    "    #         # Add ticks to the x-axis\n",
    "    #         y_min, y_max = axes[i].get_ylim()\n",
    "    #         for x in learned_tunings_pearson_corr_NREM_low_lfp_feature_value:\n",
    "    #             axes[i].plot([x, x], [y_min+0.01, y_min + 0.02], '|-', color=colors[0], linewidth=0.25)\n",
    "    #         for x in learned_tunings_pearson_corr_NREM_high_lfp_feature_value:\n",
    "    #             axes[i].plot([x, x], [y_min+0.06, y_min + 0.07], '|-', color=colors[1], linewidth=0.25)\n",
    "            \n",
    "\n",
    "    #         # Calculate and display medians\n",
    "    #         median_corr_low_lfp_feature = median_LT_PF_pearson_corr_NREM_low_lfp_feature_value['data']\n",
    "    #         median_corr_low_lfp_feature_pvalue = median_LT_PF_pearson_corr_NREM_low_lfp_feature_value['p_value']\n",
    "\n",
    "    #         median_corr_high_lfp_feature = median_LT_PF_pearson_corr_NREM_high_lfp_feature_value['data']\n",
    "    #         median_corr_high_lfp_feature_pvalue = median_LT_PF_pearson_corr_NREM_high_lfp_feature_value['p_value']\n",
    "\n",
    "    #         def get_pval_statement(pvalue):\n",
    "    #             if pvalue < 0.0001:\n",
    "    #                 pvalue_statement = 'P<1e-4'\n",
    "    #             else:\n",
    "    #                 pvalue_statement = f'P={pvalue:.4f}'\n",
    "    #             return pvalue_statement\n",
    "\n",
    "    #         axes[i].axvline(median_corr_low_lfp_feature, color=colors[0], linestyle='dashed', label=f'Median {epoch} low {current_lfp_feature}={median_corr_low_lfp_feature:.2f},{get_pval_statement(median_corr_low_lfp_feature_pvalue)}')\n",
    "    #         axes[i].axvline(median_corr_high_lfp_feature, color=colors[1], linestyle='dashed', label=f'Median {epoch} high {current_lfp_feature}={median_corr_high_lfp_feature:.2f},{get_pval_statement(median_corr_high_lfp_feature_pvalue)}')\n",
    "            \n",
    "    #         axes[i].legend(fontsize=6)\n",
    "\n",
    "    #         # Perform rank-sum test\n",
    "    #         learned_tunings_pearson_corr_NREM_low_lfp_feature_value = learned_tunings_pearson_corr_NREM_low_lfp_feature_value[~np.isnan(learned_tunings_pearson_corr_NREM_low_lfp_feature_value)]\n",
    "    #         learned_tunings_pearson_corr_NREM_high_lfp_feature_value = learned_tunings_pearson_corr_NREM_high_lfp_feature_value[~np.isnan(learned_tunings_pearson_corr_NREM_high_lfp_feature_value)]\n",
    "    #         statistic, p_value = ranksums(learned_tunings_pearson_corr_NREM_low_lfp_feature_value,\n",
    "    #                                     learned_tunings_pearson_corr_NREM_high_lfp_feature_value)\n",
    "\n",
    "\n",
    "    #         # Add line with p-value above the plot\n",
    "    #         axes[i].text((median_corr_low_lfp_feature + median_corr_high_lfp_feature) / 2, 0.2, f'{get_pval_statement(p_value)}', ha='center', fontsize=6)\n",
    "\n",
    "\n",
    "\n",
    "    #         # plot the histogram of PBEs' coincident slow wave amplitude\n",
    "    #         ax_inset = axes[i].inset_axes([0.2, 0.4, 0.3, 0.2])\n",
    "\n",
    "    #         curr_feature_values = epoch_PBEs_lfp_feature_value[epoch][current_lfp_feature]\n",
    "    #         curr_feature_threshold = lfp_feature_threshold[epoch][current_lfp_feature]\n",
    "\n",
    "    #         sns.histplot(curr_feature_values, bins=30, cumulative=False, kde=True, stat='density', color='black', edgecolor='black', ax = ax_inset)\n",
    "    #         ax_inset.axvline(x=curr_feature_threshold, color='black', linestyle = '--')\n",
    "    #         ylim_inset = ax_inset.get_ylim()\n",
    "    #         ax_inset.text(curr_feature_threshold, ylim_inset[1], f'median = {curr_feature_threshold:.2f}', ha='center', fontsize=6)\n",
    "\n",
    "    #         # ax_inset.axhline(y=num_PBEs_NREM_low_vs_high_swa[epoch]['low_swa'], color='black', linestyle = '--')\n",
    "\n",
    "    #         ax_inset.set_title(f'NREM PBEs {current_lfp_feature}', fontsize=6)\n",
    "    #         # ax_inset.set_xlim([0, np.max(epoch_PBEs_lfp_feature_value[epoch][current_lfp_feature])])\n",
    "\n",
    "    #         # ax_inset.set_xticks(range(0, int(np.ceil(np.max(epoch_PBEs_lfp_feature_value[epoch]))), 2))\n",
    "    #         ax_inset.set_ylabel('cumulative number of PBEs', fontsize=6)\n",
    "    #         ax_inset.tick_params(labelsize=6)\n",
    "    #         ax_inset.set_facecolor('none')\n",
    "    #         # ax_inset.get_legend().remove()\n",
    "\n",
    "\n",
    "    #     filename = f'{session_name}.learned_tunings_low_versus_high_{current_lfp_feature}_matched_particpation_NREM_delta_spindle_spectrogram.svg'\n",
    "    #     file_path = os.path.join(session_dataset_path, filename)\n",
    "    #     plt.savefig(file_path)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing Learned Tunings Based on PBEs During Low vs High Values of certain LFP features, Including Ripple Power, Slow Wave Amplitude, and Spindle Power.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# import numpy as np\n",
    "# data_dir = '/home/kouroshmaboudi/Documents/Learned_tuning_Python/Datasets/Achilles_11012013'\n",
    "\n",
    "# filename = 'Achilles_11012013.learned_tunings_NRE_low_versus_high_swa_matched_particpation.npy'\n",
    "# file_path = os.path.join(data_dir, filename)\n",
    "\n",
    "# learned_tunings_NREM_low_versus_high_swa_matched_particpation = np.load(file_path, allow_pickle=True).item() # use this to read the files\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "    #------------------------------------------------------------------------------------------------------\n",
    "    # Learned tunings during Non-REM versus Quiet Wake ripples, with equalized number of firing time bins for each unit across different epochs or conditions\n",
    "    # The frequency of PBEs ocurring during NREM ('N') or quiet wake ('W')\n",
    "\n",
    "    PBEs_epoch_indices_all_off_periods = [None for _ in range(2)]\n",
    "    num_PBEs_all_off_periods = dict()\n",
    "\n",
    "    PBEs_epoch_indices_NREM = [None for _ in range(2)]\n",
    "    num_PBEs_NREM = dict()\n",
    "\n",
    "    PBEs_epoch_indices_QW = [None for _ in range(2)]\n",
    "    num_PBEs_QW = dict()\n",
    "\n",
    "    time_bin_duration=0.02\n",
    "    num_PF_shuffles = 10000\n",
    "\n",
    "\n",
    "    for i, epoch in enumerate(['pre', 'post']):\n",
    "        if epoch == 'pre':\n",
    "            epoch_idx = 0\n",
    "            epoch_duration = epochs[0,1] - epochs[0,0]\n",
    "        elif epoch == 'post':\n",
    "            epoch_idx = 2\n",
    "            epoch_duration = 4*60*60\n",
    "\n",
    "            \n",
    "        if_inside_epoch = (PBEs_peak_time >= epochs[epoch_idx,0]) & (PBEs_peak_time <= epochs[epoch_idx,0]+epoch_duration)\n",
    "        \n",
    "        # Indices of PBEs ocurring during each epoch, regardless of their corresponding brain state\n",
    "        PBEs_epoch_indices_all_off_periods[i] = np.where(if_inside_epoch)[0]\n",
    "        num_PBEs_all_off_periods[epoch] = PBEs_epoch_indices_all_off_periods[i].shape[0]\n",
    "\n",
    "\n",
    "        # Indices of PBEs ocurring during each epoch, limited to specific brain state, e.g. NREM or quiet wake (QW) \n",
    "        PBEs_epoch_indices_NREM[i] = np.where(np.logical_and(if_inside_epoch, brain_state_first_letter == 'N'))[0]\n",
    "        num_PBEs_NREM[epoch] = PBEs_epoch_indices_NREM[i].shape[0]\n",
    "\n",
    "        PBEs_epoch_indices_QW[i] =  np.where(np.logical_and(if_inside_epoch, brain_state_first_letter == 'Q'))[0]\n",
    "        num_PBEs_QW[epoch] = PBEs_epoch_indices_QW[i].shape[0]\n",
    "\n",
    "\n",
    "    learned_tunings_matched_particpation = {\n",
    "        'pre':  {'all_offline': None, 'NREM': None, 'QW': None},\n",
    "        'post': {'all_offline': None, 'NREM': None, 'QW': None}\n",
    "    }\n",
    "\n",
    "    # calculate learned tunings within each epoch *using all offline PBEs*, regardless of their corresponding brain state   \n",
    "    learned_tunings_matched_participation_all = calculate_learned_tuning_matched_participation(PBEs, spikes, PBEs_epoch_indices_all_off_periods, L_ratios, time_bin_duration)\n",
    "    print('calculation of learned tunings based on all offline PBEs within each epoch was completed! ')\n",
    "\n",
    "    learned_tunings_matched_particpation['pre']['all_offline'] = learned_tunings_matched_participation_all[:,:,0]\n",
    "    learned_tunings_matched_particpation['post']['all_offline'] = learned_tunings_matched_participation_all[:,:,1]\n",
    "\n",
    "    # learned tunings based on *NREM PBEs* within each epoch\n",
    "    learned_tunings_matched_participation_NREM_all = calculate_learned_tuning_matched_participation(PBEs, spikes, PBEs_epoch_indices_NREM, L_ratios, time_bin_duration)\n",
    "    print('calculation of learned tunings based on NREM PBEs within each epoch was completed! ')\n",
    "\n",
    "    learned_tunings_matched_particpation['pre']['NREM'] = learned_tunings_matched_participation_NREM_all[:,:,0]\n",
    "    learned_tunings_matched_particpation['post']['NREM'] = learned_tunings_matched_participation_NREM_all[:,:,1]\n",
    "\n",
    "    # learned tunnigs based on *QW PBEs* within each epoch\n",
    "    learned_tunings_matched_participation_QW_all = calculate_learned_tuning_matched_participation(PBEs, spikes, PBEs_epoch_indices_QW, L_ratios, time_bin_duration)\n",
    "    print('calculation of learned tunings based on QW PBEs within each epoch was completed! ')\n",
    "\n",
    "    learned_tunings_matched_particpation['pre']['QW'] = learned_tunings_matched_participation_QW_all[:,:,0]\n",
    "    learned_tunings_matched_particpation['post']['QW'] = learned_tunings_matched_participation_QW_all[:,:,1]\n",
    "\n",
    "\n",
    "    filename = f'{session_name}.learned_tunings_NREM_vs_QW_matched_firing_PRE_and_POST.npy'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "    np.save(file_path, learned_tunings_matched_particpation)\n",
    "\n",
    "\n",
    "\n",
    "    # Plot the distributions of learned tunings for NREM and QW periods within each epoch\n",
    "    colors = sns.color_palette(\"husl\", 2) # Set color palette\n",
    "    sns.set_style('whitegrid') # Set style and context\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",
    "    fig, axes = plt.subplots(1,2, figsize = (10, 4))\n",
    "\n",
    "    for i, epoch in enumerate(['pre', 'post']):\n",
    "        \n",
    "        # learned tunings regardless of the brain state\n",
    "        num_PBEs_all = num_PBEs_all_off_periods[epoch]\n",
    "        learned_tuning_place_field_pearson_corr_all_offline, _, median_LT_PF_pearson_corr_all_offline = calculate_place_field_fidelity_of_learned_tuning(learned_tunings_matched_particpation[epoch]['all_offline'][active_units], place_fields_uni[active_units], num_PF_shuffles)\n",
    "        \n",
    "        # learned tunings separate for each brain state\n",
    "        learned_tuning_place_field_pearson_corr_NREM, _, median_LT_PF_pearson_corr_NREM = calculate_place_field_fidelity_of_learned_tuning(learned_tunings_matched_particpation[epoch]['NREM'][active_units], place_fields_uni[active_units], num_PF_shuffles)\n",
    "        learned_tuning_place_field_pearson_corr_QW, _, median_LT_PF_pearson_corr_QW = calculate_place_field_fidelity_of_learned_tuning(learned_tunings_matched_particpation[epoch]['QW'][active_units], place_fields_uni[active_units], num_PF_shuffles)\n",
    "        \n",
    "\n",
    "        # Plot the distributions using ECDFs with overlaid ticks\n",
    "        sns.ecdfplot(learned_tuning_place_field_pearson_corr_NREM, ax = axes[i], label=f'{epoch} NREM', color=colors[0], linewidth = 2)\n",
    "        sns.ecdfplot(learned_tuning_place_field_pearson_corr_QW, ax = axes[i], label=f'{epoch} QW', color=colors[1], linewidth = 2)\n",
    "        sns.ecdfplot(learned_tuning_place_field_pearson_corr_all_offline, ax = axes[i], label=f'{epoch} all offline', color='black', linewidth = 2)\n",
    "\n",
    "        axes[i].set_xlim([-1,1])\n",
    "        axes[i].set_xlabel(f'{epoch} LT fidelity', fontsize=10)\n",
    "        axes[i].set_ylabel('Proportion of units', fontsize=10)\n",
    "        axes[i].tick_params(labelsize=8)\n",
    "        axes[i].legend(fontsize=8)\n",
    "\n",
    "\n",
    "        # Add ticks to the x-axis\n",
    "        y_min, y_max = axes[i].get_ylim()\n",
    "        for x in learned_tuning_place_field_pearson_corr_NREM:\n",
    "            axes[i].plot([x, x], [y_min+0.01, y_min + 0.02], '|-', color=colors[0], linewidth=0.25)\n",
    "        for x in learned_tuning_place_field_pearson_corr_QW:\n",
    "            axes[i].plot([x, x], [y_min+0.06, y_min + 0.07], '|-', color=colors[1], linewidth=0.25)\n",
    "        \n",
    "\n",
    "        # Calculate and display medians\n",
    "        median_nrem = median_LT_PF_pearson_corr_NREM['data']\n",
    "        median_nrem_pvalue = median_LT_PF_pearson_corr_NREM['p_value']\n",
    "\n",
    "        median_qw = median_LT_PF_pearson_corr_QW['data']\n",
    "        median_qw_pvalue = median_LT_PF_pearson_corr_QW['p_value']\n",
    "\n",
    "        median_offline = median_LT_PF_pearson_corr_all_offline['data']\n",
    "        median_offline_pvalue = median_LT_PF_pearson_corr_all_offline['p_value']\n",
    "\n",
    "\n",
    "        def get_pval_statement(pvalue):\n",
    "            if pvalue < 0.0001:\n",
    "                pvalue_statement = 'P<1e-4'\n",
    "            else:\n",
    "                pvalue_statement = f'P={pvalue:.4f}'\n",
    "            return pvalue_statement\n",
    "\n",
    "\n",
    "        axes[i].axvline(median_nrem, color=colors[0], linestyle='dashed', label=f'Median {epoch} NREM={median_nrem:.2f},{get_pval_statement(median_nrem_pvalue)}')\n",
    "        axes[i].axvline(median_qw, color=colors[1], linestyle='dashed', label=f'Median {epoch} QW={median_qw:.2f},{get_pval_statement(median_qw_pvalue)}')\n",
    "        axes[i].axvline(median_offline, color='black', linestyle='dashed', label=f'Median {epoch} all offline={median_offline:.2f},{get_pval_statement(median_offline_pvalue)}')\n",
    "\n",
    "        axes[i].legend(fontsize=6)\n",
    "\n",
    "        # Perform rank-sum test\n",
    "        statistic, p_value = ranksums(learned_tuning_place_field_pearson_corr_NREM,\n",
    "                                    learned_tuning_place_field_pearson_corr_QW)\n",
    "\n",
    "        # Add line with p-value above the plot\n",
    "        axes[i].text((median_nrem + median_qw) / 2, 0.2, f'p-value = {p_value:.4f}', ha='center', fontsize=6)\n",
    "\n",
    "\n",
    "        # Bar plot of frequency of ripples in each brain states\n",
    "        ax_inset = axes[i].inset_axes([0.2, 0.4, 0.2, 0.2])\n",
    "\n",
    "\n",
    "        ax_inset.bar([0, 1],[num_PBEs_NREM[epoch], num_PBEs_QW[epoch]], color=[colors[0], colors[1]]) #unique_strings, \n",
    "        ax_inset.set_xticks([0, 1], ['NREM', 'QW'])\n",
    "        ax_inset.set_ylabel('number of PBEs', fontsize=6)\n",
    "        ax_inset.tick_params(labelsize=6)\n",
    "        # axes[i].tight_layout()\n",
    "\n",
    "\n",
    "    filename = f'{session_name}.learned_tunings_offline_NREM_vs_QW_matched_firing_PRE_and_POST_2.svg'\n",
    "    file_path = os.path.join(session_dataset_path, filename)\n",
    "    plt.savefig(file_path)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'RatS-Day2-2020-11-27_10-22-29'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "session_name"
   ]
  },
  {
   "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
}