{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quadrature\n",
    "\n",
    "This is a notebook that illustrates differences in the quadrature methods which are central to the evaluation of\n",
    "\n",
    "$$\n",
    "\\int_0^{\\pi} f(\\theta) \\sin \\theta \\; \\mathrm{d}\\theta = \\int_{-1}^{1} f(\\cos \\theta) \\; \\mathrm{d}\\cos \\theta.\n",
    "$$\n",
    "\n",
    "Quadrature is used to compute the projection onto the associated Legendre functions in `torch-harmonics`.\n",
    "\n",
    "In order to illustrate how interpolation and quadrature affect the error in the computation of the SHT, this notebook contains example for both errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_theta = 80"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../torch_harmonics\")\n",
    "\n",
    "from quadrature import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test interpolation\n",
    "\n",
    "we first assess the interpolation onto the quadrature nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# interpolation - careful - this breaks if points align (for non-periodic signals)\n",
    "def interpolate(t, tq, f):\n",
    "    j = np.searchsorted(t, tq) - 1\n",
    "    d = torch.from_numpy( (tq - t[j]) / np.diff(t)[j] )\n",
    "    j = torch.as_tensor(j)\n",
    "    interp = torch.lerp(torch.as_tensor(f[j]), torch.as_tensor(f[j+1]), d)\n",
    "    # interp = f[j] + (f[j+1] - f[j]) * (tq - t[j]) / np.diff(t)[j]\n",
    "    # print(d)\n",
    "    # print(f[j+1] - f[j])\n",
    "    # print(j)\n",
    "    return interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "cost_lg, wlg = legendre_gauss_weights(n_theta, -1, 1)\n",
    "# cost_lg, wlg = lobatto_weights(n_theta, -1, 1)\n",
    "tq = np.flip(np.arccos(cost_lg))\n",
    "teq = np.linspace(0, np.pi, n_theta, dtype=np.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8288 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(teq, '.')\n",
    "plt.plot(tq, '.')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "test interpolation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda t : np.cos(4*t)\n",
    "# f = lambda t : 1 / (1 + 25 * (2*(t-np.pi/2)/np.pi)**2)\n",
    "# f = lambda t : 1 / (1 + 25 * np.cos(t)**2)\n",
    "# f = lambda t : t**5 - 3*t**2 - 2*t + 1.0\n",
    "\n",
    "interp = interpolate(teq, tq, f(teq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.9905,  0.9602,  0.9065,  0.8309,  0.7351,  0.6214,  0.4926,  0.3519,\n",
       "         0.2026,  0.0484, -0.1070, -0.2597, -0.4062, -0.5428, -0.6661, -0.7733,\n",
       "        -0.8617, -0.9292, -0.9741, -0.9953, -0.9923, -0.9652, -0.9147, -0.8418,\n",
       "        -0.7486, -0.6370, -0.5100, -0.3706, -0.2221, -0.0681,  0.0875,  0.2411,\n",
       "         0.3888,  0.5272,  0.6527,  0.7625,  0.8537,  0.9243,  0.9724,  0.9968,\n",
       "         0.9968,  0.9724,  0.9243,  0.8537,  0.7625,  0.6527,  0.5272,  0.3888,\n",
       "         0.2411,  0.0875, -0.0681, -0.2221, -0.3706, -0.5100, -0.6370, -0.7486,\n",
       "        -0.8418, -0.9147, -0.9652, -0.9923, -0.9953, -0.9741, -0.9292, -0.8617,\n",
       "        -0.7733, -0.6661, -0.5428, -0.4062, -0.2597, -0.1070,  0.0484,  0.2026,\n",
       "         0.3519,  0.4926,  0.6214,  0.7351,  0.8309,  0.9065,  0.9602,  0.9905],\n",
       "       dtype=torch.float64)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "interp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 1.        ,  0.98737534,  0.94982013,  0.88828261,  0.80431657,\n",
       "        0.70004207,  0.578092  ,  0.4415455 ,  0.29385027,  0.13873553,\n",
       "       -0.01988219, -0.1779979 , -0.33161928, -0.4768675 , -0.61007515,\n",
       "       -0.72787881, -0.82730403, -0.90584039, -0.9615049 , -0.99289207,\n",
       "       -0.9992094 , -0.98029737, -0.9366335 , -0.86932027, -0.78005731,\n",
       "       -0.67109843, -0.54519477, -0.40552532, -0.25561663, -0.0992538 ,\n",
       "        0.05961513,  0.21697881,  0.36886392,  0.51143548,  0.64109364,\n",
       "        0.75456462,  0.84898336,  0.92196585,  0.97166934,  0.99683884,\n",
       "        0.99683884,  0.97166934,  0.92196585,  0.84898336,  0.75456462,\n",
       "        0.64109364,  0.51143548,  0.36886392,  0.21697881,  0.05961513,\n",
       "       -0.0992538 , -0.25561663, -0.40552532, -0.54519477, -0.67109843,\n",
       "       -0.78005731, -0.86932027, -0.9366335 , -0.98029737, -0.9992094 ,\n",
       "       -0.99289207, -0.9615049 , -0.90584039, -0.82730403, -0.72787881,\n",
       "       -0.61007515, -0.4768675 , -0.33161928, -0.1779979 , -0.01988219,\n",
       "        0.13873553,  0.29385027,  0.4415455 ,  0.578092  ,  0.70004207,\n",
       "        0.80431657,  0.88828261,  0.94982013,  0.98737534,  1.        ])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(teq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 21996 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plt.plot(teq, f(teq), '.-', label=\"reference\")\n",
    "plt.plot(tq, f(tq), '.-', label=\"reference\")\n",
    "# plt.plot(tq, interp, '.-', label=\"interpolated\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test quadrature with associated Legendre polynomials\n",
    "\n",
    "let us test different quadrature modes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def precompute_legpoly(m_max, l_max, x):\n",
    "    \"\"\"\n",
    "    Computes the values of P^m_n(\\cos \\theta) at the positions specified by x (theta)\n",
    "    The resulting tensor has shape (m_max, l_max, len(x))\n",
    "    \"\"\"\n",
    "\n",
    "    # compute the tensor P^m_n:\n",
    "    pct = np.zeros((m_max, l_max, len(x)), dtype=np.float64)\n",
    "\n",
    "    sinx = np.sin(x)\n",
    "    cosx = np.cos(x)\n",
    "\n",
    "    a = lambda m, l: np.sqrt((4*l**2 - 1) / (l**2 - m**2))\n",
    "    b = lambda m, l: -1 * np.sqrt((2*l+1)/(2*l-3)) * np.sqrt(((l-1)**2 - m**2)/(l**2 - m**2))\n",
    "\n",
    "    # start by populating the diagonal and the second higher diagonal\n",
    "    amm = np.sqrt( 1. / (4 * np.pi) )\n",
    "    pct[0,0,:] = amm\n",
    "    pct[0,1,:] = a(0, 1) * cosx * amm\n",
    "    for m in range(1, min(m_max, l_max)):\n",
    "        pct[m,m,:] = -1*np.sqrt( (2*m+1) / (2*m) ) * pct[m-1,m-1,:] * sinx\n",
    "        if m + 1 < l_max:\n",
    "            pct[m,m+1,:] = a(m, m+1) * cosx * pct[m,m,:]\n",
    "\n",
    "    # fill the remaining values on the upper triangle\n",
    "    for m in range(0, m_max):\n",
    "        for l in range(m+2, l_max):\n",
    "            pct[m,l,:] = a(m,l) * cosx * pct[m,l-1,:] + b(m,l) * pct[m,l-2,:]\n",
    "\n",
    "    return torch.from_numpy(pct)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us plot the Legendre polynomials:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38316 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "m = 0\n",
    "\n",
    "pct = np.sqrt(2 * np.pi) * precompute_legpoly(n_theta, n_theta, teq)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "for l in range(6):\n",
    "    ax.plot(np.cos(teq), pct[0, l])\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def project(t, w, f, mmax=None):\n",
    "    m = 0\n",
    "    if mmax == None:\n",
    "        mmax = len(t)\n",
    "\n",
    "    weights = torch.from_numpy(w)\n",
    "    pct = np.sqrt(2 * np.pi) * precompute_legpoly(mmax, mmax, t)\n",
    "    weights = torch.einsum('mlk,k->mlk', pct, weights)\n",
    "\n",
    "    proj = torch.einsum('...k,lk->...l', torch.as_tensor(f), weights[m])\n",
    "    rec = torch.einsum('...l, lk->...k', proj, pct[m] )\n",
    "    return rec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let us compare the accuracy of the different projection methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31760 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t = np.linspace(0, np.pi, n_theta)\n",
    "plt.plot(t, f(t), label=\"reference\")\n",
    "\n",
    "for quadrature in [legendre_gauss_weights, lobatto_weights, clenshaw_curtiss_weights, fejer2_weights]:\n",
    "    cost, wq = quadrature(n_theta, -1, 1)\n",
    "    tq = np.flip(np.arccos(cost))\n",
    "\n",
    "    out = project(tq, wq, f(tq))\n",
    "\n",
    "    plt.plot(tq, out, '.-', label=quadrature.__name__)\n",
    "\n",
    "plt.legend(loc='lower left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27308 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "for quadrature in [legendre_gauss_weights, lobatto_weights, clenshaw_curtiss_weights]:\n",
    "    cost, wq = quadrature(n_theta, -1, 1)\n",
    "    tq = np.flip(np.arccos(cost))\n",
    "\n",
    "    out = project(tq, wq, f(tq))\n",
    "    # print(np.abs(out - f(tq)))\n",
    "\n",
    "    plt.semilogy(tq, out - f(tq), '.-', label=quadrature.__name__)\n",
    "\n",
    "plt.legend(loc='lower left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us now add interpolation into the mix to evaluate performance with interpolation taken into account. For this particular case, we will assume that the data is given to us on an equidistant grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31444 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t = np.linspace(0, np.pi, n_theta)\n",
    "ref = f(t)\n",
    "plt.plot(t, ref, label=\"reference\")\n",
    "\n",
    "for quadrature in [legendre_gauss_weights, lobatto_weights, clenshaw_curtiss_weights]:\n",
    "    cost, wq = quadrature(n_theta, -1, 1)\n",
    "    tq = np.flip(np.arccos(cost))\n",
    "\n",
    "    if quadrature == lobatto_weights or quadrature == legendre_gauss_weights:\n",
    "        f_interp = interpolate(t, tq, ref)\n",
    "        mmax = len(tq)\n",
    "    else:\n",
    "        f_interp = ref\n",
    "        mmax = len(tq)\n",
    "\n",
    "    out = project(tq, wq, f_interp, mmax=mmax)\n",
    "\n",
    "    plt.plot(tq, out, '.-', label=quadrature.__name__)\n",
    "\n",
    "plt.legend(loc='lower left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "again, let us plot the overall error, this time including the interpolation error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32624 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "t = np.linspace(0, np.pi, n_theta)\n",
    "ref = f(t)\n",
    "\n",
    "fig, ax = plt.subplots(2, 1)\n",
    "\n",
    "for quadrature in [legendre_gauss_weights, lobatto_weights, clenshaw_curtiss_weights]:\n",
    "    cost, wq = quadrature(n_theta, -1, 1)\n",
    "    tq = np.flip(np.arccos(cost))\n",
    "\n",
    "    if quadrature == lobatto_weights or quadrature == legendre_gauss_weights:\n",
    "        f_interp = interpolate(t, tq, ref)\n",
    "        mmax = len(tq)\n",
    "    else:\n",
    "        f_interp = ref\n",
    "        mmax = len(tq)\n",
    "\n",
    "    out = project(tq, wq, f_interp, mmax=mmax)\n",
    "\n",
    "    ax[0].semilogy(tq, out - f(tq), '.-', label=quadrature.__name__)\n",
    "    ax[1].semilogy(tq, f_interp - f(tq), '.-', label=quadrature.__name__)\n",
    "\n",
    "ax[0].set_title(\"Projection error after interpolation\")\n",
    "ax[1].set_title(\"Interpolation error\")\n",
    "# ax[0].legend(loc='lower left')\n",
    "# ax[1].legend(loc='lower left')\n",
    "plt.legend()\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "we can see that the interpolation dominates when we interpolate the solution. For this reason, it is reasonable t choose Clenshaw-Curtiss quadrature in scenarios where we expect the interpolation error to dominate."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
