training/auto_notebooks/supervised/ReadySetGo-v0.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/neurogym/ngym_usage/blob/master/supervised/auto_notebooks/supervised/ReadySetGo-v0.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install packages if on Colab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment following lines to install\n",
    "# ! pip install gym   # Install gym\n",
    "# ! git clone https://github.com/gyyang/neurogym.git  # Install neurogym\n",
    "# %cd neurogym/\n",
    "# ! pip install -e .\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.decomposition import PCA\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "import gym\n",
    "import neurogym as ngym\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "\n",
    "envid = 'ReadySetGo-v0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_modelpath(envid):\n",
    "    # Make a local file directories\n",
    "    path = Path('.') / 'files'\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    path = path / envid\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    return path\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(Net, self).__init__()\n",
    "        self.lstm = nn.LSTM(input_size, hidden_size)\n",
    "        self.linear = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out, hidden = self.lstm(x)\n",
    "        x = self.linear(out)\n",
    "        return x, out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Supervised training networks.\n",
    "\n",
    "Save network in a path determined by environment ID.\n",
    "\n",
    "Args:\n",
    "    envid: str, environment ID.\n",
    "\"\"\"\n",
    "modelpath = get_modelpath(envid)\n",
    "config = {\n",
    "    'dt': 100,\n",
    "    'hidden_size': 64,\n",
    "    'lr': 1e-2,\n",
    "    'batch_size': 16,\n",
    "    'seq_len': 100,\n",
    "    'envid': envid,\n",
    "}\n",
    "\n",
    "env_kwargs = {'dt': config['dt']}\n",
    "config['env_kwargs'] = env_kwargs\n",
    "\n",
    "# Save config\n",
    "with open(modelpath / 'config.json', 'w') as f:\n",
    "    json.dump(config, f)\n",
    "\n",
    "# Make supervised dataset\n",
    "dataset = ngym.Dataset(\n",
    "    envid, env_kwargs=env_kwargs, batch_size=config['batch_size'],\n",
    "    seq_len=config['seq_len'])\n",
    "env = dataset.env\n",
    "act_size = env.action_space.n\n",
    "# Train network\n",
    "net = Net(input_size=env.observation_space.shape[0],\n",
    "          hidden_size=config['hidden_size'],\n",
    "          output_size=act_size)\n",
    "net = net.to(device)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=config['lr'])\n",
    "\n",
    "print('Training task ', envid)\n",
    "\n",
    "running_loss = 0.0\n",
    "for i in range(2000):\n",
    "    inputs, labels = dataset()\n",
    "    inputs = torch.from_numpy(inputs).type(torch.float).to(device)\n",
    "    labels = torch.from_numpy(labels.flatten()).type(torch.long).to(device)\n",
    "\n",
    "    # zero the parameter gradients\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # forward + backward + optimize\n",
    "    outputs, _ = net(inputs)\n",
    "\n",
    "    loss = criterion(outputs.view(-1, act_size), labels)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # print statistics\n",
    "    running_loss += loss.item()\n",
    "    if i % 200 == 199:\n",
    "        print('{:d} loss: {:0.5f}'.format(i + 1, running_loss / 200))\n",
    "        running_loss = 0.0\n",
    "        torch.save(net.state_dict(), modelpath / 'net.pth')\n",
    "\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_test_timing(env):\n",
    "    \"\"\"Infer timing of environment for testing.\"\"\"\n",
    "    timing = {}\n",
    "    for period in env.timing.keys():\n",
    "        period_times = [env.sample_time(period) for _ in range(100)]\n",
    "        timing[period] = np.median(period_times)\n",
    "    return timing\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run network after training for analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Run trained networks for analysis.\n",
    "\n",
    "Args:\n",
    "    envid: str, Environment ID\n",
    "\n",
    "Returns:\n",
    "    activity: a list of activity matrices, each matrix has shape (\n",
    "    N_time, N_neuron)\n",
    "    info: pandas dataframe, each row is information of a trial\n",
    "    config: dict of network, training configurations\n",
    "\"\"\"\n",
    "modelpath = get_modelpath(envid)\n",
    "with open(modelpath / 'config.json') as f:\n",
    "    config = json.load(f)\n",
    "\n",
    "env_kwargs = config['env_kwargs']\n",
    "\n",
    "# Run network to get activity and info\n",
    "# Environment\n",
    "env = gym.make(envid, **env_kwargs)\n",
    "env.timing = infer_test_timing(env)\n",
    "env.reset(no_step=True)\n",
    "\n",
    "# Instantiate the network and print information\n",
    "with torch.no_grad():\n",
    "    net = Net(input_size=env.observation_space.shape[0],\n",
    "              hidden_size=config['hidden_size'],\n",
    "              output_size=env.action_space.n)\n",
    "    net = net.to(device)\n",
    "    net.load_state_dict(torch.load(modelpath / 'net.pth'))\n",
    "\n",
    "    perf = 0\n",
    "    num_trial = 100\n",
    "\n",
    "    activity = list()\n",
    "    info = pd.DataFrame()\n",
    "\n",
    "    for i in range(num_trial):\n",
    "        env.new_trial()\n",
    "        ob, gt = env.ob, env.gt\n",
    "        inputs = torch.from_numpy(ob[:, np.newaxis, :]).type(torch.float)\n",
    "        action_pred, hidden = net(inputs)\n",
    "\n",
    "        # Compute performance\n",
    "        action_pred = action_pred.detach().numpy()\n",
    "        choice = np.argmax(action_pred[-1, 0, :])\n",
    "        correct = choice == gt[-1]\n",
    "\n",
    "        # Log trial info\n",
    "        trial_info = env.trial\n",
    "        trial_info.update({'correct': correct, 'choice': choice})\n",
    "        info = info.append(trial_info, ignore_index=True)\n",
    "\n",
    "        # Log stimulus period activity\n",
    "        activity.append(np.array(hidden)[:, 0, :])\n",
    "\n",
    "    print('Average performance', np.mean(info['correct']))\n",
    "\n",
    "activity = np.array(activity)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### General analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_average_activity(activity, info, config):\n",
    "    # Load and preprocess results\n",
    "    plt.figure(figsize=(1.2, 0.8))\n",
    "    t_plot = np.arange(activity.shape[1]) * config['dt']\n",
    "    plt.plot(t_plot, activity.mean(axis=0).mean(axis=-1))\n",
    "\n",
    "analysis_average_activity(activity, info, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_conditions(info):\n",
    "    \"\"\"Get a list of task conditions to plot.\"\"\"\n",
    "    conditions = info.columns\n",
    "    # This condition's unique value should be less than 5\n",
    "    new_conditions = list()\n",
    "    for c in conditions:\n",
    "        try:\n",
    "            n_cond = len(pd.unique(info[c]))\n",
    "            if 1 < n_cond < 5:\n",
    "                new_conditions.append(c)\n",
    "        except TypeError:\n",
    "            pass\n",
    "        \n",
    "    return new_conditions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_activity_by_condition(activity, info, config):\n",
    "    conditions = get_conditions(info)\n",
    "    for condition in conditions:\n",
    "        values = pd.unique(info[condition])\n",
    "        plt.figure(figsize=(1.2, 0.8))\n",
    "        t_plot = np.arange(activity.shape[1]) * config['dt']\n",
    "        for value in values:\n",
    "            a = activity[info[condition] == value]\n",
    "            plt.plot(t_plot, a.mean(axis=0).mean(axis=-1), label=str(value))\n",
    "        plt.legend(title=condition, loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "\n",
    "analysis_activity_by_condition(activity, info, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_example_units_by_condition(activity, info, config):\n",
    "    conditions = get_conditions(info)\n",
    "    if len(conditions) < 1:\n",
    "        return\n",
    "\n",
    "    example_ids = np.array([0, 1])    \n",
    "    for example_id in example_ids:        \n",
    "        example_activity = activity[:, :, example_id]\n",
    "        fig, axes = plt.subplots(\n",
    "                len(conditions), 1,  figsize=(1.2, 0.8 * len(conditions)),\n",
    "                sharex=True)\n",
    "        for i, condition in enumerate(conditions):\n",
    "            ax = axes[i]\n",
    "            values = pd.unique(info[condition])\n",
    "            t_plot = np.arange(activity.shape[1]) * config['dt']\n",
    "            for value in values:\n",
    "                a = example_activity[info[condition] == value]\n",
    "                ax.plot(t_plot, a.mean(axis=0), label=str(value))\n",
    "            ax.legend(title=condition, loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "            ax.set_ylabel('Activity')\n",
    "            if i == len(conditions) - 1:\n",
    "                ax.set_xlabel('Time (ms)')\n",
    "            if i == 0:\n",
    "                ax.set_title('Unit {:d}'.format(example_id + 1))\n",
    "\n",
    "analysis_example_units_by_condition(activity, info, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analysis_pca_by_condition(activity, info, config):\n",
    "    # Reshape activity to (N_trial x N_time, N_neuron)\n",
    "    activity_reshape = np.reshape(activity, (-1, activity.shape[-1]))\n",
    "    pca = PCA(n_components=2)\n",
    "    pca.fit(activity_reshape)\n",
    "    \n",
    "    conditions = get_conditions(info)\n",
    "    for condition in conditions:\n",
    "        values = pd.unique(info[condition])\n",
    "        fig = plt.figure(figsize=(2.5, 2.5))\n",
    "        ax = fig.add_axes([0.2, 0.2, 0.7, 0.7])\n",
    "        for value in values:\n",
    "            # Get relevant trials, and average across them\n",
    "            a = activity[info[condition] == value].mean(axis=0)\n",
    "            a = pca.transform(a)  # (N_time, N_PC)\n",
    "            plt.plot(a[:, 0], a[:, 1], label=str(value))\n",
    "        plt.legend(title=condition, loc='center left', bbox_to_anchor=(1.0, 0.5))\n",
    "    \n",
    "        plt.xlabel('PC 1')\n",
    "        plt.ylabel('PC 2')\n",
    "\n",
    "analysis_pca_by_condition(activity, info, config)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}