Functional

Sources

FieldPupil = Callable[[Field], Field] module-attribute

__all__ = ['point_source', 'objective_point_source', 'plane_wave', 'gaussian_plane_wave', 'generic_field'] module-attribute

gaussian_plane_wave(shape, dx, spectrum, waist, power=1.0, amplitude=1.0, kykx=(0.0, 0.0), pupil=None, scalar=True)

Generates plane wave of given power with a Gaussian intensity profile (as opposed to a totally flat plane wave). Can also be given pupil and kykx vector to control the angle of the plane wave.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).	required
dx	ScalarLike | Float[Array, wv] | Float[Array, 'wv 2']	The spacing (pixel size) of the samples of the Field in units of distance. In the simplest case, you can pass a scalar value to define this spacing which creates a square pixel. Each sample (pixel) of the Field has a height and width. Spacing is also allowed to vary with wavelength of the spectrum. To choose a different spacing per wavelength (but still square), you can pass a 1D array of spacings of the same length as the number of wavelengths in the spectrum of the Field. To create a non-square spacing, you must always pass a 2D array of shape (wavelengths 2) where the last dimension has length 2 and defines the height and width of the spacing in units of distance. To createa a non-square spacing you must always include the wavelengths dimension even if there is only a single wavelength in the spectrum (i.e. a (1 2) shaped array).	required
spectrum	Spectrum | ScalarLike | Float[Array, wv] | tuple[Float[Array, wv], Float[Array, wv]]	The Spectrum of the Field to be created. This can be specified either as a single float value representing a wavelength in units of distance for a monochromatic field, a 1D array of wavelengths for a chromatic field that has the same intensity in all wavelengths, or a tuple of two 1D arrays where the first array represents the wavelengths and the second array is a unitless array of weights that define the spectral density (the relative intensity of each wavelength in the spectrum). This second array of spectral density will automatically be normalized to sum to 1.	required
waist	ScalarLike	The size of the waist (twice the standard deviation) as a scalar in units of distance. This waist defines the width of the 2D Gaussian amplitude profile of this beam.	required
power	ScalarLike | None	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.	1.0
amplitude	ScalarLike | Float[Array, 3]	The amplitude of the electric field. For scalar Fields this doesnt do anything but scale the field (which will be undone if power is not None), but it is required for vectorial Fields to set the polarization.	1.0
kykx	ArrayLike | tuple[int, int]	Defines the orientation of the plane wave. Should be a tuple or an array of shape (2,) in the format [ky kx]. We assume that these are wave vectors, i.e. that they have already been multiplied by 2 * pi / wavelength.	(0.0, 0.0)
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.	None
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.	True

Source code in src/chromatix/functional/sources.py

def gaussian_plane_wave(
    shape: tuple[int, int],
    dx: ScalarLike | Float[Array, "wv"] | Float[Array, "wv 2"],
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    waist: ScalarLike,
    power: ScalarLike | None = 1.0,
    amplitude: ScalarLike | Float[Array, "3"] = 1.0,
    kykx: ArrayLike | tuple[int, int] = (0.0, 0.0),
    pupil: FieldPupil | None = None,
    scalar: bool = True,
) -> Field:
    """
    Generates plane wave of given ``power`` with a Gaussian intensity profile
    (as opposed to a totally flat plane wave). Can also be given ``pupil`` and
    ``kykx`` vector to control the angle of the plane wave.

    Args:
        shape: The shape (height and width) of the ``Field`` to be created in
            number of samples (pixels).
        dx: The spacing (pixel size) of the samples of the ``Field`` in units
            of distance. In the simplest case, you can pass a scalar value
            to define this spacing which creates a square pixel. Each sample
            (pixel) of the ``Field`` has a height and width. Spacing is also
            allowed to vary with wavelength of the spectrum. To choose a
            different spacing per wavelength (but still square), you can pass a
            1D array of spacings of the same length as the number of wavelengths
            in the spectrum of the ``Field``. To create a non-square spacing,
            you must always pass a 2D array of shape `(wavelengths 2)` where
            the last dimension has length 2 and defines the height and width of
            the spacing in units of distance. To createa a non-square spacing
            you must always include the `wavelengths` dimension even if there
            is only a single wavelength in the spectrum (i.e. a `(1 2)` shaped
            array).
        spectrum: The
            [``Spectrum``](core.md#chromatix.core.spectrum.Spectrum.build)
            of the ``Field`` to be created. This can be specified either as a
            single float value representing a wavelength in units of distance
            for a monochromatic field, a 1D array of wavelengths for a chromatic
            field that has the same intensity in all wavelengths, or a tuple
            of two 1D arrays where the first array represents the wavelengths
            and the second array is a unitless array of weights that define the
            spectral density (the relative intensity of each wavelength in the
            spectrum). This second array of spectral density will automatically
            be normalized to sum to 1.
        waist: The size of the waist (twice the standard deviation) as a
            scalar in units of distance. This waist defines the width of the 2D
            Gaussian amplitude profile of this beam.
        power: The total power that the result should be normalized to, defaults
            to 1.0. If ``None``, no normalization occurs.
        amplitude: The amplitude of the electric field. For scalar ``Field``s
            this doesnt do anything but scale the field (which will be undone
            if ``power`` is not ``None``), but it is required for vectorial
            ``Field``s to set the polarization.
        kykx: Defines the orientation of the plane wave. Should be a tuple or an
            array of shape `(2,)` in the format `[ky kx]`. We assume that these
            are wave vectors, i.e. that they have already been multiplied by ``2
            * pi / wavelength``.
        pupil: A function that applies a pupil to the field if provided.
            Defaults to ``None`` in which case the field is unchanged.
        scalar: Whether the result should be ``ScalarField`` (if True) or
            ``VectorField`` (if False). Defaults to True.
    """
    spectrum = Spectrum.build(spectrum)
    field = Field.empty(shape, dx, spectrum, scalar)
    amplitude = jnp.atleast_1d(jnp.asarray(amplitude))
    if scalar:
        assert_axis_dimension(amplitude, -1, 1)
    else:
        assert_axis_dimension(amplitude, -1, 3)
    kykx = jnp.atleast_1d(jnp.asarray(kykx))
    u = amplitude * jnp.exp(1j * jnp.sum(kykx * field.grid, axis=-1))
    # There's no spectral dependence so we need to manually put in the spectral axis
    # hence the ones_like term.
    field = field.replace(u=u * jnp.ones_like(field.u))
    field = gaussian_pupil(field, waist)
    if pupil is not None:
        field = pupil(field)
    if power is not None:
        field = field * jnp.sqrt(power / field.power)
    return field

generic_field(dx, spectrum, amplitude, phase, power=1.0, pupil=None, scalar=True)

Generates field with arbitrary phase and amplitude. You can likely use the appropriate constructor for the type of Field you want rather than this function, or just use Field.build.

Parameters:

Name	Type	Description	Default
dx	ScalarLike | Float[Array, wv] | Float[Array, 'wv 2']	The spacing (pixel size) of the samples of the Field in units of distance. In the simplest case, you can pass a scalar value to define this spacing which creates a square pixel. Each sample (pixel) of the Field has a height and width. Spacing is also allowed to vary with wavelength of the spectrum. To choose a different spacing per wavelength (but still square), you can pass a 1D array of spacings of the same length as the number of wavelengths in the spectrum of the Field. To create a non-square spacing, you must always pass a 2D array of shape (wavelengths 2) where the last dimension has length 2 and defines the height and width of the spacing in units of distance. To createa a non-square spacing you must always include the wavelengths dimension even if there is only a single wavelength in the spectrum (i.e. a (1 2) shaped array).	required
spectrum	Spectrum | ScalarLike | Float[Array, wv] | tuple[Float[Array, wv], Float[Array, wv]]	The Spectrum of the Field to be created. This can be specified either as a single float value representing a wavelength in units of distance for a monochromatic field, a 1D array of wavelengths for a chromatic field that has the same intensity in all wavelengths, or a tuple of two 1D arrays where the first array represents the wavelengths and the second array is a unitless array of weights that define the spectral density (the relative intensity of each wavelength in the spectrum). This second array of spectral density will automatically be normalized to sum to 1.	required
amplitude	ArrayLike	The amplitude of the field with shape (... height width) for scalar monochromatic fields, shape (... height width wavelengths) for scalar fields with multiple wavelengths, shape (height width 3) for vectorial monochromatic fields, or shape (... height width wavelengths 3) for vectorial fields with multiple wavelengths.	required
phase	ArrayLike	The phase of the field with shape (... height width) for scalar monochromatic fields, shape (... height width wavelengths) for scalar fields with multiple wavelengths, shape (height width 3) for vectorial monochromatic fields, or shape (... height width wavelengths 3) for vectorial fields with multiple wavelengths.	required
power	ScalarLike | None	The total power (a scalar value) that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.	1.0
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.	None
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.	True

Source code in src/chromatix/functional/sources.py

def generic_field(
    dx: ScalarLike | Float[Array, "wv"] | Float[Array, "wv 2"],
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    amplitude: ArrayLike,
    phase: ArrayLike,
    power: ScalarLike | None = 1.0,
    pupil: FieldPupil | None = None,
    scalar: bool = True,
) -> Field:
    """
    Generates field with arbitrary ``phase`` and ``amplitude``.
    You can likely use the appropriate constructor for the type
    of ``Field`` you want rather than this function, or just use
    [``Field.build``](field.md#chromatix.core.field.Field.build).

    Args:
        dx: The spacing (pixel size) of the samples of the ``Field`` in units
            of distance. In the simplest case, you can pass a scalar value
            to define this spacing which creates a square pixel. Each sample
            (pixel) of the ``Field`` has a height and width. Spacing is also
            allowed to vary with wavelength of the spectrum. To choose a
            different spacing per wavelength (but still square), you can pass a
            1D array of spacings of the same length as the number of wavelengths
            in the spectrum of the ``Field``. To create a non-square spacing,
            you must always pass a 2D array of shape `(wavelengths 2)` where
            the last dimension has length 2 and defines the height and width of
            the spacing in units of distance. To createa a non-square spacing
            you must always include the `wavelengths` dimension even if there
            is only a single wavelength in the spectrum (i.e. a `(1 2)` shaped
            array).
        spectrum: The
            [``Spectrum``](core.md#chromatix.core.spectrum.Spectrum.build)
            of the ``Field`` to be created. This can be specified either as a
            single float value representing a wavelength in units of distance
            for a monochromatic field, a 1D array of wavelengths for a chromatic
            field that has the same intensity in all wavelengths, or a tuple
            of two 1D arrays where the first array represents the wavelengths
            and the second array is a unitless array of weights that define the
            spectral density (the relative intensity of each wavelength in the
            spectrum). This second array of spectral density will automatically
            be normalized to sum to 1.
        amplitude: The amplitude of the field with shape `(... height
            width)` for scalar monochromatic fields, shape `(... height width
            wavelengths)` for scalar fields with multiple wavelengths, shape
            `(height width 3)` for vectorial monochromatic fields, or shape
            `(... height width wavelengths 3)` for vectorial fields with
            multiple wavelengths.
        phase: The phase of the field with shape `(... height width)` for scalar
            monochromatic fields, shape `(... height width wavelengths)` for
            scalar fields with multiple wavelengths, shape `(height width 3)`
            for vectorial monochromatic fields, or shape `(... height width
            wavelengths 3)` for vectorial fields with multiple wavelengths.
        power: The total power (a scalar value) that the result should be
            normalized to, defaults to 1.0. If ``None``, no normalization
            occurs.
        pupil: A function that applies a pupil to the field if provided.
            Defaults to ``None`` in which case the field is unchanged.
        scalar: Whether the result should be ``ScalarField`` (if ``True``) or
            ``VectorField`` (if ``False``). Defaults to ``True``.
    """
    spectrum = Spectrum.build(spectrum)
    assert_equal_shape([amplitude, phase])
    match (scalar, isinstance(spectrum, MonoSpectrum)):
        case (False, False):
            assert amplitude.ndim >= 4, (
                "Amplitude must have at least 4 dimensions: (height width wavelengths 3)"
            )
            assert phase.ndim >= 4, (
                "Phase must have at least 4 dimensions: (height width wavelengths 3)"
            )
            assert_axis_dimension(amplitude, -1, 3)
            assert_axis_dimension(phase, -1, 3)
        case (False, True):
            assert amplitude.ndim >= 3, (
                "Amplitude must have at least 3 dimensions: (height width 3)"
            )
            assert phase.ndim >= 3, (
                "Phase must have at least 3 dimensions: (height width 3)"
            )
            assert_axis_dimension(amplitude, -1, 3)
            assert_axis_dimension(phase, -1, 3)
        case (True, False):
            assert amplitude.ndim >= 3, (
                "Amplitude must have at least 3 dimensions: (height width wavelengths)"
            )
            assert phase.ndim >= 3, (
                "Phase must have at least 3 dimensions: (height width wavelengths)"
            )
        case (True, True):
            assert amplitude.ndim >= 2, (
                "Amplitude must have at least 2 dimensions: (height width)"
            )
            assert phase.ndim >= 2, (
                "Phase must have at least 2 dimensions: (height width)"
            )
    u = jnp.asarray(amplitude) * jnp.exp(1j * jnp.asarray(phase))
    field = Field.build(u, dx, spectrum)
    if pupil is not None:
        field = pupil(field)
    if power is not None:
        field = field * jnp.sqrt(power / field.power)
    return field

objective_point_source(shape, dx, spectrum, z, f, n, NA, power=1.0, amplitude=1.0, offset=(0.0, 0.0), scalar=True)

Generates field due to a point source defocused by an amount z away from the focal plane, just after passing through a thin lens with focal length f and numerical aperture NA.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).	required
dx	ScalarLike | Float[Array, wv] | Float[Array, 'wv 2']	The spacing (pixel size) of the samples of the Field in units of distance. In the simplest case, you can pass a scalar value to define this spacing which creates a square pixel. Each sample (pixel) of the Field has a height and width. Spacing is also allowed to vary with wavelength of the spectrum. To choose a different spacing per wavelength (but still square), you can pass a 1D array of spacings of the same length as the number of wavelengths in the spectrum of the Field. To create a non-square spacing, you must always pass a 2D array of shape (wavelengths 2) where the last dimension has length 2 and defines the height and width of the spacing in units of distance. To createa a non-square spacing you must always include the wavelengths dimension even if there is only a single wavelength in the spectrum (i.e. a (1 2) shaped array).	required
spectrum	Spectrum | ScalarLike | Float[Array, wv] | tuple[Float[Array, wv], Float[Array, wv]]	The Spectrum of the Field to be created. This can be specified either as a single float value representing a wavelength in units of distance for a monochromatic field, a 1D array of wavelengths for a chromatic field that has the same intensity in all wavelengths, or a tuple of two 1D arrays where the first array represents the wavelengths and the second array is a unitless array of weights that define the spectral density (the relative intensity of each wavelength in the spectrum). This second array of spectral density will automatically be normalized to sum to 1.	required
z	ScalarLike | Float[Array, z]	How far away the point source is from the focal plane of the lens in units of distance.	required
f	ScalarLike	Focal length of the objective lens in units of distance.	required
n	ScalarLike	Refractive index.	required
NA	ScalarLike	The numerical aperture of the objective lens.	required
power	ScalarLike | None	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.	1.0
amplitude	ScalarLike | Float[Array, 3]	The amplitude of the electric field. For scalar Fields this doesnt do anything but scale the field (which will be undone if power is not None), but it is required for vectorial Fields to set the polarization.	1.0
offset	Float[Array, 2] | tuple[float, float]	The offset of the point source in the plane in units of distance. Should be a tuple or an array of shape (2,) in the order y x.	(0.0, 0.0)
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.	True

Source code in src/chromatix/functional/sources.py

def objective_point_source(
    shape: tuple[int, int],
    dx: ScalarLike | Float[Array, "wv"] | Float[Array, "wv 2"],
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    z: ScalarLike | Float[Array, "z"],
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike,
    power: ScalarLike | None = 1.0,
    amplitude: ScalarLike | Float[Array, "3"] = 1.0,
    offset: Float[Array, "2"] | tuple[float, float] = (0.0, 0.0),
    scalar: bool = True,
) -> ScalarField | VectorField:
    """
    Generates field due to a point source defocused by an amount ``z`` away from
    the focal plane, just after passing through a thin lens with focal length
    ``f`` and numerical aperture ``NA``.

    Args:
        shape: The shape (height and width) of the ``Field`` to be created in
            number of samples (pixels).
        dx: The spacing (pixel size) of the samples of the ``Field`` in units
            of distance. In the simplest case, you can pass a scalar value
            to define this spacing which creates a square pixel. Each sample
            (pixel) of the ``Field`` has a height and width. Spacing is also
            allowed to vary with wavelength of the spectrum. To choose a
            different spacing per wavelength (but still square), you can pass a
            1D array of spacings of the same length as the number of wavelengths
            in the spectrum of the ``Field``. To create a non-square spacing,
            you must always pass a 2D array of shape `(wavelengths 2)` where
            the last dimension has length 2 and defines the height and width of
            the spacing in units of distance. To createa a non-square spacing
            you must always include the `wavelengths` dimension even if there
            is only a single wavelength in the spectrum (i.e. a `(1 2)` shaped
            array).
        spectrum: The
            [``Spectrum``](core.md#chromatix.core.spectrum.Spectrum.build)
            of the ``Field`` to be created. This can be specified either as a
            single float value representing a wavelength in units of distance
            for a monochromatic field, a 1D array of wavelengths for a chromatic
            field that has the same intensity in all wavelengths, or a tuple
            of two 1D arrays where the first array represents the wavelengths
            and the second array is a unitless array of weights that define the
            spectral density (the relative intensity of each wavelength in the
            spectrum). This second array of spectral density will automatically
            be normalized to sum to 1.
        z: How far away the point source is from the focal plane of the lens in
            units of distance.
        f: Focal length of the objective lens in units of distance.
        n: Refractive index.
        NA: The numerical aperture of the objective lens.
        power: The total power that the result should be normalized to, defaults
            to 1.0. If ``None``, no normalization occurs.
        amplitude: The amplitude of the electric field. For scalar ``Field``s
            this doesnt do anything but scale the field (which will be undone
            if ``power`` is not ``None``), but it is required for vectorial
            ``Field``s to set the polarization.
        offset: The offset of the point source in the plane in units of
            distance. Should be a tuple or an array of shape `(2,)` in the order
            `y x`.
        scalar: Whether the result should be ``ScalarField`` (if True) or
            ``VectorField`` (if False). Defaults to True.
    """
    spectrum = Spectrum.build(spectrum)
    field = Field.empty(shape, dx, spectrum, scalar)
    # If scalar, last axis should 1, else 3.
    amplitude = jnp.atleast_1d(jnp.asarray(amplitude))
    if scalar:
        assert_axis_dimension(amplitude, -1, 1)
    else:
        assert_axis_dimension(amplitude, -1, 3)
    z = jnp.atleast_1d(jnp.asarray(z)).squeeze()
    if z.size > 1:
        z = _broadcast_1d_to_innermost_batch(z, field.spatial_dims)
        field = field.expand_dims()
    offset = jnp.atleast_1d(jnp.asarray(offset))
    L = jnp.sqrt(field.broadcasted_wavelength * f / n)
    phase = -jnp.pi * (z / f) * l2_sq_norm(field.grid - offset) / L**2
    u = amplitude * -1j / L**2 * jnp.exp(1j * phase)
    field = field.replace(u=u)
    D = 2 * f * NA / n
    field = circular_pupil(field, D)  # type: ignore
    if power is not None:
        field = field * jnp.sqrt(power / field.power)
    return field

plane_wave(shape, dx, spectrum, power=1.0, amplitude=1.0, kykx=(0.0, 0.0), pupil=None, scalar=True)

Generates plane wave of given power, as exp(1j) at each location of the field. Can also be given pupil and kykx vector to control the angle of the plane wave. If a kykx wave vector is provided, the plane wave is constructed as exp(1j * jnp.sum(kykx * field.grid, axis=-1)).

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).	required
dx	ScalarLike | Float[Array, wv] | Float[Array, 'wv 2']	The spacing (pixel size) of the samples of the Field in units of distance. In the simplest case, you can pass a scalar value to define this spacing which creates a square pixel. Each sample (pixel) of the Field has a height and width. Spacing is also allowed to vary with wavelength of the spectrum. To choose a different spacing per wavelength (but still square), you can pass a 1D array of spacings of the same length as the number of wavelengths in the spectrum of the Field. To create a non-square spacing, you must always pass a 2D array of shape (wavelengths 2) where the last dimension has length 2 and defines the height and width of the spacing in units of distance. To createa a non-square spacing you must always include the wavelengths dimension even if there is only a single wavelength in the spectrum (i.e. a (1 2) shaped array).	required
spectrum	Spectrum | ScalarLike | Float[Array, wv] | tuple[Float[Array, wv], Float[Array, wv]]	The Spectrum of the Field to be created. This can be specified either as a single float value representing a wavelength in units of distance for a monochromatic field, a 1D array of wavelengths for a chromatic field that has the same intensity in all wavelengths, or a tuple of two 1D arrays where the first array represents the wavelengths and the second array is a unitless array of weights that define the spectral density (the relative intensity of each wavelength in the spectrum). This second array of spectral density will automatically be normalized to sum to 1.	required
power	ScalarLike | None	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.	1.0
amplitude	ScalarLike | Float[Array, 3]	The amplitude of the electric field. For scalar Fields this doesnt do anything but scale the field (which will be undone if power is not None), but it is required for vectorial Fields to set the polarization.	1.0
kykx	ArrayLike | tuple[float, float]	Defines the orientation of the plane wave. Should be a tuple or an array of shape (2,) in the format [ky kx]. We assume that these are wave vectors, i.e. that they have already been multiplied by 2 * pi / wavelength.	(0.0, 0.0)
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.	None
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.	True

