Wavefronts

Wavefront

Bases: Base

A simple class to hold the state of some wavefront as it is transformed and propagated throughout an optical system. All wavefronts assume square arrays.

Attributes:

Name	Type	Description
wavelength	(float, meters)	The wavelength of the Wavefront.
amplitude	(Array, power)	The electric field amplitude of the Wavefront.
phase	(Array, radians)	The electric field phase of the Wavefront.
pixel_scale	(float, meters / pixel or radians / pixel)	The pixel scale of the phase and amplitude arrays. If units='Cartesian' then the pixel scale is in meters/pixel, else if units='Angular' then the pixel scale is in radians/pixel.
plane	str	The current plane type of wavefront, can be 'Pupil', 'Focal' or 'Intermediate'.
units	str	The current units of the wavefront, can be 'Cartesian' or 'Angular'.

Source code in src/dLux/wavefronts.py

class Wavefront(Base):
    """
    A simple class to hold the state of some wavefront as it is transformed and
    propagated throughout an optical system. All wavefronts assume square arrays.

    Attributes
    ----------
    wavelength : float, meters
        The wavelength of the `Wavefront`.
    amplitude : Array, power
        The electric field amplitude of the `Wavefront`.
    phase : Array, radians
        The electric field phase of the `Wavefront`.
    pixel_scale : float, meters/pixel or radians/pixel
        The pixel scale of the phase and amplitude arrays. If `units='Cartesian'` then
        the pixel scale is in meters/pixel, else if `units='Angular'` then the pixel
        scale is in radians/pixel.
    plane : str
        The current plane type of wavefront, can be 'Pupil', 'Focal' or 'Intermediate'.
    units : str
        The current units of the wavefront, can be 'Cartesian' or 'Angular'.
    """

    wavelength: float
    pixel_scale: float
    amplitude: Array
    phase: Array
    plane: str
    units: str

    def __init__(
        self: Wavefront, npixels: int, diameter: float, wavelength: float
    ):
        """
        Parameters
        ----------
        npixels : int
            The number of pixels that represent the `Wavefront`.
        diameter : float, meters
            The total diameter of the `Wavefront`.
        wavelength : float, meters
            The wavelength of the `Wavefront`.
        """
        self.wavelength = np.asarray(wavelength, float)
        self.pixel_scale = np.asarray(diameter / npixels, float)
        self.amplitude = (
            np.ones((npixels, npixels), dtype=float) / npixels**2
        )
        self.phase = np.zeros((npixels, npixels), dtype=float)

        # Always initialised in Pupil plane with Cartesian Coords
        self.plane = "Pupil"
        self.units = "Cartesian"

    ####################
    # Getter Functions #
    ####################
    @property
    def diameter(self: Wavefront) -> Array:
        """
        Returns the current wavefront diameter calculated using the pixel scale and
        number of pixels.

        Returns
        -------
        diameter : Array, meters or radians
            The current diameter of the wavefront.
        """
        return self.npixels * self.pixel_scale

    @property
    def npixels(self: Wavefront) -> int:
        """
        Returns the side length of the arrays currently representing the wavefront.
        Taken from the last axis of the amplitude array.

        Returns
        -------
        pixels : int
            The number of pixels that represent the `Wavefront`.
        """
        return self.amplitude.shape[-1]

    @property
    def real(self: Wavefront) -> Array:
        """
        Returns the real component of the `Wavefront`.

        Returns
        -------
        wavefront : Array
            The real component of the `Wavefront` phasor.
        """
        return self.amplitude * np.cos(self.phase)

    @property
    def imaginary(self: Wavefront) -> Array:
        """
        Returns the imaginary component of the `Wavefront`.

        Returns
        -------
        wavefront : Array
            The imaginary component of the `Wavefront` phasor.
        """
        return self.amplitude * np.sin(self.phase)

    @property
    def phasor(self: Wavefront) -> Array:
        """
        The electric field phasor described by this Wavefront in complex form.

        Returns
        -------
        field : Array
            The electric field phasor of the wavefront.
        """
        return self.amplitude * np.exp(1j * self.phase)

    @property
    def psf(self: Wavefront) -> Array:
        """
        Calculates the Point Spread Function (PSF), i.e. the squared modulus
        of the complex wavefront.

        Returns
        -------
        psf : Array
            The PSF of the wavefront.
        """
        return self.amplitude**2

    @property
    def coordinates(self: Wavefront) -> Array:
        """
        Returns the physical positions of the wavefront pixels in meters.

        Returns
        -------
        coordinates : Array
            The coordinates of the centers of each pixel representing the
            wavefront.
        """
        return dlu.pixel_coords(self.npixels, self.diameter)

    @property
    def wavenumber(self: Wavefront) -> Array:
        """
        Returns the wavenumber of the wavefront (2 * pi / wavelength).

        Returns
        -------
        wavenumber : Array, 1/meters
            The wavenumber of the wavefront.
        """
        return 2 * np.pi / self.wavelength

    @property
    def fringe_size(self: Wavefront) -> Array:
        """
        Returns the size of the fringes in angular units.

        TODO Units check from focal plane
        Returns
        -------
        fringe_size : Array, radians
            The size of the linear diffraction fringe of the wavefront.
        """
        return self.wavelength / self.diameter

    @property
    def ndim(self: Wavefront) -> int:
        """
        Returns the number of 'dimensions' of the wavefront. This is used to track the
        vectorised version of the wavefront returned from vmapping.

        Returns
        -------
        ndim : int
            The 'dimensionality' of dimensions of the wavefront.
        """
        return self.pixel_scale.ndim

    #################
    # Magic Methods #
    #################
    def __add__(self: Wavefront, other: Any) -> Wavefront:
        """
        Adds the input 'other' to the wavefront. If the input is a numeric type, it is
        treated as an OPD, else if it is an optical layer, it will be applied to the
        wavefront.

        Parameters
        ----------
        other : Any
            The input to add to the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The output wavefront.
        """
        # None Type
        if other is None:
            return self

        # Some Optical Layer
        if isinstance(other, OpticalLayer()):
            return other.apply(self)

        # Array based inputs - Defaults to OPD
        if isinstance(other, (Array, float, int)):
            return self.add_opd(other)

        # Other
        else:
            raise TypeError(
                "Can only add an array or OpticalLayer to "
                f"Wavefront. Got: {type(other)}."
            )

    def __iadd__(self: Wavefront, other: Any) -> Wavefront:
        """
        Provides the += operator for the wavefront, calling the __add__ method.

        Parameters
        ----------
        other : Any
            The input to add to the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The output wavefront.
        """
        return self.__add__(other)

    def __mul__(self: Wavefront, other: Any) -> Wavefront:
        """
        Multiplies the input 'other' to the wavefront. If the input is a numeric type,
        it is treated as an array of transmission values and is multiplied by the
        wavefront amplitude, unless it is a complex number, in which case it will be
        multiplied with the wavefront phasor. If it is an optical layer, it will be
        applied to the wavefront.

        Parameters
        ----------
        other : Any
            The input to multiply with the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The output wavefront.
        """
        # None Type, return None
        if other is None:
            return self

        # Some Optical Layer, apply it
        if isinstance(other, OpticalLayer()):
            return other.apply(self)

        # Array based inputs
        if isinstance(other, (Array, float, int)):
            # Complex array - Multiply the phasors
            if isinstance(other, Array) and other.dtype.kind == "c":
                phasor = self.phasor * other
                return self.set(
                    ["amplitude", "phase"], [np.abs(phasor), np.angle(phasor)]
                )

            # Scalar array - Multiply amplitude
            else:
                return self.multiply("amplitude", other)

        # Other
        else:
            raise TypeError(
                "Can only multiply Wavefront by array or "
                f"OpticalLayer. Got: {type(other)}."
            )

    def __imul__(self: Wavefront, other: Any) -> Wavefront:
        """
        Provides the *= operator for the wavefront, calling the __mul__ method.

        Parameters
        ----------
        other : Any
            The input to multiply with the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The output wavefront.
        """
        return self.__mul__(other)

    ###################
    # Adder Functions #
    ###################
    def add_opd(self: Wavefront, opd: Array) -> Wavefront:
        """
        Adds an optical path difference (OPD) to the wavefront.

        Parameters
        ----------
        opd : Array, meters
            The opd to add to the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The new wavefront with the phases updated according to the supplied opd.
        """
        return self.add("phase", self.wavenumber * opd)

    def add_phase(self: Wavefront, phase: Array) -> Wavefront:
        """
        Adds a phase to the wavefront.

        Parameters
        ----------
        phase : Array, radians
            The phase to be added to the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The new wavefront with updated phases.
        """
        # Add this extra None check to allow PhaseOptics to have a None phase
        # and still be able to be 'added' to it, making this the phase
        # equivalent of `wf += opd` -> `wf = wf.add_phase(phase)`
        if phase is not None:
            return self.add("phase", phase)
        return self

    ###################
    # Other Functions #
    ###################
    def tilt(self: Wavefront, angles: Array) -> Wavefront:
        """
        Tilts the wavefront by the (x, y) angles.

        Parameters
        ----------
        angles : Array, radians
            The (x, y) angles by which to tilt the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The tilted wavefront.
        """
        if not isinstance(angles, Array) or angles.shape != (2,):
            raise ValueError("angles must be an array of shape (2,).")
        opd = -(angles[:, None, None] * self.coordinates).sum(0)
        return self.add_opd(opd)

    def normalise(self: Wavefront) -> Wavefront:
        """
        Normalises the total power of the wavefront to 1.

        Returns
        -------
        wavefront : Wavefront
            The normalised wavefront.
        """
        return self.divide("amplitude", np.linalg.norm(self.amplitude))

    def _to_field(self: Wavefront, complex: bool = False) -> Array:
        """
        Returns the wavefront in either (amplitude, phase) or (real, imaginary) form.

        Parameters
        ----------
        complex : bool = False
            Whether to return the wavefront in (real, imaginary) form.

        Returns
        -------
        field : Array
            The wavefront in either (amplitude, phase) or (real, imaginary) form.
        """
        if complex:
            return np.array([self.real, self.imaginary])
        return np.array([self.amplitude, self.phase])

    def _to_amplitude_phase(self: Wavefront, field: Array) -> Array:
        """
        Transforms the input field in (real, imaginary) to (amplitude, phase) form.

        Parameters
        ----------
        field : Array
            The wavefront field in (real, imaginary) form.

        Returns
        -------
        field : Array
            The wavefront field in (amplitude, phase) form.
        """
        amplitude = np.hypot(field[0], field[1])
        phase = np.arctan2(field[1], field[0])
        return np.array([amplitude, phase])

    def flip(self: Wavefront, axis: tuple) -> Wavefront:
        """
        Flips the wavefront along the specified axes. Note we use 'ij' indexing, so
        axis 0 is the y-axis and axis 1 is the x-axis.

        Parameters
        ----------
        axis : tuple
            The axes along which to flip the wavefront.

        Returns
        -------
        wavefront : Wavefront
            The new flipped wavefront.
        """
        field = self._to_field()
        flipper = vmap(np.flip, (0, None))
        amplitude, phase = flipper(field, axis)
        return self.set(["amplitude", "phase"], [amplitude, phase])

    def scale_to(
        self: Wavefront,
        npixels: int,
        pixel_scale: Array,
        complex: bool = False,
    ) -> Wavefront:
        """
        Interpolated the wavefront to a given npixels and pixel_scale. Can be done on
        the real and imaginary components by passing in complex=True.

        Parameters
        ----------
        npixels : int
            The number of pixels  to interpolate to.
        pixel_scale: Array
            The pixel scale to interpolate to.
        complex : bool = False
            Whether to rotate the real and imaginary representation of the wavefront as
            opposed to the amplitude and phase representation.

        Returns
        -------
        wavefront : Wavefront
            The new interpolated wavefront.
        """
        # Get field in either (amplitude, phase) or (real, imaginary)
        field = self._to_field(complex=complex)

        # Scale the field
        scale_fn = vmap(dlu.scale, (0, None, None))
        field = scale_fn(field, npixels, pixel_scale / self.pixel_scale)

        # Cast back to (amplitude, phase) if needed
        if complex:
            field = self._to_amplitude_phase(field)

        # Return new wavefront
        return self.set(
            ["amplitude", "phase", "pixel_scale"],
            [field[0], field[1], pixel_scale],
        )

    def rotate(
        self: Wavefront, angle: Array, order: int = 1, complex: bool = False
    ) -> Wavefront:
        """
        Rotates the wavefront by a given angle via interpolation. Can be done on the
        real and imaginary components by passing in complex=True.

        Parameters
        ----------
        angle : Array, radians
            The angle by which to rotate the wavefront in a clockwise
            direction.
        order : int = 1
            The interpolation order to use.
        complex : bool = False
            Whether to rotate the real and imaginary representation of the wavefront as
            opposed to the amplitude and phase representation.

        Returns
        -------
        wavefront : Wavefront
            The new wavefront rotated by angle in the clockwise direction.
        """
        # Get field in either (amplitude, phase) or (real, imaginary)
        field = self._to_field(complex=complex)

        # Rotate the field
        rotator = vmap(dlu.rotate, (0, None, None))
        field = rotator(field, angle, order)

        # Cast back to (amplitude, phase) if needed
        if complex:
            field = self._to_amplitude_phase(field)

        # Return new wavefront
        return self.set(["amplitude", "phase"], [field[0], field[1]])

    def resize(self: Wavefront, npixels: int) -> Wavefront:
        """
        Resizes the wavefront via a zero-padding or cropping operation.

        Parameters
        ----------
        npixels : int
            The size to resize the wavefront to.

        Returns
        -------
        wavefront : Wavefront
            The resized wavefront.
        """
        field = self._to_field()
        amplitude, phase = vmap(dlu.resize, (0, None))(field, npixels)
        return self.set(["amplitude", "phase"], [amplitude, phase])

    #########################
    # Propagation Functions #
    #########################
    def _prep_prop(self: Wavefront, focal_length) -> tuple:
        """
        Determines the propagation direction, output plane and output units.

        Parameters
        ----------
        focal_length : Union[float, None]
            The focal length of the propagation.

        Returns
        -------
        inverse : bool
            Whether the propagation is inverse or not.
        plane : str
            The output plane of the propagation.
        units : str
            The output units of the propagation.
        """
        # Determine propagation direction, output plane and output units
        if self.plane == "Pupil":
            inverse = False
            plane = "Focal"
            if focal_length is None:
                units = "Angular"
            else:
                units = "Cartesian"
        else:
            if focal_length is not None and self.units == "Angular":
                raise ValueError(
                    "focal_length can not be specific when"
                    "propagating from a Focal plane with angular units."
                )
            inverse = True
            plane = "Pupil"
            units = "Cartesian"

        return inverse, plane, units

    def propagate_FFT(
        self: Wavefront,
        focal_length: float = None,
        pad: int = 2,
    ) -> Wavefront:
        """
        Propagates the wavefront by performing a Fast Fourier Transform.

        Parameters
        ----------
        focal_length : float = None
            The focal length of the propagation. If None, the output pixel scale has
            units of radians, else meters.
        pad : int = 2
            The padding factory to apply to the input wavefront before the FFT.

        Returns
        -------
        wavefront : Wavefront
            The propagated wavefront.
        """
        inverse, plane, units = self._prep_prop(focal_length)

        # Calculate
        phasor, pixel_scale = dlu.FFT(
            self.phasor,
            self.wavelength,
            self.pixel_scale,
            focal_length,
            pad,
            inverse,
        )

        # Return new wavefront
        return self.set(
            ["amplitude", "phase", "pixel_scale", "plane", "units"],
            [np.abs(phasor), np.angle(phasor), pixel_scale, plane, units],
        )

    # TODO: Class method this?
    def _MFT(
        self: Wavefront,
        phasor: Array,
        wavelength: float,
        pixel_scale: float,
        *args: tuple,
    ) -> Array:
        """
        Simple alias for the MFT function to allow for vectorisation over phasors,
        wavelengths, pixel_scales, etc.

        Parameters
        ----------
        phasor : Array
            The phasor to propagate.
        wavelength : float
            The wavelength of the wavefront.
        pixel_scale : float
            The pixel scale of the wavefront.
        args : tuple
            The propagation arguments.

        Returns
        -------
        phasor : Array
            The propagated phasor.
        """
        return dlu.MFT(phasor, wavelength, pixel_scale, *args)

    def propagate(
        self: Wavefront,
        npixels: int,
        pixel_scale: float,
        focal_length: float = None,
        shift: Array = np.zeros(2),
        pixel: bool = True,
    ) -> Wavefront:
        """
        Propagates the wavefront by performing an MFT, allowing for the output pixel
        scale and npixels to be specified.

        Parameters
        ----------
        npixels : int
            The number of pixels in the output plane.
        pixel_scale : float, meters/pixel or radians/pixel
            The pixel scale of the output plane.
        focal_length : float = None
            The focal length of the propagation. If None, the propagation is angular
            and pixel_scale_out is taken in as radians/pixel, else meters/pixel.
        shift : Array = np.zeros(2)
            The shift in the center of the output plane.
        pixel : bool = True
            Should the shift be taken in units of pixels, or pixel scale.

        Returns
        -------
        wavefront : Wavefront
            The propagated wavefront.
        """
        inverse, plane, units = self._prep_prop(focal_length)

        # Enforce array so output can be vectorised by vmap
        pixel_scale = np.asarray(pixel_scale, float)

        # Calculate
        # Using a self._MFT here allows for broadband wavefronts to define
        # vectorised propagation fn over phasors, wavels, px_scales, etc.
        # It also makes the code muuuuch nicer to read
        args = (npixels, pixel_scale, focal_length, shift, pixel, inverse)
        phasor = self._MFT(
            self.phasor, self.wavelength, self.pixel_scale, *args
        )

        # Update
        return self.set(
            ["amplitude", "phase", "pixel_scale", "plane", "units"],
            [np.abs(phasor), np.angle(phasor), pixel_scale, plane, units],
        )

    # # TODO: Class method this?
    # def _fresnel(self, phasor, wavelength, pixel_scale, focal_shift, *args):
    #     return dlu.fresnel_MFT(phasor, wavelength, pixel_scale, *args)

    def propagate_fresnel(
        self: Wavefront,
        npixels: int,
        pixel_scale: float,
        focal_length: float,
        focal_shift: float = 0.0,
        shift: Array = np.zeros(2),
        pixel: bool = True,
    ) -> Wavefront:
        """
        Propagates the phasor using a Far-Field Fresnel propagation. This allows for
        psfs to be better modelled a few wavelengths from the focal plane.

        Parameters
        ----------
        npixels : int
            The number of pixels in the output plane.
        pixel_scale : float, meters/pixel or radians/pixel
            The pixel scale of the output plane.
        focal_length : float
            The focal length of the propagation.
        focal_shift: float, meters = 0.
            The shift from focus to propagate to.
        shift : Array = np.zeros(2)
            The shift in the center of the output plane.
        pixel : bool = True
            Should the shift be taken in units of pixels, or pixel scale.

        Returns
        -------
        wavefront : Wavefront
            The propagated wavefront.
        """
        # TODO: Try inverse propagation to see if it works, it probably will
        if self.plane == "Pupil":
            inverse = False
        else:
            inverse = True
        plane = "Intermediate"
        units = "Cartesian"

        # We can't fresnel from a focal plane
        if self.plane != "Pupil":
            raise ValueError(
                "Can only do an fresnel propagation from a Pupil plane, "
                f"current plane is {self.plane}."
            )

        # Calculate
        phasor = dlu.fresnel_MFT(
            self.phasor,
            self.wavelength,
            self.pixel_scale,
            npixels,
            pixel_scale,
            focal_length,
            focal_shift,
            shift,
            pixel,
            inverse,
        )

        # Update
        return self.set(
            ["amplitude", "phase", "pixel_scale", "plane", "units"],
            [np.abs(phasor), np.angle(phasor), pixel_scale, plane, units],
        )

