train_sensorium_shifter_table_1&2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from nnfabrik.builder import get_data, get_model, get_trainer\n",
    "\n",
    "import movi\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from torchvision import transforms\n",
    "from collections import OrderedDict\n",
    "import itertools\n",
    "from functools import partial\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/yuchen/2023-neurips-mouse\n"
     ]
    }
   ],
   "source": [
    "%cd /home/yuchen/2023-neurips-mouse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Subset, DataLoader, ConcatDataset\n",
    "from mouse_model.data_utils import MouseDatasetSeg\n",
    "from mouse_model.evaluation import cor_in_time\n",
    "\n",
    "\n",
    "def load_train_val_ds():\n",
    "    ds_list = [MouseDatasetSeg(file_id=args.file_id, seg_idx=i, data_split=\"train\", vid_type=\"vid_shift_mean\", \n",
    "                           seq_len=args.seq_len, predict_offset=1, smoothing=None, behav_prod=True) for i in range(args.split_range)]\n",
    "    train_ds, val_ds = [], []\n",
    "    for ds in ds_list:\n",
    "        train_ratio = 0.8\n",
    "        train_ds_len = int(len(ds) * train_ratio)\n",
    "        train_ds.append(Subset(ds, np.arange(0, train_ds_len, 1)))\n",
    "        val_ds.append(Subset(ds, np.arange(train_ds_len, len(ds), 1)))\n",
    "    train_ds = ConcatDataset(train_ds)\n",
    "    val_ds = ConcatDataset(val_ds)\n",
    "    return train_ds, val_ds\n",
    "\n",
    "def load_test_ds():\n",
    "    test_ds = [MouseDatasetSeg(file_id=args.file_id, seg_idx=i, data_split=\"test\", vid_type=\"vid_shift_mean\", \n",
    "                               seq_len=args.seq_len, predict_offset=1, smoothing=None, behav_prod=False) \n",
    "               for i in range(args.split_range)]\n",
    "    test_ds = ConcatDataset(test_ds)\n",
    "    return test_ds\n",
    "\n",
    "def normalize_movie(movie):\n",
    "    \"\"\"Normalize the range of gray levels in a movie\"\"\"\n",
    "    norm_movie = movie.astype(float)\n",
    "    norm_movie -= norm_movie.min()\n",
    "    if not np.isclose(norm_movie.max(), 0):\n",
    "        norm_movie /= norm_movie.max()\n",
    "    return norm_movie\n",
    "\n",
    "def normalize_movie_neg_pos_1(movie):\n",
    "    \"\"\"Normalize the range of gray levels in a movie\"\"\"\n",
    "    movie = normalize_movie(movie)\n",
    "    movie = (movie - 0.5) * 2\n",
    "    return movie\n",
    "\n",
    "def smoothing_with_np_conv(nsp, size=int(2000/48), return_fr=False):\n",
    "    np_conv_res = []\n",
    "    for i in range(nsp.shape[1]):\n",
    "        np_conv_res.append(np.convolve(nsp[:, i], np.ones(size)/size, mode=\"same\"))        \n",
    "    np_conv_res = np.transpose(np.array(np_conv_res))\n",
    "    if return_fr:\n",
    "        return np_conv_res/0.016\n",
    "    return np_conv_res\n",
    "\n",
    "# https://stackoverflow.com/questions/52201081/compute-product-of-all-combinations-of-a-list\n",
    "def get_all_comb_prod(array):\n",
    "    res = []\n",
    "    combinations_with_r = partial(lambda r: itertools.combinations(array, r = r))\n",
    "    for r in map(combinations_with_r, range(1, len(array) + 1)):\n",
    "        for j in r:\n",
    "            res.append(np.prod(j))       \n",
    "    return res\n",
    "\n",
    "class MouseDatasetSeg(Dataset):\n",
    "    \n",
    "    def __init__(self, file_id, seg_idx, data_split=\"train\", vid_type=\"vid_shift_mean\", \n",
    "                 seq_len=1, predict_offset=1, smoothing=None, behav_prod=False):\n",
    "        \n",
    "        self.seq_len = seq_len\n",
    "        self.predict_offset = predict_offset\n",
    "        self.behav_prod = behav_prod\n",
    "        \n",
    "        data_dir = \"{}/{}\".format(ROOT_DIR_48ms_3_segment, file_id)\n",
    "        print(data_dir)\n",
    "        \n",
    "        # firing rate below 3 Hz\n",
    "        bad_neuron_index_dict = {\"070921_J553RT\": [1,  20,  21,  25,  27,  30,  31,  35,  39,  45,  46,  49,  50,  \n",
    "                                                   51,  52,  53,  54,  55,  57,  58,  59,  64,  65,  68,  75,  77,  \n",
    "                                                   79,  81,  83,  84,  87,  88,  90,  91,  94,  95,  97, 104, 105, 107], \n",
    "                                 \"101521_J559NC\": [1,  4,  9, 22, 25, 27, 28, 29, 43, 45, 46, 47, 51, 53], \n",
    "                                 \"110421_J569LT\": [0,  1,  2,  5,  6,  8, 11, 12, 13, 16, 18, 22, 24, 26, 28, 30, 41, \n",
    "                                                   44, 47, 49]}\n",
    "        \n",
    "        all_nsp = np.load(\"{}/{}_nsp_seg_{}.npy\".format(data_dir, data_split, seg_idx))\n",
    "        good_nsp = []\n",
    "        for i in range(all_nsp.shape[1]):\n",
    "            if i not in bad_neuron_index_dict[file_id]:\n",
    "                good_nsp.append(all_nsp[:, i])\n",
    "        self.nsp = np.transpose(np.array(good_nsp))\n",
    "        \n",
    "            \n",
    "        self.images = np.load(\"{}/{}_{}_seg_{}.npy\".format(data_dir, data_split, vid_type, seg_idx))\n",
    "        self.images = normalize_movie(self.images)\n",
    "        self.images = np.expand_dims(self.images, axis=1)\n",
    "        \n",
    "#         behav_key_list = ['speed', 'gz', 'pitch', 'roll', 'phi', 'th', 'eyerad']\n",
    "        behav_key_list = ['eyerad','speed',  'th', 'phi']\n",
    "        behavior_var_list = [np.load(\"{}/{}_{}_seg_{}.npy\".format(data_dir, data_split, behav_key, seg_idx)) \n",
    "                             for behav_key in behav_key_list]\n",
    "        \n",
    "        behavior_var_list[0] = [0 if math.isnan(x) else x for x in behavior_var_list[0]]\n",
    "        temp_lst = []\n",
    "        for i in range(len(behavior_var_list[0])):\n",
    "            behavior_var_list[0][i] = ((behavior_var_list[0][i])**2) * np.pi\n",
    "            if i == 0:\n",
    "                temp_lst.append(0)\n",
    "            else:\n",
    "                temp_lst.append((behavior_var_list[0][i] - behavior_var_list[0][i-1])/48)\n",
    "\n",
    "        behavior_var_list.insert(1, temp_lst)\n",
    "\n",
    "        for i in range(len(behavior_var_list)):\n",
    "            behavior_var_list[i] = [0 if math.isnan(x) else x for x in behavior_var_list[i]]\n",
    "            behavior_var_list[i] = (behavior_var_list[i])/np.std(behavior_var_list[i])\n",
    "   \n",
    "        self.behavior_var = np.stack(behavior_var_list, axis=1)\n",
    "        self.behavior_var = np.nan_to_num(self.behavior_var)\n",
    "\n",
    "    def __len__(self):\n",
    "        len_block = self.seq_len + self.predict_offset\n",
    "        return self.images.shape[0] - len_block + 1\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        current_frame = self.images[idx:(idx+self.seq_len)]\n",
    "\n",
    "        if self.behav_prod:\n",
    "            current_behavior_var = []\n",
    "            for i in range(idx, idx+self.seq_len):\n",
    "                behav_var_frame = np.delete(self.behavior_var[i], 1)\n",
    "                current_behavior_var.append(get_all_comb_prod(behav_var_frame))\n",
    "            current_behavior_var = torch.tensor(current_behavior_var, dtype=torch.float)\n",
    "        else:\n",
    "            current_behavior_var = torch.tensor(self.behavior_var[idx:(idx+self.seq_len)], dtype=torch.float)\n",
    "        \n",
    "        current_behavior_var = current_behavior_var.squeeze().cuda()\n",
    "#         print(current_behavior_var)\n",
    "        neural_spikes = torch.tensor(self.nsp[idx+self.seq_len-1+self.predict_offset], dtype=torch.float).cuda()\n",
    "\n",
    "        current_frame = current_frame.squeeze()\n",
    "        \n",
    "        # for vision + behavior\n",
    "        current_frame = np.array([current_frame, \n",
    "                                  np.repeat(current_behavior_var[0].item(), 60*80).reshape(60,-1),\n",
    "                                 np.repeat(current_behavior_var[1].item(), 60*80).reshape(60,-1), \n",
    "                                 np.repeat(current_behavior_var[2].item(), 60*80).reshape(60,-1)])\n",
    "        current_frame = torch.tensor(current_frame, dtype=torch.float).cuda()\n",
    "        return {\"images\": current_frame, \"responses\": neural_spikes, \n",
    "                'behavior': current_behavior_var[:3], 'pupil_center':current_behavior_var[3:] }\n",
    "        \n",
    "        # for vision\n",
    "        # current_frame = np.array([current_frame]) \n",
    "        # current_frame = torch.tensor(current_frame, dtype=torch.float).cuda()\n",
    "        # return {\"images\": current_frame, \"responses\": neural_spikes, 'pupil_center':current_behavior_var[3:]}\n",
    "    \n",
    "    \n",
    "def get_dataloaders_one_file(train_ratio=0.8, batch_size=256):\n",
    "    \n",
    "    ds_list = [MouseDatasetSeg(file_id=args.file_id, seg_idx=i, data_split=\"train\", vid_type=args.vid_ty, \n",
    "                           seq_len=args.seq_len, predict_offset=1, smoothing=None, behav_prod=False) \n",
    "               for i in range(args.split_range)]\n",
    "    print('current split_range is: ' + str(args.split_range))\n",
    "    print('current vid_type is: ' + str(args.vid_ty))\n",
    "    train_ds, val_ds = [], []\n",
    "#     dl1 = TensorDataset()\n",
    "    for ds in ds_list:\n",
    "        train_ratio = 0.8\n",
    "        train_ds_len = int(len(ds) * train_ratio)\n",
    "        train_ds.append(Subset(ds, np.arange(0, train_ds_len, 1)))\n",
    "        val_ds.append(Subset(ds, np.arange(train_ds_len, len(ds), 1)))\n",
    "    train_ds = ConcatDataset(train_ds)\n",
    "    val_ds = ConcatDataset(val_ds)\n",
    "\n",
    "    test_ds = [MouseDatasetSeg(file_id=args.file_id, seg_idx=i, data_split=\"test\", vid_type=args.vid_ty, \n",
    "                               seq_len=args.seq_len, predict_offset=1, smoothing=None, behav_prod=False)  \n",
    "               for i in range(10)]\n",
    "    test_ds = ConcatDataset(test_ds)\n",
    "\n",
    "\n",
    "    dataloaders = OrderedDict()\n",
    "    dataloaders['train'] = OrderedDict()\n",
    "    dataloaders['train'][args.file_id] = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
    "    dataloaders['validation'] = OrderedDict()\n",
    "    dataloaders['validation'][args.file_id] = DataLoader(val_ds, batch_size=batch_size, shuffle=False)\n",
    "    dataloaders['test'] = OrderedDict()\n",
    "    dataloaders['test'][args.file_id] = DataLoader(test_ds, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "    return dataloaders\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "sensorium baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n",
      "110421_J569LT\n",
      "seed: 0\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "current split_range is: 10\n",
      "current vid_type is: vid_mean\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 25\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mseed: \u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mstr\u001b[39m(args\u001b[38;5;241m.\u001b[39mseed))\n\u001b[1;32m     23\u001b[0m ROOT_DIR_48ms_3_segment \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m---> 25\u001b[0m dataloaders \u001b[38;5;241m=\u001b[39m \u001b[43mget_dataloaders_one_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     27\u001b[0m model_fn \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124msensorium.models.stacked_core_full_gauss_readout\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m     28\u001b[0m model_config \u001b[38;5;241m=\u001b[39m {\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpad_input\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     29\u001b[0m   \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mstack\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     30\u001b[0m   \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mlayers\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;241m4\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     45\u001b[0m   \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mshifter\u001b[39m\u001b[38;5;124m'\u001b[39m: \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     46\u001b[0m }\n",
      "Cell \u001b[0;32mIn[3], line 167\u001b[0m, in \u001b[0;36mget_dataloaders_one_file\u001b[0;34m(train_ratio, batch_size)\u001b[0m\n\u001b[1;32m    164\u001b[0m train_ds \u001b[38;5;241m=\u001b[39m ConcatDataset(train_ds)\n\u001b[1;32m    165\u001b[0m val_ds \u001b[38;5;241m=\u001b[39m ConcatDataset(val_ds)\n\u001b[0;32m--> 167\u001b[0m test_ds \u001b[38;5;241m=\u001b[39m [MouseDatasetSeg(file_id\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mfile_id, seg_idx\u001b[38;5;241m=\u001b[39mi, data_split\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtest\u001b[39m\u001b[38;5;124m\"\u001b[39m, vid_type\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mvid_ty, \n\u001b[1;32m    168\u001b[0m                            seq_len\u001b[38;5;241m=\u001b[39margs\u001b[38;5;241m.\u001b[39mseq_len, predict_offset\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, smoothing\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, behav_prod\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)  \n\u001b[1;32m    169\u001b[0m            \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m10\u001b[39m)]\n\u001b[1;32m    170\u001b[0m test_ds \u001b[38;5;241m=\u001b[39m ConcatDataset(test_ds)\n\u001b[1;32m    173\u001b[0m dataloaders \u001b[38;5;241m=\u001b[39m OrderedDict()\n",
      "Cell \u001b[0;32mIn[3], line 167\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    164\u001b[0m train_ds \u001b[38;5;241m=\u001b[39m ConcatDataset(train_ds)\n\u001b[1;32m    165\u001b[0m val_ds \u001b[38;5;241m=\u001b[39m ConcatDataset(val_ds)\n\u001b[0;32m--> 167\u001b[0m test_ds \u001b[38;5;241m=\u001b[39m [\u001b[43mMouseDatasetSeg\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfile_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfile_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mseg_idx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mi\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata_split\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtest\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvid_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvid_ty\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m    168\u001b[0m \u001b[43m                           \u001b[49m\u001b[43mseq_len\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mseq_len\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpredict_offset\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msmoothing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbehav_prod\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m  \n\u001b[1;32m    169\u001b[0m            \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m10\u001b[39m)]\n\u001b[1;32m    170\u001b[0m test_ds \u001b[38;5;241m=\u001b[39m ConcatDataset(test_ds)\n\u001b[1;32m    173\u001b[0m dataloaders \u001b[38;5;241m=\u001b[39m OrderedDict()\n",
      "Cell \u001b[0;32mIn[3], line 87\u001b[0m, in \u001b[0;36mMouseDatasetSeg.__init__\u001b[0;34m(self, file_id, seg_idx, data_split, vid_type, seq_len, predict_offset, smoothing, behav_prod)\u001b[0m\n\u001b[1;32m     83\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnsp \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mtranspose(np\u001b[38;5;241m.\u001b[39marray(good_nsp))\n\u001b[1;32m     86\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mimages \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mload(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m_\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m_seg_\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m.npy\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(data_dir, data_split, vid_type, seg_idx))\n\u001b[0;32m---> 87\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mimages \u001b[38;5;241m=\u001b[39m \u001b[43mnormalize_movie\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimages\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     88\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mimages \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mexpand_dims(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mimages, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     90\u001b[0m \u001b[38;5;66;03m#         behav_key_list = ['speed', 'gz', 'pitch', 'roll', 'phi', 'th', 'eyerad']\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[3], line 30\u001b[0m, in \u001b[0;36mnormalize_movie\u001b[0;34m(movie)\u001b[0m\n\u001b[1;32m     28\u001b[0m norm_movie \u001b[38;5;241m=\u001b[39m movie\u001b[38;5;241m.\u001b[39mastype(\u001b[38;5;28mfloat\u001b[39m)\n\u001b[1;32m     29\u001b[0m norm_movie \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m norm_movie\u001b[38;5;241m.\u001b[39mmin()\n\u001b[0;32m---> 30\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m np\u001b[38;5;241m.\u001b[39misclose(\u001b[43mnorm_movie\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m, \u001b[38;5;241m0\u001b[39m):\n\u001b[1;32m     31\u001b[0m     norm_movie \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m=\u001b[39m norm_movie\u001b[38;5;241m.\u001b[39mmax()\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m norm_movie\n",
      "File \u001b[0;32m~/anaconda3/envs/sensorium2023/lib/python3.9/site-packages/numpy/core/_methods.py:41\u001b[0m, in \u001b[0;36m_amax\u001b[0;34m(a, axis, out, keepdims, initial, where)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_amax\u001b[39m(a, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, out\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, keepdims\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     40\u001b[0m           initial\u001b[38;5;241m=\u001b[39m_NoValue, where\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[0;32m---> 41\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mumr_maximum\u001b[49m\u001b[43m(\u001b[49m\u001b[43ma\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maxis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkeepdims\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minitial\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mwhere\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for subj in ['110421_J569LT']:\n",
    "# for subj in ['110421_J569LT','101521_J559NC', '070921_J553RT']:\n",
    "    for i in [10]:\n",
    "        for seed in [0]:\n",
    "            class Args:\n",
    "                seed = 0\n",
    "                file_id = subj\n",
    "                epochs = 100\n",
    "                batch_size = 256\n",
    "                l1_weight_behav=6.2591\n",
    "                l1_weight_comb=3.9773\n",
    "                seq_len = 1\n",
    "                split_range = 10\n",
    "                vid_ty = 'vid_mean'\n",
    "                best_val_path = None\n",
    "                best_train_path = None\n",
    "\n",
    "            args=Args()\n",
    "\n",
    "            print(i)\n",
    "            print(subj)\n",
    "            print('seed: ' + str(args.seed))\n",
    "            ROOT_DIR_48ms_3_segment = \"/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms\"\n",
    "\n",
    "            dataloaders = get_dataloaders_one_file()\n",
    "\n",
    "            model_fn = 'sensorium.models.stacked_core_full_gauss_readout'\n",
    "            model_config = {'pad_input': False,\n",
    "              'stack': -1,\n",
    "              'layers': 4,\n",
    "              'input_kern': 9,\n",
    "              'gamma_input': 6.3831,\n",
    "              'gamma_readout': 0.0076,\n",
    "              'hidden_kern': 7,\n",
    "              'hidden_channels': 64,\n",
    "              'depth_separable': True,\n",
    "              'grid_mean_predictor': {'type': None,\n",
    "               'input_dimensions': 2,\n",
    "               'hidden_layers': 1,\n",
    "               'hidden_features': 30,\n",
    "               'final_tanh': True},\n",
    "              'init_sigma': 0.1,\n",
    "              'init_mu_range': 0.3,\n",
    "              'gauss_type': 'full',\n",
    "              'shifter': False,\n",
    "            }\n",
    "\n",
    "            model = get_model(model_fn=model_fn,\n",
    "                              model_config=model_config,\n",
    "                              dataloaders=dataloaders,\n",
    "                              seed=0,)\n",
    "\n",
    "            trainer_fn = \"sensorium.training.standard_trainer\"\n",
    "\n",
    "            trainer_config = {'max_iter': 100,\n",
    "                             'verbose': False,\n",
    "                             'lr_decay_steps': 4,\n",
    "                             'avg_loss': False,\n",
    "                             'lr_init': 0.009,\n",
    "                             }\n",
    "\n",
    "            trainer = get_trainer(trainer_fn=trainer_fn, \n",
    "                                 trainer_config=trainer_config)\n",
    "\n",
    "            validation_score, trainer_output, state_dict = trainer(model, dataloaders, seed=args.seed)\n",
    "            print(validation_score)\n",
    "\n",
    "            # weight = \"/hdd/yuchen/new_data_sensorium_sota_model_\" + str(subj) + \"_split\" + str(i) + '_'+ str(args.vid_ty) +\"seed\" + str(args.seed) + \"head.pth\"\n",
    "            # torch.save(model.state_dict(), weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "current split_range is: 10\n",
      "current vid_type is: vid_mean\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/110421_J569LT\n",
      "R2 0.358379\n",
      "MSE 0.100978\n",
      "mean corr, 0.441+-0.181\n",
      "max corr 0.769153\n",
      "min corr 0.027994\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "current split_range is: 10\n",
      "current vid_type is: vid_mean\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/101521_J559NC\n",
      "R2 0.745344\n",
      "MSE 0.097545\n",
      "mean corr, 0.487+-0.146\n",
      "max corr 0.730575\n",
      "min corr 0.020470\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "current split_range is: 10\n",
      "current vid_type is: vid_mean\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms/070921_J553RT\n",
      "R2 0.803191\n",
      "MSE 0.069641\n",
      "mean corr, 0.540+-0.138\n",
      "max corr 0.788901\n",
      "min corr 0.153343\n"
     ]
    }
   ],
   "source": [
    "class Args:\n",
    "                seed = 0\n",
    "                file_id = subj\n",
    "                epochs = 100\n",
    "                batch_size = 256\n",
    "                l1_weight_behav=6.2591\n",
    "                l1_weight_comb=3.9773\n",
    "                seq_len = 1\n",
    "                split_range = 10\n",
    "                vid_ty = 'vid_mean'\n",
    "                best_val_path = None\n",
    "                best_train_path = None\n",
    "\n",
    "args=Args()\n",
    "\n",
    "ROOT_DIR_48ms_3_segment = \"/hdd/aiwenxu/mouse-data-10-segment-split-70-30-48ms\"\n",
    "\n",
    "for subj in ['110421_J569LT','101521_J559NC', '070921_J553RT']:\n",
    "    args.file_id = subj\n",
    "    dataloaders = get_dataloaders_one_file()\n",
    "    model_fn = 'sensorium.models.stacked_core_full_gauss_readout'\n",
    "    model_config = {'pad_input': False,\n",
    "              'stack': -1,\n",
    "              'layers': 4,\n",
    "              'input_kern': 9,\n",
    "              'gamma_input': 6.3831,\n",
    "              'gamma_readout': 0.0076,\n",
    "              'hidden_kern': 7,\n",
    "              'hidden_channels': 64,\n",
    "              'depth_separable': True,\n",
    "              'grid_mean_predictor': {'type': None,\n",
    "               'input_dimensions': 2,\n",
    "               'hidden_layers': 1,\n",
    "               'hidden_features': 30,\n",
    "               'final_tanh': True},\n",
    "              'init_sigma': 0.1,\n",
    "              'init_mu_range': 0.3,\n",
    "              'gauss_type': 'full',\n",
    "              'shifter': True,\n",
    "    }\n",
    "\n",
    "\n",
    "    model = get_model(model_fn=model_fn,\n",
    "                                  model_config=model_config,\n",
    "                                  dataloaders=dataloaders,\n",
    "                                  seed=0,)\n",
    "    pat = \"/hdd/yuchen/new_data_sensorium_sota_model_{}_split10_vid_meanseed0vision.pth\".format(args.file_id)\n",
    "    model.load_state_dict(torch.load(pat))\n",
    "\n",
    "    i='test'\n",
    "    model.eval()\n",
    "    correlations,label,pred = get_correlations(model, dataloaders, tier=i, device=\"cuda:0\", as_dict=False, per_neuron=False)\n",
    "    cor_array = cor_in_time(pred, label)\n",
    "#     print(\"best val model on test dataset, {:.3f}+-{:.3f}\".format(np.mean(cor_array), np.std(cor_array)))\n",
    "    pred = smoothing_with_np_conv(pred)\n",
    "    label = smoothing_with_np_conv(label)\n",
    "    print(\"R2\", \"{:.6f}\".format(r2_score(label.T, pred.T)))\n",
    "    print(\"MSE\", \"{:.6f}\".format(mean_squared_error(label, pred)))\n",
    "    cor_array = cor_in_time(pred, label)\n",
    "    print(\"mean corr, {:.3f}+-{:.3f}\".format(np.mean(cor_array), np.std(cor_array)))\n",
    "    print(\"max corr\", \"{:.6f}\".format(np.max(cor_array)))\n",
    "    print(\"min corr\", \"{:.6f}\".format(np.min(cor_array)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R2 0.358379\n",
      "MSE 0.100978\n",
      "mean corr, 0.441+-0.181\n",
      "max corr 0.769153\n",
      "min corr 0.027994\n"
     ]
    }
   ],
   "source": [
    "from neuralpredictors.training import eval_state, device_state\n",
    "from neuralpredictors.measures.np_functions import corr, fev\n",
    "\n",
    "def model_predictions(model, dataloader, data_key, device=\"cpu\"):\n",
    "\n",
    "    target, output = torch.empty(0), torch.empty(0)\n",
    "    for batch in dataloader:\n",
    "        images, responses = (\n",
    "            batch[:2]\n",
    "            if not isinstance(batch, dict)\n",
    "#             else (batch[\"inputs\"], batch[\"targets\"])\n",
    "            else (batch[\"images\"], batch[\"responses\"])\n",
    "        )\n",
    "        batch_kwargs = batch._asdict() if not isinstance(batch, dict) else batch\n",
    "\n",
    "        with torch.no_grad():\n",
    "            with device_state(model, device):\n",
    "                output = torch.cat(\n",
    "                    (\n",
    "                        output,\n",
    "                        (\n",
    "                            model(images.to(device), data_key=data_key, **batch_kwargs)\n",
    "                            .detach()\n",
    "                            .cpu()\n",
    "                        ),\n",
    "                    ),\n",
    "                    dim=0,\n",
    "                )\n",
    "            target = torch.cat((target, responses.detach().cpu()), dim=0)\n",
    "\n",
    "    return target.numpy(), output.numpy()\n",
    "\n",
    "\n",
    "def get_correlations(\n",
    "    model, dataloaders, tier=None, device=\"cpu\", as_dict=False, per_neuron=True, **kwargs\n",
    "):\n",
    "\n",
    "    correlations = {}\n",
    "    dl = dataloaders[tier] if tier is not None else dataloaders\n",
    "    target_, output_ = [], []\n",
    "    \n",
    "    for k, v in dl.items():\n",
    "        target, output = model_predictions(\n",
    "            dataloader=v, model=model, data_key=k, device=device\n",
    "        )\n",
    "\n",
    "        correlations[k] = corr(target, output, axis=0)\n",
    "\n",
    "        if np.any(np.isnan(correlations[k])):\n",
    "            warnings.warn(\n",
    "                \"{}% NaNs , NaNs will be set to Zero.\".format(\n",
    "                    np.isnan(correlations[k]).mean() * 100\n",
    "                )\n",
    "            )\n",
    "        correlations[k][np.isnan(correlations[k])] = 0\n",
    "\n",
    "    if not as_dict:\n",
    "        correlations = (\n",
    "            np.hstack([v for v in correlations.values()])\n",
    "            if per_neuron\n",
    "            else np.mean(np.hstack([v for v in correlations.values()]))\n",
    "        )\n",
    "    return correlations,target,output\n",
    "\n",
    "from sklearn.metrics import r2_score, mean_squared_error\n",
    "# for i in ['train', 'validation', 'test']:\n",
    "for i in ['test']:\n",
    "    model.eval()\n",
    "    correlations,label,pred = get_correlations(model, dataloaders, tier=i, device=\"cuda:0\", as_dict=False, per_neuron=False)\n",
    "    cor_array = cor_in_time(pred, label)\n",
    "#     print(\"best val model on test dataset, {:.3f}+-{:.3f}\".format(np.mean(cor_array), np.std(cor_array)))\n",
    "    pred = smoothing_with_np_conv(pred)\n",
    "    label = smoothing_with_np_conv(label)\n",
    "    print(\"R2\", \"{:.6f}\".format(r2_score(label.T, pred.T)))\n",
    "    print(\"MSE\", \"{:.6f}\".format(mean_squared_error(label, pred)))\n",
    "    cor_array = cor_in_time(pred, label)\n",
    "    print(\"mean corr, {:.3f}+-{:.3f}\".format(np.mean(cor_array), np.std(cor_array)))\n",
    "    print(\"max corr\", \"{:.6f}\".format(np.max(cor_array)))\n",
    "    print(\"min corr\", \"{:.6f}\".format(np.min(cor_array)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}