Elements

Sources

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

__all__ = ['PointSource', 'ObjectivePointSource', 'PlaneWave', 'GenericField'] module-attribute

GenericField

Bases: Module

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.

Attributes:

Name	Type	Description
dx	ScalarLike	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	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.
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.
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.
power	ScalarLike	The total power (a scalar value) that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.

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

class GenericField(eqx.Module):
    """
    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).

    Attributes:
        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``.
    """

    dx: ScalarLike
    spectrum: (
        Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]]
    )
    amplitude: ArrayLike
    phase: ArrayLike
    power: ScalarLike
    pupil: FieldPupil | None = eqx.field(static=True)
    scalar: bool = eqx.field(static=True)

    def __init(
        self,
        dx: ScalarLike,
        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,
    ):
        """
        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``.
        """
        self.dx = dx
        self.spectrum = spectrum
        self.amplitude = amplitude
        self.phase = phase
        self.power = power
        self.pupil = pupil
        self.scalar = scalar

    def __call__(self) -> ScalarField | VectorField:
        return generic_field(
            self.dx,
            self.spectrum,
            self.amplitude,
            self.phase,
            self.power,
            self.pupil,
            self.scalar,
        )

amplitude instance-attribute

dx instance-attribute

phase instance-attribute

power instance-attribute

pupil = eqx.field(static=True) class-attribute instance-attribute

scalar = eqx.field(static=True) class-attribute instance-attribute

spectrum instance-attribute

__call__()

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

def __call__(self) -> ScalarField | VectorField:
    return generic_field(
        self.dx,
        self.spectrum,
        self.amplitude,
        self.phase,
        self.power,
        self.pupil,
        self.scalar,
    )

__init(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	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/elements/sources.py

def __init(
    self,
    dx: ScalarLike,
    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,
):
    """
    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``.
    """
    self.dx = dx
    self.spectrum = spectrum
    self.amplitude = amplitude
    self.phase = phase
    self.power = power
    self.pupil = pupil
    self.scalar = scalar

ObjectivePointSource

Bases: Module

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.

Attributes:

Name	Type	Description
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).
dx	ScalarLike | Float[Array, 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).
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.
f	ScalarLike	Focal length of the objective lens in units of distance.
n	ScalarLike	Refractive index.
NA	ScalarLike	The numerical aperture of the objective lens.
power	ScalarLike | None	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.
amplitude	ScalarLike	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.
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.
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.

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

class ObjectivePointSource(eqx.Module):
    """
    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``.

    Attributes:
        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.
        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.
    """

    shape: tuple[int, int] = eqx.field(static=True)
    dx: ScalarLike | Float[Array, "2"]
    spectrum: (
        Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]]
    )
    f: ScalarLike
    n: ScalarLike
    NA: ScalarLike
    power: ScalarLike | None
    amplitude: ScalarLike
    offset: Float[Array, "2"] | tuple[float, float]
    scalar: bool = eqx.field(static=True)

    def __init__(
        self,
        shape: tuple[int, int],
        dx: ScalarLike,
        spectrum: Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]],
        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,
    ):
        """
        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.
            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.
        """
        self.shape = shape
        self.dx = dx
        self.spectrum = spectrum
        self.f = f
        self.n = n
        self.NA = NA
        self.power = power
        self.amplitude = amplitude
        self.offset = offset
        self.scalar = scalar

    def __call__(self, z: ScalarLike | Float[Array, "z"]) -> 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:
            z: How far away the point source is from the focal plane of the lens
                in units of distance.
        """
        return objective_point_source(
            self.shape,
            self.dx,
            self.spectrum,
            z,
            self.f,
            self.n,
            self.NA,
            self.power,
            self.amplitude,
            self.offset,
            self.scalar,
        )

NA = NA instance-attribute

amplitude = amplitude instance-attribute

dx = dx instance-attribute

f = f instance-attribute

n = n instance-attribute

offset = offset instance-attribute

power = power instance-attribute

scalar = scalar class-attribute instance-attribute

shape = shape class-attribute instance-attribute

spectrum = spectrum instance-attribute

__call__(z)

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
z	ScalarLike | Float[Array, z]	How far away the point source is from the focal plane of the lens in units of distance.	required

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

def __call__(self, z: ScalarLike | Float[Array, "z"]) -> 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:
        z: How far away the point source is from the focal plane of the lens
            in units of distance.
    """
    return objective_point_source(
        self.shape,
        self.dx,
        self.spectrum,
        z,
        self.f,
        self.n,
        self.NA,
        self.power,
        self.amplitude,
        self.offset,
        self.scalar,
    )

__init__(shape, dx, spectrum, 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	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
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/elements/sources.py

def __init__(
    self,
    shape: tuple[int, int],
    dx: ScalarLike,
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    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,
):
    """
    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.
        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.
    """
    self.shape = shape
    self.dx = dx
    self.spectrum = spectrum
    self.f = f
    self.n = n
    self.NA = NA
    self.power = power
    self.amplitude = amplitude
    self.offset = offset
    self.scalar = scalar

PlaneWave

Bases: Module

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)).

Attributes:

Name	Type	Description
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).
dx	ScalarLike	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	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.
power	ScalarLike	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.
amplitude	ScalarLike	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.
kykx	Float[Array, 2] | 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.
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.

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

class PlaneWave(eqx.Module):
    """
    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))``.

    Attributes:
        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.
    """

    shape: tuple[int, int] = eqx.field(static=True)
    dx: ScalarLike
    spectrum: (
        Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]]
    )
    power: ScalarLike
    amplitude: ScalarLike
    kykx: Float[Array, "2"] | tuple[float, float]
    pupil: FieldPupil | None = eqx.field(static=True)
    scalar: bool = eqx.field(static=True)

    def __init__(
        self,
        shape: tuple[int, int],
        dx: ScalarLike,
        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,
    ):
        """
        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.
        """
        self.shape = shape
        self.dx = dx
        self.spectrum = spectrum
        self.power = power
        self.amplitude = amplitude
        self.kykx = kykx
        self.pupil = pupil
        self.scalar = scalar

    def __call__(self) -> ScalarField | VectorField:
        return plane_wave(
            self.shape,
            self.dx,
            self.spectrum,
            self.power,
            self.amplitude,
            self.kykx,
            self.pupil,
            self.scalar,
        )

amplitude = amplitude instance-attribute

dx = dx instance-attribute

kykx = kykx instance-attribute

power = power instance-attribute

pupil = pupil class-attribute instance-attribute

scalar = scalar class-attribute instance-attribute

shape = shape class-attribute instance-attribute

spectrum = spectrum instance-attribute

__call__()

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

def __call__(self) -> ScalarField | VectorField:
    return plane_wave(
        self.shape,
        self.dx,
        self.spectrum,
        self.power,
        self.amplitude,
        self.kykx,
        self.pupil,
        self.scalar,
    )

