Utilities

General

center_crop(u, crop_length)

Symmetrically crops u with lengths specified per axis in crop_length, which should be iterable with same size as u.ndims.

Source code in src/chromatix/utils/utils.py

def center_crop(u: Array, crop_length: Sequence[int]) -> Array:
    """
    Symmetrically crops ``u`` with lengths specified per axis in
    ``crop_length``, which should be iterable with same size as ``u.ndims``.
    """
    crop_length = [0 if length is None else length for length in crop_length]
    crop = tuple([slice(n, size - n) for size, n in zip(u.shape, crop_length)])
    return u[crop]

center_pad(u, pad_width, cval=0)

Symmetrically pads u with lengths specified per axis in pad_width, which should be an iterable of integers and have the same length as u.ndims.

Source code in src/chromatix/utils/utils.py

def center_pad(u: ArrayLike, pad_width: Sequence[int], cval: float = 0) -> Array:
    """
    Symmetrically pads ``u`` with lengths specified per axis in ``pad_width``,
    which should be an iterable of integers and have the same length as
    ``u.ndims``.
    """
    pad = [(n, n) for n in pad_width]
    return jnp.pad(u, pad, constant_values=cval)

create_grid(shape, spacing)

Creates a 2D grid of vertical and horizontal coordinates with the specified shape and spacing, with the origin in the center of the grid.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape of the grid, described as a tuple of integers of the form (H W).	required
spacing	ScalarLike	The spacing of each pixel in the grid, either a single float for square pixels or an array of shape (2 1) for non-square pixels.	required

Returns: The grid as an array of shape (2 H W).

Source code in src/chromatix/utils/utils.py

def create_grid(shape: tuple[int, int], spacing: ScalarLike) -> Array:
    """
    Creates a 2D grid of vertical and horizontal coordinates with the specified
    ``shape`` and ``spacing``, with the origin in the center of the grid.

    Args:
        shape: The shape of the grid, described as a tuple of
            integers of the form ``(H W)``.
        spacing: The spacing of each pixel in the grid, either a single float
            for square pixels or an array of shape `(2 1)` for non-square
            pixels.
    Returns:
        The grid as an array of shape ``(2 H W)``.
    """
    half_size = jnp.array(shape) / 2
    spacing = jnp.atleast_1d(spacing)
    if spacing.size == 1:
        spacing = jnp.concatenate([spacing, spacing])
    assert spacing.size == 2, "Spacing must be either single float or have shape (2,)"
    spacing = rearrange(spacing, "d -> d 1 1", d=2)
    # @copypaste(Field): We must use meshgrid instead of mgrid here
    # in order to be jittable
    grid = jnp.meshgrid(
        jnp.linspace(-half_size[0], half_size[0] - 1, num=shape[0]) + 0.5,
        jnp.linspace(-half_size[1], half_size[1] - 1, num=shape[1]) + 0.5,
        indexing="ij",
    )
    grid = spacing * jnp.array(grid)
    return grid

gaussian_kernel(sigma, truncate=4.0, shape=None)

Creates ND Gaussian kernel of given sigma.

If shape is not provided, then the shape of the kernel is automatically calculated using the given truncation (the same truncation for each dimension) and sigma. The number of dimensions is determined by the length of sigma, which should be a 1D array.

If shape is provided, then truncate is ignored and the result will have the provided shape. The provided shape must be odd in all dimensions to ensure that there is a center pixel.

Parameters:

Name	Type	Description	Default
sigma	Sequence[float]	A 1D array whose length is the number of dimensions specifying the standard deviation of the Gaussian distribution in each dimension.	required
truncate	float	If shape is not provided, then this float is the number of standard deviations for which to calculate the Gaussian. This is then used to determine the shape of the kernel in each dimension.	4.0
shape	Sequence[int] | None	If provided, determines the shape of the kernel. This will cause truncate to be ignored.	None

Returns:

Type	Description
Array	The ND Gaussian kernel.

Source code in src/chromatix/utils/utils.py

def gaussian_kernel(
    sigma: Sequence[float], truncate: float = 4.0, shape: Sequence[int] | None = None
) -> Array:
    """
    Creates ND Gaussian kernel of given ``sigma``.

    If ``shape`` is not provided, then the shape of the kernel is automatically
    calculated using the given truncation (the same truncation for each
    dimension) and ``sigma``. The number of dimensions is determined by the
    length of ``sigma``, which should be a 1D array.

    If ``shape`` is provided, then ``truncate`` is ignored and the result will
    have the provided ``shape``. The provided ``shape`` must be odd in all
    dimensions to ensure that there is a center pixel.

    Args:
        sigma: A 1D array whose length is the number of dimensions specifying
            the standard deviation of the Gaussian distribution in each
            dimension.
        truncate: If ``shape`` is not provided, then this float is the number
            of standard deviations for which to calculate the Gaussian. This is
            then used to determine the shape of the kernel in each dimension.
        shape: If provided, determines the ``shape`` of the kernel. This will
            cause ``truncate`` to be ignored.

    Returns:
        The ND Gaussian kernel.
    """
    _sigma = np.atleast_1d(np.array(sigma))
    if shape is not None:
        _shape = np.atleast_1d(np.array(shape))
        assert np.all(_shape % 2 != 0), "Shape must be odd in all dimensions"
        radius = ((_shape - 1) / 2).astype(np.int16)
    else:
        radius = (truncate * _sigma + 0.5).astype(np.int16)

    x = jnp.mgrid[tuple(slice(-r, r + 1) for r in radius)]
    phi = jnp.exp(-0.5 * jnp.sum((x.T / _sigma) ** 2, axis=-1))  # type: ignore
    return phi / phi.sum()

