{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "W tym notatniku przyjrzymy się technikom stosowanym do partycjonowania dużych grafów w procesie uczenia się. Wiemy już, jak podzielić graf na podzbiory treningowe, testowe i walidacyjne, ale co zrobić, gdy graf jest zbyt duży, aby go przetworzyć? W takich przypadkach stosuje się techniki partycjonowania grafów.\n",
    "\n",
    "Będziemy pracować z grafem \"Karate Club\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as th\n",
    "import networkx as nx\n",
    "import torch_geometric as pyg\n",
    "import visualizations as viz\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from importlib import reload"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = pyg.datasets.KarateClub()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[34, 34], edge_index=[2, 156], y=[34], train_mask=[34])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_nx = pyg.utils.to_networkx(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Próbkowanie sąsiedztwa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Próbkowanie bezpośredniego sąsiedztwa dla jednego wierzchołka - max. 3 bezpośrednich sąsiadów, wielkość batcha max. 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[4, 34], edge_index=[2, 3], y=[4], train_mask=[4], n_id=[4], e_id=[3], num_sampled_nodes=[2], num_sampled_edges=[1], input_id=[1], batch_size=1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "neighbor_sampler1 = pyg.loader.NeighborLoader(\n",
    "    graph,\n",
    "    num_neighbors=[3],\n",
    "    batch_size=4,\n",
    "    input_nodes=[0],\n",
    "    directed=True\n",
    ")\n",
    "\n",
    "batch1 = next(iter(neighbor_sampler1))\n",
    "batch1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Numery wierzchołków w macierzy sąsiedztwa \n",
      " tensor([[1, 2, 3],\n",
      "        [0, 0, 0]])\n",
      "Oryginalne numery wierzchołków tensor([ 0,  3, 10,  1])\n"
     ]
    }
   ],
   "source": [
    "# PyG numeruje wierzchołki w podgrafie sąsiedztwa ZAWSZE od 0\n",
    "print(\"Numery wierzchołków w macierzy sąsiedztwa \\n\", batch1.edge_index)\n",
    "print(\"Oryginalne numery wierzchołków\", batch1.n_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19780 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig_batch1 = viz.plot_batch_with_renaming(graph_nx, batch1, seed=222);\n",
    "fig_batch1.savefig(\"./images/rys_5.18.tif\", format=\"tif\", pil_kwargs={\"compression\": \"tiff_lzw\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Próbkowanie sąsiedztwa odległego o dwa kroki dla jednego wierzchołka. Wielkość batcha - max. 10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[9, 34], edge_index=[2, 12], y=[9], train_mask=[9], n_id=[9], e_id=[12], num_sampled_nodes=[3], num_sampled_edges=[2], input_id=[1], batch_size=1)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "neighbor_sampler2 = pyg.loader.NeighborLoader(\n",
    "    graph,\n",
    "    num_neighbors=[3, 3],\n",
    "    batch_size=10,\n",
    "    input_nodes=[0],\n",
    "    directed=True,\n",
    ")\n",
    "\n",
    "batch2 = next(iter(neighbor_sampler2))\n",
    "batch2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32496 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig_batch2 = viz.plot_batch_with_renaming(graph_nx, batch2, seed=222,);\n",
    "fig_batch2.savefig(\"./images/rys_5.19.tif\", format=\"tif\", pil_kwargs={\"compression\": \"tiff_lzw\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ta sama próbka ucząca będzie wyglądać inaczej, jeśli użyjemy wszystkich krawędzi, a nie tylko realnie odwiedzonych w procesie próbkowania."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 46380 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "neighbor_sampler2a = pyg.loader.NeighborLoader(\n",
    "    graph,\n",
    "    num_neighbors=[3, 3],\n",
    "    batch_size=10,\n",
    "    input_nodes=[0],\n",
    "    directed=False,\n",
    ")\n",
    "\n",
    "batch2a = next(iter(neighbor_sampler2a))\n",
    "\n",
    "fig_batch2a = viz.plot_batch_with_renaming(graph_nx, batch2a, seed=222, figsize=(12,12));\n",
    "fig_batch2a.savefig(\"./images/rys_5.20.tif\", format=\"tif\", pil_kwargs={\"compression\": \"tiff_lzw\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Próbkowanie bezpośredniego sąsiedztwa dla kilku wierzchołków  - max. 3 bezpośrednich sąsiadów, wielkość batcha max. 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[8, 34], edge_index=[2, 6], y=[8], train_mask=[8], n_id=[8], e_id=[6], num_sampled_nodes=[2], num_sampled_edges=[1], input_id=[2], batch_size=2)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "neighbor_sampler3 = pyg.loader.NeighborLoader(\n",
    "    graph,\n",
    "    num_neighbors=[3],\n",
    "    batch_size=10,\n",
    "    input_nodes=[0, 33],\n",
    "    directed=True,\n",
    ")\n",
    "\n",
    "batch3 = next(iter(neighbor_sampler3))\n",
    "batch3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32052 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig_batch3 = viz.plot_batch_with_renaming(graph_nx, batch3, seed=222);\n",
    "fig_batch3.savefig(\"./images/rys_5.21.tif\", format=\"tif\", pil_kwargs={\"compression\": \"tiff_lzw\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Różnica pomiędzy próbkowaniem `directed = True` a `directed = False`:\n",
    "\n",
    "1. `directed = True` - próbkowanie sąsiedztwa wierzchołka w jednym kierunku. Polecane, gdy ilość używanych warstw splotu jest taka sama, jak ilość skoków w grafie.\n",
    "2. `directed = False` - próbkowanie sąsiedztwa wierzchołka w obu kierunkach. Polecane, gdy używane będą głębsze sieci splotowe - mające więcej warstw niż próbkowanych skoków sąsiedztwa."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[8, 34], edge_index=[2, 14], y=[8], train_mask=[8], n_id=[8], e_id=[14], input_id=[2], batch_size=2)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pyg.seed_everything(123)\n",
    "neighbor_sampler4 = pyg.loader.NeighborLoader(\n",
    "    graph,\n",
    "    num_neighbors=[3],\n",
    "    batch_size=10,\n",
    "    input_nodes=[0, 33],\n",
    "    directed=False,)\n",
    "\n",
    "batch4 = next(iter(neighbor_sampler4))\n",
    "batch4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35628 characters]",
      "text/plain": [
       "<Figure size 472.441x472.441 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "viz.plot_batch_with_renaming(graph_nx, batch4);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wyszukiwanie społeczności"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyG zawiera implementację "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing METIS partitioning...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "cluster_data = pyg.loader.ClusterData(graph, num_parts=3, recursive=True)\n",
    "cluster_loader = pyg.loader.ClusterLoader(cluster_data, batch_size=1, shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data(x=[11, 34], y=[11], train_mask=[11], edge_index=[2, 30])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 257896 characters]",
      "text/plain": [
       "<Figure size 1200x1200 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from matplotlib.gridspec import GridSpec\n",
    "\n",
    "fig = plt.figure(figsize=(12, 12))\n",
    "gs = GridSpec(2, 2, figure=fig)\n",
    "\n",
    "ax1 = fig.add_subplot(gs[0, 0])\n",
    "ax2 = fig.add_subplot(gs[0, 1])\n",
    "ax3 = fig.add_subplot(gs[1, :])\n",
    "\n",
    "viz.plot_batch_with_renaming(graph_nx, cluster_data[0], name_attr=None, title=\"Społ. 0\", ax=ax1, figsize=(10,10), font_size=16, node_size=500)\n",
    "\n",
    "viz.plot_batch_with_renaming(graph_nx, cluster_data[1], name_attr=None, title=\"Społ. 1\", ax=ax2, figsize=(10,10), font_size=16, node_size=500)\n",
    "viz.plot_batch_with_renaming(graph_nx, cluster_data[2], name_attr=None, title=\"Społ. 2\", ax=ax3, figsize=(10,10), font_size=16, node_size=500)\n",
    "plt.tight_layout()\n",
    "fig.savefig(\"./images/rys_5.23.tif\", format=\"tif\", pil_kwargs={\"compression\": \"tiff_lzw\"})"
   ]
  }
 ],
 "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
}
