Skip to content

Operations

General

fourier_convolution(image, kernel, *, axes=(0, 1), fast_fft_shape=True, mode='same')

Fourier convolution in n dimensions over the specified axes of an Array.

The convolution dimensions are determined by the axes argument. For example, the default axes to perform convolution over are (0, 1), or the first two axes of the input, which will perform a 2D convolution. The kernel and image should have the same number of dimensions.

This function computes the convolution kernel * image by employing the Fourier convolution theorem. The inputs are padded appropriately to avoid circular convolutions.

By default, the inputs are further padded to the nearest power of 2 that is larger than the padded input shape for faster FFT performance. If the input shape causes the difference between padded and unpadded to be too large (causing either memory or performance issues), this extra padding can be disabled.

Parameters:

Name Type Description Default
image ArrayLike

The input to be convolved.

 required
kernel ArrayLike

The convolution kernel.

 required
fast_fft_shape bool

Determines whether inputs should be further padded for increased FFT performance. Defaults to True.

 True
mode str

A string that determines whether to crop the result of the convolution to the same shape as image. Should be either "same" or "full". Defaults to "same".

 'same'
Source code in src/chromatix/ops/ops.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def fourier_convolution(
    image: ArrayLike,
    kernel: ArrayLike,
    *,
    axes: tuple[int, ...] = (0, 1),
    fast_fft_shape: bool = True,
    mode: str = "same",
) -> Array:
    """
    Fourier convolution in n dimensions over the specified axes of an ``Array``.

    The convolution dimensions are determined by the `axes` argument. For
    example, the default axes to perform convolution over are (0, 1), or the
    first two axes of the input, which will perform a 2D convolution. The
    ``kernel`` and ``image`` should have the same number of dimensions.

    This function computes the convolution ``kernel * image`` by employing the
    Fourier convolution theorem. The inputs are padded appropriately to avoid
    circular convolutions.

    By default, the inputs are further padded to the nearest power of 2 that
    is larger than the padded input shape for faster FFT performance. If the
    input shape causes the difference between padded and unpadded to be too
    large (causing either memory or performance issues), this extra padding can
    be disabled.

    Args:
        image: The input to be convolved.
        kernel: The convolution kernel.
        fast_fft_shape: Determines whether inputs should be further padded for
            increased FFT performance. Defaults to ``True``.
        mode: A string that determines whether to crop the result of the
            convolution to the same shape as ``image``. Should be either
            ``"same"`` or ``"full"``. Defaults to ``"same"``.
    """
    axes = list(axes)
    for i in range(len(axes) - 1):
        assert axes[i + 1] == (axes[i] + 1), "Axes to convolve over must be contiguous"
    assert image.ndim == kernel.ndim, (
        f"Input ({image.ndim}D) and kernel ({kernel.ndim}D) must have same number of dimensions"
    )
    for i in range(len(axes)):
        if axes[i] < 0:
            axes[i] = image.ndim + axes[i]
    # Get padded shape to prevent circular convolution
    padded_shape = [
        k1 + k2
        for k1, k2 in zip(
            image.shape[axes[0] : axes[-1] + 1], kernel.shape[axes[0] : axes[-1] + 1]
        )
    ]
    if fast_fft_shape:
        fast_shape = [next_order(k) for k in padded_shape]
    else:
        fast_shape = padded_shape
    # Save memory with rfft if inputs are not complex
    is_complex = (image.dtype.kind == "c") or (kernel.dtype.kind == "c")
    # output_shape = image.shape[axes[0]:axes[-1] + 1] if mode == "same" else fast_shape
    if is_complex:
        fft = partial(jnp.fft.fftn, s=fast_shape, axes=axes)
        ifft = partial(jnp.fft.ifftn, s=fast_shape, axes=axes)
        # ifft = partial(jnp.fft.ifftn, s=output_shape, axes=axes)
    else:
        fft = partial(jnp.fft.rfftn, s=fast_shape, axes=axes)
        ifft = partial(jnp.fft.irfftn, s=fast_shape, axes=axes)
        # ifft = partial(jnp.fft.irfftn, s=output_shape, axes=axes)
    conv = ifft(fft(image) * fft(kernel))
    # Remove padding
    if mode == "same":
        conv = conv[
            tuple(
                [
                    slice((k - 1) // 2, (k - 1) // 2 + i) if idx in axes else slice(i)
                    for idx, (i, k) in enumerate(zip(image.shape, kernel.shape))
                ]
            )
        ]
    return conv

Filters

__all__ = ['high_pass_filter', 'gaussian_filter'] module-attribute

gaussian_filter(data, sigma, axes=(1, 2), kernel_shape=None)

Performs a Gaussian filter on data.

Parameters:

Name Type Description Default
data Array

The input to be Gaussian filtered.

 required
sigma Sequence[float]

The standard deviation of the Gaussian kernel.

 required
kernel_shape Sequence[int] | None

The shape of the kernel. If not provided, the shape will be determined by the required shape of a Gaussian kernel truncated to 4.0 * sigma.

 None

Returns:

Type Description
Array

The Gaussian filtered array.
Source code in src/chromatix/ops/filters.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def gaussian_filter(
    data: Array,
    sigma: Sequence[float],
    axes: tuple[int, int] = (1, 2),
    kernel_shape: Sequence[int] | None = None,
) -> Array:
    """
    Performs a Gaussian filter on ``data``.

    Args:
        data: The input to be Gaussian filtered.
        sigma: The standard deviation of the Gaussian kernel.
        kernel_shape: The shape of the kernel. If not provided, the shape will
            be determined by the required shape of a Gaussian kernel truncated
            to ``4.0 * sigma``.

    Returns:
        The Gaussian filtered array.
    """
    assert len(axes) == len(sigma), (
        "Must specify same number of axes to convolve as elements in sigma"
    )
    kernel = gaussian_kernel(sigma, shape=kernel_shape)
    kernel = _broadcast_2d_to_spatial(kernel, data.ndim)
    return fourier_convolution(data, kernel, axes=axes)

high_pass_filter(data, sigma, axes=(1, 2), kernel_shape=None)

Performs a high pass filter on data.

The high pass filter is constructed as the difference between a delta kernel and a Gaussian kernel with standard deviation sigma.

Parameters:

Name Type Description Default
data ArrayLike

The input to be high pass filtered.

 required
sigma Sequence[float]

The standard deviation of the Gaussian kernel, which sets the low pass filter. The result of this low pass will be subtracted from the input.

 required
kernel_shape Sequence[int] | None

The shape of the kernel. If not provided, the shape will be determined by the required shape of a Gaussian kernel truncated to 4.0 * sigma.

 None

Returns:

Type Description
Array

The high pass filtered array.
Source code in src/chromatix/ops/filters.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def high_pass_filter(
    data: ArrayLike,
    sigma: Sequence[float],
    axes: tuple[int, int] = (1, 2),
    kernel_shape: Sequence[int] | None = None,
) -> Array:
    """
    Performs a high pass filter on ``data``.

    The high pass filter is constructed as the difference between a
    delta kernel and a Gaussian kernel with standard deviation ``sigma``.

    Args:
        data: The input to be high pass filtered.
        sigma: The standard deviation of the Gaussian kernel, which sets the
            low pass filter. The result of this low pass will be subtracted
            from the input.
        kernel_shape: The shape of the kernel. If not provided, the shape will
            be determined by the required shape of a Gaussian kernel truncated
            to ``4.0 * sigma``.

    Returns:
        The high pass filtered array.
    """
    assert len(axes) == len(sigma), (
        "Must specify same number of axes to convolve as elements in sigma"
    )
    low_pass_kernel = gaussian_kernel(sigma, shape=kernel_shape)
    # NOTE(gj): 1e-3 effectively gives delta kernel
    delta_kernel = gaussian_kernel((1e-3,) * len(sigma), shape=low_pass_kernel.shape)
    kernel = delta_kernel - low_pass_kernel
    if axes[0] >= 0 and axes[1] >= 0:
        _axes = (axes[0] - data.ndim, axes[1] - data.ndim)
        kernel = _broadcast_2d_to_spatial(kernel, _axes)
        kernel = rearrange(kernel, "... -> " + ("1 " * axes[0]) + "...")
    else:
        kernel = _broadcast_2d_to_spatial(kernel, axes)
    return fourier_convolution(data, kernel, axes=axes)

Noise

approximate_shot_noise(key, image)

Approximates Poisson shot noise using a Gaussian for differentiability.

Source code in src/chromatix/ops/noise.py 
 9
10
11
12
13
14
15
@custom_jvp
def approximate_shot_noise(key: PRNGKey, image: ArrayLike) -> Array:
    """
    Approximates Poisson shot noise using a Gaussian for differentiability.
    """
    noisy = image + jnp.sqrt(image) * random.normal(key, image.shape)
    return jnp.maximum(noisy, 0.0)

approximate_shotnoise_jvp(primals, tangents)

Custom gradient for approximate_shot_noise.

This is necessary to fix an instability when the input image is 0.

Source code in src/chromatix/ops/noise.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@approximate_shot_noise.defjvp
def approximate_shotnoise_jvp(primals: tuple, tangents: tuple) -> tuple:
    """
    Custom gradient for ``approximate_shot_noise``.

    This is necessary to fix an instability when the input ``image`` is 0.
    """
    key, image = primals
    _, image_dot = tangents
    primal_out = approximate_shot_noise(key, image)
    # We define the gradient to be zero if image=0
    # we just add eta as we multiply by zero later anyway
    noise_grad = jnp.ones_like(image) + random.normal(key, image.shape) / (
        2 * jnp.sqrt(image) + 1e-6
    )
    # maximum operation, abs to get rid of -0
    tangent_out = image_dot * jnp.abs(noise_grad) * (primal_out != 0)
    return primal_out, tangent_out

shot_noise(key, image)

Simulates Poisson shot noise whose gradient is approximated using the gradient of a Gaussian, just as if the simulation had been approximate_shot_noise instead.

Source code in src/chromatix/ops/noise.py 
38
39
40
41
42
43
44
45
46
@custom_jvp
def shot_noise(key: PRNGKey, image: ArrayLike) -> Array:
    """
    Simulates Poisson shot noise whose gradient is approximated using
    the gradient of a Gaussian, just as if the simulation had been
    ``approximate_shot_noise`` instead.
    """
    noisy = random.poisson(key, image, image.shape)
    return jnp.float32(noisy)

shotnoise_jvp(primals, tangents)

Custom gradient for shot_noise.

Because the Poisson distribution computed in shot_noise cannot be differentiated, this function computes the gradient as if the forward pass had been approximate_shot_noise.

Source code in src/chromatix/ops/noise.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
@shot_noise.defjvp
def shotnoise_jvp(primals: Tuple, tangents: Tuple) -> Tuple:
    """
    Custom gradient for ``shot_noise``.

    Because the Poisson distribution computed in ``shot_noise`` cannot be
    differentiated, this function computes the gradient as if the forward pass
    had been ``approximate_shot_noise``.
    """
    key, image = primals
    _, image_dot = tangents
    primal_out = shot_noise(key, image)
    # We define the gradient to be zero if image=0
    # we just add eta as we multiply by zero later anyway
    noise_grad = jnp.ones_like(image) + random.normal(key, image.shape) / (
        2 * jnp.sqrt(image) + 1e-6
    )
    # maximum operation, abs to get rid of -0)
    tangent_out = image_dot * jnp.abs(noise_grad) * (primal_out != 0)
    return primal_out, tangent_out

Quantization

__all__ = ['binarize', 'binarize_jvp', 'quantize', 'quantize_jvp'] module-attribute

binarize(x, threshold=0.5)

Binarize each pixel of amplitude mask to be either 0 or 1.

Optionally takes a threshold to choose the threshold at which values are set to 1. Defaults to 0.5.

The gradient of binarize is approximated with the gradient of the sigmoid function as in [1].

[1] Eybposh, M. Hossein, et al. "Optimization of time-multiplexed computer-generated holograms with surrogate gradients." Emerging Digital Micromirror Device Based Systems and Applications XIV. SPIE, 2022.

Parameters:

Name Type Description Default
x Array

Input to binarize to 0 or 1.

 required
threshold float

Threshold above which values are set to 1. Defaults to 0.5.

 0.5
Source code in src/chromatix/ops/quantization.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
@jax.custom_jvp
def binarize(x: Array, threshold: float = 0.5) -> Array:
    """
    Binarize each pixel of amplitude mask to be either 0 or 1.

    Optionally takes a ``threshold`` to choose the threshold at which values
    are set to 1. Defaults to 0.5.

    The gradient of ``binarize`` is approximated with the gradient of the
    sigmoid function as in [1].

    [1] Eybposh, M. Hossein, et al. "Optimization of time-multiplexed
    computer-generated holograms with surrogate gradients." Emerging Digital
    Micromirror Device Based Systems and Applications XIV. SPIE, 2022.

    Args:
        x: Input to binarize to 0 or 1.
        threshold: Threshold above which values are set to 1. Defaults to 0.5.
    """
    return (x > threshold) * 1.0

binarize_jvp(primals, tangents)

Custom gradient for binarize.

We approximate the gradient of binarize with the gradient of the sigmoid function.

Source code in src/chromatix/ops/quantization.py 
30
31
32
33
34
35
36
37
38
39
40
@binarize.defjvp
def binarize_jvp(primals: tuple, tangents: tuple) -> tuple:
    """
    Custom gradient for ``binarize``.

    We approximate the gradient of ``binarize`` with the gradient of the
    sigmoid function.
    """
    sig = 1 / (1 + jnp.exp((-1 * primals[0]) + primals[1]))
    out_tangents = tangents[0] * (1 / 2.0) * (1 - sig**2 - (1 - sig) ** 2)
    return binarize(*primals), out_tangents

quantize(x, bit_depth, range=None)

Quantize the input x to the specified bit_depth. Surrogate gradient approach [1] is used to adjust the bit depth differentiably.

[1] Eybposh, M. Hossein, et al. "Optimization of time-multiplexed computer-generated holograms with surrogate gradients." Emerging Digital Micromirror Device Based Systems and Applications XIV. SPIE, 2022.

Parameters:

Name Type Description Default
x Array

Input to quantize.

 required
bit_depth float

Number of bits. This parameter does NOT represent the number of digitization levels, but the bit depth.

 required
range tuple[int, int] | None

Range to quantize to, provided as (minimum, maximum). If not provided, the range of the values in x will be used.

 None

Returns:

Type Description
Array

The quantized input.
Source code in src/chromatix/ops/quantization.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
@jax.custom_jvp
def quantize(x: Array, bit_depth: float, range: tuple[int, int] | None = None) -> Array:
    """
    Quantize the input ``x`` to the specified ``bit_depth``. Surrogate gradient
    approach [1] is used to adjust the bit depth differentiably.

    [1] Eybposh, M. Hossein, et al. "Optimization of time-multiplexed
    computer-generated holograms with surrogate gradients." Emerging Digital
    Micromirror Device Based Systems and Applications XIV. SPIE, 2022.

    Args:
        x: Input to quantize.
        bit_depth: Number of bits. This parameter does NOT represent the number of
            digitization levels, but the bit depth.
        range: Range to quantize to, provided as ``(minimum, maximum)``. If not
            provided, the range of the values in ``x`` will be used.

    Returns:
        The quantized input.
    """
    if range is None:
        x_min = x.min()
        x_max = (x - x_min).max()
    else:
        x_min, x_max = range
    y = jnp.round(((x - x_min) / x_max) * ((2**bit_depth) - 1)).astype(jnp.float32)
    return (y / ((2**bit_depth) - 1.0)) * x_max + x_min

quantize_jvp(primals, tangents)

Custom gradient for quantize.

We approximate the gradient of quantize with surrogate gradients, as described in [1] (see definition of quantize).

Source code in src/chromatix/ops/quantization.py 
72
73
74
75
76
77
78
79
80
@quantize.defjvp
def quantize_jvp(primals: tuple, tangents: tuple) -> tuple:
    """
    Custom gradient for ``quantize``.

    We approximate the gradient of ``quantize`` with surrogate gradients, as
    described in [1] (see definition of ``quantize``).
    """
    return quantize(*primals), tangents[0]

Resample

InterpolatingPlaneResampler

Bases: Resampler

Source code in src/chromatix/ops/resample.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class InterpolatingPlaneResampler(Resampler):
    out_shape: tuple[int, int] = eqx.field(static=True)
    out_spacing: ScalarLike | Float[Array, "2"]
    resampling_method: str = eqx.field(static=True)

    def __init__(
        self,
        out_shape: tuple[int, int],
        out_spacing: ScalarLike | Float[Array, "2"],
        resampling_method: str = "linear",
    ):
        self.out_shape = out_shape
        self.out_spacing = out_spacing
        self.resampling_method = resampling_method

    def __call__(
        self, resample_input: Float[Array, "h w ..."], in_spacing: Float[Array, "2"]
    ) -> Array:
        in_spacing = jnp.atleast_1d(jnp.asarray(in_spacing).squeeze())
        assert in_spacing.size == 2, (
            "Input spacing is an array of shape (2,) representing pixel size in (y x)"
        )
        _in_shape, _out_shape = (
            jnp.asarray(resample_input.shape),
            jnp.asarray(self.out_shape),
        )
        scale = in_spacing / self.out_spacing
        translation = -0.5 * (_in_shape * scale - _out_shape)
        # NOTE(dd): Because scale_and_translate expects shape to have same
        # number of dimensions as input, we have to extend the shape with
        # any channel/ vectorial dimensions here
        # extended_shape = out_shape + x.shape
        resample_output = scale_and_translate(
            resample_input,
            self.out_shape,
            (0, 1),
            scale,
            translation,
            method=self.resampling_method,
        )
        resample_output = resample_output / jnp.prod(scale)
        return resample_output

out_shape = out_shape class-attribute instance-attribute

out_spacing = out_spacing instance-attribute

resampling_method = resampling_method class-attribute instance-attribute

__call__(resample_input, in_spacing)

Source code in src/chromatix/ops/resample.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
def __call__(
    self, resample_input: Float[Array, "h w ..."], in_spacing: Float[Array, "2"]
) -> Array:
    in_spacing = jnp.atleast_1d(jnp.asarray(in_spacing).squeeze())
    assert in_spacing.size == 2, (
        "Input spacing is an array of shape (2,) representing pixel size in (y x)"
    )
    _in_shape, _out_shape = (
        jnp.asarray(resample_input.shape),
        jnp.asarray(self.out_shape),
    )
    scale = in_spacing / self.out_spacing
    translation = -0.5 * (_in_shape * scale - _out_shape)
    # NOTE(dd): Because scale_and_translate expects shape to have same
    # number of dimensions as input, we have to extend the shape with
    # any channel/ vectorial dimensions here
    # extended_shape = out_shape + x.shape
    resample_output = scale_and_translate(
        resample_input,
        self.out_shape,
        (0, 1),
        scale,
        translation,
        method=self.resampling_method,
    )
    resample_output = resample_output / jnp.prod(scale)
    return resample_output

__init__(out_shape, out_spacing, resampling_method='linear')

Source code in src/chromatix/ops/resample.py 
38
39
40
41
42
43
44
45
46
def __init__(
    self,
    out_shape: tuple[int, int],
    out_spacing: ScalarLike | Float[Array, "2"],
    resampling_method: str = "linear",
):
    self.out_shape = out_shape
    self.out_spacing = out_spacing
    self.resampling_method = resampling_method

PoolingPlaneDownsampler

Bases: Resampler

Source code in src/chromatix/ops/resample.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class PoolingPlaneDownsampler(Resampler):
    out_shape: tuple[int, int] = eqx.field(static=True)
    out_spacing: ScalarLike | Float[Array, "2"]

    def __init__(
        self, out_shape: tuple[int, int], out_spacing: ScalarLike | Float[Array, "2"]
    ):
        self.out_shape = out_shape
        self.out_spacing = out_spacing

    def __call__(
        self, resample_input: Float[Array, "h w ..."], in_spacing: Float[Array, "2"]
    ) -> Array:
        return reduce(
            resample_input,
            "(h hf) (w wf) ... -> h w ...",
            "sum",
            h=self.out_shape[0],
            w=self.out_shape[1],
        )

out_shape = out_shape class-attribute instance-attribute

out_spacing = out_spacing instance-attribute

__call__(resample_input, in_spacing)

Source code in src/chromatix/ops/resample.py 
21
22
23
24
25
26
27
28
29
30
def __call__(
    self, resample_input: Float[Array, "h w ..."], in_spacing: Float[Array, "2"]
) -> Array:
    return reduce(
        resample_input,
        "(h hf) (w wf) ... -> h w ...",
        "sum",
        h=self.out_shape[0],
        w=self.out_shape[1],
    )

__init__(out_shape, out_spacing)

Source code in src/chromatix/ops/resample.py 
15
16
17
18
19
def __init__(
    self, out_shape: tuple[int, int], out_spacing: ScalarLike | Float[Array, "2"]
):
    self.out_shape = out_shape
    self.out_spacing = out_spacing

init_plane_resample(out_shape, out_spacing, resampling_method='linear')

Returns a function that resamples 2D planes to the specified output shape and spacing. These functions are instances of Resamplers in Chromatix.

The returned function is allowed to be jitted because the shape of the output will no longer depend on the input of this function.

Multiple resampling_methods are supported: either 'pooling' which uses sum pooling (for downsampling only) or any method supported by jax.image.scale_and_translate ('linear', 'cubic', 'lanczos3', or 'lanczos5').

The input may have any number of dimensions after the first two, but the returned function assumes that the 2D planes to be downsampled are contained in the first two axes. Any other dimensions are treated as batch dimensions, i.e. resampling is parallelized across those dimensions. In order to add arbitrary batch dimensions before the first two dimensions, use jax.vmap.

Parameters:

Name Type Description Default
out_shape tuple[int, ...]

A tuple representing the number of samples (pixels) to which the incoming plane should be resampled in the format (height width).

 required
out_spacing ScalarLike | Float[Array, 2]

Either a scalar or a 1D array of size 2 (in the format (height width)) representing the spacing between samples in units of distance. A scalar value represents square pixels, which is typically what you will want to use.

 required
resampling_method str

A string representing the type of resampling method to initialize. Can be either "linear", "cubic", "lanczos3", or "lanczos5" for arbitrary interpolation, or "pooling" for a sum pooling downsampling. Defaults to "linear".

 'linear'

Returns: A Resampler, which is a callable that actually performs the resampling.

Source code in src/chromatix/ops/resample.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def init_plane_resample(
    out_shape: tuple[int, ...],
    out_spacing: ScalarLike | Float[Array, "2"],
    resampling_method: str = "linear",
) -> Resampler:
    """
    Returns a function that resamples 2D planes to the specified output shape
    and spacing. These functions are instances of ``Resampler``s in Chromatix.

    The returned function is allowed to be jitted because the shape of the
    output will no longer depend on the input of this function.

    Multiple ``resampling_methods`` are supported: either `'pooling'` which
    uses sum pooling (for downsampling only) or any method supported by
    ``jax.image.scale_and_translate`` (`'linear'`, `'cubic'`, `'lanczos3'`,
    or `'lanczos5'`).

    The input may have any number of dimensions after the first two, but
    the returned function assumes that the 2D planes to be downsampled are
    contained in the first two axes. Any other dimensions are treated as batch
    dimensions, i.e. resampling is parallelized across those dimensions. In
    order to add arbitrary batch dimensions before the first two dimensions,
    use ``jax.vmap``.

    Args:
        out_shape: A tuple representing the number of samples (pixels) to
            which the incoming plane should be resampled in the format `(height
            width)`.
        out_spacing: Either a scalar or a 1D array of size 2 (in the format
            `(height width)`) representing the spacing between samples in units
            of distance. A scalar value represents square pixels, which is
            typically what you will want to use.
        resampling_method: A string representing the type of resampling
            method to initialize. Can be either ``"linear"``, ``"cubic"``,
            ``"lanczos3"``, or ``"lanczos5"`` for arbitrary interpolation,
            or ``"pooling"`` for a sum pooling downsampling. Defaults to
            ``"linear"``.
    Returns:
        A [``Resampler``](core.md#chromatix.core.base.Resampler), which is a
        callable that actually performs the resampling.
    """
    assert len(out_shape) == 2, "Shape must be tuple of form (height width)"
    assert np.atleast_1d(np.asarray(out_spacing).squeeze()).size <= 2, (
        "Spacing is either a float or array of shape (2,) for non-square pixels"
    )
    if resampling_method == "pool":
        return PoolingPlaneDownsampler(out_shape, out_spacing)
    else:
        return InterpolatingPlaneResampler(
            out_shape, out_spacing, resampling_method=resampling_method
        )