coordinates: Array property

Returns the physical positions of the wavefront pixels in meters.

Returns:

Name	Type	Description
coordinates	Array	The coordinates of the centers of each pixel representing the wavefront.

diameter: Array property

Returns the current wavefront diameter calculated using the pixel scale and number of pixels.

Returns:

Name	Type	Description
diameter	(Array, meters or radians)	The current diameter of the wavefront.

fringe_size: Array property

Returns the size of the fringes in angular units.

TODO Units check from focal plane

Returns:

Name	Type	Description
fringe_size	(Array, radians)	The size of the linear diffraction fringe of the wavefront.

imaginary: Array property

Returns the imaginary component of the Wavefront.

Returns:

Name	Type	Description
wavefront	Array	The imaginary component of the Wavefront phasor.

ndim: int property

Returns the number of 'dimensions' of the wavefront. This is used to track the vectorised version of the wavefront returned from vmapping.

Returns:

Name	Type	Description
ndim	int	The 'dimensionality' of dimensions of the wavefront.

npixels: int property

Returns the side length of the arrays currently representing the wavefront. Taken from the last axis of the amplitude array.

Returns:

Name	Type	Description
pixels	int	The number of pixels that represent the Wavefront.

phasor: Array property

The electric field phasor described by this Wavefront in complex form.

