Link_prediction_Pytorch_Geometric.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5bjYrUL9V7LV"
   },
   "source": [
    "# Link prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "57O237suPDmh"
   },
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import networkx as nx\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "from torch_geometric.datasets import Planetoid\n",
    "from torch_geometric.datasets import CitationFull\n",
    "from torch_geometric.datasets import HeterophilousGraphDataset\n",
    "from torch_geometric.datasets import Reddit\n",
    "from torch_geometric.datasets import PPI\n",
    "\n",
    "from torch_geometric.utils import to_networkx, from_networkx\n",
    "\n",
    "import torch\n",
    "from torch_geometric.transforms import RandomLinkSplit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOAD DATASET\n",
    "dataset = Planetoid(root='/tmp/PubMed', name='PubMed') ## PUBMED\n",
    "# dataset = CitationFull(root='/tmp/Cora', name='Cora')  ## CORA\n",
    "# dataset = HeterophilousGraphDataset(root=\"./\", name='amazon_ratings') ## AMAZON RATINGS\n",
    "# dataset = PPI(\"./\")\n",
    "data = dataset[0]\n",
    "G_BC = to_networkx(data, to_undirected=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of nodes: 19717\n",
      "Number of edges: 44324\n"
     ]
    }
   ],
   "source": [
    "labels = []\n",
    "for i in range(len(data.y)):\n",
    "    l = data.y[i].squeeze().numpy().tolist()\n",
    "    labels.append(l)\n",
    "\n",
    "group_dict = {i: labels[i] for i in range(len(labels))}\n",
    "\n",
    "for user_id, groups in group_dict.items():\n",
    "    nx.set_node_attributes(G_BC, {user_id: groups}, 'group_belonging')\n",
    "\n",
    "print(\"Number of nodes:\", G_BC.number_of_nodes())\n",
    "print(\"Number of edges:\", G_BC.number_of_edges())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def predict_link(u, v, embeddings):\n",
    "    \"\"\"\n",
    "    Computes the normalized probability for an existing link between two nodes u and v based on the input\n",
    "    embeddings.\n",
    "    :param u: a node in the graph\n",
    "    :param v: a node in the graph\n",
    "    :param embeddings: trained embeddings\n",
    "    :return: sigmoid normalized probability for the existence of a link\n",
    "    \"\"\"\n",
    "    embedding1 = embeddings[u]\n",
    "    embedding2 = embeddings[v]\n",
    "    \n",
    "    # Compute inner product (dot product)\n",
    "    dot_product = np.dot(embedding1, embedding2)\n",
    "\n",
    "    # Normalize by sigmoid function\n",
    "    link_probability = 1/(1 + np.exp(-dot_product))\n",
    "    return link_probability\n",
    "\n",
    "\n",
    "def link_predictions(embeddings, edges, y_true):\n",
    "    \"\"\"\n",
    "    Computes the ROC-AUC score for a given set of test edges based on the trained embeddings.\n",
    "    :param embeddings: a models trained embeddings\n",
    "    :param edges: test edges\n",
    "    :param y_true: the labels for edges (1=true, 0=false)\n",
    "    :return: the ROC-AUC score from predictions\n",
    "    \"\"\"\n",
    "    predictions = []\n",
    "    for edge in edges:\n",
    "        predictions.append(predict_link(edge[0], edge[1], embeddings))\n",
    "    return roc_auc_score(y_true, predictions) \n",
    "\n",
    "\n",
    "def train_test_split_graph(G):\n",
    "    \"\"\"\n",
    "    Splits a Graph into a test and train set randomly to 80-20. The test split is balanced with negative edges sampled from random vertex pairs that have no edges between them. \n",
    "    While removing edges randomly, it makes sure that no vertex is isolated.\n",
    "    :param G: a networkx graph to be split\n",
    "    :return: the train-test split as torch geometrics graphs\n",
    "    \"\"\"\n",
    "    data = from_networkx(G)\n",
    "    data.y = data.group_belonging\n",
    "    data.x = torch.arange(G.number_of_nodes()).unsqueeze(1)\n",
    "    \n",
    "    transform = RandomLinkSplit(num_val=0, num_test=0.5, is_undirected=True, add_negative_train_samples=False)\n",
    "    train_data, _, test_data = transform(data)\n",
    "    return train_data, test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOAD EMBEDDING\n",
    "from gensim.models.keyedvectors import KeyedVectors\n",
    "\n",
    "PATH = '/Users/silviaarellanogarcia/Documents/MSc MACHINE LEARNING/Advanced Machine Learning/Project/embeddings_deepwalk/'\n",
    "phi_model_YT = KeyedVectors.load_word2vec_format(PATH + 'model_pubmed_80_40.embedding')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the embedding vectors\n",
    "embeddings = phi_model_YT.vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.20333679,  0.30727208,  0.06666798, ...,  0.03066145,\n",
       "         0.11071127, -0.22095709],\n",
       "       [ 0.16608925,  0.01392533, -0.19145992, ...,  0.05858718,\n",
       "         0.268196  ,  0.05441923],\n",
       "       [-0.12973866, -0.1092654 ,  0.10813908, ..., -0.13678004,\n",
       "         0.14862283,  0.04702898],\n",
       "       ...,\n",
       "       [ 0.5484147 , -0.26757112,  0.2658604 , ..., -0.23587406,\n",
       "         0.12957697, -0.09191801],\n",
       "       [ 0.67656016,  0.40571168, -0.24498807, ...,  0.43837348,\n",
       "        -0.07032686, -0.10148539],\n",
       "       [ 0.748623  ,  0.11105119,  0.41657597, ...,  0.04104827,\n",
       "         0.03264922,  0.17897221]], dtype=float32)"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-31T12:12:50.302871Z",
     "start_time": "2023-12-31T12:12:50.301129Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# LINK PREDICTION\n",
    "train_data, test_data = train_test_split_graph(G_BC)\n",
    "\n",
    "# Prepare edges\n",
    "test_edges = test_data.edge_label_index.numpy().T\n",
    "y_true = test_data.edge_label.numpy()\n",
    "\n",
    "roc_auc = link_predictions(embeddings, test_edges, y_true)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.49776102820359625"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "roc_auc"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}