Utils

cax.utils.render

Utilities for rendering.

rgba_to_rgb(array)

Convert a premultiplied RGBA image to RGB by alpha compositing over white.

RGBA arrays in CAX are premultiplied: colour is already scaled by alpha, so a pixel holds the light it emits and compositing over white is rgb + (1 - alpha). Targets from get_emoji_array and the RGBA channels of a neural cellular automaton state follow this convention.

The function assumes the last dimension encodes channels and that the input is normalized to the range [0, 1] with shape (..., 4). The output preserves the input shape except for the channel dimension, which becomes 3.

Parameters:

Name	Type	Description	Default
array	Array	Premultiplied RGBA image with shape (..., 4) and values in [0, 1].	required

Returns:

Type	Description
Array	RGB image with shape (..., 3) and values in [0, 1].

Source code in src/cax/utils/render.py

def rgba_to_rgb(array: Array) -> Array:
    """Convert a premultiplied RGBA image to RGB by alpha compositing over white.

    RGBA arrays in CAX are premultiplied: colour is already scaled by alpha, so a pixel
    holds the light it emits and compositing over white is ``rgb + (1 - alpha)``.
    Targets from ``get_emoji_array`` and the RGBA channels of a neural cellular
    automaton state follow this convention.

    The function assumes the last dimension encodes channels and that the input is
    normalized to the range ``[0, 1]`` with shape ``(..., 4)``. The output preserves the
    input shape except for the channel dimension, which becomes ``3``.

    Args:
        array: Premultiplied RGBA image with shape ``(..., 4)`` and values in
            ``[0, 1]``.

    Returns:
        RGB image with shape ``(..., 3)`` and values in ``[0, 1]``.

    """
    if array.shape[-1] != 4:
        raise ValueError(
            f"Expected an RGBA array with 4 channels, got {array.shape[-1]}"
        )
    rgb, alpha = array[..., :-1], array[..., -1:]
    alpha = jnp.clip(alpha, min=0.0, max=1.0)
    return (1.0 - alpha) + rgb

rgb_to_hsv(rgb)

Convert RGB to HSV.

Input and output are in the range [0, 1] and use channel-last layout.

Parameters:

Name	Type	Description	Default
rgb	Array	RGB image with shape (..., 3).	required

Returns:

Type	Description
Array	HSV image with shape (..., 3).

Source code in src/cax/utils/render.py

def rgb_to_hsv(rgb: Array) -> Array:
    """Convert RGB to HSV.

    Input and output are in the range ``[0, 1]`` and use channel-last layout.

    Args:
        rgb: RGB image with shape ``(..., 3)``.

    Returns:
        HSV image with shape ``(..., 3)``.

    """
    input_shape = rgb.shape
    rgb = rgb.reshape(-1, 3)
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]

    maxc = jnp.maximum(jnp.maximum(r, g), b)
    minc = jnp.minimum(jnp.minimum(r, g), b)
    v = maxc
    deltac = maxc - minc

    s = jnp.where(maxc != 0, deltac / maxc, 0)

    deltac = jnp.where(deltac == 0, 1, deltac)  # Avoid division by zero

    rc = (maxc - r) / deltac
    gc = (maxc - g) / deltac
    bc = (maxc - b) / deltac

    h = jnp.where(
        r == maxc, bc - gc, jnp.where(g == maxc, 2.0 + rc - bc, 4.0 + gc - rc)
    )

    h = jnp.where(minc == maxc, 0.0, h)
    h = (h / 6.0) % 1.0

    hsv = jnp.stack([h, s, v], axis=-1)
    return hsv.reshape(input_shape)

hsv_to_rgb(hsv)

Convert HSV to RGB.

Input and output are in the range [0, 1] and use channel-last layout.

Parameters:

Name	Type	Description	Default
hsv	Array	HSV image with shape (..., 3).	required

Returns:

Type	Description
Array	RGB image with shape (..., 3).

Source code in src/cax/utils/render.py

