examples/03_decoder.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit and analyze decoder models\n",
    "The next step of the BehaveNet pipeline is using simultaneously recorded neural activity to predict behavior. Specifically, we can predict either of our compressed descriptions of behavior: the convolutional autoencoder latents or the ARHMM states.\n",
    "\n",
    "We use linear models or feedforward deep networks to predict the state or latents for a given frame given a window of neural activity. We then compare our predictions to baseline (aka chance) performance. We can also use the convolutional autoencoder to convert the predicted latents into a full predicted behavioral video and compare to the original behavior.\n",
    "\n",
    "\n",
    "<br>\n",
    "\n",
    "### Contents\n",
    "* [Decoding discrete states](#Decoding-discrete-states)\n",
    "* [Decoding continuous latents](#Decoding-continuous-latents)\n",
    "* [Assess decoding performance](#Assess-decoding-performance)\n",
    "* [Plot true vs predicted latents](#Plot-true-vs-predicted-latents)\n",
    "* [Make real vs predicted movies](#Make-real-vs-predicted-movies)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import scipy\n",
    "import scipy.io as sio\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "\n",
    "matplotlib.rcParams.update({'font.size': 22})\n",
    "from behavenet import get_user_dir, make_dir_if_not_exists\n",
    "from behavenet.data.utils import get_transforms_paths\n",
    "from behavenet.data.utils import load_labels_like_latents\n",
    "from behavenet.fitting.utils import get_expt_dir\n",
    "from behavenet.fitting.utils import get_session_dir\n",
    "from behavenet.fitting.utils import get_best_model_version\n",
    "from behavenet.fitting.utils import get_lab_example\n",
    "from behavenet.plotting.arhmm_utils import *\n",
    "\n",
    "save_outputs = True  # true to save figures/movies to user's figure directory\n",
    "format = 'png'  # figure format ('png' | 'jpeg' | 'pdf'); movies saved as mp4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoding discrete (ARHMM) states\n",
    "\n",
    "First copy the example json files ``decoding_data.json``, ``decoding_ae_model.json``, ``decoding_arhmm_model.json``, ``decoding_training.json`` and ``decoding_compute.json`` into your ``.behavenet`` directory, ``cd`` to the ``behavenet`` directory in the terminal, and run:\n",
    "\n",
    "```console\n",
    "$: python behavenet/fitting/decoder_grid_search.py --data_config ~/.behavenet/decoding_data.json --model_config ~/.behavenet/decoding_arhmm_model.json --training_config ~/.behavenet/decoding_training.json --compute_config ~/.behavenet/decoding_compute.json\n",
    "```\n",
    "\n",
    "[Back to contents](#Contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoding continuous states (AE latents)\n",
    "\n",
    "```console\n",
    "$: python behavenet/fitting/decoder_grid_search.py --data_config ~/.behavenet/decoding_data.json --model_config ~/.behavenet/decoding_ae_model.json --training_config ~/.behavenet/decoding_training.json --compute_config ~/.behavenet/decoding_compute.json\n",
    "```\n",
    "\n",
    "[Back to contents](#Contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assess decoding performance\n",
    "We want to examine how our predictions of both discrete states and continuous states compare to a baseline chance performance.\n",
    "\n",
    "[Back to contents](#Contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from behavenet.fitting.utils import get_subdirs\n",
    "\n",
    "# set model info\n",
    "hparams = {\n",
    "    'data_dir': get_user_dir('data'),\n",
    "    'save_dir': get_user_dir('save'),\n",
    "    'model_class': 'neural-arhmm',\n",
    "    'experiment_name': 'grid_search',\n",
    "    'model_type': 'mlp',\n",
    "    'noise_type': 'gaussian',\n",
    "    'n_arhmm_lags': 1,\n",
    "    'transitions': 'stationary',\n",
    "    'arhmm_experiment_name': 'state_number_search',\n",
    "    'ae_experiment_name': 'ae-example',\n",
    "    'n_ae_latents': 9,\n",
    "    'neural_arhmm_experiment_name':'grid_search',\n",
    "    'n_arhmm_states': 4,\n",
    "    'n_max_lags': 8,\n",
    "    'rng_seed_data': 0,\n",
    "    'trial_splits': '8;1;1;0'\n",
    "}\n",
    "\n",
    "hparams['neural_arhmm_experiment_name'] = hparams['experiment_name']\n",
    "hparams['neural_arhmm_model_type'] = hparams['model_type']\n",
    "\n",
    "get_lab_example(hparams, 'musall', 'vistrained')\n",
    "sess_idx = 0 \n",
    "hparams['session_dir'], sess_ids = get_session_dir(hparams)\n",
    "hparams['expt_dir'] = get_expt_dir(hparams)\n",
    "\n",
    "\n",
    "## Get discrete chance performance (accuracy of always predicting the most common training state)\n",
    "_, states_file = get_transforms_paths('states', hparams, sess_ids[sess_idx])\n",
    "with open(states_file, 'rb') as f:\n",
    "    all_states = pickle.load(f)\n",
    "most_common_train_state = scipy.stats.mode(np.concatenate([all_states['states'][i] for i in all_states['trials']['train']])).mode[0]\n",
    "\n",
    "all_test_states = np.concatenate([all_states['states'][i][hparams['n_max_lags']:-hparams['n_max_lags']] for i in all_states['trials']['test']])\n",
    "chance_arhmm_performance = (all_test_states==0).sum()/all_test_states.shape[0]\n",
    "\n",
    "## Get discrete chance performance (accuracy of always predicting the most common training state)\n",
    "_, states_file = get_transforms_paths('neural_arhmm_predictions', hparams, sess_ids[sess_idx])\n",
    "with open(states_file, 'rb') as f:\n",
    "    all_state_predictions = pickle.load(f)\n",
    "all_test_state_predictions = np.concatenate([np.argmax(all_state_predictions['predictions'][i][hparams['n_max_lags']:-hparams['n_max_lags']],axis=1) for i in all_state_predictions['trials']['test']])\n",
    "decoding_arhmm_performance = (all_test_states==all_test_state_predictions).sum()/all_test_states.shape[0]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set model info\n",
    "sess_idx = 0\n",
    "hparams = {\n",
    "    'data_dir': get_user_dir('data'),\n",
    "    'save_dir': get_user_dir('save'),\n",
    "    'model_class': 'neural-ae',\n",
    "    'ae_model_class': 'ae',\n",
    "    'ae_model_type': 'conv',\n",
    "    'ae_experiment_name': 'ae-example',\n",
    "    'n_ae_latents': 9,\n",
    "    'experiment_name': 'grid_search',\n",
    "    'model_type': 'mlp',\n",
    "    'n_max_lags': 8,\n",
    "    'rng_seed_data': 0,\n",
    "    'trial_splits': '8;1;1;0'\n",
    "}\n",
    "\n",
    "hparams['neural_ae_experiment_name'] = hparams['experiment_name']\n",
    "hparams['neural_ae_model_type'] = hparams['model_type']\n",
    "\n",
    "get_lab_example(hparams, 'musall', 'vistrained')\n",
    "\n",
    "hparams['session_dir'], sess_ids = get_session_dir(hparams)\n",
    "expt_dir = get_expt_dir(hparams)\n",
    "\n",
    "## Get discrete chance performance (accuracy of always predicting the most common training state)\n",
    "_, latents_file = get_transforms_paths('ae_latents', hparams, sess_ids[sess_idx])\n",
    "with open(latents_file, 'rb') as f:\n",
    "    all_latents = pickle.load(f)\n",
    "mean_ae_latents = np.mean(np.concatenate([all_latents['latents'][i] for i in all_latents['trials']['train']]),axis=0)\n",
    "\n",
    "all_test_latents = np.concatenate([all_latents['latents'][i][hparams['n_max_lags']:-hparams['n_max_lags']] for i in all_latents['trials']['test']])\n",
    "chance_ae_performance = np.mean((all_test_latents-mean_ae_latents)**2)\n",
    "\n",
    "## Get discrete chance performance (accuracy of always predicting the most common training state)\n",
    "_, latent_predictions_file = get_transforms_paths('neural_ae_predictions', hparams, sess_ids[sess_idx])\n",
    "with open(latent_predictions_file, 'rb') as f:\n",
    "    all_latent_predictions = pickle.load(f)\n",
    "all_test_latent_predictions = np.concatenate([all_latent_predictions['predictions'][i][hparams['n_max_lags']:-hparams['n_max_lags']] for i in all_latents['trials']['test']])\n",
    "decoding_ae_performance = np.mean((all_test_latents-all_test_latent_predictions)**2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,2, figsize=(10,10))\n",
    "\n",
    "bar = axes[0].bar([0,1], [chance_arhmm_performance, decoding_arhmm_performance])\n",
    "bar[0].set_color('#355C7D')\n",
    "bar[1].set_color('#F67280')\n",
    "bar = axes[1].bar([0,1], [chance_ae_performance, decoding_ae_performance])\n",
    "bar[0].set_color('#355C7D')\n",
    "bar[1].set_color('#F67280')\n",
    "\n",
    "axes[0].set_xticks([0,1])\n",
    "axes[0].set_xticklabels(['Chance','Decoding'])\n",
    "\n",
    "axes[1].set_xticks([0,1])\n",
    "axes[1].set_xticklabels(['Chance','Decoding'])\n",
    "\n",
    "axes[0].set_ylabel('Fraction correct')\n",
    "axes[1].set_ylabel('Mean Squared Error')\n",
    "\n",
    "axes[0].set_title('ARHMM State Decoding')\n",
    "axes[1].set_title('AE Latent Decoding')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot true vs predicted latents\n",
    "\n",
    "[Back to contents](#Contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_real_vs_predicted(\n",
    "        latents, latents_predicted, save_file=None, xtick_locs=None,  frame_rate=None, format='png'):\n",
    "    \"\"\"Plot real and sampled latents overlaying real and (potentially sampled) states.\n",
    "    Parameters\n",
    "    ----------\n",
    "    latents : :obj:`np.ndarray`\n",
    "        shape (n_frames, n_latents)\n",
    "    latents_samp : :obj:`np.ndarray`\n",
    "        shape (n_frames, n_latents)\n",
    "    states : :obj:`np.ndarray`\n",
    "        shape (n_frames,)\n",
    "    states_samp : :obj:`np.ndarray`\n",
    "        shape (n_frames,) if :obj:`latents_samp` are not conditioned on :obj:`states`, otherwise\n",
    "        shape (0,)\n",
    "    save_file : :obj:`str`\n",
    "        full save file (path and filename)\n",
    "    xtick_locs : :obj:`array-like`, optional\n",
    "        tick locations in bin values for plot\n",
    "    frame_rate : :obj:`float`, optional\n",
    "        behavioral video framerate; to properly relabel xticks\n",
    "    format : :obj:`str`, optional\n",
    "        any accepted matplotlib save format, e.g. 'png' | 'pdf' | 'jpeg'\n",
    "    Returns\n",
    "    -------\n",
    "    :obj:`matplotlib.figure.Figure`\n",
    "        matplotlib figure handle\n",
    "    \"\"\"\n",
    "\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(10, 8))\n",
    "\n",
    "        \n",
    "    spc = 1.1 * abs(latents.max())\n",
    "    n_latents = latents.shape[1]\n",
    "    plotting_latents = latents + spc * np.arange(n_latents)\n",
    "    plotting_predicted_latents = latents_predicted + spc * np.arange(n_latents)\n",
    "    ymin = min(-spc - 1, np.min(plotting_latents))\n",
    "    ymax = max(spc * n_latents, np.max(plotting_latents))\n",
    "    ax.plot(plotting_latents, '-k', lw=3, label='AE Latents')\n",
    "    ax.plot(plotting_predicted_latents, '-g', lw=3, label='Predicted AE latents')\n",
    "    ax.set_ylim([ymin, ymax])\n",
    "\n",
    "    ax.set_yticks([])\n",
    "\n",
    "    ax.set_xlabel('Time (bins)')\n",
    "\n",
    "    if xtick_locs is not None:\n",
    "        ax.set_xticks(xtick_locs)\n",
    "        if frame_rate is not None:\n",
    "            ax.set_xticklabels((np.asarray(xtick_locs) / frame_rate).astype('int'))\n",
    "            ax.set_xlabel('Time (sec)')\n",
    "    handles, labels = plt.gca().get_legend_handles_labels()\n",
    "    by_label = dict(zip(labels, handles))\n",
    "    plt.legend(by_label.values(), by_label.keys(),loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "\n",
    "    if save_file is not None:\n",
    "        make_dir_if_not_exists(save_file)\n",
    "        plt.savefig(save_file, dpi=300, format=format)\n",
    "\n",
    "   # return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# user params\n",
    "get_best_version = True  # False when looking at multiple models w/in a tt expt\n",
    "dtype = 'test'  # data type to draw trials from: 'train' | 'val' | 'test'\n",
    "sess_idx = 0  # when using a multisession, this determines which session is used\n",
    "max_frames = 200\n",
    "\n",
    "# define which arhmm states to plot (must already be fit)\n",
    "n_arhmm_states = [2, 4]\n",
    "\n",
    "# set model info\n",
    "hparams = {\n",
    "    'data_dir': get_user_dir('data'),\n",
    "    'save_dir': get_user_dir('save'),\n",
    "    'experiment_name': 'grid_search',\n",
    "    'model_class': 'neural-ae',\n",
    "    'model_type': 'mlp',\n",
    "    'ae_experiment_name': 'ae-example',\n",
    "    'n_ae_latents': 9,\n",
    "    'ae_model_class': 'ae',\n",
    "    'ae_model_type': 'conv',\n",
    "    'rng_seed_data': 0,\n",
    "    'trial_splits': '8;1;1;0'\n",
    "}\n",
    "\n",
    "hparams['neural_ae_experiment_name'] = hparams['experiment_name']\n",
    "hparams['neural_ae_model_type'] = hparams['model_type']\n",
    "\n",
    "get_lab_example(hparams, 'musall', 'vistrained')\n",
    "\n",
    "xtick_locs = [0, 30, 60, 90, 120, 150, 180]\n",
    "frame_rate = 30\n",
    "n_trials = 20\n",
    "           \n",
    "for n_states in n_arhmm_states:\n",
    "        \n",
    "    hparams['n_arhmm_states'] = n_states\n",
    "    hparams['session_dir'], sess_ids = get_session_dir(hparams)\n",
    "    hparams['expt_dir'] = get_expt_dir(hparams)\n",
    "\n",
    "    # get version/model\n",
    "    if get_best_version:\n",
    "        version = get_best_model_version(\n",
    "            hparams['expt_dir'], measure='val_loss', best_def='min')[0]\n",
    "    else:\n",
    "        _, version = experiment_exists(hparams, which_version=True)\n",
    "\n",
    "    # load model\n",
    "    model_file = os.path.join(\n",
    "        hparams['expt_dir'], 'version_%i' % version, 'best_val_model.pt')\n",
    "    with open(model_file, 'rb') as f:\n",
    "        hmm = pickle.load(f)\n",
    "        \n",
    "            \n",
    "    # load latents\n",
    "    _, latents_file = get_transforms_paths('ae_latents', hparams, sess_ids[sess_idx])\n",
    "    with open(latents_file, 'rb') as f:\n",
    "        all_latents = pickle.load(f)\n",
    "\n",
    "        \n",
    "    # load latent predictions\n",
    "    _, latents_file = get_transforms_paths('neural_ae_predictions', hparams, sess_ids[sess_idx])\n",
    "    with open(latents_file, 'rb') as f:\n",
    "        predicted_latents = pickle.load(f)\n",
    "\n",
    "    # choose which trials to plot\n",
    "    np.random.seed(0)\n",
    "    trial_vec = np.random.choice(\n",
    "        np.arange(0, len(all_latents['trials'][dtype])), size=(n_trials,), \n",
    "        replace=False)\n",
    "    \n",
    "model_name = str(\n",
    "    'latent_prediction_D=%02i' % (\n",
    "    hparams['n_ae_latents']))\n",
    "\n",
    "if save_outputs:\n",
    "    save_file = os.path.join(\n",
    "        get_user_dir('fig'), hparams['model_class'], model_name + '.' + format)\n",
    "else:\n",
    "    save_file = None\n",
    "\n",
    "plot_real_vs_predicted(all_latents['latents'][trial_vec[0]], predicted_latents['predictions'][trial_vec[0]], save_file=save_file, xtick_locs=None,  frame_rate=30, format='png')\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "behavenet",
   "language": "python",
   "name": "behavenet"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}