Returns:

Name	Type	Description
field	Array	The electric field phasor of the wavefront.

psf: Array property

Calculates the Point Spread Function (PSF), i.e. the squared modulus of the complex wavefront.

Returns:

Name	Type	Description
psf	Array	The PSF of the wavefront.

real: Array property

Returns the real component of the Wavefront.

Returns:

Name	Type	Description
wavefront	Array	The real component of the Wavefront phasor.

wavenumber: Array property

Returns the wavenumber of the wavefront (2 * pi / wavelength).

Returns:

Name	Type	Description
wavenumber	(Array, 1 / meters)	The wavenumber of the wavefront.

__add__(other)

Adds the input 'other' to the wavefront. If the input is a numeric type, it is treated as an OPD, else if it is an optical layer, it will be applied to the wavefront.

Parameters:

Name	Type	Description	Default
other	Any	The input to add to the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The output wavefront.

Source code in src/dLux/wavefronts.py

def __add__(self: Wavefront, other: Any) -> Wavefront:
    """
    Adds the input 'other' to the wavefront. If the input is a numeric type, it is
    treated as an OPD, else if it is an optical layer, it will be applied to the
    wavefront.

    Parameters
    ----------
    other : Any
        The input to add to the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The output wavefront.
    """
    # None Type
    if other is None:
        return self

    # Some Optical Layer
    if isinstance(other, OpticalLayer()):
        return other.apply(self)

    # Array based inputs - Defaults to OPD
    if isinstance(other, (Array, float, int)):
        return self.add_opd(other)

    # Other
    else:
        raise TypeError(
            "Can only add an array or OpticalLayer to "
            f"Wavefront. Got: {type(other)}."
        )

