Filipowicz_etal_MBMF_Bound_Computing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compute simulation and subject bounds for Daw 2-step task\n",
    "\n",
    "This following scripts compute the empirical information bottleneck, predictive information, and distance from the bound for all of the simulations and all of the subjects from the Kool et al., 2016, PLoS Comp. Biol. (https://doi.org/10.1371/journal.pcbi.1005090) assesment of model-based and model-free strategies in this task environment (their code and data can be found here: https://github.com/wkool/tradeoffs).\n",
    "\n",
    "The results from these processing steps are used in 'Filipowicz_etal_EIB_MBMF.ipynb' for the main analyses and comparisons.\n",
    "\n",
    "## 0) Loading necessary packages and helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using matplotlib backend: MacOSX\n",
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab\n",
    "%matplotlib inline\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "from scipy.stats import spearmanr, entropy\n",
    "import scipy.io as spio\n",
    "\n",
    "import glob, pickle\n",
    "\n",
    "import sys, csv\n",
    "\n",
    "\n",
    "######################\n",
    "## HELPER FUNCTIONS ##\n",
    "######################\n",
    "\n",
    "#Functions to creat labels\n",
    "def trial_label(x):\n",
    "    as_strings = [num.astype(str) for num in x]\n",
    "    return \"\".join(as_strings)\n",
    "\n",
    "def make_features(trials_data):\n",
    "    labeled_data = np.apply_along_axis(trial_label, 0, trials_data)\n",
    "    combos = np.unique(labeled_data)\n",
    "    string_to_index = dict(zip(combos, np.arange(len(combos))))\n",
    "    map_to_index = np.vectorize(lambda x: string_to_index[x])\n",
    "    mapped_data = map_to_index(labeled_data)\n",
    "    return mapped_data\n",
    "\n",
    "def get_marginal(x):\n",
    "    \"\"\"\n",
    "    Helper function to compute and return marginal probability distribution for a 1d vector (x)\n",
    "    \"\"\"\n",
    "    px = np.array([np.sum(x==xi) for xi in np.sort(np.unique(x))])/len(x)\n",
    "    return(px)\n",
    "\n",
    "def get_joint(x, y):\n",
    "    \"\"\"\n",
    "    Computes joint probability distribution between 1d vectors x and y\n",
    "    \"\"\"\n",
    "    #  set up dictionary for joint distribution (x-->y-->freq)\n",
    "    joint_x_y = {}\n",
    "    \n",
    "    for x_un in np.unique(x):\n",
    "        joint_x_y[x_un] = dict(zip(np.unique(y), np.zeros(len(np.unique(y)))))\n",
    "        \n",
    "#    populate dictionary \n",
    "    for trial, x_val in enumerate(x):\n",
    "        y_val = y[trial]\n",
    "        joint_x_y[x_val][y_val] += 1\n",
    "        \n",
    "#   normalize to make distirbution  \n",
    "    joint_sum = sum(sum(list(c.values())) for c in list(joint_x_y.values()))\n",
    "    \n",
    "    for key1 in joint_x_y:\n",
    "        for key2 in joint_x_y[key1]:\n",
    "            joint_x_y[key1][key2] /= joint_sum\n",
    "            \n",
    "    return(joint_x_y)\n",
    "\n",
    "def mutual_inf(x, y):\n",
    "    \"\"\"\n",
    "    Calculates the mutual information I(x;y)\n",
    "    Assuming x,y are both [n x 1] dimensional\n",
    "    \"\"\"  \n",
    "#     Calculate marginal distributions\n",
    "    px = get_marginal(x)\n",
    "    py = get_marginal(y)\n",
    "    \n",
    "    \n",
    "    joint_x_y = get_joint(x,y)\n",
    "# calculate mutual information\n",
    "    mi = 0\n",
    "    \n",
    "    for n_x, x_un in enumerate(np.unique(x)):\n",
    "        pxi = px[n_x] # p(x)\n",
    "        \n",
    "        for n_y, y_un in enumerate(np.unique(y)):\n",
    "            pyi = py[n_y] # p(y)            \n",
    "            \n",
    "            joint_i = joint_x_y[x_un][y_un] # P(x,y)\n",
    "            \n",
    "            if ((pxi == 0) or (pyi == 0) or (joint_i ==0 )):\n",
    "                continue\n",
    "            else:\n",
    "                mi += joint_i * np.log2(joint_i/(pxi*pyi))\n",
    "                \n",
    "    return mi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/alsfilip/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:57: DeprecationWarning: Calling np.sum(generator) is deprecated, and in the future will give a different result. Use np.sum(np.fromiter(generator)) or the python sum builtin instead.\n"
     ]
    }
   ],
   "source": [
    "# Compute/extract complexity, predictive accuracy, % correct, RT, main effect, and interaction for each subject\n",
    "# Load all data - containes trial by trial responses/stimuli and also subject by subject interaction values and previous fits from Kool et al.\n",
    "mbmf_subs = pd.read_csv('./data/MBMF_subject_data.csv')\n",
    "subID = pd.unique(mbmf_subs['subject'])\n",
    "\n",
    "# Get subject fits - here subject name is by index rather than ID but matches the order of the subject IDs in the file above\n",
    "mbmf_fits = pd.read_csv('./data/w_fits.csv')\n",
    "\n",
    "subInfo = {}\n",
    "for i in np.arange(len(subID)):\n",
    "    # Sub dictionary for each subject\n",
    "    subInfo[subID[i]] = {}\n",
    "    \n",
    "    # Get subset of data with subject data\n",
    "    sdat = mbmf_subs.loc[mbmf_subs['subject']==subID[i]]\n",
    "    \n",
    "    # Get trial info needed to compute complexity\n",
    "    R1 = np.array(sdat['R1'].loc[sdat['R1'] >= 0]-1)\n",
    "    S2 = np.array(sdat['S2'].loc[sdat['R1'] >= 0]-1)\n",
    "    Rw = np.array(sdat['Rw'].loc[sdat['R1'] >= 0])\n",
    "    R1_star = np.array(sdat['R1_star'].loc[sdat['R1'] >= 0]-1)\n",
    "    \n",
    "    # Get first-step RT info\n",
    "    rt1 = np.array(sdat['rt1'].loc[sdat['R1'] >= 0])\n",
    "    \n",
    "    # Compute complexity for different feature sizes\n",
    "    # All past features\n",
    "    all_stack = np.stack((R1,S2,Rw,R1_star),axis=1)\n",
    "    all_features = np.sum(all_stack*(2**np.arange(4)),axis=1)\n",
    "    subInfo[subID[i]]['Ipast_all'] = mutual_inf(all_features[:-1],R1[1:])\n",
    "    \n",
    "    # Reduced features\n",
    "    # No R1\n",
    "    noR1_features = np.sum(all_stack[:,1:4]*(2**np.arange(3)),axis=1)\n",
    "    subInfo[subID[i]]['Ipast_noR1'] = mutual_inf(noR1_features[:-1],R1[1:])\n",
    "    \n",
    "    # No S2\n",
    "    noS2_features = np.sum(all_stack[:,[0,2,3]]*(2**np.arange(3)),axis=1)\n",
    "    subInfo[subID[i]]['Ipast_noS2'] = mutual_inf(noS2_features[:-1],R1[1:])\n",
    "    \n",
    "    # No Rw\n",
    "    noRw_features = np.sum(all_stack[:,[0,1,3]]*(2**np.arange(3)),axis=1)\n",
    "    subInfo[subID[i]]['Ipast_noRw'] = mutual_inf(noRw_features[:-1],R1[1:])\n",
    "    \n",
    "    # No R1_star\n",
    "    noR1_star_features = np.sum(all_stack[:,0:3]*(2**np.arange(3)),axis=1)\n",
    "    subInfo[subID[i]]['Ipast_noR1_star'] = mutual_inf(noR1_star_features[:-1],R1[1:])\n",
    "    \n",
    "    # Predictive info and proportion correct\n",
    "    # Ifuture\n",
    "    subInfo[subID[i]]['Ifuture'] = mutual_inf(R1,R1_star)\n",
    "    \n",
    "    # Proportion correct\n",
    "    subInfo[subID[i]]['PropCorr'] = np.mean(Rw)\n",
    "    \n",
    "    # Prop rich choices\n",
    "    subInfo[subID[i]]['PropRich'] = np.mean(R1==R1_star)\n",
    "    \n",
    "    # Extract main effect reported by Kool et al\n",
    "    subInfo[subID[i]]['ME'] = sdat['meffect'].iloc[0]\n",
    "    \n",
    "    # Extract interaction reported by Kool et al\n",
    "    subInfo[subID[i]]['Int'] = sdat['interaction'].iloc[0]\n",
    "    \n",
    "    # Get mean log response times on first response\n",
    "    subInfo[subID[i]]['Mean_logRT'] = np.mean(np.log(rt1[rt1>0]))\n",
    "\n",
    "    # Get fit strategy mixing coefficient for stochastic model fits\n",
    "    subInfo[subID[i]]['w_stoch'] = mbmf_fits['w_stoch'].iloc[i]\n",
    "    \n",
    "    # Get fit strategy mixing coefficient for deterministic model fits\n",
    "    subInfo[subID[i]]['w_det'] = mbmf_fits['w_det'].iloc[i]\n",
    "\n",
    "# Save individual subject info for use in the main analysis\n",
    "outfile = open('./data/sub_cx_behav.pkl','wb')\n",
    "pickle.dump(subInfo,outfile)\n",
    "outfile.close()"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}