Part 1 - Knowledge Graph Construction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5319f2-b647-4bcc-90b2-8f748ed39eac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import calendar\n",
    "import gzip\n",
    "import io\n",
    "import json\n",
    "import pickle\n",
    "import os\n",
    "import re\n",
    "import tarfile\n",
    "from datetime import date\n",
    "from typing import Dict, List, Union\n",
    "\n",
    "import graphistry\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import requests\n",
    "import seaborn as sns\n",
    "import umap\n",
    "from graphistry import Plottable\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sklearn.cluster import DBSCAN, KMeans\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from torch import Tensor\n",
    "\n",
    "from graphml_class.clean import clean_graph\n",
    "from graphml_class.palette import CATEGORICAL_PALETTE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c866f3f-adce-447d-ab09-f0165b3bd20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make Seaborn pretty :)\n",
    "sns.set(style='white', context='poster')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca312d8d-ed3d-4f0d-811f-58d5c810a8c1",
   "metadata": {},
   "source": [
    "Part 1: Knowledge Graph Construction\n",
    "====================================\n",
    "\n",
    "In this section of the course, we will cover _knowledge graph construction_ or how to construct knowledge graphs from both _natural networks_ and _structural networks_. In the lecture I described _knowledge graph construction_ as the process of building a knowledge graph from raw data using ETL, data transformations or Natural Language Processing (NLP).\n",
    "\n",
    "There are two main types of networks in common use: simple and heterogeneous networks. We're going to start out building a simple network where nodes are _academic papers_ and edges are _citations between papers_.\n",
    "\n",
    "<center><img src=\"images/Graphs-vs-Heterogeneous-Graphs-2000px.png\" width=\"1000px\" /></center>\n",
    "\n",
    "There are two categories of data from which we can build networks: _natural networks_ and _structural networks_. We can also transform _existing networks_ that are already formatted and easy to work with.\n",
    "\n",
    "<center><img src=\"images/ETL-in-Natural-and-Structural-Graphs.jpg\" width=\"860px\" /></center>\n",
    "\n",
    "Structural networks are just as important as natural networks because...\n",
    "\n",
    "<center><img src=\"images/Raw-Data-are-Often-Not-Networks.png\" width=\"960px\" /></center>\n",
    "<center>Slide on Network Construction from <a href=\"https://scholar.google.com/citations?user=Q_kKkIUAAAAJ&hl=en&oi=ao\">Jure Leskovec's</a> <a href=\"https://web.stanford.edu/class/cs224w/\">Stanford CS224W class</a> when it was still called <i>Network Analysis</i>. Today it is called <i>Machine Learning with Graphs</i></center>\n",
    "\n",
    "<br />\n",
    "\n",
    "Given time, we'll be building both _simple_ and _heterogeneous networks_ from both _natural_ and _structural graphs_. We'll start with the former.\n",
    "\n",
    "# Section Textbook\n",
    "\n",
    "An excellent resource for the first two parts of this course, **Knowledge Graph Construction** and **Network Science** is the [Network Science (CC4063 / CC4095)](https://www.dcc.fc.up.pt/~pribeiro/aulas/ns2122/) class taught by [Pedro Ribeiro](https://www.dcc.fc.up.pt/~pribeiro/) at the [Center for Research in Advanced Computing Systems](https://cracs.fc.up.pt/), part of the [Computer Science Department](https://www.dcc.fc.up.pt/site/) of the [University of Porto](https://www.up.pt/portal/en/).\n",
    "\n",
    "We will be using [Section 10: Network Construction](https://www.dcc.fc.up.pt/~pribeiro/aulas/ns2122/handouts.html#construction) during this part of the course, and specifically the slides for that section: [Network Construction (selected slides from J. Leskovec and L. Lacasa)](https://www.dcc.fc.up.pt/~pribeiro/aulas/ns2122/10_netconstruction.pdf)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33276594-d0c7-4a44-8637-a5426107134e",
   "metadata": {},
   "source": [
    "# Setting up Graphistry\n",
    "\n",
    "First let's setup a network visualization tool to help us evaluate what we are building. Throughout this part of the course we will be using `pygraphistry` and [Graphistry Hub](https://hub.graphistry.com/) [https://hub.graphistry.com/](https://hub.graphistry.com/) to visualize networks. Both are free for personal use and are powerful for visualizing networks large and small.\n",
    "\n",
    "You can [signup](https://hub.graphistry.com/accounts/signup/) for a Graphistry account at [https://hub.graphistry.com/accounts/signup/](https://hub.graphistry.com/accounts/signup/). <b>You should use a username/password/email to get the required credentials</b>, although after that you can login with your Github or Google account.\n",
    "\n",
    "<center><img src=\"images/graphistry_hub_registration.png\" /></center>\n",
    "\n",
    "Retain and use your credentials in the login form and in the environment variables in the next cell below. You should set the `GRAPHISTRY_USERNAME` and `GRAPHISTRY_PASSWORD` variables in the `env/graphistry.env` file, and then restart this docker container to pickup the new values.\n",
    "\n",
    "<center><img src=\"images/graphistry_hub_homepage.png\" /></center>\n",
    "\n",
    "## Graphistry Streamlit Examples\n",
    "\n",
    "Graphistry can be used with [Streamlit]([streamlit](https://streamlit.io/)) to build simple, all-Python web applications with forms driving interactive network visualizations. Check out [graphistry/graph-app-kit](https://github.com/graphistry/graph-app-kit/tree/master/src/python/views/) for a number of very-simple-to-build [`streamlit`](https://pypi.org/project/streamlit/) applications leveraging [pygraphistry](https://pygraphistry.readthedocs.io/) ([graphistry/pygraphistry](https://github.com/graphistry/pygraphistry)).\n",
    "\n",
    "Graphistry's [UMAP support](https://umap-learn.readthedocs.io/en/latest/plotting.html) does something similar to what we do below when construction a KNN vector network, but using GPUs in **a few lines of code**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94962741-ff18-417f-a218-52e6a61f5a0e",
   "metadata": {},
   "source": [
    "# Knowledge Graph Construction via Extract-Transform-Load (ETL) for a Simple, Natural Graph\n",
    "\n",
    "We are going to start out by building a knowledge graph from an existing edge list and then add properties to it in a process called [Extract-Transform-Load or ETL](https://aws.amazon.com/what-is/etl/) (some peope like [ELT](https://www.techtarget.com/searchdatamanagement/definition/Extract-Load-Transform-ELT), it's all the same to me). A common method of building an enterprise knowledge graph as a _property graph_ is to take a company's data on a platform like [PySpark](https://spark.apache.org/docs/3.3.1/api/python/index.html#:~:text=PySpark%20is%20an%20interface%20for,data%20in%20a%20distributed%20environment.) [a Python interface to [Apache Spark](https://spark.apache.org/). [Databricks](https://www.databricks.com/glossary/pyspark) is a popular [Platform as a Service (PaaS)](https://www.techtarget.com/searchcloudcomputing/definition/Platform-as-a-Service-PaaS#:~:text=Platform%20as%20a%20service%20(PaaS,software%20on%20its%20own%20infrastructure.) solution for Spark and [MLFlow](https://mlflow.org/), a [machine-learning operations (MLOps)](https://www.databricks.com/glossary/mlops) platform] and ETL it into nodes and edges, each of which have properties. Then this data is loaded into a graph database like [Neo4j](https://neo4j.com/developer/cypher/) for querying [[Neo4j Community Edition](https://neo4j.com/licensing/#neo4j-community-edition:~:text=About%20Neo4j%20Licenses-,Neo4j%20Community%20Edition,use%20Community%20Edition%2C%20whether%20or%20not%20your%20software%20is%20closed%20source.,-Download%20Neo4j%20%E2%86%92) is running on your machine right now via [docker compose](https://docs.docker.com/compose/) if you're running this course]. We will demonstrate that process below using [NetworkX](https://networkx.org/) and [PyData](https://pydata.org/) [Pandas](https://pandas.pydata.org/docs/). We will also perform Natural Language Processing to create label for our nodes to use in Part 3 - Graph Analytics and Part 4 - Graph Machine Learning.\n",
    "\n",
    "After that we will use text embeddings [sentence encoders] to build a K-Nearest-Neighbors network of papers as related by the semantics [meaning] of their abstracts. For now, think about ETL :)\n",
    "\n",
    "## High-Energy Physics Theory Citation Network\n",
    "\n",
    "We'll be using the [High-energy physics theory citation network](https://snap.stanford.edu/data/cit-HepTh.html) from [Stanford SNAP](https://snap.stanford.edu/index.html). SNAP has many large network datasets available in the [Stanford Large Network Dataset Collection](https://snap.stanford.edu/data/).\n",
    "\n",
    "The dataset includes the following files, which we will combine:\n",
    "\n",
    "* [Citation graph edge list](https://snap.stanford.edu/data/cit-HepTh.txt.gz) contains node ID pairs. Node IDs are standard paper identifiers. This will build the core structure of our network.\n",
    "* [Paper metadata](cit-HepTh-abstracts.tar.gz) including abstracts. This will add propertis to our network.\n",
    "* [Publishing dates on arXiv](https://snap.stanford.edu/data/cit-HepTh-dates.txt.gz) will make our citation network a temporal [Directed-Acyclic-Graph (DAG)](https://en.wikipedia.org/wiki/Directed_acyclic_graph) since one paper can't cite another before it is written and there are no reciprocal edges. While we don't focus on this, it does affect our analysis.\n",
    "\n",
    "## Dataset Citation\n",
    "\n",
    "```\n",
    "Paper: hep-th/0002031\n",
    "From: Maulik K. Parikh \n",
    "Date: Fri, 4 Feb 2000 17:04:51 GMT   (10kb)\n",
    "\n",
    "Title: Confinement and the AdS/CFT Correspondence\n",
    "Authors: D. S. Berman and Maulik K. Parikh\n",
    "Comments: 12 pages, 1 figure, RevTeX\n",
    "Report-no: SPIN-1999/25, UG-1999/42\n",
    "Journal-ref: Phys.Lett. B483 (2000) 271-276\n",
    "\\\\\n",
    "  We study the thermodynamics of the confined and unconfined phases of\n",
    "superconformal Yang-Mills in finite volume and at large N using the AdS/CFT\n",
    "correspondence. We discuss the necessary conditions for a smooth phase\n",
    "crossover and obtain an N-dependent curve for the phase boundary.\n",
    "\\\\\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c191a7-0372-4723-9909-cefd2afc61a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment variable setup\n",
    "GRAPHISTRY_USERNAME = os.getenv(\"GRAPHISTRY_USERNAME\")\n",
    "GRAPHISTRY_PASSWORD = os.getenv(\"GRAPHISTRY_PASSWORD\")\n",
    "\n",
    "# GRAPHISTRY_USERNAME = \"rjurney\"\n",
    "# GRAPHISTRY_PASSWORD = \"my_password\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e92082-faad-40a2-836e-ca7f98c5f081",
   "metadata": {},
   "outputs": [],
   "source": [
    "graphistry.register(\n",
    "    api=3,\n",
    "    username=GRAPHISTRY_USERNAME,\n",
    "    password=GRAPHISTRY_PASSWORD,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09972cd9-b909-445f-8b78-94d1978b4125",
   "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": "c23d5a59-b4b5-4854-83d9-47a317b06058",
   "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": "markdown",
   "id": "46b3cf16-524f-49a0-9865-eb5924e357d1",
   "metadata": {},
   "source": [
    "## `NetworkX` on PyPi is `networkx` in code is `nx`.\n",
    "\n",
    "The convention we used above is to load NetworkX via `import networkx as nx` so we can use the shorthand `nx` to call its classes and algorithms.\n",
    "\n",
    "## Numeric Node IDs\n",
    "\n",
    "What follows is a demonstration of _knowledge graph construction_, which we covered in the lecture. A node/edge list was provided, but the IDs are not sequential... which a network sampling tool I hope to use called [littleballoffur](https://github.com/benedekrozemberczki/littleballoffur) requires. In fact many graph libraries require sequential IDs. We must transform the graph IDs, create a mapping back and forth and annotate the nodes with properties for both IDs.\n",
    "\n",
    "## Build a Directional Graph (nx.DiGraph) from a CSV\n",
    "\n",
    "The edge list is a `#` commented, space-delimited CSV. We will parse it, assign sequential IDs and build a [`nx.DiGraph`](https://networkx.org/documentation/stable/reference/classes/digraph.html).\n",
    "\n",
    "### Download the Citation Edge List\n",
    "\n",
    "First, we download the edge list and build the structure of the network: `(paper)-cited->(paper)`. Note that we cache the edge list so you can edit the code without having to re-download the data.\n",
    "\n",
    "The edge list is located at [https://snap.stanford.edu/data/cit-HepTh.txt.gz](https://snap.stanford.edu/data/cit-HepTh.txt.gz) and is stored in `data/cit-HepTh.txt.gz`. We will read the file in its compressed state via the `gzip` builtin library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a4b88f4-deaf-444a-b90f-10bd9b7dad4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download and load edges (citations) from `cit-HepTh.txt.gz`\n",
    "edge_path = \"data/cit-HepTh.txt.gz\"\n",
    "gzip_content = None\n",
    "\n",
    "if os.path.exists(edge_path):\n",
    "    print(f\"Using existing citation graph edge file {edge_path}\")\n",
    "    gzip_content = open(edge_path, \"rb\")\n",
    "else:\n",
    "    print(\"Fetching citation graph edge file ...\")\n",
    "    response = requests.get(f\"https://snap.stanford.edu/{edge_path}\")\n",
    "    gzip_content = io.BytesIO(response.content)\n",
    "\n",
    "    print(\"Writing edge list to file {edge_path}\")\n",
    "    with open(edge_path, \"wb\") as f:\n",
    "        f.write(response.content)\n",
    "        print(f\"Wrote downloaded edge file to {edge_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1ba635-1606-4e15-9b2f-795b79e44753",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the top 10 lines of our gzip text file\n",
    "!zcat data/cit-HepTh.txt.gz | head -10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88106b1c-c9a1-4ebe-82e2-a54d2b31df9e",
   "metadata": {},
   "source": [
    "### Graph and Identifier Setup\n",
    "\n",
    "We create directional a [`nx.DiGraph`](https://networkx.org/documentation/stable/reference/classes/digraph.html) because citations are inherently directional: from citer to cited. Note that whether we model them this way or not, citation graphs are temporal networks. The citing paper's publishing date must fall after cited paper's publishing date. We'll load publishing dates below.\n",
    "\n",
    "We need to setup `file_to_net` and `net_to_file` dictionaries to map back and forth between the file format's identifiers and or own sequential identifiers we'll be assigning starting with 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78b66727-06b4-4c72-b02a-1caf2a66242a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a directed graph\n",
    "G = nx.DiGraph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c3f84c-c441-46f1-988b-7bca8c561a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to create sequential IDs starting from 0 for littleballoffur and DGL\n",
    "file_to_net: Dict[int, int] = {}\n",
    "net_to_file: Dict[int, int] = {}\n",
    "current_idx = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "181ae24c-b54d-4a5a-b286-c5f0288eaab4",
   "metadata": {},
   "source": [
    "### From Text to Graph\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd66174-f00d-448d-be65-b68787ea7adb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decompress the gzip content and build the edge list for our network\n",
    "print(\"Building network structure ...\")\n",
    "\n",
    "# Note we reuse the `gzip_content` variable from the download cell. This is a weird way to do it :)\n",
    "with gzip.GzipFile(fileobj=gzip_content) as f:\n",
    "\n",
    "    # Iterate through the lines, using the `line_number` as an `edge_id` below.\n",
    "    # They won't quite start at 0 owing to comments, but that's ok in the case of edges.\n",
    "    for line_number, line in enumerate(f):\n",
    "        line = line.decode(\"utf-8\")\n",
    "\n",
    "        # Ignore comment lines that start with '#'\n",
    "        if not line.startswith(\"#\"):\n",
    "            # Source (citing), desstination (cited) papers\n",
    "            citing_key, cited_key = line.strip().split(\"\\t\")\n",
    "\n",
    "            # The edge list makes the paper ID an int, stripping 0001001 to 1001, for example\n",
    "            citing_key, cited_key = int(citing_key), int(cited_key)\n",
    "\n",
    "            # If the either of the paper IDs don't exist, make one\n",
    "            for key in [citing_key, cited_key]:\n",
    "                if key not in file_to_net:\n",
    "                    # Build up an index that maps back and forth\n",
    "                    file_to_net[key] = current_idx\n",
    "                    net_to_file[current_idx] = key\n",
    "\n",
    "                    # Bump the current ID\n",
    "                    current_idx += 1\n",
    "\n",
    "            # print(f\"Citing key: {citing_key}, Cited key: {cited_key}\")\n",
    "            # print(f\"Mapped key: {file_to_net[citing_key]}, Mapped key: {file_to_net[cited_key]}\")\n",
    "\n",
    "            G.add_edge(file_to_net[citing_key], file_to_net[cited_key], edge_id=line_number)\n",
    "\n",
    "            # Conditionally set the keys on the nodes\n",
    "            G.nodes[file_to_net[citing_key]][\"file_id\"] = citing_key\n",
    "            G.nodes[file_to_net[citing_key]][\"sequential_id\"] = file_to_net[citing_key]\n",
    "\n",
    "            G.nodes[file_to_net[cited_key]][\"file_id\"] = cited_key\n",
    "            G.nodes[file_to_net[cited_key]][\"sequential_id\"] = file_to_net[cited_key]\n",
    "\n",
    "print(f\"Network built! It contains {G.number_of_nodes():,} nodes and {G.number_of_edges():,} edges.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26805bf9-3b99-40fb-8757-9970a235fa6b",
   "metadata": {},
   "source": [
    "## Persist the `file_to_net` and `net_to_file` Maps\n",
    "\n",
    "We need to persist these indexes as we will load and use them in part 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cabae7b-8402-4677-82b5-cded122ebf26",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(file_to_net.items())[0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b8058b4-02fb-4038-834a-4bf935b9e008",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"data/citation/file_to_net.pkl\", \"wb\") as f:\n",
    "    pickle.dump(file_to_net, f)\n",
    "with open(\"data/citation/file_to_net.json\", \"w\") as f:\n",
    "    json.dump(file_to_net, f)\n",
    "\n",
    "with open(\"data/citation/net_to_file.pkl\", \"wb\") as f:\n",
    "    pickle.dump(net_to_file, f)\n",
    "with open(\"data/citation/net_to_file.json\", \"w\") as f:\n",
    "    json.dump(net_to_file, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39a5d993-5545-4886-966a-e622d2409215",
   "metadata": {},
   "source": [
    "## Node Properties from Abstract Metadata\n",
    "\n",
    "In addition to the edge list, SNAP provides the paper's essential metadata in another file, which we will load to provide node properties and text embeddings for citation graph.\n",
    "\n",
    "We are going to perform the following steps:\n",
    "\n",
    "1) Download and cache the metadata to `data/cit-HepTh-abstracts.tar.gz`. Note that 1 file corresponds to one paper node's metadata.\n",
    "2) Process the tarball file where one file corresponds to one node ID in the original file. See why we made or mappings `file_to_net` and `net_to_file`?\n",
    "3) Assign node properties by parsing the fields of the record using traditional information extraction with regular expressions.\n",
    "4) Use a sentence transformer paraphrase model to summarize the entire textual record and enable node comparison for journal label creation.\n",
    "\n",
    "### Downloading the Abstract Metadata\n",
    "\n",
    "Another file containing node metadata, including the abstracts, for about 90% of nodes in this network is provided at [https://snap.stanford.edu/data/cit-HepTh-abstracts.tar.gz](https://snap.stanford.edu/data/cit-HepTh-abstracts.tar.gz), which we will save to `data/cit-HepTh-abstracts.tar.gz`. This code works just like the edge list download code above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4609227a-1af3-41fd-9921-7715c11d9685",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the abstracts from `cit-HepTh-abstracts.tar.gz`\n",
    "print(\"Fetching paper abstracts ...\")\n",
    "abstract_path = \"data/cit-HepTh-abstracts.tar.gz\"\n",
    "abstract_gzip_content = None\n",
    "\n",
    "if os.path.exists(abstract_path):\n",
    "    print(f\"Using existing paper abstracts file {abstract_path}\")\n",
    "    with open(abstract_path, \"rb\") as f:\n",
    "        abstract_gzip_content = io.BytesIO(f.read())\n",
    "else:\n",
    "    print(\"Downloading paper abbstracts ...\")\n",
    "    abstract_response = requests.get(f\"https://snap.stanford.edu/{abstract_path}\")\n",
    "    abstract_gzip_content = io.BytesIO(abstract_response.content)\n",
    "\n",
    "    print(f\"Downloading abstract file to {abstract_path}\")\n",
    "    with open(abstract_path, \"wb\") as f:\n",
    "        f.write(abstract_response.content)\n",
    "        print(f\"Wrote downloaded abstract file to {abstract_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28ee4892-b011-4b3b-bb18-78633a297a3b",
   "metadata": {},
   "source": [
    "### Manually Parsing Node Metadata\n",
    "\n",
    "As a first pass let's use regular expressions in the Python `re` builtin library to extract each paper's fields so we can assign them as properties to our `nx.DiGraph` nodes.\n",
    "\n",
    "Here is what a couple of **test documents** look like. This is corresponds to two files in our abstract tarball."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce09eff3-3189-45c5-b139-3d346bf97e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = [\n",
    "\"\"\"------------------------------------------------------------------------------\n",
    "\\\\\n",
    "Paper: hep-th/9612115\n",
    "From: Asato Tsuchiya <tsuchiya@theory.kek.jp>\n",
    "Date: Wed, 11 Dec 1996 17:38:56 +0900   (20kb)\n",
    "Date (revised): Tue, 31 Dec 1996 01:06:34 +0900\n",
    "\n",
    "Title: A Large-N Reduced Model as Superstring\n",
    "Authors: N. Ishibashi, H. Kawai, Y. Kitazawa and A. Tsuchiya\n",
    "Comments: 29 pages, Latex, a footnote and references added, eq.(3.52)\n",
    "corrected, minor corrections\n",
    "Report-no: KEK-TH-503, TIT/HEP-357\n",
    "Journal-ref: Nucl.Phys. B498 (1997) 467-491\n",
    "\\\\\n",
    "A matrix model which has the manifest ten-dimensional N=2 super Poincare\n",
    "invariance is proposed. Interactions between BPS-saturated states are analyzed\n",
    "to show that massless spectrum is the same as that of type IIB string theory.\n",
    "It is conjectured that the large-N reduced model of ten-dimensional super\n",
    "Yang-Mills theory can be regarded as a constructive definition of this model\n",
    "and therefore is equivalent to superstring theory.\n",
    "\\\\\n",
    "\"\"\",\n",
    "    \"\"\"------------------------------------------------------------------------------\n",
    "\\\\\n",
    "Paper: hep-th/9711029\n",
    "From: John Schwarz <jhs@theory.caltech.edu>\n",
    "Date: Wed, 5 Nov 1997 17:30:55 GMT   (20kb)\n",
    "Date (revised v2): Thu, 6 Nov 1997 23:52:45 GMT   (21kb)\n",
    "\n",
    "Title: The Status of String Theory\n",
    "Author: John H. Schwarz\n",
    "Comments: 16 pages, latex, two figures; minor corrections, references added\n",
    "Report-no: CALT-68-2140\n",
    "\\\\\n",
    "There have been many remarkable developments in our understanding of\n",
    "superstring theory in the past few years, a period that has been described as\n",
    "``the second superstring revolution.'' Several of them are discussed here. The\n",
    "presentation is intended primarily for the benefit of nonexperts.\n",
    "\\\\\n",
    "\"\"\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44f48a8b-95b3-404e-8658-7c9fd4002acd",
   "metadata": {},
   "source": [
    "### Structured Information Extraction with a Regex Helper\n",
    "\n",
    "Our extract function was created through trial and error using the [Pythex Regex Editor](https://pythex.org/). Paste the test documents where it says `Your test data` and try a couple of the patterns such as `r\"\"` above that where it says `Your regular expression`. It will show you where the patterns match in your data. A new section displays the matches within the text and a window on the right shows the text your regular expression will return via the list the `match.groups()` command returns.\n",
    "\n",
    "A few cycles of this and we have a clean extraction. In practice, I used more test records than this, which I've spared you in the interest of time :) Regular expressions are difficult to learn, but there are resources and **ChatGPT-4 is quite capable at writing regex!** It wrote many of the ones below.\n",
    "\n",
    "<center><img src=\"images/Pythex-Regex-Helper.png\" width=\"1000px\" /></center>\n",
    "\n",
    "I used Pythex to help write the `extract_paper_info(record)` method below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b66d7e1-9e7a-4f51-adc2-7d26538b5dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_paper_info(record):\n",
    "    \"\"\"Extract structured information from the text of academic paper text records using regular expressions.\n",
    "\n",
    "    Note: I was written wholly or in part by ChatGPT4 on May 23, 2023.\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize an empty dictionary to hold the information\n",
    "    info = {}\n",
    "\n",
    "    # Match \"Paper\" field\n",
    "    paper_match = re.search(r\"Paper:\\s*(.*)\", record)\n",
    "    if paper_match:\n",
    "        info[\"Paper\"] = paper_match.group(1)\n",
    "\n",
    "    # # Match \"From\" field\n",
    "    # from_match = re.search(r\"From:\\s*(.*)\", record)\n",
    "    # if from_match:\n",
    "    #     info['From'] = from_match.group(1)\n",
    "\n",
    "    # Match \"From\" field\n",
    "    from_match = re.search(r\"From:\\s*([^<]*)<\", record)\n",
    "    if from_match:\n",
    "        info[\"From\"] = from_match.group(1).strip()\n",
    "\n",
    "    # Match \"Date\" field\n",
    "    date_match = re.search(r\"Date:\\s*(.*)(\\s*)(\\(\\d+kb\\))\", record)\n",
    "    if date_match:\n",
    "        info[\"Date\"] = date_match.group(1).strip()\n",
    "\n",
    "    # Match \"Title\" field\n",
    "    title_match = re.search(r\"Title:\\s*(.*)\", record)\n",
    "    if title_match:\n",
    "        info[\"Title\"] = title_match.group(1)\n",
    "\n",
    "    # Match \"Authors\" field\n",
    "    authors_match = re.search(r\"Authors:\\s*(.*)\", record)\n",
    "    if authors_match:\n",
    "        info[\"Authors\"] = authors_match.group(1)\n",
    "\n",
    "    # Match \"Comments\" field\n",
    "    comments_match = re.search(r\"Comments:\\s*(.*)\", record)\n",
    "    if comments_match:\n",
    "        info[\"Comments\"] = comments_match.group(1)\n",
    "\n",
    "    # Match \"Report-no\" field\n",
    "    report_no_match = re.search(r\"Report-no:\\s*(.*)\", record)\n",
    "    if report_no_match:\n",
    "        info[\"Report-no\"] = report_no_match.group(1)\n",
    "\n",
    "    # Match \"Journal-ref\" field\n",
    "    journal_ref_match = re.search(r\"Journal-ref:\\s*(.*)\", record)\n",
    "    if journal_ref_match:\n",
    "        info[\"Journal-ref\"] = journal_ref_match.group(1)\n",
    "\n",
    "    # Extract \"Abstract\" field\n",
    "    abstract_pattern = r\"Journal-ref:[^\\\\\\\\]*\\\\\\\\[\\n\\s]*(.*?)(?=\\\\\\\\)\"\n",
    "    abstract_match = re.search(abstract_pattern, record, re.DOTALL)\n",
    "    if abstract_match:\n",
    "        abstract = abstract_match.group(1)\n",
    "        abstract = abstract.replace(\"\\n\", \" \").replace(\"  \", \" \")\n",
    "        info[\"Abstract\"] = abstract.strip()\n",
    "\n",
    "    return info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe9f0c15-5080-4875-9c9a-9312ee9e5ab6",
   "metadata": {},
   "source": [
    "### Testing Our Information Extraction\n",
    "\n",
    "To develop the above I created the unit tests below. Inspect the values so you agree they work :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e92f08f8-aa75-44a7-8e9c-b941b7bf0d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "(doc1, doc2) = docs\n",
    "\n",
    "paper_info_one = extract_paper_info(doc1)\n",
    "# Get the paper ID part of the \"Paper\" field\n",
    "paper_id = int(paper_info_one.get(\"Paper\", \"\").split(\"/\")[-1])\n",
    "assert paper_info_one[\"Paper\"] == \"hep-th/9612115\"\n",
    "assert paper_id == 9612115\n",
    "\n",
    "paper_info_two = extract_paper_info(doc2)\n",
    "# Get the paper ID part of the \"Paper\" field\n",
    "paper_id = int(paper_info_two.get(\"Paper\", \"\").split(\"/\")[-1])\n",
    "assert paper_info_two[\"Paper\"] == \"hep-th/9711029\"\n",
    "assert paper_id == 9711029"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5b416ba-a7e3-4a84-826d-f1192665af1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "paper_info_one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4017f2fb-0321-4210-8eb7-68d4a849f55f",
   "metadata": {},
   "outputs": [],
   "source": [
    "paper_info_two"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1786f68-848a-442a-9e96-ccd70debf27a",
   "metadata": {},
   "source": [
    "### Setting Node Properties\n",
    "\n",
    "With our information extraction function tested, we are ready to loop through the abstract metadata tarball's files `G.nodes()` and assign the extract function's fields as node fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc34b911-f8a9-4c70-861a-6c11f7f0e72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use me :)\n",
    "# nx.set_node_attributes(G, feature_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18b0860-2b80-4903-9d5b-1bc7c557d7e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "hit_count, miss_count, matches = 0, 0, 0\n",
    "all_abstracts: List[str] = []\n",
    "abstracts: Dict[int, str] = {}\n",
    "paper_ids: List[int] = []\n",
    "# Decompress the gzip content, then work through the abstract files in the tarball\n",
    "with gzip.GzipFile(fileobj=abstract_gzip_content) as f:\n",
    "    with tarfile.open(fileobj=f, mode=\"r|\") as tar:\n",
    "        for member in tar:\n",
    "            abstract_file = tar.extractfile(member)\n",
    "            if abstract_file:\n",
    "                content = abstract_file.read().decode(\"utf-8\")\n",
    "\n",
    "                paper_id = int(os.path.basename(member.name).split(\".\")[0])\n",
    "\n",
    "                # We can also parse and use those values directly or embed field-wise\n",
    "                paper_info = extract_paper_info(content)\n",
    "                if paper_info:\n",
    "                    abstract_paper_id = paper_info.get(\"Paper\", \"\").split(\"/\")[-1]\n",
    "                    if paper_id != int(abstract_paper_id):\n",
    "                        matches += 1\n",
    "                        print(f\"Paper ID {paper_id} != {abstract_paper_id}\")\n",
    "\n",
    "                    # Get the paper ID part of the \"Paper\" field\n",
    "                    if paper_id in file_to_net and file_to_net[paper_id] in G:\n",
    "                        for field, value in paper_info.items():\n",
    "                            G.nodes[file_to_net[paper_id]][field] = value\n",
    "\n",
    "                        abstracts[paper_id] = content\n",
    "                        all_abstracts.append(content)\n",
    "                        paper_ids.append(paper_id)\n",
    "\n",
    "                        hit_count += 1\n",
    "\n",
    "                    else:\n",
    "                        # Add isolated nodes if paper_id isn't in G\n",
    "                        miss_count += 1\n",
    "                        # We could do this for some use cases to create isolated nodes. Not all graphs are connected. See Part 2, Network Science.\n",
    "                        # G.add_node(file_to_net[paper_id], **paper_info)\n",
    "\n",
    "# Now `G` is a property graph representing the \"High-energy physics theory citation network\" dataset\n",
    "print(f\"Added metadata to {hit_count:,} nodes, {miss_count:,} were unknown.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7767a17-0955-4a8f-a34e-9996fcb54b90",
   "metadata": {},
   "source": [
    "## Temporal Networks\n",
    "\n",
    "Our citation graph is a temporal network. Temporal networks have a determined sequence in which nodes are added to the graph - just as real networks evolve. Network science and graph machine learning for temporal networks that don't take time into account can result in incorrect analyses and inaccurate machine learning inference.\n",
    "\n",
    "<center><img src=\"images/Temporal-vs-static-networks-A-The-sequence-of-contacts-among-three-nodes-capturing_W640.jpg\" width=\"800px\" /></center>\n",
    "<center>Image source: <a href=\"https://arxiv.org/abs/1607.06168\">The fundamental advantages of temporal networks, Li et al., 2016</a></center>\n",
    "\n",
    "### Downloading Publishing Dates\n",
    "\n",
    "The timestamps are available at [https://snap.stanford.edu/data/cit-HepTh-dates.txt.gz](https://snap.stanford.edu/data/cit-HepTh-dates.txt.gz) and are downloaded to `data/cit-HepTh-dates.txt.gz`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d35f49a5-d300-4c10-8cef-ccbd664257fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download and load edges (citations) from `cit-HepTh.txt.gz`\n",
    "dates_path = \"data/cit-HepTh-dates.txt.gz\"\n",
    "date_gzip_content = None\n",
    "\n",
    "if os.path.exists(dates_path):\n",
    "    print(f\"Using existing paper dates file {dates_path}\")\n",
    "    date_gzip_content = open(dates_path, \"rb\")\n",
    "else:\n",
    "    print(\"Downloading paper publishing dates ...\")\n",
    "    date_response = requests.get(f\"https://snap.stanford.edu/{dates_path}\")\n",
    "    date_gzip_content = io.BytesIO(date_response.content)\n",
    "\n",
    "    with open(dates_path, \"wb\") as f:\n",
    "        f.write(date_response.content)\n",
    "        print(\"Wrote downloaded publishing dates file to {dates_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4dc8390-c714-4c75-b702-5f3c243e1d0d",
   "metadata": {},
   "source": [
    "### Adding a Temporal Property\n",
    "\n",
    "Publishing dates go under the `Published` node property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acaca8e5-e0c1-4fbc-ba9a-a514218acccf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decompress the gzip content and add a \"published\" date property to our nodes\n",
    "print(\"Adding publising dates ...\")\n",
    "hit_count = 0\n",
    "with gzip.GzipFile(fileobj=date_gzip_content) as f:\n",
    "    for line in f:\n",
    "        line = line.decode(\"utf-8\")\n",
    "        # Ignore lines that start with '#'\n",
    "        if not line.startswith(\"#\"):\n",
    "            paper_id, iso_date = line.strip().split(\"\\t\")\n",
    "\n",
    "            # The edge list makes the paper ID an int, stripping 0001001 to 1001, for example\n",
    "            paper_id = int(paper_id)\n",
    "\n",
    "            if paper_id in file_to_net and file_to_net[paper_id] in G:\n",
    "                # Add a UTC timestamp for the data\n",
    "                G.nodes[file_to_net[paper_id]][\"Published\"] = calendar.timegm(\n",
    "                    date.fromisoformat(iso_date).timetuple()\n",
    "                )\n",
    "                hit_count += 1\n",
    "\n",
    "print(f\"Added Published dates to {hit_count:,} nodes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72a45ded-092d-49e7-bcbd-55018997b7e7",
   "metadata": {},
   "source": [
    "### Test our Network Build\n",
    "\n",
    "Let's make sure everything built as expected..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a6d03a-4bbe-4997-b3fd-6ee1f728c927",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What is the first node?\n",
    "test_node = G.nodes[0]\n",
    "\n",
    "assert test_node[\"sequential_id\"] == 0\n",
    "assert test_node[\"file_id\"] == 1001\n",
    "\n",
    "print(json.dumps(test_node, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5228d39-a3d1-40ef-ab76-f1bf99b79c37",
   "metadata": {},
   "source": [
    "## Abstract Embeddings\n",
    "\n",
    "In addition to parsing the data, we will embed the entire record. First we create a utility to embed string columns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0043881c-c4d7-4357-b284-7e65ffb2063f",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_model = SentenceTransformer(\"sentence-transformers/all-MiniLM-L6-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa37dfea-c8da-429d-83df-011c3fbc755e",
   "metadata": {},
   "outputs": [],
   "source": [
    "paraphrase_model = SentenceTransformer(\"sentence-transformers/paraphrase-MiniLM-L6-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01bb1116-6b21-45f7-94f2-49fe6af97658",
   "metadata": {},
   "outputs": [],
   "source": [
    "def embed_paper_info(\n",
    "    records: Union[str, List[str]], convert_to_tensor=True\n",
    ") -> Union[np.ndarray, Tensor]:\n",
    "    if records and isinstance(records, str):\n",
    "        records = [records]\n",
    "    return paraphrase_model.encode(records, convert_to_tensor=convert_to_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a38e4d-40f6-4484-be82-743f496cfdf8",
   "metadata": {},
   "source": [
    "Then we embed all the nodes' metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0734266-b939-4438-8711-f675bfa93479",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed the abstracts for GNN features. Embedding is a generic approach for retrieval as well.\n",
    "# Note: NetworkX can't save lists in GEXF format, so we'll JSONize the list & save the embeddings separately.\n",
    "embedded_abstracts: np.ndarray = None\n",
    "if os.path.exists(\"data/embedded_abstracts.npy\"):\n",
    "    embedded_abstracts = np.load(\"data/embedded_abstracts.npy\")\n",
    "else:\n",
    "    embedded_abstracts = embed_paper_info(all_abstracts, convert_to_tensor=False)\n",
    "    np.save(\"data/embedded_abstracts.npy\", embedded_abstracts)\n",
    "\n",
    "node_embedding_dict: Dict[int, List[float]] = {}\n",
    "if os.path.exists(\"data/node_embedding_dict.json.gz\"):\n",
    "    node_embedding_dict = json.load(\n",
    "        gzip.GzipFile(\"data/node_embedding_dict.json.gz\", \"r\"),\n",
    "        # encoding=\"utf-8\",\n",
    "    )\n",
    "else:\n",
    "    for paper_id, emb in zip(paper_ids, embedded_abstracts):\n",
    "        assert emb.shape == (384,)\n",
    "\n",
    "        # Gephi assumes a list of floats is a time series, so we need to convert to a string\n",
    "        emb_list = emb.tolist()\n",
    "        G.nodes[file_to_net[paper_id]][\"Embedding-JSON\"] = json.dumps(emb_list)\n",
    "\n",
    "        node_embedding_dict[paper_id] = emb_list\n",
    "\n",
    "# Write the mapping from paper ID to embedding to JSON.\n",
    "# Note: All JSON keys are strings. We will have to int(key) to read the data back.\n",
    "json.dump(\n",
    "    node_embedding_dict,\n",
    "    io.TextIOWrapper(\n",
    "        gzip.GzipFile(\"data/node_embedding_dict.json.gz\", \"w\"),\n",
    "        encoding=\"utf-8\",\n",
    "    ),\n",
    "    indent=4,\n",
    "    sort_keys=True,\n",
    ")\n",
    "\n",
    "# Write the entire network using GEXF format - the date has to be in UTC format for this to work.\n",
    "nx.write_gexf(G, path=\"data/physics_embeddings.gexf.gz\", prettyprint=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6a32b3-319e-43ac-ab12-a4d906018c2d",
   "metadata": {},
   "source": [
    "# Label Making and K-Nearest-Neighbors (KNN) Graph Building\n",
    "\n",
    "Now we are going to build a pandas `DataFrame` or `pd.DataFrame` of our nodes so we can create clean labels for our journals. These will serve as labels for our machine-learning tasks using this network.\n",
    "\n",
    "In this section, we will also demonstrate another method of building a network - K-nearest-neighbors construction.\n",
    "\n",
    "## Building a Node `pd.DataFrame`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65ea0cd4-6b8d-40b6-a3d9-af30157e999b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract nodes and their attributes into a list of dictionaries\n",
    "node_data = [{**{\"node\": node}, **attr} for node, attr in G.nodes(data=True)]\n",
    "\n",
    "# Convert the list of dictionaries into a DataFrame\n",
    "node_df = pd.DataFrame(node_data)\n",
    "\n",
    "# Cleanup\n",
    "node_df.fillna(\"\", inplace=True)\n",
    "\n",
    "node_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e83edc4-87f5-483b-874b-e4715569dbb0",
   "metadata": {},
   "source": [
    "## Clustering `Journal-ref`\n",
    "\n",
    "Taking a look at the field representing the journal a paper appeared in, we have a problem if we want to use this field as a label for a categorical classification... this is a fuzzy string problem, not a regular expression problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d92b9e53-2e5b-4896-998f-5661722acb21",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_df[\"Journal-ref\"].sample(n=10).head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "253ab1d4-98e8-458e-80eb-1226342d7a6d",
   "metadata": {},
   "source": [
    "### Embedding `Title`, `Abstract` and `Journal-ref`\n",
    "\n",
    "That's ok! We will embed the `Journal-ref` field and cluster it to arrive at our class labels for each journal. This will give us a head start on presenting network construction using K-Nearest-Neighbors in the next section :)\n",
    "\n",
    "#### Simplify `Journal-ref` for Clustering\n",
    "\n",
    "What if we remove all the dates entirely? This should help us find titles :) While I am demonstrating vector operations, we might just use a Bag-of-Words model :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7d78186-7b6a-47f3-b87f-3a98918b8d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove all non-text characters from Journal-ref to make it cluster better\n",
    "node_df[\"Journal-ref-Letters\"] = node_df[\"Journal-ref\"].str.replace(r\"[^a-zA-Z.]\", \"\", regex=True).str.strip().str[:-1]\n",
    "node_df[\"Journal-ref-Letters\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38c5e31a-855a-4b3c-af6c-1b570eb2c799",
   "metadata": {},
   "outputs": [],
   "source": [
    "def embed_cache_column(input_col: str, model_name: str, model: SentenceTransformer, df: pd.DataFrame) -> np.ndarray:\n",
    "    \"\"\"Given an input column, model name, SentenceTransformer model and a pd.DataFrame, produce an embedding.\"\"\"\n",
    "\n",
    "    base_path = \"data/embedded_journals\"\n",
    "    file_name = f\"{input_col}-{model_name}-Embedding.npy\"\n",
    "    file_path = os.path.join(base_path, file_name)\n",
    "    \n",
    "    embedding: np.ndarray = None\n",
    "    if os.path.exists(file_path):\n",
    "        embedding = np.load(file_path)\n",
    "        print(f\"Cachce hit for {file_path}\")\n",
    "    else:\n",
    "        print(f\"Cache miss for {file_path}, saved\")\n",
    "        embedding = model.encode(df[column].tolist())\n",
    "        np.save(file_path, embedding)\n",
    "        print(f\"Embedded {column} using {model_name} model\")\n",
    "\n",
    "    return embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45bf4e54-40a5-4dc9-b77b-86e4ff42d17a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1. Embed the relevant text columns Journal-ref and cluster them to produce journal class labels.\n",
    "paraphrase_model = SentenceTransformer(\"sentence-transformers/paraphrase-MiniLM-L6-v2\")\n",
    "all_model = SentenceTransformer(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "\n",
    "# Embed these columns\n",
    "paraphrase_embeddings: Dict[str, np.ndarray] = {}\n",
    "all_embeddings: Dict[str, np.ndarray] = {}\n",
    "\n",
    "# Generate embeddings, but do cache\n",
    "for column in [\"Title\", \"Abstract\", \"Journal-ref\", \"Journal-ref-Letters\"]:\n",
    "\n",
    "    # Paraphrase model\n",
    "    paraphrase_embeddings[column] = embed_cache_column(column, \"Paraphrase\", paraphrase_model, node_df)\n",
    "    node_df[f\"{column}-Paraphrase-Embedding\"] = paraphrase_embeddings[column].tolist()\n",
    "\n",
    "    # All model\n",
    "    all_embeddings[column] = embed_cache_column(column, \"All\", all_model, node_df)\n",
    "    node_df[f\"{column}-All-Embedding\"] = all_embeddings[column].tolist()   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27846474-5bb9-4617-b223-d69074c81673",
   "metadata": {},
   "source": [
    "#### Assign Embeddings back to `node_df` `pd.DataFrame`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52417cab-bb35-4667-90ff-4d605c920025",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's go with Journal-ref-Letters-Embedding below, less variance\n",
    "node_df[[\"Journal-ref\", \"Journal-ref-Letters\", \"Journal-ref-Letters-All-Embedding\", \"Journal-ref-Letters-Paraphrase-Embedding\"]].sample(30).head(30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39622cf8-0814-43fe-86a6-b1680754610d",
   "metadata": {},
   "source": [
    "### Clustering `Journal-ref-Letters-Embedding`\n",
    "\n",
    "Now we will use [UMAP](https://umap-learn.readthedocs.io/en/latest/) via the [umap-learn](https://umap-learn.readthedocs.io/en/latest/basic_usage.html) PyPi library to reduce our data to 2 dimensions and then [DBSCAN](https://en.wikipedia.org/wiki/DBSCAN) [which workds with 2D data and infers circles aroud centrioids] via [scikit-learn]() to cluster it into journal names. Our final step is to nominate a real journal name for each cluster, which we will do manually as this course is about graph ML and this is NLP :)\n",
    "\n",
    "Note that while normally you need to [standardize data](https://scikit-learn.org/stable/auto_examples/preprocessing/plot_all_scaling.html#standardscaler) to give it a normal distribution with something like [sklearn.preprocessing.StandardScaler](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.StandardScaler.html), you do not need to do this when sentence encoding data with a sentence transformer as we have done.\n",
    "\n",
    "#### KMeans - Did Not Work :)\n",
    "\n",
    "KMeans didn't work, which is probably because I didn't know how to tune it effectively, but DBSCAN is fairly automatic and produced good clusters, so I used that. If you think about it though, KMeans starts with random centroids... and that is why it is splitting journals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0920893c-b620-4bac-bffa-32c419d1b880",
   "metadata": {},
   "outputs": [],
   "source": [
    "km = KMeans(\n",
    "    n_clusters=15,\n",
    "    init=\"k-means++\",\n",
    "    max_iter=200,\n",
    ")\n",
    "\n",
    "# Apply to both embeddings\n",
    "paraphrase_class_scores = km.fit_transform(node_df[\"Journal-ref-Letters-Paraphrase-Embedding\"].tolist())\n",
    "paraphrase_classes = np.argmax(paraphrase_class_scores, axis=1)\n",
    "node_df[\"Journal-ref-Letters-Paraphrase-KMeans\"] = paraphrase_classes.tolist()\n",
    "\n",
    "all_class_scores = km.fit_transform(node_df[\"Journal-ref-Letters-All-Embedding\"].tolist())\n",
    "all_classes = np.argmax(all_class_scores, axis=1)\n",
    "node_df[\"Journal-ref-Letters-All-KMeans\"] = all_classes.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a5a83c5-dbbe-4ccf-b81d-07a57e33d5a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.unique(paraphrase_classes, return_counts=True), paraphrase_classes.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b814ead0-4090-420a-8956-24d9b9bad9f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.unique(all_classes, return_counts=True), all_classes.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a52bee-c522-4553-9529-e36c91fc3c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    node_df[[\"Journal-ref\", \"Journal-ref-Letters\", \"Journal-ref-Letters-Paraphrase-KMeans\", \"Journal-ref-Letters-All-KMeans\"]]\n",
    "    .sample(30)\n",
    "    .sort_values(by=\"Journal-ref-Letters-Paraphrase-KMeans\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b1b937d-67e5-458f-a8d0-079d941d6159",
   "metadata": {},
   "source": [
    "#### KMeans Did Not Work 2.0\n",
    "\n",
    "Although a nice way to begin, KMeans didn't work.\n",
    "\n",
    "**Q: Do any of you have more experience with KMeans than I do that can make it work to produce good journal clusters? :)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ac9ad4a-14b1-4e8c-bb06-309b309f8b2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Dimension Reduction with UMAP\n",
    "reducer = umap.UMAP()\n",
    "reducer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d769fd55-8f03-4f09-a539-c277e2dc0eff",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduced_paraphrase_embeddings = reducer.fit_transform(node_df[\"Journal-ref-Letters-Paraphrase-Embedding\"].tolist())\n",
    "reduced_paraphrase_embeddings, reduced_paraphrase_embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e2d63b-b61b-4e28-98a1-cb920f01dfb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduced_all_embeddings = reducer.fit_transform(node_df[\"Journal-ref-Letters-All-Embedding\"].tolist())\n",
    "reduced_all_embeddings, reduced_all_embeddings.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "105337a3-7a0a-4c2d-a11e-2e5ec61b93a4",
   "metadata": {},
   "source": [
    "#### Clustering with DBSCAN\n",
    "\n",
    "ChatGPT and I say:\n",
    "\n",
    "> DBSCAN groups together closely packed 2-dimensional data points based on a specified distance measure and minimum number of points. It starts with an arbitrary point, expands clusters from suitable points, and identifies noise points that don't belong to any cluster. This makes it good for finding arbitrary shaped clusters and outliers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dca51b0b-2fcb-4aab-b6eb-c8b610b01e15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Clustering with DBSCAN - you can search for the best hyperparameters\n",
    "dbscan = DBSCAN(eps=0.9, min_samples=500)#, metric=\"l1\")\n",
    "dbscan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80dc8a60-a110-45d1-a87a-021dc5137a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "paraphrase_clusters = dbscan.fit_predict(reduced_paraphrase_embeddings)\n",
    "node_df[\"Journal-ref-Letters-Paraphrase-DBSCAN\"] = paraphrase_clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19d9d243-8b51-4c6c-aa78-bcd773843a25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count the clusters and their sizes\n",
    "np.unique(paraphrase_clusters, return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9cf21f9-8ea6-41a2-90db-51a5cc5ff9c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_clusters = dbscan.fit_predict(reduced_all_embeddings)\n",
    "node_df[\"Journal-ref-Letters-All-DBSCAN\"] = all_clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01088c9b-c7ea-481d-ab01-69c95a11d75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count the clusters and their sizes\n",
    "np.unique(all_clusters, return_counts=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c7bba40-1f37-40c0-8dde-9387c28ed274",
   "metadata": {},
   "source": [
    "#### Viewing DBSCAN Clusters\n",
    "\n",
    "Let's take a look comparing our two embeddings' clusters... first we need to drop the many records without Journal labels. Those make for bad positive labels... but could be a good source of negative labels, by which we could avoid negative sampling!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d935182e-e1d4-42fc-892a-f0b3508f0619",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_eval_df = node_df[[\n",
    "    \"Journal-ref\",\n",
    "    \"Journal-ref-Letters\",\n",
    "    \"Journal-ref-Letters-Paraphrase-DBSCAN\",\n",
    "    \"Journal-ref-Letters-All-DBSCAN\"\n",
    "]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35aaee4c-01a0-41e5-aafd-109de9824c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace all-whitespace strings with empty strings\n",
    "cluster_eval_df.replace(r\"^\\s*$\", \"\", regex=True, inplace=True)\n",
    "\n",
    "# Identify rows where any field is an empty string\n",
    "empty_rows = cluster_eval_df[cluster_eval_df.eq(\"\").any(axis=1)].index\n",
    "\n",
    "# Drop those rows\n",
    "cluster_eval_df.drop(empty_rows, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34a2cab6-e936-4716-a3d3-bc3ceb093043",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    cluster_eval_df\n",
    "    .sample(60)\n",
    "    .dropna(axis=0)\n",
    "    .sort_values(by=\"Journal-ref-Letters-All-DBSCAN\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25db28a3-a002-48c9-b073-037738379e70",
   "metadata": {},
   "source": [
    "#### Judging DBSCAN Clusters...\n",
    "\n",
    "These clusters look pretty good! That was a lot of trouble, though.\n",
    "\n",
    "**Q: Do any of you know how to create node labels for journals quicker than that? What about a bag of words model?**\n",
    "\n",
    "### Nominating Labels for our Clusters\n",
    "\n",
    "Before we even evaluate our clusters by displaying them alongside our `Journal-ref` and `Journal-ref-Letters` in the `node_df` `pd.DataFrame`, I want to perform a trick for labeling clusters using a `GROUP BY`, [MapReduce](https://en.wikipedia.org/wiki/MapReduce) or [split-apply-combine strategy] as you prefer :)\n",
    "\n",
    "We will group by the cluster ID, take [pd.Series.value_counts](https://pandas.pydata.org/docs/reference/api/pandas.Series.value_counts.html) [a `pd.Series` is a `pd.DataFrame` column] and assign the top value count as the label for the cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b97be23-179f-4283-b858-ee93ad15af4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign our canonical clusters to the DataFrame\n",
    "node_df[\"Journal-ref-DBSCAN\"] = node_df[\"Journal-ref-Letters-All-DBSCAN\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "959ac38c-9830-44a8-a3df-73fcffd609d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame(node_df[[\"Journal-ref\", \"Journal-ref-Letters\", \"Journal-ref-DBSCAN\"]].groupby(\"Journal-ref-DBSCAN\").value_counts())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35832642-1143-4251-9366-bf3503c08e2b",
   "metadata": {},
   "source": [
    "#### ChatGPT and Me :)\n",
    "\n",
    "I could not get the above code to preview all the clusters... so I enlisted help.\n",
    "\n",
    "Note: this code was written by ChatGPT-4 after 4 attempts. It shows that the `Journal-ref-Count`s are too low to nominate a good name. This led me to allow `.` and ` ` [spaces] in the find/replace above that looked like:\n",
    "\n",
    "```python\n",
    "node_df[\"Journal-ref-Letters\"] = node_df[\"Journal-ref\"].str.replace(r\"[^a-zA-Z]\", \"\", regex=True).str.strip()\n",
    "```\n",
    "\n",
    "The `r\"[^a-zA-Z]\"` now looks like `r\"[^a-zA-Z. ]\"` to include periods and spaces:\n",
    "\n",
    "```python\n",
    "node_df[\"Journal-ref-Letters\"] = node_df[\"Journal-ref\"].str.replace(r\"[^a-zA-Z. ]\", \"\", regex=True).str.strip()\n",
    "```\n",
    "\n",
    "And then things looked good!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed603631-3eed-4e1c-b6fc-3ab284dfb6ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_top_n_values(group, n=5):\n",
    "    result = {}\n",
    "    for column in [\"Journal-ref\", \"Journal-ref-Letters\"]:\n",
    "        top_values = group[column].value_counts().nlargest(n)\n",
    "        result[f\"{column}-Value\"] = top_values.index.tolist()\n",
    "        result[f\"{column}-Count\"] = top_values.values.tolist()\n",
    "    return pd.Series(result)\n",
    "\n",
    "# Grouping by 'Journal-ref-DBSCAN' and applying the function\n",
    "top_values_df = node_df.groupby(\"Journal-ref-DBSCAN\").apply(get_top_n_values).reset_index()\n",
    "\n",
    "top_values_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d527bed-0587-413f-8fa5-dcd1b7f52ee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_values_df['Journal-ref-Letters-Cluster'] = top_values_df['Journal-ref-Letters-Value'].str.get(0)\n",
    "\n",
    "journal_label_df = top_values_df[[\"Journal-ref-DBSCAN\", \"Journal-ref-Letters-Cluster\"]]\n",
    "\n",
    "journal_label_df.rename(\n",
    "    columns={\"Journal-ref-DBSCAN\": \"Journal-ref-DBSCAN-2\"},\n",
    "    inplace=True,\n",
    ")\n",
    "\n",
    "journal_label_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a6ab0fe-1b04-4098-9bdf-5858d7d8dbdf",
   "metadata": {},
   "source": [
    "### Plot our Clusters\n",
    "\n",
    "Hey, it took me a few tries but the clusters look good! So do our labels. Woohoo! Now we can classify our nodes. Sometimes node labels aren't cleanly available, and you have to work to get training data for graph machine learning. In Part 4 of this class we will use these labels with [PyG](https://pyg.org) to classify nodes with missing journal entries into the journals they belong to. Fun! :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128102e0-1194-45e0-b7fd-3b1b16046b33",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data = pd.DataFrame(data=reduced_all_embeddings, columns=['UMAP 1', 'UMAP 2'])\n",
    "plot_data['Cluster ID'] = all_clusters  # Add cluster labels to the DataFrame\n",
    "\n",
    "# Step 2: Plot the data\n",
    "# Initialize the plot\n",
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "# Use seaborn's scatterplot function to plot UMAP dimensions,\n",
    "# coloring the points by their cluster ID.\n",
    "# The 'palette' argument specifies the colors to use for each cluster.\n",
    "sns.scatterplot(\n",
    "    x='UMAP 1',  # X-axis: first dimension from UMAP\n",
    "    y='UMAP 2',  # Y-axis: second dimension from UMAP\n",
    "    hue='Cluster ID',  # Color by cluster ID\n",
    "    palette=sns.color_palette(\"hsv\", len(plot_data['Cluster ID'].unique())),  # Use a color palette with enough colors\n",
    "    data=plot_data,  # Data source\n",
    "    legend=\"full\",  # Display a legend\n",
    "    alpha=0.5  # Make points semi-transparent to see overlapping points\n",
    ")\n",
    "\n",
    "# Display the plot\n",
    "plt.title('Clusters in 2D UMAP space')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a478083-5ebc-4727-852d-6cbad06c5d3d",
   "metadata": {},
   "source": [
    "### Assign Class Labels Back to the DataFrame\n",
    "\n",
    "Our last step is to assign both a numeric class number and our nominated raw `Journal-ref` label to represent our clusters in the original `node_df` node list `pd.DataFrame`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e11a1a19-f2cf-48d8-b2d1-b4b80ba39150",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_df = node_df.merge(\n",
    "    journal_label_df, \n",
    "    left_on='Journal-ref-DBSCAN', \n",
    "    right_on='Journal-ref-DBSCAN-2', \n",
    "    how='left'\n",
    ")\n",
    "\n",
    "# Rename the merged column to the desired name\n",
    "node_df.rename(\n",
    "    columns={\n",
    "        \"Journal-ref-Letters-Cluster\": 'Journal-ref-Label'},\n",
    "    inplace=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7330d55-42e6-4246-99fa-56fb157f91f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_df[[\"Title\", \"Journal-ref\", \"Journal-ref-Label\", \"Journal-ref-DBSCAN\"]].head(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "281a1b55-101f-4ccc-b6b0-f3323d3f10c7",
   "metadata": {},
   "source": [
    "### Save `node_df` to Parquet Format\n",
    "\n",
    "We want to use this `pd.DataFrame` for new and improved node features for use to learn node classification, so we will save it in Parquet format. I wrote a popular [guide to Python and Parquet format](https://blog.datasyndrome.com/python-and-parquet-performance-e71da65269ce) you may want to check out :)\n",
    "\n",
    "#### Fix a Small Problem with Empty Strings `\"\"` in the `Published` Column\n",
    "\n",
    "I got a PyArrow error when I tried to save, so I am filling [imputing] that column with 0s.\n",
    "\n",
    "```python\n",
    "ArrowInvalid: (\"Could not convert '' with type str: tried to convert to double\", 'Conversion failed for column Published with type object')\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fc8d08e-67c0-4afb-bf51-493e2463efe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_df[\"Published\"] = node_df[\"Published\"].replace(\"\", \"0\").astype(int)\n",
    "node_df[\"Published\"].head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d1c6873-66bf-4429-b33d-0a9172b54e34",
   "metadata": {},
   "source": [
    "#### Lots of Zeros\n",
    "\n",
    "Hmmmm how many zeros? Enough we may need to find another temporal network or look again at how we parse dates above. Not now :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5906f8d4-4ac2-4426-8111-08878c23018e",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_df[\"Published\"].where(node_df[\"Published\"] == 0).count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "048fb6a3-cd35-4723-912e-b667270e2c35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure we didn't drop any nodes...\n",
    "len(node_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd330fd1-3513-4efe-b210-8a3df9af4aae",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_df = node_df[[\n",
    "    \"node\",\n",
    "    \"file_id\",\n",
    "    \"sequential_id\",\n",
    "    \"Paper\",\n",
    "    \"From\",\n",
    "    \"Date\",\n",
    "    \"Title\",\n",
    "    \"Authors\",\n",
    "    \"Comments\",\n",
    "    \"Report-no\",\n",
    "    \"Published\",\n",
    "    \"Journal-ref\",\n",
    "    \"Journal-ref-Label\",\n",
    "    \"Journal-ref-DBSCAN\",\n",
    "    \"Abstract\",\n",
    "]].fillna(\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9021933-6f40-4944-a8fd-97acf2726247",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_df.to_parquet(\"data/cit-HepTh-df.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3abbb5d2-1331-4985-b011-5c3127a4d776",
   "metadata": {},
   "source": [
    "### Adding Journal Label to our NetworkX `nx.DiGraph`\n",
    "\n",
    "Let's add the journal labels back to our `G` `nx.DiGraph` and save it again in [GEXF format](https://gexf.net/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b25cee0-4e69-432e-bb57-cdb0f77be31d",
   "metadata": {},
   "outputs": [],
   "source": [
    "networkx_df = output_df[[\"node\", \"Journal-ref-DBSCAN\", \"Journal-ref-Label\"]]\n",
    "\n",
    "networkx_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5aa63eb-a60e-40bc-bec0-64903250faf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the DataFrame to a dictionary where the key is the 'node' column\n",
    "networkx_dict = networkx_df.set_index('node').T.to_dict('dict')\n",
    "\n",
    "# Now iterate through this dictionary and update the attributes of your nodes\n",
    "for node, attributes in networkx_dict.items():\n",
    "    G.nodes[node].update(attributes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca082263-d553-42dd-beec-f1bfe21f1f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "G.nodes[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62e795e2-3c2a-4050-a286-33a095df52b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This cell saves the network for Part 2\n",
    "nx.write_gexf(G, path=\"data/physics_labeled.gexf.gz\", prettyprint=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "234b6250-5cc4-47d3-ae3d-302b8a7e4e09",
   "metadata": {},
   "source": [
    "### Wrapping up NetworkX Network Construction\n",
    "\n",
    "And that completes building our High Energy Physics `nx.DiGraph` :) We can now perform citation link prediction and journal prediction for papers without journal metadata. Cool! :)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b81e7a54-75ee-4efc-a117-c72ca39b0435",
   "metadata": {},
   "source": [
    "## Building K-Nearest-Neighbor Networks\n",
    "\n",
    "Next up... we are going to use the sentence encoded abstracts in `node_df[\"Abstracts-Paraphrase-Embedding\"]` to create a KNN Network.\n",
    "\n",
    "### Remove Null Abstracts\n",
    "\n",
    "First remove all the null string `Abstract` rows, or we will get lots of points with 0 distance from lots of others - poor edges indeed :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ee46dd-a0c0-4246-ab89-675494665237",
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_df = node_df.copy()[[\"node\", \"Title\", \"Authors\", \"Published\", \"Abstract\", \"Journal-ref\", \"Journal-ref-DBSCAN\", \"Journal-ref-Label\", \"Abstract-Paraphrase-Embedding\"]]\n",
    "viz_df = viz_df[viz_df[\"Abstract\"].str.strip() != \"\"]\n",
    "\n",
    "len(viz_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9417097e-70ee-471b-bc04-3162744ca71e",
   "metadata": {},
   "source": [
    "### Create a Zero-Index to Refer to Below\n",
    "\n",
    "Now that we have filtered our `viz_df` `pd.DataFrame` to filer out some rows... we need a new `zero_index`. We will need to convert `viz_df` to a dict to reference it when building our `KNN_G` `nx.Graph` using `enumerate()` so let's create a column now with a `zero_index`..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e042ab5-b2bd-4ea8-bae2-ae1955ef1b03",
   "metadata": {},
   "outputs": [],
   "source": [
    "enumerated_index = range(0, len(viz_df.index))\n",
    "list(enumerated_index)[0:10]\n",
    "viz_df[\"zero_index\"] = enumerated_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb0a48f-99c4-4b24-b184-53d07f25d0cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "526edad5-1ce6-4d33-a8f9-bb05d37c1389",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neighbors import NearestNeighbors\n",
    "\n",
    "\n",
    "# Determines the edges per node - how far out to match in similarity of the embedding vector space\n",
    "RADIUS = 3\n",
    "\n",
    "# Convert the embedding column to a 2D numpy array\n",
    "embeddings = np.vstack(viz_df[\"Abstract-Paraphrase-Embedding\"])\n",
    "\n",
    "# Step 1: Compute nearest neighbors\n",
    "knn = NearestNeighbors(radius=RADIUS, algorithm='auto')\n",
    "knn.fit(embeddings)\n",
    "\n",
    "# Use radius_neighbors instead of kneighbors\n",
    "distances, indices = knn.radius_neighbors(embeddings, sort_results=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1131473-7253-4a55-a58c-6e77c5c93a29",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23751c71-dea7-4ac4-a160-49cc923e80d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_dict = viz_df.drop(columns=[\"Abstract-Paraphrase-Embedding\"], axis=1).set_index(\"zero_index\", drop=False).to_dict(orient=\"index\")\n",
    "viz_dict[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d504a29-6bf4-4c06-bbc8-302a45c06ddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "viz_df.iloc[indices[0]][[\"node\", \"Title\", \"Abstract\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71ccee65-81f7-4a87-bc5e-9f6f2435d60c",
   "metadata": {},
   "outputs": [],
   "source": [
    "distances, distances.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da5d8d65-3896-4636-b8c5-dbdb50323269",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices, indices.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7861338-a06b-490c-9749-282adb2783e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 2: Create a graph from nearest neighbor relationships\n",
    "KNN_G = nx.Graph()\n",
    "i = 0\n",
    "while i < 10:\n",
    "    for i, neighbors in enumerate(indices):\n",
    "        for j, neighbor in enumerate(neighbors[1:]):  # Skip the point itself (neighbors[0] is always the point itself)\n",
    "            \n",
    "            # Add the node properties if we are adding an edge we haven't seen before\n",
    "            if not i in KNN_G.nodes():\n",
    "                KNN_G.add_node(i, **viz_dict[i])\n",
    "            if not neighbor in KNN_G.nodes():\n",
    "                KNN_G.add_node(neighbor, **viz_dict[neighbor])\n",
    "            \n",
    "            KNN_G.add_edge(i, neighbor, weight=distances[i][j])  # Add an edge between the point and each neighbor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5604516-a6be-4b2d-bb21-f912b854a6d5",
   "metadata": {},
   "source": [
    "### Verify that Edges Have Weights\n",
    "\n",
    "We set weights using `i` and `j` because the `indices` `np.ndarray` is the same size and structure of the corresponding `distances` `np.ndarray`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9568aed1-fea5-4a22-9401-b8b9c607d816",
   "metadata": {},
   "outputs": [],
   "source": [
    "(u, v) = list(KNN_G.edges())[0]\n",
    "\n",
    "KNN_G.get_edge_data(u, v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54155cd1-511b-4b23-a276-546d34760b78",
   "metadata": {},
   "source": [
    "### Check One Node's Properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d590b2-db78-485d-9611-c4a338531c1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I had to try 0-3 to find a real one. Look, the zero_index is correct.\n",
    "print(json.dumps(KNN_G.nodes[3], indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cea4883-9a7a-443a-8aa8-1e678feba8cc",
   "metadata": {},
   "source": [
    "### How Big is Our Network?\n",
    "\n",
    "I had to play with this to get useful structure out of it. Click arond in the rendered network below and see if you agree it is useful :) Remember that it might work well for your data - high energy physics papers are a complex topic that a domain-specific model like [sentence-transformers/paraphrase-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/paraphrase-MiniLM-L6-v2) may not work well with... I just got it working man, I didn't tune it. In practice you have to tweek these things but this should get you started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5fa345c-3438-4994-b0a0-6351bf9c5583",
   "metadata": {},
   "outputs": [],
   "source": [
    "# How big is our network?\n",
    "KNN_G.number_of_nodes(), KNN_G.number_of_edges()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd103be8-6c12-4332-8ee6-c4c599afbfa5",
   "metadata": {},
   "source": [
    "Now KNN_G is a graph where each node represents a row in df, and edges connect each node to its k-nearest neighbors.\n",
    "\n",
    "### Visualizing KNN Networks in Graphistry\n",
    "\n",
    "Now we will use Graphistry to visualize this network, checking if the edges make sense by comparing nodes' `Title`s and `Abstract`s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6501a5f9-4b3c-4808-88f4-afe5a7d7e9c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = (\n",
    "    graphistry.bind(\n",
    "        # Always for networkx.Graph/DiGraph\n",
    "        source=\"src\",\n",
    "        destination=\"dst\",\n",
    "        node=\"nodeid\",\n",
    "        # Above always for networkx.Graph/DiGraph\n",
    "        point_title=\"Title\",\n",
    "        point_label=\"Title\",\n",
    "        edge_weight=\"weight\",\n",
    "    )\n",
    "    .addStyle(\n",
    "        page={\n",
    "            \"title\": \"KNN Network 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": "22859a61-dbdb-4216-9c2f-202e85ffa4f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pygraphistry has API to encode point color by our clusters... but it doesn't work :)\n",
    "g2: Plottable = g.encode_point_color(\n",
    "    \"Journal-ref-Label\",\n",
    "    as_categorical=True,\n",
    "    palette=CATEGORICAL_PALETTE,\n",
    "    default_mapping=\"#CCC\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3c2acb8-920a-4e53-9d7a-3ddb93280e8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "KNN_G_small = KNN_G.subgraph(\n",
    "    np.random.choice(KNN_G.nodes(), size=1500, p=list(nx.pagerank(KNN_G).values()))\n",
    ")\n",
    "KNN_G_small.number_of_nodes(), KNN_G_small.number_of_edges()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b099f6-4623-4fa6-bfca-e8f2afafaf25",
   "metadata": {},
   "outputs": [],
   "source": [
    "g2.plot(KNN_G_small)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e81d12b-ad39-4602-b888-0762ebe1afd1",
   "metadata": {},
   "source": [
    "### Scaling KNN with Approximate Nearest Neighbors: A-KNN\n",
    "\n",
    "\n",
    "Both popular search engines have Elasticsearch and OpenSearch have A-KNN capabilities as does a very easy to use vector search engine \n",
    "\n",
    "Another option is the Python A-KNN library [PyNNDescent](https://github.com/lmcinnes/pynndescent). Another tool is `faiss` discussed below. Here is a [good tutorial on KNN and FAISS](https://towardsdatascience.com/comprehensive-guide-to-approximate-nearest-neighbors-algorithms-8b94f057d6b6).\n",
    "\n",
    "#### Bottom Line: KNN Can Scale\n",
    "\n",
    "Using one of the tools above or many others, you can make KNN scale via A-KNN with a little engineering or help from a data engineer :)\n",
    "\n",
    "# Final Exercise: LSH Vector Networks\n",
    "\n",
    "I ran out of time, so as an exercise, I want you to look at the code we used to generate the KNN network and make it more scalable using `faiss`. :)\n",
    "\n",
    "## Building More Scalable Locality Sensitive Hashing (LSH) Networks\n",
    "\n",
    "This works similarly, but [Locality Sensitive Hashing (LSH)](https://en.wikipedia.org/wiki/Locality-sensitive_hashing) is a more scalable algorithm that hashes data points into buckets such that similar data points are near one another in the same bucket. It is a general purpose operator for creating pairs of nodes in graph machine learning. \n",
    "\n",
    "### FAISS\n",
    "\n",
    "FAISS [[intro here](https://engineering.fb.com/2017/03/29/data-infrastructure/faiss-a-library-for-efficient-similarity-search/)] is a popular vector search engine from Facebook with an LSH feature... L2 below.\n",
    "\n",
    "### ChatGPT Takes a Swing...\n",
    "\n",
    "Here is what me and ChatGPT wrote as a first pass:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d01cd5-c38e-4f55-bc43-0441f4d8a01b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import faiss\n",
    "# import pandas as pd\n",
    "# import numpy as np\n",
    "# import networkx as nx\n",
    "\n",
    "# # Assume df is your DataFrame and 'sentence_embedding' is the column with embeddings\n",
    "# df = ...  # your DataFrame\n",
    "\n",
    "# # Assume embeddings are stored as lists, so we'll convert to a numpy array\n",
    "# # Also, FAISS requires data to be in float32\n",
    "# embedding_matrix = np.array(node_df[\"Abstract-Paraphrase-Embedding\"].tolist(), dtype=\"float32\")\n",
    "\n",
    "# # Build the FAISS index\n",
    "# dimension = embedding_matrix.shape[1]  # assuming embedding_matrix is your matrix of embeddings\n",
    "# index = faiss.IndexFlatL2(dimension)\n",
    "# index.add(embedding_matrix)\n",
    "\n",
    "# # Perform the KNN search\n",
    "# k = 6  # Search for 5 nearest neighbors + 1 because the query set is the same as the database\n",
    "# distances, indices = index.search(embedding_matrix, k)\n",
    "\n",
    "# # Construct the network\n",
    "# G = nx.Graph()\n",
    "\n",
    "# # Iterate through each row in the DataFrame\n",
    "# for idx, neighbors in enumerate(indices):\n",
    "#     for neighbor, distance in zip(neighbors[1:], distances[idx][1:]):  # Skip the first neighbor as it will be the node itself\n",
    "#         G.add_edge(df.index[idx], df.index[neighbor], weight=distance)\n",
    "\n",
    "# # Now G is your K-Nearest-Neighbor network!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd2fe5e-3bba-461f-9e3f-585d59ddaef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_df.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "142db9b9-0bc1-4f22-9ebc-60c32ce12077",
   "metadata": {},
   "source": [
    "# Big Data and Large Knowledge Graphs with PySpark and GraphFrames\n",
    "\n",
    "I'll be honest with you... I am a big data, large knowledge graph specialist. All the data we have used so far is very small to me. While it is difficult to cover scalable methods for everything, I wanted to introduce you to a tool for PySpark called [GraphFrames](https://graphframes.github.io/graphframes/docs/_site/index.html). It has some powerful utility methods like [connected components](https://graphframes.github.io/graphframes/docs/_site/user-guide.html#connected-components) that many teams use to perform tasks like merging nodes during entity resolution - node deduplication. It can also perform property graph [motif finding](https://graphframes.github.io/graphframes/docs/_site/user-guide.html#motif-finding) for networks with billions of nodes and edges.\n",
    "\n",
    "GraphFrames are graphs [created from node / edge lists](https://graphframes.github.io/graphframes/docs/_site/user-guide.html#creating-graphframes) which are [pyspark.sql.DataFrames](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/dataframe.html). This means you get the power of PySpark to write arbitrary graph operations on large datasets thanks to [Spark](https://spark.apache.org/docs/latest/).\n",
    "\n",
    "## Imperative vs Declarative\n",
    "\n",
    "SQL is a declarative language. PySpark is an imperative API, with a declarative SQL interface if you prefer. We'll be using the `pyspark.sql.DataFrame` and `pyspark.sql.functions` APIs to process data step-by-step. This is different from Pandas, which is fairly declarative in its APIs. Keep that in mind below... we get to carefully control how the dataflow works, which can be time consuming compared to a single Pandas command that does more than what Spark `DataFrames` do - MapReduce and iterate :)\n",
    "\n",
    "## Running Checks\n",
    "\n",
    "It is easy to mess up an algorithm when building knowledge graphs in PySpark / GraphFrames. You will notice that throughout  the script below, I check and print counts as I go. I encourage you to do this as well, or your scripts will seem to run but will produce bad knowledge graphs that produce bad answers.\n",
    "\n",
    "## Running GraphFrames\n",
    "\n",
    "To run GraphFrames locally from a shell, you can import the package with the `spark-shell` or `pyspark` `--packages` argument:\n",
    "\n",
    "```bash\n",
    "pyspark --packages graphframes:graphframes:0.8.3-spark3.5-s_2.12\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64efde8-74bd-4b49-b0aa-d4654037a2ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is actually already set in Docker, just reminding you Java is needed\n",
    "os.environ[\"JAVA_HOME\"] = \"/usr/lib/jvm/java-17-openjdk-amd64\"\n",
    "\n",
    "# Setup PySpark to use the GraphFrames jar package from maven central\n",
    "os.environ[\"PYSPARK_SUBMIT_ARGS\"] = \"--packages graphframes:graphframes:0.8.3-spark3.5-s_2.12 pyspark-shell\"\n",
    "\n",
    "# Show all the rows in a pd.DataFrame\n",
    "pd.set_option('display.max_columns', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c5950fb-f74a-4c7f-bcbf-f0650aabc852",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyspark.sql.functions as F\n",
    "from pyspark.sql import SparkSession\n",
    "\n",
    "# Initialize a SparkSession. You can configre SparkSession via: .config(\"spark.some.config.option\", \"some-value\")\n",
    "spark = (\n",
    "    SparkSession.builder\n",
    "    .appName(\"Big Graph Builder\")  # Set app name\n",
    "    .getOrCreate()  # Get or create the SparkSession\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf969b1-ee0a-4510-aeb7-51b146c871c8",
   "metadata": {},
   "source": [
    "## Loading the `stats.meta.stackexchange.com` Data\n",
    "\n",
    "I wanted to load the `stats.stackexchange.com` dataset, but it might have been too big for your notebook computer so instead we get [stats.meta.stackexchange.com](https://stats.meta.stackexchange.com/) :) We will load the Posts, Users and Votes and then build a knowledge graph using GraphFrames."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dccb1a8d-34b9-4f44-ad03-04cc1db6c00f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the Posts...\n",
    "posts_df = spark.read.parquet(\"data/stats.meta.stackexchange.com/Posts.parquet\").withColumn(\"Type\", F.lit(\"Post\"))\n",
    "\n",
    "print(f\"\\nTotal stats.meta.stackexchange.com Posts: {posts_df.count():,}\\n\")\n",
    "\n",
    "# Do not .sample(0.1), use filter F.rand\n",
    "posts_df.filter(F.rand() < 0.1).limit(10).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b2d5580-4a63-4248-adc9-ff8798badacf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pretty print the post column names\n",
    "print(\"\\nAll unique columns of posts:\\n\")\n",
    "pd.DataFrame({\"Column\": pd.Series(sorted(posts_df.columns))})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f53570b-c538-4db4-8797-bad4395eaa66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the Users...\n",
    "users_df = spark.read.parquet(\"data/stats.meta.stackexchange.com/Users.parquet\").withColumn(\"Type\", F.lit(\"User\"))\n",
    "print(f\"\\nTotal Users: {users_df.count():,}\\n\")\n",
    "\n",
    "users_df.limit(10).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8b762de-4491-45ac-8f00-e71da188c9cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the Votes...\n",
    "votes_df = spark.read.parquet(\"data/stats.meta.stackexchange.com/Votes.parquet\").withColumn(\"Type\", F.lit(\"Vote\"))\n",
    "print(f\"Total Votes: {votes_df.count():,}\")\n",
    "\n",
    "votes_df.limit(10).toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91774abc-81d3-4cf5-989b-a9ca84b3f69d",
   "metadata": {},
   "source": [
    "## Building Node and Edge Lists\n",
    "\n",
    "One downside of GraphFrames is that it uses a single node and edge type - so we need to combine the fields of all three types into one `node_df` `pyspark.sql.DataFrame`. The benefit is that you can work with billions of these nodes. It is painful but neccessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59b61657-1f35-40c9-806f-4004b62ef166",
   "metadata": {},
   "outputs": [],
   "source": [
    "posts_df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b728baee-4903-4c91-aeb3-76234ae4614c",
   "metadata": {},
   "outputs": [],
   "source": [
    "posts_df.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf2437a-8a98-47d5-b460-23f28fffbcc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get all the unique columns to use when filling all DataFrames with all columns\n",
    "all_columns = set(\n",
    "    list(zip(posts_df.columns, posts_df.schema)) + list(zip(users_df.columns, users_df.schema)) + list(zip(votes_df.columns, votes_df.schema))\n",
    ")\n",
    "all_columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e748bc81-2ba2-48ba-a1ee-94af9c02f06d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_missing_columns(df, all_columns):\n",
    "    \"\"\"Add any missing columns from any DataFrame among several we want to merge.\"\"\"\n",
    "    for col_name, schema_field in all_columns:\n",
    "        if col_name not in df.columns:\n",
    "            df = df.withColumn(col_name, F.lit(None).cast(schema_field.dataType))\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f92b92b3-a5bd-4c05-8857-9d8b15e7bbf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now apply this function to each of your DataFrames to get a consistent schema\n",
    "posts_df = add_missing_columns(posts_df, all_columns)\n",
    "users_df = add_missing_columns(users_df, all_columns)\n",
    "votes_df = add_missing_columns(votes_df, all_columns)\n",
    "\n",
    "assert set(posts_df.columns) == set(users_df.columns) == set(votes_df.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baadff73-32b7-45cb-847a-89cee0b47165",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now UNION the three DataFrames to get our nodes\n",
    "nodes_df = posts_df.unionByName(users_df).unionByName(votes_df)\n",
    "print(f\"Total nodes: {nodes_df.count():,}\")\n",
    "print(f\"Total distinct nodes: {nodes_df.distinct().count():,}\")\n",
    "\n",
    "# Set the nodes_df.Id column to lowercase nodes_df.id for GraphFrames\n",
    "nodes_df = nodes_df.withColumnRenamed(\"Id\", \"id\")\n",
    "\n",
    "nodes_df.limit(10).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fa9762e-c1e1-436d-ac09-52ddd5f52d62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The 'Type' column distinguishes between node types\n",
    "nodes_df.groupby(\"Type\").count().toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aadc053e-2233-43fd-bd8f-f7cd8a087e10",
   "metadata": {},
   "source": [
    "### Storing or Caching `nodes_df`\n",
    "\n",
    "It is a good idea to either store [for later] or to cache the nodes before creating a `GraphFrame` for performance reasons."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab7f83a1-b913-43a8-bd7d-095bfcb8f37b",
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes_df.write.mode(\"overwrite\").parquet(\"data/stats.meta.stackexchange.com/Nodes.parquet\")\n",
    "\n",
    "# Helps performance of GraphFrames' algorithms\n",
    "nodes_df = nodes_df.cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a03721df-d0be-44d0-9023-69ab02a676ce",
   "metadata": {},
   "source": [
    "### Computing Edges from Node ID References\n",
    "\n",
    "Now we can use references to and from node types to create edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a08b2152-9fdc-4b4f-ac44-1c1473fb4f21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pd.DataFrame({\"Columns\": pd.Series(sorted(nodes_df.columns))})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4a1108f-9edc-40d3-95fe-27cc0d92b4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# What relatonships do we want to add to our graph? How about Vote--VotedFor-->Post\n",
    "votes_df = nodes_df[nodes_df.Type == \"Vote\"]\n",
    "print(f\"\\nStill have total votes: {votes_df.count():,}\\n\")\n",
    "\n",
    "# Start our edges with voted edges - 'src', 'dst' and 'relationship' are standard edge fields in GraphFrames\n",
    "edges_df = votes_df.select([\n",
    "    F.col(\"Id\").alias(\"src\"),\n",
    "    F.col(\"PostId\").alias(\"dst\"),\n",
    "    F.lit(\"VotedFor\").alias(\"relationship\"),\n",
    "])\n",
    "print(f\"Total edges: {edges_df.count():,}\")\n",
    "\n",
    "# A fast way to sample\n",
    "edges_df.filter(F.rand() < 0.01).limit(10).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06ac4a60-1a78-4134-97eb-1a2d61a8eb22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now link questions to answers...\n",
    "posts_df.limit(4).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b06428e-11d6-4710-8af2-93642a801602",
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted(posts_df.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21a7f23b-bf28-47a9-ad79-0ffc93c4cc2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do the questions look ok? Questions have NO parent ID and DO have a Title\n",
    "questions_df = posts_df[posts_df.ParentId.isNull()]\n",
    "print(f\"\\nTotal questions: {questions_df.count():,}\\n\")\n",
    "\n",
    "questions_df[[\"Title\"]].show(10, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c88b15d1-96e9-4da8-9000-1e8e2be2c648",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Answers DO have a ParentId parent post and no Title\n",
    "answers_df = posts_df[posts_df.ParentId.isNotNull()]\n",
    "print(f\"\\nTotal answers: {answers_df.count():,}\\n\")\n",
    "\n",
    "answers_df.sample(0.10).limit(100).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3fb2511-8b21-4ce6-ad79-d8cbba0cecf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a answer[Post]--Answered-->question[Post] edge\n",
    "answered_df = answers_df.select([\n",
    "    F.col(\"Id\").alias(\"src\"),\n",
    "    F.col(\"ParentId\").alias(\"dst\"),\n",
    "    F.lit(\"Answered\").alias(\"relationship\")\n",
    "])\n",
    "print(f\"\\nTotal Answered edges: {answered_df.count():,}\\n\")\n",
    "answered_df.show()\n",
    "\n",
    "# Note: it probably makes sense to turn Posts into Question and Answers node Types, but for now..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16a054fc-91c8-43c1-a698-fecef2d96825",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the Answered edges to the VotedFor edges :)\n",
    "print(f\"Edges before Answered added to VotedFor: {edges_df.count():,}\")\n",
    "edges_df = edges_df.unionByName(answered_df)\n",
    "print(f\"Edges after Answered  added to VotedFor: {edges_df.count():,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "350b252b-82ac-4146-a929-2bbac5b9ebb9",
   "metadata": {},
   "source": [
    "## Stratified Sample with `pyspark.sql.DataFrame.sampleBy()`\n",
    "\n",
    "[pyspark.sql.DataFrame.sampleBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.sampleBy.html) enables a [stratified sample](https://www.sciencedirect.com/topics/mathematics/stratified-sampling#:~:text=Stratified%20sampling%20designs%20involve%20partitioning,samples%20independently%20from%20each%20stratum.) of records spread across the labels of a column in specific proportions. This lets us see examples of each type of edge as a **sanity check**. You must sanity check your algorithms in PySpark :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "074a0f27-d03b-47e6-89be-fd6e33cb7e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform an evenly stratified sample by edge 'relationship'\n",
    "SAMPLE_RATE = 0.01\n",
    "\n",
    "edge_count = edges_df.count()\n",
    "\n",
    "sample_dict = edges_df.groupby(\"relationship\").count().toPandas().to_dict()\n",
    "sample_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "847d5548-a0d3-436e-81b7-5d4edf8289aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the relationship types from a GROUP BY or SELECT / DISTINCT\n",
    "relationship_types = edges_df.groupby(\"relationship\").count().toPandas()[\"relationship\"].tolist()\n",
    "relationship_types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c871aec-7ca2-45ca-9509-15cc208ec867",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_rates = {key:0.001 for key in relationship_types}\n",
    "sample_rates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a192cd9d-b553-4af2-979c-598e2534661d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Overkill, but do note how I frequently wrap a python Dict with a Series, then display it pretty like as a DataFrame :)\n",
    "pd.DataFrame(pd.Series(sample_rates), columns=[\"Proportions\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bec776d-f200-48b5-a9a6-99263bff9164",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ok, look we can keep at this and get a sample... you get the idea\n",
    "edges_df.sampleBy(\"relationship\", sample_rates).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79195663-7e05-4de1-92eb-c94dae5cbe8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets' add a User--Posted-->Post\n",
    "posts_df = nodes_df[nodes_df.Type == \"Post\"]\n",
    "print(f\"\\nStill have total posts: {posts_df.count():,}\\n\")\n",
    "\n",
    "# Create Posted edges\n",
    "posted_df = posts_df.select([\n",
    "    F.col(\"OwnerUserId\").alias(\"src\"),\n",
    "    F.col(\"Id\").alias(\"dst\"),\n",
    "    F.lit(\"Posted\").alias(\"relationship\"),\n",
    "])\n",
    "# Add the Posted edges to the existing set of edges\n",
    "edges_df = edges_df.unionByName(posted_df)\n",
    "\n",
    "print(f\"Total edges: {edges_df.count():,}\")\n",
    "edges_df.filter(F.rand() < 0.01).limit(10).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c062791-3a9b-411c-b6ef-7074afa0761c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's add Users-Asked->question[Post]\n",
    "asked_df = questions_df.select([\n",
    "    F.col(\"OwnerUserId\").alias(\"src\"),\n",
    "    F.col(\"Id\").alias(\"dst\"),\n",
    "    F.lit(\"Asked\").alias(\"relationship\")\n",
    "])\n",
    "print(f\"Total questions asked: {asked_df.count():,}\")\n",
    "\n",
    "edges_df = edges_df.unionByName(asked_df)\n",
    "print(f\"Total edges now: {edges_df.count():,}\")\n",
    "\n",
    "edges_df.sample(0.001).limit(20).toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f77f63c-83e7-438c-b381-f2399a66474f",
   "metadata": {},
   "source": [
    "### Cleaning Nodes and Edges\n",
    "\n",
    "It is a good idea to deduplicate your node and edge lists."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77b3c80b-1b37-4add-b96e-414a2c104d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Node count: {nodes_df.count():,}\")\n",
    "print(f\"Edge count: {edges_df.count():,}\")\n",
    "\n",
    "nodes_df = nodes_df.distinct()\n",
    "edges_df = edges_df.distinct()\n",
    "\n",
    "print(f\"Distinct node count: {nodes_df.count():,}\")\n",
    "print(f\"Distinct edge count: {edges_df.count():,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffe8f287-00a4-4486-8fcc-3b509cbecead",
   "metadata": {},
   "source": [
    "## Creating a GraphFrame\n",
    "\n",
    "Let's create a `graphframes.GraphFrame` using the instructions in the [Quickstart](https://graphframes.github.io/graphframes/docs/_site/quick-start.html) and then test some operations from the [User Guide](https://graphframes.github.io/graphframes/docs/_site/user-guide.html). There are terse [API Docs](https://graphframes.github.io/graphframes/docs/_site/api/python/graphframes.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba8e726a-2826-447d-aa93-8a80e3c48937",
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphframes import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02964dbd-170f-41e3-a8ce-cbab029434c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = GraphFrame(nodes_df, edges_df)\n",
    "g"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5328a7a7-76f1-45c9-9921-71ae59eab65e",
   "metadata": {},
   "source": [
    "## Finding Motifs\n",
    "\n",
    "Can we think of a graph pattern to search for? What about voting rings? People who vote together for many of the same posts? I think we have the data for that...\n",
    "\n",
    "This motif stinks, I'll think of another one and get back to you. Read the GraphFrames paper [GraphFrames: An Integrated API for Mixing Graph and\n",
    "Relational Queries](https://people.eecs.berkeley.edu/~matei/papers/2016/grades_graphframes.pdf) for more reasonable motifs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdc6794e-bc57-4dc0-b19c-48b1c23f26f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can we think of a motif pattern to find? That is my hobby :D\n",
    "# (UserA)--[Asked]-->(Questions [Posts])<--[Answered]--(UserB);\n",
    "# (UserB)--[VotedFor]-->(Answers [Posts])<--Answered--(UserA);\n",
    "paths = (\n",
    "    g.find(\"(a)-[ab]->(b); (c)-[bc]->(b); (c)-[cd]->(d); (e)-[de]->(d)\")\n",
    "    # (a)-[ab]->(b)\n",
    "    .filter(\"a.Type == 'User'\")\n",
    "    #.filter(\"ab.relationship == 'Asked'\")\n",
    "    .filter(\"b.Type == 'Post'\")\n",
    "    #.filter(F.col(\"b.ParentId\").isNull())  # Post (b) is a question\n",
    "    # (b)<-[bc]-(c)\n",
    "    # .filter(\"c.Type == 'User'\")\n",
    "    # .filter(\"bc.relationship == 'Answered'\")\n",
    "    # (c)-[cd]->(d)\n",
    "        \n",
    "    # .filter(\"ab.relationship == 'Posted'\")\n",
    "    # .filter(\"b.Type == 'Post'\")\n",
    "    # .filter(\"bc.relationship == 'Answered'\")\n",
    "    # .filter(\"c.Type == 'Post'\")\n",
    ")\n",
    "paths.show()"
   ]
  },
  {
   "cell_type": "raw",
   "id": "760d2bfd-746f-4a41-b428-ec790eec78ed",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9423ae3a-10be-44c8-96cb-33b6e47f0889",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is an example of post-path match processing to get a complex, motif-based answer\n",
    "# Q: What pairs of nodes have the most common neighbors?\n",
    "unique_paths = paths.groupby([F.col(\"a.id\").alias(\"a_id\"), F.col(\"c.id\").alias(\"c_id\")]).count()\n",
    "unique_paths.orderBy(F.col(\"count\").desc()).limit(20).toPandas()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7812abb9-243f-4c61-bcda-efa8a8851277",
   "metadata": {},
   "source": [
    "### Real-World Network Motifs and Higher-Order Networks\n",
    "\n",
    "Remember these slides? This is a real-world use of network motif matching to find sources of risk for banks in terms of money laundering and terrorism funding occurring in financial networks of people and companies.\n",
    "\n",
    "<center><img src=\"images/Multiple-Path-Indirect-Ownership-Motif.jpg\" width=\"800px\" /></center>\n",
    "<center>A well known pattern to hide Ultimate Beneficial Ownership (UBO) of a company</center>\n",
    "\n",
    "<br />\n",
    "\n",
    "<center><img src=\"images/PySpark-GraphFrames-Motif-Search-Python-Code.jpg\" width=\"800px\" /></center>\n",
    "<center>The PySpark / GraphFrames pseudo-code that implements [most of] the motif above</center>\n",
    "\n",
    "<br />\n",
    "\n",
    "The next motif is more complicated... it uses Sentence Transformers with PySpark - [[here's how](https://stackoverflow.com/questions/72398129/creating-a-sentence-transformer-model-in-spark-mllib)] to perform what I think is a new type of network motif - a _semantic network motif_.\n",
    "\n",
    "<center><img src=\"images/Corrupt-Incorporation-Services-Motif.jpg\" width=\"800px\" /></center>\n",
    "<center>A more complex <i>semantic network motif</i> using Sentence Transformers to do fuzzy string matching of officer names</center>\n",
    "\n",
    "<br />\n",
    "\n",
    "While I don't have time to demonstrate it at present, once you have network motifs, it is possible to group by all members of a motif match and combine their edges into a new, higher-order node with new semantics. If you link or [cluster](https://arxiv.org/abs/1612.08447) these new nodes, they form a [higher-order network](https://arxiv.org/abs/2104.11329). If the motif is a pattern important to your problem domain, nodes in this higher-order network might exist in a _solution space_ rather than your graph's original _problem space_. For example, if you define risk motifs in a financial network, the way they cluster can show you centers of risk. Pattern matching is a different approach than graph machine learning but can yield similarly powerful results.\n",
    "\n",
    "<center><img src=\"images/Higher-Order-Networks-Using-Edge-Projection-via-Property-Graphlet-Minors.jpg\" width=\"800px\" /></center>\n",
    "<center>Forming a higher-order network using network motif clustering</center>\n",
    "\n",
    "For more information, check out Stanfor SNAP's page on [Higher-order organization of complex networks](https://snap.stanford.edu/higher-order/). The diagrams of overlapping motifs are really cool:\n",
    "\n",
    "<center><img src=\"images/Leskovec-Cutting-Motif-Clusters.png\" width=\"600px\" /></center>\n",
    "\n",
    "<br /><br />\n",
    "\n",
    "<center><img src=\"images/Leskovec-Motif-Cluster-Efficiency.png\" width=\"600px\" /></center>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fe44ee0-a487-429a-b90a-c8d39d32dc82",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}