def hsv_to_rgb(hsv: Array) -> Array:
    """Convert HSV to RGB.

    Input and output are in the range ``[0, 1]`` and use channel-last layout.

    Args:
        hsv: HSV image with shape ``(..., 3)``.

    Returns:
        RGB image with shape ``(..., 3)``.

    """
    input_shape = hsv.shape
    hsv = hsv.reshape(-1, 3)
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]

    i = jnp.floor(h * 6.0).astype(jnp.int32)
    f = (h * 6.0) - i
    p = v * (1.0 - s)
    q = v * (1.0 - s * f)
    t = v * (1.0 - s * (1.0 - f))

    i = i % 6

    rgb = jnp.zeros_like(hsv)
    rgb = jnp.where((i == 0)[..., None], jnp.stack([v, t, p], axis=-1), rgb)
    rgb = jnp.where((i == 1)[..., None], jnp.stack([q, v, p], axis=-1), rgb)
    rgb = jnp.where((i == 2)[..., None], jnp.stack([p, v, t], axis=-1), rgb)
    rgb = jnp.where((i == 3)[..., None], jnp.stack([p, q, v], axis=-1), rgb)
    rgb = jnp.where((i == 4)[..., None], jnp.stack([t, p, v], axis=-1), rgb)
    rgb = jnp.where((i == 5)[..., None], jnp.stack([v, p, q], axis=-1), rgb)

    rgb = jnp.where(s[..., None] == 0.0, jnp.full_like(rgb, v[..., None]), rgb)

    return rgb.reshape(input_shape)

clip_and_uint8(frame)

Clip a floating-point image to [0, 1] and convert to uint8.

Parameters:

Name	Type	Description	Default
frame	Array	Image-like array with values expected in or near [0, 1].	required

Returns:

Type	Description
Array	Array of dtype uint8 with values in [0, 255].

Source code in src/cax/utils/render.py

def clip_and_uint8(frame: Array) -> Array:
    """Clip a floating-point image to ``[0, 1]`` and convert to ``uint8``.

    Args:
        frame: Image-like array with values expected in or near ``[0, 1]``.

    Returns:
        Array of dtype ``uint8`` with values in ``[0, 255]``.

    """
    frame = jnp.clip(frame, min=0.0, max=1.0)
    return (frame * 255).astype(jnp.uint8)

render_array_with_channels_to_rgb(array)

Render an array with channels as an RGB image.

This function processes an input array and converts it into an RGB image based on the number of channels present in the array. The conversion logic is as follows: - If the array has 1 channel, it is repeated across the RGB channels to produce a grayscale image. - If the array has 2 channels, the first channel is interpreted as hue and the second as saturation. These are converted to RGB using a fixed brightness value, resulting in a colorful representation. - If the array has 3 or more channels, the last three channels are used directly as the RGB values.

Parameters:

Name	Type	Description	Default
array	Array	Input array with shape (..., C) and values in [0, 1].	required

Returns:

Type	Description
Array	RGB array with shape (..., 3) and values in [0, 1].

Source code in src/cax/utils/render.py

def render_array_with_channels_to_rgb(array: Array) -> Array:
    """Render an array with channels as an RGB image.

    This function processes an input array and converts it into an RGB image based on
    the number of channels present in the array. The conversion logic is as follows:
    - If the array has 1 channel, it is repeated across the RGB channels to produce a
        grayscale image.
    - If the array has 2 channels, the first channel is interpreted as hue and the
        second as saturation. These are converted to RGB using a fixed brightness value,
        resulting in a colorful representation.
    - If the array has 3 or more channels, the last three channels are used directly as
        the RGB values.

    Args:
        array: Input array with shape ``(..., C)`` and values in ``[0, 1]``.

    Returns:
        RGB array with shape ``(..., 3)`` and values in ``[0, 1]``.

    """
    num_channels = array.shape[-1]

    if num_channels == 1:
        # 1 channel
        rgb = jnp.repeat(array, 3, axis=-1)
    elif num_channels == 2:
        # 2 channels
        hue = array[..., 0:1]  # Use the first channel as hue
        saturation = array[..., 1:2]  # and the second as saturation
        value = jnp.ones_like(hue)  # Use full brightness
        hsv = jnp.concatenate([hue, saturation, value], axis=-1)
        rgb = hsv_to_rgb(hsv)
    else:
        # 3 channels or more
        rgb = array[..., -3:]

    return rgb

render_array_with_channels_to_rgba(array)

Render an array with channels as an RGBA image.

