analysis_behavior.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Behavioral analyses\n",
    "\n",
    "Extract behaviorally relevant data from the `events.tsv` files and provide them in Pandas dataframe format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:07.422899Z",
     "start_time": "2021-05-28T13:46:07.032216Z"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.439376Z",
     "start_time": "2021-05-28T13:46:07.425232Z"
    }
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import os\n",
    "import os.path as op\n",
    "import shutil\n",
    "from tempfile import mkdtemp\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pingouin\n",
    "import scipy\n",
    "import seaborn as sns\n",
    "from scipy import stats\n",
    "\n",
    "from utils import (\n",
    "    BIDS_ROOT,\n",
    "    TASK_NAME_MAP,\n",
    "    add_binned_outcomes_to_df,\n",
    "    events_to_behav_data,\n",
    "    get_df_bnt,\n",
    "    task_not_present_for_subject,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get sample description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.491334Z",
     "start_time": "2021-05-28T13:46:10.441699Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "participants_tsv = op.join(BIDS_ROOT, \"participants.tsv\")\n",
    "participants_tsv = pd.read_csv(participants_tsv, sep=\"\\t\")\n",
    "\n",
    "nsubjs = len(participants_tsv)\n",
    "nfemale = int((participants_tsv[\"sex\"] == \"f\").sum())\n",
    "mean_age = participants_tsv[\"age\"].mean().round(1)\n",
    "sd_age = participants_tsv[\"age\"].std().round(1)\n",
    "age_range = (participants_tsv[\"age\"].min(), participants_tsv[\"age\"].max())\n",
    "\n",
    "print(f\"{nsubjs} participants, {nfemale} female\")\n",
    "print(\n",
    "    f\"mean age {mean_age} ± {sd_age} years, range {age_range[0]}-{age_range[1]} years\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get behavioral data\n",
    "\n",
    "- Need to specify `BIDS_ROOT` in `utils.py`, a path to the bids directory with the participant data\n",
    "- Will collect all `events.tsv` files and save them into a temporary directory\n",
    "- Will read the `events.tsv` files, prune them, calculate some summaries, and clean error rows\n",
    "- Will append all data and save as a single dataframe\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.500095Z",
     "start_time": "2021-05-28T13:46:10.493307Z"
    }
   },
   "outputs": [],
   "source": [
    "subjects = list(range(1, 41))\n",
    "\n",
    "tasks = [\"ActiveFixed\", \"ActiveVariable\", \"YokedFixed\", \"YokedVariable\", \"description\"]\n",
    "\n",
    "# Where to save the behavioral data\n",
    "beh_fname = \"behavioral_data.csv\"\n",
    "beh_fpath = op.join(BIDS_ROOT, \"code\", beh_fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.555455Z",
     "start_time": "2021-05-28T13:46:10.504132Z"
    }
   },
   "outputs": [],
   "source": [
    "if op.exists(beh_fpath):\n",
    "    print(\n",
    "        'Data already exists at \"{}\"\\nSkipping the \"get\" procedure ...'.format(\n",
    "            beh_fpath\n",
    "        )\n",
    "    )\n",
    "else:\n",
    "    # where to temporarily save behavioral data\n",
    "    tmp_beh_dir = mkdtemp()\n",
    "    os.makedirs(tmp_beh_dir, exist_ok=True)\n",
    "\n",
    "    # copy all events.tsv files to the new behavioral dir\n",
    "    for root, dirs, files in os.walk(BIDS_ROOT):\n",
    "        for file in files:\n",
    "            if (\n",
    "                \"sourcedata\" not in root\n",
    "                and tmp_beh_dir not in root\n",
    "                and \"events.tsv\" in file\n",
    "            ):\n",
    "                print(file)\n",
    "                fpath = op.join(root, file)\n",
    "                dest = op.join(tmp_beh_dir, file)\n",
    "                shutil.copyfile(fpath, dest)\n",
    "\n",
    "    # Append all dataframe into a single one\n",
    "    dfs = list()\n",
    "    for subj_task_tuple in itertools.product(subjects, tasks):\n",
    "\n",
    "        # If this combination of subject and task is not present,\n",
    "        # skip to the next iteration of the loop\n",
    "        if task_not_present_for_subject(*subj_task_tuple):\n",
    "            continue\n",
    "\n",
    "        # Load df\n",
    "        fname = \"sub-{:02}_task-{}_events.tsv\".format(*subj_task_tuple)\n",
    "        fpath = op.join(tmp_beh_dir, fname)\n",
    "        df = events_to_behav_data(fpath)\n",
    "\n",
    "        # Add subjects and task column\n",
    "        df[\"subject\"] = subj_task_tuple[0]\n",
    "        df[\"task\"] = TASK_NAME_MAP[subj_task_tuple[1]]\n",
    "\n",
    "        # Reorder columns\n",
    "        cols = list(df)\n",
    "        cols.insert(0, cols.pop(cols.index(\"subject\")))\n",
    "        cols.insert(1, cols.pop(cols.index(\"task\")))\n",
    "        cols.insert(2, cols.pop(cols.index(\"n_samples\")))\n",
    "        df = df.loc[:, cols]\n",
    "\n",
    "        dfs.append(df)\n",
    "\n",
    "    df = pd.concat(dfs, ignore_index=True)\n",
    "\n",
    "    # Add BNT data\n",
    "    df_bnt = get_df_bnt(BIDS_ROOT)\n",
    "    tmp = df.merge(df_bnt[[\"subject\", \"bnt_quartile\"]], on=\"subject\")\n",
    "    pd.testing.assert_frame_equal(df, tmp.loc[:, \"subject\":\"timestamp_outcome\"])\n",
    "    df = tmp.copy()\n",
    "\n",
    "    # Add binned outcomes and orth binned outcomes\n",
    "    df = add_binned_outcomes_to_df(df)\n",
    "\n",
    "    # Save DF\n",
    "    df.to_csv(beh_fpath, na_rep=\"n/a\", index=False)\n",
    "\n",
    "    # Delete the temporary behavioral data dir\n",
    "    shutil.rmtree(tmp_beh_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.877265Z",
     "start_time": "2021-05-28T13:46:10.559120Z"
    }
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv(beh_fpath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Difficulty of active versus yoked sequences\n",
    "\n",
    "- One argument against the yoking scheme is based on half of participants being \"yoked-to-self\", and the other half being \"yoked-to-other\"\n",
    "- For yoked-to-self, \"active\" versus \"yoked\" sequences are matched in terms of difficulty\n",
    "- However, how is this for yoked-to-other?\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a \"yoked_to_self\" (True/False) column to the overall behavioral data\n",
    "participants_tsv[\"yoked_to_self\"] = (\n",
    "    participants_tsv[\"participant_id\"] == participants_tsv[\"yoked_to\"]\n",
    ")\n",
    "yoked_to_self = [\n",
    "    int(i)\n",
    "    for i in participants_tsv[participants_tsv[\"yoked_to_self\"] == True][\n",
    "        \"participant_id\"\n",
    "    ].str[4:]\n",
    "]\n",
    "\n",
    "df[\"yoked_to_self\"] = False\n",
    "df.loc[df[\"subject\"].isin(yoked_to_self), \"yoked_to_self\"] = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmpdf = df[\n",
    "    (df[\"yoked_to_self\"] == False) & (df[\"task\"].isin([\"AF\", \"AV\", \"YF\", \"YV\"]))\n",
    "][[\"subject\", \"task\", \"trial\", \"exp_ev0\", \"exp_ev1\"]].drop_duplicates()\n",
    "tmpdf[\"exp_ev_diff\"] = np.abs(tmpdf[\"exp_ev0\"] - tmpdf[\"exp_ev1\"])\n",
    "tmpdf = (\n",
    "    tmpdf.groupby([\"subject\", \"task\"])\n",
    "    .mean()\n",
    "    .reset_index()[[\"subject\", \"task\", \"exp_ev_diff\"]]\n",
    ")\n",
    "tmpdf[\"sampling\"] = tmpdf[\"task\"].map(\n",
    "    {\"AF\": \"active\", \"AV\": \"active\", \"YF\": \"yoked\", \"YV\": \"yoked\"}\n",
    ")\n",
    "tmpdf[\"stopping\"] = tmpdf[\"task\"].map(\n",
    "    {\"AF\": \"fixed\", \"AV\": \"variable\", \"YF\": \"fixed\", \"YV\": \"variable\"}\n",
    ")\n",
    "\n",
    "tmpdf.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "order = [\"active\", \"yoked\"]\n",
    "hue = \"stopping\"\n",
    "hue_order = [\"fixed\", \"variable\"]\n",
    "\n",
    "sns.swarmplot(\n",
    "    x=\"sampling\",\n",
    "    y=\"exp_ev_diff\",\n",
    "    data=tmpdf,\n",
    "    order=order,\n",
    "    hue=hue,\n",
    "    hue_order=hue_order,\n",
    "    ax=ax,\n",
    "    dodge=True,\n",
    ")\n",
    "sns.pointplot(\n",
    "    x=\"sampling\",\n",
    "    y=\"exp_ev_diff\",\n",
    "    data=tmpdf,\n",
    "    order=order,\n",
    "    hue=hue,\n",
    "    hue_order=hue_order,\n",
    "    ax=ax,\n",
    "    ci=68,\n",
    "    join=True,\n",
    ")\n",
    "ax.set(ylabel=\"Mean expected value difference\", title=\"Error bars are SEM\")\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "lh = plt.legend(handles[2:], labels[2:])\n",
    "lh.set_title(hue)\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fstats = pingouin.mixed_anova(\n",
    "    data=tmpdf,\n",
    "    dv=\"exp_ev_diff\",\n",
    "    within=\"sampling\",\n",
    "    between=\"stopping\",\n",
    "    subject=\"subject\",\n",
    ")\n",
    "fstats"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# bonus money earned\n",
    "\n",
    "see also:\n",
    "- https://github.com/sappelhoff/sp_experiment/blob/1c750597fbc7556d503f86eaaa440efae15091ab/sp_experiment/define_settings.py#L127\n",
    "- https://github.com/sappelhoff/sp_experiment/blob/1c750597fbc7556d503f86eaaa440efae15091ab/sp_experiment/utils.py#L67-L141\n",
    "\n",
    "do:\n",
    "\n",
    "- for each subj:\n",
    "- calculate points per task\n",
    "- multiply with `exchange_rate` and round up to integer\n",
    "- do for all tasks, and pay out the sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.882280Z",
     "start_time": "2021-05-28T13:46:10.879514Z"
    }
   },
   "outputs": [],
   "source": [
    "exchange_rate = 0.005"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:10.918675Z",
     "start_time": "2021-05-28T13:46:10.884980Z"
    }
   },
   "outputs": [],
   "source": [
    "tmp = df.drop_duplicates(subset=[\"subject\", \"task\", \"trial\"])\n",
    "assert len(tmp) == 40 * 3 * 100  # 40 subjs, 3 tasks, 100 trials each\n",
    "tmp = tmp[[\"subject\", \"task\", \"trial\", \"fin_outcome\"]]\n",
    "tmp.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.089190Z",
     "start_time": "2021-05-28T13:46:10.920579Z"
    }
   },
   "outputs": [],
   "source": [
    "all_payoffs = []\n",
    "for subj in tmp[\"subject\"].unique():\n",
    "\n",
    "    total_money = []\n",
    "    tmp_subj = tmp[tmp[\"subject\"] == subj]\n",
    "\n",
    "    for task in tmp_subj[\"task\"].unique():\n",
    "\n",
    "        tmp_subj_task = tmp_subj[tmp_subj[\"task\"] == task]\n",
    "        assert len(tmp_subj_task) == 100\n",
    "\n",
    "        points = np.sum(tmp_subj_task[\"fin_outcome\"])\n",
    "        money = int(np.ceil(points * exchange_rate))\n",
    "        total_money.append(money)\n",
    "\n",
    "    all_payoffs.append(np.sum(total_money))\n",
    "\n",
    "assert len(all_payoffs) == 40"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.101204Z",
     "start_time": "2021-05-28T13:46:11.093473Z"
    }
   },
   "outputs": [],
   "source": [
    "print(\n",
    "    f\"Participants earned {np.mean(all_payoffs):.2f} ± {np.std(all_payoffs):.2f} € on average\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-22T09:16:08.969284Z",
     "start_time": "2021-01-22T09:16:08.944128Z"
    }
   },
   "source": [
    "# Descriptives about \"error trials\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.134714Z",
     "start_time": "2021-05-28T13:46:11.104740Z"
    }
   },
   "outputs": [],
   "source": [
    "tmp = df[[\"subject\", \"task\", \"trial\", \"nerrors\"]]\n",
    "tmp = tmp.drop_duplicates(subset=[\"subject\", \"task\", \"trial\"])\n",
    "tmp = tmp[tmp[\"task\"] != \"DESC\"]\n",
    "\n",
    "tmp.groupby([\"subject\"])[\"nerrors\"].sum().agg([\"mean\", \"median\", \"sem\"]).round(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.142997Z",
     "start_time": "2021-05-28T13:46:11.136595Z"
    }
   },
   "outputs": [],
   "source": [
    "# in percent out of 100 active and 100 yoked trials\n",
    "perc_restarted = np.round(((6.5 / 200) * 100))\n",
    "print(\n",
    "    f\"On average {perc_restarted}% of trials per participant were restarted due to lack of fixation,\\n\"\n",
    "    f\"or failure to draw a sample within 3 seconds\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What was the average number of samples?\n",
    "\n",
    "NOTE: Need to calculate \"sample\" column +1 because it is 0-indexed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.183017Z",
     "start_time": "2021-05-28T13:46:11.146046Z"
    }
   },
   "outputs": [],
   "source": [
    "df_nsamples = df[df[\"task\"] == \"AV\"].drop_duplicates([\"subject\", \"trial\"])\n",
    "df_nsamples = df_nsamples[[\"subject\", \"task\", \"trial\", \"n_samples\"]]\n",
    "df_nsamples = df_nsamples.reset_index(drop=True)\n",
    "df_nsamples.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.881636Z",
     "start_time": "2021-05-28T13:46:11.185662Z"
    }
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"talk\"):\n",
    "    fig, ax = plt.subplots(figsize=(10, 5))\n",
    "    sns.histplot(\n",
    "        data=df_nsamples, x=\"n_samples\", ax=ax, bins=19, discrete=True, kde=True\n",
    "    )\n",
    "    ax.set(xticks=np.arange(20), xlabel=\"Number of samples\", ylabel=\"Count\")\n",
    "\n",
    "    sns.despine(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:11.894324Z",
     "start_time": "2021-05-28T13:46:11.883489Z"
    }
   },
   "outputs": [],
   "source": [
    "mean_samples = df_nsamples[\"n_samples\"].mean()\n",
    "print(\n",
    "    f\"mean {mean_samples:.1f}, SD {df_nsamples['n_samples'].std():.1f}\",\n",
    ")\n",
    "print(f\"median {np.median(df_nsamples['n_samples']):.1f}\")\n",
    "print(f\"mode {scipy.stats.mode(df_nsamples['n_samples'])[0][0]:.1f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compute mean sample position for each number / probability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: We'll need to drop all trials with only one sample for this analysis\n",
    "# ... because we cannot tell whether the \"relative\" sample position of a single sample is 0 or 1.\n",
    "(\n",
    "    df[df[\"task\"].isin([\"AF\", \"AV\"])].drop_duplicates([\"subject\", \"task\", \"trial\"])[\n",
    "        \"n_samples\"\n",
    "    ]\n",
    "    < 2\n",
    ").sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute mean positions of outcomes\n",
    "_tasks = [\"AF\", \"AV\"]\n",
    "_dfs = []\n",
    "for sub in np.unique(df[\"subject\"]):\n",
    "    for i, task in enumerate(_tasks):\n",
    "\n",
    "        if task_not_present_for_subject(sub, task):\n",
    "            continue\n",
    "\n",
    "        _data = df[\n",
    "            (df[\"task\"] == task) & (df[\"n_samples\"] > 1) & (df[\"subject\"] == sub)\n",
    "        ][[\"n_samples\", \"sample\", \"outcome\"]]\n",
    "        arr = _data.to_numpy(dtype=float)\n",
    "\n",
    "        # normalize sample position by sample length\n",
    "        # z = (x - x.min) / (x.max - x.min)\n",
    "        # note: samples are zero-indexed, so +1 first\n",
    "        arr[:, 1] += 1\n",
    "        arr[:, 1] = (arr[:, 1] - 1) / (arr[:, 0] - 1)\n",
    "\n",
    "        _df = pd.DataFrame(arr[:, 1:], columns=[\"sample_pos_norm\", \"value\"])\n",
    "        _df[\"control\"] = {\"AF\": \"partial\", \"AV\": \"full\"}[task]\n",
    "        _df[\"subject\"] = sub\n",
    "        _dfs.append(_df)\n",
    "\n",
    "df_values = pd.concat(_dfs)\n",
    "df_values = df_values.astype({\"value\": int})\n",
    "\n",
    "# form means in each subj\n",
    "df_values = df_values.groupby([\"subject\", \"control\", \"value\"]).mean().reset_index()\n",
    "df_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each task, trial, option, outcome ...\n",
    "# --> compute the \"probability\" (frequency) of that outcome occurring\n",
    "data = df[df[\"task\"].isin([\"AF\", \"AV\"])][\n",
    "    [\"subject\", \"task\", \"trial\", \"n_samples\", \"sample\", \"action\", \"outcome\"]\n",
    "]\n",
    "data = data[data[\"n_samples\"] > 1]\n",
    "\n",
    "# go through each subj / task / trial separately\n",
    "grps = []\n",
    "for meta, grp in data.groupby([\"subject\", \"task\", \"trial\"]):\n",
    "\n",
    "    # go through each option (left/right) separately\n",
    "    grp = grp.copy()\n",
    "    grp[\"outcome_freq\"] = 0.0\n",
    "    for _meta, _grp in grp.groupby(\"action\"):\n",
    "        # calculate \"experienced probability\" (frequency)\n",
    "        # specific for the option (left/right)\n",
    "        arr = _grp[\"outcome\"].value_counts().reset_index().to_numpy(dtype=float)\n",
    "        arr[:, -1] = arr[:, -1] / len(_grp)\n",
    "        assert arr[:, 1].sum() == 1\n",
    "\n",
    "        # add the calculated frequencies to the df\n",
    "        def _insert_val(row, arr=arr):\n",
    "            out = row[\"outcome\"]\n",
    "            freq = row[\"outcome_freq\"]\n",
    "            if out in arr[:, 0]:\n",
    "                freq = arr[:, 1][np.nonzero(arr[:, 0] == out)[0][0]]\n",
    "            return freq\n",
    "\n",
    "        grp[\"outcome_freq\"] = grp.apply(lambda row: _insert_val(row), axis=1)\n",
    "\n",
    "    grps.append(grp)\n",
    "\n",
    "_data = pd.concat(grps)\n",
    "\n",
    "# sanity check that we only added a frequency column without changing anything else\n",
    "pd.testing.assert_frame_equal(\n",
    "    data,\n",
    "    _data[[\"subject\", \"task\", \"trial\", \"n_samples\", \"sample\", \"action\", \"outcome\"]],\n",
    ")\n",
    "data = _data.copy()\n",
    "del _data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bin the frequencies into nbins evenly spaced bins\n",
    "# e.g., nbins=5 --> array([0.  , 0.25, 0.5 , 0.75, 1.  ])\n",
    "# if right=False:\n",
    "# bin 1 = 0-0.24999999\n",
    "# bin 2 = 0.25-0.4999999\n",
    "# bin 3 = 0.5-0.7499999999\n",
    "# bin 4 = 0.75-0.9999999\n",
    "# bin 5 = 1.0-infinity\n",
    "#\n",
    "# if right=True:\n",
    "# bin 1 = 0-0.25\n",
    "# bin 2 = 0.251-0.5\n",
    "# bin 3 = 0.51-0.75\n",
    "# bin 4 = 0.751-1.0\n",
    "# bin 5 = empty\n",
    "nbins = 6\n",
    "bins = np.linspace(0, 1, nbins)\n",
    "\n",
    "data[\"outcome_freq_bin\"] = np.digitize(data[\"outcome_freq\"], bins=bins, right=False)\n",
    "\n",
    "datacopy = data.copy()\n",
    "datacopy[\"outcome_freq_bin\"] = np.digitize(\n",
    "    datacopy[\"outcome_freq\"], bins=bins, right=True\n",
    ")\n",
    "\n",
    "\n",
    "# sanity check bins (see above)\n",
    "assert data[\"outcome_freq\"].min() > 0\n",
    "assert np.allclose(\n",
    "    data[data[\"outcome_freq_bin\"] == nbins][\"outcome_freq\"].to_numpy(), 1.0\n",
    ")\n",
    "assert len(datacopy[datacopy[\"outcome_freq_bin\"] == nbins]) == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Which one to pick? right open or closed?\n",
    "data = datacopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute mean positions of freq bins\n",
    "_tasks = [\"AF\", \"AV\"]\n",
    "_dfs = []\n",
    "for sub in np.unique(df[\"subject\"]):\n",
    "    for i, task in enumerate(_tasks):\n",
    "\n",
    "        if task_not_present_for_subject(sub, task):\n",
    "            continue\n",
    "\n",
    "        _data = data[(data[\"task\"] == task) & (data[\"subject\"] == sub)][\n",
    "            [\"n_samples\", \"sample\", \"outcome_freq_bin\"]\n",
    "        ]\n",
    "        arr = _data.to_numpy(dtype=float)\n",
    "\n",
    "        # normalize sample position by sample length\n",
    "        # z = (x - x.min) / (x.max - x.min)\n",
    "        # note: samples are zero-indexed, so +1 first\n",
    "        arr[:, 1] += 1\n",
    "        arr[:, 1] = (arr[:, 1] - 1) / (arr[:, 0] - 1)\n",
    "\n",
    "        _df = pd.DataFrame(arr[:, 1:], columns=[\"sample_pos_norm\", \"freq_bin\"])\n",
    "        _df[\"control\"] = {\"AF\": \"partial\", \"AV\": \"full\"}[task]\n",
    "        _df[\"subject\"] = sub\n",
    "        _dfs.append(_df)\n",
    "\n",
    "df_freqs = pd.concat(_dfs)\n",
    "df_freqs = df_freqs.astype({\"freq_bin\": int})\n",
    "\n",
    "# form means in each subj\n",
    "# NOTE: subj 20 doesn't have samples with low sample_pos_norm (i.e., empty first bin)\n",
    "df_freqs = df_freqs.groupby([\"subject\", \"control\", \"freq_bin\"]).mean().reset_index()\n",
    "df_freqs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"talk\"):\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(8, 4))\n",
    "    ax1, ax2 = axs\n",
    "\n",
    "    ci = 68\n",
    "    scale = 0.75\n",
    "    dodge = True\n",
    "    sns.pointplot(\n",
    "        data=df_values.groupby([\"subject\", \"control\", \"value\"])\n",
    "        .mean()\n",
    "        .reset_index(),  # df_values,\n",
    "        x=\"value\",\n",
    "        y=\"sample_pos_norm\",\n",
    "        hue=\"control\",\n",
    "        ci=ci,\n",
    "        ax=ax1,\n",
    "        scale=scale,\n",
    "        dodge=dodge,\n",
    "    )\n",
    "    sns.pointplot(\n",
    "        data=df_freqs,\n",
    "        x=\"freq_bin\",\n",
    "        y=\"sample_pos_norm\",\n",
    "        hue=\"control\",\n",
    "        ci=ci,\n",
    "        ax=ax2,\n",
    "        scale=scale,\n",
    "        dodge=dodge,\n",
    "    )\n",
    "\n",
    "    ax1.axhline(0.5, c=\"black\", lw=0.5, ls=\"--\", zorder=0)\n",
    "    ax2.axhline(0.5, c=\"black\", lw=0.5, ls=\"--\", zorder=0)\n",
    "\n",
    "    ax1.set(ylim=(0.4, 0.6), xlabel=\"sample value\")\n",
    "    ax2.set(\n",
    "        ylim=(0.4, 0.6),\n",
    "        xlabel=\"frequency (binned)\\nwithin option\",\n",
    "        # xticklabels=[f\"{_:.1f}\" for _ in bins[1:]],\n",
    "        xticklabels=[\"0.1\", \"0.3\", \"0.5\", \"0.7\", \"0.9\"],\n",
    "    )\n",
    "    ax1.set_ylabel(\"mean sample position\\n(normalized by trial length)\")\n",
    "    ax2.get_legend().remove()\n",
    "    ax2.set_ylabel(\"mean sample position\\n(normalized by trial length)\")\n",
    "    ax1.legend(frameon=False, title=None)\n",
    "    ax2.legend(frameon=False, title=None)\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "sns.despine(fig)\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# slope of \"sample value\"\n",
    "for control in [\"partial\", \"full\"]:\n",
    "    _d = df_values[df_values[\"control\"] == control]\n",
    "    val = _d[\"value\"].to_numpy()\n",
    "    pos = _d[\"sample_pos_norm\"].to_numpy()\n",
    "    print(control)\n",
    "    display(pingouin.regression.linear_regression(X=val, y=pos).round(4))\n",
    "    print(\"------\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = sns.lmplot(\n",
    "    x=\"value\",\n",
    "    y=\"sample_pos_norm\",\n",
    "    col=\"control\",\n",
    "    x_jitter=0.25,\n",
    "    y_jitter=0.025,\n",
    "    data=df_values,\n",
    "    line_kws={\"color\": \"red\"},\n",
    ")\n",
    "grid.fig.suptitle(\"sample values\\nindividual dots are jittered\", y=1.1)\n",
    "grid.fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What were the experienced and true accuracies in the conditions?\n",
    "\n",
    "- experienced accuracy = accuracy judged according to what was seen\n",
    "- true accuracy = accuracy judged based on the underlying, true, distributions\n",
    "\n",
    "NOTE: Need to remove \"ambiguous trials\", where one of the following conditions was met:\n",
    "\n",
    "- only one option was sampled\n",
    "- the expected values for both options were equal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate accuracies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:14.033830Z",
     "start_time": "2021-05-28T13:46:11.896539Z"
    }
   },
   "outputs": [],
   "source": [
    "# Calculate an accuracy dataframe for each accuracy type ... to eventually concatenate them\n",
    "dfs = list()\n",
    "for accuracy_type in [\"experienced\", \"true\"]:\n",
    "\n",
    "    tmp_df = df.copy()\n",
    "\n",
    "    if accuracy_type == \"experienced\":\n",
    "        ev_to_base_on0 = \"exp_ev0\"\n",
    "        ev_to_base_on1 = \"exp_ev1\"\n",
    "    elif accuracy_type == \"true\":\n",
    "        ev_to_base_on0 = \"true_ev0\"\n",
    "        ev_to_base_on1 = \"true_ev1\"\n",
    "\n",
    "    # ambiguous rows, where the experienced EVs are equal\n",
    "    ev_equal_rows = (tmp_df[\"exp_ev0\"] == tmp_df[\"exp_ev1\"]).to_numpy()\n",
    "    tmp_df[\"evs_equal\"] = ev_equal_rows\n",
    "\n",
    "    # ambiguous rows, where only one option was sampled\n",
    "    # NOTE: this will produce wrong data for DESC task\n",
    "    # (see https://github.com/sappelhoff/sp_experiment/issues/8)\n",
    "    only_left_sampled = (\n",
    "        tmp_df.groupby([\"subject\", \"trial\", \"task\"])[\"action\"].mean() == 0\n",
    "    )\n",
    "    only_right_sampled = (\n",
    "        tmp_df.groupby([\"subject\", \"trial\", \"task\"])[\"action\"].mean() == 1\n",
    "    )\n",
    "    single_option_sampled = np.logical_or(only_left_sampled, only_right_sampled)\n",
    "\n",
    "    # merge the single option rows into tmp_df\n",
    "    single_option_sampled = single_option_sampled.reset_index()\n",
    "    single_option_sampled = single_option_sampled.rename(\n",
    "        {\"action\": \"single_option_sampled\"}, axis=1\n",
    "    )\n",
    "    tmp_df = pd.merge(\n",
    "        tmp_df,\n",
    "        single_option_sampled,\n",
    "        on=[\"subject\", \"trial\", \"task\"],\n",
    "        validate=\"many_to_one\",\n",
    "    )\n",
    "\n",
    "    # NOTE: For description task, trials are never ambiguous for cases were only a single option was sampled\n",
    "    # because these cases were replaced with descriptions of true underlying distributions\n",
    "    # rather than descriptions of the experienced distributions\n",
    "    # Still, we mark them as ambiguous to prevent analyzing data from two different sources (experienced vs true)\n",
    "    # For each subject, overwrite \"single_option_sampled\" column for DESC task with the\n",
    "    # corresponding data from the active task (either AV or AF, depending on subject)\n",
    "    for subj in range(1, 41):\n",
    "        idx_active = (tmp_df[\"subject\"] == subj) & (\n",
    "            ~tmp_df[\"task\"].isin([\"DESC\", \"YV\", \"YF\"])\n",
    "        )\n",
    "        idx_desc = (tmp_df[\"subject\"] == subj) & (tmp_df[\"task\"] == \"DESC\")\n",
    "        data = (\n",
    "            tmp_df[idx_active]\n",
    "            .drop_duplicates(subset=\"trial\")[\"single_option_sampled\"]\n",
    "            .to_numpy()\n",
    "        )\n",
    "        tmp_df.loc[idx_desc, \"single_option_sampled\"] = data\n",
    "\n",
    "    # ambiguous rows overall (equal EVs, or single option sampled)\n",
    "    tmp_df[\"ambiguous_trial\"] = np.logical_or(\n",
    "        tmp_df[\"evs_equal\"], tmp_df[\"single_option_sampled\"]\n",
    "    ).to_numpy()\n",
    "\n",
    "    # Group data by subject, task, and trial ... taking the mean over samples\n",
    "    # --> the mean should leave \"fin_action\", \"ambiguous_trial\" and EVs unchanged, ...\n",
    "    # because they are the same over samples\n",
    "    tmp_df = (\n",
    "        tmp_df.groupby([\"subject\", \"task\", \"trial\"])[\n",
    "            [ev_to_base_on0, ev_to_base_on1, \"ambiguous_trial\", \"fin_action\"]\n",
    "        ]\n",
    "        .mean()\n",
    "        .reset_index()\n",
    "    )\n",
    "\n",
    "    # Remove ambiguous trials\n",
    "    tmp_df = tmp_df[~tmp_df[\"ambiguous_trial\"]]\n",
    "\n",
    "    # Add column when correct choice was made\n",
    "    right_better = (tmp_df[ev_to_base_on0] < tmp_df[ev_to_base_on1]).to_numpy()\n",
    "    right_selected = (tmp_df[\"fin_action\"] == 1).to_numpy()\n",
    "    tmp_df[\"correct_choice\"] = right_better == right_selected\n",
    "\n",
    "    # Calculate accuracy as mean correct choices over trials per subject and task\n",
    "    accuracy_df = (\n",
    "        tmp_df.groupby([\"subject\", \"task\"])[\"correct_choice\"].mean().reset_index()\n",
    "    )\n",
    "    accuracy_df[\"accuracy_type\"] = accuracy_type\n",
    "    dfs.append(accuracy_df)\n",
    "\n",
    "accuracy_df = pd.concat(dfs, ignore_index=True)\n",
    "accuracy_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:14.043299Z",
     "start_time": "2021-05-28T13:46:14.036083Z"
    }
   },
   "outputs": [],
   "source": [
    "fname_beh_acc = op.join(BIDS_ROOT, \"code\", \"beh_accs.csv\")\n",
    "accuracy_df.to_csv(fname_beh_acc, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity check for yoking scheme: Comparing \"yoked to self\" and \"yoked to other\"\n",
    "\n",
    "in \"yoked\", watching the stream of a different person may be different from watching our own stream\n",
    "\n",
    "... are the accuracies different between these groups?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:14.055416Z",
     "start_time": "2021-05-28T13:46:14.046753Z"
    }
   },
   "outputs": [],
   "source": [
    "# get list of subj ids who were yoked to themselves\n",
    "yoked_to_self = participants_tsv[\"participant_id\"][\n",
    "    participants_tsv[\"participant_id\"] == participants_tsv[\"yoked_to\"]\n",
    "].to_list()\n",
    "yoked_to_self = [int(i[-2:]) for i in yoked_to_self]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:14.068724Z",
     "start_time": "2021-05-28T13:46:14.057176Z"
    }
   },
   "outputs": [],
   "source": [
    "# add a column to accuracy df\n",
    "accuracy_df[\"yoked_to\"] = \"n/a\"\n",
    "accuracy_df.loc[accuracy_df[\"subject\"].isin(yoked_to_self), \"yoked_to\"] = \"self\"\n",
    "accuracy_df.loc[~accuracy_df[\"subject\"].isin(yoked_to_self), \"yoked_to\"] = \"other\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:14.244545Z",
     "start_time": "2021-05-28T13:46:14.071418Z"
    }
   },
   "outputs": [],
   "source": [
    "# print out a summary\n",
    "accuracy_df.groupby([\"task\", \"yoked_to\", \"accuracy_type\"])[\"correct_choice\"].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:15.327487Z",
     "start_time": "2021-05-28T13:46:14.249505Z"
    }
   },
   "outputs": [],
   "source": [
    "# visualize as 2x2x2 catplot\n",
    "tmp = accuracy_df[\n",
    "    (accuracy_df[\"task\"] != \"DESC\") & (accuracy_df[\"accuracy_type\"] == \"experienced\")\n",
    "]\n",
    "\n",
    "tmp.insert(\n",
    "    2, \"sampling\", [\"active\" if i == \"A\" else \"yoked\" for i in tmp[\"task\"].str[0]]\n",
    ")\n",
    "tmp.insert(\n",
    "    3, \"stopping\", [\"fixed\" if i == \"F\" else \"variable\" for i in tmp[\"task\"].str[1]]\n",
    ")\n",
    "\n",
    "grid = sns.catplot(\n",
    "    kind=\"point\",\n",
    "    ci=68,\n",
    "    dodge=True,\n",
    "    x=\"sampling\",\n",
    "    y=\"correct_choice\",\n",
    "    hue=\"stopping\",\n",
    "    data=tmp,\n",
    "    col=\"yoked_to\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:15.368143Z",
     "start_time": "2021-05-28T13:46:15.331892Z"
    }
   },
   "outputs": [],
   "source": [
    "tmp.groupby(\"yoked_to\")[\"correct_choice\"].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:15.421707Z",
     "start_time": "2021-05-28T13:46:15.371333Z"
    }
   },
   "outputs": [],
   "source": [
    "xname = \"self\"\n",
    "yname = \"other\"\n",
    "x = tmp[tmp[\"yoked_to\"] == xname][\"correct_choice\"].to_numpy()\n",
    "y = tmp[tmp[\"yoked_to\"] == yname][\"correct_choice\"].to_numpy()\n",
    "model = pingouin.ttest(x, y, paired=False)\n",
    "\n",
    "print(\n",
    "    f\"ttest\\nMean accuracies\\n{xname}: {np.round(x.mean(), 2)}, {yname}: {np.round(y.mean(),2)}\"\n",
    ")\n",
    "print(\n",
    "    f\"t({model['dof'][0]})={np.round(model['T'][0], 3)}, p={np.round(model['p-val'][0], 3)}\"\n",
    ")\n",
    "display(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:22:50.218491Z",
     "start_time": "2021-05-28T13:22:50.214580Z"
    }
   },
   "source": [
    "### Use rpy2 for mixed anova\n",
    "\n",
    "see: https://www.marsja.se/r-from-python-rpy2-tutorial/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rpy2.robjects as ro\n",
    "import rpy2.robjects.packages as rpackages\n",
    "from rpy2.robjects import pandas2ri\n",
    "from rpy2.robjects.conversion import localconverter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:17.127148Z",
     "start_time": "2021-05-28T13:46:17.107827Z"
    }
   },
   "outputs": [],
   "source": [
    "# convert pandas DF to R data.frame\n",
    "with localconverter(ro.default_converter + pandas2ri.converter):\n",
    "    r_from_pd_df = ro.conversion.py2rpy(tmp)\n",
    "\n",
    "r_from_pd_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:19.475286Z",
     "start_time": "2021-05-28T13:46:17.129459Z"
    }
   },
   "outputs": [],
   "source": [
    "afex = rpackages.importr(\"afex\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:19.603380Z",
     "start_time": "2021-05-28T13:46:19.477605Z"
    }
   },
   "outputs": [],
   "source": [
    "## see: https://cran.r-project.org/web/packages/afex/afex.pdf\n",
    "model = afex.aov_ez(\n",
    "    id=\"subject\",\n",
    "    dv=\"correct_choice\",\n",
    "    between=[\"stopping\", \"yoked_to\"],\n",
    "    within=\"sampling\",\n",
    "    data=r_from_pd_df,\n",
    "    check_contrasts=True,  # 3-way interaction is unaffected by this, \"True\" is recommended\n",
    "    type=3,  # using so-called type 3 sums of squares\n",
    "    print_formula=True,\n",
    ")\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:19.612868Z",
     "start_time": "2021-05-28T13:46:19.605484Z"
    }
   },
   "outputs": [],
   "source": [
    "# add sampling and stopping cols to DF\n",
    "sampstop_data = np.array([[i[0], i[1]] for i in accuracy_df[\"task\"].to_list()])\n",
    "sampstop_data[(sampstop_data == \"D\") | (sampstop_data == \"E\")] = \"\"\n",
    "accuracy_df[\"sampling\"] = sampstop_data[:, 0]\n",
    "accuracy_df[\"stopping\"] = sampstop_data[:, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:19.634026Z",
     "start_time": "2021-05-28T13:46:19.615126Z"
    }
   },
   "outputs": [],
   "source": [
    "# save data for publication plots\n",
    "plotdir = op.join(BIDS_ROOT, \"code\", \"publication_plots\")\n",
    "os.makedirs(plotdir, exist_ok=True)\n",
    "\n",
    "fname = \"beh_accuracies.csv\"\n",
    "fname = op.join(plotdir, fname)\n",
    "\n",
    "tmp = accuracy_df[accuracy_df[\"accuracy_type\"] == \"experienced\"]\n",
    "tmp = tmp[tmp[\"task\"] != \"DESC\"]\n",
    "tmp[\"sampling\"] = tmp[\"sampling\"].map({\"A\": \"Active\", \"Y\": \"Yoked\"})\n",
    "tmp[\"stopping\"] = tmp[\"stopping\"].map({\"V\": \"Variable\", \"F\": \"Fixed\"})\n",
    "\n",
    "tmp.to_csv(fname, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:20.236716Z",
     "start_time": "2021-05-28T13:46:19.636297Z"
    }
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"notebook\"):\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    data = tmp\n",
    "    order = [\"Active\", \"Yoked\"]\n",
    "    hue_order = [\"Fixed\", \"Variable\"]\n",
    "\n",
    "    sns.pointplot(\n",
    "        x=\"sampling\",\n",
    "        y=\"correct_choice\",\n",
    "        hue=\"stopping\",\n",
    "        data=data,\n",
    "        ci=68,\n",
    "        dodge=0.2,\n",
    "        order=order,\n",
    "        hue_order=hue_order,\n",
    "        ax=ax,\n",
    "        markers=\"o\",\n",
    "    )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=\"sampling\",\n",
    "        y=\"correct_choice\",\n",
    "        hue=\"stopping\",\n",
    "        data=data,\n",
    "        order=order,\n",
    "        hue_order=hue_order,\n",
    "        ax=ax,\n",
    "        dodge=True,\n",
    "        size=4,\n",
    "    )\n",
    "\n",
    "    # add legend\n",
    "    # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "    from matplotlib.lines import Line2D\n",
    "\n",
    "    legend_elements = [\n",
    "        Line2D([0], [0], color=sns.color_palette()[0], marker=\"o\", label=\"Fixed\"),\n",
    "        Line2D([0], [0], color=sns.color_palette()[1], marker=\"o\", label=\"Variable\"),\n",
    "    ]\n",
    "\n",
    "    ax.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower left\",\n",
    "        framealpha=1,\n",
    "        bbox_to_anchor=(0, 1),\n",
    "        ncol=2,\n",
    "        title=\"Stopping\",\n",
    "    )\n",
    "\n",
    "    ax.set_ylim((0.5, 1.0))\n",
    "    ax.set_ylabel(\"accuracy\")\n",
    "\n",
    "    # https://stackoverflow.com/a/51157346/5201771\n",
    "    # 2-4 3-5\n",
    "    for i_dots, (idx0, idx1) in enumerate([(2, 4), (3, 5)]):\n",
    "        locs1 = ax.get_children()[idx0].get_offsets()\n",
    "        locs2 = ax.get_children()[idx1].get_offsets()\n",
    "\n",
    "        # Need to sort locs, so data corresponds\n",
    "        sort_idxs_list = []\n",
    "        sampling = order[i_dots]\n",
    "        for stopping in hue_order:\n",
    "            arr = data[(data[\"sampling\"] == sampling) & (data[\"stopping\"] == stopping)][\n",
    "                \"correct_choice\"\n",
    "            ].to_numpy()\n",
    "            sort_idxs_list += [np.argsort(arr)]\n",
    "\n",
    "        locs2_sorted = locs2[sort_idxs_list[1].argsort()][sort_idxs_list[0]]\n",
    "\n",
    "        for i in range(locs1.shape[0]):\n",
    "            _x = [locs1[i, 0], locs2_sorted[i, 0]]\n",
    "            _y = [locs1[i, 1], locs2_sorted[i, 1]]\n",
    "            ax.plot(_x, _y, color=\"black\", alpha=0.1)\n",
    "\n",
    "    # ax.set_yticks(np.arange(0.6, 0.95, 0.02))\n",
    "    # ax.grid('on')\n",
    "    ax.set_xlabel(\"Sampling\")\n",
    "    ax.set_ylabel(\"Accuracy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## statistical testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:20.319625Z",
     "start_time": "2021-05-28T13:46:20.238207Z"
    }
   },
   "outputs": [],
   "source": [
    "# Summary statistics\n",
    "accuracy_df.groupby([\"task\", \"accuracy_type\"])[\"correct_choice\"].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:21.423389Z",
     "start_time": "2021-05-28T13:46:20.321641Z"
    }
   },
   "outputs": [],
   "source": [
    "# Mean +- SEM for specific selection\n",
    "\n",
    "\n",
    "def sns_ci(a):\n",
    "    \"\"\"Helper to get seaborn 68% ci (~SEM).\"\"\"\n",
    "    it = sns.utils.ci(sns.algorithms.bootstrap(a), 68)\n",
    "    return np.mean([np.abs(a.mean() - it[0]), np.abs(a.mean() - it[1])])\n",
    "\n",
    "\n",
    "accuracy_df[\n",
    "    (accuracy_df[\"accuracy_type\"] == \"experienced\") & (accuracy_df[\"task\"] != \"DESC\")\n",
    "].groupby([\"task\", \"accuracy_type\"]).agg(\n",
    "    {\"correct_choice\": [np.mean, scipy.stats.sem, sns_ci]}\n",
    ").round(\n",
    "    3\n",
    ").reset_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### mixed anova"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:21.683938Z",
     "start_time": "2021-05-28T13:46:21.425399Z"
    }
   },
   "outputs": [],
   "source": [
    "accuracy_type = \"experienced\"\n",
    "\n",
    "tt = accuracy_df.loc[\n",
    "    (accuracy_df[\"task\"] != \"DESC\") & (accuracy_df[\"accuracy_type\"] == accuracy_type),\n",
    "    :,\n",
    "]\n",
    "\n",
    "# calculate model\n",
    "model = pingouin.mixed_anova(\n",
    "    data=tt,\n",
    "    dv=\"correct_choice\",\n",
    "    within=\"sampling\",\n",
    "    between=\"stopping\",\n",
    "    subject=\"subject\",\n",
    ").round(3)\n",
    "\n",
    "display(model)\n",
    "\n",
    "# calculate posthocs\n",
    "\n",
    "stats = pingouin.pairwise_ttests(\n",
    "    data=tt,\n",
    "    padjust=\"bonf\",\n",
    "    dv=\"correct_choice\",\n",
    "    within=\"sampling\",\n",
    "    between=\"stopping\",\n",
    "    subject=\"subject\",\n",
    "    within_first=False,\n",
    "    effsize=\"cohen\",\n",
    ").round(3)\n",
    "\n",
    "display(stats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## single subj scatterplot\n",
    "\n",
    "with **active** task on x-axis and **yoked** task on y-axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:21.709367Z",
     "start_time": "2021-05-28T13:46:21.686651Z"
    }
   },
   "outputs": [],
   "source": [
    "# Collapse df to wide format\n",
    "dd = pd.pivot_table(\n",
    "    accuracy_df,\n",
    "    index=[\"subject\", \"accuracy_type\"],\n",
    "    values=\"correct_choice\",\n",
    "    columns=\"task\",\n",
    ")\n",
    "dd = dd.reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:21.717808Z",
     "start_time": "2021-05-28T13:46:21.712080Z"
    }
   },
   "outputs": [],
   "source": [
    "def _label_point(x, y, val, ax):\n",
    "    \"\"\"Label points in a plot: https://stackoverflow.com/a/46028674/5201771.\"\"\"\n",
    "    a = pd.concat({\"x\": x, \"y\": y, \"val\": val}, axis=1)\n",
    "    for i, point in a.iterrows():\n",
    "        ax.text(point[\"x\"], point[\"y\"], str(int(point[\"val\"])), fontsize=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:21.726138Z",
     "start_time": "2021-05-28T13:46:21.720217Z"
    }
   },
   "outputs": [],
   "source": [
    "# idxs for FIXED vs VARIABLE tasks\n",
    "variable_idx = dd[\"subject\"].isin(range(2, 41, 2)).to_numpy()\n",
    "fixed_idx = dd[\"subject\"].isin(range(1, 41, 2)).to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:22.544927Z",
     "start_time": "2021-05-28T13:46:21.728181Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot single subject accuracies\n",
    "fig, axs = plt.subplots(1, 2, figsize=(10, 6), sharex=True, sharey=False)\n",
    "\n",
    "for ax, taskset, xyidx in zip(\n",
    "    axs, [[\"AF\", \"YF\"], [\"AV\", \"YV\"]], [fixed_idx, variable_idx]\n",
    "):\n",
    "    x, y = taskset\n",
    "    legend = \"brief\" if x == \"AF\" else False\n",
    "    sns.scatterplot(\n",
    "        x=x,\n",
    "        y=y,\n",
    "        hue=\"accuracy_type\",\n",
    "        data=dd,\n",
    "        ax=ax,\n",
    "        legend=legend,\n",
    "        palette=sns.color_palette()[2:4],\n",
    "        markers=\".\",\n",
    "        hue_order=[\"true\", \"experienced\"],\n",
    "    )\n",
    "\n",
    "    _label_point(dd.loc[xyidx, x], dd.loc[xyidx, y], dd.loc[xyidx, \"subject\"], ax)\n",
    "\n",
    "    if legend:\n",
    "        ax.legend(loc=2)\n",
    "\n",
    "\n",
    "axs.flat[0].set_xlim((0.5, 1))\n",
    "axs.flat[0].set_ylim((0.5, 1))\n",
    "axs.flat[1].set_ylim((0.5, 1))\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.suptitle(\"Single subject accuracies\", y=1.02)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Is \"numeracy\" correlated with n_samples or accuracy?\n",
    "\n",
    "- According to Peters2012 --> more numeracy, more n_samples\n",
    "\n",
    "scoring according to:\n",
    "- Cokely, E. T., Galesic, M., Schulz, E., Ghazal, S., & Garcia-Retamero, R. (2012). Measuring Risk Literacy: The Berlin Numeracy Test. Judgment and Decision Making. https://doi.org/10.1037/t45862-000\n",
    "\n",
    "Dividing sample into 4 quartiles through \"adaptive scoring\":\n",
    "\n",
    "- quartile 1 --> got 1 wrong, 2a wrong \n",
    "- quartile 2 --> got 1 wrong, 2a right\n",
    "- quartile 3 --> got 1 right, 2b wrong, 3 wrong\n",
    "- quartile 4 --> got 1 right, 2b right ... OR 2b wrong, but 3 right\n",
    "\n",
    "The 4th quartile has the highest numeracy skill\n",
    "\n",
    "\n",
    "The adaptive scoring questions 1, 2a, 2b, and 3 correspond to q1, q4, q2, q3 in our data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:22.561434Z",
     "start_time": "2021-05-28T13:46:22.546557Z"
    }
   },
   "outputs": [],
   "source": [
    "df_bnt = get_df_bnt(BIDS_ROOT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:22.931566Z",
     "start_time": "2021-05-28T13:46:22.563034Z"
    }
   },
   "outputs": [],
   "source": [
    "# sanity check: high quartiles should also have an overall higher number\n",
    "# of correctly answered questions\n",
    "fig, ax = plt.subplots()\n",
    "sns.pointplot(x=\"bnt_quartile\", y=\"bnt_n_correct\", data=df_bnt, ci=68, ax=ax)\n",
    "ax.set_title(\"Number of correctly answererd questions \" \"for each scoring quartile\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:22.940708Z",
     "start_time": "2021-05-28T13:46:22.933456Z"
    }
   },
   "outputs": [],
   "source": [
    "# Sanity check 2: quartiles 1+2 should be around 50% of all cases if this\n",
    "# sample is comparable to the sample the original BNT study was based on\n",
    "np.sum(df_bnt[\"bnt_quartile\"] <= 2) / len(df_bnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:23.242898Z",
     "start_time": "2021-05-28T13:46:22.943461Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "sns.histplot(x=\"bnt_quartile\", data=df_bnt, ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## numeracy and n_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:23.272266Z",
     "start_time": "2021-05-28T13:46:23.244735Z"
    }
   },
   "outputs": [],
   "source": [
    "tmp = df_nsamples.groupby(\"subject\").mean().reset_index()\n",
    "tmp = tmp.merge(df_bnt[[\"subject\", \"bnt_quartile\"]], on=\"subject\")\n",
    "tmp = df_nsamples.merge(df_bnt[[\"subject\", \"bnt_quartile\"]], on=\"subject\")\n",
    "tmp.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:25.881204Z",
     "start_time": "2021-05-28T13:46:23.277665Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "sns.pointplot(\n",
    "    x=\"bnt_quartile\",\n",
    "    y=\"n_samples\",\n",
    "    data=tmp,\n",
    "    ci=68,\n",
    "    ax=ax,\n",
    "    estimator=np.mean,\n",
    "    color=\"blue\",\n",
    ")\n",
    "sns.pointplot(\n",
    "    x=\"bnt_quartile\",\n",
    "    y=\"n_samples\",\n",
    "    data=tmp,\n",
    "    ci=68,\n",
    "    ax=ax,\n",
    "    estimator=np.median,\n",
    "    color=\"red\",\n",
    ")\n",
    "sns.swarmplot(x=\"bnt_quartile\", y=\"n_samples\", data=tmp, ax=ax, size=0.5, color=\"black\")\n",
    "ax.set_title(\"Mean (blue) and Median (red)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## numeracy and accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:25.892361Z",
     "start_time": "2021-05-28T13:46:25.882769Z"
    }
   },
   "outputs": [],
   "source": [
    "tmp = accuracy_df.merge(df_bnt[[\"subject\", \"bnt_quartile\"]], on=\"subject\")\n",
    "tmp = tmp[tmp[\"task\"] != \"DESC\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:27.308942Z",
     "start_time": "2021-05-28T13:46:25.894778Z"
    }
   },
   "outputs": [],
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), sharex=True, sharey=True)\n",
    "\n",
    "for ax, acctype in zip([ax1, ax2], [\"experienced\", \"true\"]):\n",
    "\n",
    "    sns.pointplot(\n",
    "        x=\"bnt_quartile\",\n",
    "        y=\"correct_choice\",\n",
    "        data=tmp[tmp[\"accuracy_type\"] == acctype],\n",
    "        ci=68,\n",
    "        dodge=True,\n",
    "        hue=\"task\",\n",
    "        hue_order=[\"AV\", \"AF\", \"YV\", \"YF\"],\n",
    "        markers=[\"o\", \"x\", \"o\", \"x\"],\n",
    "        linestyles=[\"-\", \"--\", \"-\", \"--\"],\n",
    "        palette=[\n",
    "            sns.color_palette()[0],\n",
    "            sns.color_palette()[0],\n",
    "            sns.color_palette()[1],\n",
    "            sns.color_palette()[1],\n",
    "        ],\n",
    "        ax=ax,\n",
    "    )\n",
    "    ax.set_title(acctype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:27.341487Z",
     "start_time": "2021-05-28T13:46:27.311458Z"
    }
   },
   "outputs": [],
   "source": [
    "# Correlate accuracy (averaged over active/yoked per subj) and BNT\n",
    "bnt_acc_corr_df = (\n",
    "    tmp.groupby(\"subject\")[[\"correct_choice\", \"bnt_quartile\"]].mean().reset_index()\n",
    ")\n",
    "\n",
    "pingouin.correlation.corr(\n",
    "    bnt_acc_corr_df[\"correct_choice\"],\n",
    "    bnt_acc_corr_df[\"bnt_quartile\"],\n",
    "    alternative=\"greater\",  # for pingouin <0.4 this should be `tail=\"one-sided\",`\n",
    "    method=\"kendall\",\n",
    ").round(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use rpy2 for mixed anova\n",
    "\n",
    "see: https://www.marsja.se/r-from-python-rpy2-tutorial/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.529510Z",
     "start_time": "2021-05-28T13:46:29.506666Z"
    }
   },
   "outputs": [],
   "source": [
    "tmp = accuracy_df.merge(df_bnt[[\"subject\", \"bnt_quartile\"]], on=\"subject\")\n",
    "tmp = tmp[tmp[\"task\"] != \"DESC\"]\n",
    "\n",
    "acctype = \"experienced\"\n",
    "tmp = tmp[tmp[\"accuracy_type\"] == acctype]\n",
    "\n",
    "tmp.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.544592Z",
     "start_time": "2021-05-28T13:46:29.531865Z"
    }
   },
   "outputs": [],
   "source": [
    "for col in [\"stopping\", \"sampling\", \"task\", \"accuracy_type\"]:\n",
    "    tmp[col] = tmp[col].astype(\"category\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.567505Z",
     "start_time": "2021-05-28T13:46:29.552103Z"
    }
   },
   "outputs": [],
   "source": [
    "# convert pandas DF to R data.frame\n",
    "with localconverter(ro.default_converter + pandas2ri.converter):\n",
    "    r_from_pd_df = ro.conversion.py2rpy(tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.580638Z",
     "start_time": "2021-05-28T13:46:29.569346Z"
    }
   },
   "outputs": [],
   "source": [
    "r_from_pd_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.746430Z",
     "start_time": "2021-05-28T13:46:29.583027Z"
    }
   },
   "outputs": [],
   "source": [
    "afex = rpackages.importr(\"afex\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.818168Z",
     "start_time": "2021-05-28T13:46:29.750625Z"
    }
   },
   "outputs": [],
   "source": [
    "model = afex.aov_ez(\n",
    "    id=\"subject\",\n",
    "    dv=\"correct_choice\",\n",
    "    between=[\"stopping\", \"bnt_quartile\"],\n",
    "    within=\"sampling\",\n",
    "    data=r_from_pd_df,\n",
    ")\n",
    "print(\"acc type: {}\".format(acctype))\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# find stereotypic sampling pattern\n",
    "\n",
    "- for example, \"piecewise\": `a b a b a b` --> if sequence_length - 1 switches in sequence and sequence length > 4\n",
    "- or, \"comprehensive\": `a a a a b b b b` --> if single switch in sequence and sequence length > 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.910059Z",
     "start_time": "2021-05-28T13:46:29.819948Z"
    }
   },
   "outputs": [],
   "source": [
    "# Calculate switches using \"action\" and \"trial\" column\n",
    "# A switch happens when *within* a trial, the participant\n",
    "# changes from sampling one option to the other option\n",
    "# The first sample at the \"other option\" is considered the\n",
    "# switch trial.\n",
    "dat = zip(df[\"action\"].to_numpy(), df[\"trial\"].to_numpy())\n",
    "\n",
    "# initialize a \"previous trial\" variable for preventing\n",
    "# counting switches across trial boundaries\n",
    "prev_trial = -9999999\n",
    "\n",
    "# Collect booleans for each trial whether it was a switch\n",
    "switches = list()\n",
    "for act, trial in dat:\n",
    "    # if we go on to a new trial in the data, an action is\n",
    "    # never counted as a switch\n",
    "    if trial != prev_trial:\n",
    "        prev_act = act\n",
    "        prev_trial = trial\n",
    "    # Else, an action is a switch if the action is different\n",
    "    # from the action before\n",
    "    switches.append(act != prev_act)\n",
    "    prev_act = act\n",
    "\n",
    "df[\"switch\"] = switches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:29.932049Z",
     "start_time": "2021-05-28T13:46:29.912024Z"
    }
   },
   "outputs": [],
   "source": [
    "_df = df[df[\"task\"].isin([\"AF\", \"AV\"])][\n",
    "    [\"subject\", \"task\", \"trial\", \"sample\", \"action\", \"switch\"]\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:32.936901Z",
     "start_time": "2021-05-28T13:46:29.934120Z"
    }
   },
   "outputs": [],
   "source": [
    "# classify each trial into \"0=comprehensive\" (aaabbb), \"1=piecewise\"(ababab), ...\n",
    "# \"np.nan=invalid\"(seq length <4, or only one option), or \"2=other\"\n",
    "#\n",
    "# The <4 requirement is important to distinguish \"ab\" trials: these are\n",
    "# ambiguous: comprehensive or piecewise.\n",
    "#\n",
    "# when strict=True, the mean of actions must be 0.5, so only aabb, aaabbb, ...\n",
    "# and abab, ababab, ... are valid\n",
    "\n",
    "comprehensive_tolerance = 2\n",
    "strict = False\n",
    "\n",
    "classification = dict(subject=[], task=[], trial=[], classification=[])\n",
    "for meta, grp in _df.groupby([\"subject\", \"task\", \"trial\"]):\n",
    "    n_switches = grp[\"switch\"].sum()\n",
    "    n_samples = len(grp)\n",
    "\n",
    "    # too short sequences are invalid\n",
    "    if (len(grp[\"action\"].unique()) != 2) or (n_samples < 4):\n",
    "        classi = np.nan\n",
    "\n",
    "    # exactly one switch means \"comprehensive\"\n",
    "    elif (n_samples > 1) and n_switches == 1:\n",
    "        classi = 0\n",
    "\n",
    "        # Check within tolerance for equal samples from each option\n",
    "        # to exclude patterns like aaaaaaaab\n",
    "        diff_0_1 = np.abs(np.diff(grp[\"action\"].value_counts().to_numpy()))\n",
    "        if diff_0_1 > comprehensive_tolerance:\n",
    "            classi = 2\n",
    "\n",
    "        # In the strict case, left and right must be samples evenly\n",
    "        if strict and not np.allclose(grp[\"action\"].mean(), 0.5):\n",
    "            classi = 2\n",
    "\n",
    "    # switching each times meas \"piecewise\"\n",
    "    elif (n_samples >= 2) and n_switches == (n_samples - 1):\n",
    "        classi = 1\n",
    "\n",
    "        # In the strict case, left and right must be samples evenly\n",
    "        if strict and not np.allclose(grp[\"action\"].mean(), 0.5):\n",
    "            classi = 2\n",
    "\n",
    "    # all other are \"other\" (mixed)\n",
    "    else:\n",
    "        classi = 2\n",
    "\n",
    "    subj, task, trl = meta\n",
    "    classification[\"subject\"] += [subj]\n",
    "    classification[\"task\"] += [task]\n",
    "    classification[\"trial\"] += [trl]\n",
    "    classification[\"classification\"] += [classi]\n",
    "\n",
    "# make DF\n",
    "classification_df = pd.DataFrame(classification)\n",
    "assert len(classification_df) == 4000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:32.951572Z",
     "start_time": "2021-05-28T13:46:32.938381Z"
    }
   },
   "outputs": [],
   "source": [
    "n_na = len(classification_df) - len(classification_df.dropna())\n",
    "print(f\"Need to drop {n_na} trials for different reasons:\")\n",
    "print(\n",
    "    \"\\nonly one option chosen, or only one sample taken\"\n",
    "    \"\\n(strictly speaking the latter reason is the same as the former)\"\n",
    ")\n",
    "\n",
    "classification_df_na = classification_df.copy()\n",
    "classification_df = classification_df.dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:33.135932Z",
     "start_time": "2021-05-28T13:46:32.953573Z"
    }
   },
   "outputs": [],
   "source": [
    "# Summarize the proportion of used strategies per subject\n",
    "data = {\"subject\": [], \"task\": [], \"comprehensive\": [], \"piecewise\": [], \"other\": []}\n",
    "for meta, grp in classification_df.groupby([\"subject\", \"task\"]):\n",
    "    subj, task = meta\n",
    "\n",
    "    for i, style in enumerate([\"comprehensive\", \"piecewise\", \"other\"]):\n",
    "        try:\n",
    "            a = (grp[\"classification\"].value_counts() / len(grp))[i]\n",
    "        except KeyError:\n",
    "            a = 0.0\n",
    "\n",
    "        data[style] += [a]\n",
    "\n",
    "    data[\"subject\"] += [subj]\n",
    "    data[\"task\"] += [task]\n",
    "\n",
    "\n",
    "class_summary_df = pd.DataFrame(data)\n",
    "\n",
    "# stereotypic sampling is either comprehensive or piecewise\n",
    "class_summary_df[\"stereotypic\"] = (\n",
    "    class_summary_df[\"comprehensive\"] + class_summary_df[\"piecewise\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:34.641663Z",
     "start_time": "2021-05-28T13:46:33.139183Z"
    }
   },
   "outputs": [],
   "source": [
    "summary_stat_to_use = np.median  # np.median or np.mean\n",
    "\n",
    "# NOTE: np.median might give warnings for the cases where all subjs are excluded\n",
    "\n",
    "# plot tradeoff: What % of subj did use stereotypic sampling on **at least** % or trials?\n",
    "fig, axs = plt.subplots(1, 3, figsize=(10, 5), sharey=True)\n",
    "n = 100\n",
    "cutoffs = np.linspace(1, 0, n)\n",
    "nsubjs_norm = np.interp(np.arange(0, 41), (0, 40), (0, 1))\n",
    "\n",
    "for ipartition, partition in enumerate([[\"AF\"], [\"AV\"], [\"AF\", \"AV\"]]):\n",
    "\n",
    "    _df_to_work_on = class_summary_df[class_summary_df[\"task\"].isin(partition)]\n",
    "    ax = axs.flat[ipartition]\n",
    "\n",
    "    plot_data = np.full((3, n), np.nan)\n",
    "    for i, cutoff in enumerate(cutoffs):\n",
    "\n",
    "        nsubjs_in_data = len(_df_to_work_on[_df_to_work_on[\"other\"] < cutoff])\n",
    "        min_stereo = _df_to_work_on[_df_to_work_on[\"other\"] < cutoff][\n",
    "            \"stereotypic\"\n",
    "        ].min()\n",
    "\n",
    "        summary_stereo = summary_stat_to_use(\n",
    "            _df_to_work_on[_df_to_work_on[\"other\"] < cutoff][\"stereotypic\"]\n",
    "        )\n",
    "\n",
    "        plot_data[0, i] = nsubjs_in_data\n",
    "        plot_data[1, i] = min_stereo\n",
    "        plot_data[2, i] = summary_stereo\n",
    "\n",
    "    # plot\n",
    "    ax.plot(\n",
    "        cutoffs,\n",
    "        nsubjs_norm[plot_data[0, :].astype(int)],\n",
    "        color=\"r\",\n",
    "        label=\"subjects retained\",\n",
    "    )\n",
    "    ax.plot(cutoffs, plot_data[1, :], color=\"k\", label=\"minimum stereotypic sampling\")\n",
    "    ax.plot(\n",
    "        cutoffs,\n",
    "        plot_data[2, :],\n",
    "        color=\"m\",\n",
    "        label=f\"{summary_stat_to_use.__name__} stereotypic sampling\",\n",
    "    )\n",
    "\n",
    "    xlab = \"cutoff 'other'\"\n",
    "    if ipartition == 1:\n",
    "        xlab += \"\\n(subjects need to have less than this proportion of 'other' sampling styles,\"\n",
    "        xlab += \"\\ni.e., NOT 'comprehensive' or 'piecewise', to be retained)\"\n",
    "    ax.set(xlabel=xlab, ylabel=\"proportion\", title=partition)\n",
    "    ax.grid(\"on\")\n",
    "\n",
    "    if ipartition == 0:\n",
    "        leg = ax.legend()\n",
    "\n",
    "\n",
    "fig.suptitle(\n",
    "    \"How many subject sample stereotypically?\\n(i.e., 'comprehensive':aaabbb or 'piecewise':ababab)\",\n",
    "    y=1.05,\n",
    ")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:34.681318Z",
     "start_time": "2021-05-28T13:46:34.643552Z"
    }
   },
   "outputs": [],
   "source": [
    "cutoffs = [0.3, 0.3, 0.4]\n",
    "\n",
    "for ipartition, partition in enumerate([[\"AF\"], [\"AV\"], [\"AF\", \"AV\"]]):\n",
    "\n",
    "    _df_to_work_on = class_summary_df[class_summary_df[\"task\"].isin(partition)]\n",
    "    cutoff = cutoffs[ipartition]\n",
    "\n",
    "    nsubjs_in_data = len(_df_to_work_on[_df_to_work_on[\"other\"] < cutoff])\n",
    "    min_stereo = _df_to_work_on[_df_to_work_on[\"other\"] < cutoff][\"stereotypic\"].min()\n",
    "    summary_stereo = summary_stat_to_use(\n",
    "        _df_to_work_on[_df_to_work_on[\"other\"] < cutoff][\"stereotypic\"]\n",
    "    )\n",
    "\n",
    "    print(f\"\\n\\n{partition}\")\n",
    "    print(\n",
    "        f\"{nsubjs_in_data} subjects sampled at least {min_stereo*100:.2f}% of trials in stereotypic form.\"\n",
    "    )\n",
    "    print(\n",
    "        f\"{nsubjs_in_data} subjects sampled {summary_stereo*100:.2f}% of trials in stereotypic form on average ({summary_stat_to_use.__name__}).\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:34.704658Z",
     "start_time": "2021-05-28T13:46:34.684462Z"
    }
   },
   "outputs": [],
   "source": [
    "all_av_stereo = summary_stat_to_use(\n",
    "    (class_summary_df[class_summary_df[\"task\"] == \"AV\"][\"stereotypic\"]) * 100\n",
    ").round(2)\n",
    "\n",
    "all_af_stereo = summary_stat_to_use(\n",
    "    (class_summary_df[class_summary_df[\"task\"] == \"AF\"][\"stereotypic\"]) * 100\n",
    ").round(2)\n",
    "\n",
    "all_stereo = summary_stat_to_use((class_summary_df[\"stereotypic\"]) * 100).round(2)\n",
    "\n",
    "print(\n",
    "    f\"Over all subjects,\\n{all_av_stereo}% in AV,\\n{all_af_stereo}% in AF,\\n{all_stereo}% overall,\"\n",
    "    f\"\\nsampled in stereotypic form ...\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decision weights via logisitic regression\n",
    "\n",
    "Two vectors are needed:\n",
    "- `X`: the sample outcomes from both options (one option \"flipped\" such that 1=9, 2=8, ..., 9=1)\n",
    "    - this \"flip\" is equivalent to first mean-centering the vector of outcomes, and then flipping the signs\n",
    "    - e.g., 1, 2, 3, 4, 5, 6, 7, 8, 9 --> (mean-center) --> -4, -3, -2, -1, 0, 1, 2, 3, 4 --> (flip)  --> 4, 3, 2, 1, 0, -1, -2, -3, -4\n",
    "- `y`: the final choices for options (left: 0, right: 1)\n",
    "\n",
    "Obviously, `y` will have fewer entries than `X`. \n",
    "We need to repeat the `y` value of each trial for each sample in that trial.\n",
    "Finally, we do this for early, mid, and late samples separately\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:34.722622Z",
     "start_time": "2021-05-28T13:46:34.707746Z"
    }
   },
   "outputs": [],
   "source": [
    "def classify_timing(df, first_last_n):\n",
    "    \"\"\"Add a new column 'timing' to the `df` for early/mid/late samples.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    df : pandas.DataFrame\n",
    "        The behavioral data containing `n_samples` and `sample` column.\n",
    "    first_last_n : int\n",
    "        Controls the N *first* samples in a sequence and N *last* samples in a\n",
    "        sequence to be classified into \"early\" and \"late\" respectively. The\n",
    "        samples in between are *mid* samples..\n",
    "        Sequences that are 2*first_last_n samples long or shorter have\n",
    "        *mixed* timing and can later be dropped from analysis.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    df : pandas.DataFrame\n",
    "        A copy of the data with a new column \"timing\".\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    `first_last_n` should be picked low enough such that not too many trials\n",
    "    will end up with each sample classified as \"mixed\". At the same time,\n",
    "    `first_last_n` should be high enough to make the number of \"early\", \"mid\",\n",
    "    and \"late\" samples per trial as balanced as possible.\n",
    "\n",
    "    \"\"\"\n",
    "    df = df.copy()\n",
    "    df[\"timing\"] = np.nan\n",
    "\n",
    "    # NOTE: nth \"sample\" is 0-indexed ... n_samples is a count\n",
    "    # \"how many steps is *this* sample away from the sequence end?\" ...\n",
    "    # `1` is the end, due to zero-indexing\n",
    "    sample_diff = (df[\"n_samples\"] - df[\"sample\"]).to_numpy()\n",
    "\n",
    "    # \"how large does the `sample_diff` have to be so that the sample can be classified\n",
    "    # as early?\"\n",
    "    early_thresh = ((df[\"n_samples\"] + 1) - first_last_n).to_numpy()\n",
    "\n",
    "    # If there are `first_last_n` or less steps until the sequence end, we have\n",
    "    # a late sample\n",
    "    df.loc[sample_diff <= first_last_n, \"timing\"] = \"late\"\n",
    "\n",
    "    # If there are many steps left, we have an early sample\n",
    "    df.loc[sample_diff >= early_thresh, \"timing\"] = \"early\"\n",
    "\n",
    "    # The remaining rows are \"mid\" samples ...\n",
    "    df.loc[df[\"timing\"].isna(), \"timing\"] = \"mid\"\n",
    "\n",
    "    # ... EXCEPT those rows that are of sampling sequences that are too short\n",
    "    # to be divided into early, mid, and late\n",
    "    df.loc[(df[\"n_samples\"] <= first_last_n * 2), \"timing\"] = \"mixed\"\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:34.755820Z",
     "start_time": "2021-05-28T13:46:34.727555Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_log_reg_coef(df, subject, flip, selection):\n",
    "    \"\"\"Predict final choice by outcome history.\n",
    "\n",
    "    Using a logistic regression.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    df : pandas.DataFrame\n",
    "        The data, with columns [subject, sampling, stopping, trial, sample,\n",
    "        timing, action, outcome, fin_action], and no \"DESC\" task.\n",
    "    subject : int\n",
    "        The subject id.\n",
    "    flip : 0 | 1\n",
    "        The \"side\" to flip (0=left 1=right). Relates to whether an outcome\n",
    "        was sampled left or right. For the sake of the logistic regression,\n",
    "        we will \"flip\" the outcomes on one side such that 1 becomes 9, 2\n",
    "        becomes 8, ... and 9 becomes 1.\n",
    "    selection : all | timing | task | timing_task\n",
    "        On which part(s) of the data to calculate the logistic regression.\n",
    "        \"all\" will calculate over all data. \"timing\" will calculate three\n",
    "        separate logistic regressions (one for early, mid, and late each).\n",
    "        \"task\" will calculate two separate regressions (one for active,\n",
    "        one for yoked). Finally, \"timing_task\" will calculate six separate\n",
    "        regressions: early, mid, late each for active, and yoked.\n",
    "\n",
    "    Notes\n",
    "    -----\n",
    "    If you need a stopping\" factor next to \"sampling\" and are wondering why\n",
    "    this is not supplied with the `selection` parameter, remember that\n",
    "    \"stopping\" was a between factor, so this is implicit. For example, the\n",
    "    data from sub-01 is always \"fixed\", sub-02 always \"variable\", and so on.\n",
    "    You can use df[[\"subject\", \"task\"]] to see the mapping.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    coef : dict\n",
    "        The coefficient(s)\n",
    "\n",
    "    \"\"\"\n",
    "    # work on one subj\n",
    "    subj_df = df[df[\"subject\"] == subject]\n",
    "    subj_df = subj_df.reset_index(drop=True)\n",
    "\n",
    "    outcome = subj_df[\"outcome\"].to_numpy()\n",
    "    flip_idxs = (subj_df[\"action\"] == flip).to_numpy()\n",
    "\n",
    "    # \"10 - outcome\" flips outcome so that 1=9, 2=8, ... 9=1\n",
    "    outcome_flipped = 10 - outcome[flip_idxs]\n",
    "\n",
    "    # Add column to df\n",
    "    subj_df.insert(len(subj_df.columns), \"outcome_flipped\", outcome)\n",
    "    subj_df.loc[flip_idxs, \"outcome_flipped\"] = outcome_flipped\n",
    "\n",
    "    # sanity check non-flipped outcomes equal\n",
    "    test = subj_df[subj_df[\"action\"] != flip][[\"outcome\", \"outcome_flipped\"]].to_numpy()\n",
    "    np.testing.assert_array_equal(test[:, 0], test[:, 1])\n",
    "\n",
    "    # sanity check flipped outcomes + original outcomes == 10\n",
    "    test = subj_df[subj_df[\"action\"] == flip][[\"outcome\", \"outcome_flipped\"]].to_numpy()\n",
    "    np.testing.assert_allclose(test[:, 0] + test[:, 1], 10)\n",
    "\n",
    "    if selection == \"all\":\n",
    "        sel = {\"all\": np.arange(0, subj_df.shape[0])}\n",
    "    elif selection == \"timing\":\n",
    "        sel = {\n",
    "            \"early\": subj_df[\"timing\"] == \"early\",\n",
    "            \"mid\": subj_df[\"timing\"] == \"mid\",\n",
    "            \"late\": subj_df[\"timing\"] == \"late\",\n",
    "        }\n",
    "    elif selection == \"task\":\n",
    "        sel = {\n",
    "            \"active\": subj_df[\"sampling\"] == \"active\",\n",
    "            \"yoked\": subj_df[\"sampling\"] == \"yoked\",\n",
    "        }\n",
    "    elif selection == \"timing_task\":\n",
    "        sel = {\n",
    "            \"active/early\": (subj_df[\"sampling\"] == \"active\")\n",
    "            & (subj_df[\"timing\"] == \"early\"),\n",
    "            \"active/mid\": (subj_df[\"sampling\"] == \"active\")\n",
    "            & (subj_df[\"timing\"] == \"mid\"),\n",
    "            \"active/late\": (subj_df[\"sampling\"] == \"active\")\n",
    "            & (subj_df[\"timing\"] == \"late\"),\n",
    "            \"yoked/early\": (subj_df[\"sampling\"] == \"yoked\")\n",
    "            & (subj_df[\"timing\"] == \"early\"),\n",
    "            \"yoked/mid\": (subj_df[\"sampling\"] == \"yoked\")\n",
    "            & (subj_df[\"timing\"] == \"mid\"),\n",
    "            \"yoked/late\": (subj_df[\"sampling\"] == \"yoked\")\n",
    "            & (subj_df[\"timing\"] == \"late\"),\n",
    "        }\n",
    "    else:\n",
    "        raise ValueError(f\"What is {selection}?\")\n",
    "\n",
    "    coef = {}\n",
    "    for selkey, selval in sel.items():\n",
    "        # IV and DV for linear regression\n",
    "        X = subj_df.loc[selval, \"outcome_flipped\"].to_numpy()\n",
    "        X = np.atleast_2d(X).T\n",
    "        y = subj_df.loc[selval, \"fin_action\"].to_numpy()\n",
    "\n",
    "        # first output is the coef ot the intercept\n",
    "        _, this_coef = pingouin.logistic_regression(\n",
    "            X, y, coef_only=True, random_state=42\n",
    "        )\n",
    "\n",
    "        # \"save\"\n",
    "        coef[selkey] = this_coef\n",
    "\n",
    "    return coef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:35.536218Z",
     "start_time": "2021-05-28T13:46:35.346026Z"
    }
   },
   "outputs": [],
   "source": [
    "# prep df\n",
    "first_last_n = 2\n",
    "tmp_logreg = classify_timing(df, first_last_n=first_last_n)\n",
    "\n",
    "\n",
    "tmp_logreg = tmp_logreg[tmp_logreg[\"task\"] != \"DESC\"]\n",
    "\n",
    "tmp_logreg.loc[:, \"sampling\"] = np.array([\"yoked\", \"active\"])[\n",
    "    tmp_logreg[\"task\"].str.startswith(\"A\").to_numpy(dtype=int)\n",
    "]\n",
    "tmp_logreg.loc[:, \"stopping\"] = np.array([\"variable\", \"fixed\"])[\n",
    "    tmp_logreg[\"task\"].str.endswith(\"F\").to_numpy(dtype=int)\n",
    "]\n",
    "\n",
    "cols = [\n",
    "    \"subject\",\n",
    "    \"sampling\",\n",
    "    \"stopping\",\n",
    "    \"trial\",\n",
    "    \"sample\",\n",
    "    \"timing\",\n",
    "    \"action\",\n",
    "    \"outcome\",\n",
    "    \"fin_action\",\n",
    "]\n",
    "\n",
    "tmp_logreg = tmp_logreg[cols]\n",
    "tmp_logreg.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## early mid late ... WITH sampling/stopping division\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:37.207184Z",
     "start_time": "2021-05-28T13:46:35.538423Z"
    }
   },
   "outputs": [],
   "source": [
    "flip = 0\n",
    "selection = \"timing_task\"\n",
    "subjdfs = []\n",
    "for subj in range(1, 41):\n",
    "    # calculate\n",
    "    slope = get_log_reg_coef(tmp_logreg, subj, flip, selection)\n",
    "\n",
    "    colnames = {\n",
    "        \"all\": [\"data_used\"],\n",
    "        \"timing\": [\"timing\"],\n",
    "        \"task\": [\"sampling\"],\n",
    "        \"timing_task\": [\"sampling\", \"timing\"],\n",
    "    }[selection]\n",
    "\n",
    "    # \"save\"\n",
    "    tmp_logreg_subjdf = pd.DataFrame(\n",
    "        np.array([i.split(\"/\") for i in list(slope.keys())])\n",
    "    )\n",
    "    tmp_logreg_subjdf.columns = colnames\n",
    "    tmp_logreg_subjdf[\"slope\"] = np.array(list(slope.values()))\n",
    "    tmp_logreg_subjdf[\"subject\"] = subj\n",
    "    tmp_logreg_subjdf = tmp_logreg_subjdf[[\"subject\", *colnames, \"slope\"]]\n",
    "    subjdfs.append(tmp_logreg_subjdf)\n",
    "\n",
    "decision_weight_df = pd.concat(subjdfs)\n",
    "decision_weight_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:37.237013Z",
     "start_time": "2021-05-28T13:46:37.209277Z"
    }
   },
   "outputs": [],
   "source": [
    "# merge the stopping factor on to decision_weight_df\n",
    "tmp_stopping = df.drop_duplicates(subset=[\"subject\"], keep=\"first\")[\n",
    "    [\"subject\", \"task\"]\n",
    "].reset_index(drop=True)\n",
    "tmp_stopping[\"stopping\"] = tmp_stopping[\"task\"].map({\"AF\": \"fixed\", \"AV\": \"variable\"})\n",
    "\n",
    "decision_weight_df = decision_weight_df.merge(\n",
    "    tmp_stopping[[\"subject\", \"stopping\"]], on=[\"subject\"]\n",
    ")\n",
    "decision_weight_df = decision_weight_df[\n",
    "    [\"subject\", \"sampling\", \"stopping\", \"timing\", \"slope\"]\n",
    "]\n",
    "decision_weight_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:37.253412Z",
     "start_time": "2021-05-28T13:46:37.238976Z"
    }
   },
   "outputs": [],
   "source": [
    "# save data for publication plots\n",
    "plotdir = op.join(BIDS_ROOT, \"code\", \"publication_plots\")\n",
    "os.makedirs(plotdir, exist_ok=True)\n",
    "\n",
    "fname = \"decision_weights_logreg_tasks.csv\"\n",
    "fname = op.join(plotdir, fname)\n",
    "\n",
    "decision_weight_df.to_csv(fname, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:37.876316Z",
     "start_time": "2021-05-28T13:46:37.256278Z"
    }
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"talk\", font_scale=1):\n",
    "    fig, ax = plt.subplots()\n",
    "    sns.lineplot(\n",
    "        x=\"timing\",\n",
    "        y=\"slope\",\n",
    "        hue=\"stopping\",\n",
    "        hue_order=[\"fixed\", \"variable\"],\n",
    "        style=\"sampling\",\n",
    "        style_order=[\"active\", \"yoked\"],\n",
    "        data=decision_weight_df,\n",
    "        ci=68,\n",
    "        ax=ax,\n",
    "        legend=False,\n",
    "    )\n",
    "\n",
    "    ax.axhline(0, color=\"black\")\n",
    "    ax.set_ylabel(\"decision weight\")\n",
    "    ax.set_xlabel(\"sample position\")\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",
    "    legend1 = fig.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"lower right\",\n",
    "        bbox_to_anchor=(0.95, 0.175),\n",
    "        framealpha=1,\n",
    "        frameon=False,\n",
    "        ncol=2,\n",
    "        title=\"sampling\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `late - early`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:39.532238Z",
     "start_time": "2021-05-28T13:46:37.877870Z"
    }
   },
   "outputs": [],
   "source": [
    "flip = 0\n",
    "selection = \"timing_task\"\n",
    "subjdfs = []\n",
    "for subj in range(1, 41):\n",
    "    # calculate\n",
    "    coef = get_log_reg_coef(tmp_logreg, subj, flip, selection)\n",
    "\n",
    "    colnames = {\n",
    "        \"all\": [\"data_used\"],\n",
    "        \"timing\": [\"timing\"],\n",
    "        \"task\": [\"sampling\"],\n",
    "        \"timing_task\": [\"sampling\", \"timing\"],\n",
    "    }[selection]\n",
    "\n",
    "    # \"save\"\n",
    "    tmp_subjdf = pd.DataFrame(np.array([i.split(\"/\") for i in list(coef.keys())]))\n",
    "    tmp_subjdf.columns = colnames\n",
    "    tmp_subjdf[\"slope\"] = np.array(list(coef.values()))\n",
    "    tmp_subjdf[\"subject\"] = subj\n",
    "    tmp_subjdf = tmp_subjdf[[\"subject\", *colnames, \"slope\"]]\n",
    "    subjdfs.append(tmp_subjdf)\n",
    "\n",
    "decision_weight_df = pd.concat(subjdfs)\n",
    "\n",
    "# sampling column must be capitalized to match with \"tmp\" df that we want\n",
    "# to merge on this df\n",
    "decision_weight_df[\"sampling\"] = decision_weight_df[\"sampling\"].str.capitalize()\n",
    "\n",
    "decision_weight_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:39.562282Z",
     "start_time": "2021-05-28T13:46:39.533764Z"
    }
   },
   "outputs": [],
   "source": [
    "# get stopping and task factors back\n",
    "tmp = pd.read_csv(op.join(BIDS_ROOT, \"participants.tsv\"), sep=\"\\t\")\n",
    "tmp[\"subject\"] = [int(i[-2:]) for i in tmp[\"participant_id\"]]\n",
    "tmp = tmp[[\"subject\", \"stopping\"]]\n",
    "tmp[\"stopping\"] = tmp[\"stopping\"].str.capitalize()\n",
    "\n",
    "decision_weight_df = decision_weight_df.merge(tmp, on=[\"subject\"])\n",
    "\n",
    "decision_weight_df[\"task\"] = (\n",
    "    decision_weight_df[\"sampling\"].str[0] + decision_weight_df[\"stopping\"].str[0]\n",
    ").str.upper()\n",
    "\n",
    "\n",
    "decision_weight_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:39.624802Z",
     "start_time": "2021-05-28T13:46:39.564710Z"
    }
   },
   "outputs": [],
   "source": [
    "# subtract second from first --> first - second\n",
    "first = \"late\"\n",
    "second = \"early\"\n",
    "colname = f\"slope_{first}_minus_{second}\"\n",
    "tmp_dwdf = decision_weight_df[decision_weight_df[\"timing\"] == first].merge(\n",
    "    decision_weight_df[decision_weight_df[\"timing\"] == second],\n",
    "    on=[\"subject\", \"task\"],\n",
    "    suffixes=(\"_\" + first, \"_\" + second),\n",
    ")\n",
    "tmp_dwdf.insert(0, colname, tmp_dwdf[f\"slope_{first}\"] - tmp_dwdf[f\"slope_{second}\"])\n",
    "\n",
    "tmp_dwdf[\"sampling\"] = [\n",
    "    \"Active\" if i.startswith(\"A\") else \"Yoked\" for i in tmp_dwdf[\"task\"]\n",
    "]\n",
    "tmp_dwdf[\"stopping\"] = [\n",
    "    \"Fixed\" if i.endswith(\"F\") else \"Variable\" for i in tmp_dwdf[\"task\"]\n",
    "]\n",
    "\n",
    "tmp_dwdf.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### correlate recency with numeracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:39.648910Z",
     "start_time": "2021-05-28T13:46:39.627240Z"
    }
   },
   "outputs": [],
   "source": [
    "bnt_recency_corr_df = bnt_acc_corr_df.merge(\n",
    "    tmp_dwdf.groupby([\"subject\"])[colname].mean().reset_index(), on=\"subject\"\n",
    ")\n",
    "bnt_recency_corr_df = bnt_recency_corr_df.rename(columns={colname: \"recency\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:39.672506Z",
     "start_time": "2021-05-28T13:46:39.652281Z"
    }
   },
   "outputs": [],
   "source": [
    "pingouin.correlation.corr(\n",
    "    bnt_recency_corr_df[\"recency\"],\n",
    "    bnt_recency_corr_df[\"bnt_quartile\"],\n",
    "    alternative=\"greater\",  # for pingouin <0.4 this should be `tail=\"one-sided\",`\n",
    "    method=\"kendall\",\n",
    ").round(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:39.684782Z",
     "start_time": "2021-05-28T13:46:39.675277Z"
    }
   },
   "outputs": [],
   "source": [
    "# save data for publication plots\n",
    "plotdir = op.join(BIDS_ROOT, \"code\", \"publication_plots\")\n",
    "os.makedirs(plotdir, exist_ok=True)\n",
    "\n",
    "fname = \"beh_recency_logreg.csv\"\n",
    "fname = op.join(plotdir, fname)\n",
    "\n",
    "tmp_dwdf.to_csv(fname, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:40.227997Z",
     "start_time": "2021-05-28T13:46:39.688141Z"
    }
   },
   "outputs": [],
   "source": [
    "with sns.plotting_context(\"notebook\", font_scale=1.3):\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    data = tmp_dwdf\n",
    "    order = [\"Active\", \"Yoked\"]\n",
    "    hue_order = [\"Fixed\", \"Variable\"]\n",
    "\n",
    "    sns.pointplot(\n",
    "        x=\"sampling\",\n",
    "        y=colname,\n",
    "        hue=\"stopping\",\n",
    "        data=data,\n",
    "        dodge=True,\n",
    "        ci=68,\n",
    "        order=order,\n",
    "        hue_order=hue_order,\n",
    "        ax=ax,\n",
    "        markers=\"o\",\n",
    "    )\n",
    "\n",
    "    sns.swarmplot(\n",
    "        x=\"sampling\",\n",
    "        y=colname,\n",
    "        hue=\"stopping\",\n",
    "        data=data,\n",
    "        order=order,\n",
    "        hue_order=hue_order,\n",
    "        ax=ax,\n",
    "        dodge=True,\n",
    "        size=3,\n",
    "    )\n",
    "\n",
    "    # add legend\n",
    "    # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html\n",
    "    from matplotlib.lines import Line2D\n",
    "\n",
    "    legend_elements = [\n",
    "        Line2D([0], [0], color=sns.color_palette()[0], marker=\"o\", label=\"Fixed\"),\n",
    "        Line2D([0], [0], color=sns.color_palette()[1], marker=\"o\", label=\"Variable\"),\n",
    "    ]\n",
    "\n",
    "    ax.legend(\n",
    "        handles=legend_elements,\n",
    "        loc=\"upper left\",\n",
    "        prop={\"size\": 13},\n",
    "        framealpha=1,\n",
    "    )\n",
    "\n",
    "    ax.set_ylabel(f\"slope difference\\n{first}-{second}\")\n",
    "\n",
    "    # https://stackoverflow.com/a/51157346/5201771\n",
    "    # 2-4 3-5\n",
    "    for i_dots, (idx0, idx1) in enumerate([(2, 4), (3, 5)]):\n",
    "        locs1 = ax.get_children()[idx0].get_offsets()\n",
    "        locs2 = ax.get_children()[idx1].get_offsets()\n",
    "\n",
    "        # Need to sort locs, so data corresponds\n",
    "        sort_idxs_list = []\n",
    "        sampling = order[i_dots]\n",
    "        for stopping in hue_order:\n",
    "            arr = data[(data[\"sampling\"] == sampling) & (data[\"stopping\"] == stopping)][\n",
    "                colname\n",
    "            ].to_numpy()\n",
    "            sort_idxs_list += [np.argsort(arr)]\n",
    "\n",
    "        locs2_sorted = locs2[sort_idxs_list[1].argsort()][sort_idxs_list[0]]\n",
    "\n",
    "        for i in range(locs1.shape[0]):\n",
    "            _x = [locs1[i, 0], locs2_sorted[i, 0]]\n",
    "            _y = [locs1[i, 1], locs2_sorted[i, 1]]\n",
    "            ax.plot(_x, _y, color=\"black\", alpha=0.1)\n",
    "\n",
    "    ax.axhline(0, color=\"black\")\n",
    "    ax.set_title(\"Regression slopes over weights\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:40.329132Z",
     "start_time": "2021-05-28T13:46:40.229679Z"
    }
   },
   "outputs": [],
   "source": [
    "model = pingouin.mixed_anova(\n",
    "    data=tmp_dwdf, dv=colname, within=\"sampling\", subject=\"subject\", between=\"stopping\"\n",
    ")\n",
    "model.round(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-05-28T13:46:40.481376Z",
     "start_time": "2021-05-28T13:46:40.330972Z"
    }
   },
   "outputs": [],
   "source": [
    "stats = pingouin.pairwise_ttests(\n",
    "    data=tmp_dwdf,\n",
    "    dv=colname,\n",
    "    within=\"sampling\",\n",
    "    between=\"stopping\",\n",
    "    subject=\"subject\",\n",
    "    padjust=\"bonf\",\n",
    "    within_first=False,\n",
    "    effsize=\"cohen\",\n",
    ")\n",
    "\n",
    "display(stats.round(3))"
   ]
  },
  {
   "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": "592.1px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}