{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "W tym notebooku przyjrzymy się, jak łatwo jest łączyć ze sobą arbitralnie wybrane elementy modyfikujące poszczególne fazy przekazywania wiadomości w grafowych sieciach neuronowych. Podobnie, jak poprzednio wykorzystamy bibliotekę `PyTorch Geometric` oraz zbiór danych `Cora`.\n",
    "\n",
    "Zaimplementujemy warstwę splotu, która:\n",
    "\n",
    "1. W fazie tworzenia wiadomości - normalizuje reprezentację wierzchołka za pomocą pierwiastka z iloczynu stopni wierzchołków sąsiednich oraz przekształca go macierzą wag. Mechanizm podobny do sieci GCN.\n",
    "2. W fazie agregacji sąsiadów - stosuje sumowanie macierzy reprezentacji sąsiadów. Mechanizm podobny do sieci GraphSAGE.\n",
    "3. W fazie aktualizacji reprezentacji wykorzystuje sieć MLP - mechanizm podobny do sieci GIN."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch as th\n",
    "import torch.nn as nn\n",
    "import torch_geometric as pyg\n",
    "import torch_geometric.nn as pyg_nn\n",
    "import torch_geometric.transforms as pyg_t\n",
    "import torch_geometric.datasets as pyg_datasets\n",
    "import torch_geometric.data as pyg_data\n",
    "import torch_geometric.utils as pyg_utils\n",
    "\n",
    "import utilities as utils\n",
    "from base_models import NodesClassifierModel, NaiveGraphConv\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[2708, 1433], edge_index=[2, 10556], y=[2708], train_mask=[2708], val_mask=[2708], test_mask=[2708])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = pyg_datasets.Planetoid(root=Path(os.pardir, \"../data\", \"Planetoid\"), name='Cora', transform=pyg_t.NormalizeFeatures())\n",
    "cora = dataset[0]\n",
    "cora"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementacja warstwy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SampleConv(pyg_nn.MessagePassing):\n",
    "    def __init__(\n",
    "        self, \n",
    "        dim_in: int, \n",
    "        dim_out: int, \n",
    "        mlp: nn.Module,\n",
    "        aggr: str = 'sum'):\n",
    "        super().__init__(aggr=aggr)\n",
    "        self.dim_in = dim_in\n",
    "        self.dim_out = dim_out\n",
    "        self.W = nn.Linear(dim_in, dim_out)\n",
    "        self.mlp = mlp\n",
    "        self.reset_parameters()\n",
    "        \n",
    "    def reset_parameters(self):\n",
    "        super().reset_parameters()\n",
    "        pyg_nn.inits.glorot(self.W.weight)\n",
    "        pyg_nn.inits.zeros(self.W.bias)\n",
    "        pyg_nn.inits.reset(self.mlp)\n",
    "        \n",
    "        \n",
    "    def message(self, x_j, norm_coef: th.Tensor):\n",
    "        # Wiadomość od sąsiada - poprzednia reprezentacja * współczynnik normalizacji\n",
    "        h_u_norm = x_j * norm_coef\n",
    "        return super().message(h_u_norm)\n",
    "    \n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor):\n",
    "        if not pyg_utils.contains_self_loops(edge_index):\n",
    "            edge_index, _ = pyg_utils.add_self_loops(edge_index, num_nodes=x.size(0))\n",
    "        # Krok 1: wyliczenie współczynników normalizacji\n",
    "        norm_coef = self.calculate_norm_coefs(x, edge_index)\n",
    "        \n",
    "        hw = self.W(x)\n",
    "        \n",
    "        # Rozpoczęcie przekazywania wiadomości\n",
    "        hv = self.propagate(edge_index, x=hw, norm_coef=norm_coef)\n",
    "        # Krok 3: przekształcenie przez sieć MLP\n",
    "        hv_mlp = self.mlp(hv)\n",
    "        return hv_mlp\n",
    "\n",
    "    def calculate_norm_coefs(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja wyliczająca współczynniki normalizacji, wykorzystujące stopnie wierzchołków \n",
    "        u i v.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : th.Tensor\n",
    "            Macierz cech wierzchołków.\n",
    "        edge_index : th.Tensor\n",
    "            Macierz sąsiedztwa grafu.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Współczynniki normalizacji.\n",
    "        \"\"\"\n",
    "        degree_u = pyg_utils.degree(edge_index[0], x.size(0), dtype=x.dtype)\n",
    "        degree_v = pyg_utils.degree(edge_index[1], x.size(0), dtype=x.dtype)\n",
    "        norm_coef = 1 / th.sqrt(degree_u[edge_index[0]] * degree_v[edge_index[1]])\n",
    "        norm_coef = norm_coef.unsqueeze(1)\n",
    "        return norm_coef"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sprawdzenie, czy implementacja działa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2708, 7])\n"
     ]
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    mlp = nn.Sequential(nn.Linear(16, 16), nn.ReLU(), nn.Linear(16, 7))\n",
    "    conv = SampleConv(cora.num_features, 16, mlp)\n",
    "    out = conv(cora.x, cora.edge_index)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Klasyfikacja wierzchołków\n",
    "\n",
    "Wykorzystamy tak napisaną sieć do klasyfikacji wierzchołków."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "mlp1 = nn.Sequential(\n",
    "    nn.Linear(256, 128),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(128, 64),\n",
    "    nn.ReLU()\n",
    ")\n",
    "\n",
    "mlp2 = nn.Sequential(\n",
    "    nn.Linear(64, 32),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(32, 16),\n",
    ")\n",
    "\n",
    "sample_conv1 = SampleConv(dim_in=cora.num_features, dim_out=256, mlp=mlp1, aggr='sum')\n",
    "sample_conv2 = SampleConv(dim_in=64, dim_out=64, mlp=mlp2, aggr='sum')\n",
    "\n",
    "classifier = NodesClassifierModel(\n",
    "    gnn_convs=(sample_conv1, sample_conv2), \n",
    "    lin_sizes=[16, 7])\n",
    "\n",
    "optim = th.optim.Adam(classifier.parameters(), lr=0.01)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "with th.no_grad():\n",
    "    classifier.eval()\n",
    "    yhat_before_training = classifier(cora.x, cora.edge_index)\n",
    "\n",
    "epoch_losses = []\n",
    "val_losses = []\n",
    "\n",
    "for epoch in tqdm(range(15)):\n",
    "    classifier.train()\n",
    "    optim.zero_grad()\n",
    "    out = classifier(cora.x, cora.edge_index)\n",
    "    if out.isnan().any():\n",
    "        raise ValueError(\"NaNs in the output.\")\n",
    "    loss = loss_fn(out[cora.train_mask], cora.y[cora.train_mask])\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    epoch_losses.append(loss.item())\n",
    "\n",
    "    classifier.eval()\n",
    "    with th.no_grad():\n",
    "        val_loss = loss_fn(out[cora.val_mask], cora.y[cora.val_mask])\n",
    "        val_losses.append(val_loss.item())\n",
    "\n",
    "with th.no_grad():\n",
    "    classifier.eval()\n",
    "    yhat_after_training = classifier(cora.x, cora.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 41692 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "utils.plot_results(\n",
    "    epoch_losses,\n",
    "    val_losses,\n",
    "    plot_name=\"naive_gin.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu Przykładowa sieć przed trenowaniem:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.00      0.00      0.00        61\n",
      "           1       0.00      0.00      0.00        36\n",
      "           2       0.00      0.00      0.00        78\n",
      "           3       0.00      0.00      0.00       158\n",
      "           4       0.00      0.00      0.00        81\n",
      "           5       0.11      1.00      0.20        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.11       500\n",
      "   macro avg       0.02      0.14      0.03       500\n",
      "weighted avg       0.01      0.11      0.02       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu Przykładowa sieć po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.61      0.75      0.68        61\n",
      "           1       0.00      0.00      0.00        36\n",
      "           2       0.57      1.00      0.73        78\n",
      "           3       0.00      0.00      0.00       158\n",
      "           4       0.36      0.86      0.51        81\n",
      "           5       0.62      0.60      0.61        57\n",
      "           6       0.54      0.69      0.61        29\n",
      "\n",
      "    accuracy                           0.50       500\n",
      "   macro avg       0.39      0.56      0.45       500\n",
      "weighted avg       0.32      0.50      0.38       500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "utils.val_scores_before_after_train(\n",
    "    cora.y[cora.val_mask],\n",
    "    yhat_before_training[cora.val_mask],\n",
    "    yhat_after_training[cora.val_mask],\n",
    "    \"Przykładowa sieć\"\n",
    ")"
   ]
  }
 ],
 "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
}