__init__(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	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/elements/sources.py

def __init__(
    self,
    shape: tuple[int, int],
    dx: ScalarLike,
    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,
):
    """
    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.
    """
    self.shape = shape
    self.dx = dx
    self.spectrum = spectrum
    self.power = power
    self.amplitude = amplitude
    self.kykx = kykx
    self.pupil = pupil
    self.scalar = scalar

PointSource

Bases: Module

Generates field due to point source a distance z away.

Can also be given pupil.

Attributes:

Name	Type	Description
shape	tuple[int, int]	The shape (height and width) of the Field to be created in number of samples (pixels).
dx	ScalarLike | Float[Array, 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).
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.
n	ScalarLike	Refractive index.
power	ScalarLike | None	The total power that the result should be normalized to, defaults to 1.0. If None, no normalization occurs.
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.
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.
pupil	FieldPupil | None	A function that applies a pupil to the field if provided. Defaults to None in which case the field is unchanged.
scalar	bool	Whether the result should be ScalarField (if True) or VectorField (if False). Defaults to True.
epsilon	float	Value added to denominators for numerical stability when z is 0.

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

class PointSource(eqx.Module):
    """
    Generates field due to point source a distance ``z`` away.

    Can also be given ``pupil``.

    Attributes:
        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.
        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.
    """

    shape: tuple[int, int] = eqx.field(static=True)
    dx: ScalarLike | Float[Array, "2"]
    spectrum: (
        Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]]
    )
    n: ScalarLike
    power: ScalarLike | None
    amplitude: ScalarLike | Float[Array, "3"]
    offset: Float[Array, "2"] | tuple[float, float]
    pupil: FieldPupil | None = eqx.field(static=True)
    scalar: bool = eqx.field(static=True)
    epsilon: float = eqx.field(static=True)

    def __init__(
        self,
        shape: tuple[int, int],
        dx: ScalarLike | Float[Array, "2"],
        spectrum: Spectrum
        | ScalarLike
        | Float[Array, "wv"]
        | tuple[Float[Array, "wv"], Float[Array, "wv"]],
        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``.

        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.
            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.
        """
        self.shape = shape
        self.dx = dx
        self.spectrum = (spectrum,)
        self.n = n
        self.power = power
        self.amplitude = amplitude
        self.offset = offset
        self.pupil = pupil
        self.scalar = scalar
        self.epsilon = epsilon

    def __call__(self, z: ScalarLike | Float[Array, "z"]) -> ScalarField | VectorField:
        """
        Generates field due to point source a distance ``z`` away.

        Args:
            z: How far away the point source is in units of distance.
        """
        return point_source(
            self.shape,
            self.dx,
            self.spectrum,
            z,
            self.n,
            self.power,
            self.amplitude,
            self.offset,
            self.pupil,
            self.scalar,
            self.epsilon,
        )

amplitude = amplitude instance-attribute

dx = dx instance-attribute

epsilon = epsilon class-attribute instance-attribute

n = n instance-attribute

offset = offset instance-attribute

power = power instance-attribute

pupil = pupil class-attribute instance-attribute

scalar = scalar class-attribute instance-attribute

shape = shape class-attribute instance-attribute

spectrum = (spectrum,) instance-attribute

__call__(z)

Generates field due to point source a distance z away.

Parameters:

Name	Type	Description	Default
z	ScalarLike | Float[Array, z]	How far away the point source is in units of distance.	required

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

def __call__(self, z: ScalarLike | Float[Array, "z"]) -> ScalarField | VectorField:
    """
    Generates field due to point source a distance ``z`` away.

    Args:
        z: How far away the point source is in units of distance.
    """
    return point_source(
        self.shape,
        self.dx,
        self.spectrum,
        z,
        self.n,
        self.power,
        self.amplitude,
        self.offset,
        self.pupil,
        self.scalar,
        self.epsilon,
    )

__init__(shape, dx, spectrum, 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.

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, 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
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/elements/sources.py

def __init__(
    self,
    shape: tuple[int, int],
    dx: ScalarLike | Float[Array, "2"],
    spectrum: Spectrum
    | ScalarLike
    | Float[Array, "wv"]
    | tuple[Float[Array, "wv"], Float[Array, "wv"]],
    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``.

    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.
        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.
    """
    self.shape = shape
    self.dx = dx
    self.spectrum = (spectrum,)
    self.n = n
    self.power = power
    self.amplitude = amplitude
    self.offset = offset
    self.pupil = pupil
    self.scalar = scalar
    self.epsilon = epsilon

Lenses

__all__ = ['ThinLens', 'FFLens', 'DFLens'] module-attribute

DFLens

Bases: Module

Applies a thin lens placed a distance d after the incoming Field. This element returns the Field a distance f after the lens.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Attributes:

Name	Type	Description
d	ScalarLike	How far away the lens is from the incoming Field in units of distance.
f	ScalarLike	Focal length of the lens in units of distance.
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.
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.

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

class DFLens(eqx.Module):
    """
    Applies a thin lens placed a distance ``d`` after the incoming ``Field``.
    This element returns the ``Field`` a distance ``f`` after the lens.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Attributes:
        d: How far away the lens is from the incoming ``Field`` in units of distance.
        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``.
    """

    d: ScalarLike
    f: ScalarLike
    n: ScalarLike
    NA: ScalarLike | None
    inverse: bool = eqx.field(static=True)

    def __init__(
        self,
        d: ScalarLike,
        f: ScalarLike,
        n: ScalarLike,
        NA: ScalarLike | None = None,
        inverse: bool = False,
    ):
        """
        Applies a thin lens placed a distance ``d`` after the incoming ``Field``.
        This element returns the ``Field`` a distance ``f`` after the lens.

        This element can be placed after any element that returns a ``Field`` or
        before any element that accepts a ``Field``.

        Args:
            d: How far away the lens is from the incoming ``Field`` in units
                of distance.
            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``.
        """
        self.d = d
        self.f = f
        self.n = n
        self.NA = NA
        self.inverse = inverse

    def __call__(self, field: Field) -> Field:
        return cf.df_lens(field, self.d, self.f, self.n, self.NA, inverse=self.inverse)

NA = NA instance-attribute

d = d instance-attribute

f = f instance-attribute

inverse = inverse class-attribute instance-attribute

n = n instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    return cf.df_lens(field, self.d, self.f, self.n, self.NA, inverse=self.inverse)

__init__(d, f, n, NA=None, inverse=False)

Applies a thin lens placed a distance d after the incoming Field. This element returns the Field a distance f after the lens.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Parameters:

Name	Type	Description	Default
d	ScalarLike	How far away the lens is from the incoming Field in units of distance.	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

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

def __init__(
    self,
    d: ScalarLike,
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike | None = None,
    inverse: bool = False,
):
    """
    Applies a thin lens placed a distance ``d`` after the incoming ``Field``.
    This element returns the ``Field`` a distance ``f`` after the lens.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Args:
        d: How far away the lens is from the incoming ``Field`` in units
            of distance.
        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``.
    """
    self.d = d
    self.f = f
    self.n = n
    self.NA = NA
    self.inverse = inverse

FFLens

Bases: Module

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

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Attributes:

Name	Type	Description
f	ScalarLike	Focal length of the lens.
n	ScalarLike	The refractive index of the surrounding medium (assumed to be the same incoming and exiting).
NA	ScalarLike | None	If provided, the NA of the lens. By default, no pupil is applied to the incoming Field.
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.

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

class FFLens(eqx.Module):
    """
    Applies a thin lens placed a distance ``f`` after the incoming ``Field``.
    This element returns the ``Field`` a distance ``f`` after the lens.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Attributes:
        f: Focal length of the lens.
        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``.
    """

    f: ScalarLike
    n: ScalarLike
    NA: ScalarLike | None
    inverse: bool = eqx.field(static=True)

    def __init__(
        self,
        f: ScalarLike,
        n: ScalarLike,
        NA: ScalarLike | None = None,
        inverse: bool = False,
    ):
        """
        Applies a thin lens placed directly after the incoming ``Field``.
        This element returns the ``Field`` directly after the lens.

        This element can be placed after any element that returns a ``Field`` or
        before any element that accepts a ``Field``.

        Args:
            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``.
        """
        self.f = f
        self.n = n
        self.NA = NA
        self.inverse = inverse

    def __call__(self, field: Field) -> Field:
        return cf.ff_lens(field, self.f, self.n, self.NA, inverse=self.inverse)

NA = NA instance-attribute

f = f instance-attribute

inverse = inverse class-attribute instance-attribute

n = n instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    return cf.ff_lens(field, self.f, self.n, self.NA, inverse=self.inverse)

__init__(f, n, NA=None, inverse=False)

Applies a thin lens placed directly after the incoming Field. This element returns the Field directly after the lens.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Parameters:

Name	Type	Description	Default
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

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

def __init__(
    self,
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike | None = None,
    inverse: bool = False,
):
    """
    Applies a thin lens placed directly after the incoming ``Field``.
    This element returns the ``Field`` directly after the lens.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Args:
        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``.
    """
    self.f = f
    self.n = n
    self.NA = NA
    self.inverse = inverse

ThinLens

Bases: Module

Applies a thin lens placed directly after the incoming Field. This element returns the Field directly after the lens.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Attributes:

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

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

class ThinLens(eqx.Module):
    """
    Applies a thin lens placed directly after the incoming ``Field``.
    This element returns the ``Field`` directly after the lens.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Attributes:
        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``.
    """

    f: ScalarLike
    n: ScalarLike
    NA: ScalarLike | None

    def __init__(self, f: ScalarLike, n: ScalarLike, NA: ScalarLike | None = None):
        """
        Applies a thin lens placed directly after the incoming ``Field``.
        This element returns the ``Field`` directly after the lens.

        This element can be placed after any element that returns a ``Field`` or
        before any element that accepts a ``Field``.

        Args:
            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``.
        """
        self.f = f
        self.n = n
        self.NA = NA

    def __call__(self, field: Field) -> Field:
        return cf.thin_lens(field, self.f, self.n, self.NA)

NA = NA instance-attribute

f = f instance-attribute

n = n instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    return cf.thin_lens(field, self.f, self.n, self.NA)

__init__(f, n, NA=None)

Applies a thin lens placed directly after the incoming Field. This element returns the Field directly after the lens.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Parameters:

Name	Type	Description	Default
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

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

def __init__(self, f: ScalarLike, n: ScalarLike, NA: ScalarLike | None = None):
    """
    Applies a thin lens placed directly after the incoming ``Field``.
    This element returns the ``Field`` directly after the lens.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Args:
        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``.
    """
    self.f = f
    self.n = n
    self.NA = NA

Phase Masks

__all__ = ['SeidelAberrations', 'ZernikeAberrations', 'PhaseMask', 'SpatialLightModulator'] module-attribute

PhaseMask

Bases: Module

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

This element can be placed after any element that returns a Field or before any element that accepts a Field.

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.

The phase can be optimized (pixel by pixel).

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.

Attributes:

Name	Type	Description
phase	Float[Array, 'h w']	The phase mask to apply as a 2D array of phase values (in radians) of shape (height width).

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

class PhaseMask(eqx.Module):
    """
    Perturbs ``field`` by ``phase`` (**in radians**), i.e. ``field * exp(1j * phase)``.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    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.

    The ``phase`` can be optimized (pixel by pixel).

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

    Attributes:
        phase: The phase mask to apply as a 2D array of phase values (in
            radians) of shape ``(height width)``.
    """

    phase: Float[Array, "h w"]

    def __init__(self, phase: Float[Array, "h w"]):
        self.phase = phase

    def __call__(self, field: Field) -> Field:
        """Applies ``phase`` mask to incoming ``Field``."""
        return phase_change(field, self.phase)

phase = phase instance-attribute

__call__(field)

Applies phase mask to incoming Field.

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

def __call__(self, field: Field) -> Field:
    """Applies ``phase`` mask to incoming ``Field``."""
    return phase_change(field, self.phase)

__init__(phase)

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

def __init__(self, phase: Float[Array, "h w"]):
    self.phase = phase

SeidelAberrations

Bases: Module

Applies Seidel phase polynomial to an incoming Field.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

This element handles multi-wavelength Fields by assuming that the first wavelength in the spectrum of the Field is the central wavelength for which the phase was calculated, and modulates the phase by the ratio of other wavelengths in the spectrum to the central wavelength appropriately.

Attributes:

Name	Type	Description
coefficients	Float[Array, 5]	The Seidel coefficients. Should have shape (5,).
f	ScalarLike	The focal length of the system's objective lens in units of distance.
n	ScalarLike	The refractive index of the medium.
NA	ScalarLike	The numerical aperture. The applied phase will be 0 outside NA.

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

class SeidelAberrations(eqx.Module):
    """
    Applies Seidel phase polynomial to an incoming ``Field``.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    This element handles multi-wavelength ``Field``s by assuming that the first
    wavelength in the ``spectrum`` of the ``Field`` is the central wavelength
    for which the ``phase`` was calculated, and modulates the ``phase`` by the
    ratio of other wavelengths in the ``spectrum`` to the central wavelength
    appropriately.

    Attributes:
        coefficients: The Seidel coefficients. Should have shape `(5,)`.
        f: The focal length of the system's objective lens in units of distance.
        n: The refractive index of the medium.
        NA: The numerical aperture. The applied phase will be 0 outside NA.
    """

    coefficients: Float[Array, "5"]
    f: ScalarLike
    n: ScalarLike
    NA: ScalarLike

    def __init__(
        self,
        coefficients: Float[Array, "5"],
        f: ScalarLike,
        n: ScalarLike,
        NA: ScalarLike,
    ):
        self.coefficients = coefficients
        self.f = f
        self.n = n
        self.NA = NA

    def __call__(self, field: Field, u: ScalarLike, v: ScalarLike) -> Field:
        """
        Applies ``phase`` mask to incoming ``Field``.

        Args:
            field: The complex field to be perturbed.
            u: The horizontal position of the object field point in normalized
                coordinates from 0 to +/- 1. A value of 0 represents the center
                coordinate in the plane while a value of 1 represents the farthest
                point from the center. Positive values go right and negative values
                go left.
            v: The vertical position of the object field point in normalized
                coordinates from 0 to +/- 1. A value of 0 represents the center
                coordinate in the plane while a value of 1 represents the farthest
                point from the center. Positive values go down and negative values
                go up.
        """
        phase = seidel_aberrations(
            field.spatial_shape,
            field.central_dx,
            field.central_wavelength,
            self.n,
            self.f,
            self.NA,
            self.coefficients,
            self.u,
            self.v,
        )
        return phase_change(field, phase)

NA = NA instance-attribute

coefficients = coefficients instance-attribute

f = f instance-attribute

n = n instance-attribute

__call__(field, u, v)

Applies phase mask to incoming Field.

Parameters:

Name	Type	Description	Default
field	Field	The complex field to be perturbed.	required
u	ScalarLike	The horizontal position of the object field point in normalized coordinates from 0 to +/- 1. A value of 0 represents the center coordinate in the plane while a value of 1 represents the farthest point from the center. Positive values go right and negative values go left.	required
v	ScalarLike	The vertical position of the object field point in normalized coordinates from 0 to +/- 1. A value of 0 represents the center coordinate in the plane while a value of 1 represents the farthest point from the center. Positive values go down and negative values go up.	required

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

def __call__(self, field: Field, u: ScalarLike, v: ScalarLike) -> Field:
    """
    Applies ``phase`` mask to incoming ``Field``.

    Args:
        field: The complex field to be perturbed.
        u: The horizontal position of the object field point in normalized
            coordinates from 0 to +/- 1. A value of 0 represents the center
            coordinate in the plane while a value of 1 represents the farthest
            point from the center. Positive values go right and negative values
            go left.
        v: The vertical position of the object field point in normalized
            coordinates from 0 to +/- 1. A value of 0 represents the center
            coordinate in the plane while a value of 1 represents the farthest
            point from the center. Positive values go down and negative values
            go up.
    """
    phase = seidel_aberrations(
        field.spatial_shape,
        field.central_dx,
        field.central_wavelength,
        self.n,
        self.f,
        self.NA,
        self.coefficients,
        self.u,
        self.v,
    )
    return phase_change(field, phase)

__init__(coefficients, f, n, NA)

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

def __init__(
    self,
    coefficients: Float[Array, "5"],
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike,
):
    self.coefficients = coefficients
    self.f = f
    self.n = n
    self.NA = NA

SpatialLightModulator

Bases: Module

Simulates a spatial light modulator (SLM) applied to an incoming Field.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

This means that this element acts as if the SLM is phase only and transmits a Field, rather than reflecting it.

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.

This element also handles the limited phase modulation range of an SLM, controlled by phase_range.

This element also approximates the simulation of the surface of the SLM by interpolating the phase with the given shape (which can be the number of pixels in the SLM) to the shape of the incoming Field, which is assumed to have a shape larger than the given SLM shape. The order of the interpolation performed can be controlled with interpolation_order, but currently only orders of 0 (nearest neighbor) and 1 (linear interpolation) are supported.

Attributes:

Name	Type	Description
phase	Float[Array, 'h w']	The phase mask to apply as a 2D array of phase values (in radians) of shape (height width).
shape	tuple[int, int]	The shape of the SLM in number of pixels, provided as (height width).
spacing	ScalarLike	The pitch of the SLM pixels in units of distance.
phase_range	Float[Array, 2] | tuple[float, float]	The phase range in radians that the SLM can simulate, provided as (min, max).
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.
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.

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

class SpatialLightModulator(eqx.Module):
    """
    Simulates a spatial light modulator (SLM) applied to an incoming ``Field``.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    This means that this element acts as if the SLM is phase only and transmits
    a ``Field``, rather than reflecting it.

    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.

    This element also handles the limited phase modulation range of an SLM,
    controlled by ``phase_range``.

    This element also approximates the simulation of the surface of the SLM by
    interpolating the phase with the given ``shape`` (which can be the number of
    pixels in the SLM) to the shape of the incoming ``Field``, which is assumed
    to have a shape larger than the given SLM ``shape``. The order of the
    interpolation performed can be controlled with ``interpolation_order``, but
    currently only orders of 0 (nearest neighbor) and 1 (linear interpolation)
    are supported.

    Attributes:
        phase: The phase mask to apply as a 2D array of phase values (in
            radians) of shape ``(height width)``.
        shape: The shape of the SLM in number of pixels, provided as `(height
            width)`.
        spacing: The pitch of the SLM pixels in units of distance.
        phase_range: The phase range in radians that the SLM can simulate,
            provided as `(min, max)`.
        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.
    """

    phase: Float[Array, "h w"]
    shape: tuple[int, int]
    spacing: ScalarLike
    phase_range: Float[Array, "2"] | tuple[float, float]
    num_bits: int | None = None
    interpolation_order: int = 0

    def __init__(
        self,
        phase: Float[Array, "h w"],
        shape: tuple[int, int],
        spacing: ScalarLike,
        phase_range: Float[Array, "2"] | tuple[float, float],
        num_bits: int | None = None,
        interpolation_order: int = 0,
    ):
        self.phase = phase
        assert_rank(
            self.phase, 2, custom_message="Phase must be a 2D array of shape (H W)"
        )
        assert self.phase.shape == shape, (
            f"Phase array must have same shape as SLM: expected {shape} got {self.phase.shape}"
        )
        self.shape = shape
        self.spacing = spacing
        self.phase_range = phase_range
        self.num_bits = num_bits
        self.interpolation_order

    def __call__(self, field: Field) -> Field:
        """Applies simulated SLM ``phase`` mask to incoming ``Field``."""
        return interpolated_phase_change(
            field, self.phase, self.phase_range, self.num_bits, self.interpolation_order
        )

interpolation_order = 0 class-attribute instance-attribute

num_bits = num_bits class-attribute instance-attribute

phase = phase instance-attribute

phase_range = phase_range instance-attribute

shape = shape instance-attribute

spacing = spacing instance-attribute

__call__(field)

Applies simulated SLM phase mask to incoming Field.

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

def __call__(self, field: Field) -> Field:
    """Applies simulated SLM ``phase`` mask to incoming ``Field``."""
    return interpolated_phase_change(
        field, self.phase, self.phase_range, self.num_bits, self.interpolation_order
    )

__init__(phase, shape, spacing, phase_range, num_bits=None, interpolation_order=0)

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

def __init__(
    self,
    phase: Float[Array, "h w"],
    shape: tuple[int, int],
    spacing: ScalarLike,
    phase_range: Float[Array, "2"] | tuple[float, float],
    num_bits: int | None = None,
    interpolation_order: int = 0,
):
    self.phase = phase
    assert_rank(
        self.phase, 2, custom_message="Phase must be a 2D array of shape (H W)"
    )
    assert self.phase.shape == shape, (
        f"Phase array must have same shape as SLM: expected {shape} got {self.phase.shape}"
    )
    self.shape = shape
    self.spacing = spacing
    self.phase_range = phase_range
    self.num_bits = num_bits
    self.interpolation_order

ZernikeAberrations

Bases: Module

Applies Zernike aberrations to an incoming Field.

This element can be placed after any element that returns a Field or before any element that accepts a Field.

This element handles multi-wavelength Fields by assuming that the first wavelength in the spectrum of the Field is the central wavelength for which the phase was calculated, and modulates the phase by the ratio of other wavelengths in the spectrum to the central wavelength appropriately.

Attributes:

Name	Type	Description
coefficients	Float[Array, m]	The Zernike coefficients as a 1D array of the same length as the ansi_indices.
f	ScalarLike	The focal length of the system's objective lens in units of distance.
n	ScalarLike	The refractive index of the medium.
NA	ScalarLike	The numerical aperture. The applied phase will be 0 outside NA.
ansi_indices	Sequence[int]	Linear Zernike indices according to ANSI numbering.
coefficients	Float[Array, m]	Weight coefficients for the Zernike polynomials.
normalize	bool	Whether to normalize the Zernike coefficients. Defaults to True.

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

class ZernikeAberrations(eqx.Module):
    """
    Applies Zernike aberrations to an incoming ``Field``.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    This element handles multi-wavelength ``Field``s by assuming that the first
    wavelength in the ``spectrum`` of the ``Field`` is the central wavelength
    for which the ``phase`` was calculated, and modulates the ``phase`` by the
    ratio of other wavelengths in the ``spectrum`` to the central wavelength
    appropriately.

    Attributes:
        coefficients: The Zernike coefficients as a 1D array of the same length as the
            ``ansi_indices``.
        f: The focal length of the system's objective lens in units of distance.
        n: The refractive index of the medium.
        NA: The numerical aperture. The applied phase will be 0 outside NA.
        ansi_indices: Linear Zernike indices according to ANSI numbering.
        coefficients: Weight coefficients for the Zernike polynomials.
        normalize: Whether to normalize the Zernike coefficients. Defaults to
            ``True``.
    """

    coefficients: Float[Array, "m"]
    f: ScalarLike
    n: ScalarLike
    NA: ScalarLike
    ansi_indices: Sequence[int]
    normalize: bool

    def __init__(
        self,
        coefficients: Float[Array, "m"],
        f: ScalarLike,
        n: ScalarLike,
        NA: ScalarLike,
        ansi_indices: Sequence[int],
        normalize: bool = True,
    ):
        self.coefficients = coefficients
        self.f = f
        self.n = n
        self.NA = NA
        self.ansi_indices = ansi_indices
        self.normalize = normalize

    def __call__(self, field: Field) -> Field:
        """Applies ``phase`` mask to incoming ``Field``."""
        phase = zernike_aberrations(
            field.spatial_shape,
            field.central_dx,
            field.central_wavelength,
            self.n,
            self.f,
            self.NA,
            self.ansi_indices,
            self.coefficients,
            self.normalize,
        )
        return phase_change(field, phase)

NA = NA instance-attribute

ansi_indices = ansi_indices instance-attribute

coefficients = coefficients instance-attribute

f = f instance-attribute

n = n instance-attribute

normalize = normalize instance-attribute

__call__(field)

Applies phase mask to incoming Field.

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

def __call__(self, field: Field) -> Field:
    """Applies ``phase`` mask to incoming ``Field``."""
    phase = zernike_aberrations(
        field.spatial_shape,
        field.central_dx,
        field.central_wavelength,
        self.n,
        self.f,
        self.NA,
        self.ansi_indices,
        self.coefficients,
        self.normalize,
    )
    return phase_change(field, phase)

__init__(coefficients, f, n, NA, ansi_indices, normalize=True)

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

def __init__(
    self,
    coefficients: Float[Array, "m"],
    f: ScalarLike,
    n: ScalarLike,
    NA: ScalarLike,
    ansi_indices: Sequence[int],
    normalize: bool = True,
):
    self.coefficients = coefficients
    self.f = f
    self.n = n
    self.NA = NA
    self.ansi_indices = ansi_indices
    self.normalize = normalize

Amplitude Masks

__all__ = ['AmplitudeMask'] module-attribute

AmplitudeMask

Bases: Module

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

This element can be placed after any element that returns a Field or before any element that accepts a Field.

Attributes:

Name	Type	Description
amplitude	Float[Array, 'h w']	The amplitude mask to apply as a 2D array of amplitude values of shape (height width).
is_binary	bool	Binarize (make 0 or 1) the amplitude mask if True. Note that if this is False, amplitude values do not need to be in any range (i.e. they are allowed to be more than 1 which would add energy to the Field).

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

class AmplitudeMask(eqx.Module):
    """
    Perturbs ``field`` by ``amplitude``, i.e. ``field * amplitude``.

    This element can be placed after any element that returns a ``Field`` or
    before any element that accepts a ``Field``.

    Attributes:
        amplitude: The amplitude mask to apply as a 2D array of amplitude values
            of shape ``(height width)``.
        is_binary: Binarize (make 0 or 1) the amplitude mask if ``True``. Note
            that if this is ``False``, amplitude values do not need to be in
            any range (i.e. they are allowed to be more than 1 which would add
            energy to the ``Field``).
    """

    amplitude: Float[Array, "h w"]
    is_binary: bool = eqx.field(static=True)

    def __init__(self, amplitude: Float[Array, "h w"], is_binary: bool):
        self.amplitude = amplitude
        self.is_binary = is_binary

    def __call__(self, field: Field) -> Field:
        """Applies ``amplitude`` mask to incoming ``Field``."""
        if self.is_binary:
            amplitude = binarize(self.amplitude)
        else:
            amplitude = self.amplitude
        return amplitude_change(field, amplitude)

amplitude = amplitude instance-attribute

is_binary = is_binary class-attribute instance-attribute

__call__(field)

Applies amplitude mask to incoming Field.

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

def __call__(self, field: Field) -> Field:
    """Applies ``amplitude`` mask to incoming ``Field``."""
    if self.is_binary:
        amplitude = binarize(self.amplitude)
    else:
        amplitude = self.amplitude
    return amplitude_change(field, amplitude)

__init__(amplitude, is_binary)

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

def __init__(self, amplitude: Float[Array, "h w"], is_binary: bool):
    self.amplitude = amplitude
    self.is_binary = is_binary

Samples

__all__ = ['ThinSample', 'ClearThinSample', 'MultisliceThickSample', 'ClearMultisliceThickSample'] module-attribute

ClearMultisliceThickSample

Bases: Sample, Scattering, Volume

Perturbs incoming ScalarField as if it went through a clear but scattering thick sample. The thick sample is modeled as being made of many thin slices each of a given thickness. The dn_stack contains the 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.

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.

Attributes:

Name	Type	Description
dn	Float[Array, 'd h w']	The sample's isotropic refractive index change per voxel for each slice as a 3D array of shape (depth height width) where depth is the total number of slices.
n	ScalarLike	Average refractive index of the sample.
dx	ScalarLike | Float[Array, 2]	The spacing of each pixel of the sample as a scalar in units of distance or as a 1D array of shape (2,) in the format (y x).
thickness	ScalarLike	How far to propagate for each slice in units of distance.
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.

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

class ClearMultisliceThickSample(Sample, Scattering, Volume, strict=True):
    """
    Perturbs incoming ``ScalarField`` as if it went through a clear but
    scattering thick sample. The thick sample is modeled as being made of many
    thin slices each of a given thickness. The ``dn_stack`` contains the 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``.

    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``.

    Attributes:
        dn: The sample's isotropic refractive index change per voxel for
            each slice as a 3D array of shape ``(depth height width)`` where
            ``depth`` is the total number of slices.
        n: Average refractive index of the sample.
        dx: The spacing of each pixel of the sample as a scalar in units of
            distance or as a 1D array of shape `(2,)` in the format `(y x)`.
        thickness: How far to propagate for each slice in units of distance.
        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.
    """

    dn: Float[Array, "d h w"]
    n: ScalarLike
    dx: ScalarLike | Float[Array, "2"]
    thickness: ScalarLike
    NA: ScalarLike | None

    def __init__(
        self,
        dn: Float[Array, "d h w"],
        n: ScalarLike,
        dx: ScalarLike | Float[Array, "2"],
        thickness: ScalarLike,
        NA: ScalarLike | None = None,
    ):
        self.dn = dn
        self.n = n
        self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
        self.thickness = thickness
        self._verify_scalar_thickness()
        self.NA = NA

    def __call__(
        self,
        field: ScalarField,
        kykx: Float[Array, "2"] = (0.0, 0.0),
        pad_width: int = 0,
        reverse_propagate_distance: ScalarLike | None = None,
        return_stack: bool = False,
        remove_evanescent: bool = False,
        bandlimit: bool = False,
    ) -> ScalarField:
        """
        Args:
            field: The complex field to be perturbed.
            kykx: If provided, defines the orientation of the propagation.
                Should be an array of shape `(2,)` in the format ``ky kx``.
            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.
            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``.
        """
        self._verify_matching_spacing(field)
        return multislice_thick_sample(
            field,
            jnp.zeros_like(self.dn),
            self.dn,
            self.n,
            self.thickness,
            pad_width,
            NA=self.NA,
            kykx=kykx,
            reverse_propagate_distance=reverse_propagate_distance,
            return_stack=return_stack,
            remove_evanescent=remove_evanescent,
            bandlimit=bandlimit,
        )

    @property
    def shape(self) -> tuple[int, int, int]:
        return self.dn.shape

    @property
    def ndim(self) -> int:
        return self.dn.ndim

    @property
    def grid(self) -> Array:
        return grid(self.shape, self.dx)

    @property
    def num_planes(self) -> int:
        return self.dn.shape[0]

NA = NA instance-attribute

dn = dn instance-attribute

dx = _broadcast_dx_to_grid(dx, 1).squeeze() instance-attribute

grid property

n = n instance-attribute

ndim property

num_planes property

shape property

thickness = thickness instance-attribute

__call__(field, kykx=(0.0, 0.0), pad_width=0, reverse_propagate_distance=None, return_stack=False, remove_evanescent=False, bandlimit=False)

Parameters:

Name	Type	Description	Default
field	ScalarField	The complex field to be perturbed.	required
kykx	Float[Array, 2]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format ky kx.	(0.0, 0.0)
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.	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

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

def __call__(
    self,
    field: ScalarField,
    kykx: Float[Array, "2"] = (0.0, 0.0),
    pad_width: int = 0,
    reverse_propagate_distance: ScalarLike | None = None,
    return_stack: bool = False,
    remove_evanescent: bool = False,
    bandlimit: bool = False,
) -> ScalarField:
    """
    Args:
        field: The complex field to be perturbed.
        kykx: If provided, defines the orientation of the propagation.
            Should be an array of shape `(2,)` in the format ``ky kx``.
        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.
        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``.
    """
    self._verify_matching_spacing(field)
    return multislice_thick_sample(
        field,
        jnp.zeros_like(self.dn),
        self.dn,
        self.n,
        self.thickness,
        pad_width,
        NA=self.NA,
        kykx=kykx,
        reverse_propagate_distance=reverse_propagate_distance,
        return_stack=return_stack,
        remove_evanescent=remove_evanescent,
        bandlimit=bandlimit,
    )

__init__(dn, n, dx, thickness, NA=None)

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

def __init__(
    self,
    dn: Float[Array, "d h w"],
    n: ScalarLike,
    dx: ScalarLike | Float[Array, "2"],
    thickness: ScalarLike,
    NA: ScalarLike | None = None,
):
    self.dn = dn
    self.n = n
    self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
    self.thickness = thickness
    self._verify_scalar_thickness()
    self.NA = NA

ClearThinSample

Bases: Sample, Scattering

Perturbs an incoming ScalarField as if it went through a clear thin sample object with a given refractive index change dn from the refractive index of the medium and of a given thickness in the same units as the spectrum of the incoming ScalarField.

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

Returns a ScalarField with the result of the perturbation.

Attributes:

Name	Type	Description
dn	Float[Array, 'h w']	The sample's isotropic refractive index change as a 2D array of shape (height width).
dx	ScalarLike | Float[Array, 2]	The spacing of each pixel of the sample as a scalar in units of distance or as a 1D array of shape (2,) in the format (y x).
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).

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