__iadd__(other)

Provides the += operator for the wavefront, calling the add method.

Parameters:

Name	Type	Description	Default
other	Any	The input to add to the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The output wavefront.

Source code in src/dLux/wavefronts.py

def __iadd__(self: Wavefront, other: Any) -> Wavefront:
    """
    Provides the += operator for the wavefront, calling the __add__ method.

    Parameters
    ----------
    other : Any
        The input to add to the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The output wavefront.
    """
    return self.__add__(other)

__imul__(other)

Provides the *= operator for the wavefront, calling the mul method.

Parameters:

Name	Type	Description	Default
other	Any	The input to multiply with the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The output wavefront.

Source code in src/dLux/wavefronts.py

def __imul__(self: Wavefront, other: Any) -> Wavefront:
    """
    Provides the *= operator for the wavefront, calling the __mul__ method.

    Parameters
    ----------
    other : Any
        The input to multiply with the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The output wavefront.
    """
    return self.__mul__(other)

__init__(npixels, diameter, wavelength)

Parameters:

Name	Type	Description	Default
npixels	int	The number of pixels that represent the Wavefront.	required
diameter	(float, meters)	The total diameter of the Wavefront.	required
wavelength	(float, meters)	The wavelength of the Wavefront.	required

Source code in src/dLux/wavefronts.py

