{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyzing the gradients of the SHT\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from torch_harmonics.quadrature import legendre_gauss_weights, clenshaw_curtiss_weights\n",
    "from torch_harmonics.legendre import legpoly, clm\n",
    "from torch_harmonics import RealSHT, InverseRealSHT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem setting\n",
    "\n",
    "We consider the simple problem of fitting the spectral coefficients $\\theta$ such that\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\mathcal{L}\n",
    "&= ||\\mathop{\\mathrm{ISHT}}[\\theta] - u^*||^2_{S^2} \\\\\n",
    "&\\approx \\sum_j \\omega_j (\\mathop{\\mathrm{ISHT}}[\\theta](x_j) - u^*(x_j))^2 \\\\\n",
    "&= (S \\, \\theta - u^*)^T \\mathop{\\mathrm{diag}}(\\omega) \\, (S \\, \\theta - u^*) \\\\\n",
    "&= L\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "is minimized.\n",
    "\n",
    "The Vandermonde matrix $S$, which is characterized by $\\mathop{\\mathrm{ISHT}}[\\theta] = S \\theta$ realizes the action of the discrete SHT.\n",
    "\n",
    "The necessary condition for a minimizer of $L$ is\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "& \\nabla_\\theta L = S^T \\mathop{\\mathrm{diag}}(\\omega) \\, (S \\, \\theta - u^*) = 0 \\\\\n",
    "\\Leftrightarrow \\quad & S^T \\mathop{\\mathrm{diag}}(\\omega) \\, S \\; \\theta = S^T \\mathop{\\mathrm{diag}}(\\omega) \\, u^*.\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "On the Gaussian grid, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(12.5664)\n"
     ]
    }
   ],
   "source": [
    "nlat = 64\n",
    "nlon = 2*nlat\n",
    "grid = \"equiangular\"\n",
    "\n",
    "# for quadrature and plotting\n",
    "if grid == \"legendre-gauss\":\n",
    "    lmax = mmax = nlat\n",
    "    xq, wq = legendre_gauss_weights(nlat)\n",
    "elif grid ==\"equiangular\":\n",
    "    lmax = mmax = nlat//2\n",
    "    xq, wq = clenshaw_curtiss_weights(nlat)\n",
    "\n",
    "sht = RealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid)\n",
    "isht = InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid)\n",
    "\n",
    "lat = torch.arccos(xq)\n",
    "omega = torch.pi * wq.float() / nlat\n",
    "omega = omega.reshape(-1, 1)\n",
    "\n",
    "print(nlon*omega.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## analytic  gradients of the SHT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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:41:34--  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   174KB/s    in 0.3s    \n",
      "\n",
      "2025-03-20 22:41:35 (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\n",
    "\n",
    "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=(nlat,nlon)).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7cefc14e5090>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 32512 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torch_harmonics.plotting import plot_sphere\n",
    "\n",
    "plot_sphere(data, projection=\"orthographic\", cmap=\"turbo\", colorbar=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the analytic gradients are of the discrete SHT can be determined using some linear algebra. If the analysis step is $\\theta = S^\\dagger \\mathop{\\mathrm{diag}}(\\omega) u$, we can determine that $J_u \\theta = \\mathop{\\mathrm{diag}}(\\omega) S$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-5.1482e-12, -5.1482e-12, -5.1482e-12,  ..., -5.1482e-12,\n",
       "         -5.1482e-12, -5.1482e-12],\n",
       "        [-1.0307e-10, -6.7286e-11, -3.3095e-11,  ..., -2.1765e-10,\n",
       "         -1.7852e-10, -1.4024e-10],\n",
       "        [-4.1769e-10, -5.1147e-10, -5.8795e-10,  ..., -3.6529e-11,\n",
       "         -1.7920e-10, -3.0670e-10],\n",
       "        ...,\n",
       "        [ 7.5371e-10,  7.6736e-10,  7.8559e-10,  ...,  7.4207e-10,\n",
       "          7.4091e-10,  7.4484e-10],\n",
       "        [-5.9529e-10, -5.9526e-10, -5.9613e-10,  ..., -6.0131e-10,\n",
       "         -5.9828e-10, -5.9629e-10],\n",
       "        [ 3.3697e-10,  3.3697e-10,  3.3697e-10,  ...,  3.3697e-10,\n",
       "          3.3697e-10,  3.3697e-10]], dtype=torch.float64)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u = data.clone()\n",
    "u.requires_grad_(True)\n",
    "\n",
    "theta = sht(u)\n",
    "\n",
    "# direction for backward pass\n",
    "gin = torch.randn_like(theta, requires_grad=False)\n",
    "\n",
    "# gradient computed using autodiff\n",
    "theta.backward(gin)\n",
    "\n",
    "# print(u.grad)\n",
    "\n",
    "# analytic gradient computed using ISHT\n",
    "mweights = torch.ones(1, gin.shape[-1], dtype=gin.dtype)\n",
    "mweights[..., 1:] *= 0.5\n",
    "ugrad = omega * isht(mweights * gin)\n",
    "\n",
    "(u.grad - ugrad) / u.grad.max().abs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7cefb44402e0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 136016 characters]",
      "text/plain": [
       "<Figure size 1440x432 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(layout='constrained', figsize=(20, 6))\n",
    "subfigs = fig.subfigures(1, 3)\n",
    "\n",
    "# spectral fitting\n",
    "plot_sphere(u.grad.detach(), fig=subfigs[0], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"autograd\")\n",
    "plot_sphere(ugrad.detach(), fig=subfigs[1], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"analytic\")\n",
    "plot_sphere((ugrad - u.grad).detach() / u.grad.max().abs(), fig=subfigs[2], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"error\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## using Newton method to find the SH weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 0, loss: 494.6172995679299\n",
      "iter: 1, loss: 0.008023963421976575\n",
      "iter: 2, loss: 0.0080239639324424\n",
      "iter: 3, loss: 0.00802396362808806\n",
      "iter: 4, loss: 0.008023963341038724\n",
      "iter: 5, loss: 0.008023963360998972\n",
      "iter: 6, loss: 0.008023964004906828\n",
      "iter: 7, loss: 0.008023963539076105\n",
      "iter: 8, loss: 0.008023963004288193\n",
      "iter: 9, loss: 0.008023963616464676\n",
      "iter: 10, loss: 0.008023963441370515\n",
      "iter: 11, loss: 0.008023963158389706\n",
      "iter: 12, loss: 0.008023963874548495\n",
      "iter: 13, loss: 0.008023963932492268\n",
      "iter: 14, loss: 0.00802396306221282\n",
      "iter: 15, loss: 0.008023964035021834\n",
      "iter: 16, loss: 0.008023963549934635\n",
      "iter: 17, loss: 0.008023962661971371\n",
      "iter: 18, loss: 0.008023964071981958\n",
      "iter: 19, loss: 0.008023963595639045\n",
      "iter: 20, loss: 0.008023962643703104\n",
      "iter: 21, loss: 0.00802396378849733\n",
      "iter: 22, loss: 0.008023962831258453\n",
      "iter: 23, loss: 0.008023963407114714\n",
      "iter: 24, loss: 0.008023963982757057\n",
      "iter: 25, loss: 0.008023963384387687\n",
      "iter: 26, loss: 0.00802396333344522\n",
      "iter: 27, loss: 0.008023963970227685\n",
      "iter: 28, loss: 0.008023963614997916\n",
      "iter: 29, loss: 0.008023963503094355\n",
      "iter: 30, loss: 0.008023963176716418\n",
      "iter: 31, loss: 0.008023962684003006\n",
      "iter: 32, loss: 0.008023963188028484\n",
      "iter: 33, loss: 0.008023963627171175\n",
      "iter: 34, loss: 0.008023963529420412\n",
      "iter: 35, loss: 0.008023963566632819\n",
      "iter: 36, loss: 0.008023964042094543\n",
      "iter: 37, loss: 0.008023963302599049\n",
      "iter: 38, loss: 0.008023963310401236\n",
      "iter: 39, loss: 0.008023963388689038\n"
     ]
    }
   ],
   "source": [
    "lr = 1.0\n",
    "theta = nn.Parameter(torch.randn(lmax, lmax, dtype=torch.complex64))\n",
    "optim = torch.optim.SGD([theta], lr=lr)\n",
    "\n",
    "for iter in range(40):\n",
    "    optim.zero_grad(set_to_none=True)\n",
    "    loss = torch.sum(0.5*omega*(isht(theta) - data)**2)\n",
    "    loss.backward()\n",
    "\n",
    "    # action of the Hessian\n",
    "    with torch.no_grad():\n",
    "        for m in range(1,mmax):\n",
    "            theta.grad[:,m].mul_(0.5)\n",
    "\n",
    "    optim.step()\n",
    "\n",
    "    print(f\"iter: {iter}, loss: {loss}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "what's the best possible loss? $\\theta^* = (S^T \\mathop{\\mathrm{diag}}(\\omega) \\, S)^{-1} S^T \\mathop{\\mathrm{diag}}(\\omega) u^* = \\mathop{\\mathrm{SHT}}[u^*]$ gives us the global minimizer for this problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7cef9e48c070>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 316496 characters]",
      "text/plain": [
       "<Figure size 1440x864 with 12 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(layout='constrained', figsize=(20, 12))\n",
    "subfigs = fig.subfigures(2, 3)\n",
    "\n",
    "# spectral fitting\n",
    "plot_sphere(isht(theta).detach(), fig=subfigs[0,0], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"Fit\")\n",
    "plot_sphere(data, fig=subfigs[0,1], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"Ground truth\")\n",
    "plot_sphere((isht(theta) - data).detach(), fig=subfigs[0,2], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"residual\")\n",
    "\n",
    "# sht(u)\n",
    "plot_sphere(isht(sht(data)).detach(), fig=subfigs[1,0], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"isht(sht(u))\")\n",
    "plot_sphere(data, fig=subfigs[1,1], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"Ground truth\")\n",
    "plot_sphere((isht(sht(data)) - data).detach(), fig=subfigs[1,2], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"residual\")"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