class ClearThinSample(Sample, Scattering, strict=True):
    """
    Perturbs an incoming ``ScalarField`` as if it went through a clear
    thin sample object with a given refractive index change ``dn`` from the
    refractive index of the medium and of a given ``thickness`` in the same
    units as the spectrum of the incoming ``ScalarField``.

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

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

    Attributes:
        dn: The sample's isotropic refractive index change as a 2D array of
            shape ``(height width)``.
        dx: The spacing of each pixel of the sample as a scalar in units of
            distance or as a 1D array of shape `(2,)` in the format `(y x)`.
        thickness: Thickness in units of distance as a scalar or array
            broadcastable to ``(height width)`` (thickness at each sample
            location).
    """

    dn: Float[Array, "h w"]
    dx: ScalarLike | Float[Array, "2"]
    thickness: ScalarLike | Float[Array, "h w"]

    def __init__(
        self,
        dn: Float[Array, "h w"],
        dx: ScalarLike | Float[Array, "2"],
        thickness: ScalarLike | Float[Array, "h w"],
    ):
        self.dn = dn
        self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
        self.thickness = thickness

    def __call__(self, field: ScalarField) -> ScalarField:
        self._verify_matching_spacing(field)
        return thin_sample(field, jnp.zeros_like(self.dn), self.dn, self.thickness)

    @property
    def shape(self) -> tuple[int, int]:
        return self.dn.shape

    @property
    def ndim(self) -> int:
        return self.dn.ndim

    @property
    def grid(self) -> Array:
        return grid(self.shape, self.dx)

