{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ins and Outs of Polynomials\n",
    "\n",
    "This document serves as a reference for how prysm is set up to work with polynomials, in the context of OPD or surface figure error.  Much of what differentiates prysm's API in this area has to do with the fact that it [expects the grid to exist at the user level](./how-prysm-works.ipynb#Grids), but there are some deep and consequential implementation differences, too.  Before we get into those, we will create a working grid and a mask for visualization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from prysm.coordinates import make_xy_grid, cart_to_polar\n",
    "from prysm.geometry import circle\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "x, y = make_xy_grid(256, diameter=2)\n",
    "r, t = cart_to_polar(x, y)\n",
    "mask = circle(1,r) == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a long document, so you may wish to search for your preferred polynomial flavor:\n",
    "\n",
    "- [Hopkins](#Hopkins)\n",
    "- [Zernike](#Zernike)\n",
    "- [Jacobi](#Jacobi)\n",
    "- [Chebyshev](#Chebyshev)\n",
    "- [Legendre](#Legendre)\n",
    "- [Qs](#Qs)\n",
    "\n",
    "Note that all polynomial types allow evaluation for arbitrary order.\n",
    "\n",
    "## Hopkins\n",
    "\n",
    "The simplest polynomials are Hopkins':\n",
    "\n",
    "$$ \\text{OPD} = W_{abc} \\left[\\cos\\left(a\\cdot\\theta\\right) \\cdot \\rho^b \\cdot H^c \\right]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "for some set of coefficients.  The usage of this should not be surprising, for $W_{131}$, coma one can write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import hopkins\n",
    "cma = hopkins(1, 3, 1, r, t, 1)\n",
    "cma[mask]=np.nan\n",
    "plt.imshow(cma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we defined our grid to have a radius of 1, but often you may hold two copies of r, one which is normalized by some reference radius for polynomial evaluation, and one which is not for pupil geometry evaluation.  There is no further complexity in using Hopkins' polynomials."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Zernike\n",
    "\n",
    "prysm has a fairly granular implementation of Zernike polynomials, and expects its users to assemble the pieces to synthesize higher order functionality.  The basic building block is the `zernike_nm` function, which takes azimuthal and radial orders n and m, as in $Z_n^m$.  For example, to compute the equivalent \"primary coma\" Zernike mode as the hopkins example, one would:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import zernike_nm\n",
    "cmaZ = zernike_nm(3,1, r,t, norm=True)\n",
    "cmaZ[mask]=np.nan\n",
    "plt.imshow(cmaZ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the terms can be orthonormalized (given unit RMS) or not, based on the `norm` kwarg.  The order `m` can be negative to give access to the sinusoidal terms instead of cosinusoidal.  If you wish to work with a particular ordering scheme, prysm supports Fringe, Noll, and ANSI out of the box, all of which start counting at 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import noll_to_nm, fringe_to_nm, ansi_j_to_nm\n",
    "\n",
    "n, m = fringe_to_nm(9)\n",
    "sphZ = zernike_nm(n, m, r, t, norm=False)\n",
    "sphZ[mask]=np.nan\n",
    "plt.imshow(sphZ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These functions are not iterator-aware and should be used with, say, a list comprehension."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you wish to compute Zernikes much more quickly, the underlying implementation in prysm allows the work in computing lower order terms to be used to compute the higher order terms.  The Zernike polynomials are fundamentally two \"pieces\" which get multiplied.  The radial basis is where much of the work lives, and most programs that do not type out closed form solutions use Rodrigues' technique to compute the radial basis:\n",
    "\n",
    "$$\n",
    "R_n^m (\\rho) = \\sum_{k=0}^{\\frac{n-m}{2}} \\frac{(-1)^k (n-k)!}{k!(\\frac{n+m}{2}-k)!(\\frac{n-m}{2}-k)!}\\rho^{n-2k} \\tag{1}\n",
    "$$\n",
    "\n",
    "prysm does not do this, and instead uses the fact that the radial polynomial is a Jacobi polynomial under a change-of-basis:\n",
    "\n",
    "$$\n",
    "R_n^m (\\rho) = P_\\frac{n-m}{2}^\\left(0,|m|\\right)\\left(2\\rho^2 - 1\\right) \\tag{2}\n",
    "$$\n",
    "\n",
    "And the jacobi polynomials can be computed using a recurrence relation:\n",
    "$$\n",
    "a \\cdot P_n^{(\\alpha,\\beta)} = b \\cdot x \\cdot P_{n-1}^{(\\alpha,\\beta)} - c \\cdot P_{n-2}^{(\\alpha,\\beta)} \\tag{3}\n",
    "$$\n",
    "\n",
    "In other words, for a given $m$, you can compute $R$ for $n=3$ from $R$ for $n=2$ and $n=1$, and so on until you reach the highest value of N.  Because the sum in the Rodrigues formulation is increasingly large as $n,m$ grow, it has worse than linear time complexity.  Because the recurrrence in Eq. (3) does not change as $n,m$ grow it _does_ have linear time complexity.\n",
    "\n",
    "The use of this recurrence relation is hidden from the user in the `zernike_nm` function, and the recurrence relation is for a so-called auxiliary polynomial ($R$), so the Zernike polynomials themselves are not useful for this recurrence.  You _can_ make use of it by calling the `zernike_nm_sequence` function, a naming that will become familiar by the end of this reference guide.  Consider the first 16 Fringe Zernikes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import zernike_nm_sequence\n",
    "\n",
    "nms = [fringe_to_nm(i) for i in range(1,36)]\n",
    "\n",
    "# zernike_nm_sequence returns a generator\n",
    "%timeit polynomials = list(zernike_nm_sequence(nms, r, t)) # implicit norm=True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the timing to not using the sequence flavored version:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "for n, m in nms:\n",
    "    zernike_nm(n, m, r, t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The sequence function returns a generator to leave the user in control of their memory usage.  If you wished to compute 1,000 Zernike polynomials, this would avoid holding them all in memory at once while still improving performance.  These is no benefit other than performance and plausibly reduced memory usage to the `_sequence` version of the function.  A side benefit to the recurrence relation is that it is numerically stable to higher order than Rodrigues' expression, so you can compute higher order Zernike polynomials without numerical errors.  This is an especially useful property for using lower-precision data types like float32, since they will suffer from numerical imprecision earlier."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Jacobi \n",
    "Of course, because the Zernike polynomials are related to them you also have access to the Jacobi polynomials:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import jacobi, jacobi_sequence\n",
    "\n",
    "x_ = x[0,:] # not required to be 1D, just for example\n",
    "plt.plot(x_, jacobi(3,0,0,x_))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shape may be familiar as the Zernike flavor of coma across one axis."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chebyshev\n",
    "\n",
    "Both types of Chevyshev polynomials are supported.  They are both just special cases of Jacobi polynomials:\n",
    "\n",
    "$$ \\text{cheby1} \\equiv P_n^\\left(-0.5,-0.5\\right)(x) \\quad / \\quad P_n^\\left(-0.5,-0.5\\right)(1)$$\n",
    "$$ \\text{cheby2} \\equiv P_n^\\left(0.5,0.5\\right)(x) \\quad / \\quad P_n^\\left(0.5,0.5\\right)(1)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import cheby1, cheby2, cheby1_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x_, cheby1(3,x_), x_, cheby2(3,x_))\n",
    "plt.legend(['first kind', 'second kind'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most typical use of these polynomials in optics are as an orthogonal basis over some rectangular aperture.  The calculation is separable in x and y, so it can be reduced from scaling by $2(N\\cdot M)$ to just $N+M$.  prysm will compute the mode for one column of x and one row of y, then broadcast to 2D to assemble the 'image'.  This introduces three new functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import cheby1_2d_sequence, mode_1d_to_2d, sum_of_xy_modes # or cheby2_..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# orders 1, 2, 3 in x and 4, 5, 6 in y\n",
    "ns = [1, 2, 3]\n",
    "ms = [4, 5, 6]\n",
    "modesx, modesy = cheby1_2d_sequence(ns, ms, x, y)\n",
    "plt.plot(x_, modesx[0]) # modes are 1D\n",
    "plt.title('a single mode, 1D')\n",
    "plt.figure()\n",
    "# and can be expanded to 2D\n",
    "plt.imshow(mode_1d_to_2d(modesx[0], x, y))\n",
    "plt.title('a single mode, 2D')\n",
    "\n",
    "Wx = [1]*len(modesx)\n",
    "Wy = [1]*len(modesy)\n",
    "im = sum_of_xy_modes(modesx, modesy, x, y, Wx, Wy)\n",
    "plt.imshow(im)\n",
    "plt.title('a sum of modes')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a final note, there is no reason you can't just use the cheby1/cheby2 functions with 2D arrays, it is only slower:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "im = cheby2(3, y)\n",
    "plt.imshow(im)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Legendre\n",
    "\n",
    "These polynomials are just a special case of Jacobi polynomials:\n",
    "\n",
    "$$ \\text{legendre} \\equiv P_n^\\left(0,0\\right)(x) $$\n",
    "\n",
    "Usage follows from the [Chebyshev](#Chebyshev) exactly, except the functions are prefixed by `legendre`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import legendre, legendre_sequence, legendre_2d_sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Qs\n",
    "\n",
    "Qs are Greg Forbes' Q polynomials, $Q\\text{bfs}$, $Q\\text{con}$, and $Q_n^m$.  Qbfs and Qcon polynomials are radial only, and replace the 'standard' asphere equation.  The implementation of all three of these also uses a recurrence relation, although it is more complicated and outside the scope of this reference guide.  Each includes the leading prefix from the papers:\n",
    "\n",
    "- $\\rho^2(1-\\rho^2)$ for $Q\\text{bfs}$,\n",
    "- $\\rho^4$ for $Q\\text{con}$,\n",
    "- the same as $Q\\text{bfs}$ for $Q_n^m$ when $m=0$ or $\\rho^m \\cos\\left(m\\theta\\right)$ for $m\\neq 0$\n",
    "\n",
    "The $Q_n^m$ implementation departs from the papers in order to have a more Zernike-esque flavor.  Instead of having $a,b$ coefficients and $a$ map to $\\cos$ and $b$ to $\\sin$, this implementation uses the sign of $m$, with $\\cos$ for $m>0$ and $\\sin$ for $m<0$.\n",
    "\n",
    "There are six essential functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prysm.polynomials import (\n",
    "    Qbfs, Qbfs_sequence,\n",
    "    Qcon, Qcon_sequence,\n",
    "    Q2d, Q2d_sequence,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Qbfs(2,r)\n",
    "p[mask]=np.nan\n",
    "plt.imshow(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Qcon(2,r)\n",
    "p[mask]=np.nan\n",
    "plt.imshow(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = Q2d(2, 2, r, t) # cosine term\n",
    "p[mask]=np.nan\n",
    "plt.imshow(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p2 = Q2d(2, -2, r, t) # sine term\n",
    "p2[mask]=np.nan\n",
    "plt.imshow(p2)"
   ]
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}