{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Derivatives of $\\theta$ and $\\phi$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1D FFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grid setup\n",
    "N = 128\n",
    "L = 1\n",
    "x = np.linspace(0, L, N, endpoint=False)\n",
    "\n",
    "# Function definition\n",
    "f = lambda x: np.sin(4 * np.pi * x) \n",
    "f_prime = lambda x: 4 * np.pi * np.cos(4 * np.pi * x)\n",
    "\n",
    "# Compute the discrete functions\n",
    "f_h = f(x)\n",
    "f_prime_h = f_prime(x)\n",
    "\n",
    "# FFT Derivative\n",
    "k = np.fft.fftfreq(N, d=L/(2 * np.pi * N))\n",
    "f_prime_fft_h = np.fft.ifft(1j * k * np.fft.fft(f_h)).real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95036 characters]",
      "text/plain": [
       "<Figure size 1000x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot results \n",
    "fig, axs = plt.subplots(1, 2, figsize=(10, 4))\n",
    "axs[0].plot(x, f_h, label=r'$f (x)$')\n",
    "axs[0].plot(x, f_prime_h, alpha=0.9, label=r'$g (x)$')\n",
    "axs[0].scatter(x, f_prime_fft_h, marker=\".\", alpha=0.9, label=r'$g_{fft} (x)$')\n",
    "axs[0].set_xlabel(r\"$X$\")\n",
    "axs[0].set_ylabel(r\"$u(x)$\")\n",
    "axs[0].legend()\n",
    "\n",
    "axs[1].semilogy(x, np.abs(f_prime_h - f_prime_fft_h), label=r'$\\epsilon = |g(x) - g_{fft}(x)|$')\n",
    "axs[1].set_xlabel(r\"$X$\")\n",
    "axs[1].set_ylabel(r\"$\\epsilon$\")\n",
    "\n",
    "axs[1].legend()\n",
    "plt.tight_layout()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2D FFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up grid and domain\n",
    "N_phi = 128  # grid size\n",
    "N_theta = N_phi // 2 + 1  # grid size\n",
    "L_theta = 1  # domain size\n",
    "L_phi = 1  # domain size\n",
    "\n",
    "mesh_theta = np.linspace(0, L_theta, N_theta, endpoint=False)\n",
    "mesh_phi = np.linspace(0, L_phi, N_phi, endpoint=False)\n",
    "mesh = np.meshgrid(mesh_theta, mesh_phi, indexing=\"ij\")\n",
    "\n",
    "k_theta = np.fft.fftfreq(N_theta, d=L_theta/(2 * np.pi * N_theta))\n",
    "k_phi = np.fft.fftfreq(N_phi, d=L_phi/(2 * np.pi * N_phi))\n",
    "Ktheta, Kphi = np.meshgrid(k_theta, k_phi, indexing='ij')\n",
    "\n",
    "g = lambda mesh: np.sin(4 * np.pi * mesh[0]) * np.sin(4 * np.pi * mesh[1])\n",
    "g_prime_theta = lambda mesh: (\n",
    "    4 * np.pi * np.cos(4 * np.pi * mesh[0]) * np.sin(4 * np.pi * mesh[1]))\n",
    "g_prime_phi = lambda mesh: (\n",
    "    4 * np.pi * np.cos(4 * np.pi * mesh[1]) * np.sin(4 * np.pi * mesh[0]))\n",
    "g_h = g(mesh)\n",
    "g_prime_theta_h = g_prime_theta(mesh)\n",
    "g_prime_phi_h = g_prime_phi(mesh)\n",
    "\n",
    "g_prime_fft2_theta_h = np.fft.ifft2(1j * Ktheta * np.fft.fft2(g_h)).real\n",
    "g_prime_fft2_phi_h = np.fft.ifft2(1j * Kphi * np.fft.fft2(g_h)).real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 264848 characters]",
      "text/plain": [
       "<Figure size 1000x400 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean error phi:  1.6363934567153887e-14\n",
      "Mean error theta:  2.9112296192913396e-14\n"
     ]
    }
   ],
   "source": [
    "fig, axs = plt.subplots(2, 3, figsize=(10, 4), constrained_layout=True)\n",
    "\n",
    "axs[0, 0].set_title(r'$\\frac{\\partial f}{\\partial \\theta}$')\n",
    "g_prime_theta_bar = axs[0,0].imshow(g_prime_theta_h, cmap=\"turbo\")\n",
    "\n",
    "axs[0, 1].set_title(r'$\\frac{\\partial f}{\\partial \\theta}$, fft')\n",
    "g_prime_fft2_theta_bar = axs[0, 1].imshow(g_prime_fft2_theta_h, cmap=\"turbo\")\n",
    "\n",
    "axs[0, 2].set_title(r'$\\epsilon_\\theta$')\n",
    "im_err_theta = axs[0, 2].imshow(np.abs(g_prime_fft2_theta_h - g_prime_theta_h), cmap=\"inferno\")\n",
    "\n",
    "axs[1, 0].set_title(r'$\\frac{\\partial f}{\\partial \\phi}$')\n",
    "g_prime_phi_bar = axs[1, 0].imshow(g_prime_phi_h, cmap=\"turbo\")\n",
    "\n",
    "axs[1, 1].set_title(r'$\\frac{\\partial f}{\\partial \\phi}$, fft')\n",
    "g_prime_fft2_phi_bar = axs[1, 1].imshow(g_prime_fft2_phi_h, cmap=\"turbo\")\n",
    "\n",
    "axs[1, 2].set_title(r'$\\epsilon_\\phi$')\n",
    "im_err_phi = axs[1, 2].imshow(np.abs(g_prime_fft2_phi_h - g_prime_phi_h), cmap=\"inferno\")\n",
    "\n",
    "# Add colorbars, one per column\n",
    "fig.colorbar(g_prime_theta_bar, ax=axs[0, 0], shrink=0.5)\n",
    "fig.colorbar(g_prime_phi_bar, ax=axs[1, 0], shrink=0.6)\n",
    "\n",
    "fig.colorbar(g_prime_fft2_theta_bar, ax=axs[0, 1], shrink=0.6)\n",
    "fig.colorbar(g_prime_fft2_phi_bar, ax=axs[1, 1], shrink=0.6)\n",
    "\n",
    "\n",
    "fig.colorbar(im_err_theta, ax=axs[0, 2], shrink=0.6)\n",
    "fig.colorbar(im_err_phi, ax=axs[1, 2], shrink=0.6)\n",
    "#plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Statistics \n",
    "print(\"Mean error phi: \", np.mean(np.abs(g_prime_fft2_theta_h - g_prime_theta_h)))\n",
    "print(\"Mean error theta: \", np.mean(np.abs(g_prime_fft2_phi_h - g_prime_phi_h)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SHT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "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\n",
    "from torch_harmonics import RealVectorSHT, InverseRealVectorSHT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlon = 128\n",
    "nlat = nlon // 2 + 1\n",
    "grid = \"equiangular\"\n",
    "lmax = nlon // 2 + 1\n",
    "mmax = nlon // 2 + 1\n",
    "\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",
    "mesh_theta = np.linspace(0, np.pi, nlat, endpoint=False)\n",
    "mesh_phi = np.linspace(0, 2 * np.pi, nlon, endpoint=False)\n",
    "mesh_polar = np.meshgrid(mesh_theta, mesh_phi, indexing=\"ij\")\n",
    "\n",
    "l = torch.arange(lmax)\n",
    "m = torch.arange(mmax)\n",
    "ls = (l).reshape(-1, 1)\n",
    "ms = (m).reshape(1, -1)\n",
    "\n",
    "Ktheta, Kphi = np.meshgrid(ls, ms, indexing='ij')\n",
    "\n",
    "e = lambda mesh: np.sin(4 * np.pi * mesh[0]) * np.sin(4 * np.pi * mesh[1])\n",
    "e_prime_theta = lambda mesh: (\n",
    "    4 * np.pi * np.cos(4 * np.pi * mesh[0]) * np.sin(4 * np.pi * mesh[1]))\n",
    "e_prime_phi = lambda mesh: (\n",
    "    4 * np.pi * np.cos(4 * np.pi * mesh[1]) * np.sin(4 * np.pi * mesh[0]))\n",
    "\n",
    "e_h = e(mesh_polar)\n",
    "e_prime_theta_h = e_prime_theta(mesh_polar)\n",
    "e_prime_phi_h = e_prime_phi(mesh_polar)\n",
    "\n",
    "e_prime_fft2_theta_h = isht(1j * torch.tensor(Ktheta) * sht(torch.tensor(e_h)))\n",
    "e_prime_fft2_phi_h = isht(1j * torch.tensor(Kphi) * sht(torch.tensor(e_h)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 524392 characters]",
      "text/plain": [
       "<Figure size 1000x400 with 12 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean error theta:  11.095667342870255\n",
      "Mean error phi:  0.9485484019412784\n"
     ]
    }
   ],
   "source": [
    "fig, axs = plt.subplots(2, 3, figsize=(10,4), constrained_layout=True)\n",
    "\n",
    "e_prime_theta_bar = axs[0, 0].imshow(e_prime_theta_h, cmap=\"turbo\")\n",
    "axs[0, 0].set_title(r'$\\frac{\\partial f}{\\partial \\theta}$')\n",
    "\n",
    "e_prime_fft2_theta_bar = axs[0, 1].imshow(e_prime_fft2_theta_h, cmap=\"turbo\")\n",
    "axs[0, 1].set_title(r'$\\frac{\\partial f}{\\partial \\theta}$ sht')\n",
    "\n",
    "axs[0, 2].set_title(r'$\\epsilon_\\theta$')\n",
    "im_err_theta =axs[0, 2].imshow(np.abs(e_prime_fft2_theta_h.detach().numpy() - e_prime_theta_h), cmap=\"inferno\")\n",
    "\n",
    "e_prime_phi_bar = axs[1, 0].imshow(e_prime_phi_h, cmap=\"turbo\")\n",
    "axs[1, 0].set_title(r'$\\frac{\\partial f}{\\partial \\phi}$')\n",
    "\n",
    "e_prime_fft2_phi_bar = axs[1, 1].imshow(e_prime_fft2_phi_h, cmap=\"turbo\")\n",
    "axs[1, 1].set_title(r'$\\frac{\\partial f}{\\partial \\phi}$ sht')\n",
    "\n",
    "axs[1, 2].set_title(r'$\\epsilon_\\phi$')\n",
    "im_err_phi = axs[1, 2].imshow(np.abs(e_prime_fft2_phi_h.detach().numpy() - e_prime_phi_h), cmap=\"inferno\")\n",
    "\n",
    "fig.colorbar(e_prime_theta_bar, ax=axs[0, 0], shrink=0.6)\n",
    "fig.colorbar(e_prime_phi_bar, ax=axs[1, 0], shrink=0.6)\n",
    "\n",
    "fig.colorbar(e_prime_fft2_theta_bar, ax=axs[0, 1], shrink=0.6)\n",
    "fig.colorbar(e_prime_fft2_phi_bar, ax=axs[1, 1], shrink=0.6)\n",
    "\n",
    "fig.colorbar(im_err_theta, ax=axs[0, 2], shrink=0.6)\n",
    "fig.colorbar(im_err_phi, ax=axs[1, 2], shrink=0.6)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "print(\"Mean error theta: \", np.mean(np.abs(e_prime_fft2_theta_h.detach().numpy() - e_prime_theta_h)))\n",
    "print(\"Mean error phi: \", np.mean(np.abs(e_prime_fft2_phi_h.detach().numpy() - e_prime_phi_h)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