def __init__(
    self: Wavefront, npixels: int, diameter: float, wavelength: float
):
    """
    Parameters
    ----------
    npixels : int
        The number of pixels that represent the `Wavefront`.
    diameter : float, meters
        The total diameter of the `Wavefront`.
    wavelength : float, meters
        The wavelength of the `Wavefront`.
    """
    self.wavelength = np.asarray(wavelength, float)
    self.pixel_scale = np.asarray(diameter / npixels, float)
    self.amplitude = (
        np.ones((npixels, npixels), dtype=float) / npixels**2
    )
    self.phase = np.zeros((npixels, npixels), dtype=float)

    # Always initialised in Pupil plane with Cartesian Coords
    self.plane = "Pupil"
    self.units = "Cartesian"

__mul__(other)

Multiplies the input 'other' to the wavefront. If the input is a numeric type, it is treated as an array of transmission values and is multiplied by the wavefront amplitude, unless it is a complex number, in which case it will be multiplied with the wavefront phasor. If it is an optical layer, it will be applied to the wavefront.

Parameters:

Name	Type	Description	Default
other	Any	The input to multiply with the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The output wavefront.

Source code in src/dLux/wavefronts.py

def __mul__(self: Wavefront, other: Any) -> Wavefront:
    """
    Multiplies the input 'other' to the wavefront. If the input is a numeric type,
    it is treated as an array of transmission values and is multiplied by the
    wavefront amplitude, unless it is a complex number, in which case it will be
    multiplied with the wavefront phasor. If it is an optical layer, it will be
    applied to the wavefront.

    Parameters
    ----------
    other : Any
        The input to multiply with the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The output wavefront.
    """
    # None Type, return None
    if other is None:
        return self

    # Some Optical Layer, apply it
    if isinstance(other, OpticalLayer()):
        return other.apply(self)

    # Array based inputs
    if isinstance(other, (Array, float, int)):
        # Complex array - Multiply the phasors
        if isinstance(other, Array) and other.dtype.kind == "c":
            phasor = self.phasor * other
            return self.set(
                ["amplitude", "phase"], [np.abs(phasor), np.angle(phasor)]
            )

        # Scalar array - Multiply amplitude
        else:
            return self.multiply("amplitude", other)

    # Other
    else:
        raise TypeError(
            "Can only multiply Wavefront by array or "
            f"OpticalLayer. Got: {type(other)}."
        )

