Link Prediction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import hnswlib\n",
    "import networkx as nx\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse as sp\n",
    "\n",
    "\"\"\"\n",
    "Disclaimer: functions defined from lines 15 to 36 in this file come from \n",
    "tkipf/gae original repository on Graph Autoencoders. Moreover, the\n",
    "mask_test_edges function is borrowed from philipjackson's mask_test_edges \n",
    "pull request on this same repository.\n",
    "\"\"\"\n",
    "\n",
    "def sparse_to_tuple(sparse_mx):\n",
    "    if not sp.isspmatrix_coo(sparse_mx):\n",
    "        sparse_mx = sparse_mx.tocoo()\n",
    "    coords = np.vstack((sparse_mx.row, sparse_mx.col)).transpose()\n",
    "    values = sparse_mx.data\n",
    "    shape = sparse_mx.shape\n",
    "    return coords, values, shape\n",
    "\n",
    "def preprocess_graph(adj):\n",
    "    adj = sp.coo_matrix(adj)\n",
    "    adj_ = adj + sp.eye(adj.shape[0])\n",
    "    degree_mat_inv_sqrt = sp.diags(np.power(np.array(adj_.sum(1)), -0.5).flatten())\n",
    "    adj_normalized = adj_.dot(degree_mat_inv_sqrt).transpose().dot(degree_mat_inv_sqrt)\n",
    "    return sparse_to_tuple(adj_normalized)\n",
    "\n",
    "def construct_feed_dict(adj_normalized, adj, features, placeholders):\n",
    "    # Construct feed dictionary\n",
    "    feed_dict = dict()\n",
    "    feed_dict.update({placeholders['features']: features})\n",
    "    feed_dict.update({placeholders['adj']: adj_normalized})\n",
    "    feed_dict.update({placeholders['adj_orig']: adj})\n",
    "    return feed_dict\n",
    "\n",
    "def mask_test_edges(adj, test_percent=1., val_percent=0.):\n",
    "    \"\"\" Randomly removes some edges from original graph to create\n",
    "    test and validation sets for link prediction task\n",
    "    :param adj: complete sparse adjacency matrix of the graph\n",
    "    :param test_percent: percentage of edges in test set\n",
    "    :param val_percent: percentage of edges in validation set\n",
    "    :return: train incomplete adjacency matrix, validation and test sets\n",
    "    \"\"\"\n",
    "    # Remove diagonal elements\n",
    "    adj = adj - sp.dia_matrix((adj.diagonal()[None, :], [0]), shape=adj.shape)\n",
    "    adj.eliminate_zeros()\n",
    "    # Check that diag is zero:\n",
    "    #assert adj.diagonal().sum() == 0\n",
    "\n",
    "    edges_positive, _, _ = sparse_to_tuple(adj)\n",
    "    # Filtering out edges from lower triangle of adjacency matrix\n",
    "    edges_positive = edges_positive[edges_positive[:,1] > edges_positive[:,0],:]\n",
    "    # val_edges, val_edges_false, test_edges, test_edges_false = None, None, None, None\n",
    "\n",
    "    # number of positive (and negative) edges in test and val sets:\n",
    "    num_test = int(np.floor(edges_positive.shape[0] / (100. / test_percent)))\n",
    "    num_val = 0\n",
    "\n",
    "    # sample positive edges for test and val sets:\n",
    "    edges_positive_idx = np.arange(edges_positive.shape[0])\n",
    "    np.random.shuffle(edges_positive_idx)\n",
    "    val_edge_idx = edges_positive_idx[:num_val]\n",
    "    test_edge_idx = edges_positive_idx[num_val:(num_val + num_test)]\n",
    "    test_edges = edges_positive[test_edge_idx] # positive test edges\n",
    "    val_edges = edges_positive[val_edge_idx] # positive val edges\n",
    "    train_edges = np.delete(edges_positive, np.hstack([test_edge_idx, val_edge_idx]), axis = 0) # positive train edges\n",
    "\n",
    "    # the above strategy for sampling without replacement will not work for\n",
    "    # sampling negative edges on large graphs, because the pool of negative\n",
    "    # edges is much much larger due to sparsity, therefore we'll use\n",
    "    # the following strategy:\n",
    "    # 1. sample random linear indices from adjacency matrix WITH REPLACEMENT\n",
    "    # (without replacement is super slow). sample more than we need so we'll\n",
    "    # probably have enough after all the filtering steps.\n",
    "    # 2. remove any edges that have already been added to the other edge lists\n",
    "    # 3. convert to (i,j) coordinates\n",
    "    # 4. swap i and j where i > j, to ensure they're upper triangle elements\n",
    "    # 5. remove any duplicate elements if there are any\n",
    "    # 6. remove any diagonal elements\n",
    "    # 7. if we don't have enough edges, repeat this process until we get enough\n",
    "    positive_idx, _, _ = sparse_to_tuple(adj) # [i,j] coord pairs for all true edges\n",
    "    positive_idx = positive_idx[:,0]*adj.shape[0] + positive_idx[:,1] # linear indices\n",
    "    test_edges_false = np.empty((0,2),dtype='int64')\n",
    "    idx_test_edges_false = np.empty((0,),dtype='int64')\n",
    "\n",
    "    while len(test_edges_false) < len(test_edges):\n",
    "        # step 1:\n",
    "        idx = np.random.choice(adj.shape[0]**2, 2*(num_test - len(test_edges_false)), replace = True)\n",
    "        # step 2:\n",
    "        idx = idx[~np.in1d(idx, positive_idx, assume_unique = True)]\n",
    "        idx = idx[~np.in1d(idx, idx_test_edges_false, assume_unique = True)]\n",
    "        # step 3:\n",
    "        rowidx = idx // adj.shape[0]\n",
    "        colidx = idx % adj.shape[0]\n",
    "        coords = np.vstack((rowidx,colidx)).transpose()\n",
    "        # step 4:\n",
    "        lowertrimask = coords[:,0] > coords[:,1]\n",
    "        coords[lowertrimask] = coords[lowertrimask][:,::-1]\n",
    "        # step 5:\n",
    "        coords = np.unique(coords, axis = 0) # note: coords are now sorted lexicographically\n",
    "        np.random.shuffle(coords) # not anymore\n",
    "        # step 6:\n",
    "        coords = coords[coords[:,0] != coords[:,1]]\n",
    "        # step 7:\n",
    "        coords = coords[:min(num_test, len(idx))]\n",
    "        test_edges_false = np.append(test_edges_false, coords, axis = 0)\n",
    "        idx = idx[:min(num_test, len(idx))]\n",
    "        idx_test_edges_false = np.append(idx_test_edges_false, idx)\n",
    "\n",
    "    val_edges_false = np.empty((0,2), dtype = 'int64')\n",
    "    idx_val_edges_false = np.empty((0,), dtype = 'int64')\n",
    "#     while len(val_edges_false) < len(val_edges):\n",
    "#         # step 1:\n",
    "#         idx = np.random.choice(adj.shape[0]**2, 2*(num_val - len(val_edges_false)), replace = True)\n",
    "#         # step 2:\n",
    "#         idx = idx[~np.in1d(idx, positive_idx, assume_unique = True)]\n",
    "#         idx = idx[~np.in1d(idx, idx_test_edges_false, assume_unique = True)]\n",
    "#         idx = idx[~np.in1d(idx, idx_val_edges_false, assume_unique = True)]\n",
    "#         # step 3:\n",
    "#         rowidx = idx // adj.shape[0]\n",
    "#         colidx = idx % adj.shape[0]\n",
    "#         coords = np.vstack((rowidx,colidx)).transpose()\n",
    "#         # step 4:\n",
    "#         lowertrimask = coords[:,0] > coords[:,1]\n",
    "#         coords[lowertrimask] = coords[lowertrimask][:,::-1]\n",
    "#         # step 5:\n",
    "#         coords = np.unique(coords, axis = 0) # note: coords are now sorted lexicographically\n",
    "#         np.random.shuffle(coords) # not any more\n",
    "#         # step 6:\n",
    "#         coords = coords[coords[:,0] != coords[:,1]]\n",
    "#         # step 7:\n",
    "#         coords = coords[:min(num_val, len(idx))]\n",
    "#         val_edges_false = np.append(val_edges_false, coords, axis = 0)\n",
    "#         idx = idx[:min(num_val, len(idx))]\n",
    "#         idx_val_edges_false = np.append(idx_val_edges_false, idx)\n",
    "\n",
    "    # sanity checks:\n",
    "#     train_edges_linear = train_edges[:,0]*adj.shape[0] + train_edges[:,1]\n",
    "#     test_edges_linear = test_edges[:,0]*adj.shape[0] + test_edges[:,1]\n",
    "#     assert not np.any(np.in1d(idx_test_edges_false, positive_idx))\n",
    "#     assert not np.any(np.in1d(idx_val_edges_false, positive_idx))\n",
    "#     assert not np.any(np.in1d(val_edges[:,0]*adj.shape[0]+val_edges[:,1], train_edges_linear))\n",
    "#     assert not np.any(np.in1d(test_edges_linear, train_edges_linear))\n",
    "#     assert not np.any(np.in1d(val_edges[:,0]*adj.shape[0]+val_edges[:,1], test_edges_linear))\n",
    "\n",
    "    # Re-build adj matrix\n",
    "    data = np.ones(train_edges.shape[0])\n",
    "    adj_train = sp.csr_matrix((data, (train_edges[:, 0], train_edges[:, 1])), shape=adj.shape)\n",
    "    adj_train = adj_train + adj_train.T\n",
    "    return adj_train, val_edges, val_edges_false, test_edges, test_edges_false"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading Embeddings and creating Indexer for NN search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "IDList = []                                # List of paper IDs\n",
    "NNList = []                                # List of list, NNList[i]: NNs to paper whose id is IDList[i]\n",
    "embeddings = []                            # Embeddings read from the input file\n",
    "\n",
    "with open('./data/dblp_Abstract_2Thresholded_USE_Trans_Embeddings.json', 'r') as file:\n",
    "    for line in file:\n",
    "        data = json.loads(line)\n",
    "        paperID = data['id']\n",
    "        embedding = data['embedding']\n",
    "        IDList.append(paperID)\n",
    "        embeddings.append(embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "numElements = len(IDList)\n",
    "dimension = len(embeddings[0])\n",
    "embeddings = np.asarray(embeddings)\n",
    "data_labels = np.arange(numElements)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = hnswlib.Index(space='cosine', dim=dimension)  # the space can be changed - keeps the data, alters the distance function.\n",
    "\n",
    "# Increase the total capacity (max_elements), so that it will handle the new data\n",
    "p.load_index(\"./models/USETranshnswlib.bin\", max_elements = numElements)\n",
    "labels, _ = p.knn_query(embeddings, k = 4)\n",
    "del p\n",
    "del embeddings\n",
    "del data_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Adjacency List for Node Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating Citation Adjacency List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "adjList = defaultdict(set)                          # Convert set to list later for node2vec, set: to handle duplicates\n",
    "with open('./data/dblp_AIpapers2Thresholded.json', 'r') as file:\n",
    "    for line in file:\n",
    "        data = json.loads(line)\n",
    "        paperID = data['id']\n",
    "        references = data.get('references', [])\n",
    "        for referencedPaper in references:\n",
    "            adjList[paperID].add(referencedPaper)\n",
    "            adjList[referencedPaper].add(paperID)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Augmenting Adj List with FastText NNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "nnToKeep = 3\n",
    "id = 0\n",
    "for label in labels:\n",
    "    paperID = IDList[id]\n",
    "    label = [IDList[index] for index in label if index != id]\n",
    "    if (len(label) > nnToKeep):\n",
    "        del label[nnToKeep:]\n",
    "    for referencedPaper in label:\n",
    "        adjList[paperID].add(referencedPaper)\n",
    "        adjList[referencedPaper].add(paperID)\n",
    "    id += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating NetworkX Graph and reporting graph statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of nodes:  475839 . Number of edges:  6760367 . Avg Degree:  28.414514152896253\n"
     ]
    }
   ],
   "source": [
    "adjList = {key: list(values) for key, values in adjList.items()}\n",
    "G = nx.from_dict_of_lists(adjList)\n",
    "\n",
    "nnodes = G.number_of_nodes()\n",
    "avgDegree = sum(d for n, d in G.degree()) / float(nnodes)\n",
    "print('Number of nodes: ', nnodes, '. Number of edges: ', G.number_of_edges(), '. Avg Degree: ', avgDegree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "adj_sparse = nx.to_scipy_sparse_matrix(G)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adj_sparse.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Constructing new graph\n"
     ]
    }
   ],
   "source": [
    "# Perform train-test split\n",
    "adj_train, val_edges, val_edges_false, test_edges, test_edges_false = mask_test_edges(adj_sparse)\n",
    "print('Constructing new graph')\n",
    "G_train = nx.from_scipy_sparse_matrix(adj_train) # new graph object with only non-hidden edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "G_train.number_of_nodes() == G.number_of_nodes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = list(G.nodes())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for adjV in adjList[u]:\n",
    "    if adjV == v:\n",
    "        print('Found')\n",
    "    \n",
    "## paperID mapping: G.nodes() are paperIDs one to one mapped with G_train.nodes which are 0 to |V| - 1\n",
    "## test edges and test edges negative are integer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Node2Vec Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Computing transition probabilities:   1%|          | 5509/475839 [22:03<43:16:25,  3.02it/s] IOPub message rate exceeded.\n",
      "The notebook server will temporarily stop sending output\n",
      "to the client in order to avoid crashing it.\n",
      "To change this limit, set the config variable\n",
      "`--NotebookApp.iopub_msg_rate_limit`.\n",
      "\n",
      "Current values:\n",
      "NotebookApp.iopub_msg_rate_limit=1000.0 (msgs/sec)\n",
      "NotebookApp.rate_limit_window=3.0 (secs)\n",
      "\n",
      "Computing transition probabilities:   3%|▎         | 13006/475839 [43:04<32:26:49,  3.96it/s] "
     ]
    }
   ],
   "source": [
    "from node2vec import Node2Vec\n",
    "walkLength = 8\n",
    "node2vec = Node2Vec(G_train, walk_length = walkLength)#, workers = 12, temp_folder = './data/tmp_data')\n",
    "          "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = node2vec.fit()  # returns a gensim wv model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outFileName = './models/node2vec_USE_2Citation_Embeddings_WL_' + str(walkLength) + '_NN_' + str(nnToKeep) + '.kv'\n",
    "model.wv.save_word2vec_format(outFileName)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Edge Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getEdgeEmbedding(embedding1, embedding2, policy='Hadamard'):\n",
    "    if (policy=='Hadamard'):\n",
    "        return embedding1 * embedding2\n",
    "    elif (policy=='Avg'):\n",
    "        return (embedding1 + embedding2) / 2\n",
    "def average(lis):\n",
    "    return (sum(lis) / len(lis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2VecMode = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = []\n",
    "\n",
    "embeddingDict = dict()\n",
    "\n",
    "if (word2VecMode):\n",
    "    embeddingFileName = './data/dblpAbstract_2Thresholded_FT_Embeddings.json'\n",
    "    with open(embeddingFileName, 'r') as file:\n",
    "        for line in file:\n",
    "            data = json.loads(line)\n",
    "            embeddingDict[data['id']] = data['embedding']\n",
    "            \n",
    "edges = [*test_edges, *test_edges_false]\n",
    "for edge in edges :\n",
    "    if (word2VecMode):\n",
    "        paperID1 = nodes[edge[0]]\n",
    "        paperID2 = nodes[edge[1]]\n",
    "        embedding1 = np.asarray(embeddingDict[paperID1])\n",
    "        embedding2 = np.asarray(embeddingDict[paperID2])\n",
    "    else:\n",
    "        u = str(edge[0])\n",
    "        v = str(edge[1])\n",
    "        if u in model.wv.vocab:\n",
    "            embedding1 = np.asarray(model.wv[u])\n",
    "        else:\n",
    "            embedding1 = np.asarray([0] * 128)\n",
    "        if v in model.wv.vocab:\n",
    "            embedding2 = np.asarray(model.wv[v])\n",
    "        else:\n",
    "            embedding2 = np.asarray([0] * 128)\n",
    "    edgeEmbedding =  getEdgeEmbedding(embedding1, embedding2)\n",
    "    X.append(edgeEmbedding)\n",
    "\n",
    "Y = np.asarray([1] * len(test_edges) + [0] * len(test_edges_false))\n",
    "\n",
    "del embeddingDict    \n",
    "X = np.asarray(X)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neural_network import MLPClassifier\n",
    "from sklearn.svm import LinearSVC\n",
    "from sklearn.gaussian_process.kernels import RBF\n",
    "from sklearn.ensemble import RandomForestClassifier, BaggingClassifier\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.multiclass import OneVsRestClassifier\n",
    "\n",
    "\n",
    "names = [\n",
    " \"Random Forest\", \"Neural Net\", \"Logistic Regression\", \"Linear SVC\" ]\n",
    "\n",
    "classifiers = [\n",
    "    RandomForestClassifier(verbose=True, n_jobs = -1),\n",
    "    MLPClassifier(verbose=True, early_stopping=True),\n",
    "    LogisticRegression(n_jobs=-1),\n",
    "    OneVsRestClassifier(BaggingClassifier(LinearSVC(),n_jobs = -1))]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import precision_recall_fscore_support\n",
    "from sklearn import svm\n",
    "from sklearn.model_selection import StratifiedKFold, KFold\n",
    "\n",
    "kfold = KFold(n_splits=5, shuffle=True)\n",
    "for name, clf in zip(names, classifiers):\n",
    "    precScores = []\n",
    "    recallScores = []\n",
    "    f1Scores = []\n",
    "    count = 1\n",
    "    for train_index, test_index in kfold.split(X, Y):\n",
    "        X_train, X_test = X[train_index], X[test_index]\n",
    "        y_train, y_test = Y[train_index], Y[test_index]\n",
    "        print('Fitting: ', count)\n",
    "        clf.fit(X_train, y_train)\n",
    "        print('count ', count)\n",
    "        y_pred = clf.predict(X_test)\n",
    "        prec, recall, fscore, _ = precision_recall_fscore_support(y_test, y_pred, average='weighted')\n",
    "        precScores.append(prec)\n",
    "        recallScores.append(recall)\n",
    "        f1Scores.append(fscore)\n",
    "        count += 1\n",
    "    print('Name', name,'. Avg Precision: ', average(precScores), '. Avg Recall: ', average(recallScores), '. Avg F-1 Score: ', average(f1Scores) )\n"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}