plotting/power-StimOffOn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Plot STIM OFF vs. ON power (Figures 2a, S1, & S2).\"\"\"\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",
    "import pte_stats\n",
    "import pte_decode\n",
    "from matplotlib import pyplot as plt\n",
    "from statannotations.stats.StatTest import StatTest\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",
    "alpha = 0.05\n",
    "two_sided = True\n",
    "\n",
    "FOLDERPATH_ANALYSIS = \"Path_to\\\\Project\\\\Analysis\"\n",
    "FOLDERPATH_FIGURES = os.path.join(os.path.dirname(os.getcwd()), \"figures\")\n",
    "\n",
    "def permutation_onesample():\n",
    "    \"\"\"Wrapper for StatTest with permutation one-sample test.\"\"\"\n",
    "\n",
    "    def _stat_test(x, y):\n",
    "        if isinstance(x, pd.Series):\n",
    "            x = x.to_numpy()\n",
    "        if isinstance(y, pd.Series):\n",
    "            y = y.to_numpy()\n",
    "        diff = x - y\n",
    "        return pte_stats.permutation_onesample(\n",
    "            data_a=diff, data_b=0, n_perm=100_000, two_tailed=True\n",
    "        )\n",
    "\n",
    "    return StatTest(\n",
    "        func=_stat_test,\n",
    "        alpha=0.05,\n",
    "        test_long_name=\"Permutation Test\",\n",
    "        test_short_name=\"Perm. Test\"\n",
    "    )\n",
    "\n",
    "lfreq = 4.0\n",
    "hfreq = 40.0\n",
    "\n",
    "fbands = {\n",
    "    \"alpha\": (8, 12),\n",
    "    \"low_beta\": (12, 20),\n",
    "    \"high_beta\": (20, 30)\n",
    "}\n",
    "\n",
    "fband_func = np.mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WHOLE CORTEX POWER (Figure 2a; also Figure S2 when adjusting the data loaded)\n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_pow_standard_whole-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "freqs = copy.deepcopy(data[\"freqs\"][0])\n",
    "freqs_delta = np.unique(np.subtract(freqs[1:], freqs[:-1]))[0]\n",
    "data[\"freqs\"] = [freqs] * len(data[\"ch_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "pow_off = np.array(data[\"power-multitaper\"][(data[\"stim\"] == \"Off\") & (data[\"ch_regions\"] == \"cortex\")].to_list())\n",
    "pow_on = np.array(data[\"power-multitaper\"][(data[\"stim\"] == \"On\") & (data[\"ch_regions\"] == \"cortex\")].to_list())\n",
    "\n",
    "np.random.seed(44)\n",
    "\n",
    "lfreq_i = freqs.index(lfreq)\n",
    "hfreq_i = freqs.index(hfreq)\n",
    "plot_freqs = freqs[lfreq_i:hfreq_i+1]\n",
    "\n",
    "pow_off = pow_off[:, lfreq_i:hfreq_i+1]\n",
    "pow_on = pow_on[:, lfreq_i:hfreq_i+1]\n",
    "\n",
    "fig, axis = pte_decode.lineplot_prediction_compare(\n",
    "    x_1=pow_off.T,\n",
    "    x_2=pow_on.T,\n",
    "    times=np.arange(lfreq, hfreq+0.5, 0.5),\n",
    "    data_labels=[\"OFF therapy\", \"ON STN-DBS\"],\n",
    "    x_label=\"Frequency (Hz)\",\n",
    "    y_label=\"Power (% total)\",\n",
    "    two_tailed=two_sided,\n",
    "    paired_x1x2=True,\n",
    "    n_perm=n_perm,\n",
    "    correction_method=\"cluster_pvals\",\n",
    "    title=\"Cortex\",\n",
    "    colour=[\"#DF4A4A\", \"#71BCAD\"]\n",
    ")\n",
    "\n",
    "for fband_name, fband_lims in fbands.items():\n",
    "    stat, pval = pte_stats.permutation_onesample(\n",
    "        data_a=fband_func(pow_off[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1),\n",
    "        data_b=fband_func(pow_on[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1),\n",
    "        n_perm=n_perm, two_tailed=two_sided\n",
    "    )\n",
    "    if pval < alpha:\n",
    "        axis.axvspan(fband_lims[0], fband_lims[1], color=[0.8, 0.8, 0.8], alpha=0.3, linewidth=0, zorder=0)\n",
    "        prefix = \"[SIGNIFICANT]\"\n",
    "    else:\n",
    "        prefix = \"\"\n",
    "    print(f\"{prefix} {fband_name}: stat={stat}; p-val={pval}\")\n",
    "\n",
    "axis.spines['top'].set_visible(False)\n",
    "axis.spines['right'].set_visible(False)\n",
    "axis.set_box_aspect(0.8)\n",
    "axis.legend(handlelength=0.5, handletextpad=0.5, labelspacing=0, borderaxespad=0, loc=\"lower left\")\n",
    "leg = axis.get_legend()\n",
    "leg.legend_handles[0].set_linewidth(1.5)\n",
    "leg.legend_handles[1].set_linewidth(1.5)\n",
    "leg.set_frame_on(False)\n",
    "axis.tick_params(\"x\", pad=2, size=2)\n",
    "axis.tick_params(\"y\", pad=2, size=2)\n",
    "axis.set_xticks(np.arange(5, 45, 5))\n",
    "axis.xaxis.labelpad = 0\n",
    "axis.yaxis.labelpad = 1\n",
    "axis.set_yticks([1, 2, 3])\n",
    "\n",
    "fig.set_size_inches(1.6, 1.6)\n",
    "\n",
    "fig.savefig(os.path.join(FOLDERPATH_FIGURES, \"Manuscript_Power_Stim_whole_cortex.pdf\"), dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# STN POWER (Figure 2a; also Figure S2 when adjusting the data loaded)\n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_pow_standard_whole-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "freqs = copy.deepcopy(data[\"freqs\"][0])\n",
    "freqs_delta = np.unique(np.subtract(freqs[1:], freqs[:-1]))[0]\n",
    "data[\"freqs\"] = [freqs] * len(data[\"ch_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "pow_off = np.array(data[\"power-multitaper\"][(data[\"stim\"] == \"Off\") & (data[\"ch_regions\"] == \"STN\")].to_list())\n",
    "pow_on = np.array(data[\"power-multitaper\"][(data[\"stim\"] == \"On\") & (data[\"ch_regions\"] == \"STN\")].to_list())\n",
    "\n",
    "np.random.seed(44)\n",
    "\n",
    "lfreq_i = freqs.index(lfreq)\n",
    "hfreq_i = freqs.index(hfreq)\n",
    "plot_freqs = freqs[lfreq_i:hfreq_i+1]\n",
    "\n",
    "pow_off = pow_off[:, lfreq_i:hfreq_i+1]\n",
    "pow_on = pow_on[:, lfreq_i:hfreq_i+1]\n",
    "\n",
    "fig, axis = pte_decode.lineplot_prediction_compare(\n",
    "    x_1=pow_off.T,\n",
    "    x_2=pow_on.T,\n",
    "    times=np.arange(lfreq, hfreq+0.5, 0.5),\n",
    "    data_labels=[\"OFF therapy\", \"ON STN-DBS\"],\n",
    "    x_label=\"Frequency (Hz)\",\n",
    "    y_label=\"Power (% total)\",\n",
    "    two_tailed=two_sided,\n",
    "    paired_x1x2=True,\n",
    "    n_perm=n_perm,\n",
    "    correction_method=\"cluster_pvals\",\n",
    "    title=\"STN\",\n",
    "    colour=[\"#DF4A4A\", \"#71BCAD\"]\n",
    ")\n",
    "\n",
    "for fband_name, fband_lims in fbands.items():\n",
    "    stat, pval = pte_stats.permutation_onesample(\n",
    "        data_a=fband_func(pow_off[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1),\n",
    "        data_b=fband_func(pow_on[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1),\n",
    "        n_perm=n_perm, two_tailed=True\n",
    "    )\n",
    "    if pval < alpha:\n",
    "        axis.axvspan(fband_lims[0], fband_lims[1], color=[0.8, 0.8, 0.8], alpha=0.3, linewidth=0, zorder=0)\n",
    "        prefix = \"[SIGNIFICANT]\"\n",
    "    else:\n",
    "        prefix = \"\"\n",
    "    print(f\"{prefix} {fband_name}: stat={stat}; p-val={pval}\")\n",
    "\n",
    "axis.spines['top'].set_visible(False)\n",
    "axis.spines['right'].set_visible(False)\n",
    "axis.set_box_aspect(0.8)\n",
    "axis.legend(handlelength=0.5, handletextpad=0.5, labelspacing=0, borderaxespad=0, loc=\"lower left\")\n",
    "leg = axis.get_legend()\n",
    "leg.legend_handles[0].set_linewidth(1.5)\n",
    "leg.legend_handles[1].set_linewidth(1.5)\n",
    "leg.legend_handles[2].set_color([0.5, 0.5, 0.5])\n",
    "leg.set_frame_on(False)\n",
    "axis.set_xticks(np.arange(5, 45, 5))\n",
    "axis.tick_params(\"x\", pad=2, size=2)\n",
    "axis.tick_params(\"y\", pad=2, size=2)\n",
    "axis.set_ylim((0.2, 3.1))\n",
    "axis.xaxis.labelpad = 0\n",
    "axis.yaxis.labelpad = 1\n",
    "axis.set_yticks([1, 2, 3])\n",
    "\n",
    "fig.set_size_inches(1.58, 1.58)\n",
    "\n",
    "fig.savefig(os.path.join(FOLDERPATH_FIGURES, \"Manuscript_Power_Stim_STN.pdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# STN POWER OFF (Channels with low vs. high beta peaks; Figure S1a)\n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_pow_standard_channels-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "freqs = copy.deepcopy(data[\"freqs\"][0])\n",
    "freqs_delta = np.unique(np.subtract(freqs[1:], freqs[:-1]))[0]\n",
    "data[\"freqs\"] = [freqs] * len(data[\"ch_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "pow_low = []\n",
    "pow_high = []\n",
    "peak_channels_low = {\n",
    "    \"006\": [\"01-05\", \"01-06\", \"01-07\", \"05-07\", \"06-07\", \"06-08\", \"07-08\"],\n",
    "    \"007\": [\"01-05\", \"05-08\"],\n",
    "    \"009\": [\"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "    \"012\": [\"01-02\", \"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"02-08\", \"03-08\", \"04-08\"],\n",
    "    \"014\": [\"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"03-08\", \"04-08\"],\n",
    "    \"017\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"05-08\"],\n",
    "    \"019\": [\"05-08\"],\n",
    "    \"022\": [\"01-07\", \"05-07\", \"06-07\", \"06-08\"],\n",
    "    \"023\": [\"01-02\", \"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"02-08\", \"03-08\", \"04-08\"],\n",
    "    \"025\": [\"01-07\", \"05-07\", \"06-07\", \"07-08\"],\n",
    "    \"026\": [\"01-04\", \"02-04\", \"03-04\", \"03-08\", \"04-08\"],\n",
    "    \"027\": [\"05-07\"],\n",
    "}\n",
    "peak_channels_high = {\n",
    "    \"006\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "    \"007\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "    \"009\": [\"01-05\", \"05-06\", \"05-07\"],\n",
    "    \"012\": [\"01-02\", \"01-03\", \"01-04\"],\n",
    "    \"014\": [\"01-02\", \"01-03\", \"01-04\", \"02-04\", \"03-04\", \"02-08\", \"03-08\", \"04-08\"],\n",
    "    \"017\": [\"01-05\", \"05-06\", \"05-07\", \"05-08\"],\n",
    "    \"019\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-08\"],\n",
    "    \"022\": [\"05-07\"],\n",
    "    \"023\": [\"01-02\", \"02-04\"],\n",
    "    \"025\": [\"05-06\", \"06-07\"],\n",
    "    \"026\": [\"01-02\", \"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"02-08\", \"03-08\"],\n",
    "    \"027\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "}\n",
    "for sub in peak_channels_low.keys():\n",
    "    sub_pow_low = []\n",
    "    sub_pow_high = []\n",
    "    sub_pow = data[(data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)]\n",
    "    for ch in peak_channels_low[sub]:\n",
    "        if sub_pow[\"ch_names\"][sub_pow[\"stim\"] == \"Off\"].str.contains(ch).any():\n",
    "            sub_pow_low.append(data[\"power-multitaper\"][(data[\"stim\"] == \"Off\") &\n",
    "                                                        (data[\"ch_names\"].str.contains(ch)) &\n",
    "                                                        (data[\"ch_regions\"] == \"STN\") &\n",
    "                                                        (data[\"sub\"] == sub)].to_list())\n",
    "    for ch in peak_channels_high[sub]:\n",
    "        if sub_pow[\"ch_names\"][sub_pow[\"stim\"] == \"Off\"].str.contains(ch).any():\n",
    "            sub_pow_high.append(data[\"power-multitaper\"][(data[\"stim\"] == \"Off\") &\n",
    "                                                         (data[\"ch_names\"].str.contains(ch)) &\n",
    "                                                         (data[\"ch_regions\"] == \"STN\") &\n",
    "                                                         (data[\"sub\"] == sub)].to_list())\n",
    "    pow_low.append(np.mean(sub_pow_low, axis=0)[0])\n",
    "    pow_high.append(np.mean(sub_pow_high, axis=0)[0])\n",
    "pow_low = np.array(pow_low)\n",
    "pow_high = np.array(pow_high)\n",
    "\n",
    "np.random.seed(44)\n",
    "\n",
    "lfreq_i = freqs.index(lfreq)\n",
    "hfreq_i = freqs.index(hfreq)\n",
    "plot_freqs = freqs[lfreq_i:hfreq_i+1]\n",
    "\n",
    "pow_low = pow_low[:, lfreq_i:hfreq_i+1]\n",
    "pow_high = pow_high[:, lfreq_i:hfreq_i+1]\n",
    "\n",
    "fig, axis = pte_decode.lineplot_prediction_compare(\n",
    "    x_1=pow_low.T,\n",
    "    x_2=pow_high.T,\n",
    "    times=np.arange(lfreq, hfreq+0.5, 0.5),\n",
    "    data_labels=[\"Low beta\", \"High beta\"],\n",
    "    x_label=\"Frequency (Hz)\",\n",
    "    y_label=\"Power (% total)\",\n",
    "    two_tailed=True,\n",
    "    paired_x1x2=True,\n",
    "    n_perm=n_perm,\n",
    "    correction_method=\"cluster_pvals\",\n",
    "    title=\"STN (OFF therapy)\",\n",
    "    colour=[\"#CF80E7\", \"#CC9800\"]\n",
    ")\n",
    "\n",
    "for fband_name, fband_lims in fbands.items():\n",
    "    stat, pval = pte_stats.permutation_onesample(\n",
    "        data_a=fband_func(pow_low[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1),\n",
    "        data_b=fband_func(pow_high[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1),\n",
    "        n_perm=n_perm, two_tailed=True\n",
    "    )\n",
    "    if pval < alpha:\n",
    "        axis.axvspan(fband_lims[0], fband_lims[1], color=[0.8, 0.8, 0.8], alpha=0.3, linewidth=0, zorder=0)\n",
    "        prefix = \"[SIGNIFICANT]\"\n",
    "    else:\n",
    "        prefix = \"\"\n",
    "    print(f\"{prefix} {fband_name}: stat={stat}; p-val={pval}\")\n",
    "\n",
    "axis.spines['top'].set_visible(False)\n",
    "axis.spines['right'].set_visible(False)\n",
    "axis.set_box_aspect(0.8)\n",
    "axis.legend(handlelength=0.5, handletextpad=0.5, labelspacing=0, borderaxespad=0, loc=\"lower left\")\n",
    "leg = axis.get_legend()\n",
    "leg.legend_handles[0].set_linewidth(1.5)\n",
    "leg.legend_handles[1].set_linewidth(1.5)\n",
    "leg.legend_handles[2].set_color([0.5, 0.5, 0.5])\n",
    "leg.set_frame_on(False)\n",
    "axis.set_xticks(np.arange(5, 45, 5))\n",
    "axis.tick_params(\"x\", pad=2, size=2)\n",
    "axis.tick_params(\"y\", pad=2, size=2)\n",
    "axis.set_ylim((0.2, 3.1))\n",
    "axis.xaxis.labelpad = 0\n",
    "axis.yaxis.labelpad = 1\n",
    "axis.set_yticks([1, 2, 3])\n",
    "\n",
    "fig.set_size_inches(1.58, 1.58)\n",
    "\n",
    "fig.savefig(os.path.join(FOLDERPATH_FIGURES, \"Manuscript_Power_Stim_OFF_low_vs_high_beta_channels.pdf\"), dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# STN POWER OFF vs. ON (Channels with low beta peaks; Figure S1b)\n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_pow_standard_channels-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "freqs = copy.deepcopy(data[\"freqs\"][0])\n",
    "freqs_delta = np.unique(np.subtract(freqs[1:], freqs[:-1]))[0]\n",
    "data[\"freqs\"] = [freqs] * len(data[\"ch_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "pow_off = []\n",
    "pow_on = []\n",
    "peak_channels = {\n",
    "    \"006\": [\"01-05\", \"01-06\", \"01-07\", \"05-07\", \"06-07\", \"06-08\", \"07-08\"],\n",
    "    \"007\": [\"01-05\", \"05-08\"],\n",
    "    \"009\": [\"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "    \"012\": [\"01-02\", \"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"02-08\", \"03-08\", \"04-08\"],\n",
    "    \"014\": [\"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"03-08\", \"04-08\"],\n",
    "    \"017\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"05-08\"],\n",
    "    \"019\": [\"05-08\"],\n",
    "    \"022\": [\"01-07\", \"05-07\", \"06-07\", \"06-08\"],\n",
    "    \"023\": [\"01-02\", \"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"02-08\", \"03-08\", \"04-08\"],\n",
    "    \"025\": [\"01-07\", \"05-07\", \"06-07\", \"07-08\"],\n",
    "    \"026\": [\"01-04\", \"02-04\", \"03-04\", \"03-08\", \"04-08\"],\n",
    "    \"027\": [\"05-07\"],\n",
    "}\n",
    "for sub in peak_channels.keys():\n",
    "    sub_pow_off = []\n",
    "    sub_pow_on = []\n",
    "    for ch in peak_channels[sub]:\n",
    "        sub_pow = data[(data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)]\n",
    "        if sub_pow[\"ch_names\"][sub_pow[\"stim\"] == \"Off\"].str.contains(ch).any():\n",
    "            sub_pow_off.append(data[\"power-multitaper\"][(data[\"stim\"] == \"Off\") & (data[\"ch_names\"].str.contains(ch)) & (data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)].to_list())\n",
    "        if sub_pow[\"ch_names\"][sub_pow[\"stim\"] == \"On\"].str.contains(ch).any():\n",
    "            sub_pow_on.append(data[\"power-multitaper\"][(data[\"stim\"] == \"On\") & (data[\"ch_names\"].str.contains(ch)) & (data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)].to_list())\n",
    "    pow_off.append(np.mean(sub_pow_off, axis=0)[0])\n",
    "    pow_on.append(np.mean(sub_pow_on, axis=0)[0])\n",
    "pow_off = np.array(pow_off)\n",
    "pow_on = np.array(pow_on)\n",
    "\n",
    "np.random.seed(44)\n",
    "\n",
    "lfreq_i = freqs.index(lfreq)\n",
    "hfreq_i = freqs.index(hfreq)\n",
    "plot_freqs = freqs[lfreq_i:hfreq_i+1]\n",
    "\n",
    "pow_off = pow_off[:, lfreq_i:hfreq_i+1]\n",
    "pow_on = pow_on[:, lfreq_i:hfreq_i+1]\n",
    "\n",
    "for fband_name, fband_lims in fbands.items():\n",
    "    pow_off_band = fband_func(pow_off[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1)\n",
    "    pow_on_band = fband_func(pow_on[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1)\n",
    "    plot_df = pd.DataFrame.from_dict({\n",
    "        \"Power (% total)\": pow_off_band.tolist() + pow_on_band.tolist(),\n",
    "        \"Stimulation\": [\"OFF\"] * len(pow_off_band) + [\"ON\"] * len(pow_on_band),\n",
    "        \"Subject\": list(peak_channels.keys()) * 2\n",
    "    })\n",
    "\n",
    "    fig, ax = pte_decode.violinplot_results(\n",
    "        data=plot_df,\n",
    "        outpath=None,\n",
    "        x=\"Stimulation\",\n",
    "        y=\"Power (% total)\",\n",
    "        hue=None,\n",
    "        order=[\"OFF\", \"ON\"],\n",
    "        hue_order=None,\n",
    "        stat_test=permutation_onesample(),\n",
    "        alpha=0.05,\n",
    "        add_lines=\"Subject\",\n",
    "        title=f\"Low beta channels; {fband_name}\",\n",
    "        swarm_size=3,\n",
    "        colours=[\"#DF4A4A\", \"#71BCAD\"],\n",
    "        figsize=[1.3, 2],\n",
    "        show=True\n",
    "    )\n",
    "\n",
    "    fig.savefig(os.path.join(FOLDERPATH_FIGURES, f\"Manuscript_Power_Stim_STN_lowbeta_channels_{fband_name}.pdf\"), dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# STN POWER OFF vs. ON (Channels with high beta peaks; Figure S1b)\n",
    "data = pd.read_pickle(os.path.join(\n",
    "    FOLDERPATH_ANALYSIS, \"task-Rest_acq-multi_run-multi_pow_standard_channels-StimOffOn_multi_sub.pkl\")\n",
    ")\n",
    "freqs = copy.deepcopy(data[\"freqs\"][0])\n",
    "freqs_delta = np.unique(np.subtract(freqs[1:], freqs[:-1]))[0]\n",
    "data[\"freqs\"] = [freqs] * len(data[\"ch_types\"])\n",
    "data = pd.DataFrame.from_dict(data)\n",
    "\n",
    "pow_off = []\n",
    "pow_on = []\n",
    "peak_channels = {\n",
    "    \"006\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "    \"007\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "    \"009\": [\"01-05\", \"05-06\", \"05-07\"],\n",
    "    \"012\": [\"01-02\", \"01-03\", \"01-04\"],\n",
    "    \"014\": [\"01-02\", \"01-03\", \"01-04\", \"02-04\", \"03-04\", \"02-08\", \"03-08\", \"04-08\"],\n",
    "    \"017\": [\"01-05\", \"05-06\", \"05-07\", \"05-08\"],\n",
    "    \"019\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-08\"],\n",
    "    \"022\": [\"05-07\"],\n",
    "    \"023\": [\"01-02\", \"02-04\"],\n",
    "    \"025\": [\"05-06\", \"06-07\"],\n",
    "    \"026\": [\"01-02\", \"01-03\", \"01-04\", \"02-03\", \"02-04\", \"03-04\", \"02-08\", \"03-08\"],\n",
    "    \"027\": [\"01-05\", \"01-06\", \"01-07\", \"05-06\", \"05-07\", \"06-07\", \"05-08\", \"06-08\", \"07-08\"],\n",
    "}\n",
    "for sub in peak_channels.keys():\n",
    "    sub_pow_off = []\n",
    "    sub_pow_on = []\n",
    "    for ch in peak_channels[sub]:\n",
    "        sub_pow = data[(data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)]\n",
    "        if sub_pow[\"ch_names\"][sub_pow[\"stim\"] == \"Off\"].str.contains(ch).any():\n",
    "            sub_pow_off.append(data[\"power-multitaper\"][(data[\"stim\"] == \"Off\") & (data[\"ch_names\"].str.contains(ch)) & (data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)].to_list())\n",
    "        if sub_pow[\"ch_names\"][sub_pow[\"stim\"] == \"On\"].str.contains(ch).any():\n",
    "            sub_pow_on.append(data[\"power-multitaper\"][(data[\"stim\"] == \"On\") & (data[\"ch_names\"].str.contains(ch)) & (data[\"ch_regions\"] == \"STN\") & (data[\"sub\"] == sub)].to_list())\n",
    "    pow_off.append(np.mean(sub_pow_off, axis=0)[0])\n",
    "    pow_on.append(np.mean(sub_pow_on, axis=0)[0])\n",
    "pow_off = np.array(pow_off)\n",
    "pow_on = np.array(pow_on)\n",
    "\n",
    "np.random.seed(44)\n",
    "\n",
    "lfreq_i = freqs.index(lfreq)\n",
    "hfreq_i = freqs.index(hfreq)\n",
    "plot_freqs = freqs[lfreq_i:hfreq_i+1]\n",
    "\n",
    "pow_off = pow_off[:, lfreq_i:hfreq_i+1]\n",
    "pow_on = pow_on[:, lfreq_i:hfreq_i+1]\n",
    "\n",
    "for fband_name, fband_lims in fbands.items():\n",
    "    pow_off_band = fband_func(pow_off[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1)\n",
    "    pow_on_band = fband_func(pow_on[:, freqs.index(fband_lims[0]):freqs.index(fband_lims[1])+1], 1)\n",
    "    plot_df = pd.DataFrame.from_dict({\n",
    "        \"Power (% total)\": pow_off_band.tolist() + pow_on_band.tolist(),\n",
    "        \"Stimulation\": [\"OFF\"] * len(pow_off_band) + [\"ON\"] * len(pow_on_band),\n",
    "        \"Subject\": list(peak_channels.keys()) * 2\n",
    "    })\n",
    "\n",
    "    fig, ax = pte_decode.violinplot_results(\n",
    "        data=plot_df,\n",
    "        outpath=None,\n",
    "        x=\"Stimulation\",\n",
    "        y=\"Power (% total)\",\n",
    "        hue=None,\n",
    "        order=[\"OFF\", \"ON\"],\n",
    "        hue_order=None,\n",
    "        stat_test=permutation_onesample(),\n",
    "        alpha=0.05,\n",
    "        add_lines=\"Subject\",\n",
    "        title=f\"High beta channels; {fband_name}\",\n",
    "        swarm_size=3,\n",
    "        colours=[\"#DF4A4A\", \"#71BCAD\"],\n",
    "        figsize=[1.3, 2],\n",
    "        show=True\n",
    "    )\n",
    "\n",
    "    fig.savefig(os.path.join(FOLDERPATH_FIGURES, f\"Manuscript_Power_Stim_STN_highbeta_channels_{fband_name}.pdf\"), dpi=300)"
   ]
  }
 ],
 "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
}