This function processes an input array and converts it into an RGBA image based on the number of channels present in the array. The conversion logic is as follows: - If the array has 1 channel, it is repeated across the RGBA channels. - If the array has 2 channels, the first channel is used for RGB, and the second for alpha. - If the array has 3 channels, the first channel is interpreted as hue and the second as saturation. These are converted to RGB using a fixed brightness value, and the last channel is used as the alpha channel. - If the array has 4 or more channels, the last four channels are used directly as RGBA.

The result is premultiplied, as every RGBA array in CAX is (see rgba_to_rgb): the colour built from one to three channels is scaled by the alpha before it is returned, and four or more channels are taken to be premultiplied already.

Parameters:

Name	Type	Description	Default
array	Array	Input array with shape (..., C) and values in [0, 1].	required

Returns:

Type	Description
Array	Premultiplied RGBA array with shape (..., 4) and values in [0, 1].

Source code in src/cax/utils/render.py

def render_array_with_channels_to_rgba(array: Array) -> Array:
    """Render an array with channels as an RGBA image.

    This function processes an input array and converts it into an RGBA image based on
    the number of channels present in the array. The conversion logic is as follows:
    - If the array has 1 channel, it is repeated across the RGBA channels.
    - If the array has 2 channels, the first channel is used for RGB, and the second for
        alpha.
    - If the array has 3 channels, the first channel is interpreted as hue and the
        second as saturation. These are converted to RGB using a fixed brightness value,
        and the last channel is used as the alpha channel.
    - If the array has 4 or more channels, the last four channels are used directly as
        RGBA.

    The result is premultiplied, as every RGBA array in CAX is (see ``rgba_to_rgb``):
    the colour built from one to three channels is scaled by the alpha before it is
    returned, and four or more channels are taken to be premultiplied already.

    Args:
        array: Input array with shape ``(..., C)`` and values in ``[0, 1]``.

    Returns:
        Premultiplied RGBA array with shape ``(..., 4)`` and values in ``[0, 1]``.

    """
    num_channels = array.shape[-1]

    if num_channels == 1:
        # 1 channel
        rgb = jnp.repeat(array, 3, axis=-1)
        alpha = array
    elif num_channels == 2:
        # 2 channels
        rgb = jnp.repeat(array[..., 0:1], 3, axis=-1)
        alpha = array[..., 1:2]
    elif num_channels == 3:
        # 3 channels
        hue = array[..., 0:1]  # Use the first channel as hue
        saturation = array[..., 1:2]  # and the second as saturation
        value = jnp.ones_like(hue)  # Use full brightness
        hsv = jnp.concatenate([hue, saturation, value], axis=-1)
        rgb = hsv_to_rgb(hsv)
        alpha = array[..., 2:3]  # Use the last channel as alpha
    else:
        # 4 or more channels
        return array[..., -4:]

    return jnp.concatenate([rgb * alpha, alpha], axis=-1)

pixel_grid(resolution, *, low=0.0, high=1.0)

Build a square grid of pixel-center coordinates.

Parameters:

Name	Type	Description	Default
resolution	int	Number of pixels along each side.	required
low	float	Coordinate of the first pixel along each axis.	0.0
high	float	Coordinate of the last pixel along each axis.	1.0

Returns:

Type	Description
Array	Array with shape (resolution, resolution, 2) of (x, y) coordinates.

Source code in src/cax/utils/render.py

def pixel_grid(resolution: int, *, low: float = 0.0, high: float = 1.0) -> Array:
    """Build a square grid of pixel-center coordinates.

    Args:
        resolution: Number of pixels along each side.
        low: Coordinate of the first pixel along each axis.
        high: Coordinate of the last pixel along each axis.

    Returns:
        Array with shape ``(resolution, resolution, 2)`` of ``(x, y)`` coordinates.

    """
    x = jnp.linspace(low, high, resolution)
    y = jnp.linspace(low, high, resolution)
    return jnp.stack(jnp.meshgrid(x, y), axis=-1)

nearest_point(grid, points)

Find the nearest of points for every grid pixel.

Parameters:

Name	Type	Description	Default
grid	Array	Pixel coordinates with shape (resolution, resolution, 2).	required
points	Array	Point coordinates with shape (num_points, 2).	required

