Propagation

This module contains the core diffraction engine of dLux, with the core code to propagate phasors is various different ways.

FFT

Calculates the Fast Fourier Transform (FFT) of the input phasor.

Parameters:

Name	Type	Description	Default
phasor	Array[complex]	The input phasor.	required
wavelength	(float, meters)	The wavelength of the input phasor.	required
pixel_scale	(float, meters / pixel)	The pixel scale of the input phasor.	required
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 amount to pad the input array by before propagation. Note this function does not automatically crop the output.	2
inverse	bool = False	Is this a forward or inverse FFT.	False

Returns:

Name	Type	Description
phasor	Array[complex]	The propagated phasor.

Source code in src/dLux/utils/propagation.py

def FFT(
    phasor: Array,
    wavelength: float,
    pixel_scale: float,
    focal_length: float = None,
    pad: int = 2,
    inverse: bool = False,
) -> Array:
    """
    Calculates the Fast Fourier Transform (FFT) of the input phasor.

    Parameters
    ----------
    phasor : Array[complex]
        The input phasor.
    wavelength : float, meters
        The wavelength of the input phasor.
    pixel_scale : float, meters/pixel
        The pixel scale of the input phasor.
    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 amount to pad the input array by before propagation. Note this function
        does not automatically crop the output.
    inverse : bool = False
        Is this a forward or inverse FFT.

    Returns
    -------
    phasor : Array[complex]
        The propagated phasor.
    """
    npixels = phasor.shape[-1]

    # Calculate the output pixel scale
    fringe_size = wavelength / (pixel_scale * npixels)
    new_pixel_scale = fringe_size / pad
    if focal_length is not None:
        new_pixel_scale *= focal_length

    # Pad the input array
    npixels = (npixels * (pad - 1)) // 2
    phasor = np.pad(phasor, npixels)

    # Perform the FFT
    if inverse:
        phasor = np.fft.fft2(np.fft.ifftshift(phasor)) / phasor.shape[-1]
    else:
        phasor = np.fft.fftshift(np.fft.ifft2(phasor)) * phasor.shape[-1]

    return phasor, new_pixel_scale

MFT

Propagates a phasor using a Matrix Fourier Transform (MFT), allowing for output pixel scale and a shift to be specified.

TODO: Add link to Soumer et al. 2007(?), which describes the MFT.

Parameters:

Name	Type	Description	Default
phasor	Array	The input phasor.	required
wavelength	(float, meters)	The wavelength of the input phasor.	required
pixel_scale_in	(float, meters / pixel, radians / pixel)	The pixel scale of the input plane.	required
npixels_out	int	The number of pixels in the output plane.	required
pixel_scale_out	(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
phasor	Array	The propagated phasor.

Source code in src/dLux/utils/propagation.py

def MFT(
    phasor: Array,
    wavelength: float,
    pixel_scale_in: float,
    npixels_out: int,
    pixel_scale_out: float,
    focal_length: float = None,
    shift: Array = np.zeros(2),
    pixel: bool = True,
    inverse: bool = False,
) -> Array:
    """
    Propagates a phasor using a Matrix Fourier Transform (MFT), allowing for output
    pixel scale and a shift to be specified.

    TODO: Add link to Soumer et al. 2007(?), which describes the MFT.

    Parameters
    ----------
    phasor : Array
        The input phasor.
    wavelength : float, meters
        The wavelength of the input phasor.
    pixel_scale_in : float, meters/pixel, radians/pixel
        The pixel scale of the input plane.
    npixels_out : int
        The number of pixels in the output plane.
    pixel_scale_out : 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
    -------
    phasor : Array
        The propagated phasor.
    """
    # Get parameters
    npixels_in = phasor.shape[-1]
    if not pixel:
        shift /= pixel_scale_out

    # Alias the transfer matrix function
    get_tf_mat = lambda s: transfer_matrix(
        wavelength,
        npixels_in,
        pixel_scale_in,
        npixels_out,
        pixel_scale_out,
        s,
        focal_length,
        0.0,
        inverse,
    )

    # Get transfer matrices and propagate
    x_mat, y_mat = vmap(get_tf_mat)(shift)
    phasor = (y_mat.T @ phasor) @ x_mat

    # Normalise
    nfringes = calc_nfringes(
        wavelength,
        npixels_in,
        pixel_scale_in,
        npixels_out,
        pixel_scale_out,
        focal_length,
    )
    phasor *= np.exp(
        np.log(nfringes) - (np.log(npixels_in) + np.log(npixels_out))
    )

    return phasor

fresnel_MFT

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

NOTE: This does have an 'inverse' parameter, however behaviour is not guaranteed to be correct when inverse=True.

Parameters:

Name	Type	Description	Default
phasor	Array[complex]	The input phasor.	required
wavelength	(float, meters)	The wavelength of the input phasor.	required
pixel_scale_in	(float, meters / pixel, radians / pixel)	The pixel scale of the input plane.	required
npixels_out	int	The number of pixels in the output plane.	required
pixel_scale_out	(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.	required
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
inverse	bool	Is this a forward or inverse propagation.	False

Returns:

Name	Type	Description
phasor	Array[complex]	The propagated phasor.

Source code in src/dLux/utils/propagation.py

def fresnel_MFT(
    phasor: Array,
    wavelength: float,
    pixel_scale_in: float,
    npixels_out: int,
    pixel_scale_out: float,
    focal_length: float,
    focal_shift: float,
    shift: Array = np.zeros(2),
    pixel: bool = True,
    inverse: bool = False,
) -> Array:
    """
    Propagates the phasor using a Far-Field Fresnel propagation. This allows for psfs
    to be better modelled a few wavelengths from the focal plane.

    NOTE: This does have an 'inverse' parameter, however behaviour is not
    guaranteed to be correct when `inverse=True`.

    Parameters
    ----------
    phasor : Array[complex]
        The input phasor.
    wavelength : float, meters
        The wavelength of the input phasor.
    pixel_scale_in : float, meters/pixel, radians/pixel
        The pixel scale of the input plane.
    npixels_out : int
        The number of pixels in the output plane.
    pixel_scale_out : 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
        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.
    inverse: bool = False
        Is this a forward or inverse propagation.

    Returns
    -------
    phasor : Array[complex]
        The propagated phasor.
    """
    # Calculate phase factors
    first_factor, second_factor = fresnel_phase_factors(
        wavelength,
        phasor.shape[-1],
        pixel_scale_in,
        npixels_out,
        pixel_scale_out,
        focal_length,
        focal_shift,
    )

    # Propagate
    phasor *= first_factor
    phasor = MFT(
        phasor,
        wavelength,
        pixel_scale_in,
        npixels_out,
        pixel_scale_out,
        focal_length,
        shift,
        pixel,
        inverse,
    )

    phasor *= second_factor
    return phasor