notebooks/2_validation_moas.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generation of test dataset\n",
    "\n",
    "Here, I explore [DrugMechDB](https://sulab.github.io/DrugMechDB/) which could ideally provide us with true positive paths / MOAs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "from pubchempy import get_compounds\n",
    "from collections import Counter"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load in the DrugMechDB file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "KG_DIR = '../data/kg'\n",
    "VALIDATION_DIR = '../data/validation'\n",
    "MAPPING_DIR = '../data/mappings'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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>n1</th>\n",
       "      <th>e1</th>\n",
       "      <th>n2</th>\n",
       "      <th>e2</th>\n",
       "      <th>n3</th>\n",
       "      <th>e3</th>\n",
       "      <th>n4</th>\n",
       "      <th>e4</th>\n",
       "      <th>n5</th>\n",
       "      <th>e5</th>\n",
       "      <th>n6</th>\n",
       "      <th>e6</th>\n",
       "      <th>n7</th>\n",
       "      <th>e7</th>\n",
       "      <th>n8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>imatinib</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>BCR/ABL</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>CML (ph+)</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>imatinib</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>c-Kit</td>\n",
       "      <td>UP_REGULATES</td>\n",
       "      <td>Cellular proliferation</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Systemic mast cell disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         n1        e1       n2            e2                      n3      e3  \\\n",
       "0  imatinib  INHIBITS  BCR/ABL        CAUSES               CML (ph+)     NaN   \n",
       "1  imatinib  INHIBITS    c-Kit  UP_REGULATES  Cellular proliferation  CAUSES   \n",
       "\n",
       "                           n4   e4   n5   e5   n6   e6   n7   e7   n8  \n",
       "0                         NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "1  Systemic mast cell disease  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_names_df = pd.read_excel(f'{VALIDATION_DIR}/indication_MOA_paths.xlsx', sheet_name='paths')\n",
    "node_names_df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_type_df = pd.read_excel(f'{VALIDATION_DIR}/indication_MOA_paths.xlsx', sheet_name='metapaths')\n",
    "node_ids_df = pd.read_excel(f'{VALIDATION_DIR}/indication_MOA_paths.xlsx', sheet_name='node_ids')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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>n1</th>\n",
       "      <th>n2</th>\n",
       "      <th>n3</th>\n",
       "      <th>n4</th>\n",
       "      <th>n5</th>\n",
       "      <th>n6</th>\n",
       "      <th>n7</th>\n",
       "      <th>n8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>MESH:D000068877</td>\n",
       "      <td>UniProt:P00519</td>\n",
       "      <td>MESH:D015464</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>MESH:D000068877</td>\n",
       "      <td>UniProt:P10721</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>MESH:D034721</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                n1              n2            n3            n4   n5   n6   n7  \\\n",
       "0  MESH:D000068877  UniProt:P00519  MESH:D015464           NaN  NaN  NaN  NaN   \n",
       "1  MESH:D000068877  UniProt:P10721    GO:0008283  MESH:D034721  NaN  NaN  NaN   \n",
       "\n",
       "    n8  \n",
       "0  NaN  \n",
       "1  NaN  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_ids_df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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>n1</th>\n",
       "      <th>e1</th>\n",
       "      <th>n2</th>\n",
       "      <th>e2</th>\n",
       "      <th>n3</th>\n",
       "      <th>e3</th>\n",
       "      <th>n4</th>\n",
       "      <th>e4</th>\n",
       "      <th>n5</th>\n",
       "      <th>e5</th>\n",
       "      <th>n6</th>\n",
       "      <th>e6</th>\n",
       "      <th>n7</th>\n",
       "      <th>e7</th>\n",
       "      <th>n8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Protein</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Protein</td>\n",
       "      <td>UP_REGULATES</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     n1        e1       n2            e2                  n3      e3       n4  \\\n",
       "0  Drug  INHIBITS  Protein        CAUSES             Disease     NaN      NaN   \n",
       "1  Drug  INHIBITS  Protein  UP_REGULATES  Biological Process  CAUSES  Disease   \n",
       "\n",
       "    e4   n5   e5   n6   e6   n7   e7   n8  \n",
       "0  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "1  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_type_df.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get all drug-bp pairs\n",
    "\n",
    "From the main file, each row is a valide drug-bp path. So each path is considered individually and only paths that connect a drug to a biological process are selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:00<00:00, 2801.43it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "101"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drug_bp_pairs = []\n",
    "\n",
    "for idx, row in tqdm(node_type_df.iterrows(), total=node_type_df.shape[0]):\n",
    "    try:\n",
    "        assert row['n1'] == 'Drug'\n",
    "        drug_node = {\n",
    "            'idx': node_ids_df.iloc[idx]['n1'], \n",
    "            'name': node_names_df.iloc[idx]['n1']\n",
    "        }\n",
    "    except AssertionError:\n",
    "        continue\n",
    "\n",
    "    for _index in range(1, 8):\n",
    "        if row[f'n{_index}'] == 'Biological Process':\n",
    "            drug_bp_pairs.append((drug_node, {\n",
    "                'idx': node_ids_df.iloc[idx][f'n{_index}'], \n",
    "                'name': node_names_df.iloc[idx][f'n{_index}']\n",
    "            }))\n",
    "\n",
    "len(drug_bp_pairs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Harmonization of nodes to consistent ontology"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For compounds nodes, we will ground them to pubchem nodes and see overlap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 101/101 [01:07<00:00,  1.49it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "94"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drug2mech_harmonized_edges = []\n",
    "drug_id_mapping = {}\n",
    "\n",
    "for cmp_node, bp_node in tqdm(drug_bp_pairs):\n",
    "    compound = get_compounds(cmp_node['name'], 'name')\n",
    "    if len(compound) == 0:\n",
    "        continue\n",
    "\n",
    "    cidx = 'pubchem.compound:' + str(compound[0].cid)\n",
    "    drug_id_mapping[cmp_node['name']] = cidx\n",
    "    \n",
    "    drug2mech_harmonized_edges.append({\n",
    "        'source': cidx,\n",
    "        'source_node_type': 'Compound',\n",
    "        'target': bp_node['idx'],\n",
    "        'target_node_type': 'Biological Process',\n",
    "        'edge_type': 'induces'\n",
    "    })\n",
    "\n",
    "len(drug2mech_harmonized_edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:1983</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0001659</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>pubchem.compound:156391</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>pubchem.compound:4539</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006260</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    source source_node_type      target    target_node_type  \\\n",
       "0    pubchem.compound:5291         Compound  GO:0008283  Biological Process   \n",
       "2    pubchem.compound:1983         Compound  GO:0001659  Biological Process   \n",
       "3    pubchem.compound:2244         Compound  GO:0007596  Biological Process   \n",
       "5  pubchem.compound:156391         Compound  GO:0006954  Biological Process   \n",
       "6    pubchem.compound:4539         Compound  GO:0006260  Biological Process   \n",
       "\n",
       "  edge_type  \n",
       "0   induces  \n",
       "2   induces  \n",
       "3   induces  \n",
       "5   induces  \n",
       "6   induces  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df = pd.DataFrame(drug2mech_harmonized_edges)\n",
    "validation_df.drop_duplicates(inplace=True)\n",
    "validation_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(85, 5)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now, get those which are connected to our KG:\n",
    "\n",
    "We know they're already connected if both the drug and BP are already in the KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:3553</td>\n",
       "      <td>Gene</td>\n",
       "      <td>upregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:203068</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   source source_node_type           target target_node_type  \\\n",
       "0  pubchem.compound:10607         Compound    ncbigene:3553             Gene   \n",
       "1  pubchem.compound:10607         Compound  ncbigene:203068             Gene   \n",
       "\n",
       "       edge_type  \n",
       "0    upregulates  \n",
       "1  downregulates  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "combined_kg = pd.read_csv(f'{KG_DIR}/enriched_kg.tsv', sep='\\t')\n",
    "combined_kg.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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</th>\n",
       "      <th>target</th>\n",
       "      <th>edge_type</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target_node_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>ncbigene:1861</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>ncbigene:6477</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          source         target  edge_type source_node_type target_node_type\n",
       "0  ncbigene:6622  ncbigene:1861  interacts             Gene             Gene\n",
       "1  ncbigene:6622  ncbigene:6477  interacts             Gene             Gene"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "small_kg = pd.read_csv(f'{KG_DIR}/small_enriched_kg.tsv', sep='\\t')\n",
    "small_kg.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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</th>\n",
       "      <th>target</th>\n",
       "      <th>edge_type</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target_node_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>ncbigene:6477</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>GO:0043065</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          source         target     edge_type source_node_type  \\\n",
       "0  ncbigene:6622  ncbigene:6477     interacts             Gene   \n",
       "1  ncbigene:6622     GO:0043065  participates             Gene   \n",
       "\n",
       "     target_node_type  \n",
       "0                Gene  \n",
       "1  Biological Process  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xsmall_kg = pd.read_csv(f'{KG_DIR}/xsmall_enriched_kg.tsv', sep='\\t')\n",
    "xsmall_kg.head(2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to keep track of items for the full, Small, and XSmall variants of MoA-net:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "variants = {'full': dict(),\n",
    "            'small': dict(),\n",
    "            'xsmall': dict()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "variants['full']['bps'] = set(combined_kg[combined_kg['target'].str.startswith('GO')]['target'].tolist())\n",
    "variants['full']['drugs'] = set(combined_kg[combined_kg['source'].str.startswith('pubchem')]['source'].tolist())\n",
    "\n",
    "variants['small']['bps'] = set(small_kg[small_kg['target'].str.startswith('GO')]['target'].tolist())\n",
    "variants['small']['drugs'] = set(small_kg[small_kg['source'].str.startswith('pubchem')]['source'].tolist())\n",
    "\n",
    "variants['xsmall']['bps'] = set(xsmall_kg[xsmall_kg['target'].str.startswith('GO')]['target'].tolist())\n",
    "variants['xsmall']['drugs'] = set(xsmall_kg[xsmall_kg['source'].str.startswith('pubchem')]['source'].tolist())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see which are already within our KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>78</th>\n",
       "      <td>pubchem.compound:2145</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   source source_node_type      target    target_node_type  \\\n",
       "78  pubchem.compound:2145         Compound  GO:0008283  Biological Process   \n",
       "\n",
       "   edge_type  \n",
       "78   induces  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df.loc[(validation_df['source'].isin(variants['full']['drugs'])) & (validation_df['target'].isin(variants['full']['bps']))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>78</th>\n",
       "      <td>pubchem.compound:2145</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   source source_node_type      target    target_node_type  \\\n",
       "78  pubchem.compound:2145         Compound  GO:0008283  Biological Process   \n",
       "\n",
       "   edge_type  \n",
       "78   induces  "
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df.loc[(validation_df['source'].isin(variants['small']['drugs'])) & (validation_df['target'].isin(variants['small']['bps']))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>78</th>\n",
       "      <td>pubchem.compound:2145</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   source source_node_type      target    target_node_type  \\\n",
       "78  pubchem.compound:2145         Compound  GO:0008283  Biological Process   \n",
       "\n",
       "   edge_type  \n",
       "78   induces  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df.loc[(validation_df['source'].isin(variants['xsmall']['drugs'])) & (validation_df['target'].isin(variants['xsmall']['bps']))]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get all the proteins and respective PPIs from the KGs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9301"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_sets = set(combined_kg[combined_kg['target'].str.startswith('ncbigene')]['target'].tolist())\n",
    "source_sets = set(combined_kg[combined_kg['source'].str.startswith('ncbigene')]['source'].tolist())\n",
    "variants['full']['proteins'] = target_sets.union(source_sets)\n",
    "len(variants['full']['proteins'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7446"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_sets = set(small_kg[small_kg['target'].str.startswith('ncbigene')]['target'].tolist())\n",
    "source_sets = set(small_kg[small_kg['source'].str.startswith('ncbigene')]['source'].tolist())\n",
    "variants['small']['proteins'] = target_sets.union(source_sets)\n",
    "len(variants['small']['proteins'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3176"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_sets = set(xsmall_kg[xsmall_kg['target'].str.startswith('ncbigene')]['target'].tolist())\n",
    "source_sets = set(xsmall_kg[xsmall_kg['source'].str.startswith('ncbigene')]['source'].tolist())\n",
    "variants['xsmall']['proteins'] = target_sets.union(source_sets)\n",
    "len(variants['xsmall']['proteins'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need a mapping file to map DrugMechDB proteins to NCBI terminology:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open json file\n",
    "with open(f'{MAPPING_DIR}/uniprot2ncbi.json', 'r') as f:\n",
    "    uniprot2ncbi = json.load(f)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, go through the DrugMechDB df again, and for each row, get the path between the drug and BP.\n",
    "\n",
    "I know this is pretty redundant, but I will just do it for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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>n1</th>\n",
       "      <th>e1</th>\n",
       "      <th>n2</th>\n",
       "      <th>e2</th>\n",
       "      <th>n3</th>\n",
       "      <th>e3</th>\n",
       "      <th>n4</th>\n",
       "      <th>e4</th>\n",
       "      <th>n5</th>\n",
       "      <th>e5</th>\n",
       "      <th>n6</th>\n",
       "      <th>e6</th>\n",
       "      <th>n7</th>\n",
       "      <th>e7</th>\n",
       "      <th>n8</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Protein</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Protein</td>\n",
       "      <td>UP_REGULATES</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Protein</td>\n",
       "      <td>UP_REGULATES</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Pathway</td>\n",
       "      <td>ASSOCIATED_WITH</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Drug</td>\n",
       "      <td>INHIBITS</td>\n",
       "      <td>Protein</td>\n",
       "      <td>PRODUCES</td>\n",
       "      <td>Compound Class</td>\n",
       "      <td>CAUSES</td>\n",
       "      <td>Disease</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     n1        e1       n2               e2                  n3      e3  \\\n",
       "0  Drug  INHIBITS  Protein           CAUSES             Disease     NaN   \n",
       "1  Drug  INHIBITS  Protein     UP_REGULATES  Biological Process  CAUSES   \n",
       "2  Drug  INHIBITS  Protein     UP_REGULATES  Biological Process  CAUSES   \n",
       "3  Drug  INHIBITS  Pathway  ASSOCIATED_WITH             Disease     NaN   \n",
       "4  Drug  INHIBITS  Protein         PRODUCES      Compound Class  CAUSES   \n",
       "\n",
       "        n4   e4   n5   e5   n6   e6   n7   e7   n8  \n",
       "0      NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "1  Disease  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "2  Disease  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "3      NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "4  Disease  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  "
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node_type_df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 123/123 [00:00<00:00, 1794.88it/s]\n"
     ]
    }
   ],
   "source": [
    "validation_moas = dict()\n",
    "drug2prot_relation = dict()\n",
    "\n",
    "for idx, row in tqdm(node_type_df.iterrows(), total=node_type_df.shape[0]):\n",
    "    path_nodes = list()\n",
    "    try:\n",
    "        assert row['n1'] == 'Drug' and row['n2'] == 'Protein'\n",
    "        drug_node = node_names_df.iloc[idx]['n1']\n",
    "        prot_node = node_ids_df.iloc[idx][f'n2']\n",
    "        assert drug_node in drug_id_mapping and prot_node in uniprot2ncbi\n",
    "        drug_node = drug_id_mapping[drug_node]\n",
    "        prot_node = uniprot2ncbi[prot_node]\n",
    "        path_nodes.append(prot_node)\n",
    "\n",
    "        drug2prot_relation[(drug_node, prot_node)] = node_names_df.iloc[idx]['e1']\n",
    "        \n",
    "    except AssertionError:\n",
    "        continue\n",
    "\n",
    "    for _index in range(2, 8):\n",
    "        if row[f'n{_index}'] == 'Protein':\n",
    "            try:\n",
    "                prot_node = uniprot2ncbi[node_ids_df.iloc[idx][f'n{_index}']]\n",
    "                path_nodes.append(prot_node)\n",
    "            except KeyError:\n",
    "                pass\n",
    "        elif row[f'n{_index}'] == 'Biological Process':\n",
    "            bp_node = node_ids_df.iloc[idx][f'n{_index}']\n",
    "            validation_moas[(drug_node, bp_node)] = path_nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(validation_moas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have 48 MOAs which map to ontologies in our KG. That means 48 Drug-BP pairs which can go in our test dataset.\n",
    "\n",
    "Next, I check which of these drug-BP pairs are actually connected to our KG in some way.\n",
    "\n",
    "We will map the drug-protein relation using the following dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "neg_relations = ['INHIBITS', 'DECREASES_EXPRESSION', 'DISRUPTED_IN', 'DISRUPTS', 'DOWN_REGULATES', 'REDUCED_IN', 'REDUCES']\n",
    "\n",
    "pos_relations = ['ACTIVATES', 'ACTIVATED_BY', 'INCREASED_BY', 'INCREASES', 'INCREASES_EXPRESSION', 'UPREGULATED_IN',\n",
    "            'UP_REGULATES', 'ELEVATED_IN', 'STIMULATES']\n",
    "\n",
    "relation_map = {rel: 'upregulates' for rel in pos_relations}\n",
    "for rel in neg_relations:\n",
    "    relation_map[rel] = 'downregulates'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we will see whether connections between drug-BP pairs can be made through the PPIs in the network, since only one pair was fully present in the KG.\n",
    "\n",
    "We have to do it for the full, Small, and XSmall variants of MoA-net, so let's loop through."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in variants:\n",
    "    variants[key]['connected_pairs'] = dict()\n",
    "    variants[key]['other_edges'] = {'source': [],\n",
    "                                    'source_node_type': [],\n",
    "                                    'target': [],\n",
    "                                    'target_node_type': [],\n",
    "                                    'edge_type': []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "for kg in variants.keys():\n",
    "\n",
    "    for key, val in validation_moas.items():\n",
    "        if key[1].startswith('MESH'):\n",
    "            continue\n",
    "        # if both drug and GO BP in graph, great!\n",
    "        if key[0] in variants[kg]['drugs'] and key[1] in variants[kg]['bps']:\n",
    "            variants[kg]['connected_pairs'][key] = val\n",
    "        # if drug in graph but not GO BP, see if the prot connected to GO BP is:\n",
    "        elif key[0] in variants[kg]['drugs'] and len(val) > 0 and val[-1] in variants[kg]['proteins']:\n",
    "            # add the last connection to KG\n",
    "            variants[kg]['other_edges']['source'].append(val[-1])\n",
    "            variants[kg]['other_edges']['source_node_type'].append('Gene')\n",
    "            variants[kg]['other_edges']['target'].append(key[1])\n",
    "            variants[kg]['other_edges']['target_node_type'].append('Biological Process')\n",
    "            variants[kg]['other_edges']['edge_type'].append('participates')\n",
    "            variants[kg]['connected_pairs'][key] = val\n",
    "        # if GO BP in graph but not drug, see if the prot connected to drug is:\n",
    "        elif key[1] in variants[kg]['bps'] and len(val) > 0 and val[0] in variants[kg]['proteins']:\n",
    "            # add the connection to KG\n",
    "            variants[kg]['other_edges']['source'].append(key[0])\n",
    "            variants[kg]['other_edges']['source_node_type'].append('Compound')\n",
    "            variants[kg]['other_edges']['target'].append(val[0])\n",
    "            variants[kg]['other_edges']['target_node_type'].append('Gene')\n",
    "            variants[kg]['other_edges']['edge_type'].append(relation_map[drug2prot_relation[(key[0], val[0])]])\n",
    "            variants[kg]['connected_pairs'][key] = val\n",
    "        # if neither are in the KG, see if the protein(s) connected to them are:\n",
    "        elif len(val) > 0 and val[0] in variants[kg]['proteins'] and val[-1] in variants[kg]['proteins']:\n",
    "            # add BOTH connections to KG\n",
    "            variants[kg]['other_edges']['source'].append(val[-1])\n",
    "            variants[kg]['other_edges']['source_node_type'].append('Gene')\n",
    "            variants[kg]['other_edges']['target'].append(key[1])\n",
    "            variants[kg]['other_edges']['target_node_type'].append('Biological Process')\n",
    "            variants[kg]['other_edges']['edge_type'].append('participates')\n",
    "\n",
    "            variants[kg]['other_edges']['source'].append(key[0])\n",
    "            variants[kg]['other_edges']['source_node_type'].append('Compound')\n",
    "            variants[kg]['other_edges']['target'].append(val[0])\n",
    "            variants[kg]['other_edges']['target_node_type'].append('Gene')\n",
    "            variants[kg]['other_edges']['edge_type'].append(relation_map[drug2prot_relation[(key[0], val[0])]])\n",
    "\n",
    "            variants[kg]['connected_pairs'][key] = val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(variants['full']['connected_pairs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(variants['small']['connected_pairs'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(variants['xsmall']['connected_pairs'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Luckily, all 48 of them are! Now we can create the final KG as well as the final test set.\n",
    "\n",
    "First the validation / test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "for kg in variants:\n",
    "    variants[kg]['drug2process_edges'] = {'source': [],\n",
    "                        'source_node_type': [],\n",
    "                        'target': [],\n",
    "                        'target_node_type': [],\n",
    "                        'edge_type': []}\n",
    "\n",
    "    for key in variants[kg]['connected_pairs'].keys():\n",
    "        variants[kg]['drug2process_edges']['source'].append(key[0])\n",
    "        variants[kg]['drug2process_edges']['source_node_type'].append('Compound')\n",
    "        variants[kg]['drug2process_edges']['target'].append(key[1])\n",
    "        variants[kg]['drug2process_edges']['target_node_type'].append('Biological Process')\n",
    "        variants[kg]['drug2process_edges']['edge_type'].append('induces')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:156391</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pubchem.compound:5745</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:5745</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0050900</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    source source_node_type      target    target_node_type  \\\n",
       "0    pubchem.compound:5291         Compound  GO:0008283  Biological Process   \n",
       "1    pubchem.compound:2244         Compound  GO:0007596  Biological Process   \n",
       "2  pubchem.compound:156391         Compound  GO:0006954  Biological Process   \n",
       "3    pubchem.compound:5745         Compound  GO:0006954  Biological Process   \n",
       "4    pubchem.compound:5745         Compound  GO:0050900  Biological Process   \n",
       "\n",
       "  edge_type  \n",
       "0   induces  \n",
       "1   induces  \n",
       "2   induces  \n",
       "3   induces  \n",
       "4   induces  "
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df = pd.DataFrame(variants['full']['drug2process_edges'])\n",
    "validation_df.drop_duplicates(inplace=True)\n",
    "validation_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(validation_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_df.to_csv(f'{KG_DIR}/drugmechdb_triples.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:156391</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pubchem.compound:5745</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:5745</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0050900</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    source source_node_type      target    target_node_type  \\\n",
       "0    pubchem.compound:5291         Compound  GO:0008283  Biological Process   \n",
       "1    pubchem.compound:2244         Compound  GO:0007596  Biological Process   \n",
       "2  pubchem.compound:156391         Compound  GO:0006954  Biological Process   \n",
       "3    pubchem.compound:5745         Compound  GO:0006954  Biological Process   \n",
       "4    pubchem.compound:5745         Compound  GO:0050900  Biological Process   \n",
       "\n",
       "  edge_type  \n",
       "0   induces  \n",
       "1   induces  \n",
       "2   induces  \n",
       "3   induces  \n",
       "4   induces  "
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df = pd.DataFrame(variants['small']['drug2process_edges'])\n",
    "validation_df.drop_duplicates(inplace=True)\n",
    "validation_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "48"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(validation_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_df.to_csv(f'{KG_DIR}/drugmechdb_triples_small.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0008283</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:5745</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0050900</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pubchem.compound:5743</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:2200</td>\n",
       "      <td>Compound</td>\n",
       "      <td>GO:0034776</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>induces</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  source source_node_type      target    target_node_type  \\\n",
       "0  pubchem.compound:5291         Compound  GO:0008283  Biological Process   \n",
       "1  pubchem.compound:2244         Compound  GO:0007596  Biological Process   \n",
       "2  pubchem.compound:5745         Compound  GO:0050900  Biological Process   \n",
       "3  pubchem.compound:5743         Compound  GO:0006954  Biological Process   \n",
       "4  pubchem.compound:2200         Compound  GO:0034776  Biological Process   \n",
       "\n",
       "  edge_type  \n",
       "0   induces  \n",
       "1   induces  \n",
       "2   induces  \n",
       "3   induces  \n",
       "4   induces  "
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df = pd.DataFrame(variants['xsmall']['drug2process_edges'])\n",
    "validation_df.drop_duplicates(inplace=True)\n",
    "validation_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(validation_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_df.to_csv(f'{KG_DIR}/drugmechdb_triples_xsmall.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, the final KGs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5156</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:5742</td>\n",
       "      <td>Gene</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>participates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5742</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ncbigene:5743</td>\n",
       "      <td>Gene</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>participates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:156391</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5743</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    source source_node_type         target  \\\n",
       "0    pubchem.compound:5291         Compound  ncbigene:5156   \n",
       "1            ncbigene:5742             Gene     GO:0007596   \n",
       "2    pubchem.compound:2244         Compound  ncbigene:5742   \n",
       "3            ncbigene:5743             Gene     GO:0006954   \n",
       "4  pubchem.compound:156391         Compound  ncbigene:5743   \n",
       "\n",
       "     target_node_type      edge_type  \n",
       "0                Gene  downregulates  \n",
       "1  Biological Process   participates  \n",
       "2                Gene  downregulates  \n",
       "3  Biological Process   participates  \n",
       "4                Gene  downregulates  "
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kg_data = pd.DataFrame(variants['full']['other_edges'])\n",
    "\n",
    "kg_data.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:3553</td>\n",
       "      <td>Gene</td>\n",
       "      <td>upregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:203068</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:54658</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:7153</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:10607</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:7277</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   source source_node_type           target target_node_type  \\\n",
       "0  pubchem.compound:10607         Compound    ncbigene:3553             Gene   \n",
       "1  pubchem.compound:10607         Compound  ncbigene:203068             Gene   \n",
       "2  pubchem.compound:10607         Compound   ncbigene:54658             Gene   \n",
       "3  pubchem.compound:10607         Compound    ncbigene:7153             Gene   \n",
       "4  pubchem.compound:10607         Compound    ncbigene:7277             Gene   \n",
       "\n",
       "       edge_type  \n",
       "0    upregulates  \n",
       "1  downregulates  \n",
       "2  downregulates  \n",
       "3  downregulates  \n",
       "4  downregulates  "
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Append to the existing KG\n",
    "final_kg_df = pd.concat([combined_kg, kg_data], ignore_index=True)\n",
    "final_kg_df.drop_duplicates(inplace=True)\n",
    "final_kg_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_kg_df.to_csv(f'{KG_DIR}/final_kg.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Final Small KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5156</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:5742</td>\n",
       "      <td>Gene</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>participates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5742</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ncbigene:5743</td>\n",
       "      <td>Gene</td>\n",
       "      <td>GO:0006954</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>participates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:156391</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5743</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                    source source_node_type         target  \\\n",
       "0    pubchem.compound:5291         Compound  ncbigene:5156   \n",
       "1            ncbigene:5742             Gene     GO:0007596   \n",
       "2    pubchem.compound:2244         Compound  ncbigene:5742   \n",
       "3            ncbigene:5743             Gene     GO:0006954   \n",
       "4  pubchem.compound:156391         Compound  ncbigene:5743   \n",
       "\n",
       "     target_node_type      edge_type  \n",
       "0                Gene  downregulates  \n",
       "1  Biological Process   participates  \n",
       "2                Gene  downregulates  \n",
       "3  Biological Process   participates  \n",
       "4                Gene  downregulates  "
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kg_data = pd.DataFrame(variants['small']['other_edges'])\n",
    "\n",
    "kg_data.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "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</th>\n",
       "      <th>target</th>\n",
       "      <th>edge_type</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target_node_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>ncbigene:1861</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>ncbigene:6477</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>GO:0043065</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ncbigene:3308</td>\n",
       "      <td>ncbigene:6767</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>ncbigene:3308</td>\n",
       "      <td>GO:0006457</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          source         target     edge_type source_node_type  \\\n",
       "0  ncbigene:6622  ncbigene:1861     interacts             Gene   \n",
       "1  ncbigene:6622  ncbigene:6477     interacts             Gene   \n",
       "2  ncbigene:6622     GO:0043065  participates             Gene   \n",
       "3  ncbigene:3308  ncbigene:6767     interacts             Gene   \n",
       "4  ncbigene:3308     GO:0006457  participates             Gene   \n",
       "\n",
       "     target_node_type  \n",
       "0                Gene  \n",
       "1                Gene  \n",
       "2  Biological Process  \n",
       "3                Gene  \n",
       "4  Biological Process  "
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Append to the existing KG\n",
    "final_kg_df = pd.concat([small_kg, kg_data], ignore_index=True)\n",
    "final_kg_df.drop_duplicates(inplace=True)\n",
    "final_kg_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_kg_df.to_csv(f'{KG_DIR}/small_kg.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finall XSmall KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "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</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target</th>\n",
       "      <th>target_node_type</th>\n",
       "      <th>edge_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>pubchem.compound:5291</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5156</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:5742</td>\n",
       "      <td>Gene</td>\n",
       "      <td>GO:0007596</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>participates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>pubchem.compound:2244</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:5742</td>\n",
       "      <td>Gene</td>\n",
       "      <td>downregulates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ncbigene:301</td>\n",
       "      <td>Gene</td>\n",
       "      <td>GO:0050900</td>\n",
       "      <td>Biological Process</td>\n",
       "      <td>participates</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>pubchem.compound:5745</td>\n",
       "      <td>Compound</td>\n",
       "      <td>ncbigene:301</td>\n",
       "      <td>Gene</td>\n",
       "      <td>upregulates</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  source source_node_type         target    target_node_type  \\\n",
       "0  pubchem.compound:5291         Compound  ncbigene:5156                Gene   \n",
       "1          ncbigene:5742             Gene     GO:0007596  Biological Process   \n",
       "2  pubchem.compound:2244         Compound  ncbigene:5742                Gene   \n",
       "3           ncbigene:301             Gene     GO:0050900  Biological Process   \n",
       "4  pubchem.compound:5745         Compound   ncbigene:301                Gene   \n",
       "\n",
       "       edge_type  \n",
       "0  downregulates  \n",
       "1   participates  \n",
       "2  downregulates  \n",
       "3   participates  \n",
       "4    upregulates  "
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "kg_data = pd.DataFrame(variants['xsmall']['other_edges'])\n",
    "\n",
    "kg_data.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "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</th>\n",
       "      <th>target</th>\n",
       "      <th>edge_type</th>\n",
       "      <th>source_node_type</th>\n",
       "      <th>target_node_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>ncbigene:6477</td>\n",
       "      <td>interacts</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Gene</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>ncbigene:6622</td>\n",
       "      <td>GO:0043065</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>ncbigene:3308</td>\n",
       "      <td>GO:0006457</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>ncbigene:2885</td>\n",
       "      <td>GO:0007165</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>ncbigene:5295</td>\n",
       "      <td>GO:0050821</td>\n",
       "      <td>participates</td>\n",
       "      <td>Gene</td>\n",
       "      <td>Biological Process</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          source         target     edge_type source_node_type  \\\n",
       "0  ncbigene:6622  ncbigene:6477     interacts             Gene   \n",
       "1  ncbigene:6622     GO:0043065  participates             Gene   \n",
       "2  ncbigene:3308     GO:0006457  participates             Gene   \n",
       "3  ncbigene:2885     GO:0007165  participates             Gene   \n",
       "4  ncbigene:5295     GO:0050821  participates             Gene   \n",
       "\n",
       "     target_node_type  \n",
       "0                Gene  \n",
       "1  Biological Process  \n",
       "2  Biological Process  \n",
       "3  Biological Process  \n",
       "4  Biological Process  "
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Append to the existing KG\n",
    "final_kg_df = pd.concat([xsmall_kg, kg_data], ignore_index=True)\n",
    "final_kg_df.drop_duplicates(inplace=True)\n",
    "final_kg_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_kg_df.to_csv(f'{KG_DIR}/xsmall_kg.tsv', sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nesy_venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}