grid_spatial_to_pupil(grid, f, NA, n)

Source code in src/chromatix/utils/utils.py

def grid_spatial_to_pupil(
    grid: Array, f: ScalarLike, NA: ScalarLike, n: ScalarLike
) -> Array:
    R = f * NA / n  # pupil radius
    return grid / R

l1_norm(a, axis=-1)

Sum absolute value, i.e. |x| + |y|.

Source code in src/chromatix/utils/utils.py

def l1_norm(a: Array, axis: int | tuple[int, ...] = -1) -> Array:
    """Sum absolute value, i.e. `|x| + |y|`."""
    return jnp.sum(jnp.abs(a), axis=axis)

l2_norm(a, axis=-1)

Square root of l2_sq_norm, i.e. sqrt(x**2 + y**2).

Source code in src/chromatix/utils/utils.py

def l2_norm(a: Array, axis: int | tuple[int, ...] = -1) -> Array:
    """Square root of ``l2_sq_norm``, i.e. `sqrt(x**2 + y**2)`."""
    return jnp.sqrt(jnp.sum(a**2, axis=axis))

l2_sq_norm(a, axis=-1)

Sum of squares, i.e. x**2 + y**2.

Source code in src/chromatix/utils/utils.py

def l2_sq_norm(a: Array, axis: int | tuple[int, ...] = -1) -> Array:
    """Sum of squares, i.e. `x**2 + y**2`."""
    return jnp.sum(a**2, axis=axis)

linf_norm(a, axis=-1)

Max absolute value, i.e. max(|x|, |y|).

Source code in src/chromatix/utils/utils.py

def linf_norm(a: Array, axis: int | tuple[int, ...] = -1) -> Array:
    """Max absolute value, i.e. `max(|x|, |y|)`."""
    return jnp.max(jnp.abs(a), axis=axis)

matvec(x, y)

Implements batched matrix - vector multiplication. Mostly used in polarization calculations. Example [..., N, M] x [...., M] -> [...., N]

Source code in src/chromatix/utils/utils.py

def matvec(x: Array, y: Array) -> Array:
    """Implements batched matrix - vector multiplication.
    Mostly used in polarization calculations.
    Example [..., N, M] x [...., M] -> [...., N]"""
    return jnp.matmul(x, y[..., None]).squeeze(-1)

next_order(val)

Source code in src/chromatix/utils/utils.py

def next_order(val: int) -> int:
    return int(2 ** np.ceil(np.log2(val)))

outer(x, y, in_axis=-1)

Calculates batched outer product (Numpy flattens input matrices) Includes additional in_axis for which axis to use. Output axes will always be last two.

Source code in src/chromatix/utils/utils.py

def outer(x: Array, y: Array, in_axis: int = -1) -> Array:
    """Calculates batched outer product (Numpy flattens input matrices)
    Includes additional in_axis for which axis to use.
    Output axes will always be last two.
    """
    _x = jnp.moveaxis(x, in_axis, -1)
    _y = jnp.moveaxis(y, in_axis, -1)
    return _x[..., None, :] * _y[..., :, None]

rotate_grid(grid, rotation)

Rotates a 2D grid (an array of shape (2 H W)) by rotation radians. Positive rotations are assumed to be in the counter-clockwise direction.

Source code in src/chromatix/utils/utils.py

def rotate_grid(grid: Array, rotation: ScalarLike) -> Array:
    """
    Rotates a 2D grid (an array of shape ``(2 H W)``) by ``rotation`` radians.
    Positive rotations are assumed to be in the counter-clockwise direction.
    """
    rotation = jnp.array(
        [
            [jnp.cos(rotation), -jnp.sin(rotation)],
            [jnp.sin(rotation), jnp.cos(rotation)],
        ]
    )
    grid = jnp.einsum("ij, ihw -> jhw", rotation, grid)
    return grid

sigmoid_taper(shape, width)

Source code in src/chromatix/utils/utils.py

def sigmoid_taper(shape: tuple[int, int], width: float) -> Array:
    dist = jnp.asarray(
        distance_transform_edt(np.pad(np.ones((shape[0] - 2, shape[1] - 2)), 1))
    )
    taper = 2 * (nn.sigmoid(dist / width) - 0.5)  # type: ignore - it's an array!
    return taper

Initializers

__all__ = ['axicon_phase', 'flat_phase', 'microlens_array_amplitude_and_phase', 'hexagonal_microlens_array_amplitude_and_phase', 'rectangular_microlens_array_amplitude_and_phase', 'circular_phase', 'linear_phase', 'sawtooth_phase', 'sinusoid_phase', 'potato_chip', 'seidel_aberrations', 'zernike_aberrations', 'defocused_ramps'] module-attribute

axicon_phase(shape, spacing, wavelength, n_axicon, slope_angle, n_medium=1.0)

Source code in src/chromatix/utils/initializers.py

