{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import spatial_maps as sp\n",
    "import septum_mec.analysis.data_processing as dp\n",
    "import septum_mec.analysis.registration\n",
    "import expipe\n",
    "import os\n",
    "import pathlib\n",
    "import scipy\n",
    "import scipy.signal\n",
    "import numpy as np\n",
    "import exdir\n",
    "import pandas as pd\n",
    "import optogenetics as og\n",
    "import quantities as pq\n",
    "import shutil\n",
    "from distutils.dir_util import copy_tree\n",
    "import elephant as el\n",
    "import neo\n",
    "from scipy.signal import find_peaks\n",
    "from scipy.interpolate import interp1d\n",
    "from matplotlib import mlab\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from tqdm._tqdm_notebook import tqdm_notebook\n",
    "tqdm_notebook.pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_loader = dp.Data()\n",
    "actions = data_loader.actions\n",
    "project = data_loader.project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = pathlib.Path('output/stimulus-spike-lfp-response')\n",
    "(output / 'data').mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "identify_neurons = actions['identify-neurons']\n",
    "# sessions = pd.read_csv(identify_neurons.data_path('sessions'))\n",
    "units = pd.read_csv(identify_neurons.data_path('units'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lim(action_id):\n",
    "    stim_times = data_loader.stim_times(action_id)\n",
    "    if stim_times is None:\n",
    "        return [0, np.inf]\n",
    "    stim_times = np.array(stim_times)\n",
    "    return [stim_times.min(), stim_times.max()]\n",
    "\n",
    "def get_mask(lfp, lim):\n",
    "    return (lfp.times >= lim[0]) & (lfp.times <= lim[1])\n",
    "\n",
    "def zscore(a):\n",
    "    return (a - a.mean(0)) / a.std(0)\n",
    "\n",
    "def compute_stim_freq(action_id):\n",
    "    stim_times = data_loader.stim_times(action_id)\n",
    "    if stim_times is None:\n",
    "        return np.nan\n",
    "    stim_times = np.array(stim_times)\n",
    "    return 1 / np.mean(np.diff(stim_times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def signaltonoise(a, axis=0, ddof=0):\n",
    "    a = np.asanyarray(a)\n",
    "    m = a.mean(axis)\n",
    "    sd = a.std(axis=axis, ddof=ddof)\n",
    "    return np.where(sd == 0, 0, m / sd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_energy(p, f, f1, f2):\n",
    "    if np.isnan(f1) or np.all(np.isnan(p)):\n",
    "        return np.nan\n",
    "    mask = (f > f1) & (f < f2)\n",
    "    df = f[1] - f[0]\n",
    "    return np.sum(p[mask]) * df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_theta_peak(p, f, f1, f2):\n",
    "    if np.all(np.isnan(p)):\n",
    "        return np.nan, np.nan\n",
    "    mask = (f > f1) & (f < f2)\n",
    "    p_m = p[mask]\n",
    "    f_m = f[mask]\n",
    "    peaks, _ = find_peaks(p_m)\n",
    "    idx = np.argmax(p_m[peaks])\n",
    "    return f_m[peaks[idx]], p_m[peaks[idx]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_half_width(p, f, m_p, m_f):\n",
    "    if np.isnan(m_p):\n",
    "        return np.nan, np.nan\n",
    "    m_p_half = m_p / 2\n",
    "    half_p = p - m_p_half\n",
    "    idx_f = np.where(f <= m_f)[0].max()\n",
    "    idxs_p1, = np.where(np.diff(half_p[:idx_f + 1] > 0) == 1)\n",
    "    if len(idxs_p1) == 0:\n",
    "        return np.nan, np.nan\n",
    "    m1 = idxs_p1.max()\n",
    "    idxs_p2, = np.where(np.diff(half_p[idx_f:] > 0) == 1)\n",
    "    m2 = idxs_p2.min() + idx_f\n",
    "    assert p[m1] < m_p_half < p[m1+1], (p[m1], m_p_half, p[m1+1])\n",
    "    assert p[m2] > m_p_half > p[m2+1], (p[m2], m_p_half, p[m2+1])\n",
    "    \n",
    "    f1 = interp1d([half_p[m1], half_p[m1 + 1]], [f[m1], f[m1 + 1]])(0)\n",
    "    f2 = interp1d([half_p[m2], half_p[m2 + 1]], [f[m2], f[m2 + 1]])(0)\n",
    "    return f1, f2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# p = np.load('debug_p.npy')\n",
    "# f = np.load('debug_f.npy')\n",
    "# compute_half_width(p, f, 0.01038941, 30.30187709636872)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.plot(f, p)\n",
    "# plt.xlim(29.9,30.6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_stim_peak(p, f, s_f):\n",
    "    if np.isnan(s_f):\n",
    "        return np.nan\n",
    "    return interp1d(f, p)(s_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_spike_lfp_coherence(anas, sptr, NFFT):\n",
    "\n",
    "    sigs, freqs = el.sta.spike_field_coherence(anas, sptr, **{'nperseg': NFFT})\n",
    "    return sigs, freqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def butter_bandpass(lowcut, highcut, fs, order=5):\n",
    "    nyq = 0.5 * fs\n",
    "    low = lowcut / nyq\n",
    "    high = highcut / nyq\n",
    "    b, a = scipy.signal.butter(order, [low, high], btype='band')\n",
    "    return b, a\n",
    "\n",
    "\n",
    "def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):\n",
    "    b, a = butter_bandpass(lowcut, highcut, fs, order=order)\n",
    "    y = scipy.signal.filtfilt(b, a, data)\n",
    "    return y\n",
    "\n",
    "def compute_spike_phase_func(lfp, times, return_degrees=False):\n",
    "    x_a = hilbert(lfp)\n",
    "    x_phase = np.angle(x_a)\n",
    "    if return_degrees:\n",
    "        x_phase = x_phase * 180 / np.pi\n",
    "    return interp1d(times, x_phase)\n",
    "\n",
    "\n",
    "def vonmises_kde(data, kappa=100, n_bins=100):\n",
    "    from scipy.special import i0\n",
    "    bins = np.linspace(-np.pi, np.pi, n_bins)\n",
    "    x = np.linspace(-np.pi, np.pi, n_bins)\n",
    "    # integrate vonmises kernels\n",
    "    kde = np.exp(kappa * np.cos(x[:, None] - data[None, :])).sum(1) / (2 * np.pi * i0(kappa))\n",
    "    kde /= np.trapz(kde, x=bins)\n",
    "    return bins, kde\n",
    "\n",
    "\n",
    "def spike_phase_score(phase_bins, density):\n",
    "    import math\n",
    "    import pycircstat as pc\n",
    "    ang = pc.mean(phase_bins, w=density)\n",
    "    vec_len = pc.resultant_vector_length(phase_bins, w=density)\n",
    "    # ci_lim = pc.mean_ci_limits(head_angle_bins, w=rate)\n",
    "    return ang, vec_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_clean_lfp(anas, width=500, threshold=2):\n",
    "    anas = np.array(anas)\n",
    "    idxs, = np.where(abs(anas) > threshold)\n",
    "    for idx in idxs:\n",
    "        anas[idx-width:idx+width] = 0 # TODO AR model prediction\n",
    "    return anas, idxs\n",
    "\n",
    "\n",
    "def compute_clean_spikes(spikes, idxs, times, width=500):\n",
    "\n",
    "    for idx in idxs:\n",
    "        t0 = times[idx-width]\n",
    "        stop = idx + width\n",
    "        if stop > len(times) - 1:\n",
    "            stop = len(times) - 1 \n",
    "        t1 = times[stop]\n",
    "        mask = (spikes > t0) & (spikes < t1)\n",
    "        spikes = spikes[~mask]\n",
    "    spikes = spikes[spikes <= times[-1]]\n",
    "    return spikes\n",
    "\n",
    "\n",
    "def prepare_spike_lfp(anas, sptr, t_start, t_stop):\n",
    "\n",
    "    t_start = t_start * pq.s if t_start is not None else 0 * pq.s\n",
    "    sampling_rate = anas.sampling_rate\n",
    "    units = anas.units\n",
    "    times = anas.times\n",
    "    if t_start is not None and t_stop is not None:\n",
    "        t_stop = t_stop * pq.s\n",
    "        mask = (times > t_start) & (times < t_stop)\n",
    "        anas = np.array(anas)[mask,:]\n",
    "        times = times[mask]\n",
    "\n",
    "    best_channel = np.argmax(signaltonoise(anas))\n",
    "#     best_channel = np.random.choice(anas.shape[1])\n",
    "    \n",
    "    cleaned_anas, idxs = compute_clean_lfp(anas[:, best_channel])\n",
    "    \n",
    "    cleaned_anas = neo.AnalogSignal(\n",
    "        signal=cleaned_anas * units, sampling_rate=sampling_rate, t_start=t_start\n",
    "    )\n",
    "    \n",
    "    spike_units = sptr.units\n",
    "    spike_times = sptr.times\n",
    "    spike_times = compute_clean_spikes(spike_times, idxs, times)\n",
    "\n",
    "    sptr = neo.SpikeTrain(\n",
    "        spike_times[(spike_times > t_start) & (spike_times < times[-1])], units=spike_units,\n",
    "        t_start=t_start, t_stop=times[-1]\n",
    "    )\n",
    "\n",
    "    return cleaned_anas, sptr, best_channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_spike_phase_func(lfp, times, return_degrees=False):\n",
    "    from scipy.fftpack import next_fast_len\n",
    "    x_a = scipy.signal.hilbert(\n",
    "        lfp, next_fast_len(len(lfp)))[:len(lfp)]\n",
    "#     x_a = hilbert(lfp)\n",
    "    x_phase = np.angle(x_a, deg=return_degrees)\n",
    "    return interp1d(times, x_phase)\n",
    "\n",
    "\n",
    "def compute_spike_phase(lfp, spikes, flim=[6,10]):\n",
    "    \n",
    "    sample_rate = lfp.sampling_rate.magnitude\n",
    "    \n",
    "    # sometimes the position is recorded after LFP recording is ended\n",
    "    times = np.arange(lfp.shape[0]) / sample_rate\n",
    "     \n",
    "    spikes = np.array(spikes)\n",
    "    spikes = spikes[(spikes > times.min()) & (spikes < times.max())]\n",
    "    \n",
    "    filtered_lfp = butter_bandpass_filter(\n",
    "        lfp.magnitude.ravel(), *flim, fs=sample_rate, order=3)\n",
    "\n",
    "    spike_phase_func = compute_spike_phase_func(filtered_lfp, times)\n",
    "    \n",
    "    return spike_phase_func(spikes), filtered_lfp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.figure(figsize=(16,9))\n",
    "# # lfp = data_loader.lfp('1833-200619-2', 6)\n",
    "# lfp = data_loader.lfp('1834-220319-3', 6)\n",
    "# # lfp = data_loader.lfp('1849-010319-4', 6)\n",
    "# times = np.arange(lfp.shape[0]) / lfp.sampling_rate.magnitude\n",
    "# clean_lfp, _ = compute_clean_lfp(lfp.magnitude[:, 0], threshold=2)\n",
    "# plt.plot(times,lfp[:,0])\n",
    "# plt.plot(times,clean_lfp)\n",
    "\n",
    "# plt.figure(figsize=(16,9))\n",
    "# plt.psd(lfp[:,0].ravel(), Fs=1000, NFFT=10000)\n",
    "# plt.psd(clean_lfp, Fs=1000, NFFT=10000)\n",
    "# plt.xlim(0,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.figure(figsize=(16,9))\n",
    "\n",
    "# plt.plot(times,lfp[:,0])\n",
    "# # plt.plot(clean_lfp*100)\n",
    "# plt.plot(times[:-1], np.diff(lfp[:,0].magnitude.ravel()))\n",
    "# plt.xlim(64.5,65.5)\n",
    "# # plt.ylim(-250,250)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # action_id_0, channel_0, unit_0 = '1833-200619-1', 6, 163\n",
    "# # action_id_1, channel_1, unit_1 = '1833-200619-2', 6, 28\n",
    "# action_id_0, channel_0, unit_0 = '1834-220319-3', 2, 46\n",
    "# action_id_1, channel_1, unit_1 = '1834-220319-4', 2, 60\n",
    "# lfp_0 = data_loader.lfp(action_id_0, channel_0)\n",
    "# lfp_1 = data_loader.lfp(action_id_1, channel_1)\n",
    "\n",
    "# sample_rate_0 = lfp_0.sampling_rate\n",
    "# sample_rate_1 = lfp_1.sampling_rate\n",
    "\n",
    "# lim_0 = get_lim(action_id_0)\n",
    "# lim_1 = get_lim(action_id_1)\n",
    "\n",
    "# sptrs_0 = data_loader.spike_trains(action_id_0, channel_0)\n",
    "\n",
    "# sptrs_1 = data_loader.spike_trains(action_id_1, channel_1)\n",
    "\n",
    "# cleaned_lfps_0, sptr_0, best_channel_0 = prepare_spike_lfp(lfp_0, sptrs_0[unit_0], *lim_0)\n",
    "\n",
    "# cleaned_lfps_1, sptr_1, best_channel_1 = prepare_spike_lfp(lfp_1, sptrs_1[unit_1], *lim_1)\n",
    "\n",
    "# coher_0, freq_0 = compute_spike_lfp_coherence(cleaned_lfps_0, sptr_0, 4096)\n",
    "\n",
    "# coher_1, freq_1 = compute_spike_lfp_coherence(cleaned_lfps_1, sptr_1, 4096)\n",
    "\n",
    "# spike_phase_0, filtered_lfp_0 = compute_spike_phase(cleaned_lfps_0, sptrs_0[unit_0], flim=[6,10])\n",
    "\n",
    "# spike_phase_1, filtered_lfp_1 = compute_spike_phase(cleaned_lfps_1, sptrs_1[unit_1], flim=[6,10])\n",
    "\n",
    "# spike_phase_1_stim, filtered_lfp_1_stim = compute_spike_phase(cleaned_lfps_1, sptrs_1[unit_1], flim=[10.5,11.5])\n",
    "\n",
    "# plt.figure()\n",
    "# plt.plot(freq_0, coher_0.ravel())\n",
    "# plt.plot(freq_1, coher_1.ravel())\n",
    "# plt.xlim(0,20)\n",
    "\n",
    "# plt.figure()\n",
    "# bins_0, kde_0 = vonmises_kde(spike_phase_0, 100)\n",
    "# ang_0, vec_len_0 = spike_phase_score(bins_0, kde_0)\n",
    "# plt.polar(bins_0, kde_0, color='b')\n",
    "# plt.polar([ang_0, ang_0], [0, vec_len_0], color='b')\n",
    "\n",
    "# bins_1, kde_1 = vonmises_kde(spike_phase_1, 100)\n",
    "# ang_1, vec_len_1 = spike_phase_score(bins_1, kde_1)\n",
    "# plt.polar(bins_1, kde_1, color='r')\n",
    "# plt.polar([ang_1, ang_1], [0, vec_len_1], color='r')\n",
    "\n",
    "# bins_1_stim, kde_1_stim = vonmises_kde(spike_phase_1_stim, 100)\n",
    "# ang_1_stim, vec_len_1_stim = spike_phase_score(bins_1_stim, kde_1_stim)\n",
    "# plt.polar(bins_1_stim, kde_1_stim, color='k')\n",
    "# plt.polar([ang_1_stim, ang_1_stim], [0, vec_len_1_stim], color='k')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO fix artefact stuff from phase precession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "NFFT = 8192\n",
    "theta_band_f1, theta_band_f2 = 6, 10 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "coherence_data, freqency_data = {}, {}\n",
    "theta_kde_data, theta_bins_data = {}, {}\n",
    "stim_kde_data, stim_bins_data = {}, {}\n",
    "\n",
    "def process(row):\n",
    "    action_id = row['action']\n",
    "    channel_group = row['channel_group']\n",
    "    unit_name = row['unit_name']\n",
    "    \n",
    "    name = f'{action_id}_{channel_group}_{unit_name}'\n",
    "    \n",
    "    lfp = data_loader.lfp(action_id, channel_group) # TODO consider choosing strongest stim response\n",
    "    \n",
    "    sptr = data_loader.spike_train(action_id, channel_group, unit_name)\n",
    "       \n",
    "    lim = get_lim(action_id)\n",
    "    \n",
    "    cleaned_lfp, sptr, best_channel = prepare_spike_lfp(lfp, sptr, *lim)\n",
    "       \n",
    "    p_xys, freq = compute_spike_lfp_coherence(cleaned_lfp, sptr, NFFT=NFFT)\n",
    "    \n",
    "    p_xy = p_xys.magnitude.ravel()\n",
    "    freq = freq.magnitude\n",
    "    \n",
    "    theta_f, theta_p_max = find_theta_peak(p_xy, freq, theta_band_f1, theta_band_f2)\n",
    "    \n",
    "    theta_energy = compute_energy(p_xy, freq, theta_band_f1, theta_band_f2) # theta band 6 - 10 Hz\n",
    "    \n",
    "    theta_half_f1, theta_half_f2 = compute_half_width(p_xy, freq, theta_p_max, theta_f)\n",
    "    \n",
    "    theta_half_width = theta_half_f2 - theta_half_f1\n",
    "    \n",
    "    theta_half_energy = compute_energy(p_xy, freq, theta_half_f1, theta_half_f2) # theta band 6 - 10 Hz\n",
    "    \n",
    "    theta_spike_phase, _ = compute_spike_phase(cleaned_lfp, sptr, flim=[theta_band_f1, theta_band_f2])\n",
    "    theta_bins, theta_kde = vonmises_kde(theta_spike_phase)\n",
    "    theta_ang, theta_vec_len = spike_phase_score(theta_bins, theta_kde)\n",
    "    theta_kde_data.update({name: theta_kde})\n",
    "    theta_bins_data.update({name: theta_bins})\n",
    "\n",
    "    # stim\n",
    "    \n",
    "    stim_freq = compute_stim_freq(action_id)\n",
    "    \n",
    "    stim_p_max = compute_stim_peak(p_xy, freq, stim_freq)\n",
    "    \n",
    "    stim_half_f1, stim_half_f2 = compute_half_width(p_xy, freq, stim_p_max, stim_freq)\n",
    "    stim_half_width = stim_half_f2 - stim_half_f1\n",
    "    \n",
    "    stim_energy = compute_energy(p_xy, freq, stim_half_f1, stim_half_f2)\n",
    "    \n",
    "    if np.isnan(stim_freq):\n",
    "        stim_spike_phase, stim_bins, stim_kde, stim_ang, stim_vec_len = [np.nan] * 5\n",
    "    else:\n",
    "        stim_spike_phase, _ = compute_spike_phase(cleaned_lfp, sptr, flim=[stim_freq - .5, stim_freq + .5])\n",
    "        stim_bins, stim_kde = vonmises_kde(stim_spike_phase)\n",
    "        stim_ang, stim_vec_len = spike_phase_score(stim_bins, stim_kde)\n",
    "        stim_kde_data.update({name: stim_kde})\n",
    "        stim_bins_data.update({name: stim_bins})\n",
    "        \n",
    "    coherence_data.update({name: p_xy})\n",
    "    freqency_data.update({name: freq})\n",
    "    \n",
    "    result = pd.Series({\n",
    "        'best_channel': best_channel,\n",
    "        'theta_freq': theta_f,\n",
    "        'theta_peak': theta_p_max,\n",
    "        'theta_energy': theta_energy,\n",
    "        'theta_half_f1': theta_half_f1, \n",
    "        'theta_half_f2': theta_half_f2,\n",
    "        'theta_half_width': theta_half_width,\n",
    "        'theta_half_energy': theta_half_energy,\n",
    "        'theta_ang': theta_ang, \n",
    "        'theta_vec_len': theta_vec_len,\n",
    "        'stim_freq': stim_freq,\n",
    "        'stim_p_max': stim_p_max,\n",
    "        'stim_half_f1': stim_half_f1, \n",
    "        'stim_half_f2': stim_half_f2,\n",
    "        'stim_half_width': stim_half_width,\n",
    "        'stim_energy': stim_energy,\n",
    "        'stim_ang': stim_ang, \n",
    "        'stim_vec_len': stim_vec_len\n",
    "    })\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "199c002aa9604451a318e3f471c9d892",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1284), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mikkel/.virtualenvs/expipe/lib/python3.6/site-packages/scipy/signal/spectral.py:1577: RuntimeWarning: invalid value encountered in true_divide\n",
      "  Cxy = np.abs(Pxy)**2 / Pxx / Pyy\n",
      "/home/mikkel/.virtualenvs/expipe/lib/python3.6/site-packages/ipykernel_launcher.py:28: RuntimeWarning: invalid value encountered in true_divide\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "results = units.merge(\n",
    "    units.progress_apply(process, axis=1), \n",
    "    left_index=True, right_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coher, freqs = {}, {}\n",
    "# for i, row in tqdm(units.iterrows(), total=len(units)):\n",
    "#     action_id = row['action']\n",
    "#     channel_group = row['channel_group']\n",
    "#     unit_name = row['unit_name']\n",
    "    \n",
    "#     name = f'{action_id}_{channel_group}_{unit_name}'\n",
    "    \n",
    "#     lfp = data_loader.lfp(action_id, channel_group) # TODO consider choosing strongest stim response\n",
    "    \n",
    "#     sptr = data_loader.spike_train(action_id, channel_group, unit_name)\n",
    "       \n",
    "#     lim = get_lim(action_id)\n",
    "\n",
    "#     p_xys, freq, clean_lfp = compute_spike_lfp(lfp, sptr, *lim, NFFT=NFFT)\n",
    "    \n",
    "#     snls = signaltonoise(clean_lfp)\n",
    "#     best_channel = np.argmax(snls)\n",
    "#     snl = snls[best_channel]\n",
    "#     p_xy = p_xys[:,best_channel].magnitude\n",
    "#     freq = freq.magnitude\n",
    "    \n",
    "#     coher.update({name: p_xy})\n",
    "#     freqs.update({name: freq})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'frequency_data' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-26-64114f1ea8bd>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcoherence_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_feather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'data'\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'coherence.feather'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfrequency_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_feather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'data'\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'freqs.feather'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtheta_kde_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_feather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'data'\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'theta_kde.feather'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtheta_bins_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_feather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'data'\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'theta_bins.feather'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDataFrame\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstim_kde_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_feather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'data'\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;34m'stim_kde.feather'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'frequency_data' is not defined"
     ]
    }
   ],
   "source": [
    "pd.DataFrame(coherence_data).to_feather(output / 'data' / 'coherence.feather')\n",
    "pd.DataFrame(frequency_data).to_feather(output / 'data' / 'freqs.feather')\n",
    "pd.DataFrame(theta_kde_data).to_feather(output / 'data' / 'theta_kde.feather')\n",
    "pd.DataFrame(theta_bins_data).to_feather(output / 'data' / 'theta_bins.feather')\n",
    "pd.DataFrame(stim_kde_data).to_feather(output / 'data' / 'stim_kde.feather')\n",
    "pd.DataFrame(stim_bins_data).to_feather(output / 'data' / 'stim_bins.feather')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save to expipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action = project.require_action(\"stimulus-spike-lfp-response\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action.modules['parameters'] = {\n",
    "    'NFFT': NFFT,\n",
    "    'theta_band_f1': theta_band_f1,\n",
    "    'theta_band_f2': theta_band_f2\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "action.data['results'] = 'results.csv'\n",
    "results.to_csv(action.data_path('results'), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "copy_tree(output, str(action.data_path()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "septum_mec.analysis.registration.store_notebook(action, \"10-calculate-stimulus-spike-lfp-response.ipynb\")"
   ]
  },
  {
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}