Source code in src/chromatix/functional/sources.py

def plane_wave(
    shape: tuple[int, int],
    dx: ScalarLike | Float[Array, "wv"] | Float[Array, "wv 2"],
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    power: ScalarLike | None = 1.0,
    amplitude: ScalarLike | Float[Array, "3"] = 1.0,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    pupil: FieldPupil | None = None,
    scalar: bool = True,
) -> ScalarField | VectorField:
    """
    Generates plane wave of given ``power``, as ``exp(1j)`` at each location of
    the field. Can also be given ``pupil`` and ``kykx`` vector to control the
    angle of the plane wave. If a ``kykx`` wave vector is provided, the plane
    wave is constructed as ``exp(1j * jnp.sum(kykx * field.grid, axis=-1))``.

    Args:
        shape: The shape (height and width) of the ``Field`` to be created in
            number of samples (pixels).
        dx: The spacing (pixel size) of the samples of the ``Field`` in units
            of distance. In the simplest case, you can pass a scalar value
            to define this spacing which creates a square pixel. Each sample
            (pixel) of the ``Field`` has a height and width. Spacing is also
            allowed to vary with wavelength of the spectrum. To choose a
            different spacing per wavelength (but still square), you can pass a
            1D array of spacings of the same length as the number of wavelengths
            in the spectrum of the ``Field``. To create a non-square spacing,
            you must always pass a 2D array of shape `(wavelengths 2)` where
            the last dimension has length 2 and defines the height and width of
            the spacing in units of distance. To createa a non-square spacing
            you must always include the `wavelengths` dimension even if there
            is only a single wavelength in the spectrum (i.e. a `(1 2)` shaped
            array).
        spectrum: The
            [``Spectrum``](core.md#chromatix.core.spectrum.Spectrum.build)
            of the ``Field`` to be created. This can be specified either as a
            single float value representing a wavelength in units of distance
            for a monochromatic field, a 1D array of wavelengths for a chromatic
            field that has the same intensity in all wavelengths, or a tuple
            of two 1D arrays where the first array represents the wavelengths
            and the second array is a unitless array of weights that define the
            spectral density (the relative intensity of each wavelength in the
            spectrum). This second array of spectral density will automatically
            be normalized to sum to 1.
        power: The total power that the result should be normalized to, defaults
            to 1.0. If ``None``, no normalization occurs.
        amplitude: The amplitude of the electric field. For scalar ``Field``s
            this doesnt do anything but scale the field (which will be undone
            if ``power`` is not ``None``), but it is required for vectorial
            ``Field``s to set the polarization.
        kykx: Defines the orientation of the plane wave. Should be a tuple or an
            array of shape `(2,)` in the format `[ky kx]`. We assume that these
            are wave vectors, i.e. that they have already been multiplied by ``2
            * pi / wavelength``.
        pupil: A function that applies a pupil to the field if provided.
            Defaults to ``None`` in which case the field is unchanged.
        scalar: Whether the result should be ``ScalarField`` (if True) or
            ``VectorField`` (if False). Defaults to True.
    """
    spectrum = Spectrum.build(spectrum)
    field = Field.empty(shape, dx, spectrum, scalar)
    # If scalar, last axis should 1, else 3.
    amplitude = jnp.atleast_1d(amplitude)
    if scalar:
        assert_axis_dimension(amplitude, -1, 1)
    else:
        assert_axis_dimension(amplitude, -1, 3)
    kykx = jnp.atleast_1d(jnp.asarray(kykx))
    u = amplitude * jnp.exp(1j * jnp.sum(kykx * field.grid, axis=-1))
    # NOTE(dd/2025-10-02): There's no vectorial dependence on the grid so we
    # need to make sure to match the right shape, hence the multiplication
    # by the ones_like term to ensure we broadcast properly in the case of a
    # vectorial Field.
    field = field.replace(u=u * jnp.ones_like(field.u))
    if pupil is not None:
        field = pupil(field)
    if power is not None:
        field = field * jnp.sqrt(power / field.power)
    return field

point_source(shape, dx, spectrum, z, n, power=1.0, amplitude=1.0, offset=(0.0, 0.0), pupil=None, scalar=True, epsilon=float(np.finfo(np.float32).eps))

Generates field due to point source a distance z away. Can also be given pupil.

Warning

This function is numerically unstable at z = 0, so an epsilon is applied.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).	required
dx	ScalarLike | Float[Array, wv] | Float[Array, 'wv 2']	The spacing (pixel size) of the samples of the Field in units of distance. In the simplest case, you can pass a scalar value to define this spacing which creates a square pixel. Each sample (pixel) of the Field has a height and width. Spacing is also allowed to vary with wavelength of the spectrum. To choose a different spacing per wavelength (but still square), you can pass a 1D array of spacings of the same length as the number of wavelengths in the spectrum of the Field. To create a non-square spacing, you must always pass a 2D array of shape (wavelengths 2) where the last dimension has length 2 and defines the height and width of the spacing in units of distance. To createa a non-square spacing you must always include the wavelengths dimension even if there is only a single wavelength in the spectrum (i.e. a (1 2) shaped array).	required
spectrum	Spectrum | ScalarLike | Float[Array, wv] | tuple[Float[Array, wv], Float[Array, wv]]	The Spectrum of the Field to be created. This can be specified either as a single float value representing a wavelength in units of distance for a monochromatic field, a 1D array of wavelengths for a chromatic field that has the same intensity in all wavelengths, or a tuple of two 1D arrays where the first array represents the wavelengths and the second array is a unitless array of weights that define the spectral density (the relative intensity of each wavelength in the spectrum). This second array of spectral density will automatically be normalized to sum to 1.	required
z	ScalarLike | Float[Array, z]	How far away the point source is in units of distance.	required
n	ScalarLike	Refractive index.	required
power	ScalarLike | None	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.	1.0
amplitude	ScalarLike | Float[Array, 3]	The amplitude of the electric field. For scalar Fields this doesnt do anything but scale the field (which will be undone if power is not None), but it is required for vectorial Fields to set the polarization.	1.0
offset	Float[Array, 2] | tuple[float, float]	The offset of the point source in the plane in units of distance. Should be a tuple or an array of shape (2,) in the order y x.	(0.0, 0.0)
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.	None
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.	True
epsilon	float	Value added to denominators for numerical stability when z is 0.	float(eps)

Source code in src/chromatix/functional/sources.py

def point_source(
    shape: tuple[int, int],
    dx: ScalarLike | Float[Array, "wv"] | Float[Array, "wv 2"],
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    z: ScalarLike | Float[Array, "z"],
    n: ScalarLike,
    power: ScalarLike | None = 1.0,
    amplitude: ScalarLike | Float[Array, "3"] = 1.0,
    offset: Float[Array, "2"] | tuple[float, float] = (0.0, 0.0),
    pupil: FieldPupil | None = None,
    scalar: bool = True,
    epsilon: float = float(np.finfo(np.float32).eps),
) -> ScalarField | VectorField:
    """
    Generates field due to point source a distance ``z`` away. Can also be given
    ``pupil``.

    !!! warning
        This function is numerically unstable at z = 0, so an epsilon is applied.

    Args:
        shape: The shape (height and width) of the ``Field`` to be created in
            number of samples (pixels).
        dx: The spacing (pixel size) of the samples of the ``Field`` in units
            of distance. In the simplest case, you can pass a scalar value
            to define this spacing which creates a square pixel. Each sample
            (pixel) of the ``Field`` has a height and width. Spacing is also
            allowed to vary with wavelength of the spectrum. To choose a
            different spacing per wavelength (but still square), you can pass a
            1D array of spacings of the same length as the number of wavelengths
            in the spectrum of the ``Field``. To create a non-square spacing,
            you must always pass a 2D array of shape `(wavelengths 2)` where
            the last dimension has length 2 and defines the height and width of
            the spacing in units of distance. To createa a non-square spacing
            you must always include the `wavelengths` dimension even if there
            is only a single wavelength in the spectrum (i.e. a `(1 2)` shaped
            array).
        spectrum: The
            [``Spectrum``](core.md#chromatix.core.spectrum.Spectrum.build)
            of the ``Field`` to be created. This can be specified either as a
            single float value representing a wavelength in units of distance
            for a monochromatic field, a 1D array of wavelengths for a chromatic
            field that has the same intensity in all wavelengths, or a tuple
            of two 1D arrays where the first array represents the wavelengths
            and the second array is a unitless array of weights that define the
            spectral density (the relative intensity of each wavelength in the
            spectrum). This second array of spectral density will automatically
            be normalized to sum to 1.
        z: How far away the point source is in units of distance.
        n: Refractive index.
        power: The total power that the result should be normalized to, defaults
            to 1.0. If ``None``, no normalization occurs.
        amplitude: The amplitude of the electric field. For scalar ``Field``s
            this doesnt do anything but scale the field (which will be undone
            if ``power`` is not ``None``), but it is required for vectorial
            ``Field``s to set the polarization.
        offset: The offset of the point source in the plane in units of
            distance. Should be a tuple or an array of shape `(2,)` in the
            order `y x`.
        pupil: A function that applies a pupil to the field if provided.
            Defaults to ``None`` in which case the field is unchanged.
        scalar: Whether the result should be ``ScalarField`` (if ``True``) or
            ``VectorField`` (if ``False``). Defaults to ``True``.
        epsilon: Value added to denominators for numerical stability when z is 0.
    """
    spectrum = Spectrum.build(spectrum)
    field = Field.empty(shape, dx, spectrum, scalar)
    # If scalar, last axis should 1, else 3.
    amplitude = jnp.atleast_1d(amplitude)
    if scalar:
        assert_axis_dimension(amplitude, -1, 1)
    else:
        assert_axis_dimension(amplitude, -1, 3)
    z = jnp.atleast_1d(jnp.asarray(z)).squeeze()
    if z.size > 1:
        z = _broadcast_1d_to_innermost_batch(z, field.spatial_dims)
        field = field.expand_dims()
    offset = jnp.atleast_1d(jnp.asarray(offset))
    L = jnp.sqrt(field.broadcasted_wavelength * jnp.abs(z) / n)
    L_sq = jnp.sign(z) * jnp.fmax(L**2, epsilon)
    phase = jnp.pi * l2_sq_norm(field.grid - offset) / L_sq
    u = amplitude * -1j / L_sq * jnp.exp(1j * phase)
    field = field.replace(u=u)
    if pupil is not None:
        field = pupil(field)
    if power is not None:
        field = field * jnp.sqrt(power / field.power)
    return field

Lenses

__all__ = ['thin_lens', 'ff_lens', 'df_lens', 'microlens_array', 'hexagonal_microlens_array', 'rectangular_microlens_array', 'thick_plano_convex_lens', 'thick_plano_convex_ff_lens', 'high_na_ff_lens'] module-attribute

df_lens(field, d, f, n, NA=None, inverse=False)

Applies a thin lens placed a distance d after the incoming Field.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
d	ScalarLike	Distance from the incoming Field to the lens.	required
f	ScalarLike	Focal length of the lens in units of distance.	required
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.	None
inverse	bool	Whether the field is passing forwards or backwards through the lens. If True, the phase of the lens is conjugated. Defaults to False.	False

Returns:

Type	Description
Field	The Field propagated a distance f after the lens.

Source code in src/chromatix/functional/lenses.py

def df_lens(
    field: Field,
    d: ScalarLike,
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike | None = None,
    inverse: bool = False,
) -> Field:
    """
    Applies a thin lens placed a distance ``d`` after the incoming ``Field``.

    Args:
        field: The ``Field`` to which the lens will be applied.
        d: Distance from the incoming ``Field`` to the lens.
        f: Focal length of the lens in units of distance.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        NA: If provided, the NA of the lens. By default, no pupil is applied
            to the incoming ``Field``.
        inverse: Whether the field is passing forwards or backwards through
            the lens. If ``True``, the phase of the lens is conjugated.
            Defaults to ``False``.

    Returns:
        The ``Field`` propagated a distance ``f`` after the lens.
    """
    if NA is not None:
        D = 2 * f * NA / n  # Expression for NA yields width of pupil
        field = circular_pupil(field, D)

    if inverse:
        # if inverse, propagate over negative distance
        f = -d
        d = -f
    field = optical_fft(field, f, n)

    # Phase factor due to distance d from lens
    L = jnp.sqrt(jnp.complex64(field.broadcasted_wavelength * f / n))  # Lengthscale L
    phase = jnp.pi * (1 - d / f) * l2_sq_norm(field.grid) / jnp.abs(L) ** 2
    return field * jnp.exp(1j * phase)

ff_lens(field, f, n, NA=None, inverse=False)

Applies a thin lens placed a distance f after the incoming Field.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
f	ScalarLike	Focal length of the lens in units of distance.	required
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.	None
inverse	bool	Whether the field is passing forwards or backwards through the lens. If True, the phase of the lens is conjugated. Defaults to False.	False

Returns:

Type	Description
Field	The Field propagated a distance f to and after the lens.

Source code in src/chromatix/functional/lenses.py

def ff_lens(
    field: Field,
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike | None = None,
    inverse: bool = False,
) -> Field:
    """
    Applies a thin lens placed a distance ``f`` after the incoming ``Field``.

    Args:
        field: The ``Field`` to which the lens will be applied.
        f: Focal length of the lens in units of distance.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        NA: If provided, the NA of the lens. By default, no pupil is applied
            to the incoming ``Field``.
        inverse: Whether the field is passing forwards or backwards through
            the lens. If ``True``, the phase of the lens is conjugated.
            Defaults to ``False``.

    Returns:
        The ``Field`` propagated a distance ``f`` to and after the lens.
    """
    # Pupil
    if NA is not None:
        D = 2 * f * NA / n  # Expression for NA yields width of pupil
        field = circular_pupil(field, D)
    if inverse:
        # if inverse, propagate over negative distance
        f = -f
    return optical_fft(field, f, n)

hexagonal_microlens_array(field, f, n, num_lenses_per_side, radius, separation, block_between=False)

Applies a microlens array of hexagonally arranged microlenses placed immediately in the plane of the incoming Field.

Warning

If you have recently used this function prior to it being documented, note that the arguments have changed.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
f	ScalarLike	A scalar value defining the focal length of each lens in units of distance.	required
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
num_lenses_per_side	int	The number of lenses on each outer side of the hexagon (e.g. setting this number to 4 will create 37 microlenses).	required
radius	Array	A scalar value defining the radius of each microlens in units of distance.	required
separation	ScalarLike	A scalar value defining how far apart the center of each microlens is from its neighbors in units of distance.	required
block_between	bool	If True, will mask out the Field in the spaces between the microlenses. For example, this is useful to suppress background from light that is not focused by the microlenses in the PSF of a Fourier light-field microscope. Defaults to False, in which case no blocking of light occurs.	False

Returns:

Type	Description
Field	The Field immediately after the microlens array.

Source code in src/chromatix/functional/lenses.py

def hexagonal_microlens_array(
    field: Field,
    f: ScalarLike,
    n: ScalarLike,
    num_lenses_per_side: int,
    radius: Array,
    separation: ScalarLike,
    block_between: bool = False,
) -> Field:
    """
    Applies a microlens array of hexagonally arranged microlenses placed
    immediately in the plane of the incoming ``Field``.

    !!!warning
        If you have recently used this function prior to it being documented,
        note that the arguments have changed.

    Args:
        field: The ``Field`` to which the lens will be applied.
        f: A scalar value defining the focal length of each lens in units of
            distance.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        num_lenses_per_side: The number of lenses on each outer side of the
            hexagon (e.g. setting this number to 4 will create 37 microlenses).
        radius: A scalar value defining the radius of each microlens in units
            of distance.
        separation: A scalar value defining how far apart the center of each
            microlens is from its neighbors in units of distance.
        block_between: If ``True``, will mask out the ``Field`` in the spaces
            between the microlenses. For example, this is useful to suppress
            background from light that is not focused by the microlenses in the
            PSF of a Fourier light-field microscope. Defaults to ``False``, in
            which case no blocking of light occurs.

    Returns:
        The ``Field`` immediately after the microlens array.
    """
    amplitude, phase = hexagonal_microlens_array_amplitude_and_phase(
        field.spatial_shape,
        field.central_dx,
        field.central_wavelength,
        n,
        f,
        num_lenses_per_side,
        radius,
        separation,
    )
    field = phase_change(field, phase)
    if block_between:
        field = amplitude_change(field, amplitude)
    return field

high_na_ff_lens(field, f, n, NA, output_shape=None, output_dx=None)

Applies a high NA lens placed a distance f after the incoming Field.

Warning

This function assumes that the incoming Field contains only a single wavelength and has a square shape.

Parameters:

Name	Type	Description	Default
field	ScalarField | VectorField	The Field to which the lens will be applied.	required
f	float	Focal length of the lens.	required
n	float	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
NA	float	The NA of the lens.	required
output_shape	tuple[int, int] | None	The shape of the camera (in pixels). If not provided, the output shape will be the same as the shape of the incoming field.	None
output_dx	ScalarLike | None	The pixel pitch of the camera (in units of distance). If not provided, the output spacing will be the same as the spacing of the incoming field.	None

Returns:

Type	Description
ScalarField | VectorField	The Field propagated a distance f to and after the lens.

Source code in src/chromatix/functional/lenses.py

def high_na_ff_lens(
    field: ScalarField | VectorField,
    f: float,
    n: float,
    NA: float,
    output_shape: tuple[int, int] | None = None,
    output_dx: ScalarLike | None = None,
) -> ScalarField | VectorField:
    """
    Applies a high NA lens placed a distance ``f`` after the incoming ``Field``.

    !!!warning
        This function assumes that the incoming ``Field`` contains only a single
        wavelength and has a square shape.

    Args:
        field: The ``Field`` to which the lens will be applied.
        f: Focal length of the lens.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        NA: The NA of the lens.
        output_shape: The shape of the camera (in pixels). If not provided, the
            output shape will be the same as the shape of the incoming field.
        output_dx: The pixel pitch of the camera (in units of distance). If not
            provided, the output spacing will be the same as the spacing of the
            incoming field.

    Returns:
        The ``Field`` propagated a distance ``f`` to and after the lens.
    """
    if not isinstance(field, Vector):
        spherical_u = field.u
    else:
        spherical_u = cartesian_to_spherical(field, n, NA, f)
    if output_dx is None:
        output_dx = field.central_dx
    if output_shape is None:
        output_shape = field.spatial_shape
    # TODO: This only works for single wavelength so far?
    # TODO: What about non-square cases?
    fov_out = output_shape[0] * output_dx
    zoom_factor = 2 * NA * fov_out / ((field.shape[1] - 1) * field.spectrum.wavelength)
    # Correction factors
    s_grid = field.f_grid * field.spectrum.wavelength / n
    sz_sq = 1 - NA**2 * l2_sq_norm(s_grid)
    sz = jnp.sqrt(jnp.maximum(sz_sq, 0.0))
    k = 2 * jnp.pi * n / field.spectrum.wavelength
    defocus = jnp.where(sz != 0.0, jnp.exp(1j * k * sz * f) / sz, 0.0)
    # Create zoomed field
    u = zoomed_fft(
        x=spherical_u * defocus,
        k_start=-zoom_factor * jnp.pi,
        k_end=zoom_factor * jnp.pi,
        output_shape=output_shape,
        include_end=True,
        axes=field.spatial_dims,
    )
    output_dx = output_dx * jnp.ones_like(field.dx)
    return field.replace(u=u, dx=output_dx)

microlens_array(field, fs, n, centers, radii, block_between=False)

Applies a microlens array of arbitrary positioned microlenses placed immediately in the plane of the incoming Field.

Warning

If you have recently used this function prior to it being documented, note that the arguments have changed.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
fs	Float[Array, m]	A 1D array of shape (lenses) defining the focal lengths of each lens in units of distance.	required
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
centers	Float[Array, m]	A 2D array of shape (lenses 2) defining the center position in units of distance (in y x order) for each lens of the microlens array.	required
radii	Float[Array, m]	A 1D array of shape (lenses) defining the radius of each microlens in units of distance.	required

Returns:

Type	Description
Field	The Field immediately after the microlens array.

Source code in src/chromatix/functional/lenses.py