def axicon_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n_axicon: ScalarLike,
    slope_angle: ScalarLike,
    n_medium: ScalarLike = 1.0,
) -> Array:
    dn = jnp.asarray(n_axicon - n_medium)
    grid = create_grid(shape, spacing)
    thickness = jnp.sin(slope_angle) * l2_norm(grid, axis=0)
    phase = 2.0 * jnp.pi * dn * thickness / jnp.asarray(wavelength)
    return phase

circular_phase(shape, spacing, shift, w)

Source code in src/chromatix/utils/initializers.py

def circular_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    shift: ScalarLike,
    w: ScalarLike,
) -> Array:
    grid = create_grid(shape, spacing)
    phase = l2_sq_norm(grid, axis=0)
    phase = phase * (phase <= (w / 2) ** 2)
    phase = jnp.asarray(shift) * phase
    return phase

defocused_ramps(shape, spacing, wavelength, n, f, NA, num_ramps=6, delta=[2374.0] * 6, defocus=[-50.0, 150.0, -100.0, 50.0, -150.0, 100.0])

Computes the "defocused ramps" phase mask as described in [1].

This phase mask is intended to be used in a 4f microscope to produce a number of "pencil" beams in the resulting PSF. The resulting PSF produces multiple subimages of the sample on the camera that are projections of the sample along different angles and through different axial ranges, intended to be used for 3D snapshot microscopy.

The name describes the fact that the phase mask consists of multiple phase ramps around a central flat region, combined with a bowl of defocus within each phase ramp to defocus the pencil beam that results from that arm of the phase mask.

[1]: Deb et al. "FourierNets enable the design of highly non-local optical encoders for computational imaging." NeurIPS, 2022.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape of the phase mask, described as a tuple of integers of the form (H W).	required
spacing	ScalarLike	The spacing of each pixel in the phase mask.	required
wavelength	ScalarLike	The wavelength to compute the phase mask for.	required
n	ScalarLike	Refractive index.	required
f	ScalarLike	The focal distance (should be in same units as wavelength).	required
NA	ScalarLike	The numerical aperture. Phase will be 0 outside of this NA.	required
num_ramps	int	Sets the number of "pencil" beams or "ramps". The number of pencil beams will be num_ramps + 1, because of the central flat region of the phase mask.	6
delta	Sequence[float]	Controls the "slope" of each phase ramp. Higher values move the resulting pencil further away from the center of the field.	[2374.0] * 6
defocus	Sequence[float]	Controls the defocus of each pencil axially (should be in same units as wavelength).	[-50.0, 150.0, -100.0, 50.0, -150.0, 100.0]

Source code in src/chromatix/utils/initializers.py

def defocused_ramps(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    f: ScalarLike,
    NA: ScalarLike,
    num_ramps: int = 6,
    delta: Sequence[float] = [2374.0] * 6,
    defocus: Sequence[float] = [-50.0, 150.0, -100.0, 50.0, -150.0, 100.0],
) -> Array:
    """
    Computes the "defocused ramps" phase mask as described in [1].

    This phase mask is intended to be used in a 4f microscope to produce a
    number of "pencil" beams in the resulting PSF. The resulting PSF produces
    multiple subimages of the sample on the camera that are projections of the
    sample along different angles and through different axial ranges, intended
    to be used for 3D snapshot microscopy.

    The name describes the fact that the phase mask consists of multiple phase
    ramps around a central flat region, combined with a bowl of defocus within
    each phase ramp to defocus the pencil beam that results from that arm of
    the phase mask.

    [1]: Deb et al. "FourierNets enable the design of highly non-local optical
        encoders for computational imaging." NeurIPS, 2022.

    Args:
        shape: The shape of the phase mask, described as a tuple of
            integers of the form (H W).
        spacing: The spacing of each pixel in the phase mask.
        wavelength: The wavelength to compute the phase mask for.
        n: Refractive index.
        f: The focal distance (should be in same units as ``wavelength``).
        NA: The numerical aperture. Phase will be 0 outside of this NA.
        num_ramps: Sets the number of "pencil" beams or "ramps". The number of
            pencil beams will be ``num_ramps + 1``, because of the central
            flat region of the phase mask.
        delta: Controls the "slope" of each phase ramp. Higher values move the
            resulting pencil further away from the center of the field.
        defocus: Controls the defocus of each pencil axially (should be in
            same units as ``wavelength``).
    """
    grid = create_grid(shape, spacing)
    # Normalize coordinates from -1 to 1 within radius R
    grid = grid_spatial_to_pupil(grid, f, NA, n)
    l2_sq_grid = jnp.sum(grid**2, axis=0)
    theta = jnp.arctan2(*grid)
    edges = jnp.linspace(-jnp.pi, jnp.pi, num_ramps + 1)
    centers = (edges[:-1] + edges[1:]) / 2
    flat_region_edge = (num_ramps + 1) ** -0.5
    defocus_center = (flat_region_edge + 1) / 2.0
    phase = jnp.zeros(shape)

    def ramp(center, theta_bounds, delta_ramp, ramp_defocus):
        # Calculate distances along and across current ramp
        ramp_parallel_distance = grid[0] * jnp.sin(center) + grid[1] * jnp.cos(center)
        ramp_perpendicular_distance = grid[0] * jnp.cos(center) - grid[1] * jnp.sin(
            center
        )
        # Select coordinates for current ramp
        ramp_mask = (
            (theta >= theta_bounds[0])
            & (theta < theta_bounds[1])
            & (ramp_parallel_distance > flat_region_edge)
            & (l2_sq_grid < 1)
        )
        # Create ramp
        phase = ramp_mask * delta_ramp * ramp_perpendicular_distance
        # Create defocus within ramp
        ramp_quadratic = (grid[1] - jnp.cos(center) * defocus_center) ** 2 + (
            grid[0] - jnp.sin(center) * defocus_center
        ) ** 2
        phase += ramp_mask * (ramp_defocus * ramp_quadratic)
        phase -= ramp_mask * jnp.where(ramp_mask > 0, phase, 0).mean()  # type: ignore
        return phase

    for ramp_idx in range(num_ramps):
        phase += ramp(
            centers[ramp_idx],
            edges[ramp_idx : (ramp_idx + 2)],
            delta[ramp_idx],
            defocus[ramp_idx],
        )
    phase *= l2_sq_grid < 1
    return phase

