analysis_erp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Event related potentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:17:58.588752Z",
     "start_time": "2021-06-02T09:17:58.251543Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:02.397239Z",
     "start_time": "2021-06-02T09:17:59.321558Z"
    }
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import os\n",
    "import os.path as op\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import mne\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pingouin\n",
    "import seaborn as sns\n",
    "from matplotlib.colors import ListedColormap\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "from utils import (\n",
    "    BIDS_ROOT,\n",
    "    arr2df,\n",
    "    dict2arr,\n",
    "    get_erps_by_dict,\n",
    "    split_dict,\n",
    "    split_key_col,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:02.402828Z",
     "start_time": "2021-06-02T09:18:02.399243Z"
    }
   },
   "outputs": [],
   "source": [
    "# IO: Where to find the data\n",
    "# Where to find epochs\n",
    "fname_epo_template = op.join(\n",
    "    BIDS_ROOT, \"derivatives\", \"sub-{0:02}\", \"sub-{0:02}_epochs-epo.fif.gz\"\n",
    ")\n",
    "\n",
    "\n",
    "# Pack all names in a dict\n",
    "name_templates = dict()\n",
    "name_templates[\"epochs\"] = fname_epo_template\n",
    "\n",
    "subjects = range(1, 41)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:02.418673Z",
     "start_time": "2021-06-02T09:18:02.404872Z"
    }
   },
   "outputs": [],
   "source": [
    "# save data for publication plots\n",
    "plotdir = op.join(BIDS_ROOT, \"code\", \"publication_plots\")\n",
    "os.makedirs(plotdir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:03.394910Z",
     "start_time": "2021-06-02T09:18:03.352879Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot_topos(\n",
    "    erp_dict, vminmax, tmin=-0.1, tmax=0.8, step=0.1, extrapolate=\"head\", sphere=None\n",
    "):\n",
    "    \"\"\"Plot a topography over time.\"\"\"\n",
    "    # Figure settings\n",
    "    times = np.arange(tmin, tmax, step)\n",
    "\n",
    "    # Set up figure\n",
    "    n_keys = len(erp_dict)\n",
    "    fig, axs = plt.subplots(n_keys, len(times) + 1, figsize=(15, 2 * n_keys))\n",
    "    axs = np.atleast_2d(axs)\n",
    "\n",
    "    for ii, key in enumerate(erp_dict):\n",
    "        # Drop NaNs\n",
    "        arr = np.asarray(erp_dict[key])\n",
    "        n_nans = pd.isnull(arr).sum()\n",
    "        if n_nans > 0:\n",
    "            print(\"dropping {} nans from {} before plotting\".format(n_nans, key))\n",
    "            arr = arr[pd.notnull(arr)]\n",
    "\n",
    "        # plot colorbar only on first iteration\n",
    "        colorbar = True  # if ii == 0 else False\n",
    "\n",
    "        with mne.utils.use_log_level(0):\n",
    "            # Form grand average\n",
    "            ga = mne.grand_average(list(arr))\n",
    "\n",
    "            # plot\n",
    "            mne.viz.plot_evoked_topomap(\n",
    "                ga,\n",
    "                times=times,\n",
    "                vmin=-vminmax,\n",
    "                vmax=vminmax,\n",
    "                axes=axs[ii, :],\n",
    "                colorbar=colorbar,\n",
    "                extrapolate=extrapolate,\n",
    "                sphere=sphere,\n",
    "            )\n",
    "        axs[ii, 0].set_ylabel(key)\n",
    "\n",
    "    fig.suptitle(\"Sample ERP\")\n",
    "\n",
    "    return fig, axs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:03.626476Z",
     "start_time": "2021-06-02T09:18:03.580209Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot_single_erp(\n",
    "    erp_dict,\n",
    "    chn_grp,\n",
    "    title,\n",
    "    colors=None,\n",
    "    linestyles=None,\n",
    "    ylim=None,\n",
    "    cmap=None,\n",
    "    ax=None,\n",
    "    special_legend=False,\n",
    "):\n",
    "    \"\"\"Plot a single panel ERP.\"\"\"\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(1, 1)\n",
    "    else:\n",
    "        fig = ax.get_figure()\n",
    "\n",
    "    if colors is None:\n",
    "        legend = False\n",
    "    elif special_legend:\n",
    "        legend = False\n",
    "    else:\n",
    "        legend = True\n",
    "\n",
    "    # Drop NaNs\n",
    "    for key in erp_dict:\n",
    "        arr = np.asarray(erp_dict[key])\n",
    "        n_nans = pd.isnull(arr).sum()\n",
    "        if n_nans > 0:\n",
    "            print(\"dropping {} nans from {} before plotting\".format(n_nans, key))\n",
    "            arr = arr[pd.notnull(arr)]\n",
    "            erp_dict[key] = list(arr)\n",
    "\n",
    "    # plot\n",
    "    with mne.utils.use_log_level(0):\n",
    "        mne.viz.plot_compare_evokeds(\n",
    "            erp_dict,\n",
    "            ylim=ylim,\n",
    "            picks=chn_grp,\n",
    "            combine=\"mean\",\n",
    "            axes=ax,\n",
    "            show_sensors=\"lower left\",\n",
    "            legend=legend,\n",
    "            ci=0.68,\n",
    "            truncate_xaxis=False,\n",
    "            truncate_yaxis=False,\n",
    "            title=title,\n",
    "            colors=colors,\n",
    "            linestyles=linestyles,\n",
    "            cmap=cmap,\n",
    "        )\n",
    "\n",
    "    if special_legend:\n",
    "        # make ordering: stopping*sampling\n",
    "        ordering = zip(\n",
    "            [\"variable\", \"variable\", \"fixed\", \"fixed\"],\n",
    "            [\"active\", \"yoked\", \"active\", \"yoked\"],\n",
    "        )\n",
    "        legend_elements = []\n",
    "        for stopping, sampling in ordering:\n",
    "            legend_elements.append(\n",
    "                Line2D(\n",
    "                    [0],\n",
    "                    [0],\n",
    "                    color=colors[stopping],\n",
    "                    linestyle=linestyles[sampling],\n",
    "                    label=f\"{sampling}/{stopping}\",\n",
    "                )\n",
    "            )\n",
    "\n",
    "        ax.legend(\n",
    "            handles=legend_elements,\n",
    "            loc=\"best\",\n",
    "            prop={\"size\": 13},\n",
    "            framealpha=1,\n",
    "        )\n",
    "\n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define analysis parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:05.153119Z",
     "start_time": "2021-06-02T09:18:05.144185Z"
    }
   },
   "outputs": [],
   "source": [
    "# preprocessing settings\n",
    "preproc_settings = {\n",
    "    \"crop\": (0.6, 1.6),\n",
    "    \"tshift\": -0.8,\n",
    "    \"smooth\": False,\n",
    "    \"baseline\": (None, 0),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:05.313012Z",
     "start_time": "2021-06-02T09:18:05.309135Z"
    }
   },
   "outputs": [],
   "source": [
    "# Electrode groups for summarizing across channels\n",
    "p3_grp = [\"Cz\", \"C1\", \"C2\", \"CPz\", \"CP1\", \"CP2\", \"CP3\", \"CP4\", \"Pz\", \"P1\", \"P2\"]\n",
    "\n",
    "p1_n1_grp = [\"O2\", \"PO4\", \"PO8\", \"PO10\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:05.499581Z",
     "start_time": "2021-06-02T09:18:05.497058Z"
    }
   },
   "outputs": [],
   "source": [
    "# Time windows for ERP components\n",
    "p3_window = (0.3, 0.6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:05.725875Z",
     "start_time": "2021-06-02T09:18:05.722927Z"
    }
   },
   "outputs": [],
   "source": [
    "# define single panel plot scaling\n",
    "single_panel_ylim = dict(eeg=[-2, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:05.951340Z",
     "start_time": "2021-06-02T09:18:05.945808Z"
    }
   },
   "outputs": [],
   "source": [
    "# Define factors\n",
    "sampling_styles = (\"active\", \"yoked\")\n",
    "stopping_styles = (\"fixed\", \"variable\")\n",
    "timings = (\"early\", \"mid\", \"late\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:18:11.390355Z",
     "start_time": "2021-06-02T09:18:06.394692Z"
    }
   },
   "outputs": [],
   "source": [
    "# Get an ERP for extracting channel names and times array\n",
    "# it's the same for all ERPs and will be used below\n",
    "tmp = get_erps_by_dict({\"all\": list()}, name_templates, [1], **preproc_settings)\n",
    "ch_names = tmp[\"all\"][0].ch_names\n",
    "times = tmp[\"all\"][0].times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample ERP left - right contrast"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## double subtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:02.906324Z",
     "start_time": "2021-06-02T09:18:11.392497Z"
    }
   },
   "outputs": [],
   "source": [
    "# Make SIDED overall_erps = soverall_erps: DOUBLE SUBTRACTION\n",
    "# i.e., left - right, see \"subtract\" param for get_erps_by_dict\n",
    "# double_subtraction=True !\n",
    "try:\n",
    "    ssoverall_erps\n",
    "    print(\"Data already loaded\")\n",
    "except NameError:\n",
    "    subtract = (\"left\", \"right\")\n",
    "\n",
    "    ssoverall_erps = {\n",
    "        \"{0}/{1}\".format(*ii): list()\n",
    "        for ii in itertools.product(sampling_styles, stopping_styles)\n",
    "    }\n",
    "    ssoverall_erps = get_erps_by_dict(\n",
    "        ssoverall_erps,\n",
    "        name_templates,\n",
    "        subjects,\n",
    "        subtract=subtract,\n",
    "        double_subtraction=True,\n",
    "        **preproc_settings\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:02.911207Z",
     "start_time": "2021-06-02T09:19:02.908551Z"
    }
   },
   "outputs": [],
   "source": [
    "p1_window = (0.08, 0.13)\n",
    "n1_window = (0.14, 0.20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:03.266851Z",
     "start_time": "2021-06-02T09:19:02.913210Z"
    }
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(context=\"talk\", font_scale=1):\n",
    "\n",
    "    colors = {\"variable\": sns.color_palette()[1], \"fixed\": sns.color_palette()[0]}\n",
    "    linestyles = {\"active\": \"-\", \"yoked\": \"--\"}\n",
    "    custom_grp = [\"opo\"]\n",
    "    fig, ax = plot_single_erp(\n",
    "        ssoverall_erps,\n",
    "        custom_grp,\n",
    "        \"{}\".format(custom_grp),\n",
    "        colors,\n",
    "        linestyles,\n",
    "        ylim=None,\n",
    "        special_legend=True,\n",
    "    )\n",
    "\n",
    "    # P1 window\n",
    "    p1_window = (0.08, 0.13)\n",
    "    ax.axvspan(*p1_window, color=\"black\", alpha=0.1)\n",
    "\n",
    "    # N1 from 140ms to 200ms\n",
    "    n1_window = (0.14, 0.20)\n",
    "    ax.axvspan(*n1_window, color=\"black\", alpha=0.1)\n",
    "\n",
    "    # improve legend\n",
    "    leg = ax.get_legend()\n",
    "    leg.get_texts()[0].set_text(\"Full control\")\n",
    "    leg.get_texts()[1].set_text(\"Yoked\")\n",
    "    leg.get_texts()[2].set_text(\"Partial control\")\n",
    "    leg.get_texts()[3].set_text(\"Yoked\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:03.270497Z",
     "start_time": "2021-06-02T09:19:03.268163Z"
    }
   },
   "outputs": [],
   "source": [
    "special_group = [\"opo\"]\n",
    "special_ch_names = ssoverall_erps[\"active/fixed\"][0].ch_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:03.443805Z",
     "start_time": "2021-06-02T09:19:03.271936Z"
    }
   },
   "outputs": [],
   "source": [
    "ssoverall_arr = dict2arr(ssoverall_erps)\n",
    "\n",
    "sp1_df = arr2df(\n",
    "    ssoverall_arr,\n",
    "    special_group,\n",
    "    p1_window,\n",
    "    ssoverall_erps.keys(),\n",
    "    special_ch_names,\n",
    "    times,\n",
    ")\n",
    "\n",
    "sn1_df = arr2df(\n",
    "    ssoverall_arr,\n",
    "    special_group,\n",
    "    n1_window,\n",
    "    ssoverall_erps.keys(),\n",
    "    special_ch_names,\n",
    "    times,\n",
    ")\n",
    "\n",
    "\n",
    "sp1_df = split_key_col(sp1_df, [\"sampling\", \"stopping\"])\n",
    "sn1_df = split_key_col(sn1_df, [\"sampling\", \"stopping\"])\n",
    "\n",
    "\n",
    "# Get \"subject\" column from \"subj_repr\" by making use of knowledge of which\n",
    "# subj was in which stopping condition\n",
    "tmp = pd.read_csv(op.join(BIDS_ROOT, \"participants.tsv\"), sep=\"\\t\")\n",
    "fixed_ids = sorted(\n",
    "    [int(i[-2:]) for i in tmp[tmp[\"stopping\"] == \"fixed\"][\"participant_id\"]]\n",
    ")\n",
    "variable_ids = sorted(\n",
    "    [int(i[-2:]) for i in tmp[tmp[\"stopping\"] == \"variable\"][\"participant_id\"]]\n",
    ")\n",
    "\n",
    "sp1_df[\"subject\"] = pd.concat(\n",
    "    [\n",
    "        sp1_df[sp1_df[\"stopping\"] == \"fixed\"][\"subj_repr\"].map(\n",
    "            dict(zip(range(20), fixed_ids))\n",
    "        ),\n",
    "        sp1_df[sp1_df[\"stopping\"] == \"variable\"][\"subj_repr\"].map(\n",
    "            dict(zip(range(20), variable_ids))\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "sn1_df[\"subject\"] = pd.concat(\n",
    "    [\n",
    "        sn1_df[sn1_df[\"stopping\"] == \"fixed\"][\"subj_repr\"].map(\n",
    "            dict(zip(range(20), fixed_ids))\n",
    "        ),\n",
    "        sn1_df[sn1_df[\"stopping\"] == \"variable\"][\"subj_repr\"].map(\n",
    "            dict(zip(range(20), variable_ids))\n",
    "        ),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# reorder cols\n",
    "sp1_df = sp1_df[[\"key\", \"sampling\", \"stopping\", \"subject\", \"subj_repr\", \"mean_amp\"]]\n",
    "sn1_df = sn1_df[[\"key\", \"sampling\", \"stopping\", \"subject\", \"subj_repr\", \"mean_amp\"]]\n",
    "\n",
    "\n",
    "sp1_df[\"sampling\"] = sp1_df[\"sampling\"].str.capitalize()\n",
    "sp1_df[\"stopping\"] = sp1_df[\"stopping\"].str.capitalize()\n",
    "sn1_df[\"sampling\"] = sn1_df[\"sampling\"].str.capitalize()\n",
    "sn1_df[\"stopping\"] = sn1_df[\"stopping\"].str.capitalize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:03.448108Z",
     "start_time": "2021-06-02T09:19:03.445195Z"
    }
   },
   "outputs": [],
   "source": [
    "groups = {\"p1\": special_group, \"n1\": special_group}\n",
    "windows = {\"p1\": p1_window, \"n1\": n1_window}\n",
    "erps = {\"p1\": sp1_df, \"n1\": sn1_df}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:03.957389Z",
     "start_time": "2021-06-02T09:19:03.450113Z"
    }
   },
   "outputs": [],
   "source": [
    "scale = 1\n",
    "errwidth = 3\n",
    "capsize = 0.2\n",
    "\n",
    "labelpad = 12\n",
    "\n",
    "swarmsize = 2\n",
    "\n",
    "with sns.plotting_context(\"talk\", font_scale=1):\n",
    "    fig, axs = plt.subplots(\n",
    "        2, 1, figsize=(3, 6), sharex=True, gridspec_kw={\"height_ratios\": [1, 1]}\n",
    "    )\n",
    "\n",
    "for ii, use in enumerate((\"p1\", \"n1\")):\n",
    "    group = groups[use]\n",
    "    window = windows[use]\n",
    "    data = erps[use]\n",
    "    colname = \"mean_amp\"\n",
    "    order = [\"Active\", \"Yoked\"]\n",
    "    hue_order = [\"Fixed\", \"Variable\"]\n",
    "\n",
    "    with sns.plotting_context(\"talk\", font_scale=1):\n",
    "        ax = axs[ii]\n",
    "\n",
    "        sns.pointplot(\n",
    "            x=\"sampling\",\n",
    "            y=colname,\n",
    "            hue=\"stopping\",\n",
    "            data=data,\n",
    "            dodge=0.2,\n",
    "            ci=68,\n",
    "            order=order,\n",
    "            hue_order=hue_order,\n",
    "            ax=ax,\n",
    "            markers=\"o\",\n",
    "            scale=scale,\n",
    "            errwidth=errwidth,\n",
    "            capsize=capsize,\n",
    "        )\n",
    "\n",
    "        sns.swarmplot(\n",
    "            x=\"sampling\",\n",
    "            y=colname,\n",
    "            hue=\"stopping\",\n",
    "            data=data,\n",
    "            order=order,\n",
    "            hue_order=hue_order,\n",
    "            ax=ax,\n",
    "            dodge=True,\n",
    "            size=swarmsize,\n",
    "        )\n",
    "\n",
    "        # add legend\n",
    "        # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "        legend_elements = [\n",
    "            Line2D([0], [0], color=sns.color_palette()[0], marker=\"o\", label=\"Fixed\"),\n",
    "            Line2D(\n",
    "                [0], [0], color=sns.color_palette()[1], marker=\"o\", label=\"Variable\"\n",
    "            ),\n",
    "        ]\n",
    "\n",
    "        ax.legend(\n",
    "            handles=legend_elements,\n",
    "            loc=\"lower left\",\n",
    "            frameon=False,\n",
    "            bbox_to_anchor=(1, 0.5),\n",
    "        )\n",
    "\n",
    "        ax.set_ylabel(\"Mean amplitude\")\n",
    "\n",
    "        # https://stackoverflow.com/a/51157346/5201771\n",
    "        # 2-4 3-5\n",
    "        for i_dots, (idx0, idx1) in enumerate([(2, 4), (3, 5)]):\n",
    "            locs1 = ax.get_children()[idx0].get_offsets()\n",
    "            locs2 = ax.get_children()[idx1].get_offsets()\n",
    "\n",
    "            # Need to sort locs, so data corresponds\n",
    "            sort_idxs_list = []\n",
    "            sampling = order[i_dots]\n",
    "            for stopping in hue_order:\n",
    "                arr = data[\n",
    "                    (data[\"sampling\"] == sampling) & (data[\"stopping\"] == stopping)\n",
    "                ][colname].to_numpy()\n",
    "                sort_idxs_list += [np.argsort(arr)]\n",
    "\n",
    "            locs2_sorted = locs2[sort_idxs_list[1].argsort()][sort_idxs_list[0]]\n",
    "\n",
    "            for i in range(locs1.shape[0]):\n",
    "                _x = [locs1[i, 0], locs2_sorted[i, 0]]\n",
    "                _y = [locs1[i, 1], locs2_sorted[i, 1]]\n",
    "                ax.plot(_x, _y, color=\"black\", alpha=0.1, linewidth=0.75)\n",
    "\n",
    "        ax.axhline(0, color=\"black\", linewidth=0.75, linestyle=\"--\")\n",
    "        ax.set_title(\n",
    "            f\"{use}\\nMean amplitudes {' - '.join(subtract)}\\nwindow_seconds: {window}, channels: {group}\"\n",
    "        )\n",
    "\n",
    "        # settings\n",
    "        ax.set_title(\"\")\n",
    "        ax.set_ylabel(\"µV\")\n",
    "        if ii == 0:\n",
    "            ax.set_xlabel(\"\")\n",
    "            leg = ax.get_legend()\n",
    "            leg.get_texts()[0].set_text(\"Partial control\")\n",
    "            leg.get_texts()[1].set_text(\"Full control\")\n",
    "        else:\n",
    "            ax.get_legend().remove()\n",
    "\n",
    "        ax.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"])\n",
    "        ax.set_xlabel(ax.get_xlabel(), labelpad=labelpad)\n",
    "\n",
    "sns.despine(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T09:19:04.196918Z",
     "start_time": "2021-06-02T09:19:03.958872Z"
    }
   },
   "outputs": [],
   "source": [
    "# analyze in mixed model\n",
    "for use in (\"p1\", \"n1\"):\n",
    "    tmp = erps[use]\n",
    "    stats = pingouin.mixed_anova(\n",
    "        data=tmp,\n",
    "        dv=\"mean_amp\",\n",
    "        within=\"sampling\",\n",
    "        subject=\"subject\",\n",
    "        between=\"stopping\",\n",
    "    )\n",
    "    print(use)\n",
    "    display(stats.round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample ERPs early/mid/late AF/AV/YF/YV SPLIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:11.127180Z",
     "start_time": "2021-06-02T10:55:22.105572Z"
    }
   },
   "outputs": [],
   "source": [
    "condi_early_mid_late_erps = {\n",
    "    \"{}/{}/{}\".format(*ii): list()\n",
    "    for ii in itertools.product(sampling_styles, stopping_styles, timings)\n",
    "}\n",
    "condi_early_mid_late_erps = get_erps_by_dict(\n",
    "    condi_early_mid_late_erps, name_templates, subjects, **preproc_settings\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:26.525730Z",
     "start_time": "2021-06-02T10:56:26.486875Z"
    }
   },
   "outputs": [],
   "source": [
    "# also convert to array\n",
    "condi_early_mid_late_arr = dict2arr(condi_early_mid_late_erps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:26.795606Z",
     "start_time": "2021-06-02T10:56:26.791580Z"
    }
   },
   "outputs": [],
   "source": [
    "split_condi_early_mid_late_erps = split_dict(condi_early_mid_late_erps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Topo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:31.449006Z",
     "start_time": "2021-06-02T10:56:27.571800Z"
    }
   },
   "outputs": [],
   "source": [
    "condi = \"active/fixed\"\n",
    "fig, axs = plot_topos(split_condi_early_mid_late_erps[condi], vminmax=2)\n",
    "fig.suptitle(\"sample ERP {}\".format(condi))\n",
    "display(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:39.555147Z",
     "start_time": "2021-06-02T10:56:35.385964Z"
    }
   },
   "outputs": [],
   "source": [
    "condi = \"active/variable\"\n",
    "fig, axs = plot_topos(split_condi_early_mid_late_erps[condi], vminmax=2)\n",
    "fig.suptitle(\"sample ERP {}\".format(condi))\n",
    "display(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:43.648279Z",
     "start_time": "2021-06-02T10:56:39.556381Z"
    }
   },
   "outputs": [],
   "source": [
    "condi = \"yoked/fixed\"\n",
    "fig, axs = plot_topos(split_condi_early_mid_late_erps[condi], vminmax=2)\n",
    "fig.suptitle(\"sample ERP {}\".format(condi))\n",
    "display(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:47.225200Z",
     "start_time": "2021-06-02T10:56:43.649725Z"
    }
   },
   "outputs": [],
   "source": [
    "condi = \"yoked/variable\"\n",
    "fig, axs = plot_topos(split_condi_early_mid_late_erps[condi], vminmax=2)\n",
    "fig.suptitle(\"sample ERP {}\".format(condi))\n",
    "display(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## P3 ERP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:47.229639Z",
     "start_time": "2021-06-02T10:56:47.226586Z"
    }
   },
   "outputs": [],
   "source": [
    "cmap = np.array(sns.color_palette(\"Set1\"))[[0, 4, -1], :]\n",
    "cmap = sns.color_palette(cmap)\n",
    "cmap = ListedColormap(cmap.as_hex())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Overview plot**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:49.154171Z",
     "start_time": "2021-06-02T10:56:47.230693Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 2, figsize=(10, 8), sharex=True, sharey=True)\n",
    "\n",
    "for i, key in enumerate(\n",
    "    [\"active/fixed\", \"yoked/fixed\", \"active/variable\", \"yoked/variable\"]\n",
    "):\n",
    "\n",
    "    erp_dict = split_condi_early_mid_late_erps[key]\n",
    "\n",
    "    colors = np.linspace(0, 1, len(erp_dict.keys()))\n",
    "\n",
    "    ax = axs.flat[i]\n",
    "    fig, ax = plot_single_erp(\n",
    "        erp_dict, p3_grp, \"p3\", colors, cmap=cmap, ylim=single_panel_ylim, ax=ax\n",
    "    )\n",
    "\n",
    "    ax.axvspan(*p3_window, color=\"black\", alpha=0.1)\n",
    "\n",
    "    ax.set_title(ax.get_title() + \" \" + key)\n",
    "\n",
    "display(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:50.055643Z",
     "start_time": "2021-06-02T10:56:49.155362Z"
    }
   },
   "outputs": [],
   "source": [
    "include_yoked_lines_in_plot = False\n",
    "\n",
    "with sns.plotting_context(context=\"talk\", font_scale=1):\n",
    "\n",
    "    fig, axs = plt.subplots(2, 1, figsize=(6, 8), sharex=True, sharey=True)\n",
    "    fig.tight_layout(h_pad=3)\n",
    "\n",
    "    for i, (key1, key2) in enumerate(\n",
    "        [(\"active/fixed\", \"yoked/fixed\"), (\"active/variable\", \"yoked/variable\")]\n",
    "    ):\n",
    "\n",
    "        erp_dict = split_condi_early_mid_late_erps[key1]\n",
    "        erp_dict2 = split_condi_early_mid_late_erps[key2]\n",
    "\n",
    "        colors = np.linspace(0, 1, len(erp_dict.keys()))\n",
    "\n",
    "        ax = axs.flat[i]\n",
    "        fig, ax = plot_single_erp(\n",
    "            erp_dict, p3_grp, \"p3\", colors, cmap=cmap, ylim=single_panel_ylim, ax=ax\n",
    "        )\n",
    "\n",
    "        # optionally plot the yoked lines as well\n",
    "        if include_yoked_lines_in_plot:\n",
    "            fig, ax = plot_single_erp(\n",
    "                erp_dict2,\n",
    "                p3_grp,\n",
    "                \"p3\",\n",
    "                colors,\n",
    "                cmap=cmap,\n",
    "                ylim=single_panel_ylim,\n",
    "                ax=ax,\n",
    "                linestyles=dict(zip(erp_dict2.keys(), (\"--\", \"--\", \"--\"))),\n",
    "            )\n",
    "\n",
    "        ax.axvspan(*p3_window, color=\"black\", alpha=0.1)\n",
    "\n",
    "        ax.set_title({0: \"Partial control\", 1: \"Full control\"}[i])\n",
    "        if i == 0:\n",
    "            ax.set_xlabel(\"\")\n",
    "\n",
    "\n",
    "# remove duplicate legends\n",
    "objs = fig.get_children()\n",
    "# objs[0].remove()  # cannot remove\n",
    "# objs[1].remove()  # upper axes\n",
    "# objs[2].remove()  # lower axes\n",
    "objs[3].remove()  # upper montage plot\n",
    "objs[4].remove()  # upper cbar\n",
    "# objs[5].remove()  # lower montage plot\n",
    "# objs[6].remove()  # lower cbar\n",
    "\n",
    "display(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the **MEANS** from the conditions (given a window)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:50.138322Z",
     "start_time": "2021-06-02T10:56:50.056914Z"
    }
   },
   "outputs": [],
   "source": [
    "# Extract p3 means into dataframe format\n",
    "p3conditiming_df = arr2df(\n",
    "    condi_early_mid_late_arr,\n",
    "    p3_grp,\n",
    "    p3_window,\n",
    "    condi_early_mid_late_erps.keys(),\n",
    "    ch_names,\n",
    "    times,\n",
    ")\n",
    "\n",
    "# rename key\n",
    "p3conditiming_df = p3conditiming_df.rename(columns={\"key\": \"sampling-stopping-timing\"})\n",
    "p3conditiming_df[\"sampling-stopping-timing\"]\n",
    "\n",
    "# split up colnames\n",
    "tmp = pd.DataFrame(\n",
    "    [ii.split(\"/\") for ii in p3conditiming_df[\"sampling-stopping-timing\"]]\n",
    ")\n",
    "tmp.columns = [\"sampling\", \"stopping\", \"timing\"]\n",
    "p3conditiming_df = pd.concat([p3conditiming_df, tmp], axis=1)\n",
    "\n",
    "# Need to fix subj_repr to be different between stopping fixed and variable\n",
    "vals = p3conditiming_df.loc[p3conditiming_df[\"stopping\"] == \"fixed\"][\n",
    "    \"subj_repr\"\n",
    "].to_numpy()\n",
    "p3conditiming_df.loc[p3conditiming_df[\"stopping\"] == \"fixed\", \"subj_repr\"] = vals + 20\n",
    "\n",
    "# Turn subj_repr into subject column\n",
    "_all_subjs = []\n",
    "for _cond in p3conditiming_df[\"sampling-stopping-timing\"].unique():\n",
    "    # get current data to work on\n",
    "    _curdat = p3conditiming_df[p3conditiming_df[\"sampling-stopping-timing\"] == _cond]\n",
    "\n",
    "    # turn subj reprs into idx 0 to 19, assert everything is ordered\n",
    "    _subj_reprs = _curdat[\"subj_repr\"].to_numpy()\n",
    "    _subj_reprs2, idx = np.unique(_subj_reprs, return_index=True)\n",
    "    np.testing.assert_array_equal(_subj_reprs, _subj_reprs2)\n",
    "\n",
    "    # according to the study setup, \"fixed\" participants were sub 1, 3, 5, ...\n",
    "    # and \"variable\" participants were sub 2, 4, 6, ...\n",
    "    if \"fixed\" in _cond:\n",
    "        assert \"variable\" not in _cond\n",
    "        _subjs = np.arange(1, 40, 2)[idx]\n",
    "    elif \"variable\" in _cond:\n",
    "        assert \"fixed\" not in _cond\n",
    "        _subjs = np.arange(2, 41, 2)[idx]\n",
    "    else:\n",
    "        raise ValueError(f\"Unexpected _cond: '{_cond}'\")\n",
    "\n",
    "    _all_subjs += list(_subjs)\n",
    "\n",
    "\n",
    "# add subject column to df\n",
    "if \"subject\" in p3conditiming_df.columns:\n",
    "    np.testing.assert_array_equal(p3conditiming_df[\"subject\"], _all_subjs)\n",
    "else:\n",
    "    p3conditiming_df.insert(0, \"subject\", _all_subjs)\n",
    "\n",
    "# sanity check that the subj_repr to subject mapping is unique\n",
    "assert p3conditiming_df[[\"subject\", \"subj_repr\"]].drop_duplicates().shape[0] == 40\n",
    "\n",
    "p3conditiming_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:51.081585Z",
     "start_time": "2021-06-02T10:56:50.140308Z"
    }
   },
   "outputs": [],
   "source": [
    "grid = sns.catplot(\n",
    "    x=\"timing\",\n",
    "    y=\"mean_amp\",\n",
    "    data=p3conditiming_df,\n",
    "    kind=\"point\",\n",
    "    ci=68,\n",
    "    col=\"sampling\",\n",
    "    row=\"stopping\",\n",
    "    color=\"black\",\n",
    ")\n",
    "\n",
    "grid.fig.suptitle(\"window: {}s\\n{}\".format(p3_window, p3_grp), y=1.05)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Early - late \"accumulation\"  P3 ERP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:51.095548Z",
     "start_time": "2021-06-02T10:56:51.083129Z"
    }
   },
   "outputs": [],
   "source": [
    "df_accumul = p3conditiming_df[p3conditiming_df[\"timing\"].isin([\"early\", \"late\"])].copy()\n",
    "\n",
    "# capitalize factor levels\n",
    "df_accumul.loc[:, \"sampling\"] = df_accumul[\"sampling\"].str.capitalize().to_list()\n",
    "df_accumul.loc[:, \"stopping\"] = df_accumul[\"stopping\"].str.capitalize().to_list()\n",
    "\n",
    "df_accumul.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:51.629843Z",
     "start_time": "2021-06-02T10:56:51.097346Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"notebook\", font_scale=1):\n",
    "    fig, axs = plt.subplots(2, 2, sharex=True, sharey=True)\n",
    "\n",
    "    # one plot per task\n",
    "    for i, (meta, grp) in enumerate(df_accumul.groupby([\"sampling\", \"stopping\"])):\n",
    "        ax = axs.flat[i]\n",
    "\n",
    "        task = meta[0][0] + meta[1][0]\n",
    "\n",
    "        sns.pointplot(x=\"timing\", y=\"mean_amp\", data=grp, ax=ax, color=\"black\")\n",
    "        sns.swarmplot(x=\"timing\", y=\"mean_amp\", data=grp, ax=ax, size=3)\n",
    "\n",
    "        ax.set_title(task)\n",
    "\n",
    "    fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:51.644856Z",
     "start_time": "2021-06-02T10:56:51.631001Z"
    }
   },
   "outputs": [],
   "source": [
    "# compute difference first - second\n",
    "first = \"late\"\n",
    "second = \"early\"\n",
    "a = df_accumul[df_accumul[\"timing\"] == first][\n",
    "    [\"subject\", \"sampling\", \"stopping\"]\n",
    "].reset_index(drop=True)\n",
    "b = df_accumul[df_accumul[\"timing\"] == second][\n",
    "    [\"subject\", \"sampling\", \"stopping\"]\n",
    "].reset_index(drop=True)\n",
    "pd.testing.assert_frame_equal(a, b)\n",
    "\n",
    "mean_amp_diff = (\n",
    "    df_accumul[df_accumul[\"timing\"] == first][\"mean_amp\"].to_numpy()\n",
    "    - df_accumul[df_accumul[\"timing\"] == second][\"mean_amp\"].to_numpy()\n",
    ")\n",
    "\n",
    "colname = f\"mean_amp_diff_{first}_minus_{second}\"\n",
    "a[colname] = mean_amp_diff\n",
    "df_accumul = a.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:51.653678Z",
     "start_time": "2021-06-02T10:56:51.646513Z"
    }
   },
   "outputs": [],
   "source": [
    "fname = \"eeg_accumulation.csv\"\n",
    "fname = op.join(plotdir, fname)\n",
    "\n",
    "df_accumul.to_csv(fname, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:51.960052Z",
     "start_time": "2021-06-02T10:56:51.655296Z"
    }
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"notebook\", font_scale=1.3):\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    data = df_accumul\n",
    "    order = [\"Active\", \"Yoked\"]\n",
    "    hue_order = [\"Fixed\", \"Variable\"]\n",
    "\n",
    "    sns.pointplot(\n",
    "        x=\"sampling\",\n",
    "        y=colname,\n",
    "        hue=\"stopping\",\n",
    "        data=data,\n",
    "        dodge=True,\n",
    "        ci=68,\n",
    "        order=order,\n",
    "        hue_order=hue_order,\n",
    "        ax=ax,\n",
    "        markers=\"o\",\n",
    "    )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=\"sampling\",\n",
    "        y=colname,\n",
    "        hue=\"stopping\",\n",
    "        data=data,\n",
    "        order=order,\n",
    "        hue_order=hue_order,\n",
    "        ax=ax,\n",
    "        dodge=True,\n",
    "        size=3,\n",
    "    )\n",
    "\n",
    "    # add legend\n",
    "    # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "    legend_elements = [\n",
    "        Line2D([0], [0], color=sns.color_palette()[0], marker=\"o\", label=\"Fixed\"),\n",
    "        Line2D([0], [0], color=sns.color_palette()[1], marker=\"o\", label=\"Variable\"),\n",
    "    ]\n",
    "\n",
    "    ax.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower right\",\n",
    "        prop={\"size\": 13},\n",
    "        framealpha=1,\n",
    "    )\n",
    "\n",
    "    ax.set_ylabel(f\"amplidute difference\\n{first}-{second}\")\n",
    "\n",
    "    # https://stackoverflow.com/a/51157346/5201771\n",
    "    # 2-4 3-5\n",
    "    for i_dots, (idx0, idx1) in enumerate([(2, 4), (3, 5)]):\n",
    "        locs1 = ax.get_children()[idx0].get_offsets()\n",
    "        locs2 = ax.get_children()[idx1].get_offsets()\n",
    "\n",
    "        # Need to sort locs, so data corresponds\n",
    "        sort_idxs_list = []\n",
    "        sampling = order[i_dots]\n",
    "        for stopping in hue_order:\n",
    "            arr = data[(data[\"sampling\"] == sampling) & (data[\"stopping\"] == stopping)][\n",
    "                colname\n",
    "            ].to_numpy()\n",
    "            sort_idxs_list += [np.argsort(arr)]\n",
    "\n",
    "        locs2_sorted = locs2[sort_idxs_list[1].argsort()][sort_idxs_list[0]]\n",
    "\n",
    "        for i in range(locs1.shape[0]):\n",
    "            _x = [locs1[i, 0], locs2_sorted[i, 0]]\n",
    "            _y = [locs1[i, 1], locs2_sorted[i, 1]]\n",
    "            ax.plot(_x, _y, color=\"black\", alpha=0.1)\n",
    "\n",
    "    ax.axhline(0, color=\"black\")\n",
    "    ax.set_title(f\"Mean P3 amplitudes from window: {p3_window}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Calc the stats**\n",
    "\n",
    "anova + posthoc tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:56:53.009769Z",
     "start_time": "2021-06-02T10:56:52.914466Z"
    }
   },
   "outputs": [],
   "source": [
    "model = pingouin.mixed_anova(\n",
    "    data=df_accumul,\n",
    "    dv=colname,\n",
    "    within=\"sampling\",\n",
    "    between=\"stopping\",\n",
    "    subject=\"subject\",\n",
    ")\n",
    "\n",
    "model.round(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:57:16.901721Z",
     "start_time": "2021-06-02T10:57:16.773968Z"
    }
   },
   "outputs": [],
   "source": [
    "stats = pingouin.pairwise_ttests(\n",
    "    data=df_accumul,\n",
    "    dv=colname,\n",
    "    within=\"sampling\",\n",
    "    between=\"stopping\",\n",
    "    subject=\"subject\",\n",
    "    padjust=\"bonf\",\n",
    "    within_first=False,\n",
    "    effsize=\"cohen\",\n",
    ")\n",
    "display(stats.round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "T-tests for each condition > 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:57:40.255979Z",
     "start_time": "2021-06-02T10:57:40.169488Z"
    }
   },
   "outputs": [],
   "source": [
    "for meta, grp in df_accumul.groupby([\"sampling\", \"stopping\"]):\n",
    "    stats = pingouin.ttest(grp[colname], y=0)\n",
    "    print(f\"\\n{meta}, pval = {stats['p-val'][0] * 4:.3f} (bonf. corrected)\")\n",
    "\n",
    "    display(stats.round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Correlations with Accuracy and BNT**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-06-02T10:57:43.627375Z",
     "start_time": "2021-06-02T10:57:43.053308Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load behavioral accuracy data\n",
    "fname_beh_acc = \"beh_accs.csv\"\n",
    "accuracy_df = pd.read_csv(fname_beh_acc)\n",
    "\n",
    "# also load BNT data, as part of the overall behavioral DF\n",
    "fname_beh_overall = \"behavioral_data.csv\"\n",
    "beh_df = pd.read_csv(fname_beh_overall)\n",
    "\n",
    "# Drop DESC task for now\n",
    "accuracy_df = accuracy_df[accuracy_df[\"task\"] != \"DESC\"].reset_index(drop=True)\n",
    "beh_df = beh_df[beh_df[\"task\"] != \"DESC\"].reset_index(drop=True)\n",
    "\n",
    "# get sampling and stopping from task\n",
    "accuracy_df[\"sampling\"] = accuracy_df[\"task\"].str[0].map({\"A\": \"active\", \"Y\": \"yoked\"})\n",
    "accuracy_df[\"stopping\"] = (\n",
    "    accuracy_df[\"task\"].str[1].map({\"F\": \"fixed\", \"V\": \"variable\"})\n",
    ")\n",
    "\n",
    "beh_df[\"sampling\"] = beh_df[\"task\"].str[0].map({\"A\": \"active\", \"Y\": \"yoked\"})\n",
    "beh_df[\"stopping\"] = beh_df[\"task\"].str[1].map({\"F\": \"fixed\", \"V\": \"variable\"})\n",
    "\n",
    "# work on \"experienced\" accuracy\n",
    "accuracy_type = \"experienced\"\n",
    "accuracy_df = accuracy_df[accuracy_df[\"accuracy_type\"] == accuracy_type]\n",
    "\n",
    "# work on behavioral df to make it focussed on bnt only\n",
    "bnt_df = beh_df[[\"subject\", \"stopping\", \"sampling\", \"bnt_quartile\"]].drop_duplicates()\n",
    "\n",
    "# make a df for doing the correlations\n",
    "_tmp4corr = df_accumul.copy()\n",
    "_tmp4corr[\"sampling\"] = _tmp4corr[\"sampling\"].str.lower()\n",
    "_tmp4corr[\"stopping\"] = _tmp4corr[\"stopping\"].str.lower()\n",
    "_tmp4corr = _tmp4corr.merge(accuracy_df, on=[\"subject\", \"sampling\", \"stopping\"])\n",
    "_tmp4corr = _tmp4corr.merge(bnt_df, on=[\"subject\", \"sampling\", \"stopping\"])\n",
    "\n",
    "print(\"------------\\nBNT\\n\")\n",
    "# Correlate CPP and BNT\n",
    "for this_task in _tmp4corr[\"task\"].unique():\n",
    "\n",
    "    print(this_task)\n",
    "    _curdat = _tmp4corr[_tmp4corr[\"task\"] == this_task]\n",
    "    stat = pingouin.correlation.corr(\n",
    "        _curdat[colname], _curdat[\"bnt_quartile\"], method=\"kendall\", tail=\"two-sided\"\n",
    "    )\n",
    "    display(stat)\n",
    "\n",
    "print(\"------------\\nACCURACY\\n\")\n",
    "# correlate CPP and accuracy\n",
    "for this_task in _tmp4corr[\"task\"].unique():\n",
    "\n",
    "    print(this_task)\n",
    "    _curdat = _tmp4corr[_tmp4corr[\"task\"] == this_task]\n",
    "    stat = pingouin.correlation.corr(\n",
    "        _curdat[colname], _curdat[\"correct_choice\"], method=\"pearson\", tail=\"two-sided\"\n",
    "    )\n",
    "    display(stat)"
   ]
  },
  {
   "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.7.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "1267.2px",
    "left": "21px",
    "top": "111.483px",
    "width": "443.967px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}