dn = dn instance-attribute

dx = _broadcast_dx_to_grid(dx, 1).squeeze() instance-attribute

grid property

ndim property

shape property

thickness = thickness instance-attribute

__call__(field)

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

def __call__(self, field: ScalarField) -> ScalarField:
    self._verify_matching_spacing(field)
    return thin_sample(field, jnp.zeros_like(self.dn), self.dn, self.thickness)

__init__(dn, dx, thickness)

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

def __init__(
    self,
    dn: Float[Array, "h w"],
    dx: ScalarLike | Float[Array, "2"],
    thickness: ScalarLike | Float[Array, "h w"],
):
    self.dn = dn
    self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
    self.thickness = thickness

MultisliceThickSample

Bases: Sample, Absorbing, Scattering, Volume

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.

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.

Attributes:

Name	Type	Description
absorption	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.
dn	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.
n	ScalarLike	Average refractive index of the sample.
dx	ScalarLike | Float[Array, 2]	The spacing of each pixel of the sample as a scalar in units of distance or as a 1D array of shape (2,) in the format (y x).
thickness	ScalarLike	How far to propagate for each slice in units of distance.
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.

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

class MultisliceThickSample(Sample, Absorbing, Scattering, Volume, strict=True):
    """
    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``.

    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``.

    Attributes:
        absorption: 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: 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.
        dx: The spacing of each pixel of the sample as a scalar in units of
            distance or as a 1D array of shape `(2,)` in the format `(y x)`.
        thickness: How far to propagate for each slice in units of distance.
        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.
    """

    absorption: Float[Array, "d h w"]
    dn: Float[Array, "d h w"]
    n: ScalarLike
    dx: ScalarLike | Float[Array, "2"]
    thickness: ScalarLike
    NA: ScalarLike | None

    def __init__(
        self,
        absorption: Float[Array, "d h w"],
        dn: Float[Array, "d h w"],
        n: ScalarLike,
        dx: ScalarLike | Float[Array, "2"],
        thickness: ScalarLike,
        NA: ScalarLike | None = None,
    ):
        self.absorption = absorption
        self.dn = dn
        self.n = n
        self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
        self.thickness = thickness
        self._verify_scalar_thickness()
        self.NA = NA

    def __call__(
        self,
        field: ScalarField,
        kykx: Float[Array, "2"] = (0.0, 0.0),
        pad_width: int = 0,
        reverse_propagate_distance: ScalarLike | None = None,
        return_stack: bool = False,
        remove_evanescent: bool = False,
        bandlimit: bool = False,
    ) -> ScalarField:
        """
        Args:
            field: The complex field to be perturbed.
            kykx: If provided, defines the orientation of the propagation.
                Should be an array of shape `(2,)` in the format ``ky kx``.
            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.
            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``.
        """
        self._verify_matching_spacing(field)
        return multislice_thick_sample(
            field,
            self.absorption,
            self.dn,
            self.n,
            self.thickness,
            pad_width,
            NA=self.NA,
            kykx=kykx,
            reverse_propagate_distance=reverse_propagate_distance,
            return_stack=return_stack,
            remove_evanescent=remove_evanescent,
            bandlimit=bandlimit,
        )

    @property
    def shape(self) -> tuple[int, int, int]:
        return self.absorption.shape

    @property
    def ndim(self) -> int:
        return self.absorption.ndim

    @property
    def grid(self) -> Array:
        return grid(self.shape, self.dx)

    @property
    def num_planes(self) -> int:
        return self.absorption.shape[0]