flat_phase(shape, *args, value=0.0)

Computes a flat mask (one with constant value).

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape of the mask, described as a tuple of integers of the form (H W).	required
value	ScalarLike	The constant value to use for the mask, defaults to 0.	0.0

Source code in src/chromatix/utils/initializers.py

def flat_phase(shape: tuple[int, int], *args, value: ScalarLike = 0.0) -> Array:
    """
    Computes a flat mask (one with constant value).

    Args:
        shape: The shape of the mask, described as a tuple of
            integers of the form (H W).
        value: The constant value to use for the mask, defaults to 0.
    """
    return jnp.full(shape, value)

hexagonal_microlens_array_amplitude_and_phase(shape, spacing, wavelength, n, f, num_lenses_per_side, radius, separation)

Source code in src/chromatix/utils/initializers.py

@no_type_check
def hexagonal_microlens_array_amplitude_and_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    f: ScalarLike,
    num_lenses_per_side: ScalarLike,
    radius: ScalarLike,
    separation: ScalarLike,
) -> tuple[Array, Array]:
    hex_distance = num_lenses_per_side - 1
    unit_hex_coordinates = []
    q_basis = np.array([0, 1])
    r_basis = np.array([np.sqrt(3) / 2, 1 / 2])
    for q in range(-hex_distance, hex_distance + 1):
        for r in range(
            max(-hex_distance, -q - hex_distance),
            min(hex_distance, -q + hex_distance) + 1,
        ):
            unit_hex_coordinates.append(q_basis * q + r_basis * r)
    unit_hex_coordinates = np.array(unit_hex_coordinates).T
    hex_coordinates = unit_hex_coordinates * separation
    return microlens_array_amplitude_and_phase(
        shape,
        spacing,
        wavelength,
        n,
        jnp.ones(hex_coordinates.shape[1]) * f,
        hex_coordinates,
        jnp.ones(hex_coordinates.shape[1]) * radius,
    )

linear_phase(shape, spacing, wavelength, n_mask, max_thickness, rotation=0.0, n_medium=1.0)

Source code in src/chromatix/utils/initializers.py

def linear_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n_mask: ScalarLike,
    max_thickness: ScalarLike,
    rotation: ScalarLike = 0.0,
    n_medium: ScalarLike = 1.0,
) -> Array:
    dn = jnp.asarray(n_mask - n_medium)
    grid = create_grid(shape, spacing)
    grid = rotate_grid(grid, rotation)
    phase = grid[1] - grid[1].min()
    phase = (
        2
        * jnp.pi
        * dn
        * jnp.asarray(max_thickness)
        * (phase / phase.max())
        / jnp.asarray(wavelength)
    )
    return phase

microlens_array_amplitude_and_phase(shape, spacing, wavelength, n, fs, centers, radii)

Source code in src/chromatix/utils/initializers.py

@no_type_check
def microlens_array_amplitude_and_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    fs: Array,
    centers: Array,
    radii: Array,
) -> tuple[Array, Array]:
    phase = jnp.zeros(shape)
    amplitude = jnp.zeros(shape)
    grid = create_grid(shape, spacing)

    @no_type_check
    def _place_mask(
        i: int, centers_amplitude_and_phase: tuple[Array, Array]
    ) -> tuple[Array, Array]:
        centers, amplitude, phase = centers_amplitude_and_phase
        center = centers[:, i]
        squared_distance = l2_sq_norm(
            grid - center[:, jnp.newaxis, jnp.newaxis], axis=0
        )
        L = wavelength * fs[i] / n
        mask = jnp.squeeze(squared_distance) < (radii[i] ** 2)
        amplitude += mask
        phase += mask * jnp.squeeze(squared_distance / L)
        return centers, amplitude, phase

    centers, amplitude, phase = jax.lax.fori_loop(
        0, centers.shape[1], _place_mask, (centers, amplitude, phase)
    )
    phase *= -jnp.pi
    amplitude = jnp.clip(amplitude, 0.0, 1.0)
    return amplitude, phase

potato_chip(shape, spacing, wavelength, n, f, NA, d=50.0, C0=-146.7)

Computes the "potato chip" phase mask described by [1].

Also known as the "helical focus" phase mask, this phase mask was designed to produce an extended helical PSF for 3D snapshot microscopy.

