DiffractionBlurGenerator#

class deepinv.physics.generator.DiffractionBlurGenerator(psf_size, num_channels=1, zernike_index=tuple(range(4, 12)), fc=0.2, max_zernike_amplitude=0.15, pupil_size=(256, 256), apodize=False, random_rotate=False, index_convention='noll', device='cpu', dtype=torch.float32, rng=None)[source]#

Bases: PSFGenerator

Diffraction limited blur generator.

Generates 2D diffraction PSFs in optics using Zernike decomposition of the phase mask (Fresnel/Fraunhoffer diffraction theory).

Zernike polynomials are a sequence of orthogonal polynomials defined on the unit disk. They are commonly used in optical systems to describe wavefront aberrations.

The PSF \(h(\theta)\) is defined as the squared magnitude of the Fourier transform of the pupil function \(p_{\theta} = \exp(- i 2 \pi \phi_{\theta})\):

\[h(\theta) = \left| \mathcal{F} \left( \exp \left( - i 2 \pi \phi_\theta \right) \right) \right|^2,\]

where \(\phi_\theta : \mathbb{R}^2 \to \mathbb{R}\) is the wavefront error function (or aberration function), expressed in the Zernike basis:

\[\phi_\theta= \sum_{k \in K} \theta_k z_k\]

where \(K\) is set of indices of Zernike polynomials \(z_k\) used in the decomposition (equal to zernike_index). See Lakshminarayanan and Fleck[1] or this link or deepinv.physics.generator.Zernike for more details.

In the ideal diffraction-limited case (i.e., no aberrations), the PSF corresponds to the Airy pattern.

The Zernike polynomials \(z_k\) are indexed using the 'noll' or 'ansi' convention (defined by index_convention parameter). Conversion from the two conventions to the standard radial-angular indexing is done internally (see wikipedia page).

Parameters:

  • psf_size (tuple) – the shape H x W of the generated PSF in 2D

  • num_channels (int) – number of images channels. Defaults to 1.

  • zernike_index (tuple[int, ...], tuple[tuple[int, int], ...]) –

    activated Zernike coefficients in the following index_convention convention. It can be either:

    • a tuple of int corresponding to the Noll or ANSI indices, in which case the index_convention parameter is required to interpret them correctly.

    • a tuple of tuple[int, int] corresponding to the standard radial-angular indexing \((n,m)\). In this case, the index_convention parameter is ignored.

    Defaults to (4, 5, 6, 7, 8, 9, 10, 11), correspond to radial order n from 2 to 3 (included) and the spherical aberration. These correspond to the following aberrations: defocus, astigmatism, coma, trefoil and spherical aberration.

  • fc (float) – cutoff frequency (NA/emission_wavelength) * pixel_size. Should be in [0, 0.25] to respect the Shannon-Nyquist sampling theorem, defaults to 0.2.

  • max_zernike_amplitude (float) – maximum amplitude of the Zernike coefficients, defaults to 0.15. The amplitude of each Zernike coefficient is sampled uniformly in [-max_zernike_amplitude/2, max_zernike_amplitude/2].

  • pupil_size (tuple[int]) – pixel size used to synthesize the super-resolved pupil. The higher the more precise, defaults to (256, 256). If a single int is given, a square pupil is considered.

  • apodize (bool) – whether to apodize the PSF to reduce ringing artifacts, defaults to False.

  • random_rotate (bool) – whether to randomly rotate the PSF, defaults to False.

  • index_convention (str) – the convention for the Zernike polynomials indexing. Can be either 'noll' (default) or 'ansi'.

  • device (str, torch.device) – device where the tensors are allocated and processed, defaults to 'cpu'.

  • dtype (torch.dtype) – data type of the tensors, defaults to torch.float32.

  • rng (torch.Generator) – pseudo random number generator for reproducibility. Defaults to None.

Examples:

>>> from deepinv.physics.generator import DiffractionBlurGenerator
>>> generator = DiffractionBlurGenerator((5, 5), num_channels=3)
>>> print("\n".join(generator.zernike_polynomials)) # list of Zernike polynomials used
Zernike(n = 2, m = 0) -- Defocus
Zernike(n = 2, m = -2) -- Oblique Astigmatism
Zernike(n = 2, m = 2) -- Vertical Astigmatism
Zernike(n = 3, m = -1) -- Vertical Coma
Zernike(n = 3, m = 1) -- Horizontal Coma
Zernike(n = 3, m = -3) -- Vertical Trefoil
Zernike(n = 3, m = 3) -- Oblique Trefoil
Zernike(n = 4, m = 0) -- Primary Spherical
>>> blur = generator.step()  # dict_keys(['filter', 'coeff', 'pupil'])
>>> print(blur['filter'].shape)
torch.Size([1, 3, 5, 5])

References:

generate_angles(batch_size)[source]#

Generates random angles for the rotation of the PSF.

Parameters:

batch_size (int) – batch_size.

Returns:

a tensor of shape (batch_size,) angles in degrees.

Return type:

Tensor

generate_coeff(batch_size)[source]#

Generates random coefficients of the decomposition in the Zernike polynomials.

Parameters:

batch_size (int) – batch_size.

Returns:

a tensor of shape (batch_size, len(zernike_index)) coefficients in the Zernike decomposition.

Return type:

Tensor

step(batch_size=1, coeff=None, angle=None, seed=None, **kwargs)[source]#

Generate a batch of PFS with a batch of Zernike coefficients

Parameters:

  • batch_size (int) – batch_size.

  • coeff (torch.Tensor) – batch_size x len(zernike_index) coefficients of the Zernike decomposition (default is None)

  • angle (torch.Tensor) – batch_size angles in degree to rotate the PSF (defaults is None)

  • seed (int) – the seed for the random number generator.

Returns:

dictionary with keys

  • filter: tensor of size (batch_size x num_channels x psf_size[0] x psf_size[1]) batch of PSFs,

  • coeff: list of sampled Zernike coefficients in this realization,

  • pupil: the pupil function,

  • angle: the random rotation angle in degrees if random_rotate is True, nothing otherwise.

Return type:

dict

property zernike_polynomials: list[str]#

List of Zernike polynomials used in the decomposition, with the corresponding aberration if available.

Examples using DiffractionBlurGenerator:#

Tour of blur operators

Tour of blur operators