BayesianGLM_Behavior.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9199e51a",
   "metadata": {},
   "source": [
    "# Bayesian analysis of emotion-mediated memory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89921b81-06ad-4757-8349-56998e1d31b1",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "In this notebook I will conduct (somewhat) speedy Bayesian mixed-effects logistic regression to test how word-level features predict memory, with subject-level random effects. \n",
    "\n",
    "\n",
    "Notes\n",
    "___\n",
    "\n",
    "\n",
    "- Bambi needs up-to-date xarray (for 'unify_chunks' method), while the data I pickled used xarray 0.13.0 \n",
    "\n",
    "\n",
    "- Bambi can't multi-process sampling across multiple chains if you import joblib, so need to make sure to only load in csv data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f89f70c3-f2ac-49d9-93d0-6b592fea3dc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import arviz as az\n",
    "import bambi as bmb\n",
    "import argparse\n",
    "from os.path import join\n",
    "import pandas as pd\n",
    "from pandas.api.types import CategoricalDtype\n",
    "import seaborn as sns \n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.backends.backend_pdf import PdfPages\n",
    "from Bayesian_model_utils import run_model, plot_res, print_latex_table, plot_predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "930778eb-3d0f-47c3-9bcf-a9c50d0e2f7e",
   "metadata": {},
   "source": [
    "## How do arousal and valence predict recall performance?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba60cd5-e768-4e2b-8cdf-427495b6dacc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, let's build our model and plot the priors to see if they're sensible \n",
    "\n",
    "# Old dataset of pts with Hipp or Amy electrodes \n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all.csv', low_memory=False)\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "# Center valence using mean of whole database to capture relative polarity\n",
    "behav_df.valence = behav_df.valence - 0.5\n",
    "\n",
    "behav_df = behav_df[['recalled', 'arousal', 'valence', 'subject']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['arousal', 'valence'] \n",
    "Intx = ['arousal:valence']\n",
    "rand_effect = ['subject']\n",
    "\n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Set some terms \n",
    "rand_term = [f'(1|{x})' for x in rand_effect]\n",
    "formula = f'{y} ~ 1+'+'+'.join(rand_term)+'+'+'+'.join(X)+'+'+'+'.join(Intx)\n",
    "model_fam = 'bernoulli'\n",
    "priors=None\n",
    "categorical=None\n",
    "\n",
    "# construct the model \n",
    "model = bmb.Model(formula=formula, \n",
    "              data=behav_df[rand_effect + [y] + X],\n",
    "             family=model_fam,\n",
    "             priors=priors,\n",
    "             categorical=categorical)\n",
    "\n",
    "model.build()\n",
    "\n",
    "model.plot_priors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ede9821-1035-40b6-8b59-76a84ffed8d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make the graphical model: \n",
    "\n",
    "model.graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4645b6f1-b496-461b-9710-b84576e1e55b",
   "metadata": {},
   "source": [
    "Our model formulation looks sensible. Let's do it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c034824-571e-457a-a92c-644d0e55927a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Model 1: Effect of arousal and continuous, linear valence on recall \n",
    "\n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all.csv', low_memory=False)\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "# Center valence using mean of whole database to capture relative polarity\n",
    "behav_df.valence = behav_df.valence - 0.5\n",
    "\n",
    "behav_df = behav_df[['recalled', 'arousal', 'valence', 'subject']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['arousal', 'valence'] \n",
    "Intx = ['arousal:valence']\n",
    "rand_effect = ['subject']\n",
    "\n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Run the model \n",
    "run_model(behav_df, y, X, Intx, rand_effect, \n",
    "          chains=4, cores=4, tune=1000, draws=1000, \n",
    "          label = label, rand_slopes=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bc4b863-e0b1-42f8-82a5-e201c415de19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 2: Effect of arousal and binned valence on recall \n",
    "\n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all.csv', low_memory=False)\n",
    "\n",
    "behav_df = behav_df.drop(columns=['Unnamed: 0'])\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "behav_df = behav_df[['recalled', 'arousal', 'CV', 'subject']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['arousal', 'CV'] \n",
    "Intx = ['arousal:CV']\n",
    "rand_effect = ['subject']\n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Run the model \n",
    "run_model(behav_df, y, X, Intx, rand_effect, categorical=categorical, \n",
    "          chains=4, cores=4, tune=1000, draws=1000, \n",
    "          label = label, rand_slopes=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18546c50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 3: Effect of arousal and squared valence on recall \n",
    "\n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all.csv', low_memory=False)\n",
    "\n",
    "behav_df = behav_df.drop(columns=['Unnamed: 0'])\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "# Center valence using mean of whole database to capture relative polarity\n",
    "behav_df.valence = behav_df.valence - 0.5\n",
    "\n",
    "behav_df['valence_squared'] = behav_df.valence**2\n",
    "\n",
    "behav_df = behav_df[['recalled', 'arousal', 'valence_squared', 'subject']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['arousal', 'valence_squared'] \n",
    "Intx = ['arousal:valence_squared']\n",
    "rand_effect = ['subject']\n",
    "categorical =['CV'] \n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Run the model \n",
    "run_model(behav_df, y, X, Intx, rand_effect, categorical=categorical, \n",
    "          chains=4, cores=4, tune=1000, draws=1000, \n",
    "          label = label, rand_slopes=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ded18234-6cd7-4de6-bf90-10063119f3db",
   "metadata": {},
   "source": [
    "Let's now do model comparison using leave-one-out cross-validation and determine which formulation of valence is best for prediction. \n",
    "\n",
    "Source:\n",
    "https://arviz-devs.github.io/arviz/api/generated/arviz.compare.html\n",
    "\n",
    "https://arviz-devs.github.io/arviz/api/generated/arviz.plot_compare.html\n",
    "\n",
    "https://bambinos.github.io/bambi/main/notebooks/model_comparison.html?highlight=waic\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ef2d7df-4688-4c81-b6c1-5b99913f4a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model comparison: use PSI-LOO\n",
    "\n",
    "# Load each model\n",
    "save_dir = '/home1/salman.qasim/Salman_Project/FR_Emotion/BayesModels'\n",
    "linear_valence = az.from_netcdf(f'{save_dir}/recalled_arousal_valence_allFRpts_model')\n",
    "valence_squared = az.from_netcdf(f'{save_dir}/recalled_arousal_valence_squared_allFRpts_model')\n",
    "binned_valence = az.from_netcdf(f'{save_dir}/recalled_arousal_CV_allFRpts_model')\n",
    "\n",
    "models = {\"linear_valence\": linear_valence, \n",
    "          \"valence_squared\": valence_squared,\n",
    "         \"binned_valence\": binned_valence}\n",
    "\n",
    "df_compare = az.compare(models, ic='loo')\n",
    "plot_df = df_compare.reset_index().rename(columns={'index':'model'})[['model', 'loo']]\n",
    "\n",
    "c_string = '|c'*plot_df.shape[1] + '|'\n",
    "print(plot_df.to_latex(index=False, \n",
    "                                        column_format=c_string).replace(\"\\\\\\n\", \"\\\\ \\hline\\n\"))  \n",
    "az.plot_compare(df_compare, insample_dev=True)\n",
    "plt.savefig(f'{save_dir}/behav_valence_comparison.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd4ca9ac-dcf0-4d38-9cd2-1ff3678cb64e",
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_compare(df_compare, insample_dev=True)\n",
    "plt.savefig(f'{save_dir}/arousal_valence_comparison.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548a46ba-b64d-48af-95d8-b90c5a040c5f",
   "metadata": {},
   "source": [
    "## How do arousal and valence interact with direct-brain stimulation ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3493d21-2ec3-42d7-ba0c-f42d0a94cbfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 1: Effect of stimulation to Hipp only on recall as a function of arousal and valence: \n",
    "\n",
    "stim_behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/all_stim_ev.csv')\n",
    "\n",
    "# Get only encoding events\n",
    "stim_behav_df = stim_behav_df[stim_behav_df.type=='WORD']\n",
    "\n",
    "# Get only hippocampal stim:\n",
    "stim_behav_df = stim_behav_df[(stim_behav_df.stim_reg=='Hipp')]\n",
    " \n",
    "# Center valence using mean of whole database to capture relative polarity\n",
    "stim_behav_df['valence'] = stim_behav_df['valence'] - 0.5\n",
    "\n",
    "cat_type = CategoricalDtype(categories=['right', 'left'], ordered=True)\n",
    "stim_behav_df.stim_hemi = stim_behav_df.stim_hemi.astype(cat_type)\n",
    "\n",
    "stim_behav_df = stim_behav_df[['recalled', 'valence', 'arousal', 'subject', 'is_stim', 'stim_hemi']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['is_stim', 'arousal', 'valence', 'stim_hemi'] \n",
    "Intx = ['is_stim:arousal', 'is_stim:valence']\n",
    "\n",
    "rand_effect = ['subject']\n",
    "categorical = ['is_stim', 'stim_hemi']\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_HippOnly' \n",
    "\n",
    "# Drop nan data\n",
    "stim_behav_df = stim_behav_df.dropna(subset=X)\n",
    "\n",
    "run_model(stim_behav_df, y, X, Intx, rand_effect, chains=4, cores=4, tune=2000, draws=2000, target_accept=0.95,\n",
    "          rand_slopes=False, categorical=categorical, label=label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de1402b6-134e-494e-a82a-0d9e54d2e6cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 2: Effect of stimulation to MTL only on recall as a function of arousal and valence: \n",
    "\n",
    "stim_behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/all_stim_ev.csv')\n",
    "\n",
    "# Get only encoding events\n",
    "stim_behav_df = stim_behav_df[stim_behav_df.type=='WORD']\n",
    "\n",
    "# Get only hippocampal stim:\n",
    "stim_behav_df = stim_behav_df[(stim_behav_df.stim_reg=='MTL')]\n",
    " \n",
    "# Center valence using mean of whole database to capture relative polarity\n",
    "stim_behav_df['valence'] = stim_behav_df['valence'] - 0.5\n",
    "\n",
    "cat_type = CategoricalDtype(categories=['right', 'left'], ordered=True)\n",
    "stim_behav_df.stim_hemi = stim_behav_df.stim_hemi.astype(cat_type)\n",
    "\n",
    "stim_behav_df = stim_behav_df[['recalled', 'valence', 'arousal', 'subject', 'is_stim', 'stim_hemi']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['is_stim', 'arousal', 'valence', 'stim_hemi'] \n",
    "Intx = ['is_stim:arousal', 'is_stim:valence']\n",
    "\n",
    "rand_effect = ['subject']\n",
    "categorical = ['is_stim', 'stim_hemi']\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_MTLOnly' \n",
    "\n",
    "# Drop nan data\n",
    "stim_behav_df = stim_behav_df.dropna(subset=X)\n",
    "\n",
    "run_model(stim_behav_df, y, X, Intx, rand_effect, chains=4, cores=4, tune=2000, draws=2000, target_accept=0.95,\n",
    "          rand_slopes=False, categorical=categorical, label=label)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548a46ba-b64d-48af-95d8-b90c5a040c5f",
   "metadata": {},
   "source": [
    "## How do arousal and valence interact with direct-brain stimulation ?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b5492aa-56aa-4cd8-96b4-a246dafcdbb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 1: Account for both BDI and BAI scores\n",
    "\n",
    "# First, load the Beck Scores: \n",
    "BeckScores = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/Beck_Scores.csv')\n",
    "BeckScores.rename(columns={'Subject Code':'subject'}, inplace=True)\n",
    "\n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all_patients_revision.csv', memory_map=True)\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "# add in the Beck Scores \n",
    "behav_df = behav_df.merge(BeckScores, on='subject')\n",
    "\n",
    "behav_df['valence'] = behav_df['valence'] - 0.5\n",
    "\n",
    "behav_df = behav_df[['recalled', 'BDI', 'BAI', 'subject']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['BDI', 'BAI']\n",
    "Intx= ['BDI:BAI']\n",
    "\n",
    "rand_effect = ['subject']\n",
    "\n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Run the model \n",
    "run_model(behav_df, y, X, Intx, rand_effect, chains=4, cores=4, tune=1000, draws=1000, \n",
    "          label = label, rand_slopes=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c485d90-bad8-4691-8151-0bb37f6a171d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 2: Account for just BDI and interaction with arousal\n",
    "\n",
    "# First, load the Beck Scores: \n",
    "BeckScores = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/Beck_Scores.csv')\n",
    "BeckScores.rename(columns={'Subject Code':'subject'}, inplace=True)\n",
    "\n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all_patients_revision.csv', memory_map=True)\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "# add in the Beck Scores \n",
    "behav_df = behav_df.merge(BeckScores, on='subject')\n",
    "\n",
    "# make squared valence\n",
    "behav_df['valence'] = behav_df['valence'] - 0.5\n",
    "\n",
    "behav_df = behav_df[['recalled', 'subject', 'BDI', 'arousal']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['BDI', 'arousal']\n",
    "Intx= ['BDI:arousal']\n",
    "\n",
    "rand_effect = ['subject']\n",
    "\n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Run the model \n",
    "run_model(behav_df, y, X, Intx, rand_effect, chains=4, cores=4, tune=1000, draws=1000, \n",
    "          label = label, rand_slopes=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d539fb0-f728-44e3-82d7-32ed53c8e439",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model 2: Account for just BDI and interaction with valence\n",
    "\n",
    "# First, load the Beck Scores: \n",
    "BeckScores = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/Beck_Scores.csv')\n",
    "BeckScores.rename(columns={'Subject Code':'subject'}, inplace=True)\n",
    "\n",
    "behav_df = pd.read_csv('/home1/salman.qasim/Salman_Project/FR_Emotion/events_all_patients_revision.csv', memory_map=True)\n",
    "\n",
    "# Get only encoding events\n",
    "behav_df = behav_df[behav_df.type=='WORD']\n",
    "\n",
    "# add in the Beck Scores \n",
    "behav_df = behav_df.merge(BeckScores, on='subject')\n",
    "\n",
    "# make squared valence\n",
    "behav_df['valence'] = (behav_df.valence - 0.5)\n",
    "\n",
    "behav_df = behav_df[['recalled', 'subject', 'BDI', 'valence']]\n",
    "\n",
    "y = 'recalled'\n",
    "X = ['BDI', 'valence']\n",
    "Intx= ['BDI:valence']\n",
    "\n",
    "rand_effect = ['subject']\n",
    "\n",
    "# Drop nan data\n",
    "behav_df = behav_df.dropna(subset=X)\n",
    "\n",
    "label = (f\"{y}\" + \"_{}\"*len(X)).format(*X) + '_allFRpts'\n",
    "\n",
    "# Run the model \n",
    "run_model(behav_df, y, X, Intx, rand_effect, chains=4, cores=4, tune=1000, draws=1000, \n",
    "          label = label, rand_slopes=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6527a04e-7b66-4947-bfcb-d6e399091390",
   "metadata": {},
   "source": [
    "Is BDI more predictive of memory when taking word arousal or word valence into account?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef5cda1-23c1-42df-b9d0-574649701502",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model comparison: use PSI-LOO\n",
    "\n",
    "# Load each model\n",
    "save_dir = '/home1/salman.qasim/Salman_Project/FR_Emotion/BayesModels'\n",
    "BDI_arousal = az.from_netcdf(f'{save_dir}/recalled_BDI_arousal_allFRpts_model')\n",
    "BDI_valence = az.from_netcdf(f'{save_dir}/recalled_BDI_valence_squared_allFRpts_model')\n",
    "\n",
    "models = {\"BDI_arousal\": BDI_arousal, \n",
    "          \"BDI_valence\": BDI_valence}\n",
    "\n",
    "df_compare = az.compare(models, ic='loo')\n",
    "plot_df = df_compare.reset_index().rename(columns={'index':'model'})[['model', 'loo']]\n",
    "\n",
    "c_string = '|c'*plot_df.shape[1] + '|'\n",
    "print(plot_df.to_latex(index=False, \n",
    "                                        column_format=c_string).replace(\"\\\\\\n\", \"\\\\ \\hline\\n\"))  \n",
    "az.plot_compare(df_compare, insample_dev=True)\n",
    "plt.savefig(f'{save_dir}/BDI_arousal_valence_comparison.pdf')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.10 64-bit",
   "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.9.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}