zernpy

Python package for:

• Calculating real-valued Zernike polynomials based on exact (factorial) and recursive equations;
• Converting between radial (n) + azimuthal (m) orders, OSA, Noll, and Fringe indices;
• Fitting phase profiles with Zernike polynomials;
• Generating and visualizing wavefronts formed by single or sums of polynomials;
• Computing 2D PSF kernels in the imaging plane of an optical system for individual polynomials and polynomial sums.

Features

• Real-valued Zernike polynomial computation in analytical and recursive forms;
• Derivatives computation (radial and azimuthal);
• Orders / OSA / Noll / Fringe index conversions;
• Wavefront fitting;
• 2D PSF kernel simulation (image plane);
• Optional Numba acceleration of computations;
• Visualization (plotting) utilities.

Full API documentation: sklykov.github.io/zernpy
The recursive form of the equations is used by default. It becomes valuable for high-order polynomials (n > 40) due to improved numerical stability.

Setup instructions

Installation

pip install zernpy

To upgrade:

pip install -U zernpy

Requirements

The package requires numpy, scipy and matplotlib with no specific version constraints. The optional dependency is numba>=0.57.1.
Tests are run with pytest. Linting and formatting use ruff; non-strict type checking uses mypy.

Examples

Minimal Example

from zernpy import ZernPol
from math import pi
zp = ZernPol(m=0, n=4)  # Spherical polynomial initialization
# Similar init. forms: ZernPol(osa=12), ZernPol(noll=11), ZernPol(fringe=9)
indices = zp.get_indices()  # returns tuple: ((m, n), OSA, Noll, Fringe) orders / indices
naming = zp.get_polynomial_name()  # returns str with polynomial name (up to 7th order)
# Coordinates can be provided as real numbers or NumPy arrays
r = 0.5; theta = pi*0.25  # as example of a single polar point
value = zp.polynomial_value(r, theta)  # polynomial value(-s) for radial coordinates
value_r = zp.radial(r)  # radial polynomial value(-s)
value_dr = zp.radial_dr(r)  # derivative on r of radial polynomial value(-s)
value_ang = zp.triangular(theta)  # angular polynomial value(-s)
value_ang_dth = zp.triangular_dtheta(theta)  # derivative on theta of angular value(-s)
normalization = zp.normf()  # OSA or Var(ZP)=1 normalization factor

Example of a few useful static methods

from zernpy import ZernPol
# Notations conversion
m, n = ZernPol.index2orders(osa_index=10)  # Get azimuthal, radial orders. Same for noll_index, fringe_index
noll = ZernPol.osa2noll(10)  # Also available: noll2osa, osa2fringe, fringe2osa
ZernPol.plot_profile(ZernPol(fringe=11))  # interactive plotting of a single polynomial
# Sum of Zernike polynomials as a surface
zerns = ZernPol(osa=3), ZernPol(osa=6)
zerns_sum_surface = ZernPol.gen_zernikes_surface(coefficients=[0.1, -0.1], polynomials=zerns)
ZernPol.plot_zernikes_surface(zerns_sum_surface)  # interactive matplotlib plot

Note
Images below are visible on GitHub but not on the PyPI project page.

Fitting Zernike polynomials to a 2D phase profile

There are methods implemented for fitting a provided phase profile to a provided set of polynomials. Minimal example:

import matplotlib.pyplot as plt
from zernpy import ZernPol, generate_phases_image, fit_polynomials
polynomials = (ZernPol(m=-1, n=5), ZernPol(m=3, n=3), ZernPol(m=0, n=2))
pols_coeffs = (-0.114, 0.245, 0.403)
phases_image = generate_phases_image(polynomials=polynomials, polynomials_amplitudes=pols_coeffs,
                                     img_height=401, img_width=401)
plt.figure("Initial Phase Profile"); plt.imshow(phases_image, cmap="jet")
plt.axis("off"); plt.tight_layout()
polynomials_amplitudes, cropped_img = fit_polynomials(phases_image, polynomials, return_cropped_image=True)
plt.figure("Cropped Input Phase Profile"); plt.imshow(cropped_img, cmap="jet")
plt.axis("off"); plt.tight_layout()