Returns:

Type	Description
tuple[Array, Array]	A (min_distance_sq, index) tuple of (resolution, resolution) arrays: the squared distance to, and the index of, the nearest point per pixel.

Source code in src/cax/utils/render.py

def nearest_point(grid: Array, points: Array) -> tuple[Array, Array]:
    """Find the nearest of `points` for every grid pixel.

    Args:
        grid: Pixel coordinates with shape ``(resolution, resolution, 2)``.
        points: Point coordinates with shape ``(num_points, 2)``.

    Returns:
        A ``(min_distance_sq, index)`` tuple of ``(resolution, resolution)`` arrays:
            the squared distance to, and the index of, the nearest point per pixel.

    """
    distance_sq = jnp.sum(
        (grid[:, :, None, :] - points[None, None, :, :]) ** 2, axis=-1
    )
    return jnp.min(distance_sq, axis=-1), jnp.argmin(distance_sq, axis=-1)

soft_disk_mask(min_distance_sq, radius)

Anti-aliased disk coverage from squared distances to the nearest point.

Parameters:

Name	Type	Description	Default
min_distance_sq	Array	Squared distance to the nearest point per pixel.	required
radius	float	Disk radius in the grid's coordinate space.	required

Returns:

Type	Description
Array	Coverage in [0, 1]: one at the point, falling to zero at the disk edge.

Source code in src/cax/utils/render.py

def soft_disk_mask(min_distance_sq: Array, radius: float) -> Array:
    """Anti-aliased disk coverage from squared distances to the nearest point.

    Args:
        min_distance_sq: Squared distance to the nearest point per pixel.
        radius: Disk radius in the grid's coordinate space.

    Returns:
        Coverage in ``[0, 1]``: one at the point, falling to zero at the disk edge.

    """
    return jnp.clip(1.0 - min_distance_sq / (radius**2), 0.0, 1.0)

hex_to_square(array)

Resample a triangular-lattice array onto a square pixel grid.

A triangular lattice is stored in an ordinary square array whose axes stand for the lattice vectors (1, 0) and (1/2, sqrt(3)/2) rather than for a Cartesian frame. Drawn directly such an array leans over, because the viewer reads its axes as perpendicular when they are sixty degrees apart. This maps each output pixel back through the basis and samples there, so what is drawn is the lattice as it actually sits in the plane.

The lattice is treated as periodic, matching the wrap-around a cellular automaton on a torus already assumes.

Parameters:

Name	Type	Description	Default
array	Array	Values on a triangular lattice, with shape (..., height, width, channels). Leading axes are treated as batch.	required

Returns:

Type	Description
Array	An array of the same shape, holding the lattice resampled onto square pixels.

Source code in src/cax/utils/render.py

def hex_to_square(array: Array) -> Array:
    """Resample a triangular-lattice array onto a square pixel grid.

    A triangular lattice is stored in an ordinary square array whose axes stand for the
    lattice vectors `(1, 0)` and `(1/2, sqrt(3)/2)` rather than for a Cartesian frame.
    Drawn directly such an array leans over, because the viewer reads its axes as
    perpendicular when they are sixty degrees apart. This maps each output pixel back
    through the basis and samples there, so what is drawn is the lattice as it actually
    sits in the plane.

    The lattice is treated as periodic, matching the wrap-around a cellular automaton on
    a torus already assumes.

    Args:
        array: Values on a triangular lattice, with shape `(..., height, width,
            channels)`. Leading axes are treated as batch.

    Returns:
        An array of the same shape, holding the lattice resampled onto square pixels.

    """
    return _resample(array, jnp.linalg.inv(HEX_BASIS))

square_to_hex(array)

Resample a square-pixel array onto a triangular lattice.

The inverse of hex_to_square, and what a picture needs before it is placed on a triangular lattice: written in directly it would be sheared, and a shape that is no longer itself is no longer sustained by a rule that was tuned to it.

Parameters:

Name	Type	Description	Default
array	Array	Values on square pixels, with shape (..., height, width, channels). Leading axes are treated as batch.	required

Returns:

Type	Description
Array	An array of the same shape, holding the picture resampled onto the lattice.

Source code in src/cax/utils/render.py

