{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import torch\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\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from plotting import plot_sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(12.5664)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nlat = 32\n",
    "nlon = 2*nlat\n",
    "grid = \"legendre-gauss\"\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-1\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).repeat(1, nlon)\n",
    "\n",
    "omega.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forming the Vandermonde matrix\n",
    "nmodes = int(lmax * (lmax+1) / 2)\n",
    "e = torch.zeros(nmodes, lmax, mmax, dtype=torch.complex64)\n",
    "\n",
    "midx = lambda l,m : l*(l+1)//2 + m\n",
    "\n",
    "for l in range(lmax):\n",
    "    for m in range(l+1):\n",
    "        e[midx(l,m), l, m] = 1.0\n",
    "\n",
    "vdm = isht(e) # nmodes x nlat x nlon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let us visualize the some entries in the Vandermonde matrix. In principle, it should contain the Spherical Harmonics in ascending order:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 158752 characters]",
      "text/plain": [
       "<Figure size 864x576 with 21 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_lmax = 6\n",
    "\n",
    "fig = plt.figure(layout='constrained', figsize=(12, 8))\n",
    "subfigs = fig.subfigures(plt_lmax, plt_lmax)\n",
    "\n",
    "for l in range(plt_lmax):\n",
    "    for m in range(l+1):\n",
    "        plot_sphere(vdm[midx(l,m)], fig=subfigs[l, m], title=f\"l={l}, m={m}\", lat=(lat - np.pi/2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "now, let us compute the gramian matrix $S^T \\mathop{\\mathrm{diag}}(\\omega) S$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "gramian = torch.einsum(\"iqr,jqr,qr->ij\", vdm, vdm, omega)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plt.pcolormesh(xlg, np.arange(0, nmodes), , cmap=\"plasma\")\n",
    "# plt.xlabel(\"x\")\n",
    "# plt.ylabel(\"l,m mode\")\n",
    "# plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f7c3704c250>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 9856 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(gramian, interpolation=\"nearest\", cmap=\"plasma\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(2.0000)\n",
      "tensor(23276)\n",
      "tensor(-1.2321e-06)\n",
      "tensor(198430)\n"
     ]
    }
   ],
   "source": [
    "print(torch.max(gramian))\n",
    "print(torch.argmax(gramian))\n",
    "print(torch.min(gramian))\n",
    "print(torch.argmin(gramian))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.0000079"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.linalg.cond(gramian)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f7c36f91930>]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 18292 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(np.diag(gramian))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forming the Vandermonde matrix\n",
    "nmodes = int(lmax * (lmax+1) / 2)\n",
    "\n",
    "\n",
    "e = torch.zeros(lmax, mmax, lmax, mmax, dtype=torch.complex64)\n",
    "for l in range(lmax):\n",
    "    for m in range(l+1):\n",
    "        e[l, m, l, m] = 1.0\n",
    "\n",
    "Sinv = isht(e).cfloat() # l x m x nlat x nlon\n",
    "\n",
    "e = torch.zeros(nlat, nlon, nlat, nlon, dtype=torch.float32)\n",
    "for lat in range(nlat):\n",
    "    for lon in range(nlon):\n",
    "        e[lat, lon, lat, lon] = 1.0\n",
    "\n",
    "S = sht(e) # nlat x nlon x l x m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = torch.einsum(\"lmop,lmqr->lopqr\", Sinv, Sinv.conj())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7f7c1e67ad10>"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31620 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_sphere(out[3, 29, 16].abs(), colorbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS = torch.einsum(\"lmqr, oplm->qrop\", Sinv, S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f7c36dff370>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 49928 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(SS.reshape(nlat*nlon, nlat*nlon).abs(), interpolation=\"nearest\", cmap=\"plasma\")\n",
    "plt.colorbar()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7f7c3f10ea40>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 9764 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(SS[:,0,:,0].abs(), interpolation=\"nearest\", cmap=\"plasma\")\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x7f7c3f04f700>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 13708 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 23296 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "field = torch.zeros(nlat, nlon)\n",
    "field[16,3] = 1.0\n",
    "\n",
    "plot_sphere(field, colorbar=True)\n",
    "plot_sphere(isht(sht(field)), colorbar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
