BehavioralFit3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/anaconda3/lib/python3.6/site-packages/matplotlib/cbook/deprecation.py:107: MatplotlibDeprecationWarning: The mpl_toolkits.axes_grid module was deprecated in version 2.1. Use mpl_toolkits.axes_grid1 and mpl_toolkits.axisartist provies the same functionality instead.\n",
      "  warnings.warn(message, mplDeprecation, stacklevel=1)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np \n",
    "import pandas as pd\n",
    "import pystan\n",
    "import stan_utility\n",
    "import scipy\n",
    "import _pickle as cPickle\n",
    "\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "# to render plots with white background\n",
    "import seaborn as sns\n",
    "sns.set_style('white')\n",
    "sns.set_context('notebook', font_scale=1.2)\n",
    "\n",
    "from mpl_toolkits.axes_grid.anchored_artists import AnchoredText\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# TrialStart = 0, 1. much more TrialStart == 0 than ==1 \n",
    "# First lick: Reward Enter\n",
    "# Next lick: Reward Exit (after they havne't licked for a while) -- not super reliable "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def inv_logit(arr):\n",
    "    '''Elementwise inverse logit (logistic) function.'''\n",
    "    return 1 / (1 + np.exp(-arr))\n",
    "\n",
    "def phi_approx(arr):\n",
    "    '''Elementwise fast approximation of the cumulative unit normal. \n",
    "    For details, see Bowling et al. (2009). \"A logistic approximation \n",
    "    to the cumulative normal distribution.\"'''\n",
    "    return inv_logit(0.07056 * arr ** 3 + 1.5976 * arr)\n",
    "\n",
    "\n",
    "# def softmax(x):\n",
    "#     \"\"\"\n",
    "#     Compute softmax values for each sets of scores in x.\n",
    "    \n",
    "#     Rows are scores for each class. \n",
    "#     Columns are predictions (samples).\n",
    "#     \"\"\"\n",
    "#     scoreMatExp = np.exp(np.asarray(x))\n",
    "#     return scoreMatExp / scoreMatExp.sum(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# data = pd.read_csv('./data/dan_table.csv')\n",
    "data = pd.read_csv('./data/rachel_table.csv', names = ['gcamp', 'TrialStart', 'NosePokeEnter', 'NosePokeExit', 'LeverPresentation',\n",
    "                                                       'IpsLeverPress', 'ConLeverPress', 'CSplus', 'CSneg', 'RewardEnter',\n",
    "                                                       'RewardExit', 'MouseID', 'RecordLoc', 'Session'])\n",
    "\n",
    "data_g = pd.read_csv('./data/dan_table.csv')\n",
    "\n",
    "data['gcamp'] = data_g['gcamp']\n",
    "\n",
    "data = data[(data['RecordLoc'] != 5) ]\n",
    "\n",
    "# Somehow the length of data is different for different RecordLoc... I'm going to use DMS since I'll want to \n",
    "# model the neural data from that area later \n",
    "#         1 = VTA/SNc -> Dorsomedial Striatum\n",
    "#         2 = VTA/SNc -> Nucleus Accumbens\n",
    "#         3 = VTA/ SNc -> Dorsolateral Striatum (forgot i had these)\n",
    "#         4 = VTA/ SNc -> Dorsomedial Striatum Cell Bodies\n",
    "#         5 = GFP controls (DMS and NAc together)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# pull out trial by trial info \n",
    "\n",
    "data_t = []\n",
    "\n",
    "for mouse_id in np.unique(data['MouseID']):\n",
    "    data_small = data[data['MouseID'] == mouse_id]\n",
    "\n",
    "    trial_starts = np.where(data_small['TrialStart'] == 1)[0]\n",
    "    num_trials = len(trial_starts)\n",
    "\n",
    "    for t_i, t in enumerate(trial_starts):\n",
    "\n",
    "        if t_i == (num_trials - 1):\n",
    "            trial_end = len(data_small)\n",
    "        else:\n",
    "            trial_end = trial_starts[t_i + 1]\n",
    "\n",
    "        trial_df = data_small[t:trial_end];\n",
    "\n",
    "#         lever_press = np.where(trial_df['LeverPresentation'] == 1)[0][0]\n",
    "        if sum(trial_df.IpsLeverPress.values) >  sum(trial_df.ConLeverPress.values):\n",
    "            action = 'Ips'\n",
    "        else:\n",
    "            action = 'Con'\n",
    "\n",
    "        if sum(trial_df.CSplus.values) >  sum(trial_df.CSneg.values):\n",
    "            reward = 1\n",
    "        else:\n",
    "            reward = 0\n",
    "\n",
    "        sess = np.unique(trial_df['Session'])[0]\n",
    "        recordloc = np.unique(trial_df['RecordLoc'])[0]\n",
    "        data_t.append([ mouse_id - 1, t_i, action, reward, t , sess, recordloc]) #  lever_press])\n",
    "\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "df_trials = pd.DataFrame(data_t, columns = ('MouseID', 'Trial', 'Action', 'Reward', 'TrialStart',\n",
    "                                             'Session', 'RecordLoc'))\n",
    "                \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "NS = len(np.unique(df_trials['MouseID']))\n",
    "\n",
    "NT = max(df_trials.groupby('MouseID').count()['Trial'])\n",
    "NT_all = df_trials.groupby('MouseID').count()['Trial']\n",
    "\n",
    "c_str = np.empty([NS, NT], 'str')\n",
    "r = np.empty([NS, NT], int)\n",
    "r.fill(-1) # init with nan to ensure that shorter trials dont' get false data \n",
    "for i, s in enumerate(np.unique(df_trials['MouseID'])):\n",
    "    c_str[i, :NT_all[s]] = (df_trials[df_trials['MouseID'] == s])['Action'].values\n",
    "    r[i, :NT_all[s]] = (df_trials[df_trials['MouseID'] == s])['Reward'].values\n",
    "    \n",
    "c = np.empty([NS, NT], int)\n",
    "c.fill(-1)\n",
    "\n",
    "c[c_str == 'C'] = 0; \n",
    "c[c_str == 'I'] = 1;  \n",
    "\n",
    "\n",
    "standata = {'NS':NS,'NT':NT,'NT_all':NT_all, \n",
    "           'r':r, 'c':c}\n",
    "\n",
    "# pystan.stan_rdump(standata, 'rl_behavioral.data.R')\n",
    "# pickle.dump( standata, open( \"standata.p\", \"wb\" ) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# intercept for side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit = pystan.stan(file = 'stan/qlearning_intercept_side.stan',data=standata, iter=1000, warmup=250, chains=4, \n",
    "                  control=dict(adapt_delta = 0.9999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit = pystan.stan(file = 'stan/qlearning2_checkfit.stan',data=standata, iter=1000, warmup=250, chains=4, \n",
    "#                   control=dict(adapt_delta = 0.9999))\n",
    "\n",
    "# fit = sm.sampling( data=standata, iter=1000, warmup=250, chains=4, \n",
    "#                  seed = 4, control=dict(adapt_delta = 0.9999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stan_utility.check_treedepth(fit); stan_utility.check_energy(fit); stan_utility.check_div(fit)\n",
    "# print(fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## Save summary file.\n",
    "# summary = fit.summary()\n",
    "# summary = pd.DataFrame(summary['summary'], columns=summary['summary_colnames'], index=summary['summary_rownames'])\n",
    "# summary.to_csv(os.path.join('data/', 'summary_full_seed4.csv'))\n",
    "\n",
    "# # Save contents of StanFit.\n",
    "# extract = fit.extract()\n",
    "# for k, v in standata.items(): extract[k] = v\n",
    "# with open(os.path.join('data/' 'StanFit_full_seed4.pickle'), 'wb') as fn: cPickle.dump(extract, fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# intercept for last action\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fit2 = pystan.stan(file = 'stan/qlearning_intercept_last.stan',data=standata, iter=1000, warmup=250, chains=4, \n",
    "#                   control=dict(adapt_delta = 0.99))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stan_utility.check_treedepth(fit2); stan_utility.check_energy(fit2); stan_utility.check_div(fit2)\n",
    "print(fit2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ## Save summary file.\n",
    "# summary = fit2.summary()\n",
    "# summary = pd.DataFrame(summary['summary'], columns=summary['summary_colnames'], index=summary['summary_rownames'])\n",
    "# summary.to_csv(os.path.join('data/', 'summary_int_pc.csv'))\n",
    "\n",
    "# # Save contents of StanFit.\n",
    "# extract = fit2.extract()\n",
    "# for k, v in standata.items(): extract[k] = v\n",
    "# with open(os.path.join('data/' 'StanFit_int_pc.pickle'), 'wb') as fn: cPickle.dump(extract, fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis of fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/anaconda3/lib/python3.6/site-packages/ipykernel/__main__.py:3: FutureWarning: from_csv is deprecated. Please use read_csv(...) instead. Note that some of the default arguments are different, so please refer to the documentation for from_csv when changing your function calls\n",
      "  app.launch_new_instance()\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "extract = pickle.load( open( \"data/StanFit_int_pc.pickle\", \"rb\" ) )\n",
    "summary = pd.DataFrame.from_csv(\"data/summary_int_pc.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>mean</th>\n",
       "      <th>se_mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>2.5%</th>\n",
       "      <th>25%</th>\n",
       "      <th>50%</th>\n",
       "      <th>75%</th>\n",
       "      <th>97.5%</th>\n",
       "      <th>n_eff</th>\n",
       "      <th>Rhat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>betam</th>\n",
       "      <td>0.991831</td>\n",
       "      <td>0.001914</td>\n",
       "      <td>0.104817</td>\n",
       "      <td>0.785806</td>\n",
       "      <td>0.926501</td>\n",
       "      <td>0.990275</td>\n",
       "      <td>1.058405</td>\n",
       "      <td>1.204639</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000840</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alpham</th>\n",
       "      <td>0.280245</td>\n",
       "      <td>0.002688</td>\n",
       "      <td>0.125308</td>\n",
       "      <td>0.015506</td>\n",
       "      <td>0.205782</td>\n",
       "      <td>0.283441</td>\n",
       "      <td>0.357970</td>\n",
       "      <td>0.521426</td>\n",
       "      <td>2173.0</td>\n",
       "      <td>0.999823</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>staym</th>\n",
       "      <td>0.946621</td>\n",
       "      <td>0.001909</td>\n",
       "      <td>0.104576</td>\n",
       "      <td>0.742629</td>\n",
       "      <td>0.883670</td>\n",
       "      <td>0.945385</td>\n",
       "      <td>1.008465</td>\n",
       "      <td>1.163876</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999867</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betasd</th>\n",
       "      <td>0.335896</td>\n",
       "      <td>0.001608</td>\n",
       "      <td>0.088088</td>\n",
       "      <td>0.207180</td>\n",
       "      <td>0.273766</td>\n",
       "      <td>0.323315</td>\n",
       "      <td>0.381547</td>\n",
       "      <td>0.542784</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000864</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphasd</th>\n",
       "      <td>0.344023</td>\n",
       "      <td>0.003802</td>\n",
       "      <td>0.133464</td>\n",
       "      <td>0.142716</td>\n",
       "      <td>0.250423</td>\n",
       "      <td>0.325971</td>\n",
       "      <td>0.416741</td>\n",
       "      <td>0.669016</td>\n",
       "      <td>1232.0</td>\n",
       "      <td>1.003453</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>staysd</th>\n",
       "      <td>0.340301</td>\n",
       "      <td>0.001565</td>\n",
       "      <td>0.085730</td>\n",
       "      <td>0.214879</td>\n",
       "      <td>0.280670</td>\n",
       "      <td>0.328428</td>\n",
       "      <td>0.383207</td>\n",
       "      <td>0.550083</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.001274</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[0]</th>\n",
       "      <td>1.379926</td>\n",
       "      <td>0.001465</td>\n",
       "      <td>0.080263</td>\n",
       "      <td>1.228311</td>\n",
       "      <td>1.326660</td>\n",
       "      <td>1.378630</td>\n",
       "      <td>1.432049</td>\n",
       "      <td>1.537987</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000165</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[1]</th>\n",
       "      <td>1.177025</td>\n",
       "      <td>0.001428</td>\n",
       "      <td>0.078231</td>\n",
       "      <td>1.026434</td>\n",
       "      <td>1.124913</td>\n",
       "      <td>1.177907</td>\n",
       "      <td>1.226203</td>\n",
       "      <td>1.331797</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999350</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[2]</th>\n",
       "      <td>0.930249</td>\n",
       "      <td>0.001186</td>\n",
       "      <td>0.064948</td>\n",
       "      <td>0.807819</td>\n",
       "      <td>0.886897</td>\n",
       "      <td>0.929059</td>\n",
       "      <td>0.972513</td>\n",
       "      <td>1.057845</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999029</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[3]</th>\n",
       "      <td>1.098245</td>\n",
       "      <td>0.001162</td>\n",
       "      <td>0.063654</td>\n",
       "      <td>0.976399</td>\n",
       "      <td>1.055412</td>\n",
       "      <td>1.098440</td>\n",
       "      <td>1.139053</td>\n",
       "      <td>1.223629</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999604</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[4]</th>\n",
       "      <td>0.564534</td>\n",
       "      <td>0.001271</td>\n",
       "      <td>0.069606</td>\n",
       "      <td>0.428373</td>\n",
       "      <td>0.517964</td>\n",
       "      <td>0.563651</td>\n",
       "      <td>0.610798</td>\n",
       "      <td>0.704753</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000897</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[5]</th>\n",
       "      <td>0.583091</td>\n",
       "      <td>0.001250</td>\n",
       "      <td>0.068450</td>\n",
       "      <td>0.451537</td>\n",
       "      <td>0.537015</td>\n",
       "      <td>0.582412</td>\n",
       "      <td>0.630098</td>\n",
       "      <td>0.715114</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999084</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[6]</th>\n",
       "      <td>1.415278</td>\n",
       "      <td>0.001924</td>\n",
       "      <td>0.105392</td>\n",
       "      <td>1.206970</td>\n",
       "      <td>1.341734</td>\n",
       "      <td>1.418100</td>\n",
       "      <td>1.487329</td>\n",
       "      <td>1.622447</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999615</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[7]</th>\n",
       "      <td>1.116062</td>\n",
       "      <td>0.002254</td>\n",
       "      <td>0.123449</td>\n",
       "      <td>0.886945</td>\n",
       "      <td>1.031003</td>\n",
       "      <td>1.111392</td>\n",
       "      <td>1.199705</td>\n",
       "      <td>1.358343</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999433</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[8]</th>\n",
       "      <td>0.954083</td>\n",
       "      <td>0.002532</td>\n",
       "      <td>0.138672</td>\n",
       "      <td>0.688488</td>\n",
       "      <td>0.859469</td>\n",
       "      <td>0.952385</td>\n",
       "      <td>1.045171</td>\n",
       "      <td>1.232408</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999676</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[9]</th>\n",
       "      <td>0.684312</td>\n",
       "      <td>0.003164</td>\n",
       "      <td>0.148146</td>\n",
       "      <td>0.418406</td>\n",
       "      <td>0.582348</td>\n",
       "      <td>0.672727</td>\n",
       "      <td>0.778691</td>\n",
       "      <td>1.011478</td>\n",
       "      <td>2193.0</td>\n",
       "      <td>1.000923</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[10]</th>\n",
       "      <td>0.970828</td>\n",
       "      <td>0.002655</td>\n",
       "      <td>0.145434</td>\n",
       "      <td>0.685972</td>\n",
       "      <td>0.873628</td>\n",
       "      <td>0.968891</td>\n",
       "      <td>1.065741</td>\n",
       "      <td>1.259182</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999559</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>betas[11]</th>\n",
       "      <td>1.040339</td>\n",
       "      <td>0.002859</td>\n",
       "      <td>0.156585</td>\n",
       "      <td>0.736354</td>\n",
       "      <td>0.935134</td>\n",
       "      <td>1.039184</td>\n",
       "      <td>1.144083</td>\n",
       "      <td>1.350519</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999773</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[0]</th>\n",
       "      <td>0.335759</td>\n",
       "      <td>0.001952</td>\n",
       "      <td>0.106924</td>\n",
       "      <td>0.139748</td>\n",
       "      <td>0.261489</td>\n",
       "      <td>0.331041</td>\n",
       "      <td>0.405564</td>\n",
       "      <td>0.556481</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000006</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[1]</th>\n",
       "      <td>0.313749</td>\n",
       "      <td>0.002093</td>\n",
       "      <td>0.114653</td>\n",
       "      <td>0.097160</td>\n",
       "      <td>0.238824</td>\n",
       "      <td>0.308004</td>\n",
       "      <td>0.385863</td>\n",
       "      <td>0.547910</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[2]</th>\n",
       "      <td>0.349050</td>\n",
       "      <td>0.002336</td>\n",
       "      <td>0.127975</td>\n",
       "      <td>0.106704</td>\n",
       "      <td>0.259944</td>\n",
       "      <td>0.344313</td>\n",
       "      <td>0.430415</td>\n",
       "      <td>0.625239</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999509</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[3]</th>\n",
       "      <td>0.737345</td>\n",
       "      <td>0.002746</td>\n",
       "      <td>0.150379</td>\n",
       "      <td>0.472788</td>\n",
       "      <td>0.631962</td>\n",
       "      <td>0.729918</td>\n",
       "      <td>0.831182</td>\n",
       "      <td>1.065803</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000549</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[4]</th>\n",
       "      <td>0.543439</td>\n",
       "      <td>0.003798</td>\n",
       "      <td>0.208033</td>\n",
       "      <td>0.180177</td>\n",
       "      <td>0.402940</td>\n",
       "      <td>0.522610</td>\n",
       "      <td>0.665781</td>\n",
       "      <td>1.012842</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.001456</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[5]</th>\n",
       "      <td>0.429557</td>\n",
       "      <td>0.003762</td>\n",
       "      <td>0.206070</td>\n",
       "      <td>0.065418</td>\n",
       "      <td>0.287863</td>\n",
       "      <td>0.417346</td>\n",
       "      <td>0.557378</td>\n",
       "      <td>0.874793</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000158</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[6]</th>\n",
       "      <td>0.217940</td>\n",
       "      <td>0.002135</td>\n",
       "      <td>0.116931</td>\n",
       "      <td>-0.000741</td>\n",
       "      <td>0.135297</td>\n",
       "      <td>0.215179</td>\n",
       "      <td>0.295095</td>\n",
       "      <td>0.454695</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000213</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[7]</th>\n",
       "      <td>0.347727</td>\n",
       "      <td>0.004526</td>\n",
       "      <td>0.216380</td>\n",
       "      <td>-0.045205</td>\n",
       "      <td>0.204184</td>\n",
       "      <td>0.327800</td>\n",
       "      <td>0.483695</td>\n",
       "      <td>0.821103</td>\n",
       "      <td>2286.0</td>\n",
       "      <td>1.000367</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[8]</th>\n",
       "      <td>-0.123617</td>\n",
       "      <td>0.005477</td>\n",
       "      <td>0.220396</td>\n",
       "      <td>-0.518081</td>\n",
       "      <td>-0.281584</td>\n",
       "      <td>-0.131120</td>\n",
       "      <td>0.016973</td>\n",
       "      <td>0.337861</td>\n",
       "      <td>1619.0</td>\n",
       "      <td>0.999767</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[9]</th>\n",
       "      <td>-0.022117</td>\n",
       "      <td>0.008111</td>\n",
       "      <td>0.328156</td>\n",
       "      <td>-0.692413</td>\n",
       "      <td>-0.234400</td>\n",
       "      <td>-0.000705</td>\n",
       "      <td>0.205762</td>\n",
       "      <td>0.569548</td>\n",
       "      <td>1637.0</td>\n",
       "      <td>1.000958</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[10]</th>\n",
       "      <td>0.218687</td>\n",
       "      <td>0.004461</td>\n",
       "      <td>0.244318</td>\n",
       "      <td>-0.242149</td>\n",
       "      <td>0.058077</td>\n",
       "      <td>0.218161</td>\n",
       "      <td>0.369450</td>\n",
       "      <td>0.725681</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999060</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>alphas[11]</th>\n",
       "      <td>0.008869</td>\n",
       "      <td>0.004476</td>\n",
       "      <td>0.207651</td>\n",
       "      <td>-0.380493</td>\n",
       "      <td>-0.135531</td>\n",
       "      <td>0.004154</td>\n",
       "      <td>0.152844</td>\n",
       "      <td>0.427773</td>\n",
       "      <td>2152.0</td>\n",
       "      <td>1.000453</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[0]</th>\n",
       "      <td>1.286290</td>\n",
       "      <td>0.000753</td>\n",
       "      <td>0.041229</td>\n",
       "      <td>1.205779</td>\n",
       "      <td>1.258368</td>\n",
       "      <td>1.286300</td>\n",
       "      <td>1.313091</td>\n",
       "      <td>1.367689</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999884</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[1]</th>\n",
       "      <td>1.088134</td>\n",
       "      <td>0.000750</td>\n",
       "      <td>0.041099</td>\n",
       "      <td>1.007790</td>\n",
       "      <td>1.059463</td>\n",
       "      <td>1.088687</td>\n",
       "      <td>1.115770</td>\n",
       "      <td>1.170416</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999611</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[2]</th>\n",
       "      <td>0.584186</td>\n",
       "      <td>0.000653</td>\n",
       "      <td>0.035784</td>\n",
       "      <td>0.516296</td>\n",
       "      <td>0.559552</td>\n",
       "      <td>0.583430</td>\n",
       "      <td>0.607596</td>\n",
       "      <td>0.655401</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.998944</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[3]</th>\n",
       "      <td>0.860108</td>\n",
       "      <td>0.000675</td>\n",
       "      <td>0.036970</td>\n",
       "      <td>0.786962</td>\n",
       "      <td>0.834661</td>\n",
       "      <td>0.860965</td>\n",
       "      <td>0.884563</td>\n",
       "      <td>0.931428</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999373</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[4]</th>\n",
       "      <td>0.857619</td>\n",
       "      <td>0.000771</td>\n",
       "      <td>0.042206</td>\n",
       "      <td>0.771756</td>\n",
       "      <td>0.828983</td>\n",
       "      <td>0.857818</td>\n",
       "      <td>0.886119</td>\n",
       "      <td>0.941352</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>1.000088</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[5]</th>\n",
       "      <td>0.582339</td>\n",
       "      <td>0.000769</td>\n",
       "      <td>0.042102</td>\n",
       "      <td>0.500756</td>\n",
       "      <td>0.553389</td>\n",
       "      <td>0.582123</td>\n",
       "      <td>0.610419</td>\n",
       "      <td>0.664364</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999826</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[6]</th>\n",
       "      <td>1.138377</td>\n",
       "      <td>0.000905</td>\n",
       "      <td>0.049556</td>\n",
       "      <td>1.041427</td>\n",
       "      <td>1.104558</td>\n",
       "      <td>1.138832</td>\n",
       "      <td>1.172468</td>\n",
       "      <td>1.234085</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999702</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[7]</th>\n",
       "      <td>1.247312</td>\n",
       "      <td>0.001170</td>\n",
       "      <td>0.064088</td>\n",
       "      <td>1.123204</td>\n",
       "      <td>1.203998</td>\n",
       "      <td>1.246351</td>\n",
       "      <td>1.290635</td>\n",
       "      <td>1.376133</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.998835</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[8]</th>\n",
       "      <td>0.846683</td>\n",
       "      <td>0.001241</td>\n",
       "      <td>0.067993</td>\n",
       "      <td>0.704295</td>\n",
       "      <td>0.801611</td>\n",
       "      <td>0.847420</td>\n",
       "      <td>0.892267</td>\n",
       "      <td>0.975140</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999537</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[9]</th>\n",
       "      <td>0.766284</td>\n",
       "      <td>0.001394</td>\n",
       "      <td>0.076368</td>\n",
       "      <td>0.616830</td>\n",
       "      <td>0.715681</td>\n",
       "      <td>0.765879</td>\n",
       "      <td>0.816835</td>\n",
       "      <td>0.913771</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999038</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[10]</th>\n",
       "      <td>1.473008</td>\n",
       "      <td>0.001533</td>\n",
       "      <td>0.083956</td>\n",
       "      <td>1.314673</td>\n",
       "      <td>1.415101</td>\n",
       "      <td>1.472399</td>\n",
       "      <td>1.529529</td>\n",
       "      <td>1.639523</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.999342</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>stay[11]</th>\n",
       "      <td>0.642435</td>\n",
       "      <td>0.001552</td>\n",
       "      <td>0.084990</td>\n",
       "      <td>0.476009</td>\n",
       "      <td>0.584796</td>\n",
       "      <td>0.641904</td>\n",
       "      <td>0.698949</td>\n",
       "      <td>0.803311</td>\n",
       "      <td>3000.0</td>\n",
       "      <td>0.998985</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lp__</th>\n",
       "      <td>-16051.968579</td>\n",
       "      <td>0.157836</td>\n",
       "      <td>5.018586</td>\n",
       "      <td>-16062.813482</td>\n",
       "      <td>-16055.103814</td>\n",
       "      <td>-16051.640716</td>\n",
       "      <td>-16048.506824</td>\n",
       "      <td>-16042.959519</td>\n",
       "      <td>1011.0</td>\n",
       "      <td>1.002175</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    mean   se_mean        sd          2.5%           25%  \\\n",
       "betam           0.991831  0.001914  0.104817      0.785806      0.926501   \n",
       "alpham          0.280245  0.002688  0.125308      0.015506      0.205782   \n",
       "staym           0.946621  0.001909  0.104576      0.742629      0.883670   \n",
       "betasd          0.335896  0.001608  0.088088      0.207180      0.273766   \n",
       "alphasd         0.344023  0.003802  0.133464      0.142716      0.250423   \n",
       "staysd          0.340301  0.001565  0.085730      0.214879      0.280670   \n",
       "betas[0]        1.379926  0.001465  0.080263      1.228311      1.326660   \n",
       "betas[1]        1.177025  0.001428  0.078231      1.026434      1.124913   \n",
       "betas[2]        0.930249  0.001186  0.064948      0.807819      0.886897   \n",
       "betas[3]        1.098245  0.001162  0.063654      0.976399      1.055412   \n",
       "betas[4]        0.564534  0.001271  0.069606      0.428373      0.517964   \n",
       "betas[5]        0.583091  0.001250  0.068450      0.451537      0.537015   \n",
       "betas[6]        1.415278  0.001924  0.105392      1.206970      1.341734   \n",
       "betas[7]        1.116062  0.002254  0.123449      0.886945      1.031003   \n",
       "betas[8]        0.954083  0.002532  0.138672      0.688488      0.859469   \n",
       "betas[9]        0.684312  0.003164  0.148146      0.418406      0.582348   \n",
       "betas[10]       0.970828  0.002655  0.145434      0.685972      0.873628   \n",
       "betas[11]       1.040339  0.002859  0.156585      0.736354      0.935134   \n",
       "alphas[0]       0.335759  0.001952  0.106924      0.139748      0.261489   \n",
       "alphas[1]       0.313749  0.002093  0.114653      0.097160      0.238824   \n",
       "alphas[2]       0.349050  0.002336  0.127975      0.106704      0.259944   \n",
       "alphas[3]       0.737345  0.002746  0.150379      0.472788      0.631962   \n",
       "alphas[4]       0.543439  0.003798  0.208033      0.180177      0.402940   \n",
       "alphas[5]       0.429557  0.003762  0.206070      0.065418      0.287863   \n",
       "alphas[6]       0.217940  0.002135  0.116931     -0.000741      0.135297   \n",
       "alphas[7]       0.347727  0.004526  0.216380     -0.045205      0.204184   \n",
       "alphas[8]      -0.123617  0.005477  0.220396     -0.518081     -0.281584   \n",
       "alphas[9]      -0.022117  0.008111  0.328156     -0.692413     -0.234400   \n",
       "alphas[10]      0.218687  0.004461  0.244318     -0.242149      0.058077   \n",
       "alphas[11]      0.008869  0.004476  0.207651     -0.380493     -0.135531   \n",
       "stay[0]         1.286290  0.000753  0.041229      1.205779      1.258368   \n",
       "stay[1]         1.088134  0.000750  0.041099      1.007790      1.059463   \n",
       "stay[2]         0.584186  0.000653  0.035784      0.516296      0.559552   \n",
       "stay[3]         0.860108  0.000675  0.036970      0.786962      0.834661   \n",
       "stay[4]         0.857619  0.000771  0.042206      0.771756      0.828983   \n",
       "stay[5]         0.582339  0.000769  0.042102      0.500756      0.553389   \n",
       "stay[6]         1.138377  0.000905  0.049556      1.041427      1.104558   \n",
       "stay[7]         1.247312  0.001170  0.064088      1.123204      1.203998   \n",
       "stay[8]         0.846683  0.001241  0.067993      0.704295      0.801611   \n",
       "stay[9]         0.766284  0.001394  0.076368      0.616830      0.715681   \n",
       "stay[10]        1.473008  0.001533  0.083956      1.314673      1.415101   \n",
       "stay[11]        0.642435  0.001552  0.084990      0.476009      0.584796   \n",
       "lp__       -16051.968579  0.157836  5.018586 -16062.813482 -16055.103814   \n",
       "\n",
       "                     50%           75%         97.5%   n_eff      Rhat  \n",
       "betam           0.990275      1.058405      1.204639  3000.0  1.000840  \n",
       "alpham          0.283441      0.357970      0.521426  2173.0  0.999823  \n",
       "staym           0.945385      1.008465      1.163876  3000.0  0.999867  \n",
       "betasd          0.323315      0.381547      0.542784  3000.0  1.000864  \n",
       "alphasd         0.325971      0.416741      0.669016  1232.0  1.003453  \n",
       "staysd          0.328428      0.383207      0.550083  3000.0  1.001274  \n",
       "betas[0]        1.378630      1.432049      1.537987  3000.0  1.000165  \n",
       "betas[1]        1.177907      1.226203      1.331797  3000.0  0.999350  \n",
       "betas[2]        0.929059      0.972513      1.057845  3000.0  0.999029  \n",
       "betas[3]        1.098440      1.139053      1.223629  3000.0  0.999604  \n",
       "betas[4]        0.563651      0.610798      0.704753  3000.0  1.000897  \n",
       "betas[5]        0.582412      0.630098      0.715114  3000.0  0.999084  \n",
       "betas[6]        1.418100      1.487329      1.622447  3000.0  0.999615  \n",
       "betas[7]        1.111392      1.199705      1.358343  3000.0  0.999433  \n",
       "betas[8]        0.952385      1.045171      1.232408  3000.0  0.999676  \n",
       "betas[9]        0.672727      0.778691      1.011478  2193.0  1.000923  \n",
       "betas[10]       0.968891      1.065741      1.259182  3000.0  0.999559  \n",
       "betas[11]       1.039184      1.144083      1.350519  3000.0  0.999773  \n",
       "alphas[0]       0.331041      0.405564      0.556481  3000.0  1.000006  \n",
       "alphas[1]       0.308004      0.385863      0.547910  3000.0  0.999450  \n",
       "alphas[2]       0.344313      0.430415      0.625239  3000.0  0.999509  \n",
       "alphas[3]       0.729918      0.831182      1.065803  3000.0  1.000549  \n",
       "alphas[4]       0.522610      0.665781      1.012842  3000.0  1.001456  \n",
       "alphas[5]       0.417346      0.557378      0.874793  3000.0  1.000158  \n",
       "alphas[6]       0.215179      0.295095      0.454695  3000.0  1.000213  \n",
       "alphas[7]       0.327800      0.483695      0.821103  2286.0  1.000367  \n",
       "alphas[8]      -0.131120      0.016973      0.337861  1619.0  0.999767  \n",
       "alphas[9]      -0.000705      0.205762      0.569548  1637.0  1.000958  \n",
       "alphas[10]      0.218161      0.369450      0.725681  3000.0  0.999060  \n",
       "alphas[11]      0.004154      0.152844      0.427773  2152.0  1.000453  \n",
       "stay[0]         1.286300      1.313091      1.367689  3000.0  0.999884  \n",
       "stay[1]         1.088687      1.115770      1.170416  3000.0  0.999611  \n",
       "stay[2]         0.583430      0.607596      0.655401  3000.0  0.998944  \n",
       "stay[3]         0.860965      0.884563      0.931428  3000.0  0.999373  \n",
       "stay[4]         0.857818      0.886119      0.941352  3000.0  1.000088  \n",
       "stay[5]         0.582123      0.610419      0.664364  3000.0  0.999826  \n",
       "stay[6]         1.138832      1.172468      1.234085  3000.0  0.999702  \n",
       "stay[7]         1.246351      1.290635      1.376133  3000.0  0.998835  \n",
       "stay[8]         0.847420      0.892267      0.975140  3000.0  0.999537  \n",
       "stay[9]         0.765879      0.816835      0.913771  3000.0  0.999038  \n",
       "stay[10]        1.472399      1.529529      1.639523  3000.0  0.999342  \n",
       "stay[11]        0.641904      0.698949      0.803311  3000.0  0.998985  \n",
       "lp__       -16051.640716 -16048.506824 -16042.959519  1011.0  1.002175  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'NS' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-9d55c422ea4e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msns\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mJointGrid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Alphas\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"Betas\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdummy\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0ms\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mNS\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0msns\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkdeplot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0malphas\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0max\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0max_marg_x\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlegend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'NS' is not defined"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZ8AAAGfCAYAAACA4t+UAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAGx1JREFUeJzt3X9sVfX9x/FXoS2sMFpl1NvYdhWS\nyYwdiDMZ0lClhMqWS7Fjazsg0/BjbgudcyLUkS4wtlIpQ1eLwQBzOKREBK4FtAOy4koQmT+6qgOc\n0lmQdmAxci9l9Mf5/sHXq5UfXqTnfdvL85GQcO79NPd9bLxPzu3t50Y5juMIAABDfcI9AADg6kN8\nAADmiA8AwBzxAQCYIz4AAHPEBwBgjvgAAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcA\nYI74AADMER8AgDniAwAwR3wAAOaIDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADm\niA8AwBzxAQCYIz4AAHPEBwBgjvgAAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74\nAADMER8AgDniAwAwR3wAAOaIDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8A\nwBzxAQCYIz4AAHPEBwBgjvgAAMwRHwCAOeIDADAXHe4Bvqz29nY1NTWFewwACInH41F0dK99yu12\nvfa/RFNTk7KyssI9BgCEZNeuXUpOTg73GD1GlOM4TriH+DK48gHQm3Dl01WvjQ8AoPfiDQcAAHPE\nBwBgjvgAAMwRHwCAObP4OI6j+fPna/Xq1Re8v6amRl6vV9nZ2SosLJTf77caDQBgzCQ+7777rn78\n4x/rhRdeuOD9LS0tKioqUnl5uaqrq5WSkqKysjKL0QAAYWASn3Xr1ik3N1cTJ0684P21tbVKT09X\nWlqaJKmgoEBVVVXiXeAAEJlM4lNcXKzJkydf9P6mpiZ5PJ7gscfjkd/vVyAQuOjXtLe368iRI2pv\nb+/WWQEgXK6m57Ue8YaDzs7OC97ep8/Fx/tkex12OQAQKa6m57UeEZ+kpCQdP348eNzc3Kz4+HjF\nxcWFcSoAgFt6RHwyMjJUV1enhoYGSVJlZSWbhgJABAtbfOrr65WTkyNJGjx4sEpKSlRYWKiJEyfq\n0KFDmjdvXrhGAwC4zHSL1SVLlgT/np6eLp/PFzzOzMxUZmam5TgAgDDpES+7AQCuLsQHAGCO+AAA\nzBEfAIA54gMAMEd8AADmiA8AwBzxAQCYIz4AAHPEBwBgjvgAAMwRHwCAOeIDADBHfAAA5ogPAMAc\n8QEAmCM+AABzxAcAYI74AADMER8AgDniAwAwR3wAAOaIDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEf\nAIA54gMAMEd8AADmiA8AwBzxAQCYIz4AAHPEBwBgjvgAAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEA\nmCM+AABzxAcAYI74AADMER8AgDniAwAwR3wAAOaIDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA5\nk/jU1NTI6/UqOztbhYWF8vv9563ZsWOHvF6vcnJyNH36dL3//vsWowEAwsD1+LS0tKioqEjl5eWq\nrq5WSkqKysrKuqw5c+aM5s6dq8cff1w+n09ZWVlavHix26MBAMLE9fjU1tYqPT1daWlpkqSCggJV\nVVXJcZzgmo6ODjmOo1OnTkmSAoGA+vXr5/ZoAIAwiXb7AZqamuTxeILHHo9Hfr9fgUBAAwcOlCQN\nGDBACxcuVH5+vhISEtTZ2an169e7PRoAIExcv/Lp7Oy88AP3+fShDx48qIqKCm3fvl21tbW67777\nNGfOnC5XRwCAyOF6fJKSknT8+PHgcXNzs+Lj4xUXFxe8rba2VqNGjVJqaqokaerUqXrnnXd08uRJ\nt8cDAISB6/HJyMhQXV2dGhoaJEmVlZXKysrqsuamm27S/v37deLECUnSzp07lZycrGuvvdbt8QAA\nYeD6z3wGDx6skpISFRYWqq2tTampqSotLVV9fb0WLFggn8+n0aNHa8aMGZo+fbpiYmIUHx+vFStW\nuD0aACBMopxe+oOVI0eOKCsrS7t27VJycnK4xwGAK3Y1Pa+xwwEAwBzxAQCYIz4AAHPEBwBgjvgA\nAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74AADMER8AgDniAwAwR3wAAOaIDwDA\nHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8AwBzxAQCYIz4AAHPEBwBgjvgAAMwR\nHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74AADMER8AgDniAwAwR3wAAOaIDwDAHPEB\nAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8AwBzxAQCYIz4AAHPEBwBgziQ+NTU18nq9\nys7OVmFhofx+/3lrDh48qOnTp2vy5MnKzc3Vm2++aTEaACAMXI9PS0uLioqKVF5erurqaqWkpKis\nrKzLmtbWVs2YMUMzZ87Uli1b9LOf/UwPPvig26MBAMLE9fjU1tYqPT1daWlpkqSCggJVVVXJcZzg\nmj179iglJUWZmZmSpKysLD366KNujwYACBPX49PU1CSPxxM89ng88vv9CgQCwdsOHz6sIUOG6OGH\nH1Zubq7uvfdedXR0uD0aACBMXI9PZ2fnhR+4z6cP3d7ert27dysvL0+bNm3StGnTNHv2bJ09e9bt\n8QAAYeB6fJKSknT8+PHgcXNzs+Lj4xUXFxe8LTExUUOHDtWIESMkSePHj1dHR4caGxvdHg8AEAau\nxycjI0N1dXVqaGiQJFVWViorK6vLmrFjx+ro0aPBd7jt379fUVFRSk5Odns8AEAYRLv9AIMHD1ZJ\nSYkKCwvV1tam1NRUlZaWqr6+XgsWLJDP59OQIUNUUVGhhQsXqrW1VbGxsSovL1e/fv3cHg8AEAau\nx0eSMjMzg+9k+0RCQoJ8Pl/w+LbbbtOzzz5rMQ4AIMzY4QAAYI74AADMER8AgDniAwAwR3wAAOaI\nDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIC5kONz6tQpSec+e2fLli2qqalxayYAQIQLKT5btmwJ\nbgxaVlampUuXasGCBVq5cqWrwwEAIlNI8VmzZo0qKirU3t6ujRs3qqKiQpWVlXrmmWfcng8AEIFC\n+kiFpqYmjR49Wq+88opiYmI0cuRISZ++FAcAwOUIKT6JiYnau3evNm/erNtvv12SVFVVpZSUFFeH\nAwBEppDi88ADD6iwsFD9+vXT008/rb1796q4uFiPP/642/MBACJQSPEZN26c9u3bp6ioKEVFRSkp\nKUl///vfNXDgQLfnAwBEoJDi09HRob/+9a9qbm6W4ziSpLa2Nr333ntasmSJqwMCACJPSPF5+OGH\ntXv3biUkJKitrU0DBgzQoUOHNGnSJLfnAwBEoJDi87e//U0bN27UiRMn9Oc//1mPPfaYnn76ab3y\nyituzwcAiEAh/Z5PVFSUUlNTNWzYMP3rX/+SJOXn5+u1115zdTgAQGQKKT4pKSl64403FB8fr9On\nT+v48eMKBAI6c+aM2/MBACJQSC+7zZw5U/fee6+2bt2qKVOmKC8vT9HR0Ro7dqzb8wEAIlBI8bnr\nrruUnp6uxMRE3X///Ro2bJgCgYByc3Pdng8AEIFCetlt9uzZuv766xUTEyNJ8nq9ys/P1z333OPm\nbACACHXRK58jR47oqaeekiTt3btXixcv7nK/3+/Xe++95+pwAIDIdNH4JCcnKyYmRh999JEcx1Eg\nEOhyf2xsrJYvX+76gACAyHPJn/nMmzdPkjRs2DDNnDnTZCAAQOQL6Wc+M2fOVGNjox599FEVFRXp\n448/1ubNm92eDQAQoUKKz549e5Sbm6vGxkZVV1crEAho6dKlWr16tdvzAQAiUEjxKSsr02OPPaZl\ny5apb9++SkpK0urVq7Vu3Tq35wMARKCQ4tPY2KjRo0dLOrfVjiQNHz5cH3/8sXuTAQAiVkjxGTZs\nmF544YUut7300ksaOnSoK0MBACJbSDscPPTQQ5o1a5aee+45tba26v7779eePXtUUVHh9nwAgAgU\nUnxuvfVWbdu2TVu3blVqaqoSExP1q1/9SikpKW7PBwCIQCHFR5KSkpI0a9YsSVJLS4uuvfZa14YC\nAES2S/7Mp7OzU6tWrdIvf/lLbdu2TSdPntTkyZM1ZswYTZw4UQ0NDUZjAgAiySXjU1paqk2bNsnj\n8eiJJ57QjBkzdOONN+r555/X6NGjtWTJEqs5AQAR5JIvu23fvl3PPfecEhMTlZeXp4kTJ+qpp57S\noEGD9OCDD+rOO++0mhMAEEEueeXj9/uVmJgoSUpLS1P//v01aNAgSVJcXJza29vdnxAAEHEuGZ9P\nfqH0E598ns8nHMfp/okAABHvki+7OY6jt99+OxiZjo6OLsfEBwDwZVwyPq2tred9VPZnjz9/ZQQA\nQCguGZ8DBw5YzQEAuIqEtLcbAADdifgAAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmDOJT01Njbxe\nr7Kzs1VYWCi/33/RtTt37tSoUaMsxgIAhInr8WlpaVFRUZHKy8tVXV2tlJQUlZWVXXBtQ0ODSktL\n2bYHACKc6/Gpra1Venq60tLSJEkFBQWqqqo6LzCtra2aO3eu5s+f7/ZIAIAwcz0+TU1N8ng8wWOP\nxyO/369AINBlXXFxsfLy8nTjjTe6PRIAIMxcj09nZ+eFH7jPpw+9bt06RUdHa8qUKW6PAwDoAS65\nsWh3SEpKUl1dXfC4ublZ8fHxiouLC962efNmnTlzRjk5OWprawv+/cknn9R1113n9ogAAGOuxycj\nI0OlpaVqaGhQWlqaKisrlZWV1WXNxo0bg38/cuSIvF6vfD6f26MBAMLE9ZfdBg8erJKSEhUWFmri\nxIk6dOiQ5s2bp/r6euXk5Lj98ACAHsj1Kx9JyszMVGZmZpfbEhISLnh1k5ycrNdff91iLABAmLDD\nAQDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8AwBzxAQCYIz4AAHPEBwBgjvgA\nAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74AADMER8AgDniAwAwR3wAAOaIDwDA\nHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8AwBzxAQCYIz4AAHPEBwBgjvgAAMwR\nHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74AADMER8AgDniAwAwR3wAAOaIDwDAHPEB\nAJgjPgAAc8QHAGDOJD41NTXyer3Kzs5WYWGh/H7/eWt8Pp8mTZqknJwc5efnq76+3mI0AEAYuB6f\nlpYWFRUVqby8XNXV1UpJSVFZWVmXNe+9956WLl2qVatWyefz6ac//anmzJnj9mgAgDBxPT61tbVK\nT09XWlqaJKmgoEBVVVVyHCe4JjY2VosXL1ZiYqIk6eabb9aJEyd09uxZt8cDAIRBtNsP0NTUJI/H\nEzz2eDzy+/0KBAIaOHCgJCk5OVnJycmSJMdxVFJSonHjxik2Ntbt8QAAYeB6fDo7Oy94e58+5190\nnT59WvPnz1dTU5NWrVrl9mgAgDBx/WW3pKQkHT9+PHjc3Nys+Ph4xcXFdVn3wQcfKD8/X3379tXa\ntWs1aNAgt0cDAISJ6/HJyMhQXV2dGhoaJEmVlZXKysrqsuajjz7StGnTNGHCBC1fvlz9+/d3eywA\nQBi5/rLb4MGDVVJSosLCQrW1tSk1NVWlpaWqr6/XggUL5PP5tH79eh07dkw7duzQjh07gl/71FNP\n6ZprrnF7RACAsSjns28760WOHDmirKws7dq1K/hmBQDoza6m5zV2OAAAmCM+AABzxAcAYI74AADM\nER8AgDniAwAwR3wAAOaIDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8AwBzx\nAQCYIz4AAHPEBwBgjvgAAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74AADMER8A\ngDniAwAwR3wAAOaIDwDAHPEBAJgjPgAAc8QHAGCO+AAAzBEfAIA54gMAMEd8AADmiA8AwBzxAQCY\nIz4AAHPEBwBgjvgAAMwRHwCAOeIDADBHfAAA5ogPAMAc8QEAmCM+AABzxAcAYI74AADMmcSnpqZG\nXq9X2dnZKiwslN/v/1JrAACRwfX4tLS0qKioSOXl5aqurlZKSorKysouew0AIHK4Hp/a2lqlp6cr\nLS1NklRQUKCqqio5jnNZawAAkSPa7QdoamqSx+MJHns8Hvn9fgUCAQ0cODDkNZ/X0dER/FoA6Ok8\nHo+io11/yu01XP8v0dnZecHb+/Tpc1lrPu/48eOSpKlTp17BdABgY9euXUpOTg73GD2G6/FJSkpS\nXV1d8Li5uVnx8fGKi4u7rDWfd/PNN2vdunUaMmSI+vbt687wANBNPvvqzqXW7Nq1K6S1vZ3r8cnI\nyFBpaakaGhqUlpamyspKZWVlXfaaz+vfv7++/e1vuzk6AJiKjo6+aq6OohyDn+rv3r1by5YtU1tb\nm1JTU1VaWqrGxkYtWLBAPp/vomsSEhLcHg0AEAYm8QEA4LPY4QAAYI74AADM9fj4XG1b84RyLj6f\nT5MmTVJOTo7y8/NVX18fhkm7x+V873bu3KlRo0YZTtf9QjnfgwcPavr06Zo8ebJyc3P15ptvhmHS\n7hHK+e7YsUNer1c5OTmaPn263n///TBM2j0cx9H8+fO1evXqC94fSc9VV8zpwT788EPnO9/5jnP4\n8GHHcRznkUcecX7zm99c9preIpRzeffdd50xY8Y4zc3NjuM4Tk1NjZOZmWk7aDe5nO/d4cOHnfHj\nxzsjR460G7CbhXK+p0+fdsaMGePU1NQ4juM4O3bscLKzs40n7R6hnG9ra6szYsQIp6GhwXEcx/nT\nn/7kzJo1y3jS7vHvf//bmT59uvOtb33LWbVq1Xn3R9JzVXfo0Vc+V9vWPKGcS2xsrBYvXqzExERJ\n537f6cSJEzp79mw4Rr4ioX7vWltbNXfuXM2fPz8MU3afUM53z549SklJUWZmpiQpKytLjz76aDjG\nvWKhnG9HR4ccx9GpU6ckSYFAQP369QvHuFds3bp1ys3N1cSJEy94fyQ9V3WHHr3Xg1tb8/RUoZxL\ncnJy8PcAHMdRSUmJxo0bp9jY2LDMfCVC/d4VFxcrLy9PN954YzjG7DahnO/hw4c1ZMgQPfzwwzpw\n4IAGDRqkuXPnhmvkKxLK+Q4YMEALFy5Ufn6+EhIS1NnZqfXr14dr5CtSXFwsSXr55ZcveH8kPVd1\nhx595ePW1jw91eWcy+nTp/WLX/xC77//vhYvXuz2aK4I5XzXrVun6OhoTZkyxWos14Ryvu3t7dq9\ne7fy8vK0adMmTZs2TbNnz+6VV7ahnO/BgwdVUVGh7du3q7a2Vvfdd5/mzJkTkVcDkfRc1R169Fkn\nJSUF93CTLr41zxet6S1CPZcPPvhA+fn56tu3r9auXatBgwZZj9otQjnfzZs3q76+Xjk5OZo9e7bO\nnDmjnJwcNTc3h2PkKxLK+SYmJmro0KEaMWKEJGn8+PHq6OhQY2Oj+bxXKpTzra2t1ahRo5Samirp\n3F6N77zzjk6ePGk+r9si6bmqO/To+GRkZKiurk4NDQ2SdNGteb5oTW8Ryrl89NFHmjZtmiZMmKDl\ny5erf//+YZi0e4Ryvhs3btTWrVvl8/n05JNPqn///vL5fLruuuvCMPGVCeV8x44dq6NHjwbf4bZ/\n/35FRUX1yi1XQjnfm266Sfv379eJEycknXtHY3Jysq699lrrcV0XSc9V3SJc73QIVU1NjeP1ep27\n7rrLmT17tnPy5Ennn//8pzNp0qRLrumtvuh8V6xY4QwfPtyZNGlSlz8tLS1hnvzLCeX7+4nGxsZe\n/W43xwntfF955RVnypQpzve+9z3n7rvvdvbv3x/Gia9MKOf7l7/8xbnrrrscr9frTJs2zTl06FAY\nJ75y8+bNC77bLZKfq64U2+sAAMz16JfdAACRifgAAMwRHwCAOeIDADBHfAAA5ogPIsqiRYs0fPhw\nHThwIHjbvn37dMstt4T09ePGjdOLL77o1ngA/h/xQcQ4ffq0tm7dqtzcXK1duzbc4wC4BOKDiLFt\n2zYNGzZMP/nJT7Rt2za1tLSct2bfvn3Kzs7W73//e40aNUp33nmnNm7c2GXNq6++qrvvvlu33HKL\npk6dqmPHjkmS/ve//+m3v/2tJkyYoJEjR2rcuHF6/vnnJZ3bt2vx4sW6/fbbNXr0aM2YMUP/+c9/\n3D9poJciPogYlZWV+uEPf6ivf/3ruvXWW7Vhw4YLrmtoaNCZM2e0d+9ePfLII1q0aJFeffXV4P3/\n+Mc/tHLlStXW1qqjo0NPPPGEJGnNmjV666239Oyzz+q1117TPffco4ULF6qtrU07duzQ3r179eKL\nL+qll15SYmKi/vjHP5qcN9AbER9EhLfeektHjhzRd7/7XUnSj370I61fv17t7e3nrY2NjdX8+fPV\nr18/3XbbbZowYYK2bt0avL+goECJiYkaMGCA7rjjjuAnaxYUFGjFihX66le/qmPHjukrX/lKcEv8\nfv366dixY3ruued09OhR/e53v9OyZctsTh7ohXr05/kAoaqsrFQgENC4ceMknfusow8//FDV1dX6\n2te+1mXtkCFDuuwk7PF4urxE9tldwmNiYtTR0SFJ8vv9WrRokd544w0lJyfrhhtukHTuJbc77rhD\nxcXF2rBhg5YtW6brr79eDz300NW9cSRwCcQHvZ7f79fWrVu1cuVKfeMb3wjevmbNGq1du1YPPPBA\nl/UnT55UW1ubYmJiJJ37iIrPfsjXxRQXFyslJUUVFRWKiYnRW2+9Fbxiamxs1De/+U2tX79efr9f\nzzzzjO6//369+uqrvfKD/gC38bIber3nn39eiYmJGjNmjIYMGRL8k5eXpzfeeEP19fVd1p8+fVoV\nFRVqa2vT3r17tWvXLuXk5Hzh4/j9fvXv3199+/bVf//7X/3hD3+QdO4D4F5++WX9/Oc/19GjRzVg\nwAANGjRIAwcOVHQ0/74DLoT/M9DrbdiwQV6v97zb09LSNGrUqPPedv3Jz2oyMjKUkJCgJUuW6Oab\nb/7Cx/n1r3+tBQsW6NZbb9U111yjH/zgB3r77bf17rvv6vvf/74OHz6svLw8BQIB3XDDDaqoqLhq\nP6US+CJ8pAKuKvv27dN9992n119/PdyjAFc1/lkGADBHfAAA5njZDQBgjisfAIA54gMAMEd8AADm\niA8AwBzxAQCYIz4AAHP/B8BoSUsDraBAAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x432 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dummy = pd.DataFrame([[0, 0, 0]], columns = ('Alphas', 'Betas', 'Stay Intercept'))\n",
    "alphas = extract['alphas']\n",
    "betas = extract['betas']\n",
    "g = sns.JointGrid(\"Alphas\", \"Betas\", dummy) \n",
    "\n",
    "for s in np.arange(NS):\n",
    "    \n",
    "    sns.kdeplot(alphas[:,s], ax=g.ax_marg_x, legend=False)\n",
    "    sns.kdeplot(betas[:,s], ax=g.ax_marg_y, vertical=True, legend = False) \n",
    "    g.ax_joint.plot(alphas[:,s], betas[:,s], 'o', ms=5, alpha = 0.5, label = 'Mouse {:d}'.format(s))\n",
    "    \n",
    "\n",
    "g.ax_joint.legend(bbox_to_anchor=(1.2, 1), loc=2, borderaxespad=0.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy = pd.DataFrame([[0, 0, 0]], columns = ('Alphas', 'Betas', 'Stay'))\n",
    "alphas = extract['stay']\n",
    "betas = extract['betas']\n",
    "g = sns.JointGrid(\"Stay\", \"Betas\", dummy) \n",
    "\n",
    "for s in np.arange(NS):\n",
    "    \n",
    "    sns.kdeplot(alphas[:,s], ax=g.ax_marg_x, legend=False)\n",
    "    sns.kdeplot(betas[:,s], ax=g.ax_marg_y, vertical=True, legend = False) \n",
    "    g.ax_joint.plot(alphas[:,s], betas[:,s], 'o', ms=5, alpha = 0.5, label = 'Mouse {:d}'.format(s))\n",
    "    \n",
    "\n",
    "g.ax_joint.legend(bbox_to_anchor=(1.2, 1), loc=2, borderaxespad=0.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'NS' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-ec3743613071>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0ms\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mNS\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m     \u001b[0mplt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mextract\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'stay'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m100\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0malpha\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.8\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m'Mouse {:d}'\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mplt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlegend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbbox_to_anchor\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1.05\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloc\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mborderaxespad\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mplt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtitle\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Stay Intercept Distributions for Each Subject'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'NS' is not defined"
     ]
    }
   ],
   "source": [
    "for s in np.arange(NS):\n",
    "    plt.hist(extract['stay'][:,s], 100, alpha = 0.8, label = 'Mouse {:d}'.format(s))\n",
    "\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)\n",
    "plt.title('Stay Intercept Distributions for Each Subject')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 37.6 s, sys: 6.99 s, total: 44.6 s\n",
      "Wall time: 1min 4s\n"
     ]
    }
   ],
   "source": [
    "%%time \n",
    "# get the predictions: \n",
    "N_samples = np.shape(extract['betas'])[0]\n",
    "\n",
    "\n",
    "Q_i = np.zeros((N_samples, NS, (NT + 1)), float)\n",
    "Q_c = np.zeros((N_samples, NS, (NT + 1)), float)\n",
    "delta = np.zeros((N_samples, NS, NT), float)\n",
    "\n",
    "for s in np.arange(NS):\n",
    "\n",
    "    for t in np.arange(extract['NT_all'][s]):\n",
    "\n",
    "        r_i = extract['r'][s, t]\n",
    "        c_i = extract['c'][s, t]\n",
    "\n",
    "        if c_i == 0: # contra choice\n",
    "            Q_choice = Q_c; \n",
    "            Q_other = Q_i; \n",
    "        elif c_i == 1: # ipsa choice\n",
    "            Q_choice = Q_i; \n",
    "            Q_other = Q_c; \n",
    "            \n",
    "        delta[:, s, t] = r_i - Q_choice[:, s, t] \n",
    "        \n",
    "        \n",
    "        # Q of the choice is updated, the other one is not. \n",
    "        Q_choice[:, s, t + 1] = Q_choice[:, s, t] + phi_approx(extract['alphas'][:,s]) * delta[:, s, t]\n",
    "        Q_other[:, s, t + 1] = Q_other[: ,s, t]\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "_Q_i = np.mean(Q_i, 0)\n",
    "_Q_c = np.mean(Q_c, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 11.4 s, sys: 13 s, total: 24.4 s\n",
      "Wall time: 29.5 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "prob_ipsa = np.zeros((N_samples, NS, NT))\n",
    "prob_contra = np.zeros((N_samples, NS, NT))\n",
    "\n",
    "\n",
    "pc = np.zeros((NS, NT))\n",
    "pc[:, 1:NT_all[s]] = c[:, :(NT_all[s] - 1)] * 2 - 1\n",
    "\n",
    "for s in np.arange(NS):\n",
    "\n",
    "    betas_i = extract['betas'][:,s]\n",
    "    stay_i = extract['stay'][:,s]\n",
    "    \n",
    "    intercept_pc = stay_i[:, np.newaxis] * pc[np.newaxis, s, :NT_all[s]]\n",
    "    \n",
    "    ipsa = betas_i[:, np.newaxis]  * Q_i[:, s, :NT_all[s]] \n",
    "    contra = betas_i[:, np.newaxis]  * Q_c[:, s, :NT_all[s]] \n",
    "    prob_ipsa[:, s, :NT_all[s]] = 1 / (1 + np.exp(contra - ipsa + intercept_pc) )\n",
    "    prob_contra[:, s, :NT_all[s]] = 1 / (1 + np.exp(ipsa - contra + intercept_pc) )\n",
    "\n",
    "\n",
    "_prob_ipsa = np.mean(prob_ipsa, axis = 0)\n",
    "_prob_contra = np.mean(prob_contra, axis = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "_prob_choice = np.zeros((NS, NT))\n",
    "\n",
    "_prob_choice[c == 0] = _prob_contra[c == 0]\n",
    "_prob_choice[c == 1] = _prob_ipsa[c == 1]\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculating the pseudo R^2, as defined as: \n",
    "# psuedo_r = 1 - log(L_c) / log(L_null)\n",
    "# for L_null, we will define as: sum(log(total_ipsi/total_trials)) + sum(log(total_contra/total_trials))\n",
    "stayed = (c[:, :-1] == c[:, 1:]).astype(int)\n",
    "\n",
    "psuedo_r = []\n",
    "for s in np.arange(NS):\n",
    "\n",
    "    num_stayed = sum(stayed[s, :(NT_all[s] - 1)] == 1)\n",
    "    num_left = sum(stayed[s, :(NT_all[s] - 1)] == 0)\n",
    "    log_L_null = np.log(0.5) +  (np.log(num_stayed/(num_left + num_stayed)) * num_stayed) + (np.log(num_left/(num_left + num_stayed)) * num_left)\n",
    "\n",
    "    log_L_c = sum(np.log(_prob_choice[s, :NT_all[s]]))\n",
    "    \n",
    "    if log_L_null == 0: \n",
    "        import pdb; pdb.set_trace()\n",
    "\n",
    "    psuedo_r.append(1 - (log_L_c / log_L_null))\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-0.44077625054711489,\n",
       " -0.38011248509652917,\n",
       " -0.053432797497120266,\n",
       " -0.0556136175509232,\n",
       " -0.14217907020695342,\n",
       " -0.068337334067650968,\n",
       " -0.47417766477859424,\n",
       " -0.56423908219890451,\n",
       " -0.11518552071363897,\n",
       " -0.20398542490310745,\n",
       " -0.66808357770355831,\n",
       " -0.14029069600393518]"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psuedo_r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4525"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NT_all[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-2972.3149498006478"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "log_L_c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psuedo_r_old = [0.23240674515566684,\n",
    " 0.16206812812064064,\n",
    " 0.088361106266638845,\n",
    " 0.14987018108276606,\n",
    " 0.041754443245826112,\n",
    " 0.036427641632915453,\n",
    " 0.21168087024654814,\n",
    " 0.16402156651479727,\n",
    " 0.086523879833013795,\n",
    " 0.049997310358773972,\n",
    " 0.14134013612596175,\n",
    " 0.10030866571132502]\n",
    "\n",
    "plt.bar(np.arange(NS), psuedo_r, label = 'with intercept \\nfor previous choice', alpha = 0.8)\n",
    "\n",
    "plt.bar(np.arange(NS), psuedo_r_old, label = 'original')\n",
    "\n",
    "plt.ylabel('psuedo_r')\n",
    "plt.xlabel('mice')\n",
    "\n",
    "plt.legend(bbox_to_anchor=(1.2, 1), loc=2, borderaxespad=0.)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## debug code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mini = {'NS':3, 'NT':4, 'NT_all':[1, 2, 4], \n",
    "                 'r':r[:3, :4], 'c':c[:3,:4]}\n",
    "debug = sm.sampling(data=mini, iter=2, warmup=1, chains=1)\n",
    "extract = debug.extract()\n",
    "for k, v in mini.items(): extract[k] = v\n",
    "\n",
    "extract['prob_c']\n",
    "extract['c_hat']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:anaconda3]",
   "language": "python",
   "name": "conda-env-anaconda3-py"
  },
  "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.6.5"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}