11_v3_GraphTab.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import pickle\n",
    "import torch\n",
    "import copy\n",
    "import torch.nn          as nn\n",
    "import numpy             as np\n",
    "import pandas            as pd\n",
    "import matplotlib.pyplot as plt \n",
    "import seaborn           as sns\n",
    "\n",
    "from tqdm                   import tqdm\n",
    "from torch.utils.data       import Dataset, DataLoader\n",
    "from torch_geometric.loader import DataLoader as PyG_Dataloader\n",
    "\n",
    "from config import (\n",
    "    PATH_TO_FEATURES,\n",
    "    PATH_SUMMARY_DATASETS\n",
    ")\n",
    "\n",
    "IMGS_PATH = 'imgs'\n",
    "NOTEBOOKS_SUMMARY_FOLDER = 'v3/'\n",
    "WITHOUT_MISSING_FOLDER = '/without_missing/'\n",
    "\n",
    "torch.manual_seed(42)\n",
    "sns.set_theme(style=\"white\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Experiments on the `GraphTab` approach\n",
    "\n",
    "For this notebook we will use the updated cell-line gene-gene interaction dataset created in [`07_v3_graph_dataset.ipynb`](07_v3_graph_dataset.ipynb). The resulting dictionary lies in `~/datasets_for_model_building/summary_datasets/v3/cl_graphs_dict.pkl`. For gene selection we used a `combined_score` threshold of `700`. The bi-modal network approach we are using in this notebook is as follows:\n",
    "\n",
    "- replacing the cell-line branch by a GNN and \n",
    "- having the drug branch using tabular input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'PATH_SUMMARY_DATASETS' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb Cell 3'\u001b[0m in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000002?line=0'>1</a>\u001b[0m \u001b[39m# Reading the drug response matrix.\u001b[39;00m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000002?line=1'>2</a>\u001b[0m \u001b[39m# This dataset got developed in `15_summary_datasets.ipynb`\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000002?line=2'>3</a>\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mopen\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39m{\u001b[39;00mPATH_SUMMARY_DATASETS\u001b[39m}\u001b[39;00m\u001b[39m{\u001b[39;00mNOTEBOOKS_SUMMARY_FOLDER\u001b[39m}\u001b[39;00m\u001b[39mdrug_response_matrix__gdsc2.pkl\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39mrb\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mas\u001b[39;00m f: \n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000002?line=3'>4</a>\u001b[0m     drug_cl \u001b[39m=\u001b[39m pickle\u001b[39m.\u001b[39mload(f)\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000002?line=4'>5</a>\u001b[0m \u001b[39mprint\u001b[39m(drug_cl\u001b[39m.\u001b[39mshape)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'PATH_SUMMARY_DATASETS' is not defined"
     ]
    }
   ],
   "source": [
    "# Reading the drug response matrix.\n",
    "# This dataset got developed in `15_summary_datasets.ipynb`\n",
    "with open(f'{PATH_SUMMARY_DATASETS}{NOTEBOOKS_SUMMARY_FOLDER}drug_response_matrix__gdsc2.pkl', 'rb') as f: \n",
    "    drug_cl = pickle.load(f)\n",
    "print(drug_cl.shape)\n",
    "print(\"Number of unique cell lines:\", len(drug_cl.CELL_LINE_NAME.unique()))\n",
    "print(\"Number of unique drug ids:\", len(drug_cl.DRUG_ID.unique()))\n",
    "print(\"Number of unique drug names:\", len(drug_cl.DRUG_NAME.unique()))\n",
    "drug_cl.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of cell-lines/graphs: 983\n",
      "Data(x=[696, 4], edge_index=[2, 7794])\n"
     ]
    }
   ],
   "source": [
    "# The following are the gene-gene interaction graphs per cell-line were the gene-gene tuples\n",
    "# have a `combined_score` > 0.70*1_000=700.\n",
    "# This dataset got developed in `07_v3_graph_dataset.ipynb`\n",
    "with open(f'{PATH_SUMMARY_DATASETS}{NOTEBOOKS_SUMMARY_FOLDER}cl_graphs_dict.pkl', 'rb') as f:\n",
    "    cl_graphs_v3 = pd.read_pickle(f)\n",
    "print(f\"Number of cell-lines/graphs: {len(list(cl_graphs_v3.keys()))}\")\n",
    "print(cl_graphs_v3['22RV1'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pre-process\n",
    "\n",
    "In this sub-section we are going to\n",
    "- [x] select only the cell-line in the gene-gene graph dictionary which are in the drug-response-matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of cell-lines/graphs: 732\n",
      "Data(x=[696, 4], edge_index=[2, 7794])\n"
     ]
    }
   ],
   "source": [
    "cl_graphs = dict((cl, cl_graphs_v3[cl]) for cl in drug_cl.CELL_LINE_NAME.unique().tolist())\n",
    "print(f\"Number of cell-lines/graphs: {len(list(cl_graphs.keys()))}\")\n",
    "print(cl_graphs['22RV1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed for 0 (0.00 %) out of all 732 cell-lines.\n",
      "All cell-lines succeeded according to this issue: \n",
      "https://github.com/pyg-team/pytorch_geometric/issues/4588\n"
     ]
    }
   ],
   "source": [
    "fails = []\n",
    "cls = len(list(cl_graphs.keys()))\n",
    "for cl, G in cl_graphs.items():\n",
    "    if not (G.edge_index.max() < G.num_nodes):\n",
    "        fails.append(cl)\n",
    "print(f\"Failed for {len(fails)} ({100*len(fails)/cls:2.2f} %) out of all {cls} cell-lines.\")\n",
    "del cls\n",
    "\n",
    "if len(fails) > 0:\n",
    "    for cl, G in tqdm(cl_graphs.items()):\n",
    "        mapping = {}\n",
    "        mapped_edge_index = []\n",
    "        for (src, dst) in G.edge_index.t().tolist():\n",
    "            if src not in mapping:\n",
    "                mapping[src] = len(mapping)\n",
    "            if dst not in mapping:\n",
    "                mapping[dst] = len(mapping)\n",
    "            mapped_edge_index.append([mapping[src], mapping[dst]])\n",
    "        edge_index = torch.tensor(mapped_edge_index).t().contiguous()\n",
    "        cl_graphs[cl].edge_index = edge_index\n",
    "\n",
    "for cl, G in cl_graphs.items():\n",
    "    assert G.edge_index.max() < G.num_nodes, f'FAIL for cell-line: {cl}'\n",
    "print(\"All cell-lines succeeded according to this issue: \")\n",
    "print(\"https://github.com/pyg-team/pytorch_geometric/issues/4588\")  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(152, 257)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>DRUG_ID</th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>...</th>\n",
       "      <th>246</th>\n",
       "      <th>247</th>\n",
       "      <th>248</th>\n",
       "      <th>249</th>\n",
       "      <th>250</th>\n",
       "      <th>251</th>\n",
       "      <th>252</th>\n",
       "      <th>253</th>\n",
       "      <th>254</th>\n",
       "      <th>255</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1073</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>810</th>\n",
       "      <td>1910</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1562</th>\n",
       "      <td>1913</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2314</th>\n",
       "      <td>1634</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3044</th>\n",
       "      <td>2045</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 257 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      DRUG_ID  0  1  2  3  4  5  6  7  8  ...  246  247  248  249  250  251  \\\n",
       "1        1073  0  0  0  0  0  0  0  0  0  ...    0    0    0    0    0    0   \n",
       "810      1910  1  1  0  0  0  0  0  0  0  ...    1    0    0    0    0    1   \n",
       "1562     1913  0  1  1  0  1  0  0  0  0  ...    0    0    0    0    1    1   \n",
       "2314     1634  0  0  0  0  0  0  0  0  0  ...    1    0    0    0    0    1   \n",
       "3044     2045  0  1  0  0  0  0  0  0  0  ...    0    0    0    0    0    0   \n",
       "\n",
       "      252  253  254  255  \n",
       "1       0    0    0    1  \n",
       "810     0    0    0    1  \n",
       "1562    0    0    0    0  \n",
       "2314    0    0    0    1  \n",
       "3044    0    1    0    0  \n",
       "\n",
       "[5 rows x 257 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open(f'{PATH_SUMMARY_DATASETS}drug_smiles_fingerprints_matrix.pkl', 'rb') as f:\n",
    "    drug_name_smiles = pickle.load(f)\n",
    "# drug_name_smiles.set_index(['drug_name'], inplace=True)\n",
    "print(drug_name_smiles.shape)\n",
    "# TODO: Note that yet there are some DRUG_NAME's which have >1 DRUG_ID\n",
    "drug_name_smiles.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Dataset summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We got three datasets now: \n",
    "- Drug response matrix\n",
    "- fingerprint matrix\n",
    "- cell-line graph dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(91991, 5)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>CELL_LINE_NAME</th>\n",
       "      <th>DRUG_ID</th>\n",
       "      <th>DRUG_NAME</th>\n",
       "      <th>DATASET</th>\n",
       "      <th>LN_IC50</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>3441054</th>\n",
       "      <td>22RV1</td>\n",
       "      <td>1003</td>\n",
       "      <td>Camptothecin</td>\n",
       "      <td>GDSC2</td>\n",
       "      <td>-3.142631</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3445814</th>\n",
       "      <td>QGP-1</td>\n",
       "      <td>1003</td>\n",
       "      <td>Camptothecin</td>\n",
       "      <td>GDSC2</td>\n",
       "      <td>-0.534857</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3446843</th>\n",
       "      <td>RC-K8</td>\n",
       "      <td>1003</td>\n",
       "      <td>Camptothecin</td>\n",
       "      <td>GDSC2</td>\n",
       "      <td>-3.324229</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3447648</th>\n",
       "      <td>RCC-JW</td>\n",
       "      <td>1003</td>\n",
       "      <td>Camptothecin</td>\n",
       "      <td>GDSC2</td>\n",
       "      <td>-2.573593</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3440830</th>\n",
       "      <td>CHP-134</td>\n",
       "      <td>1003</td>\n",
       "      <td>Camptothecin</td>\n",
       "      <td>GDSC2</td>\n",
       "      <td>-3.817771</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        CELL_LINE_NAME  DRUG_ID     DRUG_NAME DATASET   LN_IC50\n",
       "3441054          22RV1     1003  Camptothecin   GDSC2 -3.142631\n",
       "3445814          QGP-1     1003  Camptothecin   GDSC2 -0.534857\n",
       "3446843          RC-K8     1003  Camptothecin   GDSC2 -3.324229\n",
       "3447648         RCC-JW     1003  Camptothecin   GDSC2 -2.573593\n",
       "3440830        CHP-134     1003  Camptothecin   GDSC2 -3.817771"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Drug response matrix holding the ln(IC50) values for each cell-line drug tuple.\n",
    "drug_response_matrix = copy.deepcopy(drug_cl)\n",
    "print(drug_response_matrix.shape)\n",
    "drug_response_matrix.sort_values(['DRUG_ID']).head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(152, 257)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>DRUG_ID</th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>...</th>\n",
       "      <th>246</th>\n",
       "      <th>247</th>\n",
       "      <th>248</th>\n",
       "      <th>249</th>\n",
       "      <th>250</th>\n",
       "      <th>251</th>\n",
       "      <th>252</th>\n",
       "      <th>253</th>\n",
       "      <th>254</th>\n",
       "      <th>255</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>21273</th>\n",
       "      <td>1003</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>94088</th>\n",
       "      <td>1004</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25052</th>\n",
       "      <td>1006</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39836</th>\n",
       "      <td>1010</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>61065</th>\n",
       "      <td>1011</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 257 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "       DRUG_ID  0  1  2  3  4  5  6  7  8  ...  246  247  248  249  250  251  \\\n",
       "21273     1003  0  0  0  0  0  0  0  0  0  ...    0    0    1    0    1    0   \n",
       "94088     1004  1  0  0  0  0  0  0  0  1  ...    0    1    0    1    0    1   \n",
       "25052     1006  0  0  0  0  0  0  0  0  0  ...    0    0    0    0    0    1   \n",
       "39836     1010  1  0  0  0  0  0  0  0  0  ...    0    0    0    1    0    0   \n",
       "61065     1011  0  1  0  0  0  0  0  0  0  ...    1    0    0    0    1    1   \n",
       "\n",
       "       252  253  254  255  \n",
       "21273    0    0    0    0  \n",
       "94088    0    0    0    0  \n",
       "25052    0    0    0    0  \n",
       "39836    0    0    0    1  \n",
       "61065    0    0    0    1  \n",
       "\n",
       "[5 rows x 257 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The drug matrix holding the fingerprints for the drugs in the drug response matrix.\n",
    "fingerprints = copy.deepcopy(drug_name_smiles)\n",
    "print(fingerprints.shape)\n",
    "fingerprints.sort_values(['DRUG_ID']).head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of drug fingerprints: 152\n",
      "Length of each SMILES fingerprint: 256\n"
     ]
    }
   ],
   "source": [
    "# Drugs as dictionary.\n",
    "fingerprints_dict = fingerprints.set_index('DRUG_ID').T.to_dict('list')\n",
    "print(\"Number of drug fingerprints:\", len(fingerprints_dict.keys()))\n",
    "print(\"Length of each SMILES fingerprint:\", len(fingerprints_dict[1003]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of cell-lines/graphs: 732\n",
      "Data(x=[696, 4], edge_index=[2, 7794])\n"
     ]
    }
   ],
   "source": [
    "# The cell-line graphs holding for each cell-line in the drug-response-matrix the corresponding \n",
    "# graph with the cell-line level gene features.\n",
    "cell_line_graphs = copy.deepcopy(cl_graphs)\n",
    "print(f\"Number of cell-lines/graphs: {len(list(cell_line_graphs.keys()))}\")\n",
    "print(cell_line_graphs['22RV1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed for 0 (0.00 %) out of all 732 cell-lines.\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "fails = []\n",
    "cls = len(list(cell_line_graphs.keys()))\n",
    "for cl, G in cell_line_graphs.items():\n",
    "    if not(G.edge_index.max() < G.num_nodes):\n",
    "        fails.append(cl)\n",
    "print(f\"Failed for {len(fails)} ({100*len(fails)/cls:2.2f} %) out of all {cls} cell-lines.\")\n",
    "del cls\n",
    "\n",
    "c = 0\n",
    "sum_per = []\n",
    "for k, v in cell_line_graphs.items():\n",
    "    if v.x.isnan().any(): \n",
    "        sum_per.append(v.x.isnan().sum())\n",
    "        c += 1\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Build PyTorch dataset\n",
    "\n",
    "In this subsection we are going to create the dataset which holds the correct drug, cell-line graph and corresponding `ln(IC50)` value for a given index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import Dataset\n",
    "\n",
    "class GraphTabDataset(Dataset): \n",
    "    def __init__(self, cl_graphs, drugs, drug_response_matrix):\n",
    "        super().__init__()\n",
    "\n",
    "        # SMILES fingerprints of the drugs and cell-line graphs.\n",
    "        self.drugs = drugs\n",
    "        self.cell_line_graphs = cl_graphs\n",
    "\n",
    "        # Lookup datasets for the response values.\n",
    "        drug_response_matrix.reset_index(drop=True, inplace=True)\n",
    "        self.cell_lines = drug_response_matrix['CELL_LINE_NAME']\n",
    "        self.drug_ids = drug_response_matrix['DRUG_ID']\n",
    "        self.drug_names = drug_response_matrix['DRUG_NAME']\n",
    "        self.ic50s = drug_response_matrix['LN_IC50']\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ic50s)\n",
    "\n",
    "    def __getitem__(self, idx: int):\n",
    "        \"\"\"\n",
    "        Returns a tuple of cell-line, drug and the corresponding ln(IC50)\n",
    "        value for a given index.\n",
    "\n",
    "        Args:\n",
    "            idx (`int`): Index to specify the row in the drug response matrix.  \n",
    "        Returns:\n",
    "            `Tuple[torch_geometric.data.data.Data, np.ndarray, np.float64]`:\n",
    "            Tuple of a cell-line graph, drug SMILES fingerprint and the \n",
    "            corresponding ln(IC50) value.\n",
    "        \"\"\"\n",
    "        return (self.cell_line_graphs[self.cell_lines.iloc[idx]], \n",
    "                self.drugs[self.drug_ids.iloc[idx]],\n",
    "                self.ic50s.iloc[idx])\n",
    "\n",
    "    def print_dataset_summary(self):\n",
    "        print(f\"GraphTabDataset Summary\")\n",
    "        print(f\"{23*'='}\")\n",
    "        print(f\"# observations : {len(self.ic50s)}\")\n",
    "        print(f\"# cell-lines   : {len(np.unique(self.cell_lines))}\")\n",
    "        print(f\"# drugs        : {len(np.unique(self.drug_names))}\")\n",
    "        print(f\"# genes        : {self.cell_line_graphs[next(iter(self.cell_line_graphs))].x.shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GraphTabDataset Summary\n",
      "=======================\n",
      "# observations : 91991\n",
      "# cell-lines   : 732\n",
      "# drugs        : 152\n",
      "# genes        : 696\n"
     ]
    }
   ],
   "source": [
    "graph_tab_dataset = GraphTabDataset(cl_graphs=cell_line_graphs,\n",
    "                                    drugs=fingerprints_dict,\n",
    "                                    drug_response_matrix=drug_response_matrix)\n",
    "graph_tab_dataset.print_dataset_summary()                                    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set hyperparameters\n",
    "\n",
    "In this subsection we are setting the hyperparameters for the model building."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self, batch_size, lr, train_ratio, val_ratio, num_epochs):\n",
    "        self.BATCH_SIZE = batch_size\n",
    "        self.LR = lr\n",
    "        self.TRAIN_RATIO = train_ratio\n",
    "        self.TEST_VAL_RATIO = 1-self.TRAIN_RATIO\n",
    "        self.VAL_RATIO = val_ratio\n",
    "        self.NUM_EPOCHS = num_epochs\n",
    "        self.RANDOM_SEED = 12345      \n",
    "\n",
    "args = Args(batch_size=1_000, \n",
    "            lr=0.0001, \n",
    "            train_ratio=0.8, \n",
    "            val_ratio=0.5, \n",
    "            num_epochs=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create pytorch geometric `DataLoader` datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full     shape: (91991, 5)\n",
      "train    shape: (73592, 5)\n",
      "test_val shape: (18399, 5)\n",
      "test     shape: (9199, 5)\n",
      "val      shape: (9200, 5)\n",
      "\n",
      "train_dataset:\n",
      "GraphTabDataset Summary\n",
      "=======================\n",
      "# observations : 73592\n",
      "# cell-lines   : 732\n",
      "# drugs        : 152\n",
      "# genes        : 696\n",
      "\n",
      "\n",
      "test_dataset:\n",
      "GraphTabDataset Summary\n",
      "=======================\n",
      "# observations : 9199\n",
      "# cell-lines   : 732\n",
      "# drugs        : 152\n",
      "# genes        : 696\n",
      "\n",
      "\n",
      "val_dataset:\n",
      "GraphTabDataset Summary\n",
      "=======================\n",
      "# observations : 9200\n",
      "# cell-lines   : 732\n",
      "# drugs        : 152\n",
      "# genes        : 696\n"
     ]
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "from torch_geometric.data import Batch\n",
    "from torch_geometric.loader import DataLoader as PyG_DataLoader\n",
    "\n",
    "\n",
    "def create_datasets(drm, cl_graphs, drug_mat, args):\n",
    "    print(f\"Full     shape: {drm.shape}\")\n",
    "    train_set, test_val_set = train_test_split(drm, \n",
    "                                               test_size=args.TEST_VAL_RATIO, \n",
    "                                               random_state=args.RANDOM_SEED,\n",
    "                                               stratify=drm['CELL_LINE_NAME'])\n",
    "    test_set, val_set = train_test_split(test_val_set,\n",
    "                                         test_size=args.VAL_RATIO,\n",
    "                                         random_state=args.RANDOM_SEED,\n",
    "                                         stratify=test_val_set['CELL_LINE_NAME'])\n",
    "    print(f\"train    shape: {train_set.shape}\")\n",
    "    print(f\"test_val shape: {test_val_set.shape}\")\n",
    "    print(f\"test     shape: {test_set.shape}\")\n",
    "    print(f\"val      shape: {val_set.shape}\")\n",
    "\n",
    "    train_dataset = GraphTabDataset(cl_graphs=cl_graphs, drugs=drug_mat, drug_response_matrix=train_set)\n",
    "    test_dataset = GraphTabDataset(cl_graphs=cl_graphs, drugs=drug_mat, drug_response_matrix=test_set)\n",
    "    val_dataset = GraphTabDataset(cl_graphs=cl_graphs, drugs=drug_mat, drug_response_matrix=val_set)\n",
    "\n",
    "    print(\"\\ntrain_dataset:\")\n",
    "    train_dataset.print_dataset_summary()\n",
    "    print(\"\\n\\ntest_dataset:\")\n",
    "    test_dataset.print_dataset_summary()\n",
    "    print(\"\\n\\nval_dataset:\")\n",
    "    val_dataset.print_dataset_summary()\n",
    "\n",
    "    # TODO: try out different `num_workers`.\n",
    "    train_loader = PyG_DataLoader(dataset=train_dataset, batch_size=args.BATCH_SIZE, shuffle=True)\n",
    "    test_loader = PyG_DataLoader(dataset=test_dataset, batch_size=args.BATCH_SIZE, shuffle=True)\n",
    "    val_loader = PyG_DataLoader(dataset=val_dataset, batch_size=args.BATCH_SIZE, shuffle=True)\n",
    "\n",
    "    return train_loader, test_loader, val_loader\n",
    "\n",
    "train_loader, test_loader, val_loader = create_datasets(drug_response_matrix, cl_graphs, fingerprints_dict, args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of batches per dataset:\n",
      "  train : 74\n",
      "  test  : 10\n",
      "  val   : 10\n"
     ]
    }
   ],
   "source": [
    "# for batch_cell_graph, batch_drugs, batch_ic50s in train_loader\n",
    "print(\"Number of batches per dataset:\")\n",
    "print(f\"  train : {len(train_loader)}\")\n",
    "print(f\"  test  : {len(test_loader)}\")\n",
    "print(f\"  val   : {len(val_loader)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 1:\n",
      "=======\n",
      "Number of cell-line graphs in the batch: 1000\n",
      "Gene-gene interaction topology per batch: DataBatch(x=[696000, 4], edge_index=[2, 7794000], batch=[696000], ptr=[1001])\n",
      "Number of ln(IC50) targets per batch: torch.Size([1000])\n",
      "Step 2:\n",
      "=======\n",
      "Number of cell-line graphs in the batch: 1000\n",
      "Gene-gene interaction topology per batch: DataBatch(x=[696000, 4], edge_index=[2, 7794000], batch=[696000], ptr=[1001])\n",
      "Number of ln(IC50) targets per batch: torch.Size([1000])\n",
      "Step 3:\n",
      "=======\n",
      "Number of cell-line graphs in the batch: 1000\n",
      "Gene-gene interaction topology per batch: DataBatch(x=[696000, 4], edge_index=[2, 7794000], batch=[696000], ptr=[1001])\n",
      "Number of ln(IC50) targets per batch: torch.Size([1000])\n",
      "... step 10\n",
      "... step 20\n",
      "... step 30\n",
      "... step 40\n",
      "... step 50\n",
      "... step 60\n",
      "... step 70\n",
      "Step 74:\n",
      "=======\n",
      "Number of cell-line graphs in the batch: 592\n",
      "Gene-gene interaction topology per batch: DataBatch(x=[412032, 4], edge_index=[2, 4614048], batch=[412032], ptr=[593])\n",
      "Number of ln(IC50) targets per batch: torch.Size([592])\n"
     ]
    }
   ],
   "source": [
    "for step, data in enumerate(train_loader):\n",
    "    if (step > 2) & (step < len(train_loader)-1):\n",
    "        if step % 10 == 0: \n",
    "            print(\"... step\", step) \n",
    "        continue\n",
    "    else:    \n",
    "        cl_graphs, drugs, targets = data\n",
    "        print(f'Step {step + 1}:')\n",
    "        print(f'=======')\n",
    "        print(\"Number of cell-line graphs in the batch:\", cl_graphs.num_graphs)\n",
    "        print(\"Gene-gene interaction topology per batch:\", cl_graphs)\n",
    "        print(\"Number of ln(IC50) targets per batch:\", targets.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model development\n",
    "\n",
    "In this subsection we are going to train and test the GraphTab model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from time import sleep\n",
    "from sklearn.metrics import r2_score, mean_absolute_error\n",
    "from scipy.stats import pearsonr\n",
    "\n",
    "class BuildModel():\n",
    "    def __init__(self, model, criterion, optimizer, num_epochs, \n",
    "        train_loader, test_loader, val_loader, device):\n",
    "        self.train_losses = []\n",
    "        self.test_losses = []\n",
    "        self.val_losses = []\n",
    "        self.train_loader = train_loader\n",
    "        self.test_loader = test_loader\n",
    "        self.val_loader = val_loader\n",
    "        self.num_epochs = num_epochs\n",
    "        self.model = model\n",
    "        self.criterion = criterion\n",
    "        self.optimizer = optimizer\n",
    "        self.device = device\n",
    "\n",
    "    def train(self, loader): \n",
    "        train_epoch_losses, val_epoch_losses = [], []\n",
    "        train_epoch_rmse, val_epoch_rmse = [], []\n",
    "        train_epoch_mae, val_epoch_mae = [], []\n",
    "        train_epoch_r2, val_epoch_r2 = [], []\n",
    "        train_epoch_pcorr, val_epoch_pcorr = [], []\n",
    "        y_true, y_pred = [], []\n",
    "        all_batch_losses = [] # TODO: this is just for monitoring\n",
    "        n_batches = len(loader)\n",
    "\n",
    "        self.model = self.model.float() # TODO: maybe remove\n",
    "        for epoch in range(self.num_epochs):\n",
    "            self.model.train()\n",
    "            batch_losses = []\n",
    "            for i, data in enumerate(tqdm(loader, desc='Iteration')):\n",
    "                sleep(0.01)\n",
    "                cell, drug, ic50s = data\n",
    "                drug = torch.stack(drug, 0).transpose(1, 0) # Note that this is only neede when geometric \n",
    "                                                            # Dataloader is used and no collate.\n",
    "                cell, drug, ic50s = cell.to(self.device), drug.to(self.device), ic50s.to(self.device)\n",
    "\n",
    "                self.optimizer.zero_grad()\n",
    "\n",
    "                # Models predictions of the ic50s for a batch of cell-lines and drugs\n",
    "                preds = self.model(cell, drug.float()).unsqueeze(1)\n",
    "                loss = self.criterion(preds, ic50s.view(-1, 1).float()) # =train_loss\n",
    "                batch_losses.append(loss)\n",
    "\n",
    "                \n",
    "                y_true.append(ic50s.view(-1, 1))\n",
    "                y_pred.append(preds)             \n",
    "\n",
    "                loss.backward()\n",
    "                self.optimizer.step()\n",
    "\n",
    "            all_batch_losses.append(batch_losses) # TODO: this is just for monitoring\n",
    "            total_epoch_loss = sum(batch_losses)\n",
    "            train_epoch_losses.append(total_epoch_loss / n_batches)\n",
    "\n",
    "            y_true = torch.cat(y_true, dim=0)\n",
    "            y_pred = torch.cat(y_pred, dim=0)\n",
    "            train_mse = train_epoch_losses[-1]\n",
    "            train_epoch_rmse.append(torch.sqrt(train_mse))\n",
    "            train_epoch_mae.append(mean_absolute_error(y_true.cpu(), y_pred.cpu()))\n",
    "            train_epoch_r2.append(r2_score(y_true.cpu(), y_pred.cpu()))\n",
    "            train_epoch_pcorr.append(pearsonr(y_true.cpu().numpy().flatten(), y_pred.cpu().numpy().flatten()))\n",
    "                     \n",
    "            mse, rmse, mae, r2, pcorr = self.validate(self.val_loader)\n",
    "            val_epoch_losses.append(mse)\n",
    "            val_epoch_rmse.append(rmse)\n",
    "            val_epoch_mae.append(mae)\n",
    "            val_epoch_r2.append(r2)\n",
    "            val_epoch_pcorr.append(pcorr)\n",
    "\n",
    "            print(\"=====Epoch \", epoch)\n",
    "            print(f\"Train      | MSE: {train_mse:2.5f}\")\n",
    "            print(f\"Validation | MSE: {mse:2.5f}\")\n",
    "\n",
    "        performance_stats = {\n",
    "            'train': {\n",
    "                'mse': train_epoch_losses,\n",
    "                'rmse': train_epoch_rmse,\n",
    "                'mae': train_epoch_mae,\n",
    "                'r2': train_epoch_r2,\n",
    "                'pcorr': train_epoch_pcorr\n",
    "            },\n",
    "            'val': {\n",
    "                'mse': val_epoch_losses,\n",
    "                'rmse': val_epoch_rmse,\n",
    "                'mae': val_epoch_mae,\n",
    "                'r2': val_epoch_r2,\n",
    "                'pcorr': val_epoch_pcorr\n",
    "            }            \n",
    "        }\n",
    "\n",
    "        return performance_stats           \n",
    "\n",
    "    def validate(self, loader):\n",
    "        self.model.eval()\n",
    "        y_true, y_pred = [], []\n",
    "        total_loss = 0\n",
    "        with torch.no_grad():\n",
    "            for data in tqdm(loader, desc='Iter', position=0, leave=True):\n",
    "                sleep(0.01)\n",
    "                cl, dr, ic50 = data\n",
    "                dr = torch.stack(dr, 0).transpose(1, 0)\n",
    "\n",
    "                preds = self.model(cl, dr.float()).unsqueeze(1)\n",
    "                ic50 = ic50.to(self.device)\n",
    "                total_loss += self.criterion(preds, ic50.view(-1,1).float())\n",
    "                # total_loss += F.mse_loss(preds, ic50.view(-1, 1).float(), reduction='sum')\n",
    "                y_true.append(ic50.view(-1, 1))\n",
    "                y_pred.append(preds)\n",
    "        \n",
    "        y_true = torch.cat(y_true, dim=0)\n",
    "        y_pred = torch.cat(y_pred, dim=0)\n",
    "        mse = total_loss / len(loader)\n",
    "        rmse = torch.sqrt(mse)\n",
    "        mae = mean_absolute_error(y_true.cpu(), y_pred.cpu())\n",
    "        r2 = r2_score(y_true.cpu(), y_pred.cpu())\n",
    "        pearson_corr_coef, _ = pearsonr(y_true.cpu().numpy().flatten(), \n",
    "                                        y_pred.cpu().numpy().flatten())\n",
    "\n",
    "        return mse, rmse, mae, r2, pearson_corr_coef"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cpu\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload\n",
    "\n",
    "from torch_geometric.nn import Sequential, GCNConv, global_mean_pool, global_max_pool\n",
    "\n",
    "\n",
    "class GraphTab_v1(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GraphTab_v1, self).__init__()\n",
    "\n",
    "        # Cell-line graph branch. Obtains node embeddings.\n",
    "        # https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.sequential.Sequential\n",
    "        self.cell_emb = Sequential('x, edge_index, batch', \n",
    "            [\n",
    "                (GCNConv(in_channels=4, out_channels=256), 'x, edge_index -> x1'), # TODO: GATConv() vs GCNConv()\n",
    "                nn.ReLU(inplace=True),\n",
    "                ## nn.BatchNorm1d(num_features=128),\n",
    "                ## nn.Dropout(self.dropout_p),\n",
    "                (GCNConv(in_channels=256, out_channels=256), 'x1, edge_index -> x2'),\n",
    "                nn.ReLU(inplace=True),\n",
    "                (global_mean_pool, 'x2, batch -> x3'), \n",
    "                # Start embedding\n",
    "                nn.Linear(256, 128),\n",
    "                nn.BatchNorm1d(128),\n",
    "                nn.ReLU(),\n",
    "                nn.Dropout(p=0.1),\n",
    "                nn.Linear(128, 128),\n",
    "                nn.ReLU()\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.drug_emb = nn.Sequential(\n",
    "            nn.Linear(256, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.1),\n",
    "            nn.Linear(128, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU()          \n",
    "        )\n",
    "\n",
    "        self.fcn = nn.Sequential(\n",
    "            nn.Linear(2*128, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ELU(),\n",
    "            nn.Dropout(p=0.1),\n",
    "            nn.Linear(128, 64),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ELU(),\n",
    "            nn.Dropout(p=0.1),\n",
    "            nn.Linear(64, 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, cell, drug):\n",
    "        drug_emb = self.drug_emb(drug)\n",
    "        cell_emb = self.cell_emb(cell.x.float(), cell.edge_index, cell.batch)\n",
    "        concat = torch.cat([cell_emb, drug_emb], -1)\n",
    "        y_pred = self.fcn(concat)\n",
    "        y_pred = y_pred.reshape(y_pred.shape[0])\n",
    "        return y_pred\n",
    "\n",
    "torch.manual_seed(args.RANDOM_SEED)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"device: {device}\")\n",
    "\n",
    "model = GraphTab_v1().to(device)\n",
    "loss_func = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(params=model.parameters(), lr=args.LR) # TODO: include weight_decay of lr\n",
    "\n",
    "build_model = BuildModel(model=model,\n",
    "                         criterion=loss_func,\n",
    "                         optimizer=optimizer,\n",
    "                         num_epochs=3,\n",
    "                         train_loader=train_loader,\n",
    "                         test_loader=test_loader,\n",
    "                         val_loader=val_loader, \n",
    "                         device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration:   0%|          | 0/74 [00:00<?, ?it/s]"
     ]
    }
   ],
   "source": [
    "performance_stats = build_model.train(build_model.train_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Only With a Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'mse': [tensor(14.3344, requires_grad=True)], 'rmse': [tensor(3.7861, requires_grad=True)], 'mae': [3.3293848], 'r2': [-1.195433379156095], 'pcorr': [(0.0004950899936630969, 0.9944485685845041)]}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch_geometric.nn import Sequential, GCNConv, global_mean_pool, global_max_pool\n",
    "\n",
    "\n",
    "class GraphTab_v1(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GraphTab_v1, self).__init__()\n",
    "\n",
    "        # Cell-line graph branch. Obtains node embeddings.\n",
    "        # https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.sequential.Sequential\n",
    "        self.cell_emb = Sequential('x, edge_index, batch', \n",
    "            [\n",
    "                (GCNConv(in_channels=4, out_channels=256), 'x, edge_index -> x1'), # TODO: GATConv() vs GCNConv()\n",
    "                nn.ReLU(inplace=True),\n",
    "                ## nn.BatchNorm1d(num_features=128),\n",
    "                ## nn.Dropout(self.dropout_p),\n",
    "                (GCNConv(in_channels=256, out_channels=256), 'x1, edge_index -> x2'),\n",
    "                nn.ReLU(inplace=True),\n",
    "                (global_mean_pool, 'x2, batch -> x3'), \n",
    "                # Start embedding\n",
    "                nn.Linear(256, 128),\n",
    "                nn.BatchNorm1d(128),\n",
    "                nn.ReLU(),\n",
    "                nn.Dropout(p=0.1),\n",
    "                nn.Linear(128, 128),\n",
    "                nn.ReLU()\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.drug_emb = nn.Sequential(\n",
    "            nn.Linear(256, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(p=0.1),\n",
    "            nn.Linear(128, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU()          \n",
    "        )\n",
    "\n",
    "        self.fcn = nn.Sequential(\n",
    "            nn.Linear(2*128, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ELU(),\n",
    "            nn.Dropout(p=0.1),\n",
    "            nn.Linear(128, 64),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ELU(),\n",
    "            nn.Dropout(p=0.1),\n",
    "            nn.Linear(64, 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, cell, drug):\n",
    "        drug_emb = self.drug_emb(drug)\n",
    "        cell_emb = self.cell_emb(cell.x.float(), cell.edge_index, cell.batch)\n",
    "        concat = torch.cat([cell_emb, drug_emb], -1)\n",
    "        y_pred = self.fcn(concat)\n",
    "        y_pred = y_pred.reshape(y_pred.shape[0])\n",
    "        return y_pred\n",
    "\n",
    "\n",
    "model = GraphTab_v1()\n",
    "optimizer = torch.optim.Adam(params=model.parameters(), lr=0.0001)\n",
    "\n",
    "checkpoint = torch.load(f'{PATH_SUMMARY_DATASETS}GraphTab/v3/model_performance')\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "epoch = checkpoint['epoch']\n",
    "train_perf = checkpoint['train_performances']\n",
    "val_perf = checkpoint['val_performances']\n",
    "print(train_perf)\n",
    "\n",
    "# model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'drug_response_matrix' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb Cell 30'\u001b[0m in \u001b[0;36m<cell line: 3>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000033?line=0'>1</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtorch_geometric\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mloader\u001b[39;00m \u001b[39mimport\u001b[39;00m DataLoader\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000033?line=2'>3</a>\u001b[0m sample \u001b[39m=\u001b[39m drug_response_matrix\u001b[39m.\u001b[39msample(\u001b[39m1_000\u001b[39m)\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000033?line=3'>4</a>\u001b[0m train_set, test_val_set \u001b[39m=\u001b[39m train_test_split(sample, test_size\u001b[39m=\u001b[39m\u001b[39m0.8\u001b[39m, random_state\u001b[39m=\u001b[39m\u001b[39m42\u001b[39m)\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v3_GraphTab.ipynb#ch0000033?line=4'>5</a>\u001b[0m sample_dataset \u001b[39m=\u001b[39m GraphTabDataset(cl_graphs\u001b[39m=\u001b[39mcl_graphs, drugs\u001b[39m=\u001b[39mfingerprints_dict, drug_response_matrix\u001b[39m=\u001b[39mtrain_set)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'drug_response_matrix' is not defined"
     ]
    }
   ],
   "source": [
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "sample = drug_response_matrix.sample(1_000)\n",
    "train_set, test_val_set = train_test_split(sample, test_size=0.8, random_state=42)\n",
    "sample_dataset = GraphTabDataset(cl_graphs=cl_graphs, drugs=fingerprints_dict, drug_response_matrix=train_set)\n",
    "print(\"\\ntrain_dataset:\")\n",
    "sample_dataset.print_dataset_summary()\n",
    "sample_loader = DataLoader(dataset=sample_dataset, batch_size=2, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.48it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:17<00:00,  6.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  0\n",
      "Train      | MSE: 13.98202\n",
      "Validation | MSE: 11.01190\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.67it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:16<00:00,  6.74it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  1\n",
      "Train      | MSE: 11.09647\n",
      "Validation | MSE: 9.97083\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.81it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:24<00:00,  6.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  2\n",
      "Train      | MSE: 8.56703\n",
      "Validation | MSE: 8.11504\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.24it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:33<00:00,  5.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  3\n",
      "Train      | MSE: 6.49640\n",
      "Validation | MSE: 7.55463\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.81it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:23<00:00,  6.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  4\n",
      "Train      | MSE: 6.88781\n",
      "Validation | MSE: 7.54679\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.33it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:20<00:00,  6.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  5\n",
      "Train      | MSE: 7.12907\n",
      "Validation | MSE: 7.46634\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.67it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:20<00:00,  6.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  6\n",
      "Train      | MSE: 6.83407\n",
      "Validation | MSE: 10.60364\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.90it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:22<00:00,  6.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  7\n",
      "Train      | MSE: 6.86842\n",
      "Validation | MSE: 7.34576\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.91it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:20<00:00,  6.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  8\n",
      "Train      | MSE: 6.24524\n",
      "Validation | MSE: 7.07446\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.86it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:14<00:00,  6.85it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  9\n",
      "Train      | MSE: 6.77860\n",
      "Validation | MSE: 7.24841\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.56it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:09<00:00,  7.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  10\n",
      "Train      | MSE: 6.18162\n",
      "Validation | MSE: 6.95470\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.57it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:09<00:00,  7.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  11\n",
      "Train      | MSE: 6.43730\n",
      "Validation | MSE: 6.98433\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.45it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:08<00:00,  7.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  12\n",
      "Train      | MSE: 5.98261\n",
      "Validation | MSE: 7.09261\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.22it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:08<00:00,  7.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  13\n",
      "Train      | MSE: 6.42925\n",
      "Validation | MSE: 7.40595\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.47it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:05<00:00,  7.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  14\n",
      "Train      | MSE: 6.39245\n",
      "Validation | MSE: 7.20384\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.56it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:07<00:00,  7.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  15\n",
      "Train      | MSE: 6.58797\n",
      "Validation | MSE: 7.04990\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.19it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:19<00:00,  6.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  16\n",
      "Train      | MSE: 6.48583\n",
      "Validation | MSE: 7.16788\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:07<00:00, 13.69it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:25<00:00,  6.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  17\n",
      "Train      | MSE: 6.09398\n",
      "Validation | MSE: 7.05967\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 15.39it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:22<00:00,  6.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  18\n",
      "Train      | MSE: 6.40673\n",
      "Validation | MSE: 7.07927\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 100/100 [00:06<00:00, 16.17it/s]\n",
      "Iter: 100%|██████████| 920/920 [02:18<00:00,  6.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=====Epoch  19\n",
      "Train      | MSE: 6.35017\n",
      "Validation | MSE: 7.07385\n"
     ]
    }
   ],
   "source": [
    "train_losses, val_losses = build_model.train(sample_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 1:\n",
      "=======\n",
      "Number of graphs in the current batch: 10\n",
      "DataBatch(x=[8580, 4], edge_index=[2, 831260], batch=[8580], ptr=[11])\n",
      "tensor([ 4.5438,  4.4568, -0.9637, -2.6823, -3.9976,  3.0180, -2.8586,  4.7727,\n",
      "        -2.9860, -1.4190], dtype=torch.float64)\n",
      "\n",
      "Step 2:\n",
      "=======\n",
      "Number of graphs in the current batch: 10\n",
      "DataBatch(x=[8580, 4], edge_index=[2, 831260], batch=[8580], ptr=[11])\n",
      "tensor([ 3.1815,  2.9466,  0.6865,  4.0391,  0.9185,  4.0514,  2.4978,  6.8512,\n",
      "        -3.0862,  4.6164], dtype=torch.float64)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for step, data in enumerate(train_loader):\n",
    "    cell, drug, ic50 = data\n",
    "    print(f'Step {step + 1}:')\n",
    "    print(f'=======')\n",
    "    print(f'Number of graphs in the current batch: {cell.num_graphs}')\n",
    "    print(cell)\n",
    "    print(ic50)\n",
    "    print()\n",
    "\n",
    "    if step == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 304,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "index 858000 is out of bounds for dimension 0 with size 858000",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[1;32m/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v1_GraphTab.ipynb Cell 55'\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v1_GraphTab.ipynb#ch0000061?line=0'>1</a>\u001b[0m \u001b[39m# 858000\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/cwoest/Documents/Academics/Data_Science_UP/master_thesis/material/GNN-material/11_v1_GraphTab.ipynb#ch0000061?line=1'>2</a>\u001b[0m cell\u001b[39m.\u001b[39;49mx[\u001b[39m858000\u001b[39;49m]\n",
      "\u001b[0;31mIndexError\u001b[0m: index 858000 is out of bounds for dimension 0 with size 858000"
     ]
    }
   ],
   "source": [
    "# 858000\n",
    "cell.x[858000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 305,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 4.2288, -1.0000,  2.0000,  0.0000], dtype=torch.float64)"
      ]
     },
     "execution_count": 305,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cell.x[858000-1]"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "c78b81650a0bd32063743affb6953ff71b1a0dba806fbca9e2db842718495748"
  },
  "kernelspec": {
   "display_name": "Python 3.10.4 ('master-thesis-log')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}