NA = NA instance-attribute

absorption = absorption instance-attribute

dn = dn instance-attribute

dx = _broadcast_dx_to_grid(dx, 1).squeeze() instance-attribute

grid property

n = n instance-attribute

ndim property

num_planes property

shape property

thickness = thickness instance-attribute

__call__(field, kykx=(0.0, 0.0), pad_width=0, reverse_propagate_distance=None, return_stack=False, remove_evanescent=False, bandlimit=False)

Parameters:

Name	Type	Description	Default
field	ScalarField	The complex field to be perturbed.	required
kykx	Float[Array, 2]	If provided, defines the orientation of the propagation. Should be an array of shape (2,) in the format ky kx.	(0.0, 0.0)
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.	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

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

def __call__(
    self,
    field: ScalarField,
    kykx: Float[Array, "2"] = (0.0, 0.0),
    pad_width: int = 0,
    reverse_propagate_distance: ScalarLike | None = None,
    return_stack: bool = False,
    remove_evanescent: bool = False,
    bandlimit: bool = False,
) -> ScalarField:
    """
    Args:
        field: The complex field to be perturbed.
        kykx: If provided, defines the orientation of the propagation.
            Should be an array of shape `(2,)` in the format ``ky kx``.
        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.
        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``.
    """
    self._verify_matching_spacing(field)
    return multislice_thick_sample(
        field,
        self.absorption,
        self.dn,
        self.n,
        self.thickness,
        pad_width,
        NA=self.NA,
        kykx=kykx,
        reverse_propagate_distance=reverse_propagate_distance,
        return_stack=return_stack,
        remove_evanescent=remove_evanescent,
        bandlimit=bandlimit,
    )

__init__(absorption, dn, n, dx, thickness, NA=None)

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

def __init__(
    self,
    absorption: Float[Array, "d h w"],
    dn: Float[Array, "d h w"],
    n: ScalarLike,
    dx: ScalarLike | Float[Array, "2"],
    thickness: ScalarLike,
    NA: ScalarLike | None = None,
):
    self.absorption = absorption
    self.dn = dn
    self.n = n
    self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
    self.thickness = thickness
    self._verify_scalar_thickness()
    self.NA = NA

ThinSample

Bases: Sample, Absorbing, Scattering

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

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.

Attributes:

Name	Type	Description
absorption	Float[Array, 'h w']	The sample's absorption defined as a 2D array of shape (height width).
dn	Float[Array, 'h w']	The sample's isotropic refractive index change as a 2D array of shape (height width).
dx	ScalarLike | Float[Array, 2]	The spacing of each pixel of the sample as a scalar in units of distance or as a 1D array of shape (2,) in the format (y x).
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).

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

