notebooks/4_train_test_splits.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train - Test - Validation Splits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import os.path as osp\n",
    "import os\n",
    "import networkx as nx\n",
    "from copy import deepcopy"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load in the final KGs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "KG_DIR = '../data/kg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "kg = pd.read_csv(osp.join(KG_DIR, 'final_kg.tsv'), sep='\\t')\n",
    "kg.drop_duplicates(inplace=True)\n",
    "\n",
    "kg_protclass = pd.read_csv(osp.join(KG_DIR, 'final_kg_subclassed.tsv'), sep='\\t')\n",
    "kg_protclass.drop_duplicates(subset=['source', 'target'], inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "drug_bp_pairs = kg.loc[kg['edge_type'] == 'induces']\n",
    "\n",
    "drug_bp_pairs_protclass = kg_protclass.loc[kg_protclass['edge_type'] == 'induces']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 1622 positive drug-BP pairs in the final KG\n",
      "There are 1622 positive drug-BP pairs in the protein-classed KG\n"
     ]
    }
   ],
   "source": [
    "print(f\"There are {len(drug_bp_pairs)} positive drug-BP pairs in the final KG\")\n",
    "\n",
    "print(f\"There are {len(drug_bp_pairs_protclass)} positive drug-BP pairs in the protein-classed KG\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load in the DrugMechDB pairs that go in the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48 additional drug-BP pairs come from DrugMechDB, constituting 2.9593094944512948%\n"
     ]
    }
   ],
   "source": [
    "dm_db_pairs = pd.read_csv(osp.join(KG_DIR, 'drugmechdb_triples.tsv'), sep='\\t')\n",
    "dm_db_pairs.drop_duplicates(inplace=True)\n",
    "\n",
    "print(f\"{len(dm_db_pairs)} additional drug-BP pairs come from DrugMechDB, constituting {len(dm_db_pairs)/len(drug_bp_pairs)*100}%\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specify which pairs are able to be matched:\n",
    "\n",
    "MARS / MARS can match pairs as long as they are 4 or less hops apart (our hyperparameter setting) and not connected via an inverse _CtBP edge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = nx.DiGraph()\n",
    "\n",
    "for i, row in kg.iterrows():\n",
    "    if row['edge_type'] == 'induces':\n",
    "        continue\n",
    "    src_id = row['source']\n",
    "    trgt_id = row['target']\n",
    "    if src_id not in G.nodes:\n",
    "        G.add_node(src_id, type=row['source_node_type'])\n",
    "    if trgt_id not in G.nodes:\n",
    "        G.add_node(trgt_id, type=row['target_node_type'])\n",
    "    G.add_edge(src_id, trgt_id, type=row['edge_type'])\n",
    "\n",
    "Gp = nx.DiGraph()\n",
    "\n",
    "for i, row in kg_protclass.iterrows():\n",
    "    if row['edge_type'] == 'induces':\n",
    "        continue\n",
    "    src_id = row['source']\n",
    "    trgt_id = row['target']\n",
    "    if src_id not in Gp.nodes:\n",
    "        Gp.add_node(src_id, type=row['source_node_type'])\n",
    "    if trgt_id not in Gp.nodes:\n",
    "        Gp.add_node(trgt_id, type=row['target_node_type'])\n",
    "    Gp.add_edge(src_id, trgt_id, type=row['edge_type'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the unmatched pairs for the regular, full KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "unmatched_pairs = set()\n",
    "unmatched_dm_db_pairs = set()\n",
    "\n",
    "path_lengths = dict()\n",
    "\n",
    "for i, row in drug_bp_pairs.iterrows():\n",
    "    if not nx.has_path(G, row['source'], row['target']):\n",
    "        unmatched_pairs.add(i)\n",
    "    elif 1 > nx.shortest_path_length(G, row['source'], row['target']) > 4:\n",
    "        unmatched_pairs.add(i)\n",
    "    else:\n",
    "        path_lengths[f\"{(row['source'], row['target'])}\"] = nx.shortest_path_length(G, row['source'], row['target'])\n",
    "\n",
    "for i, row in dm_db_pairs.iterrows():\n",
    "    if not nx.has_path(G, row['source'], row['target']):\n",
    "        unmatched_dm_db_pairs.add(i)\n",
    "    elif 1 > nx.shortest_path_length(G, row['source'], row['target']) > 4:\n",
    "        unmatched_dm_db_pairs.add(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "26 drug-BP pairs in the final KG are not connected via 4 hops or less\n",
      "0 drug-BP pairs from DrugMechDB are not connected via 4 hops or less\n"
     ]
    }
   ],
   "source": [
    "print(f\"{len(unmatched_pairs)} drug-BP pairs in the final KG are not connected via 4 hops or less\")\n",
    "print(f\"{len(unmatched_dm_db_pairs)} drug-BP pairs from DrugMechDB are not connected via 4 hops or less\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the unmatched pairs for the protein-subclassed KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "unmatched_pairs_p = set()\n",
    "unmatched_dm_db_pairs_p = set()\n",
    "\n",
    "path_lengths_p = dict()\n",
    "\n",
    "for i, row in drug_bp_pairs_protclass.iterrows():\n",
    "    if not row['source'] in Gp.nodes or not row['target'] in Gp.nodes:\n",
    "        unmatched_pairs_p.add(i)\n",
    "    elif not nx.has_path(Gp, row['source'], row['target']):\n",
    "        unmatched_pairs_p.add(i)\n",
    "    elif 1 > nx.shortest_path_length(Gp, row['source'], row['target']) > 4:\n",
    "        unmatched_pairs_p.add(i)\n",
    "    else:\n",
    "        path_lengths_p[f\"{(row['source'], row['target'])}\"] = nx.shortest_path_length(Gp, row['source'], row['target'])\n",
    "\n",
    "for i, row in dm_db_pairs.iterrows():\n",
    "    if not row['source'] in Gp.nodes or not row['target'] in Gp.nodes:\n",
    "        unmatched_dm_db_pairs_p.add(i)\n",
    "    elif not nx.has_path(Gp, row['source'], row['target']):\n",
    "        unmatched_dm_db_pairs_p.add(i)\n",
    "    elif 1 > nx.shortest_path_length(Gp, row['source'], row['target']) > 4:\n",
    "        unmatched_dm_db_pairs_p.add(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "693 drug-BP pairs in the protein-classed KG are not connected via 4 hops or less\n",
      "33 drug-BP pairs from DrugMechDB are not connected via 4 hops or less\n"
     ]
    }
   ],
   "source": [
    "print(f\"{len(unmatched_pairs_p)} drug-BP pairs in the protein-classed KG are not connected via 4 hops or less\")\n",
    "print(f\"{len(unmatched_dm_db_pairs_p)} drug-BP pairs from DrugMechDB are not connected via 4 hops or less\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_299297/2091281768.py:2: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  drug_bp_pairs.drop(list(unmatched_pairs), inplace=True)\n",
      "/tmp/ipykernel_299297/2091281768.py:6: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  drug_bp_pairs_protclass.drop(list(unmatched_pairs_p), inplace=True)\n"
     ]
    }
   ],
   "source": [
    "# Remove unmatched pairs from the final KG\n",
    "drug_bp_pairs.drop(list(unmatched_pairs), inplace=True)\n",
    "drug_bp_pairs.reset_index(drop=True, inplace=True)\n",
    "\n",
    "# Remove unmatched pairs from the protein-classed KG\n",
    "drug_bp_pairs_protclass.drop(list(unmatched_pairs_p), inplace=True)\n",
    "drug_bp_pairs_protclass.reset_index(drop=True, inplace=True)\n",
    "\n",
    "# Remove unmatched pairs from DrugMechDB\n",
    "# make copies of it for multiple KG variants:\n",
    "dm_db_pairs_protclass = deepcopy(dm_db_pairs)\n",
    "\n",
    "dm_db_pairs.drop(list(unmatched_dm_db_pairs), inplace=True)\n",
    "dm_db_pairs.reset_index(drop=True, inplace=True)\n",
    "\n",
    "dm_db_pairs_protclass.drop(list(unmatched_dm_db_pairs_p), inplace=True)\n",
    "dm_db_pairs_protclass.reset_index(drop=True, inplace=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add their categories (path lengths) in:\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the Splits"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the original PoLo example, they use different proportions for the splits, but let's go with something most similar to their Hetionet example, in which they do an approximate 60/20/20% split.\n",
    "\n",
    "Note that the DrugMechDB examples need to be in the test set.\n",
    "\n",
    "First, we'll exclude the subset of DrugMechDB examples that are in the test set.\n",
    "\n",
    "Then, we'll split the remaining examples into train, validation, and test sets, with the test set accumulating to 20% with the DrugMechDB examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_positives = len(drug_bp_pairs) + len(dm_db_pairs)\n",
    "\n",
    "total_positives_pc = len(drug_bp_pairs_protclass) + len(dm_db_pairs_protclass)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we need to get the following numbers from the KG positive examples for train, validation, and test sets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(986, 329, 281)\n",
      "(566, 189, 174)\n"
     ]
    }
   ],
   "source": [
    "proportions = round(0.6 * total_positives), round(0.2 * total_positives), round(0.2 * total_positives) - len(dm_db_pairs)\n",
    "print(proportions)\n",
    "\n",
    "proportions_pc = round(0.6 * total_positives_pc), round(0.2 * total_positives_pc), round(0.2 * total_positives_pc) - len(dm_db_pairs_protclass)\n",
    "print(proportions_pc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write a function which separates the dataframe into train, val and test sets of defined sizes\n",
    "def train_test_split(df, train_size, val_size, test_size):\n",
    "    df = df.sample(frac=1, random_state=7).reset_index(drop=True)\n",
    "    train = df[:train_size]\n",
    "    val = df[train_size:train_size+val_size]\n",
    "    test = df[train_size+val_size:train_size+val_size+test_size]\n",
    "    return train, val, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, val, test = train_test_split(drug_bp_pairs, proportions[0], proportions[1], proportions[2])\n",
    "\n",
    "train_pc, val_pc, test_pc = train_test_split(drug_bp_pairs_protclass, proportions_pc[0], proportions_pc[1], proportions_pc[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check it did what we want:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(986, 329, 281)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train), len(val), len(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(566, 189, 174)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_pc), len(val_pc), len(test_pc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No overlap?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_pairs = {(row['source'], row['target']) for i, row in train.iterrows()}\n",
    "test_pairs = {(row['source'], row['target']) for i, row in test.iterrows()}\n",
    "val_pairs = {(row['source'], row['target']) for i, row in val.iterrows()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "set()"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_pairs & test_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "set()"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_pairs & val_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "set()"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_pairs & val_pairs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Overlap in the protein-classed one?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_pairs = {(row['source'], row['target']) for i, row in train_pc.iterrows()}\n",
    "test_pairs = {(row['source'], row['target']) for i, row in test_pc.iterrows()}\n",
    "val_pairs = {(row['source'], row['target']) for i, row in val_pc.iterrows()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "set()"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_pairs & test_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "set()"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_pairs & val_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "set()"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_pairs & val_pairs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Good, no overlap. Add the DrugMechDB examples to the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = pd.concat([test, dm_db_pairs]).sample(frac=1, random_state=7).reset_index(drop=True)\n",
    "\n",
    "test_pc = pd.concat([test_pc, dm_db_pairs_protclass]).sample(frac=1, random_state=7).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "329\n",
      "189\n"
     ]
    }
   ],
   "source": [
    "print(len(test))\n",
    "print(len(test_pc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take the test and validation sets out the KG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "kg = kg.loc[kg['edge_type'] != 'induces']\n",
    "kg_mars = pd.concat([kg, train]).sample(frac=1, random_state=7).reset_index(drop=True)\n",
    "\n",
    "kg_protclass = kg_protclass.loc[kg_protclass['edge_type'] != 'induces']\n",
    "kg_mars_pc = pd.concat([kg_protclass, train_pc]).sample(frac=1, random_state=7).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(kg_mars.loc[kg_mars['edge_type'] == 'induces']) == len(train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(kg_mars_pc.loc[kg_mars_pc['edge_type'] == 'induces']) == len(train_pc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write everything to files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "SPLITS_DIR = osp.join(KG_DIR, 'splits')\n",
    "MOA_NET = osp.join(SPLITS_DIR, 'MoA-net')\n",
    "MOA_NET_PROTCLASSED = osp.join(SPLITS_DIR, 'MoA-net-protclass')\n",
    "\n",
    "# create the above directories if they do not exist:\n",
    "for dir in [SPLITS_DIR, MOA_NET, MOA_NET_PROTCLASSED]:\n",
    "    if not osp.exists(dir):\n",
    "        os.makedirs(dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "kg.to_csv(osp.join(MOA_NET, 'kg_no_cmp_bp.tsv'), sep='\\t', index=False)\n",
    "kg_mars.to_csv(osp.join(MOA_NET, 'kg_with_train_smpls.tsv'), sep='\\t', index=False)\n",
    "\n",
    "train.to_csv(osp.join(MOA_NET, 'train.tsv'), sep='\\t', index=False)\n",
    "val.to_csv(osp.join(MOA_NET, 'dev.tsv'), sep='\\t', index=False)\n",
    "test.to_csv(osp.join(MOA_NET, 'test.tsv'), sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "kg_protclass.to_csv(osp.join(MOA_NET_PROTCLASSED, 'kg_no_cmp_bp.tsv'), sep='\\t', index=False)\n",
    "kg_mars_pc.to_csv(osp.join(MOA_NET_PROTCLASSED, 'kg_with_train_smpls.tsv'), sep='\\t', index=False)\n",
    "\n",
    "train_pc.to_csv(osp.join(MOA_NET_PROTCLASSED, 'train.tsv'), sep='\\t', index=False)\n",
    "val_pc.to_csv(osp.join(MOA_NET_PROTCLASSED, 'dev.tsv'), sep='\\t', index=False)\n",
    "test_pc.to_csv(osp.join(MOA_NET_PROTCLASSED, 'test.tsv'), sep='\\t', index=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's make an ablation test set for MoA-net to ensure the improved metrics aren't just coming from a smaller test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_test = test.sample(100, random_state=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_test.to_csv(osp.join(MOA_NET, 'sub_test.tsv'), sep='\\t', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "moanet_env",
   "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
}