{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Efficient) Spherical Harmonic Transform\n",
    "\n",
    "## Preliminaries\n",
    "\n",
    "### Spherical harmonics transform\n",
    "\n",
    "The truncated series expansion of a function $f$ defined on the surface of a sphere can be written as\n",
    "\n",
    "$$\n",
    "f(\\theta, \\lambda) = \\sum_{m=-M}^{M} \\exp(im\\lambda) \\sum_{n=|m|}^{M} F_n^m \\bar{P}_n^m (\\cos \\theta),\n",
    "$$\n",
    "\n",
    "where $\\theta$ is the colatitude, $\\lambda$ the longitude, $\\bar{P}_n^m$ the normalized, associated Legendre polynomials and $F_n^m$, the expansion coefficient associated to the mode $(m,n)$.\n",
    "\n",
    "A direct spherical harmonic transform can be accomplished by a Fourier transform\n",
    "\n",
    "$$\n",
    "F^m(\\theta) = \\frac{1}{2 \\pi} \\int_{0}^{2\\pi} f(\\theta, \\lambda) \\exp(-im\\lambda) \\; \\mathrm{d}\\lambda\n",
    "$$\n",
    "\n",
    "in longitude and a Legendre transform\n",
    "\n",
    "$$\n",
    "F_n^m = \\frac{1}{2} \\int_{-1}^1 F^m(\\theta) \\bar{P}_n^m(\\cos \\theta) \\; \\mathrm{d} \\cos \\theta\n",
    "$$\n",
    "\n",
    "in latitude."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage\n",
    "\n",
    "### Preparation\n",
    "\n",
    "Let us begin by loading some libraries that we will be needing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.fft as fft\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "\n",
    "cmap = 'turbo'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should also load the `torch-sht`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_harmonics import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If a GPU is available we want to use it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test data\n",
    "\n",
    "Let us load some test-data in 2D for testing. We use the MOLA digital elevation map of MARS, courtesy of NASA:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/bbonev/.zshenv:export:2: not valid in this context: :/usr/local/cuda-11.7/lib64\n",
      "--2022-12-14 14:46:11--  https://astropedia.astrogeology.usgs.gov/download/Mars/GlobalSurveyor/MOLA/thumbs/Mars_MGS_MOLA_DEM_mosaic_global_1024.jpg\n",
      "Resolving astropedia.astrogeology.usgs.gov (astropedia.astrogeology.usgs.gov)... 137.227.239.81, 2001:49c8:c000:122d::81\n",
      "Connecting to astropedia.astrogeology.usgs.gov (astropedia.astrogeology.usgs.gov)|137.227.239.81|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 \n",
      "Length: 55192 (54K) [image/jpeg]\n",
      "Saving to: ‘./data/mola_topo.jpg’\n",
      "\n",
      "./data/mola_topo.jp 100%[===================>]  53.90K   165KB/s    in 0.3s    \n",
      "\n",
      "2022-12-14 14:46:12 (165 KB/s) - ‘./data/mola_topo.jpg’ saved [55192/55192]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://astropedia.astrogeology.usgs.gov/download/Mars/GlobalSurveyor/MOLA/thumbs/Mars_MGS_MOLA_DEM_mosaic_global_1024.jpg -O ./data/mola_topo.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us load and visualize the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import imageio.v3 as iio\n",
    "\n",
    "img = iio.imread('./data/mola_topo.jpg')\n",
    "#convert to grayscale\n",
    "data = np.dot(img[...,:3]/255, [0.299, 0.587, 0.114])\n",
    "# interpolate onto 512x1024 grid:\n",
    "data = nn.functional.interpolate(torch.from_numpy(data).unsqueeze(0).unsqueeze(0), size=(512,1024)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 297356 characters]",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_theta = data.shape[0]\n",
    "n_lambda = data.shape[1]\n",
    "lon = np.linspace(-np.pi, np.pi, n_lambda)\n",
    "lat = np.linspace(np.pi/2., -np.pi/2., n_theta)\n",
    "Lon, Lat = np.meshgrid(lon, lat)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "ax.pcolormesh(Lon, Lat, data, cmap=cmap)\n",
    "# ax.set_title(\"Elevation map of mars\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want the data on the GPU. Moreover, `torch_harmonics` expects the last two dimensions to be the spatial dimensions. To this end, we permute the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = data.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spherical Harmonic Transform:\n",
    "\n",
    "To allow for composability, `torch_harmonics` provides it's transform in the form of PyTorch NN modules. As such we begin by initializing them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "sht = RealSHT(n_theta, n_lambda, grid=\"equiangular\").to(device)\n",
    "isht = InverseRealSHT(n_theta, n_lambda, grid=\"equiangular\").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "coeffs = sht(signal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us plot the logarithm of the absolute coefficients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 102928 characters]",
      "text/plain": [
       "<Figure size 432x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots(1,1,figsize=(6,5))\n",
    "plt.imshow(coeffs.abs().log().cpu(), cmap=cmap)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can transform the the solution back and plot the error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 350896 characters]",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "img = ax.pcolormesh(Lon, Lat, (isht(coeffs).cpu() - data).abs(), cmap=cmap, vmin=0, vmax=0.01)\n",
    "ax.set_title(\"Interpolation error\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.colorbar(img, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Differentiability\n",
    "\n",
    "`torch_harmonics` is differentiable, so let's use this functionality. We can perform a simple fit of spectral coefficients by using a simple model which will use 100 modes to approximate our original model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SpectralModel(\n",
       "  (isht): InverseRealSHT(\n",
       "    n_theta=80, n_lambda=160,\n",
       "     h_out=512, w_out=1024,\n",
       "     grid=equiangular\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class SpectralModel(nn.Module):\n",
    "    def __init__(self, n_modes, out_dims):\n",
    "        super().__init__()\n",
    "        self.coeffs = nn.Parameter(torch.randn(n_modes, n_modes+1, dtype=torch.complex128))\n",
    "        self.isht = InverseRealSHT(n_modes, 2*n_modes, h_out=out_dims[0], w_out=out_dims[1], grid=\"equiangular\").to(device)\n",
    "\n",
    "    def forward(self):\n",
    "        return self.isht(self.coeffs)\n",
    "\n",
    "sh_model = SpectralModel(80, (n_theta, n_lambda)).to(device)\n",
    "sh_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now, let us train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iteration: 0 loss: 489.0419848135532\n",
      "iteration: 10 loss: 50.536812313055094\n",
      "iteration: 20 loss: 18.2867276705511\n",
      "iteration: 30 loss: 7.539545962856474\n",
      "iteration: 40 loss: 2.596605806971006\n",
      "iteration: 50 loss: 0.9838933916680342\n",
      "iteration: 60 loss: 0.3480041213186319\n",
      "iteration: 70 loss: 0.13015090955555858\n",
      "iteration: 80 loss: 0.04624785018116071\n",
      "iteration: 90 loss: 0.016876006231934998\n",
      "iteration: 100 loss: 0.006266858247602376\n",
      "iteration: 110 loss: 0.0026305888565850166\n",
      "iteration: 120 loss: 0.0011998036288189506\n",
      "iteration: 130 loss: 0.000736531425067511\n",
      "iteration: 140 loss: 0.0005687002004316918\n",
      "iteration: 150 loss: 0.0005063323205870124\n",
      "iteration: 160 loss: 0.00048439627165954875\n",
      "iteration: 170 loss: 0.00047673294540564255\n",
      "iteration: 180 loss: 0.00047400688336708045\n",
      "iteration: 190 loss: 0.00047300015334753086\n",
      "iteration: 200 loss: 0.0004726731414181341\n",
      "iteration: 210 loss: 0.0004725495755762753\n",
      "iteration: 220 loss: 0.00047250564161943495\n",
      "iteration: 230 loss: 0.0004724907328314938\n",
      "iteration: 240 loss: 0.0004724854040734426\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(sh_model.parameters(), lr = 1e-1)\n",
    "\n",
    "for iter in range(250):\n",
    "\n",
    "    loss = (sh_model() - signal).pow(2).mean()\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if iter % 10 == 0:\n",
    "        print(f'iteration: {iter} loss: {loss.item()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let us plot the approximation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 219864 characters]",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "ax.pcolormesh(Lon, Lat, sh_model().detach().cpu(), cmap=cmap)\n",
    "# ax.set_title(\"Elevation map of mars\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not bad! Let us again plot the error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 170264 characters]",
      "text/plain": [
       "<Figure size 720x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(10, 5))\n",
    "ax = fig.add_subplot(1, 1, 1, projection='mollweide')\n",
    "img = ax.pcolormesh(Lon, Lat, (sh_model().detach() - signal).abs().cpu(), cmap=cmap, vmin=0, vmax=1)\n",
    "ax.set_title(\"Approximation error\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.colorbar(img, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This concludes our introduction to `torch-harmonics`. If you are enjoying this library, make sure to cite us:"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 64-bit",
   "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.10.6 (main, Nov 14 2022, 16:10:14) [GCC 11.3.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
