demo_spherical_vs_graph.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [DeepSphere]: a spherical convolutional neural network\n",
    "[DeepSphere]: https://github.com/SwissDataScienceCenter/DeepSphere\n",
    "\n",
    "[Nathanaël Perraudin](https://perraudin.info), [Michaël Defferrard](http://deff.ch), Tomasz Kacprzak, Raphael Sgier\n",
    "\n",
    "# Demo: spherical harmonics vs graph Fourier modes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from time import time\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import healpy as hp\n",
    "import pygsp as pg\n",
    "import itertools\n",
    "\n",
    "from deepsphere import utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.figsize'] = (17, 5)  # (9, 4) for matplotlib notebook\n",
    "\n",
    "nside = 16\n",
    "cm = plt.cm.RdBu_r\n",
    "cm.set_under('w')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 Spherical harmonics\n",
    "\n",
    "* Spherical harmonics are indexed by degree (or angular frequency) $\\ell$ and order $m \\in [0, \\ell]$.\n",
    "* $l=0$ is the constant, $l=1$ is the two monopoles, $l=2$ is the three dipoles, etc.\n",
    "* Coefficients are commonly defined as $a_{\\ell m}$.\n",
    "* A spherical harmonic takes value $Y_\\ell^m(\\theta, \\phi)$ in angular direction $(\\theta, \\phi)$.\n",
    "* In HEALPix, as we only deal with maps of real numbers, the representation is complex and the relation $a_{\\ell, m}^* = a_{\\ell, -m}$ holds. As such, to generate all $2\\ell + 1$ orders, one needs to e.g. transform $a_{\\ell m} = 1 + 1j$ and $a_{\\ell m} = 1 - 1j$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_spherical_harmonic(nside, l, m):\n",
    "\n",
    "    lmax = l\n",
    "    idx = hp.sphtfunc.Alm.getidx(lmax, l, m)\n",
    "    size = hp.sphtfunc.Alm.getsize(lmax, mmax=lmax)\n",
    "    print('{} spherical harmonics for l in [0, {}]'.format(size, lmax))\n",
    "    print('l={}, m={} is at index {}'.format(l, m, idx))\n",
    "\n",
    "    alm = np.zeros(size, dtype=np.complex128)\n",
    "    alm[idx] = 1\n",
    "\n",
    "    map = hp.sphtfunc.alm2map(alm, nside, lmax, verbose=False)\n",
    "    hp.mollview(map, title=\"Spherical harmonic of degree l={} and order m={}\".format(l,m), cmap=cm)\n",
    "    #hp.cartview(map, title=f\"Spherical harmonic l={l}, m={m}\")\n",
    "    with utils.HiddenPrints():\n",
    "        hp.graticule();\n",
    "plot_spherical_harmonic(nside, l=2, m=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the spherical harmonics up to $\\ell_{max}$ and plot them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "harmonics = utils.compute_spherical_harmonics(nside, lmax=18)\n",
    "harmonics.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_harmonics(harmonics, title=''):\n",
    "    n_harmonics = harmonics.shape[1]\n",
    "    l, m = 0, 0\n",
    "    for idx in range(n_harmonics):\n",
    "        hp.mollview(harmonics[:, idx], \n",
    "                    title='{}: l={}, m={}'.format(title, l, m),\n",
    "                    nest=True,\n",
    "                    sub=(np.sqrt(n_harmonics), np.sqrt(n_harmonics), idx+1),\n",
    "                    max=np.max(np.abs(harmonics)),\n",
    "                    min=-np.max(np.abs(harmonics)),\n",
    "                    cbar=False,\n",
    "                    cmap=cm)\n",
    "        \n",
    "        m += 1\n",
    "        if m > l:\n",
    "            l += 1\n",
    "            m = -l\n",
    "    with utils.HiddenPrints():\n",
    "        hp.graticule();\n",
    "plot_harmonics(harmonics[:, :16], 'Spherical harmonic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 Graph Fourier modes\n",
    "\n",
    "The graph Fourier modes are the eigenvectors of the graph Laplacian.\n",
    "$$L = U \\Lambda U^T$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "graph.compute_fourier_basis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The weighted adjacency matrix is very sparse. Distance between pixels is not constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2)\n",
    "axes[0].imshow(graph.W.toarray())\n",
    "axes[1].hist(graph.W[graph.W>0].T,30);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fourier modes of the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_harmonics(graph.U[:, :16], 'Eigenvector')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The eigenvalues are clearly organized in groups, which corresponds to angular frequencies $\\ell$ of the spherical harmonics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(graph.e[:50], '.-')\n",
    "idx = 0\n",
    "for l in range(7):\n",
    "    plt.text(idx, graph.e[idx] + 0.005, 'l = {}'.format(l))\n",
    "    idx += 2*l + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Fourier modes capture the spherical geometry.\n",
    "\n",
    "The eigenvectors can be used to embed any graph in a 3D space.\n",
    "In our case, that should produce a sphere.\n",
    "That is [Laplacian eigenmaps](http://www2.imm.dtu.dk/projects/manifold/Papers/Laplacian.pdf), an algorithm that preserves local distances.\n",
    "\n",
    "Question: are the graph representation and the embedding actually equivalent?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10, 10))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.scatter(graph.U[:, 1], graph.U[:, 2], graph.U[:, 3], c=graph.d);\n",
    "\n",
    "# The spherical harmonics obviously capture the geometry too.\n",
    "# ax.scatter(harmonics[:, 1], harmonics[:, 2], harmonics[:, 3], c=graph.d);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Spherical harmonics on the graph\n",
    "\n",
    "Todo:\n",
    "* Smoothness of rotated spherical harmonics should be constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combinatorial Laplacian for the spherical harmonics.\n",
    "graph = utils.healpix_graph(nside, lap_type='combinatorial', nest=True, dtype=np.float64)\n",
    "graph.compute_fourier_basis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us try to re-order the spherical harmonic with respect of the graph frequencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind = np.argsort(np.diag(harmonics.T @ graph.L @ harmonics))\n",
    "harmonics_sort = harmonics[:,ind]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Spherical harmonics are not exactly orthogonal on the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C_euclidean = harmonics_sort.T @ harmonics_sort\n",
    "C_graph = (harmonics_sort.T @ graph.L @ harmonics_sort) #@ np.diag(1/(graph.e[:harmonics_sort.shape[1]] +0.0001))\n",
    "\n",
    "fig, axes = plt.subplots(1, 2)\n",
    "axes[0].imshow(C_euclidean)\n",
    "axes[1].imshow(C_graph)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2)\n",
    "axes[0].plot(np.diag(C_euclidean))\n",
    "axes[1].plot(np.diag(C_graph), label='diag(F^T L F)')\n",
    "axes[1].plot(graph.e[:harmonics_sort.shape[1]],label='graph eigenvalues')\n",
    "axes[1].legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 Correspondance of the subspaces\n",
    "\n",
    "* Are the subspaces equivalent?\n",
    "* Is the projection on the subspaces similar?\n",
    "* TODO: what is `win = hp.sphtfunc.pixwin(nside=32)`?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Graph eigenvectors vs sampled spherical harmonics\n",
    "\n",
    "That is not an exact representation of the correspondance as in the case of non-uniform sampling, the integrals (i.e., the SHT) are not equal to the dot products with the sampled harmonic functions (which aren't orthogonal and hence don't form a basis)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmax = 10\n",
    "n_harmonics = np.sum(np.arange(1, 2*lmax+2, 2))\n",
    "print('{} harmonics for lmax = {}'.format(n_harmonics, lmax))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "C = harmonics[:, :n_harmonics].T @ harmonics[:, :n_harmonics]\n",
    "print(np.linalg.norm(C, ord='fro'))\n",
    "\n",
    "C = graph.U[:, :n_harmonics].T @ graph.U[:, :n_harmonics]\n",
    "print(np.linalg.norm(C, ord='fro'))\n",
    "\n",
    "C = harmonics[:, :n_harmonics].T @ graph.U[:, :n_harmonics]\n",
    "print(np.linalg.norm(C, ord='fro'))\n",
    "\n",
    "plt.imshow(np.abs(C), cmap=plt.cm.gist_heat_r)\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 The sampled spherical harmonics don't form an orthonormal basis\n",
    "\n",
    "* The spherical harmonics are not exactly orthonormal (they don't form a basis).\n",
    "* They're not orthonormal as they are a sampling of the true eigenfunctions.\n",
    "* If the sampling is not uniform, the dot product is not equal to the integral."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_non_orthogonality(nside, lmax):\n",
    "\n",
    "    harmonics = utils.compute_spherical_harmonics(nside, lmax=lmax)\n",
    "    graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "    n_harmonics = np.cumsum(np.arange(1, 2*lmax+2, 2))\n",
    "    graph.compute_fourier_basis(n_eigenvectors=n_harmonics[-1])\n",
    "\n",
    "    C = harmonics.T @ harmonics\n",
    "    fig, axes = plt.subplots(2, 1, sharex=True)\n",
    "    axes[0].plot(np.sum(C**2, axis=1))\n",
    "    axes[0].plot(C.diagonal())\n",
    "\n",
    "    # Hence the rows and columns of the correlation matrix don't sum up to one.\n",
    "    C = harmonics.T @ graph.U[:, :harmonics.shape[1]]\n",
    "    axes[1].plot(np.sum(C**2, axis=0))\n",
    "    axes[1].plot(np.sum(C**2, axis=1));\n",
    "\n",
    "show_non_orthogonality(nside=16, lmax=8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 Percentage of energy in the corresponding subspaces\n",
    "\n",
    "TODO: the graph should be constructed (if that's even possible) such that the Fourier modes converge to the spherical harmonics while we sample more (i.e., nside increases)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmax = 20\n",
    "nsides = [4, 8, 16]\n",
    "\n",
    "# Normalized corresponds better than combinatorial. Less difference between nsides.\n",
    "# lap_type='combinatorial'\n",
    "lap_type='normalized'\n",
    "\n",
    "def compute_orhtogonality(nside, lmax, lap_type='normalized'):\n",
    "    n_harmonics = np.cumsum(np.arange(1, 2*lmax+2, 2))\n",
    "\n",
    "    harmonics = utils.compute_spherical_harmonics(nside, lmax=lmax)\n",
    "\n",
    "    graph = utils.healpix_graph(nside, lap_type=lap_type, nest=True, dtype=np.float64)\n",
    "    n_eigenvectors = min(n_harmonics[-1], graph.N)\n",
    "    graph.compute_fourier_basis(n_eigenvectors)\n",
    "\n",
    "    C = harmonics.T @ graph.U\n",
    "    \n",
    "    return C, n_harmonics\n",
    "\n",
    "def compute_energy(C, lmax):\n",
    "    \"\"\"Compute the percentage energy in each subspace (diagonal blocks).\"\"\"\n",
    "    energy_in = np.empty(lmax+1)\n",
    "    start = 0\n",
    "    for ell in range(lmax + 1):\n",
    "        n_harmonics = 2*ell + 1\n",
    "        end = start + n_harmonics\n",
    "        energy_square = np.sum(C[start:end, start:end]**2)\n",
    "        # col_sum = np.sum(C[start:end, :]**2)  # how each harmonic are distributed on all eigenvectors\n",
    "        # row_sum = np.sum(C[:, start:end]**2)  # how each eigenvector are distributed on all harmonics\n",
    "        # np.testing.assert_allclose(col_sum, n_harmonics, rtol=1e-2)\n",
    "        # np.testing.assert_allclose(row_sum, n_harmonics, rtol=1e-2)\n",
    "        energy_in[ell] = energy_square / n_harmonics\n",
    "        start = end\n",
    "    return energy_in\n",
    "\n",
    "energy = dict()\n",
    "\n",
    "for nside in nsides:\n",
    "    C, _ = compute_orhtogonality(nside, lmax, lap_type)\n",
    "    energy[nside] = compute_energy(C, lmax)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5, 2))\n",
    "for nside in nsides:\n",
    "    ax.plot(range(lmax+1), 100*energy[nside], '.-', label=rf'$N_{{side}}={nside}$')\n",
    "ax.legend()\n",
    "ax.set_ylabel('correspondance [%]')\n",
    "ax.set_xlabel('degree $\\ell$')\n",
    "ax.set_xticks(range(0, lmax + 1, lmax//10))\n",
    "fig.tight_layout();\n",
    "# fig.savefig(os.path.join(pathfig, 'subspace_alignment.pdf'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4 SHT of the graph eigenvectors\n",
    "\n",
    "That is a better way to compare, as it uses the real spherical harmonic transform (SHT)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nsides = [4, 8, 16, 32]  # slow\n",
    "nsides = [4, 8]\n",
    "\n",
    "spectral_content = dict()\n",
    "\n",
    "for nside in nsides:\n",
    "\n",
    "    lmax = 3 * nside - 1\n",
    "\n",
    "    n_harmonics = np.cumsum(np.arange(1, 2*lmax+2, 2))[-1]\n",
    "    graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "    graph.compute_fourier_basis(n_eigenvectors=n_harmonics)\n",
    "\n",
    "    cl = np.empty((n_harmonics, lmax+1))\n",
    "    for i in range(n_harmonics):\n",
    "        eigenvector = hp.reorder(graph.U[:, i], n2r=True)\n",
    "        # alm = hp.sphtfunc.map2alm(eigenvector)\n",
    "        cl[i] = hp.sphtfunc.anafast(eigenvector, lmax=lmax)\n",
    "\n",
    "    spectral_content[nside] = np.empty((lmax+1, lmax+1))\n",
    "    start = 0\n",
    "    for ell in range(lmax+1):\n",
    "        end = start + (2 * ell + 1)\n",
    "        spectral_content[nside][ell] = np.sum(cl[start:end,:], axis=0)\n",
    "        start = end\n",
    "\n",
    "fig1, axes = plt.subplots(1, len(nsides))\n",
    "fig2, ax2 = plt.subplots()\n",
    "\n",
    "for ax, (nside, sc) in zip(axes, spectral_content.items()):\n",
    "    sc = sc / sc[0, 0]\n",
    "    im = ax.imshow(sc, cmap=plt.cm.gist_heat_r)\n",
    "    fig1.colorbar(im, ax=ax)\n",
    "    ax.set_title(rf'$N_{{side}}={nside}$')\n",
    "\n",
    "    energy_in = np.diag(sc)\n",
    "    ax2.plot(energy_in, label=rf'$N_{{side}}={nside}$')\n",
    "\n",
    "ax2.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: should it sum to one?\n",
    "i = 10\n",
    "eigenvector = hp.reorder(graph.U[:, i], n2r=True)\n",
    "np.linalg.norm(eigenvector)\n",
    "cl = hp.sphtfunc.anafast(eigenvector, lmax=lmax)\n",
    "np.linalg.norm(cl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 Action of the graph Laplacian\n",
    "\n",
    "We don't care if the Fourier modes are not exactly the spherical harmonics, but we want the Laplacian operator to be close.\n",
    "\n",
    "The Laplacian is a second order operator on the sphere (see [Laplace's spherical harmonics](https://en.wikipedia.org/wiki/Spherical_harmonics#Laplace's_spherical_harmonics)):\n",
    "\n",
    "$$\n",
    "r^2 \\nabla^2 Y_\\ell^m(\\theta, \\phi) = -\\ell (\\ell + 1) Y_\\ell^m(\\theta, \\phi).\n",
    "$$\n",
    "\n",
    "Pluging in the graph Laplacian $L = -\\nabla^2$ and setting the radius $r=1$ gives\n",
    "\n",
    "$$\n",
    "L Y_\\ell^m(\\theta, \\phi) = \\ell (\\ell + 1) Y_\\ell^m(\\theta, \\phi)\n",
    "\\\\\n",
    "L a_\\ell^m = \\ell (\\ell + 1) a_\\ell^m\n",
    "$$\n",
    "\n",
    "So multiplying a map with the Laplacian $L$ should affect its spectrum in a predictable way.\n",
    "\n",
    "TODO:\n",
    "* a normalization factor might be missing in the Laplacian\n",
    "* ripple in the error at low-frequency\n",
    "    * Laplacian is approximated by finite differences => greater disc\n",
    "    * influence of Nside?\n",
    "    * due to adjacency being too local? More than 8 neighbors. Get coords of pixels with `hp.pixelfunc.pix2ang` or `hp.query_disc`.\n",
    "    * construct with angular distance\n",
    "    * alternative: check that it gets worse with less neighbors.\n",
    "* average across many realizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folder = 'data'\n",
    "name = 'COM_CMB_IQU-smica_1024_R2.02_full.fits'\n",
    "if not os.path.exists(os.path.join(folder, name)):\n",
    "    url = 'https://irsa.ipac.caltech.edu/data/Planck/release_2/all-sky-maps/maps/component-maps/cmb/' + name\n",
    "    utils.download(url, folder, name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 32\n",
    "npix = 12 * nside**2  # hp.nside2npix\n",
    "print('Nside = {}, Npix = {}'.format(nside, npix))\n",
    "\n",
    "graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "\n",
    "if True:\n",
    "    map1_n, _, _ = hp.read_map('data/COM_CMB_IQU-smica_1024_R2.02_full.fits', field=(0, 1, 3), nest=True)\n",
    "    map1_n = hp.ud_grade(map1_n, nside_out=nside, order_in='NESTED')\n",
    "else:\n",
    "    map1_n = np.random.normal(size=hp.nside2npix(nside))\n",
    "    # map1_n -= map1_n.mean()\n",
    "\n",
    "map1_r = hp.reorder(map1_n, n2r=True)\n",
    "\n",
    "hp.mollview(map1_r)\n",
    "\n",
    "cl1 = hp.sphtfunc.anafast(map1_r)\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(cl1);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ratio between the SHT of the map before and after multiplication by the Laplacian. The frequency response should be equal to $\\ell (\\ell + 1)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map2_n = graph.L @ map1_n\n",
    "\n",
    "map2_r = hp.reorder(map2_n, n2r=True)\n",
    "\n",
    "cl2 = hp.sphtfunc.anafast(map2_r)\n",
    "cl12 = hp.sphtfunc.anafast(map1_r, map2_r)\n",
    "\n",
    "ratio1 = cl2 / cl1\n",
    "plt.plot(ratio1, label='Empirical response')\n",
    "\n",
    "ratio2 = cl12 / cl1\n",
    "plt.plot(ratio2, label='Empirical response')\n",
    "#plt.plot(ratio[2:])\n",
    "#plt.plot(area * ratio)\n",
    "\n",
    "# The scaling issue (the asymptotic value of the error) is probably related to the area of the pixels.\n",
    "area = 4*np.pi / npix\n",
    "print('Area of a pixel: {:.2e}'.format(area))\n",
    "\n",
    "l = np.arange(cl1.shape[0])\n",
    "ideal = l*(l+1) * area / 4\n",
    "plt.plot(l, ideal, label='Ideal response: $\\ell(\\ell+1)$')\n",
    "\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(cl2);\n",
    "plt.plot(cl12);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "error = ratio1 / ideal\n",
    "plt.semilogy(error);\n",
    "\n",
    "error = ratio2 / ideal\n",
    "plt.semilogy(error);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 Action of graph filters\n",
    "\n",
    "We don't care if the Fourier modes are not exactly the spherical harmonics, but we want the Laplacian operator to be close. Or, at least, the filtering operations to be close.\n",
    "\n",
    "Todo:\n",
    "* Measure mean and variance of the difference between a graph filter and the ideal SHT filter at all locations.\n",
    "\n",
    "Functions:\n",
    "* `healpy.sphtfunc.map2alm` => spherical harmonic transform (SHT)\n",
    "* `healpy.sphtfunc.almxfl` => arbitrary smoothing\n",
    "* `healpy.sphtfunc.smoothing` => Gaussian smoothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 256\n",
    "\n",
    "map, _, _ = hp.read_map('data/COM_CMB_IQU-smica_1024_R2.02_full.fits', field=(0, 1, 3), nest=True)\n",
    "map = hp.ud_grade(map, nside_out=nside, order_in='NESTED')\n",
    "map = hp.reorder(map, inp='NESTED', out='RING')\n",
    "# map = np.random.normal(size=hp.nside2npix(nside))\n",
    "\n",
    "lmax = 3 * nside - 1\n",
    "\n",
    "print(\"We'll compute {} frequencies\".format(lmax))\n",
    "hp.mollview(map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm = hp.sphtfunc.map2alm(map)\n",
    "plt.plot(np.abs(alm))\n",
    "\n",
    "l, m = hp.Alm.getlm(lmax=lmax)\n",
    "print(list(zip(l, m))[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl = hp.sphtfunc.anafast(map)\n",
    "plt.plot(cl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "map = graph.L @ map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alm = hp.sphtfunc.map2alm(map)\n",
    "plt.plot(np.abs(alm))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl = hp.sphtfunc.anafast(map)\n",
    "plt.plot(cl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Discrete low-pass filtering using the graph eigenvectors and the spherical harmonics.\n",
    "Note that taking the dot product between the sampled spherical harmonics and the map to filter is not equivalent to performing the filtering with the real SHT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_filtering(nside=16, n_signals=100):\n",
    "    n_harmonics = 16\n",
    "    n_pixels = harmonics.shape[0]\n",
    "\n",
    "    graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "    graph.compute_fourier_basis(n_eigenvectors=n_harmonics)\n",
    "\n",
    "    signals = np.random.uniform(size=(n_pixels, n_signals))\n",
    "    eigenvectors = graph.U\n",
    "    harmonics_ = harmonics[:, :n_harmonics]\n",
    "\n",
    "    signals_sphere = harmonics_ @ harmonics_.T @ signals\n",
    "    signals_graph = eigenvectors @ eigenvectors.T @ signals\n",
    "\n",
    "    hp.mollview(signals[:, 0], nest=True)\n",
    "    hp.mollview(signals_sphere[:, 0], nest=True)\n",
    "    hp.mollview(signals_graph[:, 0], nest=True)\n",
    "\n",
    "    error = signals_graph - signals_sphere\n",
    "    hp.mollview(error[:, 0], nest=True)\n",
    "    print(np.linalg.norm(error, ord='fro'))\n",
    "\n",
    "compare_filtering()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7 Spherical power spectrum\n",
    "\n",
    "The angular power spectrum is given by\n",
    "$$ \\hat{C}_\\ell = \\frac{1}{2\\ell + 1} \\sum_m |\\hat{a}_{\\ell m}|^2 $$\n",
    "As such, $\\hat{C}_\\ell$ is the expected variance of the $\\hat{a}_{\\ell m}$ at order $\\ell$.\n",
    "It also implies the standard result that the total power at the angular wavenumber $\\ell$ is $(2\\ell + 1) \\hat{C}_\\ell$ because there are $2\\ell + 1$ modes for each $\\ell$.\n",
    "\n",
    "Todo:\n",
    "* speed\n",
    "* accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psd = hp.sphtfunc.anafast(map, lmax=lmax)\n",
    "plt.semilogy(psd);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8 Filtering speed\n",
    "\n",
    "TODO:\n",
    "* filtering with tensorflow / numpy / pytorch\n",
    "* Similarly for a part of the sphere\n",
    "* speed and accuracy w.r.t. polynomial order\n",
    "\n",
    "Functions:\n",
    "* `healpy.sphtfunc.almxfl` => arbitrary smoothing\n",
    "* `healpy.sphtfunc.smoothing` => Gaussian smoothing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 256\n",
    "\n",
    "map, _, _ = hp.read_map('data/COM_CMB_IQU-smica_1024_R2.02_full.fits', field=(0, 1, 3), nest=True)\n",
    "map = hp.ud_grade(map, nside_out=nside, order_in='NESTED')\n",
    "map = hp.reorder(map, inp='NESTED', out='RING')\n",
    "\n",
    "hp.mollview(map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Smoothing by filtering with spherical harmonics, i.e., do an element-wise multiplication in the spectral domain after a spherical transform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smooth = hp.sphtfunc.smoothing(map, sigma=0.01, verbose=False)\n",
    "hp.mollview(smooth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Smoothing by filtering on graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "graph.estimate_lmax()\n",
    "\n",
    "map = hp.reorder(map, inp='RING', out='NESTED')\n",
    "\n",
    "filter = pg.filters.Heat(graph, tau=30)\n",
    "filter = filter.approximate('Chebyshev', order=10)\n",
    "filter.plot()\n",
    "\n",
    "smooth = filter.filter(map)\n",
    "hp.mollview(smooth, nest=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the speed of both approaches.\n",
    "\n",
    "We limit OpenMP (used by HEALPix) to use a single core as graph convolutions are not parallelized yet. While the mutliplication of sparse matrices with dense vectors can be parallelized, it is not implemented by scipy. It's hopefully implemented in tensorflow or pytorch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsides = [64, 128, 256, 512, 1024, 2048]  # Max 1024 for the real Planck map.\n",
    "\n",
    "# Polynomial orders for graph filtering.\n",
    "orders = [5, 15]\n",
    "\n",
    "# Largest order (angular frequency) for the spherical harmonic transform.\n",
    "# lmax = lmax * nsides. HEALPix default is 3*nside-1.\n",
    "lmax = [2, 3]\n",
    "\n",
    "# Number of OpenMP threads for spherical harmonics.\n",
    "# TODO: set before starting jupyter.\n",
    "os.environ['OMP_NUM_THREADS'] = '1'\n",
    "\n",
    "times_graph = np.zeros((len(nsides), len(orders)))\n",
    "times_sphere = np.zeros((len(nsides), len(lmax)))\n",
    "times_part = np.zeros((len(nsides), len(orders), 3))\n",
    "\n",
    "for i, nside in enumerate(nsides):\n",
    "    print('Timing nside: {}'.format(nside))\n",
    "    # map = hp.ud_grade(map_cmb, nside_out=nside, order_in='NESTED')\n",
    "    map = np.random.normal(size=hp.nside2npix(nside))\n",
    "\n",
    "    # Filtering on the graph. Need the nested ordering.\n",
    "    graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "    graph.estimate_lmax()\n",
    "    for j, order in enumerate(orders):\n",
    "        filter = pg.filters.Heat(graph, tau=30).approximate('Chebyshev', order=order)\n",
    "        t = time()\n",
    "        smooth = filter.filter(map)\n",
    "        times_graph[i, j] = time() - t\n",
    "    print('  * Graph full: done')\n",
    "        \n",
    "    for k, o in enumerate([1,2,4]):\n",
    "        # Filtering on the graph. Need the nested ordering.\n",
    "        graph = utils.healpix_graph(nside, lap_type='normalized', nest=True, dtype=np.float64, indexes=np.arange((nside//o)**2))\n",
    "        graph.estimate_lmax()\n",
    "        for j, order in enumerate(orders):\n",
    "            filter = pg.filters.Heat(graph, tau=30).approximate('Chebyshev', order=order)\n",
    "            t = time()\n",
    "            smooth = filter.filter(map[:(nside//o)**2])\n",
    "            times_part[i, j, k] = time() - t\n",
    "    print('  * Graph part: done')\n",
    "\n",
    "\n",
    "    # Filtering with the spherical harmonics. Need the ring ordering.\n",
    "    map = hp.reorder(map, inp='NESTED', out='RING')\n",
    "    for j, lm in enumerate(lmax):\n",
    "        t = time()\n",
    "        hp.sphtfunc.smoothing(map, sigma=0.01, verbose=False, lmax=lm*nside)\n",
    "        times_sphere[i, j] = time() - t\n",
    "\n",
    "    print('  * Libsharp: done')\n",
    "\n",
    "# Reset OpenMP to use all available cores.\n",
    "del os.environ['OMP_NUM_THREADS']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save results for the plot in the paper.\n",
    "np.savez('results/filtering_speed.npz', nsides=nsides, orders=orders,\n",
    "         lmax=lmax, times_sphere=times_sphere, times_graph=times_graph, times_part=times_part)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "npix = [hp.nside2npix(nside) for nside in nsides]\n",
    "ax.loglog(npix, times_graph, '.-')\n",
    "ax.loglog(npix, times_sphere, ':')\n",
    "labels = ['Graph, polynomial order {}'.format(order) for order in orders]\n",
    "labels += ['Spherical harmonics, $\\ell_{{max}}$ = {}$N_{{side}}$'.format(lm) for lm in lmax]\n",
    "ax.legend(labels)\n",
    "for i, nside in enumerate(nsides):\n",
    "    ax.text(npix[i], times_sphere[i, -1] * 1.8, 'Nside = {}'.format(nside), horizontalalignment='center')\n",
    "ax.set_ylim(0.8 * times_graph.min(), 3 * times_sphere.max())\n",
    "ax.set_xlabel('Number of pixels')\n",
    "ax.set_ylabel('Processing time [s]')\n",
    "ax.set_title('Single core performance');\n",
    "#ax.ticklabel_format(style='sci', axis='x')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.semilogy(orders, times_graph.T, '.-')\n",
    "labels = ['Filtering with graphs, Nside = {}'.format(nside) for nside in nsides]\n",
    "ax.legend(labels)\n",
    "ax.set_xlabel('Polynomial order')\n",
    "ax.set_xticks(orders)\n",
    "ax.set_ylabel('Processing time [s]');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.semilogy(lmax, times_sphere.T, '.-')\n",
    "labels = ['Filtering with spherical harmonics, Nside = {}'.format(nside) for nside in nsides]\n",
    "ax.legend(labels)\n",
    "ax.set_xlabel('lmax')\n",
    "ax.set_xticks(lmax)\n",
    "ax.set_ylabel('Processing time [s]');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9 Convergence when number of pixels goes to infinity\n",
    "\n",
    "As a graph representation of the sphere is not dependant on the sampling, we can compute spherical harmonics from random samplings of the sphere. As the sampling increases, the eigenvectors approach the spherical harmonics.\n",
    "\n",
    "Do the eigenvalues converge to the orders $\\ell$, i.e. do the stairs become flatter, when increasing $N_{pix}$?\n",
    "No, they mostly stay the same. We are not in the Belkin & Niyogy setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_harmonics = 16\n",
    "n_points = [100, 500, 1000]\n",
    "fig = plt.figure()\n",
    "\n",
    "for i, n in enumerate(n_points):\n",
    "    ax = fig.add_subplot(1, 3, i+1, projection='3d')\n",
    "    graph = pg.graphs.Sphere(nb_pts=n)\n",
    "    graph.compute_fourier_basis(n_eigenvectors=n_harmonics)\n",
    "    graph.plot_signal(graph.U[:, 1], edges=False, ax=ax)\n",
    "    _ = ax.axis('off')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9 Part of sphere\n",
    "\n",
    "Tweak the part graph such that the eigenvectors computed on parts of the sphere are as if we would have restricted the global eigenvectors to a part of sphere.\n",
    "\n",
    "Todo: Mais tu peux essayer sur une grille normale. Le plus simple pour vérifier si c'est possible, c'est de vérifier s'il est nécessaire d'avoir une matrice circulante pour avoir des exponentielles complexes comme vecteurs propres."
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}