{
 "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+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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting the Spherical Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from plotting import plot_sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 224872 characters]",
      "text/plain": [
       "<Figure size 864x576 with 36 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_lmax = 8\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}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's also make an animated version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 864x576 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "frame=0\n",
      "frame=10\n",
      "frame=20\n",
      "frame=30\n",
      "frame=40\n",
      "frame=50\n",
      "frame=60\n",
      "frame=70\n",
      "frame=80\n",
      "frame=90\n",
      "frame=100\n",
      "frame=110\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 5432 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "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], 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
}
