graph-transe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Translational embeddings (TransE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters\n",
    "Define hyper-parameters inclduing the evaluation information:\n",
    "\n",
    "This is the basic set of parameters for a TransE model\n",
    " \n",
    "Unless you have a powerful GPU or use a different dataset than the one provided, set the number of epochs to a low number (5 or 10) for this tutorial; when you want to use and apply the model for a prediction task later, train for more epochs (ideally, on a good GPU)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "org_id = '9606' # 4932 - yeast, 9606 - human\n",
    "\n",
    "embeddings_size = 50\n",
    "batch_size = 32\n",
    "margin = 0.1\n",
    "reg_norm = 1\n",
    "optimizer = 'adam'\n",
    "loss_function = 'mse'\n",
    "epochs = 100\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kulmanm/KAUST/CBRC/ontology-tutorial/venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py:112: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  \"Converting sparse IndexedSlices to a dense Tensor of unknown shape. \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100\n",
      "26023/26023 [==============================] - 291s 11ms/step - loss: 114.0954 - val_loss: 18.6479\n",
      "Epoch 2/100\n",
      "26023/26023 [==============================] - 265s 10ms/step - loss: 21.8254 - val_loss: 4.5449\n",
      "Epoch 3/100\n",
      "26023/26023 [==============================] - 257s 10ms/step - loss: 5.7751 - val_loss: 1.1735\n",
      "Epoch 4/100\n",
      "26023/26023 [==============================] - 249s 10ms/step - loss: 1.5106 - val_loss: 0.4180\n",
      "Epoch 5/100\n",
      "26023/26023 [==============================] - 246s 9ms/step - loss: 0.7002 - val_loss: 0.2671\n",
      "Epoch 6/100\n",
      "26023/26023 [==============================] - 241s 9ms/step - loss: 0.5357 - val_loss: 0.2256\n",
      "Epoch 7/100\n",
      "26023/26023 [==============================] - 243s 9ms/step - loss: 0.3710 - val_loss: 0.2029\n",
      "Epoch 8/100\n",
      "26023/26023 [==============================] - 244s 9ms/step - loss: 0.2403 - val_loss: 0.1865\n",
      "Epoch 9/100\n",
      "26023/26023 [==============================] - 246s 9ms/step - loss: 0.2016 - val_loss: 0.1826\n",
      "Epoch 10/100\n",
      "26023/26023 [==============================] - 244s 9ms/step - loss: 0.1744 - val_loss: 0.1659\n",
      "Epoch 11/100\n",
      "26023/26023 [==============================] - 253s 10ms/step - loss: 0.1606 - val_loss: 0.1716\n",
      "Epoch 12/100\n",
      "26023/26023 [==============================] - 250s 10ms/step - loss: 0.1478 - val_loss: 0.1606\n",
      "Epoch 13/100\n",
      "26023/26023 [==============================] - 248s 10ms/step - loss: 0.1374 - val_loss: 0.1585\n",
      "Epoch 14/100\n",
      "26023/26023 [==============================] - 251s 10ms/step - loss: 0.1292 - val_loss: 0.1526\n",
      "Epoch 15/100\n",
      "26023/26023 [==============================] - 249s 10ms/step - loss: 0.1248 - val_loss: 0.1452\n",
      "Epoch 16/100\n",
      "26023/26023 [==============================] - 250s 10ms/step - loss: 0.1192 - val_loss: 0.1451\n",
      "Epoch 17/100\n",
      "26023/26023 [==============================] - 245s 9ms/step - loss: 0.1159 - val_loss: 0.1392\n",
      "Epoch 18/100\n",
      "26023/26023 [==============================] - 254s 10ms/step - loss: 0.1117 - val_loss: 0.1362\n",
      "Epoch 19/100\n",
      "26023/26023 [==============================] - 243s 9ms/step - loss: 0.1084 - val_loss: 0.1430\n",
      "Epoch 20/100\n",
      "26023/26023 [==============================] - 242s 9ms/step - loss: 0.1063 - val_loss: 0.1303\n",
      "Epoch 21/100\n",
      "26023/26023 [==============================] - 241s 9ms/step - loss: 0.1027 - val_loss: 0.1310\n",
      "Epoch 22/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.1005 - val_loss: 0.1301\n",
      "Epoch 23/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0972 - val_loss: 0.1293\n",
      "Epoch 24/100\n",
      "26023/26023 [==============================] - 239s 9ms/step - loss: 0.0967 - val_loss: 0.1260\n",
      "Epoch 25/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0952 - val_loss: 0.1284\n",
      "Epoch 26/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0933 - val_loss: 0.1253\n",
      "Epoch 27/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0923 - val_loss: 0.1241\n",
      "Epoch 28/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0911 - val_loss: 0.1248\n",
      "Epoch 29/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0903 - val_loss: 0.1251\n",
      "Epoch 30/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0894 - val_loss: 0.1236\n",
      "Epoch 31/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0869 - val_loss: 0.1202\n",
      "Epoch 32/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0868 - val_loss: 0.1275\n",
      "Epoch 33/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0845 - val_loss: 0.1199\n",
      "Epoch 34/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0838 - val_loss: 0.1222\n",
      "Epoch 35/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0836 - val_loss: 0.1241\n",
      "Epoch 36/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0824 - val_loss: 0.1200\n",
      "Epoch 37/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0814 - val_loss: 0.1178\n",
      "Epoch 38/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0813 - val_loss: 0.1193\n",
      "Epoch 39/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0807 - val_loss: 0.1190\n",
      "Epoch 40/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0786 - val_loss: 0.1232\n",
      "Epoch 41/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0790 - val_loss: 0.1181\n",
      "Epoch 42/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0780 - val_loss: 0.1187\n",
      "Epoch 43/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0778 - val_loss: 0.1189\n",
      "Epoch 44/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0773 - val_loss: 0.1141\n",
      "Epoch 45/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0755 - val_loss: 0.1182\n",
      "Epoch 46/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0765 - val_loss: 0.1200\n",
      "Epoch 47/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0739 - val_loss: 0.1122\n",
      "Epoch 48/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0755 - val_loss: 0.1175\n",
      "Epoch 49/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0753 - val_loss: 0.1155\n",
      "Epoch 50/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0746 - val_loss: 0.1156\n",
      "Epoch 51/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0740 - val_loss: 0.1166\n",
      "Epoch 52/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0728 - val_loss: 0.1139\n",
      "Epoch 53/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0735 - val_loss: 0.1127\n",
      "Epoch 54/100\n",
      "26023/26023 [==============================] - 239s 9ms/step - loss: 0.0723 - val_loss: 0.1092\n",
      "Epoch 55/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0726 - val_loss: 0.1086\n",
      "Epoch 56/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0730 - val_loss: 0.1132\n",
      "Epoch 57/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0716 - val_loss: 0.1139\n",
      "Epoch 58/100\n",
      "26023/26023 [==============================] - 239s 9ms/step - loss: 0.0706 - val_loss: 0.1095\n",
      "Epoch 59/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0706 - val_loss: 0.1115\n",
      "Epoch 60/100\n",
      "26023/26023 [==============================] - 239s 9ms/step - loss: 0.0708 - val_loss: 0.1083\n",
      "Epoch 61/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0701 - val_loss: 0.1095\n",
      "Epoch 62/100\n",
      "26023/26023 [==============================] - 239s 9ms/step - loss: 0.0713 - val_loss: 0.1093\n",
      "Epoch 63/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0693 - val_loss: 0.1097\n",
      "Epoch 64/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0697 - val_loss: 0.1096\n",
      "Epoch 65/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0689 - val_loss: 0.1111\n",
      "Epoch 66/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0678 - val_loss: 0.1131\n",
      "Epoch 67/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0691 - val_loss: 0.1123\n",
      "Epoch 68/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0681 - val_loss: 0.1139\n",
      "Epoch 69/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0681 - val_loss: 0.1119\n",
      "Epoch 70/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0677 - val_loss: 0.1099\n",
      "Epoch 71/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0677 - val_loss: 0.1092\n",
      "Epoch 72/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0675 - val_loss: 0.1105\n",
      "Epoch 73/100\n",
      "26023/26023 [==============================] - 240s 9ms/step - loss: 0.0672 - val_loss: 0.1119\n",
      "Epoch 74/100\n",
      "26023/26023 [==============================] - 240s 9ms/step - loss: 0.0660 - val_loss: 0.1090\n",
      "Epoch 75/100\n",
      "26023/26023 [==============================] - 243s 9ms/step - loss: 0.0665 - val_loss: 0.1057\n",
      "Epoch 76/100\n",
      "26023/26023 [==============================] - 240s 9ms/step - loss: 0.0662 - val_loss: 0.1100\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 77/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0661 - val_loss: 0.1110\n",
      "Epoch 78/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0671 - val_loss: 0.1129\n",
      "Epoch 79/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0663 - val_loss: 0.1056\n",
      "Epoch 80/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0658 - val_loss: 0.1070\n",
      "Epoch 81/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0662 - val_loss: 0.1098\n",
      "Epoch 82/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0656 - val_loss: 0.1081\n",
      "Epoch 83/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0663 - val_loss: 0.1077\n",
      "Epoch 84/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0657 - val_loss: 0.1054\n",
      "Epoch 85/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0648 - val_loss: 0.1060\n",
      "Epoch 86/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0649 - val_loss: 0.1063\n",
      "Epoch 87/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0646 - val_loss: 0.1039\n",
      "Epoch 88/100\n",
      "26023/26023 [==============================] - 240s 9ms/step - loss: 0.0637 - val_loss: 0.1075\n",
      "Epoch 89/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0656 - val_loss: 0.1061\n",
      "Epoch 90/100\n",
      "26023/26023 [==============================] - 240s 9ms/step - loss: 0.0639 - val_loss: 0.1073\n",
      "Epoch 91/100\n",
      "26023/26023 [==============================] - 236s 9ms/step - loss: 0.0649 - val_loss: 0.1061\n",
      "Epoch 92/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0632 - val_loss: 0.1111\n",
      "Epoch 93/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0642 - val_loss: 0.1079\n",
      "Epoch 94/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0640 - val_loss: 0.1063\n",
      "Epoch 95/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0640 - val_loss: 0.1053\n",
      "Epoch 96/100\n",
      "26023/26023 [==============================] - 238s 9ms/step - loss: 0.0640 - val_loss: 0.1035\n",
      "Epoch 97/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0635 - val_loss: 0.1069\n",
      "Epoch 98/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0641 - val_loss: 0.1054\n",
      "Epoch 99/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0633 - val_loss: 0.1083\n",
      "Epoch 100/100\n",
      "26023/26023 [==============================] - 237s 9ms/step - loss: 0.0628 - val_loss: 0.1053\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7f1c74245a20>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import Input\n",
    "from tensorflow.keras import optimizers\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, CSVLogger\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from rdflib import Graph\n",
    "import math\n",
    "import os\n",
    "\n",
    "\n",
    "class TransE(tf.keras.Model):\n",
    "    \n",
    "    def __init__(self, nb_classes, nb_relations, embedding_size, margin=0.1, reg_norm=1):\n",
    "        super(TransE, self).__init__()\n",
    "        self.nb_classes = nb_classes\n",
    "        self.nb_relations = nb_relations\n",
    "        self.margin = margin\n",
    "        self.reg_norm = 1\n",
    "        bound = 6 / math.sqrt(embedding_size)\n",
    "        cls_weights = np.random.uniform(low=-bound, high=bound, size=(nb_classes, embedding_size))\n",
    "        rel_weights = np.random.uniform(low=-bound, high=bound, size=(nb_relations, embedding_size))\n",
    "        self.cls_embeddings = tf.keras.layers.Embedding(\n",
    "            nb_classes,\n",
    "            embedding_size,\n",
    "            input_length=1,\n",
    "            weights=[cls_weights,])\n",
    "        self.rel_embeddings = tf.keras.layers.Embedding(\n",
    "            nb_relations,\n",
    "            embedding_size,\n",
    "            input_length=1,\n",
    "            weights=[rel_weights,])\n",
    "    \n",
    "    def reg(self, x):\n",
    "        res = tf.abs(tf.norm(x, axis=1) - self.reg_norm)\n",
    "        res = tf.reshape(res, [-1, 1])\n",
    "        return res\n",
    "    \n",
    "        \n",
    "    def call(self, input):\n",
    "        pos, neg = input\n",
    "        ph, pl, pt = pos[:, 0], pos[:, 1], pos[:, 2]\n",
    "        nh, nl, nt = neg[:, 0], neg[:, 1], neg[:, 2]\n",
    "        nh = self.cls_embeddings(nh)\n",
    "        nl = self.rel_embeddings(nl)\n",
    "        nt = self.cls_embeddings(nt)\n",
    "        ph = self.cls_embeddings(ph)\n",
    "        pl = self.rel_embeddings(pl)\n",
    "        pt = self.cls_embeddings(pt)\n",
    "        \n",
    "        pos = tf.subtract(tf.add(ph, pl), pt)\n",
    "        neg = tf.subtract(tf.add(nh, nl), nt)\n",
    "        \n",
    "        self._loss = tf.norm(pos, axis=1) - tf.norm(neg, axis=1) + self.margin\n",
    "        self._loss = tf.reduce_sum(tf.nn.relu(self._loss))\n",
    "        self._loss = (tf.reshape(self._loss, [-1, 1]) + self.reg(ph)\n",
    "                      + self.reg(ph) + self.reg(nh) + self.reg(nh))\n",
    "        return self._loss\n",
    "\n",
    "    \n",
    "class Generator(object):\n",
    "\n",
    "    def __init__(self, data, triple_set, nb_classes, batch_size):\n",
    "        self.data = data\n",
    "        self.triple_set = triple_set\n",
    "        self.nb_classes = nb_classes\n",
    "        self.size = len(data)\n",
    "        self.start = 0\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        if self.start < self.size:\n",
    "            end = self.start + self.batch_size\n",
    "            batch = self.data[self.start:end]\n",
    "            neg_triples = []\n",
    "            for h, l, t in batch:\n",
    "                while True:\n",
    "                    hn = h\n",
    "                    tn = t\n",
    "                    if np.random.choice([False, True]):\n",
    "                        hn = np.random.randint(0, self.nb_classes)\n",
    "                    else:\n",
    "                        tn = np.random.randint(0, self.nb_classes)\n",
    "                    if (hn, l, tn) not in self.triple_set:\n",
    "                        neg_triples.append((hn, l, tn))\n",
    "                        break\n",
    "            self.start = end\n",
    "            batch = np.hstack(batch).reshape(len(batch), 3)\n",
    "            neg_triples = np.hstack(neg_triples).reshape(len(neg_triples), 3)\n",
    "            labels = np.zeros((len(batch), 1), dtype='float32')\n",
    "            return ([batch, neg_triples], labels)\n",
    "        else:\n",
    "            self.start = 0\n",
    "            return self.__next__()\n",
    "            #raise StopIteration()\n",
    "\n",
    "\n",
    "def load_data(filename=f'data/train/{org_id}.plain.nt', format='nt'):\n",
    "    g = Graph()\n",
    "    g.parse(filename, format=format)\n",
    "    classes = {}\n",
    "    relations = {}\n",
    "    data = []\n",
    "    for h, l, t in g:\n",
    "        h, l, t = str(h), str(l), str(t)\n",
    "        if h not in classes:\n",
    "            classes[h] = len(classes)\n",
    "        if t not in classes:\n",
    "            classes[t] = len(classes)\n",
    "        if l not in relations:\n",
    "            relations[l] = len(relations)\n",
    "        data.append((classes[h], relations[l], classes[t]))\n",
    "    return data, classes, relations\n",
    "\n",
    "def load_valid_data(data_file, classes, relations):\n",
    "    data = []\n",
    "    rel = f'http://interacts'\n",
    "    with open(data_file, 'r') as f:\n",
    "        for line in f:\n",
    "            it = line.strip().split()\n",
    "            id1 = f'http://{it[0]}'\n",
    "            id2 = f'http://{it[1]}'\n",
    "            if id1 not in classes or id2 not in classes or rel not in relations:\n",
    "                continue\n",
    "            data.append((classes[id1], relations[rel], classes[id2]))\n",
    "    return data\n",
    "\n",
    "data, classes, relations = load_data()\n",
    "train_steps = int(math.ceil(len(data) / batch_size))\n",
    "train_generator = Generator(data, set(data), len(classes), batch_size)\n",
    "\n",
    "valid_data = load_valid_data(f'data/valid/{org_id}.protein.links.v11.0.txt', classes, relations)\n",
    "valid_generator = Generator(valid_data, set(data), len(classes), batch_size)\n",
    "valid_steps = int(math.ceil(len(valid_data) / batch_size))\n",
    "\n",
    "pos_input = Input(shape=(3,), dtype='int32')\n",
    "neg_input = Input(shape=(3,), dtype='int32')\n",
    "transe_model = TransE(len(classes), len(relations), embeddings_size)\n",
    "out = transe_model([pos_input, neg_input])\n",
    "model = tf.keras.Model(inputs=[pos_input, neg_input], outputs=out)\n",
    "model.compile(optimizer=optimizer, loss=loss_function)\n",
    "\n",
    "model.fit_generator(\n",
    "    train_generator,\n",
    "    steps_per_epoch=train_steps,\n",
    "    epochs=epochs,\n",
    "    validation_data=valid_generator,\n",
    "    validation_steps=valid_steps,\n",
    "    workers=4,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save embeddings:\n",
    "\n",
    "The embeddings provided in the data package have been trained for 100 epochs and may provide better results; you may want to skip this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_embeddings = transe_model.cls_embeddings.get_weights()[0]\n",
    "rel_embeddings = transe_model.rel_embeddings.get_weights()[0]\n",
    "cls_embeddings = list(cls_embeddings)\n",
    "rel_embeddings = list(rel_embeddings)\n",
    "\n",
    "if not os.path.exists('data/transe'):\n",
    "    os.makedirs('data/transe')\n",
    "\n",
    "df = pd.DataFrame({'classes': list(classes.keys()), 'embeddings': cls_embeddings})\n",
    "df.to_pickle(f'data/transe/{org_id}_cls_embeddings.pkl')\n",
    "\n",
    "df = pd.DataFrame({'relations': list(relations.keys()), 'embeddings': rel_embeddings})\n",
    "df.to_pickle(f'data/transe/{org_id}_rel_embeddings.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation results on test set:\n",
    " * Hits@k, k $\\in$ {10, 100}\n",
    " * Mean rank\n",
    " * ROC AUC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Proteins:  6229\n",
      "0.06 0.32 1125.44 0.82\n",
      "0.13 0.40 1074.69 0.83\n"
     ]
    }
   ],
   "source": [
    "from scipy.stats import rankdata\n",
    "org_id = '4932'\n",
    "\n",
    "def load_test_data(data_file, classes, relations):\n",
    "    data = []\n",
    "    rel = f'http://interacts'\n",
    "    with open(data_file, 'r') as f:\n",
    "        for line in f:\n",
    "            it = line.strip().split()\n",
    "            id1 = f'http://{it[0]}'\n",
    "            id2 = f'http://{it[1]}'\n",
    "            if id1 not in classes or id2 not in classes or rel not in relations:\n",
    "                continue\n",
    "            data.append((id1, rel, id2))\n",
    "    return data\n",
    "\n",
    "def compute_rank_roc(ranks, n_prots):\n",
    "    auc_x = list(ranks.keys())\n",
    "    auc_x.sort()\n",
    "    auc_y = []\n",
    "    tpr = 0\n",
    "    sum_rank = sum(ranks.values())\n",
    "    for x in auc_x:\n",
    "        tpr += ranks[x]\n",
    "        auc_y.append(tpr / sum_rank)\n",
    "    auc_x.append(n_prots)\n",
    "    auc_y.append(1)\n",
    "    auc = np.trapz(auc_y, auc_x) / n_prots\n",
    "    return auc\n",
    "\n",
    "# Load embeddings from saved files\n",
    "cls_df = pd.read_pickle(f'data/transe/{org_id}_cls_embeddings.pkl')\n",
    "rel_df = pd.read_pickle(f'data/transe/{org_id}_rel_embeddings.pkl')\n",
    "classes = {v: k for k, v in enumerate(cls_df['classes'])}\n",
    "relations = {v: k for k, v in enumerate(rel_df['relations'])}\n",
    "\n",
    "nb_classes = len(cls_df)\n",
    "nb_relations = len(rel_df)\n",
    "embeds_list = cls_df['embeddings'].values\n",
    "rembeds_list = rel_df['embeddings'].values\n",
    "size = len(embeds_list[0])\n",
    "embeds = np.zeros((nb_classes, size), dtype=np.float32)\n",
    "for i, emb in enumerate(embeds_list):\n",
    "    embeds[i, :] = emb\n",
    "\n",
    "proteins = {}\n",
    "for k, v in classes.items():\n",
    "    if not k.startswith('http://GO:'):\n",
    "        proteins[k] = v\n",
    "\n",
    "print('Proteins: ', len(proteins))\n",
    "\n",
    "prot_index = list(proteins.values())\n",
    "prot_embeds = embeds[prot_index, :]\n",
    "prot_dict = {v: k for k, v in enumerate(prot_index)}\n",
    "    \n",
    "rsize = len(rembeds_list[0])\n",
    "rembeds = np.zeros((nb_relations, rsize), dtype=np.float32)\n",
    "for i, emb in enumerate(rembeds_list):\n",
    "    rembeds[i, :] = emb\n",
    "\n",
    "# Load training data to computed filtered rank\n",
    "train_data = load_test_data(f'data/train/{org_id}.protein.links.v11.0.txt', classes, relations)\n",
    "valid_data = load_test_data(f'data/valid/{org_id}.protein.links.v11.0.txt', classes, relations)\n",
    "trlabels = {}\n",
    "for c, r, d in train_data:\n",
    "    c, r, d = prot_dict[classes[c]], relations[r], prot_dict[classes[d]]\n",
    "    if r not in trlabels:\n",
    "        trlabels[r] = np.ones((len(prot_embeds), len(prot_embeds)), dtype=np.int32)\n",
    "    trlabels[r][c, d] = 1000\n",
    "for c, r, d in valid_data:\n",
    "    c, r, d = prot_dict[classes[c]], relations[r], prot_dict[classes[d]]\n",
    "    if r not in trlabels:\n",
    "        trlabels[r] = np.ones((len(prot_embeds), len(prot_embeds)), dtype=np.int32)\n",
    "    trlabels[r][c, d] = 1000\n",
    "\n",
    "\n",
    "# Load test data and compute ranks for each protein\n",
    "test_data = load_test_data(f'data/test/{org_id}.protein.links.v11.0.txt', classes, relations)\n",
    "top1 = 0\n",
    "top10 = 0\n",
    "top100 = 0\n",
    "mean_rank = 0\n",
    "ftop1 = 0\n",
    "ftop10 = 0\n",
    "ftop100 = 0\n",
    "fmean_rank = 0\n",
    "labels = {}\n",
    "preds = {}\n",
    "ranks = {}\n",
    "franks = {}\n",
    "eval_data = test_data\n",
    "n = len(eval_data)\n",
    "for c, r, d in eval_data:\n",
    "    c, r, d = prot_dict[classes[c]], relations[r], prot_dict[classes[d]]\n",
    "    if r not in labels:\n",
    "        labels[r] = np.zeros((len(prot_embeds), len(prot_embeds)), dtype=np.int32)\n",
    "    if r not in preds:\n",
    "        preds[r] = np.zeros((len(prot_embeds), len(prot_embeds)), dtype=np.float32)\n",
    "    labels[r][c, d] = 1\n",
    "    ec = prot_embeds[c, :]\n",
    "    er = rembeds[r, :]\n",
    "    ec += er\n",
    "\n",
    "    # Compute distance\n",
    "    dst = np.linalg.norm(prot_embeds - ec.reshape(1, -1), axis=1)\n",
    "    res = dst.flatten()\n",
    "\n",
    "    preds[r][c, :] = res\n",
    "    index = rankdata(res, method='average')\n",
    "    rank = index[d]\n",
    "    if rank == 1:\n",
    "        top1 += 1\n",
    "    if rank <= 10:\n",
    "        top10 += 1\n",
    "    if rank <= 100:\n",
    "        top100 += 1\n",
    "    mean_rank += rank\n",
    "    if rank not in ranks:\n",
    "        ranks[rank] = 0\n",
    "    ranks[rank] += 1\n",
    "\n",
    "    # Filtered rank\n",
    "    index = rankdata((res * trlabels[r][c, :]), method='average')\n",
    "    rank = index[d]\n",
    "    if rank == 1:\n",
    "        ftop1 += 1\n",
    "    if rank <= 10:\n",
    "        ftop10 += 1\n",
    "    if rank <= 100:\n",
    "        ftop100 += 1\n",
    "    fmean_rank += rank\n",
    "\n",
    "    if rank not in franks:\n",
    "        franks[rank] = 0\n",
    "    franks[rank] += 1\n",
    "top1 /= n\n",
    "top10 /= n\n",
    "top100 /= n\n",
    "mean_rank /= n\n",
    "ftop1 /= n\n",
    "ftop10 /= n\n",
    "ftop100 /= n\n",
    "fmean_rank /= n\n",
    "\n",
    "rank_auc = compute_rank_roc(ranks, len(proteins))\n",
    "frank_auc = compute_rank_roc(franks, len(proteins))\n",
    "\n",
    "print(f'{top10:.2f} {top100:.2f} {mean_rank:.2f} {rank_auc:.2f}')\n",
    "print(f'{ftop10:.2f} {ftop100:.2f} {fmean_rank:.2f} {frank_auc:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}