add_opd(opd)

Adds an optical path difference (OPD) to the wavefront.

Parameters:

Name	Type	Description	Default
opd	(Array, meters)	The opd to add to the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The new wavefront with the phases updated according to the supplied opd.

Source code in src/dLux/wavefronts.py

def add_opd(self: Wavefront, opd: Array) -> Wavefront:
    """
    Adds an optical path difference (OPD) to the wavefront.

    Parameters
    ----------
    opd : Array, meters
        The opd to add to the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The new wavefront with the phases updated according to the supplied opd.
    """
    return self.add("phase", self.wavenumber * opd)

add_phase(phase)

Adds a phase to the wavefront.

Parameters:

Name	Type	Description	Default
phase	(Array, radians)	The phase to be added to the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The new wavefront with updated phases.

Source code in src/dLux/wavefronts.py

def add_phase(self: Wavefront, phase: Array) -> Wavefront:
    """
    Adds a phase to the wavefront.

    Parameters
    ----------
    phase : Array, radians
        The phase to be added to the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The new wavefront with updated phases.
    """
    # Add this extra None check to allow PhaseOptics to have a None phase
    # and still be able to be 'added' to it, making this the phase
    # equivalent of `wf += opd` -> `wf = wf.add_phase(phase)`
    if phase is not None:
        return self.add("phase", phase)
    return self

flip(axis)

Flips the wavefront along the specified axes. Note we use 'ij' indexing, so axis 0 is the y-axis and axis 1 is the x-axis.

Parameters:

Name	Type	Description	Default
axis	tuple	The axes along which to flip the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The new flipped wavefront.

Source code in src/dLux/wavefronts.py

def flip(self: Wavefront, axis: tuple) -> Wavefront:
    """
    Flips the wavefront along the specified axes. Note we use 'ij' indexing, so
    axis 0 is the y-axis and axis 1 is the x-axis.

    Parameters
    ----------
    axis : tuple
        The axes along which to flip the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The new flipped wavefront.
    """
    field = self._to_field()
    flipper = vmap(np.flip, (0, None))
    amplitude, phase = flipper(field, axis)
    return self.set(["amplitude", "phase"], [amplitude, phase])

normalise()

Normalises the total power of the wavefront to 1.

Returns:

Name	Type	Description
wavefront	Wavefront	The normalised wavefront.

Source code in src/dLux/wavefronts.py

def normalise(self: Wavefront) -> Wavefront:
    """
    Normalises the total power of the wavefront to 1.

    Returns
    -------
    wavefront : Wavefront
        The normalised wavefront.
    """
    return self.divide("amplitude", np.linalg.norm(self.amplitude))

propagate(npixels, pixel_scale, focal_length=None, shift=np.zeros(2), pixel=True)

Propagates the wavefront by performing an MFT, allowing for the output pixel scale and npixels to be specified.

Parameters:

Name	Type	Description	Default
npixels	int	The number of pixels in the output plane.	required
pixel_scale	(float, meters / pixel or radians / pixel)	The pixel scale of the output plane.	required
focal_length	float = None	The focal length of the propagation. If None, the propagation is angular and pixel_scale_out is taken in as radians/pixel, else meters/pixel.	None
shift	Array = np.zeros(2)	The shift in the center of the output plane.	zeros(2)
pixel	bool = True	Should the shift be taken in units of pixels, or pixel scale.	True

Returns:

Name	Type	Description
wavefront	Wavefront	The propagated wavefront.

Source code in src/dLux/wavefronts.py

