bdanskin_analysis_all.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aefcfb32-f7d0-4f00-8956-8304ed217bcd",
   "metadata": {},
   "source": [
    "# All analysis code for B. Danskin et al. 2023\n",
    "Clears and reloads between each analysis type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "122781cc-81de-49f2-95fd-0a0c1912b80b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "511c9409-f000-4c2b-9f02-af2e8d463f88",
   "metadata": {},
   "source": [
    "# Behavior analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c42ad73-cae8-4e27-b997-91b98de6f4a1",
   "metadata": {},
   "source": [
    "## General import functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "40e2835e-9553-4d45-b5bb-ceccf2a79daf",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -f\n",
    "import sys\n",
    "from os.path import dirname, join as pjoin\n",
    "from os import listdir\n",
    "sys.path.append('C:\\\\jupyter_notebooks\\\\Danskin_SciAdv_2023\\\\py_code') # set local directory\n",
    "import numpy as np\n",
    "import pickle\n",
    "import random\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "\n",
    "import h5py\n",
    "import scipy.io as sio\n",
    "from scipy import stats\n",
    "import statsmodels as sm # import statsmodels.api as sm\n",
    "from scipy.optimize import minimize, basinhopping, curve_fit\n",
    "from sklearn.model_selection import KFold\n",
    "from tqdm.notebook import tqdm, trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ed1e7523-7d7d-40cb-acec-c68e25e33ba3",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Directories\n",
    "%autoreload\n",
    "import bdanskin as BD\n",
    "\n",
    "project_dir = 'C:\\\\jupyter_notebooks\\\\Danskin_SciAdv_2023' # local directory\n",
    "behavior_dir = pjoin(project_dir, 'hattori_datasets_behavior')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a40da41a-7861-4676-bc15-d1374981edb1",
   "metadata": {},
   "source": [
    "## Logistic regresssion analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b5e9e28-514b-4e36-93cd-8fa3933c66a4",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44c3e90c-0a27-47d4-a769-b08ef7acd80e",
   "metadata": {},
   "source": [
    "#### pandas apply fit logistic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ed22ccad-9c98-4232-8be7-dd82c16d9303",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegressionCV\n",
    "from sklearn.linear_model import LogisticRegression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "17848a12-fdca-4e9e-8695-ddee7c31ec4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def modelfit_logistic_pd(row, reg_type):\n",
    "    n_back  = 10\n",
    "    a = row['a']\n",
    "    R = row['R']\n",
    "    n_choice = sum((a==1)|(a==2))\n",
    "    \n",
    "    # Fit RUC\n",
    "    glm_mat, glm_target = BD.logistic_prepare_predictors(a,R,n_back,'RUC')\n",
    "    if reg_type=='none':\n",
    "        log_reg = LogisticRegression(solver='newton-cg', penalty='none', n_jobs=-1, random_state=0)\n",
    "    else: \n",
    "        log_reg = LogisticRegressionCV(solver='saga', cv=5, penalty=reg_type, n_jobs=-1, random_state=0, refit=True)\n",
    "    fit_mdl = log_reg.fit(glm_mat,glm_target)  \n",
    "    coef = fit_mdl.coef_[0]\n",
    "    intercept = fit_mdl.intercept_[0]\n",
    "    loglik,AIC,p_logit = BD.logistic_calc_loglik(glm_mat, glm_target,coef, intercept, n_back)\n",
    "    raw_weights = np.append(coef,intercept)\n",
    "    \n",
    "    row['mdl_hist'] = n_back\n",
    "    row['mdl_penalty'] = reg_type\n",
    "    row['RUC_weights'] = raw_weights\n",
    "    row['RUC_AIC'] = AIC\n",
    "    row['RUC_loglik'] = loglik\n",
    "    row['RUC_norm_loglik'] = loglik/n_choice\n",
    "    \n",
    "    # Fit RC\n",
    "    glm_mat, glm_target = BD.logistic_prepare_predictors(a,R,n_back,'RC')\n",
    "    if reg_type=='none':\n",
    "        log_reg = LogisticRegression(solver='newton-cg', penalty='none', n_jobs=-1, random_state=0)\n",
    "    else: \n",
    "        log_reg = LogisticRegressionCV(solver='saga', cv=5, penalty=reg_type, n_jobs=-1, random_state=0, refit=True)\n",
    "    fit_mdl = log_reg.fit(glm_mat,glm_target)  \n",
    "    coef = fit_mdl.coef_[0]\n",
    "    intercept = fit_mdl.intercept_[0]\n",
    "    loglik,AIC,p_logit = BD.logistic_calc_loglik(glm_mat, glm_target,coef, intercept, n_back)\n",
    "    raw_weights = np.append(coef,intercept)\n",
    "    \n",
    "    row['RC_weights'] = raw_weights\n",
    "    row['RC_AIC'] = AIC\n",
    "    row['RC_loglik'] = loglik\n",
    "    row['RC_norm_loglik'] = loglik/n_choice\n",
    "    \n",
    "    # Fit RU\n",
    "    glm_mat, glm_target = BD.logistic_prepare_predictors(a,R,n_back,'RU')\n",
    "    if reg_type=='none':\n",
    "        log_reg = LogisticRegression(solver='newton-cg', penalty='none', n_jobs=-1, random_state=0)\n",
    "    else: \n",
    "        log_reg = LogisticRegressionCV(solver='saga', cv=5, penalty=reg_type, n_jobs=-1, random_state=0, refit=True)\n",
    "    fit_mdl = log_reg.fit(glm_mat,glm_target)  \n",
    "    coef = fit_mdl.coef_[0]\n",
    "    intercept = fit_mdl.intercept_[0]\n",
    "    loglik,AIC,p_logit = BD.logistic_calc_loglik(glm_mat, glm_target,coef, intercept, n_back)\n",
    "    raw_weights = np.append(coef,intercept)\n",
    "    \n",
    "    row['RU_weights'] = raw_weights\n",
    "    row['RU_AIC'] = AIC\n",
    "    row['RU_loglik'] = loglik\n",
    "    row['RU_norm_loglik'] = loglik/n_choice\n",
    "    \n",
    "    # Fit R\n",
    "    glm_mat, glm_target = BD.logistic_prepare_predictors(a,R,n_back,'R')\n",
    "    if reg_type=='none':\n",
    "        log_reg = LogisticRegression(solver='newton-cg', penalty='none', n_jobs=-1, random_state=0)\n",
    "    else: \n",
    "        log_reg = LogisticRegressionCV(solver='saga', cv=5, penalty=reg_type, n_jobs=-1, random_state=0, refit=True)\n",
    "    fit_mdl = log_reg.fit(glm_mat,glm_target)  \n",
    "    coef = fit_mdl.coef_[0]\n",
    "    intercept = fit_mdl.intercept_[0]\n",
    "    loglik,AIC,p_logit = BD.logistic_calc_loglik(glm_mat, glm_target,coef, intercept, n_back)\n",
    "    raw_weights = np.append(coef,intercept)\n",
    "    \n",
    "    row['R_weights'] = raw_weights\n",
    "    row['R_AIC'] = AIC\n",
    "    row['R_loglik'] = loglik\n",
    "    row['R_norm_loglik'] = loglik/n_choice\n",
    "\n",
    "    return row"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48cbdf06-429d-4996-a9bb-b38ec80c344f",
   "metadata": {},
   "source": [
    "### Load Imaging and simulation datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7885a455-9ef4-40d7-ae10-6aab31973437",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data\n",
    "imaging_beh_df = pd.read_pickle(pjoin(behavior_dir,'all_behavior_df.pkl'))\n",
    "\n",
    "# 10000 trial simulation\n",
    "sim_beh_df = pd.read_pickle(pjoin(behavior_dir, 'rl_simulation_full_base_all_imaging_10000.pkl'))\n",
    "sim_beh_df = sim_beh_df[['Area','Mouse','Date','RL_Model','a','R']]\n",
    "\n",
    "full_sim = sim_beh_df.loc[sim_beh_df['RL_Model']=='rl_full'].reset_index(drop=True)\n",
    "base_sim = sim_beh_df.loc[sim_beh_df['RL_Model']=='rl_base'].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b72cf5e-e32c-4934-98e3-75d73bad299e",
   "metadata": {},
   "source": [
    "### Fit logistic weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "8e86cc23-79fa-46b8-80c3-709db744d999",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0b766479e444c2c8f903854961f8616",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "# Fit imaging data\n",
    "beh_df = imaging_beh_df.iloc[:2]\n",
    "reg_df = beh_df.progress_apply(modelfit_logistic_pd,axis=1, args=['none'])\n",
    "reg_df = reg_df.drop(columns=['a','R'])\n",
    "reg_df.index.rename('session_id', inplace=True)\n",
    "reg_df.to_pickle(pjoin(behavior_dir, 'logstic_imaging_10_hist.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "49fa2e85-8a0d-4c4c-891c-fd6028d83ddf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee8fe3e78ca74149a9a76f2c2842b5c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "888b15b0281448f5b10618d93705d45a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "# Fit simulation data\n",
    "beh_df = full_sim\n",
    "reg_df = beh_df.progress_apply(modelfit_logistic_pd,axis=1, args=['none'])\n",
    "reg_df = reg_df.drop(columns=['a','R'])\n",
    "reg_df.index.rename('session_id', inplace=True)\n",
    "reg_df.to_pickle(pjoin(behavior_dir, 'logstic_full_rl_10_hist.pkl'))\n",
    "\n",
    "beh_df = base_sim\n",
    "reg_df = beh_df.progress_apply(modelfit_logistic_pd,axis=1, args=['none'])\n",
    "reg_df = reg_df.drop(columns=['a','R'])\n",
    "reg_df.index.rename('session_id', inplace=True)\n",
    "reg_df.to_pickle(pjoin(behavior_dir, 'logstic_base_rl_10_hist.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d21b0710-e736-460e-b141-cc8b91d53293",
   "metadata": {},
   "source": [
    "## Behlogits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe74f36a-cf3f-4cf9-8ad5-5963e4a60f7c",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b7e4484-0e9c-4b27-b66e-a921fb15a1c4",
   "metadata": {},
   "source": [
    "#### pandas apply model fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "75e26624-a2df-4a73-be0d-1c10d75fc19a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def modelfit_behlogit(row, n_back, mdl_type):\n",
    "    rng = np.random.default_rng(220727)\n",
    "    beh_func = eval('BD.logit_'+mdl_type)\n",
    "    # Prepare history matrices\n",
    "    a,R = row[['a','R']]\n",
    "    a_hist,R_hist,uR_hist = BD.prepare_hist_matrix(a,R,n_back) # New code to simplify fits\n",
    "    n_choice_trials = sum((a==1)|(a==2))\n",
    "    \n",
    "    # given different model types\n",
    "    if (mdl_type=='exp_r') | (mdl_type=='hyp_r'):\n",
    "        n_params = 3\n",
    "        initial_params = [1,1,0]\n",
    "        bounds = [[0, np.inf],[-np.inf,np.inf],[-np.inf, np.inf]]\n",
    "    elif (mdl_type=='exp_ru') | (mdl_type=='hyp_ru') | (mdl_type=='hyp_r_exp_u') | (mdl_type=='exp_r_hyp_u'):\n",
    "        n_params = 5\n",
    "        initial_params = [1,1,-1,1,0]\n",
    "        bounds = [[0, np.inf],[-np.inf,np.inf],[-np.inf, np.inf],[-np.inf, np.inf],[-np.inf, np.inf]]\n",
    "    elif (mdl_type=='exp_rc') | (mdl_type=='hyp_rc') | (mdl_type=='hyp_r_exp_c'):\n",
    "        n_params = 5\n",
    "        initial_params = [1,1,-1,1,0]\n",
    "        bounds = [[0, np.inf],[-np.inf,np.inf],[-np.inf, np.inf],[-np.inf, np.inf],[-np.inf, np.inf]]\n",
    "    elif (mdl_type=='exp_ruc') | (mdl_type=='hyp_ruc')| (mdl_type=='hyp_r_exp_uc'):\n",
    "        n_params = 7\n",
    "        initial_params = [1,1,-1,1,1,1,0]\n",
    "        bounds = [[0, np.inf],[-np.inf,np.inf],[-np.inf,0],[-np.inf,np.inf],\n",
    "                  [0, np.inf],[-np.inf,np.inf],[-np.inf, np.inf]]\n",
    "    else:\n",
    "        print('unable to parse mdl_type')\n",
    "\n",
    "    # Run k-fold cross validation\n",
    "    cv_iter = 0\n",
    "    loglik_vec = np.zeros(10)\n",
    "    norm_loglik_vec = np.zeros(10)\n",
    "    aic_vec = np.zeros(10)\n",
    "    pa_vec = np.zeros(10)\n",
    "    kf = KFold(n_splits=10, shuffle=False)#, random_state=0)\n",
    "    kf.get_n_splits(np.arange(0,len(a)))\n",
    "    for train_index, test_index in kf.split(np.arange(0,len(a))):\n",
    "        lik_model = minimize(beh_func, initial_params, bounds=bounds, method='L-BFGS-B',\n",
    "                         args=(a[train_index], a_hist[train_index,:], R_hist[train_index,:],\n",
    "                               uR_hist[train_index,:], 1))\n",
    "        temp_params = lik_model.x\n",
    "        p_logit,nloglik = beh_func(temp_params, a[test_index],a_hist[test_index,:], \n",
    "                                   R_hist[test_index,:], uR_hist[test_index,:], 0)\n",
    "        pa_vec[cv_iter] = BD.PA_logit(a[test_index],p_logit)\n",
    "        aic_vec[cv_iter] = 2*n_params + 2*nloglik\n",
    "        loglik_vec[cv_iter] = -nloglik\n",
    "        norm_loglik_vec[cv_iter] = -nloglik/len(test_index)\n",
    "        cv_iter+=1\n",
    "        \n",
    "    # Fit on full session\n",
    "    minimizer_kwargs = {\"method\":\"L-BFGS-B\",\"bounds\":bounds, \n",
    "                        \"args\":(a, a_hist, R_hist, uR_hist, 1)}\n",
    "    with np.errstate(divide='ignore',invalid='ignore'):\n",
    "        lik_model = basinhopping(beh_func, initial_params, minimizer_kwargs=minimizer_kwargs, niter=10)\n",
    "\n",
    "    temp_params = lik_model.x\n",
    "    p_logit,nloglik = beh_func(temp_params,a,a_hist,R_hist,uR_hist, 0)\n",
    "    full_PA = BD.PA_logit(a,p_logit)\n",
    "    full_AIC = 2*n_params + 2*nloglik \n",
    "\n",
    "    row['mdl_type'] = mdl_type\n",
    "    row['mdl_hist'] = n_back\n",
    "    row['full_session_AIC'] = full_AIC\n",
    "    row['full_session_pa'] = full_PA\n",
    "    row['full_session_loglik'] = -nloglik\n",
    "    row['full_session_norm_loglik'] = -nloglik/n_choice_trials\n",
    "    row['full_session_betaR'] = temp_params[0]\n",
    "    row['full_session_tauR'] = np.exp(temp_params[1])\n",
    "\n",
    "    if (mdl_type=='exp_r') | (mdl_type=='hyp_r'):\n",
    "        row['full_session_betaU'] = np.nan\n",
    "        row['full_session_tauU'] = np.nan\n",
    "        row['full_session_betaC'] = np.nan\n",
    "        row['full_session_tauC'] = np.nan\n",
    "        row['full_session_const'] = temp_params[2]\n",
    "    elif (mdl_type=='exp_ru') | (mdl_type=='hyp_ru') | (mdl_type=='hyp_r_exp_u') | (mdl_type=='exp_r_hyp_u'):\n",
    "        row['full_session_betaU'] = temp_params[2]\n",
    "        row['full_session_tauU'] = np.exp(temp_params[3])\n",
    "        row['full_session_betaC'] = np.nan\n",
    "        row['full_session_tauC'] = np.nan\n",
    "        row['full_session_const'] = temp_params[4]\n",
    "    elif (mdl_type=='exp_rc') | (mdl_type=='hyp_rc') | (mdl_type=='hyp_r_exp_c'):\n",
    "        row['full_session_betaU'] = np.nan\n",
    "        row['full_session_tauU'] = np.nan\n",
    "        row['full_session_betaC'] = temp_params[2]\n",
    "        row['full_session_tauC'] = np.exp(temp_params[3])\n",
    "        row['full_session_const'] = temp_params[4]\n",
    "    elif (mdl_type=='exp_ruc') | (mdl_type=='hyp_ruc') | (mdl_type=='hyp_r_exp_uc'):\n",
    "        row['full_session_betaU'] = temp_params[2]\n",
    "        row['full_session_tauU'] = np.exp(temp_params[3])\n",
    "        row['full_session_betaC'] = temp_params[4]\n",
    "        row['full_session_tauC'] = np.exp(temp_params[5])\n",
    "        row['full_session_const'] = temp_params[6]\n",
    "\n",
    "    row['CV_mean_loglik'] = np.nanmean(loglik_vec)\n",
    "    row['CV_mean_norm_loglik'] = np.nanmean(norm_loglik_vec)\n",
    "    row['CV_mean_aic'] = np.nanmean(aic_vec)\n",
    "    row['CV_mean_pa'] = np.nanmean(pa_vec)\n",
    "    \n",
    "    return row"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c285db1c-b553-41c0-8190-ddb4ccb65eab",
   "metadata": {},
   "source": [
    "#### pandas apply constrained model fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d12f5623-9cee-479b-9dd1-195fd96152fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def modelfit_behlogit_constrained_ru(row, n_back, mdl_type):\n",
    "    rng = np.random.default_rng(220727)\n",
    "    beh_func = eval('BD.logit_'+mdl_type)\n",
    "    # Prepare history matrices\n",
    "    a,R = row[['a','R']]\n",
    "    a_hist,R_hist,uR_hist = BD.prepare_hist_matrix(a,R,n_back) # New code to simplify fits\n",
    "    n_choice_trials = sum((a==1)|(a==2))\n",
    "    \n",
    "    # given different model types\n",
    "    if (mdl_type=='exp_ru') | (mdl_type=='hyp_ru') | (mdl_type=='hyp_r_exp_u') | (mdl_type=='exp_r_hyp_u'):\n",
    "        n_params = 5\n",
    "        initial_params = [1,1,-1,1,0]\n",
    "        bounds = [[0, np.inf],[0,np.inf],[-np.inf, 0],[0, np.inf],[-np.inf, np.inf]]\n",
    "    else:\n",
    "        print('wrong model type provided')\n",
    "\n",
    "    # Run k-fold cross validation\n",
    "    cv_iter = 0\n",
    "    loglik_vec = np.zeros(10)\n",
    "    norm_loglik_vec = np.zeros(10)\n",
    "    aic_vec = np.zeros(10)\n",
    "    pa_vec = np.zeros(10)\n",
    "    kf = KFold(n_splits=10, shuffle=False)#, random_state=0)\n",
    "    kf.get_n_splits(np.arange(0,len(a)))\n",
    "    for train_index, test_index in kf.split(np.arange(0,len(a))):\n",
    "        lik_model = minimize(beh_func, initial_params, bounds=bounds, method='L-BFGS-B',\n",
    "                         args=(a[train_index], a_hist[train_index,:], R_hist[train_index,:],\n",
    "                               uR_hist[train_index,:], 1))\n",
    "        temp_params = lik_model.x\n",
    "        p_logit,nloglik = beh_func(temp_params, a[test_index],a_hist[test_index,:], \n",
    "                                   R_hist[test_index,:], uR_hist[test_index,:], 0)\n",
    "        pa_vec[cv_iter] = BD.PA_logit(a[test_index],p_logit)\n",
    "        aic_vec[cv_iter] = 2*n_params + 2*nloglik\n",
    "        loglik_vec[cv_iter] = -nloglik\n",
    "        norm_loglik_vec[cv_iter] = -nloglik/len(test_index)\n",
    "        cv_iter+=1\n",
    "        \n",
    "    # Fit on full session\n",
    "    minimizer_kwargs = {\"method\":\"L-BFGS-B\",\"bounds\":bounds, \n",
    "                        \"args\":(a, a_hist, R_hist, uR_hist, 1)}\n",
    "    with np.errstate(divide='ignore',invalid='ignore'):\n",
    "        lik_model = basinhopping(beh_func, initial_params, minimizer_kwargs=minimizer_kwargs, niter=10)\n",
    "\n",
    "    temp_params = lik_model.x\n",
    "    p_logit,nloglik = beh_func(temp_params,a,a_hist,R_hist,uR_hist, 0)\n",
    "    full_PA = BD.PA_logit(a,p_logit)\n",
    "    full_AIC = 2*n_params + 2*nloglik \n",
    "\n",
    "    row['mdl_type'] = mdl_type\n",
    "    row['mdl_hist'] = n_back\n",
    "    row['full_session_AIC'] = full_AIC\n",
    "    row['full_session_pa'] = full_PA\n",
    "    row['full_session_loglik'] = -nloglik\n",
    "    row['full_session_norm_loglik'] = -nloglik/n_choice_trials\n",
    "    row['full_session_betaR'] = temp_params[0]\n",
    "    row['full_session_tauR'] = np.exp(temp_params[1])\n",
    "\n",
    "    if (mdl_type=='exp_r') | (mdl_type=='hyp_r'):\n",
    "        row['full_session_betaU'] = np.nan\n",
    "        row['full_session_tauU'] = np.nan\n",
    "        row['full_session_betaC'] = np.nan\n",
    "        row['full_session_tauC'] = np.nan\n",
    "        row['full_session_const'] = temp_params[2]\n",
    "    elif (mdl_type=='exp_ru') | (mdl_type=='hyp_ru') | (mdl_type=='hyp_r_exp_u') | (mdl_type=='exp_r_hyp_u'):\n",
    "        row['full_session_betaU'] = temp_params[2]\n",
    "        row['full_session_tauU'] = np.exp(temp_params[3])\n",
    "        row['full_session_betaC'] = np.nan\n",
    "        row['full_session_tauC'] = np.nan\n",
    "        row['full_session_const'] = temp_params[4]\n",
    "    elif (mdl_type=='exp_rc') | (mdl_type=='hyp_rc') | (mdl_type=='hyp_r_exp_c'):\n",
    "        row['full_session_betaU'] = np.nan\n",
    "        row['full_session_tauU'] = np.nan\n",
    "        row['full_session_betaC'] = temp_params[2]\n",
    "        row['full_session_tauC'] = np.exp(temp_params[3])\n",
    "        row['full_session_const'] = temp_params[4]\n",
    "    elif (mdl_type=='exp_ruc') | (mdl_type=='hyp_ruc') | (mdl_type=='hyp_r_exp_uc'):\n",
    "        row['full_session_betaU'] = temp_params[2]\n",
    "        row['full_session_tauU'] = np.exp(temp_params[3])\n",
    "        row['full_session_betaC'] = temp_params[4]\n",
    "        row['full_session_tauC'] = np.exp(temp_params[5])\n",
    "        row['full_session_const'] = temp_params[6]\n",
    "\n",
    "    row['CV_mean_loglik'] = np.nanmean(loglik_vec)\n",
    "    row['CV_mean_norm_loglik'] = np.nanmean(norm_loglik_vec)\n",
    "    row['CV_mean_aic'] = np.nanmean(aic_vec)\n",
    "    row['CV_mean_pa'] = np.nanmean(pa_vec)\n",
    "    \n",
    "    return row"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bad1166-a1e3-4462-9902-36e2346d28c7",
   "metadata": {},
   "source": [
    "### Load Imaging and simulation datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ec1de93d-2076-4aed-8735-f005e3bf2046",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load data\n",
    "imaging_beh_df = pd.read_pickle(pjoin(behavior_dir,'all_behavior_df.pkl'))\n",
    "\n",
    "# 10000 trial simulation\n",
    "sim_beh_df = pd.read_pickle(pjoin(behavior_dir, 'rl_simulation_full_base_all_imaging_10000.pkl'))\n",
    "sim_beh_df = sim_beh_df[['Area','Mouse','Date','RL_Model','a','R']]\n",
    "\n",
    "full_sim = sim_beh_df.loc[sim_beh_df['RL_Model']=='rl_full'].reset_index(drop=True)\n",
    "base_sim = sim_beh_df.loc[sim_beh_df['RL_Model']=='rl_base'].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6074e5fe-b46d-480f-bd50-136b53bce0ca",
   "metadata": {},
   "source": [
    "### Fit Behlogit models (15-history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "c4b55851-92a3-47b4-a3e3-45b458ff671e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "110852a8c2074c4d972de910f26cb405",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "# Fit imaging data\n",
    "beh_df = imaging_beh_df#.iloc[:3]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r_hyp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_rc']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r_exp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_rc']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_imaging_15_hist.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 252,
   "id": "715e9d53-3217-44cb-9553-f3d0675c84a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5c400c7c92ab4b65aaf208d2cdb19651",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "480687e00206433ca0ecc3582420ec8f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hyperbolic (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ebd4dea84cf944038c39fc95ba42595a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c336b29632d34b0692eebfd7a6b659b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Fit imaging data\n",
    "beh_df = imaging_beh_df # .iloc[:10]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit_constrained_ru,axis=1,args=([15,'exp_ru']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit_constrained_ru,axis=1,args=([15,'exp_r_hyp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit_constrained_ru,axis=1,args=([15,'hyp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit_constrained_ru,axis=1,args=([15,'hyp_r_exp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_imaging_15_hist_constrained.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "d79bc36a-edad-4826-8031-f312f1fe64f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c8121baad9654ae899e6999a84c0681a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e2a3fdd31f944a9e811193fa9819cf70",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "383f15201c2e467b92299cc93c6a2a45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd84c1492e6844b5ac00b7bbb60f2c84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hyperbolic (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b73540613924317aa01608cd596d79a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "635b340ff0c442a09f352c113bdeefa9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7333c21f92e046cf9bd52ec084be0428",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c30c4c112a184989882e65a6fd5f204d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fee05b17e6bd42bfabdf152295c2431c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1795ba1b479f462d8adc1dd9170d1adb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9d7b36015cfe4a2eb3b94c1dbeb0f93a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "337639333c5f425390d35d3a38f873c6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hyperbolic (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ede28db3be7f4ee79b1ce44a42b742e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "323a86fe9e74455db27edd575b515b62",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e9ed10c8cce645e7933d65b0650f6069",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "20561dd9e7fd4f1480d2b607d09f72ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Fit simulation data\n",
    "beh_df = full_sim.copy() # .iloc[:10]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r_hyp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_rc']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r_exp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_rc']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_rl_full_15_hist.pkl'))\n",
    "\n",
    "beh_df = base_sim.copy()  # .iloc[:10]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r_hyp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_rc']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_ru']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r_exp_u']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_rc']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_rl_base_15_hist.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e655c9b-61fb-4370-9f35-5bf6e1517ee3",
   "metadata": {},
   "source": [
    "### Fit Behlogit models (compare history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "42b0f04e-0a95-4a99-a0cc-80a0632116eb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1dd171fd62ab420dacf43c856c0dfd0d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a94de8bf8a64f079fc0fd88789d28a9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14d1dbd2c7da42a58118bc64380e6f4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ca75e800a4b47e7991189b8a2cd88f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b1842a67e4bd4e09b0f2a4b314bf3ad1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hyperbolic (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a2a0aac8da740098a15cff066f4cfd7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd12e198d1e843439dbb47c5a96f028d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0301af064d2f4bbca96997cac80e6081",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3eff71d04e9b435794283ce7b09a7a7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e505b3159f8c45219f567eed82c50a05",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "# Fit imaging data\n",
    "beh_df = imaging_beh_df#.iloc[:10]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([5,'exp_r']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([10,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([20,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([30,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([5,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([10,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([20,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([30,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_imaging_compare_hist.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "50ed8f66-c986-43f0-9956-cd923cdbaba0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fafe9c8536d84ff68d612ab309180a6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "374caa92fe02456da2705c3bec029127",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "998f7ef4b1b44701b7a83dfdc5a14d5a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9a3a5156812403681928ce219eb0783",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a46ad72793ab4a889498b51089f7397e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hyperbolic (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d6d15b29108404ca8f3d25a1251124c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e61e1da3dd994888a9bc451c313e9b58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29a6b1daf86243c88aa0ad2bca91c3ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fddc18fb8494731a7ece3d9befefb53",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "391245de7dce4aa3bb7dc82f8e5c9a98",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exponential (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14befaa71b8f4574892b774d8593aee3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "33e62353ed414a03aa515677c0e23715",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "45dbe5f3c6ed4f579a9ec06928032959",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5446a991bf574f7c9c730a0a549429ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0091238d6ae842288108da7e32b4a427",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hyperbolic (15-hist)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "866c0d837b7d42a597a469b4e07ee748",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "620ec01bc8884cb4911c2f16c760ece4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "389f070c08484d41b3d48d56ef773ca6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcecc0ae69a74746a20827c09b6faf52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4877770bdb9c49dbaaf344f04ed4cc6d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/80 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Fit simulation data\n",
    "beh_df = full_sim.copy() #.iloc[:10]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([5, 'exp_r']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([10,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([20,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([30,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([5,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([10,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([20,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([30,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_rl_full_compare_hist.pkl'))\n",
    "\n",
    "\n",
    "beh_df = base_sim.copy()  #.iloc[:10]\n",
    "print('exponential (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([5, 'exp_r']))\n",
    "behlogit_all = temp_df.drop(columns=['a','R'])\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([10,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([20,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([30,'exp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "print('hyperbolic (15-hist)')\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([5,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([10,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([15,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([20,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "temp_df = beh_df.progress_apply(modelfit_behlogit,axis=1,args=([30,'hyp_r']))\n",
    "behlogit_all = behlogit_all.append(temp_df.drop(columns=['a','R']), ignore_index=False)\n",
    "\n",
    "behlogit_all.to_pickle(pjoin(behavior_dir, 'behlogit_rl_base_compare_hist.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22375f6d-ad61-47fd-817f-1dc988f76f16",
   "metadata": {},
   "source": [
    "# Imaging analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6123dd4-9149-49bd-ab16-6fa57907f87f",
   "metadata": {},
   "source": [
    "## General import functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "13666063-97b8-42d7-8ed6-42356d10a6ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -f\n",
    "import sys\n",
    "from os.path import dirname, join as pjoin\n",
    "from os import listdir\n",
    "sys.path.append('C:\\\\jupyter_notebooks\\\\Danskin_SciAdv_2023\\\\py_code') # set local directory\n",
    "import numpy as np\n",
    "import pickle\n",
    "import random\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "\n",
    "import h5py\n",
    "import scipy.io as sio\n",
    "from scipy import stats\n",
    "import statsmodels as sm # import statsmodels.api as sm\n",
    "from scipy.optimize import minimize, basinhopping, curve_fit\n",
    "from sklearn.model_selection import KFold\n",
    "from tqdm.notebook import tqdm, trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "50bf8c39-dde7-4563-a8e0-e18b7045a25f",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Directories\n",
    "%autoreload\n",
    "import bdanskin as BD \n",
    "\n",
    "project_dir = 'C:\\\\jupyter_notebooks\\\\Danskin_SciAdv_2023' # local directory\n",
    "imaging_dir = pjoin(project_dir, 'hattori_datasets_xarray')\n",
    "cellfits_dir = pjoin(project_dir, 'hattori_datasets_xarray_cellfits')\n",
    "behavior_dir = pjoin(project_dir, 'hattori_datasets_behavior')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d11982ef-2162-41d7-b59b-3f12782a8ab2",
   "metadata": {},
   "source": [
    "## CV model fit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "936c7f91-6b33-4ac4-aabb-2b89ae2b5fc7",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f7fa77-b4f2-4748-86dc-4b83685cf13f",
   "metadata": {},
   "source": [
    "#### cross-validation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "71a32918-7351-4dab-9693-079f3f064cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_cv(y, c0, a_hist_c, R_hist_c, uR_hist_c, mdl_dict):\n",
    "    n_obs = len(y)  \n",
    "    mdl_name = mdl_dict['mdl_name']\n",
    "    fit_func = mdl_dict['fit_func']\n",
    "    n_params = mdl_dict['n_params']\n",
    "    n_pvals = mdl_dict['n_pvals']\n",
    "    initial_params = mdl_dict['initial_params']\n",
    "    bounds = mdl_dict['bounds']\n",
    "    \n",
    "    # Cross-validation\n",
    "    n_fold = 10\n",
    "    rsq_vec = np.zeros(n_fold)\n",
    "    snr_vec = np.zeros(n_fold)\n",
    "    aic_vec = np.zeros(n_fold)\n",
    "    ll_vec = np.zeros(n_fold)\n",
    "    norm_ll_vec = np.zeros(n_fold)\n",
    "    param_mat = np.zeros([n_fold,n_params])\n",
    "    pval_mat = np.zeros([n_fold, n_pvals])\n",
    "    kf = KFold(n_splits=10, shuffle=False)#, random_state=0)\n",
    "    kf.get_n_splits(np.arange(0,n_obs))\n",
    "    cv_iter = 0\n",
    "    for train_index, test_index in kf.split(np.arange(0,n_obs)):\n",
    "        n_test = len(test_index)\n",
    "        # Simple fit\n",
    "        lik_model = minimize(fit_func, initial_params, bounds=bounds, method='L-BFGS-B',\n",
    "                     args=(y[train_index], c0[train_index], a_hist_c[train_index,:],\n",
    "                             R_hist_c[train_index,:], uR_hist_c[train_index,:], 1))\n",
    "        temp_params = lik_model.x\n",
    "        p_values = estimate_pval(temp_params, y[train_index], c0[train_index],\n",
    "                                 a_hist_c[train_index,:],\n",
    "                                 R_hist_c[train_index,:],\n",
    "                                 uR_hist_c[train_index,:], mdl_dict)\n",
    "        \n",
    "        _,sse,rsq,snr = fit_func(temp_params, y[test_index], c0[test_index],\n",
    "                                 a_hist_c[test_index,:],\n",
    "                                 R_hist_c[test_index,:], \n",
    "                                 uR_hist_c[test_index,:], 0)\n",
    "        loglik, norm_loglik = BD.calc_mdl_loglik(sse, n_test)\n",
    "        aic = BD.calc_aic_loglik(loglik, n_params)\n",
    "        \n",
    "        rsq_vec[cv_iter] = rsq\n",
    "        snr_vec[cv_iter] = snr\n",
    "        aic_vec[cv_iter] = aic\n",
    "        ll_vec[cv_iter] = loglik\n",
    "        norm_ll_vec[cv_iter] = norm_loglik\n",
    "        param_mat[cv_iter,:] = temp_params\n",
    "        pval_mat[cv_iter] = p_values\n",
    "        cv_iter+=1\n",
    "    \n",
    "    mean_params = np.nanmean(param_mat,0)\n",
    "    std_params = np.nanstd(param_mat,0)\n",
    "    mean_pvals = stats.gmean(pval_mat,0)\n",
    "    out_dict = {'mean_params': mean_params, 'std_params': std_params, \n",
    "                'mean_pvals': mean_pvals, 'mean_rsq': np.nanmean(rsq_vec),\n",
    "                'mean_snr': np.nanmean(snr_vec), 'mean_aic': np.nanmean(aic_vec),\n",
    "                'mean_loglik': np.nanmean(ll_vec), 'mean_norm_loglik': np.nanmean(norm_ll_vec)}\n",
    "    return out_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0059e2cf-3782-4514-9da4-f85dc941151f",
   "metadata": {},
   "source": [
    "#### calculate p-values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "c1fa83b8-b4d9-47c9-adc9-2e8d944d777e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimate_pval(temp_params, y, c0, a_hist_c, R_hist_c, uR_hist_c, mdl_dict):\n",
    "    mdl_name = mdl_dict['mdl_name']\n",
    "    y = y.values\n",
    "    if (mdl_name=='null'):\n",
    "        predictors = pd.DataFrame({'c0': c0})\n",
    "    elif (mdl_name=='Rp1'):\n",
    "        past_rc = np.multiply(a_hist_c, R_hist_c)\n",
    "        rc_weights = past_rc[:,-1]\n",
    "        predictors = pd.DataFrame({'RCp1': rc_weights,\n",
    "                                       'c0': c0})\n",
    "    elif (mdl_name=='exp_r'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        rc_weights = BD.exp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        predictors = pd.DataFrame({'exp_rc': rc_weights,\n",
    "                                       'c0': c0})\n",
    "    elif (mdl_name=='hyp_r'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        rc_weights = BD.hyp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        predictors = pd.DataFrame({'hyp_rc': rc_weights,\n",
    "                                       'c0': c0})\n",
    "    elif (mdl_name=='exp_ru') | (mdl_name=='exp_ru_con'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.exp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.exp_filter(a_hist_c, R_hist_c, uc_tau,'ur')\n",
    "        predictors = pd.DataFrame({'exp_rc': rc_weights, 'exp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    elif (mdl_name=='hyp_ru') | (mdl_name=='hyp_ru_con'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.hyp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.hyp_filter(a_hist_c, R_hist_c, uc_tau,'ur')\n",
    "        predictors = pd.DataFrame({'hyp_rc': rc_weights, 'hyp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    elif (mdl_name=='exp_r_hyp_u'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.exp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.hyp_filter(a_hist_c, R_hist_c, uc_tau,'c')\n",
    "        predictors = pd.DataFrame({'exp_rc': rc_weights, 'hyp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    elif (mdl_name=='hyp_r_exp_u'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.hyp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.exp_filter(a_hist_c, R_hist_c, uc_tau,'c')\n",
    "        predictors = pd.DataFrame({'hyp_rc': rc_weights, 'exp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "       \n",
    "    X2 = sm.add_constant(predictors) # Add constant term to predictor dataframe\n",
    "    est = sm.OLS(y, X2).fit() # Fit linear regression with ordinary least squares\n",
    "    p_values = est.pvalues\n",
    "    return p_values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374b6175-5c15-46b8-9a8a-ca8c3eadcb56",
   "metadata": {},
   "source": [
    "#### pandas apply cv fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "9504799c-40a1-4b60-a132-cc511cd38013",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_fit_CV_cell(column,c0,a_hist_c,R_hist_c,uR_hist_c,mdl_name):\n",
    "    # mdl_name: 'exp_r' 'hyp_ruc' 'exp_r_con'\n",
    "    rng = np.random.default_rng(220804)\n",
    "    y = column\n",
    "    n_obs = len(y)\n",
    "    fit_func = eval('BD.reg_' + mdl_name + '_c0')\n",
    "    \n",
    "    if (mdl_name=='null'):\n",
    "        n_params = 2; n_pvals = 2\n",
    "        initial_params = [0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='Rp1'):\n",
    "        n_params = 3; n_pvals = 3\n",
    "        initial_params = [0.01, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='exp_r') | (mdl_name=='hyp_r'):\n",
    "        n_params = 4; n_pvals = 3\n",
    "        initial_params = [0.01, 1, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.log(100),np.log(100)],[-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='exp_ru') | (mdl_name=='hyp_ru'):\n",
    "        n_params = 6; n_pvals = 4\n",
    "        initial_params = [0.01, 1, -0.01, 1, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.log(100),np.log(100)],[-np.inf,np.inf],[-np.log(100),np.log(100)],\n",
    "                  [-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='exp_r_hyp_u') | (mdl_name=='hyp_r_exp_u'):\n",
    "        n_params = 6; n_pvals = 4\n",
    "        initial_params = [0.01, 1, -0.01, 1, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.log(100),np.log(100)],[-np.inf,0],[-np.log(100),np.log(100)],\n",
    "                  [-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    else:\n",
    "        print('unable to parse mdl_type')\n",
    "    \n",
    "    # Fit with cross-validation\n",
    "    mdl_dict = {'mdl_name': mdl_name, 'fit_func': fit_func, 'n_params': n_params, 'n_pvals': n_pvals,\n",
    "                'initial_params': initial_params, 'bounds': bounds}\n",
    "    cv_out_dict = fit_cv(y, c0, a_hist_c, R_hist_c, uR_hist_c, mdl_dict)\n",
    "    temp_params = cv_out_dict['mean_params']\n",
    "    std_params = cv_out_dict['std_params']\n",
    "    p_values = cv_out_dict['mean_pvals']\n",
    "    \n",
    "    column['mdl_type'] = mdl_name\n",
    "    column['cv_mean_rsq'] = cv_out_dict['mean_rsq']\n",
    "    column['cv_mean_snr'] = cv_out_dict['mean_snr']\n",
    "    column['cv_mean_aic'] = cv_out_dict['mean_aic']\n",
    "    column['cv_mean_loglik'] = cv_out_dict['mean_loglik']\n",
    "    column['cv_mean_norm_loglik'] = cv_out_dict['mean_norm_loglik']\n",
    "    if (mdl_name=='null'):\n",
    "        column['cv_mean_beta_RC'] =  np.nan\n",
    "        column['cv_mean_tau_RC'] =  np.nan\n",
    "        column['cv_mean_beta_UC'] = np.nan\n",
    "        column['cv_mean_tau_UC'] = np.nan\n",
    "        column['cv_mean_beta_C0'] = temp_params[0]\n",
    "        column['cv_mean_beta_0'] = temp_params[1]\n",
    "        column['cv_std_beta_RC'] =  np.nan\n",
    "        column['cv_std_tau_RC'] =  np.nan\n",
    "        column['cv_std_beta_UC'] =  np.nan\n",
    "        column['cv_std_tau_UC'] =  np.nan\n",
    "        column['cv_std_beta_C0'] =  std_params[0]\n",
    "        column['cv_std_beta_0'] =  std_params[1]\n",
    "        column['cv_gmean_p_beta_RC'] = np.nan\n",
    "        column['cv_gmean_p_beta_UC'] = np.nan\n",
    "        column['cv_gmean_p_beta_C0'] = p_values[1]\n",
    "        column['cv_gmean_p_beta_0'] = p_values[0]\n",
    "    elif (mdl_name=='Rp1'):\n",
    "        column['cv_mean_beta_RC'] =  temp_params[0]\n",
    "        column['cv_mean_tau_RC'] =  np.nan\n",
    "        column['cv_mean_beta_UC'] = np.nan\n",
    "        column['cv_mean_tau_UC'] = np.nan\n",
    "        column['cv_mean_beta_C0'] = temp_params[1]\n",
    "        column['cv_mean_beta_0'] = temp_params[2]\n",
    "        column['cv_std_beta_RC'] =  std_params[0]\n",
    "        column['cv_std_tau_RC'] =  np.nan\n",
    "        column['cv_std_beta_UC'] =  np.nan\n",
    "        column['cv_std_tau_UC'] =  np.nan\n",
    "        column['cv_std_beta_C0'] =  std_params[1]\n",
    "        column['cv_std_beta_0'] =  std_params[2]\n",
    "        column['cv_gmean_p_beta_RC'] = p_values[1]\n",
    "        column['cv_gmean_p_beta_UC'] = np.nan\n",
    "        column['cv_gmean_p_beta_C0'] = p_values[2]\n",
    "        column['cv_gmean_p_beta_0'] = p_values[0]\n",
    "    elif (mdl_name=='exp_r') | (mdl_name=='hyp_r'):\n",
    "        column['cv_mean_beta_RC'] =  temp_params[0]\n",
    "        column['cv_mean_tau_RC'] =  np.exp(temp_params[1])\n",
    "        column['cv_mean_beta_UC'] = np.nan\n",
    "        column['cv_mean_tau_UC'] = np.nan\n",
    "        column['cv_mean_beta_C0'] = temp_params[2]\n",
    "        column['cv_mean_beta_0'] = temp_params[3]\n",
    "        column['cv_std_beta_RC'] =  std_params[0]\n",
    "        column['cv_std_tau_RC'] =  np.exp(std_params[1])\n",
    "        column['cv_std_beta_UC'] =  np.nan\n",
    "        column['cv_std_tau_UC'] =  np.nan\n",
    "        column['cv_std_beta_C0'] =  std_params[2]\n",
    "        column['cv_std_beta_0'] =  std_params[3]\n",
    "        column['cv_gmean_p_beta_RC'] = p_values[1]\n",
    "        column['cv_gmean_p_beta_UC'] = np.nan\n",
    "        column['cv_gmean_p_beta_C0'] = p_values[2]\n",
    "        column['cv_gmean_p_beta_0'] = p_values[0]\n",
    "    elif (mdl_name=='exp_ru') | (mdl_name=='hyp_ru') | (mdl_name=='exp_r_hyp_u') | (mdl_name=='hyp_r_exp_u'):\n",
    "        column['cv_mean_beta_RC'] =  temp_params[0]\n",
    "        column['cv_mean_tau_RC'] =  np.exp(temp_params[1])\n",
    "        column['cv_mean_beta_UC'] = temp_params[2]\n",
    "        column['cv_mean_tau_UC'] = np.exp(temp_params[3])\n",
    "        column['cv_mean_beta_C0'] = temp_params[4]\n",
    "        column['cv_mean_beta_0'] = temp_params[5]\n",
    "        column['cv_std_beta_RC'] =  std_params[0]\n",
    "        column['cv_std_tau_RC'] =  np.exp(std_params[1])\n",
    "        column['cv_std_beta_UC'] = std_params[2]\n",
    "        column['cv_std_tau_UC'] = np.exp(std_params[3])\n",
    "        column['cv_std_beta_C0'] = std_params[4]\n",
    "        column['cv_std_beta_0'] = std_params[5]\n",
    "        column['cv_gmean_p_beta_RC'] = p_values[1]\n",
    "        column['cv_gmean_p_beta_UC'] = p_values[2]\n",
    "        column['cv_gmean_p_beta_C0'] = p_values[3]\n",
    "        column['cv_gmean_p_beta_0'] = p_values[0]\n",
    "    return column"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0a13371-f21e-4bb2-a168-6a2a0eb7bc13",
   "metadata": {},
   "source": [
    "### Process Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2243b1fa-870a-4d31-b8a7-f8e9a90fcbd4",
   "metadata": {},
   "source": [
    "#### Estimate across sessions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "3a16e20a-f157-46a2-8f5f-69f5e1461dd6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2023964dd95b4e8480d1ceca2c7b3105",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checking analysis for RH055 on 180715\n",
      "null models,  Checking analysis for RH055 on 180807\n",
      "null models,  "
     ]
    }
   ],
   "source": [
    "sessions_all = BD.get_sessions('all')\n",
    "for ss in trange(len(sessions_all)):\n",
    "    mouse = sessions_all[ss][:5]\n",
    "    date = sessions_all[ss][6:12]\n",
    "    session = sessions_all[ss][6:]\n",
    "    print('Checking analysis for {} on {}'.format(mouse, date))\n",
    "\n",
    "    ds_2s = xr.open_dataset(pjoin(imaging_dir,'preprocess_data_2s','{}_ready_2s.nc'.format(sessions_all[ss])))\n",
    "\n",
    "    # Prepare behavior information\n",
    "    a = np.array(ds_2s.a)\n",
    "    R = np.array(ds_2s.R)\n",
    "    n_back = 15\n",
    "    a_hist,R_hist,uR_hist = BD.prepare_hist_matrix(a,R,n_back)\n",
    "\n",
    "    # Choice metrics\n",
    "    choice_trials = (a==1)|(a==2)\n",
    "    left_trials = (a==1)+0\n",
    "    left_trials[a==0]=-1\n",
    "\n",
    "    # NaN check to remove choice trials, necessary for some PPC sessions\n",
    "    temp_da = ds_2s.spks.isel(trial=choice_trials)#,cell=slice(0,3))\n",
    "    neuron_df = ds_2s.spks.to_pandas().reset_index(drop = True)\n",
    "    neuron_df_choice = temp_da.to_pandas().reset_index(drop = True)\n",
    "    neuron_isnan = np.isnan(neuron_df_choice.loc[:,1])\n",
    "    if sum(neuron_isnan)>0:\n",
    "        print('NaN data choice trials detected. Using only choice trials with data')\n",
    "        choice_trials[np.isnan(neuron_df.loc[:,1])]=False\n",
    "        temp_da = ds_2s.spks.isel(trial=choice_trials)#,cell=slice(0,3))\n",
    "        neuron_df_choice = temp_da.to_pandas().reset_index(drop = True)\n",
    "\n",
    "    # Full session\n",
    "    a_hist_c = a_hist[choice_trials,:]\n",
    "    R_hist_c = R_hist[choice_trials,:]\n",
    "    uR_hist_c = uR_hist[choice_trials,:]\n",
    "    c0 = left_trials[choice_trials]\n",
    "    \n",
    "    print('null models, ', end =' ')\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'null']))\n",
    "    regression_temp = reg_mdl.iloc[-22:].T\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'Rp1']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "    \n",
    "    print('exponential models,' , end =' ')\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'exp_r']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'exp_ru']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "    \n",
    "    print('hyperbolic models,' , end =' ')\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'hyp_r']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'hyp_ru']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "\n",
    "    print('mixed models,' , end =' ')\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                                     args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'exp_r_hyp_u']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df_choice.apply(apply_fit_CV_cell, axis=0,\n",
    "                                     args=([c0, a_hist_c, R_hist_c, uR_hist_c, 'hyp_r_exp_u']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-22:].T, ignore_index=False)\n",
    "    print('done')\n",
    "    \n",
    "    Save outputs\n",
    "    output_multi = regression_temp.reset_index().set_index(['mdl_type','cell'])  \n",
    "    output_xarray = output_multi.to_xarray()\n",
    "    save_filename = pjoin(cellfits_dir,'ready_cellfits','cv_compare_mdl','sse_{}_15hist.nc'.format(sessions_all[ss]))\n",
    "    output_xarray.to_netcdf(save_filename) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd29825d-abd3-4f74-848f-6361417b37ec",
   "metadata": {},
   "source": [
    "## Halves tau estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15e68204-be19-4b89-bc29-3253d3795bcd",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99c62934-cef2-40bd-b1fa-15e8e2bcb8bb",
   "metadata": {},
   "source": [
    "#### calculate p-values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "8697d916-d719-472c-b006-7d771354308a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimate_pval(temp_params, y, c0, a_hist_c, R_hist_c, uR_hist_c, mdl_name):\n",
    "    #mdl_name = mdl_dict['mdl_name']\n",
    "    y = y.values\n",
    "    if (mdl_name=='null'):\n",
    "        predictors = pd.DataFrame({'c0': c0})\n",
    "    elif (mdl_name=='Rp1'):\n",
    "        past_rc = np.multiply(a_hist_c, R_hist_c)\n",
    "        rc_weights = past_rc[:,-1]\n",
    "        predictors = pd.DataFrame({'RCp1': rc_weights,\n",
    "                                       'c0': c0})\n",
    "    elif (mdl_name=='exp_r'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        rc_weights = BD.exp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        predictors = pd.DataFrame({'exp_rc': rc_weights,\n",
    "                                       'c0': c0})\n",
    "    elif (mdl_name=='hyp_r'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        rc_weights = BD.hyp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        predictors = pd.DataFrame({'hyp_rc': rc_weights,\n",
    "                                       'c0': c0})\n",
    "    elif (mdl_name=='exp_ru') | (mdl_name=='exp_ru_con'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.exp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.exp_filter(a_hist_c, R_hist_c, uc_tau,'ur')\n",
    "        predictors = pd.DataFrame({'exp_rc': rc_weights, 'exp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    elif (mdl_name=='hyp_ru') | (mdl_name=='hyp_ru_con'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.hyp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.hyp_filter(a_hist_c, R_hist_c, uc_tau,'ur')\n",
    "        predictors = pd.DataFrame({'hyp_rc': rc_weights, 'hyp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    elif (mdl_name=='exp_r_hyp_u'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.exp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.hyp_filter(a_hist_c, R_hist_c, uc_tau,'c')\n",
    "        predictors = pd.DataFrame({'exp_rc': rc_weights, 'hyp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    elif (mdl_name=='hyp_r_exp_u'):\n",
    "        rc_tau = np.exp(temp_params[1])\n",
    "        uc_tau = np.exp(temp_params[3])\n",
    "        rc_weights = BD.hyp_filter(a_hist_c, R_hist_c, rc_tau,'r')\n",
    "        uc_weights = BD.exp_filter(a_hist_c, R_hist_c, uc_tau,'c')\n",
    "        predictors = pd.DataFrame({'hyp_rc': rc_weights, 'exp_uc': uc_weights,\n",
    "                                   'c0': c0})\n",
    "    # print(len(y), len(predictors))    \n",
    "    X2 = sm.add_constant(predictors) # Add constant term to predictor dataframe\n",
    "    est = sm.OLS(y, X2).fit() # Fit linear regression with ordinary least squares\n",
    "    p_values = est.pvalues\n",
    "    return p_values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f134e9c8-2999-49bb-a202-ad9d4fec8da4",
   "metadata": {},
   "source": [
    "#### pandas apply (decay)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "ca8409f9-bd2e-4efb-abda-6c279621bf14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_fit_cell(column,c0,a_hist_c,R_hist_c,uR_hist_c,half_id,mdl_name):\n",
    "    # mdl_name: 'exp_r' 'hyp_ruc' 'exp_r_con'\n",
    "    rng = np.random.default_rng(220804)\n",
    "    y = column\n",
    "    n_obs = len(y)\n",
    "    fit_func = eval('BD.reg_' + mdl_name + '_c0')\n",
    "    \n",
    "    if (mdl_name=='null'):\n",
    "        n_params = 2; n_pvals = 2\n",
    "        initial_params = [0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='Rp1'):\n",
    "        n_params = 3; n_pvals = 3\n",
    "        initial_params = [0.01, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='exp_r') | (mdl_name=='hyp_r'):\n",
    "        n_params = 4; n_pvals = 3\n",
    "        initial_params = [0.01, 1, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.log(100),np.log(100)],[-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='exp_ru') | (mdl_name=='hyp_ru'):\n",
    "        n_params = 6; n_pvals = 4\n",
    "        initial_params = [0.01, 1, -0.01, 1, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.log(100),np.log(100)],[-np.inf,np.inf],[-np.log(100),np.log(100)],\n",
    "                  [-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    elif (mdl_name=='exp_r_hyp_u') | (mdl_name=='hyp_r_exp_u'):\n",
    "        n_params = 6; n_pvals = 4\n",
    "        initial_params = [0.01, 1, -0.01, 1, 0.01, 0.01]\n",
    "        bounds = [[-np.inf,np.inf],[-np.log(100),np.log(100)],[-np.inf,0],[-np.log(100),np.log(100)],\n",
    "                  [-np.inf,np.inf],[-np.inf,np.inf]]\n",
    "    else:\n",
    "        print('unable to parse mdl_type')\n",
    "    \n",
    "    # Simple fit\n",
    "    lik_model = minimize(fit_func, initial_params, bounds=bounds, method='L-BFGS-B',\n",
    "                 args=(y, c0, a_hist_c, R_hist_c, uR_hist_c, 1))\n",
    "    temp_params = lik_model.x\n",
    "    p_values = estimate_pval(temp_params, y, c0, a_hist_c, R_hist_c, uR_hist_c, mdl_name)\n",
    "    _,sse,rsq,snr = fit_func(temp_params, y, c0, a_hist_c, R_hist_c, uR_hist_c, 0)\n",
    "    loglik, norm_loglik = BD.calc_mdl_loglik(sse, n_obs)\n",
    "    aic = BD.calc_aic_loglik(loglik, n_params)\n",
    "    \n",
    "    column['half'] = half_id\n",
    "    column['mdl_type'] = mdl_name\n",
    "    column['rsq'] = rsq\n",
    "    column['snr'] = snr\n",
    "    column['aic'] = aic\n",
    "    column['loglik'] = loglik\n",
    "    column['norm_loglik'] = norm_loglik\n",
    "    if (mdl_name=='null'):\n",
    "        column['beta_RC'] =  np.nan\n",
    "        column['tau_RC'] =  np.nan\n",
    "        column['beta_UC'] = np.nan\n",
    "        column['tau_UC'] = np.nan\n",
    "        column['beta_C'] = np.nan\n",
    "        column['tau_C'] = np.nan\n",
    "        column['beta_C0'] = temp_params[0]\n",
    "        column['beta_0'] = temp_params[1]\n",
    "        column['p_beta_RC'] = np.nan\n",
    "        column['p_beta_UC'] = np.nan\n",
    "        column['p_beta_C'] = np.nan\n",
    "        column['p_beta_C0'] = p_values[1]\n",
    "        column['p_beta_0'] = p_values[0]\n",
    "    elif (mdl_name=='Rp1'):\n",
    "        column['beta_RC'] =  temp_params[0]\n",
    "        column['tau_RC'] =  np.nan\n",
    "        column['beta_UC'] = np.nan\n",
    "        column['tau_UC'] = np.nan\n",
    "        column['beta_C'] = np.nan\n",
    "        column['tau_C'] = np.nan\n",
    "        column['beta_C0'] = temp_params[1]\n",
    "        column['beta_0'] = temp_params[2]\n",
    "        column['p_beta_RC'] = p_values[1]\n",
    "        column['p_beta_UC'] = np.nan\n",
    "        column['p_beta_C'] = np.nan\n",
    "        column['p_beta_C0'] = p_values[2]\n",
    "        column['p_beta_0'] = p_values[0]\n",
    "    elif (mdl_name=='exp_r') | (mdl_name=='hyp_r'):\n",
    "        column['beta_RC'] =  temp_params[0]\n",
    "        column['tau_RC'] =  np.exp(temp_params[1])\n",
    "        column['beta_UC'] = np.nan\n",
    "        column['tau_UC'] = np.nan\n",
    "        column['beta_C'] = np.nan\n",
    "        column['tau_C'] = np.nan\n",
    "        column['beta_C0'] = temp_params[2]\n",
    "        column['beta_0'] = temp_params[3]\n",
    "        column['p_beta_RC'] = p_values[1]\n",
    "        column['p_beta_UC'] = np.nan\n",
    "        column['p_beta_C'] = np.nan\n",
    "        column['p_beta_C0'] = p_values[2]\n",
    "        column['p_beta_0'] = p_values[0]\n",
    "    elif (mdl_name=='exp_ru') | (mdl_name=='hyp_ru') | (mdl_name=='exp_r_hyp_u') | (mdl_name=='hyp_r_exp_u'):\n",
    "        column['beta_RC'] =  temp_params[0]\n",
    "        column['tau_RC'] =  np.exp(temp_params[1])\n",
    "        column['beta_UC'] = temp_params[2]\n",
    "        column['tau_UC'] = np.exp(temp_params[3])\n",
    "        column['beta_C'] = np.nan\n",
    "        column['tau_C'] = np.nan \n",
    "        column['beta_C0'] = temp_params[4]\n",
    "        column['beta_0'] = temp_params[5]\n",
    "        column['p_beta_RC'] = p_values[1]\n",
    "        column['p_beta_UC'] = p_values[2]\n",
    "        column['p_beta_C'] = np.nan\n",
    "        column['p_beta_C0'] = p_values[3]\n",
    "        column['p_beta_0'] = p_values[0]\n",
    "    \n",
    "    return column"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06c9ccba-4be6-4584-8f26-3428daa644df",
   "metadata": {},
   "source": [
    "#### iterate across pandas apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "c8a7e650-6979-488d-a72e-6974a77e3ad6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_models_half(neuron_df, c0, a_hist_c, R_hist_c, half_id):\n",
    "    print('null models, ', end =' ')\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'null']))\n",
    "    regression_temp = reg_mdl.iloc[-20:].T\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'Rp1']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "    \n",
    "    print('exponential models,' , end =' ')\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'exp_r']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'exp_ru']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'exp_r_hyp_u']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "\n",
    "    print('hyperbolic models,' , end =' ')\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'hyp_r']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'hyp_ru']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "    reg_mdl = neuron_df.apply(apply_fit_cell, axis=0,\n",
    "                              args=([c0, a_hist_c, R_hist_c, uR_hist_c, half_id, 'hyp_r_exp_u']))\n",
    "    regression_temp = regression_temp.append(reg_mdl.iloc[-20:].T, ignore_index=False)\n",
    "\n",
    "    print('done')\n",
    "    \n",
    "    return regression_temp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec5365f3-9e03-41cc-8ce4-63cc2f5e6e59",
   "metadata": {},
   "source": [
    "### Process Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "818be8ee-dce6-49f9-a394-5a35c0ddefb0",
   "metadata": {},
   "source": [
    "#### Estimate across sessions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "d9dca8fb-8570-41b9-b72b-2e0953cd767d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f83267399844126afc04ac1b6583768",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checking analysis for RH055 on 180715\n",
      "null models,  exponential models, hyperbolic models, done\n",
      "null models,  exponential models, hyperbolic models, done\n",
      "null models,  exponential models, hyperbolic models, done\n",
      "Checking analysis for RH055 on 180807\n",
      "null models,  exponential models, hyperbolic models, done\n",
      "null models,  exponential models, hyperbolic models, done\n",
      "null models,  exponential models, hyperbolic models, done\n"
     ]
    }
   ],
   "source": [
    "sessions_all = BD.get_sessions('all')\n",
    "for ss in trange(len(sessions_all)):\n",
    "    mouse = sessions_all[ss][:5]\n",
    "    date = sessions_all[ss][6:12]\n",
    "    session = sessions_all[ss][6:]\n",
    "    print('Checking analysis for {} on {}'.format(mouse, date))\n",
    "\n",
    "    ds_2s = xr.open_dataset(pjoin(imaging_dir,'preprocess_data_2s','{}_ready_2s.nc'.format(sessions_all[ss])))\n",
    "\n",
    "    # Prepare behavior information\n",
    "    a = np.array(ds_2s.a)\n",
    "    R = np.array(ds_2s.R)\n",
    "    n_back = 15\n",
    "    a_hist,R_hist,uR_hist = BD.prepare_hist_matrix(a,R,n_back)\n",
    "    a_cleaned = a.copy()+0.\n",
    "    a_cleaned[a>2] = 0\n",
    "    a_cleaned[a==2] = -1\n",
    "\n",
    "    # Choice metrics\n",
    "    choice_trials = (a==1)|(a==2)\n",
    "    left_trials = (a==1)+0\n",
    "    left_trials[a==0]=-1\n",
    "\n",
    "    # NaN check to remove choice trials, necessary for some PPC sessions\n",
    "    temp_da = ds_2s.spks.isel(trial=choice_trials)#,cell=slice(0,3))\n",
    "    neuron_df = ds_2s.spks.to_pandas().reset_index(drop = True)\n",
    "    neuron_df_choice = temp_da.to_pandas().reset_index(drop = True)\n",
    "    neuron_isnan = np.isnan(neuron_df_choice.loc[:,1])\n",
    "    if sum(neuron_isnan)>0:\n",
    "        print('NaN data choice trials detected. Using only choice trials with data')\n",
    "        choice_trials[np.isnan(neuron_df.loc[:,1])]=False\n",
    "        temp_da = ds_2s.spks.isel(trial=choice_trials)#,cell=slice(0,3))\n",
    "        neuron_df_choice = temp_da.to_pandas().reset_index(drop = True)\n",
    "\n",
    "    # Full session\n",
    "    a_hist_c = a_hist[choice_trials,:]\n",
    "    R_hist_c = R_hist[choice_trials,:]\n",
    "    uR_hist_c = uR_hist[choice_trials,:]\n",
    "    c0 = left_trials[choice_trials]\n",
    "    regression_output = fit_models_half(neuron_df_choice, c0, a_hist_c, R_hist_c, 0)\n",
    "      \n",
    "    # Select Half of choice trials\n",
    "    choice_inds = np.where(choice_trials)[0]\n",
    "    num_choice = sum(choice_trials)\n",
    "    half_choices = np.floor(num_choice/2).astype(int) \n",
    "    sub_trials_1 = choice_inds[:half_choices]\n",
    "    sub_trials_2 = choice_inds[half_choices:]\n",
    "    # Half 1\n",
    "    a_hist_c = a_hist[sub_trials_1,:]\n",
    "    R_hist_c = R_hist[sub_trials_1,:]\n",
    "    uR_hist_c = uR_hist[sub_trials_1,:]\n",
    "    c0 = a_cleaned[sub_trials_1]\n",
    "    temp_da = ds_2s.spks.isel(trial=sub_trials_1)#,cell=slice(0,3))\n",
    "    neuron_df_choice = temp_da.to_pandas().reset_index(drop = True)\n",
    "    regression_temp = fit_models_half(neuron_df_choice, c0, a_hist_c, R_hist_c, 1)\n",
    "    regression_output = regression_output.append(regression_temp)\n",
    "    \n",
    "    # Half 2\n",
    "    a_hist_c = a_hist[sub_trials_2,:]\n",
    "    R_hist_c = R_hist[sub_trials_2,:]\n",
    "    uR_hist_c = uR_hist[sub_trials_2,:]\n",
    "    c0 = a_cleaned[sub_trials_2]\n",
    "    temp_da = ds_2s.spks.isel(trial=sub_trials_2)#,cell=slice(0,3))\n",
    "    neuron_df_choice = temp_da.to_pandas().reset_index(drop = True)\n",
    "    regression_temp = fit_models_half(neuron_df_choice, c0, a_hist_c, R_hist_c, 2)\n",
    "    regression_output = regression_output.append(regression_temp)\n",
    "    \n",
    "    # Save outputs\n",
    "    output_multi = regression_output.reset_index().set_index(['half','mdl_type','cell'])  \n",
    "    output_xarray = output_multi.to_xarray()\n",
    "    save_filename = pjoin(cellfits_dir,'ready_cellfits','halves_compare_mdl','sse_{}_15hist.nc'.format(sessions_all[ss]))\n",
    "    # output_xarray.to_netcdf(save_filename) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21abe203-a72b-4722-8316-8d7d5d4570fe",
   "metadata": {},
   "source": [
    "## Quasihyperbolic analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ab8cd7-618c-4a05-9527-62d605d91d83",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ffb799f-012f-46ba-ba4a-05179c31808a",
   "metadata": {},
   "source": [
    "#### load data function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "91ebc2dc-ee1e-4435-9217-12522cfa733d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_cell_data(session):\n",
    "    save_filename = pjoin(cellfits_dir,'ready_cellfits','halves_compare_mdl',\n",
    "                          'sse_{}_15hist.nc'.format(session))\n",
    "    \n",
    "    r_xarray = xr.open_dataset(save_filename)\n",
    "    crit_bool = ( (r_xarray.sel(mdl_type='exp_ru',half=1).p_beta_RC<0.05) &\n",
    "                  (r_xarray.sel(mdl_type='exp_ru',half=2).p_beta_RC<0.05) &\n",
    "                  (r_xarray.sel(mdl_type='exp_ru',half=0).tau_RC<99)&\n",
    "                  (r_xarray.sel(mdl_type='exp_ru',half=0).tau_RC>0.01) )\n",
    "    \n",
    "    out_pd = r_xarray.sel(mdl_type='exp_ru',half=0).to_dataframe()[['beta_RC','tau_RC']]\n",
    "    out_pd['crit'] = crit_bool\n",
    "    return out_pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9b1c3ae-81a6-489e-a030-ed6ecff2d7a3",
   "metadata": {},
   "source": [
    "#### modelfit weighted neusum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "804e6d5f-2de2-4f6e-a608-82e4e6c8ca56",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for a given session\n",
    "def modelfit_quasilogit(a,R,n_tau,drawn_tau):\n",
    "    n_back = 15\n",
    "    # Prepare history matrices\n",
    "    a_hist,R_hist,_ = BD.prepare_hist_matrix(a,R,n_back) \n",
    "    \n",
    "    n_draw_iter = np.floor(n_draw/n_tau).astype(int); #print(n_tau,n_draw_iter)\n",
    "    drawn_tau_mat = np.reshape(drawn_tau[:n_draw_iter*n_tau], (n_draw_iter,n_tau))\n",
    "    initialize_bool = True\n",
    "    for dd in range(n_draw_iter):\n",
    "        tau_vec = tau_vec = drawn_tau_mat[dd,:]\n",
    "    \n",
    "        # Prepare the model initial guesses and parameters\n",
    "        filter_initguess = np.random.rand(n_tau)\n",
    "        filter_initguess[filter_initguess<1e-3] = 1e-3 # to ensure no 0 value\n",
    "        init_params = np.append(0,filter_initguess) # add constant\n",
    "        bounds = np.vstack([np.append(-np.inf,0*np.ones(len(tau_vec))),\n",
    "                            np.inf*np.ones(len(init_params))]).T     \n",
    "\n",
    "        n_params = len(filter_initguess)\n",
    "        # Run k-fold cross validation\n",
    "        cv_iter = 0\n",
    "        loglik_vec = np.zeros(10)\n",
    "        aic_vec = np.zeros(10)\n",
    "        pa_vec = np.zeros(10)\n",
    "        kf = KFold(n_splits=10, shuffle=False)#, random_state=0)\n",
    "        kf.get_n_splits(np.arange(0,len(a)))\n",
    "        for train_index, test_index in kf.split(np.arange(0,len(a))):  \n",
    "            minimizer_kwargs = {'method':'L-BFGS-B', 'bounds':bounds, 'args':(\n",
    "                tau_vec, a[train_index], a_hist[train_index,:], R_hist[train_index,:], 1)}\n",
    "            with np.errstate(divide='ignore',invalid='ignore'):\n",
    "                try:\n",
    "                    lik_model = basinhopping(BD.logit_quasi_fit, init_params, minimizer_kwargs=minimizer_kwargs, niter=3)\n",
    "                except:# Necessary to stop weird x0 error\n",
    "                    init_params = 0.5*np.ones(len(init_params))\n",
    "                    lik_model = basinhopping(BD.logit_quasi_fit, init_params, minimizer_kwargs=minimizer_kwargs, niter=3)\n",
    "            mdl_weights = lik_model.x\n",
    "            p_logit, _, _, nloglik = BD.logit_quasi_fit(mdl_weights, tau_vec,\n",
    "                                                   a[test_index], \n",
    "                                                   a_hist[test_index,:], \n",
    "                                                   R_hist[test_index,:], 0)\n",
    "            pa_vec[cv_iter],_ = BD.PA_LL_logit(a[test_index],p_logit)\n",
    "            aic_vec[cv_iter] = 2*n_params + 2*nloglik\n",
    "            loglik_vec[cv_iter] = -nloglik\n",
    "            cv_iter+=1\n",
    "        \n",
    "        temp_series = pd.Series({'n_tau': n_tau, 'draw_iter': dd, \n",
    "                                 'loglik': loglik_vec.mean(), 'loglik_trial': loglik_vec.mean()/len(test_index), \n",
    "                                 'pa': pa_vec.mean(),\n",
    "                                 'aic': aic_vec.mean()})\n",
    "        if initialize_bool:\n",
    "            temp_df = pd.DataFrame(temp_series).T\n",
    "            initialize_bool = False\n",
    "        else:\n",
    "            temp_df = temp_df.append(temp_series, ignore_index=True)\n",
    "    \n",
    "    return temp_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "674f5884-3bae-4c4f-a80a-93bbabfd304b",
   "metadata": {},
   "source": [
    "### Load behavior fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "4074941b-c736-4be4-83a5-1b75e511cd25",
   "metadata": {},
   "outputs": [],
   "source": [
    "file_savename = pjoin(behavior_dir, 'behlogit_exp_hyp_history.nc')\n",
    "behlogit_expert = xr.open_dataset(file_savename)\n",
    "behlogit_fit = behlogit_expert.sel(mdl_name=['exp_r','hyp_r'],mdl_history=15).drop('mdl_history')[[\n",
    "    'CV_mean_loglik', 'CV_mean_aic', 'CV_mean_pa']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7533449b-169c-42dc-affe-43ec8e326fe2",
   "metadata": {},
   "source": [
    "### Fit weighted quasi-hyperbolic to behavior, given tau"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "08b5dc83-27f9-47e3-9d68-85fd04daf77f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running for area RSC\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04beafafd1644f72a6ae294251f84af9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "n_draw = 150\n",
    "max_n_tau = 15 # max number of tau to draw from distribution, iterated up to this number\n",
    "\n",
    "# Load data\n",
    "beh_df = pd.read_pickle(pjoin(behavior_dir,'all_behavior_df.pkl'))\n",
    "\n",
    "initialize_df = True # True False\n",
    "for area in ['RSC']:#,'PPC','ALM','M2','S1']: # ,,'V1']\n",
    "    print('Running for area {}'.format(area))\n",
    "    generating_area = area\n",
    "    sessions_all = BD.get_sessions(area)\n",
    "    output_all = pd.DataFrame([])\n",
    "    for ss in range(len(sessions_all)):\n",
    "        out_pd = load_cell_data(sessions_all[ss])\n",
    "        output_all = output_all.append(out_pd,ignore_index=True)\n",
    "    tau_distribution = output_all.loc[output_all['crit'],'tau_RC'].values\n",
    "    np.random.seed(211021) # n_draw_iteration\n",
    "    drawn_tau = np.random.choice(tau_distribution,n_draw)\n",
    "    for nn in trange(max_n_tau): \n",
    "        n_tau = nn+1 # number of tau used to produce QH\n",
    "        # print('Running with {} taus'.format(n_tau))\n",
    "        for session_id in range(len(beh_df)):\n",
    "            # Prepare behavior\n",
    "            n_back = 15 \n",
    "            a = beh_df.loc[session_id,'a']\n",
    "            R = beh_df.loc[session_id,'R']\n",
    "            # Run CV fit\n",
    "            cv_fit_df = modelfit_quasilogit(a,R,n_tau,drawn_tau)\n",
    "            # Add fit metadata\n",
    "            cv_fit_df['generating_area']= generating_area\n",
    "            cv_fit_df['session_id']= session_id\n",
    "            # Add fit comparisons\n",
    "            ll_exp = behlogit_fit.sel(mdl_name='exp_r').isel(session_id=session_id).CV_mean_loglik.item()\n",
    "            ll_hyp = behlogit_fit.sel(mdl_name='hyp_r').isel(session_id=session_id).CV_mean_loglik.item()\n",
    "            aic_exp = behlogit_fit.sel(mdl_name='exp_r').isel(session_id=session_id).CV_mean_aic.item()\n",
    "            aic_hyp = behlogit_fit.sel(mdl_name='hyp_r').isel(session_id=session_id).CV_mean_aic.item()\n",
    "            pa_exp = behlogit_fit.sel(mdl_name='exp_r').isel(session_id=session_id).CV_mean_pa.item()\n",
    "            pa_hyp = behlogit_fit.sel(mdl_name='hyp_r').isel(session_id=session_id).CV_mean_pa.item()\n",
    "            cv_fit_df['loglik_dexp'] = cv_fit_df['loglik'] - ll_exp\n",
    "            cv_fit_df['loglik_dhyp'] = cv_fit_df['loglik'] - ll_hyp\n",
    "            cv_fit_df['aic_dexp'] = cv_fit_df['aic'] - aic_exp\n",
    "            cv_fit_df['aic_dhyp'] = cv_fit_df['aic'] - aic_hyp\n",
    "            cv_fit_df['pa_dexp'] = cv_fit_df['pa'] - pa_exp\n",
    "            cv_fit_df['pa_dhyp'] = cv_fit_df['pa'] - pa_hyp\n",
    "            \n",
    "            if initialize_df:\n",
    "                performance_cv_df = cv_fit_df\n",
    "                initialize_df = False\n",
    "            else:\n",
    "                performance_cv_df = performance_cv_df.append(cv_fit_df, ignore_index=True)\n",
    "                \n",
    "    output_xarray = performance_cv_df.set_index(['generating_area','n_tau','session_id','draw_iter']).to_xarray()\n",
    "    \n",
    "    file_savename = pjoin(behavior_dir,'quasihyp_subsample_cv.nc')\n",
    "    output_xarray.to_netcdf(file_savename) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c40a829-33a8-491d-bb2b-285c0413f07e",
   "metadata": {},
   "source": [
    "# Opto analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dc2e2e2-fefa-4528-82ea-e39f1faafe67",
   "metadata": {},
   "source": [
    "## General import functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4343a58c-c2d2-4437-aa68-8cd74f7b9fc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reset -f\n",
    "import sys\n",
    "from os.path import dirname, join as pjoin\n",
    "from os import listdir\n",
    "sys.path.append('C:\\\\jupyter_notebooks\\\\Danskin_SciAdv_2023\\\\py_code') # set local directory\n",
    "import numpy as np\n",
    "import pickle\n",
    "import random\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "\n",
    "import h5py\n",
    "import scipy.io as sio\n",
    "from scipy import stats\n",
    "import statsmodels as sm # import statsmodels.api as sm\n",
    "from scipy.optimize import minimize, basinhopping, curve_fit\n",
    "from sklearn.model_selection import KFold\n",
    "from tqdm.notebook import tqdm, trange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3e5352a3-37ea-4254-be5a-68cd40904b9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Directories\n",
    "%autoreload\n",
    "import bdanskin as BD \n",
    "\n",
    "project_dir = 'C:\\\\jupyter_notebooks\\\\Danskin_SciAdv_2023' # local directory\n",
    "opto_dir = pjoin(project_dir, 'chr2_optogenetics')\n",
    "behavior_dir = pjoin(project_dir, 'hattori_datasets_behavior')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5ae1931-28ed-47a9-a152-e2c32ac91b4f",
   "metadata": {},
   "source": [
    "## Winstay, loseswitch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3a8afe1-95a0-4843-b3c7-35a348e2225b",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79b328d5-959e-4b44-bc3e-df95a17af2b5",
   "metadata": {},
   "source": [
    "#### pandas winstay, loseswitch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "319a7d13-307d-462a-8314-b2862e645963",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ws_ls_opto_metrics_pd(row):\n",
    "    inac,a,R = row[['Inac Type','a','R']]\n",
    "    # print(a, R)\n",
    "    if inac=='ra':\n",
    "        opto = row['Opto']\n",
    "    else:\n",
    "        opto = row['Opto Shift']\n",
    "        \n",
    "    pre_opto = np.append(opto[1:],0)\n",
    "    post_opto = np.append(0,opto[:-1])  \n",
    "  \n",
    "    choice = (a==1)+0\n",
    "    choice[a==2] = -1\n",
    "    previous_choice = np.append(0, choice[:-1])\n",
    "    stay_vec = ((choice==previous_choice) & choice!=0)+0\n",
    "    switch_vec = ((choice!=previous_choice) & (choice!=0) & (previous_choice!=0))+0\n",
    "    post_win = np.append(0, R[:-1])\n",
    "    post_lose = np.append(0, R[:-1]==0)\n",
    "    post_win[choice==0] = 0\n",
    "    post_lose[choice==0] = 0\n",
    "    post_lose[previous_choice==0] = 0\n",
    "    \n",
    "    p_switch_ctrl = sum(switch_vec  & (opto==0))/sum((choice!=0) & (previous_choice!=0) & (opto==0))\n",
    "    p_switch_opto = sum(switch_vec  & (opto==1))/sum((choice!=0) & (previous_choice!=0) & (opto==1))\n",
    "    p_switch_preopto = sum(switch_vec  & (pre_opto==1))/sum((choice!=0) & (previous_choice!=0) & (pre_opto==1))\n",
    "    p_switch_postopto = sum(switch_vec  & (post_opto==1))/sum((choice!=0) & (previous_choice!=0) & (post_opto==1))\n",
    "\n",
    "    p_stay_ctrl = sum(stay_vec & (opto==0))/sum((choice!=0) & (previous_choice!=0) & (opto==0))\n",
    "    p_stay_opto = sum(stay_vec & (opto==1))/sum((choice!=0) & (previous_choice!=0) & (opto==1))\n",
    "    p_stay_preopto = sum(stay_vec & (pre_opto==1))/sum((choice!=0) & (previous_choice!=0) & (pre_opto==1))\n",
    "    p_stay_postopto = sum(stay_vec & (post_opto==1))/sum((choice!=0) & (previous_choice!=0) & (post_opto==1))\n",
    "\n",
    "    ws_ctrl = sum(post_win & stay_vec & (opto==0)) / sum(post_win & (opto==0))\n",
    "    ws_opto = sum(post_win & stay_vec & (opto==1)) / sum(post_win & (opto==1))\n",
    "    ws_preopto = sum(post_win & stay_vec & (pre_opto==0)) / sum(post_win & (pre_opto==0))\n",
    "    ws_postopto = sum(post_win & stay_vec & (post_opto==0)) / sum(post_win & (post_opto==0))\n",
    "\n",
    "    ws_ctrl = sum(post_win & stay_vec & (opto==0)) / sum(post_win & (opto==0))\n",
    "    ws_opto = sum(post_win & stay_vec & (opto==1)) / sum(post_win & (opto==1))\n",
    "    ws_preopto = sum(post_win & stay_vec & (pre_opto==0)) / sum(post_win & (pre_opto==0))\n",
    "    ws_postopto = sum(post_win & stay_vec & (post_opto==0)) / sum(post_win & (post_opto==0))\n",
    "\n",
    "    ls_ctrl = sum(post_lose & switch_vec & (opto==0)) / sum(post_lose & (opto==0))\n",
    "    ls_opto = sum(post_lose & switch_vec & (opto==1)) / sum(post_lose & (opto==1))\n",
    "    ls_preopto = sum(post_lose & switch_vec & (pre_opto==0)) / sum(post_lose & (pre_opto==0))\n",
    "    ls_postopto = sum(post_lose & switch_vec & (post_opto==0)) / sum(post_lose & (post_opto==0))\n",
    "    \n",
    "    row['p_stay_ctrl'] = p_stay_ctrl\n",
    "    row['p_stay_preopto'] = p_stay_preopto\n",
    "    row['p_stay_opto'] = p_stay_opto\n",
    "    row['p_stay_postopto'] = p_stay_postopto\n",
    "    row['p_switch_ctrl'] = p_switch_ctrl\n",
    "    row['p_switch_preopto'] = p_switch_preopto\n",
    "    row['p_switch_opto'] = p_switch_opto\n",
    "    row['p_switch_postopto'] = p_switch_postopto\n",
    "    row['ws_ctrl'] = ws_ctrl\n",
    "    row['ws_preopto'] = ws_preopto\n",
    "    row['ws_opto'] = ws_opto\n",
    "    row['ws_postopto'] = ws_postopto\n",
    "    row['ls_ctrl'] = ls_ctrl\n",
    "    row['ls_preopto'] = ls_preopto\n",
    "    row['ls_opto'] = ls_opto\n",
    "    row['ls_postopto'] = ls_postopto\n",
    "    row['ws_stay_ctrl'] = ws_ctrl/p_stay_ctrl\n",
    "    row['ws_stay_preopto'] = ws_preopto/p_stay_preopto\n",
    "    row['ws_stay_opto'] = ws_opto/p_stay_opto\n",
    "    row['ws_stay_postopto'] = ws_postopto/p_stay_postopto\n",
    "    row['ls_switch_ctrl'] = ls_ctrl/p_switch_ctrl\n",
    "    row['ls_switch_preopto'] = ls_preopto/p_switch_preopto\n",
    "    row['ls_switch_opto'] = ls_opto/p_switch_opto\n",
    "    row['ls_switch_postopto'] = ls_postopto/p_switch_postopto\n",
    "    \n",
    "    row['ws_stay_ctrl_prime'] = ws_ctrl/np.mean([ws_ctrl, 1-ls_ctrl])\n",
    "    row['ws_stay_preopto_prime'] = ws_preopto/np.mean([ws_preopto, 1-ls_preopto])\n",
    "    row['ws_stay_opto_prime'] = ws_opto/np.mean([ws_opto, 1-ls_opto])\n",
    "    row['ws_stay_postopto_prime'] = ws_postopto/np.mean([ws_postopto, 1-ls_postopto])\n",
    "    row['ls_switch_ctrl_prime'] = ls_ctrl/np.mean([ls_ctrl, 1-ws_ctrl])\n",
    "    row['ls_switch_preopto_prime'] = ls_preopto/np.mean([ls_preopto, 1-ws_preopto])\n",
    "    row['ls_switch_opto_prime'] = ls_opto/np.mean([ls_opto, 1-ws_opto])\n",
    "    row['ls_switch_postopto_prime'] = ls_postopto/np.mean([ls_postopto, 1-ws_postopto])\n",
    "\n",
    "    return row"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ee6f801-8266-48c0-89eb-7682150dd659",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cc8e0011-4a99-4dfd-a9ca-6d4444ae1841",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load session summary\n",
    "chr2_session_summary = pd.read_pickle(pjoin(opto_dir,'ChR2_session_summary.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66d02bc1-6cd0-4995-8fdb-67d467af33be",
   "metadata": {},
   "source": [
    "### Iterate across model fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f00e58a4-7ee4-4f63-821b-cc8af4084c1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa1635021ff14f4bb0036971ead46d84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/255 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "\n",
    "fit_df = chr2_session_summary.progress_apply(ws_ls_opto_metrics_pd, axis=1)\n",
    "chr2_ws_ls = fit_df.drop(columns=['a','R','Opto','Opto Shift'])\n",
    "\n",
    "chr2_ws_ls.to_pickle(pjoin(opto_dir, 'ChR2_wsls.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d327fa20-2a2c-4427-a938-706940f80609",
   "metadata": {},
   "source": [
    "## Logistic regression"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a90572c-07bb-49c2-8536-48ca3794b8d3",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1db8f55c-9075-4ea6-bccf-b4e6e60bee9b",
   "metadata": {},
   "source": [
    "#### pandas logistic modelfit (subsample 1000x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "4e94dfa6-9f44-407f-9617-dbe8ad2639e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegressionCV\n",
    "from sklearn.linear_model import LogisticRegression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "63ecc2d6-2a22-4ff9-ab0e-42855f8d02c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def modelfit_opto_logistic_subsample_pd(row, mdl_type, reg_type):\n",
    "    rng = np.random.default_rng(220727)\n",
    "    # rng = np.random.default_rng(100000) \n",
    "    \n",
    "    n_back = 5\n",
    "    inac,a,R = row[['Inac Type','a','R']]\n",
    "    if inac=='ra':\n",
    "        Opto = row['Opto']\n",
    "    else:\n",
    "        Opto = row['Opto Shift']   \n",
    "    n_choice = sum((a==1)|(a==2))\n",
    "    choice_trials = (a==1)|(a==2)\n",
    "    \n",
    "    predictor_df = BD.prepare_opto_predictor_df(a,R,Opto)\n",
    "    glm_mat, glm_target, opto_param,_  = BD.prepare_opto_glm_5(predictor_df, mdl_type)\n",
    "    \n",
    "    # Subsample\n",
    "    n_iterations = 1000\n",
    "    sample_size = 0.9\n",
    "    opto_inds = np.where(opto_param)[0]\n",
    "    ctrl_inds = np.where(opto_param==0)[0]\n",
    "    n_opto = len(opto_inds)\n",
    "    n_per_iteration = np.floor(n_opto*sample_size).astype(int)\n",
    "    n_per_test = n_opto-n_per_iteration\n",
    "    \n",
    "    for nn in range(n_iterations):\n",
    "        perm_ctrl = rng.permutation(ctrl_inds)\n",
    "        perm_opto = rng.permutation(opto_inds)\n",
    "        ctrl_train = perm_ctrl[:n_per_iteration]\n",
    "        ctrl_test = perm_ctrl[n_per_iteration:n_per_iteration+n_per_test]\n",
    "        opto_train = perm_opto[:n_per_iteration]\n",
    "        opto_test = perm_opto[n_per_iteration:n_per_iteration+n_per_test]\n",
    "        \n",
    "        train_inds = np.concatenate([opto_train, ctrl_train])\n",
    "        train_inds.sort()\n",
    "        test_inds = np.concatenate([opto_test, ctrl_test])\n",
    "        test_inds.sort()\n",
    "        n_test = len(test_inds)\n",
    "        \n",
    "        # Estimate regression\n",
    "        if reg_type=='none':\n",
    "            log_reg = LogisticRegression(solver='newton-cg', penalty='none', n_jobs=-1, random_state=0)\n",
    "        else: \n",
    "            log_reg = LogisticRegressionCV(solver='saga', cv=5, penalty=reg_type, n_jobs=-1, random_state=0, refit=True)\n",
    "        fit_mdl = log_reg.fit(glm_mat[train_inds,:],glm_target[train_inds])    \n",
    "        \n",
    "        coef = fit_mdl.coef_[0]\n",
    "        intercept = fit_mdl.intercept_[0]\n",
    "        temp_coef = coef.copy()\n",
    "        temp_coef[-1] = temp_coef[-1]+intercept\n",
    "        mdl_coef = np.append(temp_coef, intercept)\n",
    "\n",
    "        loglik_train, pa_train = BD.logistic_calc_loglik_pa(glm_mat[train_inds,:], glm_target[train_inds], mdl_coef)\n",
    "        loglik_test, pa_test = BD.logistic_calc_loglik_pa(glm_mat[test_inds,:], glm_target[test_inds], mdl_coef)\n",
    "        oll_0, cll_0, opa_0, cpa_0 = BD.logistic_opto_calc_loglik(glm_mat[train_inds,:], glm_target[train_inds], mdl_coef)\n",
    "        oll_1, cll_1, opa_1, cpa_1 = BD.logistic_opto_calc_loglik(glm_mat[test_inds,:], glm_target[test_inds], mdl_coef)\n",
    "            \n",
    "        if nn==0:\n",
    "            all_coef_mat = mdl_coef.copy()\n",
    "            ll_train_all= loglik_train.copy()\n",
    "            pa_train_all = pa_train.copy()\n",
    "            ll_test_all= loglik_test.copy()\n",
    "            pa_test_all = pa_test.copy()\n",
    "            ll_train_ctrl = cll_0.copy()\n",
    "            ll_train_opto = oll_0.copy()\n",
    "            pa_train_ctrl = cpa_0.copy()\n",
    "            pa_train_opto = opa_0.copy()\n",
    "            ll_test_ctrl = cll_1.copy()\n",
    "            ll_test_opto = oll_1.copy()\n",
    "            pa_test_ctrl = cpa_1.copy()\n",
    "            pa_test_opto = opa_1.copy()\n",
    "\n",
    "        else:\n",
    "            all_coef_mat = np.vstack((all_coef_mat, mdl_coef))\n",
    "            ll_train_all = np.append(ll_train_all, loglik_train)\n",
    "            pa_train_all = np.append(pa_train_all, pa_train)\n",
    "            ll_test_all = np.append(ll_test_all, loglik_test)\n",
    "            pa_test_all = np.append(pa_test_all, pa_test)\n",
    "            ll_train_ctrl = np.append(ll_train_ctrl, cll_0)\n",
    "            ll_train_opto = np.append(ll_train_opto, oll_0)\n",
    "            pa_train_ctrl = np.append(pa_train_ctrl, cpa_0)\n",
    "            pa_train_opto = np.append(pa_train_opto, opa_0)\n",
    "            ll_test_ctrl = np.append(ll_test_ctrl, cll_1)\n",
    "            ll_test_opto = np.append(ll_test_opto, oll_1)\n",
    "            pa_test_ctrl = np.append(pa_test_ctrl, cpa_1)\n",
    "            pa_test_opto = np.append(pa_test_opto, opa_1)\n",
    "    \n",
    "    ll_train_all[ll_train_all==-np.inf] = np.nan\n",
    "    ll_test_all[ll_test_all==-np.inf] = np.nan\n",
    "    ll_train_ctrl[ll_train_ctrl==-np.inf] = np.nan\n",
    "    ll_train_opto[ll_train_opto==-np.inf] = np.nan\n",
    "    ll_test_ctrl[ll_test_ctrl==-np.inf] = np.nan\n",
    "    ll_test_opto[ll_test_opto==-np.inf] = np.nan\n",
    "    frac_opto_neg = calc_bootstrap_onetail(all_coef_mat, mdl_type)\n",
    "    \n",
    "    row['mdl_type'] = mdl_type\n",
    "    row['mdl_hist'] = n_back\n",
    "    row['mdl_penalty'] = reg_type\n",
    "    row['mdl_subsample'] = n_iterations\n",
    "    row['mean_weights'] = np.nanmean(all_coef_mat,0)\n",
    "    row['median_weights'] = np.nanmedian(all_coef_mat,0)\n",
    "    row['boot_opto_neg'] = frac_opto_neg\n",
    "    row['loglik_train'] = np.nanmean(ll_train_all)\n",
    "    row['loglik_train_ctrl'] = np.nanmean(ll_train_ctrl)\n",
    "    row['loglik_train_opto'] = np.nanmean(ll_train_opto)\n",
    "    row['loglik_test'] = np.nanmean(ll_test_all)\n",
    "    row['loglik_test_ctrl'] = np.nanmean(ll_test_ctrl)\n",
    "    row['loglik_test_opto'] = np.nanmean(ll_test_opto)\n",
    "    row['pa_train'] = np.nanmean(pa_train_all)\n",
    "    row['pa_train_ctrl'] = np.nanmean(pa_train_ctrl)\n",
    "    row['pa_train_opto'] = np.nanmean(pa_train_opto)\n",
    "    row['pa_test'] = np.nanmean(pa_test_all)\n",
    "    row['pa_test_ctrl'] = np.nanmean(pa_test_ctrl)\n",
    "    row['pa_test_opto'] = np.nanmean(pa_test_opto)\n",
    "    \n",
    "    # Full session fit\n",
    "    if reg_type=='none':\n",
    "        log_reg = LogisticRegression(solver='newton-cg', penalty='none', n_jobs=-1, random_state=0)\n",
    "    else: \n",
    "        log_reg = LogisticRegressionCV(solver='saga', cv=5, penalty=reg_type, n_jobs=-1, random_state=0, refit=True)\n",
    "    fit_mdl = log_reg.fit(glm_mat,glm_target)    \n",
    "\n",
    "    coef = fit_mdl.coef_[0]\n",
    "    intercept = fit_mdl.intercept_[0]\n",
    "    temp_coef = coef.copy()\n",
    "    temp_coef[-1] = temp_coef[-1]+intercept\n",
    "    mdl_coef = np.append(temp_coef, intercept)\n",
    "    loglik_full, pa_full = BD.logistic_calc_loglik_pa(glm_mat, glm_target, mdl_coef)\n",
    "\n",
    "    row['full_weights'] = mdl_coef\n",
    "    row['loglik_full'] = loglik_full\n",
    "    row['pa_full'] = pa_full\n",
    "    return row"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8581e8e2-1c63-40ce-b450-a4688de6a751",
   "metadata": {},
   "source": [
    "#### calculate opto stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "87455e08-136b-463a-aa3f-04b9edb0a648",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_bootstrap_onetail(coef_mat, mdl_type):\n",
    "    # Assumed 5-history model\n",
    "    if (mdl_type=='RU')|(mdl_type=='RC'):\n",
    "        delta_weights = coef_mat[:,slice(10,20)] - coef_mat[:,slice(0,10)]\n",
    "        delta_weights = np.column_stack((delta_weights, (coef_mat[:,-2] - coef_mat[:,-1])))\n",
    "        mean_bool_delta = np.mean(delta_weights<0,axis=0)\n",
    "    elif (mdl_type=='RUC'):\n",
    "        delta_weights = coef_mat[:,slice(15,30)] - coef_mat[:,slice(0,15)]\n",
    "        delta_weights = np.column_stack((delta_weights, (coef_mat[:,-2] - coef_mat[:,-1])))\n",
    "        mean_bool_delta = np.mean(delta_weights<0,axis=0)\n",
    "        \n",
    "    return mean_bool_delta"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3170c35-4381-4a49-825a-25a3a1e9f4ee",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "7e2deb27-3232-48bd-98e0-2c8f78834251",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load session summary\n",
    "chr2_session_summary = pd.read_pickle(pjoin(opto_dir,'ChR2_session_summary.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f81ba221-98a1-48a5-b333-b64b317accc1",
   "metadata": {},
   "source": [
    "### Iterate across model fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "565ec8e6-8893-4c4f-8c06-c6e4c6fce80f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fitting RU model...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0c49327ccd714f6c9db0c0ca6702b5ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tqdm.pandas()\n",
    "temp_df = chr2_session_summary#.iloc[:3]\n",
    "\n",
    "# modelfit_opto_logistic_subsample_pd(row, mdl_type, reg_type)\n",
    "print('Fitting RU model...')\n",
    "fit_df = temp_df.progress_apply(modelfit_opto_logistic_subsample_pd,axis=1, args=['RU','none'])\n",
    "output_reg = fit_df.drop(columns=['a','R','Opto','Opto Shift'])\n",
    "fit_df = temp_df.progress_apply(modelfit_opto_logistic_subsample_pd,axis=1, args=['RU','l1'])\n",
    "output_reg = output_reg.append(fit_df.drop(columns=['a','R','Opto','Opto Shift']), ignore_index=False)\n",
    "fit_df = temp_df.progress_apply(modelfit_opto_logistic_subsample_pd,axis=1, args=['RU','l2'])\n",
    "output_reg = output_reg.append(fit_df.drop(columns=['a','R','Opto','Opto Shift']), ignore_index=False)\n",
    "\n",
    "print('  done!')\n",
    "output_reg.to_pickle(pjoin(opto_dir, 'logistic_ChR2_l1l2_subsample1000.pkl'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fba7902a-0448-47b1-85d5-e065583851db",
   "metadata": {},
   "source": [
    "## Behlogits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a1a6994-b6a1-49b9-90dc-5693822fdd0a",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2811d1ca-641b-409d-9ddb-16a593131641",
   "metadata": {},
   "source": [
    "#### behlogits local (extra bias terms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f29add98-37b4-4fbe-89e3-551fa4b6aa95",
   "metadata": {},
   "outputs": [],
   "source": [
    "def logit_exp_r_w(params,a,a_hist,R_hist,uR_hist,session_mat,fit_bool):\n",
    "    n_back = np.shape(a_hist)[1]\n",
    "    beta_rc = params[0]\n",
    "    tau_rc = params[1]\n",
    "    const = params[2:]\n",
    "    \n",
    "    past_r = np.multiply(a_hist,R_hist) # already flipped\n",
    "    past_ur = np.multiply(a_hist,uR_hist)\n",
    "    g_rc = np.flip(np.exp(-np.arange(0,n_back)/tau_rc)) \n",
    "    g_rc[np.isnan(g_rc)]=0.  # Catch to remove nans\n",
    "    if sum(g_rc)!=0:\n",
    "        g_rc = g_rc/sum(g_rc) \n",
    "\n",
    "    # convert the sum into a probability, calculate negative log likelihood\n",
    "    sum_weights = (beta_rc*(np.dot(past_r,g_rc))  + np.dot(session_mat,const))\n",
    "    p_logit = 1/(1+np.exp(-sum_weights))\n",
    "    nloglik = -(np.nansum(np.log(p_logit[a==1]))+np.nansum(np.log(1-p_logit[a==0])))\n",
    "    \n",
    "    # If this is for fitting, return only nloglik\n",
    "    if fit_bool:\n",
    "        return nloglik\n",
    "    else:\n",
    "        return p_logit, nloglik\n",
    "    \n",
    "def logit_hyp_r_w(params,a,a_hist,R_hist,uR_hist,session_mat,fit_bool):\n",
    "    n_back = np.shape(a_hist)[1]\n",
    "    beta_rc = params[0]\n",
    "    tau_rc = params[1]\n",
    "    const = params[2:]\n",
    "    \n",
    "    past_r = np.multiply(a_hist,R_hist) # already flipped\n",
    "    past_ur = np.multiply(a_hist,uR_hist)\n",
    "    g_rc = np.flip(1/(1+np.arange(0,n_back)/tau_rc))\n",
    "    g_rc[np.isnan(g_rc)]=0.  # Catch to remove nans\n",
    "    if sum(g_rc)!=0:\n",
    "        g_rc = g_rc/sum(g_rc) \n",
    "    \n",
    "    # convert the sum into a probability, calculate negative log likelihood\n",
    "    sum_weights = (beta_rc*(np.dot(past_r,g_rc)) + np.dot(session_mat,const))\n",
    "    p_logit = 1/(1+np.exp(-sum_weights))\n",
    "    nloglik = -(np.nansum(np.log(p_logit[a==1]))+np.nansum(np.log(1-p_logit[a==0])))\n",
    "    \n",
    "    # If this is for fitting, return only nloglik\n",
    "    if fit_bool:\n",
    "        return nloglik\n",
    "    else:\n",
    "        return p_logit, nloglik"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3748ee8-34aa-44c4-9fb8-0b86c83c43fb",
   "metadata": {},
   "source": [
    "#### prepare design matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "28672f86-5d02-45b0-82c3-112db32a3198",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_opto_predictor_df(a,R,Opto,area):\n",
    "    a_hist, R_hist, uR_hist = BD.prepare_hist_matrix(a,R,10)\n",
    "\n",
    "    # Codify the dependent variable: choice\n",
    "    dv = (a==1).astype(int)\n",
    "    dv = dv.reshape([len(dv),1])\n",
    "\n",
    "    opto_on = np.where(Opto)[0]\n",
    "    choice_trials = ((a==1) | (a==2)) \n",
    "    opto_choice_trials = ((a==1) | (a==2)) & Opto\n",
    "    opto_inds = np.where(opto_choice_trials)[0]\n",
    "    \n",
    "    opto_choice_trials = opto_choice_trials.reshape([len(opto_choice_trials),1])\n",
    "    \n",
    "    trials_to_use = np.zeros(len(a))\n",
    "    trials_to_use[opto_inds]=1\n",
    "    trials_to_use = trials_to_use.astype(bool)\n",
    "    \n",
    "    if area=='Control':\n",
    "        trial_type = np.zeros([len(a),1])\n",
    "    else:\n",
    "        trial_type = np.ones([len(a),1])     \n",
    "\n",
    "    temp_df = pd.DataFrame(np.hstack([dv[trials_to_use], \n",
    "                                      trial_type[trials_to_use],\n",
    "                                      a_hist[trials_to_use,:],\n",
    "                                      R_hist[trials_to_use,:],\n",
    "                                      uR_hist[trials_to_use,:]]),\n",
    "                  columns=['DV','opto',\n",
    "                           'ap10','ap9','ap8','ap7','ap6','ap5','ap4','ap3','ap2','ap1',\n",
    "                           'Rp10','Rp9','Rp8','Rp7','Rp6','Rp5','Rp4','Rp3','Rp2','Rp1',\n",
    "                           'uRp10','uRp9','uRp8','uRp7','uRp6','uRp5','uRp4','uRp3','uRp2','uRp1'])\n",
    "\n",
    "    return temp_df  \n",
    "\n",
    "def prepare_concatenated_df_by_mouse(mouse_sub_df):\n",
    "    n_sessions = len(mouse_sub_df)\n",
    "    for nn in range(n_sessions):\n",
    "        temp_row = mouse_sub_df.iloc[nn]\n",
    "        area,inac,a,R = temp_row[['Area','Inac Type','a','R']]\n",
    "        if inac=='ra':\n",
    "            Opto = temp_row['Opto']\n",
    "        else:\n",
    "            Opto = temp_row['Opto Shift'] \n",
    "    \n",
    "        temp_df = prepare_opto_predictor_df(a,R,Opto,area)\n",
    "        temp_df['ss_'+str(nn)] = np.ones(len(temp_df))\n",
    "        \n",
    "        if nn==0:\n",
    "            out_df = temp_df.copy()\n",
    "        else:\n",
    "            out_df = out_df.append(temp_df)\n",
    "            \n",
    "    out_df = out_df.fillna(0)\n",
    "    out_df['s_all'] = 1\n",
    "    return out_df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e438d5e3-cf2d-4395-ab44-5c50b7c47040",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_RUC(predictor_df,n_back):\n",
    "    # Prepare matrices\n",
    "    if n_back==10:\n",
    "        temp_a_hist = predictor_df[['ap10','ap9','ap8','ap7','ap6','ap5','ap4','ap3','ap2','ap1']].to_numpy()\n",
    "        temp_R_hist = predictor_df[['Rp10','Rp9','Rp8','Rp7','Rp6','Rp5','Rp4','Rp3','Rp2','Rp1']].to_numpy()\n",
    "        temp_uR_hist = predictor_df[['uRp10','uRp9','uRp8','uRp7','uRp6','uRp5','uRp4','uRp3','uRp2','uRp1']].to_numpy()\n",
    "    elif n_back==5:\n",
    "        temp_a_hist = predictor_df[['ap5','ap4','ap3','ap2','ap1']].to_numpy()\n",
    "        temp_R_hist = predictor_df[['Rp5','Rp4','Rp3','Rp2','Rp1']].to_numpy()\n",
    "        temp_uR_hist = predictor_df[['uRp5','uRp4','uRp3','uRp2','uRp1']].to_numpy()\n",
    "    \n",
    "    session_mat = predictor_df.iloc[:,32:].to_numpy()\n",
    "    glm_target = predictor_df['DV'].values\n",
    "    return glm_target, temp_a_hist, temp_R_hist, temp_uR_hist, session_mat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5187a788-afd6-4e34-ada7-f9fb05cc1c87",
   "metadata": {},
   "source": [
    "#### modelfit function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c9b22294-b01d-4eea-9ceb-afde2e298550",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_subsample_behlogit(temp_beh, area, fit_method, n_iterations, n_back):\n",
    "    rng = np.random.default_rng(220202)\n",
    "    # opt_catch = True\n",
    "    predictor_df = prepare_concatenated_df_by_mouse(temp_beh)    \n",
    "    glm_target, a_hist, R_hist, uR_hist, session_mat = generate_RUC(predictor_df,n_back)\n",
    "    \n",
    "    # Model specs\n",
    "    n_bias = np.shape(session_mat)[1]\n",
    "    n_params = 2+n_bias\n",
    "    initial_params = np.append([1,2],np.zeros(n_bias))\n",
    "    bounds = np.array([[0,np.inf],[0,20]])\n",
    "    bias_bounds = np.zeros([n_bias,2]); bias_bounds[:,0] = -np.inf; bias_bounds[:,1] = np.inf; \n",
    "    bounds = np.vstack([bounds,bias_bounds])\n",
    "    \n",
    "    # Subsample\n",
    "    sample_size = 0.9\n",
    "    opto_bool = predictor_df['opto'].values\n",
    "    opto_inds = np.where(opto_bool)[0]\n",
    "    ctrl_inds = np.where(opto_bool==0)[0]\n",
    "    n_condition = np.min([len(opto_inds),len(ctrl_inds)])\n",
    "    n_per_iteration = np.floor(n_condition*sample_size).astype(int)\n",
    "    n_per_test = n_condition-n_per_iteration\n",
    "    print(len(opto_inds), len(ctrl_inds),n_per_iteration)\n",
    "    \n",
    "    for nn in range(n_iterations):\n",
    "        perm_ctrl = rng.permutation(ctrl_inds)\n",
    "        perm_opto = rng.permutation(opto_inds)\n",
    "        ctrl_train = perm_ctrl[:n_per_iteration]\n",
    "        ctrl_test = perm_ctrl[n_per_iteration:n_per_iteration+n_per_test]\n",
    "        opto_train = perm_opto[:n_per_iteration]\n",
    "        opto_test = perm_opto[n_per_iteration:]\n",
    "        \n",
    "        # Fit on opto trials\n",
    "        if fit_method=='basinhopping':\n",
    "            minimizer_kwargs = {\"method\":'L-BFGS-B', \"bounds\":bounds,\n",
    "                                \"args\": (glm_target[opto_train], a_hist[opto_train,:],\n",
    "                                         R_hist[opto_train,:], uR_hist[opto_train,:], \n",
    "                                         session_mat[opto_train,:], 1)}\n",
    "            \n",
    "            # Fit exponential\n",
    "            beh_func = logit_exp_r_w\n",
    "            with np.errstate(divide='ignore',invalid='ignore'):\n",
    "                lik_model = basinhopping(beh_func, initial_params, minimizer_kwargs=minimizer_kwargs, niter=10)\n",
    "            opto_exp_train_params = lik_model.x\n",
    "            _,opto_exp_nloglik = beh_func(opto_exp_train_params, glm_target[opto_test],\n",
    "                                          a_hist[opto_test,:], R_hist[opto_test,:],\n",
    "                                          uR_hist[opto_test,:], session_mat[opto_test,:],0)\n",
    "            \n",
    "            # Fit hyperbolic\n",
    "            beh_func = logit_hyp_r_w\n",
    "            with np.errstate(divide='ignore',invalid='ignore'):\n",
    "                lik_model = basinhopping(beh_func, initial_params, minimizer_kwargs=minimizer_kwargs, niter=10)\n",
    "            opto_hyp_train_params = lik_model.x\n",
    "            _,opto_hyp_nloglik = beh_func(opto_hyp_train_params, glm_target[opto_test],\n",
    "                                          a_hist[opto_test,:], R_hist[opto_test,:],\n",
    "                                          uR_hist[opto_test,:], session_mat[opto_test,:],0)\n",
    "        else:\n",
    "            # Fit exponential\n",
    "            beh_func = logit_exp_r_w\n",
    "            lik_model = minimize(beh_func, initial_params, bounds=bounds, method=fit_method,\n",
    "                                 args=(glm_target[opto_train], a_hist[opto_train,:],\n",
    "                                       R_hist[opto_train,:], uR_hist[opto_train,:],\n",
    "                                       session_mat[opto_train,:], 1))\n",
    "            opto_exp_train_params = lik_model.x\n",
    "            _,opto_exp_nloglik = beh_func(opto_exp_train_params, glm_target[opto_test],\n",
    "                                          a_hist[opto_test,:], R_hist[opto_test,:],\n",
    "                                          uR_hist[opto_test,:], session_mat[opto_test,:],0)\n",
    "            \n",
    "            # Fit hyperbolic\n",
    "            beh_func = logit_hyp_r_w\n",
    "            lik_model = minimize(beh_func, initial_params, bounds=bounds, method=fit_method,\n",
    "                                 args=(glm_target[opto_train], a_hist[opto_train,:],\n",
    "                                       R_hist[opto_train,:], uR_hist[opto_train,:],\n",
    "                                       session_mat[opto_train,:], 1))\n",
    "            opto_hyp_train_params = lik_model.x\n",
    "            _,opto_hyp_nloglik = beh_func(opto_hyp_train_params, glm_target[opto_test],\n",
    "                                          a_hist[opto_test,:], R_hist[opto_test,:],\n",
    "                                          uR_hist[opto_test,:], session_mat[opto_test,:],0)\n",
    "        # Fit on ctrl trials\n",
    "        if fit_method=='basinhopping':\n",
    "            minimizer_kwargs = {\"method\":'L-BFGS-B', \"bounds\":bounds,\n",
    "                                \"args\": (glm_target[ctrl_train], a_hist[ctrl_train,:],\n",
    "                                         R_hist[ctrl_train,:], uR_hist[ctrl_train,:], \n",
    "                                         session_mat[ctrl_train,:], 1)}\n",
    "            \n",
    "            # Fit exponential\n",
    "            beh_func = logit_exp_r_w\n",
    "            with np.errstate(divide='ignore',invalid='ignore'):\n",
    "                lik_model = basinhopping(beh_func, initial_params, minimizer_kwargs=minimizer_kwargs, niter=10)\n",
    "            ctrl_exp_train_params = lik_model.x\n",
    "            _,ctrl_exp_nloglik = beh_func(ctrl_exp_train_params, glm_target[ctrl_test],\n",
    "                                          a_hist[ctrl_test,:], R_hist[ctrl_test,:],\n",
    "                                          uR_hist[ctrl_test,:], session_mat[ctrl_test,:],0)\n",
    "            \n",
    "            # Fit hyperbolic\n",
    "            beh_func = logit_hyp_r_w\n",
    "            with np.errstate(divide='ignore',invalid='ignore'):\n",
    "                lik_model = basinhopping(beh_func, initial_params, minimizer_kwargs=minimizer_kwargs, niter=10)\n",
    "            ctrl_hyp_train_params = lik_model.x\n",
    "            _,ctrl_hyp_nloglik = beh_func(ctrl_hyp_train_params, glm_target[ctrl_test],\n",
    "                                          a_hist[ctrl_test,:], R_hist[ctrl_test,:],\n",
    "                                          uR_hist[ctrl_test,:], session_mat[ctrl_test,:],0)\n",
    "        else:\n",
    "            # Fit exponential\n",
    "            beh_func = logit_exp_r_w\n",
    "            lik_model = minimize(beh_func, initial_params, bounds=bounds, method=fit_method,\n",
    "                                 args=(glm_target[ctrl_train], a_hist[ctrl_train,:],\n",
    "                                       R_hist[ctrl_train,:], uR_hist[ctrl_train,:],\n",
    "                                       session_mat[ctrl_train,:], 1))\n",
    "            ctrl_exp_train_params = lik_model.x\n",
    "            _,ctrl_exp_nloglik = beh_func(ctrl_exp_train_params, glm_target[ctrl_test],\n",
    "                                          a_hist[ctrl_test,:], R_hist[ctrl_test,:],\n",
    "                                          uR_hist[ctrl_test,:], session_mat[ctrl_test,:],0)\n",
    "            \n",
    "            # Fit hyperbolic\n",
    "            beh_func = logit_hyp_r_w\n",
    "            lik_model = minimize(beh_func, initial_params, bounds=bounds, method=fit_method,\n",
    "                                 args=(glm_target[ctrl_train], a_hist[ctrl_train,:],\n",
    "                                       R_hist[ctrl_train,:], uR_hist[ctrl_train,:],\n",
    "                                       session_mat[ctrl_train,:], 1))\n",
    "            ctrl_hyp_train_params = lik_model.x\n",
    "            _,ctrl_hyp_nloglik = beh_func(ctrl_hyp_train_params, glm_target[ctrl_test],\n",
    "                                          a_hist[ctrl_test,:], R_hist[ctrl_test,:],\n",
    "                                          uR_hist[ctrl_test,:], session_mat[ctrl_test,:],0)\n",
    "       \n",
    "        opto_delta_loglik = (-opto_hyp_nloglik) - (-opto_exp_nloglik)\n",
    "        ctrl_delta_loglik = (-ctrl_hyp_nloglik) - (-ctrl_exp_nloglik)\n",
    "        \n",
    "        if nn==0:\n",
    "            opto_exp_param_mat = opto_exp_train_params.copy()\n",
    "            ctrl_exp_param_mat = ctrl_exp_train_params.copy()\n",
    "            opto_hyp_param_mat = opto_hyp_train_params.copy()\n",
    "            ctrl_hyp_param_mat = ctrl_hyp_train_params.copy()\n",
    "            opto_dll_mat = opto_delta_loglik\n",
    "            ctrl_dll_mat = ctrl_delta_loglik\n",
    "        else:\n",
    "            opto_exp_param_mat = np.vstack((opto_exp_param_mat, opto_exp_train_params))\n",
    "            ctrl_exp_param_mat = np.vstack((ctrl_exp_param_mat, ctrl_exp_train_params))\n",
    "            opto_hyp_param_mat = np.vstack((opto_hyp_param_mat, opto_hyp_train_params))\n",
    "            ctrl_hyp_param_mat = np.vstack((ctrl_hyp_param_mat, ctrl_hyp_train_params))\n",
    "            opto_dll_mat = np.vstack((opto_dll_mat, opto_delta_loglik))\n",
    "            ctrl_dll_mat = np.vstack((ctrl_dll_mat, ctrl_delta_loglik))\n",
    "            \n",
    "    fit_df = pd.DataFrame({'n_iter': n_iterations, 'mdl_hist': n_back, 'fit_method': fit_method,\n",
    "                           'opt_catch': opt_catch,\n",
    "                           'median_opto_dll': np.nanmedian(opto_dll_mat,axis=0),\n",
    "                           'median_ctrl_dll': np.nanmedian(ctrl_dll_mat,axis=0),\n",
    "                           'mean_opto_dll': np.nanmean(opto_dll_mat,axis=0),\n",
    "                           'mean_ctrl_dll': np.nanmean(ctrl_dll_mat,axis=0),\n",
    "                           'median_opto_dll_trial': np.nanmedian(opto_dll_mat,axis=0)/(n_per_test),\n",
    "                           'median_ctrl_dll_trial': np.nanmedian(ctrl_dll_mat,axis=0)/(n_per_test),\n",
    "                           'median_opto_exp_params': [np.nanmedian(opto_exp_param_mat,axis=0)],\n",
    "                           'median_ctrl_exp_params': [np.nanmedian(ctrl_exp_param_mat,axis=0)],\n",
    "                           'median_opto_hyp_params': [np.nanmedian(opto_hyp_param_mat,axis=0)],\n",
    "                           'median_ctrl_hyp_params': [np.nanmedian(ctrl_hyp_param_mat,axis=0)],\n",
    "                          })\n",
    "\n",
    "    return fit_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "558a0977-be5c-4a3d-b744-8f4403f995eb",
   "metadata": {},
   "source": [
    "#### across-area modelfit iterable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9d5ee5ef-7b2e-4c77-90cb-36ae318e5f1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_subsample_given_area(beh_df, area, m_all, fit_method, n_iterations, n_back,opt_catch):\n",
    "    fit_df = pd.DataFrame([])\n",
    "    # For given area, estimate all sessions for mouse\n",
    "    for mm in m_all:\n",
    "        temp_beh = beh_df.loc[(beh_df['Mouse']==mm) & \n",
    "                      ((beh_df['Area']==area) | (beh_df['Area']=='Control'))]\n",
    "        temp_df = fit_subsample_behlogit(temp_beh, area, fit_method, n_iterations, n_back,opt_catch)\n",
    "        temp_df['Mouse'] = mm\n",
    "        temp_df['Area'] = area\n",
    "        \n",
    "        fit_df = fit_df.append(temp_df)\n",
    "    return fit_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c34a86e-ca26-4fcc-9316-cb6e716c9e54",
   "metadata": {},
   "source": [
    "### Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "c7391eb7-a5ce-4c0e-a23d-2c089267bc14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load session summary\n",
    "chr2_session_summary = pd.read_pickle(pjoin(opto_dir,'ChR2_session_summary.pkl'))\n",
    "\n",
    "beh_df = chr2_session_summary.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ae767f5f-f2a3-4e91-9129-1d39b9eaad5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['BD032' 'BD033' 'BD109' 'BD110' 'BD111' 'BD112' 'BD115' 'BD116' 'BD118']\n"
     ]
    }
   ],
   "source": [
    "# Mouse IDs\n",
    "rsc_m_id = beh_df.loc[(beh_df['Area']=='RSCb'),'Mouse'].unique()\n",
    "ppc_m_id = beh_df.loc[(beh_df['Area']=='PPCb'),'Mouse'].unique()\n",
    "m2_m_id = beh_df.loc[(beh_df['Area']=='M2b'),'Mouse'].unique()\n",
    "loh_m_id = beh_df.loc[(beh_df['Area']=='Control'),'Mouse'].unique()\n",
    "\n",
    "# Group intersection\n",
    "loh_rsc_m_id = np.intersect1d(loh_m_id, rsc_m_id)\n",
    "# rsc_to_use = np.isin(rsc_m_id,loh_rsc_m_id)\n",
    "loh_to_use = np.isin(loh_m_id, loh_rsc_m_id)\n",
    "loh_rsc_m_id = loh_m_id[loh_to_use]\n",
    "\n",
    "loh_ppc_m_id = np.intersect1d(loh_m_id, ppc_m_id)\n",
    "# ppc_to_use = np.isin(ppc_m_id,loh_ppc_m_id)\n",
    "loh_to_use = np.isin(loh_m_id, loh_ppc_m_id)\n",
    "loh_ppc_m_id = loh_m_id[loh_to_use]\n",
    "\n",
    "loh_m2_m_id = np.intersect1d(loh_m_id, m2_m_id)\n",
    "loh_to_use = np.isin(loh_m_id, loh_m2_m_id)\n",
    "loh_m2_m_id = loh_m_id[loh_to_use]\n",
    "\n",
    "print(loh_ppc_m_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f32a6327-a236-4459-8751-7d0c27a01909",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Iterate across areas "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "cf772b23-94fe-4a59-a34b-a8bc32d9c365",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "292 199 179\n",
      "174 68 61\n",
      "66 63 56\n",
      "143 194 128\n",
      "163 103 92\n",
      "142 124 111\n",
      "227 427 204\n",
      "92 241 82\n",
      "101 333 90\n",
      "310 199 179\n",
      "147 68 61\n",
      "176 194 158\n",
      "295 103 92\n",
      "90 124 81\n",
      "121 427 108\n",
      "156 241 140\n",
      "96 333 86\n",
      "151 179 135\n",
      "341 199 179\n",
      "235 68 61\n",
      "85 63 56\n",
      "99 103 89\n",
      "107 427 96\n",
      "57 57 51\n",
      "187 241 168\n",
      "89 333 80\n",
      "264 179 161\n"
     ]
    }
   ],
   "source": [
    "temp_df = fit_subsample_given_area(beh_df, 'RSCb', loh_rsc_m_id, 'basinhopping', 30, 10)\n",
    "behlogit_out = temp_df.copy()\n",
    "temp_df = fit_subsample_given_area(beh_df, 'PPCb', loh_ppc_m_id, 'basinhopping',30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "temp_df = fit_subsample_given_area(beh_df, 'M2b', loh_m2_m_id, 'basinhopping',30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "temp_df = fit_subsample_given_area(beh_df, 'RSCb', loh_rsc_m_id, 'basinhopping', 30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "temp_df = fit_subsample_given_area(beh_df, 'PPCb', loh_ppc_m_id, 'basinhopping',30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "temp_df = fit_subsample_given_area(beh_df, 'M2b', loh_m2_m_id, 'basinhopping',30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "\n",
    "temp_df = fit_subsample_given_area(beh_df, 'RSCb', loh_rsc_m_id, 'Nelder-Mead', 30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "temp_df = fit_subsample_given_area(beh_df, 'PPCb', loh_ppc_m_id, 'Nelder-Mead',30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "temp_df = fit_subsample_given_area(beh_df, 'M2b', loh_m2_m_id, 'Nelder-Mead',30, 10)\n",
    "behlogit_out = behlogit_out.append(temp_df)\n",
    "\n",
    "local_dir = 'C://jupyter_notebooks//processed_data'\n",
    "pickle.dump(behlogit_out,open(pjoin(opto_dir, 'opto_behlogit_RA_10hist.pkl'),'wb'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}