def square_to_hex(array: Array) -> Array:
    """Resample a square-pixel array onto a triangular lattice.

    The inverse of `hex_to_square`, and what a picture needs before it is placed on a
    triangular lattice: written in directly it would be sheared, and a shape that is no
    longer itself is no longer sustained by a rule that was tuned to it.

    Args:
        array: Values on square pixels, with shape `(..., height, width, channels)`.
            Leading axes are treated as batch.

    Returns:
        An array of the same shape, holding the picture resampled onto the lattice.

    """
    return _resample(array, HEX_BASIS)

render_states(cs, states, **kwargs)

Render every state of a trajectory, one frame at a time.

Rendering a particle system costs one (resolution^2, num_particles) array per frame, which is far larger than the frame it produces. Vectorizing over a whole trajectory asks for all of them at once — terabytes for a long run — and survives only where the compiler happens to fuse the intermediate away, so the same notebook runs on an accelerator and dies on a CPU. Scanning over the trajectory bounds the peak at a single frame whatever the backend, and costs nothing: the frames are independent.

Parameters:

Name	Type	Description	Default
cs	_Renderable	The complex system, whose render draws one state.	required
states	Any	A trajectory: a pytree whose leaves have the time steps on axis 0.	required
**kwargs	Any	Forwarded to cs.render (resolution, particle_radius, ...).	{}

Returns:

Type	Description
Array	The rendered frames, with shape (num_steps, resolution, resolution, 3).

Source code in src/cax/utils/render.py

def render_states(cs: _Renderable, states: Any, **kwargs: Any) -> Array:
    """Render every state of a trajectory, one frame at a time.

    Rendering a particle system costs one `(resolution^2, num_particles)` array per
    frame, which is far larger than the frame it produces. Vectorizing over a whole
    trajectory asks for all of them at once — terabytes for a long run — and survives
    only where the compiler happens to fuse the intermediate away, so the same notebook
    runs on an accelerator and dies on a CPU. Scanning over the trajectory bounds the
    peak at a single frame whatever the backend, and costs nothing: the frames are
    independent.

    Args:
        cs: The complex system, whose `render` draws one state.
        states: A trajectory: a pytree whose leaves have the time steps on axis 0.
        **kwargs: Forwarded to `cs.render` (`resolution`, `particle_radius`, ...).

    Returns:
        The rendered frames, with shape `(num_steps, resolution, resolution, 3)`.

    """

    def render_fn(cs: _Renderable, state: Any) -> Array:
        return cs.render(state, **kwargs)

    return nnx.scan(
        render_fn,
        in_axes=(nnx.StateAxes({...: nnx.Carry}), 0),
        out_axes=0,
    )(cs, states)

cax.utils.emoji

Utilities for emojis.

get_image_from_url(url)

Fetch an image from a given URL.

Parameters:

Name	Type	Description	Default
url	str	The URL of the image to fetch.	required

Returns:

Type	Description
Image	The fetched image as a PIL Image object.

Raises:

Type	Description
ConnectionError	If the download fails — most commonly because the machine is offline. The original error is chained.

Source code in src/cax/utils/emoji.py

def get_image_from_url(url: str) -> Image:
    """Fetch an image from a given URL.

    Args:
        url: The URL of the image to fetch.

    Returns:
        The fetched image as a PIL Image object.

    Raises:
        ConnectionError: If the download fails — most commonly because the machine
            is offline. The original error is chained.

    """
    try:
        with urlopen(url, timeout=_FETCH_TIMEOUT_S) as response:
            image_data = response.read()
    except (URLError, TimeoutError) as error:
        raise ConnectionError(
            f"Could not download {url}. Emoji images are fetched from the network at "
            f"call time; check the connection and retry."
        ) from error

    image_pil = PIL.Image.open(io.BytesIO(image_data))
    return image_pil

get_emoji_filename(emoji)

Build the Noto Emoji filename for an emoji.

Noto names a glyph after the codepoints that spell it, in lowercase hexadecimal, joined by underscores. Sequences are spelled out in full, so the zero-width joiner of a glyph like 👨‍💻 is part of the name, while the variation selector that merely asks for an emoji presentation is not.

Parameters:

Name	Type	Description	Default
emoji	str	The emoji character or sequence.	required

