ZERNIPAX: A Fast and Accurate Zernike Polynomial Calculator in Python

TL;DR

This work introduces ZERNIPAX, a fast and accurate Python package for evaluating Zernike polynomials on the unit disk by harnessing Jacobi polynomials and their stable recursion. Implemented with JAX, it provides CPU and GPU pathways that dramatically reduce computation time while maintaining high numerical accuracy, especially for high-order modes. The approach explicitly addresses numerical stability issues inherent in direct evaluation and leverages automatic differentiation for integration with optimization workflows. The results demonstrate substantial speedups over existing open-source tools and enable real-time or optimization-driven applications in optics, astrophysics, and plasma physics. The release emphasizes open-source availability and practicality for high-fidelity simulations across diverse scientific domains.

Abstract

Zernike polynomials serve as an orthogonal basis on the unit disc, and have proven to be effective in optics simulations, astrophysics, and more recently in plasma simulations. Unlike Bessel functions, Zernike polynomials are inherently finite and smooth at the disc center (r=0), ensuring continuous differentiability along the axis. This property makes them particularly suitable for simulations, requiring no additional handling at the origin. We developed ZERNIPAX, an open-source Python package capable of utilizing CPU/GPUs, leveraging Google's JAX package and available on GitHub as well as the Python software repository PyPI. Our implementation of the recursion relation between Jacobi polynomials significantly improves computation time compared to alternative methods by use of parallel computing while still performing more accurately for high-mode numbers.

Figures (5)

• Figure 1: Accuracy of Jacobi recursion relation in ZERNIPAX (left) compared to direct polynomial evaluation (right) of a) radial part of Zernike polynomials b) first derivative c) second derivative d) third derivative. For the error, $\max_{x\in(0,1)}|Z_{nm}(x)-\Tilde{Z}_{nm}(x)|$, we evaluate both methods at 100 linearly spaced radial points for each $(n,m)$ mode in $n\in[0,50]$ and $m\in[0,50]$, and take the maximum absolute value difference with the exact calculation (with mpmath 100 significant digit precision), $\Tilde{Z}_{nm}$, for the error of that mode. The results, $Z_{nm}$, have been obtained using the CPU version of ZERNIPAX with 64-bit precision on a 2.8 GHz Intel Cascade Lake CPU.
• Figure 2: Comparison of the accuracies with 3 open-source codes, namely ZERN, ZERNPY and ZERNIKE, respectively. For the error, $\max_{x\in(0,1)}|Z_{nm}(x)-\Tilde{Z}_{nm}(x)|$, we evaluate every code at 100 linearly spaced radial points for each $(n,m)$ mode in $n\in[0,50]$ and $m\in[0,50]$, and take the maximum absolute value difference with the exact calculation (with mpmath 100 significant digit precision), $\Tilde{Z}_{nm}$, for the error of that mode. The results, $Z_{nm}$, have been obtained using a 2.8 GHz Intel Cascade Lake CPU with 64-bit precision.
• Figure 3: Execution time of the full set of radial Zernike polynomials for both 100 and 1000 linearly spaced points along radial direction a) for resolutions $n\in[10, 100]$ in log scale b) for resolutions $n\in[10,30]$ on a linear scale for packages ZERN, ZERNIKE and both CPU and GPU-optimized versions of ZERNIPAX. Packages ZERN and ZERNIKE can only use CPU resources. ZERNIPAX results don't include the compilation time which only affects the first run. The results have been obtained using a 2.8 GHz Intel Cascade Lake CPU and 80GB memory Nvidia A100 GPU with 64-bit precision.
• Figure 4: The corresponding bit precision of different mpmath decimal point precisions shown for range $[8,100]$. 64-bit precision can be achieved by setting dps to 18. The minimum value chosen for this figure dps=8 has 30-bit and the maximum value dps=100 has 336-bit precision.
• Figure 5: The maximum difference between Zernike polynomials up to $n,m$=100 calculated with lower dps and 200 dps, $\Tilde{Z}_{nm}(x)$. For the log scale, 0's are set to $2^{-53}$.