roi_image_encoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0678cb22-26ed-4b31-8f9c-01df86152ef5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/hdd/yuchen\n"
     ]
    }
   ],
   "source": [
    "%cd /hdd/yuchen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e6efe314-dd30-40f3-a29e-7f72efaa9c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import glob\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset, BatchSampler\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tqdm.notebook import tqdm\n",
    "import random\n",
    "import nltk, re, string, collections\n",
    "import scipy \n",
    "import IProgress \n",
    "from ipywidgets import FloatProgress\n",
    "from nilearn import connectome\n",
    "from sklearn.preprocessing import Normalizer, OrdinalEncoder, OneHotEncoder, StandardScaler\n",
    "from torch.nn import Linear, ReLU, Dropout\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv, TopKPooling, BatchNorm, global_mean_pool\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.loader import DataLoader\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7784657d-f060-432a-af91-671b7d40eae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42\n",
    "random.seed(seed)\n",
    "os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3c7c0ca5-ed95-47b9-9e61-18f8916c7749",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9df388782104dc9ae321f0e8b4e2e7a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "featname = 'init_latent'\n",
    "use_stim = 'each'\n",
    "\n",
    "subject= 'subj01'\n",
    "topdir = '/hdd/yuchen/cs292/nsdfeat'\n",
    "\n",
    "savedir = f'{topdir}/subjfeat/'\n",
    "featdir = f'{topdir}/15dim/{featname}/'\n",
    "\n",
    "nsd_expdesign = scipy.io.loadmat('/hdd/yuchen/nsd/nsddata/experiments/nsd/nsd_expdesign.mat')\n",
    "\n",
    "# Note that most of them are 1-base index!\n",
    "sharedix = nsd_expdesign['sharedix'] -1 \n",
    "stims = np.load(f'/hdd/yuchen/mrifeat/{subject}/{subject}_stims.npy')\n",
    "\n",
    "feats = []\n",
    "tr_idx = np.zeros(len(stims))\n",
    "\n",
    "ct = 0\n",
    "\n",
    "for idx, s in tqdm(enumerate(stims)): \n",
    "    if s in sharedix:\n",
    "        tr_idx[idx] = 0\n",
    "    else:\n",
    "        tr_idx[idx] = 1    \n",
    "    if os.path.exists(f'{featdir}/{s:06}.npy'):\n",
    "        feat = np.load(f'{featdir}/{s:06}.npy')\n",
    "        feats.append(feat)\n",
    "    else: \n",
    "        ct += 1\n",
    "\n",
    "feats = np.stack(feats)    \n",
    "\n",
    "feats_tr = feats[tr_idx==1]\n",
    "feats_te = feats[tr_idx==0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c6307b28-f18f-498c-8d3b-ca442f5d6d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "savedir = '/hdd/yuchen/mrifeat/subj01/'\n",
    "roi = 'V1'\n",
    "\n",
    "betas_tr = np.load(f'{savedir}/{subject}_{roi}_betas_tr.npy')\n",
    "betas_te = np.load(f'{savedir}/{subject}_{roi}_betas_te.npy')\n",
    "betas_tr_avg = np.load(f'{savedir}/{subject}_{roi}_betas_ave_tr.npy')\n",
    "betas_te_avg = np.load(f'{savedir}/{subject}_{roi}_betas_ave_te.npy')\n",
    "\n",
    "num_trials_tr, num_neurons = betas_tr.shape\n",
    "num_trials_te_avg, _ = betas_te_avg.shape\n",
    "roi_neuron = {'V1': num_neurons, \n",
    "              'V2':0, 'V3':0, 'V4':0}\n",
    "\n",
    "for roi in ['V2', 'V3', 'V4']:\n",
    "    original_array = np.load(f'{savedir}/{subject}_{roi}_betas_tr.npy')\n",
    "    roi_neuron[roi] = original_array.shape[1]\n",
    "    betas_tr = np.hstack([betas_tr,original_array])\n",
    "    original_array = np.load(f'{savedir}/{subject}_{roi}_betas_ave_te.npy')\n",
    "    betas_te_avg = np.hstack([betas_te_avg,original_array])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8a77f2f3-0935-46b6-892f-56e6fc69d1cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(982, 10631)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "betas_te_avg.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "85e190db-651f-49c5-a1bb-c8fe78786040",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'V1': 4308, 'V2': 2916, 'V3': 2096, 'V4': 1311}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "roi_neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fb74d48f-1e1c-490d-a6cf-9970435e5319",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9883c16927004758aa2bed9f3c39f618",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "stims = np.load(f'/hdd/yuchen/mrifeat/{subject}/{subject}_stims_ave.npy')\n",
    "\n",
    "feats_avg = []\n",
    "tr_idx = np.zeros(len(stims))\n",
    "\n",
    "ct = 0\n",
    "\n",
    "for idx, s in tqdm(enumerate(stims)): \n",
    "    if s in sharedix:\n",
    "        tr_idx[idx] = 0\n",
    "    else:\n",
    "        tr_idx[idx] = 1    \n",
    "    if os.path.exists(f'{featdir}/{s:06}.npy'):\n",
    "        # feat = np.load(f'{featdir}/{s:06}.npy')[0]\n",
    "        feat = np.load(f'{featdir}/{s:06}.npy')\n",
    "        \n",
    "        feats_avg.append(feat)\n",
    "    else: \n",
    "        # print(f'{s:06}.npy does not exist')\n",
    "        # feats_avg.append('empty')\n",
    "        feats_avg.append(np.array([0]*768*15))\n",
    "        ct += 1\n",
    "\n",
    "feats_avg = np.stack(feats_avg)    \n",
    "\n",
    "feats_tr_avg = feats_avg[tr_idx==1]\n",
    "feats_te_avg = feats_avg[tr_idx==0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d5879e77-75d1-484e-907c-e9c9202de0ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feats_tr (24980, 6400)\n",
      "feats_te (2770, 6400)\n",
      "betas_tr (24980, 10631)\n",
      "betas_te (2770, 4308)\n",
      "feats_tr_avg (8859, 6400)\n",
      "feats_te_avg (982, 6400)\n",
      "betas_tr_avg (8859, 4308)\n",
      "betas_te_avg (982, 10631)\n"
     ]
    }
   ],
   "source": [
    "print('feats_tr', feats_tr.shape)\n",
    "print('feats_te', feats_te.shape)\n",
    "print('betas_tr', betas_tr.shape)\n",
    "print('betas_te', betas_te.shape)\n",
    "\n",
    "print('feats_tr_avg', feats_tr_avg.shape)\n",
    "print('feats_te_avg', feats_te_avg.shape)\n",
    "print('betas_tr_avg', betas_tr_avg.shape)\n",
    "print('betas_te_avg', betas_te_avg.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4f2033d7-b18d-4f15-af59-86fb5ff57919",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'V1': 4308, 'V2': 2916, 'V3': 2096, 'V4': 1311}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "roi_neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fafd8421-9704-4e84-80fc-a32cda7ade94",
   "metadata": {},
   "outputs": [],
   "source": [
    "betas_tr_ = (betas_tr - betas_tr.mean()) / betas_tr.std()\n",
    "betas_tr_avg_ = (betas_tr_avg - betas_tr_avg.mean()) / betas_tr_avg.std()\n",
    "betas_te_ = (betas_te - betas_te.mean()) / betas_te.std()\n",
    "betas_te_avg_ = (betas_te_avg - betas_te_avg.mean()) / betas_te_avg.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "be9eae0f-16ca-40e2-8067-0d9561959226",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "59add481-fe29-4b65-a549-64865040f28f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((24980, 10631), (982, 10631))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "betas_tr_.shape, betas_te_avg_.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e180a496-c301-4b91-b54a-7573d02ff5ca",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## goal: 0.753467549664157"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "389b5d54-aacc-4db9-9c1f-98ff7af7823d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse, os\n",
    "import numpy as np\n",
    "from himalaya.backend import set_backend\n",
    "from himalaya.ridge import RidgeCV\n",
    "from himalaya.scoring import correlation_score\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "alpha = [0.000001,0.00001,0.0001,0.001,0.01, 0.1, 1]\n",
    "\n",
    "ridge = RidgeCV(alphas=alpha)\n",
    "\n",
    "preprocess_pipeline = make_pipeline(\n",
    "        StandardScaler(with_mean=True, with_std=True),\n",
    "    )\n",
    "pipeline = make_pipeline(\n",
    "        preprocess_pipeline,\n",
    "        ridge,\n",
    ")    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7914f467-bc01-4c6b-bd24-312df879bfb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.753467549664157"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = betas_tr_\n",
    "Y = feats_tr\n",
    "X_te = betas_te_avg_\n",
    "pipeline.fit(X, Y)\n",
    "scores = pipeline.predict(X_te)\n",
    "from sklearn.metrics import mean_squared_error\n",
    "mean_squared_error(scores,feats_te_avg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da88f155-b5ed-4fa5-9e1e-b399a746db76",
   "metadata": {},
   "source": [
    "## gcn best: 0.55396"
   ]
  },
  {
   "cell_type": "raw",
   "id": "89005bec-c620-4078-a34b-7a55dc0c1e04",
   "metadata": {},
   "source": [
    "ideas inspired by:\n",
    "https://www.mdpi.com/2076-3425/12/10/1394\n",
    "- Decoding Visual fMRI Stimuli from Human Brain Based on Graph Convolutional Neural Network\n",
    "https://www.bendsawyer.com/wp-content/uploads/2023/07/Saeidi-et-al-Decoding-Task-Based-fMRI-Data-with-Graph-Neural-Networks-Considering-Individual-Differences.pdf\n",
    "- Decoding Task-Based fMRI Data with Graph Neural Networks, Considering Individual Differences\n",
    "\n",
    "code adapted from:\n",
    "https://github.com/krzysztoffiok/gnn-classification-pipeline/blob/master/gnn_fw/models.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981fa705-082b-4bbe-9751-a5d25b12ba7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "betas_tr_temp = betas_tr_[:, :roi_neuron['V1']]\n",
    "# for roi in [2,3,4]:\n",
    "for roi in [2,3]:\n",
    "    temp = betas_tr_[:, sum(list(roi_neuron.values())[:roi-1]):sum(list(roi_neuron.values())[:roi])]\n",
    "    # if roi == 'V3': temp = betas_tr_[:, roi_neuron[pre_roi]:]\n",
    "    if roi == 3: temp = betas_tr_[:, sum(list(roi_neuron.values())[:roi-1]):]\n",
    "    zeros = np.zeros((betas_tr_.shape[0], roi_neuron['V1']))\n",
    "    zeros[:, :temp.shape[-1]] = temp\n",
    "    betas_tr_temp = np.hstack([betas_tr_temp, zeros])\n",
    "\n",
    "betas_te_avg_temp = betas_te_avg_[:, :roi_neuron['V1']]\n",
    "# for roi in [2,3,4]:\n",
    "for roi in [2,3]:\n",
    "    # pre_roi = 'V' + str(int(roi[-1])-1)\n",
    "    # temp = betas_te_avg_[:, roi_neuron[pre_roi]:roi_neuron[pre_roi]+roi_neuron[roi]]\n",
    "    # if roi == 'V3': temp = betas_tr_[:, roi_neuron[pre_roi]:]\n",
    "    temp = betas_te_avg_[:, sum(list(roi_neuron.values())[:roi-1]):sum(list(roi_neuron.values())[:roi])]\n",
    "    if roi == 3: temp = betas_te_avg_[:, sum(list(roi_neuron.values())[:roi-1]):]\n",
    "        \n",
    "    zeros = np.zeros((betas_te_avg_.shape[0], roi_neuron['V1']))\n",
    "    zeros[:, :temp.shape[-1]] = temp\n",
    "    betas_te_avg_temp = np.hstack([betas_te_avg_temp, zeros])\n",
    "\n",
    "betas_tr_ = betas_tr_temp\n",
    "betas_te_avg_ = betas_te_avg_temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7785af68-f500-4e10-8b37-0643b35dd7d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_roi = 3\n",
    "# num_roi = 4\n",
    "\n",
    "num_neurons = betas_tr_.shape[-1]\n",
    "num_pad = (num_neurons // num_roi + 1) * num_roi - num_neurons\n",
    "\n",
    "betas_tr = np.pad(betas_tr_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_tr_avg = np.pad(betas_tr_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te = np.pad(betas_te_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te_avg = np.pad(betas_te_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "73e6d383-91be-41ef-bd97-0ff64bb0bc83",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/nilearn/connectome/connectivity_matrices.py:509: FutureWarning: The default strategy for standardize is currently 'zscore' which incorrectly uses population std to calculate sample zscores. The new strategy 'zscore_sample' corrects this behavior by using the sample std. In release 0.13, the default strategy will be replaced by the new strategy and the 'zscore' option will be removed. Please use 'zscore_sample' instead.\n",
      "  covariances_std = [\n"
     ]
    }
   ],
   "source": [
    "num_neurons_each_roi = int(betas_te_avg.shape[-1]/num_roi)\n",
    "\n",
    "conn_measure = connectome.ConnectivityMeasure(kind='correlation')\n",
    "\n",
    "function_con = betas_tr.reshape(-1, num_roi, num_neurons_each_roi)\n",
    "\n",
    "connectivity = conn_measure.fit_transform([np.mean(function_con, axis=-1)])\n",
    "for i in range(len(connectivity[0])):\n",
    "    connectivity[0][i, i] = 0  # setting the i-th row to zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "de5db6dc-10a8-4b04-b2f1-1fd87895f868",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 3, 3)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "connectivity.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "40ee759e-4e0e-4dd2-903a-0ef1e6ffe576",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 6)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adjmatrix = np.where(abs(connectivity[0]) > 0.5)\n",
    "edge_index = np.array(adjmatrix)  # edge index matrix of shape [2, num_edges]\n",
    "\n",
    "edge_index.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "33455771-3966-47d2-bd97-bf35f967253b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6,)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_attr = []\n",
    "for idx in range(edge_index.shape[1]):\n",
    "    edge_attr.append(connectivity[0][edge_index[0,idx], edge_index[1,idx]])\n",
    "edge_attr = np.array(edge_attr)\n",
    "edge_attr.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "25b4e514-17c6-4b5a-a4f8-9efd2ad5506f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 1, 1, 2, 2],\n",
       "       [1, 2, 0, 2, 0, 1]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "da977361-3ae4-48a8-b19f-f0d554014529",
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_index_temp = []\n",
    "edge_attr_temp = []\n",
    "for i in range(edge_attr.shape[0]):\n",
    "    edge = edge_index.T[i]\n",
    "    if edge[0] < edge[1]:\n",
    "        edge_index_temp.append(edge)\n",
    "        edge_attr_temp.append(edge_attr[i])\n",
    "edge_attr = np.array(edge_attr_temp)\n",
    "edge_index = np.array(edge_index_temp).T\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ea2ce094-8b66-4597-90f3-f9d4ad3bcd5d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 1],\n",
       "       [1, 2, 2]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "edge_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "de0db021-bc55-46cb-9f18-c0515471db6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trial_tr, num_trials_te = feats_tr.shape[0], feats_te_avg.shape[0]\n",
    "\n",
    "train_dataset = []\n",
    "\n",
    "for i in range(num_trial_tr):\n",
    "    train_dataset.append(Data(x=torch.tensor(betas_tr.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_tr[i], dtype=torch.float)\n",
    "            ))\n",
    "    \n",
    "test_dataset = []\n",
    "for i in range(num_trials_te):\n",
    "    test_dataset.append(Data(x=torch.tensor(betas_te_avg.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_te_avg[i], dtype=torch.float)\n",
    "            ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "15bd82d5-b3c7-4f5f-900f-92aef341a1a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(24980, 982)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_dataset), len(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "20b580f4-8cac-42d4-9d7b-f7814758fddf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Data(x=[3, 4309], edge_index=[2, 3], edge_attr=[3], y=[6400]),\n",
       " Data(x=[3, 4309], edge_index=[2, 3], edge_attr=[3], y=[6400]))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[0], test_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3710c1eb-b263-4a19-94b7-e79b44b1c557",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9ae53258-f5a4-48f6-a535-6399098bcde3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv, TopKPooling, BatchNorm,ChebConv, SAGEConv,GATConv\n",
    "from torch_geometric.nn import global_mean_pool, global_max_pool,global_add_pool\n",
    "\n",
    "class GNN(torch.nn.Module):\n",
    "    def __init__(self, num_neurons_each_roi):\n",
    "        super(GNN, self).__init__()\n",
    "        \n",
    "        self.conv1 = ChebConv(num_neurons_each_roi, 512, K=4)\n",
    "\n",
    "        self.conv = ChebConv(512, 512, K=3)\n",
    "        self.bn = BatchNorm(512)\n",
    "\n",
    "        nnconv_channel = 16\n",
    "        self.nnconv = nn.Sequential(\n",
    "            nn.Conv1d(1, nnconv_channel, kernel_size=1),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv1d(nnconv_channel, nnconv_channel, kernel_size=3, padding = 1),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "        self.downsample = nn.Conv1d(nnconv_channel, 1, kernel_size=1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(512, 1024),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(1024, 6400)\n",
    "        )\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch,edge_attr = data.x, data.edge_index, data.batch, data.edge_attr\n",
    "        batch_size = batch.max() + 1\n",
    "\n",
    "        x = F.relu(self.conv1(x, edge_index, edge_attr))\n",
    "        \n",
    "        for i in range(1):\n",
    "            a = self.conv(x, edge_index, edge_attr)\n",
    "            a = self.bn(a)\n",
    "            x = x + F.leaky_relu(a)\n",
    "        \n",
    "        x = global_mean_pool(x, batch) # print(x.shape) torch.Size([64, 256])\n",
    "\n",
    "        x = x.unsqueeze(1)\n",
    "        x = self.nnconv(x)\n",
    "        for j in range(2):\n",
    "            x = x + self.net(x)\n",
    "        x = self.downsample(x)\n",
    "        x = x.flatten(start_dim=1)\n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "a7044222-7cce-44de-891b-481e01a5d892",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 0.62051 val loss 0.58267\n",
      "[epoch 1] train loss 0.59657 val loss 0.57291\n",
      "[epoch 2] train loss 0.58292 val loss 0.56874\n",
      "[epoch 3] train loss 0.57191 val loss 0.57412\n",
      "[epoch 4] train loss 0.56174 val loss 0.5609\n",
      "[epoch 5] train loss 0.55355 val loss 0.57046\n",
      "[epoch 6] train loss 0.54591 val loss 0.55605\n",
      "[epoch 7] train loss 0.53832 val loss 0.55433\n",
      "[epoch 8] train loss 0.53103 val loss 0.55641\n",
      "[epoch 9] train loss 0.52431 val loss 0.56152\n",
      "[epoch 10] train loss 0.51754 val loss 0.56442\n",
      "[epoch 11] train loss 0.51065 val loss 0.56317\n",
      "[epoch 12] train loss 0.50398 val loss 0.56185\n",
      "[epoch 13] train loss 0.4973 val loss 0.55834\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:1'\n",
    "model = GNN(num_neurons_each_roi=num_neurons_each_roi).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        batch = batch.to(device)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "        pred = model(batch)\n",
    "        \n",
    "        loss = loss_fn(pred, feats)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "\n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            batch = batch.to(device)\n",
    "            batch_num = batch.batch.max() + 1\n",
    "            feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "            pred = model(batch)\n",
    "            loss = loss_fn(pred, feats)\n",
    "            total_loss_val += loss.item()\n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_gcn_roi.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdfaf03c-30b6-454a-80e3-299a04ad6f4b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f2b68f06-46bf-4db2-ab9c-18b433cf8b61",
   "metadata": {},
   "source": [
    "num_roi=8: 0.5654079392552376\n",
    "num_roi=16: 0.5733139477670193\n",
    "num_roi=4: 0.56504\n",
    "    conv512: 0.5668724402785301\n",
    "    conv256: 0.5616156719624996\n",
    "        linear256->6400: 0.5727455839514732\n",
    "        GCNimproved=True: 0.562708642333746\n",
    "        linear256->256->6400: 0.5671156793832779\n",
    "        linear256->1024->6400: 0.5647310167551041\n",
    "        linear256->512->1024->6400: 0.5680358745157719\n",
    "        linear256->512->512->6400: 0.5688614509999752\n",
    "        ChebConv(K=2) 0.5641; (K=3) 0.5607; (K=4) 0.5610; (K=3&4) 0.5597; (K=4&3) 0.5596\n",
    "        SAGEConv and GATConv: worse \n",
    "            nnconv.T: 0.5561986267566681 (2layers of nnconv.T is worse)\n",
    "            nnconv: 2skip: 0.5567543022334576 (0.5583899803459644)\n",
    "                replace relu with leakyrelu: 0.557660523802042\n",
    "                replace nnconv's relu with leakyrelu: 0.5604802779853344\n",
    "                replace gcn's second and all nnconv's relu with leakyrelu: 0.5566234886646271\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "3f27f073-3a5d-4b98-9395-93b119dfb157",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = 'cuda:1'\n",
    "model = GNN(num_neurons_each_roi=num_neurons_each_roi).to(device)\n",
    "model.load_state_dict(torch.load('/hdd/yuchen/modelweights_visual_gcn_roi.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "fbb9afec-12d7-4c29-9c00-106dab3c4e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "lst = []\n",
    "lst_gt = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch in test_dataloader:\n",
    "        batch = batch.to(device)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        pred = model(batch)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "        \n",
    "        lst.append(pred.cpu().detach().numpy())\n",
    "        lst_gt.append(feats.cpu().detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7be9ee0a-8e9a-447d-9283-2f3162c84c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "lst = np.vstack(lst)\n",
    "lst_gt = np.vstack(lst_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "15e3273d-fd0a-4d17-bfbb-0834707b3943",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.55396"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import mean_squared_error\n",
    "mean_squared_error(lst,lst_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "19838a3b-5d2c-4686-a5fa-a6a57bdbcecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('/hdd/yuchen/temp_lst_latent_init_roi.npy', lst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66a1391-bb64-4f3c-a9e9-37dbe0c732a6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "4ac24f94-050d-4299-a0d2-fbde9ba50a0e",
   "metadata": {},
   "source": [
    "## adapted text encoder best: 0.5579"
   ]
  },
  {
   "cell_type": "raw",
   "id": "40a49a21-91f4-463d-976f-1223b3ab13af",
   "metadata": {},
   "source": [
    "idea inspired by:\n",
    "https://arxiv.org/abs/2210.01769\n",
    "- Mind Reader: Reconstructing complex images from brain activities\n",
    "\n",
    "code adapted from\n",
    "https://github.com/sklin93/mind-reader/blob/main/modules.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "80b8067b-d8e8-42d0-88de-71dd6cef735f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, betas, feats):\n",
    "        self.betas = torch.tensor(betas, dtype=torch.float32)\n",
    "        self.feats = torch.tensor(feats, dtype=torch.float32)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.betas)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {'betas': self.betas[idx], 'feats': self.feats[idx]}\n",
    "\n",
    "train_dataset = CustomDataset(betas_tr_, feats_tr)\n",
    "test_dataset = CustomDataset(betas_te_avg_, feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "68dd501c-457b-417f-969d-09fe479dcf41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([10631]), torch.Size([6400]))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset[0]['betas'].shape, train_dataset[0]['feats'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1006689e-817d-4dcf-9b1f-977fbf3ca882",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256  \n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "for batch in train_dataloader:\n",
    "    break\n",
    "num_neurons = batch['betas'].shape[-1]\n",
    "num_features = batch['feats'].shape[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "1ca98026-8aee-4cbe-b9f0-f3f8b1035bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self, num_neurons, num_features):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        \n",
    "        self.num_neurons = num_neurons\n",
    "        channel = 64\n",
    "        self.conv1d = nn.Conv1d(1, channel, kernel_size=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv1d(channel, channel, kernel_size=11, padding = 5),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv1d(channel, channel, kernel_size=11, padding = 5),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.downsample = nn.Conv1d(channel, 1, kernel_size=1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(num_neurons, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 6400)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1d(x)\n",
    "        for i in range(2):\n",
    "            x = x + self.net(x)\n",
    "        x = self.downsample(x)\n",
    "        x = x.reshape(-1, self.num_neurons)\n",
    "        x = self.fc(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3db402b6-4067-4701-bd80-a4d69b1639f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 0.63639 val loss 0.59706\n",
      "[epoch 1] train loss 0.60864 val loss 0.57652\n",
      "[epoch 2] train loss 0.59012 val loss 0.56902\n",
      "[epoch 3] train loss 0.57534 val loss 0.5643\n",
      "[epoch 4] train loss 0.56173 val loss 0.56871\n",
      "[epoch 5] train loss 0.55002 val loss 0.55794\n",
      "[epoch 6] train loss 0.54059 val loss 0.56048\n",
      "[epoch 7] train loss 0.53047 val loss 0.56575\n",
      "[epoch 8] train loss 0.52317 val loss 0.56198\n",
      "[epoch 9] train loss 0.5157 val loss 0.56166\n",
      "[epoch 10] train loss 0.50986 val loss 0.56431\n",
      "[epoch 11] train loss 0.50329 val loss 0.56392\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:1'\n",
    "model = NeuralNetwork(num_neurons = num_neurons, num_features = num_features).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        betas = batch['betas'].to(device).unsqueeze(1)\n",
    "        feats = batch['feats'].to(device)\n",
    "        \n",
    "        pred = model(betas)\n",
    "        loss = loss_fn(pred, feats)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "\n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            betas = batch['betas'].to(device).unsqueeze(1)\n",
    "            feats = batch['feats'].to(device)\n",
    "            pred = model(betas)\n",
    "            loss = loss_fn(pred, feats)\n",
    "            total_loss_val += loss.item()\n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_textnn.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "raw",
   "id": "ff6d852a-a8ad-4bad-9855-c3be8ea5f0cd",
   "metadata": {},
   "source": [
    "lr0.001: 0.56826; lr0.0005: 0.56561\n",
    "    2 resblock: 0.56314; 1resblock: 0.56571\n",
    "        conv_channnel32: 0.56449; conv_channnel128: 0.56253 (very slow)\n",
    "        remove last conv in each resblock: 0.5635\n",
    "        kernal=5: 0.56329; kernal=11: 0.56288; kernal=13: 0.5633823126554489\n",
    "            num_neurons->512&1024: 0.5636\n",
    "            num_neurons->256: 0.5682\n",
    "            num_neurons->1024: 0.5634\n",
    "            num_neurons->512&512: 0.56112\n",
    "            num_neurons->512&256: bad\n",
    "                conv without padding: kernal=11: 0.56968; kernal=5: 0.5665; kernal=3: 0.56189\n",
    "                gelu in resblock: 0.5616\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "3f962a50-af0d-4541-88d0-4ce9015ceae1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5579410791397095"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_loss_val"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef17732b-2938-453c-beb0-aa959a2be6a6",
   "metadata": {},
   "source": [
    "## VAE best: 0.5665474534034729"
   ]
  },
  {
   "cell_type": "raw",
   "id": "dd18e84a-129e-4a46-88fa-5f701f027cd5",
   "metadata": {},
   "source": [
    "Idea inspired by: https://arxiv.org/pdf/2302.12971.pdf\n",
    "\"BrainCLIP: Bridging Brain and Visual-Linguistic Representation Via CLIP for Generic Natural Visual Stimulus Decoding\"\n",
    "\n",
    "encoder of the VAE reconstructed from appendix figure 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b6917528-52f5-4737-bbd0-59b19e9ad93a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, betas, feats):\n",
    "        self.betas = torch.tensor(betas, dtype=torch.float32)\n",
    "        self.feats = torch.tensor(feats, dtype=torch.float32)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.betas)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {'betas': self.betas[idx], 'feats': self.feats[idx]}\n",
    "\n",
    "train_dataset = CustomDataset(betas_tr_, feats_tr)\n",
    "test_dataset = CustomDataset(betas_te_avg_, feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "863db648-1fab-4513-975c-78b4b4ee688e",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256  \n",
    "train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "for batch in train_dataloader:\n",
    "    break\n",
    "num_neurons = batch['betas'].shape[-1]\n",
    "num_features = batch['feats'].shape[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6b0f84d3-bc11-4611-a621-9ed9b5c95fc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "class FBL(nn.Module):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super(FBL, self).__init__()\n",
    "        self.fc = nn.Linear(in_features, out_features)\n",
    "        self.bn = nn.BatchNorm1d(out_features)\n",
    "        self.act = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc(x)\n",
    "        x = self.bn(x)\n",
    "        return self.act(x)\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dims, z_dim):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.ln1 = nn.Linear(input_dim, 512)\n",
    "        self.relu = nn.LeakyReLU()\n",
    "        self.ln2 = nn.Linear(512, 256)\n",
    "        self.fc_mu = nn.Linear(256, z_dim)\n",
    "        self.fc_log_var = nn.Linear(256, z_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.ln1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.ln2(x)\n",
    "        mu = self.fc_mu(x)\n",
    "        log_var = self.fc_log_var(x)\n",
    "        return mu, log_var\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dims, output_dim):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.fbl1 = FBL(z_dim, hidden_dims)\n",
    "        self.fbl2 = FBL(hidden_dims, hidden_dims)\n",
    "        self.fbl3 = FBL(hidden_dims, hidden_dims)\n",
    "        self.final_fc = nn.Linear(hidden_dims, output_dim)\n",
    "\n",
    "    def forward(self, z):\n",
    "        z = self.fbl1(z)\n",
    "        for i in range(1):\n",
    "            z = self.fbl2(z)\n",
    "        z = self.fbl3(z)\n",
    "        return self.final_fc(z)\n",
    "\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, fMRI_input_dim, hidden_dims, z_dim, output_dim):\n",
    "        super(VAE, self).__init__()\n",
    "        self.fMRI_encoder = Encoder(fMRI_input_dim, hidden_dims, z_dim)\n",
    "        self.decoder = Decoder(z_dim, hidden_dims, output_dim)\n",
    "\n",
    "    def reparameterize(self, mu, log_var):\n",
    "        std = torch.exp(0.5*log_var)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps*std\n",
    "\n",
    "    def forward(self, fMRI):\n",
    "        fMRI_mu, fMRI_log_var = self.fMRI_encoder(fMRI)\n",
    "        z_fMRI = self.reparameterize(fMRI_mu, fMRI_log_var)\n",
    "        fMRI_embedding = self.decoder(z_fMRI)\n",
    "\n",
    "        return fMRI_embedding, fMRI_mu, fMRI_log_var, z_fMRI\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f81de1d7-370a-4d48-be07-e842cff6fae6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 1.27908 val loss 0.62805\n",
      "[epoch 1] train loss 1.02074 val loss 0.61877\n",
      "[epoch 2] train loss 0.98401 val loss 0.60876\n",
      "[epoch 3] train loss 0.95482 val loss 0.60088\n",
      "[epoch 4] train loss 0.93094 val loss 0.59462\n",
      "[epoch 5] train loss 0.91274 val loss 0.58854\n",
      "[epoch 6] train loss 0.89968 val loss 0.58721\n",
      "[epoch 7] train loss 0.88945 val loss 0.58078\n",
      "[epoch 8] train loss 0.87884 val loss 0.57848\n",
      "[epoch 9] train loss 0.87066 val loss 0.57868\n",
      "[epoch 10] train loss 0.8622 val loss 0.5821\n",
      "[epoch 11] train loss 0.85776 val loss 0.57713\n",
      "[epoch 12] train loss 0.84802 val loss 0.57883\n",
      "[epoch 13] train loss 0.84132 val loss 0.57566\n",
      "[epoch 14] train loss 0.83617 val loss 0.57621\n",
      "[epoch 15] train loss 0.83102 val loss 0.57499\n",
      "[epoch 16] train loss 0.82325 val loss 0.57161\n",
      "[epoch 17] train loss 0.81833 val loss 0.58257\n",
      "[epoch 18] train loss 0.81233 val loss 0.57486\n",
      "[epoch 19] train loss 0.80495 val loss 0.57503\n",
      "[epoch 20] train loss 0.80231 val loss 0.57576\n",
      "[epoch 21] train loss 0.79543 val loss 0.57136\n",
      "[epoch 22] train loss 0.7904 val loss 0.57142\n",
      "[epoch 23] train loss 0.78489 val loss 0.56956\n",
      "[epoch 24] train loss 0.77968 val loss 0.5709\n",
      "[epoch 25] train loss 0.77652 val loss 0.56779\n",
      "[epoch 26] train loss 0.77133 val loss 0.56658\n",
      "[epoch 27] train loss 0.76637 val loss 0.56784\n",
      "[epoch 28] train loss 0.76448 val loss 0.57666\n",
      "[epoch 29] train loss 0.75796 val loss 0.57028\n",
      "[epoch 30] train loss 0.75363 val loss 0.57426\n",
      "[epoch 31] train loss 0.74956 val loss 0.56655\n",
      "[epoch 32] train loss 0.74679 val loss 0.56935\n",
      "[epoch 33] train loss 0.74073 val loss 0.57147\n",
      "[epoch 34] train loss 0.7368 val loss 0.56981\n",
      "[epoch 35] train loss 0.73365 val loss 0.56794\n",
      "[epoch 36] train loss 0.73142 val loss 0.56868\n",
      "[epoch 37] train loss 0.72793 val loss 0.56868\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:1'\n",
    "fMRI_input_dim = num_neurons  \n",
    "hidden_dims = 4096 \n",
    "z_dim = num_features  \n",
    "output_dim = num_neurons  \n",
    "\n",
    "model = VAE(fMRI_input_dim, hidden_dims, z_dim, output_dim).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0005)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        betas = batch['betas'].to(device)\n",
    "        feats = batch['feats'].to(device)\n",
    "\n",
    "        fMRI_embedding, fMRI_mu, fMRI_log_var, z_fMRI = model(betas)\n",
    "        \n",
    "        recon_loss_fMRI = loss_fn(fMRI_embedding, betas)\n",
    "        # kld_fMRI = -0.5 * torch.mean(1 + fMRI_log_var - fMRI_mu.pow(2) - fMRI_log_var.exp())\n",
    "        kld_fMRI = loss_fn(feats, z_fMRI)\n",
    "        loss = recon_loss_fMRI + kld_fMRI\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "        \n",
    "        \n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            betas = batch['betas'].to(device)\n",
    "            feats = batch['feats'].to(device)\n",
    "\n",
    "            fMRI_embedding, fMRI_mu, fMRI_log_var, z_fMRI = model(betas)\n",
    "        \n",
    "            loss = loss_fn(feats, z_fMRI)\n",
    "            total_loss_val += loss.item()\n",
    "        \n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_vae.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2f3352fd-64b4-478c-85ce-972927d6940b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5665474534034729"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_loss_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "564f0737-c3aa-4435-94e2-feb48d675741",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "2db13172-eeb2-47d5-8237-155f98c9342c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## use ventral stream data: regression: 0.697; gcn: 0.611"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d969aba0-cc4d-420a-a53d-d8bdefe7662e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6971274161290577"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import argparse, os\n",
    "import numpy as np\n",
    "from himalaya.backend import set_backend\n",
    "from himalaya.ridge import RidgeCV\n",
    "from himalaya.scoring import correlation_score\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "alpha = [0.000001,0.00001,0.0001,0.001,0.01, 0.1, 1]\n",
    "\n",
    "ridge = RidgeCV(alphas=alpha)\n",
    "\n",
    "preprocess_pipeline = make_pipeline(\n",
    "        StandardScaler(with_mean=True, with_std=True),\n",
    "    )\n",
    "pipeline = make_pipeline(\n",
    "        preprocess_pipeline,\n",
    "        ridge,\n",
    ")    \n",
    "\n",
    "X = betas_tr\n",
    "Y = feats_tr\n",
    "X_te = betas_te_avg\n",
    "pipeline.fit(X, Y)\n",
    "scores = pipeline.predict(X_te)\n",
    "from sklearn.metrics import mean_squared_error\n",
    "mean_squared_error(scores,feats_te_avg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e9e610-3c35-40d7-8fa7-4ce1f13f3ee6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "18ef401d-08f5-46c9-bde6-f35596f78e1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/nilearn/connectome/connectivity_matrices.py:509: FutureWarning: The default strategy for standardize is currently 'zscore' which incorrectly uses population std to calculate sample zscores. The new strategy 'zscore_sample' corrects this behavior by using the sample std. In release 0.13, the default strategy will be replaced by the new strategy and the 'zscore' option will be removed. Please use 'zscore_sample' instead.\n",
      "  covariances_std = [\n"
     ]
    }
   ],
   "source": [
    "num_roi = 4\n",
    "num_neurons = betas_tr_.shape[-1]\n",
    "num_pad = (num_neurons // num_roi + 1) * num_roi - num_neurons\n",
    "\n",
    "betas_tr = np.pad(betas_tr_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_tr_avg = np.pad(betas_tr_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te = np.pad(betas_te_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "betas_te_avg = np.pad(betas_te_avg_, ((0, 0), (0, num_pad)), 'constant', constant_values=0)\n",
    "\n",
    "num_neurons_each_roi = int(betas_te_avg.shape[-1]/num_roi)\n",
    "conn_measure = connectome.ConnectivityMeasure(kind='correlation')\n",
    "function_con = betas_tr.reshape(-1, num_roi, num_neurons_each_roi)\n",
    "connectivity = conn_measure.fit_transform([np.mean(function_con, axis=-1)])\n",
    "for i in range(len(connectivity[0])):\n",
    "    connectivity[0][i, i] = 0  # Setting the i-th row to zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0a9b3a32-1b49-4b9d-93ca-a3bb7fd8bd0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 4)\n"
     ]
    }
   ],
   "source": [
    "adjmatrix = np.where(abs(connectivity[0]) > 0.8)\n",
    "edge_index = np.array(adjmatrix)  # Edge index matrix of shape [2, num_edges]\n",
    "print(edge_index.shape)\n",
    "\n",
    "edge_attr = []\n",
    "for idx in range(edge_index.shape[1]):\n",
    "    edge_attr.append(connectivity[0][edge_index[0,idx], edge_index[1,idx]])\n",
    "edge_attr = np.array(edge_attr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d65ddb67-a48d-4659-bb15-d1e72657c199",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trial_tr, num_trials_te = feats_tr.shape[0], feats_te_avg.shape[0]\n",
    "\n",
    "train_dataset = []\n",
    "for i in range(num_trial_tr):\n",
    "    train_dataset.append(Data(x=torch.tensor(betas_tr.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_tr[i], dtype=torch.float)\n",
    "            ))\n",
    "    \n",
    "test_dataset = []\n",
    "for i in range(num_trials_te):\n",
    "    test_dataset.append(Data(x=torch.tensor(betas_te_avg.reshape(-1, num_roi, num_neurons_each_roi)[i], dtype=torch.float), \n",
    "            edge_index=torch.tensor(edge_index, dtype=torch.long), \n",
    "            edge_attr=torch.tensor(edge_attr, dtype=torch.float),\n",
    "            y=torch.tensor(feats_te_avg[i], dtype=torch.float)\n",
    "            ))\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0dadecd5-48d4-4df6-b59f-8a393b6e8ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GCNConv, TopKPooling, BatchNorm,ChebConv, SAGEConv,GATConv\n",
    "from torch_geometric.nn import global_mean_pool, global_max_pool,global_add_pool\n",
    "\n",
    "class GCNWithPooling(torch.nn.Module):\n",
    "    def __init__(self, num_neurons_each_roi):\n",
    "        super(GCNWithPooling, self).__init__()\n",
    "        \n",
    "        self.conv1 = ChebConv(num_neurons_each_roi, 256, K=4)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=0.1)\n",
    "        self.conv = ChebConv(256, 256, K=3)\n",
    "        self.bn = BatchNorm(256)\n",
    "\n",
    "        nnconv_channel = 16\n",
    "        self.nnconv = nn.Sequential(\n",
    "            nn.Conv1d(1, nnconv_channel, kernel_size=1),\n",
    "            nn.ReLU(), \n",
    "        )\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv1d(nnconv_channel, nnconv_channel, kernel_size=3, padding = 1),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        self.downsample = nn.Conv1d(nnconv_channel, 1, kernel_size=1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(256, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 6400)\n",
    "        )\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index, batch,edge_attr = data.x, data.edge_index, data.batch, data.edge_attr\n",
    "        batch_size = batch.max() + 1\n",
    "\n",
    "        x = F.relu(self.conv1(x, edge_index, edge_attr))\n",
    "\n",
    "        for i in range(1):\n",
    "            a = self.conv(x, edge_index, edge_attr)\n",
    "            a = self.bn(a)\n",
    "            x = x + F.relu(a)\n",
    "\n",
    "        x = global_mean_pool(x, batch) # print(x.shape) torch.Size([64, 256])\n",
    "\n",
    "        x = x.unsqueeze(1)\n",
    "        x = self.nnconv(x)\n",
    "        \n",
    "        for j in range(2):\n",
    "            x = x + self.net(x)\n",
    "        x = self.downsample(x)\n",
    "        x = x.flatten(start_dim=1)\n",
    "        x = self.fc(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2edb8c5d-afaa-40d5-ac55-50c16d2ab01f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[epoch 0] train loss 0.65145 val loss 0.62354\n",
      "[epoch 1] train loss 0.64557 val loss 0.62117\n",
      "[epoch 2] train loss 0.64482 val loss 0.62083\n",
      "[epoch 3] train loss 0.64456 val loss 0.62064\n",
      "[epoch 4] train loss 0.64436 val loss 0.62031\n",
      "[epoch 5] train loss 0.64402 val loss 0.61988\n",
      "[epoch 6] train loss 0.64311 val loss 0.61861\n",
      "[epoch 7] train loss 0.64154 val loss 0.6169\n",
      "[epoch 8] train loss 0.63933 val loss 0.61434\n",
      "[epoch 9] train loss 0.63718 val loss 0.61329\n",
      "[epoch 10] train loss 0.63523 val loss 0.61178\n",
      "[epoch 11] train loss 0.63353 val loss 0.61143\n",
      "[epoch 12] train loss 0.63188 val loss 0.61121\n",
      "[epoch 13] train loss 0.63019 val loss 0.61096\n",
      "[epoch 14] train loss 0.62867 val loss 0.61115\n",
      "[epoch 15] train loss 0.62739 val loss 0.611\n",
      "[epoch 16] train loss 0.62583 val loss 0.6112\n",
      "[epoch 17] train loss 0.62411 val loss 0.61157\n",
      "[epoch 18] train loss 0.6224 val loss 0.61154\n",
      "[epoch 19] train loss 0.62063 val loss 0.61303\n",
      "end\n"
     ]
    }
   ],
   "source": [
    "device = 'cuda:1'\n",
    "model = GCNWithPooling(num_neurons_each_roi=num_neurons_each_roi).to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.00001)\n",
    "\n",
    "num_epochs = 100\n",
    "patience = 5\n",
    "ct = 0\n",
    "best_loss_val = np.inf\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss_train = 0\n",
    "    model.train()\n",
    "    for batch in train_dataloader:\n",
    "        batch = batch.to(device)\n",
    "        batch_num = batch.batch.max() + 1\n",
    "        feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "        pred = model(batch)\n",
    "        \n",
    "        loss = loss_fn(pred, feats)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss_train += loss.item()\n",
    "    print('[epoch {}] train loss {}'.format(epoch, round(total_loss_train/len(train_dataloader), 5)), end = ' ')\n",
    "\n",
    "    model.eval()\n",
    "    total_loss_val = 0\n",
    "    with torch.no_grad():\n",
    "        for batch in test_dataloader:\n",
    "            batch = batch.to(device)\n",
    "            batch_num = batch.batch.max() + 1\n",
    "            feats = batch['y'].reshape(batch_num,-1).to(device)\n",
    "            pred = model(batch)\n",
    "            loss = loss_fn(pred, feats)\n",
    "            total_loss_val += loss.item()\n",
    "        total_loss_val = total_loss_val/len(test_dataloader)\n",
    "        print('val loss {}'.format(round(total_loss_val, 5)))\n",
    "        \n",
    "        if total_loss_val < best_loss_val:\n",
    "            ct = 0\n",
    "            best_loss_val = total_loss_val\n",
    "            torch.save(model.state_dict(), '/hdd/yuchen/modelweights_visual_gcn_test.pth')\n",
    "        else:\n",
    "            ct += 1\n",
    "            if ct > patience:\n",
    "                print('end')\n",
    "                break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34532e39-6366-4904-bbda-f6f01911acd1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad5eba82-3818-4266-861c-a63600c17451",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "797ca880-1891-448e-b634-0e2457634bfe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93a571a4-3b48-4639-92b9-a181b17d208c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4563000-839e-47b4-af2f-f217ff37c8ac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7318bf-595e-4f98-87a6-75dcf44c9e1c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbce9413-34a6-484e-a18c-3873f289ec19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "25344d33-c3d8-4f1e-8f64-d3fbcdace130",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 6400])"
      ]
     },
     "execution_count": 181,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = model(batch)\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "243a89e2-b94b-4206-9144-4b90a2f765ad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1480, 64])"
      ]
     },
     "execution_count": 151,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(dataset[0]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "569d8a18-f8a6-4e42-95d1-cbe4b1950bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "id": "bba313f1-d979-4e96-a0de-2bfbf53df258",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1480, 64])"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = model(batch)\n",
    "out.reshape(16, -1, 64).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "a538617e-4f55-4fc5-9698-0464de17d274",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 140,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_loader = DataLoader(dataset, batch_size=1)\n",
    "len(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "934214e5-79cd-4af4-ba57-3df915ea4569",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.loader import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "199989bd-ca74-4650-a787-c2cc286749a3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "512ff401-1d56-47cb-bc58-4925013398bd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "06095492-27a1-4bca-ba7a-5e1ad767b259",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import GCNConv\n",
    "from torch_geometric.datasets import Planetoid\n",
    "\n",
    "# Load a dataset\n",
    "dataset = Planetoid(root='/tmp/Cora', name='Cora')\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "54158bca-c6d6-484d-973e-d02760dee2f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "56bacd45-18fe-4cf2-8bd9-84919941d8c8",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'GlobalStorage' object has no attribute 'num_classes'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[55], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnum_classes\u001b[49m\n",
      "File \u001b[0;32m/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/torch_geometric/data/data.py:482\u001b[0m, in \u001b[0;36mData.__getattr__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m    476\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m_store\u001b[39m\u001b[38;5;124m'\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__dict__\u001b[39m:\n\u001b[1;32m    477\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[1;32m    478\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThe \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdata\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m object was created by an older version of PyG. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    479\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIf this error occurred while loading an already existing \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    480\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdataset, remove the \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mprocessed/\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m directory in the dataset\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    481\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mroot folder and try again.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 482\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mgetattr\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_store\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/hdd/yuchen/anaconda3/envs/ldm/lib/python3.8/site-packages/torch_geometric/data/storage.py:87\u001b[0m, in \u001b[0;36mBaseStorage.__getattr__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m     85\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m[key]\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[0;32m---> 87\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m(\n\u001b[1;32m     88\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m object has no attribute \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mkey\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     89\u001b[0m     ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'GlobalStorage' object has no attribute 'num_classes'"
     ]
    }
   ],
   "source": [
    "data.num_classes"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}