{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the Spherical harmonics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "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\n",
    "from torch_harmonics.legendre import legpoly, clm\n",
    "from torch_harmonics import RealSHT, InverseRealSHT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forming the Vandermonde matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlat = 60\n",
    "nlon = 2*nlat\n",
    "lmax = mmax = nlat\n",
    "\n",
    "sht = RealSHT(nlat, nlon, lmax=lmax, mmax=mmax)\n",
    "isht = InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# forming the Vandermonde matrix\n",
    "\n",
    "nmodes = int(lmax * lmax)\n",
    "e = torch.zeros(nmodes, lmax, mmax, dtype=torch.complex64)\n",
    "\n",
    "midx = lambda l,m : l*(l+1) + m\n",
    "\n",
    "for l in range(lmax):\n",
    "    for m in range(-l, l+1):\n",
    "        e[midx(l,m), l, abs(m)] = 1.0 if m >= 0 else 1.0j\n",
    "\n",
    "vdm = isht(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "abs(-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "midx(1,-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.+0.j)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e[1,1,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting the Spherical Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from torch_harmonics.plotting import plot_sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 175284 characters]",
      "text/plain": [
       "<Figure size 864x576 with 9 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_lmax = 3\n",
    "\n",
    "fig = plt.figure(layout='constrained', figsize=(12, 8))\n",
    "subfigs = fig.subfigures(plt_lmax, 2*plt_lmax-1)\n",
    "\n",
    "for l in range(plt_lmax):\n",
    "    for m in range(-l, l+1):\n",
    "        plot_sphere(vdm[midx(l,m)], fig=subfigs[l, m+plt_lmax-1], projection=\"orthographic\", title=f\"l={l}, m={m}\", central_latitude=20, central_longitude=30, cmap=\"plasma\")\n",
    "\n",
    "plt.savefig(\"spherical_harmonics.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's also make an animated version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import matplotlib.animation as animation\n",
    "\n",
    "# fig = plt.figure(layout='constrained', figsize=(12, 8), dpi=72)\n",
    "# subfigs = fig.subfigures(plt_lmax, plt_lmax)\n",
    "\n",
    "# moviewriter = animation.writers['pillow'](fps=30)\n",
    "# moviewriter.setup(fig, 'spherical_harmonics.gif', dpi=72)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for frame in range(120):\n",
    "#     # compute the rotation of the sphere\n",
    "#     lon = -120 - 3 * frame\n",
    "\n",
    "#     if frame % 10 == 0:\n",
    "#         print(f\"frame={frame}\")\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], projection=\"orthographic\", title=f\"l={l}, m={m}\", central_longitude=lon)\n",
    "\n",
    "#     plt.draw()\n",
    "#     moviewriter.grab_frame()\n",
    "\n",
    "# moviewriter.finish()"
   ]
  }
 ],
 "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
}
