publication_plots.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports and settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:30:14.825274Z",
     "start_time": "2021-05-25T14:30:11.529266Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import mne\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pdir\n",
    "import pingouin\n",
    "import seaborn as sns\n",
    "import statsmodels\n",
    "from matplotlib import transforms\n",
    "from matplotlib.lines import Line2D\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from mpl_toolkits.axes_grid1.inset_locator import (\n",
    "    BboxConnector,\n",
    "    BboxPatch,\n",
    "    TransformedBbox,\n",
    "    inset_axes,\n",
    ")\n",
    "\n",
    "from clusterperm import evaluate_significance, return_observed_clusters\n",
    "from utils import (\n",
    "    BIDS_ROOT,\n",
    "    arr2df,\n",
    "    dict2arr,\n",
    "    find_time_idxs,\n",
    "    get_erps_by_dict,\n",
    "    get_model_rdm,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# https://matplotlib.org/stable/gallery/text_labels_and_annotations/font_family_rc_sgskip.html\n",
    "plt.rcParams.update(\n",
    "    {\n",
    "        \"font.family\": \"sans-serif\",\n",
    "        \"font.sans-serif\": [\"Liberation Sans\"],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:30:14.830317Z",
     "start_time": "2021-05-25T14:30:14.827094Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# settings\n",
    "\n",
    "plotdir = os.path.join(BIDS_ROOT, \"code\", \"publication_plots\")\n",
    "\n",
    "for directory in [plotdir]:\n",
    "    os.makedirs(directory, exist_ok=True)\n",
    "\n",
    "dina4_width_inch = 8.27"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:30:14.966338Z",
     "start_time": "2021-05-25T14:30:14.959965Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# settings\n",
    "DPI = 600\n",
    "\n",
    "plotting_context = dict(context=\"talk\", font_scale=1)\n",
    "\n",
    "swarmsize = 2  # for strip and swarmplot how big the dots\n",
    "\n",
    "ci = 68  # the error bars, 68 ~ SEM\n",
    "\n",
    "axhline_args = dict(color=\"black\", linestyle=\"--\", linewidth=1)\n",
    "\n",
    "pointdodge = 0.2  # for pointplot, how much means dodge\n",
    "pointscale = 3\n",
    "pointmarkers = \".\"  # for pointplot which marker style\n",
    "pointerrwidth = 3\n",
    "pointlinewidth = axhline_args[\"linewidth\"]\n",
    "pointcapwidth = 0.1\n",
    "\n",
    "labelpad = 12\n",
    "\n",
    "sampling_order = [\"active\", \"yoked\"]\n",
    "stopping_order = [\"fixed\", \"variable\"]\n",
    "\n",
    "subj_line_settings = dict(color=\"black\", alpha=0.1, linewidth=0.75)\n",
    "\n",
    "panel_letter_kwargs = dict(\n",
    "    x=-0.2,\n",
    "    y=1.15,\n",
    "    horizontalalignment=\"center\",\n",
    "    verticalalignment=\"center\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T09:19:51.088961Z",
     "start_time": "2021-04-27T09:19:50.475463Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# start new figure\n",
    "with sns.plotting_context(**plotting_context):\n",
    "    fig, axs = plt.subplots(\n",
    "        1, 3, figsize=(dina4_width_inch + 3, 4), gridspec_kw={\"width_ratios\": [3, 3, 3]}\n",
    "    )\n",
    "    fig.tight_layout()\n",
    "\n",
    "    # plot panel letters\n",
    "    # see: https://gitter.im/matplotlib/matplotlib?at=5fb53257c6fe0131d40227bb\n",
    "    for ax, label in zip(axs.flat, (i for i in \"bcd\")):\n",
    "\n",
    "        x = -0.5\n",
    "        y = 0.3\n",
    "        offset = transforms.ScaledTranslation(x, y, fig.dpi_scale_trans)\n",
    "        transform = ax.transAxes + offset\n",
    "\n",
    "        ax.text(\n",
    "            s=label,\n",
    "            x=0,\n",
    "            y=1,\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            fontweight=\"bold\",\n",
    "            fontsize=plt.rcParams[\"font.size\"] + 2,\n",
    "            transform=transform,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:11.946636Z",
     "start_time": "2021-03-29T10:32:11.737013Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load data\n",
    "fname = \"beh_accuracies.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "beh_accuracies = pd.read_csv(fname)\n",
    "\n",
    "for factor in [\"sampling\", \"stopping\"]:\n",
    "    beh_accuracies[factor] = beh_accuracies[factor].str.lower()\n",
    "\n",
    "beh_accuracies[\"accuracy\"] = beh_accuracies[\"correct_choice\"].to_numpy()\n",
    "\n",
    "# plot\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    x = \"sampling\"\n",
    "    order = sampling_order\n",
    "    colname = \"accuracy\"\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    ax = axs[0]\n",
    "    data = beh_accuracies\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.pointplot(\n",
    "            x=x,\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            ci=ci,\n",
    "            dodge=pointdodge,\n",
    "            ax=ax,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax,\n",
    "        dodge=pointdodge,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    ax.get_legend().remove()\n",
    "    ax.set_xlabel(\"\")  # \"Sampling\"\n",
    "\n",
    "    # connect subj dots with lines\n",
    "    # https://stackoverflow.com/a/63171175/5201771\n",
    "    if pointdodge:\n",
    "        _idxs_to_pick = [(2, 4), (3, 5)]\n",
    "    else:\n",
    "        _idxs_to_pick = [(2, 3)]\n",
    "    for i_dots, (idx0, idx1) in enumerate(_idxs_to_pick):\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, **subj_line_settings)\n",
    "\n",
    "    ax.yaxis.set_major_locator(plt.MaxNLocator(5))\n",
    "    ax.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# PANEL FOR REVIEWER - NOT TO BE IN PAPER\n",
    "# --------------------------------------------\n",
    "\n",
    "# load data\n",
    "# fname = \"decision_weights.csv\"\n",
    "fname = \"decision_weights_logreg_tasks.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "decision_weights = pd.read_csv(fname)\n",
    "\n",
    "# plot\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    x = \"timing\"\n",
    "    y = \"slope\"\n",
    "# load data\n",
    "# fname = \"decision_weights.csv\"\n",
    "fname = \"decision_weights_logreg_tasks.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "decision_weights = pd.read_csv(fname)\n",
    "\n",
    "# plot\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    x = \"timing\"\n",
    "    y = \"slope\"\n",
    "\n",
    "    # cannot pass order to lineplot, so sanity check data\n",
    "    order = [\"early\", \"mid\", \"late\"]\n",
    "    assert decision_weights[\"timing\"][0:3].to_list() == order\n",
    "\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    style = \"sampling\"\n",
    "    style_order = sampling_order\n",
    "    _fig, _ax = plt.subplots()\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.lineplot(\n",
    "            x=\"timing\",\n",
    "            y=\"slope\",\n",
    "            hue=\"stopping\",\n",
    "            hue_order=hue_order,\n",
    "            style=style,\n",
    "            style_order=style_order,\n",
    "            data=decision_weights[decision_weights[\"stopping\"] == \"variable\"],\n",
    "            ci=None,\n",
    "            ax=_ax,\n",
    "            legend=False,\n",
    "            palette={\n",
    "                \"variable\": sns.color_palette()[1],\n",
    "                \"fixed\": sns.color_palette()[0],\n",
    "            },\n",
    "        )\n",
    "\n",
    "    _ax.axhline(0, **axhline_args)\n",
    "    _ax.set_ylabel(\"decision weight\")\n",
    "    _ax.set_xlabel(\"\")  # \"sample position\"\n",
    "\n",
    "    with sns.plotting_context(\n",
    "        context=plotting_context[\"context\"],\n",
    "        font_scale=plotting_context[\"font_scale\"] - 0.2,\n",
    "    ):\n",
    "        _text = _ax.text(\n",
    "            x=0.5,\n",
    "            y=0.9,\n",
    "            s=\"sample position\",\n",
    "            ha=\"center\",\n",
    "            transform=_ax.transAxes,\n",
    "        )\n",
    "\n",
    "    # add legend\n",
    "    # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "    legend_elements = []\n",
    "    for linestyle in [\"-\", \"--\"]:\n",
    "        for i, label in enumerate([\"full control\"]):\n",
    "\n",
    "            if linestyle == \"--\":\n",
    "                label = \"yoked\"\n",
    "\n",
    "            color = sns.color_palette()[i+1]\n",
    "            legend_elements.append(\n",
    "                Line2D(\n",
    "                    [0], [0], color=color, marker=None, label=label, linestyle=linestyle\n",
    "                ),\n",
    "            )\n",
    "\n",
    "    legend = _fig.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower center\",\n",
    "        bbox_to_anchor=(0.5, 0.9),\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=2,\n",
    "        title=\"\",  # \"Sampling\"\n",
    "    )\n",
    "\n",
    "    _ax.yaxis.set_major_locator(plt.MaxNLocator(5))\n",
    "    \n",
    "_fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:12.234118Z",
     "start_time": "2021-03-29T10:32:11.948174Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# load data\n",
    "# fname = \"decision_weights.csv\"\n",
    "fname = \"decision_weights_logreg_tasks.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "decision_weights = pd.read_csv(fname)\n",
    "\n",
    "# plot\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    x = \"timing\"\n",
    "    y = \"slope\"\n",
    "# load data\n",
    "# fname = \"decision_weights.csv\"\n",
    "fname = \"decision_weights_logreg_tasks.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "decision_weights = pd.read_csv(fname)\n",
    "\n",
    "# plot\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    x = \"timing\"\n",
    "    y = \"slope\"\n",
    "\n",
    "    # cannot pass order to lineplot, so sanity check data\n",
    "    order = [\"early\", \"mid\", \"late\"]\n",
    "    assert decision_weights[\"timing\"][0:3].to_list() == order\n",
    "\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    style = \"sampling\"\n",
    "    style_order = sampling_order\n",
    "    ax = axs[1]\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.lineplot(\n",
    "            x=\"timing\",\n",
    "            y=\"slope\",\n",
    "            hue=\"stopping\",\n",
    "            hue_order=hue_order,\n",
    "            style=style,\n",
    "            style_order=style_order,\n",
    "            data=decision_weights,\n",
    "            ci=68,\n",
    "            ax=ax,\n",
    "            legend=False,\n",
    "            palette={\n",
    "                \"variable\": sns.color_palette()[1],\n",
    "                \"fixed\": sns.color_palette()[0],\n",
    "            },\n",
    "        )\n",
    "\n",
    "    ax.axhline(0, **axhline_args)\n",
    "    ax.set_ylabel(\"decision weight\")\n",
    "    ax.set_xlabel(\"\")  # \"sample position\"\n",
    "\n",
    "    with sns.plotting_context(\n",
    "        context=plotting_context[\"context\"],\n",
    "        font_scale=plotting_context[\"font_scale\"] - 0.2,\n",
    "    ):\n",
    "        _text = ax.text(\n",
    "            x=0.5,\n",
    "            y=0.9,\n",
    "            s=\"sample position\",\n",
    "            ha=\"center\",\n",
    "            transform=ax.transAxes,\n",
    "        )\n",
    "\n",
    "    # add legend\n",
    "    # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "    legend_elements = []\n",
    "    for linestyle in [\"-\", \"--\"]:\n",
    "        for i, label in enumerate([\"partial control\", \"full control\"]):\n",
    "\n",
    "            if linestyle == \"--\":\n",
    "                label = \"yoked\"\n",
    "\n",
    "            color = sns.color_palette()[i]\n",
    "            legend_elements.append(\n",
    "                Line2D(\n",
    "                    [0], [0], color=color, marker=None, label=label, linestyle=linestyle\n",
    "                ),\n",
    "            )\n",
    "\n",
    "    legend = fig.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower center\",\n",
    "        bbox_to_anchor=(0.5, 1.05),\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=2,\n",
    "        title=\"\",  # \"Sampling\"\n",
    "    )\n",
    "\n",
    "    ax.yaxis.set_major_locator(plt.MaxNLocator(5))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:12.457417Z",
     "start_time": "2021-03-29T10:32:12.235248Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load data\n",
    "# fname = \"beh_recency.csv\"\n",
    "fname = \"beh_recency_logreg.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "beh_recency = pd.read_csv(fname)\n",
    "\n",
    "for factor in [\"sampling\", \"stopping\"]:\n",
    "    beh_recency[factor] = beh_recency[factor].str.lower()\n",
    "\n",
    "# reconstruct variables from df\n",
    "potential_colnames = [col for col in beh_recency.columns if col.startswith(\"slope_\")]\n",
    "colname = [col for col in potential_colnames if \"_minus_\" in col][0]\n",
    "first = colname.split(\"_\")[1]\n",
    "second = colname.split(\"_\")[3]\n",
    "\n",
    "\n",
    "# plot\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    ax = axs.flat[2]\n",
    "    x = \"sampling\"\n",
    "    order = sampling_order\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    data = beh_recency\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.pointplot(\n",
    "            x=x,\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            dodge=pointdodge,\n",
    "            ci=ci,\n",
    "            ax=ax,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax,\n",
    "        dodge=pointdodge,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    ax.axhline(0, **axhline_args)\n",
    "\n",
    "    ax.set_xlabel(\"\")  # \"Sampling\"\n",
    "    ax.set_ylabel(f\"{first} - {second}\")\n",
    "    ax.get_legend().remove()\n",
    "\n",
    "    # connect subj dots with lines\n",
    "    # https://stackoverflow.com/a/63171175/5201771\n",
    "    if pointdodge:\n",
    "        _idxs_to_pick = [(2, 4), (3, 5)]\n",
    "    else:\n",
    "        _idxs_to_pick = [(2, 3)]\n",
    "    for i_dots, (idx0, idx1) in enumerate(_idxs_to_pick):\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, **subj_line_settings)\n",
    "\n",
    "    ax.yaxis.set_major_locator(plt.MaxNLocator(5))\n",
    "    ax.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## tweaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:12.464703Z",
     "start_time": "2021-03-29T10:32:12.458816Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# make 1c y axis a bit larger than need be\n",
    "# so that 1c does not visually dominate\n",
    "ylims = axs.flat[1].get_ylim()\n",
    "axs.flat[1].set_ylim(ylims[0], ylims[1] * 1.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:12.472720Z",
     "start_time": "2021-03-29T10:32:12.466203Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# make sure \"zero line\" is aligned between 1c and 1d\n",
    "# see: https://stackoverflow.com/a/65824524/5201771\n",
    "ax1 = axs.flat[1]\n",
    "ax2 = axs.flat[-1]\n",
    "\n",
    "ax1_ylims = ax1.axes.get_ylim()\n",
    "ax1_yratio = ax1_ylims[0] / ax1_ylims[1]\n",
    "\n",
    "ax2_ylims = ax2.axes.get_ylim()\n",
    "ax2_yratio = ax2_ylims[0] / ax2_ylims[1]\n",
    "\n",
    "if ax1_yratio < ax2_yratio:\n",
    "    ax2.set_ylim(bottom=ax2_ylims[1] * ax1_yratio)\n",
    "else:\n",
    "    ax1.set_ylim(bottom=ax1_ylims[1] * ax2_yratio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:12.481313Z",
     "start_time": "2021-03-29T10:32:12.474291Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# to make the xticks \"narrower\" / squeeze them horizontally\n",
    "# at the expense of having whitespace\n",
    "ax = axs.flat[1]\n",
    "ax.set_xlim(-0.6, 2.7)\n",
    "\n",
    "# same game for the other panels\n",
    "ax = axs.flat[0]\n",
    "ax.set_xlim(-0.75, 1.75)\n",
    "\n",
    "ax = axs.flat[2]\n",
    "ax.set_xlim(-0.75, 1.75)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finish"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T10:32:12.980989Z",
     "start_time": "2021-03-29T10:32:12.483195Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "sns.despine(fig)\n",
    "fig.align_labels()\n",
    "\n",
    "# Will then be combined with fig1a.pdf to form fig1.pdf\n",
    "# fig1.pdf can then easily be converted to TIFF format\n",
    "fname = \"fig1bcd.pdf\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "fig.savefig(fname, bbox_extra_artists=[legend], bbox_inches=\"tight\")\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:18.724507Z",
     "start_time": "2021-03-29T15:01:18.304462Z"
    },
    "tags": []
   },
   "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",
    "}\n",
    "\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "    # equivalent to: fig, axs = plt.subplots(2, 2, figsize=(dina4_width_inch, 8))\n",
    "    # ... but allows for sharing y-axis across individual subplots\n",
    "    fig = plt.figure(figsize=(dina4_width_inch, 8))\n",
    "    ax1 = fig.add_subplot(2, 2, 1)\n",
    "    ax2 = fig.add_subplot(2, 2, 2)\n",
    "    ax3 = fig.add_subplot(2, 2, 3)\n",
    "    ax4 = fig.add_subplot(2, 2, 4, sharey=None)\n",
    "\n",
    "    axs = np.array(((ax1, ax2), (ax3, ax4)))\n",
    "    fig.tight_layout(h_pad=3, w_pad=3)\n",
    "\n",
    "\n",
    "# plot panel letters\n",
    "# see: https://gitter.im/matplotlib/matplotlib?at=5fb53257c6fe0131d40227bb\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    for ax, label in zip(axs.flat, (i for i in \"abcd\")):\n",
    "\n",
    "        x = -0.5\n",
    "        y = 0.3\n",
    "        offset = transforms.ScaledTranslation(x, y, fig.dpi_scale_trans)\n",
    "        transform = ax.transAxes + offset\n",
    "\n",
    "        ax.text(\n",
    "            s=label,\n",
    "            x=0,\n",
    "            y=1,\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            fontweight=\"bold\",\n",
    "            fontsize=plt.rcParams[\"font.size\"] + 2,\n",
    "            transform=transform,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:18.728807Z",
     "start_time": "2021-03-29T15:01:18.725836Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# IO: Where to find the data\n",
    "# Where to find epochs\n",
    "fname_epo_template = os.path.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-03-29T15:01:18.737982Z",
     "start_time": "2021-03-29T15:01:18.730131Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "colors = {\"variable\": sns.color_palette()[1], \"fixed\": sns.color_palette()[0]}\n",
    "linestyles = {\"active\": \"-\", \"yoked\": \"--\"}\n",
    "\n",
    "p1_n1_grp = [\"O2\", \"PO4\", \"PO8\", \"PO10\"]\n",
    "opo_grp = p1_n1_grp + [\"O1\", \"PO3\", \"PO7\", \"PO9\"]\n",
    "\n",
    "p1_window = (0.08, 0.13)\n",
    "n1_window = (0.14, 0.20)\n",
    "\n",
    "subtract = (\"left\", \"right\")\n",
    "\n",
    "sampling_styles = (\"active\", \"yoked\")\n",
    "stopping_styles = (\"fixed\", \"variable\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:23.955029Z",
     "start_time": "2021-03-29T15:01:19.052700Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "double_subtraction = True\n",
    "\n",
    "try:\n",
    "    soverall_erps\n",
    "    print(\"Data already loaded\")\n",
    "except NameError:\n",
    "    soverall_erps = {\n",
    "        \"{0}/{1}\".format(*ii): list()\n",
    "        for ii in itertools.product(sampling_styles, stopping_styles)\n",
    "    }\n",
    "    soverall_erps = get_erps_by_dict(\n",
    "        soverall_erps,\n",
    "        name_templates,\n",
    "        subjects,\n",
    "        subtract=subtract,\n",
    "        double_subtraction=double_subtraction,\n",
    "        **preproc_settings\n",
    "    )\n",
    "\n",
    "if double_subtraction:\n",
    "    # if we do a double subtraction, we need an extra dummy ERP to get\n",
    "    # channel positions\n",
    "    tmp = get_erps_by_dict({\"all\": list()}, name_templates, [1], **preproc_settings)\n",
    "    coord_dummy_erp = tmp[\"all\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:23.961575Z",
     "start_time": "2021-03-29T15:01:23.957331Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# obtain an erp data object to access the digitizied electrode positions\n",
    "# and ch_names, etc\n",
    "tmperp = soverall_erps[\"active/fixed\"][0].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:24.033229Z",
     "start_time": "2021-03-29T15:01:23.963935Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# convert ERP dict to data frame\n",
    "soverall_arr = dict2arr(soverall_erps)\n",
    "\n",
    "# array to data frame\n",
    "arr = np.column_stack(\n",
    "    list(map(np.ravel, np.meshgrid(*map(np.arange, soverall_arr.shape), indexing=\"ij\")))\n",
    "    + [soverall_arr.ravel()]\n",
    ")\n",
    "df_soverall = pd.DataFrame(\n",
    "    arr, columns=[\"task\", \"ch_names\", \"time_s\", \"subj_repr\", \"amplitude\"]\n",
    ")\n",
    "\n",
    "# make data frame pretty\n",
    "df_soverall[\"ch_names\"] = df_soverall[\"ch_names\"].map(\n",
    "    dict(zip(range(len(tmperp.ch_names)), tmperp.ch_names))\n",
    ")\n",
    "df_soverall[\"task\"] = df_soverall[\"task\"].map(\n",
    "    dict(zip(range(len(soverall_erps)), list(soverall_erps.keys())))\n",
    ")\n",
    "df_soverall[\"time_s\"] = df_soverall[\"time_s\"].map(\n",
    "    dict(zip(range(len(tmperp.times)), tmperp.times))\n",
    ")\n",
    "\n",
    "df_soverall[\"sampling\"] = [i[0] for i in df_soverall[\"task\"].str.split(\"/\")]\n",
    "df_soverall[\"stopping\"] = [i[1] for i in df_soverall[\"task\"].str.split(\"/\")]\n",
    "\n",
    "df_soverall[\"subj_repr\"] = df_soverall[\"subj_repr\"].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:24.065460Z",
     "start_time": "2021-03-29T15:01:24.034790Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if double_subtraction:\n",
    "    df_soverall_chsel = df_soverall[df_soverall[\"ch_names\"].isin([\"opo\"])]\n",
    "else:\n",
    "    df_soverall_chsel = df_soverall[df_soverall[\"ch_names\"].isin(p1_n1_grp)]\n",
    "\n",
    "amps = []\n",
    "for meta, grp in df_soverall_chsel.groupby(\"ch_names\"):\n",
    "    amps.append(grp[\"amplitude\"])\n",
    "amps = np.mean(np.stack(amps), axis=0)\n",
    "\n",
    "grp[\"amplitude\"] = amps\n",
    "grp = grp.drop(\"ch_names\", axis=1)\n",
    "df_soverall_chsel = grp.copy()\n",
    "df_soverall_chsel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.335330Z",
     "start_time": "2021-03-29T15:01:24.067302Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ci = 68  # 68  #set to \"sd\" to save plotting time\n",
    "ax = axs.flat[0]\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    hue_order = stopping_order\n",
    "\n",
    "    sns.lineplot(\n",
    "        x=\"time_s\",\n",
    "        y=\"amplitude\",\n",
    "        data=df_soverall_chsel,\n",
    "        hue=\"stopping\",\n",
    "        hue_order=hue_order,\n",
    "        style=\"sampling\",\n",
    "        ax=ax,\n",
    "        ci=ci,\n",
    "        legend=False,\n",
    "    )\n",
    "\n",
    "    # make ylims symmetric\n",
    "    maxylim = np.abs(ax.get_ylim()).max()\n",
    "    ax.set_ylim((-maxylim, maxylim))\n",
    "\n",
    "    # plot the selected electrodes\n",
    "    montage_m1 = mne.channels.make_standard_montage(\"easycap-M1\")\n",
    "    axins = inset_axes(ax, width=\"30%\", height=\"30%\", loc=\"upper left\")\n",
    "\n",
    "    if double_subtraction:\n",
    "        coord_dummy_erp.set_montage(montage_m1)\n",
    "        coord_dummy_erp.copy().pick_channels(opo_grp).plot_sensors(\n",
    "            kind=\"topomap\", title=\"\", axes=axins, show=False\n",
    "        )\n",
    "    else:\n",
    "        tmperp.set_montage(montage_m1)\n",
    "        tmperp.copy().pick_channels(p1_n1_grp).plot_sensors(\n",
    "            kind=\"topomap\", title=\"\", axes=axins, show=False\n",
    "        )\n",
    "    axins.collections[0].set_sizes([1])\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "    ax.axvline(0, ymax=0.69, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    ax.set_xlabel(\"time (s)\")\n",
    "    ax.set_ylabel(\"left - right (µV)\\n(double subtraction)\")\n",
    "\n",
    "\n",
    "# plot gray shadings\n",
    "ax.axvspan(*p1_window, ymin=0.5, ymax=1.0, color=\"black\", alpha=0.1)\n",
    "ax.axvspan(*n1_window, ymin=0.0, ymax=0.5, color=\"black\", alpha=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.749266Z",
     "start_time": "2021-03-29T15:01:42.336473Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get the correct orderings in plotting\n",
    "# NOTE: pd.unique provides in the order \"as found\"\n",
    "stopping_order = list(pd.unique(df_soverall_chsel[\"stopping\"]))\n",
    "sampling_order = list(pd.unique(df_soverall_chsel[\"sampling\"]))\n",
    "\n",
    "# The following legend is encoded according to these assumptions:\n",
    "assert stopping_order == hue_order\n",
    "assert sampling_order == [\"active\", \"yoked\"]\n",
    "\n",
    "# add legend\n",
    "# https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "legend_elements = []\n",
    "for linestyle in [\"-\", \"--\"]:\n",
    "    for i, label in enumerate([\"partial control\", \"full control\"]):\n",
    "\n",
    "        if linestyle == \"--\":\n",
    "            label = \"yoked\"\n",
    "\n",
    "        color = sns.color_palette()[i]\n",
    "        legend_elements.append(\n",
    "            Line2D(\n",
    "                [0], [0], color=color, marker=None, label=label, linestyle=linestyle\n",
    "            ),\n",
    "        )\n",
    "\n",
    "# define invisible bogus item to tweak legend\n",
    "bogus_item = Line2D([0], [0], color=\"white\", marker=None, label=\"\", linestyle=linestyle)\n",
    "legend_elements.insert(2, bogus_item)\n",
    "\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "    ax = axs.flat[0]\n",
    "\n",
    "    legend1 = fig.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower center\",\n",
    "        bbox_to_anchor=(0.5, 0.975),\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=2,\n",
    "        title=\"\",  # \"sampling\"\n",
    "    )\n",
    "\n",
    "sns.despine(ax=ax)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.756663Z",
     "start_time": "2021-03-29T15:01:42.751979Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "p3_grp = [\"Cz\", \"C1\", \"C2\", \"CPz\", \"CP1\", \"CP2\", \"CP3\", \"CP4\", \"Pz\", \"P1\", \"P2\"]\n",
    "\n",
    "p3_window = (0.3, 0.6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.766760Z",
     "start_time": "2021-03-29T15:01:42.759680Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "timings = (\"early\", \"mid\", \"late\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.781939Z",
     "start_time": "2021-03-29T15:01:42.769139Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fname = \"eeg_accumulation.csv\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "eeg_accumulation = pd.read_csv(fname)\n",
    "colname = \"mean_amp_diff_late_minus_early\"\n",
    "assert eeg_accumulation.columns[-1] == colname\n",
    "\n",
    "# make sampling and stopping levels lowercase\n",
    "eeg_accumulation[\"sampling\"] = eeg_accumulation[\"sampling\"].str.lower()\n",
    "eeg_accumulation[\"stopping\"] = eeg_accumulation[\"stopping\"].str.lower()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.806956Z",
     "start_time": "2021-03-29T15:01:42.785035Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "eeg_accumulation = eeg_accumulation.sort_values(by=[\"subject\"])\n",
    "eeg_accumulation.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:42.815329Z",
     "start_time": "2021-03-29T15:01:42.808748Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "eeg_acc_full = eeg_accumulation[\n",
    "    (eeg_accumulation[\"sampling\"] == \"active\")\n",
    "    & (eeg_accumulation[\"stopping\"] == \"variable\")\n",
    "]\n",
    "\n",
    "eeg_acc_partial = eeg_accumulation[\n",
    "    (eeg_accumulation[\"sampling\"] == \"active\")\n",
    "    & (eeg_accumulation[\"stopping\"] == \"fixed\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:48.432353Z",
     "start_time": "2021-03-29T15:01:42.817325Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "    ax = axs.flat[1]\n",
    "\n",
    "    x = \"sampling\"\n",
    "    order = sampling_order\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    data = eeg_accumulation\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.pointplot(\n",
    "            x=x,\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            dodge=pointdodge,\n",
    "            ci=ci,\n",
    "            ax=ax,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax,\n",
    "        dodge=True,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    # connect subj dots\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, **subj_line_settings)\n",
    "\n",
    "    # make ylims symmetric\n",
    "    # maxylim = np.abs(ax.get_ylim()).max()\n",
    "    # ax.set_ylim((-maxylim, maxylim))\n",
    "    ax.set_ylim(None, 5)\n",
    "\n",
    "    # plot axins\n",
    "    # plot the selected electrodes\n",
    "    _ = get_erps_by_dict({\"all\": list()}, name_templates, [1], **preproc_settings)[\n",
    "        \"all\"\n",
    "    ][0]\n",
    "\n",
    "    axins = inset_axes(\n",
    "        ax,\n",
    "        width=\"30%\",\n",
    "        height=\"30%\",\n",
    "        loc=\"upper left\",\n",
    "    )\n",
    "\n",
    "    montage_m1 = mne.channels.make_standard_montage(\"easycap-M1\")\n",
    "    _.set_montage(montage_m1)\n",
    "    _.copy().pick_channels(p3_grp).plot_sensors(\n",
    "        kind=\"topomap\", title=\"\", axes=axins, show=False\n",
    "    )\n",
    "    axins.collections[0].set_sizes([1])\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    # remove unnecessary legend\n",
    "    ax.get_legend().remove()\n",
    "\n",
    "    ax.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"])\n",
    "    ax.set_xlabel(\"\")\n",
    "    ax.set_ylabel(\"late - early samples (µV)\")\n",
    "    assert colname.endswith(\"late_minus_early\")\n",
    "\n",
    "\n",
    "sns.despine(ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:48.438498Z",
     "start_time": "2021-03-29T15:01:48.433772Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    eml_partial\n",
    "    print(\"Data already loaded\")\n",
    "except NameError:\n",
    "    eml_partial = {\"active/fixed/{0}\".format(ii): list() for ii in timings}\n",
    "    eml_partial = get_erps_by_dict(\n",
    "        eml_partial, name_templates, subjects, **preproc_settings\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:48.452803Z",
     "start_time": "2021-03-29T15:01:48.439897Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# obtain an erp data object to access the digitizied electrode positions\n",
    "# and ch_names, etc\n",
    "tmperp = eml_partial[\"active/fixed/early\"][0].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:48.868357Z",
     "start_time": "2021-03-29T15:01:48.454043Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# convert ERP dict to data frame\n",
    "eml_partial_arr = dict2arr(eml_partial)\n",
    "\n",
    "# array to data frame\n",
    "arr = np.column_stack(\n",
    "    list(\n",
    "        map(\n",
    "            np.ravel, np.meshgrid(*map(np.arange, eml_partial_arr.shape), indexing=\"ij\")\n",
    "        )\n",
    "    )\n",
    "    + [eml_partial_arr.ravel()]\n",
    ")\n",
    "df_eml_partial = pd.DataFrame(\n",
    "    arr, columns=[\"timing\", \"ch_names\", \"time_s\", \"subj_repr\", \"amplitude\"]\n",
    ")\n",
    "\n",
    "\n",
    "# make data frame pretty\n",
    "df_eml_partial[\"ch_names\"] = df_eml_partial[\"ch_names\"].map(\n",
    "    dict(zip(range(len(tmperp.ch_names)), tmperp.ch_names))\n",
    ")\n",
    "df_eml_partial[\"timing\"] = df_eml_partial[\"timing\"].map(\n",
    "    dict(zip(range(len(eml_partial)), list(eml_partial.keys())))\n",
    ")\n",
    "df_eml_partial[\"time_s\"] = df_eml_partial[\"time_s\"].map(\n",
    "    dict(zip(range(len(tmperp.times)), tmperp.times))\n",
    ")\n",
    "\n",
    "df_eml_partial[\"subj_repr\"] = df_eml_partial[\"subj_repr\"].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:01:48.960179Z",
     "start_time": "2021-03-29T15:01:48.869757Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_eml_partial_chsel = df_eml_partial[df_eml_partial[\"ch_names\"].isin(p3_grp)]\n",
    "amps = []\n",
    "for meta, grp in df_eml_partial_chsel.groupby(\"ch_names\"):\n",
    "    amps.append(grp[\"amplitude\"])\n",
    "amps = np.mean(np.stack(amps), axis=0)\n",
    "\n",
    "grp[\"amplitude\"] = amps\n",
    "grp = grp.drop(\"ch_names\", axis=1)\n",
    "df_eml_partial_chsel = grp.copy()\n",
    "df_eml_partial_chsel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:01.330242Z",
     "start_time": "2021-03-29T15:01:48.961624Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "    ax = axs.flat[2]\n",
    "\n",
    "    sns.lineplot(\n",
    "        x=\"time_s\",\n",
    "        y=\"amplitude\",\n",
    "        data=df_eml_partial_chsel,\n",
    "        hue=\"timing\",\n",
    "        palette=\"crest_r\",\n",
    "        ax=ax,\n",
    "        ci=ci,\n",
    "        legend=False,\n",
    "    )\n",
    "\n",
    "    # make ylims symmetric\n",
    "    maxylim = np.abs(ax.get_ylim()).max()\n",
    "    ax.set_ylim((-maxylim, maxylim))\n",
    "\n",
    "    # plot the selected electrodes\n",
    "    axins = inset_axes(ax, width=\"30%\", height=\"30%\", loc=\"upper left\")\n",
    "    montage_m1 = mne.channels.make_standard_montage(\"easycap-M1\")\n",
    "    tmperp.set_montage(montage_m1)\n",
    "    tmperp.copy().pick_channels(p3_grp).plot_sensors(\n",
    "        kind=\"topomap\", title=\"\", axes=axins, show=False\n",
    "    )\n",
    "    axins.collections[0].set_sizes([1])\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "    ax.axvline(0, ymax=0.69, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    # plot P3 window span\n",
    "    ax.axvspan(*p3_window, color=\"black\", alpha=0.1)\n",
    "\n",
    "    ax.set_xlabel(\"time (s)\")\n",
    "    ax.set_ylabel(\"amplitude (µV)\")\n",
    "\n",
    "sns.despine(ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:01.340309Z",
     "start_time": "2021-03-29T15:02:01.331267Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get the correct orderings in plotting\n",
    "# NOTE: pd.unique provides in the order \"as found\"\n",
    "timing_order = list(pd.unique(df_eml_partial_chsel[\"timing\"]))\n",
    "\n",
    "if timing_order[0].startswith(\"active\"):\n",
    "    timing_order = [t.split(\"/\")[-1] for t in timing_order]\n",
    "\n",
    "# The following legend is encoded according to these assumptions:\n",
    "assert tuple(timing_order) == timings\n",
    "\n",
    "# add legend\n",
    "legend_elements = []\n",
    "for i, label in enumerate(timings):\n",
    "\n",
    "    color = sns.color_palette(\"crest_r\", n_colors=3)[i]\n",
    "    legend_elements.append(\n",
    "        Line2D(\n",
    "            [0],\n",
    "            [0],\n",
    "            color=color,\n",
    "            marker=None,\n",
    "            label=label,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "ax = axs.flat[2]\n",
    "\n",
    "with sns.plotting_context(\n",
    "    context=plotting_context[\"context\"], font_scale=plotting_context[\"font_scale\"] - 0.2\n",
    "):\n",
    "    legend = fig.legend(\n",
    "        handles=legend_elements[::-1],\n",
    "        loc=\"lower right\",\n",
    "        bbox_to_anchor=(1.2, -0.05),\n",
    "        bbox_transform=ax.axes.transAxes,\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=1,\n",
    "        title=\"sample position\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:01.345569Z",
     "start_time": "2021-03-29T15:02:01.341478Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    early_mid_late_erps\n",
    "    print(\"Data already loaded\")\n",
    "except NameError:\n",
    "    early_mid_late_erps = {\"active/variable/{0}\".format(ii): list() for ii in timings}\n",
    "    early_mid_late_erps = get_erps_by_dict(\n",
    "        early_mid_late_erps, name_templates, subjects, **preproc_settings\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:01.743388Z",
     "start_time": "2021-03-29T15:02:01.347112Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# convert ERP dict to data frame\n",
    "eml_arr = dict2arr(early_mid_late_erps)\n",
    "\n",
    "# array to data frame\n",
    "arr = np.column_stack(\n",
    "    list(map(np.ravel, np.meshgrid(*map(np.arange, eml_arr.shape), indexing=\"ij\")))\n",
    "    + [eml_arr.ravel()]\n",
    ")\n",
    "df_eml = pd.DataFrame(\n",
    "    arr, columns=[\"timing\", \"ch_names\", \"time_s\", \"subj_repr\", \"amplitude\"]\n",
    ")\n",
    "\n",
    "# make data frame pretty\n",
    "df_eml[\"ch_names\"] = df_eml[\"ch_names\"].map(\n",
    "    dict(zip(range(len(tmperp.ch_names)), tmperp.ch_names))\n",
    ")\n",
    "df_eml[\"timing\"] = df_eml[\"timing\"].map(\n",
    "    dict(zip(range(len(early_mid_late_erps)), list(early_mid_late_erps.keys())))\n",
    ")\n",
    "df_eml[\"time_s\"] = df_eml[\"time_s\"].map(\n",
    "    dict(zip(range(len(tmperp.times)), tmperp.times))\n",
    ")\n",
    "\n",
    "df_eml[\"subj_repr\"] = df_eml[\"subj_repr\"].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:01.824544Z",
     "start_time": "2021-03-29T15:02:01.744543Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_eml_chsel = df_eml[df_eml[\"ch_names\"].isin(p3_grp)]\n",
    "amps = []\n",
    "for meta, grp in df_eml_chsel.groupby(\"ch_names\"):\n",
    "    amps.append(grp[\"amplitude\"])\n",
    "amps = np.mean(np.stack(amps), axis=0)\n",
    "\n",
    "grp[\"amplitude\"] = amps\n",
    "grp = grp.drop(\"ch_names\", axis=1)\n",
    "df_eml_chsel = grp.copy()\n",
    "df_eml_chsel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:14.083546Z",
     "start_time": "2021-03-29T15:02:01.825635Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "    ax = axs.flat[3]\n",
    "\n",
    "    sns.lineplot(\n",
    "        x=\"time_s\",\n",
    "        y=\"amplitude\",\n",
    "        data=df_eml_chsel,\n",
    "        hue=\"timing\",\n",
    "        palette=\"crest_r\",\n",
    "        ax=ax,\n",
    "        ci=ci,\n",
    "        legend=False,\n",
    "    )\n",
    "\n",
    "    # make ylims symmetric\n",
    "    maxylim = np.abs(ax.get_ylim()).max()\n",
    "    ax.set_ylim((-maxylim, maxylim))\n",
    "\n",
    "    # plot the selected electrodes\n",
    "    axins = inset_axes(ax, width=\"30%\", height=\"30%\", loc=\"upper left\")\n",
    "    montage_m1 = mne.channels.make_standard_montage(\"easycap-M1\")\n",
    "    tmperp.set_montage(montage_m1)    \n",
    "    tmperp.copy().pick_channels(p3_grp).plot_sensors(\n",
    "        kind=\"topomap\", title=\"\", axes=axins, show=False\n",
    "    )\n",
    "    axins.collections[0].set_sizes([1])\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "    ax.axvline(0, ymax=0.69, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    # plot P3 window span\n",
    "    ax.axvspan(*p3_window, color=\"black\", alpha=0.1)\n",
    "\n",
    "    ax.set_xlabel(\"time (s)\")\n",
    "    ax.set_ylabel(\"amplitude (µV)\")\n",
    "\n",
    "sns.despine(ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:14.091531Z",
     "start_time": "2021-03-29T15:02:14.086335Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = axs.flat[3]\n",
    "\n",
    "with sns.plotting_context(\n",
    "    context=plotting_context[\"context\"], font_scale=plotting_context[\"font_scale\"] - 0.2\n",
    "):\n",
    "    legend3 = fig.legend(\n",
    "        handles=legend_elements[::-1],\n",
    "        loc=\"lower right\",\n",
    "        bbox_to_anchor=(1.2, -0.05),\n",
    "        bbox_transform=ax.axes.transAxes,\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=1,\n",
    "        title=\"sample position\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## tweaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:14.102863Z",
     "start_time": "2021-03-29T15:02:14.093082Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# align axes labels\n",
    "fig.align_ylabels(axs=axs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:14.110789Z",
     "start_time": "2021-03-29T15:02:14.104108Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# tweak\n",
    "ax1.yaxis.set_major_locator(plt.MaxNLocator(6))\n",
    "ax2.yaxis.set_major_locator(plt.MaxNLocator(4))\n",
    "ax4.set_ylim(-2.5, 2.5)\n",
    "\n",
    "with sns.plotting_context(\n",
    "    context=plotting_context[\"context\"], font_scale=plotting_context[\"font_scale\"] - 0.2\n",
    "):\n",
    "    ax2.text(x=0.5, y=1, s=\"CPP (300-600 ms)\", transform=ax2.transAxes, ha=\"center\")\n",
    "    ax3.text(x=0.5, y=1, s=\"partial control\", transform=ax3.transAxes, ha=\"center\")\n",
    "    ax4.text(x=0.5, y=1, s=\"full control\", transform=ax4.transAxes, ha=\"center\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-29T15:02:14.116812Z",
     "start_time": "2021-03-29T15:02:14.112013Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# to make the xticks \"narrower\" / squeeze them horizontally\n",
    "# at the expense of having whitespace\n",
    "ax2.set_xlim(-0.75, 1.75)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finish"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fname = \"fig2\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "for ext in [\".pdf\", \".png\", \".tif\"]:\n",
    "    fname_save = fname + ext\n",
    "\n",
    "    dpi = DPI if ext == \".tif\" else None\n",
    "\n",
    "    fig.savefig(\n",
    "        fname_save,\n",
    "        bbox_extra_artists=[legend1, legend3],\n",
    "        bbox_inches=\"tight\",\n",
    "        dpi=dpi,\n",
    "        pil_kwargs={\"compression\": \"tiff_lzw\"},\n",
    "    )\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:31:21.651293Z",
     "start_time": "2021-05-25T14:31:20.810211Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    fig = plt.figure(figsize=(dina4_width_inch, 8))\n",
    "\n",
    "    # debugging\n",
    "    # left, bottom, width, height\n",
    "    # fake = fig.add_axes([0, 0, 1, 1])\n",
    "    # fig.set_facecolor(\"lightgray\")\n",
    "\n",
    "    ax3a_upper = fig.add_axes([0.00, 0.4, 0.3, 0.3])\n",
    "    ax3a_lower = fig.add_axes(\n",
    "        [0.00, 0.0, 0.3, 0.3], sharey=ax3a_upper, sharex=ax3a_upper\n",
    "    )\n",
    "\n",
    "    ax3b = fig.add_axes([0.65, 0.40, 0.50, 0.30])\n",
    "\n",
    "    ax3c_left = fig.add_axes([0.65, 0.00, 0.20, 0.25])\n",
    "    ax3c_right = fig.add_axes(\n",
    "        [0.95, 0.00, 0.20, 0.25], sharey=ax3c_left, sharex=ax3c_left\n",
    "    )\n",
    "\n",
    "    # cbar axes\n",
    "    cax3a_upper = fig.add_axes([0.3 + 0.005, 0.4, 0.025, 0.3])\n",
    "    cax3a_lower = fig.add_axes([0.3 + 0.005, 0.0, 0.025, 0.3])\n",
    "\n",
    "    # need to collect some objs to later form a tight box around the plot\n",
    "    extra_artists = []\n",
    "\n",
    "    # plot panel letters\n",
    "    # see: https://gitter.im/matplotlib/matplotlib?at=5fb53257c6fe0131d40227bb\n",
    "    axs = np.array([ax3a_upper, ax3a_lower, ax3b, ax3c_left, ax3c_right])\n",
    "    for ax, label in zip(axs.flat, (i for i in [\"a\", \"\", \"b\", \"c\", \"\"])):\n",
    "\n",
    "        x = -0.5\n",
    "        y = 0.3\n",
    "        offset = transforms.ScaledTranslation(x, y, fig.dpi_scale_trans)\n",
    "        transform = ax.transAxes + offset\n",
    "\n",
    "        ax.text(\n",
    "            s=label,\n",
    "            x=0,\n",
    "            y=1,\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            fontweight=\"bold\",\n",
    "            fontsize=plt.rcParams[\"font.size\"] + 2,\n",
    "            transform=transform,\n",
    "        )\n",
    "\n",
    "        if label in [\"a\", \"b\", \"c\"]:\n",
    "            extra_artists.append(ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:31:22.624188Z",
     "start_time": "2021-05-25T14:31:22.618245Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# on which results to work\n",
    "rsa_results_path = (\n",
    "    \"cv-False_flip-False_rsa-pearson_dist-euclidean_half-both_exclude-Identity_\"\n",
    "    \"mnn-FalseFalseFalse_ec-False_c-(0.6, 1.6)t--0.8s-150b-(None, 0)s-True_outcome\",    \n",
    ")\n",
    "\n",
    "# whether to get results from CPT or to try locally\n",
    "# NOTE: **expects the ARC CPT server from MPIB to be mounted**\n",
    "get_from_cpt = False\n",
    "if get_from_cpt:\n",
    "    CPT07_ip = \"141.14.159.168\"\n",
    "    CPT10_ip = \"141.14.156.202\"\n",
    "    ip_to_use = CPT07_ip\n",
    "\n",
    "    # Get data for plotting\n",
    "    rsa_path = (\n",
    "        \"/run/user/1000/gvfs/\"\n",
    "        f\"sftp:host={ip_to_use},\"\n",
    "        \"user=appelhoff/vol2/appelhoff/sp_data/derivatives/rsa_9x9\"\n",
    "    )\n",
    "\n",
    "else:\n",
    "    rsa_path = os.path.join(BIDS_ROOT, \"derivatives\", \"rsa_9x9\")\n",
    "\n",
    "# build full paths\n",
    "rsa_results_path = os.path.join(\n",
    "    rsa_path,\n",
    "    \"cv-False_flip-False_rsa-pearson_dist-euclidean_half-both_exclude-Identity_\"\n",
    "    \"mnn-FalseFalseFalse_ec-False_c-(0.6, 1.6)t--0.8s-150b-(None, 0)s-True_outcome\",\n",
    ")\n",
    "    \n",
    "results = os.path.join(rsa_results_path, \"rsa_results.csv\")\n",
    "results_first_half = results.replace(\"half-both\", \"half-first\")\n",
    "results_second_half = results.replace(\"half-both\", \"half-second\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:31:22.867751Z",
     "start_time": "2021-05-25T14:31:22.864945Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "just_plot = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:31:23.123666Z",
     "start_time": "2021-05-25T14:31:23.120656Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "effect = \"orthnumberline_length\"  # orthnumberline_length, orthextremity_length\n",
    "stat_effect_to_plot = \"interaction\"  # sampling, stopping, interaction\n",
    "\n",
    "add_extremity_to_fig3 = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:31:23.412136Z",
     "start_time": "2021-05-25T14:31:23.369267Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if add_extremity_to_fig3:\n",
    "    # add axes for extremity\n",
    "    with sns.plotting_context(**plotting_context):\n",
    "        # left bottom width height\n",
    "        ax3d = fig.add_axes([0.0, -0.5, 0.3, 0.3])\n",
    "        cax3d = fig.add_axes([0.3 + 0.005, -0.5, 0.025, 0.3])\n",
    "\n",
    "        ax3e = fig.add_axes([0.65, -0.5, 0.50, 0.30])\n",
    "\n",
    "        # plot panel letters\n",
    "        # see: https://gitter.im/matplotlib/matplotlib?at=5fb53257c6fe0131d40227bb\n",
    "        axs = np.array([ax3d, ax3e])\n",
    "        for ax, label in zip(axs.flat, (i for i in \"de\")):\n",
    "\n",
    "            x = -0.5\n",
    "            y = 0.3\n",
    "            offset = transforms.ScaledTranslation(x, y, fig.dpi_scale_trans)\n",
    "            transform = ax.transAxes + offset\n",
    "\n",
    "            ax.text(\n",
    "                s=label,\n",
    "                x=0,\n",
    "                y=1,\n",
    "                ha=\"center\",\n",
    "                va=\"center\",\n",
    "                fontweight=\"bold\",\n",
    "                fontsize=plt.rcParams[\"font.size\"] + 2,\n",
    "                transform=transform,\n",
    "            )\n",
    "\n",
    "            extra_artists.append(ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:31:23.657207Z",
     "start_time": "2021-05-25T14:31:23.618391Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def _calculations(effect, stat_effect_to_plot, rsa_results_path):\n",
    "\n",
    "    # get names from effect\n",
    "    _modelname, clusterstat = effect.split(\"_\")\n",
    "    orth = True if _modelname.startswith(\"orth\") else False\n",
    "    modelname = _modelname.lstrip(\"orth\")\n",
    "\n",
    "    print(\n",
    "        f\"Working on {effect}, (orth:{orth}, modelname:{modelname})\\nStat:{stat_effect_to_plot}\"\n",
    "    )\n",
    "\n",
    "    effect_mod = \"o\" + modelname if orth else modelname\n",
    "\n",
    "    average_effect_json = os.path.join(\n",
    "        rsa_results_path,\n",
    "        \"perm_and_2x2_outputs\",\n",
    "        effect,\n",
    "        f\"average_{effect_mod}_clusters.json\",\n",
    "    )\n",
    "    with open(average_effect_json, \"r\") as fin:\n",
    "        average_effect_clusters = json.load(fin)\n",
    "\n",
    "    print(f\"{len(average_effect_clusters)} average effect clusters found.\")\n",
    "\n",
    "    # Load permutation distribution\n",
    "    cluster_distr_file = os.path.join(\n",
    "        rsa_results_path,\n",
    "        f\"clusterperm_results_2021-01-19T22:46:03.297053/model-{modelname}_stat-length_thresh-0.05_distr.npy\",\n",
    "    )\n",
    "\n",
    "    # Get column order\n",
    "    cluster_distr_col_order_file = cluster_distr_file.replace(\n",
    "        \".npy\", \"_column_order.txt\"\n",
    "    )\n",
    "    with open(cluster_distr_col_order_file, \"r\") as fin:\n",
    "        lines = fin.readlines()\n",
    "    column_order = [line.strip() for line in lines]\n",
    "    effect_order = [\"stopping\", \"sampling\", \"Interaction\"]\n",
    "    assert column_order == [i.lower() for i in effect_order]\n",
    "\n",
    "    # Load data\n",
    "    df_rsa = pd.read_csv(results)\n",
    "    df_rsa_first_half = pd.read_csv(results_first_half)\n",
    "    df_rsa_second_half = pd.read_csv(results_second_half)\n",
    "\n",
    "    # find the rsa_method that was used\n",
    "    assert (\n",
    "        df_rsa[\"method\"].nunique() == 1\n",
    "    ), f\"more than one rsa_method detected: {df_rsa['method']}\"\n",
    "    rsa_method = df_rsa[\"method\"].unique()[0]\n",
    "\n",
    "    # find the distance metric that was used\n",
    "    assert (\n",
    "        df_rsa[\"distance_metric\"].nunique() == 1\n",
    "    ), f\"more than one distance_metric detected: {df_rsa['distance_metric']}\"\n",
    "    distance_metric = df_rsa[\"distance_metric\"].unique()[0]\n",
    "\n",
    "    print(\n",
    "        f\"Using rsa_method: >> {rsa_method}\\nwith distance_metric: >> {distance_metric}\"\n",
    "    )\n",
    "\n",
    "    thresh = 0.05\n",
    "    clusterthresh = 0.05\n",
    "\n",
    "    # model\n",
    "    model = \"o_\" + modelname if orth else modelname\n",
    "    df_rsa_model_orth = df_rsa[\n",
    "        (df_rsa[\"orth\"] == orth) & (df_rsa[\"model\"] == modelname)\n",
    "    ]\n",
    "\n",
    "    assert df_rsa_model_orth[\"model\"].nunique() == 1\n",
    "\n",
    "    # find the observed clusters in the data\n",
    "    clusters_obs, models_obs = return_observed_clusters(\n",
    "        df_rsa_model_orth, thresh, pingouin.mixed_anova\n",
    "    )\n",
    "\n",
    "    # Load cluster distribution array\n",
    "    cluster_distr = np.load(cluster_distr_file)\n",
    "\n",
    "    # turn into data frame\n",
    "    cluster_distr = pd.DataFrame(cluster_distr, columns=column_order)\n",
    "\n",
    "    # evaluate the significance of the observed clusters\n",
    "    (clustersig_threshs, clusters_obs_stats, clusters_obs_sig,) = evaluate_significance(\n",
    "        cluster_distr, clusters_obs, clusterstat, clusterthresh, models_obs\n",
    "    )\n",
    "\n",
    "    times = df_rsa[\"time_s\"].unique()\n",
    "    sig_clu_to_plot = False\n",
    "\n",
    "    # Find start and stop of largest effect\n",
    "    _ = [len(i) for i in clusters_obs_sig[stat_effect_to_plot]]\n",
    "    if len(_) > 0:\n",
    "        cluster_effect_to_plot = clusters_obs_sig[stat_effect_to_plot][np.argmax(_)]\n",
    "        cluster_effect_to_plot_times = times[cluster_effect_to_plot]\n",
    "        tstart, tstop = (\n",
    "            cluster_effect_to_plot_times[0],\n",
    "            cluster_effect_to_plot_times[-1],\n",
    "        )\n",
    "        _extent = np.round(tstart, 2), np.round(tstop, 2)\n",
    "        print(f\"Cluster extent in seconds: {_extent}\")\n",
    "        sig_clu_to_plot = True\n",
    "\n",
    "    else:\n",
    "        cluster_effect_to_plot = []\n",
    "        cluster_effect_to_plot_times = np.array([])\n",
    "        print(\"No significant cluster\")\n",
    "\n",
    "    # Find start and stop of AVERAGE effect\n",
    "    # picking the largest of all average clusters\n",
    "    which_ave_cluster = \"{}\".format(\n",
    "        np.argmax([len(val[\"cluster\"]) for key, val in average_effect_clusters.items()])\n",
    "    )\n",
    "    print(\n",
    "        \"p =\",\n",
    "        np.format_float_positional(average_effect_clusters[which_ave_cluster][\"p-val\"]),\n",
    "    )\n",
    "    tmp = times[average_effect_clusters[which_ave_cluster][\"cluster\"]]\n",
    "    tstart, tstop = tmp[0], tmp[-1]\n",
    "    _extent = np.round(tstart, 2), np.round(tstop, 2)\n",
    "    print(f\"Cluster extent in seconds: {_extent}\")\n",
    "\n",
    "    return (\n",
    "        df_rsa_model_orth,\n",
    "        df_rsa_first_half,\n",
    "        df_rsa_second_half,\n",
    "        model,\n",
    "        modelname,\n",
    "        orth,\n",
    "        clusters_obs_sig,\n",
    "        cluster_effect_to_plot,\n",
    "        average_effect_clusters,\n",
    "        rsa_method,\n",
    "        times,\n",
    "        sig_clu_to_plot,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    df_rsa_model_orth,\n",
    "    df_rsa_first_half,\n",
    "    df_rsa_second_half,\n",
    "    model,\n",
    "    modelname,\n",
    "    orth,\n",
    "    clusters_obs_sig,\n",
    "    cluster_effect_to_plot,\n",
    "    average_effect_clusters,\n",
    "    rsa_method,\n",
    "    times,\n",
    "    sig_clu_to_plot,\n",
    ") = _calculations(effect, stat_effect_to_plot, rsa_results_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 3a upper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:32:42.542969Z",
     "start_time": "2021-05-25T14:32:42.539616Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def prep_to_plot(rdm):\n",
    "    \"\"\"Remove upper triangle from rdm.\"\"\"\n",
    "    tri_idx = np.triu_indices(rdm.shape[0])\n",
    "    tmprdm = rdm.copy()\n",
    "    tmprdm = tmprdm.astype(float)\n",
    "    tmprdm[tri_idx] = np.nan\n",
    "    return tmprdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    name = model.lstrip(\"o_\")\n",
    "    rdm = get_model_rdm(name, normalize=False)\n",
    "    thismax = int(rdm.max())\n",
    "    to_plot = prep_to_plot(rdm)\n",
    "\n",
    "    im = ax3a_upper.imshow(to_plot, cmap=\"viridis\")\n",
    "\n",
    "    cbar = plt.colorbar(im, cax=cax3a_upper)\n",
    "\n",
    "    # ticklabels and labels\n",
    "    cbar.set_ticks(range(0, thismax + 2))\n",
    "    _labelmod = (1, 4, 8) if name == \"numberline\" else (0, 2, 4)\n",
    "    cbar.ax.set_yticklabels(\n",
    "        [i if i in _labelmod else \"\" for i in range(0, thismax + 2)]\n",
    "    )\n",
    "\n",
    "    _label = \"numerical distance\" if name == \"numberline\" else \"extremity distance\"\n",
    "\n",
    "    _t1 = cax3a_upper.text(\n",
    "        x=4.5,\n",
    "        y=0.5,\n",
    "        s=_label,\n",
    "        ha=\"center\",\n",
    "        va=\"center\",\n",
    "        rotation=\"vertical\",\n",
    "        color=\"black\",\n",
    "        transform=cax3a_upper.transAxes,\n",
    "    )\n",
    "\n",
    "    ax3a_upper.set_xticks(range(9))\n",
    "    ax3a_upper.set_yticks(range(9))\n",
    "\n",
    "    use_all_number_labels = True\n",
    "    if use_all_number_labels:\n",
    "        ax3a_ticklabels = range(1, 10)\n",
    "    else:\n",
    "        [i if i % 2 != 0 else \"\" for i in range(1, 10)]\n",
    "    ax3a_upper.set_xticklabels(ax3a_ticklabels)\n",
    "    ax3a_upper.set_yticklabels(ax3a_ticklabels)\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 3b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:33:02.317656Z",
     "start_time": "2021-05-25T14:32:45.497675Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    style = \"sampling\"\n",
    "    style_order = sampling_order\n",
    "\n",
    "    # Plot this group\n",
    "    sns.lineplot(\n",
    "        x=\"time_s\",\n",
    "        y=\"similarity\",\n",
    "        ci=68,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        style=style,\n",
    "        style_order=style_order,\n",
    "        legend=False,\n",
    "        data=df_rsa_model_orth,\n",
    "        ax=ax3b,\n",
    "    )\n",
    "\n",
    "    # Plot interaction effect\n",
    "    try:\n",
    "        stat_effect_significant = np.concatenate(clusters_obs_sig[stat_effect_to_plot])\n",
    "        if len(clusters_obs_sig[stat_effect_to_plot]) > 1:\n",
    "            print(\"WARNING: concatenating several clusters\")\n",
    "        else:\n",
    "            np.testing.assert_array_equal(\n",
    "                cluster_effect_to_plot, stat_effect_significant\n",
    "            )\n",
    "\n",
    "    except ValueError:\n",
    "        stat_effect_significant = []\n",
    "\n",
    "    xs = df_rsa_model_orth[\"time_s\"].unique()\n",
    "    xs = xs[stat_effect_significant]\n",
    "    ys = np.repeat(ax3b.get_ylim()[0], xs.shape[-1])\n",
    "    color = \"m\"\n",
    "    with sns.plotting_context(\"paper\"):\n",
    "        (line,) = ax3b.plot(xs, ys, marker=\"s\", color=color, linestyle=\"None\")\n",
    "\n",
    "    lab = stat_effect_to_plot\n",
    "    line.set_label(lab)\n",
    "\n",
    "    # plot general effect\n",
    "    # Get all cluster indices\n",
    "    _clusters = [d[\"cluster\"] for i, d in average_effect_clusters.items()]\n",
    "\n",
    "    # plot general numberline effect into panel\n",
    "    xs = df_rsa_model_orth[\"time_s\"].unique()\n",
    "    xs = xs[np.concatenate(_clusters)]\n",
    "    ys = np.repeat(ax3b.get_ylim()[0], xs.shape[-1])\n",
    "    color = \"k\"\n",
    "    with sns.plotting_context(\"paper\"):\n",
    "        (line,) = ax3b.plot(xs, ys, marker=\"s\", color=color, linestyle=\"None\")\n",
    "    line.set_label(f\"mean {modelname}\")\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax3b.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "    ax3b.axvline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    ax3b.set_xlabel(\"time (s)\")\n",
    "    ax3b.set_ylabel(f\"{rsa_method.capitalize()}'s r\")\n",
    "\n",
    "    sns.despine(ax=ax3b)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### legend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:33:02.962895Z",
     "start_time": "2021-05-25T14:33:02.319284Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get the correct orderings in plotting\n",
    "# NOTE: pd.unique provides in the order \"as found\"\n",
    "stopping_order = list(pd.unique(df_rsa_model_orth[\"stopping\"]))\n",
    "sampling_order = list(pd.unique(df_rsa_model_orth[\"sampling\"]))\n",
    "\n",
    "# The following legend is encoded according to these assumptions:\n",
    "assert stopping_order == hue_order\n",
    "assert sampling_order == style_order\n",
    "\n",
    "# add legend\n",
    "# https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "legend_elements = []\n",
    "for linestyle in [\"-\", \"--\"]:\n",
    "    for i, label in enumerate([\"partial control\", \"full control\"]):\n",
    "\n",
    "        if linestyle == \"--\":\n",
    "            label = \"yoked\"\n",
    "\n",
    "        color = sns.color_palette()[i]\n",
    "        legend_elements.append(\n",
    "            Line2D(\n",
    "                [0], [0], color=color, marker=None, label=label, linestyle=linestyle\n",
    "            ),\n",
    "        )\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "    legend3b = fig.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower left\",\n",
    "        bbox_to_anchor=(ax3b.get_xlim()[0], ax3b.get_ylim()[1]),\n",
    "        bbox_transform=ax3b.axes.transData,\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=2,\n",
    "        title=\"\",  # \"Sampling\"\n",
    "    )\n",
    "\n",
    "    extra_artists.append(legend3b)\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 3a lower"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:33:05.787916Z",
     "start_time": "2021-05-25T14:33:05.783003Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Determine based on what time window to plot the cluster\n",
    "\n",
    "eeg_rdm_window = \"custom\"  # largest_cluster, average_cluster, custom\n",
    "\n",
    "if eeg_rdm_window == \"largest_cluster\":\n",
    "    # Take the stat_effect_to_plot cluster\n",
    "    max_cluster = cluster_effect_to_plot.copy()\n",
    "elif eeg_rdm_window == \"average_cluster\":\n",
    "    # Get *average* cluster indices of the one we also plotted in panel b\n",
    "    max_cluster = average_effect_clusters[which_ave_cluster][\"cluster\"]\n",
    "elif eeg_rdm_window == \"custom\":\n",
    "    custom_window = (0.3, 0.6)\n",
    "    max_cluster = find_time_idxs(custom_window, times)\n",
    "    max_cluster\n",
    "else:\n",
    "    raise ValueError(\"unknown eeg_rdm_window\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:33:06.608402Z",
     "start_time": "2021-05-25T14:33:06.603635Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "subjects = range(1, 41)\n",
    "tasks = (\"AF\", \"AV\", \"YF\", \"YV\")\n",
    "\n",
    "rdms_folder_template = os.path.join(\n",
    "    rsa_results_path, \"single_subj_plots\", \"sub-{:02}_task-{}_rdm_times.npy\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:33:54.289618Z",
     "start_time": "2021-05-25T14:33:07.625935Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not just_plot:\n",
    "\n",
    "    # Go over subjects and tasks\n",
    "    # save RDMs in a large array for later averaging\n",
    "    # all_rdms is (40 x 4 x 9 x 9) containing the RDMs averaged over the\n",
    "    # max_cluster window ... half of the arrays are NaN, because each\n",
    "    # of the 40 subjs only had 2, instead of 4 tasks\n",
    "\n",
    "    all_rdms = np.full((len(subjects), len(tasks), 9, 9), np.nan)\n",
    "    did_not_find = 0\n",
    "    for isubj, subj in enumerate(subjects):\n",
    "\n",
    "        for itask, task in enumerate(tasks):\n",
    "\n",
    "            fname = rdms_folder_template.format(subj, task)\n",
    "            if not os.path.exists(fname):\n",
    "                did_not_find += 1\n",
    "                continue\n",
    "\n",
    "            rdm_times = np.load(fname)\n",
    "            rdm_average = np.mean(rdm_times[..., max_cluster], axis=-1)\n",
    "            all_rdms[isubj, itask, ...] = rdm_average\n",
    "\n",
    "    # Make RDM over all subjects and tasks\n",
    "    # 1x1 plot (1 panel)\n",
    "    subj_mean = np.nanmean(all_rdms, axis=0)\n",
    "    subj_task_mean = np.nanmean(subj_mean, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:35.767045Z",
     "start_time": "2021-05-25T14:35:35.763895Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if did_not_find == len(subjects) * 2:\n",
    "    pass  # everything as expected, 2 of 4 tasks for each subj should be skipped\n",
    "else:\n",
    "    raise ValueError(f\"unexpected number of rdm_times.npy found: {did_not_find}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:36.270200Z",
     "start_time": "2021-05-25T14:35:36.266013Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "look_only_at_full_control_for_panel_c = False\n",
    "\n",
    "if look_only_at_full_control_for_panel_c:\n",
    "    # look only at FULL CONTROL for panel C\n",
    "    print(\"Plotting panel C only with full control data!\")\n",
    "    full_control_idx = list(tasks).index(\"AV\")\n",
    "    subj_mean = np.nanmean(all_rdms[:, full_control_idx:, ...], axis=0)\n",
    "    subj_task_mean = np.nanmean(subj_mean, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:36.717459Z",
     "start_time": "2021-05-25T14:35:36.713812Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_do_normalization = False\n",
    "\n",
    "if _do_normalization:\n",
    "    # normalize distances to range 0 - 1\n",
    "    subj_task_mean = (subj_task_mean - np.nanmin(subj_task_mean)) / (\n",
    "        np.nanmax(subj_task_mean) - np.nanmin(subj_task_mean)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:37.991323Z",
     "start_time": "2021-05-25T14:35:37.287168Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    to_plot = prep_to_plot(subj_task_mean)\n",
    "\n",
    "    im = ax3a_lower.imshow(to_plot, cmap=\"viridis\")\n",
    "\n",
    "    cbar = plt.colorbar(im, cax=cax3a_lower)\n",
    "\n",
    "    # ticklabels and labels\n",
    "\n",
    "    _t2 = cax3a_lower.text(\n",
    "        x=4.55,\n",
    "        y=0.5,\n",
    "        s=\"Euclidean distance\",\n",
    "        ha=\"center\",\n",
    "        va=\"center\",\n",
    "        rotation=\"vertical\",\n",
    "        color=\"black\",\n",
    "        transform=cax3a_lower.transAxes,\n",
    "    )\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## panel 3c left+right"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:40.669718Z",
     "start_time": "2021-05-25T14:35:40.665937Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = {\n",
    "    \"first_half\": df_rsa_first_half,\n",
    "    \"second_half\": df_rsa_second_half,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:41.027454Z",
     "start_time": "2021-05-25T14:35:41.023029Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if sig_clu_to_plot:\n",
    "    # go over largest sig. cluster window\n",
    "    largest_idx = np.argmax([len(i) for i in clusters_obs_sig[stat_effect_to_plot]])\n",
    "    stat_effect_significant_largest = clusters_obs_sig[stat_effect_to_plot][largest_idx]\n",
    "    assert cluster_effect_to_plot == stat_effect_significant_largest\n",
    "else:\n",
    "    assert sig_clu_to_plot == False\n",
    "    # if no sig. cluster, go over largest average window\n",
    "    stat_effect_significant_largest = average_effect_clusters[which_ave_cluster][\n",
    "        \"cluster\"\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:41.777089Z",
     "start_time": "2021-05-25T14:35:41.722901Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dfs = {}\n",
    "for half, tmp in data.items():\n",
    "    grp = tmp[(tmp[\"orth\"] == orth) & (tmp[\"model\"] == modelname)]\n",
    "\n",
    "    # take mean across significant cluster window\n",
    "    time_idxs = grp[\"itime\"].unique()[stat_effect_significant_largest]\n",
    "\n",
    "    tmp = grp[grp[\"itime\"].isin(time_idxs)]\n",
    "    tmp = (\n",
    "        tmp.groupby([\"subject\", \"stopping\", \"sampling\", \"orth\", \"model\"])\n",
    "        .mean()\n",
    "        .reset_index()\n",
    "    )\n",
    "    tmp = tmp[[\"subject\", \"stopping\", \"sampling\", \"orth\", \"model\", \"similarity\"]]\n",
    "\n",
    "    # tmp \"save\"\n",
    "    dfs[half] = tmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:43.225435Z",
     "start_time": "2021-05-25T14:35:42.873133Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    x = \"sampling\"\n",
    "    order = sampling_order\n",
    "    colname = \"similarity\"\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    data = dfs[\"first_half\"]\n",
    "\n",
    "    # 3d\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        # plot\n",
    "        sns.pointplot(\n",
    "            x=x,\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            ci=68,\n",
    "            dodge=pointdodge,\n",
    "            ax=ax3c_left,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    # --------------------------- plot single subj dots\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax3c_left,\n",
    "        dodge=True,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    # connect subj dots\n",
    "    for i_dots, (idx0, idx1) in enumerate([(2, 4), (3, 5)]):\n",
    "        locs1 = ax3c_left.get_children()[idx0].get_offsets()\n",
    "        locs2 = ax3c_left.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",
    "            ax3c_left.plot(_x, _y, **subj_line_settings)\n",
    "\n",
    "    # ---------------------------\n",
    "\n",
    "    ax3c_left.get_legend().remove()\n",
    "    sns.despine(ax=ax3c_left)\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax3c_left.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    ax3c_left.set_xlabel(\"\")\n",
    "    ax3c_left.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"], rotation=45)\n",
    "    ax3c_left.set_ylabel(\"Pearson's r\")\n",
    "\n",
    "    ax3c_left.set_title(\"early samples\")\n",
    "\n",
    "    # 3e\n",
    "\n",
    "    data = dfs[\"second_half\"]\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        # plot\n",
    "        sns.pointplot(\n",
    "            x=\"sampling\",\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            ci=68,\n",
    "            dodge=pointdodge,\n",
    "            ax=ax3c_right,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    # --------------------------- plot single subj dots\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax3c_right,\n",
    "        dodge=True,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    # connect subj dots\n",
    "    for i_dots, (idx0, idx1) in enumerate([(2, 4), (3, 5)]):\n",
    "        locs1 = ax3c_right.get_children()[idx0].get_offsets()\n",
    "        locs2 = ax3c_right.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",
    "            ax3c_right.plot(_x, _y, **subj_line_settings)\n",
    "\n",
    "    # ---------------------------\n",
    "    ax3c_right.get_legend().remove()\n",
    "    sns.despine(ax=ax3c_right)\n",
    "\n",
    "    # Plot Zero lines\n",
    "    ax3c_right.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    ax3c_right.set_xlabel(\"\")\n",
    "    ax3c_right.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"], rotation=45)\n",
    "\n",
    "    ax3c_right.set_title(\"late samples\")\n",
    "\n",
    "    # Make ylims consistent\n",
    "    ymin = min(min(ax3c_left.get_ylim()), min(ax3c_right.get_ylim())) * 1.5\n",
    "    ymax = max(max(ax3c_left.get_ylim()), max(ax3c_right.get_ylim()))\n",
    "\n",
    "    ax3c_left.set_ylim((ymin, ymax))\n",
    "    ax3c_right.set_ylim((ymin, ymax))\n",
    "    ax3c_right.set_ylabel(\"\")\n",
    "\n",
    "    # ... but tweak ylim so that we have a negative y value\n",
    "    ylim = list(ax3c_left.get_ylim())\n",
    "    ylim[0] = -0.1\n",
    "    ax3c_left.set_ylim(ylim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:43.919990Z",
     "start_time": "2021-05-25T14:35:43.916210Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "just_plot = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## tweaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:57.870369Z",
     "start_time": "2021-05-25T14:35:57.864690Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# If figure 4, set ylim in panel b to same as it is in fig 3b\n",
    "# do same for panel c and d\n",
    "ax3b_ylim = (-0.18055231836762342, 0.3559825322474018)\n",
    "ax3c_left_ylim = (-0.1, 0.7151948522554632)\n",
    "if name == \"numberline\":\n",
    "    # assert that nothing has changes\n",
    "    _b = ax3b.get_ylim()\n",
    "    _c = ax3c_left.get_ylim()\n",
    "    assert _b == ax3b_ylim, f\"\\nPanel B\\nwant:{_b}\\ngot:{ax3b_ylim}\"\n",
    "    assert _c == ax3c_left_ylim, f\"\\nPanel C\\nwant:{_c}\\ngot:{ax3c_left_ylim}\"\n",
    "else:\n",
    "    assert name == \"extremity\"\n",
    "    ax3b.set_ylim(ax3b_ylim)\n",
    "    ax3c_left.set_ylim(ax3c_left_ylim)\n",
    "\n",
    "    # need to fix legend for panel b afterwards\n",
    "    legend3b.set_bbox_to_anchor(\n",
    "        (ax3b.get_xlim()[0], ax3b.get_ylim()[1]), transform=ax3b.transData\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:35:59.218640Z",
     "start_time": "2021-05-25T14:35:59.211937Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# to make the xticks \"narrower\" / squeeze them horizontally\n",
    "# at the expense of having whitespace\n",
    "normal = (-0.5, 1.5)\n",
    "narrow = (-0.75, 1.75)\n",
    "ax3c_left.set_xlim(normal)\n",
    "ax3c_right.set_xlim(normal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:36:00.218371Z",
     "start_time": "2021-05-25T14:36:00.214607Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# optionally keep only a part of the figure\n",
    "only_a_upper_and_b = True if name == \"extremity\" else False\n",
    "\n",
    "if name == \"extremity\" and only_a_upper_and_b:\n",
    "    ax3a_lower.remove()\n",
    "    cax3a_lower.remove()\n",
    "    ax3c_left.remove()\n",
    "    ax3c_right.remove()\n",
    "    extra_artists = extra_artists[0:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if add_extremity_to_fig3:\n",
    "\n",
    "    # get data for extremity, pick only what we need, rest assigned to _\n",
    "    (\n",
    "        df_rsa_model_orth,\n",
    "        _,\n",
    "        _,\n",
    "        model,\n",
    "        modelname,\n",
    "        orth,\n",
    "        clusters_obs_sig,\n",
    "        cluster_effect_to_plot,\n",
    "        average_effect_clusters,\n",
    "        rsa_method,\n",
    "        _,\n",
    "        _,\n",
    "    ) = _calculations(\n",
    "        effect=\"orthextremity_length\",\n",
    "        stat_effect_to_plot=\"sampling\",\n",
    "        rsa_results_path=rsa_results_path,\n",
    "    )\n",
    "\n",
    "    # need to recompute if re-run script\n",
    "    just_plot = False\n",
    "\n",
    "    with sns.plotting_context(**plotting_context):\n",
    "\n",
    "        # add panel d\n",
    "        # -----------\n",
    "        rdm = get_model_rdm(\"extremity\", normalize=False)\n",
    "        thismax = int(rdm.max())\n",
    "        to_plot = prep_to_plot(rdm)\n",
    "\n",
    "        im = ax3d.imshow(to_plot, cmap=\"viridis\")\n",
    "\n",
    "        cbar = plt.colorbar(im, cax=cax3d)\n",
    "\n",
    "        # ticklabels and labels\n",
    "        cbar.set_ticks(range(0, thismax + 2))\n",
    "        cbar.ax.set_yticklabels(\n",
    "            [i if i in (0, 2, 4) else \"\" for i in range(0, thismax + 2)]\n",
    "        )\n",
    "\n",
    "        _label = \"extremity distance\"\n",
    "\n",
    "        _t1 = cax3d.text(\n",
    "            x=4.5,\n",
    "            y=0.5,\n",
    "            s=_label,\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            rotation=\"vertical\",\n",
    "            color=\"black\",\n",
    "            transform=cax3d.transAxes,\n",
    "        )\n",
    "\n",
    "        ax3d.set_xticks(range(9))\n",
    "        ax3d.set_yticks(range(9))\n",
    "\n",
    "        use_all_number_labels = True\n",
    "        if use_all_number_labels:\n",
    "            ax3d_ticklabels = range(1, 10)\n",
    "        else:\n",
    "            [i if i % 2 != 0 else \"\" for i in range(1, 10)]\n",
    "        ax3d.set_xticklabels(ax3d_ticklabels)\n",
    "        ax3d.set_yticklabels(ax3d_ticklabels)\n",
    "\n",
    "        # add panel e\n",
    "        # -----------\n",
    "        hue = \"stopping\"\n",
    "        hue_order = stopping_order\n",
    "        style = \"sampling\"\n",
    "        style_order = sampling_order\n",
    "\n",
    "        # Plot this group\n",
    "        sns.lineplot(\n",
    "            x=\"time_s\",\n",
    "            y=\"similarity\",\n",
    "            ci=68,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            style=style,\n",
    "            style_order=style_order,\n",
    "            legend=False,\n",
    "            data=df_rsa_model_orth,\n",
    "            ax=ax3e,\n",
    "        )\n",
    "\n",
    "        # Plot interaction effect\n",
    "        try:\n",
    "            stat_effect_significant = np.concatenate(\n",
    "                clusters_obs_sig[stat_effect_to_plot]\n",
    "            )\n",
    "            if len(clusters_obs_sig[stat_effect_to_plot]) > 1:\n",
    "                print(\"WARNING: concatenating several clusters\")\n",
    "            else:\n",
    "                np.testing.assert_array_equal(\n",
    "                    cluster_effect_to_plot, stat_effect_significant\n",
    "                )\n",
    "\n",
    "        except ValueError:\n",
    "            stat_effect_significant = []\n",
    "\n",
    "        xs = df_rsa_model_orth[\"time_s\"].unique()\n",
    "        xs = xs[stat_effect_significant]\n",
    "        ys = np.repeat(ax3e.get_ylim()[0], xs.shape[-1])\n",
    "        color = \"m\"\n",
    "        with sns.plotting_context(\"paper\"):\n",
    "            (line,) = ax3e.plot(xs, ys, marker=\"s\", color=color, linestyle=\"None\")\n",
    "\n",
    "        lab = stat_effect_to_plot\n",
    "        line.set_label(lab)\n",
    "\n",
    "        # plot general effect\n",
    "        # Get all cluster indices\n",
    "        _clusters = [d[\"cluster\"] for i, d in average_effect_clusters.items()]\n",
    "\n",
    "        # plot general numberline effect into panel\n",
    "        xs = df_rsa_model_orth[\"time_s\"].unique()\n",
    "        xs = xs[np.concatenate(_clusters)]\n",
    "        ys = np.repeat(ax3e.get_ylim()[0], xs.shape[-1])\n",
    "        color = \"k\"\n",
    "        with sns.plotting_context(\"paper\"):\n",
    "            (line,) = ax3e.plot(xs, ys, marker=\"s\", color=color, linestyle=\"None\")\n",
    "        line.set_label(f\"mean {modelname}\")\n",
    "\n",
    "        # Plot Zero lines\n",
    "        ax3e.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "        ax3e.axvline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "        ax3e.set_xlabel(\"time (s)\")\n",
    "        ax3e.set_ylabel(f\"{rsa_method.capitalize()}'s r\")\n",
    "\n",
    "        sns.despine(ax=ax3e)\n",
    "\n",
    "        # adjust ylim to be same as for numberline\n",
    "        ax3e.set_ylim(ax3b_ylim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-11-18T14:21:39.495233Z",
     "start_time": "2020-11-18T14:21:39.036679Z"
    }
   },
   "source": [
    "## finish"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fname = f\"fig3_{effect}\" if not add_extremity_to_fig3 else \"fig3_combined\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "for ext in [\".pdf\", \".png\", \".tif\"]:\n",
    "    fname_save = fname + ext\n",
    "\n",
    "    dpi = DPI if ext == \".tif\" else None\n",
    "\n",
    "    fig.savefig(\n",
    "        fname_save,\n",
    "        bbox_extra_artists=extra_artists,\n",
    "        bbox_inches=\"tight\",\n",
    "        dpi=dpi,\n",
    "        pil_kwargs={\"compression\": \"tiff_lzw\"},\n",
    "    )\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-25T14:39:55.100255Z",
     "start_time": "2021-05-25T14:39:52.837619Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fname = f\"fig3_{effect}\" if not add_extremity_to_fig3 else \"fig3_combined\"\n",
    "fname = os.path.join(plotdir, fname)\n",
    "\n",
    "for ext in [\".pdf\", \".png\", \".tif\"]:\n",
    "    fname_save = fname + ext\n",
    "\n",
    "    dpi = DPI if ext == \".tif\" else None\n",
    "\n",
    "    fig.savefig(\n",
    "        fname_save,\n",
    "        bbox_extra_artists=extra_artists,\n",
    "        bbox_inches=\"tight\",\n",
    "        dpi=dpi,\n",
    "        pil_kwargs={\"compression\": \"tiff_lzw\"},\n",
    "    )\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Combined models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-27T14:28:24.934617Z",
     "start_time": "2021-04-27T14:28:24.613059Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "rdm = get_model_rdm(\"numberline\", normalize=True)\n",
    "rdm2 = get_model_rdm(\"extremity\", normalize=True)\n",
    "\n",
    "data = [rdm, rdm2, rdm + rdm2]\n",
    "fig, axs = plt.subplots(1, 3, sharex=True, sharey=True, figsize=(10, 5))\n",
    "for i, ax in enumerate(axs.flat):\n",
    "    dat = prep_to_plot(data[i])\n",
    "    ax.imshow(dat)\n",
    "fig.tight_layout()\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neurometrics figure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:39.644559Z",
     "start_time": "2021-05-02T14:07:39.641355Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "tasks = (\"AF\", \"AV\", \"YF\", \"YV\")\n",
    "model_to_use = \"both\"\n",
    "biasmax = 0.75\n",
    "kappamax = 10.0\n",
    "corr_method = \"pearson\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:40.977237Z",
     "start_time": "2021-05-02T14:07:40.191742Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    fig = plt.figure(figsize=(dina4_width_inch, 8))\n",
    "\n",
    "    width = 0.4\n",
    "    height = 0.4\n",
    "\n",
    "    horz_space = 0.3\n",
    "    vert_space = 0.1\n",
    "\n",
    "    ax1_rect = [0.0, 0.5 + vert_space, width, height]\n",
    "    ax1 = fig.add_axes(ax1_rect)\n",
    "\n",
    "    # cbar axes\n",
    "    cbar_width = 0.025\n",
    "    cbar_space = 0.005\n",
    "    cax1_rect = [\n",
    "        ax1_rect[0] + ax1_rect[2] + cbar_space,\n",
    "        ax1_rect[1],\n",
    "        cbar_width,\n",
    "        ax1_rect[-1],\n",
    "    ]\n",
    "    cax1 = fig.add_axes(cax1_rect)\n",
    "\n",
    "    ax2_rect = [\n",
    "        cbar_width + cbar_space + width + horz_space,\n",
    "        0.5 + vert_space,\n",
    "        width,\n",
    "        height,\n",
    "    ]\n",
    "    ax2 = fig.add_axes(ax2_rect)\n",
    "\n",
    "    ax3_rect = [0.0, 0.0, width, height]\n",
    "    ax3 = fig.add_axes(ax3_rect)\n",
    "\n",
    "    ax4_rect = [cbar_width + cbar_space + width + horz_space, 0.0, width, height]\n",
    "    ax4 = fig.add_axes(ax4_rect)\n",
    "\n",
    "    # need to collect some objs to later form a tight box around the plot\n",
    "    extra_artists = []\n",
    "\n",
    "    # plot panel letters\n",
    "    # see: https://gitter.im/matplotlib/matplotlib?at=5fb53257c6fe0131d40227bb\n",
    "    axs = np.array([ax1, ax2, ax3, ax4])\n",
    "    for ax, label in zip(axs.flat, (i for i in [\"a\", \"b\", \"c\", \"\"])):\n",
    "\n",
    "        x = -0.5\n",
    "        y = 0.3\n",
    "        offset = transforms.ScaledTranslation(x, y, fig.dpi_scale_trans)\n",
    "        transform = ax.transAxes + offset\n",
    "\n",
    "        ax.text(\n",
    "            s=label,\n",
    "            x=0,\n",
    "            y=1,\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            fontweight=\"bold\",\n",
    "            fontsize=plt.rcParams[\"font.size\"] + 2,\n",
    "            transform=transform,\n",
    "        )\n",
    "\n",
    "        if label in [\"a\", \"b\", \"c\", \"d\"]:\n",
    "            extra_artists.append(ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Panel 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:41.182364Z",
     "start_time": "2021-05-02T14:07:40.978436Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "tasks_to_use = \"all\"  # \"all\" or tasks\n",
    "plotdir = os.path.join(BIDS_ROOT, \"code\", \"publication_plots\")\n",
    "fname_all = os.path.join(\n",
    "    plotdir,\n",
    "    f\"nm_data_{'-'.join(tasks_to_use)}_{corr_method}_b-{biasmax}_k-{kappamax}.pickle\",\n",
    ")\n",
    "with open(fname_all, \"rb\") as fin:\n",
    "    loaddict = pickle.load(fin)\n",
    "\n",
    "bias_kappa_combis, biases, kappas, bias_0_idx, kappa_1_idx = loaddict[\"bias_kappa_info\"]\n",
    "neurometrics_data_all = dict(loaddict)\n",
    "del neurometrics_data_all[\"bias_kappa_info\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:43.565867Z",
     "start_time": "2021-05-02T14:07:43.560678Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    coef_map,\n",
    "    pval_map_subjs,\n",
    "    tval_map,\n",
    "    pval_map,\n",
    "    subj_maxima,\n",
    "    wilcoxon,\n",
    ") = neurometrics_data_all[model_to_use]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:43.849568Z",
     "start_time": "2021-05-02T14:07:43.839547Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_coefmap_to_plot = np.mean(coef_map, axis=0)\n",
    "_tmap_to_plot = tval_map[...]\n",
    "_pmap_to_plot = pval_map[...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:44.052341Z",
     "start_time": "2021-05-02T14:07:44.028673Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Make a \"mask\" for significant pixels\n",
    "# mask = np.ones_like(_pmap_to_plot)\n",
    "# mask[(_pmap_to_plot < 0.05) & (_coefmap_to_plot >= 0)] = 0.75\n",
    "# mask[(_pmap_to_plot > 0.05) | (_coefmap_to_plot < 0)] = 0.75\n",
    "\n",
    "# use B/H FDR correction\n",
    "sig, pval_corr = statsmodels.stats.multitest.fdrcorrection(\n",
    "    _pmap_to_plot.flatten(), alpha=0.001\n",
    ")\n",
    "sig = sig.reshape(_pmap_to_plot.shape)\n",
    "\n",
    "mask = sig.copy().astype(float)\n",
    "mask[mask == 0] = 0.75\n",
    "mask[_coefmap_to_plot <= 0] = 0.75"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:44.889895Z",
     "start_time": "2021-05-02T14:07:44.208794Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_mapmean_instead_of_subjmean_panel_4a = True\n",
    "\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    # plot the image\n",
    "    _ = ax1.imshow(\n",
    "        _coefmap_to_plot, origin=\"upper\", interpolation=\"nearest\", vmin=0, alpha=mask\n",
    "    )\n",
    "\n",
    "    # need to plot an additional tweak image with same settings BUT without alpha\n",
    "    # so that the colorbar renders correctly\n",
    "    _, _tweak_ax = plt.subplots()\n",
    "    im1 = _tweak_ax.imshow(\n",
    "        _coefmap_to_plot, origin=\"upper\", interpolation=\"nearest\", vmin=0\n",
    "    )\n",
    "\n",
    "    # add colorbar\n",
    "    _lab = {\"pearson\": \"Δ Pearson's r\", \"kendall_a\": \"Δ Kendall's tau\"}[corr_method]\n",
    "    cbar = plt.colorbar(im1, cax=cax1, label=_lab)\n",
    "\n",
    "    # plot \"cross\" at bias 0 and kappa 1\n",
    "    ax1.axvline(bias_0_idx, color=\"white\", linestyle=\"--\", zorder=1)\n",
    "    ax1.axhline(kappa_1_idx, color=\"white\", linestyle=\"--\", zorder=1)\n",
    "\n",
    "    # plot single subjects\n",
    "    # subj_maxima are (ycoord, xcoord, value)\n",
    "    ax1.scatter(\n",
    "        subj_maxima[..., 0:2][:, 1],\n",
    "        subj_maxima[..., 0:2][:, 0],\n",
    "        marker=\"o\",\n",
    "        facecolor=\"white\",\n",
    "        edgecolor=\"black\",\n",
    "        s=20,\n",
    "        zorder=5,\n",
    "    )\n",
    "\n",
    "    # plot single summary point\n",
    "    if plot_mapmean_instead_of_subjmean_panel_4a:\n",
    "        kappa_idx_mean, bias_idx_mean = np.unravel_index(\n",
    "            np.argmax(_coefmap_to_plot), _coefmap_to_plot.shape\n",
    "        )\n",
    "    else:\n",
    "        kappa_idx_mean, bias_idx_mean = np.around(\n",
    "            np.mean(subj_maxima[..., 0:2], axis=0)\n",
    "        ).astype(int)\n",
    "\n",
    "    ax1.scatter(\n",
    "        bias_idx_mean,\n",
    "        kappa_idx_mean,\n",
    "        marker=\"d\",\n",
    "        facecolor=\"white\",\n",
    "        edgecolor=\"black\",\n",
    "        s=75,\n",
    "        zorder=10,\n",
    "    )\n",
    "\n",
    "    # settings\n",
    "    n = _coefmap_to_plot.shape[0]\n",
    "    step = int(n / 4)\n",
    "    ticks = np.arange(0, n, step)\n",
    "    ticks = np.append(ticks, n - 1)\n",
    "\n",
    "    _ = ax1.set(\n",
    "        xticks=ticks,\n",
    "        xticklabels=biases[ticks].round(2),\n",
    "        xlabel=\"bias ($\\it{b}$)\",\n",
    "        yticks=ticks,\n",
    "        yticklabels=kappas[ticks].round(2),\n",
    "        ylabel=\"kappa ($\\it{k}$)\",\n",
    "    )\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:44.901083Z",
     "start_time": "2021-05-02T14:07:44.892704Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Edit cbar y tick labels\n",
    "labels = [item.get_text() for item in cax1.get_yticklabels()]\n",
    "labels[0] = \"<=0\"\n",
    "cax1.set_yticklabels(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test overall kappa and bias against 1 and zero respectively\n",
    "# subj_maxima are (ycoord=kappa, xcoord=bias, value=corrcoef)\n",
    "print(\"biases diff from 0\")\n",
    "subj_biases = biases[(subj_maxima[..., 1]).astype(int)]\n",
    "stats = pingouin.ttest(subj_biases, 0)\n",
    "display(stats.round(3))\n",
    "\n",
    "print(\"kappas diff from 1\")\n",
    "subj_kappas = kappas[(subj_maxima[..., 0]).astype(int)]\n",
    "stats = pingouin.ttest(subj_kappas, 1)\n",
    "display(stats.round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Panel 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:45.226232Z",
     "start_time": "2021-05-02T14:07:44.975640Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Read behavioral data to get mapping which subj was in which task\n",
    "beh_fname = \"behavioral_data.csv\"\n",
    "beh_fpath = os.path.join(BIDS_ROOT, \"code\", beh_fname)\n",
    "df = pd.read_csv(beh_fpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:45.237788Z",
     "start_time": "2021-05-02T14:07:45.228002Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Collect subject wide maxima for bias and kappa ...\n",
    "(\n",
    "    coef_map,\n",
    "    pval_map_subjs,\n",
    "    tval_map,\n",
    "    pval_map,\n",
    "    subj_maxima,\n",
    "    wilcoxon,\n",
    ") = neurometrics_data_all[model_to_use]\n",
    "\n",
    "data = {}\n",
    "\n",
    "subj_idxs = np.unique(df[\"subject\"]) - 1\n",
    "param_idxs = subj_maxima[subj_idxs, ...][:, 0:2]\n",
    "\n",
    "data[\"kappa\"] = kappas[param_idxs[:, 0].astype(int)]\n",
    "data[\"bias\"] = biases[param_idxs[:, 1].astype(int)]\n",
    "data[\"subject\"] = subj_idxs + 1\n",
    "\n",
    "nm_all_df = pd.DataFrame(data)[[\"subject\", \"bias\", \"kappa\"]]\n",
    "nm_all_df[\"model\"] = model_to_use\n",
    "\n",
    "assert not nm_all_df.isna().any().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:45.294864Z",
     "start_time": "2021-05-02T14:07:45.291137Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def eq1(X, bias, kappa):\n",
    "    \"\"\"See equation 1 from Spitzer et al. 2017, Nature Human Behavior.\"\"\"\n",
    "    dv = np.sign(X + bias) * (np.abs(X + bias) ** kappa)\n",
    "    return dv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:45.529418Z",
     "start_time": "2021-05-02T14:07:45.499777Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "numbers_rescaled = np.interp(np.arange(1, 10), (1, 9), (-1, 1))\n",
    "\n",
    "plot_number_markers_panel_4b = True\n",
    "plot_single_subj_panel_4b = False\n",
    "plot_theoretical_curves = False\n",
    "plot_theoretical_inset = True\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    # plot mean neurometric curve\n",
    "    # xs = np.linspace(-1, 1, 100)\n",
    "    xs = numbers_rescaled\n",
    "\n",
    "    if plot_mapmean_instead_of_subjmean_panel_4a:\n",
    "        bias_all_mean = biases[bias_idx_mean]\n",
    "        kappa_all_mean = kappas[kappa_idx_mean]\n",
    "    else:\n",
    "        bias_all_mean = nm_all_df[\"bias\"].mean()\n",
    "        kappa_all_mean = nm_all_df[\"kappa\"].mean()\n",
    "\n",
    "    values = eq1(\n",
    "        xs,\n",
    "        bias=bias_all_mean,\n",
    "        kappa=kappa_all_mean,\n",
    "    )\n",
    "\n",
    "    if plot_single_subj_panel_4b:\n",
    "        # need to rescale values to fit everything in plot\n",
    "        values = np.interp(values, (values.min(), values.max()), (-1, +1))\n",
    "\n",
    "    kappa_str = \"$\\it{k}$\"\n",
    "    bias_str = \"$\\it{b}$\"\n",
    "    mainline = ax2.plot(\n",
    "        xs,\n",
    "        values,\n",
    "        color=\"black\",\n",
    "        label=f\"{kappa_str}={kappa_all_mean:.2f}\\n{bias_str}={bias_all_mean:.2f}\",\n",
    "    )\n",
    "\n",
    "    # plot number markers\n",
    "    if plot_number_markers_panel_4b:\n",
    "        ax2.scatter(xs, values, s=300, color=\"k\")\n",
    "        for x, v, s in zip(xs, values, np.arange(1, 10)):\n",
    "            ax2.text(x, v, str(s), color=\"white\", ha=\"center\", va=\"center\")\n",
    "\n",
    "    # plot subj wise neurometric curves\n",
    "    if plot_single_subj_panel_4b:\n",
    "        for subject in np.unique(nm_all_df[\"subject\"]):\n",
    "\n",
    "            _tmpdf = nm_all_df[nm_all_df[\"subject\"] == subject]\n",
    "            values = eq1(\n",
    "                xs,\n",
    "                bias=_tmpdf[\"bias\"].to_list()[0],\n",
    "                kappa=_tmpdf[\"kappa\"].to_list()[0],\n",
    "            )\n",
    "\n",
    "            values_rescaled = np.interp(values, (values.min(), values.max()), (-1, +1))\n",
    "\n",
    "            ax2.plot(xs, values_rescaled, **subj_line_settings)\n",
    "\n",
    "    # plot light gray theoretical curves\n",
    "    if plot_theoretical_curves:\n",
    "        theo_lines = []\n",
    "        for _kappa, ls in zip([0.25, 1, 2], [\"-\", \"--\", \":\"]):\n",
    "            values = eq1(\n",
    "                xs,\n",
    "                bias=0,\n",
    "                kappa=_kappa,\n",
    "            )\n",
    "\n",
    "            _line, = ax2.plot(xs, values, color=\"gray\", linestyle=ls, zorder=0,\n",
    "                             label=f\"{kappa_str}={_kappa:.2f}\")\n",
    "            theo_lines.append(_line)\n",
    "            \n",
    "        leg_theo = ax2.legend(handles=theo_lines, loc=\"upper left\", frameon=False)\n",
    "        ax2.add_artist(leg_theo)\n",
    "        \n",
    "    # plot theoretical curves as inset\n",
    "    if plot_theoretical_inset:\n",
    "        ax2_2 = fig.add_axes([0.75,.775,.225,.225])\n",
    "        _xs = np.linspace(-1, 1, 100)\n",
    "        inset_lines = []\n",
    "        for _kappa, ls in zip([0.3, 1, 3], [\":\", \"-\", \"--\"]):\n",
    "            values = eq1(\n",
    "                _xs,\n",
    "                bias=0,\n",
    "                kappa=_kappa,\n",
    "            )\n",
    "\n",
    "            _line, = ax2_2.plot(_xs, values, color=\"gray\", linestyle=ls, zorder=10,\n",
    "                             label=f\"{kappa_str}={_kappa:.1f}\")        \n",
    "            inset_lines.append(_line)\n",
    "        ax2_2.set(xticks=(), yticks=())\n",
    "        ax2_2.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.25)\n",
    "        ax2_2.axvline(0, linestyle=\"--\", color=\"black\", linewidth=0.25)\n",
    "        ax2_2.legend(handles=inset_lines, frameon=False, fontsize=10)\n",
    "        ax2_2.patch.set_alpha(0.75)\n",
    "            \n",
    "    # plot zero lines\n",
    "    ax2.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "    ax2.axvline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    # settings\n",
    "    ax2.set(xlabel=\"X\", ylabel=\"$\\it{v}$\")\n",
    "\n",
    "    leg_main = ax2.legend(handles=mainline, loc=\"lower right\", frameon=False)\n",
    "    \n",
    "\n",
    "sns.despine(ax=ax2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Panel 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:45.884757Z",
     "start_time": "2021-05-02T14:07:45.881472Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# NOTE: for the 2x2 plots we load a map that had more expansive borders\n",
    "use_expansive_borders = False\n",
    "biasmax_expanded = 1.0\n",
    "kappamax_expanded = 20.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:46.766763Z",
     "start_time": "2021-05-02T14:07:46.050115Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "tasks_to_use = tasks\n",
    "\n",
    "if use_expansive_borders:\n",
    "    fname_tasks = os.path.join(\n",
    "        plotdir,\n",
    "        f\"nm_data_{'-'.join(tasks_to_use)}_{corr_method}_b-{biasmax_expanded}_k-{kappamax_expanded}.pickle\",\n",
    "    )\n",
    "else:\n",
    "    fname_tasks = os.path.join(\n",
    "        plotdir,\n",
    "        f\"nm_data_{'-'.join(tasks_to_use)}_{corr_method}_b-{biasmax}_k-{kappamax}.pickle\",\n",
    "    )\n",
    "\n",
    "with open(fname_tasks, \"rb\") as fin:\n",
    "    loaddict = pickle.load(fin)\n",
    "\n",
    "_, biases, kappas, _, _ = loaddict[\"bias_kappa_info\"]\n",
    "neurometrics_data = dict(loaddict)\n",
    "del neurometrics_data[\"bias_kappa_info\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:46.817654Z",
     "start_time": "2021-05-02T14:07:46.768411Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Collect subject wide maxima for bias and kappa ...\n",
    "\n",
    "(\n",
    "    coef_map,\n",
    "    pval_map_subjs,\n",
    "    tval_map,\n",
    "    pval_map,\n",
    "    subj_maxima,\n",
    "    wilcoxon,\n",
    ") = neurometrics_data[model_to_use]\n",
    "\n",
    "dfs = []\n",
    "for itask, task in enumerate(tasks):\n",
    "    data = {}\n",
    "\n",
    "    subj_idxs = np.unique(df[df[\"task\"] == task][\"subject\"]) - 1\n",
    "    param_idxs = subj_maxima[subj_idxs, itask, ...][:, 0:2]\n",
    "\n",
    "    data[\"kappa\"] = kappas[param_idxs[:, 0].astype(int)]\n",
    "    data[\"bias\"] = biases[param_idxs[:, 1].astype(int)]\n",
    "    data[\"subject\"] = subj_idxs + 1\n",
    "\n",
    "    _df = pd.DataFrame(data)\n",
    "    _df[\"task\"] = task\n",
    "    dfs.append(_df)\n",
    "\n",
    "nm_2x2_df = pd.concat(dfs)[[\"subject\", \"task\", \"bias\", \"kappa\"]]\n",
    "nm_2x2_df[\"sampling\"] = nm_2x2_df[\"task\"].str[0].map({\"A\": \"active\", \"Y\": \"yoked\"})\n",
    "nm_2x2_df[\"stopping\"] = nm_2x2_df[\"task\"].str[1].map({\"V\": \"variable\", \"F\": \"fixed\"})\n",
    "nm_2x2_df[\"model\"] = model_to_use\n",
    "\n",
    "assert not nm_2x2_df.isna().any().any()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:47.022521Z",
     "start_time": "2021-05-02T14:07:46.819221Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "    ax = axs.flat[2]\n",
    "\n",
    "    x = \"sampling\"\n",
    "    order = sampling_order\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    colname = \"kappa\"\n",
    "    data = nm_2x2_df\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.pointplot(\n",
    "            x=x,\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            dodge=pointdodge,\n",
    "            ci=ci,\n",
    "            ax=ax,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax,\n",
    "        dodge=True,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    # connect subj dots\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, **subj_line_settings)\n",
    "\n",
    "    # Plot line at one\n",
    "    ax.axhline(1, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    # remove unnecessary legend\n",
    "    ax.get_legend().remove()\n",
    "\n",
    "    ax.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"])\n",
    "    ax.set_xlabel(\"\")\n",
    "    ax.set_ylabel(r\"kappa ($\\it{k}$)\")\n",
    "    ax.set_title(\"kappa\")\n",
    "\n",
    "sns.despine(ax=ax3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Panel 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:47.297952Z",
     "start_time": "2021-05-02T14:07:47.024031Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(**plotting_context):\n",
    "    ax = axs.flat[3]\n",
    "\n",
    "    x = \"sampling\"\n",
    "    order = sampling_order\n",
    "    hue = \"stopping\"\n",
    "    hue_order = stopping_order\n",
    "    colname = \"bias\"\n",
    "    data = nm_2x2_df\n",
    "\n",
    "    with plt.rc_context({\"lines.linewidth\": pointlinewidth}):\n",
    "        sns.pointplot(\n",
    "            x=x,\n",
    "            order=order,\n",
    "            y=colname,\n",
    "            hue=hue,\n",
    "            hue_order=hue_order,\n",
    "            data=data,\n",
    "            dodge=pointdodge,\n",
    "            ci=ci,\n",
    "            ax=ax,\n",
    "            markers=pointmarkers,\n",
    "            scale=pointscale,\n",
    "            errwidth=pointerrwidth,\n",
    "            capsize=pointcapwidth,\n",
    "        )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=x,\n",
    "        order=order,\n",
    "        y=colname,\n",
    "        hue=hue,\n",
    "        hue_order=hue_order,\n",
    "        data=data,\n",
    "        ax=ax,\n",
    "        dodge=True,\n",
    "        size=swarmsize,\n",
    "    )\n",
    "\n",
    "    # connect subj dots\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, **subj_line_settings)\n",
    "\n",
    "    # Plot Zero line\n",
    "    ax.axhline(0, linestyle=\"--\", color=\"black\", linewidth=0.75)\n",
    "\n",
    "    # remove unnecessary legend\n",
    "    ax.get_legend().remove()\n",
    "\n",
    "    ax.set_xticklabels([\"self-\\ncontrolled\", \"yoked\"])\n",
    "    ax.set_xlabel(\"\")\n",
    "    ax.set_ylabel(\"bias ($\\it{b}$)\")\n",
    "    ax.set_title(\"bias\")\n",
    "\n",
    "sns.despine(ax=ax4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# quick stats\n",
    "for param in [\"bias\", \"kappa\"]:\n",
    "    print(\"\\n\\n\", param)\n",
    "\n",
    "    stats = pingouin.mixed_anova(\n",
    "        data=nm_2x2_df,\n",
    "        dv=param,\n",
    "        within=\"sampling\",\n",
    "        between=\"stopping\",\n",
    "        subject=\"subject\",\n",
    "    )\n",
    "\n",
    "    display(stats.round(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for within_first in [True, False]:\n",
    "    stats = pingouin.pairwise_ttests(\n",
    "        data=nm_2x2_df,\n",
    "        dv=\"bias\",\n",
    "        within=\"sampling\",\n",
    "        between=\"stopping\",\n",
    "        subject=\"subject\",\n",
    "        effsize=\"cohen\",\n",
    "        within_first=within_first,\n",
    "    )\n",
    "    display(stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test kappa in two paired ttests (one for full control, one for partial control)\n",
    "# each against yoked baseline\n",
    "for meta, grp in nm_2x2_df.groupby(\"stopping\"):\n",
    "    _x = grp[grp[\"sampling\"] == \"active\"][\"kappa\"]\n",
    "    _y = grp[grp[\"sampling\"] == \"yoked\"][\"kappa\"]\n",
    "    stats = pingouin.ttest(_x, _y, paired=True)\n",
    "    print(meta)\n",
    "    display(stats.round(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test bias in two paired ttests (one for full control, one for partial control)\n",
    "# each against yoked baseline\n",
    "for meta, grp in nm_2x2_df.groupby(\"stopping\"):\n",
    "    _x = grp[grp[\"sampling\"] == \"active\"][\"bias\"]\n",
    "    _y = grp[grp[\"sampling\"] == \"yoked\"][\"bias\"]\n",
    "    stats = pingouin.ttest(_x, _y, paired=True)\n",
    "    print(meta)\n",
    "    display(stats.round(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## legend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:47.843798Z",
     "start_time": "2021-05-02T14:07:47.824765Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# get the correct orderings in plotting\n",
    "# NOTE: pd.unique provides in the order \"as found\"\n",
    "stopping_order = list(pd.unique(nm_2x2_df[\"stopping\"]))\n",
    "sampling_order = list(pd.unique(nm_2x2_df[\"sampling\"]))\n",
    "\n",
    "# The following legend is encoded according to these assumptions:\n",
    "assert stopping_order == hue_order\n",
    "assert sampling_order == [\"active\", \"yoked\"]\n",
    "\n",
    "# add legend\n",
    "# https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "legend_elements = []\n",
    "for linestyle in [\"-\"]:\n",
    "    for i, label in enumerate([\"partial control\", \"full control\"]):\n",
    "\n",
    "        color = sns.color_palette()[i]\n",
    "        legend_elements.append(\n",
    "            Line2D(\n",
    "                [0], [0], color=color, marker=None, label=label, linestyle=linestyle\n",
    "            ),\n",
    "        )\n",
    "\n",
    "with sns.plotting_context(**plotting_context):\n",
    "\n",
    "    legend1 = fig.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower center\",\n",
    "        bbox_to_anchor=(0.5, -0.2),\n",
    "        bbox_transform=fig.transFigure,\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=2,\n",
    "        title=\"\",  # \"sampling\"\n",
    "    )\n",
    "    extra_artists.append(legend1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finish"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:48.101622Z",
     "start_time": "2021-05-02T14:07:48.097626Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if use_expansive_borders:\n",
    "    # Read back orig kappas, not to use those from panel 3/4\n",
    "    with open(fname_all, \"rb\") as fin:\n",
    "        loaddict = pickle.load(fin)\n",
    "\n",
    "    _, biases, kappas, _, _ = loaddict[\"bias_kappa_info\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:48.289211Z",
     "start_time": "2021-05-02T14:07:48.283612Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# to make the xticks \"narrower\" / squeeze them horizontally\n",
    "# at the expense of having whitespace\n",
    "ax = axs.flat[2]\n",
    "ax.set_xlim(-0.75, 1.75)\n",
    "\n",
    "ax = axs.flat[3]\n",
    "ax.set_xlim(-0.75, 1.75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-02T14:07:50.158129Z",
     "start_time": "2021-05-02T14:07:48.468900Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "also_png = True\n",
    "for ext in [\"pdf\", \"png\", \"tif\"]:\n",
    "    if not also_png and ext == \"png\":\n",
    "        continue\n",
    "\n",
    "    dpi = DPI if ext == \"tif\" else None\n",
    "\n",
    "    fname = (\n",
    "        f\"fig4_nm_{model_to_use}_{corr_method}_b-{biases.max()}_k-{kappas.max()}.{ext}\"\n",
    "    )\n",
    "    fname = os.path.join(plotdir, fname)\n",
    "\n",
    "    fig.savefig(\n",
    "        fname,\n",
    "        bbox_extra_artists=extra_artists,\n",
    "        bbox_inches=\"tight\",\n",
    "        dpi=dpi,\n",
    "        pil_kwargs={\"compression\": \"tiff_lzw\"},\n",
    "    )\n",
    "\n",
    "fig"
   ]
  },
  {
   "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": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "359.333px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}