def propagate(
    self: Wavefront,
    npixels: int,
    pixel_scale: float,
    focal_length: float = None,
    shift: Array = np.zeros(2),
    pixel: bool = True,
) -> Wavefront:
    """
    Propagates the wavefront by performing an MFT, allowing for the output pixel
    scale and npixels to be specified.

    Parameters
    ----------
    npixels : int
        The number of pixels in the output plane.
    pixel_scale : float, meters/pixel or radians/pixel
        The pixel scale of the output plane.
    focal_length : float = None
        The focal length of the propagation. If None, the propagation is angular
        and pixel_scale_out is taken in as radians/pixel, else meters/pixel.
    shift : Array = np.zeros(2)
        The shift in the center of the output plane.
    pixel : bool = True
        Should the shift be taken in units of pixels, or pixel scale.

    Returns
    -------
    wavefront : Wavefront
        The propagated wavefront.
    """
    inverse, plane, units = self._prep_prop(focal_length)

    # Enforce array so output can be vectorised by vmap
    pixel_scale = np.asarray(pixel_scale, float)

    # Calculate
    # Using a self._MFT here allows for broadband wavefronts to define
    # vectorised propagation fn over phasors, wavels, px_scales, etc.
    # It also makes the code muuuuch nicer to read
    args = (npixels, pixel_scale, focal_length, shift, pixel, inverse)
    phasor = self._MFT(
        self.phasor, self.wavelength, self.pixel_scale, *args
    )

    # Update
    return self.set(
        ["amplitude", "phase", "pixel_scale", "plane", "units"],
        [np.abs(phasor), np.angle(phasor), pixel_scale, plane, units],
    )

propagate_FFT(focal_length=None, pad=2)

Propagates the wavefront by performing a Fast Fourier Transform.

Parameters:

Name	Type	Description	Default
focal_length	float = None	The focal length of the propagation. If None, the output pixel scale has units of radians, else meters.	None
pad	int = 2	The padding factory to apply to the input wavefront before the FFT.	2

Returns:

Name	Type	Description
wavefront	Wavefront	The propagated wavefront.

Source code in src/dLux/wavefronts.py

def propagate_FFT(
    self: Wavefront,
    focal_length: float = None,
    pad: int = 2,
) -> Wavefront:
    """
    Propagates the wavefront by performing a Fast Fourier Transform.

    Parameters
    ----------
    focal_length : float = None
        The focal length of the propagation. If None, the output pixel scale has
        units of radians, else meters.
    pad : int = 2
        The padding factory to apply to the input wavefront before the FFT.

    Returns
    -------
    wavefront : Wavefront
        The propagated wavefront.
    """
    inverse, plane, units = self._prep_prop(focal_length)

    # Calculate
    phasor, pixel_scale = dlu.FFT(
        self.phasor,
        self.wavelength,
        self.pixel_scale,
        focal_length,
        pad,
        inverse,
    )

    # Return new wavefront
    return self.set(
        ["amplitude", "phase", "pixel_scale", "plane", "units"],
        [np.abs(phasor), np.angle(phasor), pixel_scale, plane, units],
    )

propagate_fresnel(npixels, pixel_scale, focal_length, focal_shift=0.0, shift=np.zeros(2), pixel=True)

Propagates the phasor using a Far-Field Fresnel propagation. This allows for psfs to be better modelled a few wavelengths from the focal plane.

Parameters:

Name	Type	Description	Default
npixels	int	The number of pixels in the output plane.	required
pixel_scale	(float, meters / pixel or radians / pixel)	The pixel scale of the output plane.	required
focal_length	float	The focal length of the propagation.	required
focal_shift	float	The shift from focus to propagate to.	0.0
shift	Array = np.zeros(2)	The shift in the center of the output plane.	zeros(2)
pixel	bool = True	Should the shift be taken in units of pixels, or pixel scale.	True

Returns:

Name	Type	Description
wavefront	Wavefront	The propagated wavefront.

Source code in src/dLux/wavefronts.py

def propagate_fresnel(
    self: Wavefront,
    npixels: int,
    pixel_scale: float,
    focal_length: float,
    focal_shift: float = 0.0,
    shift: Array = np.zeros(2),
    pixel: bool = True,
) -> Wavefront:
    """
    Propagates the phasor using a Far-Field Fresnel propagation. This allows for
    psfs to be better modelled a few wavelengths from the focal plane.

    Parameters
    ----------
    npixels : int
        The number of pixels in the output plane.
    pixel_scale : float, meters/pixel or radians/pixel
        The pixel scale of the output plane.
    focal_length : float
        The focal length of the propagation.
    focal_shift: float, meters = 0.
        The shift from focus to propagate to.
    shift : Array = np.zeros(2)
        The shift in the center of the output plane.
    pixel : bool = True
        Should the shift be taken in units of pixels, or pixel scale.

    Returns
    -------
    wavefront : Wavefront
        The propagated wavefront.
    """
    # TODO: Try inverse propagation to see if it works, it probably will
    if self.plane == "Pupil":
        inverse = False
    else:
        inverse = True
    plane = "Intermediate"
    units = "Cartesian"

    # We can't fresnel from a focal plane
    if self.plane != "Pupil":
        raise ValueError(
            "Can only do an fresnel propagation from a Pupil plane, "
            f"current plane is {self.plane}."
        )

    # Calculate
    phasor = dlu.fresnel_MFT(
        self.phasor,
        self.wavelength,
        self.pixel_scale,
        npixels,
        pixel_scale,
        focal_length,
        focal_shift,
        shift,
        pixel,
        inverse,
    )

    # Update
    return self.set(
        ["amplitude", "phase", "pixel_scale", "plane", "units"],
        [np.abs(phasor), np.angle(phasor), pixel_scale, plane, units],
    )

