Part 4 - Graph Machine Learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85b497ba-5964-42b7-94cc-ae81a6184314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import os\n",
    "import pickle\n",
    "import tarfile\n",
    "from pprint import pprint\n",
    "from typing import Dict, List, Optional\n",
    "\n",
    "import graphistry\n",
    "import dgl\n",
    "# import igraph as ig\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "# import optuna\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from dgl.data import DGLDataset\n",
    "from dgl.data.utils import download as dgl_download\n",
    "from dgl.data.utils import load_graphs, save_graphs\n",
    "from dgl.sampling import global_uniform_negative_sampling\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from graphml_class.clean import clean_graph\n",
    "from graphml_class.palette import CATEGORICAL_PALETTE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4ac7f4d-0aa0-4daa-b613-7ff400e77829",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 31337"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd0fb2c6-4032-4eff-a174-a41c09a1f3b8",
   "metadata": {},
   "source": [
    "# Graph Machine Learning by Hand\n",
    "\n",
    "The old-school method of machine learning that was common practice before graph embeddings and Graph Neural Networks was to hand engineer features of the nodes of a network by querying the network with `networkx` or a graph database like Neo4j using Cypher or Gremlin and to assign the results of those queries to nodes.\n",
    "\n",
    "What are some common features when hand engineering node features in graph machine learning tasks?\n",
    "\n",
    "## Hand Engineered Features in Graph ML\n",
    "\n",
    "* Scalable Centralities: Degree [in, out, total], Eigenvector, PageRank [a form of eigenvector centrality], Katz\n",
    "* Less Scalable Centralities: Closeness, Betweenness\n",
    "\n",
    "You can compute these methods on single edges or on multiple edges as a way of projecting a _higher-order_ network, a network composed of more complex semantics built from the base network. Graph ML worked this way until around 2015, when graph embeddings changed everything :)\n",
    "\n",
    "Check out [networkx.algorithms.centrality](https://networkx.org/documentation/stable/reference/algorithms/centrality.html) and find some that fit different aspects of networks you'd like to include in a classifier.\n",
    "\n",
    "## Loading our Citation Graph\n",
    "\n",
    "In Part 1 on knowledge graph construction, we saved our graph in GEXF format. Let's load it using the same procedure we performed in Part 2. I'm going to abbreviate the comments here, look at Part 2 for more details on what I am doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d192cd97-8a45-4c60-a7a3-42c438cca9a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "G: nx.DiGraph = nx.read_gexf(path=\"data/physics_labeled.gexf.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c24db786-8fb6-4904-a45b-7c59a9a9970a",
   "metadata": {},
   "outputs": [],
   "source": [
    "G.number_of_nodes(), G.number_of_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "356be9c2-d542-4ef5-9f9d-3cfe608bf5f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"data/citation/file_to_net.pkl\", \"rb\") as f:\n",
    "    file_to_net = pickle.load(f)\n",
    "\n",
    "# Everything ok? Yes!\n",
    "list(file_to_net.items())[0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c236eb7-c936-49cf-ae8d-2ac823f0bf65",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"data/citation/net_to_file.pkl\", \"rb\") as f:\n",
    "    net_to_file = pickle.load(f)\n",
    "\n",
    "# Everything ok here too? Yes!\n",
    "list(net_to_file.items())[0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ac6c2b-5ddd-40a3-9fe1-be9ccd6b0d8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_ids_to_int(G):\n",
    "    # Create a new directed graph\n",
    "    G_int = nx.DiGraph()\n",
    "    \n",
    "    # Create a mapping from string IDs to integer IDs\n",
    "    id_mapping = {str_id: int(str_id) for str_id in G.nodes()}\n",
    "    \n",
    "    # Copy nodes and attributes, converting IDs to integers\n",
    "    for str_id, data in G.nodes(data=True):\n",
    "        int_id = id_mapping[str_id]\n",
    "        G_int.add_node(int_id, **data)\n",
    "    \n",
    "    # Copy edges and attributes, converting IDs to integers\n",
    "    for str_id1, str_id2, data in G.edges(data=True):\n",
    "        int_id1, int_id2 = id_mapping[str_id1], id_mapping[str_id2]\n",
    "        G_int.add_edge(int_id1, int_id2, **data)\n",
    "    \n",
    "    return G_int\n",
    "\n",
    "# Convert G to use integer IDs\n",
    "G_int = convert_ids_to_int(G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f12c7bc5-de7d-42e2-8277-e0b541dd73bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "G_int.number_of_nodes(), G_int.number_of_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15cbdb80-115f-4595-9b30-1c5480a6e18c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's test or integer index now... we pickled it instead of JSONized it so it would retain its integer keys and values!\n",
    "test_id = file_to_net[9711194]\n",
    "G_int.nodes[test_id]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d18e7fff-07e6-4a58-94b7-ec9fb81a755e",
   "metadata": {},
   "source": [
    "### `G_int` --> `G`\n",
    "\n",
    "Now we can assign our new integer graph back to `G` and use it below :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cfabed8-c670-49d3-845e-678dca9633e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = G_int"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86319fce-9dae-4c59-aa67-c55a02b168d0",
   "metadata": {},
   "source": [
    "## Feature Engineering for Node Classification\n",
    "\n",
    "We are going to load our citation graph and then combine a simple encoding of the text features with manually computed features to perform node classification. We will see how our network metrics describing topology affect performance. Try not to always jump to using a network embedding, this is a useful exercise as a baseline :)\n",
    "\n",
    "Our job is to prepare a network features vector `features` that we can append to our other features before we train a classifier to perform link prediction.\n",
    "\n",
    "### Historical Review\n",
    "\n",
    "Remember these slides? Think on them again :) We are going to implement a couple of these to drive our first Journal classifier :) This slide shows numerous methods used in link prediction.\n",
    "\n",
    "<center><img src=\"images/Feature-Engineering-for-Link-Prediction.jpg\" width=\"1000px\" /></center>\n",
    "<center><i>Feature engineering for link predicton can be complicated due to the need to compute them for pairs of nodes.</i></center>\n",
    "\n",
    "<br />\n",
    "\n",
    "<center><img src=\"images/JanusGraph-Titan-Feature-Engineering-in-2014.png\" width=\"1000px\" /></center>\n",
    "<center><i>Feature engineering circa 2014 on a business graph representing relationships between companies in a property graph using <a href=\"https://titan.thinkaurelius.com/\">Titan</a> via <a href=\"https://tinkerpop.apache.org/gremlin.html\">Tinkerpop / Gremlin</a>. Titan became <a href=\"https://janusgraph.org/\">JanusGraph</a>.</i></center>\n",
    "\n",
    "<br />\n",
    "\n",
    "### Store Features by Column in `features`\n",
    "\n",
    "We will compute features for each node or pairs of nodes and then store them in the `features` `np.ndarray` by appending one column at a time. One row is a node. This is described in [NumPy: How to add an extra column to a NumPy array](https://www.w3resource.com/python-exercises/numpy/python-numpy-exercise-86.php) as looking like:\n",
    "\n",
    "<center><img src=\"images/append-column-to-matrix.png\" width=\"500px\" /></center>\n",
    "<center>Image Source: <a href=\"https://www.w3resource.com/python-exercises/numpy/python-numpy-exercise-86.php\">w3resource: NumPy: How to add an extra column to a NumPy array</a></center>\n",
    "\n",
    "#### Re-Running Feature Appends Will Make Duplicate Feature Columns\n",
    "\n",
    "Warning: If you run the cells below that do `np.append(features, my_feat_ary, axis=1)` more than once... you will get a duplicate features column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd25944-0efe-4c57-a145-1c4567cbd1aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I had to look this up, so here you are :)\n",
    "import numpy as np\n",
    "x = np.array([[10,20,30], [40,50,60]])\n",
    "y = np.array([[100], [200]])\n",
    "x, y, np.append(x, y, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a87ecd87-e335-4ab5-8928-aacc95b3f92c",
   "metadata": {},
   "source": [
    "### Initialize `features` `nd.array`\n",
    "\n",
    "You can re-run this to recalculate the features from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a566905-e7fc-4838-987f-aea0b082d5d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an inner list for each node ID\n",
    "inner_lists = [[] for x in range(G.number_of_nodes())]\n",
    "\n",
    "# The shape is 27,770 nodes long with zero feature columns\n",
    "features = np.array(inner_lists)\n",
    "features, features.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aee28f99-ca8a-4922-85cf-44517ac70359",
   "metadata": {},
   "source": [
    "### `networkx` Centrality Metrics\n",
    "\n",
    "For our first features, we will use `networkx` to compute some centrality metrics. I'm going to limit our calculations to two centralities you can use in practice on networks of most any size:\n",
    "\n",
    "* Degree Centrality (in degree, out degree, degree [both]) - this is simple and shows how well connected a node is\n",
    "* Eigenvector Centrality - measures the influence of a node in a network, where a node is considered influential if it is connected to other influential nodes.\n",
    "\n",
    "Check out this list of [networkx.centrality](https://networkx.org/documentation/stable/reference/algorithms/centrality.html) metrics for feature ideas.\n",
    "\n",
    "#### Display Metrics using `pd.Series`\n",
    "\n",
    "If you want to see a metric computer by `nx.my_function()` you can wrap the `Dict` it returns [which will ALL print in a notebok] in a `pd.Series` and it will use the node ID as the index and display the feature. It is a good idea to lay eyes on things you compute - when possible - as often something goes wrong on a first attempt :)\n",
    "\n",
    "#### Degree Centrality\n",
    "\n",
    "Degree centrality is a **local centrality** metric. It measure influence in the nearby network.\n",
    "\n",
    "NetworkX measures degree centrality as a relative value for all nodes in a network. This can complicate inference but keeps values within an interpretable range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "682b2420-ade7-42b1-af66-81f2030409f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "degree = nx.degree_centrality(G)\n",
    "\n",
    "# See - takes just 11 lines. The Dict printed over 27K rows on my screen. Try it!\n",
    "pd.Series(degree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6c2911d-ac45-4c99-b81d-3177fa73ee49",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_degree = nx.in_degree_centrality(G)\n",
    "pd.Series(in_degree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c761e8d-305e-401b-9bc2-037b6a24e0b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_degree = nx.out_degree_centrality(G)\n",
    "pd.Series(out_degree)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cea3852-a8a2-4966-a70e-29cc7210a28e",
   "metadata": {},
   "source": [
    "#### Eigenvector Centrality\n",
    "\n",
    "Eigenvector centrality is a **global centrality** metric. It measures influence within the entire network. _We always want a local and global centrality metric in our feature set_. To understand eigenvectors, check out this amazing simulation: [Eigenvectors and Eigenvalues](https://setosa.io/ev/eigenvectors-and-eigenvalues/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1e078dd-2b84-4a11-82fd-d9b75cf19530",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Eigenvector centrality is problematic in disconnected graphs, so take the largest component...\n",
    "components = nx.connected_components(G.to_undirected())\n",
    "largest_component = max(components, key=len)\n",
    "G_big_component = G.subgraph(largest_component).copy()\n",
    "\n",
    "print(f\"Largest comonent size: {len(G_big_component):,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6340a5d1-44e6-40e8-ae5e-07c7c95a27dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the eigenvectors... see \n",
    "eigenvector = nx.eigenvector_centrality_numpy(G_big_component.to_undirected())\n",
    "pd.Series(eigenvector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140187fb-ccee-49e5-8e85-15b7d1d2ce10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ... and make it undirected, to ensure it is connected\n",
    "eigenvectors = nx.eigenvector_centrality_numpy(\n",
    "    G_big_component.to_undirected()\n",
    ")\n",
    "\n",
    "for node, score in eigenvectors.items():\n",
    "    G_big_component.nodes[node][\"eigenvector\"] = score\n",
    "    G_big_component.nodes[node][\"eigenvector_log10\"] = np.log10(score + 0.000001)  # can't be zero for log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5f90b9d-b65e-46b2-8f51-18587d080081",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize eigenvector_centrality property as 0 for all nodes\n",
    "nx.set_node_attributes(G, 0.0, \"eigenvector\")\n",
    "nx.set_node_attributes(G, 0.0, \"eigenvector_log10\")\n",
    "\n",
    "# Assign the eigenvector centrality values from subgraph to original graph\n",
    "nx.set_node_attributes(G, {node: eigenvectors[node] for node in G_big_component}, \"eigenvector\")\n",
    "nx.set_node_attributes(G, {node: np.log10(G.nodes[node]['eigenvector'] + 1e-10) for node in G.nodes()}, \"eigenvector_log10\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50b11761-9631-4d56-9f10-5a1e04725c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "eigenvector_dict = nx.get_node_attributes(G, 'eigenvector')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58a519d9-b1ad-435e-8c2b-451956cb0d4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get an array for our features array\n",
    "eigenvector_array = np.array([G.nodes[node][\"eigenvector\"] for node in G.nodes()])\n",
    "pd.Series(eigenvector_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cc84c28-1df2-465c-bde0-b5e30ea0f5f7",
   "metadata": {},
   "source": [
    "### Add to `features` `np.array`\n",
    "\n",
    "That gives us four features so far that indicate a node's prominence in the network. Let's combine them to produce a four feature long feature vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01466c8a-7930-4efc-bb02-7d8817ba3487",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(list(degree.values())[0:10]), np.array(list(degree.values())[0:10]).reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ff6162b-7420-49a8-b520-524ec83e89a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare and append the degree features for addition to `features`\n",
    "for d in [degree, in_degree, out_degree, eigenvector_dict]:\n",
    "    d_ary = np.array(list(d.values())).reshape(-1, 1)\n",
    "    features = np.append(features, d_ary, axis=1)\n",
    "\n",
    "features, features.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cb6dd77-7f42-4430-90fa-04bc4fb3d179",
   "metadata": {},
   "source": [
    "### Community Detection with Louvain Modularity\n",
    "\n",
    "What about other kinds of features? One important category is the local group a node belongs to after performing an operation called \n",
    "\n",
    "One thing we didn't cover in Part 2 - Network Science, was community detection. What community a node belongs to is an important feature. We're going to cluster our network into communities and visualize them in Graphistry. Then we're going to add the cluster IDs as features in our `features` `np.array`.\n",
    "\n",
    "While there are many network clustering algorithms, the most common method of community detection is [Louvain Modularity](https://en.wikipedia.org/wiki/Louvain_method). We're going to use [nx.community.louvain_communities](https://networkx.org/documentation/stable/reference/algorithms/generated/networkx.algorithms.community.louvain.louvain_communities.html#networkx.algorithms.community.louvain.louvain_communities) because it is efficient enough to work on networks up to moderate size using one computer. You can usually find a Louvain Modularity implementation for whatever platform you use to process even large networks - big data.\n",
    "\n",
    "> The method is a greedy optimization method that appears to run in time `O(n * log(n))`.\n",
    ">\n",
    "> ...\n",
    "> \n",
    "> In the Louvain Method of community detection, first small communities are found by optimizing modularity locally on all nodes, then each small community is grouped into one node and the first step is repeated.\n",
    "\n",
    "-- Wikipedia, [Louvain method](https://en.wikipedia.org/wiki/Louvain_method)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bf507b5-97e5-4ca3-941b-9bb6735c9ef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "clusters = nx.community.louvain_communities(G, seed=SEED)\n",
    "len(clusters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d629b24e-5b15-471c-b1d1-504ef0c9a5f7",
   "metadata": {},
   "source": [
    "#### Evaluating Modularity\n",
    "\n",
    "I am curious how many clusters there are... lets plot their size in Seaborn using log scale. You can see there is a power law distribtion in cluster size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f903978-eb1a-4ce0-af0f-9e9cca7300c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get connected components and their sizes\n",
    "cluster_size = [len(c) for c in clusters]\n",
    "\n",
    "# Use seaborn to create the histogram\n",
    "sns.histplot(cluster_size, kde=True, bins=20, log_scale=True)\n",
    "plt.title(\"Histogram of Cluster Sizes\")\n",
    "plt.xlabel(\"Cluster Size\")\n",
    "plt.ylabel(\"Count\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2f43db9-4e1f-4a1b-bf8c-a85c3f52029b",
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(sorted(cluster_size, reverse=True)[0:50], indent=4, compact=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "206b3295-e0dc-4526-8384-519b2d5448f6",
   "metadata": {},
   "source": [
    "### ChatGPT and Me\n",
    "\n",
    "I needed to do consult a reference to interpret what this might mean for a citation graph... I knew it had a power law distribution with a long tail. I fed the following prompt to ChatGPT-4 Browser:\n",
    "\n",
    "#### **CHATGPT PROMPT**\n",
    "\n",
    "I have a citation graph of physics papers from arXiv, the high energy physics citation network from Stanford SNAP. I used NetworkX to compute Louvain Modularity and measured cluster sizes. Then I plotted them in seaborn on a log scale.\n",
    "\n",
    "The first 50 cluster sizes, sorted in descending order are:\n",
    "\n",
    "```\n",
    "[   2481, 2144, 1779, 1731, 1727, 1596, 1594, 1551, 1447, 1412, 1332, 1255,\n",
    "    1108, 919, 844, 835, 817, 674, 540, 495, 316, 314, 170, 157, 113, 11, 10,\n",
    "    10, 8, 8, 7, 7, 6, 6, 6, 5, 5, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4, 4]\n",
    "```\n",
    "\n",
    "Please interpret this histogram of cluster sizes and tell me what it means for a citation graph.\n",
    "\n",
    "#### Comment\n",
    "\n",
    "The output checks out... ChatGPT-4 can interpret a histogram in a specific domain fairly well. Check it out. I'm trying to demonstrate how to incorporate AI into your workflows... I did ask it to rewrite this in fewer words, but I'm printing the entire output below.\n",
    "\n",
    "#### **CHATGPT RESPONSE**\n",
    "\n",
    "The cluster sizes you've provided, derived from applying the Louvain Modularity algorithm on a citation graph of physics papers, represent the sizes of communities of papers that cite each other. The sizes are listed in descending order, indicating a few large clusters followed by many smaller ones. This pattern is typical in real-world networks and is often referred to as a scale-free or heavy-tailed distribution. Here are a few interpretations and insights based on this distribution:\n",
    "\n",
    "* Community Structure - There are a few large communities of papers (e.g., the largest has 1909 papers) which likely represent major topics or fields within high energy physics. These large clusters suggest that there are common themes that have garnered significant attention and citation in the community.\n",
    "* Long Tail - As we move down the list, the cluster sizes decrease, forming a long tail of smaller communities. This indicates a wide variety of specialized or niche topics within the field that have smaller groups of papers citing each other.\n",
    "* Power-Law Distribution - The distribution of cluster sizes resembles a power-law distribution, which is common in many natural and social phenomena. In a power-law distribution, a few items (in this case, topics or fields) are exceptionally common, but the vast majority are rare.\n",
    "* Interdisciplinary or Emerging Areas - The smaller clusters may represent interdisciplinary or emerging areas of research that have not yet amassed a large number of citations.\n",
    "* Potential for Further Exploration - Understanding the thematic content of the large clusters could provide insights into the major areas of focus in high energy physics, while exploring the smaller clusters might reveal emerging or under-explored areas of research.\n",
    "* Rich-Get-Richer Phenomenon - The distribution may also reflect a \"rich-get-richer\" phenomenon where well-cited papers continue to accumulate more citations over time, reinforcing the prominence of certain topics.\n",
    "* Validation of Louvain Modularity - The distribution of cluster sizes could also serve as a validation of the Louvain Modularity algorithm's effectiveness in detecting community structure within the citation network.\n",
    "\n",
    "### Featurizing Communities\n",
    "\n",
    "To featurize our communities, we need to assign numeric values to the clusters and append a column for cluster ID to `features`. Note, we could alternatively [one-hot-encode](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.OneHotEncoder.html) the features to give them one feature column per cluster. This would make the features sparse, which means lesss signal - something neural networks don't like but that another algorithm might not mind. We could use the [hashing trick](https://scikit-learn.org/stable/modules/generated/sklearn.feature_extraction.FeatureHasher.html) - a method of feature engineering somewhere in between one-hot-encoding and neural embeddings. We are instead going to _label encode_ them with [sklearn.preprocessing.LabelEncoder](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.LabelEncoder.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "127299fd-fe7a-4e11-9dfb-780bb0b6111e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for cluster in clusters:\n",
    "pprint(clusters[0:3], indent=4, compact=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d10360c5-59f0-4d5b-9a57-6dbe6678f675",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iterate through clusters, assigning cluster IDs, then map them to node IDs\n",
    "node_clusters: Dict[int, int] = {}\n",
    "for cluster_id, cluster in zip(range(len(clusters)), clusters):\n",
    "    for node_id in cluster:\n",
    "        node_clusters[node_id] = cluster_id\n",
    "\n",
    "node_series = pd.Series(node_clusters).sort_index()\n",
    "node_series"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ae96de9-7b2f-4d84-af08-9bae0ada6fbf",
   "metadata": {},
   "source": [
    "### Append Clusters to `features`\n",
    "\n",
    "Note, the clusters number from approximately 1-1,500. This dwarfs the previously computed, normalized centrality scores. We will be using scikit-learn's [`StandardScaler`](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.StandardScaler.html) to address this.\n",
    "\n",
    "Now we've got 5 topological features. We are going to add one more and then do a simple feature extraction of our Titles and Abstracts that might normally go with a first pass at a node classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d60ee73c-5f7f-4a2b-861e-6249f28e70b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare and append the degree features for addition to `features`\n",
    "cluster_ary = np.array(node_series).reshape(-1, 1)\n",
    "features.shape, cluster_ary.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b3c438b-603a-45eb-9f3d-d6f4508fab0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = np.append(features, d_ary, axis=1)\n",
    "features, features.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b33827e-408e-412e-b15c-e1e3a8855a71",
   "metadata": {},
   "source": [
    "## Role Discovery via RolX and ReFeX\n",
    "\n",
    "RolX role discovery is one of my favorite algorithms for networks :) You may recall from the slides how I used it to build a visually appealing, intrepretable map of the big data market. We're going to compute RolX roles to use as feature for ML tasks.\n",
    "\n",
    "### RolX\n",
    "\n",
    "You can learn more about RolX in the paper [RolX: Structural Role Extraction & Mining in Large Graphs](https://web.eecs.umich.edu/~dkoutra/papers/12-kdd-recursiverole.pdf).\n",
    "\n",
    "<center><img src=\"images/RolX-Roles-vs-Communities.png\" width=\"550px\" /></center>\n",
    "<center><i>You can see the difference between structural roles and communities. A related concept called _network motifs_ or structural patterns in networks can determine the nodes roles play in their neighborhoods.</i></center>\n",
    "\n",
    "<br /><br />\n",
    "\n",
    "<center><img src=\"images/Graph-ML-Class-Slide-for-Rolx-Roles-for-Mapping-Big-Data.png\" width=\"1000px\" /></center>\n",
    "<center><i>An example of RolX in action to filter nodes with roles I didn't care about from a visualization to tell the story of interest.</i></center>\n",
    "\n",
    "### ReFeX via `graphrole`\n",
    "\n",
    "ReFex is a role discovery algorithm similar to RolX. It was defined in [It’s Who You Know: Graph Mining Using Recursive Structural Features](https://www.andrew.cmu.edu/user/lakoglu/pubs/ReFeX.pdf).\n",
    "\n",
    "GraphRole draws from Rolx and ReFeX to do role discovery. Note that the `graphrole` implementation of RolX is a little slow... even on this network. The reality is that for large networks, you need the support of a data engineer to implement algorithms. It isn't feasible to teach a lot of large network computing skills in a five day, four hour course. RolX itself was designed to support large networks, and there are other implementations available.\n",
    "\n",
    "Therefore I have pickled the results and made them available for you to load. **Feel free to uncomment the next block of code out and run it if you like, otherwise the RolX features are computed and stored in Github.** They are loaded below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a0b75b8-2b2f-4f5f-9b19-caeca158cda7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrole import RecursiveFeatureExtractor, RoleExtractor\n",
    "\n",
    "\n",
    "# Extract features\n",
    "feature_extractor = RecursiveFeatureExtractor(G)\n",
    "rolx_features = feature_extractor.extract_features()\n",
    "\n",
    "# Calculate roles\n",
    "role_extractor = RoleExtractor(n_roles=None)\n",
    "role_extractor.extract_role_factors(rolx_features)\n",
    "\n",
    "# Can also use role percentages\n",
    "role_extractor.role_percentage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e583652-e900-457c-b8c7-3f4505244060",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_roles = role_extractor.roles\n",
    "node_role_s = pd.Series(node_roles)\n",
    "node_role_s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23fc2d18-6933-4a28-97e5-c5aecab912e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_roles() -> pd.Series:\n",
    "    \"\"\"Compute the roles unless they are already computed.\"\"\"\n",
    "    # Extract features\n",
    "    feature_extractor = RecursiveFeatureExtractor(G)\n",
    "    rolx_features = feature_extractor.extract_features()\n",
    "    \n",
    "    # Calculate roles\n",
    "    role_extractor = RoleExtractor(n_roles=None)\n",
    "    role_extractor.extract_role_factors(rolx_features)\n",
    "\n",
    "    # return pd.Series(role_extractor.role_percentages)\n",
    "    return pd.Series(role_extractor.roles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "672f829d-f178-4cf3-a424-5a4cc9a9801f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the feature_extractor and role_extractor via pickle\n",
    "node_role_path = \"data/citation/node_role_s.pkl\"\n",
    "node_role_s: Optional[pd.Series] = None\n",
    "if os.path.exists(node_role_path):\n",
    "    with open(node_role_path, \"rb\") as f:\n",
    "        node_role_s = pickle.load(f)\n",
    "else:\n",
    "    node_role_s = get_roles()\n",
    "    with open(node_role_path, \"wb\") as f:\n",
    "        pickle.dump(node_role_s, f)\n",
    "\n",
    "node_roles = node_role_s.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c95acd-b0ac-4afb-a27a-dece62fbb04f",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_role_s, node_role_s.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4a720f0-c07f-4b14-97bc-4da9d45e35aa",
   "metadata": {},
   "source": [
    "#### Label Encoding Roles with [sklearn.preprocessing.LabelEncoder](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.LabelEncoder.html)\n",
    "\n",
    "We need numeric role numbers, not strings, so we can add the `node_roles` to our `features` and start classifying. This is what scikit-learn's `LabelEncoder` does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51eccf4b-7a05-4b9e-80a3-7d98e8960453",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "\n",
    "le = LabelEncoder()\n",
    "flat_role_feature = le.fit_transform(node_role_s)\n",
    "flat_role_feature, flat_role_feature.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4576f2f-a8c8-4ee0-9dc3-7b8ded5a2ba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "role_feature = flat_role_feature.reshape([-1,1])\n",
    "role_feature, role_feature.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ab497c-c4d5-42d6-a70e-dfd838f7deaf",
   "metadata": {},
   "source": [
    "#### Add `role_feature` to `features`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c292d24b-0b69-40b2-9266-ef3d1d725e28",
   "metadata": {},
   "outputs": [],
   "source": [
    "features = np.append(features, role_feature, axis=1)\n",
    "features, features.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b51c2e-30fe-406f-a05c-cf034e47d6d0",
   "metadata": {},
   "source": [
    "## Visualizing Communities and Roles in Graphistry\n",
    "\n",
    "We're going to use `graphistry` to visualize the communities and roles we determined to see if we can interpret them the way ChatGPT did - or in some other way :) This sort of belongs in Part 2 - Network Science, but we saved it here to use as a feature for node classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61ff0c4-54d7-45e4-8c3c-11599eddf945",
   "metadata": {},
   "outputs": [],
   "source": [
    "nx.set_node_attributes(G, values=node_roles, name=\"rolx_role\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccc34712-177a-44d3-86e1-1fa1773f1978",
   "metadata": {},
   "outputs": [],
   "source": [
    "nx.set_node_attributes(G, values=node_clusters, name=\"community\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddf72b49-970c-4a6c-9a55-88da23d8c397",
   "metadata": {},
   "outputs": [],
   "source": [
    "G.nodes[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99230828-1908-4d55-92d1-17d241e7a551",
   "metadata": {},
   "outputs": [],
   "source": [
    "from littleballoffur import PageRankBasedSampler\n",
    "\n",
    "\n",
    "G_orig = G.copy()\n",
    "sampler = PageRankBasedSampler(10000)\n",
    "G_sample = nx.convert_node_labels_to_integers(\n",
    "    sampler.sample(G_orig.to_undirected())\n",
    ")\n",
    "\n",
    "G_sample.number_of_nodes(), G_sample.number_of_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54a87987-b27c-4564-a4ea-a4d7ee33e30e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment variable setup\n",
    "GRAPHISTRY_USERNAME = os.getenv(\"GRAPHISTRY_USERNAME\")\n",
    "GRAPHISTRY_PASSWORD = os.getenv(\"GRAPHISTRY_PASSWORD\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44f5cac2-d3b1-4b7a-824a-9d2decb096c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "graphistry.register(\n",
    "    api=3,\n",
    "    username=GRAPHISTRY_USERNAME,\n",
    "    password=GRAPHISTRY_PASSWORD,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9707cfa6-97ab-4048-af7b-37396dedf7ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration for Graphistry\n",
    "FAVICON_URL = \"https://graphlet.ai/assets/icons/favicon.ico\"\n",
    "LOGO = {\"url\": \"https://graphlet.ai/assets/Branding/Graphlet%20AI.svg\", \"dimensions\": {\"maxWidth\": 60, \"maxHeight\": 60}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6569bbae-f35a-4e3e-a0dc-80cb65c77573",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters for GRAPHISTRY\n",
    "GRAPHISTRY_PARAMS = {\n",
    "    \"play\": 500,\n",
    "    \"pointOpacity\": 0.7,\n",
    "    \"edgeOpacity\": 0.3,\n",
    "    \"edgeCurvature\": 0.3,\n",
    "    \"showArrows\": True,\n",
    "    \"gravity\": 0.15,\n",
    "    \"showPointsOfInterestLabel\": False,\n",
    "    \"labels\": {\n",
    "        \"shortenLabels\": False,\n",
    "    },\n",
    "    \"scalingRatio\": 0.5,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7531df39-abdf-440f-9744-139a91200884",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = (\n",
    "    graphistry.bind(\n",
    "        source=\"src\",\n",
    "        destination=\"dst\",\n",
    "        node=\"nodeid\",\n",
    "        point_title=\"Title\",\n",
    "        point_label=\"Title\",\n",
    "    )\n",
    "    .addStyle(\n",
    "        page={\n",
    "            \"title\": \"Community Detection Plot\",\n",
    "            \"favicon\": FAVICON_URL\n",
    "        },\n",
    "        logo=LOGO,\n",
    "    )\n",
    "    .settings(\n",
    "        url_params=GRAPHISTRY_PARAMS,\n",
    "        height=800,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8804c5e4-c541-4381-aa97-7d427869057a",
   "metadata": {},
   "outputs": [],
   "source": [
    "g.plot(G_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5353b7bd-be31-4a18-b11d-746cb1d896b0",
   "metadata": {},
   "source": [
    "## Classifying Nodes into Journals\n",
    "\n",
    "We are going to train a classifier and then classify the nodes without journal entries into the most common journals in the field of high energy physics to see where they most likely belong. We are going to use a simple algorithm to do this, then get more sophisticated in the next sections on _graph embeddings_ and _graph neural networks (GNNs)_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb584dae-e30d-453c-8d78-7eaa564a1f82",
   "metadata": {},
   "outputs": [],
   "source": [
    "journal_labels = np.array([\n",
    "    x[1][\"Journal-ref-DBSCAN\"] for x in list(G.nodes(data=True))\n",
    "])\n",
    "journal_labels, journal_labels.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "553ac18e-64c6-4ca2-8340-1fcb674a544a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    features,\n",
    "    journal_labels,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f79b9947-477d-474a-b955-1c1ce15b0b3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, confusion_matrix, precision_score, recall_score, f1_score\n",
    "\n",
    "\n",
    "def score_classifier(y_true, y_pred):\n",
    "    \"\"\"Run and return a pd.DataFrame of various classifier test metrics.\"\"\"\n",
    "    return {\n",
    "        \"Metric\": [\"Accuracy\", \"Precision\", \"Recall\", \"F1\"],\n",
    "        \"Score\": [\n",
    "            accuracy_score(y_true, y_pred),\n",
    "            precision_score(y_true, y_pred, average=\"weighted\"),\n",
    "            recall_score(y_true, y_pred, average=\"weighted\"),\n",
    "            f1_score(y_true, y_pred, average=\"weighted\"),\n",
    "        ],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e0ee6e-0ebc-409b-ac5e-1e61b978a58e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "\n",
    "scaler = StandardScaler()\n",
    "clf = make_pipeline(StandardScaler(), LogisticRegression())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a9a2520-c57f-409e-809e-ec260d7bab9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = clf.fit(X_train, y_train)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7642130a-c9ba-45ce-a82f-db0e82492774",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model.predict(X_test)\n",
    "\n",
    "pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "917c3f92-c1c4-4d47-8bd9-031ed6174d92",
   "metadata": {},
   "source": [
    "## Exercise: Try to improve the perormance of our model using more `networkx` metrics\n",
    "\n",
    "Take a look at the [long list of algorithms provided by networkx]() including its own [node classification methods](https://networkx.org/documentation/stable/reference/algorithms/node_classification.html). Compute several new features and add them to the `features` `np.array` as above, before re-scoring our classifier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b22669fa-e113-4f64-9fee-a65f0c2a4dd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from networkx.algorithms import node_classification\n",
    "#\n",
    "# harmonic_class = node_classification.harmonic_function(G.to_undirected(), label_name=\"Journal-ref-DBSCAN\")\n",
    "# harmonic_class_ary = np.array(harmonic_class)\n",
    "# harmonic_class_ary, harmonic_class_ary.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b83576-ab36-4dfa-90f7-03bb269f30e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# features = np.append(\n",
    "#     features,\n",
    "#     harmonic_class_ary.reshape(-1, 1),\n",
    "#     axis=1\n",
    "# )\n",
    "# features, features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaacd5c1-0d0b-4af9-868d-90dc7ad2aa33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# consist_class = node_classification.local_and_global_consistency(G.to_undirected(), label_name=\"Journal-ref-DBSCAN\")\n",
    "# consist_class_ary = np.array(consist_class)\n",
    "# consist_class_ary, consist_class_ary.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e369b4b-a062-435d-b7da-f360d9f10b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# features = np.append(\n",
    "#     features,\n",
    "#     consist_class_ary.reshape(-1, 1),\n",
    "#     axis=1\n",
    "# )\n",
    "# features, features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b8ed05-f80a-4dd5-88a2-4a64191d71a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# X_train, X_test, y_train, y_test = train_test_split(\n",
    "#     features,\n",
    "#     journal_labels,\n",
    "#     test_size=0.25,\n",
    "#     random_state=SEED\n",
    "# )\n",
    "\n",
    "# model = clf.fit(X_train, y_train)\n",
    "# y_pred = model.predict(X_test)\n",
    "\n",
    "# pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2115b871-aed8-4863-aae9-ab472710fe0e",
   "metadata": {},
   "source": [
    "## Exercise: Improve the model to improve performance\n",
    "\n",
    "SVC may not be the best model for this type of work - can you improve the classification using a different type of model? Look at [scikit-learn's supervised learning catalogue](https://scikit-learn.org/stable/supervised_learning.html) for some additional models to try.\n",
    "\n",
    "### Implement a model in XGBoost\n",
    "\n",
    "XGBoost is a popular library for building high performance classifiers. Let's implement a classifier using an [XGBClassifier](https://xgboost.readthedocs.io/en/latest/python/python_api.html#xgboost.XGBClassifier)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a2dcf4b-6c0d-4802-9f16-607da4dc8f45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from xgboost import XGBClassifier\n",
    "\n",
    "\n",
    "# # Get number of unique classes\n",
    "# n_classes = np.unique(journal_labels).size\n",
    "\n",
    "# # Replace LogisticRegression with XGBoost\n",
    "# xgb_clf = XGBClassifier(\n",
    "#     objective='multi:softmax',\n",
    "#     num_class=n_classes,  # Add this parameter\n",
    "#     random_state=SEED,\n",
    "#     n_estimators=100,\n",
    "#     learning_rate=0.1\n",
    "# )\n",
    "\n",
    "# # Fit the model - note that XGBoost doesn't require scaling\n",
    "# model = xgb_clf.fit(X_train, y_train)\n",
    "\n",
    "# # Make predictions\n",
    "# y_pred = model.predict(X_test)\n",
    "\n",
    "# # Score the model\n",
    "# pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5d4b875-4d25-4edf-a0af-05c35d96fb03",
   "metadata": {},
   "source": [
    "## Topology and Text Features\n",
    "\n",
    "Now that we've improved the topological features of our classifier, let's try adding the abstract paraphrase embedding to see how our model performs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa7213f-cf1c-477f-9db8-2b4f60544755",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "embedded_abstracts = np.load(\"data/embedded_abstracts.npy\")\n",
    "\n",
    "pca = PCA(n_components=32)\n",
    "pca_abstracts = pca.fit_transform(embedded_abstracts)\n",
    "pca_abstracts, pca_abstracts.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a951a9c-7819-4e6d-b20c-f917f70fabbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "super_features = np.append(features, pca_abstracts, axis=1)\n",
    "super_features, super_features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58f75c1f-8b1c-418d-9ad1-efa7e25315d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import RandomForestClassifier\n",
    "\n",
    "\n",
    "scaler = StandardScaler()\n",
    "clf = make_pipeline(StandardScaler(), RandomForestClassifier())\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    super_features,\n",
    "    journal_labels,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED\n",
    ")\n",
    "\n",
    "model = clf.fit(X_train, y_train)\n",
    "y_pred = model.predict(X_test)\n",
    "pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "226ac361-bc95-4fa2-b8f7-ac8f4770a725",
   "metadata": {},
   "source": [
    "### Graphs > Text\n",
    "\n",
    "Isn't it interesting that a sentence embedding of the entire abstract can tell a classifier for academic journal less than a few hand-computed network metrics?\n",
    "\n",
    "## Excrcise: Graph + Text Feature Engineering\n",
    "\n",
    "See if you can feature engineer the graph and text features together along with the model selection or model parameters to produce improved performance in our node classifier!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80a4fe12-f7a0-4134-a466-17e1524665d5",
   "metadata": {},
   "source": [
    "# Graph Machine Learning with Graph Embeddings\n",
    "\n",
    "[DeepWalk](https://arxiv.org/abs/1403.6652), Perozzi et al., 2014, was a revolution in graph machine learning. Along with [node2vec](https://snap.stanford.edu/node2vec/), Grover, A.; Leskovec, J, 2016, which came with code on Github [[aditya-grover/node2vec](https://github.com/aditya-grover/node2vec)], it removed the need to spend as much time doing feature engineering by hand. Although initially these embeddings worked on simple graphs with one type of edge and ignore node and edge properties... by efficiently encoding topology around a node, they autoomated much of the work involved in tasks like node classification/labeling and link prediction.\n",
    "\n",
    "## Sampling Our Network\n",
    "\n",
    "We will need to sample our network to compute some of these graph embeddings in a timeline manner. Remember our work from Part 3 - we will need it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2784311e-3bca-4760-bda3-8d397d84cf2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from littleballoffur import PageRankBasedSampler\n",
    "\n",
    "\n",
    "G_orig = G.copy()\n",
    "sampler = PageRankBasedSampler(5000)\n",
    "G_sample = nx.convert_node_labels_to_integers(\n",
    "    sampler.sample(G_orig.to_undirected())\n",
    ")\n",
    "\n",
    "G_sample.number_of_nodes(), G_sample.number_of_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b6efa7-f0f1-4c6d-852d-90fe416c0eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "journal_labels = np.array([x[1][\"Journal-ref-DBSCAN\"] for x in list(G_sample.nodes(data=True))])\n",
    "journal_labels, journal_labels.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb787f8f-a403-43e2-a75a-71212ef10ef7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DeepWalk was first, but it hardly state of the art. Let's see what it can do...\n",
    "from karateclub import DeepWalk\n",
    "\n",
    "\n",
    "d = DeepWalk()\n",
    "d.fit(G_sample)\n",
    "deepwalk_features = d.get_embedding()\n",
    "\n",
    "deepwalk_features, deepwalk_features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45bdd7d8-46f1-4e15-97f8-f0287a60b6c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    deepwalk_features,\n",
    "    journal_labels,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED\n",
    ")\n",
    "\n",
    "clf = RandomForestClassifier()\n",
    "model = clf.fit(X_train, y_train)\n",
    "y_pred = model.predict(X_test)\n",
    "\n",
    "pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e09dc402-203d-4ccc-a787-5b93a95136ed",
   "metadata": {},
   "source": [
    "### Node2Vec\n",
    "\n",
    "After DeepWalk came Node2Vec. Let's try it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2174f85-a25c-4fc2-8b64-1ff640b57344",
   "metadata": {},
   "outputs": [],
   "source": [
    "from karateclub import Node2Vec\n",
    "\n",
    "\n",
    "# Default: walk_number=10, walk_length=80\n",
    "d = Node2Vec(walk_number=8, walk_length=16)\n",
    "d.fit(G_sample)\n",
    "node2vec_features = d.get_embedding()\n",
    "\n",
    "node2vec_features, node2vec_features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e11e21-2bf1-4114-9e04-87642ee532da",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    node2vec_features,\n",
    "    journal_labels,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED\n",
    ")\n",
    "\n",
    "clf = RandomForestClassifier()\n",
    "model = clf.fit(X_train, y_train)\n",
    "y_pred = model.predict(X_test)\n",
    "\n",
    "pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94987cfc-73fa-4bf1-bcd0-a31ceed2c08b",
   "metadata": {},
   "source": [
    "### LaplacianEigenmaps\n",
    "\n",
    "LaplacianEigenmaps is a newer embedding, let's see how it performs..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4303e8bb-7281-4e9b-b647-f9dac1269f7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from karateclub import LaplacianEigenmaps\n",
    "\n",
    "\n",
    "d = LaplacianEigenmaps()\n",
    "d.fit(G_sample)\n",
    "laplacian_features = d.get_embedding()\n",
    "\n",
    "laplacian_features, laplacian_features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0db881e0-313e-4932-a17d-b02126a3bac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    laplacian_features,\n",
    "    journal_labels,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED\n",
    ")\n",
    "\n",
    "clf = RandomForestClassifier()\n",
    "model = clf.fit(X_train, y_train)\n",
    "y_pred = model.predict(X_test)\n",
    "\n",
    "pd.DataFrame(score_classifier(y_test, y_pred))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06fb7f2f-a3fd-432a-ba81-e9f74237b0ae",
   "metadata": {},
   "source": [
    "Before we are too hard on our classifiers, keep in mind that a text model did just as poorly as many of these. It seems likely that articles can appear in multiple physics journals and that there is some chance as to which one a given paper is accepted in.\n",
    "\n",
    "### Exercise: Compare Network Clusters with Journal Labels\n",
    "\n",
    "One way to check this is to go back to our visualization of communities above and compare them to our labels. If they don't match up, the papers may not be differentiable on that basis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88f7a665-aeb5-47e7-8a47-e343bf8645d6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f979d062-61ac-4a62-99e1-a556ba516efd",
   "metadata": {},
   "source": [
    "### Exercise: Titles, Not Abstracts\n",
    "\n",
    "What about the abstracts? Maybe they are the problem. Reduce them to 2-dimensions as we did the Journal embeddings and cluster them to see how differentiable they are. Then make an informed decision about how tenable a task our journal classification is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c5e0df-fcc7-4996-85d1-add448953376",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fdb52be0-dfc9-4dc3-aade-15a439a7311c",
   "metadata": {},
   "source": [
    "## Link Prediction\n",
    "\n",
    "Our next topic is link prediction, a common task in graph machine learning. We will use it to predict whether one paper cites another - this could be useful for missing citations or for predicting future citations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "193491ae-8b7c-433d-bd65-4f6343292ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from littleballoffur import RandomNodeEdgeSampler\n",
    "\n",
    "\n",
    "# Link prediction features are expensive to compute - sample edges from our graph\n",
    "edge_sampler = RandomNodeEdgeSampler(1000) \n",
    "sampled_edges = list(edge_sampler.sample(G.to_undirected()).edges())\n",
    "negative_sampled_edges = list(\n",
    "    map(\n",
    "        tuple, np.random.choice(G.nodes, (1000,2))\n",
    "    )\n",
    ")\n",
    "all_edges = sampled_edges + negative_sampled_edges\n",
    "\n",
    "# Set positive and negative labels for our edges\n",
    "pair_labels = [[1] * len(sampled_edges) + [0] * len(negative_sampled_edges)]\n",
    "pair_labels_ary = np.array(pair_labels).reshape([-1])\n",
    "\n",
    "# Compute some features\n",
    "jaccards = list(nx.jaccard_coefficient(G.to_undirected(), all_edges))\n",
    "ra = list(nx.resource_allocation_index(G.to_undirected(), all_edges))\n",
    "aai = list(nx.adamic_adar_index(G.to_undirected(), all_edges))\n",
    "pa = list(nx.preferential_attachment(G.to_undirected(), all_edges))\n",
    "csh = list(nx.cn_soundarajan_hopcroft(G.to_undirected(), all_edges))\n",
    "rish = list(nx.ra_index_soundarajan_hopcroft(G.to_undirected(), all_edges))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25db0df6-145a-41f7-8d70-72d2900aa670",
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_features = []\n",
    "for i, pair in enumerate(all_edges):\n",
    "    featureset = []\n",
    "    for metric in jaccards[i], ra[i], aai[i], pa[i], csh[i], rish[i]:\n",
    "        featureset.append(metric[-1])\n",
    "    pair_features.append(featureset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236386c8-787b-4fc9-a89b-3756933b20d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_features_ary = np.array(pair_features)\n",
    "pair_features_ary, pair_features_ary.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04f72eeb-adba-49cc-a1bb-9d86dfab7eb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    pair_features_ary,\n",
    "    pair_labels_ary,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED\n",
    ")\n",
    "\n",
    "clf = RandomForestClassifier()\n",
    "\n",
    "model = clf.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fcb6ee7-a4e9-474c-bc4b-888bb798e95d",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model.predict(X_test)\n",
    "score_classifier(y_test, y_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a699f91a-2fea-4adb-9c6d-7795e0665046",
   "metadata": {},
   "source": [
    "## Exercise: Link Prediction with Node Features\n",
    "\n",
    "Manual link prediction can use the same features that node classification can, if you provide features for both nodes. Improve the above link prediction by adding the previously computed node features for each node in each pair of nodes in the training and test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b18351c-95a8-4bba-8b74-4215357b2270",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8440849c-16ba-434f-947c-772c59335174",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "cdae3188-bdb7-466d-94d0-a0fd634bd53e",
   "metadata": {},
   "source": [
    "## Exercise: Link Prediction with Embeddings\n",
    "\n",
    "Pairs of embeddings can do a better job of link prediction. Append the embeddings previously computed and see if you can improve the performance of our link prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23c711c5-a6a5-4d33-b0c2-71aedc2feedc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DeepWalk was first, but it hardly state of the art. Let's see what it can do...\n",
    "from karateclub import DeepWalk\n",
    "\n",
    "\n",
    "d = DeepWalk()\n",
    "d.fit(G)\n",
    "deepwalk_features = d.get_embedding()\n",
    "\n",
    "deepwalk_features, deepwalk_features.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0a3181b-222b-433b-85d9-461798176fe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_embeddings = []\n",
    "\n",
    "for edge_pair in all_edges:\n",
    "    feature_row = np.append(\n",
    "        deepwalk_features[edge_pair[0]],\n",
    "        deepwalk_features[edge_pair[1]]\n",
    "    )\n",
    "    pair_embeddings.append(feature_row)\n",
    "\n",
    "pair_features_ary = np.stack(pair_embeddings)\n",
    "pair_features_ary, pair_features_ary.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4890ffcb-24a5-4fb1-b2f7-4fbfa31f1f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    pair_features_ary,\n",
    "    pair_labels_ary,\n",
    "    test_size=0.25,\n",
    "    random_state=SEED,\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "clf = RandomForestClassifier()\n",
    "\n",
    "model = clf.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5150bb90-c6d1-4d99-9b05-953a554b8742",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model.predict(X_test)\n",
    "score_classifier(y_test, y_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d424c16-2fdb-42bf-aed4-79ba4ee1d127",
   "metadata": {},
   "source": [
    "# Graph Machine Learning with Graph Neural Networks (GNNs)\n",
    "\n",
    "Having explored network science, we are about to dive into Graph Neural Networks (GNNs). The best introduction to GNNs is a long blog post by []() entitled [A Gentle Introduction to Graph Neural Networks](https://distill.pub/2021/gnn-intro/) which the authors have _generously_ licensed under the Creative Commons. This lets me utilize their work to explain how GNNs work while providing source code along with it to bring your theoretical understanding to a practical one.\n",
    "\n",
    "## Citation: A Gentle Introduction to Graph Neural Networks\n",
    "\n",
    "Parts of the content in Part 4 of this course are based upon: `Sanchez-Lengeling, et al., \"A Gentle Introduction to Graph Neural Networks\", Distill, 2021.` This content is cited inline. Students are encouraged to read this blog post before or after class, and to reference it if they become confused about concepts in their data science and machine learning practice. \n",
    "\n",
    "The full list of authors is:\n",
    "\n",
    "* [Benjamin Sanchez-Lengeling](https://research.google/people/106640/)\n",
    "* [Emily Reif](https://research.google/people/106150/)\n",
    "* [Adam Pearce](https://research.google/people/AdamPearce/)\n",
    "* [Alexander B. Wiltschko](https://www.linkedin.com/in/alex-wiltschko-0a7b7537/)\n",
    "\n",
    "During the course you will have access to the instructor, who understands GNNs and can elaborate further and answer any questions you may have :)\n",
    "\n",
    "## Why is there so much talk about Graph Neural Networks?\n",
    "\n",
    "Knowledge graphs are at the peak of the Gartner hype cycle and graph neural networks (GNNs) are soon to be high on the ramp because they tap and unlock the potential of enterprise knowledge graphs. Data lakes put data in one place, knowledge graphs link datasets together and graph neural networks automate business processes using data from across an enterprise. \n",
    "\n",
    "\n",
    "\n",
    "Most graph databases are fast becoming cloud-based GNN platforms:\n",
    "\n",
    "* Neo4j → [Neo4j Graph Data Science](https://neo4j.com/product/graph-data-science/)\n",
    "* TigerGraph → [Machine Learning Workbench](https://www.tigergraph.com/ml-workbench/)\n",
    "* ArangoDB → [ArrangoGraphML](https://www.arangodb.com/arangodb-for-machine-learning/)\n",
    "* Kumo → [SQL query the future](https://kumo.ai/)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e97d8554-d1db-47f0-b9cb-ed8487147a3c",
   "metadata": {},
   "source": [
    "# Graph Neural Networks (GNNs) with DGL (Deep Graph Library)\n",
    "\n",
    "[DGL or Deep Graph Library](https://dgl.ai) is the simplest way to get started with graph machine learning using graph neural networks (GNNs).\n",
    "\n",
    "First we will cover a few common operations with each major task type we covered in the lecture: node-level, edge-level, subgraph-level and graph-level.\n",
    "\n",
    "## Node-Level Tasks: Classification\n",
    "\n",
    "Node-level tasks usually involve property prediction - classifying nodes into categories or regressing one of their numeric properties. We'll cover both.\n",
    "\n",
    "As in the network science section of this course, we will start with a Text Attributed Graph (TAG) called a Citation Graph. We are going to use the [CORA dataset](https://relational.fit.cvut.cz/dataset/CORA), [described by Papers with Code](https://paperswithcode.com/dataset/cora) as:\n",
    "\n",
    "> Introduced by Andrew McCallum et al. in [Automating the Construction of Internet Portals with Machine Learning](https://doi.org/10.1023/A:1009953814988)\n",
    ">\n",
    "> The Cora dataset consists of 2708 scientific publications classified into one of seven classes. The citation network consists of 5429 links. Each publication in the dataset is described by a 0/1-valued word vector indicating the absence/presence of the corresponding word from the dictionary. The dictionary consists of 1433 unique words.\n",
    "\n",
    "### CORA Node Features: Bag of Words\n",
    "\n",
    "Note... the features for this network are a [Bag of Words](https://en.wikipedia.org/wiki/Bag-of-words_model) model: simple and _sparse_ rather than modern text representations which are _dense_, distributed representations in the form of language models or [embeddings](https://cloud.google.com/blog/topics/developers-practitioners/meet-ais-multitool-vector-embeddings). Each node has a row in the feature matrix and each of 1,433 unique words get a column with the word count. Before [Word2Vec](https://arxiv.org/abs/1301.3781) introduced text embeddings in 2013, the features for NLP problems were mostly 0s, with a few non-zero values.\n",
    "\n",
    "<center><img src=\"images/sparse_vs_dense_vectors.webp\" width=\"800px\" alt=\"Bag-of-Words (BoW) sparse vectors used in traditional NLP versus dense, embedded vector representations used in modern deep learning NLP\" /></center>\n",
    "\n",
    "The [curse of dimensionality](https://en.wikipedia.org/wiki/Curse_of_dimensionality) prevented NLP applications from realizing their modern capabilities - the more words that were added, the more dimensions the features data got and the more dimensions you add to a _sparse_ feature vector... the more all the values of that vector start to approximate the same value. They stretch out over many dimensions and look the same.\n",
    "\n",
    "Embeddings like Word2Vec related _sparse representations_ of words to the text around them by storing a middle layer of a neural network, creating _dense representations_.\n",
    "\n",
    "<center><img src=\"images/from_sparse_to_dense.webp\" width=\"700px\" alt=\"The Word2Vec's Skipgram architecture maps sparse to dense vectors via a shallow embedding technique\" /></center>\n",
    "\n",
    "These are very useful because the dimensions of the feature vector correspond to particular semantics, and because you can compare two dense vectors and get a sense of how similar the objects the represent are. This is very useful for information retreival applications like search and clustering.\n",
    "\n",
    "<center><img src=\"images/king_minus_man_plus_woman.webp\" width=\"800px\" alt=\"Given the dense embedding vector for the word 'king', if we subtract the vector for 'man' and add 'woman', we arrive at a vector very close to 'queen'.\" /></center>\n",
    "\n",
    "<br />\n",
    "\n",
    "We could use a language model or large language model (LLM) to embed the features or the original text and get better performance from our GCN. However, it is good to start simple and worry about feature engineering lately... you can spend an endless amount of time over optimizing a task nobody cares about. Make sure they want the prototype before you engineer incredible performance. A Bag of Words representation is a fine start.\n",
    "\n",
    "### CORA Classifier: Graph Convolutional Network\n",
    "\n",
    "We are going to use a neural network architecture that may be familiar to you: a convolutional neural network. The type we will employ is called a Graph Convolutional Network (GCN). Message passing occurs between nodes and the series of input messages to a node are summarized by the layers of a GCN after each round of message passing.\n",
    "\n",
    "<center><img src=\"images/gcn-decagon-overview.png\" width=\"1000px\" alt=\"Graph Neural Networks for Multirelational Link Prediction\" /><a href=\"https://snap.stanford.edu/decagon/\">Graph Neural Networks for Multirelational Link Prediction, Zitnik et al., 2018</a></center>\n",
    "\n",
    "\n",
    "There is often a big of tinkering required to make GNNs run, so even for this simple problem in DGL, we must specify our GNN architectre. It is simple enough. Let's see how it looks...\n",
    "\n",
    "Note: Figures Sources: [Dense Vectors: Capturing Meaning with Code](https://towardsdatascience.com/dense-vectors-capturing-meaning-with-code-88fc18bd94b9) by [James Briggs](https://jamescalam.medium.com/), [Graph Neural Networks for Multirelational Link Prediction, Zitnik et al., 2018](https://snap.stanford.edu/decagon/)\n",
    "\n",
    "### Building a GCN in DGL\n",
    "\n",
    "Let's build, train and evaluate our first GNN: a graph convoltional network for classifying CORA articles into categories.\n",
    "\n",
    "Note: Source for this section is the [Blitz Tutorial, Node Classification with DGL](https://docs.dgl.ai/tutorials/blitz/1_introduction.html#sphx-glr-tutorials-blitz-1-introduction-py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d27e1dac-9af7-473c-b65f-5a8d854bf81d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# DGL can also use Tensorflow or MXNet\n",
    "os.environ[\"DGLBACKEND\"] = \"pytorch\"\n",
    "import dgl\n",
    "import dgl.data\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00a5ce97-daeb-4446-9a90-6976b58d417b",
   "metadata": {},
   "source": [
    "For now we will use a pre-loaded dataset. It contains the standard CORA bag-of-word (BoW) featres. Later we will construct our own graphs to perform feature engineering on them to do more sophisticated work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f36d76-168a-4406-bb7c-4d127c5bb3db",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dgl.data.CoraGraphDataset()\n",
    "\n",
    "print(f\"Number of categories: {dataset.num_classes}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b369768-a99c-4a5c-bdc6-a78d8cbe93dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# There can be more than one graph, this dataset has just one\n",
    "g = dataset[0]\n",
    "g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c3fa7a9-8084-4f1f-a8c6-67a02e0775fe",
   "metadata": {},
   "source": [
    "`train_mask`, `val_mask` and `test_mask` are bit masks that denote the rows in the `label` and `feat` [Schemes](https://github.com/dmlc/dgl/blob/master/python/dgl/frame.py#L125) which with `DGLBACKEND=pytorch` contain DGL mappings to the [torch.Tensors](https://pytorch.org/docs/stable/tensors.html) making up the training, validation and test datasets respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48502f10-280f-4230-8939-16847c30cba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Node features\")\n",
    "print(g.ndata[\"feat\"].shape)\n",
    "\n",
    "print(\"Edge features\")\n",
    "print(g.edata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0f34d9f-c178-4dd4-8d05-d796b62f9ce2",
   "metadata": {},
   "source": [
    "### GCN Model Architecture - Diagrams, then Code\n",
    "\n",
    "The model itself is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) that uses the [dgl.nn.conv.GraphConv](https://docs.dgl.ai/generated/dgl.nn.pytorch.conv.GraphConv.html) class. \n",
    "\n",
    "<center><img src=\"images/Schematic-diagram-of-a-two-layer-GCN-model-The-dark-green-denotes-target-nodes-that-need_W640.jpg\" alt=\"Diagram of 2-layer GCN from Graph neural networks in node classification: survey and evaluation, Xiao et al., 2022\" width=\"600px\" /></center>\n",
    "\n",
    "<br />\n",
    "\n",
    "<center>Image credit: <a href=\"https://www.researchgate.net/publication/355873169_Graph_neural_networks_in_node_classification_survey_and_evaluation\">Diagram of 2-layer GCN from Graph neural networks in node classification: survey and evaluation, Xiao et al., 2022</a></center>\n",
    "\n",
    "<br />\n",
    "\n",
    "Let's dig into this diagram of our GCN before coding it in DGL.\n",
    "\n",
    "### Over Smoothing in GNNs: Too Many Layers Means Too Many Hops Sampled\n",
    "\n",
    "Note that **each layer of the GCN represents a round of message passing where nodes aggregate information from their neighbors.** This is important to know, as if you have too many layers in a GNN, you run into the [oversmoothing problem](https://towardsdatascience.com/over-smoothing-issue-in-graph-neural-network-bddc8fbc2472) where nodes start to look the same as all the other nodes.\n",
    "\n",
    "<center><img src=\"images/GNN-oversmoothing-first-layer.webp\" width=\"840px\" alt=\"First layer of GNN message passing, aggregation and summarization results in features of different colors\" /></center>\n",
    "<center>The first layer of GNN message passing, aggregation and summarization results in features represented by different colors.</center>\n",
    "<center><i>Image credit: <a href=\"https://towardsdatascience.com/over-smoothing-issue-in-graph-neural-network-bddc8fbc2472\">Over-smoothing issue in graph neural network</a> by <a href=\"https://towardsdatascience.com/over-smoothing-issue-in-graph-neural-network-bddc8fbc2472\">Anas Ait Aomar</a></i></center>\n",
    "\n",
    "<br /><br />\n",
    "\n",
    "<center><img src=\"images/GNN-oversmoothing-second-layer.webp\" width=\"1000px\" alt=\"Second layer of GNN message passing, aggregation and summarization results in features with more similar colors\" /></center>\n",
    "<center>The second layer of GNN message passing, aggregation and summarization results in features represented by more similar colors.</center>\n",
    "\n",
    "<center><i>Image credit: <a href=\"https://towardsdatascience.com/over-smoothing-issue-in-graph-neural-network-bddc8fbc2472\">Over-smoothing issue in graph neural network</a> by <a href=\"https://towardsdatascience.com/over-smoothing-issue-in-graph-neural-network-bddc8fbc2472\">Anas Ait Aomar</a></i></center>\n",
    "\n",
    "### Relu Activation Function\n",
    "\n",
    "Note how the GraphConv layers in the GCN architecture diagram above are separated by a Relu layer. Without this layer, the GCN could not learn effectively. Relu is an activation function that enables nonlinearity in neural networks - it lets them model messy data in a way that is much more powerful than a linear model. Relu is defined as `max(0, x)` which means that it maps negative values to 0 and positive values are left alone. Note that there are many derivatives of Relu that attempt to improve its performance.\n",
    "\n",
    "<center><img src=\"images/relu.png\" width=\"600px\" alt=\"Relu is max(0, x), making its plot flat when x is less than zero, and evently diagonal in a 1:1 ratio when x is greater than zero.\" /></center>\n",
    "<center>The Relu activation function: <code>max(0, x)</code></center>\n",
    "<center><i>Image Credit: <a href=\"https://medium.com/@danqing/a-practical-guide-to-relu-b83ca804f1f7\">A Practical Guide to ReLU</a> by <a href=\"https://medium.com/@danqing\">Danqing Liu</a></i></center>\n",
    "\n",
    "> we can stack as many linear classifiers as we want on top of each other, and without nonlinear functions between them, it will just be the same as one linear classifier.\n",
    ">\n",
    "> But if we put a nonlinear function between them, such as max, then this is no longer true. Now each linear layer is actually somewhat decoupled from the other ones and can do its own useful work. The max function operates as a simple if statement.\n",
    ">\n",
    "_Source: [Nonlinearity and Neural Networks](https://medium.com/unpackai/nonlinearity-and-neural-networks-2ffaaac0e6ff) by [Aravinda 加阳](https://medium.com/@aravinda-gn)_\n",
    "\n",
    "This video by [deeplizard on Youtube](https://www.youtube.com/@deeplizard) explains Relu and its significance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1526ec3-dc02-41b5-9185-7e4fb342ac24",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%HTML\n",
    "<center><iframe width=\"800\" height=\"460\" src=\"https://www.youtube.com/embed/6MmGNZsA5nI?si=sglt8BijkpykWdWP&amp;start=10\"></iframe></center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "096def14-9ac0-4b30-947a-b6ab70dcb36b",
   "metadata": {},
   "source": [
    "### Coding the Above GCN Diagram\n",
    "\n",
    "The equivalent DGL code for the GCN diagram above appears below. The graph structure and CORA BoW features are shown as the input, which feeds into one GCN layer, then a Relu activaton function, another GCN layer and finally they are mapped into the labels of our classes, in this case fields of study."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eef1325-e796-4619-82cd-a66f911ca2b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dgl.nn import GraphConv\n",
    "\n",
    "\n",
    "class GCN(nn.Module):\n",
    "    \"\"\"2-layer Graph Convolutional Network\"\"\"\n",
    "    \n",
    "    def __init__(self, in_feats, h_feats, num_classes):\n",
    "        \"\"\"Setup two GCN layers of with the input, inner and output dimensions.\"\"\"\n",
    "        super(GCN, self).__init__()\n",
    "        self.conv1 = GraphConv(in_feats, h_feats)\n",
    "        self.conv2 = GraphConv(h_feats, num_classes)\n",
    "\n",
    "    def forward(self, g, in_feat):\n",
    "        \"\"\"Operate a forward pass of the network\"\"\"\n",
    "        h = self.conv1(g, in_feat)\n",
    "        h = F.relu(h)\n",
    "        h = self.conv2(g, h)\n",
    "        return h\n",
    "\n",
    "\n",
    "# Create the model with given dimensions\n",
    "model = GCN(g.ndata[\"feat\"].shape[1], 16, dataset.num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "447bbf5b-8caf-43d3-be50-cf6254f9f13d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7393294-5d9d-4823-b76a-530f2e56c66a",
   "metadata": {},
   "source": [
    "### Training a GCN\n",
    "\n",
    "Below we define a training function that will iteratively train our GCN using message passing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35dad0c6-bc5c-43f9-b719-e4dd97eed4e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "\n",
    "\n",
    "def metrics(y_true, y_pred):\n",
    "\n",
    "    return {\n",
    "        \"accuracy\": accuracy_score(y_true, y_pred),\n",
    "        \"precision\": precision_score(y_true, y_pred, average=\"micro\"),\n",
    "        \"recall\": recall_score(y_true, y_pred, average=\"micro\"),\n",
    "        \"f1\": f1_score(y_true, y_pred, average=\"micro\"),\n",
    "    }\n",
    "\n",
    "\n",
    "def train(g, model):\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "    best_val_acc = 0\n",
    "    best_test_acc = 0\n",
    "\n",
    "    features = g.ndata[\"feat\"]\n",
    "    labels = g.ndata[\"label\"]\n",
    "    train_mask = g.ndata[\"train_mask\"]\n",
    "    val_mask = g.ndata[\"val_mask\"]\n",
    "    test_mask = g.ndata[\"test_mask\"]\n",
    "    for e in range(100):\n",
    "        # Forward\n",
    "        logits = model(g, features)\n",
    "        # print(logits.shape)\n",
    "\n",
    "        # Compute prediction\n",
    "        pred = logits.argmax(1)\n",
    "\n",
    "        # Compute loss\n",
    "        # Note that you should only compute the losses of the nodes in the training set.\n",
    "        loss = F.cross_entropy(logits[train_mask], labels[train_mask])\n",
    "\n",
    "        # Compute accuracy on training/validation/test\n",
    "        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()\n",
    "        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()\n",
    "        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()\n",
    "\n",
    "        train_scores = metrics(labels[train_mask], pred[train_mask])\n",
    "        val_scores = metrics(labels[val_mask], pred[val_mask])\n",
    "        test_scores = metrics(labels[test_mask], pred[test_mask])\n",
    "\n",
    "        # Save the best validation accuracy and the corresponding test accuracy.\n",
    "        if best_val_acc < val_acc:\n",
    "            best_val_acc = val_acc\n",
    "            best_test_acc = test_acc\n",
    "\n",
    "        # Backward\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if e % 5 == 0:\n",
    "            print(\n",
    "                f\"In epoch {e}, loss: {loss:.3f}, val acc: {val_acc:.3f} (best {best_val_acc:.3f}), test acc: {test_acc:.3f} (best {best_test_acc:.3f}),\",\n",
    "                f'val precision: {val_scores[\"precision\"]:.3f}, val recall: {val_scores[\"recall\"]:.3f}, val f1: {val_scores[\"f1\"]:.3f}'\n",
    "            )\n",
    "\n",
    "\n",
    "model = GCN(g.ndata[\"feat\"].shape[1], 16, dataset.num_classes)\n",
    "train(g, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceae7723-2ecd-464e-8232-ad17d6e1762c",
   "metadata": {},
   "source": [
    "## Graph Attention Networks (GATs)\n",
    "\n",
    "Let's try a more sophisticated architecture for node classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd5cac91-65c7-457e-99e7-e3fffc7d826e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dgl.nn import GATConv\n",
    "\n",
    "\n",
    "class GAT(nn.Module):\n",
    "    def __init__(self, in_dim, hidden_dim, out_dim, num_heads):\n",
    "        super(GAT, self).__init__()\n",
    "        self.layer1 = GATConv(in_dim, hidden_dim, num_heads=num_heads, activation=F.relu, feat_drop=0.3, attn_drop=0.1)\n",
    "        self.layer2 = GATConv(hidden_dim * num_heads, out_dim, num_heads=1)\n",
    "\n",
    "    def forward(self, g, in_feat):\n",
    "        h = self.layer1(g, in_feat)\n",
    "        h = h.view(h.shape[0], -1)\n",
    "        h = self.layer2(g, h)\n",
    "        return h.squeeze(1)\n",
    "\n",
    "\n",
    "gatconv = GAT(g.ndata[\"feat\"].shape[1], 10, dataset.num_classes, num_heads=2)\n",
    "train(g, gatconv)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a0614dd-46d7-463a-bebc-894afa922fa4",
   "metadata": {},
   "source": [
    "## GraphDataLoaders\n",
    "\n",
    "I have written a GraphDataLoader for our dataset in [`graphml_class/citation/dgl.py`](graphml_class/citation/dgl.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a0ea22c-6a01-48e9-af5d-0cac062f795c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "3abdb6ff-6efd-45ad-a760-6b68665933e6",
   "metadata": {},
   "source": [
    "# Link Prediction with GNNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "730f841f-a276-433d-9a47-279a0d63b4c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl.data\n",
    "\n",
    "dataset = dgl.data.CoraGraphDataset()\n",
    "g = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33eddec-7d29-4521-a561-905df95d374a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse as sp\n",
    "\n",
    "# Split edge set for training and testing\n",
    "u, v = g.edges()\n",
    "\n",
    "u, v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "087388f5-b85d-492f-8110-e383edd747c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "eids = np.arange(g.number_of_edges())\n",
    "eids = np.random.permutation(eids)\n",
    "\n",
    "test_size = int(len(eids) * 0.1)\n",
    "train_size = g.number_of_edges() - test_size\n",
    "\n",
    "test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]\n",
    "train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]\n",
    "\n",
    "# Find all negative edges and split them for training and testing\n",
    "adj = sp.coo_matrix((np.ones(len(u)), (u.numpy(), v.numpy())))\n",
    "adj_neg = 1 - adj.todense() - np.eye(g.number_of_nodes())\n",
    "neg_u, neg_v = np.where(adj_neg != 0)\n",
    "\n",
    "neg_eids = np.random.choice(len(neg_u), g.number_of_edges())\n",
    "test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]\n",
    "train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39af1702-dc07-497d-a0bb-4bd933b072bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_g = dgl.remove_edges(g, eids[:test_size])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e1de23-6a33-47fc-8df8-793126d52248",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dgl.nn import SAGEConv\n",
    "\n",
    "# ----------- 2. create model -------------- #\n",
    "# build a two-layer GraphSAGE model\n",
    "class GraphSAGE(nn.Module):\n",
    "    def __init__(self, in_feats, h_feats):\n",
    "        super(GraphSAGE, self).__init__()\n",
    "        self.conv1 = SAGEConv(in_feats, h_feats, 'mean')\n",
    "        self.conv2 = SAGEConv(h_feats, h_feats, 'mean')\n",
    "        # self.conv3 = SAGEConv(h_feats, h_feats, 'mean')\n",
    "\n",
    "    def forward(self, g, in_feat):\n",
    "        h = self.conv1(g, in_feat)\n",
    "        h = F.relu(h)\n",
    "        h = self.conv2(g, h)\n",
    "        # h = self.conv3(g, h)\n",
    "        return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96109427-2630-41d1-9c8c-5da3f2a94203",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.number_of_nodes())\n",
    "train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.number_of_nodes())\n",
    "\n",
    "test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.number_of_nodes())\n",
    "test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.number_of_nodes())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90e0847c-4309-4314-9f48-30ae5d5c0bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl.function as fn\n",
    "\n",
    "class DotPredictor(nn.Module):\n",
    "    def forward(self, g, h):\n",
    "        with g.local_scope():\n",
    "            g.ndata['h'] = h\n",
    "            # Compute a new edge feature named 'score' by a dot-product between the\n",
    "            # source node feature 'h' and destination node feature 'h'.\n",
    "            g.apply_edges(fn.u_dot_v('h', 'h', 'score'))\n",
    "            # u_dot_v returns a 1-element vector for each edge so you need to squeeze it.\n",
    "            return g.edata['score'][:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d84ac47-d26a-4bc3-ae8d-f24ead21716f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GraphSAGE(train_g.ndata['feat'].shape[1], 16)\n",
    "# You can replace DotPredictor with MLPPredictor.\n",
    "#pred = MLPPredictor(16)\n",
    "pred = DotPredictor()\n",
    "\n",
    "def compute_loss(pos_score, neg_score):\n",
    "    scores = torch.cat([pos_score, neg_score])\n",
    "    labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])])\n",
    "    return F.binary_cross_entropy_with_logits(scores, labels)\n",
    "\n",
    "def compute_auc(pos_score, neg_score):\n",
    "    scores = torch.cat([pos_score, neg_score]).numpy()\n",
    "    labels = torch.cat(\n",
    "        [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()\n",
    "    return roc_auc_score(labels, scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88d36ece-912e-47e3-9c3f-d6ffcba9e677",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "# ----------- 3. set up loss and optimizer -------------- #\n",
    "# in this case, loss will in training loop\n",
    "optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=0.01)\n",
    "\n",
    "# list(itertools.chain(model.parameters(), pred.parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5943e11-591c-4304-9eef-56cb59e2f9a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------- 4. training -------------------------------- #\n",
    "all_logits = []\n",
    "for e in range(100):\n",
    "    # forward\n",
    "    h = model(train_g, train_g.ndata['feat'])\n",
    "    pos_score = pred(train_pos_g, h)\n",
    "    neg_score = pred(train_neg_g, h)\n",
    "    loss = compute_loss(pos_score, neg_score)\n",
    "\n",
    "    # backward\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if e % 5 == 0:\n",
    "        print('In epoch {}, loss: {}'.format(e, loss))\n",
    "\n",
    "# ----------- 5. check results ------------------------ #\n",
    "from sklearn.metrics import roc_auc_score\n",
    "with torch.no_grad():\n",
    "    pos_score = pred(test_pos_g, h)\n",
    "    neg_score = pred(test_neg_g, h)\n",
    "    print('AUC', compute_auc(pos_score, neg_score))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44ddc839-3dca-4458-905d-ce35793f7d3c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}