{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Wstęp\n",
        "\n",
        "W tym notatniku przyjrzymy się sposobowi, w jaki można zaimplementować podejście **przekazywania wiadomości** (ang. *message passing*) w sieciach neuronowych.\n",
        "Zobaczymy trzy implementacje:\n",
        "1. \"Naiwną\", samodzielną.\n",
        "2. Z użyciem biblioteki PyTorch Geometric.\n",
        "3. Z użyciem biblioteki Deep Graph Library."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Import bibliotek"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "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 przygotowania podstawowych informacji o grafie (o strukturze identycznej, jak w poprzednich przykładach).\n",
        "Będziemy potrzebować:\n",
        "1. Macierz cech wierzchołków.\n",
        "2. Listę sąsiedztwa (krawędzie)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "X = th.tensor([\n",
        "        [1, 2, 3],\n",
        "        [4, 5, 6],\n",
        "        [7, 8, 9],\n",
        "        [10, 20, 30]\n",
        "    ]).to(th.float)\n",
        "edges = th.tensor([\n",
        "        [0, 1, 2, 3, 3, 3],\n",
        "        [3, 3, 3, 0, 1, 2]\n",
        "    ])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Wizualizacja grafu, z wykorzystaniem struktur danych z PyTorch Geometric."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "g = pyg_data.Data(\n",
        "    x=X,\n",
        "    edge_index=edges,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 20264 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": [
        "# Implementacje"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Naiwna implementacja\n",
        "\n",
        "Zaczniemy od bardzo prostej implementacji, wykorzystującej jedynie bazowe elementy z biblioteki PyTorch. Zakładamy **brak nieliniowości** i proste sumowanie jako agregację."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "class NaiveMPNN(nn.Module):\n",
        "\n",
        "    def __init__(self, dim_in: int):\n",
        "        super().__init__()\n",
        "        self.dim_in = dim_in\n",
        "\n",
        "    def forward(self, X: th.Tensor, edges: th.Tensor) -> th.Tensor:\n",
        "        h_new = th.zeros_like(X)\n",
        "        for v in range(X.size(0)):\n",
        "            v_index_as_receiver = th.where(edges[1] == v)[0]\n",
        "            neighbors = edges[0, v_index_as_receiver]\n",
        "            msg = self.message(neighbors, X)\n",
        "            agg = self.aggregate(msg)\n",
        "            new_v = self.update(X[v], agg)\n",
        "            h_new[v] = new_v\n",
        "        return h_new\n",
        "\n",
        "\n",
        "    def message(self, v_neighbors: th.Tensor, X: th.Tensor) -> th.Tensor:\n",
        "        \"\"\"Implantacja funkcji `message` - przyjmuje indeksy wierzchołków\n",
        "        sąsiadów oraz macierz cech. Dla uproszczenia - zwraca po prostu cechy\n",
        "        bez żadnych przekształceń.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        v_neighbors: th.Tensor\n",
        "            Indeksy sąsiadów v: N(v).\n",
        "        X : th.Tensor\n",
        "            Macierz cech wierzchołków\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        th.Tensor\n",
        "            Wiadomości od sąsiadów - ich przekształcone cechy. W tym wypadku:\n",
        "            bez żadnych przekształceń.\n",
        "        \"\"\"\n",
        "        return X[v_neighbors]\n",
        "\n",
        "    def aggregate(self, messages: th.Tensor) -> th.Tensor:\n",
        "        \"\"\"Implementacja funkcji `aggregate` - przyjmuje wiadomości od\n",
        "        sąsiadów i zwraca ich sumę.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        messages : th.Tensor\n",
        "            Wiadomości od sąsiadów.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        th.Tensor\n",
        "            Zsumowane wiadomości od sąsiadów.\n",
        "        \"\"\"\n",
        "        return th.sum(messages, dim=0)\n",
        "\n",
        "    def update(self, v: th.Tensor, agg: th.Tensor) -> th.Tensor:\n",
        "        \"\"\"Implementacja funkcji `update` - przyjmuje cechy wierzchołka\n",
        "        oraz zsumowane wiadomości od sąsiadów i zwraca nowe cechy wierzchołka.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        v : th.Tensor\n",
        "            Cechy wierzchołka.\n",
        "        agg : th.Tensor\n",
        "            Zsumowane wiadomości od sąsiadów.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        th.Tensor\n",
        "            Nowe cechy wierzchołka.\n",
        "        \"\"\"\n",
        "        return v + agg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Sprawdzenie w działaniu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tensor([[11., 22., 33.],\n",
              "        [14., 25., 36.],\n",
              "        [17., 28., 39.],\n",
              "        [22., 35., 48.]])"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "naive_mpnn = NaiveMPNN(dim_in=X.shape[1])\n",
        "with th.no_grad():\n",
        "    h_new_naive_mpnn = naive_mpnn(X, edges)\n",
        "h_new_naive_mpnn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Uzyskane wyniki są zgodne z uproszczonym podejściem teoretycznym opisanym w książce."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Implementacja PyG\n",
        "\n",
        "Zaimplementujemy takie samo podejście w bibliotece Pytorch Geometric, korzystając z jej mechanizmów do przetwarzania grafów."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "class MPNNPyG(pyg_nn.MessagePassing):\n",
        "\n",
        "    def __init__(self, dim_in: int, aggr='add'):\n",
        "        \"\"\"Inicjalizacja warstwy MPNN w PyG.\n",
        "        Co istotne - standardowe agregacje są zaimplementowane w PyG, więc nie musimy ponownie\n",
        "        implementować funkcji `agg` w standardowych przypadkach (suma, średnia, max, min).\"\"\"\n",
        "        super().__init__(aggr=aggr)\n",
        "\n",
        "    def forward(self, X: th.Tensor, edge_index: th.Tensor) -> th.Tensor:\n",
        "        return self.propagate(edge_index, x=X)\n",
        "\n",
        "    def message(self, x_j: th.Tensor) -> th.Tensor:\n",
        "        \"\"\"Implementacja funkcji `message` . W konwencji przyjętej w PyG\n",
        "        `x_j` odpowiada sąsiadowi - \"u\". Można też wykorzystywać argumenty takie jak\n",
        "        `x_i` - odpowiada wierzchołkowi \"v\" (czyli temu, dla którego budujemy reprezentację),\n",
        "        oraz inne, dodatkowe argumenty przekazane do funkcji propagate.\n",
        "        Te dodatkowe argumenty są dynamicznie alokowane podczas wywołania, dlatego nie \n",
        "        występują w statycznej dokumentacji funkcji.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        x_j : th.Tensor\n",
        "            Cechy sąsiada.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        th.Tensor\n",
        "            Wiadomość od sąsiada.\n",
        "        \"\"\"\n",
        "        return x_j\n",
        "\n",
        "    def update(self, aggr_out: th.Tensor, x: th.Tensor) -> th.Tensor:\n",
        "        \"\"\"Implementacja funkcji `update`. W konwencji przyjętej w PyG\n",
        "        `aggr_out` odpowiada zsumowanym wiadomościom od sąsiadów. Funkcja\n",
        "        może też przyjmować inne argumenty, takie jak `x` - cechy wierzchołka, \n",
        "        przekazane do funkcji propagate. Te dodatkowe argumenty są dynamicznie\n",
        "        alokowane podczas wywołania, dlatego nie występują w statycznej dokumentacji\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        aggr_out : th.Tensor\n",
        "            Zsumowane wiadomości od sąsiadów.\n",
        "        x : th.Tensor\n",
        "            Cechy wierzchołka - dodatkowy argument przekazany do funkcji propagate.\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        th.Tensor\n",
        "            Nowe cechy wierzchołka.\n",
        "        \"\"\"\n",
        "        return x + aggr_out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tensor([[11., 22., 33.],\n",
              "        [14., 25., 36.],\n",
              "        [17., 28., 39.],\n",
              "        [22., 35., 48.]])"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "pyg_mpnn = MPNNPyG(dim_in=X.shape[1])\n",
        "with th.no_grad():\n",
        "    h_new_pyg = pyg_mpnn(X, edges)\n",
        "h_new_pyg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Uzyskane wyniki są zgodne z uproszczonym podejściem teoretycznym opisanym w książce."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Implementacja DGL\n",
        "\n",
        "W tej sekcji przedstawimy analogiczną implementację mechanizmu *message passing* z wykorzystaniem biblioteki Deep Graph Library (DGL), bazując na jej natywnych mechanizmach przetwarzania grafów. Warto podkreślić, że DGL przyjmuje odmienny model programistyczny niż PyTorch Geometric: sieć neuronowa deleguje operacje na grafie do **obiektu grafu**, zamiast wykonywać transformacje wyłącznie po stronie warstwy/sieci.\n",
        "\n",
        "> **Uwaga (stan na styczeń 2026).** W momencie przygotowywania książki i niniejszego materiału biblioteka DGL nie była aktywnie rozwijana i nie oferowała wydań kompatybilnych z PyTorch 2.8 (wykorzystywanym w tym projekcie, m.in. przez PyTorch Geometric). Z tego powodu DGL **nie została** uwzględniona w pliku `pyproject.toml`, ponieważ nie było możliwe uzgodnienie zależności w ramach jednego środowiska.\n",
        ">\n",
        "> W konsekwencji poniższy `import dgl` w aktualnym środowisku **nie zadziała**. Aby uruchomić przykłady z DGL, należy przygotować **osobne** środowisko wirtualne z wersjami bibliotek zgodnymi z DGL.\n",
        "\n",
        "Wszystko wskazuje na to, że DGL przestanie być aktywnie rozwijana w przyszłości, niestety.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import dgl\n",
        "import dgl.function as dgl_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Krok 1: utworzenie grafu DGL"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<bound method DGLGraph.num_nodes of Graph(num_nodes=4, num_edges=6,\n",
              "      ndata_schemes={}\n",
              "      edata_schemes={})>"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dgl_graph = dgl.graph((edges[0], edges[1]))\n",
        "dgl_graph.num_nodes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "dgl_graph.ndata['X'] = X"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Krok 2: utworzenie sieci neuronowej delegującej kalkulacje do grafu"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [],
      "source": [
        "class NaiveDGLNet(nn.Module):\n",
        "\n",
        "    def __init__(self, dim_in: int):\n",
        "        super().__init__()\n",
        "        self.dim_in = dim_in\n",
        "\n",
        "    def forward(self, g: dgl.DGLGraph) -> th.Tensor:\n",
        "       with g.local_scope():\n",
        "           g.update_all(\n",
        "               message_func=dgl_fn.copy_u('X', 'm'),\n",
        "               reduce_func=dgl_fn.sum('m', 'h_N')\n",
        "           )\n",
        "           h_N = g.ndata['h_N']\n",
        "           new_h = g.ndata['X'] + h_N\n",
        "           return new_h"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tensor([[11., 22., 33.],\n",
              "        [14., 25., 36.],\n",
              "        [17., 28., 39.],\n",
              "        [22., 35., 48.]])"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "naive_dgl = NaiveDGLNet(dim_in=X.shape[1])\n",
        "with th.no_grad():\n",
        "    h_new_naive_dgl = naive_dgl(dgl_graph)\n",
        "h_new_naive_dgl"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Wyniki są zgodne z uproszczonym podejściem teoretycznym opisanym w książce."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Sprawdzenie identyczności wyników"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [],
      "source": [
        "assert th.equal(h_new_naive_dgl, h_new_naive_mpnn)\n",
        "assert th.equal(h_new_naive_dgl, h_new_pyg)"
      ]
    }
  ],
  "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
}
