GNN/Data Handling in PyG 2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "69a46b4c",
      "metadata": {
        "id": "69a46b4c"
      },
      "source": [
        "## Tutorial #15: Data Handling in PyG (part 2)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e7a6bfa0",
      "metadata": {
        "id": "e7a6bfa0"
      },
      "source": [
        "### Custom PyG dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b5560f4f",
      "metadata": {
        "id": "b5560f4f"
      },
      "source": [
        "In the first part of the notebook we will see how to create a custom dataset in PyG.\n",
        "\n",
        "The dataset we'll load is called FRANKENSTEIN, the files can be downloaded from the\n",
        "[networkrepository](http://networkrepository.com/FRANKENSTEIN.php) site. The dataset was originally presented in the paper titled [Graph Invariant Kernels](https://www.ijcai.org/Proceedings/15/Papers/528.pdf).\n",
        "\n",
        "The dataset is a collection of graphs representing molecules, but the atom symbols of the vertices are substituted with MNIST digits. Each graph is associated with a label, indicating the mutagenicity of the molecule."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b339e866",
      "metadata": {
        "id": "b339e866"
      },
      "source": [
        "The following is the README of the dataset:\n",
        "\n",
        ">FRANKENSTEIN contains the following comma separated text files:\n",
        ">\n",
        ">n: total number of nodes\n",
        ">m: total number of edges\n",
        ">N: number of graphs\n",
        ">\n",
        ">\n",
        ">*.node_attrs (n lines)\n",
        "\tmatrix of node attributes,\n",
        "\tthe comma seperated values in the i-th line is the attribute vector of the node with node_id i\n",
        ">\n",
        ">*.edges (m lines)\n",
        "\tsparse (block diagonal) adjacency matrix for all graphs,\n",
        "\teach line corresponds to (row, col) resp. (node_id, node_id)\n",
        ">\n",
        ">*.graph_labels (N lines)\n",
        "\tclass labels for all graphs in the dataset,\n",
        "\tthe value in the i-th line is the class label of the graph with graph_id i\n",
        ">\n",
        ">*.graph_idx (n lines)\n",
        "\tcolumn vector of graph identifiers for all nodes of all graphs,\n",
        "\tthe value in the i-th line is the graph_id of the node with node_id i"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d39106a1",
      "metadata": {
        "id": "d39106a1"
      },
      "source": [
        "The dataset is composed as follows:\n",
        "\n",
        "    Nr. of graphs:         4337\n",
        "    Total nr. of nodes:    73283 x 780 (weird, we'll change it later to 784)\n",
        "    Total nr. of edges:    155068"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "811a858d",
      "metadata": {
        "id": "811a858d"
      },
      "outputs": [],
      "source": [
        "!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7a9ea054",
      "metadata": {
        "id": "7a9ea054"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import os\n",
        "import pandas as pd\n",
        "from torch_geometric.data import InMemoryDataset, Data, download_url, extract_zip\n",
        "from torch_geometric.utils.convert import to_networkx\n",
        "import networkx as nx\n",
        "from tqdm import tqdm\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7d92abcf",
      "metadata": {
        "id": "7d92abcf"
      },
      "source": [
        "To create the dataset we need to convert the raw information into a ```Data``` object (a graph) in PyG.\n",
        "\n",
        "The first step is to load the csv files, this can be done manually or using some data library as Pandas:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "671effcc",
      "metadata": {
        "id": "671effcc"
      },
      "outputs": [],
      "source": [
        "raw_dir = \"tmp/raw/\"\n",
        "\n",
        "path = os.path.join(raw_dir, 'FRANKENSTEIN.node_attrs')\n",
        "node_attrs = pd.read_csv(path, sep=',', header=None)\n",
        "node_attrs.index += 1\n",
        "\n",
        "path = os.path.join(raw_dir, 'FRANKENSTEIN.edges')\n",
        "edge_index = pd.read_csv(path, sep=',', names=['source', 'target'])\n",
        "edge_index.index += 1\n",
        "\n",
        "path = os.path.join(raw_dir, 'FRANKENSTEIN.graph_idx')\n",
        "graph_idx = pd.read_csv(path, sep=',', names=['idx'])\n",
        "graph_idx.index += 1\n",
        "\n",
        "path = os.path.join(raw_dir, 'FRANKENSTEIN.graph_labels')\n",
        "graph_labels = pd.read_csv(path, sep=',', names=['label'])\n",
        "graph_labels.index += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "12c78ee6",
      "metadata": {
        "id": "12c78ee6"
      },
      "source": [
        "Graph ids go from 1 to 4337, let's extract the information for a single graph (id: 2345)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "09cfff64",
      "metadata": {
        "id": "09cfff64"
      },
      "outputs": [],
      "source": [
        "g_idx=2345\n",
        "\n",
        "node_ids = graph_idx.loc[graph_idx['idx']==g_idx].index\n",
        "\n",
        "# Node features\n",
        "attributes = node_attrs.loc[node_ids, :]\n",
        "\n",
        "# Edges info\n",
        "edges = edge_index.loc[edge_index['source'].isin(node_ids)]\n",
        "edges_ids = edges.index\n",
        "\n",
        "# Graph label\n",
        "label = graph_labels.loc[g_idx]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f8b7078b",
      "metadata": {
        "id": "f8b7078b"
      },
      "outputs": [],
      "source": [
        "print(\"Nodes:\", node_ids.shape)\n",
        "print(\"Attributes:\", attributes.shape)\n",
        "print(\"Edges:\", edges.shape)\n",
        "print(\"Label:\", label.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4123c56e",
      "metadata": {
        "id": "4123c56e"
      },
      "outputs": [],
      "source": [
        "print(\"Nodes:\", node_ids)\n",
        "print(\"Attributes:\", attributes)\n",
        "print(\"Edges:\", edges)\n",
        "print(\"Label:\", label)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1d0804e7",
      "metadata": {
        "id": "1d0804e7"
      },
      "source": [
        "At this stage the indices in the ```edges``` variable are not normalized for the single graph, e.g. they do not start from 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "01003f31",
      "metadata": {
        "id": "01003f31"
      },
      "outputs": [],
      "source": [
        "edge_idx = torch.tensor(edges.to_numpy().transpose(), dtype=torch.long)\n",
        "map_dict = {v.item():i for i,v in enumerate(torch.unique(edge_idx))}\n",
        "map_edge = torch.zeros_like(edge_idx)\n",
        "for k,v in map_dict.items():\n",
        "    map_edge[edge_idx==k] = v"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "12aa0faf",
      "metadata": {
        "id": "12aa0faf"
      },
      "outputs": [],
      "source": [
        "map_dict, map_edge, map_edge.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0d71d186",
      "metadata": {
        "id": "0d71d186"
      },
      "source": [
        "As final step we convert the ```DataFrames``` to torch tensors. The node features are basically MNIST images, therefore their size should be 784 (28x28), but for some reason the files provide vectors of length 780. To adjust this, we simply add a padding of 4 zeros at the end of the vector (it will not affect the digits representation, as we'll see later)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "62ba9bf9",
      "metadata": {
        "id": "62ba9bf9"
      },
      "outputs": [],
      "source": [
        "attrs = torch.tensor(attributes.to_numpy(), dtype=torch.float)\n",
        "pad = torch.zeros((attrs.shape[0], 4), dtype=torch.float)\n",
        "x = torch.cat((attrs, pad), dim=-1)\n",
        "\n",
        "edge_idx = map_edge.long()\n",
        "\n",
        "np_lab = label.to_numpy()\n",
        "y = torch.tensor(np_lab if np_lab[0] == 1 else [0], dtype=torch.long)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9ec8e317",
      "metadata": {
        "id": "9ec8e317"
      },
      "outputs": [],
      "source": [
        "x.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "71a27ae1",
      "metadata": {
        "id": "71a27ae1"
      },
      "source": [
        "Then we create the ```Data``` object representing the graph"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2cc06f18",
      "metadata": {
        "id": "2cc06f18"
      },
      "outputs": [],
      "source": [
        "graph = Data(x=x, edge_index=edge_idx,  y=y)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "57338892",
      "metadata": {
        "id": "57338892"
      },
      "source": [
        "Let's visualize the graph!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1b040996",
      "metadata": {
        "id": "1b040996"
      },
      "outputs": [],
      "source": [
        "vis = to_networkx(graph)\n",
        "plt.figure(1,figsize=(8,8))\n",
        "nx.draw(vis, cmap=plt.get_cmap('Set3'),node_size=70,linewidths=6)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6221b4cc",
      "metadata": {
        "id": "6221b4cc"
      },
      "source": [
        "We can also plot the vertices attributes (a.k.a. the digits):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1e114906",
      "metadata": {
        "id": "1e114906"
      },
      "outputs": [],
      "source": [
        "digit = x[5].reshape(28,28)\n",
        "plt.matshow(digit, cmap='gray')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1e184ce0",
      "metadata": {
        "id": "1e184ce0"
      },
      "source": [
        "Let's now put the process above into the ```Dataset``` class of PyG. Specifically, we are going to create an ```InMemoryDataset```. From the official PyG documentation we see that some methods need to be override:"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a467a5a8",
      "metadata": {
        "id": "a467a5a8"
      },
      "source": [
        "    torch_geometric.data.InMemoryDataset.raw_file_names(): A list of files in the raw_dir which needs to be found in order to skip the download.\n",
        "\n",
        "    torch_geometric.data.InMemoryDataset.processed_file_names(): A list of files in the processed_dir which needs to be found in order to skip the processing.\n",
        "\n",
        "    torch_geometric.data.InMemoryDataset.download(): Downloads raw data into raw_dir.\n",
        "\n",
        "    torch_geometric.data.InMemoryDataset.process(): Processes raw data and saves it into the processed_dir.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb29f522",
      "metadata": {
        "id": "eb29f522"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import os\n",
        "import pandas as pd\n",
        "from torch_geometric.data import InMemoryDataset, Data, download_url, extract_zip\n",
        "from torch_geometric.utils.convert import to_networkx\n",
        "import networkx as nx\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "class Frankenstein(InMemoryDataset):\n",
        "\n",
        "    # Base url to download the files\n",
        "    url = 'http://nrvis.com/download/data/labeled/FRANKENSTEIN.zip'\n",
        "\n",
        "    def __init__(self, root, transform=None, pre_transform=None):\n",
        "        super(Frankenstein, self).__init__(root, transform, pre_transform)\n",
        "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
        "\n",
        "    @property\n",
        "    def raw_file_names(self):\n",
        "        # List of the raw files\n",
        "        return ['FRANKENSTEIN.edges', 'FRANKENSTEIN.graph_idx',\n",
        "                'FRANKENSTEIN.graph_labels', 'FRANKENSTEIN.node_attrs']\n",
        "\n",
        "    @property\n",
        "    def processed_file_names(self):\n",
        "        return 'data.pt'\n",
        "\n",
        "    def download(self):\n",
        "        # Download the file specified in self.url and store\n",
        "        # it in self.raw_dir\n",
        "        path = download_url(self.url, self.raw_dir)\n",
        "        extract_zip(path, self.raw_dir)\n",
        "        # The zip file is removed\n",
        "        os.unlink(path)\n",
        "\n",
        "\n",
        "    def process(self):\n",
        "        # Read the files' content as Pandas DataFrame. Nodes and graphs ids\n",
        "        # are based on the file row-index, we adjust the DataFrames indices\n",
        "        # by starting from 1 instead of 0.\n",
        "\n",
        "        path = os.path.join(self.raw_dir, 'FRANKENSTEIN.node_attrs')\n",
        "        node_attrs = pd.read_csv(path, sep=',', header=None)\n",
        "        node_attrs.index += 1\n",
        "\n",
        "        path = os.path.join(self.raw_dir, 'FRANKENSTEIN.edges')\n",
        "        edge_index = pd.read_csv(path, sep=',', names=['source', 'target'])\n",
        "        edge_index.index += 1\n",
        "\n",
        "        path = os.path.join(self.raw_dir, 'FRANKENSTEIN.graph_idx')\n",
        "        graph_idx = pd.read_csv(path, sep=',', names=['idx'])\n",
        "        graph_idx.index += 1\n",
        "\n",
        "        path = os.path.join(self.raw_dir, 'FRANKENSTEIN.graph_labels')\n",
        "        graph_labels = pd.read_csv(path, sep=',', names=['label'])\n",
        "        graph_labels.index += 1\n",
        "\n",
        "\n",
        "        # In the loop we extract the nodes' embeddings, edges connectivity for\n",
        "        # and label for a graph, process the information and put it in a Data\n",
        "        # object, then we add the object to a list\n",
        "        data_list = []\n",
        "        ids_list = graph_idx['idx'].unique()\n",
        "        for g_idx in tqdm(ids_list):\n",
        "            node_ids = graph_idx.loc[graph_idx['idx']==g_idx].index\n",
        "\n",
        "            # Node features\n",
        "            attributes = node_attrs.loc[node_ids, :]\n",
        "\n",
        "            # Edges info\n",
        "            edges = edge_index.loc[edge_index['source'].isin(node_ids)]\n",
        "            edges_ids = edges.index\n",
        "\n",
        "            # Graph label\n",
        "            label = graph_labels.loc[g_idx]\n",
        "\n",
        "            # Normalize the edges indices\n",
        "            edge_idx = torch.tensor(edges.to_numpy().transpose(), dtype=torch.long)\n",
        "            map_dict = {v.item():i for i,v in enumerate(torch.unique(edge_idx))}\n",
        "            map_edge = torch.zeros_like(edge_idx)\n",
        "            for k,v in map_dict.items():\n",
        "                map_edge[edge_idx==k] = v\n",
        "\n",
        "            # Convert the DataFrames into tensors\n",
        "            attrs = torch.tensor(attributes.to_numpy(), dtype=torch.float)\n",
        "            pad = torch.zeros((attrs.shape[0], 4), dtype=torch.float)\n",
        "            x = torch.cat((attrs, pad), dim=-1)\n",
        "\n",
        "            edge_idx = map_edge.long()\n",
        "\n",
        "            np_lab = label.to_numpy()\n",
        "            y = torch.tensor(np_lab if np_lab[0] == 1 else [0], dtype=torch.long)\n",
        "\n",
        "            graph = Data(x=x, edge_index=edge_idx,  y=y)\n",
        "\n",
        "            data_list.append(graph)\n",
        "\n",
        "        # Apply the functions specified in pre_filter and pre_transform\n",
        "        if self.pre_filter is not None:\n",
        "            data_list = [data for data in data_list if self.pre_filter(data)]\n",
        "\n",
        "        if self.pre_transform is not None:\n",
        "            data_list = [self.pre_transform(data) for data in data_list]\n",
        "\n",
        "        # Store the processed data\n",
        "        data, slices = self.collate(data_list)\n",
        "        torch.save((data, slices), self.processed_paths[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "07d079a9",
      "metadata": {
        "id": "07d079a9"
      },
      "outputs": [],
      "source": [
        "import torch_geometric.transforms as T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "98a1d072",
      "metadata": {
        "id": "98a1d072"
      },
      "outputs": [],
      "source": [
        "dataset = Frankenstein(root='data', pre_transform=T.GCNNorm())"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2fc95e99",
      "metadata": {
        "id": "2fc95e99"
      },
      "source": [
        "### Open Graph Benchmark datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "99c18fa1",
      "metadata": {
        "id": "99c18fa1"
      },
      "source": [
        "Open Graph Benchmark is available as a python library, to install it just run\n",
        "\n",
        "```pip install ogb```"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "86f5eab0",
      "metadata": {
        "id": "86f5eab0"
      },
      "source": [
        "OGB allows to load a dataset in three ways: for PyG applications, for DGL (Deep Graph Library, another widely used tool for GNNs in python) and in an 'agnostic' manner. There is a naming convention to load a dataset, depending on the task an the dataset name:\n",
        "\n",
        "    ogbn-[name]: for node tasks\n",
        "    ogbg-[name]: for graph tasks\n",
        "    ogbl-[name]: for link tasks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1c2a3931",
      "metadata": {
        "id": "1c2a3931"
      },
      "outputs": [],
      "source": [
        "from ogb.nodeproppred import PygNodePropPredDataset\n",
        "#from ogb.graphproppred import PygGraphPropPredDataset\n",
        "#from ogb.linkproppred import PygLinkPropPredDataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1ec243ea",
      "metadata": {
        "id": "1ec243ea"
      },
      "outputs": [],
      "source": [
        "dataset_name = 'ogbn-arxiv'\n",
        "dataset = PygNodePropPredDataset(name = dataset_name, root='data')\n",
        "\n",
        "split_idx = dataset.get_idx_split()\n",
        "train_idx, valid_idx, test_idx = split_idx[\"train\"], split_idx[\"valid\"], split_idx[\"test\"]\n",
        "graph = dataset[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9bc1d992",
      "metadata": {
        "id": "9bc1d992"
      },
      "outputs": [],
      "source": [
        "graph"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ac064ad9",
      "metadata": {
        "id": "ac064ad9"
      },
      "source": [
        "### Benchmarking Graph Neural Networks Datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "209977d0",
      "metadata": {
        "id": "209977d0"
      },
      "outputs": [],
      "source": [
        "import torch_geometric.datasets as datasets\n",
        "\n",
        "datasets.__all__"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5a8fda25",
      "metadata": {
        "id": "5a8fda25"
      },
      "outputs": [],
      "source": [
        "datasets.GNNBenchmarkDataset.names"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "070cf962",
      "metadata": {
        "id": "070cf962"
      },
      "outputs": [],
      "source": [
        "dataset = datasets.GNNBenchmarkDataset(name='MNIST', root='data')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c0137546",
      "metadata": {
        "id": "c0137546"
      },
      "outputs": [],
      "source": [
        "dataset"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.7"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}