PsyTrack_Learning_Examples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import re\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.optimize import minimize\n",
    "from scipy.special import expit\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import psytrack_learning as psy\n",
    "from psytrack_learning.getMAP import getMAP\n",
    "from psytrack_learning.helper.helperFunctions import update_hyper, hyper_to_list\n",
    "from psytrack_learning.helper.jacHessCheck import compHess, compHess_nolog\n",
    "from psytrack_learning.helper.invBlkTriDiag import getCredibleInterval\n",
    "from psytrack_learning.hyperparameter_optimization import evd_lossfun\n",
    "from psytrack_learning.learning_rules import RewardMax, PredictMax, REINFORCE, REINFORCE_base\n",
    "from psytrack_learning.simulate_learning import reward_max, predict_max, reinforce, reinforce_base \n",
    "from psytrack_learning.simulate_learning import simulate_learning\n",
    "\n",
    "\n",
    "# Set matplotlib defaults from making files editable and consistent in Illustrator\n",
    "colors = psy.COLORS\n",
    "zorder = psy.ZORDER\n",
    "plt.rcParams['figure.dpi'] = 140\n",
    "plt.rcParams['savefig.dpi'] = 300\n",
    "plt.rcParams['savefig.facecolor'] = (1,1,1,0)\n",
    "plt.rcParams['savefig.bbox'] = \"tight\"\n",
    "plt.rcParams['font.size'] = 10\n",
    "plt.rcParams['font.family'] = 'cmu serif'\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "plt.rcParams['pdf.fonttype'] = 42\n",
    "plt.rcParams['xtick.labelsize'] = 10\n",
    "plt.rcParams['ytick.labelsize'] = 10\n",
    "plt.rcParams['axes.labelsize'] = 12\n",
    "\n",
    "# Set save path for all figures\n",
    "spath = \"/Users/nicholasroy/ActiveLearningPapers/NeurIPS20/Figures/FigureComponents/\"  # UPDATE\n",
    "sim_colors = [\"#D81159\", \"#4357AD\", \"#EE8434\", \"#CC3399\", \"#409091\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download and pre-process IBL mouse data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1) Download the [IBL dataset](https://doi.org/10.6084/m9.figshare.11636748.v7) (version 7, uploaded Feb 7, 2020).\n",
    "\n",
    "2) Update the `ibl_data_path` variable below to where the `ibl-behavioral-data-Dec2019` directory exists on your computer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ibl_data_path = '/Users/nicholasroy/ibl-behavioral-data-Dec2019/'  # --- UPDATE ---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3) We will also need to install the [ONE Light](https://github.com/int-brain-lab/ibllib/tree/master/oneibl) Python library (from the IBL) with `pip install ibllib`. This allows us to build a table of all the subject and session data contained within the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from oneibl.onelight import ONE\n",
    "import pandas as pd\n",
    "\n",
    "current_cwd = os.getcwd()\n",
    "os.chdir(ibl_data_path)\n",
    "\n",
    "# Search all sessions that have these dataset types.\n",
    "required_vars = ['_ibl_trials.choice', '_ibl_trials.contrastLeft',\n",
    "                 '_ibl_trials.contrastRight','_ibl_trials.feedbackType']\n",
    "one = ONE()\n",
    "eids = one.search(required_vars)\n",
    "\n",
    "mouseData = pd.DataFrame()\n",
    "for eid in eids:\n",
    "    lab, _, subject, date, session = eid.split(\"/\")    \n",
    "    sess_vars = {\n",
    "        \"eid\": eid,\n",
    "        \"lab\": lab,\n",
    "        \"subject\": subject,\n",
    "        \"date\": date,\n",
    "        \"session\": session,\n",
    "    }\n",
    "    mouseData = mouseData.append(sess_vars, sort=True, ignore_index=True)\n",
    "\n",
    "os.chdir(current_cwd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4) Next, we will use the table of session data to process the raw trial data below into a single CSV file, `ibl_processed.csv`, saved locally.\n",
    "\n",
    "There are several known anomalies in the raw data:\n",
    " - CSHL_002 codes left contrasts as negative right contrasts on 81 trials (these trials are corrected)\n",
    " - ZM_1084 has `feedbackType` of 0 for 3 trials (these trials are omitted)\n",
    " - DY_009, DY_010, DY_011 each have <5000 trials total (no adjustment)\n",
    " - ZM_1367, ZM_1369, ZM_1371, ZM_1372, and ZM_1743 are shown non-standard contrast values of 0.04 and 0.08 (no adjustment)\n",
    "\n",
    "_2 min_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_vars = [\"contrastLeft\", \"contrastRight\", \"choice\", \"feedbackType\", \"probabilityLeft\"]\n",
    "df = pd.DataFrame()\n",
    "\n",
    "all_mice = []\n",
    "for j, s in enumerate(mouseData[\"subject\"].unique()):\n",
    "    print(\"\\rProcessing \" + str(j+1) + \" of \" + str(len(mouseData[\"subject\"].unique())), end=\"\")\n",
    "    mouse = mouseData[mouseData[\"subject\"]==s].sort_values(['date', 'session']).reset_index()\n",
    "    for i, row in mouse.iterrows():\n",
    "        myVars = {}\n",
    "        for v in all_vars:\n",
    "            filename = \"_ibl_trials.\" + v + \".npy\"\n",
    "            var_file = os.path.join(ibl_data_path, row.eid, \"alf\", filename)\n",
    "            myVars[v] = list(np.load(var_file).flatten())\n",
    "\n",
    "        num_trials = len(myVars[v])\n",
    "        myVars['lab'] = [row.lab]*num_trials\n",
    "        myVars['subject'] = [row.subject]*num_trials\n",
    "        myVars['date'] = [row.date]*num_trials\n",
    "        myVars['session'] = [row.session]*num_trials\n",
    "\n",
    "        all_mice += [pd.DataFrame(myVars, columns=myVars.keys())]\n",
    "        \n",
    "df = pd.concat(all_mice, ignore_index=True)\n",
    "\n",
    "df = df[df['choice'] != 0]        # dump mistrials\n",
    "df = df[df['feedbackType'] != 0]  # 3 anomalous trials from ZM_1084, omit\n",
    "df.loc[np.isnan(df['contrastLeft']), \"contrastLeft\"] = 0\n",
    "df.loc[np.isnan(df['contrastRight']), \"contrastRight\"] = 0\n",
    "df.loc[df[\"contrastRight\"] < 0, \"contrastLeft\"] = np.abs(df.loc[df[\"contrastRight\"] < 0, \"contrastRight\"])\n",
    "df.loc[df[\"contrastRight\"] < 0, \"contrastRight\"] = 0  # 81 anomalous trials in CSHL_002, correct\n",
    "df[\"answer\"] = df[\"feedbackType\"] * df[\"choice\"]      # new column to indicate correct answer\n",
    "df.loc[df[\"answer\"]==1, \"answer\"] = 0\n",
    "df.loc[df[\"answer\"]==-1, \"answer\"] = 1\n",
    "df.loc[df[\"feedbackType\"]==-1, \"feedbackType\"] = 0\n",
    "df.loc[df[\"choice\"]==1, \"choice\"] = 0\n",
    "df.loc[df[\"choice\"]==-1, \"choice\"] = 1\n",
    "df.to_csv(spath+\"ibl_processed.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5) Next we do a few sanity checks on our data, to make sure everything processed correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"contrastLeft: \", np.unique(df['contrastLeft']))   # [0, 0.0625, 0.125, 0.25, 0.5, 1.0] and [0.04, 0.08]\n",
    "print(\"contrastRight: \", np.unique(df['contrastRight'])) # [0, 0.0625, 0.125, 0.25, 0.5, 1.0] and [0.04, 0.08]\n",
    "print(\"choice: \", np.unique(df['choice']))               # [0, 1]\n",
    "print(\"feedbackType: \", np.unique(df['feedbackType']))   # [0, 1]\n",
    "print(\"answer: \", np.unique(df['answer']))               # [0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6) Finally, we define a function `getMouse()` that extracts the data for a single mouse from our CSV file, and returns it in a PsyTrack compatible dictionary. We will use this function to access IBL mouse data in the figures below. Note the keyword argument and default value $p=5$ which controls the strength of the $\\tanh$ transformation on the contrast values. See Figure S3 and the STAR Methods for more details.\n",
    "\n",
    "**Note:** Once steps 1-6 have been run once, only step 6 will need to be run on subsequent uses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mouse_data_path = spath + \"ibl_processed.csv\"   # --- UPDATE if necessary ---\n",
    "MOUSE_DF = pd.read_csv(mouse_data_path)\n",
    "\n",
    "def getMouse(subject, p=5):\n",
    "    df = MOUSE_DF[MOUSE_DF['subject']==subject]   # Restrict data to the subject specified\n",
    "    \n",
    "    cL = np.tanh(p*df['contrastLeft'])/np.tanh(p)   # tanh transformation of left contrasts\n",
    "    cR = np.tanh(p*df['contrastRight'])/np.tanh(p)  # tanh transformation of right contrasts\n",
    "    cBoth = cR - cL\n",
    "    inputs = dict(cL = np.array(cL)[:, None], cR = np.array(cR)[:, None], cBoth = np.array(cBoth)[:, None])\n",
    "\n",
    "    dat = dict(\n",
    "        subject=subject,\n",
    "        lab=np.unique(df[\"lab\"])[0],\n",
    "        contrastLeft=np.array(df['contrastLeft']),\n",
    "        contrastRight=np.array(df['contrastRight']),\n",
    "        date=np.array(df['date']),\n",
    "        dayLength=np.array(df.groupby(['date','session']).size()),\n",
    "        correct=np.array(df['feedbackType']),\n",
    "        answer=np.array(df['answer']),\n",
    "        probL=np.array(df['probabilityLeft']),\n",
    "        inputs = inputs,\n",
    "        y = np.array(df['choice'])\n",
    "    )\n",
    "    \n",
    "    return dat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Athena Rat Data\n",
    "\n",
    "You can download the processed IBL data as a CSV file, `rat_behavior.csv`, from [here](https://drive.google.com/drive/u/0/folders/1eOfqq3ijdWLxvhumFioJ3LgNmLZdrBtq).\n",
    "Next, update the `rat_data_path` below to where the CSV file is locally saved.\n",
    "You can the use the `getRat` function below to extract the data for a single subject into a PsyTrack compatible dict.\n",
    "\n",
    "Note that mistrials have been omitted, as have sessions from before the reward rule was enforced (i.e. shaping sessions). The `getRat` function also has two optional parameters: `first` which will return a data set with only the first `first` trials (20,000 seems reasonable for our purposes); `cutoff` excludes sessions with fewer than `cutoff` valid trials (currently set to 50)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rat_data_path = \"~/rat_behavior.csv\"   # --- UPDATE ---\n",
    "RAT_DF = pd.read_csv(rat_data_path)\n",
    "\n",
    "RAT_DF = RAT_DF[RAT_DF[\"training_stage\"] > 2]  # Remove trials from early training\n",
    "RAT_DF = RAT_DF[~np.isnan(RAT_DF[\"choice\"])]   # Remove mistrials\n",
    "\n",
    "def getRat(subject, first=20000, cutoff=50):\n",
    "\n",
    "    df = RAT_DF[RAT_DF['subject_id']==subject]  # restrict dataset to single subject\n",
    "    df = df[:first]  # restrict to \"first\" trials of data\n",
    "    # remove sessions with fewer than \"cutoff\" valid trials\n",
    "    df = df.groupby('session').filter(lambda x: len(x) >= cutoff)   \n",
    "\n",
    "    # Normalize the stimuli to standard normal\n",
    "    s_a = (df[\"s_a\"] - np.mean(df[\"s_a\"]))/np.std(df[\"s_a\"])\n",
    "    s_b = (df[\"s_b\"] - np.mean(df[\"s_b\"]))/np.std(df[\"s_b\"])\n",
    "    \n",
    "    # Determine which trials do not have a valid previous trial (mistrial or session boundary)\n",
    "    t = np.array(df[\"trial\"])\n",
    "    prior = ((t[1:] - t[:-1]) == 1).astype(int)\n",
    "    prior = np.hstack(([0], prior))\n",
    "\n",
    "    # Calculate previous average tone value\n",
    "    s_avg = (df[\"s_a\"][:-1] + df[\"s_b\"][:-1])/2\n",
    "    s_avg = (s_avg - np.mean(s_avg))/np.std(s_avg)\n",
    "    s_avg = np.hstack(([0], s_avg))\n",
    "    s_avg = s_avg * prior  # for trials without a valid previous trial, set to 0\n",
    "\n",
    "    # Calculate previous correct answer\n",
    "    h = (df[\"correct_side\"][:-1] * 2 - 1).astype(int)   # map from (0,1) to (-1,1)\n",
    "    h = np.hstack(([0], h))\n",
    "    h = h * prior  # for trials without a valid previous trial, set to 0\n",
    "    \n",
    "    # Calculate previous choice\n",
    "    c = (df[\"choice\"][:-1] * 2 - 1).astype(int)   # map from (0,1) to (-1,1)\n",
    "    c = np.hstack(([0], c))\n",
    "    c = c * prior  # for trials without a valid previous trial, set to 0\n",
    "    \n",
    "    inputs = dict(s_a = np.array(s_a)[:, None],\n",
    "                  s_b = np.array(s_b)[:, None],\n",
    "                  s_avg = np.array(s_avg)[:, None],\n",
    "                  h = np.array(h)[:, None],\n",
    "                  c = np.array(c)[:, None])\n",
    "\n",
    "    dat = dict(\n",
    "        subject = subject,\n",
    "        inputs = inputs,\n",
    "        s_a = np.array(df['s_a']),\n",
    "        s_b = np.array(df['s_b']),\n",
    "        correct = np.array(df['hit']),\n",
    "        answer = np.array(df['correct_side']),\n",
    "        y = np.array(df['choice']),\n",
    "        dayLength=np.array(df.groupby(['session']).size()),\n",
    "    )\n",
    "    return dat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 2 | REINFORCE Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### For Athena Task + example rat\n",
    "N = 5000\n",
    "seed = 42\n",
    "mouse_name = \"CSHL_003\"\n",
    "np.random.seed(seed)\n",
    "\n",
    "orig_dat = getMouse(mouse_name, 5)\n",
    "trim_dat = psy.trim(orig_dat, END=N)\n",
    "\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "K = np.sum([i for i in weights.values()])\n",
    "X = psy.read_input(trim_dat, weights)\n",
    "answer = trim_dat[\"answer\"]\n",
    "true_sigma  = 2**np.array([-4.0, -5.0])\n",
    "true_alpha  = 2**np.array([-6.0, -8.0])\n",
    "sim_learning_rule = reinforce\n",
    "rec_learning_rule = REINFORCE\n",
    "\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1,\n",
    "                                       alpha=true_alpha, learning_rule=sim_learning_rule)\n",
    "\n",
    "dat = {'inputs': trim_dat['inputs'].copy(), 'y': y, 'answer': answer, \"correct\": r}\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule, \"rec_learning_rule\": rec_learning_rule,\n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "\n",
    "fig = psy.plot_weights(W, weights)\n",
    "# np.savez_compressed(spath+'fig2a_data_gen.npz', gen_dat=gen_dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute\n",
    "hyper_guess = {\n",
    "    'alpha': [2**-6] * K,\n",
    "    'sigma': [2**-4] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for both sigma and alpha simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": dat, \"K\": K, \"learning_rule\": rec_learning_rule,\n",
    "        \"hyper\": hyper_guess, \"weights\": weights, \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-8, \"showOpt\": True,\n",
    "       }\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA')\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "wMode_sim2, Hess, logEvd, other = getMAP(dat, opt_hyper, weights, W0=None,\n",
    "                                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "wMode_sim2 = wMode_sim2.reshape((K, -1), order=\"C\")\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, \"W_std\": W_std, \"wMode\": wMode_sim2}\n",
    "\n",
    "# np.savez_compressed(spath+'fig2a_data_rec.npz', rec_dat=rec_dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reload data\n",
    "# rec_dat = np.load(spath+'fig2a_data_rec.npz', allow_pickle=True)['rec_dat'].item()\n",
    "# gen_dat = np.load(spath+'fig2a_data_gen.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "# Plotting\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['h']]\n",
    "\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(gen_dat['W'][i], c=c, lw=0.5, zorder=2*i)\n",
    "    plt.plot(rec_dat['wMode'][i], c=c, lw=1, linestyle='--', alpha=0.5, zorder=2*i+1)\n",
    "    plt.fill_between(np.arange(len(rec_dat['wMode'][i])),\n",
    "                     rec_dat['wMode'][i] - 2 * rec_dat['W_std'][i],\n",
    "                     rec_dat['wMode'][i] + 2 * rec_dat['W_std'][i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i+1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(-6,7,2))\n",
    "plt.xlim(0,5000); plt.ylim(-6.2,6.2)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weights\")\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "# plt.savefig(spath + \"Fig2a.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recover error bars for hyperparameters\n",
    "# rec_dat = np.load(spath+'fig2a_data_rec.npz', allow_pickle=True)['rec_dat'].item()\n",
    "hess_args = rec_dat['args'].copy()\n",
    "hess_args[\"wMode\"] = rec_dat['wMode'].flatten()\n",
    "hess_args[\"learning_rule\"] = hess_args[\"learning_rule\"]\n",
    "H, g = compHess_nolog(evd_lossfun, rec_dat['res'].x, 5e-2, {\"keywords\": hess_args})\n",
    "hyp_std = np.sqrt(np.diag(np.linalg.inv(H)))\n",
    "\n",
    "rec_dat['hyp_std'] = hyp_std\n",
    "# np.savez_compressed(spath+'fig2a_data_rec.npz', rec_dat=rec_dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reload data\n",
    "# rec_dat = np.load(spath+'fig2a_data_rec.npz', allow_pickle=True)['rec_dat'].item()\n",
    "# gen_dat = np.load(spath+'fig2a_data_gen.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "# Plot recovered hyperparameters\n",
    "plt.figure(figsize=(1.5,1.25))\n",
    "sim_colors = [colors['bias'], colors['h']]\n",
    "num_std = 1.00\n",
    "K = gen_dat['K']\n",
    "\n",
    "true_sigma = gen_dat['true_sigma']\n",
    "avg_sigma = np.log2(rec_dat['opt_hyper']['sigma'])\n",
    "err_sigma = rec_dat['hyp_std'][:K]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    yerr = [[-(np.log2(2**avg_sigma[i] - err_sigma[i]*num_std) - avg_sigma[i])],\n",
    "            [np.log2(2**avg_sigma[i] + err_sigma[i]*num_std) - avg_sigma[i]]]\n",
    "    if np.isnan(yerr[0][0]):\n",
    "        yerr[0][0] = 100\n",
    "    plt.plot([i-0.3, i+0.3], [np.log2(true_sigma[i])]*2, color=c, linestyle=\"-\", lw=1.2, zorder=1, alpha=0.8)\n",
    "    plt.errorbar([i], avg_sigma[i], yerr=yerr, c=c, lw=1, marker='o', markersize=4)\n",
    "    \n",
    "true_alpha = gen_dat['true_alpha']\n",
    "avg_alpha = np.log2(rec_dat['opt_hyper']['alpha'])\n",
    "err_alpha = rec_dat['hyp_std'][K:2*K]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    yerr = [[-(np.log2(2**avg_alpha[i] - err_alpha[i]*num_std) - avg_alpha[i])],\n",
    "            [np.log2(2**avg_alpha[i] + err_alpha[i]*num_std) - avg_alpha[i]]]\n",
    "    if np.isnan(yerr[0][0]):\n",
    "        yerr[0][0] = 100\n",
    "    plt.plot([i+1.7, i+2.3], [np.log2(true_alpha[i])]*2, color=c, linestyle=\"-\", lw=1.2, zorder=1, alpha=0.8)\n",
    "    plt.errorbar([i+2], avg_alpha[i], yerr=yerr, c=c, lw=1, marker='s', markersize=4)\n",
    "    \n",
    "plt.ylim(-8.75,-3.5)\n",
    "plt.yticks(np.arange(-8,-3))\n",
    "plt.xticks([0,1,2,3])\n",
    "plt.gca().set_xticklabels([r\"$\\sigma_1$\", r\"$\\sigma_2$\",\n",
    "                           r\"$\\alpha_1$\", r\"$\\alpha_2$\"])\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "# plt.savefig(spath + \"Fig2b.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rec_dat = np.load(spath+'fig2a_data_rec.npz', allow_pickle=True)['rec_dat'].item()\n",
    "# gen_dat = np.load(spath+'fig2a_data_gen.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "wMode, Hess, logEvd, other = getMAP(gen_dat['dat'], rec_dat['opt_hyper'], gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=gen_dat['rec_learning_rule'], showOpt=0, tol=1e-12)\n",
    "wMode = wMode.reshape((gen_dat['K'], -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((gen_dat['K'], -1), order=\"C\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise = np.cumsum(E_rw, axis=1)\n",
    "learning = np.cumsum(v_rw, axis=1)\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning[i], c=c, lw=1.5, linestyle='-', alpha=0.85, zorder=10+2*i)\n",
    "    plt.plot(noise[i] - noise[i, 0], c=c, lw=1.5, linestyle='--', alpha=0.85, zorder=10+2*i+1)\n",
    "    plt.plot(learning_sim[i], c=\"gray\", lw=0.5, linestyle='-', alpha=0.25, zorder=2)\n",
    "    plt.plot(noise_sim[i] - noise_sim[i, 0], c=\"gray\", lw=0.5, linestyle='-', alpha=0.25, zorder=2)\n",
    "    plt.fill_between(np.arange(len(noise[i])),\n",
    "                     learning[i], learning_sim[i],\n",
    "                     facecolor=\"lightgray\", alpha=0.5, zorder=1)\n",
    "    plt.fill_between(np.arange(len(noise[i])),\n",
    "                     noise[i] - noise[i, 0], noise_sim[i] - noise_sim[i, 0],\n",
    "                     facecolor=\"lightgray\", alpha=0.5, zorder=1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "plt.gca().set_xticklabels([0,1000,2000,3000,4000,5000])\n",
    "plt.yticks(np.arange(-4,5,2))\n",
    "plt.xlim(0,5000); plt.ylim(-4.2,4.2)\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "# plt.savefig(spath + \"Fig2c.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 3 | Example IBL Mouse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 3 runs three models, $[RF_0, RF_1, RF_K, RF_\\beta]$ on an example mouse (mouse `CSHL_003` from the IBL data). We show fitting for the $RF_K$ model here — the other models can be fit in an analogous way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect data from manually determined training period\n",
    "outData = getMouse('CSHL_003', 5)\n",
    "dat = psy.trim(outData, START=0, END=6000)\n",
    "\n",
    "# Compute\n",
    "weights = {'bias' : 1, 'cL' : 0, 'cR' : 0, 'cBoth': 1}\n",
    "K = np.sum([weights[i] for i in weights.keys()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute\n",
    "hyper_guess = {\n",
    "    'alpha': [2**-6] * K,\n",
    "    'sigma': [2**-4] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for both sigma and alpha simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": dat, \"K\": K, \"learning_rule\": REINFORCE,\n",
    "        \"hyper\": hyper_guess, \"weights\": weights, \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA')\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, weights, W0=None,\n",
    "                                    learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, \"W_std\": W_std, \"wMode\": wMode}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(3.25,1.25))\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(wMode[i], c=c, lw=1, linestyle='-', alpha=0.85, zorder=2*i+1)\n",
    "    plt.fill_between(np.arange(len(wMode[i])), wMode[i] - 2 * W_std[i], wMode[i] + 2 * W_std[i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i)\n",
    "    \n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "plt.xticks(1000*np.arange(0,7))\n",
    "plt.yticks(np.arange(-2,3,2))\n",
    "plt.xlim(0,6000)\n",
    "plt.ylim(-3.5,3.5)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((gen_dat['K'], -1), order=\"C\")\n",
    "\n",
    "noise = np.cumsum(E_rw, axis=1)\n",
    "learning = np.cumsum(v_rw, axis=1)\n",
    "\n",
    "plt.figure(figsize=(3.25,1.25))\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning[i], c=c, lw=0.75, linestyle='-', alpha=0.85, zorder=2*i)\n",
    "    plt.fill_between(np.arange(len(noise[i])), learning[i], learning[i] + noise[i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i)\n",
    "    plt.plot(learning[i] + noise[i], c=c, lw=0.25, linestyle='-', alpha=0.5, zorder=2*i+1)\n",
    "    \n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "plt.xticks(1000*np.arange(0,7))\n",
    "plt.yticks(np.arange(-2,3,2))\n",
    "plt.xlim(0,6000)\n",
    "plt.ylim(-3.5,3.5)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 4 | IBL Population Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 4 is a population level analysis that cannot be run locally in a reasonable amount of time. See the `cluster_scripts` directory for some files for running fitting in parallel on a cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 5 | Reward Landscape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect data from manually determined training period\n",
    "outData = getMouse('CSHL_003', 5)\n",
    "dat = psy.trim(outData, START=0, END=6000)\n",
    "\n",
    "# Compute\n",
    "weights = {'bias' : 1, 'cL' : 0, 'cR' : 0, 'cBoth': 1}\n",
    "K = np.sum([weights[i] for i in weights.keys()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute\n",
    "hyper_guess = {\n",
    "    'alpha': [2**-6] * K,\n",
    "    'sigma': [2**-4] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for both sigma and alpha simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": dat, \"K\": K, \"learning_rule\": REINFORCE,\n",
    "        \"hyper\": hyper_guess, \"weights\": weights, \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA')\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, weights, W0=None,\n",
    "                                    learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, \"W_std\": W_std, \"wMode\": wMode}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prob_right(w_bias, w_stim, stim):\n",
    "    return expit(w_bias + w_stim*stim)\n",
    "\n",
    "def calculate_expected_reward(w_bias, w_stim):\n",
    "    # Contrast vals considered\n",
    "    contrast_vals = [-1.,  -0.98670389,  0.98670389, 1. ]\n",
    "    # Assume each contrast occurs with equal frequency\n",
    "    expected_reward = 0\n",
    "    for contrast in contrast_vals:\n",
    "        if contrast < 0:\n",
    "            expected_reward += (1-get_prob_right(w_bias, w_stim, contrast))/len(contrast_vals)\n",
    "        else:\n",
    "            expected_reward += get_prob_right(w_bias, w_stim, contrast)/len(contrast_vals)\n",
    "    return expected_reward\n",
    "\n",
    "\n",
    "def simulate_data(mouse, W0, alpha,  seed):\n",
    "    full_mouse_data = getMouse(mouse, 5)\n",
    "    dat = psy.trim(full_mouse_data, END=6000)\n",
    "    # Compute\n",
    "    weights = {'bias': 1, 'cL': 0, 'cR': 0, 'cBoth': 1}\n",
    "    K = np.sum([weights[i] for i in weights.keys()])\n",
    "\n",
    "    X = psy.read_input(dat, weights)\n",
    "    answer = dat[\"answer\"]\n",
    "    sigma = [0, 0]\n",
    "    W, y, r, _ = simulate_learning(X=X, answer=answer, sigma=sigma, alpha=alpha,\n",
    "                                   learning_rule=reinforce, W0 = W0, seed = seed)\n",
    "    return W, y, r\n",
    "\n",
    "\n",
    "N_bias = 200\n",
    "N_contrast = 200\n",
    "mouse = \"CSHL_003\"\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(3,1.5))\n",
    "\n",
    "w_bias = np.linspace(-2.2, 2.2, N_bias)\n",
    "w_contrast = np.linspace(-0.3,1.7, N_contrast)\n",
    "xx, yy = np.meshgrid(w_bias, w_contrast)\n",
    "z = np.zeros((N_bias, N_contrast))\n",
    "for i, w_b in enumerate(w_bias):\n",
    "    for j, w_c in enumerate(w_contrast):\n",
    "        z[i, j] = calculate_expected_reward(w_b, w_c)\n",
    "CS = plt.contourf(w_contrast,  w_bias,  z, levels = np.arange(0,11)/10,\n",
    "                  cmap = 'Greys', alpha = 0.85, vmin = 0.2, vmax = 1)\n",
    "\n",
    "N = len(wMode[0])\n",
    "inc = np.arange(0,N+1000,1000)\n",
    "inc[-1] = -1\n",
    "jump = 3\n",
    "\n",
    "plt.plot(wMode[1,::jump] ,wMode[0,::jump], lw = 1, color=\"#EF709D\")\n",
    "plt.scatter(wMode[1, inc], wMode[0, inc],\n",
    "            color = \"#EF709D\", s=15, edgecolor='w', marker=\"o\", zorder=10, lw=0.5)\n",
    "\n",
    "\n",
    "# Simulate weight trajectory according to REINFORCE with single alpha retrieved from data\n",
    "single_alpha = 0.0014743978688641268   # -- value of alpha when fitting above mouse with RF_1 \n",
    "alpha = [single_alpha, single_alpha]\n",
    "W_reinforce, _, _ = simulate_data(mouse, wMode[:,0], alpha, seed = 1)\n",
    "plt.plot(W_reinforce[1,::jump] ,W_reinforce[0,::jump], lw = 1.5, color=\"#4BB3FD\")\n",
    "plt.scatter(W_reinforce[1, inc], W_reinforce[0, inc],\n",
    "            color = \"#4BB3FD\", s=15, edgecolor='w', marker=\"o\", zorder=10, lw=0.5)\n",
    "\n",
    "\n",
    "# Simulate weight trajectory according to REINFORCE with multiple alpha retrieved from data\n",
    "multiple_alpha = opt_hyper['alpha']\n",
    "W_reinforce_mult_alpha, _, _ = simulate_data(mouse, wMode[:, 0], multiple_alpha, seed=1)\n",
    "\n",
    "plt.plot(W_reinforce_mult_alpha[1,::jump] ,W_reinforce_mult_alpha[0,::jump], lw = 1, color=colors['c'])\n",
    "plt.scatter(W_reinforce_mult_alpha[1, inc], W_reinforce_mult_alpha[0, inc],\n",
    "            color = colors['c'], s=15, edgecolor='w', marker=\"o\", zorder=10, lw=0.5)\n",
    "\n",
    "plt.xticks(np.arange(0,2))\n",
    "plt.yticks(np.arange(-2,3))\n",
    "plt.xlim((-0.3, 1.7))\n",
    "plt.ylim((-2, 2))\n",
    "plt.subplots_adjust(0,0,1,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Figure 6 | Example Athena Rat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 3 runs two models, $[RF_0, RF_\\beta]$ on an example rat (rat `W082` from the Akrami data). We show fitting for the $RF_1$ model here — the other models can be fit in an analogous way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orig_dat = getRat(\"W082\")\n",
    "dat = psy.trim(orig_dat, END=6000)\n",
    "\n",
    "weights = {\"bias\": 1, \"s_a\": 1, \"s_b\": 1}\n",
    "K = np.sum([i for i in weights.values()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute\n",
    "hyper_guess = {\n",
    "    'alpha': 2**-6,\n",
    "    'sigma': [2**-4] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for both sigma and alpha simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": dat, \"K\": K, \"learning_rule\": REINFORCE,\n",
    "        \"hyper\": hyper_guess, \"weights\": weights, \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA')\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, weights, W0=None,\n",
    "                                    learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, \"W_std\": W_std, \"wMode\": wMode}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_colors = [colors['bias'], colors['s_a'], colors['s_b']]\n",
    "\n",
    "plt.figure(figsize=(3.25,1.25))\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(wMode[i], c=c, lw=1, linestyle='-', alpha=0.85, zorder=2*i+1)\n",
    "    plt.fill_between(np.arange(len(wMode[i])), wMode[i] - 2 * W_std[i], wMode[i] + 2 * W_std[i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i)\n",
    "    \n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "plt.xticks(1000*np.arange(0,11))\n",
    "plt.yticks(np.arange(-2,3,1))\n",
    "plt.xlim(0,6000)\n",
    "plt.ylim(-2.25,2.25)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "noise = np.cumsum(E_rw, axis=1)\n",
    "learning = np.cumsum(v_rw, axis=1)\n",
    "\n",
    "plt.figure(figsize=(3.25,1.25))\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning[i], c=c, lw=0.75, linestyle='-', alpha=0.85, zorder=2*i)\n",
    "    plt.fill_between(np.arange(len(noise[i])), learning[i], learning[i] + noise[i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i)\n",
    "    plt.plot(learning[i] + noise[i], c=c, lw=0.25, linestyle='-', alpha=0.5, zorder=2*i+1)\n",
    "    \n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "plt.xticks(1000*np.arange(0,11))\n",
    "plt.yticks(np.arange(-2,3,1))\n",
    "plt.xlim(0,6000)\n",
    "plt.ylim(-2.25,2.25)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Supplement"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Figure S4\n",
    "\n",
    "Plot recovery of weights and hyperparameters of R+B model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "### Set up\n",
    "N = 6000  # would like to work with ~6000 trials, should try\n",
    "\n",
    "seed = 22  # 20\n",
    "print(\"Seed:\", seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Simulate sequence of contrasts directly (instead of grabbing from real mouse)\n",
    "p=5\n",
    "contrasts = np.array([0.0625, 0.125, 0.25, 0.5, 1.0])  # full set of contrasts\n",
    "tanh_contrasts = np.tanh(p*contrasts)/np.tanh(p)  # tanh transform\n",
    "all_contrasts = np.hstack((-tanh_contrasts[::-1], [0], tanh_contrasts))  # left contrasts are neg\n",
    "\n",
    "contrast_seq = np.random.choice(all_contrasts, size=N) # generate random sequence\n",
    "answer = np.sign(contrast_seq).astype(int)  # infer answers for each trial\n",
    "answer[answer==0] = np.sign(np.random.randn(np.sum(answer==0)))  # randomly select answer for 0 contrast\n",
    "answer[answer<0] = 0\n",
    "\n",
    "# Build data set from simulated data\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "dat = {'inputs': {'cBoth': contrast_seq.reshape(-1,1)}, 'answer': answer, 'y': np.zeros(N)}\n",
    "X = psy.read_input(dat, weights)\n",
    "K = np.sum([i for i in weights.values()])\n",
    "\n",
    "# Set true parameters of weight trajectory simulation\n",
    "true_sigma  = 2**np.array([-5.0, -6.0])  # -5, -6\n",
    "true_alpha  = 2**np.array([-6.0, -7.0])  # -6, -7\n",
    "true_base  = np.array([-0.15, 0.1])  # -0.15, 0.1, careful about which parametrization is being used\n",
    "true_hyper = {\n",
    "    'adder': true_base,\n",
    "    'alpha': true_alpha,\n",
    "    'sigma': true_sigma,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "sim_learning_rule = reinforce_base  # learning rule used in generation of weights\n",
    "rec_learning_rule = REINFORCE_base  # learning rule used for recovery of weights\n",
    "\n",
    "# Simulate the weight trajectories, also returning choices and rewards\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1, seed=seed,\n",
    "                                       alpha=true_alpha, base=true_base/true_alpha, # NOTE\n",
    "                                       learning_rule=sim_learning_rule)\n",
    "\n",
    "dat.update({\"y\": y, \"correct\": r})\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, 'true_base': true_base,\n",
    "           \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule.__name__, \"rec_learning_rule\": rec_learning_rule.__name__,\n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "np.savez_compressed(spath+'FigN1_data2.npz', gen_dat=gen_dat)\n",
    "\n",
    "\n",
    "# Display simulated weights, broken up into learning/nosie components\n",
    "fig = psy.plot_weights(W, weights)\n",
    "\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning_sim[i], c=c, lw=1, linestyle='-', alpha=0.75, zorder=1)\n",
    "    plt.plot(noise_sim[i], c=c, lw=0.5, linestyle='--', alpha=0.5, zorder=2)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "\n",
    "\n",
    "# Calculate evdience under true hyper\n",
    "wMode_true, _, logEvd, _ = getMAP(dat, true_hyper, weights, W0=None,\n",
    "                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "print(\"True Evidence:\", logEvd)\n",
    "fig = psy.plot_weights(wMode_true.reshape(K,-1), weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# gen_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-4] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for sigma, alpha, and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha', 'adder']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": gen_dat['rec_learning_rule'],\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", logEvd)\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[gen_dat['rec_learning_rule']],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "# np.savez_compressed(spath+'FigN1_data2.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "H, g = compHess_nolog(evd_lossfun, res.x, 5e-2, {\"keywords\": args})  # try 5e-2 if you get nans\n",
    "hyp_std = np.sqrt(np.diag(np.linalg.inv(H)))\n",
    "\n",
    "rec_dat.update({'hyp_std': hyp_std})\n",
    "# np.savez_compressed(spath+'FigN1_data2.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n",
    "\n",
    "print(\"Rec: \", res.x)\n",
    "print(\"Std: \", hyp_std)\n",
    "\n",
    "print(\"True:\", np.log2(true_sigma), np.log2(true_alpha))\n",
    "print(np.log2(2**res.x[:4] + hyp_std[:4]*1.96))\n",
    "print(np.log2(2**res.x[:4] - hyp_std[:4]*1.96))\n",
    "\n",
    "print(\"True:\", true_base)\n",
    "print(res.x[4:] + hyp_std[4:]*1.96)\n",
    "print(res.x[4:] - hyp_std[4:]*1.96)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "print(\"Rec: \", res.x)\n",
    "print(\"Std: \", hyp_std)\n",
    "\n",
    "print(\"True:\", np.log2(true_sigma), np.log2(true_alpha))\n",
    "print(np.log2(2**res.x[:4] + hyp_std[:4]))\n",
    "print(np.log2(2**res.x[:4] - hyp_std[:4]))\n",
    "\n",
    "print(\"True:\", true_base)\n",
    "print(res.x[4:] + hyp_std[4:])\n",
    "print(res.x[4:] - hyp_std[4:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "# rec_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['rec_dat'].item()\n",
    "# gen_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "# Plot recovered hyperparameters\n",
    "plt.figure(figsize=(3.75,1.4))\n",
    "sim_colors = [colors['bias'], colors['h']]\n",
    "num_std = 1.00\n",
    "K = gen_dat['K']\n",
    "\n",
    "\n",
    "true_sigma = gen_dat['true_sigma']\n",
    "avg_sigma = np.log2(rec_dat['opt_hyper']['sigma'])\n",
    "err_sigma = rec_dat['hyp_std'][:K]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    yerr = [[-(np.log2(2**avg_sigma[i] - err_sigma[i]*num_std) - avg_sigma[i])],\n",
    "            [np.log2(2**avg_sigma[i] + err_sigma[i]*num_std) - avg_sigma[i]]]\n",
    "    if np.isnan(yerr[0][0]):\n",
    "        yerr[0][0] = 100\n",
    "    plt.plot([i-0.3, i+0.3], [np.log2(true_sigma[i])]*2, color=\"black\", linestyle=\"-\", lw=1.2, zorder=1, alpha=0.4)\n",
    "    plt.errorbar([i], avg_sigma[i], yerr=yerr, c=c, lw=1, marker='o', markersize=4)\n",
    "    \n",
    "true_alpha = gen_dat['true_alpha']\n",
    "avg_alpha = np.log2(rec_dat['opt_hyper']['alpha'])\n",
    "err_alpha = rec_dat['hyp_std'][K:2*K]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    yerr = [[-(np.log2(2**avg_alpha[i] - err_alpha[i]*num_std) - avg_alpha[i])],\n",
    "            [np.log2(2**avg_alpha[i] + err_alpha[i]*num_std) - avg_alpha[i]]]\n",
    "    if np.isnan(yerr[0][0]):\n",
    "        yerr[0][0] = 100\n",
    "    plt.plot([i+1.7, i+2.3], [np.log2(true_alpha[i])]*2, color=\"black\", linestyle=\"-\", lw=1.2, zorder=1, alpha=0.4)\n",
    "    plt.errorbar([i+2], avg_alpha[i], yerr=yerr, c=c, lw=1, marker='s', markersize=4)\n",
    "    \n",
    "plt.ylim(-8.5,-3.5)\n",
    "plt.yticks(np.arange(-8,-3))\n",
    "# plt.axvline(4, linestyle ='-', color=\"gray\", lw=0.5)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "\n",
    "ax2 = plt.twinx()\n",
    "plt.sca(ax2)\n",
    "true_base = gen_dat['true_base']\n",
    "avg_base = rec_dat['opt_hyper']['adder']\n",
    "err_base = rec_dat['hyp_std'][-K:]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot([i+4.7, i+5.3], [true_base[i]]*2, color=\"black\", linestyle=\"-\", lw=1.2, zorder=1, alpha=0.4)\n",
    "    plt.errorbar([i+5], avg_base[i], yerr=num_std*err_base[i], c=c, lw=1, marker='^', markersize=4)\n",
    "\n",
    "plt.xticks([0,1,2,3,5,6])\n",
    "# plt.xlim(-0.5,3.5); \n",
    "plt.ylim(-0.25,0.25)\n",
    "\n",
    "plt.gca().set_xticklabels([r\"$\\sigma_1$\", r\"$\\sigma_2$\",\n",
    "                           r\"$\\alpha_1$\", r\"$\\alpha_2$\",\n",
    "                           r\"$\\beta_1$\", r\"$\\beta_2$\"])\n",
    "\n",
    "# plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN1b.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "rec_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['rec_dat'].item()\n",
    "gen_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "# Plotting\n",
    "fig = plt.figure(figsize=(3.75,1.4))\n",
    "sim_colors = [colors['bias'], colors['h']]\n",
    "\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(gen_dat['W'][i], c=\"black\", lw=0.5, alpha=0.5, zorder=2*i)\n",
    "    plt.plot(rec_dat['wMode'][i], c=c, lw=1, linestyle='-', alpha=0.85, zorder=2*i+1)\n",
    "    plt.fill_between(np.arange(len(rec_dat['wMode'][i])),\n",
    "                     rec_dat['wMode'][i] - 2 * rec_dat['W_std'][i],\n",
    "                     rec_dat['wMode'][i] + 2 * rec_dat['W_std'][i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i+1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "# plt.xticks(1000*np.arange(0,11))\n",
    "plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(-4,5,2))\n",
    "\n",
    "plt.xlim(0,len(gen_dat['W'][0])); plt.ylim(-5,5)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weights\")\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN1a.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "# rec_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['rec_dat'].item()\n",
    "# gen_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "noise = np.cumsum(rec_dat['E_rw'], axis=1)\n",
    "learning = np.cumsum(rec_dat['v_rw'], axis=1)\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.75,1.4))\n",
    "\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning[i], c=c, lw=1.0, linestyle='-', alpha=0.85, zorder=10+2*i)\n",
    "    plt.plot(learning_sim[i], c=\"black\", lw=0.5, linestyle='-', alpha=0.25, zorder=2)\n",
    "    plt.fill_between(np.arange(len(noise[i])),\n",
    "                     learning[i], learning_sim[i],\n",
    "                     facecolor=\"lightgray\", alpha=0.25, zorder=1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "# plt.xticks(1000*np.arange(0,11))\n",
    "plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(-4,5,2))\n",
    "\n",
    "plt.xlim(0,len(noise[i])); plt.ylim(-5,5)\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weight\\nComponents\")\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN1c.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "# rec_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['rec_dat'].item()\n",
    "# gen_dat = np.load(spath+'FigN1_data2.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "noise = np.cumsum(rec_dat['E_rw'], axis=1)\n",
    "learning = np.cumsum(rec_dat['v_rw'], axis=1)\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.75,1.4))\n",
    "\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(noise[i], c=c, lw=1.0, linestyle='-', alpha=0.85, zorder=10+2*i+1)\n",
    "    plt.plot(noise_sim[i], c=\"black\", lw=0.5, linestyle='-', alpha=0.25, zorder=2)\n",
    "    plt.fill_between(np.arange(len(noise[i])),\n",
    "                     noise[i], noise_sim[i],\n",
    "                     facecolor=\"lightgray\", alpha=0.25, zorder=1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "# plt.xticks(1000*np.arange(0,11))\n",
    "# plt.gca().set_xticklabels([0,1000,2000,3000,4000,5000])\n",
    "plt.yticks(np.arange(-4,5,2))\n",
    "\n",
    "plt.xlim(0,len(noise[i])); plt.ylim(-5,5)\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weight\\nComponents\")\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN1d.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Figure S5\n",
    "\n",
    "Evaluating model mismatch and recovery of the correct learning rule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "### Set up\n",
    "N = 10000  # would like to work with ~6000 trials, should try\n",
    "\n",
    "seed = 20  # 20\n",
    "print(\"Seed:\", seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Simulate sequence of contrasts directly (instead of grabbing from real mouse)\n",
    "p=5\n",
    "contrasts = np.array([0.0625, 0.125, 0.25, 0.5, 1.0])  # full set of contrasts\n",
    "tanh_contrasts = np.tanh(p*contrasts)/np.tanh(p)  # tanh transform\n",
    "all_contrasts = np.hstack((-tanh_contrasts[::-1], [0], tanh_contrasts))  # left contrasts are neg\n",
    "\n",
    "contrast_seq = np.random.choice(all_contrasts, size=N) # generate random sequence\n",
    "answer = np.sign(contrast_seq).astype(int)  # infer answers for each trial\n",
    "answer[answer==0] = np.sign(np.random.randn(np.sum(answer==0)))  # randomly select answer for 0 contrast\n",
    "answer[answer<0] = 0\n",
    "\n",
    "# Build data set from simulated data\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "dat = {'inputs': {'cBoth': contrast_seq.reshape(-1,1)}, 'answer': answer, 'y': np.zeros(N)}\n",
    "X = psy.read_input(dat, weights)\n",
    "K = np.sum([i for i in weights.values()])\n",
    "\n",
    "# Set true parameters of weight trajectory simulation\n",
    "true_sigma  = 2**np.array([-9.0, -9.0])  # -5, -6\n",
    "true_alpha  = 2**np.array([-6.0, -7.0])  # -6, -7\n",
    "true_base  = np.array([0.4, -0.2])  # -0.15, 0.1, careful about which parametrization is being used\n",
    "true_hyper = {\n",
    "    'adder': true_base,\n",
    "    'alpha': true_alpha,\n",
    "    'sigma': true_sigma,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "sim_learning_rule = reinforce_base  # learning rule used in generation of weights\n",
    "rec_learning_rule = REINFORCE_base  # learning rule used for recovery of weights\n",
    "\n",
    "# Simulate the weight trajectories, also returning choices and rewards\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1, seed=seed,\n",
    "                                       alpha=true_alpha, base=true_base/true_alpha, # NOTE\n",
    "                                       learning_rule=sim_learning_rule)\n",
    "\n",
    "dat.update({\"y\": y, \"correct\": r})\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, 'true_base': true_base,\n",
    "           \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule.__name__, \"rec_learning_rule\": rec_learning_rule.__name__,\n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "np.savez_compressed(spath+'FigN2_data.npz', gen_dat=gen_dat)\n",
    "\n",
    "\n",
    "# Display simulated weights, broken up into learning/nosie components\n",
    "fig = psy.plot_weights(W, weights)\n",
    "\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning_sim[i], c=c, lw=1, linestyle='-', alpha=0.75, zorder=1)\n",
    "    plt.plot(noise_sim[i], c=c, lw=0.5, linestyle='--', alpha=0.5, zorder=2)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "\n",
    "\n",
    "# Calculate evdience under true hyper\n",
    "wMode_true, _, logEvd, _ = getMAP(dat, true_hyper, weights, W0=None,\n",
    "                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "print(\"True Evidence:\", logEvd)\n",
    "fig = psy.plot_weights(wMode_true.reshape(K,-1), weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN2_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-16] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['alpha', 'adder']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": gen_dat['rec_learning_rule'],\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[gen_dat['rec_learning_rule']],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN2_data.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "rec_dat = np.load(spath+'FigN2_data.npz', allow_pickle=True)['rec_dat'].item()\n",
    "gen_dat = np.load(spath+'FigN2_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "# Plotting\n",
    "fig = plt.figure(figsize=(3.75,1.4))\n",
    "sim_colors = [colors['bias'], colors['h']]\n",
    "\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(gen_dat['W'][i], c=\"black\", lw=0.5, alpha=0.5, zorder=2*i)\n",
    "    plt.plot(rec_dat['wMode'][i], c=c, lw=1, linestyle='-', alpha=0.85, zorder=2*i+1)\n",
    "    plt.fill_between(np.arange(len(rec_dat['wMode'][i])),\n",
    "                     rec_dat['wMode'][i] - 2 * rec_dat['W_std'][i],\n",
    "                     rec_dat['wMode'][i] + 2 * rec_dat['W_std'][i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i+1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "# plt.xticks(1000*np.arange(0,11))\n",
    "plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(-8,9,2))\n",
    "\n",
    "plt.xlim(0,len(gen_dat['W'][0])); plt.ylim(-7.5,7.5)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weights\")\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN2a.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN2_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha', 'adder']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": gen_dat['rec_learning_rule'],\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[gen_dat['rec_learning_rule']],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN2_data_b.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "rec_dat = np.load(spath+'FigN2_data_b.npz', allow_pickle=True)['rec_dat'].item()\n",
    "gen_dat = np.load(spath+'FigN2_data_b.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "# Plotting\n",
    "fig = plt.figure(figsize=(3.75,1.4))\n",
    "sim_colors = [colors['bias'], colors['h']]\n",
    "\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(gen_dat['W'][i], c=\"black\", lw=0.5, alpha=0.5, zorder=2*i)\n",
    "    plt.plot(rec_dat['wMode'][i], c=c, lw=1, linestyle='-', alpha=0.85, zorder=2*i+1)\n",
    "    plt.fill_between(np.arange(len(rec_dat['wMode'][i])),\n",
    "                     rec_dat['wMode'][i] - 2 * rec_dat['W_std'][i],\n",
    "                     rec_dat['wMode'][i] + 2 * rec_dat['W_std'][i],\n",
    "                     facecolor=c, alpha=0.2, zorder=2*i+1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "# plt.xticks(1000*np.arange(0,11))\n",
    "# plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(-8,9,2))\n",
    "\n",
    "plt.xlim(0,len(gen_dat['W'][0])); plt.ylim(-7.5,7.5)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weights\")\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN2b.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Figure S6\n",
    "\n",
    "Fitting the model without noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "### Set up\n",
    "N = 10000  # would like to work with ~6000 trials, should try\n",
    "\n",
    "seed = 20  # 20\n",
    "print(\"Seed:\", seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Simulate sequence of contrasts directly (instead of grabbing from real mouse)\n",
    "p=5\n",
    "contrasts = np.array([0.0625, 0.125, 0.25, 0.5, 1.0])  # full set of contrasts\n",
    "tanh_contrasts = np.tanh(p*contrasts)/np.tanh(p)  # tanh transform\n",
    "all_contrasts = np.hstack((-tanh_contrasts[::-1], [0], tanh_contrasts))  # left contrasts are neg\n",
    "\n",
    "contrast_seq = np.random.choice(all_contrasts, size=N) # generate random sequence\n",
    "answer = np.sign(contrast_seq).astype(int)  # infer answers for each trial\n",
    "answer[answer==0] = np.sign(np.random.randn(np.sum(answer==0)))  # randomly select answer for 0 contrast\n",
    "answer[answer<0] = 0\n",
    "\n",
    "# Build data set from simulated data\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "dat = {'inputs': {'cBoth': contrast_seq.reshape(-1,1)}, 'answer': answer, 'y': np.zeros(N)}\n",
    "X = psy.read_input(dat, weights)\n",
    "K = np.sum([i for i in weights.values()])\n",
    "\n",
    "# Set true parameters of weight trajectory simulation\n",
    "true_sigma  = 2**np.array([-5.0, -5.0])  # -5, -6\n",
    "true_alpha  = 2**np.array([-32.0, -32.0])  # -6, -7\n",
    "true_base  = np.array([0.0, 0.0])  # -0.15, 0.1, careful about which parametrization is being used\n",
    "true_hyper = {\n",
    "    'adder': true_base,\n",
    "    'alpha': true_alpha,\n",
    "    'sigma': true_sigma,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "def no_learning(W, X, y, r, answer, i, base=None):\n",
    "    pR = expit(X[i-1] @ W[i-1])\n",
    "    return pR * (1-pR) * X[i-1] * (-1)**(answer[i-1]+1) / (1/2)\n",
    "\n",
    "sim_learning_rule = no_learning  # learning rule used in generation of weights\n",
    "rec_learning_rule = None  # learning rule used for recovery of weights\n",
    "\n",
    "# Simulate the weight trajectories, also returning choices and rewards\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1, seed=seed,\n",
    "                                       alpha=true_alpha, base=true_base/true_alpha, # NOTE\n",
    "                                       learning_rule=sim_learning_rule)\n",
    "\n",
    "dat.update({\"y\": y, \"correct\": r})\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, 'true_base': true_base,\n",
    "           \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule.__name__, \n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "np.savez_compressed(spath+'FigN3_data.npz', gen_dat=gen_dat)\n",
    "\n",
    "\n",
    "# Display simulated weights, broken up into learning/nosie components\n",
    "fig = psy.plot_weights(W, weights)\n",
    "\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning_sim[i], c=c, lw=1, linestyle='-', alpha=0.75, zorder=1)\n",
    "    plt.plot(noise_sim[i], c=c, lw=0.5, linestyle='--', alpha=0.5, zorder=2)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "\n",
    "\n",
    "# Calculate evdience under true hyper\n",
    "wMode_true, _, logEvd, _ = getMAP(dat, true_hyper, weights, W0=None,\n",
    "                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "print(\"True Evidence:\", logEvd)\n",
    "fig = psy.plot_weights(wMode_true.reshape(K,-1), weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": None,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=None,\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode}\n",
    "np.savez_compressed(spath+'FigN3_data_a1.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': 2**-7,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_a2.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_a3.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha', 'adder']\n",
    "lr = \"REINFORCE_base\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_a4.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "### Set up\n",
    "N = 10000  # would like to work with ~6000 trials, should try\n",
    "\n",
    "seed = 20  # 20\n",
    "print(\"Seed:\", seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Simulate sequence of contrasts directly (instead of grabbing from real mouse)\n",
    "p=5\n",
    "contrasts = np.array([0.0625, 0.125, 0.25, 0.5, 1.0])  # full set of contrasts\n",
    "tanh_contrasts = np.tanh(p*contrasts)/np.tanh(p)  # tanh transform\n",
    "all_contrasts = np.hstack((-tanh_contrasts[::-1], [0], tanh_contrasts))  # left contrasts are neg\n",
    "\n",
    "contrast_seq = np.random.choice(all_contrasts, size=N) # generate random sequence\n",
    "answer = np.sign(contrast_seq).astype(int)  # infer answers for each trial\n",
    "answer[answer==0] = np.sign(np.random.randn(np.sum(answer==0)))  # randomly select answer for 0 contrast\n",
    "answer[answer<0] = 0\n",
    "\n",
    "# Build data set from simulated data\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "dat = {'inputs': {'cBoth': contrast_seq.reshape(-1,1)}, 'answer': answer, 'y': np.zeros(N)}\n",
    "X = psy.read_input(dat, weights)\n",
    "K = np.sum([i for i in weights.values()])\n",
    "\n",
    "# Set true parameters of weight trajectory simulation\n",
    "true_sigma  = 2**np.array([-5.0, -5.0])  # -5, -6\n",
    "true_alpha  = 2**np.array([-7.0])  # -6, -7\n",
    "true_base  = np.array([0.0, 0.0])  # -0.15, 0.1, careful about which parametrization is being used\n",
    "true_hyper = {\n",
    "    'adder': true_base,\n",
    "    'alpha': true_alpha,\n",
    "    'sigma': true_sigma,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "\n",
    "sim_learning_rule = reinforce  # learning rule used in generation of weights\n",
    "rec_learning_rule = REINFORCE  # learning rule used for recovery of weights\n",
    "\n",
    "# Simulate the weight trajectories, also returning choices and rewards\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1, seed=seed,\n",
    "                                       alpha=true_alpha, base=true_base/true_alpha, # NOTE\n",
    "                                       learning_rule=sim_learning_rule)\n",
    "\n",
    "dat.update({\"y\": y, \"correct\": r})\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, 'true_base': true_base,\n",
    "           \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule.__name__, 'rec_learning_rule': rec_learning_rule.__name__,\n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "np.savez_compressed(spath+'FigN3_data_b.npz', gen_dat=gen_dat)\n",
    "\n",
    "\n",
    "# Display simulated weights, broken up into learning/nosie components\n",
    "fig = psy.plot_weights(W, weights)\n",
    "\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning_sim[i], c=c, lw=1, linestyle='-', alpha=0.75, zorder=1)\n",
    "    plt.plot(noise_sim[i], c=c, lw=0.5, linestyle='--', alpha=0.5, zorder=2)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "\n",
    "\n",
    "# Calculate evdience under true hyper\n",
    "wMode_true, _, logEvd, _ = getMAP(dat, true_hyper, weights, W0=None,\n",
    "                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "print(\"True Evidence:\", logEvd)\n",
    "fig = psy.plot_weights(wMode_true.reshape(K,-1), weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_b.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": None,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=None,\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode}\n",
    "np.savez_compressed(spath+'FigN3_data_b1.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_b.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': 2**-7,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_b2.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_b.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_b3.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_b.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha', 'adder']\n",
    "lr = \"REINFORCE_base\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_b4.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "### Set up\n",
    "N = 10000  # would like to work with ~6000 trials, should try\n",
    "\n",
    "seed = 20  # 20\n",
    "print(\"Seed:\", seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Simulate sequence of contrasts directly (instead of grabbing from real mouse)\n",
    "p=5\n",
    "contrasts = np.array([0.0625, 0.125, 0.25, 0.5, 1.0])  # full set of contrasts\n",
    "tanh_contrasts = np.tanh(p*contrasts)/np.tanh(p)  # tanh transform\n",
    "all_contrasts = np.hstack((-tanh_contrasts[::-1], [0], tanh_contrasts))  # left contrasts are neg\n",
    "\n",
    "contrast_seq = np.random.choice(all_contrasts, size=N) # generate random sequence\n",
    "answer = np.sign(contrast_seq).astype(int)  # infer answers for each trial\n",
    "answer[answer==0] = np.sign(np.random.randn(np.sum(answer==0)))  # randomly select answer for 0 contrast\n",
    "answer[answer<0] = 0\n",
    "\n",
    "# Build data set from simulated data\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "dat = {'inputs': {'cBoth': contrast_seq.reshape(-1,1)}, 'answer': answer, 'y': np.zeros(N)}\n",
    "X = psy.read_input(dat, weights)\n",
    "K = np.sum([i for i in weights.values()])\n",
    "\n",
    "# Set true parameters of weight trajectory simulation\n",
    "true_sigma  = 2**np.array([-5.0, -5.0])  # -5, -6\n",
    "true_alpha  = 2**np.array([-5.0, -9.0])  # -6, -7\n",
    "true_base  = np.array([0.0, 0.0])  # -0.15, 0.1, careful about which parametrization is being used\n",
    "true_hyper = {\n",
    "    'adder': true_base,\n",
    "    'alpha': true_alpha,\n",
    "    'sigma': true_sigma,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "\n",
    "sim_learning_rule = reinforce  # learning rule used in generation of weights\n",
    "rec_learning_rule = REINFORCE  # learning rule used for recovery of weights\n",
    "\n",
    "# Simulate the weight trajectories, also returning choices and rewards\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1, seed=seed,\n",
    "                                       alpha=true_alpha, base=true_base/true_alpha, # NOTE\n",
    "                                       learning_rule=sim_learning_rule)\n",
    "\n",
    "dat.update({\"y\": y, \"correct\": r})\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, 'true_base': true_base,\n",
    "           \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule.__name__, 'rec_learning_rule': rec_learning_rule.__name__,\n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "np.savez_compressed(spath+'FigN3_data_c.npz', gen_dat=gen_dat)\n",
    "\n",
    "\n",
    "# Display simulated weights, broken up into learning/nosie components\n",
    "fig = psy.plot_weights(W, weights)\n",
    "\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning_sim[i], c=c, lw=1, linestyle='-', alpha=0.75, zorder=1)\n",
    "    plt.plot(noise_sim[i], c=c, lw=0.5, linestyle='--', alpha=0.5, zorder=2)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "\n",
    "\n",
    "# Calculate evdience under true hyper\n",
    "wMode_true, _, logEvd, _ = getMAP(dat, true_hyper, weights, W0=None,\n",
    "                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "print(\"True Evidence:\", logEvd)\n",
    "fig = psy.plot_weights(wMode_true.reshape(K,-1), weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_c.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": None,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=None,\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode}\n",
    "np.savez_compressed(spath+'FigN3_data_c1.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_c.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': 2**-7,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_c2.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_c.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_c3.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_c.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha', 'adder']\n",
    "lr = \"REINFORCE_base\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_c4.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "### Set up\n",
    "N = 10000  # would like to work with ~6000 trials, should try\n",
    "\n",
    "seed = 20  # 20\n",
    "print(\"Seed:\", seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Simulate sequence of contrasts directly (instead of grabbing from real mouse)\n",
    "p=5\n",
    "contrasts = np.array([0.0625, 0.125, 0.25, 0.5, 1.0])  # full set of contrasts\n",
    "tanh_contrasts = np.tanh(p*contrasts)/np.tanh(p)  # tanh transform\n",
    "all_contrasts = np.hstack((-tanh_contrasts[::-1], [0], tanh_contrasts))  # left contrasts are neg\n",
    "\n",
    "contrast_seq = np.random.choice(all_contrasts, size=N) # generate random sequence\n",
    "answer = np.sign(contrast_seq).astype(int)  # infer answers for each trial\n",
    "answer[answer==0] = np.sign(np.random.randn(np.sum(answer==0)))  # randomly select answer for 0 contrast\n",
    "answer[answer<0] = 0\n",
    "\n",
    "# Build data set from simulated data\n",
    "weights = {\"bias\": 1, \"cBoth\": 1}\n",
    "dat = {'inputs': {'cBoth': contrast_seq.reshape(-1,1)}, 'answer': answer, 'y': np.zeros(N)}\n",
    "X = psy.read_input(dat, weights)\n",
    "K = np.sum([i for i in weights.values()])\n",
    "\n",
    "# Set true parameters of weight trajectory simulation\n",
    "true_sigma  = 2**np.array([-5.0, -5.0])  # -5, -6\n",
    "true_alpha  = 2**np.array([-5.0, -9.0])  # -6, -7\n",
    "true_base  = np.array([0.2, -0.2])  # -0.15, 0.1, careful about which parametrization is being used\n",
    "true_hyper = {\n",
    "    'adder': true_base,\n",
    "    'alpha': true_alpha,\n",
    "    'sigma': true_sigma,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "\n",
    "sim_learning_rule = reinforce_base  # learning rule used in generation of weights\n",
    "rec_learning_rule = REINFORCE_base  # learning rule used for recovery of weights\n",
    "\n",
    "# Simulate the weight trajectories, also returning choices and rewards\n",
    "W, y, r, sim_noise = simulate_learning(X=X, answer=answer, sigma=true_sigma, sigma0=1, seed=seed,\n",
    "                                       alpha=true_alpha, base=true_base/true_alpha, # NOTE\n",
    "                                       learning_rule=sim_learning_rule)\n",
    "\n",
    "dat.update({\"y\": y, \"correct\": r})\n",
    "gen_dat = {\"dat\": dat, 'true_sigma': true_sigma, 'true_alpha': true_alpha, 'true_base': true_base,\n",
    "           \"weights\": weights, \"K\": K,\n",
    "           'sim_learning_rule': sim_learning_rule.__name__, 'rec_learning_rule': rec_learning_rule.__name__,\n",
    "           \"W\": W, \"sim_noise\": sim_noise, \"seed\": seed}\n",
    "np.savez_compressed(spath+'FigN3_data_d.npz', gen_dat=gen_dat)\n",
    "\n",
    "\n",
    "# Display simulated weights, broken up into learning/nosie components\n",
    "fig = psy.plot_weights(W, weights)\n",
    "\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.25,1.25))\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning_sim[i], c=c, lw=1, linestyle='-', alpha=0.75, zorder=1)\n",
    "    plt.plot(noise_sim[i], c=c, lw=0.5, linestyle='--', alpha=0.5, zorder=2)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "\n",
    "\n",
    "# Calculate evdience under true hyper\n",
    "wMode_true, _, logEvd, _ = getMAP(dat, true_hyper, weights, W0=None,\n",
    "                         learning_rule=rec_learning_rule, showOpt=0, tol=1e-12)\n",
    "print(\"True Evidence:\", logEvd)\n",
    "fig = psy.plot_weights(wMode_true.reshape(K,-1), weights)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_d.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma']\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": None,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=None,\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode}\n",
    "np.savez_compressed(spath+'FigN3_data_d1.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_d.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': 2**-7,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_d2.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_d.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha']\n",
    "lr = \"REINFORCE\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_d3.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "gen_dat = np.load(spath+'FigN3_data_d.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "K = gen_dat['K']\n",
    "hyper_guess = {\n",
    "    'adder': [0.0] * K,\n",
    "    'alpha': [2**-7] * K,\n",
    "    'sigma': [2**-5] * K,\n",
    "    'sigInit': [2**4] * K,\n",
    "    'sigDay': None,\n",
    "}\n",
    "\n",
    "# Optimizing for alpha and baseline (adder) simultaneously\n",
    "optList = ['sigma', 'alpha', 'adder']\n",
    "lr = \"REINFORCE_base\"\n",
    "\n",
    "# List of extra arguments used by evd_lossfun in optimization of evidence\n",
    "args = {\"optList\": optList, \"dat\": gen_dat['dat'], \"K\": K, \"learning_rule\": lr,\n",
    "        \"hyper\": hyper_guess, \"weights\": gen_dat['weights'], \"update_w\": True, \"wMode\": None,\n",
    "        \"tol\": 1e-6, \"showOpt\": True,\n",
    "       }\n",
    "options = {'maxiter': 1e5}\n",
    "\n",
    "# Optimization, can also use Nelder-Mead but COBYLA is fastest and pretty reliable\n",
    "print(\"True Hyper:\", np.log2(gen_dat['true_sigma']), np.log2(gen_dat['true_alpha']), gen_dat['true_base'])\n",
    "# print(\"True Evidence:\", gen_dat['logEvd'])\n",
    "res = minimize(evd_lossfun, hyper_to_list(hyper_guess, optList, K), args=args, method='COBYLA', options=options)\n",
    "print(\"Evidence:\", -res.fun, \"  \", optList, \": \", res.x)\n",
    "\n",
    "opt_hyper = update_hyper(res.x, optList, hyper_guess, K)\n",
    "\n",
    "# Recover weights from optimal hyper\n",
    "wMode, Hess, logEvd, other = getMAP(dat, opt_hyper, gen_dat['weights'], W0=None,\n",
    "                                    learning_rule=globals()[lr],\n",
    "                                    showOpt=0, tol=1e-12)\n",
    "\n",
    "wMode = wMode.reshape((K, -1), order=\"C\")\n",
    "v_rw = other['pT']['learning_terms']['v'].T\n",
    "E_rw = other['pT']['learning_terms'][\"E_flat\"].reshape((K, -1), order=\"C\")\n",
    "\n",
    "# Recover error bars for weights\n",
    "W_std = getCredibleInterval(Hess, K)\n",
    "\n",
    "rec_dat = {\"args\": args, 'res': res, 'opt_hyper': opt_hyper, 'logEvd': logEvd,\n",
    "           \"W_std\": W_std, \"wMode\": wMode, 'v_rw': v_rw, 'E_rw': E_rw}\n",
    "np.savez_compressed(spath+'FigN3_data_d4.npz', rec_dat=rec_dat, gen_dat=gen_dat)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sim_colors = [colors['bias'], colors['h']]\n",
    "fig, axs = plt.subplots(4,4,sharex=False,sharey=False,figsize=(5, 2))\n",
    "\n",
    "def standard_adjust(ax):\n",
    "    plt.sca(ax)\n",
    "    plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "    plt.xticks(1000*np.arange(0,11,2))\n",
    "    plt.gca().set_xticklabels([])\n",
    "    plt.yticks(np.arange(-6,9,3))\n",
    "    plt.gca().set_yticklabels([])\n",
    "    plt.xlim(0,10000)\n",
    "    plt.ylim(-6.5,6.5)\n",
    "    plt.gca().spines['right'].set_visible(False)\n",
    "    plt.gca().spines['top'].set_visible(False)\n",
    "    plt.setp(ax.spines.values(), linewidth=0.5)\n",
    "    plt.setp(ax.xaxis.get_ticklines(), 'markeredgewidth', 0.5)\n",
    "    plt.setp(ax.yaxis.get_ticklines(), 'markeredgewidth', 0.5)\n",
    "\n",
    "    \n",
    "\n",
    "sim_paths = ['a', 'b', 'c', 'd']\n",
    "rec_paths = ['1', '2', '3', '4']\n",
    "\n",
    "for k in range(4):\n",
    "    for j in range(4):\n",
    "        # Reload data\n",
    "        rec_dat = np.load(spath+'FigN3_data_'+sim_paths[k]+rec_paths[j]+'.npz', allow_pickle=True)['rec_dat'].item()\n",
    "        gen_dat = np.load(spath+'FigN3_data_'+sim_paths[k]+rec_paths[j]+'.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "        # Plotting\n",
    "        plt.sca(axs[k,j])\n",
    "\n",
    "        for i, c in enumerate(sim_colors):\n",
    "            plt.plot(gen_dat['W'][i], c=\"black\", lw=0.5, alpha=0.5, zorder=2*i)\n",
    "            plt.plot(rec_dat['wMode'][i], c=c, lw=1, linestyle='-', alpha=0.85, zorder=2*i+1)\n",
    "            plt.fill_between(np.arange(len(rec_dat['wMode'][i])),\n",
    "                             rec_dat['wMode'][i] - 2 * rec_dat['W_std'][i],\n",
    "                             rec_dat['wMode'][i] + 2 * rec_dat['W_std'][i],\n",
    "                             facecolor=c, alpha=0.2, zorder=2*i+1)\n",
    "        standard_adjust(plt.gca())\n",
    "        \n",
    "        if k != 3:\n",
    "            plt.xticks([])\n",
    "        if j != 0:\n",
    "            plt.yticks([])\n",
    "            \n",
    "        if k == j:\n",
    "            axs[k,j].set_facecolor(\"#f2f2f2\")\n",
    "            \n",
    "#         if k==3 and j==0:\n",
    "#             plt.gca().set_xticklabels([0, None, None, None, None,10000])\n",
    "#             plt.gca().set_yticklabels([-6,None,0,None,6])\n",
    "\n",
    "        \n",
    "plt.subplots_adjust(0,0,1,1, wspace=0.05, hspace=0.1)\n",
    "plt.savefig(spath + \"FigN3a.pdf\")\n",
    "\n",
    "# axs[0,1].remove()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# Reload data\n",
    "rec_dat = np.load(spath+'FigN1_data.npz', allow_pickle=True)['rec_dat'].item()\n",
    "gen_dat = np.load(spath+'FigN1_data.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "noise = np.cumsum(rec_dat['E_rw'], axis=1)\n",
    "learning = np.cumsum(rec_dat['v_rw'], axis=1)\n",
    "noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "fig = plt.figure(figsize=(3.75,1.4))\n",
    "\n",
    "sim_colors = [colors['bias'], colors['cBoth']]\n",
    "for i, c in enumerate(sim_colors):\n",
    "    plt.plot(learning[i], c=c, lw=1.0, linestyle='-', alpha=0.85, zorder=10+2*i)\n",
    "    plt.plot(learning_sim[i], c=\"black\", lw=0.5, linestyle='-', alpha=0.25, zorder=2)\n",
    "    plt.fill_between(np.arange(len(noise[i])),\n",
    "                     learning[i], learning_sim[i],\n",
    "                     facecolor=\"lightgray\", alpha=0.25, zorder=1)\n",
    "\n",
    "plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "\n",
    "# plt.xticks(1000*np.arange(0,11))\n",
    "plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(-4,5,2))\n",
    "\n",
    "plt.xlim(0,len(noise[i])); plt.ylim(-5,5)\n",
    "\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "# plt.xlabel(\"Trials\"); plt.ylabel(\"Weight\\nComponents\")\n",
    "\n",
    "plt.subplots_adjust(0,0,1,1) \n",
    "plt.savefig(spath + \"FigN1c.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "sim_colors = [colors['bias'], colors['h']]\n",
    "fig, axs = plt.subplots(4,4,sharex=False,sharey=False,figsize=(5, 2))\n",
    "\n",
    "def standard_adjust(ax):\n",
    "    plt.sca(ax)\n",
    "    plt.axhline(0, color=\"black\", linestyle=\"--\", lw=0.5, alpha=0.5, zorder=0)\n",
    "    plt.xticks(1000*np.arange(0,11,2))\n",
    "    plt.gca().set_xticklabels([])\n",
    "    plt.yticks(np.arange(-6,9,3))\n",
    "    plt.gca().set_yticklabels([])\n",
    "    plt.xlim(0,10000)\n",
    "    plt.ylim(-6.5,6.5)\n",
    "    plt.gca().spines['right'].set_visible(False)\n",
    "    plt.gca().spines['top'].set_visible(False)\n",
    "    plt.setp(ax.spines.values(), linewidth=0.5)\n",
    "    plt.setp(ax.xaxis.get_ticklines(), 'markeredgewidth', 0.5)\n",
    "    plt.setp(ax.yaxis.get_ticklines(), 'markeredgewidth', 0.5)\n",
    "\n",
    "    \n",
    "\n",
    "sim_paths = ['a', 'b', 'c', 'd']\n",
    "rec_paths = ['1', '2', '3', '4']\n",
    "\n",
    "for k in range(4):\n",
    "    for j in range(4):\n",
    "        \n",
    "        if j==0:\n",
    "            axs[k,j].remove()\n",
    "            continue\n",
    "        # Reload data\n",
    "        rec_dat = np.load(spath+'FigN3_data_'+sim_paths[k]+rec_paths[j]+'.npz', allow_pickle=True)['rec_dat'].item()\n",
    "        gen_dat = np.load(spath+'FigN3_data_'+sim_paths[k]+rec_paths[j]+'.npz', allow_pickle=True)['gen_dat'].item()\n",
    "\n",
    "        noise = np.cumsum(rec_dat['E_rw'], axis=1)\n",
    "        learning = np.cumsum(rec_dat['v_rw'], axis=1)\n",
    "        noise_sim = np.cumsum(gen_dat['sim_noise'].T, axis=1)\n",
    "        learning_sim = gen_dat['W'] - noise_sim\n",
    "\n",
    "        # Plotting\n",
    "        plt.sca(axs[k,j])\n",
    "\n",
    "        for i, c in enumerate(sim_colors):\n",
    "            plt.plot(learning[i], c=c, lw=1.0, linestyle='-', alpha=0.85, zorder=10+2*i)\n",
    "            plt.plot(learning_sim[i], c=\"black\", lw=0.5, linestyle='-', alpha=0.25, zorder=2)\n",
    "            plt.fill_between(np.arange(len(noise[i])),\n",
    "                             learning[i], learning_sim[i],\n",
    "                             facecolor=\"lightgray\", alpha=0.25, zorder=1)\n",
    "        standard_adjust(plt.gca())\n",
    "        \n",
    "        if k != 3:\n",
    "            plt.xticks([])\n",
    "        if j != 1:\n",
    "            plt.yticks([])\n",
    "            \n",
    "        if k == j:\n",
    "            axs[k,j].set_facecolor(\"#f2f2f2\")\n",
    "            \n",
    "#         if k==3 and j==1:\n",
    "#             plt.gca().set_xticklabels([0, None, None, None, None,10000])\n",
    "#             plt.gca().set_yticklabels([-6,None,0,None,6])\n",
    "\n",
    "        \n",
    "plt.subplots_adjust(0,0,1,1, wspace=0.05, hspace=0.1)\n",
    "plt.savefig(spath + \"FigN3b.pdf\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def calculate_aic(logEvd, k):\n",
    "    aic = 2*k + 2*logEvd\n",
    "    return aic\n",
    "\n",
    "plt.figure(figsize=(2,1))\n",
    "\n",
    "sim_paths = ['a', 'b', 'c', 'd']\n",
    "rec_paths = ['1', '2', '3', '4']\n",
    "ks = [2,3,4,6]\n",
    "mycolors = [\"red\", \"blue\", \"green\", \"purple\"]\n",
    "mycolors = [psy.COLORS['s1'], psy.COLORS['s_avg'], psy.COLORS['c'], psy.COLORS['s2']]\n",
    "myls = [\"-\", \"--\", \":\", \"-.\"]\n",
    "\n",
    "for k in range(4):\n",
    "    model_aic = []\n",
    "    for j in range(4):\n",
    "        rec_dat = np.load(spath+'FigN3_data_'+sim_paths[k]+rec_paths[j]+'.npz', allow_pickle=True)['rec_dat'].item()\n",
    "        model_aic += [calculate_aic(-rec_dat['logEvd'], ks[j])]\n",
    "    \n",
    "    model_aic = np.array(model_aic) - model_aic[k]\n",
    "    plt.plot(model_aic, color=mycolors[k], alpha=1.0, label=str(k), marker='o', markersize=2.5, ls='-', lw=0.75)\n",
    "    plt.scatter([k], model_aic[k], marker=\"o\", edgecolor=mycolors[k], color=\"white\", s=20, zorder=100, lw=0.75)\n",
    "\n",
    "plt.plot([-5,3], [0,0], color=\"black\", linestyle=\":\", lw=1.0, alpha=0.6, zorder=0)\n",
    "plt.xticks(np.arange(0,4,1))\n",
    "plt.gca().set_xticklabels([])\n",
    "plt.yticks(np.arange(0,20,5))\n",
    "# plt.gca().set_yticklabels([])\n",
    "plt.xlim(-0.08,3.08)\n",
    "plt.ylim(-1,15)\n",
    "plt.gca().spines['right'].set_visible(False)\n",
    "plt.gca().spines['top'].set_visible(False)\n",
    "    \n",
    "plt.subplots_adjust(0,0,1,1)\n",
    "plt.savefig(spath + \"FigN3c.pdf\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true,
    "hidden": true
   },
   "source": [
    "## Figures S7-10\n",
    "\n",
    "Figures S7-10 are analyses of entire populations that cannot be run locally in a reasonable amount of time. See the `cluster_scripts` directory for some files for running fitting in parallel on a cluster."
   ]
  }
 ],
 "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"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}