{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wstęp\n",
    "\n",
    "W tym notatniku przyjrzymy się uogólnionej i uproszczonej architekturze sieci GNN. Zbudujemy jej naiwną i poglądową implementację z wykorzystaniem PyTorcha."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as th\n",
    "import torch.nn as nn\n",
    "import networkx as nx\n",
    "import torch_geometric as pyg\n",
    "import torch_geometric.data as pyg_data\n",
    "import torch_geometric.nn as pyg_nn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie grafu nieskierowanego"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Zaczniemy od zbudowania grafu - z atrybutami takimi, jak w pokazanej wizualizacji."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = pyg_data.Data(\n",
    "    x=th.tensor([\n",
    "        [1, 2, 3],\n",
    "        [4, 5, 6],\n",
    "        [7, 8, 9],\n",
    "        [10, 20, 30]\n",
    "    ]).to(th.float),\n",
    "    edge_index=th.tensor([\n",
    "        [0, 1, 2, 3, 3, 3],\n",
    "        [3, 3, 3, 0, 1, 2]\n",
    "    ]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20932 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g_nx = pyg.utils.to_networkx(g)\n",
    "\n",
    "nx.draw(g_nx, with_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modele"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naiwna implementacja\n",
    "\n",
    "\n",
    "Zaczniemy od implementacji naiwnej, która jest najbardziej intuicyjna. W tym przypadku zaimplementujemy model GNN jako funkcję, która przyjmuje macierz sąsiedztwa i macierz atrybutów wierzchołków. Będziemy samodzielnie iterować po wierzchołkach, znajdując sąsiedztwo i agregując atrybuty wierzchołków sąsiadujących, dokładnie tak, jak pokazano w pseudokodzie implementacji.\n",
    "\n",
    "Dodatkowe założenia:\n",
    "\n",
    "1. Nasza sieć będzie posiadać **wyłącznie** warstwy liniowe bez aktywacji aby łatwiej było zweryfikować wyniki obliczeń.\n",
    "2. Agregacja sąsiadów jest prostą sumą atrybutów wierzchołków sąsiadujących.\n",
    "3. Wykonujemy osadzenie (ang. *embedding*) na dwóch atrybutach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NaiveGeneralGNN(nn.Module):\n",
    "\n",
    "    def __init__(self, dim_in: int =3):\n",
    "        super().__init__()\n",
    "        self.dim_in = dim_in\n",
    "\n",
    "    def aggregate_neighbors(self, v: int, X: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Agregacja sąsiadów.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        v : int\n",
    "            Indeks bieżącego wierzchołka.\n",
    "        x : th.Tensor\n",
    "            Tensor cech sąsiadów. Wymiary: [n x dim_in]\n",
    "        edge_index : th.Tensor\n",
    "            Indeksy krawędzi. Wymiary: [2 x e].\n",
    "            Wymiar pierwszy [0, e] -> wierzchołki \"wysyłające\" wiadomość.\n",
    "            Wymiar drugi [1, e] -> wierzchołki \"odbierające\" wiadomość.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Zagregowany tensor cech sąsiadów. Wymiary: [dim_in].\n",
    "        \"\"\"\n",
    "        v_index_in_edge = (edge_index[1] == v).argwhere().squeeze()  # Pozycje, gdzie v jest odbiorcą\n",
    "        neighbors = edge_index[0, v_index_in_edge]                   # Indeksy sąsiadów\n",
    "        X_neighbors = th.atleast_2d(X[neighbors])                     # Cechy sąsiadów pod indeksami.\n",
    "                                                                     # nawet dla pojedynczego sąsiada - utrzymanie macierzy 2D\n",
    "        return th.sum(X_neighbors, dim=0)                            # Sumowanie cech sąsiadów - agregacja\n",
    "\n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja wykonująca pełne przekształcenie i zwracająca wynik w postaci nowego tensora osadzenia.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : th.Tensor\n",
    "            Cechy wierzchołków. Wymiary: [n x dim_in]\n",
    "        edge_index : th.Tensor\n",
    "            Indeksy krawędzi. Wymiary: [2 x e].\n",
    "            Wymiar pierwszy [0, e] -> wierzchołki \"wysyłające\" wiadomość.\n",
    "            Wymiar drugi [1, e] -> wierzchołki \"odbierające\" wiadomość.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Nowy tensor osadzenia. Wymiary: [n, dim_in].\n",
    "        \"\"\"\n",
    "        n = x.shape[0]                             # Liczba wierzchołków\n",
    "        h_new = th.zeros(n, self.dim_in)           # Inicjalizacja nowej macierzy osadzenia.\n",
    "        for v in range(x.shape[0]):                # Dla każdego wierzchołka\n",
    "            aggregated_neighbors = self.aggregate_neighbors(v, x, edge_index) # Agregacja sąsiadów\n",
    "            h_new[v] = x[v] + aggregated_neighbors\n",
    "        return h_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sprawdźmy, czy sieć prawidłowo wykonuje osadzenia:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[11., 22., 33.],\n",
      "        [14., 25., 36.],\n",
      "        [17., 28., 39.],\n",
      "        [22., 35., 48.]])\n"
     ]
    }
   ],
   "source": [
    "th.manual_seed(123)\n",
    "naive_gnn = NaiveGeneralGNN(dim_in=g.x.shape[1])\n",
    "with th.no_grad():\n",
    "    print(naive_gnn(g.x, g.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sprawdźmy wyniki dla analizy sąsiedztwa wierzchołka 0: `v = 0`, mającego tylko jednego sąsiada:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Zagregowane cechy sąsiadów wierzchołka 0: tensor([10., 20., 30.])\n"
     ]
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    aggregated_neighbors_v0 = naive_gnn.aggregate_neighbors(0, g.x, g.edge_index)\n",
    "print(\"Zagregowane cechy sąsiadów wierzchołka 0:\", aggregated_neighbors_v0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Jedynym sąsiadem v=0 jest wierzchołek 3 o cechach `[10, 20, 30]`, więc wszystko się zgadza."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sprawdźmy wyniki analizy sąsiedztwa wierzchołka 3: `v = 3`, mającego aż trzech sąsiadów:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Zagregowane cechy sąsiadów wierzchołka 3: tensor([12., 15., 18.])\n"
     ]
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    aggregated_neighbors_v3 = naive_gnn.aggregate_neighbors(3, g.x, g.edge_index)\n",
    "print(\"Zagregowane cechy sąsiadów wierzchołka 3:\", aggregated_neighbors_v3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wierzchołek v=3 ma trzech sąsiadów: 0, 1, 2, suma ich cech daje dokładnie takie wyniki."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replikacja w Pytorch Geometric\n",
    "\n",
    "Uzyskamy ten sam efekt za pomocą biblioteki `Pytorch Geometric` wykorzystując wbudowane w nią mechanizmy propagacji wiadomości i funkcje `message` oraz `aggregate`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import Tensor\n",
    "\n",
    "\n",
    "class PyGNaiveModel(pyg_nn.MessagePassing):\n",
    "\n",
    "    def __init__(self, dim_in: int = 3, aggr: int = 'sum'):\n",
    "        super().__init__(aggr=aggr)\n",
    "        self.dim_in = dim_in\n",
    "\n",
    "    def message(self, x_j: th.Tensor) -> th.Tensor:\n",
    "        \"\"\"Funkcja konstruująca wiadomości. We frameworku Pytorch Geometric\n",
    "        x_j - wierzchołek nadający informację (sąsiad)\n",
    "        x_i - wierzchołek odbierający informację (bieżący wierzchołek). \n",
    "        W tym przypadku nie przeprowadzamy żadnej transformacji, zwracamy cechy sąsiada po prostu.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x_j : th.Tensor\n",
    "            Cechy sąsiada. Wymiary: [1 x dim_in]\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        th.Tensor\n",
    "            Wiadomość. Wymiary: [1 x dim_in]\n",
    "        \"\"\"\n",
    "        return x_j\n",
    "\n",
    "    def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
    "        aggregated_neighbors = self.propagate(edge_index, x=x) # Funkcja propaguje wiadomości i agreguje sąsiadów\n",
    "        return x + aggregated_neighbors                        # Dodajemy osadzenie wierzchołka bieżącego do zagregowanych sąsiadów\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[11., 22., 33.],\n",
      "        [14., 25., 36.],\n",
      "        [17., 28., 39.],\n",
      "        [22., 35., 48.]])\n"
     ]
    }
   ],
   "source": [
    "naive_pyg_gnn = PyGNaiveModel(dim_in=g.x.shape[1])\n",
    "with th.no_grad():\n",
    "    print(naive_pyg_gnn(g.x, g.edge_index))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wyniki są identyczne jak w naiwnej manualnej implementacji."
   ]
  }
 ],
 "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
}