def microlens_array(
    field: Field,
    fs: Float[Array, "m"],
    n: ScalarLike,
    centers: Float[Array, "m"],
    radii: Float[Array, "m"],
    block_between: bool = False,
) -> Field:
    """
    Applies a microlens array of arbitrary positioned microlenses placed
    immediately in the plane of the incoming ``Field``.

    !!!warning
        If you have recently used this function prior to it being documented,
        note that the arguments have changed.

    Args:
        field: The ``Field`` to which the lens will be applied.
        fs: A 1D array of shape ``(lenses)`` defining the focal lengths of each
            lens in units of distance.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        centers: A 2D array of shape ``(lenses 2)`` defining the center position
            in units of distance (in `y x` order) for each lens of the microlens
            array.
        radii: A 1D array of shape ``(lenses)`` defining the radius of each
            microlens in units of distance.

    Returns:
        The ``Field`` immediately after the microlens array.
    """
    amplitude, phase = microlens_array_amplitude_and_phase(
        field.spatial_shape,
        field.central_dx,
        field.central_wavelength,
        n,
        fs,
        centers,
        radii,
    )
    field = phase_change(field, phase)
    if block_between:
        field = amplitude_change(field, amplitude)
    return field

rectangular_microlens_array(field, n, f, num_lenses_height, num_lenses_width, radius, separation, block_between=False)

Applies a microlens array of hexagonally arranged microlenses placed immediately in the plane of the incoming Field.

Warning

If you have recently used this function prior to it being documented, note that the arguments have changed.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
f	Array	A scalar value defining the focal length of each lens in units of distance.	required
num_lenses_height	int	The number of lenses on each vertical side of the rectangle.	required
num_lenses_width	int	The number of lenses on each horizontal side of the rectangle.	required
radius	Array	A scalar value defining the radius of each microlens in units of distance.	required
separation	ScalarLike	A scalar value defining how far apart the center of each microlens is from its neighbors in units of distance.	required
block_between	bool	If True, will mask out the Field in the spaces between the microlenses. For example, this is useful to suppress background from light that is not focused by the microlenses in the PSF of a Fourier light-field microscope. Defaults to False, in which case no blocking of light occurs.	False

Returns:

Type	Description
Field	The Field immediately after the microlens array.

Source code in src/chromatix/functional/lenses.py

def rectangular_microlens_array(
    field: Field,
    n: ScalarLike,
    f: Array,
    num_lenses_height: int,
    num_lenses_width: int,
    radius: Array,
    separation: ScalarLike,
    block_between: bool = False,
) -> Field:
    """
    Applies a microlens array of hexagonally arranged microlenses placed
    immediately in the plane of the incoming ``Field``.

    !!!warning
        If you have recently used this function prior to it being documented,
        note that the arguments have changed.

    Args:
        field: The ``Field`` to which the lens will be applied.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        f: A scalar value defining the focal length of each lens in units of
            distance.
        num_lenses_height: The number of lenses on each vertical side of the
            rectangle.
        num_lenses_width: The number of lenses on each horizontal side of the
            rectangle.
        radius: A scalar value defining the radius of each microlens in units
            of distance.
        separation: A scalar value defining how far apart the center of each
            microlens is from its neighbors in units of distance.
        block_between: If ``True``, will mask out the ``Field`` in the spaces
            between the microlenses. For example, this is useful to suppress
            background from light that is not focused by the microlenses in the
            PSF of a Fourier light-field microscope. Defaults to ``False``, in
            which case no blocking of light occurs.

    Returns:
        The ``Field`` immediately after the microlens array.
    """
    amplitude, phase = rectangular_microlens_array_amplitude_and_phase(
        field.spatial_shape,
        field.central_dx,
        field.central_wavelength,
        n,
        f,
        num_lenses_height,
        num_lenses_width,
        radius,
        separation,
    )
    field = phase_change(field, phase)
    if block_between:
        field = amplitude_change(field, amplitude)
    return field

thick_plano_convex_ff_lens(field, f, radius, center_thickness, n_lens, n_medium=1.0, NA=None, inverse=False, magnification=1.0)

Applies a thick plano-convex lens placed a distance f after the incoming Field. This lens includes propagation by a small distance within the lens (defined by center_thickness).

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
f	ScalarLike	The focal length of the lens in units of distance.	required
radius	ScalarLike	The radius of the spherical part of the plano-convex lens in units of distance.	required
center_thickness	ScalarLike	The maximum thickness of the plano-convex lens (i.e. the distance through the center of the lens) in units of distance.	required
n_lens	ScalarLike	The refractive index of the lens material (e.g. glass).	required
n_medium	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting). Defaults to 1.0 for air.	1.0
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.	None
inverse	bool	Whether the field is passing forwards (plano-convex) or backwards (convex-plano) through the lens. If True, the phase of the lens is conjugated. Defaults to False.	False
magnification	ScalarLike	The magnification to be applied to the propagation through the system. A magnification of greater than 1 will zoom in during the propagation (decrease the spacing of the outgoing Field) and a magnification of smaller than 1 will do the opposite. Defaults to 1.0 for no change to the spacing of the Field.	1.0

Returns:

Type	Description
Field	The Field propagated a distance f after the lens.

Source code in src/chromatix/functional/lenses.py

def thick_plano_convex_ff_lens(
    field: Field,
    f: ScalarLike,
    radius: ScalarLike,
    center_thickness: ScalarLike,
    n_lens: ScalarLike,
    n_medium: ScalarLike = 1.0,
    NA: ScalarLike | None = None,
    inverse: bool = False,
    magnification: ScalarLike = 1.0,
) -> Field:
    """
    Applies a thick plano-convex lens placed a distance ``f`` after the incoming
    ``Field``. This lens includes propagation by a small distance within the
    lens (defined by ``center_thickness``).

    Args:
        field: The ``Field`` to which the lens will be applied.
        f: The focal length of the lens in units of distance.
        radius: The radius of the spherical part of the plano-convex lens in
            units of distance.
        center_thickness: The maximum thickness of the plano-convex lens (i.e.
            the distance through the center of the lens) in units of distance.
        n_lens: The refractive index of the lens material (e.g. glass).
        n_medium: The refractive index of the surrounding medium (assumed to be
            the same incoming and exiting). Defaults to 1.0 for air.
        NA: If provided, the NA of the lens. By default, no pupil is applied
            to the incoming ``Field``.
        inverse: Whether the field is passing forwards (plano-convex) or
            backwards (convex-plano) through the lens. If ``True``, the phase of
            the lens is conjugated. Defaults to ``False``.
        magnification: The magnification to be applied to the propagation
            through the system. A magnification of greater than 1 will zoom
            in during the propagation (decrease the spacing of the outgoing
            ``Field``) and a magnification of smaller than 1 will do the
            opposite. Defaults to 1.0 for no change to the spacing of the
            ``Field``.

    Returns:
        The ``Field`` propagated a distance ``f`` after the lens.
    """
    if NA is not None:
        D = 2 * f * NA / n_medium  # Expression for NA yields width of pupil
        field = circular_pupil(field, D)
    _lens = compute_plano_convex_spherical_lens_abcd(
        f, radius, center_thickness, n_lens, n_medium, inverse
    )
    _free_space = compute_free_space_abcd(f)
    ABCD = _free_space @ _lens @ _free_space
    field = ray_transfer(field, ABCD, n_medium, magnification=magnification)
    return field

thick_plano_convex_lens(field, f, radius, center_thickness, n_lens, n_medium=1.0, NA=None, inverse=False, magnification=1.0)

Applies a thick plano-convex lens placed immediately in the plane of the incoming Field. This lens includes propagation by a small distance within the lens (defined by center_thickness).

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
f	ScalarLike	The focal length of the lens in units of distance.	required
radius	ScalarLike	The radius of the spherical part of the plano-convex lens in units of distance.	required
center_thickness	ScalarLike	The maximum thickness of the plano-convex lens (i.e. the distance through the center of the lens) in units of distance.	required
n_lens	ScalarLike	The refractive index of the lens material (e.g. glass).	required
n_medium	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting). Defaults to 1.0 for air.	1.0
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.	None
inverse	bool	Whether the field is passing forwards (plano-convex) or backwards (convex-plano) through the lens. If True, the phase of the lens is conjugated. Defaults to False.	False
magnification	ScalarLike	The magnification to be applied to the propagation through the system. A magnification of greater than 1 will zoom in during the propagation (decrease the spacing of the outgoing Field) and a magnification of smaller than 1 will do the opposite. Defaults to 1.0 for no change to the spacing of the Field.	1.0

Returns:

Type	Description
Field	The Field immediately after the lens.

Source code in src/chromatix/functional/lenses.py

def thick_plano_convex_lens(
    field: Field,
    f: ScalarLike,
    radius: ScalarLike,
    center_thickness: ScalarLike,
    n_lens: ScalarLike,
    n_medium: ScalarLike = 1.0,
    NA: ScalarLike | None = None,
    inverse: bool = False,
    magnification: ScalarLike = 1.0,
) -> Field:
    """
    Applies a thick plano-convex lens placed immediately in the plane of the
    incoming ``Field``. This lens includes propagation by a small distance
    within the lens (defined by ``center_thickness``).

    Args:
        field: The ``Field`` to which the lens will be applied.
        f: The focal length of the lens in units of distance.
        radius: The radius of the spherical part of the plano-convex lens in
            units of distance.
        center_thickness: The maximum thickness of the plano-convex lens (i.e.
            the distance through the center of the lens) in units of distance.
        n_lens: The refractive index of the lens material (e.g. glass).
        n_medium: The refractive index of the surrounding medium (assumed to be
            the same incoming and exiting). Defaults to 1.0 for air.
        NA: If provided, the NA of the lens. By default, no pupil is applied
            to the incoming ``Field``.
        inverse: Whether the field is passing forwards (plano-convex) or
            backwards (convex-plano) through the lens. If ``True``, the phase of
            the lens is conjugated. Defaults to ``False``.
        magnification: The magnification to be applied to the propagation
            through the system. A magnification of greater than 1 will zoom
            in during the propagation (decrease the spacing of the outgoing
            ``Field``) and a magnification of smaller than 1 will do the
            opposite. Defaults to 1.0 for no change to the spacing of the
            ``Field``.

    Returns:
        The ``Field`` immediately after the lens.
    """
    if NA is not None:
        D = 2 * f * NA / n_medium  # Expression for NA yields width of pupil
        field = circular_pupil(field, D)
    ABCD = compute_plano_convex_spherical_lens_abcd(
        f, radius, center_thickness, n_lens, n_medium, inverse
    )
    field = ray_transfer(field, ABCD, n_medium, magnification=magnification)
    return field

thin_lens(field, f, n, NA=None)

Applies a thin lens placed immediately in the plane of the incoming Field.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the lens will be applied.	required
f	ScalarLike	Focal length of the lens in units of distance.	required
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).	required
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.	None

Returns:

Type	Description
Field	The Field immediately after the lens.

Source code in src/chromatix/functional/lenses.py

def thin_lens(
    field: Field, f: ScalarLike, n: ScalarLike, NA: ScalarLike | None = None
) -> Field:
    """
    Applies a thin lens placed immediately in the plane of the incoming ``Field``.

    Args:
        field: The ``Field`` to which the lens will be applied.
        f: Focal length of the lens in units of distance.
        n: The refractive index of the surrounding medium (assumed to be the
            same incoming and exiting).
        NA: If provided, the NA of the lens. By default, no pupil is applied
            to the incoming ``Field``.

    Returns:
        The ``Field`` immediately after the lens.
    """
    L = jnp.sqrt(field.broadcasted_wavelength * f / n)
    phase = -jnp.pi * l2_sq_norm(field.grid) / L**2

    if NA is not None:
        D = 2 * f * NA / n  # Expression for NA yields width of pupil
        field = circular_pupil(field, D)

    return field * jnp.exp(1j * phase)

Phase Masks

__all__ = ['phase_change', 'interpolated_phase_change', 'wrap_phase', 'spectrally_modulate_phase', 'thin_prism', 'sawtooth_grating', 'sinusoid_grating', 'axicon'] module-attribute

axicon(field, n_axicon, slope_angle, n_medium=1.0)

Applies an axicon placed immediately in the plane of the incoming Field.

The steepness of the axicon can be controlled with slope_angle. To change the direction of the prism, a rotation can be applied.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the axicon will be applied.	required
n_axicon	ScalarLike	The refractive index of the axicon material (e.g. glass).	required
slope_angle	ScalarLike	The angle between the base of the axicon and the base in radians.	required
n_medium	ScalarLike	The refractive index of the surrounding medium. Defaults to 1.0 for air.	1.0

Returns:

Type	Description
Field	The perturbed Field immediately after the axicon.

Source code in src/chromatix/functional/phase_masks.py

def axicon(
    field: Field,
    n_axicon: ScalarLike,
    slope_angle: ScalarLike,
    n_medium: ScalarLike = 1.0,
) -> Field:
    """
    Applies an axicon placed immediately in the plane of the incoming ``Field``.

    The steepness of the axicon can be controlled with ``slope_angle``. To change the
    direction of the prism, a ``rotation`` can be applied.

    Args:
        field: The ``Field`` to which the axicon will be applied.
        n_axicon: The refractive index of the axicon material (e.g. glass).
        slope_angle: The angle between the base of the axicon and the base in
            radians.
        n_medium: The refractive index of the surrounding medium. Defaults to
            1.0 for air.

    Returns:
        The perturbed ``Field`` immediately after the axicon.
    """
    phase = axicon_phase(
        field.spatial_shape,
        field._dx[0, 0],
        field.spectrum[..., 0, 0].squeeze(),
        n_axicon,
        slope_angle,
        n_medium,
    )
    field = phase_change(field, phase)
    return field

interpolated_phase_change(field, phase, phase_range=None, num_bits=None, interpolation_order=0, spectrally_modulate=True)

Perturbs field by phase (in radians), i.e. field * exp(1j * phase), but with an interpolation of the phase mask to the number of samples (pixels) in the incoming field. This is useful when the phase mask itself has a different number of pixels than the field (typically the phase mask has fewer pixels than the finely-sampled field, e.g. when simulating a spatial light modulator). Assumes that the phase mask should be interpolated to the entire extent of the incoming field, i.e. that the extent of the field is exactly that of the phase mask.

Also scales the phase by the ratio of each wavelength to the central wavelength of the spectrum (i.e. the first wavelength in the array of wavelengths) by default. This scaling is necessary to achieve the proper chromatic dispersion effect after propagating a field that is not monochromatic through a phase mask (especially if the phase mask is a fine grating). Returns a new Field with the result of the perturbation.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
phase	Float[Array, 'h w']	The phase mask to apply as a 2D array of phase values (in radians) of shape (height width).	required
phase_range	tuple[float, float] | None	A tuple of two scalar values defining the minimum and maximum phase range values respectively (in radians) to which the provided phase values should be wrapped. For example, this could be (0.0, 2 * jnp.pi). This is useful when simulating devices with some allowed range of phase values or that have a baseline phase offset. Defaults to None, in which case no wrapping of the phase mask is performed.	None
num_bits	int | None	An integer value representing the number of bits to which the phase mask should be quantized, i.e. only 2**num_bits values will be allowed in the phase mask. This quantization occurs after the wrapping of the phase values to the provided phase_range, or if no phase_range is provided then to the full range of values in the provided phase.	None
interpolation_order	int	An integer defining the order of the interpolation of the phase values to the number of pixels of the field. Set to 0 by default for nearest-neighbor interpolation, but can also be set to 1 for bilinear interpolation. No higher order interpolation is supported for now.	0
spectrally_modulate	bool	Whether to perform the per-wavelength scaling of the phase based on the given spectrum of the incoming field. Assumes the phase mask is designed for the central wavelength of the spectrum (assumed to be the first wavelength in the spectrum). Set to True by default.	True

Returns:

Type	Description
Field	The perturbed Field immediately after the phase mask.

Source code in src/chromatix/functional/phase_masks.py

def interpolated_phase_change(
    field: Field,
    phase: Float[Array, "h w"],
    phase_range: tuple[float, float] | None = None,
    num_bits: int | None = None,
    interpolation_order: int = 0,
    spectrally_modulate: bool = True,
) -> Field:
    """
    Perturbs ``field`` by ``phase`` (**in radians**), i.e. ``field * exp(1j
    * phase)``, but with an interpolation of the phase mask to the number of
    samples (pixels) in the incoming field. This is useful when the phase mask
    itself has a different number of pixels than the field (typically the phase
    mask has fewer pixels than the finely-sampled field, e.g. when simulating a
    spatial light modulator). Assumes that the phase mask should be interpolated
    to the entire extent of the incoming field, i.e. that the extent of the
    field is exactly that of the phase mask.

    Also scales the phase by the ratio of each wavelength to the central
    wavelength of the spectrum (i.e. the first wavelength in the array
    of wavelengths) by default. This scaling is necessary to achieve the
    proper chromatic dispersion effect after propagating a field that is not
    monochromatic through a phase mask (especially if the phase mask is a fine
    grating). Returns a new ``Field`` with the result of the perturbation.

    Args:
        field: The complex field to be perturbed.
        phase: The phase mask to apply as a 2D array of phase values (in
            radians) of shape ``(height width)``.
        phase_range: A tuple of two scalar values defining the minimum and
            maximum phase range values respectively (in radians) to which the
            provided phase values should be wrapped. For example, this could
            be `(0.0, 2 * jnp.pi)`. This is useful when simulating devices with
            some allowed range of phase values or that have a baseline phase
            offset. Defaults to ``None``, in which case no wrapping of the phase
            mask is performed.
        num_bits: An integer value representing the number of bits to which the
            phase mask should be quantized, i.e. only `2**num_bits` values will
            be allowed in the phase mask. This quantization occurs after the
            wrapping of the phase values to the provided ``phase_range``, or if
            no ``phase_range`` is provided then to the full range of values in
            the provided ``phase``.
        interpolation_order: An integer defining the order of the interpolation
            of the phase values to the number of pixels of the field. Set to `0`
            by default for nearest-neighbor interpolation, but can also be set
            to 1 for bilinear interpolation. No higher order interpolation is
            supported for now.
        spectrally_modulate: Whether to perform the per-wavelength scaling
            of the phase based on the given spectrum of the incoming field.
            Assumes the phase mask is designed for the central wavelength of the
            spectrum (assumed to be the first wavelength in the spectrum). Set
            to ``True`` by default.

    Returns:
        The perturbed ``Field`` immediately after the phase mask.
    """
    if phase_range is not None:
        phase = wrap_phase(phase, phase_range)
    if num_bits is not None:
        phase = quantize(phase, num_bits, range=phase_range)
    # NOTE(dd/2025-08-14): Assumes the phase should be interpolated across the full
    # extent of the Field
    field_pixel_grid = jnp.meshgrid(
        jnp.linspace(0, phase.shape[0] - 1, num=field.spatial_shape[0]) + 0.5,
        jnp.linspace(0, phase.shape[1] - 1, num=field.spatial_shape[1]) + 0.5,
        indexing="ij",
    )
    phase = map_coordinates(phase, field_pixel_grid, interpolation_order)
    return phase_change(field, phase, spectrally_modulate=spectrally_modulate)

phase_change(field, phase, spectrally_modulate=True)

Perturbs field by phase (in radians), i.e. field * exp(1j * phase).

Also scales the phase by the ratio of each wavelength to the central wavelength of the spectrum (i.e. the first wavelength in the array of wavelengths) by default. This scaling is necessary to achieve the proper chromatic dispersion effect after propagating a field that is not monochromatic through a phase mask (especially if the phase mask is a fine grating). Returns a new Field with the result of the perturbation.

Common phase mask generators can be found in chromatix.utils.initializers. These generators typically assume that the phase mask is placed at the pupil plane of a system.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
phase	Float[Array, 'h w']	The phase mask to apply as a 2D array of phase values (in radians) of shape (height width).	required
spectrally_modulate	bool	Whether to perform the per-wavelength scaling of the phase based on the given spectrum of the incoming field. Assumes the phase mask is designed for the central wavelength of the spectrum (assumed to be the first wavelength in the spectrum). Set to True by default.	True

Returns:

Type	Description
Field	The perturbed Field immediately after the phase mask.

Source code in src/chromatix/functional/phase_masks.py

def phase_change(
    field: Field, phase: Float[Array, "h w"], spectrally_modulate: bool = True
) -> Field:
    """
    Perturbs ``field`` by ``phase`` (**in radians**), i.e. ``field * exp(1j * phase)``.

    Also scales the phase by the ratio of each wavelength to the central
    wavelength of the spectrum (i.e. the first wavelength in the array
    of wavelengths) by default. This scaling is necessary to achieve the
    proper chromatic dispersion effect after propagating a field that is not
    monochromatic through a phase mask (especially if the phase mask is a fine
    grating). Returns a new ``Field`` with the result of the perturbation.

    Common phase mask generators can be found in ``chromatix.utils.initializers``.
    These generators typically assume that the phase mask is placed at the pupil
    plane of a system.

    Args:
        field: The complex field to be perturbed.
        phase: The phase mask to apply as a 2D array of phase values (in
            radians) of shape ``(height width)``.
        spectrally_modulate: Whether to perform the per-wavelength scaling
            of the phase based on the given spectrum of the incoming field.
            Assumes the phase mask is designed for the central wavelength of the
            spectrum (assumed to be the first wavelength in the spectrum). Set
            to ``True`` by default.

    Returns:
        The perturbed ``Field`` immediately after the phase mask.
    """
    if phase.ndim != field.ndim:
        phase = _broadcast_2d_to_spatial(phase, field.spatial_dims)
    assert_equal_shape(
        (phase, field.u),
        dims=field.spatial_dims,
        custom_message="Phase must have same height and width as incoming ``Field``.",
    )
    if spectrally_modulate:
        phase = spectrally_modulate_phase(phase, field.spectrum)
    return field * jnp.exp(1j * phase)

