notebooks/train_autoencoder_monkey.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# append parent directory to path (../notebooks -> ..) this cell should be run once. otherwise it goes one more level up\n",
    "sys.path.append(os.path.dirname(os.getcwd()))\n",
    "os.chdir(os.path.dirname(os.getcwd()))\n",
    "\n",
    "import accelerate\n",
    "import lovely_tensors as lt\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import scipy.io as io\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import wandb\n",
    "import yaml\n",
    "from diffusers.optimization import get_scheduler\n",
    "from omegaconf import OmegaConf\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from ldns.data.monkey import get_monkey_dataloaders\n",
    "from ldns.networks import AutoEncoder, CountWrapper\n",
    "from ldns.utils.plotting_utils import *\n",
    "from ldns.losses import latent_regularizer\n",
    "\n",
    "# always run from ../ntldm\n",
    "\n",
    "\n",
    "lt.monkey_patch()\n",
    "matplotlib.rc_file('matplotlibrc') # mackelab plotting style\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/conda/envs/ldns/lib/python3.9/site-packages/torch/cuda/__init__.py:129: UserWarning: CUDA initialization: CUDA driver initialization failed, you might not have a CUDA gpu. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)\n",
      "  return torch._C._cuda_getDeviceCount() > 0\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model:\n",
      "  C_in: 182\n",
      "  C: 256\n",
      "  C_latent: 16\n",
      "  num_blocks: 4\n",
      "  num_blocks_decoder: 0\n",
      "  num_lin_per_mlp: 2\n",
      "  bidirectional: false\n",
      "dataset:\n",
      "  system_name: monkey\n",
      "  task: mc_maze\n",
      "  datapath: data/000128/sub-Jenkins/\n",
      "  signal_length: 140\n",
      "training:\n",
      "  lr: 0.001\n",
      "  num_epochs: 260\n",
      "  num_warmup_epochs: 10\n",
      "  batch_size: 64\n",
      "  random_seed: 42\n",
      "  precision: bf16\n",
      "  latent_beta: 0.001\n",
      "  latent_td_beta: 0.2\n",
      "  tk_k: 5\n",
      "  mask_prob: 0.5\n",
      "exp_name: autoencoder-monkey_z=16\n",
      "\n"
     ]
    }
   ],
   "source": [
    "cfg_yaml = \"\"\"\n",
    "model:\n",
    "  C_in: 182\n",
    "  C: 256\n",
    "  C_latent: 16\n",
    "  num_blocks: 4\n",
    "  num_blocks_decoder: 0\n",
    "  num_lin_per_mlp: 2\n",
    "  bidirectional: False\n",
    "dataset:\n",
    "  system_name: monkey\n",
    "  task: mc_maze\n",
    "  datapath: data/000128/sub-Jenkins/\n",
    "  signal_length: 140\n",
    "training:\n",
    "  lr: 0.001\n",
    "  num_epochs: 260\n",
    "  num_warmup_epochs: 10\n",
    "  batch_size: 64\n",
    "  random_seed: 42\n",
    "  precision: bf16\n",
    "  latent_beta: 0.001\n",
    "  latent_td_beta: 0.2\n",
    "  tk_k: 5\n",
    "  mask_prob: 0.5\n",
    "exp_name: autoencoder-monkey_z=16\n",
    "\"\"\"\n",
    "\n",
    "# omegaconf from yaml\n",
    "cfg = OmegaConf.create(yaml.safe_load(cfg_yaml))\n",
    "print(OmegaConf.to_yaml(cfg))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task: mc_maze, Bin width: 5 ms\n",
      "Train: 2581, Val: 71, Test: 216\n"
     ]
    }
   ],
   "source": [
    "# set seed\n",
    "torch.manual_seed(cfg.training.random_seed)\n",
    "np.random.seed(cfg.training.random_seed)\n",
    "\n",
    "train_dataloader, val_dataloader, test_dataloader = get_monkey_dataloaders(\n",
    "    cfg.dataset.task, cfg.dataset.datapath, bin_width=5, batch_size=cfg.training.batch_size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualise dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fb9ef81c9a0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAdQAAAFmCAYAAADUJWsSAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAABcSAAAXEgFnn9JSAABNzUlEQVR4nO3df1xUVf4/8JcgMiAwgEyWssYvfyCWuq0WCYpkkEopG2JmKvrZDMtdTZH4KglkGYjRWpu6Zeuv/JgmaK2ZsrVS+XN1cVKTtNEwJdNMGdDlR8j9/tGHu94ZGO4wv2dez8djHg/nzLn3njsXz3vOueee00kQBAFERERkEjdbF4CIiMgZMKASERGZAQMqERGRGTCgEhERmQEDKhERkRkwoBIREZkBAyoREZEZMKASERGZAQMqERGRGTCgEhERmQEDKhERkRkwoBIREZkBAyoREZEZMKASERGZAQMqERHZnebmZqxevRr33nsvfHx8EBYWhueffx41NTUGt9u8eTOioqLg5eWFyMhIrF+/Xi/P0aNHERcXBx8fH/To0QMLFy5EY2OjyWW2u4BaWlqKIUOGwNvbG6GhoVi+fDm4ZCsRkWtZtmwZZs+ejbFjx2LHjh3IyMjAhg0b8Pjjj7cZE4qLizF58mQkJCRgx44diIuLQ1paGt5//30xz7lz5zBq1Ch4eXlh69atmD9/PoqKivCnP/3J9EILduTgwYOCh4eH8NRTTwmffPKJsGjRIqFTp07Cq6++auuiERGRldy6dUvw9/cXnn32WUn6+++/LwAQjhw50up2ffr0EVJTUyVpqampQkREhPh+5syZQnBwsNDQ0CCmrVy5UnBzcxPOnz9vUrk7mx6SzScnJweDBw/Gxo0bAQCPPPIIfvnlFyxduhRz5syBl5eXUfsTBAFarVZ8r1Qq0alTJ7OWmYjI2ejWnR1RX1+PhoYG+Pn56dW7CoUCCoWizW1ramowZcoUTJw4UZLer18/AMDZs2fxu9/9TvJZZWUlzpw5g7y8PEl6SkoKtm7dim+//Ra9e/fGnj17MHbsWHTp0kWS59lnn8WePXvw9NNPd+h8AcBuAmpDQwPKyspa/TKWLVuGffv24eGHH25z+/r6etTX10vStFotQkJCxPfXr1+Hv7+/OYtNROR0tFotAgICLLb/nJwc5Obmtvm5v78/3njjDb30HTt2AACioqL0PquoqAAA9OnTR5IeEREBADh9+jSCg4Nx/vx5vTwqlQp+fn44ffq0Maehx24C6rlz59DY2GjwyzAUUPPz8/WCcUesXLlS/Pezzz5r8v6sraX8liy7NY5BRHS7w4cPIz8/H48++igGDBig93lLi9rPz0+S7uvrC+DXVm9beVrytTfgqT12E1DlfBnW4OhBwhrld/TviIgcy/79+5GUlITQ0FCsXbu21TzNzc0G9+Hm5iYrjynsJqBa+kSJiKhjjB17Itw2CreyshJKpVLyuaH7p7q2bNmCtLQ09OnTB7t370a3bt1azddyjNraWkl6S2NMqVSKDTbdPC35dMtpLLuJUnK+DEOysrJw/fp1yauystIiZSUiciWdOnWCm5ubrJdu8FUqlfD395e85AbU5cuXY9KkSYiOjsYXX3yBu+66q828ffv2BQBoNBpJesv7yMhI+Pj4oGfPnnp5rly5gtraWkRGRsoqV1vsJqCGh4fD3d3d4JdhiEKh0Ltocn5trFy5UnLflIiIbO+vf/0rFixYgNTUVOzevbvd+jwiIgKhoaHYtm2bJL24uBi9e/cWB6gmJCRg586daGhokORxd3dHfHy8SWXuJAj2M2tCfHw86urqcODAAfFXzgsvvIC//vWv+OGHH+Dt7W3U/qqrqyUj1cw1ytdWA5d0BwNxcBARWYJu3dm5c2fZ3b6CIKCpqUl835F698cff0RYWBi6d++OjRs3onNn6d3J8PBweHp64tSpUwgPD4dKpQIArFu3DtOnT8esWbPw2GOP4cMPP8Tq1avx/vvvi4/gfPPNNxg8eDCio6Mxb948nDlzBgsXLsSMGTNMblzZzT1UAMjOzsaoUaOQmpqKGTNm4MCBAygsLER+fr7RwZSIiMyjU6dOVn2Gf9euXairq0NlZSViY2P1Pl+7di1CQkIwcuRIrF27FmlpaQCAtLQ0NDQ0YPny5fjb3/6GsLAwbNiwQfI8a79+/VBaWooFCxYgJSUFQUFBeP755/HSSy+ZXG67aqECwPbt25GTk4PTp0+jZ8+eeO655zB//vwO7ctSLVQiIkuxh54n3bqzS5cuRrVQb58X15XqXbtqoQJAcnIykpOTbV0MIiL6P9ZuoToquwuoRESuzB7HRDCgysOASk7DHrrKiJwRA6o8dvPYDBERkSNjC9VE5moVtbUfcz2iYy9zFHe0HHK+Zzn768j1am0ovaWukyXIOWdL/x1b8pjG0r2e9na97FFrEza0xc7GuVoVAyoRERnELl957O6xGXPiYzNERMbTrTt9fX2NaqHePoWsK9W7bKESETkIW3WTs4UqDwMqERG1i/dQ28eASkQm4yNL1mGr79eYFqort2QZUInIZAykzq1laTYyjN8QERGRGbCFKkNby6bdnkbk7Nit67rY5SsPAyoRERnEgCoPn0N1Aq7UYnakVpIjldVZudI1MOe56tadQUFBsu+hNjc34+rVq+J7Z613W8OASkREErp1p0qlMiqg/vTTT+J7V6p32eVLREQGsctXHgZUF+FK3cJkOv690O342Iw87PJ1YK50f4iIrEe37rzrrruM6vK9dOmS+N7Z6l1D+JODiIjIDNjl68DYMiUia+A9VHkYUImIyCAGVHkYUImIyCAGVHkYUImIyCAGVHk4KImIiAzq1KmT+OhMey9LBNSLFy/C398fZWVlbeZZt26dGPhbe61fv17MGxwc3Gqe22d46gi2UImIyCBbtlAvXLiAxMREaLVag/nGjh2LgwcP6qX/4Q9/QE1NDcaMGQMAuHr1KqqqqlBYWIiYmBhJXlMf72FAJSIig2wRUJubm7FhwwZkZGRAznQJKpUKKpVKkvbGG2+goqICBw4cED9Tq9UAgOTkZISHh5ulrC3Y5UtERBaj1WpRXV0tedXX17e73fHjx5Geno6pU6di48aNRh/38uXLyM7OxqxZs3D//feL6Wq1Gr6+vggLCzN6n+2xWkBtbm7G6tWrce+998LHxwdhYWF4/vnnUVNTI+aJiYlptV/76NGj1iqmw1u5cqVk2jh7OaYtymUvWs7dEudv79+rsedu6e/JVt+VvV+n9hi6N9na63YhISEICAiQvPLz89s9Zq9evaDRaFBUVARvb2+jy5yTkwM3Nze8/PLLknS1Wo3AwECkpKRAqVTCx8cHEydOlMzu1FFW6/JdtmwZsrOzsWDBAjz00EM4c+YMXnzxRZw8eRKlpaUAfv1FMm/ePEyYMEGybWRkpLWKSUREOmzR5RsYGIjAwMAObXvlyhWsX78e8+fP17svqlarUVVVhZkzZ2Lu3LmoqKjA4sWLMWLECBw7dgxdu3btcJmtMpdvc3MzunXrhieffBJvvfWWmL5lyxY88cQTOHLkCPz9/dG7d298+umneOihh8xyXGefy5eIyBJ0686IiAi4u7vL2vbWrVvQaDQG8+Tk5CA3N1d2ecrKyjBy5Ejs3bsXcXFx7eZfunQpcnNzUVVVpXdf9eDBg+jcuTOGDBkipu3fvx8xMTFYuXIlZs2aJbtcuqzSQq2pqcGUKVMwceJESXq/fv0AAGfPnhUv1qBBgzp0jPr6er1++fZGhRERUftMaaFWVlZCqVRK0hQKhdnK1ppt27YhISFBL5gCQHR0tF7asGHDoFQq8dVXX5l0XKvcQ/X398cbb7yBYcOGSdJ37NgBAIiKioJarYaPjw8yMjIQFBQEhUKBMWPG4PTp07KOkZ+fr9dPHxISYuYzISJyPabcQ1UqlfD395e8LBlQq6qqcOzYMaSmpup9ptVq8be//Q0nT56UpDc3N6OxsbHVAGwMmz02c/jwYeTn5+PRRx/FgAEDoFarcePGDQQEBGD79u04f/488vLyEBsbC7VajR49etiqqETk4Ky5vmtrg48cfSELY9ZDtfWKoIcPHwYAvQYcAHh6emL27NlITk7Gpk2bxPSPPvoIdXV1GDlypEnHtklA3b9/P5KSkhAaGoq1a9cCAF555RVkZmZi+PDhAIDY2Fg8+OCDiIyMxIoVK1BQUGCLohIRuTx7nHqwpqYGp06dQnh4uKRleeLECXh6erb6jKlCoUBWVhZycnLQvXt3jBkzBidOnEBubi7GjRuH+Ph4k8pk9YC6ZcsWpKWloU+fPti9eze6desGABg4cKBe3rCwMERGRsrq187KysLcuXMlaVqtlt2+RAbY2yL1liqPNc/PXr5LZ1deXo6RI0di7dq1SEtLE9MvX75scPBpdnY2VCoV3nrrLaxatQrdunVDenq6UYOk2mKVUb4tli9fjszMTMTFxWH79u3ijeqmpiZs2rQJffr00bth3LdvXwwdOrRDD/ZylC85IkcMcm3lMdTVam/nSf+lW3dGRUUZNcr366+/Ft+7Ur1rtYkd/vrXv2LBggVITU3F7t27JaO+OnfujLy8PCxYsECyTXl5OTQajcn92kRE1HGmDEpyJVZpof74448ICwtD9+7dsXHjRnTuLO1pDg8PxyeffIJp06ZhypQpmDJlCs6fP4/Fixfjrrvuwr/+9S/Zv45uxxYqEZHxdOvOAQMG6NXbbWlqapKMonWletcq91B37dqFuro6VFZWIjY2Vu/zlj5whUKBZcuWYfz48ejatSuSk5Px6quvdiiYujp2p1mXNUeRkuuy1f9rY0b5ys3njKwSUGfMmIEZM2a0my81NbXVZ4eIiMh27HGUrz2y6qAkazN3l29Hfx2a+1el7v4s2TpyhJauI5SR5OG1tA+6defgwYON6vI9duyY+J5dvkRERP+HLVR52EIltgqISEK37vztb39rVAu1vLxcfO9K9S5bqA7EGR56tyT+MCCyDA5KkocBlYiIDGKXrzwMqA7E3lpe9tYitJdyOBJ7u4ZknxhQ5WFAJSIig9jlKw8DKnUYWzWOj9eQ5GALVR4GVCIiMogtVHkYUImIqF2u3PKUiwGVyMVw3mEiy2BAJXIxDKJkLN5DlYcBlYiIDOI9VHkYUImIyCC2UOVhQCUiIoPYQpWHAZWIiAxiC1UeBlQiIjKIAVUe122bExERmRFbqGR3+JykeTnLBPjOch6OiPdQ5eEC4y6soxWUK1dsznDulvzB4gjfjyOU0dZ0686EhAR4eHjI2vaXX35BaWmp+N4c9e7FixcxYMAA7NixA3FxcW3m02g06N27t156VFQUTp48Kb4vLS3FokWL8PXXX6N79+547rnnMH/+fJO7q9lCJSIig2zZQr1w4QISExOh1WrbzatWqwEAn332Gby9vcX02/996NAhJCUlYeLEiViyZAn27duHzMxMNDU1ISsry6SysoXqIpylG5WtC3JEjvZ3q1t3jh492qgW6ieffCK+72i929zcjA0bNiAjIwOCIODatWvYu3evwRZqdnY21q9fjwsXLrSZJzExEdXV1Th8+LCY9sILL2DVqlW4fPkyvLy8jC5rC9ft7CYiIlk6deoktlLbe+l2m2q1WlRXV0te9fX17R7z+PHjSE9Px9SpU7Fx40ZZ5VSr1Rg0aFCbnzc0NKCsrAzJycmS9JSUFNTW1mLfvn2yjtMWdvm6CEf5ZdweZzkPsozbe2IA+/l7sZdydJQpj82EhITo5cnJyUFubq7B/fTq1QsajQbBwcEoKyuTdWy1Wo2IiAg8+OCDKC8vh7+/P9LS0rBkyRJ4eHjg3LlzaGxsRJ8+fSTbRUREAABOnz6Nhx9+WNaxWsOASkROw9EDl72yxXOogYGBCAwMlJ3/6tWrqKqqQlNTE5YtW4a7774bn332GQoKCnDhwgVs2rRJvA/r5+cn2dbX1xcAUFNTY1KZGVCJiMjhde3aFaWlpejdu7fYKh4xYgQ8PT2RnZ2N7OxsNDc3G9yHqQOqGFCJiMggU1qolZWVUCqVkjSFQmG2srXw8vJqtbt27NixyM7OxldffYV77rkHAFBbWyvJ09Iy1S2nsaw2KKm+vh4eHh7ihWl5+fj4iHmOHj2KuLg4+Pj4oEePHli4cCEaGxutVUQiImqFbr3d3ut2SqUS/v7+kpclAuq3336Lv/71r6iurpak19XVAQBUKhXCw8Ph7u4OjUYjydPyPjIy0qQyWK2FevLkSTQ1NeG9995DeHi4mO7u7g4AOHfuHEaNGoXo6Ghs3boVFRUVWLRoEa5du4bVq1dbq5itsteBDh3laEP4SZ8tryH/flyPI8zle+nSJaSnp8PNzQ1PP/20mL5lyxb4+fnhvvvug0KhwPDhw1FSUoKMjAyxrMXFxVAqlRg6dKhJZbBaQFWr1ejcuTNSUlLg6emp93lBQQF8fX3x4YcfokuXLhgzZgy8vb0xe/ZsLFy4EL169bJWUfU4W8XhbOfjimx5Dfn343rsMaDW1NTg1KlTCA8Ph0qlQkxMDB566CHMnz8fdXV16N+/Pz7++GO88cYbKCoqEp+Fzc7OxqhRo5CamooZM2bgwIEDKCwsRH5+vmQCiI6wWpevWq1Gv379Wg2mALBnzx6MHTsWXbp0EdNSUlLQ3NyMPXv2tLv/+vp6vWed5MysQUREhpnS5Wsp5eXliI6Oxscffwzg1wFFJSUlePrpp/H6668jKSkJpaWlePvttzF37lxxu/j4eBQXF+P06dMYP348Nm3ahMLCQmRmZppcJqu3UBMSErB//354enpiwoQJWL58OTp37ozz58/rPRukUqng5+eH06dPt7v//Px85OXlWar4TsWaXXbsHnQtjni9HbHM1mbrFmpcXBx0J/VrLc3Pzw+vvfYaXnvtNYP7S05O1pvcwRysElAFQcDx48chCAL+8Ic/IDs7G0eOHEFeXh5OnTqFLVu2ANB/Ngj49fkgU58NIilrVhyGjtVWRdbaNIn8EeAYOnK9ba218pirrB3Zjz1+T1xtRh6rBdSPPvoIKpUKUVFRAIDhw4fjzjvvxFNPPYW9e/ca3N6VLxARETkGm06O3zIB85IlS/Diiy+iqKgIzz//vCSPn58fZs2ahYKCAoP7qq+v15sfUqvVSqa9cuXJ8elXjtDStccWCtkXS/+N6E6OP2HCBMn4FkMaGxvxwQcfiO9dqd61Sgv1hx9+wMcff4zExETJaN2W54Puuusu9OzZU+/ZoCtXrqC2tlbWs0EKhcIizzYREbk6W99DdRRWCahNTU2YOXMmFi5ciFdeeUVM37JlC9zd3REbG4uEhATs3LkTRUVF4kjg4uJiuLu7Iz4+3hrFJAuw1C9pY5ejs9Q9MUP77ej9xLbuKRtznnLzW5o5yyPneXBjnhm3h2d5O3p8a5eZAVUeq3X5zpgxA++99x4WL16M6Oho7Nu3D0uXLsWsWbPw5z//Gd988w0GDx6M6OhozJs3D2fOnMHChQsxY8YMvf8kcnE9VPtjakViqBJkV6mUo3wf9vYjgPTrzieeeMKoLt/3339ffO9K9a7VHptZtWoVwsLCsHHjRrz88ssIDg7GSy+9hAULFgAA+vXrh9LSUixYsAApKSkICgrC888/j5deeslaRSQiolawhSqPTQclWRpbqERExtOtOydNmtTmpDy6GhoasHnzZvG9K9W7XG2GyIk4S/epuburrfG9WKOL3Vbd+GyhysMHPImIiMyAXb5ktI6MOrXmL2pTj2nLMlv7uGQdjjJArIVu3Tl58mSjunw3bdokvnelepddvkREZBC7fOVhC5WIiCR0684pU6YY1ULduHGj+N6V6l22UImIyCC2UOVhQCUiIoMYUOVhQCUiIoMYUOVhQCUiIoMYUOXhc6hEZNDKlSs7PJ82kSthC5WIOrxqDrkGtlDlYUAlIiKDGFDlYUAlIrZCySAGVHkYUImIyCAGVHkYUImIyCAGVHkYUImIyCAGVHn42AwRERnUElDlvszt4sWL8Pf3R1lZmcF8DQ0NWLp0Kfr164euXbuib9++eOmll9DY2CjJFxwc3Gq5r169alI52UIlIovhsnRkqgsXLiAxMRFarbbdvHPmzMHGjRvx4osvYsiQITh69Cjy8vJw/vx5vPvuuwCAq1evoqqqCoWFhYiJiZFsb+ok/gyoRGQxDKLOw9pduc3NzdiwYQMyMjIgZ1G0n3/+GW+//TYKCgqwYMECAMBDDz0EAMjKykJ+fj5UKhXUajUAIDk5GeHh4WYtM7t8iYjIIFO6fLVaLaqrqyWv+vr6do95/PhxpKenY+rUqZLl4NpSU1OD9PR0PPbYY5L0fv36AQDOnTsHAFCr1fD19UVYWJjc05eNAZWIiAwyJaCGhIQgICBA8srPz2/3mL169YJGo0FRURG8vb3bzR8aGoqVK1eib9++kvQdO3bAw8MDffr0AfBrQA0MDERKSgqUSiV8fHwwceJEXLp0yYhvpHUMqEREZJAtBiUFBgYiODjYpH1s374d69evR3p6urhgulqtRlVVFe677z7s3LkTRUVF+PzzzzFixAjcvHnTpOPxHioRERnk5uYGNzd57S+5+SytpKQETz75JGJiYrBs2TIx/Z133kHnzp0xZMgQAEBsbCyioqIQExODDRs2YNasWR0+JgOqnTE0SbkjcZbzINdmzN9xR/LKzW/KsczBlOdQKysroVQqJWkKhcJsZWvN66+/joyMDMTFxWHHjh2S40VHR+vlHzZsGJRKJb766iuTjsuASkREBpkSUJVKpcmPo8glCALmzJmDN998E5MmTcK6devQpUsX8XOtVovi4mIMHToUAwYMENObm5vR2NgIlUpl0vEZUO2MNVt0cn7ldvSXMFum5AyM+Ts29W/emP9r/P/VuoULF+LNN9/EvHnzsHz5cr3g7unpidmzZyM5ORmbNm0S0z/66CPU1dVh5MiRJh3f5QOqvXVNGipPW5/J6T5qLY+c/RhTDkvSPaYlJwywt78Jsm8d+T/raOvO2uPUgzU1NTh16hTCw8PF50sLCgowZMgQTJgwAYcPH5bk79+/P/z8/JCVlYWcnBx0794dY8aMwYkTJ5Cbm4tx48YhPj7epDK5fEAlIiLD7DGglpeXY+TIkVi7di3S0tJQUlICQRBw5MiRVu+T7t27F3FxccjOzoZKpcJbb72FVatWoVu3bkhPT0dubq7JZeokyJmCwkRlZWUGm9K5ubnIyclBTEwM9u/fr/f5kSNH8Lvf/c7o41ZXV4tDpQHg+vXrVuvLJ8uy9ZR2jtiKdcQyG+Js56PLln/junXnvHnzZA8kqq+vR1FRkfjelepdq7RQf/vb3+LgwYN66dnZ2Thy5AgmTZoEQRBw/PhxzJs3DxMmTJDki4yMtEYxiYioFfbYQrVHVmmhtuajjz7CuHHj8MEHHyAlJQUajQa9e/fGp59+Ks6/aCq2UO2HNVsTzt5yIePZukfD0ejWnRkZGUa1UJcvXy6+d6V61yb3UOvq6vDHP/4RY8eORUpKCgCIExYPGjSoQ/usr6/Xmx9SzuoE5HxYYZIu/k2Yhi1UeWwypcWKFStQVVWFP//5z2KaWq2Gj48PMjIyEBQUBIVCgTFjxuD06dOy9pmfn683X2RISIhlToCIiEiH1VuojY2NWLFiBZ544glERESI6Wq1Gjdu3EBAQAC2b9+O8+fPIy8vD7GxsVCr1ejRo4e1i0pmxBYCkeNiC1UeqwfUbdu24ccffxTXq2vxyiuvIDMzE8OHDwfw6/yKDz74ICIjI7FixQoUFBRYu6jkRHhfleyJo93TZUCVxyYBNSoqCgMHDpSk674HgLCwMERGRsqaXzErKwtz586VpGm1Wnb7EhGZqFOnTrInvWdAtZJffvkFe/bswQsvvCBJb2pqwqZNm9CnTx+9B3Lr6upkza+oUCgsPuEyOS5HaAWQ63C0v0e2UOWx6qCkEydO4D//+Q+GDRsmSe/cuTPy8vL0uoHLy8uh0WhMnl+RiIg6zhbroToiq7ZQT5w4AeDXORV15ebmYtq0aZg6dSqmTJmC8+fPY/HixRg0aBCmTZtmzWKSC3G0e1lEtsAWqjxWDaiXL18GAMkDwy2mTp0KhUKBZcuWYfz48ejatSuSk5Px6quvwt3d3ZrFdEqGAoclBuw4yiAgey+fM+ro34aj/E05IwZUeawaUDMzM5GZmdnm56mpqUhNTbViiYiIiMzDZlMPWgOnHiQiMp5u3fniiy/Cy8tL1rZ1dXVYsmSJ+N6V6l0u30ZERAaxy1ceBlQiIjKIAVUeBlQz48AJInI2DKjyMKASEZFBDKjyuFRAXbNmDby9vQGY1oI09AiKNVum1nz8wJGe12QvARnDmL8XU/8fmPq3efvxTdmPsdzc3GRPPSg3nzPiKF+yO44UvFs4YpmJ2qJbd7788stGjfLNzs4W37tSveu6PyWIiIjMyKW6fMm+OfIMOmyVkjPjPVR5GFCJiMggBlR5GFBdmD207G7X0XJYaoAZmZel/t6sOVDIkn8v9vb/UZcrB0q5OCiJiIgkdOvOV199VXxCoj3/+c9/8P/+3/8T37tSvctBSUREZJCt10O9ePEi/P39UVZW1m7ezZs3IyoqCl5eXoiMjMT69ev18hw9ehRxcXHw8fFBjx49sHDhQjQ2NppcTnb52ilzd0Oxa9Nx2bor0NLH59+m/bPlPdQLFy4gMTERWq223bzFxcWYPHky5syZg0ceeQQ7duxAWloaPD098cQTTwAAzp07h1GjRiE6Ohpbt25FRUUFFi1ahGvXrmH16tUmlZVdvkREJKFbdy5btsyoLt/bl+nsaL3b3NyMDRs2ICMjA4Ig4Nq1a9i7dy/i4uLa3KZv374YNGgQtmzZIqZNnDgR5eXl+PbbbwEAzzzzDHbt2oWzZ8+iS5cuAIBVq1Zh9uzZ+O6779CrVy+jy9qCXb5ERGQxWq0W1dXVkld9fX272x0/fhzp6emYOnUqNm7c2G7+yspKnDlzBsnJyZL0lJQUaDQaMaDu2bMHY8eOFYNpS57m5mbs2bPHyLOTYkC1gZUrV+pNIUZEZK9MuYcaEhKCgIAAySs/P7/dY/bq1QsajQZFRUWyWscVFRUAgD59+kjSIyIiAACnT59GXV0dzp8/r5dHpVLBz88Pp0+fbvc4hvAeKhERGWSLuXwDAwMRGBgoO3/LPVY/Pz9Juq+vLwCgpqamzTwt+WpqajpaXAAuHFBtPdDDHnFwCJHt2WPd5AgTOzQ3Nxv83M3NTVYeU7hsQOU0dfrstVxErsQe/x+aElArKyuhVColaQqFwmxla9FyjNraWkl6S6tTqVSKLVPdPC35dMtpLJcNqEREJI8pAVWpVFrl6Yq+ffsCADQaDQYPHiymazQaAEBkZCR8fHzQs2dPMa3FlStXUFtbi8jISJPKwEFJRERkUMs9VLkvW4iIiEBoaCi2bdsmSS8uLkbv3r0REhICAEhISMDOnTvR0NAgyePu7o74+HiTysAWKhERGWSP91Brampw6tQphIeHQ6VSAQAWL16M6dOno1u3bnjsscfw4YcfYuvWrXj//ffF7TIzM7F582aMHj0a8+bNw5kzZ7Bw4ULMnDnTpGdQAbZQiYjIAZWXlyM6Ohoff/yxmJaWlobVq1fjH//4B8aPH4/PP/8cGzZswMSJE8U8/fr1Q2lpKf7zn/8gJSUFRUVFeP7557FixQqTy8SZknRYYxo/R17301iOWGYiV6dbd/7lL38xaqak2bNni+9daYY6dvkSEZFB9tjla4/YQiUiIgnduvOtt95C165dZW178+ZNPPfcc+J7V6p3zX4Pta1ldjQaDR599FH4+/sjKCgIs2bN0puV4saNG3juuedw5513wsfHB2PGjDF5Kiiyf3KmYuR0jYbx+yFLcoRRvvbArF2+bS2zU11djfj4eNx5551Yv349rly5gszMTHz33XfYvXu3mO/JJ5/EoUOHsGzZMvj5+SEvLw8jR47E119/Lfm1RERE1sMuX3nMElB1l9nRtWrVKvz8888oLy9HUFAQACA4OBhjxozB/v37MWzYMBw8eBB///vfsWvXLowePRoAEBsbi9DQUKxcuRKLFi0yR1HJDskZsMRBTYbx+yFLYkCVxyxt8/aW2dmzZw9iY2PFYAr8+nCtr68vdu3aJebp2rUrEhISxDwqlQojRowQ8xhSX1+vt0SQnAVpiYjIMFNWm3ElZgmo7S2zU1FRobdcjru7O0JDQ8V7pBUVFQgLC4O7u7skX0REhKz7qPn5+XpLBLXMjEFERGRpZunybW+ZHa1W2+5yOXLyEBGR9dli+TZHZJXnUA0tmdPy5cvJQ0RE1sd7qPJYJVIplcp2l8uRk8eQrKwsXL9+XfKqrKw0uexERK6O91DlsUoLtW/fvnrL5dy6dQvfffcdfv/734t59uzZg+bmZkmLVKPRyFpSR6FQWGSNPWfBxcOJqKPYQpXHKi3UhIQEfP755/jpp5/EtNLSUty4cUMc1ZuQkIDa2lrs2bNHzPPTTz/hiy++kIz8JSIi62ILVR6rtFBnzZqFN998Ew8//DBycnLw888/IzMzE6NHj8aDDz4IABg+fDji4uIwefJkLFu2DN26dUNubi78/f0xa9YsaxRTj7kmdtfdj6mT47cwZntDeS05gT0nx7cMe/9e7b18HeWqPT0clCSPVc5cpVJh7969CAoKwuTJk7Fo0SJMmDABW7ZskeQrKSnBuHHjsGDBAqSlpaFnz5747LPPOEsSERHZPU6ObwXO+mtdDkPn7srfS0cY2zoy5fs1dCxn6dEw9W/TFn+/1jqmbt25ceNGoybHnzJlivjelSbHZ0Alo3XkP7UrrQHrisx9nXjdbUu37nzvvfeMCqhPPfWU+N6V6l2uh0pERO1y5cFGcrGF6oD46910pn6HvAamY6+F/dKtO//3f/8XPj4+sra9ceMGnnzySfG9s9S7cjCgErkYVx2pSvLp1p2bN282KqBOmjRJfO9K9S67fImIyCBO7CAPA6qdM7U1Yc2uTUu2fOx1tLC5n1U2x77a2re5Bwzdzpp/m9a43m0do6PXyZSBfIbKYa2/eQZUeVz3CVwiIiIz4j1UGRxpEIS9Pj9nKkcsM7XOWZ5jtTRb3uvWrTu3bt1q1D3U1NRU8b0r3UNlQLVT9lzpmDIFojnLQabjNaDW6Nad27ZtMyqgpqSkiO9NqXdLS0uxaNEifP311+jevTuee+45zJ8/v9Vu5XXr1mH69Olt7mvdunWYNm0aACA4OBhVVVV6eX766ScEBQV1qKwA76ESEVE7bHEP9dChQ0hKSsLEiROxZMkS7Nu3D5mZmWhqakJWVpZe/rFjx+LgwYN66X/4wx9QU1ODMWPGAACuXr2KqqoqFBYWIiYmRpLX1AYXA6qdstTAFHPsW3eSf3Oz5EAsWw/QsbcWob2Ug+ybLQJqTk4OBg8ejI0bNwIAHnnkEfzyyy9YunQp5syZAy8vL0l+lUoFlUolSXvjjTdQUVGBAwcOiJ+p1WoAQHJyMsLDw81S1hYclERERAaZsnybVqtFdXW15FVfX2/weA0NDSgrK0NycrIkPSUlBbW1tdi3b1+7Zb58+TKys7Mxa9Ys3H///WK6Wq2Gr68vwsLCjPgG5OE9VCIiktCtOz/88EOj7qGOGzfOYJ6cnBzk5ua2+XlFRQX69++P4uJi/P73vxfTr1+/jsDAQLz55puYPXu2wWOkp6fj/fffR2VlpaTef+qpp7Bv3z7cd999+PTTT3Hr1i2MHTsWf/7zn3HXXXfJOse2sMvXisw1mIeoI+ytu5kch7W7fLVaLQDAz89Pku7r6wsAqKmpMbj9lStXsH79esyfP1+vEaVWq1FVVYWZM2di7ty5qKiowOLFizFixAgcO3ZM9iIArWFAJSIiu9Lc3Gzw8/YWMV+zZg1u3bqFOXPm6H32zjvvoHPnzhgyZAgAIDY2FlFRUYiJicGGDRswa9asDpebAdWKXLFl4KytIkc8L0cqKzmPyspKKJVKSZpCoTC4TUv+2tpaSXpLy1R3f7q2bduGhIQEvUFKABAdHa2XNmzYMCiVSnz11VcG99seBlSyKGetxJ31vIhaY0qXr1KpNHrsSnh4ONzd3aHRaCTpLe8jIyPb3LaqqgrHjh3D3Llz9T7TarUoLi7G0KFDMWDAADG9ubkZjY2NrQZgY3CULxERtasjI3w7SqFQYPjw4SgpKcHt42aLi4uhVCoxdOjQNrc9fPgwgF9bnbo8PT0xe/ZsvPrqq5L0jz76CHV1dRg5cqRJ5WYL1YZae97SEbsSici52eI51OzsbIwaNQqpqamYMWMGDhw4gMLCQuTn58Pb2xs1NTU4deoUwsPDJS3LEydOwNPTs9VnTBUKBbKyspCTk4Pu3btjzJgxOHHiBHJzczFu3DjEx8ebVGa2UImIyCBTnkPtqPj4eBQXF+P06dMYP348Nm3ahMLCQmRmZgIAysvLER0djY8//liy3eXLlw12MWdnZ2PlypUoLS3Fo48+itdeew3p6enYvHmzyWXmc6hERCShW3fu2rVLfGSlPbW1teI0f4Br1bvs8iWX4Kxd6c56Xo7MGW/lcD1UedjlS0REZAbs8jWRo//ylMvQhPX29h3YW3mo42y5Jqgr0607d+/ebVSX7yOPPCK+d6UuXwZUIiKSYEDtGN5DJSIig3gPVR4GVCIiMogBVR6zD0q6ePEi/P39UVZWJkn/5z//iREjRiAgIAB33nknHn/8cZw9e1aSJzs7u9VnmpYvX27uYhIRkUy2eA7VEZm1hXrhwgUkJiaKS++02L9/PxISEjBu3Dhs2rQJN2/exJIlSzBs2DCcPHkSQUFBAH5dVicuLk5vWqi7777bnMUkJ8XBSESWwRaqPGYJqM3NzdiwYQMyMjLQ2hingoIC9O/fHx988IG47M6wYcPwm9/8BuvWrUNGRgaAXwPq9OnT8cADD5ijWORiGEiJLIMBVR6zBNTjx48jPT0dzz77LEaNGoWxY8dKPr///vsxfvx4yRp2PXr0gFKpFLt9r169iqqqKgwaNKhDZaivr0d9fb0kTbelTEREZClmuYfaq1cvaDQaFBUVwdvbW+/zRYsWYcaMGZK0zz//HNevX0dUVBSAX1unALBz507cfffd8PDwwODBg/HJJ5/IKkN+fj4CAgIkr5CQEJPOi4iIeA9VLrME1MDAQAQHB8vOf/XqVTz99NPo0aMHpk2bBuC/AfXHH3/EmjVrsH37dtxxxx1ISkrCnj17zFFMIiLqAAZUeaz+2MylS5eQmJiIS5cu4dNPPxUfFk5NTUW/fv0wevRouLu7AwASExMxcOBALF68GImJidYuKhERkWxWDagnTpxAUlISamtrsXv3btx///3iZ7169UKvXr0k+T08PJCQkIDVq1e3u++srCy9Fdq1Wq3Ld/uaa+Qrp4BzDfY8UtreyianPPZW5o7ioCR5rBZQ9+7di/Hjx0OpVOLLL78U75222LVrF+rq6vD4449L0uvq6iSLx7ZFoVBAoVCYtcxERMSAKpfZ5/ItKyvDyJEjsXfvXsTFxQEAjh07hpiYGISGhqK0tBQ9evTQ227GjBnYtm0bKisrERgYCAC4efMm+vbti/j4eGzYsMHosnAuXyIi4+nWnV988YVRc/kOHz5cfO9K9a5VWqj/8z//g19++QV5eXn4/vvv8f3334ufqVQqhIeHY8GCBdi6dStGjx6NhQsX4tatWygoKMDNmzeRl5dnjWICcM61DImITOXKLU+5LB5Qz507h2PHjgEAUlJS9D6fNm0a1q1bh8jISHzxxRdYuHAhZsyYgcbGRgwfPhzvvvsuQkNDLV1MIiIik7js8m323uo0pnwcMOQ4TPm743XuOGf97ix1Xrp155dffgk/Pz9Z29bU1CA2NlZ8zy5fIiKi/8NBSfK4bAuViIhap1t37t+/36gW6rBhw8T3rlTvsoUqgzW6h+29C9pcbN315qzfs7OeF9kHtlDlMft6qERERK6IXb7kdNhaIzKNbt158OBBo7p8o6OjxfeuVO+yy5dsRjfwmSsQtrb97V3N5jiGrdlitLCtf6jYYqo/Sx3T1rc+jGWrLt/S0lIsWrQIX3/9Nbp3747nnnsO8+fPb/MYGo0GvXv31kuPiorCyZMnO7xfuRhQiYjI7hw6dAhJSUmYOHEilixZgn379iEzMxNNTU3IyspqdZuWVcs+++wzyVKit/+7I/uVi12+ZmLrX++62mr93Z5my/JY85hyj2tv19AYlmqZmXOfzsrQd++of1O6defhw4eN6vK9feGTjta7iYmJqK6uxuHDh8W0F154AatWrcLly5fh5eWlt012djbWr1+PCxcumHW/cnFQEhERGWTKeqharRbV1dWSV319vcHjNTQ0oKysDMnJyZL0lJQU1NbWYt++fa1up1arMWjQILPvVy62UMlsdO9TAsbdezLX7FCO2iowhiucozOzRA+AOf8mdOvOI0eOGNVCHTJkiME8OTk5yM3NbfPziooK9O/fH8XFxfj9738vpl+/fh2BgYF48803MXv2bL3tgoODERERgcbGRpSXl8Pf3x9paWlYsmQJPDw8OrxfuXgPlSw6GMiY/MZsbyivKwQZW3Sdu8L3agpjvidLfJeWvD7WHpSk1WoBQC+It6x4U1NTo7fN1atXUVVVhaamJixbtgx33303PvvsMxQUFODChQvYtGlTh/ZrDAZUIiKyK83NzQY/d3PTv1vZtWtXlJaWonfv3ggJCQEAjBgxAp6ensjOzkZ2dnaH9msMBlQjOOsvdWc7HzIvtoblccQyy2VKC7WyshJKpVKSplAoDO6jJX9tba0kvaUFqbs/APDy8sLDDz+slz527FhkZ2fjq6++wj333GP0fo3BgEpERBajVCqNHrsSHh4Od3d3aDQaSXrL+8jISL1tvv32W/zzn//ExIkTJcerq6sD8N+1t43drzEYUI0gZ4CNnLy25sitADLM0R93ccQyuwJr30NVKBQYPnw4SkpKkJGRIe6zuLgYSqUSQ4cO1dvm0qVLSE9Ph5ubG55++mkxfcuWLfDz88N9993Xof0agwHVTBypInCkspJxnPna8oega8nOzsaoUaOQmpqKGTNm4MCBAygsLER+fj68vb1RU1ODU6dOITw8HCqVCjExMXjooYcwf/581NXVoX///vj444/xxhtvoKioSGy1trdfU/A5VCIiMsiU51A7Kj4+HsXFxTh9+jTGjx+PTZs2obCwEJmZmQCA8vJyREdH4+OPPwbw64CikpISPP3003j99deRlJSE0tJSvP3225g7d67s/ZqCz6ESEZGEbt2pVquNeg719skVXKneZZcvuRR2GxIZj+uhysMuXyIiIjNgC5UswhbLaMnBlilRx7hyy1MutlCJiIjMgC1UksXY5xvN3RI09/54L5Uc8W+gpZVo7bGkvIcqD0f5ksmsUTE5YuVnrxx98geyPN268/jx47Kn5dNqtbj33nvF965U77KFSkREBrGFKg9bqGQRbFESOS7duvPkyZNGtVAHDBggvnelepctVLIIrlBiP/j9kKnYQpWHo3yJiIjMwOwB9eLFi/D390dZWZkkPSYmptU5H48ePSrmuXz5MiZPnoxu3bpBqVRi0qRJuHTpkrmLSE7m2WeftUjra+XKlXqrCFlKy7EscTxLfT/OzJrX3hHYYi5fR2TWLt8LFy4gMTERWq1Wki4IAo4fP4558+ZhwoQJks9a1p9ramrC6NGjUVNTg9WrV+OXX35BVlYWEhISUF5eDg8PD3MWlYiIyKzMElCbm5uxYcMGZGRktPp81NmzZ1FbW4sxY8bggQceaHUfH3zwAY4dO4avv/4a/fv3BwAMGjQIAwYMwNatWzF58mRzFNVh8L6XlC2+D0c/lj38DTnqIzqOVFZr4D1UecwyyletVuOBBx7As88+i1GjRmHs2LHYu3cv4uLiAADbtm3DhAkTcPXqVXTr1q3VfaSlpeHQoUP45ptvJOlRUVEYNGgQNm3aZLAM9fX1qK+vl6RptVqEhISI711ptBkRUUfpjvL95ptvjBrl269fP/G9K9W7ZrmH2qtXL2g0GhQVFbW6QKtarYaPjw8yMjIQFBQEhUKBMWPG4PTp02KeiooK9OnTR2/biIgISb625OfnIyAgQPK6PZgSERFZklkCamBgIIKDg9v8XK1W48aNGwgICMD27duxZs0afPvtt4iNjcUPP/wA4NdfNa2tt+fr64uamhpzFJOIiDqAg5LkscpjM6+88go+//xzFBUVITY2Fk899RT27NkDrVaLFStWAPj1PmybhXTj0z1ERLbCgCqPVSZ2GDhwoF5aWFgYIiMj8dVXXwEAlEolamtr9fLV1NTI6rvPysrC3LlzJWm691CJyP446sAle8Dvzr5YPKA2NTVh06ZN6NOnD6KjoyWf1dXVQaVSAQD69u2LY8eO6W2v0WgwdOjQdo+jUCigUCjMU2gishoGgo6z1nfHUb7yWLwvtXPnzsjLy8OCBQsk6eXl5dBoNBg5ciQAICEhARUVFTh16pSY59SpU6ioqEBCQoKli0lERGQSq9yczM3Nxf79+zF16lT84x//wJo1a5CUlIRBgwZh2rRpAICJEyeiT58+GD16NDZv3ozNmzdj9OjRuOeee5CammqNYhKRiTjDkHPiPVR5rHIPderUqVAoFFi2bBnGjx+Prl27Ijk5Ga+++irc3d0BAJ6envjHP/6BOXPmYObMmfDw8EBCQgJef/11dO7MOfyJHAG7b50Tu3zl4fJtREQkoVt3nj171qiJHcLDw8X3rlTvumzTz1zTstnD9G5EZN8cvZ5gC1UePuBJREQG2eoeamlpKYYMGQJvb2+EhoZi+fLlrc4X36KhoQFLly5Fv3790LVrV/Tt2xcvvfQSGhsbJfmCg4NbLffVq1dNKq/LtlDN9UvRUX9xEpH1sJ4w3qFDh5CUlISJEydiyZIl2LdvHzIzM9HU1ISsrKxWt5kzZw42btyIF198EUOGDMHRo0eRl5eH8+fP49133wUAXL16FVVVVSgsLERMTIxke1O7pnkPlYiIJHTrzsrKSqPuoZpjUZLExERUV1fj8OHDYtoLL7yAVatW4fLly/Dy8pLk//nnn6FSqVBQUCB5TLOgoABZWVm4cuUKVCoVPv30Uzz88MPQaDSSe73mwC5fIiKyGK1Wi+rqaslLd2UwXQ0NDSgrK0NycrIkPSUlBbW1tdi3b5/eNjU1NUhPT8djjz0mSW9Z+ebcuXMAfp1b3tfXF2FhYaacVqsYUMmlOOJzko5YZnI+Hb1/GhISorcSWH5+vsFjnTt3Do2NjXorkEVERABAqyuQhYaGYuXKlejbt68kfceOHfDw8BD3pVarERgYiJSUFCiVSvj4+GDixIm4dOmSUd9Ha1z2Hiq5Jke8l+WIZSYyhVarBQC9Fch8fX0BQPYKZNu3b8f69esxe/ZssQtbrVajqqoKM2fOxNy5c1FRUYHFixdjxIgROHbsGLp27drhcjOgEhGRQdZ+bMbQ6mOAvBXISkpK8OSTTyImJgbLli0T09955x107twZQ4YMAQDExsYiKioKMTEx2LBhA2bNmtXhcjOgkogrVxCRubU2oKm9hUxa8uuuQNbSMm1vgNTrr7+OjIwMxMXFYceOHZLj6S7SAgDDhg2DUqkUVz/rKAZUIiIyyJQWqlKpNHqUb3h4ONzd3aHRaCTpLe8jIyNb3U4QBMyZMwdvvvkmJk2ahHXr1qFLly7i51qtFsXFxRg6dCgGDBggpjc3N6OxsVFc/ayjGFBJxFYpEdkDhUKB4cOHo6SkBBkZGWKQLi4uhlKpbHNJz4ULF+LNN9/EvHnzsHz5cr3g7unpidmzZyM5ORmbNm0S0z/66CPU1dWJq591FAOqk3L0qc6IyH7YYurB7OxsjBo1CqmpqZgxYwYOHDiAwsJC5Ofnw9vbGzU1NTh16hTCw8OhUqmgVqtRUFCAIUOGYMKECZLnVwGgf//+8PPzQ1ZWFnJyctC9e3eMGTMGJ06cQG5uLsaNG4f4+HiTysyASkREdic+Ph7FxcXIycnB+PHj0bNnTxQWFmL+/PkAfl1Te+TIkVi7di3S0tJQUlICQRBw5MiRVu+T7t27F3FxccjOzoZKpcJbb72FVatWoVu3bkhPT0dubq7JZeZMSUREJKFbd164cEF23VldXY3f/OY34ntXqnfZQiUiIoO42ow8DKgOjPdJ7Yc9PnLEvw8i62KXrwymVkys2IjIkejWnVVVVUZ1+fbs2VN870pdvpzLl4iIyAzY5SuDqS1LtkyJyJHxHqo8bKESERGZAQMqERGRGbDLl4iIDGKXrzxsoRIREZkBAyoREZEZsMuXiIgMYpevPGyhEhERmQFbqC6MMzgRWZaz/B9jC1UeswfUixcvYsCAAdixYwfi4uIAGP6C4+LisHfvXgC/rn/3yiuv6OUpLCxERkaGuYvq8hz9P3lbOlqJmbvyM9f+bFEpW3JuYnsNMuY6Z3uc15msw6wB9cKFC0hMTIRWq5WkHzx4UC9vSUkJCgsLkZ6eLqap1WrExcXh1VdfleS9++67zVlMIiIiszPL5PjNzc3YsGEDMjIyIAgCrl27Ji7m2poLFy7g3nvvxeTJk/GXv/xFTA8ODsb06dOxZMkSU4sEgOuhEhF1hG7deeXKFaMmx7/jjjvE965U75plUNLx48eRnp6OqVOnYuPGje3mnz9/Pry8vLB06VIx7erVq6iqqsKgQYM6VIb6+npUV1dLXrotZSIiIksxS0Dt1asXNBoNioqK4O3tbTDvoUOH8MEHH2Dp0qXw8/MT09VqNQBg586duPvuu+Hh4YHBgwfjk08+kVWG/Px8BAQESF4hISEdPSWLWrlypeQ+i633Y00tZXa0chO5spZBSXJfrsos91ADAwMRGBgoK++yZcsQEhKCp556SpLeElB//PFHrFmzBg0NDXjzzTeRlJSEXbt2ITEx0eRyrlmzRgz4thgsYO7BGJY8B0sNHOEgDSJyVlZ9bObixYv48MMPUVRUhM6dpYdOTU1Fv379MHr0aLi7uwMAEhMTMXDgQCxevNgsAZWIiIzHx2bkMcugpNuVlZVh5MiRrQ5KeuONNzBv3jxUVVWhe/fusvY3d+5crF69GvX19Qbz1dfX6+XRarWSbl9Hujkup4Vor48fEJFj0x2UdPXqVaMGJQUFBYnvHaneNZVVW6g7d+7E8OHDWw2mu3btQl1dHR5//HFJel1dHVQqVbv7VigUUCgUZisrERGRMawWUAVBwL/+9S/88Y9/bPXzbdu2Ydu2bRg5cqR4P/bmzZv4+OOPER8fb61i2g05rU62TInIGtjlK4/V5vL9/vvvodVq0b9//1Y/X7BgAZqbmzF69Gh8+OGHKCkpQXx8PG7evIm8vDxrFZOIiOxEaWkphgwZAm9vb4SGhmL58uVo7y7l5s2bERUVBS8vL0RGRmL9+vV6eY4ePYq4uDj4+PigR48eWLhwIRobG00ur9UC6uXLlwFA0i9/u8jISHzxxRcICAjAjBkzMG3aNAQFBeHLL79EaGiotYpJREQ6bPHYzKFDh5CUlIR+/fqhpKQEkydPRmZmJgoKCtrcpri4GJMnT0ZCQoI4/W1aWhref/99Mc+5c+cwatQoeHl5YevWrZg/fz6Kiorwpz/9yeQym31Qkj3pyExJ1piH014GE1mjHPZyrkQkn27dee3aNaMGJd3+GGVHByUlJiaiuroahw8fFtNeeOEFrFq1CpcvX4aXl5feNn379sWgQYOwZcsWMW3ixIkoLy/Ht99+CwB45plnsGvXLpw9exZdunQBAKxatQqzZ8/Gd999h169ehld1hZOHVCvX78uubCVlZVQKpU2LBHZ2po1awAAf/jDH2xcEiLLUiqVHW4t6gZUY+pO3acrWtu2vUGkDQ0N8PPzQ15eHrKyssT0I0eOYOjQoSgtLcXDDz8s2aayshKhoaHYvHkznnjiCTH9gw8+QGpqKs6cOYPevXsjJCQEjzzyCFavXi3m+emnn3DHHXfg7bffxtNPPy3rPFvj1Mu31dTUSN7b68xJZH0LFiywdRGILMqcj6uYUne2tm1OTg5yc3Pb3ObcuXNobGxEnz59JOkREREAgNOnT+sF1IqKCgAwuE1wcDDOnz+vl0elUsHPzw+nT5+WdU5t4QLjRERkV1rmYb99eloA8PX1BaDfWJK7TVt5WvK1tl9jMKASEZFdaW5uNvi5m5t+6JKzTUf2awyn7vL9zW9+g6+++goDBw4U05zxPqqcexaOjOfn+Jz9HO3x/Ew5vlKpxPXr1006fn19vXgvVPdebnuT8LSUvba2VpLe0oJs7dzkbNPSMtXN05LP1Gvm1AHVzc1Nb8SWUql0+mmwnP0ceX6Oz9nP0dHPr1OnTjYtf3h4ONzd3aHRaCTpLe8jIyP1tunbt6+YZ/Dgwa1u4+Pjg549e+rt98qVK6itrW11v8Zgly8REdkVhUKB4cOHo6SkRDKRQ3FxMZRKJYYOHaq3TUREBEJDQ7Ft2zZJenFxsTi6FwASEhKwc+dONDQ0SPK4u7ubPCufU7dQiYjIMWVnZ2PUqFFITU3FjBkzcODAARQWFiI/Px/e3t6oqanBqVOnEB4eLs73vnjxYkyfPh3dunXDY489hg8//BBbt26VTOyQmZmJzZs3Y/To0Zg3bx7OnDmDhQsXYubMmSY9gwoAEJzc9evXBQDi6/r167Yuktk5+zny/Byfs5+js5+frZSUlAj33HOP0KVLFyE0NFRYvny5+NnevXsFAMLatWsl26xevVqIiIgQPD09hcjISGHDhg16+/3iiy+E+++/X/D09BR69uwpZGVlCY2NjSaX1+lbqAqFAjk5OZL3zsbZz5Hn5/ic/Ryd/fxsJTk5GcnJya1+FhcX1+q8vs888wyeeeYZg/uNjY3FoUOHzFLG2zn1TElERETWwkFJREREZsCASkREZAYMqERERGbAgEpERGQGDKhERERmwIBKRERkBgyoREREZuASAbW0tBRDhgyBt7c3QkNDsXz58lYfCLZ3zc3NWL16Ne699174+PggLCwMzz//vGQNv5iYGHTq1EnvdfToURuWXL76+np4eHjold/Hx0fMc/ToUcTFxcHHxwc9evTAwoUL0djYaMNSy1NWVtbqtWl55eXlAXDca3jx4kX4+/ujrKxMkq7RaPDoo4/C398fQUFBmDVrlt66kzdu3MBzzz2HO++8Ez4+PhgzZozJiz2bW1vn989//hMjRoxAQEAA7rzzTjz++OM4e/asJE92dnar13T58uVWPAOyNKefKenQoUNISkrCxIkTsWTJEuzbtw+ZmZloampCVlaWrYtnlGXLliE7OxsLFizAQw89hDNnzuDFF1/EyZMnUVpaCgA4fvw45s2bhwkTJki2NXUVBWs5efIkmpqa8N577yE8PFxMd3d3BwCcO3cOo0aNQnR0NLZu3YqKigosWrQI165dw+rVq21VbFl++9vf4uDBg3rp2dnZOHLkCCZNmgRBEBzyGl64cAGJiYniAs4tqqurER8fjzvvvBPr16/HlStXkJmZie+++w67d+8W8z355JM4dOgQli1bBj8/P+Tl5WHkyJH4+uuvERAQYO3T0dPW+e3fvx8JCQkYN24cNm3ahJs3b2LJkiUYNmwYTp48iaCgIACAWq1GXFwcXn31Vcn2d999t9XOgazA5MkL7VxCQoIwdOhQSVpmZqbg6+sr/Oc//7FRqYx369Ytwd/fX3j22Wcl6e+//74AQDhy5Ijw7bffCgCETz/91EalNN0777wjdO7cWaivr2/185kzZwrBwcFCQ0ODmLZy5UrBzc1NOH/+vLWKaTYffvihAED44IMPBEEQHO4a3rp1S1i7dq3QrVs3ITAwUAAg7N27V/x86dKlgre3t/DTTz+Jabt27RIACPv27RMEQRAOHDggABB27dol5rly5YrQtWtX4eWXX7baubSmvfN79NFHhXvuuUe4deuWmFZVVSW4ubkJhYWFYlrPnj2F7OxsaxadbMCpu3wbGhpQVlamNxdkSkoKamtrsW/fPhuVzHg1NTWYMmUKnnzySUl6v379AABnz56FWq0GAAwaNMjKpTMftVqNfv36wdPTs9XP9+zZg7Fjx6JLly5iWkpKCpqbm7Fnzx5rFdMs6urq8Mc//hFjx45FSkoKADjcNTx+/DjS09MxdepUbNy4Ue/zPXv2IDY2VmypAb8un+Xr64tdu3aJebp27YqEhAQxj0qlwogRI8Q8ttLe+d1///2YO3cu3Nz+W5X26NEDSqVS7Pa9evUqqqqqHOaaUsc5dUA9d+4cGhsb0adPH0l6REQEANjdPRpD/P398cYbb2DYsGGS9B07dgAAoqKioFar4ePjg4yMDAQFBUGhUNjlvShD1Go1OnfujISEBHTt2hWBgYF45plnUFtbi7q6Opw/f17veqpUKvj5+TnUeQLAihUrUFVVhT//+c9imqNdw169ekGj0aCoqAje3t56n1dUVOhdL3d3d4SGhornVFFRgbCwMLFbv0VERITNz7u981u0aBFmzJghSfv8889x/fp1REVFAfjvj6SdO3fi7rvvhoeHBwYPHoxPPvnE4uUn63LqgNpyv8PPz0+S7uvrCwB6AyMczeHDh5Gfn49HH30UAwYMgFqtxo0bNxAQEIDt27djzZo1+PbbbxEbG4sffvjB1sVtl/B/9w81Gg3GjRuHTz75BIsWLcLmzZsxZswYXL9+HYD+9QR+vaaOdD0bGxuxYsUKPPHEE+IPPAAOdw0DAwMRHBzc5udarbbd6yUnj620d366rl69iqeffho9evTAtGnTAPw3oP74449Ys2YNtm/fjjvuuANJSUkO16tChjn1oKTm5maDn9/eTeNo9u/fj6SkJISGhmLt2rUAgFdeeQWZmZkYPnw4gF+XKHrwwQcRGRmJFStWoKCgwJZFbpcgCPjoo4+gUqnEX/fDhw/HnXfeiaeeegp79+41uL0jXc9t27bhxx9/xIIFCyTpjn4NdRn6P9hyveTkcQSXLl1CYmIiLl26hE8//VT84Z6amop+/fph9OjRYis8MTERAwcOxOLFi5GYmGjLYpMZOXVAVSqVAIDa2lpJesuv3pbPHc2WLVuQlpaGPn36YPfu3ejWrRsAYODAgXp5w8LCEBkZia+++sraxTSam5sb4uLi9NLHjh0LAPjuu+8A6F9P4Ndr6kjXc9u2bYiKitK7Zo5+DXUplco2r1fPnj3FPJcvX241j6Nc0xMnTiApKQm1tbXYvXs37r//fvGzXr16oVevXpL8Hh4eSEhIsPuR6WQcx/n51wHh4eFwd3eHRqORpLe8t+fHENqyfPlyTJo0CdHR0fjiiy9w1113AQCampqwfv36Vh/LqKurg0qlsnZRjfbDDz/gnXfewffffy9Jr6urAwDcdddd6Nmzp971vHLlCmprax3mev7yyy/Ys2cPUlNTJenOcA119e3bV+963bp1C9999514vfr27YvvvvtOr6Wq0Wgc4pru3bsXMTExEAQBX375pd44h127dqG4uFhvO0e9pmSAjUcZW9zIkSOFBx54QGhubhbTMjMzBaVSKdy8edOGJTPe6tWrBQDCxIkTJY+NtAgNDRWGDRsmSfv3v/8tuLm5Ce+++661itlh58+fFwAICxculKS//vrrgru7u3D69Glh+vTpQq9evSSP1axcuVJwd3d3mMdm/v3vf7f5aIwjX8O9e/fqPVaSl5cndO3aVbhy5YqY1vLYzP79+wVBEITPP/+8zcdmXnnlFauVvz2tnV95ebng7e0tREVFCVVVVa1uN336dMHX11f4+eefxbQbN24IPXv2FKZMmWLpYpMVOX1A/eyzz4ROnToJKSkpwq5du4Ts7GyhU6dOQkFBga2LZpRLly4JXl5eQkhIiPDll18KBw8elLyuXLkirF+/XgAgTJkyRSgtLRXeeecd4a677hJ++9vfCk1NTbY+BVmmT58ueHh4CEuWLBE+/fRTITc3V+jSpYswZ84cQRAEoaKiQlAoFMLIkSOFv//978Jrr70meHp6CrNmzbJtwY2wbt06AYDwww8/6H3myNewtYBz5coVISgoSBg4cKBQUlIivPPOO0JAQIAwevRoybZxcXFCQECA8M477wglJSXCvffeK/Ts2VO4du2alc+iba2d3+DBgwUPDw9h27Ztev8nNRqNIAiCcOrUKaFr167C0KFDhR07dgjFxcXC0KFDBX9/f+HcuXM2OhuyBKcPqIIgCCUlJcI999wjdOnSRQgNDRWWL19u6yIZ7d133xUAtPlau3atIAiCsGXLFuG+++4TvL29BZVKJcycOVPyy9je1dfXC0uWLBH69OkjeHp6CuHh4UJ+fr7kwfkvvvhCuP/++wVPT0+hZ8+eQlZWltDY2GjDUhunoKBAACDU1dW1+rmjXsPWAo4gCMKJEyeEhx56SPDy8hLuuOMOYebMmUJNTY0kz7Vr14S0tDTB399f8PPzE0aPHi188803Vix9+3TP7+zZswb/T06bNk3c9t///reQmJgoBAYGCj4+PsKYMWOEEydO2OZEyGI6CYIDTmpLRERkZ5x6UBIREZG1MKASERGZAQMqERGRGTCgEhERmQEDKhERkRkwoBIREZkBAyoREZEZMKASERGZAQMqERGRGTCgEhERmQEDKhERkRkwoBIREZkBAyoREZEZMKASERGZAQMqERGRGTCgEhERmcH/B+Fnbspxm9npAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 508.5x381 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "plt.imshow(train_dataloader.dataset[0]['signal'], aspect='auto', cmap='Greys')\n",
    "plt.colorbar()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define model, Optimizer and Learning rate scheduler\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cpu\n",
      "Number of params 1.415366 M\n",
      "CountWrapper(\n",
      "  (ae_net): AutoEncoder(\n",
      "    (encoder_in): Conv1d(182, 256, kernel_size=(1,), stride=(1,))\n",
      "    (encoder): ModuleList(\n",
      "      (0-3): 4 x AutoEncoderBlock(\n",
      "        (time_mixer): FFTConv(\n",
      "          (activation): Identity()\n",
      "          (kernel): SSMKernelDPLR()\n",
      "          (drop): Identity()\n",
      "          (drop_kernel): Identity()\n",
      "        )\n",
      "        (post_tm_scale): Conv1d(256, 256, kernel_size=(1,), stride=(1,), padding=same, groups=256)\n",
      "        (channel_mixer): Sequential(\n",
      "          (0): Rearrange('b c l -> b l c')\n",
      "          (1): Linear(in_features=256, out_features=512, bias=False)\n",
      "          (2): GELU(approximate='none')\n",
      "          (3): Linear(in_features=512, out_features=256, bias=False)\n",
      "          (4): Rearrange('b l c -> b c l')\n",
      "        )\n",
      "        (norm1): InstanceNorm1d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (norm2): InstanceNorm1d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (act1): GELU(approximate='none')\n",
      "        (act2): Identity()\n",
      "      )\n",
      "    )\n",
      "    (bottleneck): Conv1d(256, 16, kernel_size=(1,), stride=(1,), groups=16)\n",
      "    (act_bottleneck): Identity()\n",
      "    (unbottleneck): Conv1d(16, 256, kernel_size=(1,), stride=(1,), groups=16)\n",
      "    (decoder): ModuleList(\n",
      "      (0): GELU(approximate='none')\n",
      "    )\n",
      "    (decoder_out): Conv1d(256, 182, kernel_size=(1,), stride=(1,))\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "\n",
    "\n",
    "ae = AutoEncoder(\n",
    "    C_in=cfg.model.C_in,\n",
    "    C=cfg.model.C,\n",
    "    C_latent=cfg.model.C_latent,\n",
    "    L=cfg.dataset.signal_length,\n",
    "    num_blocks=cfg.model.num_blocks,\n",
    "    num_blocks_decoder=cfg.model.get(\"num_blocks_decoder\", cfg.model.num_blocks),\n",
    "    num_lin_per_mlp=cfg.model.get(\"num_lin_per_mlp\", 2),  # default 2\n",
    "    bidirectional=cfg.model.get(\"bidirectional\", True),\n",
    ")\n",
    "\n",
    "print(\"Number of params\", sum(p.numel() for p in ae.parameters() if p.requires_grad)/1e6, \"M\")\n",
    "\n",
    "ae = CountWrapper(ae)\n",
    "print(ae)\n",
    "\n",
    "ae = ae.to(device)\n",
    "optimizer = torch.optim.AdamW(\n",
    "    ae.parameters(), lr=cfg.training.lr\n",
    ")  # default wd=0.01 for now\n",
    "\n",
    "num_batches = len(train_dataloader)\n",
    "lr_scheduler = get_scheduler(\n",
    "    name=\"cosine\",\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=num_batches * cfg.training.num_warmup_epochs,  # warmup for 10% of epochs\n",
    "    num_training_steps=num_batches * cfg.training.num_epochs * 1.1,  # total number of steps for scheduler to go till lr=0\n",
    ")\n",
    "\n",
    "# check if signal length is power of 2\n",
    "if cfg.dataset.signal_length & (cfg.dataset.signal_length - 1) != 0:\n",
    "    cfg.training.precision = \"no\"  # torch.fft doesnt support half if L!=2^x\n",
    "\n",
    "\n",
    "#\n",
    "\n",
    "# Accelerate setuo\n",
    "\n",
    "accelerator = accelerate.Accelerator(\n",
    "    mixed_precision=cfg.training.precision,\n",
    "    log_with=\"wandb\",\n",
    ")\n",
    "\n",
    "\n",
    "(\n",
    "    ae,\n",
    "    optimizer,\n",
    "    lr_scheduler,\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ") = accelerator.prepare(\n",
    "    ae,\n",
    "    optimizer,\n",
    "    lr_scheduler,\n",
    "    train_dataloader,\n",
    "    val_dataloader,\n",
    "    test_dataloader,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# set up losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion_poisson = nn.PoissonNLLLoss(log_input=False, full=True, reduction=\"none\")\n",
    "def compute_val_loss(net, dataloader, do_plot=False):\n",
    "    net.eval()\n",
    "    poisson_loss_total = 0\n",
    "    batch_count = 0\n",
    "\n",
    "    for batch in dataloader:\n",
    "        signal = batch[\"signal\"]\n",
    "        with torch.no_grad():\n",
    "            output_rates = net(signal)[0].cpu()\n",
    "\n",
    "        signal = signal.cpu()  # move signal to cpu\n",
    "\n",
    "        # compute pointwise l2 loss\n",
    "        poisson_loss = criterion_poisson(output_rates, signal)\n",
    "        poisson_loss_total += poisson_loss.mean().item()\n",
    "        batch_count += 1\n",
    "\n",
    "    # compute average losses over all batches\n",
    "    avg_poisson_loss = poisson_loss_total / batch_count * cfg.training.mask_prob\n",
    "    print(f\"Validation loss: {avg_poisson_loss:.4f}, mask_prob {cfg.training.mask_prob}\")\n",
    "\n",
    "    if do_plot:\n",
    "        fig, ax = plt.subplots(2, 1, figsize=(10, 2), dpi=300)\n",
    "        for row in range(2):  # plot channels 0 and 71, as an example\n",
    "            ax[row].plot(output_rates[0, 92 * (row)].cpu().clip(0, 3).numpy(), label=\"pred\")\n",
    "        ax[row].plot(\n",
    "            batch[\"signal\"][0, 92 * (row)].cpu().clip(0, 3).numpy(),\n",
    "            label=\"spikes\",\n",
    "            alpha=0.5,\n",
    "            color=\"grey\",\n",
    "        )\n",
    "        plt.legend()\n",
    "        \n",
    "        plt.show()\n",
    "        plt.close(fig)\n",
    "\n",
    "    return avg_poisson_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some plotting utils during training (and evaluation)\n",
    "\n",
    "from ldns.utils.plotting_utils import cm2inch\n",
    "from einops import rearrange\n",
    "    \n",
    "def plot_rate_traces_real(model, dataloader,figsize=(12, 5), idx=0):\n",
    "    \"\"\"\n",
    "    Plot rates and spikes.\n",
    "    \n",
    "    Args:\n",
    "        model: Trained autoencoder model\n",
    "        dataloader: Data loader\n",
    "        figsize: Figure size\n",
    "        idx: Index of the batch to plot\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    for batch in dataloader:\n",
    "        signal = batch[\"signal\"]\n",
    "        with torch.no_grad():\n",
    "            output_rates = model(signal)[0].cpu()\n",
    "\n",
    "        signal = signal.cpu()  # move signal to cpu\n",
    "        break\n",
    "\n",
    "    channels = [136, 8] # as an example\n",
    "    fig, ax = plt.subplots(len(channels), 1, figsize=cm2inch(figsize), dpi=150)\n",
    "\n",
    "    for i, channel in enumerate(channels):\n",
    "        print(batch[\"signal\"][idx, channel])\n",
    "        L = batch[\"signal\"][idx, channel].shape[0]\n",
    "        ax[i].vlines(\n",
    "            torch.arange(L),\n",
    "            torch.zeros(L),\n",
    "            torch.ones(L)*output_rates[idx, channel].cpu().max().item(),\n",
    "            # batch[\"signal\"][0, channel].cpu().numpy(),\n",
    "            # label=\"spikes\",\n",
    "            color=\"black\",\n",
    "            alpha=np.min(\n",
    "                np.stack(\n",
    "                    (np.ones(L), batch[\"signal\"][idx, channel].cpu().numpy() * 0.2),\n",
    "                    axis=1,\n",
    "                ),\n",
    "                axis=1,\n",
    "            ),\n",
    "        )\n",
    "        ax[i].plot(output_rates[idx, channel].cpu().numpy(), label=\"pred\", color=\"red\")\n",
    "        ax[i].set_title(f\"channel {channel}\")\n",
    "\n",
    "    ax[-1].legend()\n",
    "\n",
    "    fig.suptitle(\"rate traces for channels\")\n",
    "    fig.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def imshow_rates_real(model, dataloader,figsize=(12, 5), idx=0):\n",
    "    \"\"\"\n",
    "    Plot rates and spikes.\n",
    "    \n",
    "    Args:\n",
    "        model: Trained autoencoder model\n",
    "        dataloader: Data loader\n",
    "        figsize: Figure size\n",
    "        idx: Index of the batch to plot\n",
    "    \"\"\"\n",
    "\n",
    "    model.eval()\n",
    "    for batch in dataloader:\n",
    "        signal = batch[\"signal\"]\n",
    "        with torch.no_grad():\n",
    "            output_rates = model(signal)[0].cpu()\n",
    "\n",
    "        signal = signal.cpu()  # move signal to cpu\n",
    "        break\n",
    "\n",
    "    fig, ax = plt.subplots(2, 1, figsize=cm2inch(figsize), dpi=150)\n",
    "\n",
    "    im1 = ax[0].imshow(output_rates[idx].cpu().numpy(), label=\"rates\", aspect='auto', cmap='Greys')\n",
    "    im2 = ax[1].imshow(signal[idx].cpu().numpy(), label=\"rates\", aspect='auto', cmap='Greys')\n",
    "    plt.colorbar(im1, ax=ax[0])\n",
    "    plt.colorbar(im2, ax=ax[1])\n",
    "    plt.show()\n",
    "\n",
    "    # ax.set_title(f\"channel {channel}\")\n",
    "\n",
    "    ax[-1].legend()\n",
    "\n",
    "    fig.suptitle(f\"infeered rates, idx {idx}\")\n",
    "    fig.tight_layout()\n",
    "    plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training loop with validation\n",
    "# tracks reconstruction loss, latent regularization, and validation metrics\n",
    "\n",
    "rec_losses, latent_losses, total_losses, lrs, val_rate_losses = [], [], [], [], []\n",
    "avg_poisson_loss, avg_rate_loss = 0, 0\n",
    "\n",
    "# flags for enabling/disabling different evaluations, enable what you want\n",
    "do_val_loss = True  # compute validation loss every 10 epochs\n",
    "do_val_plots = False  # generate validation plots every 20 epochs\n",
    "\n",
    "# lists to store losses for plotting\n",
    "train_losses = []\n",
    "val_losses = []\n",
    "\n",
    "with tqdm(range(0, cfg.training.num_epochs)) as pbar:\n",
    "    for epoch in pbar:\n",
    "        ae.train()\n",
    "        epoch_losses = []\n",
    "\n",
    "        for i, data in enumerate(train_dataloader):\n",
    "            optimizer.zero_grad()\n",
    "            signal = data[\"signal\"]\n",
    "\n",
    "            # coordinated dropout masking\n",
    "            # randomly mask out some inputs and scale remaining ones\n",
    "            mask_prob = cfg.training.get(\"mask_prob\", 0.5) # if not specified, mask out half of the inputs\n",
    "            mask = (torch.rand_like(signal) > mask_prob).float()  \n",
    "            input_signal = signal * (mask / (1 - mask_prob))  # scale unmasked values\n",
    "\n",
    "            # forward pass\n",
    "            output_rates, z = ae(input_signal)\n",
    "            numel = signal.shape[0] * signal.shape[1] * signal.shape[2]\n",
    "\n",
    "            # compute losses\n",
    "            unmasked = (1 - mask) if mask_prob > 0 else torch.ones_like(mask)\n",
    "            poisson_loss = criterion_poisson(output_rates, signal) * unmasked\n",
    "            poisson_loss = poisson_loss.mean()\n",
    "            rec_loss = poisson_loss\n",
    "\n",
    "            # add latent regularization\n",
    "            latent_loss = latent_regularizer(z, cfg) / numel\n",
    "            loss = rec_loss + cfg.training.latent_beta * latent_loss\n",
    "\n",
    "            # optimization step\n",
    "            accelerator.backward(loss)\n",
    "            accelerator.clip_grad_norm_(ae.parameters(), 2.0)\n",
    "            optimizer.step()\n",
    "            lr_scheduler.step()\n",
    "\n",
    "            # logging\n",
    "            pbar.set_postfix(\n",
    "                **{\n",
    "                    \"rec_loss\": rec_loss.item(),\n",
    "                    \"latent_loss\": latent_loss.item(), \n",
    "                    \"total_loss\": loss.item(),\n",
    "                    \"lr\": optimizer.param_groups[0][\"lr\"],\n",
    "                    \"val_poisson_loss\": avg_poisson_loss,\n",
    "                }\n",
    "            )\n",
    "\n",
    "            # store metrics\n",
    "            rec_losses.append(rec_loss.item())\n",
    "            latent_losses.append(latent_loss.item())\n",
    "            total_losses.append(loss.item())\n",
    "            lrs.append(optimizer.param_groups[0][\"lr\"])\n",
    "            epoch_losses.append(loss.item())\n",
    "\n",
    "        # store average training loss for this epoch\n",
    "        train_losses.append(np.mean(epoch_losses))\n",
    "\n",
    "        # validation and checkpointing\n",
    "        if accelerator.is_main_process:\n",
    "            # compute validation loss every 10 epochs\n",
    "            if do_val_loss and ((epoch + 1) % 10 == 0 or epoch == cfg.training.num_epochs - 1):\n",
    "                avg_poisson_loss = compute_val_loss(ae, val_dataloader)\n",
    "                val_losses.append(avg_poisson_loss)\n",
    "\n",
    "            # generate plots and save model every 20 epochs\n",
    "            if (epoch + 1) % 20 == 0 or epoch == cfg.training.num_epochs - 1:\n",
    "                ae.eval()\n",
    "                \n",
    "                # plot training curves\n",
    "                plt.figure(figsize=cm2inch(6, 4))\n",
    "                plt.plot(train_losses, label='train loss')\n",
    "                plt.plot(np.arange(9, len(val_losses)*10, 10), val_losses, label='val loss')\n",
    "                plt.xlabel('epoch')\n",
    "                plt.ylabel('loss')\n",
    "                plt.yscale('log')\n",
    "                plt.legend()\n",
    "                plt.title('training and validation losses')\n",
    "                plt.show()\n",
    "                \n",
    "                # generate validation plots\n",
    "                if do_val_plots:\n",
    "                    plot_rate_traces_real(ae, val_dataloader, figsize=(6, 6), idx=1)\n",
    "                    imshow_rates_real(ae, val_dataloader, figsize=(6, 3), idx=1)\n",
    "                \n",
    "                # save checkpoint\n",
    "                accelerator.save_state(f\"exp/{cfg.exp_name}/epoch_{(epoch+20)//20*20}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# at the end, save the final model\n",
    "torch.save(ae.state_dict(), f\"exp/{cfg.exp_name}/model.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train the diffusion model, use the `train_diffusion_monkey.ipynb` notebook, where this trained autoencoder is used.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "timeseries",
   "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}