examples/ml-pipelines-node-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/neo4j/graph-data-science-client/blob/main/examples/ml-pipelines-node-classification.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KhH0v_bwUrXp"
   },
   "source": [
    "# Machine learning pipelines: Node classification\n",
    "\n",
    "This notebook shows the usage of GDS machine learning pipelines with the Python client and the well-known [Cora dataset](https://paperswithcode.com/dataset/cora).\n",
    "\n",
    "The task we cover here is a typical use case in graph machine learning: the classification of nodes given a graph and some node features."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X5ttqidnWPTC"
   },
   "source": [
    "## Setup\n",
    "\n",
    "We need a dedicated environment where Neo4j and GDS are available, for example a fresh AuraDS instance (which comes with GDS preinstalled) or Neo4j Desktop with a dedicated database. \n",
    "\n",
    "**Please note that we will be writing to and deleting data from Neo4j.**\n",
    "\n",
    "Once the credentials to access this environment are available, we can install the `graphdatascience` package and import the client class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "izJD_HlcAdQ0"
   },
   "outputs": [],
   "source": [
    "%pip install graphdatascience"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphdatascience import GraphDataScience"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using a local Neo4j setup, the default connection URI is `bolt://localhost:7687`; when using AuraDS, instead, the connection URI is slightly different as it uses the `neo4j+s` protocol. In this case, the client should also include the `aura_ds=True` flag to enable AuraDS-recommended settings. Check the [Neo4j GDS Client docs](https://neo4j.com/docs/graph-data-science-client/current/getting-started/) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gJLQgHK_AghP"
   },
   "outputs": [],
   "source": [
    "# Replace with the actual connection URI and credentials\n",
    "NEO4J_CONNECTION_URI = \"bolt://localhost:7687\"\n",
    "NEO4J_USERNAME = \"neo4j\"\n",
    "NEO4J_PASSWORD = \"test\"\n",
    "\n",
    "gds = GraphDataScience(NEO4J_CONNECTION_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD), database=\"neo4j\")\n",
    "\n",
    "# On AuraDS:\n",
    "#\n",
    "# NEO4J_CONNECTION_URI = \"neo4j+s://xxxxxxxx.databases.neo4j.io\"\n",
    "# NEO4J_USERNAME = \"neo4j\"\n",
    "# NEO4J_PASSWORD = \"...\"\n",
    "#\n",
    "# gds = GraphDataScience(NEO4J_CONNECTION_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD), aura_ds=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need to check that the version of the GDS library is 2.2.0 or newer, because we will use the concept of \"context\" introduced in [GDS 2.2.0](https://github.com/neo4j/graph-data-science/releases/tag/2.2.0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert gds.version() >= \"2.2.0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zOK08QpoWFKv"
   },
   "source": [
    "Finally, we import `json` to help in writing the Cypher queries used to load the data, and `numpy` and `pandas` for further data processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-ikcvI-cSDoK"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oDE2LqTXAmy8"
   },
   "source": [
    "## Loading the Cora dataset\n",
    "\n",
    "First of all, we need to load the Cora dataset on Neo4j. The latest versions of the GDS client include the Cora dataset as a ready-to-use graph (see for instance the [PyG example notebook](https://github.com/neo4j/graph-data-science-client/blob/main/examples/import-sample-export-gnn.ipynb)); alternatively, the [graph construction notebook](https://github.com/neo4j/graph-data-science-client/blob/main/examples/load-data-via-graph-construction.ipynb) shows how to project the Cora graph in memory without writing it to Neo4j. In this tutorial, anyway, we use data from CSV files and some Cypher code to run an end-to-end example, from loading the source data into Neo4j to training a model and using it for predictions.\n",
    "\n",
    "**Please note that, if you use the Cora graph loader or the graph construction method on an AuraDS instance, you cannot write the data to the Neo4j database.**\n",
    "\n",
    "The CSV files can be found at the following URIs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tOZW-I2FBLm_"
   },
   "outputs": [],
   "source": [
    "CORA_CONTENT = (\n",
    "    \"https://raw.githubusercontent.com/neo4j/graph-data-science/master/test-utils/src/main/resources/cora.content\"\n",
    ")\n",
    "CORA_CITES = (\n",
    "    \"https://raw.githubusercontent.com/neo4j/graph-data-science/master/test-utils/src/main/resources/cora.cites\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upon loading, we need to perform an additional preprocessing step to convert the `subject` field (which is a string in the dataset) into an integer, because node properties have to be numerical in order to be projected into a graph; although we could assign consecutive IDs, we assign an ID other than 0 to the first subject to later show how the class labels are represented in the model.\n",
    "\n",
    "We also select a number of nodes to be held out to test the model after it has been trained. **NOTE:** This is not related to the algorithm test/split ratio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZwPI7OHSGj1L"
   },
   "outputs": [],
   "source": [
    "SUBJECT_TO_ID = {\n",
    "    \"Neural_Networks\": 100,\n",
    "    \"Rule_Learning\": 1,\n",
    "    \"Reinforcement_Learning\": 2,\n",
    "    \"Probabilistic_Methods\": 3,\n",
    "    \"Theory\": 4,\n",
    "    \"Genetic_Algorithms\": 5,\n",
    "    \"Case_Based\": 6,\n",
    "}\n",
    "\n",
    "HOLDOUT_NODES = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now load the CSV files using the `LOAD CSV` Cypher statement and some basic data transformation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d_p_Npg3ecyz"
   },
   "outputs": [],
   "source": [
    "# Define a string representation of the SUBJECT_TO_ID map using backticks\n",
    "subject_map = json.dumps(SUBJECT_TO_ID).replace('\"', \"`\")\n",
    "\n",
    "# Cypher command to load the nodes using `LOAD CSV`, taking care of\n",
    "# converting the string `subject` field into an integer and\n",
    "# replacing the node label for the holdout nodes\n",
    "load_nodes = f\"\"\"\n",
    "    LOAD CSV FROM \"{CORA_CONTENT}\" AS row\n",
    "    WITH \n",
    "      {subject_map} AS subject_to_id,\n",
    "      toInteger(row[0]) AS extId, \n",
    "      row[1] AS subject, \n",
    "      toIntegerList(row[2..]) AS features\n",
    "    MERGE (p:Paper {{extId: extId, subject: subject_to_id[subject], features: features}})\n",
    "    WITH p LIMIT {HOLDOUT_NODES}\n",
    "    REMOVE p:Paper\n",
    "    SET p:UnclassifiedPaper\n",
    "\"\"\"\n",
    "\n",
    "# Cypher command to load the relationships using `LOAD CSV`\n",
    "load_relationships = f\"\"\"\n",
    "    LOAD CSV FROM \"{CORA_CITES}\" AS row\n",
    "    MATCH (n), (m) \n",
    "    WHERE n.extId = toInteger(row[0]) AND m.extId = toInteger(row[1])\n",
    "    MERGE (n)-[:CITES]->(m)\n",
    "\"\"\"\n",
    "\n",
    "# Load nodes and relationships on Neo4j\n",
    "gds.run_cypher(load_nodes)\n",
    "gds.run_cypher(load_relationships)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AmFOLcCunfXB"
   },
   "source": [
    "With the data loaded on Neo4j, we can now project a graph including all the nodes and the `CITES` relationship as undirected (and with `SINGLE` aggregation, to skip repeated relationships as a result of adding the inverse direction)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KdMPbSQ9OKZ6"
   },
   "outputs": [],
   "source": [
    "# Create the projected graph containing both classified and unclassified nodes\n",
    "G, _ = gds.graph.project(\n",
    "    \"cora-graph\",\n",
    "    {\"Paper\": {\"properties\": [\"features\", \"subject\"]}, \"UnclassifiedPaper\": {\"properties\": [\"features\"]}},\n",
    "    {\"CITES\": {\"orientation\": \"UNDIRECTED\", \"aggregation\": \"SINGLE\"}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can finally check the number of nodes and relationships in the newly-projected graph to make sure it has been created correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert G.node_count() == 2708\n",
    "assert G.relationship_count() == 10556"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uGWiTTrGN-Be"
   },
   "source": [
    "## Pipeline catalog basics\n",
    "\n",
    "Once the dataset has been loaded, we can define a node classification machine learning pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mP64bLMxN_IB"
   },
   "outputs": [],
   "source": [
    "# Create the pipeline\n",
    "node_pipeline, _ = gds.beta.pipeline.nodeClassification.create(\"cora-pipeline\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can check that the pipeline has actually been created with the `list` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jguecWUjQ6DV"
   },
   "outputs": [],
   "source": [
    "# List all pipelines\n",
    "gds.beta.pipeline.list()\n",
    "\n",
    "# Alternatively, get the details of a specific pipeline object\n",
    "gds.beta.pipeline.list(node_pipeline)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cil3MMlTtZ7L"
   },
   "source": [
    "## Configuring the pipeline\n",
    "\n",
    "We can now configure the pipeline. As a reminder, we need to:\n",
    "\n",
    "1. Select a subset of the available node properties to be used as features for the machine learning model\n",
    "1. Configure the train/test split and the number of folds for k-fold cross-validation _(optional)_\n",
    "1. Configure the candidate models for training\n",
    "1. Configure autotuning _(optional)_\n",
    "In this example we use Logistic Regression as a candidate model for the training, but other algorithms (such as Random Forest) are available as well. We also set some reasonable starting parameters that can be further tuned according to the needed metrics.\n",
    "\n",
    "Some hyperparameters such as `penalty` can be single values or ranges. If they are expressed as ranges, autotuning is used to search their best value.\n",
    "\n",
    "The `configureAutoTuning` method can be used to set the number of model candidates to try. Here we choose 5 to keep the training time short."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"Mark\" some node properties that will be used as features\n",
    "node_pipeline.selectFeatures([\"features\"])\n",
    "\n",
    "# If needed, change the train/test split ratio and the number of folds\n",
    "# for k-fold cross-validation\n",
    "node_pipeline.configureSplit(testFraction=0.2, validationFolds=5)\n",
    "\n",
    "# Add a model candidate to train\n",
    "node_pipeline.addLogisticRegression(maxEpochs=200, penalty=(0.0, 0.5))\n",
    "\n",
    "# Explicit set the number of trials for autotuning (default = 10)\n",
    "node_pipeline.configureAutoTuning(maxTrials=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qFIfl6yYXQgY"
   },
   "source": [
    "## Training the pipeline\n",
    "\n",
    "The configured pipeline is now ready to select and train a model. We also run a training estimate, to make sure there are enough resources to run the actual training afterwards.\n",
    "\n",
    "The Node Classification model supports several evaluation metrics. Here we use the global metric `F1_WEIGHTED`.\n",
    "\n",
    "**NOTE:** The `concurrency` parameter is explicitly set to 4 (the default value) for demonstration purposes. \n",
    "The maximum concurrency in the library is limited to 4 for Neo4j Community Edition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N5Eey0EsOXiG"
   },
   "outputs": [],
   "source": [
    "# Estimate the resources needed for training the model\n",
    "node_pipeline.train_estimate(\n",
    "    G,\n",
    "    targetNodeLabels=[\"Paper\"],\n",
    "    modelName=\"cora-pipeline-model\",\n",
    "    targetProperty=\"subject\",\n",
    "    metrics=[\"F1_WEIGHTED\"],\n",
    "    randomSeed=42,\n",
    "    concurrency=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "GVInjJHwOexk"
   },
   "outputs": [],
   "source": [
    "# Perform the actual training\n",
    "model, stats = node_pipeline.train(\n",
    "    G,\n",
    "    targetNodeLabels=[\"Paper\"],\n",
    "    modelName=\"cora-pipeline-model\",\n",
    "    targetProperty=\"subject\",\n",
    "    metrics=[\"F1_WEIGHTED\"],\n",
    "    randomSeed=42,\n",
    "    concurrency=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kxYPHz6seMEn"
   },
   "source": [
    "We can inspect the result of the training, for example to print the evaluation metrics of the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "B1pvnkbSZK9A"
   },
   "outputs": [],
   "source": [
    "# Uncomment to print all stats\n",
    "# print(stats.to_json(indent=2))\n",
    "\n",
    "# Print F1_WEIGHTED metric\n",
    "stats[\"modelInfo\"][\"metrics\"][\"F1_WEIGHTED\"][\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QBEfXROYfNES"
   },
   "source": [
    "## Using the model for prediction\n",
    "\n",
    "After training, the model is ready to classify unclassified data. \n",
    "\n",
    "One simple way to use the `predict` mode is to just stream the result of the prediction. This can be impractical when a graph is very large, so it should be only used for experimentation purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "92XksLN9Om4l"
   },
   "outputs": [],
   "source": [
    "predicted = model.predict_stream(\n",
    "    G, modelName=\"cora-pipeline-model\", includePredictedProbabilities=True, targetNodeLabels=[\"UnclassifiedPaper\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "knLyBCSUgR2-"
   },
   "source": [
    "The result of the prediction is a Pandas `DataFrame` containing the predicted class and the predicted probabilities for all the classes for each node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "pw5AmIAW7rBT"
   },
   "outputs": [],
   "source": [
    "predicted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lJ2iADIDgwmH"
   },
   "source": [
    "The order of the classes in the `predictedProbabilities` field is given in the model information, and can be used to retrieve the predicted probability for the predicted class. \n",
    "\n",
    "Please note that the order in which the classes appear in the `predictedProbabilities` field is somewhat arbitrary, so the correct way to access each probability is via the class index obtained from the model, _not_ its position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of class labels\n",
    "classes = stats[\"modelInfo\"][\"classes\"]\n",
    "print(\"Class labels:\", classes)\n",
    "\n",
    "# Calculate the confidence percentage for the predicted class\n",
    "predicted[\"confidence\"] = predicted.apply(\n",
    "    lambda row: np.floor(row[\"predictedProbabilities\"][classes.index(row[\"predictedClass\"])] * 100), axis=1\n",
    ")\n",
    "\n",
    "predicted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rL7Mi1AIc7OJ"
   },
   "source": [
    "## Adding a data preprocessing step\n",
    "\n",
    "The quality of the model can potentially be increased by adding more features or by using different features altogether. One way is to use algorithms such as FastRP that create embeddings based on both node properties and graph features, which can be added via the `addNodeProperty` pipeline method. Such properties are \"transient\", in that they are automatically created and removed by the pipeline itself.\n",
    "\n",
    "In this example we also use the `contextNodeLabels` parameter to explicitly set the types of nodes we calculate the embeddings for, and we include both the classified and the unclassified nodes. This is useful because the more nodes are used, the better the generated embeddings are. Although it may seem counterintuitive, unclassified nodes do not need to be completely unobserved during training (so, for instance, information on their neighbours can be retained). More information can be found in graph ML publications such as the [Graph Representation Learning Book](https://www.cs.mcgill.ca/~wlh/grl_book/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "F_z0gLG9QCAn"
   },
   "outputs": [],
   "source": [
    "node_pipeline_fastrp, _ = gds.beta.pipeline.nodeClassification.create(\"cora-pipeline-fastrp\")\n",
    "\n",
    "# Add a step in the pipeline that mutates the graph\n",
    "node_pipeline_fastrp.addNodeProperty(\n",
    "    \"fastRP\",\n",
    "    mutateProperty=\"embedding\",\n",
    "    embeddingDimension=512,\n",
    "    propertyRatio=1.0,\n",
    "    randomSeed=42,\n",
    "    featureProperties=[\"features\"],\n",
    "    contextNodeLabels=[\"Paper\", \"UnclassifiedPaper\"],\n",
    ")\n",
    "\n",
    "# With the node embeddings available as features, we no longer use the original raw `features`.\n",
    "node_pipeline_fastrp.selectFeatures([\"embedding\"])\n",
    "\n",
    "# Configure the pipeline as before\n",
    "node_pipeline_fastrp.configureSplit(testFraction=0.2, validationFolds=5)\n",
    "node_pipeline_fastrp.addLogisticRegression(maxEpochs=200, penalty=(0.0, 0.5))\n",
    "node_pipeline.configureAutoTuning(maxTrials=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training then proceeds as in the previous section:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "Kvr9iU4HdoTv"
   },
   "outputs": [],
   "source": [
    "# Perform the actual training\n",
    "model_fastrp, stats_fastrp = node_pipeline_fastrp.train(\n",
    "    G,\n",
    "    targetNodeLabels=[\"Paper\"],\n",
    "    modelName=\"cora-pipeline-model-fastrp\",\n",
    "    targetProperty=\"subject\",\n",
    "    metrics=[\"F1_WEIGHTED\"],\n",
    "    randomSeed=42,\n",
    "    concurrency=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `F1_WEIGHTED` metrics is better with embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "_Q3zMmpCd2o0"
   },
   "outputs": [],
   "source": [
    "print(stats_fastrp[\"modelInfo\"][\"metrics\"][\"F1_WEIGHTED\"][\"test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The classification using `predict_stream` can be run in the same way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "vgfN-ppHw4Rp"
   },
   "outputs": [],
   "source": [
    "predicted_fastrp = model_fastrp.predict_stream(\n",
    "    G,\n",
    "    modelName=\"cora-pipeline-model-fastrp\",\n",
    "    includePredictedProbabilities=True,\n",
    "    targetNodeLabels=[\"UnclassifiedPaper\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(predicted_fastrp))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PNbRL2y84F6Q"
   },
   "source": [
    "Instead of streaming the results, the prediction can be run in `mutate` mode to be more performant, especially when the predicted values are used multiple times. The predicted nodes can be retrieved using the `streamNodeProperty` method with the `UnclassifiedPaper` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "_HVufWHbxIoE"
   },
   "outputs": [],
   "source": [
    "model_fastrp.predict_mutate(\n",
    "    G,\n",
    "    mutateProperty=\"predictedClass\",\n",
    "    modelName=\"cora-pipeline-model-fastrp\",\n",
    "    predictedProbabilityProperty=\"predictedProbabilities\",\n",
    "    targetNodeLabels=[\"UnclassifiedPaper\"],\n",
    ")\n",
    "\n",
    "predicted_fastrp = gds.graph.streamNodeProperty(G, \"predictedClass\", [\"UnclassifiedPaper\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_fastrp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is useful to compare the result of classification with the original `subject` value of the test nodes, which must be retrieved from the Neo4j database since it has been excluded from the projected graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "thVR-anwmmzo"
   },
   "outputs": [],
   "source": [
    "# Retrieve node information from Neo4j using the node IDs from the prediction result\n",
    "nodes = gds.util.asNodes(predicted_fastrp.nodeId.to_list())\n",
    "\n",
    "# Create a new DataFrame containing node IDs along with node properties\n",
    "nodes_df = pd.DataFrame([(node.id, node[\"subject\"]) for node in nodes], columns=[\"nodeId\", \"subject\"])\n",
    "\n",
    "# Merge with the prediction result on node IDs, to check the predicted value\n",
    "# against the original subject\n",
    "#\n",
    "# NOTE: This could also be replaced by just appending `node[\"subject\"]` as a\n",
    "# Series since the node order would not change, but a proper merge (or join)\n",
    "# is clearer and less prone to errors.\n",
    "predicted_fastrp.merge(nodes_df, on=\"nodeId\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z27MnchLr9mi"
   },
   "source": [
    "As we can see, the prediction for all the test nodes is accurate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UdEVypaiMc5d"
   },
   "source": [
    "## Writing result back to Neo4j"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having the predicted class written back to the graph, we can now write them back to the Neo4j database.\n",
    "\n",
    "**Please note that this step is not applicable if you are running this notebook on AuraDS.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "aKuUQwxbPOaV"
   },
   "outputs": [],
   "source": [
    "gds.graph.nodeProperties.write(\n",
    "    G,\n",
    "    node_properties=[\"predictedClass\"],\n",
    "    node_labels=[\"UnclassifiedPaper\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "plRCiikGOofd"
   },
   "source": [
    "## Cleanup\n",
    "\n",
    "When the graph, the model and the pipeline are no longer needed, they should be dropped to free up memory. This only needs to be done if the Neo4j or AuraDS instance is not restarted, since a restart would clean up all the in-memory content anyway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RFWzcrX3ORXI"
   },
   "outputs": [],
   "source": [
    "model.drop()\n",
    "model_fastrp.drop()\n",
    "node_pipeline.drop()\n",
    "node_pipeline_fastrp.drop()\n",
    "\n",
    "G.drop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Neo4j database instead needs to be cleaned up explicitly if no longer useful:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(\"MATCH (n:Paper|UnclassifiedPaper) DETACH DELETE n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is good practice to close the client as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.close()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyN7QqVTzAahL3BPKHgRt8jS",
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Machine learning pipelines - Node classification example",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.5 ('venv': 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.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "d9b85228d1674a60a9aba6bdb3ec2509e82e7a19678e59e4fceb33f4c3ba222b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}