graph_model_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\lg\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\geopandas\\_compat.py:124: UserWarning: The Shapely GEOS version (3.11.1-CAPI-1.17.1) is incompatible with the GEOS version PyGEOS was compiled with (3.10.4-CAPI-1.16.2). Conversions between both will be slow.\n",
      "  warnings.warn(\n",
      "C:\\Users\\lg\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\spatialdata\\__init__.py:9: UserWarning: Geopandas was set to use PyGEOS, changing to shapely 2.0 with:\n",
      "\n",
      "\tgeopandas.options.use_pygeos = True\n",
      "\n",
      "If you intended to use PyGEOS, set the option to False.\n",
      "  _check_geopandas_using_shapely()\n"
     ]
    }
   ],
   "source": [
    "from spatialSSL.Dataloader import FullImageDatasetConstructor\n",
    "from spatialSSL.Utils import split_dataset\n",
    "from spatialSSL.Training import train\n",
    "from spatialSSL.Testing import test\n",
    "from spatialSSL.Dataset import InMemoryGraphDataset\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "'\\nimport os\\nimport scanpy as sc\\n\\nadata = sc.read(\"./data/atlas_brain_638850_CCF.h5ad\")\\n# subset the adata\\nselected_img = np.unique(adata.obs[\"section\"])[0:20]\\nsub_adata = adata[adata.obs[\"section\"].isin(selected_img)].copy()\\n\\n# specify the output directory and file\\noutput_dir = \\'./data\\'\\nos.makedirs(output_dir, exist_ok=True)\\noutput_file = os.path.join(output_dir, \\'sub_20img_adata.h5ad\\')\\n\\n# write the sub_adata into a .h5ad file\\nsub_adata.write(output_file)\\n'"
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'''\n",
    "import os\n",
    "import scanpy as sc\n",
    "\n",
    "adata = sc.read(\"./data/atlas_brain_638850_CCF.h5ad\")\n",
    "# subset the adata\n",
    "selected_img = np.unique(adata.obs[\"section\"])[0:20]\n",
    "sub_adata = adata[adata.obs[\"section\"].isin(selected_img)].copy()\n",
    "\n",
    "# specify the output directory and file\n",
    "output_dir = './data'\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "output_file = os.path.join(output_dir, 'sub_20img_adata.h5ad')\n",
    "\n",
    "# write the sub_adata into a .h5ad file\n",
    "sub_adata.write(output_file)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create an instance of Full_image_dataloader\n",
    "\n",
    " \n",
    "#file_path = \"./data/img_119670929.h5ad\"\n",
    "file_path = \"./data/subset_6img_atlas_brain.h5ad\"\n",
    "#file_path = \"./data/atlas_brain_638850_CCF.h5ad\"\n",
    "#file_path = \"./data/sub_20img_adata.h5ad\"\n",
    "\n",
    "data_constracter = FullImageDatasetConstructor(file_path=file_path,\n",
    "                                        image_col=\"section\",\n",
    "                                        label_col=\"class_label\",\n",
    "                                        include_label=False,\n",
    "                                        radius=20,\n",
    "                                        node_level = 1,\n",
    "                                       mask_method = \"cell_type\",\n",
    "                                       celltype_to_mask = \"LQ\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load the data\n",
    "data_constracter.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "Constructing Graphs:   0%|          | 0/6 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1e2691cb4a3b463e84bc1a67e8abc64d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Construct the graph\n",
    "graph_list = data_constracter.construct_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "6"
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(graph_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_loader, val_loader, test_loader = split_dataset(graph_list,split_percent=(0.6, 0.2, 0.2), batch_size=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": "<torch_geometric.loader.dataloader.DataLoader at 0x29dc3f4ee60>"
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train size: 4\n",
      "Validation size: 1\n",
      "Test size: 1\n"
     ]
    }
   ],
   "source": [
    "# Print out the size of each set to verify\n",
    "print(f\"Train size: {len(train_loader.dataset)}\")\n",
    "print(f\"Validation size: {len(val_loader.dataset)}\")\n",
    "print(f\"Test size: {len(test_loader.dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from torch import nn, optim, Tensor\n",
    "import torch\n",
    "from torch_geometric.nn import GCNConv,GATConv\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from sklearn.metrics import r2_score\n",
    "from torch.nn import LeakyReLU, Dropout\n",
    "import time\n",
    "\n",
    "class GCN(nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels, dropout_rate=0.5):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(in_channels, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GCNConv(hidden_channels, out_channels)\n",
    "        self.conv4 = GCNConv(hidden_channels, out_channels)\n",
    "        self.dropout = Dropout(dropout_rate)\n",
    "        \n",
    "        self.act = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:\n",
    "        x = self.dropout(self.act(self.conv1(x, edge_index)))\n",
    "        x = self.dropout(self.act(self.conv2(x, edge_index)))\n",
    "        x = self.dropout(self.act(self.conv3(x, edge_index)))\n",
    "        x = self.act(self.conv4(x, edge_index))  # Typically, dropout is not applied to the final layer.\n",
    "        return x\n",
    "\n",
    "class GNN(nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels, dropout_rate=0.5):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(in_channels, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GCNConv(hidden_channels, out_channels)\n",
    "        self.conv4 = GCNConv(hidden_channels, out_channels)\n",
    "        self.dropout = Dropout(dropout_rate)\n",
    "        \n",
    "        self.act = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:\n",
    "        x = self.dropout(self.act(self.conv1(x, edge_index)))\n",
    "        x = self.dropout(self.act(self.conv2(x, edge_index)))\n",
    "        x = self.dropout(self.act(self.conv3(x, edge_index)))\n",
    "        x = self.act(self.conv4(x, edge_index))  # Typically, dropout is not applied to the final layer.\n",
    "        return x\n",
    "    \n",
    "class GAT(nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels, dropout_rate=0.5):\n",
    "        super().__init__()\n",
    "        self.conv1 = GATConv(in_channels, hidden_channels)\n",
    "        self.conv2 = GATConv(hidden_channels, hidden_channels)\n",
    "        self.conv3 = GATConv(hidden_channels, out_channels)\n",
    "        self.conv4 = GATConv(hidden_channels, out_channels)\n",
    "        self.dropout = Dropout(dropout_rate)\n",
    "        self.act = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:\n",
    "        x = self.dropout(self.act(self.conv1(x, edge_index)))\n",
    "        x = self.dropout(self.act(self.conv2(x, edge_index)))\n",
    "        x = self.dropout(self.act(self.conv3(x, edge_index)))\n",
    "        x = self.act(self.conv4(x, edge_index))  # Typically, dropout is not applied to the final layer.\n",
    "        return x\n",
    "    \n",
    "from torch_geometric.nn import TopKPooling\n",
    "\n",
    "class GAT_pool(nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels, dropout_rate=0.5):\n",
    "        super().__init__()\n",
    "        self.conv1 = GATConv(in_channels, hidden_channels)\n",
    "        self.pool1 = TopKPooling(hidden_channels, ratio=0.8)\n",
    "        self.conv2 = GATConv(hidden_channels, hidden_channels)\n",
    "        self.pool2 = TopKPooling(hidden_channels, ratio=0.8)\n",
    "        self.conv3 = GATConv(hidden_channels, hidden_channels)\n",
    "        self.conv4 = GATConv(hidden_channels, out_channels)\n",
    "        self.dropout = Dropout(dropout_rate)\n",
    "        self.act = nn.LeakyReLU()\n",
    "\n",
    "    def forward(self, x: Tensor, edge_index: Tensor, mask: Tensor) -> Tensor:\n",
    "        x = self.dropout(self.act(self.conv1(x, edge_index)))\n",
    "        x, edge_index, _, batch, _, mask = self.pool1(x, edge_index, None, None, mask)\n",
    "        x = self.dropout(self.act(self.conv2(x, edge_index)))\n",
    "        x, edge_index, _, batch, _, mask = self.pool2(x, edge_index, None, batch, mask)\n",
    "        x = self.dropout(self.act(self.conv3(x, edge_index)))\n",
    "        x = self.act(self.conv4(x, edge_index))  # Typically, dropout is not applied to the final layer.\n",
    "        return x, mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Spatial training with pool\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/150, train loss: 0.8534, train r2: -1.3632,  val loss: 0.7423, val r2: -0.7248, Time: 27.489891052246094s\n",
      "Epoch 2/150, train loss: 0.7351, train r2: -1.3089,  val loss: 0.5972, val r2: -0.7231, Time: 21.353416681289673s\n",
      "Epoch 3/150, train loss: 0.6357, train r2: -1.6027,  val loss: 0.5093, val r2: -0.5802, Time: 20.558507204055786s\n",
      "Epoch 4/150, train loss: 0.5721, train r2: -1.0490,  val loss: 0.4669, val r2: -0.3823, Time: 21.35031819343567s\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[13], line 16\u001B[0m\n\u001B[0;32m     12\u001B[0m num_epochs \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m150\u001B[39m\n\u001B[0;32m     13\u001B[0m patience \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m8\u001B[39m\n\u001B[1;32m---> 16\u001B[0m \u001B[43mtrain\u001B[49m\u001B[43m(\u001B[49m\u001B[43mmodel\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mmodel\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtrain_loader\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mtrain_loader\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mval_loader\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mval_loader\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcriterion\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mcriterion\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mnum_epochs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m \u001B[49m\u001B[43mnum_epochs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mpatience\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m \u001B[49m\u001B[43mpatience\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43moptimizer\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m \u001B[49m\u001B[43moptimizer\u001B[49m\u001B[43m,\u001B[49m\u001B[43mmodel_path\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m \u001B[49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43m./models/model.pt\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m)\u001B[49m\n",
      "File \u001B[1;32m~\\PycharmProjects\\spatial_atlas_ssl\\spatialSSL\\Training.py:56\u001B[0m, in \u001B[0;36mtrain\u001B[1;34m(model, train_loader, val_loader, criterion, num_epochs, patience, optimizer, model_path, gene_expression)\u001B[0m\n\u001B[0;32m     54\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m epoch \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mrange\u001B[39m(num_epochs):\n\u001B[0;32m     55\u001B[0m     start_time \u001B[38;5;241m=\u001B[39m time\u001B[38;5;241m.\u001B[39mtime()\n\u001B[1;32m---> 56\u001B[0m     train_loss, train_r2 \u001B[38;5;241m=\u001B[39m \u001B[43mtrain_epoch\u001B[49m\u001B[43m(\u001B[49m\u001B[43mmodel\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtrain_loader\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43moptimizer\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcriterion\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mgene_expression\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtraining\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m)\u001B[49m\n\u001B[0;32m     57\u001B[0m     val_loss, val_r2 \u001B[38;5;241m=\u001B[39m train_epoch(model, val_loader, optimizer,criterion, gene_expression, training\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mFalse\u001B[39;00m)\n\u001B[0;32m     58\u001B[0m     \u001B[38;5;66;03m# scheduler.step() # Decrease learning rate by scheduler\u001B[39;00m\n",
      "File \u001B[1;32m~\\PycharmProjects\\spatial_atlas_ssl\\spatialSSL\\Training.py:27\u001B[0m, in \u001B[0;36mtrain_epoch\u001B[1;34m(model, loader, optimizer, criterion, gene_expression, training)\u001B[0m\n\u001B[0;32m     25\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m gene_expression \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[0;32m     26\u001B[0m     data \u001B[38;5;241m=\u001B[39m data\u001B[38;5;241m.\u001B[39mto(device)\n\u001B[1;32m---> 27\u001B[0m     outputs \u001B[38;5;241m=\u001B[39m \u001B[43mmodel\u001B[49m\u001B[43m(\u001B[49m\u001B[43mdata\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mx\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mfloat\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mdata\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43medge_index\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mlong\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m\n\u001B[0;32m     28\u001B[0m     loss \u001B[38;5;241m=\u001B[39m criterion(outputs[data\u001B[38;5;241m.\u001B[39mmask], data\u001B[38;5;241m.\u001B[39my\u001B[38;5;241m.\u001B[39mfloat())\n\u001B[0;32m     29\u001B[0m     targets_list\u001B[38;5;241m.\u001B[39mappend(data\u001B[38;5;241m.\u001B[39my\u001B[38;5;241m.\u001B[39mcpu()\u001B[38;5;241m.\u001B[39mdetach())\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[0;32m   1497\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[0;32m   1498\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[0;32m   1499\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[0;32m   1500\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[1;32m-> 1501\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m forward_call(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n\u001B[0;32m   1502\u001B[0m \u001B[38;5;66;03m# Do not call functions when jit is used\u001B[39;00m\n\u001B[0;32m   1503\u001B[0m full_backward_hooks, non_full_backward_hooks \u001B[38;5;241m=\u001B[39m [], []\n",
      "Cell \u001B[1;32mIn[10], line 58\u001B[0m, in \u001B[0;36mGAT.forward\u001B[1;34m(self, x, edge_index)\u001B[0m\n\u001B[0;32m     56\u001B[0m x \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mdropout(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mact(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mconv1(x, edge_index)))\n\u001B[0;32m     57\u001B[0m x \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mdropout(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mact(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mconv2(x, edge_index)))\n\u001B[1;32m---> 58\u001B[0m x \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mdropout(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mact(\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mconv3\u001B[49m\u001B[43m(\u001B[49m\u001B[43mx\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43medge_index\u001B[49m\u001B[43m)\u001B[49m))\n\u001B[0;32m     59\u001B[0m x \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mact(\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mconv4(x, edge_index))  \u001B[38;5;66;03m# Typically, dropout is not applied to the final layer.\u001B[39;00m\n\u001B[0;32m     60\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m x\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[0;32m   1497\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[0;32m   1498\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[0;32m   1499\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[0;32m   1500\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[1;32m-> 1501\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m forward_call(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n\u001B[0;32m   1502\u001B[0m \u001B[38;5;66;03m# Do not call functions when jit is used\u001B[39;00m\n\u001B[0;32m   1503\u001B[0m full_backward_hooks, non_full_backward_hooks \u001B[38;5;241m=\u001B[39m [], []\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\torch_geometric\\nn\\conv\\gat_conv.py:213\u001B[0m, in \u001B[0;36mGATConv.forward\u001B[1;34m(self, x, edge_index, edge_attr, size, return_attention_weights)\u001B[0m\n\u001B[0;32m    211\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(x, Tensor):\n\u001B[0;32m    212\u001B[0m     \u001B[38;5;28;01massert\u001B[39;00m x\u001B[38;5;241m.\u001B[39mdim() \u001B[38;5;241m==\u001B[39m \u001B[38;5;241m2\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mStatic graphs not supported in \u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mGATConv\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m--> 213\u001B[0m     x_src \u001B[38;5;241m=\u001B[39m x_dst \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mlin_src\u001B[49m\u001B[43m(\u001B[49m\u001B[43mx\u001B[49m\u001B[43m)\u001B[49m\u001B[38;5;241m.\u001B[39mview(\u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m, H, C)\n\u001B[0;32m    214\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:  \u001B[38;5;66;03m# Tuple of source and target node features:\u001B[39;00m\n\u001B[0;32m    215\u001B[0m     x_src, x_dst \u001B[38;5;241m=\u001B[39m x\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\torch\\nn\\modules\\module.py:1501\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[1;34m(self, *args, **kwargs)\u001B[0m\n\u001B[0;32m   1496\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[0;32m   1497\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[0;32m   1498\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[0;32m   1499\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[0;32m   1500\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[1;32m-> 1501\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m forward_call(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)\n\u001B[0;32m   1502\u001B[0m \u001B[38;5;66;03m# Do not call functions when jit is used\u001B[39;00m\n\u001B[0;32m   1503\u001B[0m full_backward_hooks, non_full_backward_hooks \u001B[38;5;241m=\u001B[39m [], []\n",
      "File \u001B[1;32m~\\anaconda3\\envs\\spatial_atlas_ssl\\lib\\site-packages\\torch_geometric\\nn\\dense\\linear.py:132\u001B[0m, in \u001B[0;36mLinear.forward\u001B[1;34m(self, x)\u001B[0m\n\u001B[0;32m    127\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mforward\u001B[39m(\u001B[38;5;28mself\u001B[39m, x: Tensor) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m Tensor:\n\u001B[0;32m    128\u001B[0m \u001B[38;5;250m    \u001B[39m\u001B[38;5;124mr\u001B[39m\u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[0;32m    129\u001B[0m \u001B[38;5;124;03m    Args:\u001B[39;00m\n\u001B[0;32m    130\u001B[0m \u001B[38;5;124;03m        x (torch.Tensor): The input features.\u001B[39;00m\n\u001B[0;32m    131\u001B[0m \u001B[38;5;124;03m    \"\"\"\u001B[39;00m\n\u001B[1;32m--> 132\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mF\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mlinear\u001B[49m\u001B[43m(\u001B[49m\u001B[43mx\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mweight\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbias\u001B[49m\u001B[43m)\u001B[49m\n",
      "\u001B[1;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "# Define the device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Create the model\n",
    "model = GAT(550, 550, 550).to(device) # in_channels is set to 100 as an example. Please replace it with your actual feature size.\n",
    "\n",
    "# Define loss function and optimizer\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.0005)\n",
    "#scheduler = StepLR(optimizer, step_size=80, gamma=0.1) # learning rate scheduler\n",
    "\n",
    "num_epochs = 150\n",
    "patience = 8\n",
    "\n",
    "\n",
    "train(model=model, train_loader=train_loader, val_loader=val_loader, criterion=criterion, num_epochs= num_epochs, patience = patience, optimizer= optimizer,model_path = './models/model.pt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_loss, test_r2,test_targets_list,test_outputs_list,test_celltypes_list = test(model, test_loader, criterion, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "# usage\n",
    "model = GAT(550, 550, 550)  # Replace with your actual model\n",
    "print(f'The model has {count_parameters(model):,} parameters')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Flatten the nested list into a single-dimensional list\n",
    "flattened_list = [item for sublist in np.concatenate(test_celltypes_list) for item in sublist]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "preds_test_per_cell = pd.DataFrame({'celltype': flattened_list, 'r2_score': r2_score(np.transpose(torch.cat(test_targets_list).cpu().detach().numpy()), np.transpose(torch.cat(test_outputs_list).cpu().detach().numpy()), multioutput='raw_values')})\n",
    "\n",
    "preds_test_per_gene = pd.DataFrame({'gene': data_constracter.adata.var_names, 'r2_score': r2_score(torch.cat(test_targets_list).cpu().detach().numpy(), torch.cat(test_outputs_list).cpu().detach().numpy(), multioutput='raw_values')})\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "preds_test_per_gene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# huge class inbalance problem, need to handel -> add weight?\n",
    "# add pooling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Plotting the histogram\n",
    "plt.figure(figsize=(10,6))\n",
    "plt.hist(preds_test_per_gene['r2_score'], bins=20, alpha=0.8, color='skyblue', edgecolor='black')\n",
    "plt.title('R2 Score Distribution per Gene')\n",
    "plt.xlabel('R2 Score')\n",
    "plt.ylabel('Frequency')\n",
    "plt.show()\n",
    "\n",
    "# Sorting by r2_score\n",
    "preds_test_per_gene = preds_test_per_gene.sort_values(by='r2_score', ascending=False)\n",
    "\n",
    "# Display the sorted DataFrame\n",
    "print(preds_test_per_gene)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "preds_test_per_gene\n",
    "\n",
    "# plot r2 score distribution per gene, histogram\n",
    "# density plot\n",
    "preds_test_per_gene.r2_score.plot.kde()\n",
    "#preds_test_per_gene.r2_score.plot.hist(bins=20, alpha=0.8)\n",
    "\n",
    "# plot r2 score distribution per celltype, histogram\n",
    "#preds_test_per_cell.r2_score.plot.hist(bins=10, alpha=0.8)\n",
    "\n",
    "# sort by r2 score\n",
    "preds_test_per_gene.sort_values(by='r2_score', ascending=False, inplace=True)\n",
    "preds_test_per_gene"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# set style\n",
    "sns.set_palette(\"colorblind\")\n",
    "sns.set_context(\"paper\", font_scale=1.5)\n",
    "\n",
    "# make plot larger\n",
    "plt.figure(figsize=(18,8))\n",
    "\n",
    "# rotate x-axis labels\n",
    "plt.xticks(rotation=45)\n",
    "\n",
    "# create plot of mean R2 score per celltype\n",
    "# order bars by mean R2 score\n",
    "\n",
    "# add avlue counts from adata.obs.class_label.value_counts()\n",
    "ax = sns.barplot(x=\"celltype\", y=\"r2_score\", data=preds_test_per_cell, order=preds_test_per_cell.groupby('celltype').mean().sort_values(by='r2_score', ascending=False).index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Get the frequency of each celltype\n",
    "celltype_counts = preds_test_per_cell['celltype'].value_counts().sort_values(ascending=False)\n",
    "\n",
    "plt.figure(figsize=(10,6))\n",
    "plt.bar(celltype_counts.index, celltype_counts.values, alpha=0.8, color='skyblue', edgecolor='black')\n",
    "plt.xticks(rotation='vertical')\n",
    "plt.xticks(rotation=90)\n",
    "plt.title('Celltype Distribution')\n",
    "plt.xlabel('Celltype')\n",
    "plt.ylabel('Frequency')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "sns.violinplot(x=\"celltype\", y=\"r2_score\", data=preds_test_per_cell)\n",
    "# rotate x-axis labels\n",
    "plt.xticks(rotation=90)\n",
    "plt.title('R2 Score Distribution per Gene')\n",
    "plt.xlabel('R2 Score')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "file_path = \"./data/subset_6img_atlas_brain.h5ad\"\n",
    "root_path = \"./data\"\n",
    "dataset = InMemoryGraphDataset(root= root_path,\n",
    "                               data_names = \"data_test\",\n",
    "                               file_path=file_path, image_col=\"section\", label_col=\"class_label\", include_label=False, radius=20,node_level = 1, batch_size=1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load your dataset\n",
    "import torch\n",
    "data = torch.load('./data/data_test.pt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  },
  "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}