NPX_based_analyses/extract_unit_event_spikes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### This script loads subjects with probes, parses each unit's spike times (saves as a pkl file), and collects each unit's event-related spike times (for all events) and saves that as a pkl file, too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "import time\n",
    "import gspread\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy import signal, fftpack, stats, ndimage\n",
    "import statsmodels.stats.multitest as multitest"
   ]
  },
  {
   "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,\n",
    "    get_evoked_traces,\n",
    "    get_evoked_firing_rates,\n",
    "    find_nearest_ind\n",
    ")\n",
    "from allensdk.brain_observatory.ecephys.lfp_subsampling.subsampling import remove_lfp_offset\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": [
    "Load CCF for identifying cortical areas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "mcc = MouseConnectivityCache(resolution=10)\n",
    "str_tree = mcc.get_structure_tree()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load Zap_Zip-log_exp to get metadata for experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "_gc = gspread.service_account() # need a key file to access the account\n",
    "_sh = _gc.open('Zap_Zip-log_exp') # open the spreadsheet\n",
    "_df = pd.DataFrame(_sh.sheet1.get()) # load the first worksheet\n",
    "zzmetadata = _df.T.set_index(0).T # put it in a nicely formatted dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define areas of interest to plot population activity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "areas_of_interest = {\n",
    "    'MO': [\n",
    "        'MOp1', 'MOp2/3', 'MOp5', 'MOp6a', 'MOp6b',\n",
    "        'MOs1', 'MOs2/3', 'MOs5', 'MOs6a', 'MOs6b'\n",
    "    ],\n",
    "    'ACA': [\n",
    "        'ACAd1', 'ACAd2/3', 'ACAd5', 'ACAd6a', 'ACAd6b',\n",
    "        'ACAv1', 'ACAv2/3', 'ACAv5', 'ACAv6a', 'ACAv6b'\n",
    "    ],\n",
    "    'SS': [\n",
    "        'SSp-bfd1', 'SSp-bfd2/3', 'SSp-bfd4', 'SSp-bfd5', 'SSp-bfd6a', 'SSp-bfd6b',\n",
    "        'SSp-ll1', 'SSp-ll2/3', 'SSp-ll4', 'SSp-ll5', 'SSp-ll6a', 'SSp-ll6b',\n",
    "        'SSp-tr1', 'SSp-tr2/3', 'SSp-tr4', 'SSp-tr5', 'SSp-tr6a', 'SSp-tr6b'\n",
    "    ],\n",
    "    'VIS': [\n",
    "        'VISp1', 'VISp2/3', 'VISp4', 'VISp5', 'VISp6a', 'VISp6b',\n",
    "        'VISam1', 'VISam2/3', 'VISam4', 'VISam5', 'VISam6a', 'VISam6b',\n",
    "        'VISpm1', 'VISpm2/3', 'VISpm4', 'VISpm5', 'VISpm6a', 'VISpm6b',\n",
    "        'VISrl1', 'VISrl2/3', 'VISrl4', 'VISrl5', 'VISrl6a', 'VISrl6b',\n",
    "    ],\n",
    "    'MO-TH': [\n",
    "        'AV', 'CL', 'MD', 'PO', 'RT', 'VAL', 'VPL', 'VPM', 'VM' # this is the default\n",
    "#         'AV', 'CL', 'MD', 'PO', 'PF', 'VAL', 'VPL', 'VPM', 'VM' # consider removing RT for future\n",
    "    ],\n",
    "}\n",
    "\n",
    "area_colors = {\n",
    "    'MO': 'blue',\n",
    "    'ACA': 'deepskyblue',\n",
    "    'SS': 'blueviolet',\n",
    "    'VIS': 'green',\n",
    "    'MO-TH': 'steelblue',\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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(unit_ids, all_spikes_dict):\n",
    "    burst_info = {}\n",
    "    for uid in unit_ids:\n",
    "        alluspikes = all_spikes_dict[uid]\n",
    "        preISIs = np.diff(alluspikes)[:-1]\n",
    "        postISIs = np.diff(alluspikes)[1:]\n",
    "        ## Find starts ##\n",
    "        bs_inds = np.nonzero((preISIs > 0.1) * (postISIs < 0.004))[0]\n",
    "        if len(bs_inds) == 0:\n",
    "            continue\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",
    "        ## Store burst info ##\n",
    "        burst_info[uid] = {}\n",
    "        burst_info[uid]['start_times'] = alluspikes[burst_starts]\n",
    "        burst_info[uid]['burst_spike_counts'] = np.array(burst_counts)\n",
    "    return burst_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": "markdown",
   "metadata": {},
   "source": [
    "### Load subjects from file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(r'C:\\Users\\lesliec\\OneDrive - Allen Institute\\data\\all_EEG_subjects.json') as subjects_file:\n",
    "    multi_sub_dict = json.load(subjects_file)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "## Load fewer subjects for testing script ##\n",
    "with open(r'C:\\Users\\lesliec\\OneDrive - Allen Institute\\data\\TEST_EEG_subjects.json') as subjects_file:\n",
    "    multi_sub_dict = json.load(subjects_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load EEG_exp and gather units' spike times"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "overwrite_existing_files = False\n",
    "\n",
    "event_window = [-2.0, 2.0] # find event-related spikes within this window (s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "521885: estim1\n",
      " This experiment has no probe data, not making spike times files.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "521886: estim1\n",
      " This experiment has no probe data, not making spike times files.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "521887: estim1\n",
      " This experiment has no probe data, not making spike times files.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "543393: estim1\n",
      "  Getting probe info...\n",
      "  probeB does not have area assignments, not processing.\n",
      "  NO area assignments for any probes, not analyzing.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "543394: estim1\n",
      "  Getting probe info...\n",
      "  probeB does not have area assignments, not processing.\n",
      "  NO area assignments for any probes, not analyzing.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "543395: estim1\n",
      "  Getting probe info...\n",
      "  probeB does not have area assignments, not processing.\n",
      "  NO area assignments for any probes, not analyzing.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "543396: estim_vis1\n",
      "  F:\\EEG_exp\\mouse543396\\estim_vis1_2020-09-18_12-04-46\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "546655: estim_vis\n",
      "  F:\\EEG_exp\\mouse546655\\estim_vis_2020-10-23_11-01-14\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "551400: estim_vis\n",
      "  F:\\EEG_exp\\mouse551400\\estim_vis_2021-01-22_11-07-12\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "551399: estim\n",
      "  F:\\EEG_exp\\mouse551399\\estim_2021-01-28_13-59-09\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "Body camera file not found.\n",
      "Pupil camera file not found.\n",
      "551397: estim_vis\n",
      "  F:\\EEG_exp\\mouse551397\\estim_vis_2021-02-11_10-45-23\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "569062: estim_vis\n",
      "  F:\\EEG_exp\\mouse569062\\estim_vis_2021-02-18_11-17-51\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "569068: estim_vis\n",
      "  F:\\EEG_exp\\mouse569068\\estim_vis_2021-03-04_10-51-38\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "569069: estim_vis2\n",
      "  F:\\EEG_exp\\mouse569069\\estim_vis2_2021-03-12_10-52-44\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "569064: estim_vis\n",
      "  F:\\EEG_exp\\mouse569064\\estim_vis_2021-04-08_10-28-24\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "569073a: estim_vis\n",
      "  F:\\EEG_exp\\mouse569073\\estim_vis_2021-04-15_10-27-22\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "569073b: estim\n",
      "  F:\\EEG_exp\\mouse569073\\estim_2021-04-16_10-42-44\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "571618: estim1\n",
      "  F:\\EEG_exp\\mouse571618\\estim1_2021-04-29_12-28-54\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "575102: estim_vis\n",
      "  F:\\EEG_exp\\mouse575102\\estim_vis_2021-06-03_11-25-01\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "575100: estim_vis\n",
      "  F:\\EEG_exp\\mouse575100\\estim_vis_2021-06-10_11-39-06\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "571619a: estim1\n",
      "  F:\\EEG_exp\\mouse571619\\estim1_2021-03-18_11-41-26\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "571619b: estim2\n",
      "  F:\\EEG_exp\\mouse571619\\estim2_2021-03-19_10-09-01\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "569070a: estim1\n",
      "  Getting probe info...\n",
      "  probeB does not have area assignments, not processing.\n",
      "  probeC does not have area assignments, not processing.\n",
      "  probeF does not have area assignments, not processing.\n",
      "  NO area assignments for any probes, not analyzing.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "569070b: estim2\n",
      "  Getting probe info...\n",
      "  probeB does not have area assignments, not processing.\n",
      "  probeC does not have area assignments, not processing.\n",
      "  probeF does not have area assignments, not processing.\n",
      "  NO area assignments for any probes, not analyzing.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "569072: estim_vis\n",
      "  F:\\EEG_exp\\mouse569072\\estim_vis_2021-04-22_10-26-58\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "571620: estim_vis\n",
      "  F:\\EEG_exp\\mouse571620\\estim_vis_2021-05-13_11-33-47\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "582387: estim_vis\n",
      "  F:\\EEG_exp\\mouse582387\\estim_vis_2021-07-22_09-59-47\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "586466: estim_vis\n",
      "  F:\\EEG_exp\\mouse586466\\estim_vis_2021-07-29_10-37-41\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "586468a: estim_vis\n",
      "  F:\\EEG_exp\\mouse586468\\estim_vis_2021-08-19_11-36-53\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "586468b: estim\n",
      "  F:\\EEG_exp\\mouse586468\\estim_2021-08-20_11-34-54\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical stimulation\n",
      "SomnoSuite log file not found.\n",
      "590479: estim\n",
      "  F:\\EEG_exp\\mouse590479\\estim_2021-09-21_17-24-07\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "590480: estim_vis\n",
      "  F:\\EEG_exp\\mouse590480\\estim_vis_2021-11-11_10-54-19\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "590481: estim_vis\n",
      "  F:\\EEG_exp\\mouse590481\\estim_vis_2021-11-18_10-52-12\\experiment1\\recording1\\evoked_data\\all_units_info.csv already exists, skipping analysis.\n",
      "\n",
      "Experiment type: electrical and sensory stimulation\n",
      "SomnoSuite log file not found.\n",
      "599017: estim_vis\n",
      "  Getting probe info...\n",
      "  Getting spike times...\n",
      "  Time to get unit spike times and save: 3.47 min\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for mouse_num, rec_folder in multi_sub_dict.items():\n",
    "    exp = EEGexp(rec_folder, preprocess=False, make_stim_csv=False)\n",
    "    exp_tag = exp.experiment_folder[exp.experiment_folder.find('mouse')+12:exp.experiment_folder.find(str(exp.date.year))-1]\n",
    "    print('{}: {}'.format(mouse_num, exp_tag))\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_locs = np.ones((len(probe_list)), dtype=bool)\n",
    "    probe_data = {}\n",
    "    for pbi, probei in enumerate(probe_list):\n",
    "        ## Load probe_info.json ##\n",
    "        with open(exp.ephys_params[probei]['probe_info']) as data_file:\n",
    "            data = json.load(data_file)\n",
    "        if 'area_ch' in data.keys():\n",
    "            probe_data[probei] = {}\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",
    "            unit_metrics['cell_type'] = unit_metrics['duration'].apply(lambda x: 'FS' if x <= 0.4 else 'RS')\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]['cell_type'] = unit_metrics['cell_type'].values[np.squeeze(np.argsort(peak_chs))]\n",
    "            probe_data[probei]['areas'] = unit_metrics.area.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",
    "        else:\n",
    "            print('  {} does not have area assignments, not processing.'.format(probei))\n",
    "            probe_locs[pbi] = False\n",
    "\n",
    "    if probe_locs.any():\n",
    "        pass\n",
    "    else:\n",
    "        print('  NO area assignments for any probes, not analyzing.\\n')\n",
    "        continue\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",
    "            unit_allspiketimes[unit_name] = {}\n",
    "            ## Gather unit info ##\n",
    "            unit_info = [unit_name] ## get unit_id ##\n",
    "            unit_region = pdata['areas'][unitind]\n",
    "            unit_info.extend([pdata['cell_type'][unitind], pdata['ch_depths'][pdata['chs'][unitind]], unit_region])\n",
    "            parent_region = [key for key in list(areas_of_interest.keys()) if unit_region in areas_of_interest[key]]\n",
    "            if len(parent_region) == 1:\n",
    "                unit_info.append(parent_region[0])\n",
    "            else:\n",
    "                unit_info.append('notROI')\n",
    "            all_units_info.append(unit_info)\n",
    "            \n",
    "            ## Get all and event spike times ##\n",
    "            spikesi = np.squeeze(pdata['spike_times'][pdata['spike_clusters'] == uniti])\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, columns=['unit_id', 'cell_type', 'depth', 'region', 'parent'])\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 2
}