resize(npixels)

Resizes the wavefront via a zero-padding or cropping operation.

Parameters:

Name	Type	Description	Default
npixels	int	The size to resize the wavefront to.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The resized wavefront.

Source code in src/dLux/wavefronts.py

def resize(self: Wavefront, npixels: int) -> Wavefront:
    """
    Resizes the wavefront via a zero-padding or cropping operation.

    Parameters
    ----------
    npixels : int
        The size to resize the wavefront to.

    Returns
    -------
    wavefront : Wavefront
        The resized wavefront.
    """
    field = self._to_field()
    amplitude, phase = vmap(dlu.resize, (0, None))(field, npixels)
    return self.set(["amplitude", "phase"], [amplitude, phase])

rotate(angle, order=1, complex=False)

Rotates the wavefront by a given angle via interpolation. Can be done on the real and imaginary components by passing in complex=True.

Parameters:

Name	Type	Description	Default
angle	(Array, radians)	The angle by which to rotate the wavefront in a clockwise direction.	required
order	int = 1	The interpolation order to use.	1
complex	bool = False	Whether to rotate the real and imaginary representation of the wavefront as opposed to the amplitude and phase representation.	False

Returns:

Name	Type	Description
wavefront	Wavefront	The new wavefront rotated by angle in the clockwise direction.

Source code in src/dLux/wavefronts.py

def rotate(
    self: Wavefront, angle: Array, order: int = 1, complex: bool = False
) -> Wavefront:
    """
    Rotates the wavefront by a given angle via interpolation. Can be done on the
    real and imaginary components by passing in complex=True.

    Parameters
    ----------
    angle : Array, radians
        The angle by which to rotate the wavefront in a clockwise
        direction.
    order : int = 1
        The interpolation order to use.
    complex : bool = False
        Whether to rotate the real and imaginary representation of the wavefront as
        opposed to the amplitude and phase representation.

    Returns
    -------
    wavefront : Wavefront
        The new wavefront rotated by angle in the clockwise direction.
    """
    # Get field in either (amplitude, phase) or (real, imaginary)
    field = self._to_field(complex=complex)

    # Rotate the field
    rotator = vmap(dlu.rotate, (0, None, None))
    field = rotator(field, angle, order)

    # Cast back to (amplitude, phase) if needed
    if complex:
        field = self._to_amplitude_phase(field)

    # Return new wavefront
    return self.set(["amplitude", "phase"], [field[0], field[1]])

scale_to(npixels, pixel_scale, complex=False)

Interpolated the wavefront to a given npixels and pixel_scale. Can be done on the real and imaginary components by passing in complex=True.

Parameters:

Name	Type	Description	Default
npixels	int	The number of pixels to interpolate to.	required
pixel_scale	Array	The pixel scale to interpolate to.	required
complex	bool = False	Whether to rotate the real and imaginary representation of the wavefront as opposed to the amplitude and phase representation.	False

Returns:

Name	Type	Description
wavefront	Wavefront	The new interpolated wavefront.

Source code in src/dLux/wavefronts.py

def scale_to(
    self: Wavefront,
    npixels: int,
    pixel_scale: Array,
    complex: bool = False,
) -> Wavefront:
    """
    Interpolated the wavefront to a given npixels and pixel_scale. Can be done on
    the real and imaginary components by passing in complex=True.

    Parameters
    ----------
    npixels : int
        The number of pixels  to interpolate to.
    pixel_scale: Array
        The pixel scale to interpolate to.
    complex : bool = False
        Whether to rotate the real and imaginary representation of the wavefront as
        opposed to the amplitude and phase representation.

    Returns
    -------
    wavefront : Wavefront
        The new interpolated wavefront.
    """
    # Get field in either (amplitude, phase) or (real, imaginary)
    field = self._to_field(complex=complex)

    # Scale the field
    scale_fn = vmap(dlu.scale, (0, None, None))
    field = scale_fn(field, npixels, pixel_scale / self.pixel_scale)

    # Cast back to (amplitude, phase) if needed
    if complex:
        field = self._to_amplitude_phase(field)

    # Return new wavefront
    return self.set(
        ["amplitude", "phase", "pixel_scale"],
        [field[0], field[1], pixel_scale],
    )

tilt(angles)

Tilts the wavefront by the (x, y) angles.

Parameters:

Name	Type	Description	Default
angles	(Array, radians)	The (x, y) angles by which to tilt the wavefront.	required

Returns:

Name	Type	Description
wavefront	Wavefront	The tilted wavefront.

Source code in src/dLux/wavefronts.py

def tilt(self: Wavefront, angles: Array) -> Wavefront:
    """
    Tilts the wavefront by the (x, y) angles.

    Parameters
    ----------
    angles : Array, radians
        The (x, y) angles by which to tilt the wavefront.

    Returns
    -------
    wavefront : Wavefront
        The tilted wavefront.
    """
    if not isinstance(angles, Array) or angles.shape != (2,):
        raise ValueError("angles must be an array of shape (2,).")
    opd = -(angles[:, None, None] * self.coordinates).sum(0)
    return self.add_opd(opd)