class ThinSample(Sample, Absorbing, Scattering, strict=True):
    """
    Perturbs an incoming ``ScalarField`` as if it went through a thin sample
    object with a given ``absorption``, refractive index change ``dn`` from
    the refractive index of the medium and of a given ``thickness`` in the same
    units as the spectrum of the incoming ``ScalarField``.

    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.

    Attributes:
        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)``.
        dx: The spacing of each pixel of the sample as a scalar in units of
            distance or as a 1D array of shape `(2,)` in the format `(y x)`.
        thickness: Thickness in units of distance as a scalar or array
            broadcastable to ``(height width)`` (thickness at each sample
            location).
    """

    absorption: Float[Array, "h w"]
    dn: Float[Array, "h w"]
    dx: ScalarLike | Float[Array, "2"]
    thickness: ScalarLike | Float[Array, "h w"]

    def __init__(
        self,
        absorption: Float[Array, "h w"],
        dn: Float[Array, "h w"],
        dx: ScalarLike | Float[Array, "2"],
        thickness: ScalarLike | Float[Array, "h w"],
    ):
        assert_equal_shape(
            [absorption, dn],
            custom_message="Absorption and refractive index change arrays must have same shape",
        )
        self.absorption = absorption
        self.dn = dn
        self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
        self.thickness = thickness

    def __call__(self, field: ScalarField) -> ScalarField:
        self._verify_matching_spacing(field)
        return thin_sample(field, self.absorption, self.dn, self.thickness)

    @property
    def shape(self) -> tuple[int, int]:
        return self.absorption.shape

    @property
    def ndim(self) -> int:
        return self.absorption.ndim

    @property
    def grid(self) -> Array:
        return grid(self.shape, self.dx)

absorption = absorption instance-attribute

dn = dn instance-attribute

dx = _broadcast_dx_to_grid(dx, 1).squeeze() instance-attribute

grid property

ndim property

shape property

thickness = thickness instance-attribute

__call__(field)

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

def __call__(self, field: ScalarField) -> ScalarField:
    self._verify_matching_spacing(field)
    return thin_sample(field, self.absorption, self.dn, self.thickness)

__init__(absorption, dn, dx, thickness)

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

def __init__(
    self,
    absorption: Float[Array, "h w"],
    dn: Float[Array, "h w"],
    dx: ScalarLike | Float[Array, "2"],
    thickness: ScalarLike | Float[Array, "h w"],
):
    assert_equal_shape(
        [absorption, dn],
        custom_message="Absorption and refractive index change arrays must have same shape",
    )
    self.absorption = absorption
    self.dn = dn
    self.dx = _broadcast_dx_to_grid(dx, 1).squeeze()
    self.thickness = thickness

Sensors

__all__ = ['BasicSensor'] module-attribute

BasicSensor

Bases: Sensor

Produces an intensity image from an incoming Field with shot noise. Optionally, can also accept an intensity Array if input_spacing is specified.

Attributes:

Name	Type	Description
shape	tuple[int, int]	The shape in pixels of the sensor. Should be a tuple of the form height width.
spacing	ScalarLike | Float[Array, 2]	The pitch of the sensor pixels.
shot_noise_mode	Literal['approximate', 'poisson'] | None	What type of shot noise simulation to use. Defaults to None, in which case no shot noise is simulated.
resampling_method	str | None	What kind of sampling to use when resampling the incoming Field to the pitch of the sensor. Can be either 'pooling' which uses sum pooling (for downsampling only) or any method supported by jax.image.scale_and_translate ('linear', 'cubic', 'lanczos3', or 'lanczos5'). If None, then no resampling will occur.
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.
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.

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

class BasicSensor(Sensor):
    """
    Produces an intensity image from an incoming ``Field`` with shot noise.
    Optionally, can also accept an intensity Array if ``input_spacing`` is
    specified.

    Attributes:
        shape: The shape in pixels of the sensor. Should be a tuple of the form
            `height width`.
        spacing: The pitch of the sensor pixels.
        shot_noise_mode: What type of shot noise simulation to use. Defaults to
            ``None``, in which case no shot noise is simulated.
        resampling_method: What kind of sampling to use when resampling the
            incoming ``Field`` to the pitch of the sensor. Can be either
            `'pooling'` which uses sum pooling (for downsampling only) or any
            method supported by ``jax.image.scale_and_translate`` (`'linear'`,
            `'cubic'`, `'lanczos3'`, or `'lanczos5'`). If ``None``, then no
            resampling will occur.
        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``.
    """

    shape: tuple[int, int] = eqx.field(static=True)
    spacing: ScalarLike | Float[Array, "2"]
    shot_noise_mode: Literal["approximate", "poisson"] | None = eqx.field(static=True)
    resampling_method: str | None = eqx.field(static=True)
    resampler: Resampler | None = eqx.field(static=True)
    reduce_axis: int | None = eqx.field(static=True)
    reduce_parallel_axis_name: str | None = eqx.field(static=True)

    def __init__(
        self,
        shape: tuple[int, int],
        spacing: ScalarLike | Float[Array, "2"],
        shot_noise_mode: Literal["approximate", "poisson"] | None = None,
        resampling_method: str | None = "linear",
        reduce_axis: int | None = None,
        reduce_parallel_axis_name: str | None = None,
    ):
        self.shape = shape
        self.spacing = spacing
        self.shot_noise_mode = shot_noise_mode
        self.resampling_method = resampling_method
        if self.resampling_method is not None:
            self.resampler = init_plane_resample(
                self.shape, self.spacing, self.resampling_method
            )
        else:
            self.resampler = None
        self.reduce_axis = reduce_axis
        self.reduce_parallel_axis_name = reduce_parallel_axis_name

    def __call__(
        self,
        sensor_input: Field | Float[Array, "... h w"],
        input_spacing: ScalarLike | None = None,
        resample: bool = True,
        key: PRNGKeyArray | None = None,
    ) -> Array:
        """
        Resample the given ``sensor_input`` to the pixels of the sensor and
        potentially reduce the result and add shot noise.

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

        Args:
            sensor_input: The incoming ``Field`` or intensity ``Array`` of shape
                ``(... height width)``.
            input_spacing: The spacing of the input, only required if resampling
                is desired and the input is an ``Array``.
            resample: Whether to perform resampling or not. Only matters if
                ``resampling_method`` is not ``None``. Defaults to ``True``.
            key: If provided, will be used to generate the shot noise. If
                ``shot_noise_mode`` is not ``None``, this must be provided.
        """
        if resample and self.resampling_method is not None:
            resampler = self.resampler
        else:
            resampler = None
        if self.shot_noise_mode is not None:
            key is not None, ("Must provide a PRNGKey if shot_noise_mode is not None")
            _, noise_key = jax.random.split(key)
        else:
            noise_key = None
        return basic_sensor(
            sensor_input,
            self.shot_noise_mode,
            resampler,
            self.reduce_axis,
            self.reduce_parallel_axis_name,
            input_spacing=input_spacing,
            noise_key=noise_key,
        )

    def resample(self, resample_input: Array, input_spacing: ScalarLike) -> Array:
        """
        Resample the given ``resample_input`` to the pixels of the sensor.

        Args:
            resample_input: The ``Array`` to resample of shape `(... height
                width)` (resampling will occur along the last two dimensions,
                with all other dimensions treated as batch dimensions).
            input_spacing: The spacing of the input pixels as a scalar in units
                of distance.
        """
        if self.resampling_method is not None:
            resampler = self.resampler
            for i in range(resample_input.ndim - 2):
                resampler = jax.vmap(resampler, in_axes=(0, None))
            return resampler(resample_input, input_spacing)
        else:
            return resample_input

reduce_axis = reduce_axis class-attribute instance-attribute

reduce_parallel_axis_name = reduce_parallel_axis_name class-attribute instance-attribute

resampler = eqx.field(static=True) class-attribute instance-attribute

resampling_method = resampling_method class-attribute instance-attribute

shape = shape class-attribute instance-attribute

shot_noise_mode = shot_noise_mode class-attribute instance-attribute

spacing = spacing instance-attribute

__call__(sensor_input, input_spacing=None, resample=True, key=None)

Resample the given sensor_input to the pixels of the sensor and potentially reduce the result and add shot noise.

Warning

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

Parameters:

Name	Type	Description	Default
sensor_input	Field | Float[Array, '... h w']	The incoming Field or intensity Array of shape (... height width).	required
input_spacing	ScalarLike | None	The spacing of the input, only required if resampling is desired and the input is an Array.	None
resample	bool	Whether to perform resampling or not. Only matters if resampling_method is not None. Defaults to True.	True
key	PRNGKeyArray | None	If provided, will be used to generate the shot noise. If shot_noise_mode is not None, this must be provided.	None

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

def __call__(
    self,
    sensor_input: Field | Float[Array, "... h w"],
    input_spacing: ScalarLike | None = None,
    resample: bool = True,
    key: PRNGKeyArray | None = None,
) -> Array:
    """
    Resample the given ``sensor_input`` to the pixels of the sensor and
    potentially reduce the result and add shot noise.

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

    Args:
        sensor_input: The incoming ``Field`` or intensity ``Array`` of shape
            ``(... height width)``.
        input_spacing: The spacing of the input, only required if resampling
            is desired and the input is an ``Array``.
        resample: Whether to perform resampling or not. Only matters if
            ``resampling_method`` is not ``None``. Defaults to ``True``.
        key: If provided, will be used to generate the shot noise. If
            ``shot_noise_mode`` is not ``None``, this must be provided.
    """
    if resample and self.resampling_method is not None:
        resampler = self.resampler
    else:
        resampler = None
    if self.shot_noise_mode is not None:
        key is not None, ("Must provide a PRNGKey if shot_noise_mode is not None")
        _, noise_key = jax.random.split(key)
    else:
        noise_key = None
    return basic_sensor(
        sensor_input,
        self.shot_noise_mode,
        resampler,
        self.reduce_axis,
        self.reduce_parallel_axis_name,
        input_spacing=input_spacing,
        noise_key=noise_key,
    )

