{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%env CUDA_VISIBLE_DEVICES=1\n",
    "import matplotlib.pyplot as plt\n",
    "from utils.utils import get_path,lap_eig_decay, diffusion_decay, eff_res_decay, eff_res_corr_decay\n",
    "from utils.toydata_utils import get_toy_data\n",
    "from utils.fig_utils import dist_to_color\n",
    "from utils.dist_utils import compute_laplacian, get_sknn\n",
    "import os\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import scipy.sparse.linalg\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6a8938af0e720c79"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "style_file = \"utils.style\"\n",
    "plt.style.use(style_file)\n",
    "root_path = get_path(\"data\")\n",
    "fig_path = os.path.join(root_path, \"figures\")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "49f9de0ef0a5a776"
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Eigenvector decay fig"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "1cd0336a718af695"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# get data\n",
    "seed = 0\n",
    "sigmas = [0.0, 0.1, 0.25]\n",
    "n = 1000\n",
    "d = 50\n",
    "dataset = \"toy_circle\"\n",
    "\n",
    "data = np.array([get_toy_data(dataset = dataset, n= n, d=d, seed=seed, **{\"gaussian\":{\"sigma\": sigma}}) for sigma in sigmas])"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c146f95193c189b9"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# get eigenvalues and eigenvectors\n",
    "k = 15\n",
    "n_evecs = 999\n",
    "norms = [\"sym\"]  # only use symmetric normalization for the plot\n",
    "\n",
    "evals = {}\n",
    "evecs = {}\n",
    "for norm in norms:\n",
    "    evals_norm = {}\n",
    "    evecs_norm = {}\n",
    "    for i, data_sig in enumerate(data):\n",
    "        sknn_coo = get_sknn(data_sig, k=k)\n",
    "\n",
    "        L = compute_laplacian(sknn_coo, normalization=norm)\n",
    "\n",
    "        # compute eigenvectors, following the setting in the UMAP code base\n",
    "        eigenvalues, eigenvectors = scipy.sparse.linalg.eigsh(\n",
    "                        L,\n",
    "                        n_evecs,\n",
    "                        which=\"SM\",\n",
    "                        tol=1e-4,\n",
    "                        v0=np.ones(L.shape[0]),\n",
    "                        maxiter=L.shape[0] * 5,\n",
    "                    )\n",
    "        evals_norm[sigmas[i]] = eigenvalues\n",
    "        evecs_norm[sigmas[i]] = eigenvectors\n",
    "        print(f\"Done with norm {norm} and sigma {sigmas[i]}\")\n",
    "    evals[norm] = evals_norm\n",
    "    evecs[norm] = evecs_norm"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "240cf4b56c1b8944"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# plot figure\n",
    "fig, ax = plt.subplots(ncols=5, width_ratios=(1.2, 0.8,1,1,1), figsize=(5.5, 1.4))\n",
    "\n",
    "\n",
    "letters = \"abcde\"\n",
    "\n",
    "# panel a with eigenvalues for different noise levels\n",
    "linestyles = [\"solid\", \"dashed\", \"dotted\"]\n",
    "for j, sigma in enumerate(sigmas):\n",
    "    ax[0].plot(np.arange(1, 999), evals[\"sym\"][sigma][1:], label=f\"$\\sigma = {sigma}$\", c=\"k\", linestyle=linestyles[j])\n",
    "ax[0].legend(frameon=False, loc=(0.4, 0.01), handlelength=1.0)\n",
    "ax[0].set_xscale(\"symlog\")\n",
    "ax[0].set_yscale(\"log\")\n",
    "ax[0].set_xlim(1, 1000)\n",
    "ax[0].set_ylim(None, 1.4)\n",
    "ax[0].spines['left'].set_position(('data', 1))\n",
    "ax[0].spines['bottom'].set_position(('data', 0.0002))\n",
    "\n",
    "ax[0].set_xticks([1, 10, 100, 1000])\n",
    "ax[0].set_xticklabels([1, 10, 100, 1000])\n",
    "\n",
    "ax[0].xaxis.set_minor_locator(matplotlib.ticker.LogLocator(base=10.0, subs=(0.1, 0.2, 0.3, 0.4, 0.5,  0.6, 0.7, 0.8, 0.9)))\n",
    "ax[0].set_title(\"Eigenvalue spectra\")\n",
    "ax[0].set_title(\n",
    "    letters[0],\n",
    "    ha=\"right\",\n",
    "    loc=\"left\",\n",
    "    fontweight=\"bold\",\n",
    ")\n",
    "ax[0].set_ylabel(\"Eigenvalue\")\n",
    "ax[0].set_xlabel(\"Eigenvalue index\")\n",
    "\n",
    "\n",
    "cmap = matplotlib.cm.get_cmap('tab20')\n",
    "\n",
    "# panel b with the eigenvector decay functions\n",
    "n_xvals = 100\n",
    "possible_evals = np.arange(1, 1.25* n_xvals) / n_xvals\n",
    "ax[1].plot(possible_evals, eff_res_corr_decay(possible_evals), label=\"eff res corr\", c=dist_to_color[\"eff_res\"], zorder=10, clip_on=False)\n",
    "ax[1].plot(possible_evals, eff_res_decay(possible_evals), label=\"eff res\", c=dist_to_color[\"eff_res\"], linestyle=\"dashdot\", zorder=10, clip_on=False)\n",
    "ax[1].plot(possible_evals, diffusion_decay(possible_evals, t=8), label=\"Diffusion t=8\", c=dist_to_color[\"diffusion\"], zorder=10, clip_on=False)\n",
    "ax[1].plot(possible_evals, diffusion_decay(possible_evals, t=64), label=\"Diffusion t=64\", c=dist_to_color[\"diffusion\"], linestyle=\"dashdot\", zorder=10, clip_on=False)\n",
    "\n",
    "ax[1].set_xticks([0, 0.5, 1.0])\n",
    "ax[1].set_ylim(0.001, 1)\n",
    "ax[1].set_xlim(0, 1.25)\n",
    "\n",
    "ax[1].spines['left'].set_position(('data', 0))\n",
    "ax[1].spines['bottom'].set_position(('data', 0))\n",
    "\n",
    "ax[1].set_xlabel(\"Eigenvalue\")\n",
    "ax[1].set_ylabel(\"Relative decay\")\n",
    "ax[1].set_title(\"Decay\")\n",
    "ax[1].set_title(\n",
    "    letters[1],\n",
    "    ha=\"right\",\n",
    "    loc=\"left\",\n",
    "    fontweight=\"bold\",\n",
    ")\n",
    "\n",
    "# remaining panels with the decay of each eigenvector, one for each noise level\n",
    "for i, sigma in enumerate([0.0, 0.1, 0.25]):\n",
    "    i = i+2\n",
    "    ax[i].plot(np.arange(1, 999), eff_res_corr_decay(evals[\"sym\"][sigma][1:]), label=\"Eff. resistance\", c=dist_to_color[\"eff_res\"], clip_on=False, zorder =5)\n",
    "    ax[i].plot(np.arange(1, 999), eff_res_decay(evals[\"sym\"][sigma][1:]), label=\"Eff. res. naive\", c=dist_to_color[\"eff_res\"], clip_on=False, linestyle=\"dashdot\", zorder =5)\n",
    "    ax[i].plot(np.arange(1, 999), diffusion_decay(evals[\"sym\"][sigma][1:], t=8), label=\"Diffusion $t = 8$\", c=dist_to_color[\"diffusion\"], clip_on=False, zorder =5)\n",
    "    ax[i].plot(np.arange(1, 999), diffusion_decay(evals[\"sym\"][sigma][1:], t=64), label=\"Diffusion $t = 64$\", c=dist_to_color[\"diffusion\"], clip_on=False, linestyle=\"dashdot\", zorder =5)\n",
    "    ax[i].plot(np.arange(1, 999), lap_eig_decay(evals[\"sym\"][sigma][1:], n_evecs=2), label=\"Lap. Eig. $D=2$\", c=dist_to_color[\"spectral\"], clip_on=False, zorder =5)\n",
    "    \n",
    "    ax[i].set_xscale(\"log\")\n",
    "    ax[i].set_xscale(\"log\")\n",
    "    ax[i].set_ylim(0, 1)\n",
    "    ax[i].set_xlim(1, 1000)\n",
    "    \n",
    "    ax[i].spines['left'].set_position(('data', 1))\n",
    "    ax[i].spines['bottom'].set_position(('data', 0))\n",
    "\n",
    "    if i != 0:\n",
    "        ax[i].set_yticklabels([])\n",
    "    ax[i].set_xlabel(\"Eigenvector index\")\n",
    "    ax[i].set_xticks([1, 10, 100, 1000])    \n",
    "    ax[i].set_xticklabels([1, 10, 100, 1000], zorder=6)   \n",
    "    if i == 4:\n",
    "        ax[i].legend(frameon=False, loc=(0.4,0.35))\n",
    "    ax[i].set_title(f\"$\\sigma = {sigma}$\")\n",
    "    ax[i].set_title(\n",
    "    letters[i],\n",
    "    ha=\"right\",\n",
    "    loc=\"left\",\n",
    "    fontweight=\"bold\",\n",
    ")\n",
    "fig.savefig(os.path.join(fig_path, \"fig_eig_decay_log.pdf\"))"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a95d557d19607b66"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "7a96f2364611c894"
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "conda-env-ph-py",
   "language": "python",
   "display_name": "Python [conda env:ph]"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
