{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0d642e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import sys\n",
    "sys.path += ['utils/']  \n",
    "\n",
    "from dynamical_processes import *\n",
    "from DyANE import * \n",
    "from node_embedding import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "07b98b85",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 360 # number of time frames\n",
    "weighted = True # use a weighted representation of the graph \n",
    "\n",
    "G, all_nodes = loadGraph('SFHH', T, weighted) # load the graph\n",
    "G_supra = SupraAdjacencyMatrix(G, all_nodes) # generate the supra-adjacency representation\n",
    "A_supra = nx.adjacency_matrix(G_supra) # convert it into the adjacency matrix representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "a7d3707c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Learning embeddings using the Skip-Gram model...\n",
      "Preprocessing progress: 99.67%%cessing progress: %0%Preprocessing progress: 0%00%%\n",
      "Preprocessing transition probabilities: done!\n",
      "Walking progress: 96.06%92.06%ogress: 0%0%\n",
      "Simulating random walks: done!\n",
      "Learning embeddings: done!\n"
     ]
    }
   ],
   "source": [
    "# We create the embedding with DyANE\n",
    "dim = 200\n",
    "embedding = DyANE(A_supra, ww = 15, dim = dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "9ec7433e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running the optimization for k = 1\n",
      "[========================>] 100%\r"
     ]
    }
   ],
   "source": [
    "# We create the embedding with EDRep\n",
    "res = NodeEmbedding(A_supra, dim, n_epochs = 20, walk_length = 1, k = 1, verbose = True, η = .5, sym = True)\n",
    "embedding2 = res.X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "f7b97160",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now we perform the SIR simulation\n",
    "\n",
    "λ, η = .1, 0.01\n",
    "n = len(all_nodes)\n",
    "\n",
    "s0 = np.random.choice([0,1,2], n, p = [1, 0., 0.])\n",
    "s0[np.random.choice(all_nodes, 1, replace = False)] = 1\n",
    "\n",
    "SIRresult = SIR(G, all_nodes, λ, η, s0)\n",
    "\n",
    "# Perform the reconstruction\n",
    "p = 1*n/len(G_supra.nodes)\n",
    "T = len(G)\n",
    "\n",
    "n_m = 50\n",
    "pr, pr2 = np.zeros((n_m, T)), np.zeros((n_m, T))\n",
    "\n",
    "for i in range(n_m):\n",
    "    try:\n",
    "        idx = np.random.binomial(1, p, len(G_supra.nodes)) == 1\n",
    "        Ipred = Predict(embedding, idx, SIRresult.state, G)\n",
    "        pr[i] = Ipred\n",
    "\n",
    "        Ipred = Predict(embedding2, idx, SIRresult.state, G)\n",
    "        pr2[i] = Ipred\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "b7caea8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 175256 characters]",
      "text/plain": [
       "<Figure size 1200x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize = (12,5))\n",
    "\n",
    "m1, s = np.mean(pr, axis = 0), np.std(pr, axis = 0)\n",
    "ax[0].fill_between(np.arange(T), m1 - s, m1 + s, color = 'crimson', alpha = 0.3)\n",
    "ax[0].plot(m1, color = 'crimson', linewidth = 2, label = 'DyANE', linestyle = '--')\n",
    "\n",
    "m2, s = np.mean(pr2, axis = 0), np.std(pr, axis = 0)\n",
    "ax[0].fill_between(np.arange(T), m2 - s, m2 + s, color = 'dodgerblue', alpha = 0.3)\n",
    "ax[0].plot(m2, color = 'navy', linewidth = 2, label = 'EDRep', linestyle = '-.')\n",
    "\n",
    "ax[0].plot(SIRresult.Iactive, color = 'k', linewidth = 3, label = 'Measured')\n",
    "ax[0].legend(fontsize = 15)\n",
    "ax[0].set_title('DyANE', fontsize = 26)\n",
    "ax[0].tick_params(axis='both', which='major', labelsize=20)\n",
    "ax[0].set_xlabel('time', fontsize = 24)\n",
    "ax[0].set_ylabel('n. infected', fontsize = 24)\n",
    "\n",
    "\n",
    "ax[1].scatter(m1, m2, color = 'orange', alpha = 0.7, edgecolor = 'brown', s = 80)\n",
    "ax[1].set_xlabel('Prediction DyANE', fontsize = 24)\n",
    "ax[1].yaxis.set_label_position(\"right\")\n",
    "ax[1].set_ylabel('Prediction EDRep', fontsize = 24, rotation = -90, labelpad = 30)\n",
    "ax[1].tick_params(axis='both', which='major', labelsize=20)\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "909576f5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "EDRep",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