sawtooth_grating(field, n_grating, period, thickness, rotation=0.0, n_medium=1.0)

Applies a sawtooth grating placed immediately in the plane of the incoming Field.

The grating can be varied in thickness which changes the overall scale of the phase values used to calculate the effect of the grating (i.e. the grating phase is calculated as a thin sample). To change the direction of the grating, a rotation can be applied.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the grating will be applied.	required
n_grating	ScalarLike	The refractive index of the grating material (e.g. glass).	required
period	ScalarLike	The period of the sawtooth wave defining the grating in units of distance.	required
thickness	ScalarLike	A scalar defining the thickness of the grating in units of distance.	required
rotation	ScalarLike	How much to rotate the grating in-plane (counter-clockwise in radians).	0.0
n_medium	ScalarLike	The refractive index of the surrounding medium. Defaults to 1.0 for air.	1.0

Returns:

Type	Description
Field	The perturbed Field immediately after the grating.

Source code in src/chromatix/functional/phase_masks.py

def sawtooth_grating(
    field: Field,
    n_grating: ScalarLike,
    period: ScalarLike,
    thickness: ScalarLike,
    rotation: ScalarLike = 0.0,
    n_medium: ScalarLike = 1.0,
) -> Field:
    """
    Applies a sawtooth grating placed immediately in the plane of the incoming
    ``Field``.

    The grating can be varied in thickness which changes the overall scale
    of the phase values used to calculate the effect of the grating (i.e. the
    grating phase is calculated as a thin sample). To change the direction of
    the grating, a ``rotation`` can be applied.

    Args:
        field: The ``Field`` to which the grating will be applied.
        n_grating: The refractive index of the grating material (e.g. glass).
        period: The period of the sawtooth wave defining the grating in units
            of distance.
        thickness: A scalar defining the thickness of the grating in units of
            distance.
        rotation: How much to rotate the grating in-plane (counter-clockwise
            in radians).
        n_medium: The refractive index of the surrounding medium. Defaults to
            1.0 for air.

    Returns:
        The perturbed ``Field`` immediately after the grating.
    """
    phase = sawtooth_phase(
        field.spatial_shape,
        field._dx[0, 0],
        field.spectrum[..., 0, 0].squeeze(),
        n_grating,
        period,
        thickness,
        rotation,
        n_medium,
    )
    field = phase_change(field, phase)
    return field

sinusoid_grating(field, n_grating, period, thickness, rotation=0.0, n_medium=1.0)

Applies a sinusoid grating placed immediately in the plane of the incoming Field.

The grating can be varied in thickness which changes the overall scale of the phase values used to calculate the effect of the grating (i.e. the grating phase is calculated as a thin sample). To change the direction of the grating, a rotation can be applied.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the grating will be applied.	required
n_grating	ScalarLike	The refractive index of the grating material (e.g. glass).	required
period	ScalarLike	The period of the sinusoid wave defining the grating in units of distance.	required
thickness	ScalarLike	A scalar defining the thickness of the grating in units of distance.	required
rotation	ScalarLike	How much to rotate the grating in-plane (counter-clockwise in radians).	0.0
n_medium	ScalarLike	The refractive index of the surrounding medium. Defaults to 1.0 for air.	1.0

Returns:

Type	Description
Field	The perturbed Field immediately after the grating.

Source code in src/chromatix/functional/phase_masks.py

def sinusoid_grating(
    field: Field,
    n_grating: ScalarLike,
    period: ScalarLike,
    thickness: ScalarLike,
    rotation: ScalarLike = 0.0,
    n_medium: ScalarLike = 1.0,
) -> Field:
    """
    Applies a sinusoid grating placed immediately in the plane of the incoming
    ``Field``.

    The grating can be varied in thickness which changes the overall scale
    of the phase values used to calculate the effect of the grating (i.e. the
    grating phase is calculated as a thin sample). To change the direction of
    the grating, a ``rotation`` can be applied.

    Args:
        field: The ``Field`` to which the grating will be applied.
        n_grating: The refractive index of the grating material (e.g. glass).
        period: The period of the sinusoid wave defining the grating in units
            of distance.
        thickness: A scalar defining the thickness of the grating in units of
            distance.
        rotation: How much to rotate the grating in-plane (counter-clockwise
            in radians).
        n_medium: The refractive index of the surrounding medium. Defaults to
            1.0 for air.

    Returns:
        The perturbed ``Field`` immediately after the grating.
    """
    phase = sinusoid_phase(
        field.spatial_shape,
        field._dx[0, 0],
        field.spectrum[..., 0, 0].squeeze(),
        n_grating,
        period,
        thickness,
        rotation,
        n_medium,
    )
    field = phase_change(field, phase)
    return field

spectrally_modulate_phase(phase, spectrum)

Spectrally modulates a given phase for multiple wavelengths.

Scales the phase by the ratio of each wavelength to the central wavelength of the spectrum (i.e. the first wavelength in the array of wavelengths) by default. This scaling is necessary to achieve the proper chromatic dispersion effect after propagating a field that is not monochromatic through a phase mask (especially if the phase mask is a fine grating).

Parameters:

Name	Type	Description	Default
phase	Array	The phase mask to apply as a 2D array of phase values (in radians) of shape (height width).	required
spectrum	Spectrum	The Spectrum used to modulate the phase (contains the wavelengths necessary to scale the phase values). The first wavelength in the spectrum is assumed to be the central wavelength for which the phase mask was designed.	required

Source code in src/chromatix/functional/phase_masks.py

def spectrally_modulate_phase(phase: Array, spectrum: Spectrum) -> Array:
    """
    Spectrally modulates a given ``phase`` for multiple wavelengths.

    Scales the phase by the ratio of each wavelength to the central
    wavelength of the spectrum (i.e. the first wavelength in the array
    of wavelengths) by default. This scaling is necessary to achieve the
    proper chromatic dispersion effect after propagating a field that is not
    monochromatic through a phase mask (especially if the phase mask is a fine
    grating).

    Args:
        phase: The phase mask to apply as a 2D array of phase values (in
            radians) of shape ``(height width)``.
        spectrum: The ``Spectrum`` used to modulate the phase (contains the
            wavelengths necessary to scale the phase values). The first
            wavelength in the spectrum is assumed to be the central wavelength
            for which the phase mask was designed.
    """
    return phase * spectrum.spectral_modulation

thin_prism(field, n_prism, max_thickness, rotation=0.0, n_medium=1.0)

Applies a thin prism placed immediately in the plane of the incoming Field.

The prism is applied as a linear phase ramp that starts at 0 and increases horizontally to the right to a maximum value determined by the specified max_thickness. To change the direction of the prism, a rotation can be applied.

Parameters:

Name	Type	Description	Default
field	Field	The Field to which the prism will be applied.	required
n_prism	ScalarLike	The refractive index of the prism material (e.g. glass).	required
max_thickness	ScalarLike	A scalar defining the maximum thickness of the prism in units of distance.	required
rotation	ScalarLike	How much to rotate the prism in-plane (counter-clockwise in radians).	0.0
n_medium	ScalarLike	The refractive index of the surrounding medium. Defaults to 1.0 for air.	1.0

Returns:

Type	Description
Field	The perturbed Field immediately after the prism.

Source code in src/chromatix/functional/phase_masks.py

def thin_prism(
    field: Field,
    n_prism: ScalarLike,
    max_thickness: ScalarLike,
    rotation: ScalarLike = 0.0,
    n_medium: ScalarLike = 1.0,
) -> Field:
    """
    Applies a thin prism placed immediately in the plane of the incoming ``Field``.

    The prism is applied as a linear phase ramp that starts at 0 and increases
    horizontally to the right to a maximum value determined by the specified
    ``max_thickness``. To change the direction of the prism, a ``rotation`` can
    be applied.

    Args:
        field: The ``Field`` to which the prism will be applied.
        n_prism: The refractive index of the prism material (e.g. glass).
        max_thickness: A scalar defining the maximum thickness of the prism in
            units of distance.
        rotation: How much to rotate the prism in-plane (counter-clockwise in
            radians).
        n_medium: The refractive index of the surrounding medium. Defaults to
            1.0 for air.

    Returns:
        The perturbed ``Field`` immediately after the prism.
    """
    phase = linear_phase(
        field.spatial_shape,
        field._dx[0, 0],
        field.spectrum[..., 0, 0].squeeze(),
        n_prism,
        max_thickness,
        rotation,
        n_medium,
    )
    field = phase_change(field, phase)
    return field

wrap_phase(phase, limits=(-jnp.pi, jnp.pi))

Wraps values of phase to the range given by limits.

Parameters:

Name	Type	Description	Default
phase	Float[Array, 'h w']	The phase mask to wrap (in radians) as a 2D array of shape (height width) (though this function will work on an array of any shape).	required
limits	ArrayLike | tuple[float, float]	A tuple defining the minimum and maximum value that phase will be wrapped to.	(-pi, pi)

Source code in src/chromatix/functional/phase_masks.py

