{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "W tym notatniku przeanalizujemy splot GIN (ang. *Graph Isomorphism Network*). Jest to warstwa splotu, która uzyskuje (dzięki wykorzystaniu nieliniowego mapowania siecią MLP) możliwości generalizacji porównywalne z testem izomorfizmu grafów Weisfeilera-Lehmana.\n",
    "\n",
    "Oryginalna publikacja:\n",
    "\n",
    "> Xu, K., Hu, W., Leskovec, J., & Jegelka, S. (2018). How powerful are graph neural networks?. arXiv preprint arXiv:1810.00826."
   ]
  },
  {
   "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 networkx as nx\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",
    "\n",
    "from pathlib import Path\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import utilities as utils\n",
    "from base_models import NodesClassifierModel, NaiveGraphConv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Relacja z testem Weisfeilera-Lehmana"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Przygotowanie grafów, analogicznych do tych z przykładu\n",
    "G1 = nx.Graph()\n",
    "G1.add_node('v1', feature=np.array([1., 1.]))\n",
    "G1.add_node('u1', feature=np.array([1., 10.]))\n",
    "G1.add_node('u2', feature=np.array([1., 10.]))\n",
    "G1.add_edge('v1', 'u1')\n",
    "G1.add_edge('v1', 'u2')\n",
    "\n",
    "\n",
    "G2 = nx.Graph()\n",
    "G2.add_node('v2', feature=np.array([1., 1.]))\n",
    "G2.add_node('u3', feature=np.array([1., 10.]))\n",
    "G2.add_node('u4', feature=np.array([1., 10.]))\n",
    "G2.add_node('u5', feature=np.array([1., 10.]))\n",
    "G2.add_edge('v2', 'u3')\n",
    "G2.add_edge('v2', 'u4')\n",
    "G2.add_edge('v2', 'u5')\n",
    "\n",
    "G3 = nx.Graph()\n",
    "G3.add_node('v3', feature=np.array([1., 1.]))\n",
    "G3.add_node('u6', feature=np.array([2., 12.]))\n",
    "G3.add_node('u7', feature=np.array([-1., -2.]))\n",
    "G3.add_edge('v3', 'u6')\n",
    "G3.add_edge('v3', 'u7')\n",
    "\n",
    "G4 = nx.Graph()\n",
    "G4.add_node('v4', feature=np.array([1., 1.]))\n",
    "G4.add_node('u8', feature=np.array([3., 15.]))\n",
    "G4.add_node('u9', feature=np.array([0., -3.]))\n",
    "G4.add_node('u10', feature=np.array([-2., -3.]))\n",
    "G4.add_edge('v4', 'u8')\n",
    "G4.add_edge('v4', 'u9')\n",
    "G4.add_edge('v4', 'u10')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: '402141eb1a088465fa99fc74a1bae21a',\n",
       " 2: 'c2985fd875b9d512a94cce0b5cee693b',\n",
       " 3: '898442196603a1a612bf852b67a4ca73',\n",
       " 4: 'f607d4819a2aef3d1a73aab2d09fa11e'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Wyliczenie hashy dla grafów\n",
    "graph_hashes = {}\n",
    "for i, g in enumerate([G1, G2, G3, G4]):\n",
    "    graph_hashes[i+1] = nx.weisfeiler_lehman_graph_hash(g, node_attr='feature', iterations=5)\n",
    "graph_hashes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 6,
     "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": [
    "# GIN\n",
    "\n",
    "Postać formalna warstwy GIN jest relatywnie nieskomplikowana - opiera się na dwóch podstawowych elementach:\n",
    "1. Współczynniku udziału / wpływu poprzedniej reprezentacji wierzchołka $v$ na jego nową reprezentację - $\\epsilon$;\n",
    "2. Sieci MLP, które są używane do przekształcenia reprezentacji wierzchołka.\n",
    "\n",
    "Możemy zapisać warstwę GIN jako:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\mathbf{h}_{\\mathcal{N}(v)}^{(k)} &= \\sum_{u \\in \\mathcal{N}(v)}\\mathbf{h}_u^{(k-1)} &\\triangleright \\text{Agregacja wiadomości sąsiadów poprzez sumowanie}  \\\\[1em]\n",
    "\\mathbf{h}_v^{(k)} &= \\text{MLP}^{(k)}\\left(\\left(1+\\epsilon^{(k)}\\right) \\mathbf{h}_{v}^{(k-1)} + \\mathbf{h}_{\\mathcal{N}(v)}^{(k)} \\right) &\\triangleright \\text{Aktualizacja reprezentacji wierzchołka siecią MLP}  \\\\[1em]\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementacja naiwna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NaiveGIN(NaiveGraphConv):\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            mlp: nn.Module,\n",
    "            agg_fn=pyg_nn.SumAggregation(),\n",
    "            initial_eps: float = 0.1,\n",
    "            add_self_loops: bool = True):\n",
    "        dim_in = mlp[0].in_features\n",
    "        if isinstance(mlp[-1], nn.Linear):\n",
    "            dim_out = mlp[-1].out_features\n",
    "        else:\n",
    "            dim_out = mlp[-2].out_features\n",
    "        super().__init__(\n",
    "            dim_in, \n",
    "            dim_out, \n",
    "            aggr=agg_fn, \n",
    "            sigma=nn.Identity(), \n",
    "            add_self_loops=add_self_loops)\n",
    "        self.initial_eps = initial_eps\n",
    "        self.mlp = mlp\n",
    "        self.eps = nn.Parameter(th.tensor(initial_eps))\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.eps = nn.Parameter(th.tensor(self.initial_eps))\n",
    "        for layer in self.mlp.children():\n",
    "            if isinstance(layer, nn.Linear):\n",
    "                pyg_nn.inits.glorot(layer.weight)\n",
    "                pyg_nn.inits.zeros(layer.bias)\n",
    "\n",
    "    def message(self, v_neighbors: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja budująca wiadomości od sąsiadów, implementacja abstrakcyjnej\n",
    "        funkcji `message` z architektury MPNN. Przyjmuje jako wejście\n",
    "        macierz cech sąsiadów wierzchołka v:  `u in N(v)` i buduje z niej wiadomości.\n",
    "        W tym przypadku - wiadomości od sąsiadów są ich atrybutami, nie podlegają\n",
    "        żadnym przekształceniom.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        v_neighbors : th.Tensor\n",
    "            Macierz cech sąsiadów wierzchołka v: `N(v)`.\n",
    "            Wymiar: `|N(v)| x dim_in`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Macierz wiadomości od sąsiadów.\n",
    "            Wymiar: `|N(v)| x dim_in`.\n",
    "        \"\"\"\n",
    "        return v_neighbors\n",
    "\n",
    "    def aggregate(self, m_nv: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja agregująca wiadomości od sąsiadów, implementacja abstrakcyjnej\n",
    "        funkcji `aggregate` z architektury MPNN. Przyjmuje jako wejście\n",
    "        macierz wiadomości od sąsiadów wierzchołka v i agreguje je. W przypadku\n",
    "        splotu GIN - agregacja polega na zsumowaniu wiadomości od sąsiadów.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        m_nv : th.Tensor\n",
    "            Macierz wiadomości od sąsiadów.\n",
    "            Wymiar: `|N(v)| x dim_out`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zagregowane wiadomości. Wymiar: `1 x dim_in`.\n",
    "        \"\"\"\n",
    "        return self.aggr(m_nv, dim=0)\n",
    "\n",
    "    def update(self, v: th.Tensor, agg: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Implementacja funkcji `update` z architektury MPNN. Aktualizuje\n",
    "        reprezentacje wierzchołka v na podstawie zagregowanych wiadomości\n",
    "        od sąsiadów. W przypadku splotu GIN - wykorzystuje część informacji\n",
    "        wierzchołka v ((1+epsilon) * v) i dodaje do niego zagregowane wiadomości.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        v : th.Tensor\n",
    "            Aktualna reprezentacja wierzchołka v. Wymiar: `1 x dim_in`.\n",
    "        agg : th.Tensor\n",
    "            Zagregowane wiadomości sąsiadów. Wymiar: `1 x dim_out`.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zaktualizowana reprezentacja wierzchołka v. Wymiar: `1 x dim_out`.\n",
    "        \"\"\"\n",
    "        v_eps = (1+self.eps) * v\n",
    "        up_v = v_eps + agg\n",
    "        return up_v\n",
    "            \n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja przepływu danych w warstwie splotu grafowego. Rozpoczyna cykl\n",
    "        przekazywania wiadomości, dla każdego wierzchołka w grafie: v po kolei:\n",
    "        1. Znajdując sąsiadów;\n",
    "        2. Budując wiadomości od sąsiadów;\n",
    "        3. Agregując wiadomości od sąsiadów;\n",
    "        4. Aktualizując reprezentację wierzchołka.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : th.Tensor\n",
    "            Macierz cech wierzchołków. Wymiar: N x dim_in\n",
    "\n",
    "        edge_index : th.Tensor\n",
    "            Macierz sąsiedztwa grafu. Wymiar: 2 x E\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zaktualizowane cechy wierzchołków.\n",
    "        \"\"\"\n",
    "        h_v = super().forward(x, edge_index)\n",
    "        h_v = self.mlp(h_v)\n",
    "        return h_v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sprawdzenie, czy implementacja działa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2708, 32])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    mlp = nn.Sequential(\n",
    "        nn.Linear(cora.num_features, 128),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(128, 64),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(64, 32),\n",
    "        nn.ReLU()\n",
    "    )\n",
    "    naive_gat = NaiveGIN(mlp=mlp, initial_eps=0.5)\n",
    "    naive_gat.eval()\n",
    "    h = naive_gat(cora.x, cora.edge_index)\n",
    "h.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Klasyfikacja wierzchołków"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model klasyfikacji z naiwną implementacją GIN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "⚠️ Uwaga, ze względu na fakt, iż implementacja naiwna nie jest zoptymalizowana, może być konieczne dłuższe oczekiwanie na wyniki. W zależności od posiadanego sprzętu i procesora, całość może zająć chwilę. Jeśli czas oczekiwani jest za długi - warto zmniejszyć ilość epok.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "mlp1 = nn.Sequential(\n",
    "    nn.Linear(cora.num_features, 256),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(256, 128),\n",
    "    nn.ReLU()\n",
    ")\n",
    "\n",
    "mlp2 = nn.Sequential(\n",
    "    nn.Linear(128, 64),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(64, 32),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(32, 7),\n",
    ")\n",
    "\n",
    "naive_gin1 = NaiveGIN(mlp=mlp1, initial_eps=0.1)\n",
    "naive_gin2 = NaiveGIN(mlp=mlp2, initial_eps=0.1)\n",
    "\n",
    "naive_clf = NodesClassifierModel(\n",
    "    gnn_convs=(naive_gin1, naive_gin2), \n",
    "    lin_sizes=[])\n",
    "\n",
    "optim = th.optim.Adam(naive_clf.parameters(), lr=0.01)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "with th.no_grad():\n",
    "    naive_clf.eval()\n",
    "    yhat_naive_before_training = naive_clf(cora.x, cora.edge_index)\n",
    "\n",
    "\n",
    "epoch_losses_naive = []\n",
    "val_losses_naive = []\n",
    "\n",
    "for epoch in tqdm(range(10)):\n",
    "    naive_clf.train()\n",
    "    optim.zero_grad()\n",
    "    out = naive_clf(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_naive.append(loss.item())\n",
    "\n",
    "    naive_clf.eval()\n",
    "    with th.no_grad():\n",
    "        val_loss = loss_fn(out[cora.val_mask], cora.y[cora.val_mask])\n",
    "        val_losses_naive.append(val_loss.item())\n",
    "\n",
    "with th.no_grad():\n",
    "    naive_clf.eval()\n",
    "    yhat_naive_after_training = naive_clf(cora.x, cora.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu Naiwny GIN przed trenowaniem:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.16      0.21      0.18        61\n",
      "           1       0.00      0.00      0.00        36\n",
      "           2       0.27      0.55      0.36        78\n",
      "           3       0.37      0.25      0.30       158\n",
      "           4       0.00      0.00      0.00        81\n",
      "           5       0.12      0.32      0.17        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.23       500\n",
      "   macro avg       0.13      0.19      0.14       500\n",
      "weighted avg       0.19      0.23      0.19       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu Naiwny GIN po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.70      0.61      0.65        61\n",
      "           1       0.83      0.69      0.76        36\n",
      "           2       0.85      0.81      0.83        78\n",
      "           3       0.76      0.67      0.71       158\n",
      "           4       0.67      0.85      0.75        81\n",
      "           5       0.56      0.70      0.62        57\n",
      "           6       0.63      0.66      0.64        29\n",
      "\n",
      "    accuracy                           0.72       500\n",
      "   macro avg       0.72      0.71      0.71       500\n",
      "weighted avg       0.73      0.72      0.72       500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "utils.val_scores_before_after_train(\n",
    "    cora.y[cora.val_mask],\n",
    "    yhat_naive_before_training[cora.val_mask],\n",
    "    yhat_naive_after_training[cora.val_mask],\n",
    "    \"Naiwny GIN\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44116 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "utils.plot_results(\n",
    "    epoch_losses_naive,\n",
    "    val_losses_naive,\n",
    "    plot_name=\"gin_naive.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------+--------------------------+----------------+----------+\n",
      "| Layer                            | Input Shape              | Output Shape   | #Param   |\n",
      "|----------------------------------+--------------------------+----------------+----------|\n",
      "| NodesClassifierModel             | [2708, 1433], [2, 10556] | [2708, 7]      | 410,569  |\n",
      "| ├─(gnn_convs)ModuleList          | --                       | --             | 410,569  |\n",
      "| │    └─(0)NaiveGIN               | [2708, 1433], [2, 10556] | [2708, 128]    | 400,001  |\n",
      "| │    │    └─(aggr)SumAggregation | [4, 1433]                | [1, 1433]      | --       |\n",
      "| │    │    └─(sigma)Identity      | --                       | --             | --       |\n",
      "| │    │    └─(mlp)Sequential      | [2708, 1433]             | [2708, 128]    | 400,000  |\n",
      "| │    └─(1)NaiveGIN               | [2708, 128], [2, 10556]  | [2708, 7]      | 10,568   |\n",
      "| │    │    └─(aggr)SumAggregation | [4, 1433]                | [1, 1433]      | --       |\n",
      "| │    │    └─(sigma)Identity      | --                       | --             | --       |\n",
      "| │    │    └─(mlp)Sequential      | [2708, 128]              | [2708, 7]      | 10,567   |\n",
      "| ├─(lin_layers)Sequential         | [2708, 7]                | [2708, 7]      | --       |\n",
      "+----------------------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg_nn.summary(naive_clf, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model klasyfikacji z warstwą splotu GNN z PyG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "\n",
    "mlp1 = nn.Sequential(\n",
    "    nn.Linear(cora.num_features, 256),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(256, 128),\n",
    "    nn.ReLU()\n",
    ")\n",
    "\n",
    "mlp2 = nn.Sequential(\n",
    "    nn.Linear(128, 64),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(64, 32),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(32, 7),\n",
    ")\n",
    "\n",
    "pyg_gin1 = pyg_nn.GINConv(nn=mlp1, train_eps=True, eps=0.1)\n",
    "pyg_gin2 = pyg_nn.GINConv(nn=mlp2, train_eps=True, eps=0.1)\n",
    "\n",
    "pyg_clf = NodesClassifierModel(gnn_convs=(pyg_gin1, pyg_gin2), lin_sizes=[])\n",
    "optim = th.optim.Adam(pyg_clf.parameters(), lr=0.01)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "with th.no_grad():\n",
    "    pyg_clf.eval()\n",
    "    yhat_pyg_before_training = pyg_clf(cora.x, cora.edge_index)\n",
    "\n",
    "\n",
    "epoch_losses_pyg = []\n",
    "val_losses_pyg = []\n",
    "\n",
    "for epoch in tqdm(range(10)):\n",
    "    pyg_clf.train()\n",
    "    optim.zero_grad()\n",
    "    out = pyg_clf(cora.x, cora.edge_index)\n",
    "    loss = loss_fn(out[cora.train_mask], cora.y[cora.train_mask])\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    epoch_losses_pyg.append(loss.item())\n",
    "\n",
    "    pyg_clf.eval()\n",
    "    with th.no_grad():\n",
    "        val_loss = loss_fn(out[cora.val_mask], cora.y[cora.val_mask])\n",
    "        val_losses_pyg.append(val_loss.item())\n",
    "\n",
    "with th.no_grad():\n",
    "    pyg_clf.eval()\n",
    "    yhat_pyg_after_training = pyg_clf(cora.x, cora.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wynik klasyfikacji modelu PyG GIN 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.16      1.00      0.28        81\n",
      "           5       0.00      0.00      0.00        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.16       500\n",
      "   macro avg       0.02      0.14      0.04       500\n",
      "weighted avg       0.03      0.16      0.05       500\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "Wynik klasyfikacji modelu PyG GIN po trenowaniu:\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.47      0.70      0.57        61\n",
      "           1       0.72      0.86      0.78        36\n",
      "           2       0.87      0.86      0.86        78\n",
      "           3       0.71      0.66      0.69       158\n",
      "           4       0.55      0.63      0.59        81\n",
      "           5       0.67      0.58      0.62        57\n",
      "           6       0.00      0.00      0.00        29\n",
      "\n",
      "    accuracy                           0.66       500\n",
      "   macro avg       0.57      0.61      0.59       500\n",
      "weighted avg       0.64      0.66      0.64       500\n",
      "\n"
     ]
    }
   ],
   "source": [
    "utils.val_scores_before_after_train(\n",
    "    cora.y[cora.val_mask],\n",
    "    yhat_pyg_before_training[cora.val_mask],\n",
    "    yhat_pyg_after_training[cora.val_mask],\n",
    "    \"PyG GIN\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37968 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "utils.plot_results(\n",
    "    epoch_losses_pyg,\n",
    "    val_losses_pyg,\n",
    "    plot_name=\"rys_4.14.tiff\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------+--------------------------+----------------+----------+\n",
      "| Layer                    | Input Shape              | Output Shape   | #Param   |\n",
      "|--------------------------+--------------------------+----------------+----------|\n",
      "| NodesClassifierModel     | [2708, 1433], [2, 10556] | [2708, 7]      | 410,569  |\n",
      "| ├─(gnn_convs)ModuleList  | --                       | --             | 410,569  |\n",
      "| │    └─(0)GINConv        | [2708, 1433], [2, 10556] | [2708, 128]    | 400,001  |\n",
      "| │    └─(1)GINConv        | [2708, 128], [2, 10556]  | [2708, 7]      | 10,568   |\n",
      "| ├─(lin_layers)Sequential | [2708, 7]                | [2708, 7]      | --       |\n",
      "+--------------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg_nn.summary(pyg_clf, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Alternatywna forma implementacji PyG\n",
    "\n",
    "Biblioteka PyG umożliwia tworzenie warstw splotu na dwa sposoby:\n",
    "1. Albo jako pojedyncze warstwy - każda z nich jest osobnym obiektem, który można wykorzystać w modelu (odmiennie ustawiając parametry itd.). Wtedy warstwy są łączone w modelu, tak jak w przypadku modelu z naiwną implementacją.\n",
    "2. Albo jako jeden model, który jest automatycznie konstruowany przez PyG z kilku warstw splotu zestawionych razem. \n",
    "\n",
    "\n",
    "W przedstawionym przykładzie zastosowaliśmy podejście (1) - samodzielnie opisując każdą warstwę. Poniżej przedstawiony jest przykład podejścia (2): PyG zbuduje jeden model grafowy, z dwiema warstwami GIN. Wejściem do każdej będzie:\n",
    "* macierz cech wierzchołków (wyjście z poprzedniej warstwy lub oryginalne cechy wierzchołków),\n",
    "* macierz sąsiedztwa grafu (nie ulega zmianie).\n",
    "\n",
    "Określamy specyficzne dla GIN parametry, które będą miały zastosowanie **do wszystkich warstw**.\n",
    "\n",
    "\n",
    "W przypadku architektury GIN używanie automatycznie konstruowanego modelu jest najmniej korzystne - pozbawia nas bowiem możliwości różnicowania sieci MLP w poszczególnych warstwach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------+--------------------------+----------------+----------+\n",
      "| Layer               | Input Shape              | Output Shape   | #Param   |\n",
      "|---------------------+--------------------------+----------------+----------|\n",
      "| GIN                 | [2708, 1433], [2, 10556] | [2708, 64]     | 453,504  |\n",
      "| ├─(dropout)Dropout  | [2708, 256]              | [2708, 256]    | --       |\n",
      "| ├─(act)ReLU         | [2708, 256]              | [2708, 256]    | --       |\n",
      "| ├─(convs)ModuleList | --                       | --             | 453,504  |\n",
      "| │    └─(0)GINConv   | [2708, 1433], [2, 10556] | [2708, 256]    | 432,896  |\n",
      "| │    └─(1)GINConv   | [2708, 256], [2, 10556]  | [2708, 64]     | 20,608   |\n",
      "| ├─(norms)ModuleList | --                       | --             | --       |\n",
      "| │    └─(0)Identity  | [2708, 256]              | [2708, 256]    | --       |\n",
      "| │    └─(1)Identity  | --                       | --             | --       |\n",
      "+---------------------+--------------------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "gin_model = pyg_nn.GIN(\n",
    "    in_channels=cora.num_features, # Ilość cech wejściowych - oryginalnych\n",
    "    hidden_channels=256,           # Ilość cech w warstwie ukrytej\n",
    "    num_layers=2,                  # Ilość warstw ukrytych - każda warstwa jest typu GINConv\n",
    "    out_channels=64,               # Ilość klas do przewidzenia\n",
    "    # Dodatkowe hiperparametry dla każdej warstwy\n",
    "    aggr='mean',\n",
    ")\n",
    "\n",
    "\n",
    "print(pyg_nn.summary(gin_model, cora.x, cora.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Taki model można używać identycznie, jak wcześniej."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    yhat = gin_model(cora.x, cora.edge_index)\n",
    "    classifier = NodesClassifierModel(gnn_convs=(gin_model,), lin_sizes=(64, 32, 7))\n",
    "    yhat_classifier = classifier(cora.x, cora.edge_index)"
   ]
  }
 ],
 "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
}
