{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "W tym notatniku zbadamy różne modele grafowych sieci neuronowych (GNN) i porównamy ich wydajność na zbiorze danych Cora. Wykorzystamy biblioteki takie jak PyTorch, PyTorch Geometric oraz PyTorch Lightning do budowy, trenowania i oceny modeli. Przeanalizujemy klasyczne modele GNN, modele z wielokrotną agregacją, modele z mechanizmem Jumping Knowledge oraz inne zaawansowane architektury.\n",
    "W szczególności, notatnik przedstawia działanie różnych rozszerzeń do sieci GNN, takich jak połączenia skokowe (Jumping Knowledge) oraz wiele agregacji (Multi-Aggregation). Na końcu przeprowadzimy eksperymenty i porównamy działanie różnych modeli, aby ocenić ich wydajność na zbiorze danych Cora."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch as th\n",
    "import torch_geometric as pyg\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torchmetrics as tm\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "import pytorch_lightning as pl\n",
    "from tqdm.auto import tqdm\n",
    "from pathlib import Path\n",
    "from typing import List, Literal, Tuple, Callable\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wczytanie danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = Path(os.pardir, \"data\", \"Planetoid\")\n",
    "data_path.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/filip/projects/grafowe_sieci_neuronowe/.venv/lib/python3.12/site-packages/torch_geometric/data/dataset.py:115: UserWarning: The `pre_transform` argument differs from the one used in the pre-processed version of this dataset. If you want to make use of another pre-processing technique, pass `force_reload=True` explicitly to reload the dataset.\n",
      "  self._process()\n"
     ]
    }
   ],
   "source": [
    "cora = pyg.datasets.Planetoid(root=data_path, name=\"Cora\", pre_transform=pyg.transforms.NormalizeFeatures())[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ustawienia treningu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 25\n",
    "lr = 0.01\n",
    "ncls = th.unique(cora['y']).size(0)\n",
    "dim_in = cora['x'].size(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collect_scores(yhat: th.Tensor, y: th.Tensor, ncls: int, metric: Callable):\n",
    "    return metric(\n",
    "        yhat,\n",
    "        y,\n",
    "        num_classes=ncls,\n",
    "        average='weighted',\n",
    "        task='multiclass'\n",
    "    )\n",
    "\n",
    "\n",
    "def train_and_test_model(\n",
    "        model: th.nn.Module,\n",
    "        data: pyg.data.Data,\n",
    "        epochs: int = epochs,\n",
    "        lr: float = lr):\n",
    "    optim = th.optim.Adam(model.parameters(), lr=lr)\n",
    "    train_losses = []\n",
    "    val_losses = []\n",
    "    for _ in tqdm(range(epochs)):\n",
    "        model.train()\n",
    "        optim.zero_grad()\n",
    "        out = model(data['x'], data['edge_index'])\n",
    "        loss = th.nn.functional.cross_entropy(out[data['train_mask']], data['y'][data['train_mask']])\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "        train_losses.append(loss.item())\n",
    "\n",
    "        model.eval()\n",
    "        out = model(data['x'], data['edge_index'])\n",
    "        loss = th.nn.functional.cross_entropy(out[data['val_mask']], data['y'][data['val_mask']])\n",
    "        val_losses.append(loss.item())\n",
    "\n",
    "    model.eval()\n",
    "    ncls = th.unique(data['y']).size(0)\n",
    "    yhat = model(data['x'], data['edge_index'])\n",
    "\n",
    "    scores = {}\n",
    "    for mask in ['test_mask', 'val_mask']:\n",
    "        set_name = mask.replace(\"_mask\", \"\")\n",
    "        yhat_mask = yhat[data[mask]]\n",
    "        y_mask = data['y'][data[mask]]\n",
    "        scores[set_name] = {\n",
    "            \"f1\": collect_scores(yhat_mask, y_mask, ncls, tm.functional.f1_score).item(),\n",
    "            \"roc_auc\": collect_scores(yhat_mask, y_mask, ncls, tm.functional.auroc).item()\n",
    "        }\n",
    "\n",
    "    return scores, train_losses, val_losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model klasyczny"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f41db3416cb54495bbddcb9a7514d5ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "classic_model = pyg.nn.Sequential('x, edge_index', [\n",
    "    (pyg.nn.SAGEConv(dim_in, 512, normalize=True,), 'x, edge_index -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (pyg.nn.SAGEConv(512, 256, normalize=True), 'x, edge_index -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (pyg.nn.SAGEConv(256, 128, normalize=True), 'x, edge_index -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (th.nn.Dropout(0.2), 'x -> x'),\n",
    "    (th.nn.Linear(128, 64), 'x -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (th.nn.Dropout(0.2), 'x -> x'),\n",
    "    (th.nn.Linear(64, ncls), 'x -> x'),\n",
    "])\n",
    "\n",
    "scores_classic, classic_train_losses, classic_test_losses = train_and_test_model(classic_model, cora)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "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>test</th>\n",
       "      <th>val</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>f1</th>\n",
       "      <td>0.775108</td>\n",
       "      <td>0.765161</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>roc_auc</th>\n",
       "      <td>0.934366</td>\n",
       "      <td>0.929659</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             test       val\n",
       "f1       0.775108  0.765161\n",
       "roc_auc  0.934366  0.929659"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(scores_classic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37560 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(classic_train_losses, label='train')\n",
    "plt.plot(classic_test_losses, label='val')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rozszerzenia klasycznych modeli"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wielokrotne agregacje"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Zachowanie pojedynczej warstwy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2708, 64])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sage_layer_multi = pyg.nn.SAGEConv(dim_in, 64, normalize=True, aggr=['mean', 'sum', 'max'])\n",
    "with th.no_grad():\n",
    "    out_sage_multi = sage_layer_multi(cora['x'], cora['edge_index'])\n",
    "out_sage_multi.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 4299])\n",
      "torch.Size([64])\n",
      "torch.Size([64, 1433])\n"
     ]
    }
   ],
   "source": [
    "for param in sage_layer_multi.parameters():\n",
    "    print(param.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trening modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a26744460b4c4713b8c3612ab1b83489",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "multiagg_model = pyg.nn.Sequential('x, edge_index', [\n",
    "    (pyg.nn.SAGEConv(dim_in, 128, normalize=True, aggr=['mean', 'max', 'sum', 'min']), 'x, edge_index -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (pyg.nn.SAGEConv(128, 128, normalize=True, aggr=['mean', 'max', 'sum', 'min']), 'x, edge_index -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (pyg.nn.SAGEConv(128, 64, normalize=True, aggr=['mean', 'max', 'sum', 'min']), 'x, edge_index -> x'),\n",
    "    (th.nn.ReLU(), 'x -> x'),\n",
    "    (th.nn.Dropout(0.2), 'x -> x'),\n",
    "    (th.nn.Linear(64, ncls), 'x -> x'),\n",
    "])\n",
    "\n",
    "scores_multiagg, multiagg_train_losses, multiagg_test_losses = train_and_test_model(multiagg_model, cora)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "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>test</th>\n",
       "      <th>val</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>f1</th>\n",
       "      <td>0.768343</td>\n",
       "      <td>0.740840</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>roc_auc</th>\n",
       "      <td>0.928856</td>\n",
       "      <td>0.909718</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             test       val\n",
       "f1       0.768343  0.740840\n",
       "roc_auc  0.928856  0.909718"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(scores_multiagg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34384 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(multiagg_train_losses, label='train')\n",
    "plt.plot(multiagg_test_losses, label='val')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network in graph network\n",
    "\n",
    "1. Publikacja: https://arxiv.org/abs/2111.11638\n",
    "2. Post na blogu DGL: https://www.dgl.ai/blog/2022/11/28/ngnn.html\n",
    "3. Artykuł na medium.com: https://medium.com/@hellorahulk/exploring-over-smoothing-in-graph-neural-networks-gnns-e311b8eb8a85\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NGNN(th.nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "            self, \n",
    "            gnn_module: pyg.nn.MessagePassing,\n",
    "            in_channels: int,\n",
    "            hidden_channels: int,\n",
    "            out_channels: int,\n",
    "            *args,\n",
    "            **kwargs):\n",
    "        super().__init__()\n",
    "        self.gnn = gnn_module(in_channels, hidden_channels, *args, **kwargs)\n",
    "        self.lin = th.nn.Linear(hidden_channels, out_channels)\n",
    "\n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        gnn_out = self.gnn(x, edge_index).relu()\n",
    "        lin_out = self.lin(gnn_out)\n",
    "        return lin_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    pyg.seed_everything(123)\n",
    "    nin_gnn = NGNN(pyg.nn.SAGEConv, dim_in, 512, 512, normalize=True)\n",
    "    out = nin_gnn(cora['x'], cora['edge_index'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c60e93aeaee2426f9ab52dbfd2c0c67b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "ngnn = pyg.nn.Sequential('x, edge_index', [\n",
    "    (NGNN(pyg.nn.SAGEConv, dim_in, 512, 256, normalize=True), 'x,edge_index -> x'),\n",
    "    (NGNN(pyg.nn.SAGEConv, 256, 128, 64, normalize=True), 'x,edge_index -> x'),\n",
    "    (NGNN(pyg.nn.SAGEConv, 64, 64, ncls, normalize=True), 'x,edge_index -> x'),\n",
    "])\n",
    "\n",
    "scores_ninn, ninn_train_losses, ninn_test_losses = train_and_test_model(ngnn, cora)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "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>test</th>\n",
       "      <th>val</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>f1</th>\n",
       "      <td>0.773679</td>\n",
       "      <td>0.763213</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>roc_auc</th>\n",
       "      <td>0.931426</td>\n",
       "      <td>0.922790</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             test       val\n",
       "f1       0.773679  0.763213\n",
       "roc_auc  0.931426  0.922790"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame(scores_ninn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34288 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(ninn_train_losses, label='train')\n",
    "plt.plot(ninn_test_losses, label='val')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Jumping knowledge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Własna implementacja"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class JKModel(th.nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            gnn_modules: List[pyg.nn.MessagePassing],\n",
    "            lin_sizes: Tuple[int, ...],\n",
    "            mode: Literal[\"cat\", \"max\", \"mean\"] = \"cat\",\n",
    "            dropout: float = 0.2):\n",
    "        super().__init__()\n",
    "        self.convs = th.nn.ModuleList()\n",
    "        self.jk_mode = mode\n",
    "        self.dropout = dropout\n",
    "\n",
    "        conv_out_total_size = 0\n",
    "        last_hidden_channels = 0\n",
    "\n",
    "        # Konstruowanie serii warstw splotu GNN\n",
    "        for idx, gnn in enumerate(gnn_modules):\n",
    "            # Dodaj warstwę splotu GNN do listy\n",
    "            self.convs.append(gnn)\n",
    "            # Dodaj wielkość wyjścia do całkowitej wielkości wyjścia\n",
    "            conv_out_total_size += gnn.out_channels\n",
    "\n",
    "            # Jeśli mechanizm JK ma opierać się na średniej lub maksimum:\n",
    "            # wszystkie warstwy muszą mieć taką samą liczbę kanałów\n",
    "            if idx > 0 and self.jk_mode != \"cat\" and gnn.out_channels != last_hidden_channels:\n",
    "                raise ValueError(f\"All GNNs must output the same number of channels if {mode} is used!\")\n",
    "            last_hidden_channels = gnn.out_channels\n",
    "\n",
    "        # Konstruowanie warstw liniowych\n",
    "        self._build_lin_layers(lin_sizes, conv_out_total_size, last_hidden_channels)\n",
    "\n",
    "    def _build_lin_layers(self, lin_sizes: List[int], conv_out_total_size: int, last_hidden_channels: int):\n",
    "        \"\"\"Funkcja konstruująca warstwy liniowe.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        lin_sizes : List[int]\n",
    "            Lista wielkości warstw liniowych.\n",
    "        conv_out_total_size : int\n",
    "            Całkowita wielkość wyjścia z warstw splotowych (po konkatenacji).\n",
    "        last_hidden_channels : int\n",
    "            Ostatnia liczba kanałów w warstwie splotowej.\n",
    "        \"\"\"\n",
    "        lin_input_size = conv_out_total_size if self.jk_mode == \"cat\" else last_hidden_channels\n",
    "        all_lin_sizes = [lin_input_size, *lin_sizes]\n",
    "        lin_layers = []\n",
    "        for lin_in, lin_out in zip(all_lin_sizes[:-1], all_lin_sizes[1:]):\n",
    "            if self.dropout > 0:\n",
    "                lin_layers.append(th.nn.Dropout(self.dropout))\n",
    "            lin_layers.append(th.nn.Linear(lin_in, lin_out))\n",
    "            lin_layers.append(th.nn.ReLU())\n",
    "        self.lin_layers = th.nn.Sequential(*lin_layers)\n",
    "\n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        conv_outs = []\n",
    "        last_out = x\n",
    "        # Dla każdej z k warstw splotu\n",
    "        for conv in self.convs:\n",
    "            # Przekaż ostatnie znane reprezentacje wierzchołków...\n",
    "            # ..i utwórz nowe\n",
    "            last_out = conv(last_out, edge_index)\n",
    "            \n",
    "            # Zachowaj nowe reprezentacje\n",
    "            conv_outs.append(last_out)\n",
    "        \n",
    "        # Połącz reprezentacje wierzchołków, ze wszystkich iteracji.\n",
    "        jk_input = self._prepare_jk(conv_outs)\n",
    "        # Dokonaj predykcji\n",
    "        lin_out = self.lin_layers(jk_input)\n",
    "        return lin_out\n",
    "\n",
    "    def _prepare_jk(self, conv_outs: List[th.Tensor]) -> th.Tensor:\n",
    "        # Konkatenacja poprzednich reprezentacji\n",
    "        if self.jk_mode == \"cat\":\n",
    "            conv_out_concat = th.cat(conv_outs, dim=1)\n",
    "        # Uśrednianie poprzednich reprezentacji\n",
    "        elif self.jk_mode == \"mean\":\n",
    "            conv_out_concat = th.mean(th.stack(conv_outs), dim=0)\n",
    "        # Wybór max. z poprzednich reprezentacji\n",
    "        elif self.jk_mode == \"max\":\n",
    "            conv_out_concat = th.max(th.stack(conv_outs), dim=0).values\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown JK mode {self.jk_mode}\")\n",
    "        return conv_out_concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------+--------------------------+----------------+----------+\n",
      "| Layer                    | Input Shape              | Output Shape   | #Param   |\n",
      "|--------------------------+--------------------------+----------------+----------|\n",
      "| JKModel                  | [2708, 1433], [2, 10556] | [2708, 64]     | 457,408  |\n",
      "| ├─(convs)ModuleList      | --                       | --             | 432,768  |\n",
      "| │    └─(0)SAGEConv       | [2708, 1433], [2, 10556] | [2708, 128]    | 366,976  |\n",
      "| │    └─(1)SAGEConv       | [2708, 128], [2, 10556]  | [2708, 128]    | 32,896   |\n",
      "| │    └─(2)SAGEConv       | [2708, 128], [2, 10556]  | [2708, 128]    | 32,896   |\n",
      "| ├─(lin_layers)Sequential | [2708, 384]              | [2708, 64]     | 24,640   |\n",
      "| │    └─(0)Dropout        | [2708, 384]              | [2708, 384]    | --       |\n",
      "| │    └─(1)Linear         | [2708, 384]              | [2708, 64]     | 24,640   |\n",
      "| │    └─(2)ReLU           | [2708, 64]               | [2708, 64]     | --       |\n",
      "+--------------------------+--------------------------+----------------+----------+\n",
      "torch.Size([128, 1433])\n",
      "torch.Size([128])\n",
      "torch.Size([128, 1433])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([128])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([128])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([64, 384])\n",
      "torch.Size([64])\n"
     ]
    }
   ],
   "source": [
    "jk_reimpl = JKModel(\n",
    "    [\n",
    "        pyg.nn.SAGEConv(dim_in, out_channels=128, aggr='mean'),\n",
    "        pyg.nn.SAGEConv(128, out_channels=128, aggr='mean'),\n",
    "        pyg.nn.SAGEConv(128, out_channels=128, aggr='mean')\n",
    "    ],\n",
    "    lin_sizes=[64],\n",
    "    mode=\"cat\"\n",
    ")\n",
    "\n",
    "with th.no_grad():\n",
    "    out = jk_reimpl(cora['x'], cora['edge_index'])\n",
    "\n",
    "print(pyg.nn.summary(jk_reimpl, cora['x'], cora['edge_index']))\n",
    "\n",
    "for param in jk_reimpl.parameters():\n",
    "    print(param.size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Odpowiednik w PyG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dla modeli automatycznych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "jk_gnn = pyg.nn.GraphSAGE(\n",
    "    dim_in, \n",
    "    hidden_channels=128,\n",
    "    out_channels=64, \n",
    "    aggr='mean',\n",
    "    jk='cat',\n",
    "    num_layers=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2708, 64])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    pyg.seed_everything(123)\n",
    "    out = jk_gnn(cora['x'], cora['edge_index'])\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([128, 1433])\n",
      "torch.Size([128])\n",
      "torch.Size([128, 1433])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([128])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([128])\n",
      "torch.Size([128, 128])\n",
      "torch.Size([64, 384])\n",
      "torch.Size([64])\n"
     ]
    }
   ],
   "source": [
    "for param in jk_gnn.parameters():\n",
    "    print(param.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------------+--------------------------+----------------+----------+\n",
      "| Layer                  | Input Shape              | Output Shape   | #Param   |\n",
      "|------------------------+--------------------------+----------------+----------|\n",
      "| GraphSAGE              | [2708, 1433], [2, 10556] | [2708, 64]     | 457,408  |\n",
      "| ├─(dropout)Dropout     | [2708, 128]              | [2708, 128]    | --       |\n",
      "| ├─(act)ReLU            | [2708, 128]              | [2708, 128]    | --       |\n",
      "| ├─(convs)ModuleList    | --                       | --             | 432,768  |\n",
      "| │    └─(0)SAGEConv     | [2708, 1433], [2, 10556] | [2708, 128]    | 366,976  |\n",
      "| │    └─(1)SAGEConv     | [2708, 128], [2, 10556]  | [2708, 128]    | 32,896   |\n",
      "| │    └─(2)SAGEConv     | [2708, 128], [2, 10556]  | [2708, 128]    | 32,896   |\n",
      "| ├─(norms)ModuleList    | --                       | --             | --       |\n",
      "| │    └─(0)Identity     | [2708, 128]              | [2708, 128]    | --       |\n",
      "| │    └─(1)Identity     | [2708, 128]              | [2708, 128]    | --       |\n",
      "| │    └─(2)Identity     | [2708, 128]              | [2708, 128]    | --       |\n",
      "| ├─(jk)JumpingKnowledge |                          | [2708, 384]    | --       |\n",
      "| ├─(lin)Linear          | [2708, 384]              | [2708, 64]     | 24,640   |\n",
      "+------------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg.nn.summary(jk_gnn, cora['x'], cora['edge_index']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Dla modeli budowanych samodzielnie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2708, 64])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class GNNWithJK(th.nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "        self, \n",
    "        gnns: List[pyg.nn.MessagePassing], \n",
    "        lin_sizes: Tuple[int, ...], \n",
    "        jk_mode: Literal[\"cat\", \"max\", \"mean\"] = \"cat\"):\n",
    "        super().__init__()\n",
    "        self.gnns = th.nn.ModuleList(gnns)\n",
    "        self.jk = pyg.nn.JumpingKnowledge(mode=jk_mode)\n",
    "        lin_layers = []\n",
    "        for dim_in, dim_out in zip(lin_sizes[:-1], lin_sizes[1:]):\n",
    "            lin_layers.append(th.nn.Linear(dim_in, dim_out))\n",
    "        self.lin = th.nn.Sequential(*lin_layers)\n",
    "\n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor):\n",
    "        conv_outs = []\n",
    "        last_out = x\n",
    "        for conv in self.gnns:\n",
    "            last_out = conv(last_out, edge_index).relu()\n",
    "            conv_outs.append(last_out)\n",
    "        jk_out = self.jk(conv_outs)\n",
    "        return self.lin(jk_out)\n",
    "\n",
    "\n",
    "with th.no_grad():\n",
    "    gnn_with_jk = GNNWithJK(\n",
    "        [\n",
    "            pyg.nn.SAGEConv(dim_in, out_channels=256, aggr='mean'),\n",
    "            pyg.nn.SAGEConv(256, out_channels=64, aggr='mean'),\n",
    "            pyg.nn.SAGEConv(64, out_channels=32, aggr='mean')\n",
    "        ],\n",
    "        lin_sizes=[352, 64],\n",
    "        jk_mode=\"cat\"\n",
    "    )\n",
    "    out = gnn_with_jk(cora['x'], cora['edge_index'])\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Szkolenie sieci z mechanizmem Jumping Knowledge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6cf0f7471d6e4070858fc13e911cffee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>test</th>\n",
       "      <th>val</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>f1</th>\n",
       "      <td>0.786201</td>\n",
       "      <td>0.751608</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>roc_auc</th>\n",
       "      <td>0.954501</td>\n",
       "      <td>0.942515</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             test       val\n",
       "f1       0.786201  0.751608\n",
       "roc_auc  0.954501  0.942515"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "gnn_with_jk =  GNNWithJK(\n",
    "        [\n",
    "            pyg.nn.SAGEConv(dim_in, out_channels=32, aggr='mean'),\n",
    "            pyg.nn.SAGEConv(32, out_channels=32, aggr='mean'),\n",
    "            pyg.nn.SAGEConv(32, out_channels=32, aggr='mean')\n",
    "        ],\n",
    "        lin_sizes=[96, ncls],\n",
    "        jk_mode=\"cat\"\n",
    "    )\n",
    "\n",
    "scores_gnn_with_jk, gnn_with_jk_train_losses, gnn_with_jk_test_losses = train_and_test_model(gnn_with_jk, cora, epochs=10)\n",
    "pd.DataFrame(scores_gnn_with_jk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f2aff3a8260>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 32976 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(gnn_with_jk_train_losses, label='train')\n",
    "plt.plot(gnn_with_jk_test_losses, label='val')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Eksperyment - porównanie wielu modeli\n",
    "\n",
    "Sekcja poniżej prezentuje kompleksowy eksperyment, porównujacy różne modele GNN na zbiorze danych Cora.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Przygotowanie składowych do eksperymentu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlGNN(pl.LightningModule):\n",
    "    \n",
    "    def __init__(\n",
    "        self, \n",
    "        gnn: th.nn.Module, \n",
    "        data: pyg.data.Data, \n",
    "        lr: float = 0.01):\n",
    "        super().__init__()\n",
    "        self.gnn = gnn\n",
    "        self.data = data\n",
    "        self.lr = lr\n",
    "        self.loss = th.nn.CrossEntropyLoss()\n",
    "        self.train_f1 = tm.F1Score(num_classes=ncls, average='macro', task='multiclass')\n",
    "        self.val_f1 = tm.F1Score(num_classes=ncls, average='macro', task='multiclass')\n",
    "        self.test_f1 = tm.F1Score(num_classes=ncls, average='macro', task='multiclass')\n",
    "        self.train_roc_auc = tm.AUROC(num_classes=ncls, average='macro', task='multiclass')\n",
    "        self.val_roc_auc = tm.AUROC(num_classes=ncls, average='macro', task='multiclass')\n",
    "        self.test_roc_auc = tm.AUROC(num_classes=ncls, average='macro', task='multiclass')\n",
    "        \n",
    "    def forward(self, x, edge_index):\n",
    "        return self.gnn(x, edge_index)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, edge_index, y = batch.x, batch.edge_index, batch.y\n",
    "        batch_size = batch.train_mask.sum()\n",
    "        out = self(x, edge_index)\n",
    "        loss = self.loss(out[batch.train_mask], y[batch.train_mask])\n",
    "        self.log('train_loss', loss, batch_size=batch_size)\n",
    "        f1 = self.train_f1(out[batch.train_mask], y[batch.train_mask])\n",
    "        roc_auc = self.train_roc_auc(out[batch.train_mask], y[batch.train_mask])\n",
    "        self.log('train_f1', f1, batch_size=batch_size)\n",
    "        self.log('train_roc_auc', roc_auc, batch_size=batch_size)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, edge_index, y = batch.x, batch.edge_index, batch.y\n",
    "        out = self(x, edge_index)\n",
    "        batch_size = batch.val_mask.sum()\n",
    "        loss = self.loss(out[batch.val_mask], y[batch.val_mask])\n",
    "        self.log('val_loss', loss, batch_size=batch_size)\n",
    "        f1 = self.val_f1(out[batch.val_mask], y[batch.val_mask])\n",
    "        roc_auc = self.val_roc_auc(out[batch.val_mask], y[batch.val_mask])\n",
    "        self.log('val_f1', f1, batch_size=batch_size)\n",
    "        self.log('val_roc_auc', roc_auc, batch_size=batch_size)\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, edge_index, y = batch.x, batch.edge_index, batch.y\n",
    "        out = self(x, edge_index)\n",
    "        batch_size = batch.test_mask.sum()\n",
    "        loss = self.loss(out[batch.test_mask], y[batch.test_mask])\n",
    "        self.log('test_loss', loss, batch_size=batch_size)\n",
    "        f1 = self.test_f1(out[batch.test_mask], y[batch.test_mask])\n",
    "        roc_auc = self.test_roc_auc(out[batch.test_mask], y[batch.test_mask])\n",
    "        self.log('test_f1', f1, batch_size=batch_size)\n",
    "        self.log('test_roc_auc', roc_auc, batch_size=batch_size)\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return th.optim.Adam(self.parameters(), lr=self.lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = pyg.loader.NeighborLoader(\n",
    "    cora,\n",
    "    input_nodes=cora.train_mask,\n",
    "    num_neighbors=[16, 8],\n",
    "    batch_size=256,\n",
    "    shuffle=True,\n",
    "    subgraph_type=\"induced\"\n",
    ")\n",
    "\n",
    "val_loader = pyg.loader.NeighborLoader(\n",
    "    cora,\n",
    "    input_nodes=cora.val_mask,\n",
    "    num_neighbors=[16, 8],\n",
    "    batch_size=256,\n",
    "    shuffle=False,\n",
    "    subgraph_type=\"induced\"\n",
    ")\n",
    "\n",
    "test_loader = pyg.loader.NeighborLoader(\n",
    "    cora,\n",
    "    input_nodes=cora.test_mask,\n",
    "    num_neighbors=[16, 8],\n",
    "    batch_size=256,\n",
    "    shuffle=False,\n",
    "    subgraph_type=\"induced\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(raw_model: th.nn.Module, train_loader, val_loader, test_loader, lr: float = 0.01, epochs: int = 25):\n",
    "    pyg.seed_everything(123)\n",
    "    model = PlGNN(raw_model, cora, lr)\n",
    "    trainer = pl.Trainer(max_epochs=epochs, deterministic=True)\n",
    "    trainer.fit(model, train_loader, val_loader)\n",
    "    val_scores = trainer.test(model, val_loader)\n",
    "    test_scores = trainer.test(model, test_loader)\n",
    "    return val_scores, test_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testowanie modeli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {\n",
    "    \"SAGE\": pyg.nn.Sequential(\n",
    "        \"x, edge_index\",\n",
    "        [\n",
    "            (pyg.nn.SAGEConv(dim_in, 256, normalize=True), \"x, edge_index -> x\"),\n",
    "            (th.nn.ReLU(), \"x -> x\"),\n",
    "            (pyg.nn.SAGEConv(256, 64, normalize=True), \"x, edge_index -> x\"),\n",
    "            (th.nn.ReLU(), \"x -> x\"),\n",
    "            (pyg.nn.SAGEConv(64, ncls, normalize=True), \"x, edge_index -> x\"),\n",
    "        ],\n",
    "    ),\n",
    "    \"SAGE Multiagg\": pyg.nn.Sequential(\n",
    "        \"x, edge_index\",\n",
    "        [\n",
    "            (\n",
    "                pyg.nn.SAGEConv(dim_in, 256, normalize=True, aggr=[\"sum\", \"max\"]),\n",
    "                \"x, edge_index -> x\",\n",
    "            ),\n",
    "            (th.nn.ReLU(), \"x -> x\"),\n",
    "            (\n",
    "                pyg.nn.SAGEConv(256, 64, normalize=True, aggr=[\"sum\", \"max\"]),\n",
    "                \"x, edge_index -> x\",\n",
    "            ),\n",
    "            (th.nn.ReLU(), \"x -> x\"),\n",
    "            (\n",
    "                pyg.nn.SAGEConv(64, ncls, normalize=True, aggr=[\"sum\", \"max\"]),\n",
    "                \"x, edge_index -> x\",\n",
    "            ),\n",
    "        ],\n",
    "    ),\n",
    "    \"NGNN\": pyg.nn.Sequential(\n",
    "        \"x, edge_index\",\n",
    "        [\n",
    "            (\n",
    "                NGNN(pyg.nn.SAGEConv, dim_in, 512, 256, normalize=True, aggr=['sum', 'mean']),\n",
    "                \"x, edge_index -> x\",\n",
    "            ),\n",
    "            (NGNN(pyg.nn.SAGEConv, 256, 128, 64, normalize=True,  aggr=['sum', 'mean']), \"x, edge_index -> x\"),\n",
    "            (NGNN(pyg.nn.SAGEConv, 64, 64, ncls, normalize=True,  aggr=['sum', 'mean']), \"x, edge_index -> x\"),\n",
    "        ],\n",
    "    ),\n",
    "    \"JK\": GNNWithJK(\n",
    "        [\n",
    "            pyg.nn.SAGEConv(dim_in, out_channels=256, aggr=[\"sum\"], normalize=True),\n",
    "            pyg.nn.SAGEConv(256, out_channels=128, aggr=[\"sum\"], normalize=True),\n",
    "            pyg.nn.SAGEConv(128, out_channels=128, aggr=[\"sum\"], normalize=True),\n",
    "        ],\n",
    "        lin_sizes=[512, ncls],\n",
    "        jk_mode=\"cat\",\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_test_scores = {}\n",
    "model_val_scores = {}\n",
    "\n",
    "for model_name, model in tqdm(models.items()):\n",
    "    print(f\"Training {model_name}\")\n",
    "    val_scores, test_scores = train_model(model, train_loader, val_loader, test_loader)\n",
    "    model_val_scores[model_name] = val_scores\n",
    "    model_test_scores[model_name] = test_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Porównanie wyników modeli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_scores_dataframe(scores_dict: dict) -> pd.DataFrame:\n",
    "    # Initialize an empty list to store the rows of the DataFrame\n",
    "    rows = []\n",
    "\n",
    "    # Iterate over the models and their scores\n",
    "    for model_name, scores in scores_dict.items():\n",
    "        # Flatten the dictionary and add the model name as a separate column\n",
    "        for metric, value in scores[0].items():\n",
    "            rows.append({\"model_name\": model_name, \"metric\": metric, \"value\": value})\n",
    "\n",
    "    # Create a DataFrame from the list of rows\n",
    "    df = pd.DataFrame(rows)\n",
    "    return df\n",
    "\n",
    "# Build the DataFrame using the model_val_scores dictionary\n",
    "val_scores_df = build_scores_dataframe(model_val_scores)\n",
    "test_scores_df = build_scores_dataframe(model_test_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "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>model_name</th>\n",
       "      <th>metric</th>\n",
       "      <th>value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SAGE</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>1.345788</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SAGE</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.778440</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SAGE</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.945473</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SAGE Multiagg</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>1.325065</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SAGE Multiagg</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.789885</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>SAGE Multiagg</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.949324</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>NGNN</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>1.157342</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>NGNN</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.767158</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>NGNN</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.933804</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>JK</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>0.778452</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>JK</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.780067</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>JK</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.950470</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       model_name        metric     value\n",
       "0            SAGE     test_loss  1.345788\n",
       "1            SAGE       test_f1  0.778440\n",
       "2            SAGE  test_roc_auc  0.945473\n",
       "3   SAGE Multiagg     test_loss  1.325065\n",
       "4   SAGE Multiagg       test_f1  0.789885\n",
       "5   SAGE Multiagg  test_roc_auc  0.949324\n",
       "6            NGNN     test_loss  1.157342\n",
       "7            NGNN       test_f1  0.767158\n",
       "8            NGNN  test_roc_auc  0.933804\n",
       "9              JK     test_loss  0.778452\n",
       "10             JK       test_f1  0.780067\n",
       "11             JK  test_roc_auc  0.950470"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_scores_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "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>model_name</th>\n",
       "      <th>metric</th>\n",
       "      <th>value</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SAGE</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>1.350073</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SAGE</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.758473</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SAGE</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.944750</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SAGE Multiagg</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>1.328668</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SAGE Multiagg</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.766866</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>SAGE Multiagg</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.948835</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>NGNN</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>1.168216</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>NGNN</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.746500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>NGNN</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.929269</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>JK</td>\n",
       "      <td>test_loss</td>\n",
       "      <td>0.782826</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>JK</td>\n",
       "      <td>test_f1</td>\n",
       "      <td>0.750939</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>JK</td>\n",
       "      <td>test_roc_auc</td>\n",
       "      <td>0.946713</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       model_name        metric     value\n",
       "0            SAGE     test_loss  1.350073\n",
       "1            SAGE       test_f1  0.758473\n",
       "2            SAGE  test_roc_auc  0.944750\n",
       "3   SAGE Multiagg     test_loss  1.328668\n",
       "4   SAGE Multiagg       test_f1  0.766866\n",
       "5   SAGE Multiagg  test_roc_auc  0.948835\n",
       "6            NGNN     test_loss  1.168216\n",
       "7            NGNN       test_f1  0.746500\n",
       "8            NGNN  test_roc_auc  0.929269\n",
       "9              JK     test_loss  0.782826\n",
       "10             JK       test_f1  0.750939\n",
       "11             JK  test_roc_auc  0.946713"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_scores_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_scores_df.round(4).to_excel(\"val_scores.xlsx\")\n",
    "test_scores_df.round(4).to_excel(\"test_scores.xlsx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "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>Number of Parameters</th>\n",
       "      <th>Increase [%]</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>NGNN</th>\n",
       "      <td>2452423</td>\n",
       "      <td>219.46</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SAGE Multiagg</th>\n",
       "      <td>1151367</td>\n",
       "      <td>49.98</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>JK</th>\n",
       "      <td>836103</td>\n",
       "      <td>8.91</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>SAGE</th>\n",
       "      <td>767687</td>\n",
       "      <td>0.00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               Number of Parameters  Increase [%]\n",
       "NGNN                        2452423        219.46\n",
       "SAGE Multiagg               1151367         49.98\n",
       "JK                           836103          8.91\n",
       "SAGE                         767687          0.00"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def count_parameters(model: th.nn.Module) -> int:\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "model_parameters = {name: count_parameters(model) for name, model in models.items()}\n",
    "model_parameters_df = pd.DataFrame.from_dict(model_parameters, orient='index', columns=['Number of Parameters'])\n",
    "model_parameters_df['Increase [%]'] = (\n",
    "    (model_parameters_df['Number of Parameters'] - model_parameters_df.loc['SAGE', 'Number of Parameters']) \n",
    "    / model_parameters_df.loc['SAGE', 'Number of Parameters'] * 100\n",
    ")\n",
    "display(model_parameters_df.sort_values(by='Number of Parameters', ascending=False).round(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_parameters_df.sort_values(by='Number of Parameters', ascending=False).round(2).to_excel(\"model_parameters.xlsx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_f336e_row0_col1, #T_f336e_row0_col2 {\n",
       "  background-color: #006837;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_f336e_row1_col1 {\n",
       "  background-color: #ddf191;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_f336e_row1_col2 {\n",
       "  background-color: #69be63;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_f336e_row2_col1 {\n",
       "  background-color: #f57748;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_f336e_row2_col2 {\n",
       "  background-color: #1e9a51;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_f336e_row3_col1, #T_f336e_row3_col2 {\n",
       "  background-color: #a50026;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_f336e\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_f336e_level0_col0\" class=\"col_heading level0 col0\" >test_loss</th>\n",
       "      <th id=\"T_f336e_level0_col1\" class=\"col_heading level0 col1\" >test_f1</th>\n",
       "      <th id=\"T_f336e_level0_col2\" class=\"col_heading level0 col2\" >test_roc_auc</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_f336e_level0_row0\" class=\"row_heading level0 row0\" >SAGE Multiagg</th>\n",
       "      <td id=\"T_f336e_row0_col0\" class=\"data row0 col0\" >1.328700</td>\n",
       "      <td id=\"T_f336e_row0_col1\" class=\"data row0 col1\" >0.766900</td>\n",
       "      <td id=\"T_f336e_row0_col2\" class=\"data row0 col2\" >0.948800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_f336e_level0_row1\" class=\"row_heading level0 row1\" >SAGE</th>\n",
       "      <td id=\"T_f336e_row1_col0\" class=\"data row1 col0\" >1.350100</td>\n",
       "      <td id=\"T_f336e_row1_col1\" class=\"data row1 col1\" >0.758500</td>\n",
       "      <td id=\"T_f336e_row1_col2\" class=\"data row1 col2\" >0.944800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_f336e_level0_row2\" class=\"row_heading level0 row2\" >JK</th>\n",
       "      <td id=\"T_f336e_row2_col0\" class=\"data row2 col0\" >0.782800</td>\n",
       "      <td id=\"T_f336e_row2_col1\" class=\"data row2 col1\" >0.750900</td>\n",
       "      <td id=\"T_f336e_row2_col2\" class=\"data row2 col2\" >0.946700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_f336e_level0_row3\" class=\"row_heading level0 row3\" >NGNN</th>\n",
       "      <td id=\"T_f336e_row3_col0\" class=\"data row3 col0\" >1.168200</td>\n",
       "      <td id=\"T_f336e_row3_col1\" class=\"data row3 col1\" >0.746500</td>\n",
       "      <td id=\"T_f336e_row3_col2\" class=\"data row3 col2\" >0.929300</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f2af42a1af0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "#T_d2b18_row0_col1, #T_d2b18_row1_col2 {\n",
       "  background-color: #006837;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_d2b18_row0_col2 {\n",
       "  background-color: #128a49;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "#T_d2b18_row1_col1 {\n",
       "  background-color: #e5f49b;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_d2b18_row2_col1 {\n",
       "  background-color: #fffdbc;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_d2b18_row2_col2 {\n",
       "  background-color: #a5d86a;\n",
       "  color: #000000;\n",
       "}\n",
       "#T_d2b18_row3_col1, #T_d2b18_row3_col2 {\n",
       "  background-color: #a50026;\n",
       "  color: #f1f1f1;\n",
       "}\n",
       "</style>\n",
       "<table id=\"T_d2b18\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th id=\"T_d2b18_level0_col0\" class=\"col_heading level0 col0\" >val_loss</th>\n",
       "      <th id=\"T_d2b18_level0_col1\" class=\"col_heading level0 col1\" >val_f1</th>\n",
       "      <th id=\"T_d2b18_level0_col2\" class=\"col_heading level0 col2\" >val_roc_auc</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_d2b18_level0_row0\" class=\"row_heading level0 row0\" >SAGE Multiagg</th>\n",
       "      <td id=\"T_d2b18_row0_col0\" class=\"data row0 col0\" >1.325100</td>\n",
       "      <td id=\"T_d2b18_row0_col1\" class=\"data row0 col1\" >0.789900</td>\n",
       "      <td id=\"T_d2b18_row0_col2\" class=\"data row0 col2\" >0.949300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_d2b18_level0_row1\" class=\"row_heading level0 row1\" >JK</th>\n",
       "      <td id=\"T_d2b18_row1_col0\" class=\"data row1 col0\" >0.778500</td>\n",
       "      <td id=\"T_d2b18_row1_col1\" class=\"data row1 col1\" >0.780100</td>\n",
       "      <td id=\"T_d2b18_row1_col2\" class=\"data row1 col2\" >0.950500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_d2b18_level0_row2\" class=\"row_heading level0 row2\" >SAGE</th>\n",
       "      <td id=\"T_d2b18_row2_col0\" class=\"data row2 col0\" >1.345800</td>\n",
       "      <td id=\"T_d2b18_row2_col1\" class=\"data row2 col1\" >0.778400</td>\n",
       "      <td id=\"T_d2b18_row2_col2\" class=\"data row2 col2\" >0.945500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_d2b18_level0_row3\" class=\"row_heading level0 row3\" >NGNN</th>\n",
       "      <td id=\"T_d2b18_row3_col0\" class=\"data row3 col0\" >1.157300</td>\n",
       "      <td id=\"T_d2b18_row3_col1\" class=\"data row3 col1\" >0.767200</td>\n",
       "      <td id=\"T_d2b18_row3_col2\" class=\"data row3 col2\" >0.933800</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7f2af4361700>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 37640 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 37828 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Zbieranie wyników z wszystkich modeli\n",
    "\n",
    "# Tworzenie tabeli wyników dla testów\n",
    "test_scores_df = pd.DataFrame(model_test_scores).T\n",
    "test_scores_df = pd.json_normalize(test_scores_df[0])\n",
    "test_scores_df.index = model_test_scores.keys()\n",
    "test_scores_df = test_scores_df.sort_values(by=\"test_f1\", ascending=False)\n",
    "\n",
    "# Tworzenie tabeli wyników dla walidacji\n",
    "val_scores_df = pd.DataFrame(model_val_scores).T\n",
    "val_scores_df = pd.json_normalize(val_scores_df[0])\n",
    "val_scores_df.index = model_val_scores.keys()\n",
    "val_scores_df.columns = [col.replace(\"test_\", \"val_\") for col in val_scores_df.columns]\n",
    "val_scores_df = val_scores_df.sort_values(by=\"val_f1\", ascending=False)\n",
    "\n",
    "# Wyświetlanie tabeli wyników z podświetleniem najlepszych wyników\n",
    "display(test_scores_df.round(4).style.background_gradient(subset=['test_f1', 'test_roc_auc'], cmap='RdYlGn', axis=0))\n",
    "display(val_scores_df.round(4).style.background_gradient(subset=['val_f1', 'val_roc_auc'], cmap='RdYlGn', axis=0))\n",
    "\n",
    "\n",
    "# Wizualizacja wyników testów\n",
    "test_scores_melted = test_scores_df.reset_index().melt(id_vars=\"index\", var_name=\"Metric\", value_name=\"Value\")\n",
    "test_scores_melted = test_scores_melted.rename(columns={\"index\": \"Model\"})\n",
    "\n",
    "plt.figure(figsize=(12, 6))\n",
    "sns.barplot(data=test_scores_melted, x=\"Metric\", y=\"Value\", hue=\"Model\")\n",
    "plt.title(\"Porównanie wyników testów modeli GNN na zbiorze danych Cora\")\n",
    "plt.xlabel(\"Metryka\")\n",
    "plt.ylabel(\"Wartość\")\n",
    "plt.legend(loc='best')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Wizualizacja wyników walidacji\n",
    "val_scores_melted = val_scores_df.reset_index().melt(id_vars=\"index\", var_name=\"Metric\", value_name=\"Value\")\n",
    "val_scores_melted = val_scores_melted.rename(columns={\"index\": \"Model\"})\n",
    "\n",
    "plt.figure(figsize=(12, 6))\n",
    "sns.barplot(data=val_scores_melted, x=\"Metric\", y=\"Value\", hue=\"Model\")\n",
    "plt.title(\"Porównanie wyników walidacji modeli GNN na zbiorze danych Cora\")\n",
    "plt.xlabel(\"Metryka\")\n",
    "plt.ylabel(\"Wartość\")\n",
    "plt.legend(loc='best')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