@jax.custom_jvp
def wrap_phase(
    phase: Float[Array, "h w"],
    limits: ArrayLike | tuple[float, float] = (-jnp.pi, jnp.pi),
) -> Array:
    """
    Wraps values of ``phase`` to the range given by ``limits``.

    Args:
        phase: The phase mask to wrap (in radians) as a 2D array of shape
            ``(height width)`` (though this function will work on an array of
            any shape).
        limits: A tuple defining the minimum and maximum value that ``phase``
            will be wrapped to.
    """
    phase_min, phase_max = limits
    phase = jnp.where(
        phase < phase_min,
        phase + (2 * jnp.pi * (1 + (phase_min - phase) // (2 * jnp.pi))),
        phase,
    )
    phase = jnp.where(
        phase > phase_max,
        phase - (2 * jnp.pi * (1 + (phase - phase_max) // (2 * jnp.pi))),
        phase,
    )
    return phase

wrap_phase_jvp(primals, tangents)

Source code in src/chromatix/functional/phase_masks.py

@wrap_phase.defjvp
def wrap_phase_jvp(primals: tuple, tangents: tuple) -> tuple:
    return wrap_phase(*primals), tangents[0]

Amplitude Masks

__all__ = ['amplitude_change', 'interpolated_amplitude_change'] module-attribute

amplitude_change(field, amplitude)

Perturbs field by amplitude, i.e. field * amplitude.

Returns a new Field with the result of the perturbation.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
amplitude	Float[Array, 'h w']	The amplitude mask to apply as a 2D array of amplitude values of shape (height width).	required

Source code in src/chromatix/functional/amplitude_masks.py

def amplitude_change(field: Field, amplitude: Float[Array, "h w"]) -> Field:
    """
    Perturbs ``field`` by ``amplitude``, i.e. ``field * amplitude``.

    Returns a new ``Field`` with the result of the perturbation.

    Args:
        field: The complex field to be perturbed.
        amplitude: The amplitude mask to apply as a 2D array of amplitude values
            of shape ``(height width)``.
    """
    if amplitude.ndim != field.ndim:
        amplitude = _broadcast_2d_to_spatial(amplitude, field.spatial_dims)
    assert_equal_shape(
        (amplitude, field.u),
        dims=field.spatial_dims,
        custom_message=(
            "Amplitude must have same height and width as incoming ``Field``."
        ),
    )
    return field * amplitude.astype(jnp.complex64)

interpolated_amplitude_change(field, amplitude, binary=False, num_bits=None, interpolation_order=0)

Perturbs field by amplitude, i.e. field * amplitude, but with an interpolation of the amplitude mask to the number of samples (pixels) in the incoming field. This is useful when the amplitude mask itself has a different number of pixels than the field (typically the amplitude mask has fewer pixels than the finely-sampled field, e.g. when simulating a digital micromirror device). Assumes that the amplitude mask should be interpolated to the entire extent of the incoming field, i.e. that the extent of the field is exactly that of the amplitude mask.

Returns a new Field with the result of the perturbation.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
amplitude	Float[Array, 'h w']	The amplitude mask to apply as a 2D array of amplitude values of shape (height width).	required
binary	bool	Whether the amplitude values should be binarized (0 or 1) or not. Defaults to False, in which case the amplitude values are not binarized.	False
num_bits	int | None	An integer value representing the number of bits to which the amplitude mask should be quantized, i.e. only 2**num_bits values will be allowed in the amplitude mask. This quantization occurs within the full range of values in the provided amplitude.	None
interpolation_order	int	An integer defining the order of the interpolation of the amplitude values to the number of pixels of the field. Set to 0 by default for nearest-neighbor interpolation, but can also be set to 1 for bilinear interpolation. No higher order interpolation is supported for now.	0

Source code in src/chromatix/functional/amplitude_masks.py

def interpolated_amplitude_change(
    field: Field,
    amplitude: Float[Array, "h w"],
    binary: bool = False,
    num_bits: int | None = None,
    interpolation_order: int = 0,
) -> Field:
    """
    Perturbs ``field`` by ``amplitude``, i.e. ``field * amplitude``, but with
    an interpolation of the amplitude mask to the number of samples (pixels)
    in the incoming field. This is useful when the amplitude mask itself has a
    different number of pixels than the field (typically the amplitude mask has
    fewer pixels than the finely-sampled field, e.g. when simulating a digital
    micromirror device). Assumes that the amplitude mask should be interpolated
    to the entire extent of the incoming field, i.e. that the extent of the
    field is exactly that of the amplitude mask.

    Returns a new ``Field`` with the result of the perturbation.

    Args:
        field: The complex field to be perturbed.
        amplitude: The amplitude mask to apply as a 2D array of amplitude values
            of shape ``(height width)``.
        binary: Whether the amplitude values should be binarized (0 or 1) or
            not. Defaults to ``False``, in which case the amplitude values are
            not binarized.
        num_bits: An integer value representing the number of bits to which the
            amplitude mask should be quantized, i.e. only `2**num_bits` values
            will be allowed in the amplitude mask. This quantization occurs
            within the full range of values in the provided ``amplitude``.
        interpolation_order: An integer defining the order of the interpolation
            of the amplitude values to the number of pixels of the field. Set to
            `0` by default for nearest-neighbor interpolation, but can also be
            set to 1 for bilinear interpolation. No higher order interpolation
            is supported for now.
    """
    if binary:
        amplitude = binarize(amplitude)
    elif num_bits is not None:
        amplitude = quantize(amplitude, num_bits)
    field_pixel_grid = jnp.meshgrid(
        jnp.linspace(0, amplitude.shape[0] - 1, num=field.spatial_shape[0]) + 0.5,
        jnp.linspace(0, amplitude.shape[1] - 1, num=field.spatial_shape[1]) + 0.5,
        indexing="ij",
    )
    amplitude = map_coordinates(amplitude, field_pixel_grid, interpolation_order)
    return amplitude_change(field, amplitude)

Samples

fluorescent_multislice_thick_sample(field, fluorescence_stack, dn_stack, n, thickness_per_slice, pad_width, key, num_samples=1, propagator_forward=None, propagator_backward=None, kykx=(0.0, 0.0), remove_evanescent=False, bandlimit=False)

Perturbs incoming Field as if it went through a thick, transparent, and fluorescent sample, i.e. a sample consisting of some distribution of fluorophores emitting within a clear (phase only) scattering volume (with isotropic refractive index). The thick sample is modeled as being made of many thin slices each of a given thickness. The fluorescence_stack contains the fluorescence intensities of each sample slice. The dn_stack contains the phase delay as change in refractive index from the refractive index of the medium of each sample slice.

This function simulates the incoherent light from fluorophores using a Monte-Carlo approach in which random phases are applied to the fluorescence and the resulting propagations are averaged.

A propagator_forward and a propagator_backward defining the propagation kernels for the field going forward and backward through each slice can be provided. By default, these propagator kernels is calculated inside the function. After passing through all slices, the field is propagated backwards slice by slice to compute the scattered fluorescence intensity.

Returns an Array with the result of the scattered fluorescence intensity volume.

Warning

The underlying propagation method now defaults to the angular spectrum method (ASM) with bandlimit=False and `remove_evanescent=False.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
fluorescence_stack	Float[Array, 'd h w']	The sample fluorescence amplitude for each slice defined as a 3D array of shape (depth height width).	required
dn_stack	Float[Array, 'd h w']	Sample refractive index change as a 3D array of shape(depth height width). Shape must be the same as that for fluorescence_stack.	required
n	ScalarLike	Average refractive index of the sample.	required
thickness_per_slice	ScalarLike	How far to propagate for each slice in units of distance.	required
pad_width	int	An integer defining the pad length for the propagation FFT (NOTE: should not be a jax Array, otherwise a ConcretizationError will arise when traced!). You can use padding calculator utilities from chromatix.functional.propagation to estimate the padding.	required
key	PRNGKey	A PRNGKey used to generate the random phases in each sample.	required
num_samples	int	The number of Monte-Carlo samples (random phases) to simulate.	1
propagator_forward	ArrayLike | None	The propagator kernel for the forward propagation through the sample.	None
propagator_backward	ArrayLike | None	The propagator kernel for the backward propagation through the sample.	None
kykx	ArrayLike | tuple[float, float]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format ky kx.	(0.0, 0.0)
remove_evanescent	bool	If True, removes evanescent waves. Defaults to False.	False
bandlimit	bool	Whether to bandlimit the field before propagation. Defaults to False.	False

Returns:

Type	Description
Array	The scattered fluorescence stack as a 3D array of shape ``(depth height
Array	width)`` after scattering through the provided sample.

Source code in src/chromatix/functional/samples.py

def fluorescent_multislice_thick_sample(
    field: Field,
    fluorescence_stack: Float[Array, "d h w"],
    dn_stack: Float[Array, "d h w"],
    n: ScalarLike,
    thickness_per_slice: ScalarLike,
    pad_width: int,
    key: PRNGKey,
    num_samples: int = 1,
    propagator_forward: ArrayLike | None = None,
    propagator_backward: ArrayLike | None = None,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    remove_evanescent: bool = False,
    bandlimit: bool = False,
) -> Array:
    """
    Perturbs incoming ``Field`` as if it went through a thick, transparent,
    and fluorescent sample, i.e. a sample consisting of some distribution of
    fluorophores emitting within a clear (phase only) scattering volume (with
    isotropic refractive index). The thick sample is modeled as being made
    of many thin slices each of a given thickness. The ``fluorescence_stack``
    contains the fluorescence intensities of each sample slice. The ``dn_stack``
    contains the phase delay as change in refractive index from the refractive
    index of the medium of each sample slice.

    This function simulates the incoherent light from fluorophores using a
    Monte-Carlo approach in which random phases are applied to the fluorescence
    and the resulting propagations are averaged.

    A ``propagator_forward`` and a ``propagator_backward`` defining the
    propagation kernels for the field going forward and backward through each
    slice can be provided. By default, these propagator kernels is calculated
    inside the function. After passing through all slices, the field is
    propagated backwards slice by slice to compute the scattered fluorescence
    intensity.

    Returns an ``Array`` with the result of the scattered fluorescence intensity
    volume.

    !!! warning
        The underlying propagation method now defaults to the angular spectrum
        method (ASM) with ``bandlimit=False`` and ``remove_evanescent=False`.

    Args:
        field: The complex field to be perturbed.
        fluorescence_stack: The sample fluorescence amplitude for each slice
            defined as a 3D array of shape ``(depth height width)``.
        dn_stack: Sample refractive index change as a 3D array of shape``(depth
            height width)``. Shape must be the same as that for
            ``fluorescence_stack``.
        n: Average refractive index of the sample.
        thickness_per_slice: How far to propagate for each slice in units of
            distance.
        pad_width: An integer defining the pad length for the
            propagation FFT (NOTE: should not be a `jax` ``Array``, otherwise
            a ConcretizationError will arise when traced!). You can use padding
            calculator utilities from ``chromatix.functional.propagation`` to
            estimate the padding.
        key: A ``PRNGKey`` used to generate the random phases in each sample.
        num_samples: The number of Monte-Carlo samples (random phases) to simulate.
        propagator_forward: The propagator kernel for the forward propagation through
            the sample.
        propagator_backward: The propagator kernel for the backward propagation
            through the sample.
        kykx: If provided, defines the orientation of the propagation. Should
            be an array of shape `(2,)` in the format ``ky kx``.
        remove_evanescent: If ``True``, removes evanescent waves. Defaults to
            ``False``.
        bandlimit: Whether to bandlimit the field before propagation. Defaults
            to ``False``.

    Returns:
        The scattered fluorescence stack as a 3D array of shape ``(depth height
        width)`` after scattering through the provided sample.
    """
    keys = jax.random.split(key, num=num_samples)
    original_field_shape = field.spatial_shape
    axes = field.spatial_dims
    assert_equal_shape([fluorescence_stack, dn_stack])
    field = pad(field, pad_width)
    dn_stack = center_pad(dn_stack, (0, pad_width, pad_width))
    if propagator_forward is None:
        propagator_forward = compute_asm_propagator(
            field,
            thickness_per_slice,
            n,
            kykx,
            remove_evanescent=remove_evanescent,
            bandlimit=bandlimit,
        )
    if propagator_backward is None:
        propagator_backward = compute_asm_propagator(
            field,
            -thickness_per_slice,
            n,
            kykx,
            remove_evanescent=remove_evanescent,
            bandlimit=bandlimit,
        )

    def _forward(
        i: int, field_and_fluorescence_stack: tuple[Field, Float[Array, "d h w"]]
    ) -> tuple[Field, Float[Array, "d h w"]]:
        (field, fluorescence_stack) = field_and_fluorescence_stack
        fluorescence = _broadcast_2d_to_spatial(
            fluorescence_stack[i], field.spatial_dims
        )
        dn = _broadcast_2d_to_spatial(dn_stack[i], field.spatial_dims)
        field = field * jnp.exp(
            1j * 2 * jnp.pi * dn * thickness_per_slice / field.broadcasted_wavelength
        )
        u = ifft(
            fft(fluorescence + field.u, axes=axes, shift=False) * propagator_forward,
            axes=axes,
            shift=False,
        )
        field = field.replace(u=u)
        return field, fluorescence_stack

    def _backward(
        i: int, field_and_intensity_stack: tuple[Field, Float[Array, "d h w"]]
    ) -> tuple[Field, Float[Array, "d h w"]]:
        (field, intensity_stack) = field_and_intensity_stack
        u = field.u * propagator_backward
        field = field.replace(u=u)
        field_i = field.replace(u=ifft(u, axes=axes, shift=False))
        intensity_stack = intensity_stack.at[intensity_stack.shape[0] - 1 - i].add(
            crop(field_i, pad_width).intensity
        )
        return (field, intensity_stack)

    def _sample(
        i: int, field_and_intensity_stack: tuple[Field, Float[Array, "d h w"]]
    ) -> tuple[Field, Float[Array, "d h w"]]:
        (field, intensity_stack) = field_and_intensity_stack
        random_phase_stack = jax.random.uniform(
            keys[i], fluorescence_stack.shape, minval=0, maxval=2 * jnp.pi
        )
        _fluorescence_stack = center_pad(
            fluorescence_stack * jnp.exp(1j * random_phase_stack),
            (0, pad_width, pad_width),
        )
        (field, _) = jax.lax.fori_loop(
            0, _fluorescence_stack.shape[0], _forward, (field, _fluorescence_stack)
        )
        field = field.replace(u=fft(field.u, axes=axes, shift=False))
        (field, intensity_stack) = jax.lax.fori_loop(
            0, intensity_stack.shape[0], _backward, (field, intensity_stack)
        )
        return (field, intensity_stack)

    intensity_stack = jnp.zeros((fluorescence_stack.shape[0], *original_field_shape))
    (_, intensity_stack) = jax.lax.fori_loop(
        0, num_samples, _sample, (field, intensity_stack)
    )
    intensity_stack /= num_samples
    return intensity_stack

jones_sample(field, absorption, dn, thickness)

Perturbs an incoming vectorial Field as if it went through a thin sample object with a given absorption, anisotropic refractive index change from the refractive index of the medium dn, and thickness. Ignores the incoming field in z direction.

The sample is supposed to follow the thin sample approximation, so the sample perturbation is calculated for each component in the Jones matrix as: exp(1j * 2 * pi * (dn + 1j * absorption) * thickness / lambda) where dn and absorption are allowed to vary per component of the Jones matrix, but thickness is assumed to be the same for each component of the Jones matrix.

Returns a VectorField with the result of the perturbation.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	The complex field to be perturbed.	required
absorption	Float[Array, '2 2 h w']	The sample absorption defined as a 4D array of shape (2 2 height width).	required
dn	Float[Array, '2 2 h w']	Sample refractive index change as a 4D array of shape (2 2 height width).	required
thickness	ScalarLike	Thickness at each sample location as either a scalar or an array with shape or shape broadcastable to (height width).	required

Returns:

Type	Description
VectorField | ChromaticVectorField	Field immediately after propagating through sample.

Source code in src/chromatix/functional/samples.py

def jones_sample(
    field: VectorField | ChromaticVectorField,
    absorption: Float[Array, "2 2 h w"],
    dn: Float[Array, "2 2 h w"],
    thickness: ScalarLike,
) -> VectorField | ChromaticVectorField:
    """
    Perturbs an incoming vectorial ``Field`` as if it went through a thin sample
    object with a given ``absorption``, anisotropic refractive index change from
    the refractive index of the medium ``dn``, and ``thickness``. Ignores the
    incoming field in z direction.

    The sample is supposed to follow the thin sample approximation, so the
    sample perturbation is calculated for each component in the Jones matrix as:
    ``exp(1j * 2 * pi * (dn + 1j * absorption) * thickness / lambda)`` where dn
    and absorption are allowed to vary per component of the Jones matrix, but
    thickness is assumed to be the same for each component of the Jones matrix.

    Returns a ``VectorField`` with the result of the perturbation.

    Args:
        field: The complex field to be perturbed.
        absorption: The sample absorption defined as a 4D array of shape ``(2 2
            height width)``.
        dn: Sample refractive index change as a 4D array of shape ``(2 2 height
            width)``.
        thickness: Thickness at each sample location as either a scalar or an
            array with shape or shape broadcastable to ``(height width)``.

    Returns:
        ``Field`` immediately after propagating through sample.
    """
    assert_rank(
        absorption,
        4,
        custom_message="Absorption must be array of shape ``(2 2 height width)``",
    )
    assert_rank(
        dn,
        4,
        custom_message="Refractive index must be array of shape ``(2 2 height width)``",
    )
    shape_spec = "m m h w -> m m h w " + ("1 " * (abs(field.spatial_dims[0]) - 2))
    absorption = rearrange(absorption, shape_spec, m=2)
    dn = rearrange(dn, shape_spec, m=2)
    # Thickness is the same for four elements in Jones Matrix
    sample = jnp.exp(
        1j
        * 2
        * jnp.pi
        * (dn + 1j * absorption)
        * thickness
        / field.broadcasted_wavelength
    )
    return polarizer(field, sample[0, 0], sample[0, 1], sample[1, 0], sample[1, 1])

multislice_thick_sample(field, absorption_stack, dn_stack, n, thickness_per_slice, pad_width, NA=None, propagator=None, kykx=(0.0, 0.0), reverse_propagate_distance=None, return_stack=False, remove_evanescent=False, bandlimit=False)

Perturbs incoming ScalarField as if it went through a thick sample. The thick sample is modeled as being made of many thin slices each of a given thickness. The absorption_stack and dn_stack contain the absorbance and change in isotropic refractive index from the refractive index of the medium of each sample slice. Expects that the same sample is being applied to all elements across the batch of the incoming Field.

A propagator defining the propagation kernel for the field through each slice can be provided. By default, a propagator is calculated inside the function. After passing through all slices, the field is propagated backwards to the center of the stack, or by the distances specified by reverse_propagate_distance if provided.

Returns a Field with the result of the perturbation.

Warning

The underlying propagation method now defaults to the angular spectrum method (ASM) with bandlimit=False and remove_evanescent=False.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
absorption_stack	Float[Array, 'd h w']	The sample's absorption per voxel for each slice defined as a 3D array of shape (depth height width), where depth is the total number of slices.	required
dn_stack	Float[Array, 'd h w']	The sample's isotropic refractive index change for each slice as a 3D array of shape (depth height width). Shape must be the same that for absorption_stack.	required
n	ScalarLike	Average refractive index of the sample.	required
thickness_per_slice	ScalarLike	How far to propagate for each slice.	required
pad_width	int	An integer defining the pad length for the propagation FFT (NOTE: should not be a jax Array, otherwise a ConcretizationError will arise when traced!). You can use padding calculator utilities from chromatix.functional.propagation to estimate the padding.	required
NA	ScalarLike | None	If provided, will be used to define the numerical aperture (limiting the captured frequencies) of the lens that is imaging the center of the volume. If not provided (the default case), this function will return the scattered field directly which may have undesirable high frequencies.	None
propagator	ArrayLike | None	If provided, will be used as the propagation kernel at each plane of the sample. By default, the propagation kernel is constructed automatically prior to looping through the planes of the sample.	None
kykx	ArrayLike | tuple[float, float]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format ky kx.	(0.0, 0.0)
reverse_propagate_distance	ScalarLike | None	If provided, propagates field at the end backwards by this amount from the top of the stack. By default, field is propagated backwards to the middle of the sample.	None
return_stack	bool	If True, returns the 3D stack of intermediate scattered fields at each plane of the thick sample. This 3D stack is returned as a Field where the innermost batch dimension is the number of planes in the provided dn_stack/ absorption_stack instead of the field defocused to the middle of the sample after scattering through the whole sample. If True, reverse_propagate_distance is ignored. Defaults to False.	False
remove_evanescent	bool	If True, removes evanescent waves. Defaults to False.	False
bandlimit	bool	Whether to bandlimit the field before propagation. Defaults to False.	False

Returns:

Type	Description
Field	If return_stack is False (the default), Field after
Field	propagating through sample and being focused to its center (or to any
Field	plane of the sample defined by reverse_propagate_distance) by an
Field	optical system with the specified numerical aperture. Otherwise, returns
Field	the intermediate scattered fields at each plane of the sample (where
Field	each plane is contained in the first batch dimension of the resulting
Field	Field).

Source code in src/chromatix/functional/samples.py

def multislice_thick_sample(
    field: Field,
    absorption_stack: Float[Array, "d h w"],
    dn_stack: Float[Array, "d h w"],
    n: ScalarLike,
    thickness_per_slice: ScalarLike,
    pad_width: int,
    NA: ScalarLike | None = None,
    propagator: ArrayLike | None = None,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    reverse_propagate_distance: ScalarLike | None = None,
    return_stack: bool = False,
    remove_evanescent: bool = False,
    bandlimit: bool = False,
) -> Field:
    """
    Perturbs incoming ``ScalarField`` as if it went through a thick sample. The
    thick sample is modeled as being made of many thin slices each of a given
    thickness. The ``absorption_stack`` and ``dn_stack`` contain the absorbance
    and change in isotropic refractive index from the refractive index of the
    medium of each sample slice. Expects that the same sample is being applied
    to all elements across the batch of the incoming ``Field``.

    A ``propagator`` defining the propagation kernel for the field through each
    slice can be provided. By default, a ``propagator`` is calculated inside
    the function. After passing through all slices, the field is propagated
    backwards to the center of the stack, or by the distances specified by
    ``reverse_propagate_distance`` if provided.

    Returns a ``Field`` with the result of the perturbation.

    !!! warning
        The underlying propagation method now defaults to the angular spectrum
        method (ASM) with ``bandlimit=False`` and ``remove_evanescent=False``.

    Args:
        field: The complex field to be perturbed.
        absorption_stack: The sample's absorption per voxel for each slice
            defined as a 3D array of shape ``(depth height width)``, where
            ``depth`` is the total number of slices.
        dn_stack: The sample's isotropic refractive index change for each slice
            as a 3D array of shape ``(depth height width)``. Shape must be the
            same that for ``absorption_stack``.
        n: Average refractive index of the sample.
        thickness_per_slice: How far to propagate for each slice.
        pad_width: An integer defining the pad length for the
            propagation FFT (NOTE: should not be a `jax` ``Array``, otherwise
            a ConcretizationError will arise when traced!). You can use padding
            calculator utilities from ``chromatix.functional.propagation`` to
            estimate the padding.
        NA: If provided, will be used to define the numerical aperture (limiting
            the captured frequencies) of the lens that is imaging the center of
            the volume. If not provided (the default case), this function will
            return the scattered field directly which may have undesirable high
            frequencies.
        propagator: If provided, will be used as the propagation kernel at
            each plane of the sample. By default, the propagation kernel is
            constructed automatically prior to looping through the planes of
            the sample.
        kykx: If provided, defines the orientation of the propagation. Should
            be an array of shape `(2,)` in the format ``ky kx``.
        reverse_propagate_distance: If provided, propagates field at the end
            backwards by this amount from the top of the stack. By default,
            field is propagated backwards to the middle of the sample.
        return_stack: If ``True``, returns the 3D stack of intermediate
            scattered fields at each plane of the thick sample. This 3D
            stack is returned as a ``Field`` where the innermost batch
            dimension is the number of planes in the provided ``dn_stack``/
            ``absorption_stack`` instead of the field defocused to the middle of
            the sample after scattering through the whole sample. If ``True``,
            ``reverse_propagate_distance`` is ignored. Defaults to ``False``.
        remove_evanescent: If ``True``, removes evanescent waves. Defaults to
            ``False``.
        bandlimit: Whether to bandlimit the field before propagation. Defaults
            to ``False``.

    Returns:
        If ``return_stack`` is ``False`` (the default), ``Field`` after
        propagating through sample and being focused to its center (or to any
        plane of the sample defined by ``reverse_propagate_distance``) by an
        optical system with the specified numerical aperture. Otherwise, returns
        the intermediate scattered fields at each plane of the sample (where
        each plane is contained in the first batch dimension of the resulting
        ``Field``).
    """
    assert_equal_shape([absorption_stack, dn_stack])
    field = pad(field, pad_width)
    absorption_stack = center_pad(absorption_stack, (0, pad_width, pad_width))
    dn_stack = center_pad(dn_stack, (0, pad_width, pad_width))
    if propagator is None:
        propagator = compute_asm_propagator(
            field,
            thickness_per_slice,
            n,
            kykx,
            remove_evanescent=remove_evanescent,
            bandlimit=bandlimit,
        )

    def _scatter_through_plane(i: int, u: Array) -> Array:
        absorption = absorption_stack[i]
        dn = dn_stack[i]
        field_i = field.replace(u=u)
        field_i = kernel_propagate(
            field_i,
            propagator,
        )
        field_i = thin_sample(field_i, absorption, dn, thickness_per_slice)
        return field_i.u  # pyright: ignore

    def _accumulate_field_at_each_plane(i: int, fields: Array) -> Array:
        fields = fields.at[i].set(
            _scatter_through_plane(i - 1, field.replace(u=fields[i - 1])).u  # pyright: ignore
        )
        return fields

    if return_stack:
        fields = jnp.zeros((dn_stack.shape[0] + 1,) + field.shape, dtype=field.u.dtype)
        fields = fields.at[0].set(field.u)
        fields = jax.lax.fori_loop(
            1, dn_stack.shape[0] + 1, _accumulate_field_at_each_plane, fields
        )
        field = field.replace(u=jnp.concatenate(fields[1:], axis=-5))
        return crop(field, pad_width)
    else:
        field = field.replace(
            u=jax.lax.fori_loop(0, dn_stack.shape[0], _scatter_through_plane, field.u)
        )
        # Propagate field backwards to the middle (or chosen distance) of the stack
        if reverse_propagate_distance is None:
            reverse_propagate_distance = thickness_per_slice * dn_stack.shape[0] / 2
        defocus_propagator = compute_asm_propagator(
            field, -reverse_propagate_distance, n, kykx=kykx
        )
        if NA is not None:
            # NOTE(dd/2024-12-12): @copypaste(ff_lens) Maybe eventually we should
            # just have some functions for creating masks but not applying them
            # to fields. Here we're creating a custom mask for the NA of the lens
            # imaging the desired plane of the scattering sample.
            mask = l2_sq_norm(field.f_grid) <= (
                (NA / field.broadcasted_wavelength) ** 2
            )
            mask = jnp.fft.ifftshift(mask, axes=field.spatial_dims)
            defocus_propagator *= mask
        field = kernel_propagate(field, defocus_propagator)
        return crop(field, pad_width)

polarized_multislice_thick_sample(field, potential_stack, n, thickness_per_slice, NA=1.0)

Implements a thick sample method for samples with anisotroipc refractive index (birefringent samples) per [1].

[1]: Multislice computational model for birefringent scattering, https://doi.org/10.1364/OPTICA.472077.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	Incoming vectorial field to propagate through sample.	required
potential_stack	Float[Array, 'd h w 3 3']	Scattering potential as a 5D array of shape (depth height width 3 3) where the last two axes of shape (3, 3) represent the 3 x 3 tensor representing all possible couplings of input and output polarization components of the sample at each 3D location. Note that along each row or column of these matrices we expect (z y x) order (which matches the (z y x) order of our vectorial Field components).	required
n	ScalarLike	Refractive index of medium.	required
thickness_per_slice	ScalarLike	Thickness of each slice of the sample along the optical axis in units of distance.	required
NA	ScalarLike	Numerical aperture of the imaging system. Defaults to 1.0.	1.0

Returns:

Type	Description
VectorField | ChromaticVectorField	Field just after propagating through sample, imaged through pupil
VectorField | ChromaticVectorField	with the given numerical aperture.

Source code in src/chromatix/functional/samples.py

def polarized_multislice_thick_sample(
    field: VectorField | ChromaticVectorField,
    potential_stack: Float[Array, "d h w 3 3"],
    n: ScalarLike,
    thickness_per_slice: ScalarLike,
    NA: ScalarLike = 1.0,
) -> VectorField | ChromaticVectorField:
    """Implements a thick sample method for samples with anisotroipc refractive
    index (birefringent samples) per [1].

    [1]: Multislice computational model for birefringent scattering,
    https://doi.org/10.1364/OPTICA.472077.

    Args:
        field: Incoming vectorial field to propagate through sample.
        potential_stack: Scattering potential as a 5D array of shape ``(depth
            height width 3 3)`` where the last two axes of shape `(3, 3)`
            represent the `3 x 3` tensor representing all possible couplings
            of input and output polarization components of the sample at each
            3D location. Note that along each row or column of these matrices
            we expect `(z y x)` order (which matches the `(z y x)` order of our
            vectorial ``Field`` components).
        n: Refractive index of medium.
        thickness_per_slice: Thickness of each slice of the sample along the
            optical axis in units of distance.
        NA: Numerical aperture of the imaging system. Defaults to 1.0.

    Returns:
        ``Field`` just after propagating through sample, imaged through pupil
        with the given numerical aperture.
    """

    def Q_op(u: Array) -> Array:
        # correct
        """Polarization transfer operator"""
        return crop_op(
            ifft(matvec(Q, fft(pad_op(u), axes=spatial_dims)), axes=spatial_dims)
        )

    def H_op(u: Array) -> Array:
        # correct
        """Vectorial scattering operator"""
        prefactor = jnp.where(
            kz > 0,
            -1j / 2 * jnp.exp(1j * kz * thickness_per_slice) / kz * thickness_per_slice,
            0,
        )
        return crop_op(
            ifft(
                matvec(Q, prefactor * fft(pad_op(u), axes=spatial_dims)),
                axes=spatial_dims,
            )
        )

    def P_op(u: Array) -> Array:
        """Vectorial free space operator"""
        prefactor = jnp.where(kz > 0, jnp.exp(1j * kz * thickness_per_slice), 0)
        return crop_op(
            ifft(
                matvec(Q, prefactor * fft(pad_op(u), axes=spatial_dims)),
                axes=spatial_dims,
            )
        )

    def propagate_slice(u: Array, potential_slice: Array) -> tuple[Array, None]:
        scatter_field = matvec(potential_slice, Q_op(u))
        new_field = P_op(u) + H_op(scatter_field)
        return new_field, new_field

    def pad_op(u: Array) -> Array:
        return jnp.pad(u, padding)

    def crop_op(u: Array) -> Array:
        u = u[tuple(cropping)]
        return u

    assert isinstance(field, Vector), "Must be a vectorial Field"
    if not isinstance(field, Monochromatic):
        potential_stack = potential_stack[:, :, :, jnp.newaxis, :, :]
    # Padding and cropping for circular convolution
    padded_shape = 2 * np.array(field.spatial_shape)
    pad_width = padded_shape - np.array(field.spatial_shape)
    spatial_dims = [field.ndim + d for d in field.spatial_dims]
    padding = [(0, 0) for d in range(field.ndim)]
    padding[spatial_dims[0]] = (0, pad_width[0])
    padding[spatial_dims[1]] = (0, pad_width[1])
    cropping = [slice(0, field.shape[d]) for d in range(field.ndim)]
    cropping[spatial_dims[0]] = slice(0, field.spatial_shape[0])
    cropping[spatial_dims[1]] = slice(0, field.spatial_shape[1])

    # Getting k_grid
    k_grid = jnp.fft.ifftshift(pad(field, pad_width // 2).k_grid, axes=spatial_dims)
    km = 2 * jnp.pi * n / field.broadcasted_wavelength
    kz = jnp.sqrt(
        jnp.maximum(0.0, km**2 - l2_sq_norm(k_grid))
    )  # chop off evanescent waves
    k_grid = jnp.concatenate([kz[..., jnp.newaxis], k_grid], axis=-1)

    # Getting PTFT and band limiting
    Q = (-outer(k_grid / km, k_grid / km, in_axis=-1) + jnp.eye(3)).squeeze(-3)
    Q = jnp.where(
        jnp.sum(k_grid[..., 1:, None] ** 2, axis=-2) <= NA**2 * km**2, Q, 0
    )  # Add the NA here

    # Running scan over sample
    u, intermediates = jax.lax.scan(propagate_slice, field.u, potential_stack)
    return field.replace(u=u)

thick_polarized_sample(field, potential_stack, n, thickness_per_slice, NA=1.0)

Alias for polarized_multislice_thick_sample. See polarized_multislice_thick_sample for documentation.

Warning

This alias is deprecated and will be removed in a future release. Please switch to the updated name.

Source code in src/chromatix/functional/samples.py

def thick_polarized_sample(
    field: VectorField | ChromaticVectorField,
    potential_stack: Float[Array, "d h w 3 3"],
    n: ScalarLike,
    thickness_per_slice: ScalarLike,
    NA: ScalarLike = 1.0,
) -> VectorField | ChromaticVectorField:
    """
    Alias for ``polarized_multislice_thick_sample``. See
    ``polarized_multislice_thick_sample`` for documentation.

    !!! warning
        This alias is deprecated and will be removed in a future release. Please
        switch to the updated name.
    """
    print(
        "Please switch to the updated function name `polarized_multislice_thick_sample`"
    )
    return polarized_multislice_thick_sample(
        field, potential_stack, n, thickness_per_slice, NA=NA
    )

thin_sample(field, absorption, dn, thickness)

Perturbs an incoming ScalarField as if it went through a thin sample object with a given absorption, isotropic refractive index change from the refractive index of the medium dn, and thickness.

The sample is supposed to follow the thin sample approximation, so the sample perturbation is calculated as: exp(1j * 2 * pi * (dn + 1j * absorption) * thickness / lambda).

Returns a ScalarField with the result of the perturbation.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
absorption	Float[Array, 'h w']	The sample's absorption defined as a 2D array of shape (height width).	required
dn	Float[Array, 'h w']	The sample's isotropic refractive index change as a 2D array of shape (height width).	required
thickness	ScalarLike | Float[Array, 'h w']	Thickness in units of distance as a scalar or array broadcastable to (height width) (thickness at each sample location).	required

Returns: Field immediately after propagating through sample.

Source code in src/chromatix/functional/samples.py

def thin_sample(
    field: Field,
    absorption: Float[Array, "h w"],
    dn: Float[Array, "h w"],
    thickness: ScalarLike | Float[Array, "h w"],
) -> Field:
    """
    Perturbs an incoming ``ScalarField`` as if it went through a thin sample
    object with a given ``absorption``, isotropic refractive index change from
    the refractive index of the medium ``dn``, and ``thickness``.

    The sample is supposed to follow the thin sample approximation, so the
    sample perturbation is calculated as:
    ``exp(1j * 2 * pi * (dn + 1j * absorption) * thickness / lambda)``.

    Returns a ``ScalarField`` with the result of the perturbation.

    Args:
        field: The complex field to be perturbed.
        absorption: The sample's absorption defined as a 2D array of shape
            ``(height width)``.
        dn: The sample's isotropic refractive index change as a 2D array of
            shape ``(height width)``.
        thickness: Thickness in units of distance as a scalar or array
            broadcastable to ``(height width)`` (thickness at each sample
            location).
    Returns:
        ``Field`` immediately after propagating through sample.
    """
    assert_rank(
        absorption,
        2,
        custom_message=f"Absorption must have shape {field.spatial_shape}, got {absorption.shape}",
    )
    assert_rank(
        dn,
        2,
        custom_message=f"Refractive index must have shape {field.spatial_shape}, got {dn.shape}.",
    )
    absorption = _broadcast_2d_to_spatial(absorption, field.spatial_dims)
    dn = _broadcast_2d_to_spatial(dn, field.spatial_dims)
    sample = jnp.exp(
        1j
        * 2
        * jnp.pi
        * (dn + 1j * absorption)
        * thickness
        / field.broadcasted_wavelength
    )
    return field * sample

Sensors

__all__ = ['basic_sensor'] module-attribute

basic_sensor(sensor_input, shot_noise_mode=None, resampler=None, reduce_axis=None, reduce_parallel_axis_name=None, input_spacing=None, noise_key=None)

Produces an intensity image from an incoming Field or intensity Array and simulates shot noise. In most cases, it is better and easier to use the BasicSensor element which will handle some of the state/arguments necessary for this function (especially for resampling to the pixel size of the sensor).

Warning

Assumes that the input has the same spacing at all wavelengths!

Parameters:

Name	Type	Description	Default
sensor_input	Field | Array	Either the incoming Field or an intensity Array to be sampled by the sensor with shot noise.	required
shot_noise_mode	Literal['approximate', 'poisson'] | None	What type of shot noise simulation to use. Can be either "approximate", "poisson", or None. Defaults to None, in which case no shot noise is simulated (the sensor is perfect).	None
resampler	Resampler | None	If provided, will be called to resample the incoming Field to the given shape. These are instances of Resamplers which can be created using chromatix.ops.resample.init_plane_resample.	None
reduce_axis	int | None	If provided, the result will be summed along this dimension. Useful for simulating multiple depth planes being summed at the sensor plane.	None
reduce_parallel_axis_name	str | None	If provided, will do a jax.lax.psum to perform a sum across multiple devices along the pmapped axis with this name. Useful when simulating multiple planes reaching the sensor when those planes are sharded across multiple GPUs/TPUs/etc. using jax.lax.pmap. Note that in that case you will likely want to use both reduce_axis and reduce_parallel_axis_name.	None
input_spacing	ScalarLike | None	Only needs to be provided if sensor_input is an intensity Array and not a Field. If provided, defines the spacing of the input in units of distance to be used for resampling by the sensor.	None
noise_key	PRNGKey | None	If provided, will be used to generate the shot noise. If shot_noise_mode is not None, this must be provided.	None

Returns:

Type	Description
Array	The measured intensity of the incoming Field (potentially summed
Array	across any depth/batch axis of the incoming Field) at the sensor
Array	plane.

Source code in src/chromatix/functional/sensors.py

def basic_sensor(
    sensor_input: Field | Array,
    shot_noise_mode: Literal["approximate", "poisson"] | None = None,
    resampler: Resampler | None = None,
    reduce_axis: int | None = None,
    reduce_parallel_axis_name: str | None = None,
    input_spacing: ScalarLike | None = None,
    noise_key: PRNGKey | None = None,
) -> Array:
    """
    Produces an intensity image from an incoming ``Field`` or intensity ``Array``
    and simulates shot noise. In most cases, it is better and easier to use
    the [``BasicSensor``](elements.md#chromatix.elements.sensors.BasicSensor)
    element which will handle some of the state/arguments necessary for this
    function (especially for resampling to the pixel size of the sensor).

    !!! warning
        Assumes that the input has the same spacing at all wavelengths!

    Args:
        sensor_input: Either the incoming ``Field`` or an intensity ``Array`` to
            be sampled by the sensor with shot noise.
        shot_noise_mode: What type of shot noise simulation to use. Can be
            either ``"approximate"``, ``"poisson"``, or ``None``. Defaults to
            ``None``, in which case no shot noise is simulated (the sensor is
            perfect).
        resampler: If provided, will be called to resample the
            incoming ``Field`` to the given ``shape``. These are
            instances of ``Resampler``s which can be created using
            [``chromatix.ops.resample.init_plane_resample``](ops.md#chromatix.op
            s.resample.init_plane_resample).
        reduce_axis: If provided, the result will be summed along this
            dimension. Useful for simulating multiple depth planes being summed
            at the sensor plane.
        reduce_parallel_axis_name: If provided, will do a ``jax.lax.psum `` to
            perform a sum across multiple devices along the pmapped axis with
            this name. Useful when simulating multiple planes reaching the
            sensor when those planes are sharded across multiple GPUs/TPUs/etc.
            using ``jax.lax.pmap``. Note that in that case you will likely want
            to use both ``reduce_axis`` and ``reduce_parallel_axis_name``.
        input_spacing: Only needs to be provided if ``sensor_input`` is an
            intensity ``Array`` and not a ``Field``. If provided, defines the
            spacing of the input in units of distance to be used for resampling
            by the sensor.
        noise_key: If provided, will be used to generate the shot noise. If
            ``shot_noise_mode`` is not ``None``, this must be provided.

    Returns:
        The measured intensity of the incoming ``Field`` (potentially summed
        across any depth/batch axis of the incoming ``Field``) at the sensor
        plane.
    """
    if isinstance(sensor_input, Field):
        intensity = sensor_input.intensity
        # WARNING(dd): @copypaste(Microscope) Assumes that field has same
        # spacing at all wavelengths when calculating intensity!
        input_spacing = sensor_input.central_rectangular_dx
    else:
        intensity = sensor_input
    if resampler is not None:
        assert input_spacing is not None, (
            "Must provide input_spacing for intensity array"
        )
    if (resampler is not None) and (input_spacing is not None):
        for i in range(intensity.ndim - 2):
            resampler = jax.vmap(resampler, in_axes=(0, None))
        image = resampler(intensity, jnp.atleast_1d(input_spacing))
    else:
        image = intensity
    if reduce_axis is not None:
        image = jnp.sum(image, axis=reduce_axis)
    if reduce_parallel_axis_name is not None:
        image = jax.lax.psum(image, axis_name=reduce_parallel_axis_name)
    if shot_noise_mode is not None:
        assert noise_key is not None, "Must provide a PRNG key to generate noise"
    if shot_noise_mode == "approximate":
        image = approximate_shot_noise(noise_key, image)
    elif shot_noise_mode == "poisson":
        image = shot_noise(noise_key, image)
    return image

Propagation

__all__ = ['transform_propagate', 'compute_sas_precompensation', 'transform_propagate_sas', 'transfer_propagate', 'asm_propagate', 'kernel_propagate', 'compute_transfer_propagator', 'compute_asm_propagator', 'compute_padding_transform', 'compute_padding_transfer', 'compute_padding_exact'] module-attribute

asm_propagate(field, z, n, pad_width, cval=0, absorbing_boundary=None, absorbing_boundary_width=0.65, kykx=(0.0, 0.0), remove_evanescent=False, bandlimit=False, shift_yx=(0.0, 0.0), output_dx=None, output_shape=None, use_czt=True, mode='full')

Propagate field for a distance z using angular spectrum method.

This method does not remove evanescent waves.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
z	ScalarLike | Float[Array, z]	How far to propagate as either a scalar value in units of distance or a 1D array of distances (in which case a batch dimension will be added to the resulting Field).	required
n	ScalarLike	A float that defines the (isotropic) refractive index of the medium.	required
pad_width	int	A keyword argument integer defining the pad length for the propagation FFT. Use padding calculator utilities from chromatix.functional.propagation to compute the padding. Warning The pad value hould not be a Jax array, otherwise a ConcretizationError will arise when traced!	required
cval	float	The background value to use when padding the Field. Defaults to 0 for zero padding.	0
absorbing_boundary	Literal['tukey', 'super_gaussian'] | None	An optional string that determines which absorbing boundary condition is applied (either "tukey" or "super_gaussian", for the Tukey or super Gaussian pupils respectively). Either choice will taper the propagated field to 0 at the edges to reduce aliasing at the edges due to wrapping. Defaults to None in which case no absorbing boundary is applied.	None
absorbing_boundary	Literal['tukey', 'super_gaussian'] | None	A float determining the diameter (as a percentage) of the propagated field that will be permitted without being absorbed. The edges of the field beyond this boundary will taper smoothly to 0 using the chosen boundary function.	None
kykx	ArrayLike | tuple[float, float]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format [ky, kx]. We assume that these are wave vectors, i.e. that they have already been multiplied by 2 * pi / wavelength. Defaults to (0.0, 0.0) for on-axis propagation (no tilting).	(0.0, 0.0)
remove_evanescent	bool	If True, removes evanescent waves. Defaults to False.	False
bandlimit	bool	If True, bandlimited the kernel according to "Band- Limited Angular Spectrum Method for Numerical Simulation of Free- Space Propagation in Far and Near Fields" (2009) by Matsushima and Shimobaba. Defaults to False.	False
shift_yx	ArrayLike | tuple[float, float]	If provided, defines a shift in the destination plane. Should be an array of shape [2,] in the format [y, x].	(0.0, 0.0)
output_dx	ArrayLike | None	If provided, defines a different output sampling at the output plane. Must be of the same shape as that of field.dx. For example, you could construct this as field.dx / zoom_factor where zoom_factor is how much you want to zoom in.	None
output_shape	tuple[int, int]	If provided, defines the output shape of the field. Should be a tuple of integers. If not provided and dx is provided, the output shape will default to that of the input field.	None
use_czt	bool	Whether or not to use chirp Z-transform for different output sampling. Defaults to True if output_dx or output_shape is provided, and to False if neither is provided.	True
mode	Literal['full', 'same']	Either "full" or "same". If "same", the shape of the output Field will match the shape of the incoming Field. Defaults to "full", in which case the output shape will include padding.	'full'

Source code in src/chromatix/functional/propagation.py

def asm_propagate(
    field: Field,
    z: ScalarLike | Float[Array, "z"],
    n: ScalarLike,
    pad_width: int,
    cval: float = 0,
    absorbing_boundary: Literal["tukey", "super_gaussian"] | None = None,
    absorbing_boundary_width: float = 0.65,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    remove_evanescent: bool = False,
    bandlimit: bool = False,
    shift_yx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    output_dx: ArrayLike | None = None,
    output_shape: tuple[int, int] = None,
    use_czt: bool = True,
    mode: Literal["full", "same"] = "full",
) -> Field:
    """
    Propagate ``field`` for a distance ``z`` using angular spectrum method.

    This method does not remove evanescent waves.

    Args:
        field: ``Field`` to be propagated.
        z: How far to propagate as either a scalar value in units of distance
            or a 1D array of distances (in which case a batch dimension will be
            added to the resulting ``Field``).
        n: A float that defines the (isotropic) refractive index of the medium.
        pad_width: A keyword argument integer defining the pad length for
            the propagation FFT. Use padding calculator utilities from
            ``chromatix.functional.propagation`` to compute the padding.
            !!! warning
                The pad value hould not be a Jax array, otherwise a
                ConcretizationError will arise when traced!
        cval: The background value to use when padding the Field. Defaults to 0
            for zero padding.
        absorbing_boundary: An optional string that determines which absorbing
            boundary condition is applied (either "tukey" or "super_gaussian",
            for the Tukey or super Gaussian pupils respectively). Either choice
            will taper the propagated field to 0 at the edges to reduce aliasing
            at the edges due to wrapping. Defaults to None in which case no
            absorbing boundary is applied.
        absorbing_boundary: A float determining the diameter (as a percentage)
            of the propagated field that will be permitted without being
            absorbed. The edges of the field beyond this boundary will taper
            smoothly to 0 using the chosen boundary function.
        kykx: If provided, defines the orientation of the propagation. Should
            be an array of shape `(2,)` in the format `[ky, kx]`. We assume that
            these are wave vectors, i.e. that they have already been multiplied
            by ``2 * pi / wavelength``. Defaults to `(0.0, 0.0)` for on-axis
            propagation (no tilting).
        remove_evanescent: If ``True``, removes evanescent waves. Defaults to
            False.
        bandlimit: If ``True``, bandlimited the kernel according to "Band-
            Limited Angular Spectrum Method for Numerical Simulation of Free-
            Space Propagation in Far and Near Fields" (2009) by Matsushima and
            Shimobaba. Defaults to ``False``.
        shift_yx: If provided, defines a shift in the destination
            plane. Should be an array of shape `[2,]` in the format `[y, x]`.
        output_dx: If provided, defines a different output sampling at the
            output plane. Must be of the same shape as that of ``field.dx``.
            For example, you could construct this as ``field.dx / zoom_factor``
            where ``zoom_factor`` is how much you want to zoom in.
        output_shape: If provided, defines the output shape of the field. Should be
            a tuple of integers. If not provided and ``dx`` is provided, the
            output shape will default to that of the input field.
        use_czt: Whether or not to use chirp Z-transform for different output
            sampling. Defaults to True if `output_dx` or `output_shape` is provided, and
            to False if neither is provided.
        mode: Either "full" or "same". If "same", the shape of the output
            ``Field`` will match the shape of the incoming ``Field``. Defaults
            to "full", in which case the output shape will include padding.
    """
    field = pad(field, pad_width, cval=cval)
    if output_dx is None and output_shape is None:
        # If neither output_dx nor output_shape is provided, use the default ASM propagation
        # as FFT is faster than CZT
        use_czt = False
    propagator = compute_asm_propagator(
        field,
        z,
        n,
        kykx,
        bandlimit,
        shift_yx if not use_czt else (0.0, 0.0),
        remove_evanescent=remove_evanescent,
    )
    field = kernel_propagate(
        field,
        propagator,
        absorbing_boundary=absorbing_boundary,
        absorbing_boundary_width=absorbing_boundary_width,
        output_dx=output_dx,
        output_shape=output_shape,
        shift_yx=shift_yx,
        use_czt=use_czt,
    )
    if mode == "same":
        field = crop(field, pad_width)
    return field

compute_asm_propagator(field, z, n, kykx=(0.0, 0.0), bandlimit=False, shift_yx=(0.0, 0.0), remove_evanescent=False)

Compute propagation kernel for propagation with no Fresnel approximation.

This version of the propagation kernel does not remove evanescent waves, as per the definition of the angular spectrum method. Returns an array that can be multiplied with the Fourier transform of the incoming Field, as performed by kernel_propagate.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
z	ScalarLike | Float[Array, z]	How far to propagate as either a scalar value in units of distance or a 1D array of distances (in which case a batch dimension will be added to the resulting Field).	required
n	ScalarLike	A float that defines the (isotropic) refractive index of the medium.	required
kykx	ArrayLike | tuple[float, float]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format [ky, kx]. We assume that these are wave vectors, i.e. that they have already been multiplied by 2 * pi / wavelength. Defaults to (0.0, 0.0) for on-axis propagation (no tilting).	(0.0, 0.0)
bandlimit	bool	If True, bandlimited the kernel according to "Band- Limited Angular Spectrum Method for Numerical Simulation of Free- Space Propagation in Far and Near Fields" (2009) by Matsushima and Shimobaba. Defaults to False.	False
shift_yx	ArrayLike | tuple[float, float]	If provided, defines a shift in the destination plane. Should be an array of shape [2,] in the format [y, x].	(0.0, 0.0)
remove_evanescent	bool	If True, removes evanescent waves. Defaults to False.	False

Source code in src/chromatix/functional/propagation.py

def compute_asm_propagator(
    field: Field,
    z: ScalarLike | Float[Array, "z"],
    n: ScalarLike,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    bandlimit: bool = False,
    shift_yx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    remove_evanescent: bool = False,
) -> Array:
    """
    Compute propagation kernel for propagation with no Fresnel approximation.

    This version of the propagation kernel does not remove evanescent waves,
    as per the definition of the angular spectrum method. Returns an array
    that can be multiplied with the Fourier transform of the incoming Field, as
    performed by kernel_propagate.

    Args:
        field: ``Field`` to be propagated.
        z: How far to propagate as either a scalar value in units of distance
            or a 1D array of distances (in which case a batch dimension will be
            added to the resulting ``Field``).
        n: A float that defines the (isotropic) refractive index of the medium.
        kykx: If provided, defines the orientation of the propagation. Should
            be an array of shape `(2,)` in the format `[ky, kx]`. We assume that
            these are wave vectors, i.e. that they have already been multiplied
            by ``2 * pi / wavelength``. Defaults to `(0.0, 0.0)` for on-axis
            propagation (no tilting).
        bandlimit: If ``True``, bandlimited the kernel according to "Band-
            Limited Angular Spectrum Method for Numerical Simulation of Free-
            Space Propagation in Far and Near Fields" (2009) by Matsushima and
            Shimobaba. Defaults to ``False``.
        shift_yx: If provided, defines a shift in the destination
            plane. Should be an array of shape `[2,]` in the format `[y, x]`.
        remove_evanescent: If ``True``, removes evanescent waves. Defaults to
            False.
    """
    # kykx = _broadcast_1d_to_grid(kykx, field.ndim)
    z = jnp.asarray(z)
    kykx = jnp.asarray(kykx)
    shift_yx = jnp.asarray(shift_yx)
    if z.size > 1:
        z = _broadcast_1d_to_innermost_batch(z, field.spatial_dims)
    kernel = 1 - (field.broadcasted_wavelength / n) ** 2 * l2_sq_norm(
        field.f_grid - kykx
    )
    if remove_evanescent:
        delay = jnp.sqrt(jnp.maximum(kernel, 0.0))
    else:
        delay = jnp.sqrt(jnp.complex64(kernel))
    # shift in output plane
    # shift_yx = _broadcast_1d_to_grid(shift_yx, field.ndim)
    out_shift = 2 * jnp.pi * jnp.sum(field.f_grid * shift_yx, axis=-1)
    # compute field
    phase = (
        2 * jnp.pi * (jnp.abs(z) * n / field.broadcasted_wavelength) * delay + out_shift
    )
    kernel_field = jnp.where(z >= 0, jnp.exp(1j * phase), jnp.conj(jnp.exp(1j * phase)))
    if bandlimit:
        # Table 1 of "Shifted angular spectrum method for off-axis numerical
        # propagation" (2010) by Matsushima in vectorized form
        # `z` carries trailing singletons for the spatial dims but not for the
        # trailing length-2 yx-vector axis of `field.df` / `shift_yx`. Add it so
        # the cutoffs broadcast against `field.f_grid` for multi-element `z`.
        z_bl = z[..., None] if z.ndim > 0 else z
        k_limit_p = ((shift_yx + 1 / (2 * field.df)) ** (-2) * z_bl**2 + 1) ** (
            -1 / 2
        ) / field.broadcasted_wavelength
        k_limit_n = ((shift_yx - 1 / (2 * field.df)) ** (-2) * z_bl**2 + 1) ** (
            -1 / 2
        ) / field.broadcasted_wavelength
        k0 = (1 / 2) * (
            jnp.sign(shift_yx + field.extent) * k_limit_p
            + jnp.sign(shift_yx - field.extent) * k_limit_n
        )
        k_width = (
            jnp.sign(shift_yx + field.extent) * k_limit_p
            - jnp.sign(shift_yx - field.extent) * k_limit_n
        )
        k_max = k_width / 2
        # obtain rect filter to bandlimit (Eq. 23)
        H_filter_yx = jnp.abs(field.f_grid - k0) <= k_max
        H_filter = H_filter_yx[..., 0] * H_filter_yx[..., 1]
        # apply filter
        kernel_field = kernel_field * H_filter
    return jnp.fft.ifftshift(kernel_field, axes=field.spatial_dims)

compute_padding_exact(height, wavelength, dx, z)

Automatically estimate the padding required for exact/angular wavelength propagation.

Parameters:

Name	Type	Description	Default
height	int	Height (number of pixels in the y direction) of the field. This assumes that the field is a square (height is the same as width).	required
wavelength	float	The wavelength of the field as a scalar in units of distance (assumed to be a monochromatic field).	required
dx	float	The spacing of the samples of the field as a scalar in units of distance. Assumes square pixels.	required
z	float	A float that defines how far to propagate in units of distance.	required

Source code in src/chromatix/functional/propagation.py

def compute_padding_exact(height: int, wavelength: float, dx: float, z: float) -> int:
    """
    Automatically estimate the padding required for exact/angular wavelength propagation.

    Args:
        height: Height (number of pixels in the y direction) of the field. This
            assumes that the field is a square (height is the same as width).
        wavelength: The wavelength of the field as a scalar in units of distance
            (assumed to be a monochromatic field).
        dx: The spacing of the samples of the field as a scalar in units of
            distance. Assumes square pixels.
        z: A float that defines how far to propagate in units of distance.
    """
    # TODO: works only for square fields
    D = height * dx  # height of field in real coordinates
    Nf = np.max((D / 2) ** 2 / (wavelength * z))  # Fresnel number
    M = height  # height of field in pixels
    Q = 2 * np.maximum(1.0, M / (4 * Nf))  # minimum pad ratio * 2
    scale = np.max((wavelength / (2 * dx)))
    # assert scale < 1, "Can't do exact transfer when field.dx < lambda / 2"
    Q = Q / np.sqrt(1 - scale**2)  # minimum pad ratio for exact transfer
    N = (np.ceil((Q * M) / 2) * 2).astype(int)
    pad_width = (N - M).astype(int)
    return pad_width

compute_padding_transfer(height, wavelength, dx, z)

Automatically estimate the padding required for transfer propagation.

Parameters:

Name	Type	Description	Default
height	int	Height (number of pixels in the y direction) of the field. This assumes that the field is a square (height is the same as width).	required
wavelength	float	The wavelength of the field as a scalar in units of distance (assumed to be a monochromatic field).	required
dx	float	The spacing of the samples of the field as a scalar in units of distance. Assumes square pixels.	required
z	float	A float that defines how far to propagate in units of distance.	required

Source code in src/chromatix/functional/propagation.py

def compute_padding_transfer(
    height: int, wavelength: float, dx: float, z: float
) -> int:
    """
    Automatically estimate the padding required for transfer propagation.

    Args:
        height: Height (number of pixels in the y direction) of the field. This
            assumes that the field is a square (height is the same as width).
        wavelength: The wavelength of the field as a scalar in units of distance
            (assumed to be a monochromatic field).
        dx: The spacing of the samples of the field as a scalar in units of
            distance. Assumes square pixels.
        z: A float that defines how far to propagate in units of distance.
    """
    # TODO: works only for square fields
    D = height * dx  # height of field in real coordinates
    Nf = np.max((D / 2) ** 2 / (wavelength * z))  # Fresnel number
    M = height  # height of field in pixels
    Q = 2 * np.maximum(1.0, M / (4 * Nf))  # minimum pad ratio * 2
    N = (np.ceil((Q * M) / 2) * 2).astype(int)
    pad_width = (N - M).astype(int)
    return pad_width

compute_padding_transform(height, wavelength, dx, z)

Automatically estimate the padding required for transform propagation.

Parameters:

Name	Type	Description	Default
height	int	Height (number of pixels in the y direction) of the field. This assumes that the field is a square (height is the same as width).	required
wavelength	float	The wavelength of the field as a scalar in units of distance (assumed to be a monochromatic field).	required
dx	float	The spacing of the samples of the field as a scalar in units of distance. Assumes square pixels.	required
z	float	A float that defines how far to propagate in units of distance.	required

Source code in src/chromatix/functional/propagation.py

def compute_padding_transform(
    height: int, wavelength: float, dx: float, z: float
) -> int:
    """
    Automatically estimate the padding required for transform propagation.

    Args:
        height: Height (number of pixels in the y direction) of the field. This
            assumes that the field is a square (height is the same as width).
        wavelength: The wavelength of the field as a scalar in units of distance
            (assumed to be a monochromatic field).
        dx: The spacing of the samples of the field as a scalar in units of
            distance. Assumes square pixels.
        z: A float that defines how far to propagate in units of distance.
    """
    # TODO: works only for square fields
    D = height * dx  # height of field in real coordinates
    Nf = np.max((D / 2) ** 2 / (wavelength * z))  # Fresnel number
    M = height  # height of field in pixels
    Q = 2 * np.maximum(1.0, M / (4 * Nf))  # minimum pad ratio * 2
    N = (np.ceil((Q * M) / 2) * 2).astype(int)
    pad_width = (N - M).astype(int)
    return pad_width

compute_sas_precompensation(field, z, n)

Source code in src/chromatix/functional/propagation.py

def compute_sas_precompensation(
    field: Field,
    z: ScalarLike,
    n: ScalarLike,
) -> Array:
    kz = 2 * z * jnp.pi * n / field.broadcasted_wavelength
    s = field.broadcasted_wavelength * field.f_grid / n
    s_sq = s**2
    pad_factor = 2
    L = pad_factor * field.extent
    t = L / pad_factor / jnp.abs(z) + jnp.abs(s)
    W = jnp.prod((s_sq * (2 + 1 / t**2) <= 1), axis=-1)
    H_AS = jnp.sqrt(
        jnp.maximum(0, 1 - jnp.sum(s_sq, axis=-1))
    )  # NOTE(rh): Or cast to complex? Can W be larger than the free-space limit?
    H_Fr = 1 - jnp.sum(s_sq, axis=-1) / 2
    delta_H = W * jnp.exp(1j * kz * (H_AS - H_Fr))
    delta_H = jnp.fft.ifftshift(delta_H, axes=field.spatial_dims)
    return delta_H

compute_transfer_propagator(field, z, n, kykx=(0.0, 0.0))

Compute propagation kernel for Fresnel propagation. Returns an array that can be multiplied with the Fourier transform of the incoming Field, as performed by kernel_propagate.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
z	ScalarLike | Float[Array, z]	How far to propagate as either a scalar value in units of distance or a 1D array of distances (in which case a batch dimension will be added to the resulting Field).	required
n	ScalarLike	A float that defines the (isotropic) refractive index of the medium.	required
kykx	ArrayLike | tuple[float, float]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format [ky, kx]. We assume that these are wave vectors, i.e. that they have already been multiplied by 2 * pi / wavelength. Defaults to (0.0, 0.0) for on-axis propagation (no tilting).	(0.0, 0.0)

Source code in src/chromatix/functional/propagation.py

def compute_transfer_propagator(
    field: Field,
    z: ScalarLike | Float[Array, "z"],
    n: ScalarLike,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
) -> Array:
    """
    Compute propagation kernel for Fresnel propagation.
    Returns an array that can be multiplied with the Fourier transform of the
    incoming Field, as performed by kernel_propagate.

    Args:
        field: ``Field`` to be propagated.
        z: How far to propagate as either a scalar value in units of distance
            or a 1D array of distances (in which case a batch dimension will be
            added to the resulting ``Field``).
        n: A float that defines the (isotropic) refractive index of the medium.
        kykx: If provided, defines the orientation of the propagation. Should
            be an array of shape `(2,)` in the format `[ky, kx]`. We assume that
            these are wave vectors, i.e. that they have already been multiplied
            by ``2 * pi / wavelength``. Defaults to `(0.0, 0.0)` for on-axis
            propagation (no tilting).
    """
    z = jnp.asarray(z)
    kykx = jnp.asarray(kykx)
    if z.size > 1:
        z = _broadcast_1d_to_innermost_batch(z, field.spatial_dims)
    phase = (
        -jnp.pi * field.broadcasted_wavelength / n * z * l2_sq_norm(field.f_grid - kykx)
    )
    return jnp.fft.ifftshift(jnp.exp(1j * phase), axes=field.spatial_dims)

kernel_propagate(field, propagator, absorbing_boundary=None, absorbing_boundary_width=0.65, output_dx=None, output_shape=None, shift_yx=(0.0, 0.0), use_czt=False)

Propagate an incoming Field by the given propagation kernel (propagator). This amounts to performing a Fourier convolution of the field and the propagator. Can optionally apply an absorbing boundary (a tapered pupil function) to the field after propagation.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
propagator	ArrayLike	The propagation kernel.	required
absorbing_boundary	Literal['tukey', 'super_gaussian'] | None	An optional string that determines which absorbing boundary condition is applied (either "tukey" or "super_gaussian", for the Tukey or super Gaussian pupils respectively). Either choice will taper the propagated field to 0 at the edges to reduce aliasing at the edges due to wrapping. Defaults to None in which case no absorbing boundary is applied.	None
absorbing_boundary	Literal['tukey', 'super_gaussian'] | None	A float determining the diameter (as a percentage) of the propagated field that will be permitted without being absorbed. The edges of the field beyond this boundary will taper smoothly to 0 using the chosen boundary function.	None

Source code in src/chromatix/functional/propagation.py

def kernel_propagate(
    field: Field,
    propagator: ArrayLike,
    absorbing_boundary: Literal["tukey", "super_gaussian"] | None = None,
    absorbing_boundary_width: float = 0.65,
    output_dx: ArrayLike | None = None,
    output_shape: tuple[int, int] | None = None,
    shift_yx: tuple[float, float] | Float[Array, "2"] = (0.0, 0.0),
    use_czt: bool = False,
) -> Field:
    """
    Propagate an incoming ``Field`` by the given propagation kernel
    (``propagator``). This amounts to performing a Fourier convolution of the
    ``field`` and the ``propagator``. Can optionally apply an absorbing boundary
    (a tapered pupil function) to the field after propagation.

    Args:
        field: ``Field`` to be propagated.
        propagator: The propagation kernel.
        absorbing_boundary: An optional string that determines which absorbing
            boundary condition is applied (either "tukey" or "super_gaussian",
            for the Tukey or super Gaussian pupils respectively). Either choice
            will taper the propagated field to 0 at the edges to reduce aliasing
            at the edges due to wrapping. Defaults to None in which case no
            absorbing boundary is applied.
        absorbing_boundary: A float determining the diameter (as a percentage)
            of the propagated field that will be permitted without being
            absorbed. The edges of the field beyond this boundary will taper
            smoothly to 0 using the chosen boundary function.
    """
    _boundaries = {"tukey": tukey_pupil, "super_gaussian": super_gaussian_pupil}
    assert absorbing_boundary is None or absorbing_boundary in _boundaries, (
        f"The absorbing_boundary must be None or in {_boundaries.keys()}."
    )
    axes = field.spatial_dims
    shift_yx = jnp.asarray(shift_yx)
    if output_dx is None and output_shape is None and not use_czt:
        # shifting accounted for in `propagator`
        u = jnp.fft.ifft2(jnp.fft.fft2(field.u, axes=axes) * propagator, axes=axes)
        field = shift_grid(field, shift_yx)
    else:
        if output_shape is None:
            output_shape = field.spatial_shape
        if output_dx is None:
            output_dx = field.dx
        in_field = field.u
        in_field_df = field.df
        in_field_f_grid = field.f_grid
        in_field_extent = field.extent.squeeze()
        field = shift_grid(field, shift_yx)
        field = Field.empty_like(field, dx=output_dx, shape=output_shape)
        # Scaling factor in Eq 7 of "Band-limited angular spectrum numerical
        # propagation method with selective scaling of observation window size
        # and sample number"
        alpha = field.dx / in_field_df
        # Output field in k-space
        u = jnp.fft.fftshift(
            propagator
            * jnp.fft.fft2(jnp.fft.ifftshift(in_field, axes=axes), axes=axes),
            axes=axes,
        )
        if use_czt:
            spatial_limits = field.spatial_limits
            y_min = spatial_limits[0, 0]
            y_max = spatial_limits[0, 1]
            x_min = spatial_limits[1, 0]
            x_max = spatial_limits[1, 1]
            limits_min = [y_min, x_min]
            limits_max = [y_max, x_max]
            T = in_field_extent
            for d in range(len(axes)):
                # -- chirp z-transform
                m = output_shape[d]
                a = jnp.exp(-1j * 2 * jnp.pi / T[d] * limits_min[d])
                w = jnp.exp(
                    1j * (2 * jnp.pi / T[d]) * (limits_max[d] - limits_min[d]) / (m - 1)
                )
                u = czt(x=u, m=m, a=a, w=w, axis=axes[d])
                # -- modulate
                N = (m - 1) // 2
                u = jnp.moveaxis(u, axes[d], -1)
                C = w ** (-N * jnp.arange(m)) * (a**N)
                u *= C  # applied to last dimension
                u = jnp.moveaxis(u, -1, axes[d])

            u *= jnp.prod(1 / alpha)
        else:
            # Eq 9 of "Band-limited angular spectrum numerical propagation method
            # with selective scaling of observation window size and sample number"
            # (2012)
            wn = alpha * in_field_f_grid
            f = jnp.prod(jnp.exp(-1j * jnp.pi / alpha * wn**2), axis=-1)
            B = u * jnp.prod(
                (1 / alpha) * jnp.exp(1j * jnp.pi / alpha * wn**2), axis=-1
            )
            mod_terms = jnp.prod(
                field.dx * jnp.exp(1j * jnp.pi / alpha * field.grid**2),
                axis=-1,
            )
            u = mod_terms * fftconvolve(B, f, mode="same", axes=axes)
    field = field.replace(u=u)
    if absorbing_boundary is not None:
        pupil = _boundaries[absorbing_boundary]
        absorbing_boundary_width *= field.extent[1]
        field = pupil(field, absorbing_boundary_width)
    return field

transfer_propagate(field, z, n, pad_width, cval=0, absorbing_boundary=None, absorbing_boundary_width=0.65, kykx=(0.0, 0.0), shift_yx=(0.0, 0.0), output_dx=None, output_shape=None, use_czt=True, mode='full')

Fresnel propagate field for a distance z using transfer method. This method is also called the convolutional Fresnel propagation (CV-FR) method.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
z	ScalarLike | Float[Array, z]	How far to propagate as either a scalar value in units of distance or a 1D array of distances (in which case a batch dimension will be added to the resulting Field).	required
n	ScalarLike	A float that defines the (isotropic) refractive index of the medium.	required
pad_width	int	A keyword argument integer defining the pad length for the propagation FFT. Use padding calculator utilities from chromatix.functional.propagation to compute the padding. Warning The pad value hould not be a Jax array, otherwise a ConcretizationError will arise when traced!	required
cval	float	The background value to use when padding the Field. Defaults to 0 for zero padding.	0
absorbing_boundary	Literal['tukey', 'super_gaussian'] | None	An optional string that determines which absorbing boundary condition is applied (either "tukey" or "super_gaussian", for the Tukey or super Gaussian pupils respectively). Either choice will taper the propagated field to 0 at the edges to reduce aliasing at the edges due to wrapping. Defaults to None in which case no absorbing boundary is applied.	None
absorbing_boundary	Literal['tukey', 'super_gaussian'] | None	A float determining the diameter (as a percentage) of the propagated field that will be permitted without being absorbed. The edges of the field beyond this boundary will taper smoothly to 0 using the chosen boundary function.	None
kykx	ArrayLike | tuple[float, float]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format [ky, kx]. We assume that these are wave vectors, i.e. that they have already been multiplied by 2 * pi / wavelength. Defaults to (0.0, 0.0) for on-axis propagation (no tilting).	(0.0, 0.0)
shift_yx	ArrayLike | tuple[float, float]	If provided, defines a shift in the destination plane. Should be an array of shape [2,] in the format [y, x].	(0.0, 0.0)
output_dx	ArrayLike | None	If provided, defines a different output sampling at the output plane.	None
output_shape	tuple[int, int] | None	If provided, defines the output shape of the field. Should be a tuple of integers. If not provided and dx is provided, the output shape will default to that of the input field.	None
use_czt	bool	Whether or not to use chirp Z-transform for different output sampling. Defaults to True if output_dx or output_shape is provided, and to False if neither is provided.	True
mode	Literal['full', 'same']	Either "full" or "same". If "same", the shape of the output Field will match the shape of the incoming Field. Defaults to "full", in which case the output shape will include padding.	'full'

Source code in src/chromatix/functional/propagation.py

def transfer_propagate(
    field: Field,
    z: ScalarLike | Float[Array, "z"],
    n: ScalarLike,
    pad_width: int,
    cval: float = 0,
    absorbing_boundary: Literal["tukey", "super_gaussian"] | None = None,
    absorbing_boundary_width: float = 0.65,
    kykx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    shift_yx: ArrayLike | tuple[float, float] = (0.0, 0.0),
    output_dx: ArrayLike | None = None,
    output_shape: tuple[int, int] | None = None,
    use_czt: bool = True,
    mode: Literal["full", "same"] = "full",
) -> Field:
    """
    Fresnel propagate ``field`` for a distance ``z`` using transfer method. This
    method is also called the convolutional Fresnel propagation (CV-FR) method.

    Args:
        field: ``Field`` to be propagated.
        z: How far to propagate as either a scalar value in units of distance
            or a 1D array of distances (in which case a batch dimension will be
            added to the resulting ``Field``).
        n: A float that defines the (isotropic) refractive index of the medium.
        pad_width: A keyword argument integer defining the pad length for
            the propagation FFT. Use padding calculator utilities from
            ``chromatix.functional.propagation`` to compute the padding.
            !!! warning
                The pad value hould not be a Jax array, otherwise a
                ConcretizationError will arise when traced!
        cval: The background value to use when padding the Field. Defaults to 0
            for zero padding.
        absorbing_boundary: An optional string that determines which absorbing
            boundary condition is applied (either "tukey" or "super_gaussian",
            for the Tukey or super Gaussian pupils respectively). Either choice
            will taper the propagated field to 0 at the edges to reduce aliasing
            at the edges due to wrapping. Defaults to None in which case no
            absorbing boundary is applied.
        absorbing_boundary: A float determining the diameter (as a percentage)
            of the propagated field that will be permitted without being
            absorbed. The edges of the field beyond this boundary will taper
            smoothly to 0 using the chosen boundary function.
        kykx: If provided, defines the orientation of the propagation. Should
            be an array of shape `(2,)` in the format `[ky, kx]`. We assume that
            these are wave vectors, i.e. that they have already been multiplied
            by ``2 * pi / wavelength``. Defaults to `(0.0, 0.0)` for on-axis
            propagation (no tilting).
        shift_yx: If provided, defines a shift in the destination
            plane. Should be an array of shape `[2,]` in the format `[y, x]`.
        output_dx: If provided, defines a different output sampling at the output
            plane.
        output_shape: If provided, defines the output shape of the field. Should be
            a tuple of integers. If not provided and ``dx`` is provided, the
            output shape will default to that of the input field.
        use_czt: Whether or not to use chirp Z-transform for different output
            sampling. Defaults to True if `output_dx` or `output_shape` is provided, and
            to False if neither is provided.
        mode: Either "full" or "same". If "same", the shape of the output
            ``Field`` will match the shape of the incoming ``Field``. Defaults
            to "full", in which case the output shape will include padding.
    """
    field = pad(field, pad_width, cval=cval)
    if output_dx is None and output_shape is None:
        # If neither output_dx nor output_shape is provided, use the default ASM propagation
        # as FFT is faster than CZT
        use_czt = False
    z = jnp.atleast_1d(z)
    # assert field.num_batch_dims == 0 or field.batch_dims[-1] == z.size, (
    #     "Must have no batch dimensions or innermost batch dimension must have size z"
    # )
    propagator = compute_transfer_propagator(field, z, n, kykx)
    field = kernel_propagate(
        field,
        propagator,
        absorbing_boundary=absorbing_boundary,
        absorbing_boundary_width=absorbing_boundary_width,
        output_dx=output_dx,
        output_shape=output_shape,
        shift_yx=shift_yx,
        use_czt=use_czt,
    )
    if mode == "same":
        field = crop(field, pad_width)
    return field

transform_propagate(field, z, n, pad_width, cval=0, skip_initial_phase=False, skip_final_phase=False)

Fresnel propagate field for a distance z using transform method. This method is also called the single-FFT (SFT-FR) Fresnel propagation method. Note that this method changes the sampling of the resulting field. If the distance is negative, the field is propagated back to the source inverting essentially performing an inverse.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
z	ScalarLike	How far to propagate as a scalar value in units of distance.	required
n	ScalarLike	A float that defines the (isotropic) refractive index of the medium.	required
pad_width	int | tuple[int, int]	A keyword argument integer defining the pad length for the propagation FFT. Use padding calculator utilities from chromatix.functional.propagation to compute the padding. Warning The pad value hould not be a Jax array, otherwise a ConcretizationError will arise when traced!	required
cval	float	The background value to use when padding the Field. Defaults to 0 for zero padding.	0
skip_initial_phase	bool	Whether to skip the input phase change (before Fourier transforming). Defaults to False, in which case the input phase change is not skipped.	False
skip_final_phase	bool	Whether to skip the output phase change (after Fourier transforming). Defaults to False, in which case the output phase change is not skipped.	False

Source code in src/chromatix/functional/propagation.py

def transform_propagate(
    field: Field,
    z: ScalarLike,
    n: ScalarLike,
    pad_width: int | tuple[int, int],
    cval: float = 0,
    skip_initial_phase: bool = False,
    skip_final_phase: bool = False,
) -> Field:
    """
    Fresnel propagate ``field`` for a distance ``z`` using transform method.
    This method is also called the single-FFT (SFT-FR) Fresnel propagation
    method. Note that this method changes the sampling of the resulting field.
    If the distance is negative, the field is propagated back to the source
    inverting essentially performing an inverse.

    Args:
        field: ``Field`` to be propagated.
        z: How far to propagate as a scalar value in units of distance.
        n: A float that defines the (isotropic) refractive index of the medium.
        pad_width: A keyword argument integer defining the pad length for
            the propagation FFT. Use padding calculator utilities from
            ``chromatix.functional.propagation`` to compute the padding.
            !!! warning
                The pad value hould not be a Jax array, otherwise a
                ConcretizationError will arise when traced!
        cval: The background value to use when padding the Field. Defaults to 0
            for zero padding.
        skip_initial_phase: Whether to skip the input phase change (before
            Fourier transforming). Defaults to False, in which case the input
            phase change is not skipped.
        skip_final_phase: Whether to skip the output phase change (after Fourier
            transforming). Defaults to False, in which case the output phase
            change is not skipped.
    """
    field = pad(field, pad_width, cval=cval)
    # Fourier normalization factor
    L_sq = field.broadcasted_wavelength * z / n
    # New field is optical_fft minus -1j factor
    if not skip_initial_phase:
        # Calculating input phase change (defining Q1)
        input_phase = (jnp.pi / L_sq) * l2_sq_norm(field.grid)
        field = field * jnp.exp(1j * input_phase)
    field = 1j * optical_fft(field, z, n)
    # Calculating output phase change (defining Q2)
    if not skip_final_phase:
        output_phase = (jnp.pi / L_sq) * l2_sq_norm(field.grid)
        field = field * jnp.exp(1j * output_phase)
    return crop(field, pad_width)

transform_propagate_sas(field, z, n, cval=0, skip_initial_phase=False, skip_final_phase=False)

Propagate field for a distance z using the scalable angular spectrum (SAS) method. See https://doi.org/10.1364/OPTICA.497809 It changes the pixelsize like the transform method, but it is more accurate because it precompensates the phase error. Since it uses three FFTS, it is slower than the transform method. Note that the field is automatically padded by a factor of 2, so the pixelsize is halved.

Note also that a negative propagation distance causes the code to apply the inverse propagation, i.e. propagating by a positive z and then a negative z would propagate you back to the original field. In the negative z case the order of single step Fresnel propagation and precompensation is reversed.

Parameters:

Name	Type	Description	Default
field	Field	Field to be propagated.	required
z	ScalarLike	How far to propagate as a scalar value in units of distance.	required
n	ScalarLike	A float that defines the (isotropic) refractive index of the medium.	required
cval	float	The background value to use when padding the Field. Defaults to 0 for zero padding.	0
skip_initial_phase	bool	Whether to skip the input phase change (before Fourier transforming). Defaults to False, in which case the input phase change is not skipped.	False
skip_final_phase	bool	Whether to skip the output phase change (after Fourier transforming). Defaults to False, in which case the output phase change is not skipped.	False

Source code in src/chromatix/functional/propagation.py

def transform_propagate_sas(
    field: Field,
    z: ScalarLike,
    n: ScalarLike,
    cval: float = 0,
    skip_initial_phase: bool = False,
    skip_final_phase: bool = False,
) -> Field:
    """
    Propagate ``field`` for a distance ``z`` using the scalable angular spectrum
    (SAS) method. See https://doi.org/10.1364/OPTICA.497809 It changes the
    pixelsize like the transform method, but it is more accurate because it
    precompensates the phase error. Since it uses three FFTS, it is slower
    than the transform method. Note that the field is automatically padded by a
    factor of 2, so the pixelsize is halved.

    Note also that a negative propagation distance causes the code to apply
    the inverse propagation, i.e. propagating by a positive ``z`` and then
    a negative ``z`` would propagate you back to the original ``field``. In
    the negative ``z`` case the order of single step Fresnel propagation and
    precompensation is reversed.

    Args:
        field: ``Field`` to be propagated.
        z: How far to propagate as a scalar value in units of distance.
        n: A float that defines the (isotropic) refractive index of the medium.
        cval: The background value to use when padding the Field. Defaults to 0
            for zero padding.
        skip_initial_phase: Whether to skip the input phase change (before
            Fourier transforming). Defaults to False, in which case the input
            phase change is not skipped.
        skip_final_phase: Whether to skip the output phase change (after Fourier
            transforming). Defaults to False, in which case the output phase
            change is not skipped.
    """
    # Don't change this pad_factor, only 2 is supported
    pad_factor = 2
    sz = np.array(field.spatial_shape)
    pad_width = tuple(sz // pad_factor)
    field = pad(field, pad_width, cval=cval)

    def _forward(field: Field, z: ScalarLike) -> tuple[Array, Array]:
        delta_H = compute_sas_precompensation(field, z, n)
        field = kernel_propagate(field, delta_H)
        field = transform_propagate(
            field, z, n, 0, 0, skip_initial_phase, skip_final_phase
        )
        return field.u, field.dx

    def _inverse(field: Field, z: ScalarLike) -> tuple[Array, Array]:
        field = transform_propagate(
            field, z, n, 0, 0, skip_initial_phase, skip_final_phase
        )
        delta_H = compute_sas_precompensation(field, z, n)
        field = kernel_propagate(field, delta_H)
        return field.u, field.dx

    u, dx = jax.lax.cond(z >= 0, _forward, _inverse, field, z)
    field = field.replace(u=u, dx=dx)
    return crop(field, pad_width)

Polarization

ComplexScalarLike = Complex module-attribute

__all__ = ['jones_vector', 'polarizer', 'phase_retarder', 'linear', 'left_circular', 'right_circular', 'linear_polarizer', 'left_circular_polarizer', 'right_circular_polarizer', 'universal_compensator', 'wave_plate', 'halfwave_plate', 'quarterwave_plate'] module-attribute

halfwave_plate(field, theta)

Applies a halfwave plate with angle theta to the incoming field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	incoming field.	required
theta	float	angle w.r.t. horizontal.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: outgoing field.

Source code in src/chromatix/functional/polarizers.py

def halfwave_plate(
    field: VectorField | ChromaticVectorField, theta: ScalarLike
) -> VectorField | ChromaticVectorField:
    """Applies a halfwave plate with angle theta to the incoming field.

    Args:
        field (VectorField | ChromaticVectorField): incoming field.
        theta (float): angle w.r.t. horizontal.

    Returns:
        VectorField | ChromaticVectorField: outgoing field.
    """
    return phase_retarder(field, theta, eta=jnp.pi, phi=0)

jones_vector(theta, beta)

Generates a Jones vector given by [cos(theta), sin(theta)exp(1j*beta)].

Parameters:

Name	Type	Description	Default
theta	float	Polarization angle.	required
beta	float	Relative delay between components.	required

Returns:

Name	Type	Description
Array	Array	Jones vector.

Source code in src/chromatix/functional/polarizers.py

def jones_vector(theta: ScalarLike, beta: ScalarLike) -> Array:
    """Generates a Jones vector given by [cos(theta), sin(theta)exp(1j*beta)].

    Args:
        theta (float): Polarization angle.
        beta (float): Relative delay between components.

    Returns:
        Array: Jones vector.
    """

    # Generates a Jones vector with a given beta = alpha_y - alpha_x.
    # Assumes alpha_x=0.
    return jnp.array(
        [0, jnp.sin(theta) * jnp.exp(1j * beta), jnp.cos(theta)], dtype=jnp.complex64
    )

left_circular()

Generates a Jones vector for left circularly polarized light.

Returns:

Name	Type	Description
Array	Array	Left circularly polarized Jones vector.

Source code in src/chromatix/functional/polarizers.py

def left_circular() -> Array:
    """Generates a Jones vector for left circularly polarized
    light.

    Returns:
        Array: Left circularly polarized Jones vector.
    """
    return jones_vector(jnp.pi / 4, jnp.pi / 2)

left_circular_polarizer(field)

Applies a left circular polarizer to the incoming field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	incoming field.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: outgoing field.

Source code in src/chromatix/functional/polarizers.py

def left_circular_polarizer(
    field: VectorField | ChromaticVectorField,
) -> VectorField | ChromaticVectorField:
    """Applies a left circular polarizer to the incoming field.

    Args:
        field (VectorField | ChromaticVectorField): incoming field.

    Returns:
        VectorField | ChromaticVectorField: outgoing field.
    """
    J00 = 1
    J11 = 1
    J01 = -1j
    J10 = 1j
    return polarizer(field, J00, J01, J10, J11)

linear(theta)

Generates a Jones vector for linearly polarized light with an angle $ heta$ w.r.t. to the horizontal.

Parameters:

Name	Type	Description	Default
theta	float	Angle w.r.t horizontal.	required

Returns:

Name	Type	Description
Array	Array	Linearly polarized Jones vector.

Source code in src/chromatix/functional/polarizers.py

def linear(theta: ScalarLike) -> Array:
    """Generates a Jones vector for linearly polarized
    light with an angle $\theta$ w.r.t. to the horizontal.

    Args:
        theta (float): Angle w.r.t horizontal.

    Returns:
        Array: Linearly polarized Jones vector.
    """
    return jones_vector(theta, 0)

linear_polarizer(field, angle)

Applies a linear polarizer with a given angle to the incoming field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	incoming field.	required
angle	float	angle w.r.t to the horizontal.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: outgoing field.

Source code in src/chromatix/functional/polarizers.py

def linear_polarizer(
    field: VectorField | ChromaticVectorField, angle: ScalarLike
) -> VectorField | ChromaticVectorField:
    """Applies a linear polarizer with a given angle to the incoming field.

    Args:
        field (VectorField | ChromaticVectorField): incoming field.
        angle (float): angle w.r.t to the horizontal.

    Returns:
        VectorField | ChromaticVectorField: outgoing field.
    """

    c, s = jnp.cos(angle), jnp.sin(angle)
    J00 = c**2
    J11 = s**2
    J01 = s * c
    J10 = J01
    return polarizer(field, J00, J01, J10, J11)

phase_retarder(field, theta, eta, phi)

Applies a general purpose retardation matrix with angle w.r.t horizontal theta, relative phase change eta and circularity phi.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	incoming field.	required
theta	float	angle w.r.t horizonal axis.	required
eta	float	relative phase retardation.	required
phi	float	circularity.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: outgoing field.

Source code in src/chromatix/functional/polarizers.py

def phase_retarder(
    field: VectorField | ChromaticVectorField,
    theta: ScalarLike,
    eta: ScalarLike,
    phi: ScalarLike,
) -> VectorField | ChromaticVectorField:
    """Applies a general purpose retardation matrix with angle w.r.t horizontal theta,
    relative phase change eta and circularity phi.

    Args:
        field (VectorField | ChromaticVectorField): incoming field.
        theta (float): angle w.r.t horizonal axis.
        eta (float): relative phase retardation.
        phi (float): circularity.

    Returns:
        VectorField | ChromaticVectorField: outgoing field.
    """
    s, c = jnp.sin(theta), jnp.cos(theta)
    scale = jnp.exp(-1j * eta / 2)
    J00 = scale * (c**2 + jnp.exp(1j * eta) * s**2)
    J11 = scale * (s**2 + jnp.exp(1j * eta) * c**2)
    J01 = scale * (1 - jnp.exp(1j * eta)) * jnp.exp(-1j * phi) * s * c
    J10 = scale * (1 - jnp.exp(1j * eta)) * jnp.exp(1j * phi) * s * c
    return polarizer(field, J00, J01, J10, J11)

polarizer(field, J00, J01, J10, J11)

Applies a Jones matrix with given components to the field. Note that the components here refer to the common choice of coordinate system and are inverted by us - i.e. J00 refers to Jxx.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	field to apply polarization to.	required
J00	Union[float, complex, Array]	description	required
J01	Union[float, complex, Array]	description	required
J10	Union[float, complex, Array]	description	required
J11	Union[float, complex, Array]	description	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: Field after polarizer.

Source code in src/chromatix/functional/polarizers.py

def polarizer(
    field: VectorField | ChromaticVectorField,
    J00: ComplexScalarLike,
    J01: ComplexScalarLike,
    J10: ComplexScalarLike,
    J11: ComplexScalarLike,
) -> VectorField | ChromaticVectorField:
    """Applies a Jones matrix with given components to the field.
    Note that the components here refer to the common choice of coordinate
    system and are inverted by us - i.e. J00 refers to Jxx.

    Args:
        field (VectorField | ChromaticVectorField): field to apply polarization to.
        J00 (Union[float, complex, Array]): _description_
        J01 (Union[float, complex, Array]): _description_
        J10 (Union[float, complex, Array]): _description_
        J11 (Union[float, complex, Array]): _description_

    Returns:
        VectorField | ChromaticVectorField: Field after polarizer.
    """
    assert isinstance(field, Vector), "Must be a vectorial Field"
    # Invert the axes as our order is zyx
    LP = jnp.array([[0, 0, 0], [0, J11, J10], [0, J01, J00]])
    LP = LP / jnp.linalg.norm(LP)
    return field.replace(u=matvec(LP, field.u))

quarterwave_plate(field, theta)

Applies a quarterwave plate with angle theta to the incoming field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	incoming field.	required
theta	float	angle w.r.t. horizontal.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: outgoing field.

Source code in src/chromatix/functional/polarizers.py

def quarterwave_plate(
    field: VectorField | ChromaticVectorField, theta: ScalarLike
) -> VectorField | ChromaticVectorField:
    """Applies a quarterwave plate with angle theta to the incoming field.

    Args:
        field (VectorField | ChromaticVectorField): incoming field.
        theta (float): angle w.r.t. horizontal.

    Returns:
        VectorField | ChromaticVectorField: outgoing field.
    """
    return phase_retarder(field, theta, eta=jnp.pi / 2, phi=0)

right_circular()

Generates a Jones vector for right circularly polarized light.

Returns:

Name	Type	Description
Array	Array	Right circularly polarized Jones vector.

Source code in src/chromatix/functional/polarizers.py

def right_circular() -> Array:
    """Generates a Jones vector for right circularly polarized
    light.

    Returns:
        Array: Right circularly polarized Jones vector.
    """
    return jones_vector(jnp.pi / 4, -jnp.pi / 2)

right_circular_polarizer(field)

Applies a thin RCP polarizer to the incoming Field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	The Field to which the polarizer will be applied.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	The Field directly after the polarizer.

Source code in src/chromatix/functional/polarizers.py

def right_circular_polarizer(
    field: VectorField | ChromaticVectorField,
) -> VectorField | ChromaticVectorField:
    """
    Applies a thin RCP polarizer to the incoming ``Field``.

    Args:
        field: The ``Field`` to which the polarizer will be applied.

    Returns:
        The ``Field`` directly after the polarizer.
    """
    J00 = 1
    J11 = 1
    J01 = 1j
    J10 = -1j
    return polarizer(field, J00, J01, J10, J11)

universal_compensator(field, retardance_45, retardance_0)

Applies the Universal Polarizer for the LC-PolScope to the incoming field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	Incoming field to polarize.	required
retardance_45	ScalarLike	Retardance induced at a 45 deg angle.	required
retardance_0	ScalarLike	Retardance induced at a 0 deg angle.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	The outgoing field.

Source code in src/chromatix/functional/polarizers.py

def universal_compensator(
    field: VectorField | ChromaticVectorField,
    retardance_45: ScalarLike,
    retardance_0: ScalarLike,
) -> VectorField | ChromaticVectorField:
    """Applies the Universal Polarizer for the LC-PolScope to the incoming field.

    Args:
        field: Incoming field to polarize.
        retardance_45: Retardance induced at a 45 deg angle.
        retardance_0: Retardance induced at a 0 deg angle.

    Returns:
        The outgoing field.
    """
    field = linear_polarizer(field, 0)
    field_retardance_45 = wave_plate(field, -jnp.pi / 4, retardance_45)
    field_retardance_0 = wave_plate(field_retardance_45, 0, retardance_0)
    return field_retardance_0

wave_plate(field, theta, eta)

Applies a general waveplate with angle theta and delay eta to the field.

Parameters:

Name	Type	Description	Default
field	VectorField | ChromaticVectorField	incoming field.	required
theta	float	angle w.r.t horizontal.	required
eta	float	relative delay between components.	required

Returns:

Type	Description
VectorField | ChromaticVectorField	VectorField | ChromaticVectorField: outgoing field.

Source code in src/chromatix/functional/polarizers.py

def wave_plate(
    field: VectorField | ChromaticVectorField, theta: ScalarLike, eta: ScalarLike
) -> VectorField | ChromaticVectorField:
    """Applies a general waveplate with angle theta and delay eta to the field.

    Args:
        field (VectorField | ChromaticVectorField): incoming field.
        theta (float): angle w.r.t horizontal.
        eta (float): relative delay between components.

    Returns:
        VectorField | ChromaticVectorField: outgoing field.
    """
    return phase_retarder(field, theta, eta, phi=0)

Convenience

optical_fft(field, z, n)

Computes the optical fft or ifft on an incoming Field propagated by z, assuming that the distance is in the far field. Can also be used for simulating propagation through a lens from the focal plane to the back focal plane if z is the focal length of the lens. The direction of the propagation depends on the sign of z (which is a scalar value that may be positive or negative). If z is positive an fft```will be performed, otherwise anifft(due to the1 / (lambda * z)term in the single Fourier transform Fresnel propagation, which requires this behavior). Theifftis calculated in terms of the conjugate of thefftwith appropriate normalization applied so that propagating forwards and then backwards yields the sameFieldup to numerical precision. This function also appropriately changes the sampling (dx) of the resultingField``.

Parameters:

Name	Type	Description	Default
field	Field	The Field to be propagated by fft.	required
z	ScalarLike	Scalar representing how far the Field will be propagated in units of distance.	required
n	ScalarLike	Real-valued scalar representing (isotropic) refractive index of the propagation medium.	required

Returns: The propagated Field, transformed by fft/ifft.

Source code in src/chromatix/functional/convenience.py

def optical_fft(field: Field, z: ScalarLike, n: ScalarLike) -> Field:
    """
    Computes the optical ``fft`` or ``ifft`` on an incoming ``Field`` propagated
    by ``z``, assuming that the distance is in the far field. Can also be used
    for simulating propagation through a lens from the focal plane to the back
    focal plane if ``z`` is the focal length of the lens. The direction of the
    propagation depends on the sign of ``z`` (which is a scalar value that may
    be positive or negative). If ``z`` is positive an ``fft```will be performed,
    otherwise an ``ifft`` (due to the ``1 / (lambda * z)`` term in the single
    Fourier transform Fresnel propagation, which requires this behavior).
    The ``ifft`` is calculated in terms of the conjugate of the ``fft`` with
    appropriate normalization applied so that propagating forwards and then
    backwards yields the same ``Field`` up to numerical precision. This function
    also appropriately changes the sampling (``dx``) of the resulting ``Field``.

    Args:
        field: The ``Field`` to be propagated by ``fft``.
        z: Scalar representing how far the ``Field`` will be propagated in units
            of distance.
        n: Real-valued scalar representing (isotropic) refractive index of the
            propagation medium.
    Returns:
        The propagated ``Field``, transformed by ``fft``/``ifft``.
    """
    L_sq = field.broadcasted_wavelength * z / n
    if field.spectrum.size > 1:
        shape_spec = "wv -> wv"
        for i in range(field.df.ndim - 1):
            shape_spec += " 1"
        _L_sq = rearrange(L_sq.squeeze(), shape_spec)
        du = field.df * jnp.abs(_L_sq)
    else:
        du = field.df * jnp.abs(L_sq)
    # Forward transform normalization for z >= 0
    norm_fft = (z >= 0) * -1j * jnp.prod(field.dx, axis=-1) / L_sq
    # Inverse transform normalization for z < 0
    norm_ifft = (
        (z < 0)
        * -1j  # Sign change because we take the conjugate of the input
        * (L_sq / jnp.prod(du, axis=-1))  # Inverse length scale
        / jnp.prod(
            jnp.array(
                field.spatial_shape
            )  # Due to a different norm factor for fft and ifft
        )
    )
    # Inverse transform input needs to use the conjugate
    fft_input = (norm_fft * field.u) + (norm_ifft * field.conj.u)
    fft_output = fft(fft_input, axes=field.spatial_dims, shift=True)
    u = (z >= 0) * fft_output + (z < 0) * jnp.conj(fft_output)
    return field.replace(u=u, dx=du)