baselines/tndm/tndm_reach.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "093b3cbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b228b29f-91ad-42b3-92e5-406bbadfe804",
   "metadata": {},
   "outputs": [],
   "source": [
    "%config Completer.use_jedi = False\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "\n",
    "import os\n",
    "import tensorflow as tf\n",
    "tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)\n",
    "\n",
    "use_cpu = False\n",
    "cuda_device = '0'\n",
    "\n",
    "if use_cpu:\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'\n",
    "    os.environ['CPU_ONLY'] = \"TRUE\"\n",
    "    physical_devices = tf.config.list_physical_devices('CPU')\n",
    "    tf.config.set_logical_device_configuration(\n",
    "        physical_devices[0],\n",
    "        [tf.config.LogicalDeviceConfiguration() for i in range(8)])\n",
    "    logical_devices = tf.config.list_logical_devices('CPU')\n",
    "\n",
    "    print(logical_devices)\n",
    "else:\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = cuda_device\n",
    "    os.environ['CPU_ONLY'] = \"FALSE\"\n",
    "    physical_devices = tf.config.list_physical_devices('GPU')\n",
    "    print(physical_devices)\n",
    "    \n",
    "from collections import defaultdict\n",
    "from sklearn.metrics import r2_score\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from datetime import datetime\n",
    "\n",
    "from tndm.data import DataManager\n",
    "from tndm import TNDM\n",
    "from tndm.runtime import Runtime, ModelType\n",
    "from tndm.utils import AdaptiveWeights\n",
    "from tndm.models.model_loader import ModelLoader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c459601",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e58556b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.sysconfig.get_build_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b99ed371-f805-486c-8251-6249a9b1515c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2008, 140, 182) (2008, 140, 2)\n",
      "(71, 140, 182) (71, 140, 2)\n",
      "(216, 140, 182) (216, 140, 2)\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import h5py\n",
    "\n",
    "# use teh same h5 file as in LFADS\n",
    "with h5py.File('data/baselines/lfads/data/datasets/monkey_5ms.h5', 'r') as f:\n",
    "    dataset = {key: np.array(f[key]) for key in f.keys()}\n",
    "\n",
    "\n",
    "\n",
    "# test set is combined valid and test\n",
    "neural_data = dataset['train_data'].astype('float')\n",
    "valid_neural_data = dataset['valid_data'].astype('float')\n",
    "test_neural_data = dataset['test_data'].astype('float')\n",
    "\n",
    "behavioural_data = dataset['train_beh'].astype('float')\n",
    "valid_behavioural_data = dataset['valid_beh'].astype('float')\n",
    "test_behavioural_data = dataset['test_beh'].astype('float')\n",
    "\n",
    "\n",
    "print(neural_data.shape, behavioural_data.shape)\n",
    "print(valid_neural_data.shape, valid_behavioural_data.shape)\n",
    "print(test_neural_data.shape, test_behavioural_data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "983a9d83-e8ad-4510-acec-0824b792c567",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# center behaviour at zero, using first time step (not strictly required)\n",
    "b_mean = np.mean(behavioural_data[:,0,:],axis=0)\n",
    "for i in range(2):\n",
    "    behavioural_data[:,:,i] = behavioural_data[:,:,i]-b_mean[i]\n",
    "    valid_behavioural_data[:,:,i] = valid_behavioural_data[:,:,i]-b_mean[i]\n",
    "    test_behavioural_data[:,:,i] = test_behavioural_data[:,:,i]-b_mean[i]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "068b9011-2934-4c02-aeea-101242151d6b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# model parameters\n",
    "\n",
    "# l2 regulariser for the recurrent decoder weights\n",
    "l2_reg = 1\n",
    "initial_neural_weight = 1.0 # weight of neural nll\n",
    "initial_behaviour_weight = .2 # weight of behaviour loss\n",
    "lambda_q = 100.0\n",
    "update_rate = .0005 \n",
    "dropout = .15\n",
    "seed = 0\n",
    "GRU_pre_activation = False\n",
    "var_min = 0.0001\n",
    "prior_variance = 1\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(\n",
    "    learning_rate=1e-2,\n",
    "    beta_1=0.9, \n",
    "    beta_2=0.999,\n",
    "    epsilon=1e-08)\n",
    "\n",
    "layers_settings=defaultdict(lambda: dict(\n",
    "    kernel_initializer=tf.keras.initializers.VarianceScaling(\n",
    "        scale=1.0, mode='fan_in', distribution='normal'),\n",
    "    kernel_regularizer=tf.keras.regularizers.l2(l2=0.0)\n",
    "))\n",
    "\n",
    "layers_settings['encoder'].update(dict(var_min=var_min, var_trainable=True))\n",
    "layers_settings['relevant_decoder'].update(dict(kernel_regularizer=tf.keras.regularizers.l2(l2=0),\n",
    "                                      recurrent_regularizer=tf.keras.regularizers.l2(l2=l2_reg),\n",
    "                                      original_cell=False))    \n",
    "layers_settings['irrelevant_decoder'].update(dict(kernel_regularizer=tf.keras.regularizers.l2(l2=0),\n",
    "                                      recurrent_regularizer=tf.keras.regularizers.l2(l2=l2_reg),\n",
    "                                      original_cell=False))   \n",
    "layers_settings['behavioural_dense'].update(dict(behaviour_type='causal'))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e003ee5b-e0f5-4d83-bc53-632e65f67abc",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = datetime.today().strftime(\"%y_%m_%d_%X\")\n",
    "\n",
    "spike_data_dir = \"tndm_exp\"\n",
    "\n",
    "logdir = os.path.join( spike_data_dir, 'log_l2_reg_'+str(l2_reg)+'_do_' + str(dropout)+'_no_norm_'+T)\n",
    "modeldir = os.path.join( spike_data_dir, 'model_l2_reg_'+str(l2_reg)+'_do_' + str(dropout)+'_no_norm_'+T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c537bdff-f267-4fde-b31b-d3d5f34da2e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "T0 = datetime.now()\n",
    "\n",
    "model, history = Runtime.train(\n",
    "    model_type=ModelType.TNDM,\n",
    "    adaptive_lr=dict(factor=0.95, patience=10, min_lr=1e-5),\n",
    "    model_settings=dict(\n",
    "        rel_factors=5,\n",
    "        irr_factors=5,\n",
    "        encoded_dim=64,\n",
    "        max_grad_norm=200,\n",
    "        dropout=dropout,\n",
    "        prior_variance=prior_variance,\n",
    "        GRU_pre_activation=GRU_pre_activation, #NEW\n",
    "        timestep=0.005, # hardcoded\n",
    "        seed=seed\n",
    "    ),\n",
    "    layers_settings=layers_settings,\n",
    "    optimizer=optimizer, \n",
    "    epochs=1000, \n",
    "    # logdir=logdir,\n",
    "    logdir=None,\n",
    "    train_dataset=(neural_data, behavioural_data), \n",
    "    val_dataset=(valid_neural_data, valid_behavioural_data),\n",
    "    adaptive_weights=AdaptiveWeights(\n",
    "        initial=[initial_neural_weight, initial_behaviour_weight, .0, .0, lambda_q, .0], #changed\n",
    "        update_start=[0, 0, 0, 1000, 1000, 0],\n",
    "        update_rate=[0., 0., update_rate, update_rate, 0.0, update_rate],\n",
    "        min_weight=[initial_neural_weight, initial_behaviour_weight, 0.0, 0.0, lambda_q, 0.0],#changed\n",
    "        max_weight=[initial_neural_weight, initial_behaviour_weight, 1.0, 1.0, lambda_q, 1.0],#changed\n",
    "    ),\n",
    "    batch_size=128,\n",
    "    verbose=2 # set to 2 to see the losses during training\n",
    ")\n",
    "\n",
    "model.save(modeldir)\n",
    "\n",
    "print('Training took '+str(datetime.now()-T0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "da7d3921",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.layers_settings.default_factory  = lambda : model.layers_settings.default_factory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67ce909f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(modeldir)\n",
    "\n",
    "print('Training took '+str(datetime.now()-T0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "26bbe845",
   "metadata": {},
   "outputs": [],
   "source": [
    "# manual\n",
    "modeldir = 'tndm_exp'\n",
    "from tndm.utils import CustomEncoder, upsert_empty_folder\n",
    "location = modeldir\n",
    "import json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2274b1c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "settings = model.get_settings()\n",
    "upsert_empty_folder(location)\n",
    "with open(os.path.join(location, \"settings.json\"), \"w\") as fp:\n",
    "    json.dump(settings, fp, cls=CustomEncoder)\n",
    "model.save_weights(os.path.join(location, \".weights.h5\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e97971c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "modeldir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d633ed8",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20d86842-203f-4049-b3ad-a9451d21c336",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ModelLoader.load(modeldir, model_class=TNDM)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f770ebf8-abcf-4641-ab68-e6912de205bc",
   "metadata": {},
   "source": [
    "# Latent space"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f84a787-3ee6-4f6b-bdba-263aea1cc211",
   "metadata": {},
   "source": [
    "## Training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1866305e-40b7-4687-a3d9-3783261e770b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "data = neural_data\n",
    "test_sample_mode = 'posterior_sample' #choose 'mean' for previous behaviour\n",
    "\n",
    "if test_sample_mode == 'mean':\n",
    "    log_f, b, (g0_r, mean_r, logvar_r), (g0_i, mean_i, logvar_i), (z_r, z_i) = \\\n",
    "        model(data.astype('float'), training=False, test_sample_mode=test_sample_mode)\n",
    "else:\n",
    "    batch_size_eval = 128\n",
    "    log_fs = []\n",
    "    bs = []\n",
    "    g0_rs, mean_rs, logvar_rs, z_rs = [], [], [], []\n",
    "    g0_is, mean_is, logvar_is, z_is = [], [], [], []\n",
    "    for neural_datum in tqdm(data):\n",
    "        neural_datum_batch = np.repeat(np.expand_dims(neural_datum, 0), batch_size_eval, axis=0)\n",
    "        log_f, b, (g0_r, mean_r, logvar_r), (g0_i, mean_i, logvar_i), (z_r, z_i) = \\\n",
    "            model(neural_datum_batch.astype('float'), training=False, test_sample_mode=test_sample_mode)\n",
    "        log_fs.append(np.mean(log_f, 0))\n",
    "        bs.append(np.mean(b, 0))\n",
    "        g0_rs.append(np.mean(g0_r, 0))\n",
    "        mean_rs.append(np.mean(mean_r, 0))\n",
    "        logvar_rs.append(np.mean(logvar_r, 0))\n",
    "        z_rs.append(np.mean(z_r, 0))\n",
    "        g0_is.append(np.mean(g0_i, 0))\n",
    "        mean_is.append(np.mean(mean_i, 0))\n",
    "        logvar_is.append(np.mean(logvar_i, 0))\n",
    "        z_is.append(np.mean(z_i, 0))\n",
    "    log_f = tf.stack(log_fs)\n",
    "    b = tf.stack(bs)\n",
    "    g0_r = tf.stack(g0_rs)\n",
    "    mean_r = tf.stack(mean_rs)\n",
    "    logvar_r = tf.stack(logvar_rs)\n",
    "    z_r = tf.stack(z_rs)\n",
    "    g0_i = tf.stack(g0_is)\n",
    "    mean_i = tf.stack(mean_is)\n",
    "    logvar_i = tf.stack(logvar_is)\n",
    "    z_i = tf.stack(z_is)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2702db69",
   "metadata": {},
   "source": [
    "##  test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "9d7af077",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "data = neural_data\n",
    "# data = test_neural_data\n",
    "test_sample_mode = 'mean' #choose 'mean' for previous behaviour\n",
    "\n",
    "if test_sample_mode == 'mean':\n",
    "    log_f, b, (g0_r, mean_r, logvar_r), (g0_i, mean_i, logvar_i), (z_r, z_i) = \\\n",
    "        model(data.astype('float'), training=False, test_sample_mode=test_sample_mode)\n",
    "else:\n",
    "    batch_size_eval = 128\n",
    "    log_fs = []\n",
    "    bs = []\n",
    "    g0_rs, mean_rs, logvar_rs, z_rs = [], [], [], []\n",
    "    g0_is, mean_is, logvar_is, z_is = [], [], [], []\n",
    "    for neural_datum in tqdm(data):\n",
    "        neural_datum_batch = np.repeat(np.expand_dims(neural_datum, 0), batch_size_eval, axis=0)\n",
    "        log_f, b, (g0_r, mean_r, logvar_r), (g0_i, mean_i, logvar_i), (z_r, z_i) = \\\n",
    "            model(neural_datum_batch.astype('float'), training=False, test_sample_mode=test_sample_mode)\n",
    "        log_fs.append(np.mean(log_f, 0))\n",
    "        bs.append(np.mean(b, 0))\n",
    "        g0_rs.append(np.mean(g0_r, 0))\n",
    "        mean_rs.append(np.mean(mean_r, 0))\n",
    "        logvar_rs.append(np.mean(logvar_r, 0))\n",
    "        z_rs.append(np.mean(z_r, 0))\n",
    "        g0_is.append(np.mean(g0_i, 0))\n",
    "        mean_is.append(np.mean(mean_i, 0))\n",
    "        logvar_is.append(np.mean(logvar_i, 0))\n",
    "        z_is.append(np.mean(z_i, 0))\n",
    "    log_f = tf.stack(log_fs)\n",
    "    b = tf.stack(bs)\n",
    "    g0_r = tf.stack(g0_rs)\n",
    "    mean_r = tf.stack(mean_rs)\n",
    "    logvar_r = tf.stack(logvar_rs)\n",
    "    z_r = tf.stack(z_rs)\n",
    "    g0_i = tf.stack(g0_is)\n",
    "    mean_i = tf.stack(mean_is)\n",
    "    logvar_i = tf.stack(logvar_is)\n",
    "    z_i = tf.stack(z_is)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "5703d1d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "log_f (2008, 140, 182)\n",
      "b (2008, 140, 2)\n",
      "g0_r (2008, 64)\n",
      "mean_r (2008, 64)\n",
      "logvar_r (2008, 64)\n",
      "z_r (2008, 140, 5)\n",
      "g0_i (2008, 64)\n",
      "mean_i (2008, 64)\n",
      "logvar_i (2008, 64)\n",
      "z_i (2008, 140, 5)\n",
      "ae_rates (2008, 140, 182)\n",
      "ae_latents_relevant (2008, 140, 5)\n",
      "ae_latents_irrelevant (2008, 140, 5)\n",
      "ae_behaviour (2008, 140, 2)\n",
      "gt_spikes (2008, 140, 182)\n",
      "init_states_gt_relevant (2008, 64)\n",
      "init_states_gt_irrelevant (2008, 64)\n"
     ]
    }
   ],
   "source": [
    "print(\"log_f\", log_f.shape)\n",
    "print(\"b\", b.shape)\n",
    "print(\"g0_r\", g0_r.shape)\n",
    "print(\"mean_r\", mean_r.shape)\n",
    "print(\"logvar_r\", logvar_r.shape)\n",
    "print(\"z_r\", z_r.shape)\n",
    "print(\"g0_i\", g0_i.shape)\n",
    "print(\"mean_i\", mean_i.shape)\n",
    "print(\"logvar_i\", logvar_i.shape)\n",
    "print(\"z_i\", z_i.shape)\n",
    "\n",
    "\n",
    "\n",
    "# ret_dict_test = {}\n",
    "# ret_dict_test[\"ae_rates\"] = (0.005*tf.math.exp(log_f)).numpy()\n",
    "# ret_dict_test[\"ae_latents_relevant\"] = z_r.numpy()\n",
    "# ret_dict_test[\"ae_latents_irrelevant\"] = z_i.numpy()\n",
    "# ret_dict_test[\"ae_behaviour\"] = b.numpy()\n",
    "# ret_dict_test[\"gt_spikes\"] = test_neural_data\n",
    "# ret_dict_test[\"init_states_gt_relevant\"] = g0_r.numpy()\n",
    "# ret_dict_test[\"init_states_gt_irrelevant\"] = g0_i.numpy()\n",
    "\n",
    "\n",
    "\n",
    "ret_dict_train = {}\n",
    "ret_dict_train[\"ae_rates\"] = (0.005*tf.math.exp(log_f)).numpy()\n",
    "ret_dict_train[\"ae_latents_relevant\"] = z_r.numpy()\n",
    "ret_dict_train[\"ae_latents_irrelevant\"] = z_i.numpy()\n",
    "ret_dict_train[\"ae_behaviour\"] = b.numpy()\n",
    "ret_dict_train[\"gt_spikes\"] = neural_data\n",
    "ret_dict_train[\"init_states_gt_relevant\"] = g0_r.numpy()\n",
    "ret_dict_train[\"init_states_gt_irrelevant\"] = g0_i.numpy()\n",
    "\n",
    "for key, val in ret_dict_train.items():\n",
    "    print(key, val.shape)\n",
    "\n",
    "\n",
    "\n",
    "import pickle\n",
    "with open('tndm_exp/tndm_samples.pkl', 'wb') as f:\n",
    "    pickle.dump(ret_dict_train, f)\n",
    "    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tndm",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}