Returns:

Type	Description
str	The filename, such as emoji_u1f468_200d_1f4bb.png.

Raises:

Type	Description
ValueError	If emoji is empty, or is a flag. Noto keeps flags apart from the rest, named by country code rather than by codepoint.

Source code in src/cax/utils/emoji.py

def get_emoji_filename(emoji: str) -> str:
    """Build the Noto Emoji filename for an emoji.

    Noto names a glyph after the codepoints that spell it, in lowercase hexadecimal,
    joined by underscores. Sequences are spelled out in full, so the zero-width joiner
    of a glyph like 👨‍💻 is part of the name, while the variation selector that merely
    asks for an emoji presentation is not.

    Args:
        emoji: The emoji character or sequence.

    Returns:
        The filename, such as ``emoji_u1f468_200d_1f4bb.png``.

    Raises:
        ValueError: If ``emoji`` is empty, or is a flag. Noto keeps flags apart from the
            rest, named by country code rather than by codepoint.

    """
    codepoints = [
        ord(character) for character in emoji if character != _EMOJI_PRESENTATION
    ]
    if not codepoints:
        raise ValueError("Cannot build a filename for an empty emoji.")
    if any(codepoint in _REGIONAL_INDICATORS for codepoint in codepoints):
        raise ValueError(
            f"Flags such as {emoji!r} are not available: Noto Emoji stores them apart "
            f"from the other glyphs, named by country code rather than by codepoint."
        )

    return "emoji_u" + "_".join(f"{codepoint:x}" for codepoint in codepoints) + ".png"

get_emoji(emoji) cached

Fetch and return an emoji as a PIL Image.

The glyph is downloaded from Google's Noto Emoji (PNG, 128 px) and cached in memory, so repeated calls for the same emoji fetch once. The image is returned without further processing; callers may convert to arrays or resize as needed.

Parameters:

Name	Type	Description	Default
emoji	str	The emoji character or sequence to fetch.	required

Returns:

Type	Description
Image	A PIL.Image.Image instance containing the emoji.

Raises:

Type	Description
ValueError	If the emoji has no Noto glyph under this naming scheme.
ConnectionError	If the download fails.

Source code in src/cax/utils/emoji.py

@cache
def get_emoji(emoji: str) -> Image:
    """Fetch and return an emoji as a PIL Image.

    The glyph is downloaded from Google's Noto Emoji (PNG, 128 px) and cached in memory,
    so repeated calls for the same emoji fetch once. The image is returned without
    further processing; callers may convert to arrays or resize as needed.

    Args:
        emoji: The emoji character or sequence to fetch.

    Returns:
        A ``PIL.Image.Image`` instance containing the emoji.

    Raises:
        ValueError: If the emoji has no Noto glyph under this naming scheme.
        ConnectionError: If the download fails.

    """
    filename = get_emoji_filename(emoji)
    url = (
        f"https://cdn.jsdelivr.net/gh/googlefonts/noto-emoji@{_NOTO_EMOJI_COMMIT}"
        f"/png/128/{filename}"
    )
    return get_image_from_url(url)

get_emoji_array(emoji, size, pad_width=0)

Fetch an emoji as a padded, premultiplied RGBA array.

The glyph is resized to size and framed in transparent pixels, which is what a growing cellular automaton needs: the target sits in the middle of a larger grid, so the automaton has somewhere to overshoot into and can be penalised for doing so.

Colour is premultiplied by alpha, CAX's convention for RGBA arrays: each pixel holds the light it emits, so a transparent pixel is zero in every channel and a loss on the array measures what is seen rather than the colour a PNG stores behind invisible pixels. rgba_to_rgb composites arrays in this convention.

Parameters:

Name	Type	Description	Default
emoji	str	The emoji character or sequence to fetch.	required
size	int	Width and height, in pixels, to resize the glyph to.	required
pad_width	int	Transparent pixels to add on each side.	0

Returns:

Type	Description
Array	An array of shape (size + 2 * pad_width, size + 2 * pad_width, 4) holding
Array	premultiplied RGBA values in the unit interval.

Raises:

Type	Description
ValueError	If the emoji has no Noto glyph under this naming scheme.
ConnectionError	If the download fails.

