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.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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X5ttqidnWPTC"
   },
   "source": [
    "## Setup\n",
    "\n",
    "We need an environment where Neo4j and GDS are available, for example AuraDS (which comes with GDS preinstalled) or Neo4j Desktop. \n",
    "\n",
    "Once the credentials to 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`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "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 = \"\"\n",
    "\n",
    "gds = GraphDataScience(NEO4J_CONNECTION_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using AuraDS, the connection URI is slightly different as it uses the `neo4j+s` protocol. The client should also include the `aura_ds=True` flag to enable AuraDS-recommended settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace with the actual connection URI and credentials\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": {
    "id": "zOK08QpoWFKv"
   },
   "source": [
    "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": 38,
   "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": [
    "## Load the Cora dataset\n",
    "\n",
    "First of all, we need to load the Cora dataset on Neo4j. The CSV files can be found at the following URIs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "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 also 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 class to show how the classes 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": 40,
   "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` 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": 42,
   "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": {
    "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": 43,
   "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",
    "# List a specific pipeline object\n",
    "gds.beta.pipeline.list(node_pipeline)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cil3MMlTtZ7L"
   },
   "source": [
    "## Configure 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PcTDdBnxOA0y"
   },
   "outputs": [],
   "source": [
    "# \"Mark\" some node properties that will be used as features\n",
    "node_pipeline.selectFeatures([\"features\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s8ar00tHEmt5"
   },
   "outputs": [],
   "source": [
    "# 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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_Jtfr236EO0N"
   },
   "source": [
    "Here we use Logistic Regression as an example for the training, but other algorithms (such as Random Forest) are available as well.\n",
    "\n",
    "Some hyperparameters such as `penalty` can be single values or ranges. If they are expressed as ranges, auto-tuning is used to search their best value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vA2jFvx-OHSS"
   },
   "outputs": [],
   "source": [
    "# Add a model candidate to train\n",
    "node_pipeline.addLogisticRegression(maxEpochs=1000, penalty=(0.00038, 0.00042))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qFIfl6yYXQgY"
   },
   "source": [
    "## Run the training\n",
    "\n",
    "It is now possible to train the configured models. 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",
    "    nodeLabels=[\"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",
    "    nodeLabels=[\"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": 50,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "B1pvnkbSZK9A"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7131261628836523"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": [
    "## Use the model for prediction\n",
    "\n",
    "After the model has been trained, it is possible to use it 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 used only for experimentation purposes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ryg5xlPBwE3d"
   },
   "source": [
    "In this example we use the `nodeLabels=[\"UnclassifiedPaper\"]` filter to only run prediction on the unclassified nodes. It must be noted that, when using models that have `nodePropertySteps` that use relationships (such as FastRP and other models that create embeddings, but also algorithms like PageRank), the `nodeLabels` filter should not be used because it would \"cut out\" all the linked nodes that have a different label. We will see an example of this in the following section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "92XksLN9Om4l"
   },
   "outputs": [],
   "source": [
    "predicted = model.predict_stream(\n",
    "    G, modelName=\"cora-pipeline-model\", includePredictedProbabilities=True, nodeLabels=[\"UnclassifiedPaper\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "knLyBCSUgR2-"
   },
   "source": [
    "The result of the prediction is a DataFrame containing the predicted class and the predicted probabilities for all the classes for each node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "pw5AmIAW7rBT"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>nodeId</th>\n",
       "      <th>predictedClass</th>\n",
       "      <th>predictedProbabilities</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15188</td>\n",
       "      <td>100</td>\n",
       "      <td>[0.013587384554264968, 0.0008230294746663742, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15189</td>\n",
       "      <td>1</td>\n",
       "      <td>[0.3197115910761458, 0.02600152888962292, 0.12...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15190</td>\n",
       "      <td>2</td>\n",
       "      <td>[0.0077080939234027215, 0.7666212201367124, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>15191</td>\n",
       "      <td>2</td>\n",
       "      <td>[0.003479178144017117, 0.976798336373076, 0.01...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>15192</td>\n",
       "      <td>3</td>\n",
       "      <td>[0.030846315380417298, 0.0013541278264996802, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>15193</td>\n",
       "      <td>1</td>\n",
       "      <td>[0.2743067591955607, 0.26636894112196013, 0.09...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>15194</td>\n",
       "      <td>6</td>\n",
       "      <td>[0.017862637875544374, 0.06103214395299978, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>15195</td>\n",
       "      <td>100</td>\n",
       "      <td>[0.0028305594972709534, 0.004251281491345211, ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>15196</td>\n",
       "      <td>100</td>\n",
       "      <td>[0.018424478556861662, 0.026205387730423434, 0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>15197</td>\n",
       "      <td>4</td>\n",
       "      <td>[0.01153218347682615, 0.3319040026951682, 0.00...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   nodeId  predictedClass                             predictedProbabilities\n",
       "0   15188             100  [0.013587384554264968, 0.0008230294746663742, ...\n",
       "1   15189               1  [0.3197115910761458, 0.02600152888962292, 0.12...\n",
       "2   15190               2  [0.0077080939234027215, 0.7666212201367124, 0....\n",
       "3   15191               2  [0.003479178144017117, 0.976798336373076, 0.01...\n",
       "4   15192               3  [0.030846315380417298, 0.0013541278264996802, ...\n",
       "5   15193               1  [0.2743067591955607, 0.26636894112196013, 0.09...\n",
       "6   15194               6  [0.017862637875544374, 0.06103214395299978, 0....\n",
       "7   15195             100  [0.0028305594972709534, 0.004251281491345211, ...\n",
       "8   15196             100  [0.018424478556861662, 0.026205387730423434, 0...\n",
       "9   15197               4  [0.01153218347682615, 0.3319040026951682, 0.00..."
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "lbxPNQsg4ByD"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 100]\n"
     ]
    }
   ],
   "source": [
    "classes = stats[\"modelInfo\"][\"classes\"]\n",
    "print(classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "YkkX3JPz7uxA"
   },
   "outputs": [],
   "source": [
    "# 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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "L8pUbVA7-F98"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>nodeId</th>\n",
       "      <th>predictedClass</th>\n",
       "      <th>predictedProbabilities</th>\n",
       "      <th>confidence</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15188</td>\n",
       "      <td>100</td>\n",
       "      <td>[0.013587384554264968, 0.0008230294746663742, ...</td>\n",
       "      <td>73.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15189</td>\n",
       "      <td>1</td>\n",
       "      <td>[0.3197115910761458, 0.02600152888962292, 0.12...</td>\n",
       "      <td>31.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15190</td>\n",
       "      <td>2</td>\n",
       "      <td>[0.0077080939234027215, 0.7666212201367124, 0....</td>\n",
       "      <td>76.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>15191</td>\n",
       "      <td>2</td>\n",
       "      <td>[0.003479178144017117, 0.976798336373076, 0.01...</td>\n",
       "      <td>97.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>15192</td>\n",
       "      <td>3</td>\n",
       "      <td>[0.030846315380417298, 0.0013541278264996802, ...</td>\n",
       "      <td>87.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>15193</td>\n",
       "      <td>1</td>\n",
       "      <td>[0.2743067591955607, 0.26636894112196013, 0.09...</td>\n",
       "      <td>27.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>15194</td>\n",
       "      <td>6</td>\n",
       "      <td>[0.017862637875544374, 0.06103214395299978, 0....</td>\n",
       "      <td>82.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>15195</td>\n",
       "      <td>100</td>\n",
       "      <td>[0.0028305594972709534, 0.004251281491345211, ...</td>\n",
       "      <td>89.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>15196</td>\n",
       "      <td>100</td>\n",
       "      <td>[0.018424478556861662, 0.026205387730423434, 0...</td>\n",
       "      <td>54.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>15197</td>\n",
       "      <td>4</td>\n",
       "      <td>[0.01153218347682615, 0.3319040026951682, 0.00...</td>\n",
       "      <td>46.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   nodeId  predictedClass                             predictedProbabilities  \\\n",
       "0   15188             100  [0.013587384554264968, 0.0008230294746663742, ...   \n",
       "1   15189               1  [0.3197115910761458, 0.02600152888962292, 0.12...   \n",
       "2   15190               2  [0.0077080939234027215, 0.7666212201367124, 0....   \n",
       "3   15191               2  [0.003479178144017117, 0.976798336373076, 0.01...   \n",
       "4   15192               3  [0.030846315380417298, 0.0013541278264996802, ...   \n",
       "5   15193               1  [0.2743067591955607, 0.26636894112196013, 0.09...   \n",
       "6   15194               6  [0.017862637875544374, 0.06103214395299978, 0....   \n",
       "7   15195             100  [0.0028305594972709534, 0.004251281491345211, ...   \n",
       "8   15196             100  [0.018424478556861662, 0.026205387730423434, 0...   \n",
       "9   15197               4  [0.01153218347682615, 0.3319040026951682, 0.00...   \n",
       "\n",
       "   confidence  \n",
       "0        73.0  \n",
       "1        31.0  \n",
       "2        76.0  \n",
       "3        97.0  \n",
       "4        87.0  \n",
       "5        27.0  \n",
       "6        82.0  \n",
       "7        89.0  \n",
       "8        54.0  \n",
       "9        46.0  "
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.set_option(\"display.precision\", 6)\n",
    "\n",
    "predicted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rL7Mi1AIc7OJ"
   },
   "source": [
    "## Adding a data preprocessing step\n",
    "\n",
    "The performance of the model can potentially be increased by adding more features or by using different features altogether. One way is to use models that create embeddings based on both node properties and graph features. One of such models is FastRP, which can be added via the `addNodeProperty` pipeline method.\n",
    "\n",
    "More embedding methods are available in GDS, as well as other pre-processing algorithms."
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lM51DhPWX4e3"
   },
   "source": [
    "With the node embeddings available as features, we no longer use the original raw `features`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "2uQY--MDdePK"
   },
   "outputs": [],
   "source": [
    "node_pipeline_fastrp.selectFeatures([\"embedding\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "c0krIvXTipcD"
   },
   "outputs": [],
   "source": [
    "# Configure the pipeline as before\n",
    "node_pipeline_fastrp.configureSplit(testFraction=0.2, validationFolds=5)\n",
    "\n",
    "node_pipeline_fastrp.addLogisticRegression(maxEpochs=1000, penalty=(0.00048, 0.00050))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "Kvr9iU4HdoTv"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Node Classification Train Pipeline: 100%|██████████| 100.0/100 [01:35<00:00,  1.05%/s]\n"
     ]
    }
   ],
   "source": [
    "# Perform the actual training\n",
    "model_fastrp, stats_fastrp = node_pipeline_fastrp.train(\n",
    "    G,\n",
    "    nodeLabels=[\"Paper\"],\n",
    "    modelName=\"cora-pipeline-model-fastrp\",\n",
    "    targetProperty=\"subject\",\n",
    "    metrics=[\"F1_WEIGHTED\"],\n",
    "    randomSeed=42,\n",
    "    concurrency=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "_Q3zMmpCd2o0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8652542239963541\n"
     ]
    }
   ],
   "source": [
    "print(stats_fastrp[\"modelInfo\"][\"metrics\"][\"F1_WEIGHTED\"][\"test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rhaPoK0xFdK3"
   },
   "source": [
    "## Use the new model for prediction\n",
    "\n",
    "Here we are **not** using the `nodeLabels=[\"UnclassifiedPaper\"]` parameter, because FastRP depends on neighbour nodes. When using models that have `nodePropertySteps` that use relationships (such as FastRP and other models that create embeddings, but also algorithms like PageRank), the `nodeLabels` filter should not be used because it would \"cut out\" all the linked nodes that have a different label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "vgfN-ppHw4Rp"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Node Classification Predict Pipeline: 100%|██████████| 100.0/100 [00:00<00:00, 229.18%/s]\n"
     ]
    }
   ],
   "source": [
    "predicted_fastrp = model_fastrp.predict_stream(\n",
    "    G, modelName=\"cora-pipeline-model-fastrp\", includePredictedProbabilities=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "Lp_wzQKdmTDB"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "nodeId                    2708\n",
       "predictedClass            2708\n",
       "predictedProbabilities    2708\n",
       "dtype: int64"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predicted_fastrp.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u2o_J8dzN-qQ"
   },
   "source": [
    "Since we have used no filters, the `predicted` result contains _all_ the nodes. The way to filter the nodes is via the `streamNodeProperty` method, which can be used only after the new property is written to the graph via the `mutate` mode. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PNbRL2y84F6Q"
   },
   "source": [
    "Instead of streaming the results, the prediction can be run in `mutate` mode to be more performant. The predicted nodes can be retrieved using the `streamNodeProperty` method with the `UnclassifiedPaper` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "_HVufWHbxIoE"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Node Classification Predict Pipeline: 100%|██████████| 100.0/100 [00:00<00:00, 193.51%/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "nodePropertiesWritten                                                 5416\n",
       "mutateMillis                                                             0\n",
       "postProcessingMillis                                                     0\n",
       "preProcessingMillis                                                      0\n",
       "computeMillis                                                          928\n",
       "configuration            {'graphName': 'cora-graph', 'jobId': 'e2f0947e...\n",
       "Name: 0, dtype: object"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_fastrp.predict_mutate(\n",
    "    G,\n",
    "    mutateProperty=\"predictedClass\",\n",
    "    modelName=\"cora-pipeline-model-fastrp\",\n",
    "    predictedProbabilityProperty=\"predictedProbabilities\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "dSoNCQGjGLKt"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>nodeId</th>\n",
       "      <th>propertyValue</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15188</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15189</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15190</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>15191</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>15192</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>15193</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>15194</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>15195</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>15196</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>15197</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   nodeId  propertyValue\n",
       "0   15188            100\n",
       "1   15189              1\n",
       "2   15190              2\n",
       "3   15191              2\n",
       "4   15192              3\n",
       "5   15193              3\n",
       "6   15194              4\n",
       "7   15195            100\n",
       "8   15196            100\n",
       "9   15197              4"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predicted_fastrp = gds.graph.streamNodeProperty(G, \"predictedClass\", [\"UnclassifiedPaper\"])\n",
    "\n",
    "predicted_fastrp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "thVR-anwmmzo"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>nodeId</th>\n",
       "      <th>propertyValue</th>\n",
       "      <th>subject</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15188</td>\n",
       "      <td>100</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15189</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15190</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>15191</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>15192</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>15193</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>15194</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>15195</td>\n",
       "      <td>100</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>15196</td>\n",
       "      <td>100</td>\n",
       "      <td>100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>15197</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   nodeId  propertyValue  subject\n",
       "0   15188            100      100\n",
       "1   15189              1        1\n",
       "2   15190              2        2\n",
       "3   15191              2        2\n",
       "4   15192              3        3\n",
       "5   15193              3        3\n",
       "6   15194              4        4\n",
       "7   15195            100      100\n",
       "8   15196            100      100\n",
       "9   15197              4        4"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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 holdout nodes is accurate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UdEVypaiMc5d"
   },
   "source": [
    "## Write result back to Neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "aKuUQwxbPOaV"
   },
   "outputs": [],
   "source": [
    "model_fastrp.predict_write(\n",
    "    G,\n",
    "    writeProperty=\"predictedSubject\",\n",
    "    modelName=\"cora-pipeline-model-fastrp\",\n",
    "    predictedProbabilityProperty=\"predictedProbabilities\",\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:"
   ]
  },
  {
   "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()\n",
    "gds.run_cypher(\"MATCH (n) DETACH DELETE n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is good practice to close the client as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "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.5"
  },
  "vscode": {
   "interpreter": {
    "hash": "d9b85228d1674a60a9aba6bdb3ec2509e82e7a19678e59e4fceb33f4c3ba222b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}