plotting/tde-StimOffOn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Plot time delay estimates (Figure 3d).\"\"\"\n",
    "\n",
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import copy\n",
    "import pandas as pd\n",
    "from matplotlib import pyplot as plt\n",
    "import pte_stats\n",
    "import pte_decode\n",
    "from scipy.stats import sem\n",
    "from scipy.signal import find_peaks\n",
    "\n",
    "cd_path = Path(os.getcwd()).absolute().parent\n",
    "sys.path.append(os.path.join(cd_path, \"coherence\"))\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.rc('xtick', labelsize=6)\n",
    "matplotlib.rc('ytick', labelsize=6)\n",
    "matplotlib.rc('legend', fontsize=6)\n",
    "matplotlib.rc(\"font\", size=6, family=\"Arial\")\n",
    "matplotlib.rc('axes', labelsize=7)\n",
    "matplotlib.rc('axes', titlesize=7)\n",
    "matplotlib.rcParams[\"pdf.fonttype\"] = 42\n",
    "matplotlib.rcParams[\"ps.fonttype\"] = 42\n",
    "\n",
    "prop_cycle = plt.rcParams[\"axes.prop_cycle\"]\n",
    "\n",
    "n_perm = 100_000\n",
    "two_sided = False\n",
    "alpha = 0.05\n",
    "\n",
    "FOLDERPATH_ANALYSIS = \"Path_to\\\\Project\\\\Analysis\"\n",
    "FOLDERPATH_FIGURES = os.path.join(os.path.dirname(os.getcwd()), \"figures\")\n",
    "\n",
    "ltime = -5\n",
    "htime = 51\n",
    "\n",
    "subregions_mapping = {\n",
    "    \"Motor cortex\": \"motor\",\n",
    "    \"Parietal cortex\": \"parietal\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CORTEX -> STN TIME DELAY (Figure 3d)\n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_con_tde-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "times = copy.deepcopy(data[\"timepoints\"][0])\n",
    "data[\"timepoints\"] = [times] * len(data[\"seed_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "subregion_i = 0\n",
    "fig, axis = plt.subplots(2, 1, gridspec_kw={\"height_ratios\": [2, 1]})\n",
    "hist_bins = np.arange(0, htime+1, 10)\n",
    "hist_bins[0] += 1\n",
    "tde = {}\n",
    "tde_off = {}\n",
    "tde_on = {}\n",
    "tde_off_noavg = {}\n",
    "tde_on_noavg = {}\n",
    "peaks_off = {}\n",
    "peaks_on = {}\n",
    "peaks_subs_off = {}\n",
    "peaks_subs_on = {}\n",
    "peaks_subs = {}\n",
    "n_cons_subs = {}\n",
    "n_cons_subs_off = {}\n",
    "n_cons_subs_on = {}\n",
    "peaks = {}\n",
    "n_subs = {}\n",
    "\n",
    "for subregion_name, subregion_label in subregions_mapping.items():\n",
    "    tde_off[subregion_label] = []\n",
    "    tde_off_noavg[subregion_label] = []\n",
    "    tde_on[subregion_label] = []\n",
    "    tde_on_noavg[subregion_label] = []\n",
    "    tau = []\n",
    "    peaks_off[subregion_label] = []\n",
    "    peaks_on[subregion_label] = []\n",
    "    peaks_subs_off[subregion_label] = {}\n",
    "    peaks_subs_on[subregion_label] = {}\n",
    "    peaks_subs[subregion_label] = {}\n",
    "    n_cons_subs[subregion_label] = {}\n",
    "    n_cons_subs_off[subregion_label] = {}\n",
    "    n_cons_subs_on[subregion_label] = {}\n",
    "    for sub in np.unique(data[\"sub\"]):\n",
    "        peaks_subs[subregion_label][sub] = []\n",
    "        n_cons_subs[subregion_label][sub] = 0\n",
    "        sub_off_idcs = ((data[\"stim\"] == \"Off\") & (data[\"freq_band_names\"] == \"all\")\n",
    "                        & (data[\"seed_subregions\"] == subregion_label) & (data[\"sub\"] == sub)\n",
    "                        & ~((data[\"tde-i_standard_tau_ci_80_low\"] < 0) & (data[\"tde-i_standard_tau_ci_80_high\"] > 0)))\n",
    "        tde_off_sub = np.array(data[\"tde-i_standard\"][sub_off_idcs].to_list())\n",
    "        tau_off_sub = np.array(data[\"tde-i_standard_tau\"][sub_off_idcs].to_list())\n",
    "        n_cons_off = tde_off_sub.shape[0]\n",
    "        if tde_off_sub.size > 0:\n",
    "            peaks_off_sub = []\n",
    "            for this_tde in tde_off_sub:\n",
    "                peaks_off_sub.extend(find_peaks(this_tde, height=this_tde[times.index(0)], prominence=(0.5, None), distance=5)[0])\n",
    "            tde_off_sub = tde_off_sub.mean(0)\n",
    "            tau_off_sub = tau_off_sub.mean(0)\n",
    "\n",
    "        sub_on_idcs = ((data[\"stim\"] == \"On\") & (data[\"freq_band_names\"] == \"all\")\n",
    "                        & (data[\"seed_subregions\"] == subregion_label) & (data[\"sub\"] == sub)\n",
    "                        & ~((data[\"tde-i_standard_tau_ci_80_low\"] < 0) & (data[\"tde-i_standard_tau_ci_80_high\"] > 0)))\n",
    "        tde_on_sub = np.array(data[\"tde-i_standard\"][sub_on_idcs].to_list())\n",
    "        tau_on_sub = np.array(data[\"tde-i_standard_tau\"][sub_on_idcs].to_list())\n",
    "        n_cons_on = tde_on_sub.shape[0]\n",
    "        if tde_on_sub.size > 0:\n",
    "            peaks_on_sub = []\n",
    "            for this_tde in tde_on_sub:\n",
    "                peaks_on_sub.extend(find_peaks(this_tde, height=this_tde[times.index(0)], prominence=(0.5, None), distance=5)[0])\n",
    "            tde_on_sub = tde_on_sub.mean(0)\n",
    "            tau_on_sub = tau_on_sub.mean(0)\n",
    "        \n",
    "        if tde_off_sub.size > 0 and tde_on_sub.size > 0:\n",
    "            tde_off[subregion_label].append(tde_off_sub)\n",
    "            tde_off_noavg[subregion_label].append(tde_off_sub)\n",
    "            tde_on[subregion_label].append(tde_on_sub)\n",
    "            tde_on_noavg[subregion_label].append(tde_on_sub)\n",
    "            peaks_off[subregion_label].extend(peaks_off_sub)\n",
    "            peaks_on[subregion_label].extend(peaks_on_sub)\n",
    "            peaks_subs_off[subregion_label][sub] = np.array(peaks_off_sub)\n",
    "            peaks_subs_on[subregion_label][sub] = np.array(peaks_on_sub)\n",
    "            n_cons_subs_off[subregion_label][sub] = n_cons_off\n",
    "            n_cons_subs_on[subregion_label][sub] = n_cons_on\n",
    "    \n",
    "    tde_off[subregion_label] = np.array(tde_off[subregion_label]).mean(axis=0)\n",
    "    tde_on[subregion_label] = np.array(tde_on[subregion_label]).mean(axis=0)\n",
    "    tde_off_noavg[subregion_label] = np.array(tde_off_noavg[subregion_label])\n",
    "    tde_on_noavg[subregion_label] = np.array(tde_on_noavg[subregion_label])\n",
    "    peaks_off[subregion_label] = np.array([times[peak_idx] for peak_idx in peaks_off[subregion_label]])\n",
    "    peaks_on[subregion_label] = np.array([times[peak_idx] for peak_idx in peaks_on[subregion_label]])\n",
    "    n_subs[subregion_label] = tde_off_noavg[subregion_label].shape[0]\n",
    "    \n",
    "    subregion_i += 1\n",
    "\n",
    "tde_off_baseline = tde_off[\"parietal\"]\n",
    "tde_on_baseline = tde_on[\"parietal\"]\n",
    "\n",
    "axis[0].plot(times[times.index(ltime):times.index(htime)+1],\n",
    "             tde_off[\"motor\"][times.index(ltime):times.index(htime)+1],\n",
    "             color=\"#DF4A4A\")\n",
    "axis[0].plot(times[times.index(ltime):times.index(htime)+1],\n",
    "             tde_on[\"motor\"][times.index(ltime):times.index(htime)+1],\n",
    "             color=\"#71BCAD\")\n",
    "\n",
    "peak_off_hists = {}\n",
    "peak_on_hists = {}\n",
    "for subregion_label in peaks_off.keys():\n",
    "    peak_off_hists[subregion_label] = np.histogram(\n",
    "        peaks_off[subregion_label], bins=hist_bins, range=(0, htime),\n",
    "        weights=(1 / n_subs[subregion_label]) * np.ones_like(peaks_off[subregion_label])\n",
    "    )[0]\n",
    "    peak_on_hists[subregion_label] = np.histogram(\n",
    "        peaks_on[subregion_label], bins=hist_bins, range=(0, htime),\n",
    "        weights=(1 / n_subs[subregion_label]) * np.ones_like(peaks_on[subregion_label])\n",
    "    )[0]\n",
    "axis[1].hist(np.arange(1, htime, 10), bins=hist_bins, range=(0, htime),\n",
    "                weights=peak_off_hists[\"motor\"],\n",
    "                alpha=0.7, color=\"#DF4A4A\")\n",
    "axis[1].hist(np.arange(1, htime, 10), bins=hist_bins, range=(0, htime),\n",
    "                weights=peak_on_hists[\"motor\"],\n",
    "                alpha=0.7, color=\"#71BCAD\")\n",
    "\n",
    "axis[0].spines['top'].set_visible(False)\n",
    "axis[0].spines['right'].set_visible(False)\n",
    "axis[1].spines['top'].set_visible(False)\n",
    "axis[1].spines['right'].set_visible(False)\n",
    "axis[0].set_xticks(np.arange(0, htime+1, 10))\n",
    "axis[1].set_xticks(np.arange(0, htime+1, 10))\n",
    "axis[1].set_xlim(axis[0].get_xlim())\n",
    "axis[0].plot([0, 0], axis[0].get_ylim(), color=\"k\", linestyle=\"--\", linewidth=1)\n",
    "axis[1].plot([0, 0], axis[1].get_ylim(), color=\"k\", linestyle=\"--\", linewidth=1)\n",
    "\n",
    "axis[1].set_xlabel(\"Time (ms)\")\n",
    "axis[0].set_ylabel(\"Time delay estimate\\nstrength (Z-score)\")\n",
    "axis[0].set_title(\"Motor cortex -> STN\")\n",
    "\n",
    "hist_max = np.max([peak_off_hists[\"motor\"], peak_on_hists[\"motor\"]])\n",
    "axis[1].set_ylim(hist_max/2, hist_max)\n",
    "axis[1].set_yticks(np.linspace(hist_max/4, hist_max, 4))\n",
    "axis[1].set_yticklabels([\"25\", \"50\", \"75\", \"100\"])\n",
    "axis[1].set_ylabel(\"Peak count\\n(% maximum)\")\n",
    "\n",
    "axis[0].tick_params(\"x\", pad=2, size=2)\n",
    "axis[0].tick_params(\"y\", pad=2, size=2)\n",
    "axis[1].tick_params(\"x\", pad=2, size=2)\n",
    "axis[1].tick_params(\"y\", pad=2, size=2)\n",
    "\n",
    "axis[0].xaxis.labelpad = 1\n",
    "axis[0].yaxis.labelpad = 1\n",
    "axis[1].xaxis.labelpad = 1\n",
    "axis[1].yaxis.labelpad = 1\n",
    "\n",
    "fig.set_size_inches(2.1, 2)\n",
    "\n",
    "fig.savefig(os.path.join(FOLDERPATH_FIGURES, \"Manuscript_TDE_Stim_spectrum.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find average peak time in 1-9 ms window\n",
    "peak_times = peaks_off[\"motor\"]\n",
    "peak_times_sub10 = peak_times[(peak_times > 0) & (peak_times < 10)]\n",
    "print(np.mean(peak_times_sub10))\n",
    "print(sem(peak_times_sub10))\n",
    "print(\n",
    "    \"Average motor cortex -> STN peak time below 10 ms (OFF): \"\n",
    "    f\"{np.mean(peak_times_sub10)} +/- {sem(peak_times_sub10)} ms (mean +/- SEM)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find times of significant estimates (OFF & ON)\n",
    "fig, axis = pte_decode.lineplot_prediction_compare(\n",
    "    x_1=(tde_off_noavg[\"motor\"] - tde_off_baseline)[:, times.index(ltime):times.index(htime)+1].T,\n",
    "    x_2=np.zeros_like(tde_off_noavg[\"motor\"])[:, times.index(ltime):times.index(htime)+1].T,\n",
    "    times=np.array(times[times.index(ltime):times.index(htime)+1]),\n",
    "    data_labels=[\"motor\", \"parietal\"],\n",
    "    x_label=\"Time (ms)\",\n",
    "    y_label=\"Est (A.U.)\",\n",
    "    two_tailed=two_sided,\n",
    "    paired_x1x2=False,\n",
    "    n_perm=n_perm,\n",
    "    correction_method=\"cluster_pvals\",\n",
    "    title=\"Cortex -> STN (OFF)\",\n",
    "    colour=[\"#DF4A4A\", \"k\"]\n",
    ")\n",
    "\n",
    "fig, axis = pte_decode.lineplot_prediction_compare(\n",
    "    x_1=(tde_on_noavg[\"motor\"] - tde_on_baseline)[:, times.index(ltime):times.index(htime)+1].T,\n",
    "    x_2=np.zeros_like(tde_on_noavg[\"motor\"])[:, times.index(ltime):times.index(htime)+1].T,\n",
    "    times=np.array(times[times.index(ltime):times.index(htime)+1]),\n",
    "    data_labels=[\"motor\", \"parietal\"],\n",
    "    x_label=\"Time (ms)\",\n",
    "    y_label=\"Est (A.U.)\",\n",
    "    two_tailed=two_sided,\n",
    "    paired_x1x2=False,\n",
    "    n_perm=n_perm,\n",
    "    correction_method=\"cluster_pvals\",\n",
    "    title=\"Cortex -> STN (ON)\",\n",
    "    colour=[\"#71BCAD\", \"k\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find time bins with significant number of peaks (OFF)\n",
    "n_peaks = {}\n",
    "for subregion_label in subregions_mapping.values():\n",
    "    n_peaks[subregion_label] = np.zeros((len(n_cons_subs_off[subregion_label].values()), len(hist_bins) - 1))\n",
    "    sub_idx = 0\n",
    "    for sub, n_cons in zip(peaks_subs_off[subregion_label].values(), n_cons_subs_off[subregion_label].values()):\n",
    "        if n_cons > 0:\n",
    "            bin_idx = 0\n",
    "            for start_time, end_time in zip(hist_bins[:-1], hist_bins[1:]):\n",
    "                for idx in sub:\n",
    "                    n_peaks[subregion_label][sub_idx, bin_idx] += np.count_nonzero((times[idx] >= start_time) & (times[idx] < end_time))\n",
    "                n_peaks[subregion_label][sub_idx, bin_idx] /= n_cons\n",
    "                bin_idx += 1\n",
    "        else:\n",
    "            n_peaks[subregion_label][sub_idx, :] = np.nan\n",
    "        sub_idx += 1\n",
    "\n",
    "n_peaks_pvals = {\"motor\": []}\n",
    "for subregion_label in n_peaks_pvals.keys():\n",
    "    for time_bin_idx, time_bin in enumerate(n_peaks[subregion_label].T):\n",
    "        _, pval = pte_stats.permutation_twosample(\n",
    "            data_a=np.array(time_bin[~np.isnan(time_bin)]) - np.nanmean(n_peaks[\"parietal\"], axis=0)[time_bin_idx],\n",
    "            data_b=np.zeros_like(time_bin[~np.isnan(time_bin)]),\n",
    "            n_perm=n_perm, two_tailed=two_sided\n",
    "        )\n",
    "        n_peaks_pvals[subregion_label].append(pval)\n",
    "\n",
    "n_peaks_pvals[\"motor\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find time bins with significant number of peaks (ON)\n",
    "n_peaks = {}\n",
    "for subregion_label in subregions_mapping.values():\n",
    "    n_peaks[subregion_label] = np.zeros((len(n_cons_subs_on[subregion_label].values()), len(hist_bins) - 1))\n",
    "    sub_idx = 0\n",
    "    for sub, n_cons in zip(peaks_subs_on[subregion_label].values(), n_cons_subs_on[subregion_label].values()):\n",
    "        if n_cons > 0:\n",
    "            bin_idx = 0\n",
    "            for start_time, end_time in zip(hist_bins[:-1], hist_bins[1:]):\n",
    "                for idx in sub:\n",
    "                    n_peaks[subregion_label][sub_idx, bin_idx] += np.count_nonzero((times[idx] >= start_time) & (times[idx] < end_time))\n",
    "                n_peaks[subregion_label][sub_idx, bin_idx] /= n_cons\n",
    "                bin_idx += 1\n",
    "        else:\n",
    "            n_peaks[subregion_label][sub_idx, :] = np.nan\n",
    "        sub_idx += 1\n",
    "\n",
    "n_peaks_pvals = {\"motor\": []}\n",
    "for subregion_label in n_peaks_pvals.keys():\n",
    "    for time_bin_idx, time_bin in enumerate(n_peaks[subregion_label].T):\n",
    "        _, pval = pte_stats.permutation_twosample(\n",
    "            data_a=np.array(time_bin[~np.isnan(time_bin)]) - np.nanmean(n_peaks[\"parietal\"], axis=0)[time_bin_idx],\n",
    "            data_b=np.zeros_like(time_bin[~np.isnan(time_bin)]),\n",
    "            n_perm=n_perm, two_tailed=two_sided\n",
    "        )\n",
    "        n_peaks_pvals[subregion_label].append(pval)\n",
    "\n",
    "n_peaks_pvals[\"motor\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute taus \n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_con_tde-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "times = copy.deepcopy(data[\"timepoints\"][0])\n",
    "data[\"timepoints\"] = [times] * len(data[\"seed_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "tau = {}\n",
    "tau_off = {}\n",
    "tau_on = {}\n",
    "for subregion_name, subregion_label in subregions_mapping.items():\n",
    "    tau[subregion_label] = []\n",
    "    tau_off[subregion_label] = []\n",
    "    tau_on[subregion_label] = []\n",
    "    for sub in np.unique(data[\"sub\"]):\n",
    "        sub_off_idcs = ((data[\"stim\"] == \"Off\") & (data[\"freq_band_names\"] == \"all\")\n",
    "                        & (data[\"seed_subregions\"] == subregion_label)\n",
    "                        & (data[\"tde-i_standard_tau\"] > 0) & (data[\"sub\"] == sub)\n",
    "                        & ~((data[\"tde-i_standard_tau_ci_80_low\"] < 0) & (data[\"tde-i_standard_tau_ci_80_high\"] > 0)))\n",
    "        tau_off_sub = np.array(data[\"tde-i_standard_tau\"][sub_off_idcs].to_list())\n",
    "        if tau_off_sub.size > 0:\n",
    "            tau_off_sub = tau_off_sub.mean(0)\n",
    "\n",
    "        sub_on_idcs = ((data[\"stim\"] == \"On\") & (data[\"freq_band_names\"] == \"all\")\n",
    "                        & (data[\"seed_subregions\"] == subregion_label)\n",
    "                        & (data[\"tde-i_standard_tau\"] > 0) & (data[\"sub\"] == sub)\n",
    "                        & ~((data[\"tde-i_standard_tau_ci_80_low\"] < 0) & (data[\"tde-i_standard_tau_ci_80_high\"] > 0)))\n",
    "        tau_on_sub = np.array(data[\"tde-i_standard_tau\"][sub_on_idcs].to_list())\n",
    "        if tau_on_sub.size > 0:\n",
    "            tau_on_sub = tau_on_sub.mean(0)\n",
    "        \n",
    "        if tau_off_sub.size > 0 and tau_on_sub.size > 0:\n",
    "            tau_off[subregion_label].append(tau_off_sub)\n",
    "            tau_on[subregion_label].append(tau_on_sub)\n",
    "            tau[subregion_label].append(np.mean([tau_off_sub, tau_on_sub], 0))\n",
    "    \n",
    "    tau_off[subregion_label] = np.array(tau_off)\n",
    "    tau_on[subregion_label] = np.array(tau_on)\n",
    "    tau[subregion_label] = np.array(tau[subregion_label])\n",
    "\n",
    "    _, pval = pte_stats.permutation_onesample(data_a=tau_off, data_b=tau_on, n_perm=n_perm, two_tailed=two_sided)\n",
    "    print(f\"{subregion_name} tau, OFF vs. ON: p-value = {pval}\")\n",
    "\n",
    "for subregion in tau.keys():\n",
    "    print(f\"{subregion} tau (avg. OFF & ON): {np.mean(tau[subregion])} +/- {sem(tau[subregion])} ms (mean +/- SEM)\")\n",
    "    print(f\"{subregion} tau (OFF): {np.mean(tau_off[subregion])} +/- {sem(tau_off[subregion])} ms (mean +/- SEM)\")\n",
    "    print(f\"{subregion} tau (ON): {np.mean(tau_on[subregion])} +/- {sem(tau_on[subregion])} ms (mean +/- SEM)\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pte",
   "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.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}