[1]: Broxton, Michael. "Volume reconstruction and resolution limits for three dimensional snapshot microscopy." Dissertation, Stanford University, 2017.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape of the phase mask, described as a tuple of integers of the form (H W).	required
spacing	ScalarLike	The spacing of each pixel in the phase mask.	required
wavelength	ScalarLike	The wavelength to compute the phase mask for.	required
n	ScalarLike	Refractive index.	required
f	ScalarLike	The focal distance (should be in same units as wavelength).	required
NA	ScalarLike	The numerical aperture. Phase will be 0 outside of this NA.	required
d	ScalarLike	Sets the axial extent of the PSF (should be in same units as wavelength). Defaults to 50 microns, as shown in [1]. See [1] for more details.	50.0
C0	ScalarLike	Adjusts the focus of the PSF. Set to value described in [1]. See [1] for more details.	-146.7

Source code in src/chromatix/utils/initializers.py

def potato_chip(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    f: ScalarLike,
    NA: ScalarLike,
    d: ScalarLike = 50.0,
    C0: ScalarLike = -146.7,
) -> Array:
    """
    Computes the "potato chip" phase mask described by [1].

    Also known as the "helical focus" phase mask, this phase mask was designed
    to produce an extended helical PSF for 3D snapshot microscopy.

    [1]: Broxton, Michael. "Volume reconstruction and resolution limits for
        three dimensional snapshot microscopy."
        Dissertation, Stanford University, 2017.

    Args:
        shape: The shape of the phase mask, described as a tuple of
            integers of the form (H W).
        spacing: The spacing of each pixel in the phase mask.
        wavelength: The wavelength to compute the phase mask for.
        n: Refractive index.
        f: The focal distance (should be in same units as ``wavelength``).
        NA: The numerical aperture. Phase will be 0 outside of this NA.
        d: Sets the axial extent of the PSF (should be in same units as
            ``wavelength``). Defaults to 50 microns, as shown in [1]. See [1]
            for more details.
        C0: Adjusts the focus of the PSF. Set to value described in [1]. See
            [1] for more details.
    """
    # @copypaste(Field): We must use meshgrid instead of mgrid here
    # in order to be jittable
    grid = create_grid(shape, spacing)
    # Normalize coordinates from -1 to 1 within radius R
    grid = grid_spatial_to_pupil(grid, f, NA, n)
    l2_sq_grid = jnp.sum(grid**2, axis=0)
    theta = jnp.arctan2(*grid)
    k = n / wavelength
    phase = theta * (d * jnp.sqrt(k**2 - l2_sq_grid) + C0)
    phase *= l2_sq_grid < 1
    return phase

rectangular_microlens_array_amplitude_and_phase(shape, spacing, wavelength, n, f, num_lenses_height, num_lenses_width, radius, separation)

Source code in src/chromatix/utils/initializers.py

@no_type_check
def rectangular_microlens_array_amplitude_and_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    f: ScalarLike,
    num_lenses_height: int,
    num_lenses_width: int,
    radius: ScalarLike,
    separation: ScalarLike,
) -> tuple[Array, Array]:
    unit_coordinates = np.meshgrid(
        np.arange(num_lenses_height) - num_lenses_height // 2,
        np.arange(num_lenses_width) - num_lenses_width // 2,
        indexing="ij",
    )
    unit_coordinates = np.array(unit_coordinates).reshape(
        2, num_lenses_height * num_lenses_width
    )
    coordinates = unit_coordinates * separation
    return microlens_array_amplitude_and_phase(
        shape,
        spacing,
        wavelength,
        n,
        jnp.ones(coordinates.shape[1]) * f,
        coordinates,
        jnp.ones(coordinates.shape[1]) * radius,
    )

sawtooth_phase(shape, spacing, wavelength, n_grating, period, thickness, rotation=0.0, n_medium=1.0)

Source code in src/chromatix/utils/initializers.py

def sawtooth_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n_grating: ScalarLike,
    period: ScalarLike,
    thickness: ScalarLike,
    rotation: ScalarLike = 0.0,
    n_medium: ScalarLike = 1.0,
) -> Array:
    dn = jnp.asarray(n_grating - n_medium)
    grid = create_grid(shape, spacing)
    grid = rotate_grid(grid, rotation)
    phase = grid[1] - grid[1].min()
    phase = phase % period
    phase = (
        2 * jnp.pi * dn * thickness * (phase / phase.max()) / jnp.asarray(wavelength)
    )
    return phase

seidel_aberrations(shape, spacing, wavelength, n, f, NA, coefficients, u=0, v=0)

Computes the Seidel phase polynomial described by [1]. Accounts for spatially varying aberrations by creating a different phase mask for each object field position (u,v)

[1]: Voelz, David George. Computational fourier optics: a MATLAB tutorial. Vol. 534. Bellingham, Washington: SPIE press, 2011.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape of the phase mask, described as a tuple of integers of the form (H W).	required
spacing	ScalarLike	The spacing of each pixel in the phase mask.	required
wavelength	ScalarLike	The wavelength to compute the phase mask for.	required
n	ScalarLike	Refractive index.	required
f	ScalarLike	The focal distance (should be in same units as wavelength).	required
NA	ScalarLike	The numerical aperture. Phase will be 0 outside of this NA.	required
coefficients	Sequence[float]	weight coefficients for Seidel aberrations	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.	0
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.	0

Source code in src/chromatix/utils/initializers.py

