{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyzing the gradients of the SHT\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "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 = np.arccos(xq)\n",
    "omega = torch.pi * torch.from_numpy(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": 3,
   "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",
      "--2024-09-19 22:55:03--  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)... "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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   151KB/s    in 0.4s    \n",
      "\n",
      "2024-09-19 22:55:04 (151 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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7a4199b49990>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44280 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 plotting import plot_sphere\n",
    "\n",
    "plot_sphere(data, 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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 2.4418e-10,  2.4418e-10,  2.4418e-10,  ...,  2.4418e-10,\n",
       "          2.4418e-10,  2.4418e-10],\n",
       "        [-6.0433e-10, -5.9786e-10, -5.9359e-10,  ..., -6.3639e-10,\n",
       "         -6.2369e-10, -6.1297e-10],\n",
       "        [-7.2399e-10, -7.6395e-10, -7.8275e-10,  ..., -4.9298e-10,\n",
       "         -5.8655e-10, -6.6427e-10],\n",
       "        ...,\n",
       "        [ 1.0240e-09,  1.0131e-09,  9.9555e-10,  ...,  1.0141e-09,\n",
       "          1.0246e-09,  1.0279e-09],\n",
       "        [-4.1902e-10, -4.1270e-10, -4.0478e-10,  ..., -4.2815e-10,\n",
       "         -4.2676e-10, -4.2372e-10],\n",
       "        [-1.0371e-11, -1.0371e-11, -1.0371e-11,  ..., -1.0371e-11,\n",
       "         -1.0371e-11, -1.0371e-11]], dtype=torch.float64)"
      ]
     },
     "execution_count": 5,
     "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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7a418dfbad10>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 223976 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], cmap=\"turbo\", colorbar=True, title=\"autograd\")\n",
    "plot_sphere(ugrad.detach(), fig=subfigs[1], cmap=\"turbo\", colorbar=True, title=\"analytic\")\n",
    "plot_sphere((ugrad - u.grad).detach() / u.grad.max().abs(), fig=subfigs[2], cmap=\"turbo\", colorbar=True, title=\"error\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## using Newton method to find the SH weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 0, loss: 524.3101699017082\n",
      "iter: 1, loss: 0.00802396445458801\n",
      "iter: 2, loss: 0.008023964153899235\n",
      "iter: 3, loss: 0.00802396394417556\n",
      "iter: 4, loss: 0.00802396352995326\n",
      "iter: 5, loss: 0.008023963385773996\n",
      "iter: 6, loss: 0.008023963547206097\n",
      "iter: 7, loss: 0.008023963465045698\n",
      "iter: 8, loss: 0.008023963682567867\n",
      "iter: 9, loss: 0.008023963892175393\n",
      "iter: 10, loss: 0.008023963426161781\n",
      "iter: 11, loss: 0.008023963220818102\n",
      "iter: 12, loss: 0.008023964286015772\n",
      "iter: 13, loss: 0.008023963408898455\n",
      "iter: 14, loss: 0.0080239633010406\n",
      "iter: 15, loss: 0.00802396354274379\n",
      "iter: 16, loss: 0.008023963894971519\n",
      "iter: 17, loss: 0.00802396273543675\n",
      "iter: 18, loss: 0.008023963845382277\n",
      "iter: 19, loss: 0.008023963170493507\n",
      "iter: 20, loss: 0.008023963538155878\n",
      "iter: 21, loss: 0.008023963736162949\n",
      "iter: 22, loss: 0.008023963007515712\n",
      "iter: 23, loss: 0.00802396251903999\n",
      "iter: 24, loss: 0.008023963512694849\n",
      "iter: 25, loss: 0.008023963303701073\n",
      "iter: 26, loss: 0.008023963528319265\n",
      "iter: 27, loss: 0.008023964187561856\n",
      "iter: 28, loss: 0.00802396545716888\n",
      "iter: 29, loss: 0.008023963124880929\n",
      "iter: 30, loss: 0.008023963778864577\n",
      "iter: 31, loss: 0.008023963064092091\n",
      "iter: 32, loss: 0.008023963720056796\n",
      "iter: 33, loss: 0.00802396347358365\n",
      "iter: 34, loss: 0.008023964103128693\n",
      "iter: 35, loss: 0.008023963717014568\n",
      "iter: 36, loss: 0.00802396302845848\n",
      "iter: 37, loss: 0.008023963326501782\n",
      "iter: 38, loss: 0.008023964046142596\n",
      "iter: 39, loss: 0.008023963855775787\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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7a4188b0b850>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 461176 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], cmap=\"turbo\", colorbar=True, title=\"Fit\")\n",
    "plot_sphere(data, fig=subfigs[0,1], cmap=\"turbo\", colorbar=True, title=\"Ground truth\")\n",
    "plot_sphere((isht(theta) - data).detach(), fig=subfigs[0,2], cmap=\"turbo\", colorbar=True, title=\"residual\")\n",
    "\n",
    "# sht(u)\n",
    "plot_sphere(isht(sht(data)).detach(), fig=subfigs[1,0], cmap=\"turbo\", colorbar=True, title=\"isht(sht(u))\")\n",
    "plot_sphere(data, fig=subfigs[1,1], cmap=\"turbo\", colorbar=True, title=\"Ground truth\")\n",
    "plot_sphere((isht(sht(data)) - data).detach(), fig=subfigs[1,2], 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
}