Source code in src/cax/utils/emoji.py

def get_emoji_array(emoji: str, size: int, pad_width: int = 0) -> Array:
    """Fetch an emoji as a padded, premultiplied RGBA array.

    The glyph is resized to ``size`` and framed in transparent pixels, which is what a
    growing cellular automaton needs: the target sits in the middle of a larger grid, so
    the automaton has somewhere to overshoot into and can be penalised for doing so.

    Colour is premultiplied by alpha, CAX's convention for RGBA arrays: each pixel
    holds the light it emits, so a transparent pixel is zero in every channel and a
    loss on the array measures what is seen rather than the colour a PNG stores behind
    invisible pixels. ``rgba_to_rgb`` composites arrays in this convention.

    Args:
        emoji: The emoji character or sequence to fetch.
        size: Width and height, in pixels, to resize the glyph to.
        pad_width: Transparent pixels to add on each side.

    Returns:
        An array of shape ``(size + 2 * pad_width, size + 2 * pad_width, 4)`` holding
        premultiplied RGBA values in the unit interval.

    Raises:
        ValueError: If the emoji has no Noto glyph under this naming scheme.
        ConnectionError: If the download fails.

    """
    image_pil = get_emoji(emoji).resize(
        (size, size), resample=PIL.Image.Resampling.LANCZOS
    )
    array = jnp.asarray(image_pil, dtype=jnp.float32) / 255.0
    array = array.at[..., :3].multiply(array[..., 3:])
    return jnp.pad(array, ((pad_width, pad_width), (pad_width, pad_width), (0, 0)))

cax.utils.numerics

Numerically safe primitives for differentiable simulation.

JAX propagates cotangents through both branches of jnp.where, so masking an invalid value after it has been computed leaves nan in the gradient even when the forward pass is finite — the documented "where-NaN" trap. The library-wide convention is therefore to sanitize the input of the unsafe operation, not its output: every division, norm, or singular kernel evaluated where its argument can be degenerate goes through one of these helpers (or repeats their double-where pattern inline, with a comment naming it).

safe_divide(numerator, denominator, *, where)

Divide two arrays, returning zero and a clean gradient where invalid.

Uses the double-where pattern: the denominator is replaced by one at invalid positions before dividing, so neither the forward value nor the gradient ever touches the singular point.

Parameters:

Name	Type	Description	Default
numerator	Array	Numerator array.	required
denominator	Array	Denominator array, broadcastable against numerator.	required
where	Array	Boolean mask, true where the division is valid. Broadcastable against the result; invalid positions yield zero.	required

Returns:

Type	Description
Array	numerator / denominator where where is true, zero elsewhere, with gradients that are finite everywhere the inputs are.

Source code in src/cax/utils/numerics.py

def safe_divide(numerator: Array, denominator: Array, *, where: Array) -> Array:
    """Divide two arrays, returning zero and a clean gradient where invalid.

    Uses the double-`where` pattern: the denominator is replaced by one at invalid
    positions *before* dividing, so neither the forward value nor the gradient ever
    touches the singular point.

    Args:
        numerator: Numerator array.
        denominator: Denominator array, broadcastable against `numerator`.
        where: Boolean mask, true where the division is valid. Broadcastable against
            the result; invalid positions yield zero.

    Returns:
        `numerator / denominator` where `where` is true, zero elsewhere, with gradients
            that are finite everywhere the inputs are.

    """
    denominator_safe = jnp.where(where, denominator, jnp.ones_like(denominator))
    return jnp.where(where, numerator / denominator_safe, jnp.zeros_like(numerator))

safe_norm(vector, *, axis=-1, keepdims=False)

Euclidean norm with a finite gradient at the origin.

jnp.linalg.norm differentiates to x / ||x||, which is nan at zero. This computes the same value but returns a zero gradient at the origin, which is the convention steering and force computations want: a vanished vector exerts no pull.

Parameters:

Name	Type	Description	Default
vector	Array	Input array.	required
axis	int	Axis holding the vector components.	-1
keepdims	bool	Whether the reduced axis is kept with size one.	False

Returns:

Type	Description
Array	Norm of vector along axis, with gradient zero where the norm is zero.

Source code in src/cax/utils/numerics.py

