multisub_extract_evoked_responses.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This creates epoched files for evoked EEG and units responses for multiple subjects. Designed for use with main experimental subjects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "from glob import glob\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import pickle\n",
    "\n",
    "import gspread\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append(r'C:\\Users\\lesliec\\code')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tbd_eeg.tbd_eeg.data_analysis.eegutils import EEGexp\n",
    "from tbd_eeg.tbd_eeg.data_analysis.Utilities.utilities import (\n",
    "    get_stim_events, get_evoked_traces, find_nearest_ind, qualitycheck_trials)\n",
    "from allensdk.core.mouse_connectivity_cache import MouseConnectivityCache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Developed in NPX_find_bursts_testing.ipynb, this version is faster and only returns start times and spike counts ##\n",
    "def find_bursts_indunit(spike_times):\n",
    "    \n",
    "    preISIs = np.diff(spike_times)[:-1]\n",
    "    postISIs = np.diff(spike_times)[1:]\n",
    "    ## Find starts ##\n",
    "    bs_inds = np.nonzero((preISIs > 0.1) * (postISIs < 0.004))[0]\n",
    "    if len(bs_inds) == 0:\n",
    "        return np.array([]), np.array([])\n",
    "    \n",
    "    burst_starts = bs_inds + 1 # +1 corrects for the actual spike ind\n",
    "    ## Loop through burst starts to find spikes that belong to the burst\n",
    "    burst_counts = []\n",
    "    for st_ind in bs_inds:\n",
    "        spkind = st_ind+1\n",
    "        bcount = 1\n",
    "        while (spkind < len(preISIs)) and (preISIs[spkind] < 0.004):\n",
    "            spkind += 1\n",
    "            bcount += 1\n",
    "        burst_counts.append(bcount)\n",
    "    \n",
    "    return spike_times[burst_starts], np.array(burst_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_closest_region(sunit_info, struct_tree, annot):\n",
    "    ## Finds a grey matter region above/below an unknown region ##\n",
    "    Vind = sunit_info.CCF_DV\n",
    "    vent_sip = struct_tree.get_structures_by_id([annot[sunit_info.CCF_AP, Vind, sunit_info.CCF_ML]])[0]['structure_id_path']\n",
    "    while not struct_tree.structure_descends_from(vent_sip[-1], 8):\n",
    "        Vind += 1\n",
    "        vent_sip = struct_tree.get_structures_by_id([annot[sunit_info.CCF_AP, Vind, sunit_info.CCF_ML]])[0]['structure_id_path']\n",
    "\n",
    "    Dind = sunit_info.CCF_DV\n",
    "    dors_sip = struct_tree.get_structures_by_id([annot[sunit_info.CCF_AP, Dind, sunit_info.CCF_ML]])[0]['structure_id_path']\n",
    "    while not struct_tree.structure_descends_from(dors_sip[-1], 8):\n",
    "        Dind -= 1\n",
    "        dors_sip = struct_tree.get_structures_by_id([annot[sunit_info.CCF_AP, Dind, sunit_info.CCF_ML]])[0]['structure_id_path']\n",
    "\n",
    "    if (Vind - sunit_info.CCF_DV) <= (sunit_info.CCF_DV - Dind):\n",
    "        return struct_tree.get_structures_by_id([vent_sip[-1]])[0]['acronym']\n",
    "    elif (Vind - sunit_info.CCF_DV) > (sunit_info.CCF_DV - Dind):\n",
    "        return struct_tree.get_structures_by_id([dors_sip[-1]])[0]['acronym']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_region_from_children(test_id, parent_id, struct_tree):\n",
    "    try:\n",
    "        child_ind = np.nonzero([\n",
    "            struct_tree.structure_descends_from(test_id, x) for x in struct_tree.child_ids([parent_id])[0]\n",
    "        ])[0][0]\n",
    "        return struct_tree.get_structures_by_id([struct_tree.child_ids([parent_id])[0][child_ind]])[0]['acronym']\n",
    "    except:\n",
    "        return struct_tree.get_structures_by_id([parent_id])[0]['acronym']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_parent_region(region_acronym, struct_tree):\n",
    "    areas_of_interest = {\n",
    "        'SM-TH': ['AV', 'CL', 'MD', 'PO', 'PF', 'VAL', 'VPL', 'VPM', 'VM'],\n",
    "    }\n",
    "    \n",
    "    reg_id = struct_tree.get_structures_by_acronym([region_acronym])[0]['id']\n",
    "    if struct_tree.structure_descends_from(reg_id, 567):\n",
    "        if struct_tree.structure_descends_from(reg_id, 315):\n",
    "            return get_region_from_children(reg_id, 315, struct_tree)\n",
    "        elif struct_tree.structure_descends_from(reg_id, 698):\n",
    "            return 'OLF'\n",
    "        elif struct_tree.structure_descends_from(reg_id, 1089):\n",
    "            return get_region_from_children(reg_id, 1089, struct_tree)\n",
    "        elif struct_tree.structure_descends_from(reg_id, 703):\n",
    "            return get_region_from_children(reg_id, 703, struct_tree)\n",
    "        elif struct_tree.structure_descends_from(reg_id, 477):\n",
    "            return 'STR'\n",
    "        elif struct_tree.structure_descends_from(reg_id, 803):\n",
    "            return 'PAL'\n",
    "        else:\n",
    "            return 'unassigned'\n",
    "    elif struct_tree.structure_descends_from(reg_id, 343):\n",
    "        if struct_tree.structure_descends_from(reg_id, 1129):\n",
    "            if region_acronym == 'RT':\n",
    "                return 'RT-TH'\n",
    "            elif region_acronym in areas_of_interest['SM-TH']:\n",
    "                return 'SM-TH'\n",
    "            else:\n",
    "                return 'other-TH'\n",
    "        elif struct_tree.structure_descends_from(reg_id, 1097):\n",
    "            return 'HY'\n",
    "        else:\n",
    "            return get_region_from_children(reg_id, 343, struct_tree)\n",
    "    else:\n",
    "        return 'unassigned'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_parent_region_to_df(unit_info_df, struct_tree, annot):\n",
    "    ## First, make sure all names in region column correspond to a CCF region (removes nan values) ##\n",
    "    adj_regions = unit_info_df['region'].values.copy()\n",
    "    for indi, rowi in unit_info_df.iterrows():\n",
    "        try:\n",
    "            str_info = struct_tree.get_structures_by_acronym([rowi.region])[0]\n",
    "        except KeyError:\n",
    "            if rowi.depth <= 0: # unit was placed above brain\n",
    "                new_region_id = annot[rowi.CCF_AP, np.nonzero(annot[rowi.CCF_AP, :, rowi.CCF_ML])[0][0], rowi.CCF_ML]\n",
    "                adj_regions[indi] = struct_tree.get_structures_by_id([new_region_id])[0]['acronym']\n",
    "            else:\n",
    "                Lind = rowi.CCF_ML\n",
    "                while annot[rowi.CCF_AP, rowi.CCF_DV, Lind] == 0:\n",
    "                    Lind -= 1\n",
    "                new_region_id = struct_tree.get_structures_by_id(\n",
    "                    [annot[rowi.CCF_AP, rowi.CCF_DV, Lind]])[0]['structure_id_path'][-1]\n",
    "                adj_regions[indi] = struct_tree.get_structures_by_id([new_region_id])[0]['acronym']\n",
    "    unit_info_df['adj_region'] = adj_regions\n",
    "    \n",
    "    ## Second, re-assign any non-grey matter areas to the closest region ##\n",
    "    adj_regions = unit_info_df['adj_region'].values.copy()\n",
    "    for indi, rowi in unit_info_df.iterrows():\n",
    "        reg_id = struct_tree.get_structures_by_acronym([rowi.adj_region])[0]['id']\n",
    "        if not struct_tree.structure_descends_from(reg_id, 8):\n",
    "            adj_regions[indi] = find_closest_region(rowi, struct_tree, annot)\n",
    "    unit_info_df['adj_region'] = adj_regions\n",
    "    \n",
    "    ## Finally, assign a parent region to each adjusted CCF region ##\n",
    "    parent_regions = unit_info_df['adj_region'].values.copy()\n",
    "    for indi, rowi in unit_info_df.iterrows():\n",
    "        parent_regions[indi] = get_parent_region(rowi.adj_region, struct_tree)\n",
    "    unit_info_df['parent_region'] = parent_regions\n",
    "    \n",
    "    return unit_info_df.drop('adj_region', axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load Google excel log file to get metadata for experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "_gc = gspread.service_account() # need a key file to access the account\n",
    "_sh = _gc.open('Templeton-log_exp') # open the spreadsheet\n",
    "# _sh = _gc.open('Zap_Zip-log_exp') # open the spreadsheet\n",
    "_df = pd.DataFrame(_sh.sheet1.get()) # load the first worksheet\n",
    "metadata = _df.T.set_index(0).T # put it in a nicely formatted dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load subjects.csv file to get CCF resolution"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "multisub_file = r\"C:\\Users\\lesliec\\OneDrive - Allen Institute\\data\\brain_states_subjects.csv\"\n",
    "subject_df = pd.read_csv(multisub_file, converters={'mouse': str}).astype({'analyze': bool})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load subjects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "subjects = {\n",
    "#     '689242': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse689242\\aw_iso_2023-07-20_10-52-57\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse689242\\aw_psi_2023-07-19_10-29-49\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse689242\\urethane_2023-07-21_12-32-25\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '689239': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse689239\\aw_iso_2023-08-09_11-15-42\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse689239\\aw_psi_2023-08-10_11-26-36\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse689239\\urethane_2023-08-11_11-33-46\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '689240': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse689240\\aw_iso_2023-08-18_10-03-39\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse689240\\aw_psi_2023-08-17_09-46-39\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '692644': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse692644\\aw_iso_2023-09-06_11-27-02\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse692644\\aw_psi_2023-09-07_11-51-57\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse692644\\urethane_2023-09-08_11-15-18\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '688277': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse688277\\aw_iso_2023-09-21_10-41-15\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse688277\\urethane_2023-09-22_12-04-23\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '703062': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse703062\\aw_iso_2023-10-19_13-13-25\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse703062\\urethane_2023-10-20_12-00-46\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '703063': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse703063\\aw_iso_2023-11-16_11-16-30\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse703063\\aw_psi_2023-11-15_11-08-12\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse703063\\urethane_2023-11-17_12-47-18\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '703064': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse703064\\aw_iso_2023-11-29_11-23-30\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse703064\\aw_psi_2023-11-30_12-06-43\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse703064\\urethane_2023-12-01_11-20-15\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '703065': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse703065\\aw_iso_2023-12-07_10-23-39\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse703065\\aw_psi_2023-12-06_11-04-39\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse703065\\urethane_2023-12-08_11-57-44\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '709401': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse709401\\aw_iso_2023-12-13_09-55-07\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse709401\\aw_psi_2023-12-14_11-17-30\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse709401\\urethane_2023-12-15_11-03-00\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '709400': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse709400\\aw_iso_2024-01-31_11-35-57\\experiment1\\recording1',\n",
    "#         'awake-ketanserin': r'F:\\psi_exp\\mouse709400\\aw_ket_2024-02-01_11-12-34\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse709400\\urethane_2024-02-02_11-52-38\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '709402': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse709402\\aw_iso_2024-02-15_10-26-16\\experiment1\\recording1',\n",
    "#         'awake-ketanserin': r'F:\\psi_exp\\mouse709402\\aw_ket_2024-02-14_12-32-14\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse709402\\urethane_2024-02-16_11-15-43\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '720762': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse720762\\aw_iso_2024-02-21_10-32-01\\experiment1\\recording1',\n",
    "#         'awake-ketanserin': r'F:\\psi_exp\\mouse720762\\aw_ket_2024-02-22_10-43-35\\experiment1\\recording1',\n",
    "#         'urethane': r'F:\\psi_exp\\mouse720762\\urethane_2024-02-23_11-35-21\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '724057': {\n",
    "#         'awake-sal': r'F:\\psi_exp\\mouse724057\\aw_sal_2024-04-03_10-32-05\\experiment1\\recording1',\n",
    "#         'awake-psilocybin': r'F:\\psi_exp\\mouse724057\\aw_psi_2024-04-04_10-35-18\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '730911': {\n",
    "#         'awake-sal': r'F:\\psi_exp\\mouse730911\\aw_sal_2024-05-01_11-57-16\\experiment1\\recording1',\n",
    "#         'awake-ketanserin': r'F:\\psi_exp\\mouse730911\\aw_ket_2024-05-02_10-41-41\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '730913': {\n",
    "#         'awake-sal': r'F:\\psi_exp\\mouse730913\\aw_sal_2024-04-24_10-43-30\\experiment1\\recording1',\n",
    "#         'awake-ketanserin': r'F:\\psi_exp\\mouse730913\\aw_ket_2024-04-26_10-11-34\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '735049': {\n",
    "#         'awake-sal': r'F:\\psi_exp\\mouse735049\\aw_sal_2024-05-22_11-05-25\\experiment1\\recording1',\n",
    "#         'awake-ketanserin': r'F:\\psi_exp\\mouse735049\\aw_ket_2024-05-23_11-42-08\\experiment1\\recording1',\n",
    "#     },\n",
    "#     '709399': {\n",
    "#         'awake-iso': r'F:\\psi_exp\\mouse709399\\aw_iso_2024-01-24_10-26-14\\experiment1\\recording1',\n",
    "#     },\n",
    "    '655955': {\n",
    "        'saline': r'F:\\psi_exp\\mouse655955\\estim_2022-12-13_10-05-19\\experiment1\\recording1',\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check trial quality for EEG signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "655955: saline\n",
      "Experiment type: electrical stimulation\n",
      "Stim_log saved with trial quality column.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "### Re-process quality for all trials in all subjects (updated quality function 6/4/24) ###\n",
    "for mouse, explist in subjects.items():\n",
    "    for exptype, dataloc in explist.items():\n",
    "        print('{}: {}'.format(mouse, exptype))\n",
    "        exp = EEGexp(dataloc, preprocess=False, make_stim_csv=False)\n",
    "        \n",
    "        ## Check quality for all trials ##\n",
    "        if np.any([True for x in exp.experiment_data if 'recording' in x]):\n",
    "            ## Grab exp metadata from Templeton-log_exp ##\n",
    "            exp_meta = metadata[(\n",
    "                (metadata['mouse_name'].str.contains(mouse)) &\n",
    "                (metadata['exp_name'].str.contains(os.path.basename(os.path.dirname(exp.experiment_folder))))\n",
    "            )].squeeze()\n",
    "            badchstr = exp_meta['EEG bad_channels'].replace(' ','')\n",
    "            if badchstr == 'all':\n",
    "                print(' All EEG chs are bad.')\n",
    "                bad_chs = np.arange(30)\n",
    "            else:\n",
    "                bad_chs = []\n",
    "                for char in badchstr.split(','):\n",
    "                    if char.isdecimal():\n",
    "                        bad_chs.append(int(char))\n",
    "            \n",
    "            qualitycheck_trials(exp, known_bad_chs=bad_chs)\n",
    "\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "for mouse, explist in subjects.items():\n",
    "    for exptype, dataloc in explist.items():\n",
    "        print('{}: {}'.format(mouse, exptype))\n",
    "        exp = EEGexp(dataloc, preprocess=False, make_stim_csv=False)\n",
    "        \n",
    "        ## Load stim log ##\n",
    "        stim_log = pd.read_csv(exp.stimulus_log_file).astype({'parameter': str})\n",
    "        if 'good' not in stim_log.columns:\n",
    "            if np.any([True for x in exp.experiment_data if 'recording' in x]):\n",
    "                ## Grab exp metadata from Templeton-log_exp ##\n",
    "                exp_meta = metadata[(\n",
    "                    (metadata['mouse_name'].str.contains(mouse)) &\n",
    "                    (metadata['exp_name'].str.contains(os.path.basename(os.path.dirname(exp.experiment_folder))))\n",
    "                )].squeeze()\n",
    "                badchstr = exp_meta['EEG bad_channels'].replace(' ','')\n",
    "                if badchstr == 'all':\n",
    "                    print(' All EEG chs are bad, marking all trials as \"good\".')\n",
    "                    stim_log['good'] = [True] * len(stim_log)\n",
    "                    stim_log.to_csv(exp.stimulus_log_file, index=False)\n",
    "                else:\n",
    "                    print(' Checking trial quality...')\n",
    "                    bad_chs = []\n",
    "                    for char in badchstr.split(','):\n",
    "                        if char.isdecimal():\n",
    "                            bad_chs.append(int(char))\n",
    "                    qualitycheck_trials(exp, known_bad_chs=bad_chs)\n",
    "            else:\n",
    "                print(' Experiment has no EEG, making all trials \"good\".')\n",
    "                stim_log['good'] = [True] * len(stim_log)\n",
    "                stim_log.to_csv(exp.stimulus_log_file, index=False)\n",
    "        else:\n",
    "            print(' Stim_log already has trial quality, skipping.')\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "overwrite_existing_files = False\n",
    "\n",
    "event_window = [-2.0, 2.0]\n",
    "\n",
    "apply_mask = True\n",
    "apply_hpass = True\n",
    "apply_lpass = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Process running signal and EEG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "655955: saline\n",
      "Experiment type: electrical stimulation\n",
      "  Getting event-related running...\n",
      "  Loading EEG data...\n",
      "  Getting EEG traces...\n",
      "   ...saving F:\\psi_exp\\mouse655955\\estim_2022-12-13_10-05-19\\experiment1\\recording1\\evoked_data\\event_EEGtraces.npy.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for mouse, explist in subjects.items():\n",
    "    for exptype, dataloc in explist.items():\n",
    "        print('{}: {}'.format(mouse, exptype))\n",
    "        exp = EEGexp(dataloc, preprocess=False, make_stim_csv=False)\n",
    "        \n",
    "        ## Set file names ##\n",
    "        running_file = os.path.join(exp.data_folder, 'running_signal.npy')\n",
    "        running_ts_file = os.path.join(exp.data_folder, 'running_timestamps_master_clock.npy')\n",
    "#         running_ts_file = os.path.join(exp.data_folder, 'running_timestamps.npy')\n",
    "        evoked_folder = os.path.join(exp.data_folder, 'evoked_data')\n",
    "        if not os.path.exists(evoked_folder):\n",
    "            os.mkdir(evoked_folder)\n",
    "        event_running_file = os.path.join(evoked_folder, 'event_running_speed.npy')\n",
    "        event_running_ts_file = os.path.join(evoked_folder, 'event_running_times.npy')\n",
    "        event_EEGtraces_file = os.path.join(evoked_folder, 'event_EEGtraces.npy')\n",
    "        event_EEGtraces_ts_file = os.path.join(evoked_folder, 'event_EEGtraces_times.npy')\n",
    "\n",
    "        ## Load stim log ##\n",
    "        stim_log = pd.read_csv(exp.stimulus_log_file)\n",
    "        all_event_times = stim_log['onset'].values\n",
    "            \n",
    "        ## Load running signal and get mean event speed ##\n",
    "        if not os.path.exists(event_running_file) or overwrite_existing_files:\n",
    "            print('  Getting event-related running...')\n",
    "            if os.path.exists(running_file):\n",
    "                run_signal = np.load(running_file)\n",
    "                run_timestamps = np.load(running_ts_file)\n",
    "            else:\n",
    "                print('  Loading running from sync and saving...')\n",
    "                run_signal, run_timestamps = exp.load_running()\n",
    "                np.save(running_file, run_signal, allow_pickle=False)\n",
    "                np.save(running_ts_file, run_timestamps, allow_pickle=False)\n",
    "            \n",
    "            rinds = np.arange(-int(-event_window[0] * 100), int(event_window[1] * 100))\n",
    "            event_inds = np.array([find_nearest_ind(run_timestamps, x) for x in all_event_times])\n",
    "            event_run_speed = run_signal[np.repeat([rinds], len(event_inds), axis=0).T + event_inds]\n",
    "            event_run_times = rinds / 100\n",
    "            ## Save ##\n",
    "            np.save(event_running_file, event_run_speed, allow_pickle=False)\n",
    "            np.save(event_running_ts_file, event_run_times, allow_pickle=False)\n",
    "            ## Add speed to stim_log ##\n",
    "            evinds = np.nonzero((event_run_times >= -0.5) & (event_run_times < 0.5))[0]\n",
    "            mean_speed = np.mean(event_run_speed[evinds, :], axis=0)\n",
    "            stim_log['mean_speed'] = mean_speed\n",
    "            stim_log['resting_trial'] = stim_log['mean_speed'] == 0\n",
    "            stim_log.to_csv(exp.stimulus_log_file, index=False)\n",
    "            \n",
    "        if np.any([True for x in exp.experiment_data if 'recording' in x]):\n",
    "            ## Grab exp metadata from Templeton-log_exp ##\n",
    "            exp_meta = metadata[(\n",
    "                (metadata['mouse_name'].str.contains(mouse)) &\n",
    "                (metadata['exp_name'].str.contains(os.path.basename(os.path.dirname(exp.experiment_folder))))\n",
    "            )].squeeze()\n",
    "            badchstr = exp_meta['EEG bad_channels'].replace(' ','')\n",
    "            if (not os.path.exists(event_EEGtraces_file) or overwrite_existing_files) and (badchstr != 'all'):\n",
    "                ## Load EEG data and preprocess ##\n",
    "                print('  Loading EEG data...')\n",
    "                datai, tsi = exp.load_eegdata()\n",
    "                eeg_chs = np.arange(0, datai.shape[1])\n",
    "\n",
    "                ## Mask estim artifact ##\n",
    "                if apply_mask:\n",
    "                    mask_samples = int(0.002 * exp.ephys_params['EEG']['sample_rate'])\n",
    "                    for etime in stim_log.loc[stim_log['stim_type'] == 'biphasic', 'onset'].to_numpy():\n",
    "                        val = find_nearest_ind(tsi, etime) - 2\n",
    "                        datai[val:val+mask_samples, :] = datai[val:val-mask_samples:-1, :]\n",
    "\n",
    "                ## Apply high-pass filter ##\n",
    "                if apply_hpass:\n",
    "                    hpb, hpa = signal.butter(3, 0.1/(exp.ephys_params['EEG']['sample_rate']/2), btype='highpass')\n",
    "                    datai = signal.filtfilt(hpb, hpa, datai, axis=0)\n",
    "\n",
    "                ## Get evoked traces ##\n",
    "                print('  Getting EEG traces...')\n",
    "                event_traces, event_ts = get_evoked_traces(\n",
    "                    datai, tsi, all_event_times, -event_window[0], event_window[1], exp.ephys_params['EEG']['sample_rate'])\n",
    "\n",
    "                ## Apply lowpass filter ##\n",
    "                if apply_lpass:\n",
    "                    lpb, lpa = signal.butter(3, 100/(exp.ephys_params['EEG']['sample_rate']/2), btype='low')\n",
    "                    event_traces = signal.filtfilt(lpb, lpa, event_traces, axis=0)\n",
    "\n",
    "                ## Save ##\n",
    "                print('   ...saving {}.'.format(event_EEGtraces_file))\n",
    "                np.save(event_EEGtraces_file, event_traces, allow_pickle=False)\n",
    "                np.save(event_EEGtraces_ts_file, event_ts, allow_pickle=False)\n",
    "            else:\n",
    "                print('  Not creating EEG traces file, it already exists or all EEG chs are bad.')\n",
    "        else:\n",
    "            print('  No EEG in this recording.')\n",
    "\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Process units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "655955: saline\n",
      "Experiment type: electrical stimulation\n",
      "  Getting probe info...\n",
      "  Getting spike times...\n",
      "  Adding parent region...\n",
      "  Time to get unit spike times and save: 0.99 min\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for mouse, explist in subjects.items():\n",
    "    for exptype, dataloc in explist.items():\n",
    "        print('{}: {}'.format(mouse, exptype))\n",
    "        exp = EEGexp(dataloc, preprocess=False, make_stim_csv=False)\n",
    "        \n",
    "        probe_list = [x.replace('_sorted', '') for x in exp.experiment_data if 'probe' in x]\n",
    "        if len(probe_list) == 0:\n",
    "            print(' This experiment has no probe data, not making spike times files.\\n')\n",
    "            continue\n",
    "        \n",
    "        ## Set file names ##\n",
    "        evoked_folder = os.path.join(exp.data_folder, 'evoked_data')\n",
    "        if not os.path.exists(evoked_folder):\n",
    "            os.mkdir(evoked_folder)\n",
    "        unit_info_file = os.path.join(evoked_folder, 'all_units_info.csv')\n",
    "        unit_allspiketimes_file = os.path.join(evoked_folder, 'units_allspikes.pkl')\n",
    "        unit_eventspikes_file = os.path.join(evoked_folder, 'units_event_spikes.pkl')\n",
    "        if overwrite_existing_files:\n",
    "            pass # will overwrite all subjects' files\n",
    "        else:\n",
    "            if os.path.exists(unit_info_file):\n",
    "                print('  {} already exists, skipping analysis.\\n'.format(unit_info_file))\n",
    "                continue\n",
    "\n",
    "        ## Load stim log ##\n",
    "        stim_log = pd.read_csv(exp.stimulus_log_file)\n",
    "        all_event_times = stim_log['onset'].values\n",
    "        \n",
    "        ## Get probe info ##\n",
    "        print('  Getting probe info...')\n",
    "        probe_data = {}\n",
    "        for pbi, probei in enumerate(probe_list):\n",
    "            probe_data[probei] = {}\n",
    "            ## Load probe_info.json ##\n",
    "            with open(exp.ephys_params[probei]['probe_info']) as data_file:\n",
    "                data = json.load(data_file)\n",
    "            npx_allch = np.array(data['channel'])\n",
    "            surface_ch = int(data['surface_channel'])\n",
    "            allch_z = np.array(data['vertical_pos'])\n",
    "            ref_mask = np.array(data['mask'])\n",
    "            npx_chs = np.array([x for x in npx_allch if ref_mask[x] and x <= surface_ch])\n",
    "            probe_data[probei]['ch_depths'] = allch_z[surface_ch] - allch_z\n",
    "            \n",
    "            ## Select units and get peak chs ##\n",
    "            select_units, peak_chs, unit_metrics = exp.get_probe_units(probei)\n",
    "            ## Sort units ##\n",
    "            probe_data[probei]['units'] = select_units[np.squeeze(np.argsort(peak_chs))]\n",
    "            probe_data[probei]['chs'] = peak_chs[np.squeeze(np.argsort(peak_chs))]\n",
    "            probe_data[probei]['duration'] = unit_metrics.duration.values[np.squeeze(np.argsort(peak_chs))]\n",
    "            \n",
    "            ## Load spike times and cluster ids ##\n",
    "            probe_data[probei]['spike_times'] = np.load(exp.ephys_params[probei]['spike_times'])\n",
    "            probe_data[probei]['spike_clusters'] = np.load(exp.ephys_params[probei]['spike_clusters'])\n",
    "            \n",
    "            if 'area_ch' in data.keys():\n",
    "                probe_data[probei]['areas'] = unit_metrics.area.values[np.squeeze(np.argsort(peak_chs))]\n",
    "                probe_data[probei]['CCF_coords'] = unit_metrics.ccf_coord.values[np.squeeze(np.argsort(peak_chs))]\n",
    "                \n",
    "        ## Get unit info, spikes and event-spikes, then save files ##\n",
    "        print('  Getting spike times...')\n",
    "        start = time.time()\n",
    "        all_units_info = []\n",
    "        unit_allspiketimes = {}\n",
    "        unit_eventspikestimes = {'event_window': event_window, 'event_spikes': {}, 'event_bursts': {}}\n",
    "        for probei, pdata in probe_data.items():\n",
    "            for unitind, uniti in enumerate(pdata['units']):\n",
    "                unit_name = probei[-1] + str(uniti)\n",
    "                spikesi = np.squeeze(pdata['spike_times'][pdata['spike_clusters'] == uniti])\n",
    "                if spikesi.size < 50:\n",
    "                    continue\n",
    "                unit_allspiketimes[unit_name] = {}\n",
    "                \n",
    "                ## Gather unit info ##\n",
    "                if 'areas' in pdata.keys():\n",
    "                    unit_region = pdata['areas'][unitind]\n",
    "                    unit_coords = [\n",
    "                        int(x) for x in pdata['CCF_coords'][unitind].replace('[','').replace(']','').replace(' ','').split(',')\n",
    "                    ]\n",
    "                else:\n",
    "                    unit_region = 'none'\n",
    "                    unit_coords = [-1, -1, -1]\n",
    "                all_units_info.append([\n",
    "                    unit_name, probei, pdata['chs'][unitind], pdata['ch_depths'][pdata['chs'][unitind]],\n",
    "                    pdata['duration'][unitind], unit_region, unit_coords[0], unit_coords[1], unit_coords[2]\n",
    "                ])\n",
    "\n",
    "                ## Get all and event spike times ##\n",
    "                unit_allspiketimes[unit_name]['spikes'] = spikesi\n",
    "                burstsi, burst_counts = find_bursts_indunit(spikesi)\n",
    "                unit_allspiketimes[unit_name]['bursts'] = burstsi\n",
    "                unit_allspiketimes[unit_name]['burst_counts'] = burst_counts\n",
    "                event_raster = []\n",
    "                burst_raster = []\n",
    "                burst_count_raster = []\n",
    "                for eventi in all_event_times:\n",
    "                    spikeinds = np.nonzero((spikesi >= eventi + event_window[0]) & (spikesi <= eventi + event_window[1]))[0]\n",
    "                    event_raster.append(spikesi[spikeinds] - eventi)\n",
    "                    burstinds = np.nonzero((burstsi >= eventi + event_window[0]) & (burstsi <= eventi + event_window[1]))[0]\n",
    "                    burst_raster.append(burstsi[burstinds] - eventi)\n",
    "                    burst_count_raster.append(burst_counts[burstinds])\n",
    "                unit_eventspikestimes['event_spikes'][unit_name] = event_raster\n",
    "                unit_eventspikestimes['event_bursts'][unit_name] = {'times': burst_raster, 'counts': burst_count_raster}\n",
    "\n",
    "        ## Save the data files to mouse's recordingX\\evoked_data folder ##\n",
    "        all_units_info_df = pd.DataFrame(\n",
    "            all_units_info,\n",
    "            columns=['unit_id', 'probe', 'peak_ch', 'depth', 'spike_duration', 'region', 'CCF_AP', 'CCF_DV', 'CCF_ML']\n",
    "        )\n",
    "        \n",
    "        ## Add parent region column ##\n",
    "        if len(np.unique(all_units_info_df['region'].values.astype(str))) > 1:\n",
    "            print('  Adding parent region...')\n",
    "#             sub_CCF_res = subject_df[subject_df['mouse'] == mouse]['CCF_res'].iloc[0]\n",
    "            sub_CCF_res = 25\n",
    "            mcc = MouseConnectivityCache(resolution=sub_CCF_res)\n",
    "            str_tree = mcc.get_structure_tree()\n",
    "            annot, annot_info = mcc.get_annotation_volume()\n",
    "            all_units_info_df = add_parent_region_to_df(all_units_info_df, str_tree, annot)\n",
    "        \n",
    "        all_units_info_df.to_csv(unit_info_file, index=False)\n",
    "        pickle.dump(unit_allspiketimes, open(unit_allspiketimes_file, 'wb'))\n",
    "        pickle.dump(unit_eventspikestimes, open(unit_eventspikes_file, 'wb'))\n",
    "\n",
    "        end = time.time()\n",
    "        print('  Time to get unit spike times and save: {:.2f} min\\n'.format((end-start)/60))\n",
    "        ## After each subject, delete common variables ##\n",
    "        del stim_log, probe_data, all_event_times, all_units_info, all_units_info_df, unit_allspiketimes, unit_eventspikestimes"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "len(np.unique(all_units_info_df['region'].values.astype(str)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Old subjects that have been processed"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "subjects = {\n",
    "    '631037': {\n",
    "        'awake-saline': r'F:\\psi_exp\\mouse631037\\estim_2022-12-06_09-54-04\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse631037\\urethane_2022-12-07_10-34-51\\experiment1\\recording1',\n",
    "    },\n",
    "    '654182': {\n",
    "        'awake-saline': r'F:\\psi_exp\\mouse654182\\estim_vis_2022-12-01_10-33-50\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse654182\\urethane_vis_2022-12-02_11-02-25\\experiment1\\recording1',\n",
    "    },\n",
    "    '655956': {\n",
    "        'awake-saline': r'F:\\psi_exp\\mouse655956\\estim_2022-12-15_10-07-59\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse655956\\urethane_2022-12-16_10-45-18\\experiment1\\recording1',\n",
    "    },\n",
    "    '657903': {\n",
    "        'awake-psilocybin': r'F:\\psi_exp\\mouse657903\\pilot_aw_psi_2023-01-13_12-18-22\\experiment1\\recording1',\n",
    "    },\n",
    "    '666193': {\n",
    "        'saline': r'F:\\psi_exp\\mouse666193\\pilot_aw_2023-02-15_11-44-11\\experiment1\\recording1',\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse666193\\pilot_aw_psi_2023-02-16_10-55-48\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse666193\\pilot_ur_2023-02-17_12-22-51\\experiment1\\recording1',\n",
    "    }, \n",
    "    '666194': {\n",
    "        'saline': r'F:\\psi_exp\\mouse666194\\pilot_aw_2023-02-22_12-32-58\\experiment1\\recording1',\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse666194\\pilot_aw_psi_2023-02-23_10-40-34\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse666194\\pliot_ur_2023-02-24_11-19-43\\experiment1\\recording1',\n",
    "    },\n",
    "    '666196': {\n",
    "        'saline': r'F:\\psi_exp\\mouse666196\\pilot_aw_2023-03-15_12-29-06\\experiment1\\recording1',\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse666196\\pilot_aw_psi_2023-03-16_10-21-29\\experiment1\\recording1',\n",
    "    },\n",
    "    '669118': {\n",
    "        'saline': r'F:\\psi_exp\\mouse669118\\pilot_aw_2023-03-23_12-14-39\\experiment1\\recording1',\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse669118\\pilot_aw_psi_2023-03-24_09-55-33\\experiment1\\recording1',\n",
    "    },\n",
    "    '669117': {\n",
    "        'saline': r'F:\\psi_exp\\mouse669117\\pilot_aw_2023-03-29_11-09-15\\experiment1\\recording1',\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse669117\\pilot_aw_psi_2023-03-30_11-37-07\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse669117\\pilot_ur_2023-03-31_11-51-53\\experiment1\\recording1',\n",
    "    },\n",
    "    '673449': {\n",
    "        'psilocybin1': r'F:\\psi_exp\\mouse673449\\aw_psi_2023-04-19_11-23-26\\experiment1\\recording1',\n",
    "        'psilocybin2': r'F:\\psi_exp\\mouse673449\\aw_psi_d2_2023-04-20_10-05-31\\experiment1\\recording1',\n",
    "        'saline': r'F:\\psi_exp\\mouse673449\\aw_2023-04-21_09-28-23\\experiment1\\recording1',\n",
    "    },\n",
    "    '676726': {\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse676726\\aw_psi_2023-05-03_11-08-22\\experiment1\\recording1',\n",
    "        'awake-iso': r'F:\\psi_exp\\mouse676726\\aw_iso_2023-05-04_11-02-16\\experiment1\\recording1',\n",
    "    },\n",
    "    '676727': {\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse676727\\aw_psi_2023-05-10_09-49-12\\experiment1\\recording1',\n",
    "        'awake-iso': r'F:\\psi_exp\\mouse676727\\aw_iso_2023-05-11_09-44-46\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse676727\\urethane_2023-05-12_11-35-38\\experiment1\\recording1',\n",
    "    },\n",
    "    '678912': {\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse678912\\spont_aw_psi_2023-06-22_11-42-00\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse678912\\urethane_2023-06-23_11-08-17\\experiment1\\recording1',\n",
    "    },\n",
    "    '678913': {\n",
    "        'psilocybin': r'F:\\psi_exp\\mouse678913\\spont_aw_psi_2023-06-29_12-49-40\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse678913\\urethane_2023-06-30_10-56-33\\experiment1\\recording1',\n",
    "    },\n",
    "    ## awake-iso ##\n",
    "    '551397': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse551397\\estim_vis_2021-02-11_10-45-23\\experiment1\\recording1',\n",
    "    },\n",
    "    '551399': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse551399\\estim_2021-01-28_13-59-09\\experiment1\\recording1',\n",
    "    },\n",
    "    '569062': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse569062\\estim_vis_2021-02-18_11-17-51\\experiment1\\recording1',\n",
    "    },\n",
    "    '569064': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse569064\\estim_vis_2021-04-08_10-28-24\\experiment1\\recording1',\n",
    "    },\n",
    "    '569068': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse569068\\estim_vis_2021-03-04_10-51-38\\experiment1\\recording1',\n",
    "    },\n",
    "    '569069': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse569069\\estim_vis2_2021-03-12_10-52-44\\experiment1\\recording1',\n",
    "    },\n",
    "    '569070': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse569070\\estim1_2021-04-01_10-27-33\\experiment1\\recording1',\n",
    "    },\n",
    "    '569073': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse569073\\estim_vis_2021-04-15_10-27-22\\experiment1\\recording1',\n",
    "    },\n",
    "    '571619': {\n",
    "        'awake-iso': r'F:\\ZZmanuscript_eLife\\mouse571619\\estim2_2021-03-19_10-09-01\\experiment1\\recording1',\n",
    "    },\n",
    "    '635397': {\n",
    "        'awake-iso': r'F:\\psi_exp\\mouse635397\\estim_vis_2022-08-18_12-08-15\\experiment1\\recording1',\n",
    "    },\n",
    "    '654181': {\n",
    "        'awake-iso': r'F:\\psi_exp\\mouse654181\\estim_vis_2022-11-22_09-42-58\\experiment1\\recording1',\n",
    "        'urethane': r'F:\\psi_exp\\mouse654181\\urethane_vis_2022-11-23_08-30-16\\experiment1\\recording1',\n",
    "    },\n",
    "    ## urethane ##\n",
    "    '638703': {\n",
    "        'urethane': r'F:\\psi_exp\\mouse638703\\urethane_estim_2022-10-14_12-25-20\\experiment1\\recording1',\n",
    "    },\n",
    "    '655955': {\n",
    "        'urethane': r'F:\\psi_exp\\mouse655955\\urethane_2022-12-14_10-38-00\\experiment1\\recording1',\n",
    "    },\n",
    "    '582386': {\n",
    "        'urethane': r'F:\\psi_exp\\mouse582386\\urethane_2021-07-15_11-36-58\\experiment1\\recording1',\n",
    "    },\n",
    "}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tbd_eeg",
   "language": "python",
   "name": "tbd_eeg"
  },
  "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}