psf_generator.utils.zernike

A collection of functions related to Zernike polynomials.

Functions

create_pupil_mesh(→ tuple[torch.Tensor, Ellipsis])	Create a 2D square meshgrid for the pupil function.
zernike_nl(→ torch.Tensor)	Compute the Zernike polynomial of order n and m in the polar coordinates
index_to_nl(→ tuple[int, int])	Find the [n, l]-pair given OSA index l for Zernike polynomials.
create_zernike_aberrations(→ torch.Tensor)	Create Zernike aberrations for the pupil function.
create_special_pupil(→ torch.Tensor)	Special phase masks not included in the space spanned by the Zernike polynomials.

Module Contents

psf_generator.utils.zernike.create_pupil_mesh(n_pixels: int) → tuple[torch.Tensor, Ellipsis]

Create a 2D square meshgrid for the pupil function.

Parameters

n_pixelsint

Number of pixels for the pupil function.

Returns

(kx, ky): Tuple[torch.Tensor, …]

Two Tensors that represent the 2D coordinates on the mesh.

psf_generator.utils.zernike.zernike_nl(n: int, l: int, rho: torch.float, phi: float, radius: float = 1) → torch.Tensor

Compute the Zernike polynomial of order n and m in the polar coordinates

Parameters

nint

Index n in the definition on wikipedia, positive integer.

lint

\(|l| = m\), m is the index m in the definition on wikipedia. l can be positive or negative.

rhotorch.Float

Radial distance.

phifloat

Azimuthal angle.

radiusfloat

Radius of the disk on which the Zernike polynomial is defined, default is 1.

Returns

Z: torch.Tensor

Zernike polynomial Z(rho, phi) evaluated at rho and phi given indices n and l.

psf_generator.utils.zernike.index_to_nl(index: int) → tuple[int, int]

Find the [n, l]-pair given OSA index l for Zernike polynomials.

The OSA index ‘j’ is defined as \(j = (n(n + 2) + l) / 2\).

Parameters

indexint

OSA index j.

Returns

(n, - n + 2 * l)Tuple[int, int]

Corresponding (n, l)-pair.

psf_generator.utils.zernike.create_zernike_aberrations(zernike_coefficients: torch.Tensor, n_pix_pupil: int, mesh_type: str) → torch.Tensor

Create Zernike aberrations for the pupil function.

Arbitrary Zernike aberrations can be applied to the Cartesian propagator.

How it works: - Given the Zernike coefficients as a 1D Tensor of length n_zernike, a stack of the first n_zernike Zernike polynomials are constructed. - Then, the coefficients and the polynomials are multiplied and summed accordingly to create a phase mask. Finally, we create the complex field to be multiple with the existing pupil function to add this aberration.

For the Spherical case, only the axis-symmetric Zernike polynomials (i.e. only dependent on the radius rho not the angle phi), such as _defocus_ and ‘primary spherical’, can be applied due to the axis-symmetric assumption of the spherical propagator. See Spherical_propagators.py for details.

Parameters

zernike_coefficientstorch.Tensor

1D Tensor of Zernike coefficients

n_pix_pupilint

Number of pixels of the pupil function

mesh_typestr

Choose ‘spherical’ or ‘cartesian’.

Returns

Zernike_aberrations: torch.Tensor

Of type torch.complex64.

psf_generator.utils.zernike.create_special_pupil(n_pix_pupil: int, mask=None, tophat_radius: float = 0.5) → torch.Tensor

Special phase masks not included in the space spanned by the Zernike polynomials.

The supported special phase masks are: - None <-> flat phase, Gaussian beam - vortex <-> donut beam - halfmoon-h <-> horizontal halfmoon beam - halfmoon-v <-> vertical halfmoon beam - tophat <-> tophat beam

Notes

These special masks only applies in the Cartesian case.

Parameters

n_pix_pupilint

Number of pixels on the pupil plane.

namestr

Name of the special phase mask. Valid choices: None, ‘vortex’, ‘halfmoon-h’, ‘halfmoon-v’, ‘tophat’.

tophat_radiusfloat

Radius of the tophat mask. Default is 0.5. TODO: relate to cutoff frequency of the system.

Returns

pupiltorch.Tensor

Pupil function of the special phase mask.