From c3be5f6d597db962223baa7bf7e0274a83216bc6 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 12 Jun 2024 00:49:33 -0400 Subject: [PATCH] update plotting and notebooks --- notebooks/zernike_eval.ipynb | 432 +++++++++++++++++++++++++++++++++++ zernipax/plotting.py | 36 +-- 2 files changed, 455 insertions(+), 13 deletions(-) create mode 100644 notebooks/zernike_eval.ipynb diff --git a/notebooks/zernike_eval.ipynb b/notebooks/zernike_eval.ipynb new file mode 100644 index 0000000..bca7c5c --- /dev/null +++ b/notebooks/zernike_eval.ipynb @@ -0,0 +1,432 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "# Zernike Polynomial Evaluation" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "This notebook compares different methods for evaluating the radial part of the Zernike polynomials, in terms of both speed and accuracy" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "The two primary methods we consider are direct polynomial evaluation using Horner's method and an evaluation scheme based on a recurrence relation for Jacobi polynomials.\n", + "\n", + "The radial part of the Zernike polynomials is given by\n", + "\n", + "$$\n", + " \\mathcal{R}_l^{|m|} (\\rho) = \\sum_{s=0}^{(l-|m|)/2} \\frac{(-1)^s(l-s)!}{ s!\\left( \\cfrac{l+|m|}{2} - s\\right)! \\left( \\cfrac{l-|m|}{2} - s\\right)! } \\rho^{l-2s}\n", + "$$\n", + "\n", + "Because the coefficient of rho is made up of entirely integer operations, it can be evaluated quickly and exactly to arbitrary orders (recall that python natively supports arbitrary length integer arithmetic). These coefficients can then be evaluated using Horner's method. This is done in the `zernike_radial_poly` function.\n", + "\n", + "The other approach uses the fact that the above equation can be written as\n", + "\n", + "$$\n", + " \\mathcal{R}_l^{m} (\\rho) = (-1)^{(l-m)/2} \\rho^m P_{(l-m)/2}^{m, 0} (1 - 2 \\rho^2) \\hspace{1cm}\\text{for } m\\geq0\n", + "$$\n", + "where $P_{n}^{\\alpha, \\beta}(\\rho)$ is a Jacobi polynomial. This allows us to use stable recurrence relations for the Jacobi polynomials, as is done in the `zernike_radial` function.\n", + "\n", + "The recurrence relationship for the Jacobi polynomials is,\n", + "$$\n", + " 2n(c-n)(c-2)P_{n}^{\\alpha,\\beta}(\\rho) = (c-1)[c(c-2)\\rho + (a-b)(c-2n)]P_{n-1}^{\\alpha,\\beta}(\\rho) - 2(a-1)(b-1)cP_{n-2}^{\\alpha,\\beta}(\\rho)\n", + "$$\n", + "where \n", + "$$\n", + " c = 2n + \\alpha + \\beta, \\hspace{1cm} a = n +\\alpha, \\hspace{1cm} b = n + \\beta\n", + "$$\n", + "\n", + "For the derivatives of Zernike Radial part, we will also need derivatives of Jacobi polynomials, for which there exist another relation. \n", + "$$\n", + " \\cfrac{d^k}{dx^k} P_n^{(\\alpha, \\beta)}(x) = \\cfrac{\\Gamma(\\alpha + \\beta + n + 1 + k)}{2^k\\Gamma(\\alpha + \\beta + n + 1)} P_{n-k}^{(\\alpha + k, \\beta + k)}(x)\n", + "$$\n", + "This function relates the derivatives to normal Jacobi function, and we can use above recursion relation to calculate derivatives. To further reduce the numerical inaccuracies, we can use Pochammer form instead of Gamma function. This gives us,\n", + "$$\n", + " \\cfrac{d^k}{dx^k} P_n^{(\\alpha, \\beta)}(x) = \\cfrac{(\\alpha + \\beta + n + 1)_k}{2^k} P_{n-k}^{(\\alpha + k, \\beta + k)}(x)\n", + "$$\n", + "where\n", + "$$\n", + " (\\alpha + \\beta + n + 1)_k = (\\alpha + \\beta + n + 1)(\\alpha + \\beta + n + 2) ... (\\alpha + \\beta + n + k)\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "import sys\n", + "import os\n", + "\n", + "sys.path.insert(0, os.path.abspath(\".\"))\n", + "sys.path.append(os.path.abspath(\"../\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import mpmath\n", + "import matplotlib\n", + "import matplotlib.pyplot as plt\n", + "from zernipax.basis import ZernikePolynomial\n", + "from zernipax.zernike import(\n", + " polyder_vec,\n", + " zernike_radial,\n", + " zernike_radial_old_desc,\n", + " zernike_radial_coeffs,\n", + " zernike_radial_poly,\n", + ")\n", + "from zernipax.plotting import plot_comparison" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "res = 200\n", + "basis = ZernikePolynomial(L=res, M=res, spectral_indexing=\"ansi\", sym=\"cos\")\n", + "r = np.linspace(0, 1, 100)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "Here we time the evaluation for a basis set containing 676 modes on a grid of 100 points, for derivative orders 0 through 3. (note the `block_until_ready` is needed to get [accurate timing with jax](https://jax.readthedocs.io/en/latest/async_dispatch.html))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "print(\"zernike_radial, 0th derivative\")\n", + "%timeit _ = zernike_radial(r, basis.modes[:,0], basis.modes[:,1], 0).block_until_ready()\n", + "print(\"zernike_radial, 1st derivative\")\n", + "%timeit _ = zernike_radial(r, basis.modes[:,0], basis.modes[:,1], 1).block_until_ready()\n", + "print(\"zernike_radial, 2nd derivative\")\n", + "%timeit _ = zernike_radial(r, basis.modes[:,0], basis.modes[:,1], 2).block_until_ready()\n", + "print(\"zernike_radial, 3rd derivative\")\n", + "%timeit _ = zernike_radial(r, basis.modes[:,0], basis.modes[:,1], 3).block_until_ready()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "print(\"zernike_radial_poly, 0th derivative\")\n", + "%timeit _ = zernike_radial_poly(r[:,np.newaxis], basis.modes[:,0], basis.modes[:,1], dr=0, exact=False)\n", + "print(\"zernike_radial_poly, 1st derivative\")\n", + "%timeit _ = zernike_radial_poly(r[:,np.newaxis], basis.modes[:,0], basis.modes[:,1], dr=1, exact=False)\n", + "print(\"zernike_radial_poly, 2nd derivative\")\n", + "%timeit _ = zernike_radial_poly(r[:,np.newaxis], basis.modes[:,0], basis.modes[:,1], dr=2, exact=False)\n", + "print(\"zernike_radial_poly, 3rd derivative\")\n", + "%timeit _ = zernike_radial_poly(r[:,np.newaxis], basis.modes[:,0], basis.modes[:,1], dr=3, exact=False)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "We see that the implementation using Jacobi polynomial recurrence relation is significantly faster, despite the overhead from the JAX just-in-time compiler" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "For accuracy comparison, we will also evaluate the Zernike radial polynomials in extended precision (100 digits of accuracy) and treat this as the \"true\" value." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "mpmath.mp.dps = 100\n", + "print(\"Calculate radial Zernike polynomial coefficients (exact)\")\n", + "%time c = zernike_radial_coeffs(basis.modes[:, 0], basis.modes[:, 1], exact=True)\n", + "\n", + "print(\"\\nzernike_radial_exact, 0th derivative\")\n", + "%time zt0 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in c]).T\n", + "print(\"zernike_radial_exact, 1st derivative\")\n", + "%time zt1 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in polyder_vec(c, 1, exact=True)]).T\n", + "print(\"zernike_radial_exact, 2nd derivative\")\n", + "%time zt2 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in polyder_vec(c, 2, exact=True)]).T\n", + "print(\"zernike_radial_exact, 3rd derivative\")\n", + "%time zt3 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in polyder_vec(c, 3, exact=True)]).T" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "Next we can plot the error resulting from the two evaluation methods (polynomial evaluation and jacobi recurrence relation) vs the true solution computed in exact arithmetic. We plot the max absolute error as well as the max relative error over $\\rho \\in (0,1)$ for each derivative order." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mpmath.mp.dps = 100\n", + "c = zernike_radial_coeffs(basis.modes[:, 0], basis.modes[:, 1], exact=True)\n", + "zt0 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in c]).T\n", + "\n", + "zr0 = zernike_radial(r, basis.modes[:, 0], basis.modes[:, 1], 0)\n", + "zr1 = zernike_radial(r, basis.modes[:, 0], basis.modes[:, 1], 1)\n", + "zr2 = zernike_radial(r, basis.modes[:, 0], basis.modes[:, 1], 2)\n", + "zr3 = zernike_radial(r, basis.modes[:, 0], basis.modes[:, 1], 3)\n", + "zp0 = zernike_radial_poly(\n", + " r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=0, exact=False\n", + ")\n", + "zp1 = zernike_radial_poly(\n", + " r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=1, exact=False\n", + ")\n", + "zp2 = zernike_radial_poly(\n", + " r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=2, exact=False\n", + ")\n", + "zp3 = zernike_radial_poly(\n", + " r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=3, exact=False\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Absolute Error\n", + "\n", + "### 0th derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt0, (zp0, zr0), basis, 0, \"absolute\", names=(\"Poly: \", \"Radial: \"), print_error=True)\n", + "plt.savefig(\"compare.png\", dpi=1000)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1st derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt1, (zp1, zr1), basis, 1, \"absolute\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2nd derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt2, (zp2, zr2), basis, 2, \"absolute\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3rd derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt3, (zp3, zr3), basis, 3, \"absolute\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Relative Error\n", + "\n", + "### 0th derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt0, (zp0, zr0), basis, 0, \"relative\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 1st derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt1, (zp1, zr1), basis, 1, \"relative\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2nd derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt2, (zp2, zr2), basis, 2, \"relative\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3rd derivative" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plot_comparison(zt3, (zp3, zr3), basis, 3, \"relative\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "So in addition to being faster, the evaluation using the Jacobi recurrence relation is also significantly more accurate as the mode numbers increase, keeping absolute error less than $10^{-5}$ and relative error less than $10^{-9}$, while directly evaluating the polynomial leads to errors greater than 100\\% for large $l$" + ] + } + ], + "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.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/zernipax/plotting.py b/zernipax/plotting.py index a648c69..f29b311 100644 --- a/zernipax/plotting.py +++ b/zernipax/plotting.py @@ -265,9 +265,13 @@ def plot_modes(modes, rho=100, theta=100, **kwargs): return fig, ax -def plot_comparison(exact, methods, basis, dx=0, type="absolute"): +def plot_comparison( + exact, methods, basis, dx=0, type="absolute", names=None, print_error=False +): """Plot comparison of exact and approximate methods.""" assert type in ["absolute", "relative"], "type must be 'absolute' or 'relative'" + if names is not None: + assert len(names) == len(methods), "title must have the same length as methods" N = len(methods) res = basis.L @@ -285,8 +289,9 @@ def plot_comparison(exact, methods, basis, dx=0, type="absolute"): bounds = np.logspace(-16, 0, 17) norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N) - fig, ax = plt.subplots(1, N, squeeze=True, figsize=(N * 5 + 3, 5)) + fig, ax = plt.subplots(1, N, squeeze=True, figsize=(N * 5, 4)) for i in range(N): + description = names[i] if names is not None else f"Method {i+1}:" if dx == 0: Zmn = "Z_{nm}(x)" Zmn_p = "\\tilde{Z}_{nm}(x)" @@ -294,21 +299,13 @@ def plot_comparison(exact, methods, basis, dx=0, type="absolute"): derv = "^" + str(dx) if dx > 1 else "" Zmn = "\\frac{d" + derv + "Z_{nm}(x)}{d x" + derv + "}" Zmn_p = "\\frac{d" + derv + "\\tilde{Z}_{nm}(x)}{d x" + derv + "}" + title = description + "$\\max_{x \\in (0,1)} |" + Zmn + "-" + Zmn_p if type == "absolute": c = np.max(abs(methods[i] - exact), axis=0) / np.mean(abs(exact)) - title = ( - f"Method {i+1}:" + "$\\max_{x \\in (0,1)} |" + Zmn + "-" + Zmn_p + "|$" - ) + title = title + "|$" else: c = np.max(abs(methods[i] - exact), axis=0) - title = ( - f"Method {i+1}:" - + "$\\max_{x \\in (0,1)} |" - + Zmn - + "-" - + Zmn_p - + "| / |\\bar{Z}_{lm}|$" - ) + title = title + "| / |\\bar{Z}_{lm}|$" im = ax[i].scatter( basis.modes[:, 0], basis.modes[:, 1], @@ -323,6 +320,19 @@ def plot_comparison(exact, methods, basis, dx=0, type="absolute"): ax[i].set_xlabel("$n$", fontsize=12) ax[i].set_ylabel("$m$", fontsize=12) ax[i].set_title(title, fontsize=14) + if print_error: + ax[i].text( + 0, + 45, + f"Max error: {np.max(c):.2e}", + fontsize=13, + ) + ax[i].text( + 0, + 40, + f"Mean error: {np.mean(c):.2e}", + fontsize=13, + ) # Create a separate axis for the colorbar cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7]) # [left, bottom, width, height] cbar = fig.colorbar(im, cax=cbar_ax, ticks=bounds)