examples/kg_download_quick_start/kg_download_quick_start.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8c999031",
   "metadata": {},
   "source": [
    "# Using Deep Search Knowledge Graphs with PyTorch Geometric"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "959e2f71",
   "metadata": {},
   "source": [
    "Deep Search can construct Knowledge Graphs (KGs) by parsing large collections of documents. This tutorial shows how to download these knowledge graphs and import them locally in [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric), a popular graph neural network library. \n",
    "\n",
    "### Access required\n",
    "\n",
    "The content of this notebook requires access to Deep Search capabilities which are not\n",
    "available on the public access system.\n",
    "\n",
    "[Contact us](https://ds4sd.github.io) if you are interested in exploring\n",
    "these Deep Search capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8274fd05",
   "metadata": {},
   "source": [
    "### Set notebook parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "01767c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsnotebooks.settings import KGProjectNotebookSettings\n",
    "\n",
    "# notebook settings auto-loaded from .env / env vars\n",
    "notebook_settings = KGProjectNotebookSettings()\n",
    "\n",
    "PROFILE_NAME = notebook_settings.profile  # the profile to use\n",
    "PROJECT_KEY = notebook_settings.proj_key\n",
    "KG_KEY = notebook_settings.kg_key\n",
    "BASE_DIR = \"./KG-data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be45eb77",
   "metadata": {},
   "source": [
    "### Example dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1b897f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import standard dependencies\n",
    "import os\n",
    "import ssl\n",
    "import json\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import tarfile\n",
    "from urllib.request import urlopen\n",
    "\n",
    "# Dependencies related to PyTorch Geometric\n",
    "import torch\n",
    "from torch_geometric.data import HeteroData\n",
    "from torch_geometric.transforms import ToUndirected\n",
    "\n",
    "# IPython utilities\n",
    "from IPython.display import display\n",
    "\n",
    "# Import deepsearch-toolkit\n",
    "from deepsearch.core.client import DeepSearchConfig\n",
    "from deepsearch.cps.client.api import CpsApi, CpsApiClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a282398b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create base directory if it does not exist\n",
    "if not os.path.exists(BASE_DIR):\n",
    "    os.mkdir(BASE_DIR)\n",
    "\n",
    "\n",
    "# Raise an error if the base directory is not empty\n",
    "if len(os.listdir(BASE_DIR)) > 0:\n",
    "    raise ValueError(\n",
    "        f\"BASE_DIR must be empty but found the following contents: {os.listdir(BASE_DIR)}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b6e5ca",
   "metadata": {},
   "source": [
    "# Part 1: Downloading the Knowledge Graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06dc4f79",
   "metadata": {},
   "source": [
    "### Connect to Deep Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "004e7de9",
   "metadata": {},
   "outputs": [],
   "source": [
    "api = ds.CpsApi.from_env(profile_name=PROFILE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42b1cc3f",
   "metadata": {},
   "source": [
    "### Download the knowledge graph\n",
    "\n",
    "We use an example knowledge graph based on 753 documents from arXiv related to the search phrase \"power conversion efficiency\" + \"organic\". The `PROJECT_KEY` and `KG_KEY` parameters that were set above correspond to this knowledge graph. In general, these parameters can be obtained from the API section of the UI as described in the [documentation](https://ds4sd.github.io/deepsearch-toolkit/guide/kgs/). \n",
    "\n",
    "We begin by using the API to get a URL for downloading the knowledge graph. **This step takes about 6-8 minutes due to the size of our KG.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d5c38ba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the download url using the API\n",
    "kg = api.knowledge_graphs.get(PROJECT_KEY, KG_KEY)\n",
    "download_url = kg.download()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1b786ea",
   "metadata": {},
   "source": [
    "The `download_url` can now be used to download a gzipped file that contains the contents of our knowledge graph. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f8b24fe9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌| 99.79631035475828/100 [00:41<00:00,  2.60it/s]/Users/shubham/Experiments/KGDownload-Example/deepsearch-examples/.venv/lib/python3.10/site-packages/tqdm/std.py:533: TqdmWarning: clamping frac to range [0, 1]\n",
      "  full_bar = Bar(frac,\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100.00000000004945/100 [00:41<00:00,  2.39it/s]\n"
     ]
    }
   ],
   "source": [
    "# Download the knowledge graph using urlopen\n",
    "zipped_file_path = os.path.join(BASE_DIR, \"kg_data.tar.gz\")\n",
    "context = ssl.create_default_context()\n",
    "context.check_hostname = False\n",
    "context.verify_mode = ssl.CERT_NONE\n",
    "\n",
    "with open(zipped_file_path, \"wb+\") as download_file, urlopen(\n",
    "    download_url, context=context\n",
    ") as response:\n",
    "    content_length = int(response.getheader(\"content-length\"))\n",
    "    with tqdm(total=100, position=0) as pbar:\n",
    "        for line in response:\n",
    "            download_file.write(line)\n",
    "            pbar.update((len(line) / content_length) * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "012c944a",
   "metadata": {},
   "source": [
    "### Extract the contents of the downloaded knowledge graph\n",
    "We use the `tarfile` module to unzip the contents of the downloaded `.tar.gz` file into the directory `BASE_DIR/unzipped_data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c1dada55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the unzipped KG\n",
    "unzipped_dir = os.path.join(BASE_DIR, \"unzipped_data\")\n",
    "os.mkdir(unzipped_dir)\n",
    "with tarfile.open(zipped_file_path, \"r\") as f:\n",
    "    f.extractall(path=unzipped_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc1faf6",
   "metadata": {},
   "source": [
    "### Understanding the downloaded data\n",
    "The extracted data in `BASE_DIR/unzipped_data` consists of several files in the [JSON Lines](https://jsonlines.org) format. As the downloaded KG is heterogeneous (meaning there are different types of nodes and edges), we get one `.jsonl` file for every node type. A special file named `_edges.jsonl` contains information about the edges in the KG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3f9cb932",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['_edges.jsonl',\n",
       " 'abbreviation.jsonl',\n",
       " 'affiliations.jsonl',\n",
       " 'authors.jsonl',\n",
       " 'classifications.jsonl',\n",
       " 'documents.jsonl',\n",
       " 'identifiers.jsonl',\n",
       " 'images.jsonl',\n",
       " 'material-attribute.jsonl',\n",
       " 'material-class.jsonl',\n",
       " 'material-part.jsonl',\n",
       " 'material-to-abbreviation.jsonl',\n",
       " 'material-to-material-class.jsonl',\n",
       " 'material-to-property.jsonl',\n",
       " 'material.jsonl',\n",
       " 'model-type.jsonl',\n",
       " 'property-to-property-value.jsonl',\n",
       " 'property-value-type.jsonl',\n",
       " 'property-value.jsonl',\n",
       " 'property.jsonl',\n",
       " 'publishers.jsonl',\n",
       " 'smiles.jsonl',\n",
       " 'statement.jsonl',\n",
       " 'tables.jsonl',\n",
       " 'terms.jsonl',\n",
       " 'text-type.jsonl',\n",
       " 'texts.jsonl',\n",
       " 'unit-si.jsonl',\n",
       " 'unit.jsonl',\n",
       " 'year-month.jsonl',\n",
       " 'year.jsonl']"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Get a list of all the files in the unzipped data\n",
    "files = list(os.walk(os.path.join(BASE_DIR, \"unzipped_data\")))[0][2]\n",
    "display(sorted(files))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce5950ba",
   "metadata": {},
   "source": [
    "Each line in a `.jsonl` file contains information about a node encoded in the JSON format. For example, we show the first record from the `material.jsonl` file below. This file contains information about different materials found in the documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7408bb61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A function for reading the contents of a jsonl file into a pandas dataframe\n",
    "def jsonl2df(filepath: str) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Reads the contents of a jsonl file into a Pandas DataFrame\n",
    "    :param filepath: Path to the jsonl file\n",
    "    :return dataframe: A pandas DataFrame corresponding to the data stored in the file\n",
    "    \"\"\"\n",
    "    with open(filepath, \"r\") as f:\n",
    "        data = pd.DataFrame([json.loads(line) for line in f])\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b6293224",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "_id                                           6411a23382b445f4a85ab6b9\n",
       "_hash                                 6f9f25a0a0ef3246314db75f7184489c\n",
       "_name                                                        Pb iodide\n",
       "_synonyms                                                  [Pb iodide]\n",
       "_type                                                         material\n",
       "identifiers          [{'_name': 'ent_id#64efbff3', 'type': 'ent_id'...\n",
       "labels                                                    [non-atomic]\n",
       "model                                             [IUPACNameAnnotator]\n",
       "names                [{'_name': 'chemical_name#lead iodide', 'type'...\n",
       "prov                                                      [text, data]\n",
       "flags                  {'aggregations': {'@smi': True, 'model': True}}\n",
       "coor                              [0, 8, 13, 15, 17, 21, 24, 31, 2, 1]\n",
       "doping-parameter                                                   NaN\n",
       "element-parameter                                                  NaN\n",
       "Name: 0, dtype: object"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show the first record in the materials file\n",
    "materials = jsonl2df(os.path.join(BASE_DIR, \"unzipped_data\", \"material.jsonl\"))\n",
    "display(materials.iloc[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef08de5a",
   "metadata": {},
   "source": [
    "We now take a look at the `_edges.jsonl` file. As before, each line contains information about an edge in the JSON format. Given a directed edge from `x` to `y`, we call `x` the _tail_ of this edge and `y` its _head_. We show a few interesting properties of these edges listed below.\n",
    "* `source_collection`: Type of the node at the tail of the edge\n",
    "* `target_collection`: Type of the node at the head of the edge\n",
    "* `source_hash`: Hash for the node at the tail of the edge\n",
    "* `target_hash`: Hash for the node at the head of the edge\n",
    "* `symmetric`: Whether the edge is directed (if `False`) or undirected (if `True`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "883bfa79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>source_collection</th>\n",
       "      <th>target_collection</th>\n",
       "      <th>source_hash</th>\n",
       "      <th>target_hash</th>\n",
       "      <th>symmetric</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>documents</td>\n",
       "      <td>texts</td>\n",
       "      <td>c4d297ab8e33f5419e2bf081cf8b9f2c</td>\n",
       "      <td>ac6753cb167f782854d66f49ef8a47fc</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>documents</td>\n",
       "      <td>texts</td>\n",
       "      <td>6db1ccc0bdb23fe0256571ca6e7020b0</td>\n",
       "      <td>d8843473fe754654733260c16936f051</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>documents</td>\n",
       "      <td>texts</td>\n",
       "      <td>6db1ccc0bdb23fe0256571ca6e7020b0</td>\n",
       "      <td>91d5b26959546063b595ea75f7b99bde</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>documents</td>\n",
       "      <td>texts</td>\n",
       "      <td>6db1ccc0bdb23fe0256571ca6e7020b0</td>\n",
       "      <td>4435cbe3ce0b1c8feea9a593f1e31c1d</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>documents</td>\n",
       "      <td>texts</td>\n",
       "      <td>c4d297ab8e33f5419e2bf081cf8b9f2c</td>\n",
       "      <td>8c3f5d93479d56f7a204c65fe1a3daaf</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  source_collection target_collection                       source_hash  \\\n",
       "0         documents             texts  c4d297ab8e33f5419e2bf081cf8b9f2c   \n",
       "1         documents             texts  6db1ccc0bdb23fe0256571ca6e7020b0   \n",
       "2         documents             texts  6db1ccc0bdb23fe0256571ca6e7020b0   \n",
       "3         documents             texts  6db1ccc0bdb23fe0256571ca6e7020b0   \n",
       "4         documents             texts  c4d297ab8e33f5419e2bf081cf8b9f2c   \n",
       "\n",
       "                        target_hash  symmetric  \n",
       "0  ac6753cb167f782854d66f49ef8a47fc       True  \n",
       "1  d8843473fe754654733260c16936f051       True  \n",
       "2  91d5b26959546063b595ea75f7b99bde       True  \n",
       "3  4435cbe3ce0b1c8feea9a593f1e31c1d       True  \n",
       "4  8c3f5d93479d56f7a204c65fe1a3daaf       True  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show the first few edges\n",
    "edges = jsonl2df(os.path.join(BASE_DIR, \"unzipped_data\", \"_edges.jsonl\"))\n",
    "edges = edges[\n",
    "    [\n",
    "        \"source_collection\",\n",
    "        \"target_collection\",\n",
    "        \"source_hash\",\n",
    "        \"target_hash\",\n",
    "        \"symmetric\",\n",
    "    ]\n",
    "]\n",
    "display(edges.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44b88420",
   "metadata": {},
   "source": [
    "# Part 2: Creating a PyTorch Geometric Knowledge Graph\n",
    "\n",
    "Recall that the knowledge graph downloaded above is based on 753 documents from arXiv related to the search phrase \"power conversion efficiency\" + \"organic\". Deep Search has extracted various entities and relationships from these documents using a user-defined data flow. In this section, we will use a subset of the downloaded data to create a heterogeneous knowledge graph in PyTorch Geometric.\n",
    "\n",
    "The subset we are interested in is concerned with materials and their properties. We will add the following two types of nodes from the downloaded KG to our PyTorch Geometric KG:\n",
    "1. `material`: Materials discovered in the documents\n",
    "2. `property`: Various material properties extracted from the documents\n",
    "\n",
    "We will then extract edges relating these nodes from the `_edges.jsonl` file. Let us begin by initializing an empty knowledge graph.\n",
    "\n",
    "### Initialize the KG\n",
    "\n",
    "PyTorch Geometric uses the `HeteroData` class to represent a heterogeneous knowledge graph. Below, we create an empty `HeteroData` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "396ab9c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an empty heterogeneous knowledge graph\n",
    "hetero_kg = HeteroData()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad42c72",
   "metadata": {},
   "source": [
    "### Add nodes to the KG\n",
    "Next, we add the `material` and `property` nodes to our KG as mentioned above. In this simplified example, the nodes do not have explicit features associated with them. We therefore use one-hot encoding to set `hetero_kg[nodetype].x`, as is a common practice. We also add two other attributes to each node, the `_hash` for adding edges later on and the `_name` for visualization purposes. One can also set `hetero_kg[node_type].y = ...` if node labels are available as attributes in the corresponding `.jsonl` file for the node type. We do not use any labels in this minimal example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "02fc20f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "nodetypes = {\n",
    "    \"material\": os.path.join(BASE_DIR, \"unzipped_data\", \"material.jsonl\"),\n",
    "    \"property\": os.path.join(BASE_DIR, \"unzipped_data\", \"property.jsonl\"),\n",
    "}\n",
    "\n",
    "for nodetype in nodetypes:\n",
    "    data = jsonl2df(nodetypes[nodetype])\n",
    "    hetero_kg[nodetype].x = torch.eye(data.shape[0])\n",
    "    hetero_kg[nodetype][\"_hash\"] = dict(\n",
    "        (_hash, _idx) for _idx, _hash in enumerate(data[\"_hash\"].to_list())\n",
    "    )\n",
    "    hetero_kg[nodetype][\"_name\"] = data[\"_name\"].to_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e79607cd",
   "metadata": {},
   "source": [
    "### Add edges to the KG\n",
    "\n",
    "As we only have `material` and `property` nodes in `hetero_kg`, we search for edges between these two node types in the `_edges.jsonl` file and add them to `hetero_kg`. This adds an edge between `material A` and `property B` if `B` was mentioned in the context of `A` in at least one of the arXiv documents. We can go one step further and add the value that `property B` takes for `material A` as an edge attribute[$^{[1]}$](#footnote1), but we avoid it here to simplify the exposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "21fd8c27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the relevant edges\n",
    "edges = jsonl2df(os.path.join(BASE_DIR, \"unzipped_data\", \"_edges.jsonl\"))\n",
    "edges = edges[\n",
    "    (edges.source_collection == \"material\") & (edges.target_collection == \"property\")\n",
    "]\n",
    "edges = [edges[\"source_hash\"].to_list(), edges[\"target_hash\"].to_list()]\n",
    "\n",
    "# Create the edge index\n",
    "edge_index = []\n",
    "for hash_mat, hash_prop in zip(*edges):\n",
    "    edge_index.append(\n",
    "        [\n",
    "            hetero_kg[\"material\"][\"_hash\"][hash_mat],\n",
    "            hetero_kg[\"property\"][\"_hash\"][hash_prop],\n",
    "        ]\n",
    "    )\n",
    "\n",
    "# Add edge index to the KG\n",
    "hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index = (\n",
    "    torch.tensor(edge_index).long().t()\n",
    ")\n",
    "\n",
    "# Make the graph undirected\n",
    "hetero_kg = ToUndirected()(hetero_kg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b583cb1e",
   "metadata": {},
   "source": [
    "### Summarize the created KG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6e1e2e48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of nodes\n",
      "\tmaterial -> 9040\n",
      "\tproperty -> 16581\n",
      "Total number of nodes: 25621\n",
      "\n",
      "Number of edges\n",
      "\t('material', 'mat2prop', 'property') -> 12752\n",
      "\t('property', 'rev_mat2prop', 'material') -> 12752\n",
      "Total number of edges: 25504\n"
     ]
    }
   ],
   "source": [
    "# Summarize the KG\n",
    "print(\"Number of nodes\")\n",
    "for node_type in hetero_kg.node_types:\n",
    "    print(f\"\\t{node_type} -> {hetero_kg[node_type].num_nodes}\")\n",
    "print(f\"Total number of nodes: {hetero_kg.num_nodes}\")\n",
    "\n",
    "print(\"\\nNumber of edges\")\n",
    "for edge_type in hetero_kg.edge_types:\n",
    "    print(f\"\\t{edge_type} -> {hetero_kg[edge_type].num_edges}\")\n",
    "print(f\"Total number of edges: {hetero_kg.num_edges}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a25e524",
   "metadata": {},
   "source": [
    "It is interesting to note that the constructed KG is very sparse as the number of edges is almost equal to the number of nodes. This is because Deep Search extracted several properties and materials but each material was only linked to a handful of properties in the document collection."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d122b744",
   "metadata": {},
   "source": [
    "### Visualizing the KG\n",
    "\n",
    "Note that our KG is a bipartite graph as it only contains edges between `material` and `property` nodes. We visualize a small subset of this KG by selecting four interesting `material` nodes and listing up to four randomly chosen `property` nodes that they are connected to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "90d1ec78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>material</th>\n",
       "      <th>Property0</th>\n",
       "      <th>Property1</th>\n",
       "      <th>Property2</th>\n",
       "      <th>Property3</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>perovskite/Si</td>\n",
       "      <td>band gap</td>\n",
       "      <td>energy yield</td>\n",
       "      <td>rough surface</td>\n",
       "      <td>temperature</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>O(2) Ti(1)</td>\n",
       "      <td>optical bandgap</td>\n",
       "      <td>pH value</td>\n",
       "      <td>ambient temperature</td>\n",
       "      <td>electrical conductivity</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A(1) I(3) M(1) Pb(1)</td>\n",
       "      <td>free energy surface</td>\n",
       "      <td>kinetic energy</td>\n",
       "      <td>domain size</td>\n",
       "      <td>optical bandgap</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>O(1) Zn(1)</td>\n",
       "      <td>density-of-states</td>\n",
       "      <td>high temperature</td>\n",
       "      <td>work function</td>\n",
       "      <td>symmetry-class</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               material            Property0         Property1  \\\n",
       "0         perovskite/Si             band gap      energy yield   \n",
       "1            O(2) Ti(1)      optical bandgap          pH value   \n",
       "2  A(1) I(3) M(1) Pb(1)  free energy surface    kinetic energy   \n",
       "3            O(1) Zn(1)    density-of-states  high temperature   \n",
       "\n",
       "             Property2                Property3  \n",
       "0        rough surface              temperature  \n",
       "1  ambient temperature  electrical conductivity  \n",
       "2          domain size          optical bandgap  \n",
       "3        work function           symmetry-class  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Select materials to display\n",
    "materials = [\"perovskite/Si\", \"O(2) Ti(1)\", \"A(1) I(3) M(1) Pb(1)\", \"O(1) Zn(1)\"]\n",
    "mat_idx = [hetero_kg[\"material\"][\"_name\"].index(mat) for mat in materials]\n",
    "\n",
    "# Get properties corresponding to each material\n",
    "properties = dict()\n",
    "for m_idx, material in zip(mat_idx, materials):\n",
    "    current_edges = (\n",
    "        hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[0, :] == m_idx\n",
    "    )\n",
    "    prop_idx = hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[\n",
    "        1, current_edges\n",
    "    ]\n",
    "    properties[material] = [\n",
    "        hetero_kg[\"property\"][\"_name\"][idx] for idx in prop_idx.tolist()\n",
    "    ]\n",
    "\n",
    "# Show up to four randomly chosen properties for each material\n",
    "df = pd.DataFrame()\n",
    "for mat, prop in properties.items():\n",
    "    # Restrict to four properties\n",
    "    if len(prop) > 4:\n",
    "        prop = [prop[idx] for idx in torch.randperm(len(prop)).tolist()[:4]]\n",
    "\n",
    "    # Add the row to the dataframe\n",
    "    curr_dict = dict(\n",
    "        [(\"material\", [mat])]\n",
    "        + [(f\"Property{p_idx}\", [p]) for p_idx, p in enumerate(prop)]\n",
    "    )\n",
    "    curr_df = pd.DataFrame(curr_dict)\n",
    "    df = pd.concat([df, curr_df]).reset_index(drop=True)\n",
    "\n",
    "display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d8926fb",
   "metadata": {},
   "source": [
    "Let's take another example. A `perovskite/Si` tandem is useful for making efficient solar cells. A few interesting properties for this material include its power conversion efficiency (PCE) and its band gap. Below we confirm that these properties are indeed linked to `perovskite/Si` in the KG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fc57fd5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Is perovskite/Si linked to power conversion efficiency: True\n",
      "Is perovskite/Si linked to band gap: True\n"
     ]
    }
   ],
   "source": [
    "# Find properties linked to perovskite/Si\n",
    "m_idx = hetero_kg[\"material\"][\"_name\"].index(\"perovskite/Si\")\n",
    "perovskite_edges = (\n",
    "    hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[0, :] == m_idx\n",
    ")\n",
    "prop_idx = hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[1, perovskite_edges]\n",
    "properties = [hetero_kg[\"property\"][\"_name\"][idx] for idx in prop_idx.tolist()]\n",
    "\n",
    "# Check if the desired properties are linked\n",
    "print(\n",
    "    f'Is perovskite/Si linked to power conversion efficiency: {\"power conversion efficiency\" in properties}'\n",
    ")\n",
    "print(f'Is perovskite/Si linked to band gap: {\"band gap\" in properties}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4923ea",
   "metadata": {},
   "source": [
    "This concludes the process of creating a simple PyTorch Geometric KG from the KG downloaded from Deep Search. One can now train powerful graph neural networks for various downstream tasks like node classification and link prediction using `hetero_kg`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e92255c",
   "metadata": {},
   "source": [
    "# Footnotes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67654dac",
   "metadata": {},
   "source": [
    "<a name='footnote1'></a> \n",
    "### \\[1\\] Edges with attributes\n",
    "\n",
    "Edges with attributes can be stored as additional nodes in Deep Search. For example, consider four nodes `{A, B, X, Y}` where `{A, B}` have the same node type (say `author`) and `{X, Y}` have another node type (say `paper`). The relationship between nodes of type `author` and `paper` can be encoded as nodes of type `author-to-paper`. If author `A` has written paper `X`, then there will a node `A_X` of type `author-to-paper`, and both `A` and `X` will be connected to `A_X` in the KG. \n",
    "\n",
    "Now suppose we want to add more information about the relationship between `author A` and `paper X`. For example, we may be interested in knowing the index at which `A` appears in the author list of `X`. Such information can be added as attributes of node `A_X` in Deep Search. One can then easily parse these nodes to add edge attributes in PyTorch Geometric."
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}