def safe_norm(vector: Array, *, axis: int = -1, keepdims: bool = False) -> Array:
    """Euclidean norm with a finite gradient at the origin.

    `jnp.linalg.norm` differentiates to `x / ||x||`, which is `nan` at zero. This
    computes the same value but returns a zero gradient at the origin, which is the
    convention steering and force computations want: a vanished vector exerts no pull.

    Args:
        vector: Input array.
        axis: Axis holding the vector components.
        keepdims: Whether the reduced axis is kept with size one.

    Returns:
        Norm of `vector` along `axis`, with gradient zero where the norm is zero.

    """
    squared = jnp.sum(jnp.square(vector), axis=axis, keepdims=keepdims)
    is_positive = squared > 0.0
    squared_safe = jnp.where(is_positive, squared, jnp.ones_like(squared))
    return jnp.where(is_positive, jnp.sqrt(squared_safe), jnp.zeros_like(squared))

cax.utils.dynamics

Shared dynamics primitives for particle systems on the unit torus.

Boids and Particle Life integrate the same way: exponential velocity damping, a semi-implicit Euler step, and periodic boundary conditions. These helpers hold that shared physics in one place; each system keeps its own defaults and state types.

toroidal_difference(position_1, position_2, *, period=1.0)

Minimum-image vector from position_1 to position_2 on a torus.

Applies periodic boundary conditions component-wise so each component of the result lies in [-period / 2, period / 2] — the shortest displacement on a torus of the given period.

Parameters:

Name	Type	Description	Default
position_1	Array	Start positions.	required
position_2	Array	End positions, broadcastable against position_1.	required
period	float	Length of the torus along every axis.	1.0

Returns:

Type	Description
Array	Component-wise shortest displacement from position_1 to position_2.

Source code in src/cax/utils/dynamics.py

def toroidal_difference(
    position_1: Array, position_2: Array, *, period: float = 1.0
) -> Array:
    """Minimum-image vector from `position_1` to `position_2` on a torus.

    Applies periodic boundary conditions component-wise so each component of the
    result lies in `[-period / 2, period / 2]` — the shortest displacement on a
    torus of the given period.

    Args:
        position_1: Start positions.
        position_2: End positions, broadcastable against `position_1`.
        period: Length of the torus along every axis.

    Returns:
        Component-wise shortest displacement from `position_1` to `position_2`.

    """
    difference = position_2 - position_1
    difference = jnp.where(difference > period / 2, difference - period, difference)
    difference = jnp.where(difference < -period / 2, difference + period, difference)
    return difference

damped_euler_step(position, velocity, acceleration, *, dt, friction_factor, period=1.0)

Semi-implicit Euler step with velocity damping and periodic boundaries.

The velocity is damped by friction_factor (typically 0.5 ** (dt / half_life)) and accelerated, then the position is advanced with the new velocity and wrapped onto the torus.

Parameters:

Name	Type	Description	Default
position	Array	Positions on the torus.	required
velocity	Array	Velocities.	required
acceleration	Array	Accelerations from the perception step.	required
dt	float	Time step.	required
friction_factor	float	Multiplicative velocity decay per step.	required
period	float	Length of the torus along every axis.	1.0

Returns:

Type	Description
tuple[Array, Array]	A (position, velocity) tuple after one step.

Source code in src/cax/utils/dynamics.py

def damped_euler_step(
    position: Array,
    velocity: Array,
    acceleration: Array,
    *,
    dt: float,
    friction_factor: float,
    period: float = 1.0,
) -> tuple[Array, Array]:
    """Semi-implicit Euler step with velocity damping and periodic boundaries.

    The velocity is damped by `friction_factor` (typically `0.5 ** (dt / half_life)`)
    and accelerated, then the position is advanced with the *new* velocity and wrapped
    onto the torus.

    Args:
        position: Positions on the torus.
        velocity: Velocities.
        acceleration: Accelerations from the perception step.
        dt: Time step.
        friction_factor: Multiplicative velocity decay per step.
        period: Length of the torus along every axis.

    Returns:
        A `(position, velocity)` tuple after one step.

    """
    velocity = friction_factor * velocity + dt * acceleration
    position = (position + dt * velocity) % period
    return position, velocity