def seidel_aberrations(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    f: ScalarLike,
    NA: ScalarLike,
    coefficients: Sequence[float],
    u: ScalarLike = 0,
    v: ScalarLike = 0,
) -> Array:
    """
    Computes the Seidel phase polynomial described by [1]. Accounts for spatially
    varying aberrations by creating a different phase mask for each object field
    position (u,v)

    [1]: Voelz, David George. Computational fourier optics: a MATLAB tutorial. Vol. 534.
    Bellingham, Washington: SPIE press, 2011.

    Args:
        shape: The shape of the phase mask, described as a tuple of
            integers of the form (H W).
        spacing: The spacing of each pixel in the phase mask.
        wavelength: The wavelength to compute the phase mask for.
        n: Refractive index.
        f: The focal distance (should be in same units as ``wavelength``).
        NA: The numerical aperture. Phase will be 0 outside of this NA.
        coefficients: weight coefficients for Seidel aberrations
        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.
    """
    # @copypaste(Field): We must use meshgrid instead of mgrid here
    # in order to be jittable
    grid = create_grid(shape, spacing)
    # Normalize coordinates from -1 to 1 within radius R
    grid = grid_spatial_to_pupil(grid, f, NA, n)
    Y, X = grid

    rot_angle = jnp.arctan2(v, u)

    obj_rad = jnp.sqrt(u**2 + v**2)

    X_rot = X * jnp.cos(rot_angle) + Y * jnp.sin(rot_angle)
    Y_rot = -X * jnp.sin(rot_angle) + Y * jnp.cos(rot_angle)

    pupil_radii = jnp.square(X_rot) + jnp.square(Y_rot)
    phase = (
        wavelength * coefficients[0] * jnp.square(pupil_radii)
        + wavelength * coefficients[1] * obj_rad * pupil_radii * X_rot
        + wavelength * coefficients[2] * (obj_rad**2) * jnp.square(X_rot)
        + wavelength * coefficients[3] * (obj_rad**2) * pupil_radii
        + wavelength * coefficients[4] * (obj_rad**3) * X_rot
    )

    l2_sq_grid = X**2 + Y**2

    phase *= l2_sq_grid <= 1
    phase *= 2 * jnp.pi / wavelength
    return phase

sinusoid_phase(shape, spacing, wavelength, n_grating, period, thickness, rotation=0.0, n_medium=1.0)

Source code in src/chromatix/utils/initializers.py

def sinusoid_phase(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n_grating: ScalarLike,
    period: ScalarLike,
    thickness: ScalarLike,
    rotation: ScalarLike = 0.0,
    n_medium: ScalarLike = 1.0,
) -> Array:
    dn = jnp.asarray(n_grating - n_medium)
    grid = create_grid(shape, spacing)
    grid = rotate_grid(grid, rotation)
    phase = grid[1] - grid[1].min()
    phase = jnp.sin(2 * jnp.pi * phase / period)
    phase = (
        2 * jnp.pi * dn * thickness * (phase / phase.max()) / jnp.asarray(wavelength)
    )
    return phase

zernike_aberrations(shape, spacing, wavelength, n, f, NA, ansi_indices, coefficients, normalize=True)

Computes Zernike aberrations given indices of Zernike modes and their corresponding weights.

Parameters:

Name	Type	Description	Default
shape	tuple[int, int]	The shape of the phase mask, described as a tuple of integers of the form (H W).	required
spacing	ScalarLike	The spacing of each pixel in the phase mask.	required
wavelength	ScalarLike	The wavelength to compute the phase mask for.	required
n	ScalarLike	Refractive index of the medium.	required
f	ScalarLike	The focal length of the objective (should be in same units as wavelength).	required
NA	ScalarLike	The numerical aperture of the objective. Phase will be 0 outside of this NA.	required
ansi_indices	Sequence[int]	Linear Zernike indices according to ANSI numbering.	required
coefficients	Sequence[float]	Weight coefficients for the Zernike polynomials.	required
normalize	bool	Whether to normalize the Zernike coefficients. Defaults to True.	True

Source code in src/chromatix/utils/initializers.py

