demo_spherical_convolution.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 convolution\n",
    "\n",
    "In this small notebook, we test an implementation of a spherical convolution. The general idea is to use a graph instead of the tradtionial 2 dimensional grid as a support for convolution."
   ]
  },
  {
   "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",
    "\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",
    "\n",
    "from deepsphere import utils, plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.figsize'] = (17, 5)  # (9, 4) for matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 Graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us start by constructing two small graphs on the healpix sampling scheme and visualizing them (with and without edges)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "\n",
    "ax = fig.add_subplot(121, projection='3d')\n",
    "G = utils.healpix_graph(nside=8, nest=True)\n",
    "G.plotting.update(vertex_size=10)\n",
    "G.plot(edges=False, ax=ax)\n",
    "\n",
    "ax = fig.add_subplot(122, projection='3d')\n",
    "G = utils.healpix_graph(nside=4, nest=True)\n",
    "G.plotting.update(vertex_size=20)\n",
    "G.plot(ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The healpix sampling induces an 8 nearest neighbors graph, i.e. a graph where each vertex is connected to 8 vertices. Some vertices are however connected to 7 neighbors only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in np.unique(G.d):\n",
    "    print('Number of nodes with {} neighbors: {}'.format(i, np.sum(G.d == i)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 Fourier basis\n",
    "\n",
    "Graph convolution is defined as the pointwise multiplication in the graph spectral domain. Hence it is important to verify the spectral property of the graph. Note that this operation requires the diagonalization of the Laplacian, which is very costly in computations and mermory. Nevertheless, when it comes to convolution, their exist fast methods that only require sparse matrix multiplications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = utils.healpix_graph(nside=16, lap_type='normalized', nest=True, dtype=np.float64)\n",
    "\n",
    "print('max weighted degree: {:.2f}'.format(G.dw.max()))\n",
    "print('min weighted degree: {:.2f}'.format(G.dw.min()))\n",
    "print('mean weighted degree: {:.2f}'.format(G.dw.mean()))\n",
    "print('Is the graph directed? {}'.format(G.is_directed()))\n",
    "print('Number of nodes: {}'.format(G.N))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The eigenvectors are obtained by diagonalizing the graph laplacian defined as $L=I-D^{\\frac{1}{2}}WD^{\\frac{1}{2}}$, where $W$ is the weight/adjacency matrix and $D$ the degree matrix. \n",
    "\n",
    "The Fourier basis $U$ by definition satisfies\n",
    "$$ L  = U \\Lambda U^*. $$\n",
    "Here the eigenvalues contained in the diagonal of $\\Lambda$ somehow correspond to the graph squared frequencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute all eigenvectors.\n",
    "G.compute_fourier_basis()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Mean: {}'.format(np.mean(np.abs(G.U.ravel()))))\n",
    "print('Min: {}'.format(np.min(G.U.ravel())))\n",
    "print('Max: {}'.format(np.max(G.U.ravel())))\n",
    "print('Perline: {}'.format(np.max(G.U, axis=1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us display a few Fourier modes on the healpix map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(17, 10))\n",
    "ne = 16\n",
    "cm = plt.cm.RdBu_r\n",
    "cm.set_under('w')\n",
    "\n",
    "for ind in range(ne):\n",
    "    hp.mollview(G.U[:, ind], \n",
    "                title='Eigenvector {}'.format(ind), \n",
    "                nest=True, \n",
    "                sub=(ne//4, 4, ind+1),\n",
    "                max=np.max(np.abs(G.U[:, :ne])),\n",
    "                min=-np.max(np.abs(G.U[:, :ne])),\n",
    "                cbar=False,\n",
    "                cmap=cm,\n",
    "                rot=(0,0,0))\n",
    "with utils.HiddenPrints():\n",
    "    hp.graticule();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should also check higher frequency modes as they can be more localized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IND = 3000\n",
    "fig = plt.figure(figsize=(3, 2))\n",
    "hp.mollview(G.U[:, IND], title=\"Eigenvector {}\".format(IND), nest=True, cbar=False, sub=(1, 1, 1), cmap=cm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most localized eigenvector is considered to be the one with the heighest coherence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind = np.argmax(np.max(np.abs(G.U), axis=0))\n",
    "fig = plt.figure(figsize=(3, 2))\n",
    "hp.mollview(G.U[:, ind], title=\"Eigenvector {}\".format(ind), nest=True, cbar=False, sub=(1, 1, 1), cmap=cm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This eigenvector is clearly very localized. Let us display the modulus of the Fourier eigenvector to have a more general idea about all eigenvectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.imshow(np.abs(G.U), cmap=cm, vmin=-np.max(np.abs(G.U)), vmax=np.max(np.abs(G.U)));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embedding the data into a 3D space using the eigenvectors\n",
    "\n",
    "We clearly see that the eigenvectors capture the spherical geometry."
   ]
  },
  {
   "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(G.U[:, 1], G.U[:, 2], G.U[:, 3], c=G.d);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 Convolution on graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The convolution of a signal $f$ and a kernel $k(x)$ on a graph is defined as the pointwise multiplication in the spectral domain, i.e.\n",
    "$$f_c  = U k(\\Lambda)U^*f. $$\n",
    "Here $U^*f$ is the graph Fourier transform of $f$ and $k(\\Lambda)$ is a diagonal matrix where the kernel $k$ is applied on each element of the diagonal of $\\Lambda$. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us start with the heat diffusion problem. We solve the following equation on the graph:\n",
    "$$ L f(t) = \\tau \\partial_t f(t),$$\n",
    "where $f(t): \\mathbb{R}_+ \\rightarrow \\mathbb{R}^N$ is a multivariate function depending on the time, $L$ a positive semi-definite matrix representing the Laplacian of a graph, and $\\tau$ a constant.\n",
    "\n",
    "Given the vector $f_0 = f(0)$, the solution of this equation for time $t$ can be written as:\n",
    "$$ f(t) = K_t(L) f_0, $$\n",
    "where \n",
    "$$ K_t(L) = e^{-\\tau t L}.$$\n",
    "In the equation $f(t) = K_t(L) f_0$, the kernel $K_t(x)=e^{-\\tau t x}$ can be considered as the convolution kernel and the heat diffusion problem can be solved using a simple convolution on the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "taus = [5, 10, 20, 50]\n",
    "hf = pg.filters.Heat(G, tau=taus)\n",
    "fig, ax = plt.subplots()\n",
    "hf.plot(eigenvalues=True, sum=False, ax=ax)\n",
    "ax.set_title('Filter frequency response');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def arcmin2rad(x):\n",
    "    return x / 60 / 360 * 2 * np.pi\n",
    "\n",
    "def gaussian_smoothing(sig, sigma, nest=True):\n",
    "    with utils.HiddenPrints():\n",
    "        if nest:\n",
    "            sig = hp.reorder(sig, n2r=True)\n",
    "        smooth = hp.sphtfunc.smoothing(sig, sigma=arcmin2rad(sigma))\n",
    "        if nest:\n",
    "            smooth = hp.reorder(smooth, r2n=True)\n",
    "        return smooth\n",
    "_, center = plot.get_index_equator(hp.npix2nside(G.N), radius=20)\n",
    "ind0 = center\n",
    "sig = np.zeros(G.N)\n",
    "sig[ind0] = 1\n",
    "conv = hf.analyze(sig)\n",
    "\n",
    "fig = plt.figure()\n",
    "rel_diff = []\n",
    "cm = plt.cm.seismic\n",
    "# cm = plt.cm.jet\n",
    "cm.set_under('w')\n",
    "m = 0\n",
    "for i, (tau, sigma) in enumerate(zip(taus, [315, 465, 670, 1080])):\n",
    "    smooth = gaussian_smoothing(sig, sigma, nest=True)\n",
    "    m = max(m, max(smooth))\n",
    "    hp.mollview(conv[:, i], \n",
    "                title='$t={}$'.format(tau), \n",
    "                nest=True, \n",
    "                min=-m, max=m,\n",
    "                cbar=False,\n",
    "                rot=(180,0,180),\n",
    "                sub=(2, len(taus), i+1),\n",
    "                cmap=cm)\n",
    "    \n",
    "    hp.mollview(smooth, \n",
    "                title='$\\sigma={}$'.format(sigma), \n",
    "                nest=True, \n",
    "                min=-m, max=m,\n",
    "                cbar=False,\n",
    "                rot=(180,0,180),\n",
    "                sub=(2,len(taus),i+len(taus)+1),\n",
    "                cmap=cm)\n",
    "    \n",
    "    diff = (conv[:, i]-smooth)\n",
    "    rel_diff.append(np.linalg.norm(diff)/np.linalg.norm(smooth))\n",
    "\n",
    "with utils.HiddenPrints():\n",
    "    hp.graticule();\n",
    "\n",
    "print(rel_diff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 Smoothing a Planck map\n",
    "\n",
    "Let us play with a Planck map."
   ]
  },
  {
   "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": [
    "map_cmb, map_noise, map_mask = hp.read_map('data/COM_CMB_IQU-smica_1024_R2.02_full.fits', field=(0, 1, 3), nest=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cm = plt.cm.RdBu\n",
    "cm.set_under('w')\n",
    "\n",
    "hp.mollview(map_cmb, title='cmb', nest=True, cmap=cm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us first downsample the map in order to make the smoothing operation more visible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 256\n",
    "map_cmb_lores = hp.ud_grade(map_cmb, nside_out=nside, order_in='NESTED')\n",
    "G = utils.healpix_graph(nside=nside, nest=True)\n",
    "G.estimate_lmax()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let apply our heat operator. It will smooth the map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "taus = [5, 10, 20, 50]\n",
    "hf = pg.filters.Heat(G, tau=taus)\n",
    "conv_map_lowres = hf.analyze(map_cmb_lores)\n",
    "\n",
    "fig = plt.figure(figsize=(9, 6))\n",
    "for i, tau in enumerate(taus):\n",
    "    hp.mollview(conv_map_lowres[:, i], \n",
    "                title=\"Tau: {}\".format(tau), \n",
    "                nest=True, \n",
    "                cmap=cm,\n",
    "                sub=(np.sqrt(len(taus)), np.sqrt(len(taus)), i+1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 Power spectral density\n",
    "\n",
    "Let us now compute the power spectral density on the sphere. This is going to be different than the one based on the spherical harmonics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def estimate_graph_psd(graph, signal, n_rand=10, n_point=30):\n",
    "    \"\"\"Estimate the power spectral density on graph.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    graph : pygsp.graphs.Graph\n",
    "    signal : numpy.ndarray\n",
    "        Signal whose PSD is to be estimated.\n",
    "    n_rand : int\n",
    "        Number of random signals used for the estimation.\n",
    "    n_point : int\n",
    "        Number of points at which the PSD is estimated.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    spectrum : numpy.ndarray\n",
    "        Points at which the PSD was estimated.\n",
    "    psd : numpy.ndarray\n",
    "        Power spectral density (PSD) at the evaluation points.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Define filterbank.\n",
    "    g = pg.filters.Itersine(graph, Nf=n_point, overlap=2)\n",
    "    spectrum = np.linspace(0, graph.lmax, n_point)\n",
    "    \n",
    "    # Filter signal.\n",
    "    sig_filt = g.filter(signal, method='chebyshev', order=2*n_point)\n",
    "    sig_dist = np.sum(sig_filt**2, axis=0)\n",
    "    if sig_dist.ndim > 1:\n",
    "        sig_dist = np.mean(sig_dist, axis=0).squeeze()\n",
    "    \n",
    "    # Estimate the eigenvectors by filtering random signals.\n",
    "    rand_sig = np.random.binomial(n=1, p=0.5, size=[graph.N, n_rand]) * 2 - 1\n",
    "    rand_sig_filered = g.filter(rand_sig, method='chebyshev', order=2*n_point)\n",
    "    eig_dist = np.mean(np.sum(rand_sig_filered**2, axis=0), axis=0).squeeze()\n",
    "    \n",
    "    # Compute PSD.\n",
    "    psd_values = sig_dist / eig_dist\n",
    "    \n",
    "    return spectrum, psd_values\n",
    "\n",
    "spectrum, psd = estimate_graph_psd(G, map_cmb_lores, n_rand=5, n_point=30)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.semilogy(spectrum, psd, '.-');\n",
    "ax.set_xlabel(\"laplacian's eigenvalues\");\n",
    "ax.set_ylabel('power spectral density (PSD)');"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}