Source code for lasy.optical_elements.zernike_aberrations

from scipy.constants import c

from lasy.backend import xp
from lasy.utils.zernike import zernike

from .optical_element import OpticalElement




class ZernikeAberrations(OpticalElement):
    r"""
    Class for an optic applying a set of Zernike aberrations.

    More precisely, the amplitude multiplier corresponds to:

    .. math::

        T(\boldsymbol{x}_\perp,\omega) = \exp\left( i \frac{\omega}{c}\sum_j a_j Z_j(\boldsymbol{x}_\perp)\right)

    where
    :math:`\boldsymbol{x}_\perp` is the transverse coordinate (orthogonal
    to the propagation direction). :math:`Z_j` is the j-th Zernike polynomial,
    ordered according the OSA/ANSI indexing. The Zernike polynomials are normalized
    such that their integral over the unit disk is equal to :math:`\pi`. In the above
    formula, the total phase added to the pulse is a weighted sum of these Zernike
    polynomials with weights :math:`a_j`.

    For more information see: https://en.wikipedia.org/wiki/Zernike_polynomials


    Parameters
    ----------
    pupil_coords : tuple of floats (meters)
        A tuple of floats (cgx,cgy,r) with the first two elements corresponding to the center
        point and third element the radius of the pupil on which the zernike polynomial is
        defined.
    zernike_amplitudes : dict
        A dictionary with integer keys representing the OSA/ANSI indexing of the
        individual Zernike Polynomials. The values corresponding to these keys
        are floats giving the amplitudes / weights of the relevant Zernike polynomials.

    """

    def __init__(self, pupil_coords, zernike_amplitudes):
        self.pupil_coords = pupil_coords
        self.zernike_amplitudes = zernike_amplitudes



    def amplitude_multiplier(self, x, y, omega):
        """
        Return the amplitude multiplier.

        Parameters
        ----------
        x, y, omega : ndarrays of floats
            Define points on which to evaluate the multiplier.
            These arrays need to all have the same shape.

        Returns
        -------
        multiplier : ndarray of complex numbers
            Contains the value of the multiplier at the specified points.
            This array has the same shape as the array omega.
        """
        rr = xp.sqrt(x**2 + y**2)
        phase = xp.zeros_like(rr)

        for j in list(self.zernike_amplitudes):
            phase += self.zernike_amplitudes[j] * zernike(x, y, self.pupil_coords, j)

        return xp.exp(1j * omega / c * phase)