examples/import-sample-export-gnn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "82568b27",
   "metadata": {},
   "source": [
    "# PyG integration: Sample and export"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4b7883a",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/neo4j/graph-data-science-client/blob/main/examples/import-sample-export-gnn.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "922d0fbb",
   "metadata": {},
   "source": [
    "This Jupyter notebook is hosted [here](https://github.com/neo4j/graph-data-science-client/blob/main/examples/import-sample-export-gnn.ipynb) in the Neo4j Graph Data Science Client Github repository.\n",
    "\n",
    "For a video presentation on this notebook, see the talk [GNNs at Scale With Graph Data Science Sampling and Python Client Integration](https://www.youtube.com/watch?v=c66u_wsPz-U&t=1082s) that was given at the NODES 2022 conference.\n",
    "\n",
    "The notebook exemplifies how to use the `graphdatascience` and PyTorch Geometric (PyG) Python libraries to:\n",
    "\n",
    "* Import the [CORA dataset](https://paperswithcode.com/dataset/cora) directly into GDS\n",
    "* Sample a part of CORA using the [GDS Random walk with restarts algorithm](https://neo4j.com/docs/graph-data-science/current/management-ops/projections/rwr/)\n",
    "* Export the CORA sample client side\n",
    "* Define and train a Graph Convolutional Neural Network (GCN) on the CORA sample\n",
    "* Evaluate the GCN on a test set"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c64321df",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "Running this notebook requires a Neo4j server with a recent GDS version (2.5+) installed.\n",
    "We recommend using Neo4j Desktop with GDS, or AuraDS.\n",
    "\n",
    "Also required are of course the Python libraries:\n",
    "\n",
    "* `graphdatascience` (see [docs](https://neo4j.com/docs/graph-data-science-client/current/installation/) for installation instructions)\n",
    "* PyG (see [PyG docs](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) for installation instructions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bdf33d6",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We start by importing our dependencies and setting up our GDS client connection to the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc4c3baf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install necessary dependencies\n",
    "%pip install graphdatascience torch torch_scatter torch_sparse torch_geometric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26fb276b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.nn import GCNConv\n",
    "from torch_geometric.transforms import RandomNodeSplit\n",
    "\n",
    "from graphdatascience import GraphDataScience"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce14da18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set seeds for consistent results\n",
    "random.seed(42)\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "torch.cuda.manual_seed_all(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "937f1476",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get Neo4j DB URI, credentials and name from environment if applicable\n",
    "NEO4J_URI = os.environ.get(\"NEO4J_URI\", \"bolt://localhost:7687\")\n",
    "NEO4J_AUTH = None\n",
    "NEO4J_DB = os.environ.get(\"NEO4J_DB\", \"neo4j\")\n",
    "if os.environ.get(\"NEO4J_USER\") and os.environ.get(\"NEO4J_PASSWORD\"):\n",
    "    NEO4J_AUTH = (\n",
    "        os.environ.get(\"NEO4J_USER\"),\n",
    "        os.environ.get(\"NEO4J_PASSWORD\"),\n",
    "    )\n",
    "gds = GraphDataScience(NEO4J_URI, auth=NEO4J_AUTH, database=NEO4J_DB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6aec1d31",
   "metadata": {
    "tags": [
     "verify-version"
    ]
   },
   "outputs": [],
   "source": [
    "from graphdatascience import ServerVersion\n",
    "\n",
    "assert gds.server_version() >= ServerVersion(2, 5, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1a8ab78",
   "metadata": {},
   "source": [
    "## Sampling CORA\n",
    "\n",
    "Next we use the built-in CORA loader to get the data into GDS. We will then sample it to get a smaller graph to train on.\n",
    "In a real world scenario, we would probably project data from a Neo4j database into GDS instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a575da60",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = gds.graph.load_cora()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "428ac7b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's make sure we constructed the correct graph\n",
    "print(f\"Metadata for our loaded Cora graph `G`: {G}\")\n",
    "print(f\"Node labels present in `G`: {G.node_labels()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e8b8a83",
   "metadata": {},
   "source": [
    "It's looks correct! Now let's go ahead and sample the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d116738",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use the random walk with restarts sampling algorithm with default values\n",
    "G_sample, _ = gds.graph.sample.rwr(\"cora_sample\", G, randomSeed=42, concurrency=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d66ea2d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We should have somewhere around 0.15 * 2708 ~ 406 nodes in our sample\n",
    "print(f\"Number of nodes in our sample: {G_sample.node_count()}\")\n",
    "\n",
    "# And let's see how many relationships we got\n",
    "print(f\"Number of relationships in our sample: {G_sample.relationship_count()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "722caa7b",
   "metadata": {},
   "source": [
    "## Exporting sampled CORA\n",
    "\n",
    "We can now export the topology and node properties of the sampled graph that we need to train our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3c76b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the relationship data from our sample\n",
    "sample_topology_df = gds.graph.relationships.stream(G_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8308917d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see what we got:\n",
    "display(sample_topology_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4ad967b",
   "metadata": {},
   "source": [
    "We get the right amount of rows, one for each expected relationship. However, the node IDs are rather large whereas PyG expects consecutive IDs starting from zero. We now will start to massage our data structure containing the relationships until PyG can use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b64d6ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# By using `by_rel_type` we get the topology in a format that can be used as input to several GNN frameworks:\n",
    "# {\"rel_type\": [[source_nodes], [target_nodes]]}\n",
    "\n",
    "sample_topology = sample_topology_df.by_rel_type()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09e0f763",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We should only have the \"CITES\" keys since there's only one relationship type\n",
    "print(f\"Relationship type keys: {sample_topology.keys()}\")\n",
    "print(f\"Number of  {len(sample_topology['CITES'])}\")\n",
    "\n",
    "# How many source nodes do we have?\n",
    "print(len(sample_topology[\"CITES\"][0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b87a5cb",
   "metadata": {},
   "source": [
    "Great, it looks like we have the format we need to create a PyG `edge_index` later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c52e3b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We also need to export the node properties corresponding to our node labels and features, represented by the\n",
    "# \"subject\" and \"features\" node properties respectively\n",
    "sample_node_properties = gds.graph.nodeProperties.stream(\n",
    "    G_sample,\n",
    "    [\"subject\", \"features\"],\n",
    "    separate_property_columns=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb81fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's make sure we got the data we expected\n",
    "display(sample_node_properties)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75ce6f84",
   "metadata": {},
   "source": [
    "## Constructing GCN input\n",
    "\n",
    "Now that we have all information we need client side, we can construct the PyG `Data` object we will use as training input. We will remap the node IDs to be consecutive and starting from zero. We use the ordering of node ids in `sample_node_properties` as our remapping so that the index is aligned with the node properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f7027dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In order for the node ids used in the `topology` to be consecutive and starting from zero,\n",
    "# we will need to remap them. This way they will also align with the row numbering of the\n",
    "# `sample_node_properties` data frame\n",
    "def normalize_topology_index(new_idx_to_old, topology):\n",
    "    # Create a reverse mapping based on new idx -> old idx\n",
    "    old_idx_to_new = dict((v, k) for k, v in new_idx_to_old.items())\n",
    "    return [[old_idx_to_new[node_id] for node_id in nodes] for nodes in topology]\n",
    "\n",
    "\n",
    "# We use the ordering of node ids in `sample_node_properties` as our remapping\n",
    "# The result is: [[mapped_source_nodes], [mapped_target_nodes]]\n",
    "normalized_topology = normalize_topology_index(dict(sample_node_properties[\"nodeId\"]), sample_topology[\"CITES\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d795e325",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use the ordering of node ids in `sample_node_properties` as our remapping\n",
    "edge_index = torch.tensor(normalized_topology, dtype=torch.long)\n",
    "\n",
    "# We specify the node property \"features\" as the zero-layer node embeddings\n",
    "x = torch.tensor(sample_node_properties[\"features\"], dtype=torch.float)\n",
    "\n",
    "# We specify the node property \"subject\" as class labels\n",
    "y = torch.tensor(sample_node_properties[\"subject\"], dtype=torch.long)\n",
    "\n",
    "data = Data(x=x, y=y, edge_index=edge_index)\n",
    "\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1885ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do a random split of the data so that ~10% goes into a test set and the rest used for training\n",
    "transform = RandomNodeSplit(num_test=40, num_val=0)\n",
    "data = transform(data)\n",
    "\n",
    "# We can see that our `data` object have been extended with some masks defining the split\n",
    "print(data)\n",
    "print(data.train_mask.sum().item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "531854f4",
   "metadata": {},
   "source": [
    "As a sidenote, if we had wanted to do some hyperarameter tuning, it would have been useful to keep some data for a validation set as well."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9ceb11",
   "metadata": {},
   "source": [
    "## Training and evaluating a GCN\n",
    "\n",
    "Let's now define and train the GCN using PyG and our sampled CORA as input. We adapt the CORA GCN example from the [PyG documentation](https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html#learning-methods-on-graphs).\n",
    "\n",
    "In this example we evaluate the model on a test set of the sampled CORA. Please note however, that since GCN is an inductive algorithm we could also have evaluated it on the full CORA dataset, or even another (similar) graph for that matter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74b65751",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = y.unique().shape[0]\n",
    "\n",
    "\n",
    "# Define the GCN architecture\n",
    "class GCN(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(data.num_node_features, 16)\n",
    "        self.conv2 = GCNConv(16, num_classes)\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index = data.x, data.edge_index\n",
    "\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "\n",
    "        # We use log_softmax and nll_loss instead of softmax output and cross entropy loss\n",
    "        # for reasons for performance and numerical stability.\n",
    "        # They are mathematically equivalent\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eeaef0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare training by setting up for the chosen device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Let's see what device was chosen\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff062e04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In standard PyTorch fashion we instantiate our model, and transfer it to the memory of the chosen device\n",
    "model = GCN().to(device)\n",
    "\n",
    "# Let's inspect our model architecture\n",
    "print(model)\n",
    "\n",
    "# Pass our input data to the chosen device too\n",
    "data = data.to(device)\n",
    "\n",
    "# Since hyperparameter tuning is out of scope for this small example, we initialize an\n",
    "# Adam optimizer with some fixed learning rate and weight decay\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "483d90d4",
   "metadata": {},
   "source": [
    "From inspecting the model we can see that that the output size is 7, which looks correct since Cora does indeed have 7 different paper subjects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02e8fadb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the GCN using the CORA sample represented by `data` using the standard PyTorch training loop\n",
    "model.train()\n",
    "for epoch in range(200):\n",
    "    optimizer.zero_grad()\n",
    "    out = model(data)\n",
    "    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9169bf30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the trained GCN model on our test set\n",
    "model.eval()\n",
    "pred = model(data).argmax(dim=1)\n",
    "correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()\n",
    "acc = int(correct) / int(data.test_mask.sum())\n",
    "\n",
    "print(f\"Accuracy: {acc:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a61f680a",
   "metadata": {},
   "source": [
    "The accuracy looks good. The next step would be to run the GCN model we trained our subsample on the entire Cora graph. This part is left as an exercise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "928156de",
   "metadata": {},
   "source": [
    "## Cleanup\n",
    "\n",
    "We remove the CORA graphs from the GDS graph catalog."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33523a0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = G_sample.drop()\n",
    "_ = G.drop()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}