{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting package metadata (current_repodata.json): done\n",
      "Solving environment: done\n",
      "\n",
      "# All requested packages already installed.\n",
      "\n",
      "\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Collecting package metadata (current_repodata.json): done\n",
      "Solving environment: done\n",
      "\n",
      "# All requested packages already installed.\n",
      "\n",
      "\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%conda install pyg -c pyg -c conda-forge\n",
    "%conda install pytorch torchvision torchaudio -c pytorch\n",
    "%pip install networkx seaborn -Uq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30136 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "AttributeError",
     "evalue": "'Graph' object has no attribute 'edge_index'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m/var/folders/3_/gmvd1nkx285133z5yh3chz2c0000gp/T/ipykernel_16599/2356794267.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0mnx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdraw_networkx_edge_labels\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpubmed_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0medge_labels\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlabels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0mplt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshow\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m \u001b[0mpubmed_graph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0medge_index\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m: 'Graph' object has no attribute 'edge_index'"
     ]
    }
   ],
   "source": [
    "# pubmed = pd.read_csv('~/workspace/cogtext/data/pubmed_abstracts.')\n",
    "\n",
    "# pubmed = pd.DataFrame(np.random.random((100,2)), columns=['label','pmid','vector'])\n",
    "# pubmed.head()\n",
    "\n",
    "pubmed_graph = nx.random_geometric_graph(5,5)\n",
    "\n",
    "for (u, v) in pubmed_graph.edges():\n",
    "    pubmed_graph.edges[u,v]['weight'] = np.random.randint(0,10)\n",
    "\n",
    "pos = nx.spring_layout(pubmed_graph)\n",
    "nx.draw(pubmed_graph, pos)\n",
    "labels = nx.get_edge_attributes(pubmed_graph, 'weight')\n",
    "nx.draw_networkx_edge_labels(pubmed_graph, pos,edge_labels=labels)\n",
    "plt.show()\n",
    "pubmed_graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create an edge between two nodes if their corresponding corpora share documents >= min_docs. The weight of the edge will be set to the pointwise mutual information between the two labels.\n",
    "\n",
    "log(xy) - log(x) - log(y) + log(D)\n",
    "xy is the number of articles shared between the two labels\n",
    "x is the number of articles in the x corpus\n",
    "and D is the total number of articles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.nn import Node2Vec\n",
    "\n",
    "data = \n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model = Node2Vec(data.edge_index, embedding_dim=128, \n",
    "                 walk_length=20,                        # lenght of rw\n",
    "                 context_size=10, walks_per_node=20,\n",
    "                 num_negative_samples=1, \n",
    "                 p=200, q=1,                             # bias parameters\n",
    "                 sparse=True).to(device)\n",
    "\n",
    "loader = model.loader(batch_size=128, shuffle=True, num_workers=4)\n",
    "\n",
    "for idx, (pos_rw, neg_rw) in enumerate(loader):\n",
    "    print(idx, pos_rw.shape, neg_rw.shape)\n",
    "    \n",
    "edge_tuples = [tuple(x) for x in data.edge_index.numpy().transpose()]\n",
    "G = nx.from_edgelist(edge_tuples)\n",
    "pos = nx.spring_layout(G, center=[0.5, 0.5])\n",
    "nx.set_node_attributes(G, pos, 'pos')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "fd70b4c63b8ac7010e057e1e7961ebae705e9dba34aa5a2e9dfe5bc9196414e5"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 64-bit ('pyg': conda)",
   "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.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
