{
 "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": [
    {
     "data": {
      "text/plain": [
       "tensor(12.5664)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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",
    "nlon*omega.sum()"
   ]
  },
  {
   "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",
      "--2023-10-30 18:00:14--  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   154KB/s    in 0.3s    \n",
      "\n",
      "2023-10-30 18:00:15 (154 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 0x7f49e4952380>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44276 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": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 0, loss: 453.0968931302793\n",
      "iter: 1, loss: 0.008023964326606358\n",
      "iter: 2, loss: 0.008023963388341868\n",
      "iter: 3, loss: 0.008023963340660247\n",
      "iter: 4, loss: 0.008023963596959654\n",
      "iter: 5, loss: 0.008023963735337598\n",
      "iter: 6, loss: 0.008023964260612844\n",
      "iter: 7, loss: 0.008023964042363394\n",
      "iter: 8, loss: 0.00802396368406042\n",
      "iter: 9, loss: 0.008023962714947052\n",
      "iter: 10, loss: 0.008023963489819921\n",
      "iter: 11, loss: 0.008023963701078593\n",
      "iter: 12, loss: 0.008023962923266034\n",
      "iter: 13, loss: 0.008023964198518512\n",
      "iter: 14, loss: 0.008023962813486126\n",
      "iter: 15, loss: 0.008023964110803488\n",
      "iter: 16, loss: 0.00802396403813473\n",
      "iter: 17, loss: 0.008023963786036484\n",
      "iter: 18, loss: 0.008023964195574898\n",
      "iter: 19, loss: 0.008023963516124565\n",
      "iter: 20, loss: 0.008023964508201684\n",
      "iter: 21, loss: 0.008023963767474551\n",
      "iter: 22, loss: 0.008023963648388185\n",
      "iter: 23, loss: 0.008023963972575866\n",
      "iter: 24, loss: 0.008023964038780116\n",
      "iter: 25, loss: 0.008023963707541834\n",
      "iter: 26, loss: 0.008023963269911932\n",
      "iter: 27, loss: 0.008023963391352053\n",
      "iter: 28, loss: 0.008023963414851426\n",
      "iter: 29, loss: 0.008023964147064296\n",
      "iter: 30, loss: 0.008023963760174639\n",
      "iter: 31, loss: 0.008023963924162339\n",
      "iter: 32, loss: 0.00802396360354566\n",
      "iter: 33, loss: 0.00802396407422616\n",
      "iter: 34, loss: 0.008023962918493041\n",
      "iter: 35, loss: 0.008023963622013491\n",
      "iter: 36, loss: 0.0080239635670241\n",
      "iter: 37, loss: 0.008023963871070301\n",
      "iter: 38, loss: 0.008023963587685968\n",
      "iter: 39, loss: 0.008023963496770136\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}\")\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",
    "#     with torch.no_grad():\n",
    "#         theta.add_(theta.grad, alpha=-lr)\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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7f49d214b9a0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 461168 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
}
