{
 "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": 16,
   "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": 17,
   "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": 18,
   "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": 19,
   "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",
      "--2025-01-09 00:30:29--  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... "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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   169KB/s    in 0.3s    \n",
      "\n",
      "2025-01-09 00:30:50 (169 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": 20,
   "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": 21,
   "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": 22,
   "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": 23,
   "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": 24,
   "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": 25,
   "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": 26,
   "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": 27,
   "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": 27,
     "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": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iteration: 0 loss: 1110.5501509434434\n",
      "iteration: 10 loss: 307.2719964223394\n",
      "iteration: 20 loss: 94.85549328211366\n",
      "iteration: 30 loss: 33.84012136397667\n",
      "iteration: 40 loss: 13.27256860537048\n",
      "iteration: 50 loss: 5.330296084257617\n",
      "iteration: 60 loss: 2.22780331377308\n",
      "iteration: 70 loss: 0.9787531762231616\n",
      "iteration: 80 loss: 0.4689970429107282\n",
      "iteration: 90 loss: 0.23740534161941174\n",
      "iteration: 100 loss: 0.12912743071119842\n",
      "iteration: 110 loss: 0.07353356327159222\n",
      "iteration: 120 loss: 0.04269109352645242\n",
      "iteration: 130 loss: 0.025076285913135578\n",
      "iteration: 140 loss: 0.014735300559362162\n",
      "iteration: 150 loss: 0.008633168595399724\n",
      "iteration: 160 loss: 0.005042391788189882\n",
      "iteration: 170 loss: 0.002959286307408625\n",
      "iteration: 180 loss: 0.0017676892624043142\n",
      "iteration: 190 loss: 0.0010960600006928898\n",
      "iteration: 200 loss: 0.0007237718646435973\n",
      "iteration: 210 loss: 0.0005208526650641412\n",
      "iteration: 220 loss: 0.0004121851107413985\n",
      "iteration: 230 loss: 0.00035507011897227737\n",
      "iteration: 240 loss: 0.00032563989271714995\n",
      "iteration: 250 loss: 0.000310792061015569\n",
      "iteration: 260 loss: 0.0003034701125007225\n",
      "iteration: 270 loss: 0.0002999472249394351\n",
      "iteration: 280 loss: 0.0002982967943372406\n",
      "iteration: 290 loss: 0.0003016019634782542\n",
      "iteration: 300 loss: 0.000303156822051305\n",
      "iteration: 310 loss: 0.000297139553423036\n",
      "iteration: 320 loss: 0.00030631147429905335\n",
      "iteration: 330 loss: 0.0002996713866962133\n",
      "iteration: 340 loss: 0.0003168400443074444\n",
      "iteration: 350 loss: 0.00030543049738370874\n",
      "iteration: 360 loss: 0.0003402379769750154\n",
      "iteration: 370 loss: 0.0003095132951485824\n",
      "iteration: 380 loss: 0.0002978572515053569\n",
      "iteration: 390 loss: 0.00030888117299483997\n",
      "iteration: 400 loss: 0.00030039260805842823\n",
      "iteration: 410 loss: 0.00029768920329069414\n",
      "iteration: 420 loss: 0.0002971736682496712\n",
      "iteration: 430 loss: 0.00029714712379955404\n",
      "iteration: 440 loss: 0.00029698562925891184\n",
      "iteration: 450 loss: 0.00029698897310414857\n",
      "iteration: 460 loss: 0.0002969805265735664\n",
      "iteration: 470 loss: 0.00029697392337402185\n",
      "iteration: 480 loss: 0.0002969738702713484\n",
      "iteration: 490 loss: 0.0002969738774566405\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": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 242176 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, 159476 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",
   "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
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