def zernike_aberrations(
    shape: tuple[int, int],
    spacing: ScalarLike,
    wavelength: ScalarLike,
    n: ScalarLike,
    f: ScalarLike,
    NA: ScalarLike,
    ansi_indices: Sequence[int],
    coefficients: Sequence[float],
    normalize: bool = True,
) -> Array:
    """
    Computes Zernike aberrations given indices of Zernike modes and their
    corresponding weights.

    Args:
        shape: The shape of the phase mask, described as a tuple of
            integers of the form (H W).
        spacing: The spacing of each pixel in the phase mask.
        wavelength: The wavelength to compute the phase mask for.
        n: Refractive index of the medium.
        f: The focal length of the objective (should be in same units as
            ``wavelength``).
        NA: The numerical aperture of the objective. Phase will be 0 outside of
            this 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``.
    """

    def convert_ansi_to_zernike_indices(indices):
        d = [math.sqrt(9 + 8 * ind) for ind in indices]
        n = [math.ceil((x - 3) / 2) for x in d]
        m = [2 * ind - x * (x + 2) for ind, x in zip(indices, n)]
        return tuple(zip(n, m))

    def radial_polynomial(n, m):
        """Returns a function calculating the specified radial polynomial."""

        def R(rho):
            if (n - m) % 2 == 0:
                sum = 0
                for k in range(int((n - m) / 2) + 1):
                    sum += (
                        rho ** (n - 2 * k)
                        * ((-1) ** k)
                        * comb(n - k, k)
                        * comb(n - 2 * k, (n - m) / 2 - k)
                    )
                R_nm = sum
            else:
                R_nm = 0
            return R_nm

        return R

    # @copypaste(Field): We must use meshgrid instead of mgrid here
    # in order to be jittable
    grid = create_grid(shape, spacing)
    # Normalize coordinates from -1 to 1 within radius R
    grid = grid_spatial_to_pupil(grid, f, NA, n)

    rho = l2_norm(grid, axis=0)  # radial coordinate

    mask = rho <= 1
    rho = rho * mask
    theta = jnp.arctan2(*grid) * mask  # angle coordinate

    # construct zernike bases to combine during forward pass
    zernike_polynomials = []
    zernike_indices = convert_ansi_to_zernike_indices(ansi_indices)

    for n, m in zernike_indices:
        calc_polynomial = radial_polynomial(n, m)
        R_nm = calc_polynomial(rho)

        if m == 0:
            Z = R_nm
        elif m > 0:  # 'even' Zernike polynomials
            Z = R_nm * jnp.cos(theta * abs(m))
        else:  # 'odd' Zernike polynomials
            Z = R_nm * jnp.sin(theta * abs(m))

        Z = Z * mask

        if normalize:
            if m == 0:
                Z = Z * jnp.sqrt(n + 1)
            else:
                Z = Z * jnp.sqrt(2 * (n + 1))

        zernike_polynomials.append(Z)

    zernike_polynomials = jnp.asarray(zernike_polynomials)
    zernike_polynomials = rearrange(zernike_polynomials, "b h w -> h w b")

    phase = (2 * jnp.pi / jnp.asarray(wavelength)) * jnp.dot(
        zernike_polynomials, jnp.asarray(coefficients)
    )

    return phase

FFT

fft(x, axes=(1, 2), shift=False)

Computes fft2 for input of shape (B... H W C). If shift is true, first applies ifftshift, than an fftshift to make sure everything stays centered.

Source code in src/chromatix/utils/fft.py

def fft(x: ArrayLike, axes: tuple[int, int] = (1, 2), shift: bool = False) -> Array:
    """
    Computes ``fft2`` for input of shape `(B... H W C)`.
    If shift is true, first applies ``ifftshift``, than an ``fftshift`` to
    make sure everything stays centered.
    """
    fft = partial(jnp.fft.fft2, axes=axes)
    fftshift = partial(jnp.fft.fftshift, axes=axes)
    ifftshift = partial(jnp.fft.ifftshift, axes=axes)
    if shift:
        return fftshift(fft(ifftshift(x)))
    else:
        return fft(x)

ifft(x, axes=(1, 2), shift=False)

Computes ifft2 for input of shape (B... H W C). If shift is true, first applies ifftshift, than an fftshift to make sure everything stays centered.

Source code in src/chromatix/utils/fft.py

def ifft(x: ArrayLike, axes: tuple[int, int] = (1, 2), shift: bool = False) -> Array:
    """
    Computes ``ifft2`` for input of shape `(B... H W C)`.
    If shift is true, first applies ``ifftshift``, than an ``fftshift`` to
    make sure everything stays centered.
    """
    ifft = partial(jnp.fft.ifft2, axes=axes)
    fftshift = partial(jnp.fft.fftshift, axes=axes)
    ifftshift = partial(jnp.fft.ifftshift, axes=axes)
    if shift:
        return fftshift(ifft(ifftshift(x)))
    else:
        return ifft(x)

Chirp Z-transform (CZT)

czt(x, m, a, w, axis=-1)

Chirp Z-transform (CZT) of a signal along one dimension. The CZT is a generalization of the discrete Fourier transform (DFT). The DFT samples the Z plane at uniformly-spaced points on the unit circle, whereas the CZT samples the Z plane at uniformly-spaced points on a spiral. This can be used to interpolate the DFT to any desired frequency resolution.

Bluestein's algorithm is used to compute the CZT as a convolution.

Warning

Using float32/complex64 may have numerical accuracy issues.

Parameters:

Name	Type	Description	Default
x	ArrayLike	Input signal to transform.	required
m	int	Number of samples in the output.	required
a	Complex	The starting point in the complex plane. Must lie on the unit circle for numerical stability.	required
w	Complex	The ratio between points in each step. Should lie on the unit circle.	required
axis	int	Axis along which to perform the CZT.	-1

Source code in src/chromatix/utils/czt.py