__init__(shape, spacing, shot_noise_mode=None, resampling_method='linear', reduce_axis=None, reduce_parallel_axis_name=None)

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

def __init__(
    self,
    shape: tuple[int, int],
    spacing: ScalarLike | Float[Array, "2"],
    shot_noise_mode: Literal["approximate", "poisson"] | None = None,
    resampling_method: str | None = "linear",
    reduce_axis: int | None = None,
    reduce_parallel_axis_name: str | None = None,
):
    self.shape = shape
    self.spacing = spacing
    self.shot_noise_mode = shot_noise_mode
    self.resampling_method = resampling_method
    if self.resampling_method is not None:
        self.resampler = init_plane_resample(
            self.shape, self.spacing, self.resampling_method
        )
    else:
        self.resampler = None
    self.reduce_axis = reduce_axis
    self.reduce_parallel_axis_name = reduce_parallel_axis_name

resample(resample_input, input_spacing)

Resample the given resample_input to the pixels of the sensor.

Parameters:

Name	Type	Description	Default
resample_input	Array	The Array to resample of shape (... height width) (resampling will occur along the last two dimensions, with all other dimensions treated as batch dimensions).	required
input_spacing	ScalarLike	The spacing of the input pixels as a scalar in units of distance.	required

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

def resample(self, resample_input: Array, input_spacing: ScalarLike) -> Array:
    """
    Resample the given ``resample_input`` to the pixels of the sensor.

    Args:
        resample_input: The ``Array`` to resample of shape `(... height
            width)` (resampling will occur along the last two dimensions,
            with all other dimensions treated as batch dimensions).
        input_spacing: The spacing of the input pixels as a scalar in units
            of distance.
    """
    if self.resampling_method is not None:
        resampler = self.resampler
        for i in range(resample_input.ndim - 2):
            resampler = jax.vmap(resampler, in_axes=(0, None))
        return resampler(resample_input, input_spacing)
    else:
        return resample_input

Propagation

__all__ = ['Propagate', 'KernelPropagate'] module-attribute

KernelPropagate

Bases: Module

Free space propagation with a precomputed propagation kernel.

This element takes a Field as input and outputs a Field that has been propagated by a distance that is already defined by a propagation kernel. Propagation is performed via Fourier convolution (two Fourier transforms).

Attributes:

Name	Type	Description
propagator	Array	The propagation kernel to use (can be created using e.g. compute_asm_propagator).
pad_width	int	The padding for propagation (will be used as both height and width padding). To automatically calculate the padding, use padding calculation functions from chromatix.functional. This must be passed outside of a jax.jit. Defaults to 0 (no padding), which will cause circular convolutions (edge artifacts) when propagating.
cval	ScalarLike	The value to pad with if pad_width is greater than 0. Defaults to 0.
mode	Literal['full', 'same']	Defines the cropping of the output if the method is NOT "transform". Defaults to "same", which returns a Field of the same shape, unlike the functional methods.

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

class KernelPropagate(eqx.Module):
    """
    Free space propagation with a precomputed propagation kernel.

    This element takes a ``Field`` as input and outputs a ``Field`` that has
    been propagated by a distance that is already defined by a propagation
    kernel. Propagation is performed via Fourier convolution (two Fourier
    transforms).

    Attributes:
        propagator: The propagation kernel to use (can be created using e.g.
            [`compute_asm_propagator`](functional.md#chromatix.functional.propagation.compute_asm_propagator)).
        pad_width: The padding for propagation (will be used as both height and
            width padding). To automatically calculate the padding, use padding
            calculation functions from  ``chromatix.functional``. This must be
            passed outside of a ``jax.jit``. Defaults to 0 (no padding), which
            will cause circular convolutions (edge artifacts) when propagating.
        cval: The value to pad with if ``pad_width`` is greater than 0. Defaults
            to 0.
        mode: Defines the cropping of the output if the method is NOT
            "transform". Defaults to "same", which returns a ``Field`` of the
            same shape, unlike the functional methods.
    """

    propagator: Array
    pad_width: int = eqx.field(static=True)
    cval: ScalarLike
    mode: Literal["full", "same"] = eqx.field(static=True)

    def __init__(
        self,
        propagator: Array,
        pad_width: int = 0,
        cval: float = 0.0,
        mode: Literal["full", "same"] = "same",
    ):
        self.propagator = propagator
        self.pad_width = pad_width
        self.cval = cval
        self.mode = mode

    def __call__(self, field: Field) -> Field:
        field = pad(field, self.pad_width, cval=self.cval)
        field = kernel_propagate(field, self.propagator)
        if self.mode == "same":
            field = crop(field, self.pad_width)
        return field

cval = cval instance-attribute

mode = mode class-attribute instance-attribute

pad_width = pad_width class-attribute instance-attribute

propagator = propagator instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    field = pad(field, self.pad_width, cval=self.cval)
    field = kernel_propagate(field, self.propagator)
    if self.mode == "same":
        field = crop(field, self.pad_width)
    return field

__init__(propagator, pad_width=0, cval=0.0, mode='same')

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

def __init__(
    self,
    propagator: Array,
    pad_width: int = 0,
    cval: float = 0.0,
    mode: Literal["full", "same"] = "same",
):
    self.propagator = propagator
    self.pad_width = pad_width
    self.cval = cval
    self.mode = mode

Propagate

Bases: Module

Free space propagation that can be placed after or between other elements.

This element takes a Field as input and outputs a Field that has been propagated by a distance z. Initialization of this element requires a Field to construct the propagation kernels.

For example, this element can be constructed as:

from chromatix.elements import Propagate
from chromatix.functional import plane_wave
field = plane_wave(...)
propagation = Propagate(field, z=1e3, n=1.33, method='asm')

Warning

The underlying propagation method now defaults to method=asm, bandlimit=True and remove_evanescent=False, which corresponds to bandlimited angular spectrum (BLAS) as proposed in "Band-Limited Angular Spectrum Method for Numerical Simulation of Free-Space Propagation in Far and Near Fields" (Matsumina et al., 2009).

Warning

By default this element caches the propagation kernel using the option cache_propagator. If you would like to have a propagation kernel with your own initialization, see KernelPropagate which accepts a custom propagator.

Attributes:

Name	Type	Description
z	ScalarLike | Float[Array, z]	How far to propagate as a scalar value or a 1D array in units of distance.
n	ScalarLike	Refractive index.
pad_width	int	The padding for propagation (will be used as both height and width padding). To automatically calculate the padding, use padding calculation functions from chromatix.functional. This must be passesd outside of a jax.jit. Defaults to 0 (no padding), which will cause circular convolutions (edge artifacts) when propagating.
cval	ScalarLike	The value to pad with if pad_width is greater than 0. Defaults to 0.
kykx	Float[Array, 2] | 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).
method	Literal['transform', 'transfer', 'asm']	The propagation method, which can be "transform", "transfer", or "asm". Defaults to "asm", which is propagation as a Fourier convolution (two Fourier transforms) without the Fresnel approximation.
bandlimit	bool	Whether to bandlimit the field before propagation for "asm". Defaults to True.
remove_evanescent	bool	Whether to remove evanescent waves when using the "asm" method. Defaults to False.
mode	Literal['full', 'same']	Defines the cropping of the output if the method is NOT "transform". Defaults to "same", which returns a Field of the same shape, unlike the functional methods.
cache_propagator	bool	Whether to precompute and store the propagation kernel or not. If True, none of the other attributes will be used after initializing the propagation kernel and the propagation kernel will only be computed once at initialization. Defaults to True.

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