Initially generated profile:

Cropped input profile used for fitting:

2D PSF kernel calculation

The 2D PSF kernel is calculated using scalar diffraction theory by integrating over a circular pupil with a phase distribution described by Zernike polynomials. The calculation is restricted to the focal plane (no Z-axis dependency).

Initialization and usage of the class instance (basic usage with default calculation parameters, such as the kernel size):

from zernpy import ZernPSF, ZernPol
zpsf = ZernPSF(ZernPol(m=1, n=3))  # horizontal coma
NA = 0.4; wavelength = 0.55; pixel_physical_size = 0.24*wavelength; expansion_coeff = -0.1
zpsf.set_physical_props(NA, wavelength, expansion_coeff, pixel_physical_size)
kernel = zpsf.calculate_psf_kernel(normalized=True)  # get a kernel as a squared matrix

The PSF kernel obtained for horizontal coma with the above parameters is shown below:

Check the API documentation for other available methods.

PSF kernel for several polynomials

Similarly to the code above, it's possible to calculate a PSF kernel associated with a sum profile of several polynomials:

from zernpy import ZernPSF, ZernPol 
zp1 = ZernPol(m=-1, n=3); zp2 = ZernPol(m=2, n=4); zp3 = ZernPol(m=0, n=4)
pols = (zp1, zp2, zp3); coeffs = (0.5, 0.21, 0.15)
zpsf_pic = ZernPSF(pols); zpsf_pic.set_physical_props(NA=0.65, wavelength=0.6, expansion_coeff=coeffs, 
                                                      pixel_physical_size=0.6/5.0)
zpsf_pic.calculate_psf_kernel(); zpsf_pic.plot_kernel("Sum of Polynomials Profile")

The resulting profile is:

Acceleration of kernel calculation by numba

It's possible to accelerate the calculation of a kernel by installing the numba library in the same Python environment and providing the appropriate flags in a calculation method, similar to the following code snippet:

from zernpy import force_get_psf_compilation, ZernPol, ZernPSF
force_get_psf_compilation()  # optional precompilation for saving and reusing compiled code
NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 4.6; ampl = -0.12
zp = ZernPol(m=0, n=2); zpsf = ZernPSF(zp) 
zpsf.set_physical_props(NA, wavelength, ampl, pixel_size)
zpsf.calculate_psf_kernel(accelerated=True)

Cropping kernel

By default, the kernel size is estimated empirically and may need to be increased if the PSF extends beyond its boundaries. The estimated size generally increases with the polynomial orders and the magnitudes of their coefficients. To reduce the size of kernel, it's possible to call the method crop_kernel, e.g.:

from zernpy import ZernPSF, ZernPol
zpsf = ZernPSF(zernpol=(ZernPol(m=0, n=4)))  # Spherical aberration
zpsf.set_physical_props(NA=1.25, wavelength=0.5, expansion_coeff=0.25, pixel_physical_size=0.5/5.0)
zpsf.calculate_psf_kernel(accelerated=True, verbose_info=True); zpsf.plot_kernel("Not Cropped")
zpsf.crop_kernel(min_part_of_max=0.025)  #  2.5% of max kernel value is used for a threshold
zpsf.plot_kernel("Cropped")

Originally used (estimated) kernel with size (77, 77) for Spherical aberration:

Cropped kernel with size (33, 33) for Spherical aberration:

If a warning is encountered that estimated kernel size isn't enough for kernel calculation (it can be thrown in the end of computation), it's possible to manually increase kernel size for next computation, e.g.:

zpsf.set_calculation_props(kernel_size=51)  # as an example

References

The recursive and tabular equations, along with references to the essential information about Zernike polynomials, are sourced from:

1. Honarvar Shakibaei and Paramesran 2013
2. Lakshminarayanan and Fleck 2011
3. Andersen 2018

The equations for calculation of PSF are sourced from:

1. Principles of Optics, by M. Born and E. Wolf, 4 ed., 1968
2. Open Source Articles: Lecture, Thesis.
3. Mahajan and Díaz 2013