def czt(x: ArrayLike, m: int, a: Complex, w: Complex, axis: int = -1) -> Array:
    """
    Chirp Z-transform (CZT) of a signal along one dimension. The CZT is a
    generalization of the discrete Fourier transform (DFT). The DFT samples the
    Z plane at uniformly-spaced points on the unit circle, whereas the CZT
    samples the Z plane at uniformly-spaced points on a spiral. This can be
    used to interpolate the DFT to any desired frequency resolution.

    Bluestein's algorithm is used to compute the CZT as a convolution.

    !!! warning
        Using float32/complex64 may have numerical accuracy issues.

    Args:
        x: Input signal to transform.
        m: Number of samples in the output.
        a: The starting point in the complex plane. Must lie on the unit circle
            for numerical stability.
        w: The ratio between points in each step. Should lie on the unit
            circle.
        axis: Axis along which to perform the CZT.
    """

    # TODO switch to jaxtyping
    # # check input values
    # checkify.check(m > 0, "m needs to positive")
    # axis = axis + x.ndim if axis < 0 else axis
    # checkify.check(
    #     axis < x.ndim, "axis needs to be less than the number of dimensions of x"
    # )

    # compute modulation terms
    n = x.shape[axis]
    n_czt = m + n - 1
    k = jnp.arange(n_czt)
    wk2 = w ** (k**2 / 2)
    Awk2 = a ** -k[:n] * wk2[:n]
    Fwk2 = jnp.fft.fft(1 / jnp.hstack((wk2[n - 1 : 0 : -1], wk2[:m])), n_czt)
    wk2 = wk2[:m]

    # perform CZT
    x = jnp.moveaxis(x, axis, -1)
    y = jnp.fft.ifft(jnp.fft.fft(x * Awk2, n_czt, axis=-1) * Fwk2, axis=-1)
    y = y[..., n - 1 : n + m - 1] * wk2
    y = jnp.moveaxis(y, -1, axis)
    return y

cztn(x, m, a, w, axes=(-2, -1))

Chirp Z-transform (CZT) of a signal along multiple dimensions as defined by the axes parameter. This implementation loops over the dimensions and performs the CZT along each dimension.

Warning

Using float32/complex64 may have numerical accuracy issues.

Parameters:

Name	Type	Description	Default
x	ArrayLike	Input signal to transform.	required
m	tuple[int]	Number of samples in the output. List for each dimension.	required
a	Complex[Array, m]	The starting point in the complex plane. List for each dimension.	required
w	Complex[Array, m]	The ratio between points in each step. List for each dimension.	required
axes	tuple[int]	Axes along which to perform the CZT.	(-2, -1)

Source code in src/chromatix/utils/czt.py

def cztn(
    x: ArrayLike,
    m: tuple[int],
    a: Complex[Array, "m"],
    w: Complex[Array, "m"],
    axes: tuple[int] = (-2, -1),
) -> Array:
    """
    Chirp Z-transform (CZT) of a signal along multiple dimensions as defined by
    the `axes` parameter. This implementation loops over the dimensions and
    performs the CZT along each dimension.

    !!! warning
        Using float32/complex64 may have numerical accuracy issues.

    Args:
        x: Input signal to transform.
        m: Number of samples in the output. List for each dimension.
        a: The starting point in the complex plane. List for each dimension.
        w: The ratio between points in each step. List for each dimension.
        axes: Axes along which to perform the CZT.
    """
    x_czt = x
    for d, ax in enumerate(axes):
        x_czt = czt(x_czt, a=a[d], w=w[d], m=m[d], axis=ax)
    return x_czt

zoomed_fft(x, k_start, k_end, output_shape, axes=(-2, -1), include_end=True)

Custom FFTN function that uses the Chirp Z-transform (CZT) to compute the Fourier transform of a signal. It allows to generate zoomed FFT in a region between k_start and k_end with arbitrary output shape. The usual FFT corresponds to output_shape = x.shape, k_start = 0, k_end = 2 * pi, and include_end = False.

Parameters:

Name	Type	Description	Default
x	ArrayLike	Input signal to transform.	required
k_start	float	Start of the frequency range.	required
k_end	float	End of the frequency range.	required
output_shape	tuple[int]	Desired shape of the output.	required
axes	tuple[int]	Axes along which to perform the CZT.	(-2, -1)
include_end	bool	Whether to include the end point in the frequency range.	True

Returns:

Type	Description
ArrayLike	The Fourier transform of the input signal.

Source code in src/chromatix/utils/czt.py

def zoomed_fft(
    x: ArrayLike,
    k_start: float,
    k_end: float,
    output_shape: tuple[int],
    axes: tuple[int] = (-2, -1),
    include_end: bool = True,
) -> ArrayLike:
    """
    Custom FFTN function that uses the Chirp Z-transform (CZT) to compute the
    Fourier transform of a signal. It allows to generate zoomed FFT in a region
    between k_start and k_end with arbitrary output shape. The usual FFT
    corresponds to output_shape = x.shape, k_start = 0, k_end = 2 * pi, and
    include_end = False.

    Args:
        x: Input signal to transform.
        k_start: Start of the frequency range.
        k_end: End of the frequency range.
        output_shape: Desired shape of the output.
        axes: Axes along which to perform the CZT.
        include_end: Whether to include the end point in the frequency range.

    Returns:
        The Fourier transform of the input signal.
    """
    if include_end:
        renorm = tuple(m - 1 for m in output_shape)
    else:
        renorm = output_shape
    w = tuple(jnp.exp(1j * (k_end - k_start) / n) for n in renorm)

    a = jnp.exp(-1j * k_start)
    a = tuple(a for _ in range(len(output_shape)))

    return cztn(
        x=x,
        m=output_shape,
        a=a,
        w=w,
        axes=axes,
    )

Shapes

__all__ = ['_broadcast_1d_to_channels', '_broadcast_1d_to_polarization', '_broadcast_1d_to_innermost_batch', '_broadcast_1d_to_grid', '_broadcast_2d_to_grid', '_broadcast_dx_to_grid', '_squeeze_grid_to_2d', '_broadcast_2d_to_spatial'] module-attribute