class Propagate(eqx.Module):
    """
    Free space propagation that can be placed after or between other elements.

    This element takes a ``Field`` as input and outputs a ``Field`` that has
    been propagated by a distance ``z``. Initialization of this element requires a
    ``Field`` to construct the propagation kernels.

    For example, this element can be constructed as:

    ```python
    from chromatix.elements import Propagate
    from chromatix.functional import plane_wave
    field = plane_wave(...)
    propagation = Propagate(field, z=1e3, n=1.33, method='asm')
    ```

    !!! warning
        The underlying propagation method now defaults to ``method=asm``,
        ``bandlimit=True`` and ``remove_evanescent=False``, which corresponds
        to bandlimited angular spectrum (BLAS) as proposed in "Band-Limited
        Angular Spectrum Method for Numerical Simulation of Free-Space
        Propagation in Far and Near Fields" (Matsumina et al., 2009).

    !!! warning
        By default this element caches the propagation kernel using the option
        ``cache_propagator``. If you would like to have a propagation kernel
        with your own initialization, see ``KernelPropagate`` which accepts a
        custom ``propagator``.

    Attributes:
        z: How far to propagate as a scalar value or a 1D array in units of distance.
        n: Refractive index.
        pad_width: The padding for propagation (will be used as both height and
            width padding). To automatically calculate the padding, use padding
            calculation functions from  ``chromatix.functional``. This must be
            passesd outside of a ``jax.jit``. Defaults to 0 (no padding), which
            will cause circular convolutions (edge artifacts) when propagating.
        cval: The value to pad with if ``pad_width`` is greater than 0. Defaults
            to 0.
        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).
        method: The propagation method, which can be "transform", "transfer",
            or "asm". Defaults to "asm", which is propagation as a Fourier
            convolution (two Fourier transforms) without the Fresnel
            approximation.
        bandlimit: Whether to bandlimit the field before propagation for "asm".
            Defaults to True.
        remove_evanescent: Whether to remove evanescent waves when using the
            "asm" method. Defaults to False.
        mode: Defines the cropping of the output if the method is NOT
            "transform". Defaults to "same", which returns a ``Field`` of the
            same shape, unlike the functional methods.
        cache_propagator: Whether to precompute and store the propagation kernel
            or not. If ``True``, none of the other attributes will be used after
            initializing the propagation kernel and the propagation kernel will
            only be computed once at initialization. Defaults to ``True``.
    """

    z: ScalarLike | Float[Array, "z"]
    n: ScalarLike
    pad_width: int = eqx.field(static=True)
    cval: ScalarLike
    kykx: Float[Array, "2"] | tuple[float, float]
    method: Literal["transform", "transfer", "asm"] = eqx.field(static=True)
    bandlimit: bool = eqx.field(static=True)
    remove_evanescent: bool = eqx.field(static=True)
    mode: Literal["full", "same"] = eqx.field(static=True)
    cache_propagator: bool = eqx.field(static=True)
    propagator: Array | None

    def __init__(
        self,
        field: Field,
        z: ScalarLike | Float[Array, "z"],
        n: ScalarLike,
        pad_width: int = 0,
        cval: float = 0,
        kykx: Float[Array, "2"] | tuple[float, float] = (0.0, 0.0),
        method: Literal["transform", "transfer", "asm"] = "asm",
        bandlimit: bool = True,
        remove_evanescent: bool = False,
        mode: Literal["full", "same"] = "same",
        cache_propagator: bool = True,
    ):
        self.z = z
        self.n = n
        self.pad_width = pad_width
        self.cval = cval
        self.kykx = kykx
        self.method = method
        self.bandlimit = bandlimit
        self.remove_evanescent = remove_evanescent
        self.mode = mode
        self.cache_propagator = cache_propagator
        if self.cache_propagator:
            field = pad(field, self.pad_width, cval=self.cval)
            propagator_args = (
                field,
                self.z,
                self.n,
                self.kykx,
            )
            if self.method == "transfer":
                self.propagator = compute_transfer_propagator(*propagator_args)
            elif self.method == "asm":
                self.propagator = compute_asm_propagator(
                    *propagator_args,
                    remove_evanescent=self.remove_evanescent,
                    bandlimit=self.bandlimit,
                )
            else:
                raise ValueError(
                    "Can only cache kernel for 'transfer' or 'asm' methods"
                )
        else:
            self.propagator = None

    def __call__(self, field: Field) -> Field:
        if self.cache_propagator:
            field = pad(field, self.pad_width, cval=self.cval)
            field = kernel_propagate(field, self.propagator)
            if self.mode == "same":
                field = crop(field, self.pad_width)
            return field
        if self.method == "transform":
            return transform_propagate(
                field,
                self.z,
                self.n,
                pad_width=self.pad_width,
                cval=self.cval,
            )
        elif self.method == "transfer":
            return transfer_propagate(
                field,
                self.z,
                self.n,
                pad_width=self.pad_width,
                cval=self.cval,
                kykx=self.kykx,
                mode=self.mode,
            )
        elif self.method == "asm":
            return asm_propagate(
                field,
                self.z,
                self.n,
                pad_width=self.pad_width,
                cval=self.cval,
                kykx=self.kykx,
                mode=self.mode,
                remove_evanescent=self.remove_evanescent,
                bandlimit=self.bandlimit,
            )
        else:
            raise ValueError("Method must be one of: 'transform', 'transfer', 'asm'")

bandlimit = bandlimit class-attribute instance-attribute

cache_propagator = cache_propagator class-attribute instance-attribute

cval = cval instance-attribute

kykx = kykx instance-attribute

method = method class-attribute instance-attribute

mode = mode class-attribute instance-attribute

n = n instance-attribute

pad_width = pad_width class-attribute instance-attribute

propagator instance-attribute

remove_evanescent = remove_evanescent class-attribute instance-attribute

z = z instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    if self.cache_propagator:
        field = pad(field, self.pad_width, cval=self.cval)
        field = kernel_propagate(field, self.propagator)
        if self.mode == "same":
            field = crop(field, self.pad_width)
        return field
    if self.method == "transform":
        return transform_propagate(
            field,
            self.z,
            self.n,
            pad_width=self.pad_width,
            cval=self.cval,
        )
    elif self.method == "transfer":
        return transfer_propagate(
            field,
            self.z,
            self.n,
            pad_width=self.pad_width,
            cval=self.cval,
            kykx=self.kykx,
            mode=self.mode,
        )
    elif self.method == "asm":
        return asm_propagate(
            field,
            self.z,
            self.n,
            pad_width=self.pad_width,
            cval=self.cval,
            kykx=self.kykx,
            mode=self.mode,
            remove_evanescent=self.remove_evanescent,
            bandlimit=self.bandlimit,
        )
    else:
        raise ValueError("Method must be one of: 'transform', 'transfer', 'asm'")

__init__(field, z, n, pad_width=0, cval=0, kykx=(0.0, 0.0), method='asm', bandlimit=True, remove_evanescent=False, mode='same', cache_propagator=True)

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

def __init__(
    self,
    field: Field,
    z: ScalarLike | Float[Array, "z"],
    n: ScalarLike,
    pad_width: int = 0,
    cval: float = 0,
    kykx: Float[Array, "2"] | tuple[float, float] = (0.0, 0.0),
    method: Literal["transform", "transfer", "asm"] = "asm",
    bandlimit: bool = True,
    remove_evanescent: bool = False,
    mode: Literal["full", "same"] = "same",
    cache_propagator: bool = True,
):
    self.z = z
    self.n = n
    self.pad_width = pad_width
    self.cval = cval
    self.kykx = kykx
    self.method = method
    self.bandlimit = bandlimit
    self.remove_evanescent = remove_evanescent
    self.mode = mode
    self.cache_propagator = cache_propagator
    if self.cache_propagator:
        field = pad(field, self.pad_width, cval=self.cval)
        propagator_args = (
            field,
            self.z,
            self.n,
            self.kykx,
        )
        if self.method == "transfer":
            self.propagator = compute_transfer_propagator(*propagator_args)
        elif self.method == "asm":
            self.propagator = compute_asm_propagator(
                *propagator_args,
                remove_evanescent=self.remove_evanescent,
                bandlimit=self.bandlimit,
            )
        else:
            raise ValueError(
                "Can only cache kernel for 'transfer' or 'asm' methods"
            )
    else:
        self.propagator = None

Convenience

__all__ = ['Flip', 'ScaleAndBias', 'Binarize', 'Quantize'] module-attribute

Binarize

Bases: Module

This element binarizes the incoming Field.

See chromatix.ops.quantization.binarize for more details.

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

class Binarize(eqx.Module):
    """
    This element binarizes the incoming ``Field``.

    See ``chromatix.ops.quantization.binarize`` for more details.
    """

    threshold: float

    def __init__(self, threshold: float = 0.5):
        self.threshold = threshold

    def __call__(self, field: Field) -> Field:
        return field.replace(u=binarize(field.u, self.threshold))

threshold = threshold instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    return field.replace(u=binarize(field.u, self.threshold))

__init__(threshold=0.5)

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

def __init__(self, threshold: float = 0.5):
    self.threshold = threshold

Flip

Bases: Module

This element flips the incoming Field upside down.

This is useful for situations where an upside down image is formed.

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

class Flip(eqx.Module):
    """
    This element flips the incoming ``Field`` upside down.

    This is useful for situations where an upside down image is formed.
    """

    def __call__(self, field: Field) -> Field:
        # TODO: Add in support for both axes?
        return field.replace(u=jnp.flip(field.u, axis=field.spatial_dims[0]))

__call__(field)

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

def __call__(self, field: Field) -> Field:
    # TODO: Add in support for both axes?
    return field.replace(u=jnp.flip(field.u, axis=field.spatial_dims[0]))

Quantize

Bases: Module

This element quantizes the incoming Field to the given bit_depth.

See chromatix.ops.quantization.quantize for more details.

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

class Quantize(eqx.Module):
    """
    This element quantizes the incoming ``Field`` to the given ``bit_depth``.

    See ``chromatix.ops.quantization.quantize`` for more details.
    """

    bit_depth: int = eqx.field(static=True)
    range: Float[Array, "2"] | tuple[int, int] | None = eqx.field(static=True)

    def __init__(
        self, bit_depth: int, range: Float[Array, "2"] | tuple[int, int] | None = None
    ):
        self.bit_depth = bit_depth
        self.range = range

    def __call__(self, field: Field) -> Field:
        return field.replace(u=quantize(field.u, self.bit_depth, self.range))

bit_depth = bit_depth class-attribute instance-attribute

range = range class-attribute instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    return field.replace(u=quantize(field.u, self.bit_depth, self.range))

__init__(bit_depth, range=None)

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

def __init__(
    self, bit_depth: int, range: Float[Array, "2"] | tuple[int, int] | None = None
):
    self.bit_depth = bit_depth
    self.range = range

ScaleAndBias

Bases: Module

This element applies a scale and bias to the incoming Field.

The scale and bias can either be scalars or Arrays broadcastable to the shape of the incoming Field.

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

class ScaleAndBias(eqx.Module):
    """
    This element applies a ``scale`` and ``bias`` to the incoming ``Field``.

    The ``scale`` and ``bias`` can either be scalars or ``Array``s
    broadcastable to the shape of the incoming ``Field``.
    """

    bias: ScalarLike
    scale: ScalarLike

    def __init(self, bias: ScalarLike, scale: ScalarLike):
        self.bias = bias
        self.scale = scale

    def __call__(self, field: Field) -> Field:
        return (field + self.bias) * self.scale

bias instance-attribute

scale instance-attribute

__call__(field)

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

def __call__(self, field: Field) -> Field:
    return (field + self.bias) * self.scale

__init(bias, scale)

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

def __init(self, bias: ScalarLike, scale: ScalarLike):
    self.bias = bias
    self.scale = scale