{
 "cells": [
  {
   "attachments": {},
   "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."
   ]
  },
  {
   "attachments": {},
   "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'"
   ]
  },
  {
   "attachments": {},
   "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 *"
   ]
  },
  {
   "attachments": {},
   "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')"
   ]
  },
  {
   "attachments": {},
   "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",
      "--2023-08-28 12:58:43--  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   109KB/s    in 0.5s    \n",
      "\n",
      "2023-08-28 12:58:45 (109 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"
   ]
  },
  {
   "attachments": {},
   "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()"
   ]
  },
  {
   "attachments": {},
   "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)"
   ]
  },
  {
   "attachments": {},
   "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)"
   ]
  },
  {
   "attachments": {},
   "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, 340640 characters]",
      "text/plain": [
       "<Figure size 1584x360 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(22, 5))\n",
    "ax = fig.add_subplot(1, 2, 1)\n",
    "plt.imshow(coeffs.abs().log().cpu(), cmap=cmap)\n",
    "plt.title(\"$\\log | SHT(u) |$\")\n",
    "plt.colorbar()\n",
    "\n",
    "ax = fig.add_subplot(1, 2, 2, projection='mollweide')\n",
    "plt.pcolormesh(Lon, Lat, isht(coeffs).cpu(), cmap=cmap)\n",
    "ax.set_title(\"Reconstructed elevation map\")\n",
    "ax.grid(True)\n",
    "ax.set_xticklabels([])\n",
    "ax.set_yticklabels([])\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "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, 348608 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()"
   ]
  },
  {
   "attachments": {},
   "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",
       "    nlat=512, nlon=1024,\n",
       "     lmax=120, mmax=121,\n",
       "     grid=equiangular, csphase=True\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(out_dims[0], out_dims[1], lmax=n_modes, mmax=n_modes+1, grid=\"equiangular\").to(device)\n",
    "\n",
    "    def forward(self):\n",
    "        return self.isht(self.coeffs)\n",
    "\n",
    "sh_model = SpectralModel(120, (n_theta, n_lambda)).to(device)\n",
    "sh_model"
   ]
  },
  {
   "attachments": {},
   "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: 1114.4378013113828\n",
      "iteration: 10 loss: 308.6806107370437\n",
      "iteration: 20 loss: 97.24193412864521\n",
      "iteration: 30 loss: 34.48748205511207\n",
      "iteration: 40 loss: 13.27448603501562\n",
      "iteration: 50 loss: 5.325563991891355\n",
      "iteration: 60 loss: 2.2486285268724533\n",
      "iteration: 70 loss: 0.9864678073689532\n",
      "iteration: 80 loss: 0.4463606510988175\n",
      "iteration: 90 loss: 0.20572443417422379\n",
      "iteration: 100 loss: 0.09943813502065274\n",
      "iteration: 110 loss: 0.04926881816439331\n",
      "iteration: 120 loss: 0.024924083750881523\n",
      "iteration: 130 loss: 0.012995443120258202\n",
      "iteration: 140 loss: 0.006897149705030453\n",
      "iteration: 150 loss: 0.003725578972393333\n",
      "iteration: 160 loss: 0.002063529906069868\n",
      "iteration: 170 loss: 0.0011982119677471067\n",
      "iteration: 180 loss: 0.0007490736206441505\n",
      "iteration: 190 loss: 0.0005194080521105679\n",
      "iteration: 200 loss: 0.0004041695627585252\n",
      "iteration: 210 loss: 0.00034752646960863065\n",
      "iteration: 220 loss: 0.0003202664238212154\n",
      "iteration: 230 loss: 0.00030745193490353955\n",
      "iteration: 240 loss: 0.00030156991359177326\n",
      "iteration: 250 loss: 0.0002989361504325404\n",
      "iteration: 260 loss: 0.0002977872136962757\n",
      "iteration: 270 loss: 0.0002973001713235517\n",
      "iteration: 280 loss: 0.0002971001235014028\n",
      "iteration: 290 loss: 0.00029702186077974\n",
      "iteration: 300 loss: 0.0002978981918974569\n",
      "iteration: 310 loss: 0.00029971193546475474\n",
      "iteration: 320 loss: 0.0002979065155556966\n",
      "iteration: 330 loss: 0.0003072179364963948\n",
      "iteration: 340 loss: 0.0002995392828352215\n",
      "iteration: 350 loss: 0.0002970504961957656\n",
      "iteration: 360 loss: 0.0002974129281438742\n",
      "iteration: 370 loss: 0.0002970145080492319\n",
      "iteration: 380 loss: 0.00029697935480860785\n",
      "iteration: 390 loss: 0.0002969858026193997\n",
      "iteration: 400 loss: 0.0002969800186391534\n",
      "iteration: 410 loss: 0.00029697604232545294\n",
      "iteration: 420 loss: 0.00029697424115560106\n",
      "iteration: 430 loss: 0.00029697367805708874\n",
      "iteration: 440 loss: 0.00029697367014723704\n",
      "iteration: 450 loss: 0.0002969736811532905\n",
      "iteration: 460 loss: 0.00029697363618637953\n",
      "iteration: 470 loss: 0.0002969736761192982\n",
      "iteration: 480 loss: 0.00029698408676114024\n",
      "iteration: 490 loss: 0.0003110991882101014\n"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(sh_model.parameters(), lr = 5e-2)\n",
    "\n",
    "for iter in range(500):\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()}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let us plot the approximation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 243748 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()"
   ]
  },
  {
   "attachments": {},
   "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, 161660 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()"
   ]
  },
  {
   "attachments": {},
   "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.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
