{
 "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": 7,
   "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-03-20 22:30:53--  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   174KB/s    in 0.3s    \n",
      "\n",
      "2025-03-20 22:30:54 (174 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": 8,
   "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": 9,
   "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": 10,
   "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": 11,
   "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": 12,
   "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": 13,
   "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": 14,
   "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": 15,
   "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": 15,
     "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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iteration: 0 loss: 1151.2896154798989\n",
      "iteration: 10 loss: 318.96639669728927\n",
      "iteration: 20 loss: 101.45816768663033\n",
      "iteration: 30 loss: 35.91950678733323\n",
      "iteration: 40 loss: 13.505140650990615\n",
      "iteration: 50 loss: 5.231296587249963\n",
      "iteration: 60 loss: 2.0786350484945997\n",
      "iteration: 70 loss: 0.8521032756663567\n",
      "iteration: 80 loss: 0.3662866106013982\n",
      "iteration: 90 loss: 0.1604243831943348\n",
      "iteration: 100 loss: 0.07254225429192539\n",
      "iteration: 110 loss: 0.03416170266976784\n",
      "iteration: 120 loss: 0.01661330766930938\n",
      "iteration: 130 loss: 0.008195295534645642\n",
      "iteration: 140 loss: 0.0041078305295401615\n",
      "iteration: 150 loss: 0.002102936422261028\n",
      "iteration: 160 loss: 0.0011401038639609626\n",
      "iteration: 170 loss: 0.0006832547855850787\n",
      "iteration: 180 loss: 0.0004690471929420203\n",
      "iteration: 190 loss: 0.00037255164472209997\n",
      "iteration: 200 loss: 0.0003278577921196046\n",
      "iteration: 210 loss: 0.000309609375132965\n",
      "iteration: 220 loss: 0.00030175472568891976\n",
      "iteration: 230 loss: 0.0002986564600665542\n",
      "iteration: 240 loss: 0.0002975273681647285\n",
      "iteration: 250 loss: 0.0003143859548930335\n",
      "iteration: 260 loss: 0.00030096663360024476\n",
      "iteration: 270 loss: 0.00029826281710717113\n",
      "iteration: 280 loss: 0.0002969830175221479\n",
      "iteration: 290 loss: 0.0002971419917634186\n",
      "iteration: 300 loss: 0.00029705558647964133\n",
      "iteration: 310 loss: 0.00031499472201119766\n",
      "iteration: 320 loss: 0.0003000265202913814\n",
      "iteration: 330 loss: 0.0003079234939104373\n",
      "iteration: 340 loss: 0.00030121188521898954\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iteration: 350 loss: 0.0002982996964641134\n",
      "iteration: 360 loss: 0.0002993674045253317\n",
      "iteration: 370 loss: 0.00029768720652354505\n",
      "iteration: 380 loss: 0.0003026274490778747\n",
      "iteration: 390 loss: 0.00031050489651539693\n",
      "iteration: 400 loss: 0.00029796251614888994\n",
      "iteration: 410 loss: 0.0002978342764537453\n",
      "iteration: 420 loss: 0.00032767747129473996\n",
      "iteration: 430 loss: 0.0003050815921180515\n",
      "iteration: 440 loss: 0.0002973681673583962\n",
      "iteration: 450 loss: 0.00030011220644066025\n",
      "iteration: 460 loss: 0.00029795456294238744\n",
      "iteration: 470 loss: 0.00029865453432689727\n",
      "iteration: 480 loss: 0.00029797776542376863\n",
      "iteration: 490 loss: 0.00030059723348185055\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": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 242432 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": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 159028 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
}
