Skip to content

Boids

cax.cs.boids.cs.Boids

Bases: ComplexSystem[BoidsState, Array]

Boids class.

Source code in src/cax/cs/boids/cs.py 
 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
 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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
class Boids(ComplexSystem[BoidsState, Array]):
    """Boids class."""

    def __init__(
        self,
        *,
        dt: float = 0.01,
        velocity_half_life: float = jnp.inf,
        boid_policy: BoidsPolicy,
    ):
        """Initialize Boids.

        Args:
            dt: Time step of the simulation in arbitrary time units. Smaller values
                produce smoother motion but require more steps for the same duration.
            velocity_half_life: Time constant for velocity decay due to friction. After
                this time, velocity is halved without steering input. Use jnp.inf for no
                friction. Smaller values create more damped, sluggish motion.
            boid_policy: Policy defining the behavior of the boids.

        """
        self.perceive = BoidsPerceive(
            boid_policy=boid_policy,
        )
        self.update = BoidsUpdate(
            dt=dt,
            velocity_half_life=velocity_half_life,
        )

    @override
    def _step(self, state: BoidsState, input: Array | None = None) -> BoidsState:
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        return next_state

    @nnx.jit(static_argnames=("resolution", "particle_radius"))
    @override
    def render(
        self,
        state: BoidsState,
        *,
        resolution: int = 512,
        particle_radius: float = 0.01,
    ) -> Array:
        """Render state to RGB image.

        Renders boids as triangular agents pointing in their direction of motion on a
        white background. Each boid is drawn as a filled triangle with the tip pointing
        in the velocity direction, providing visual feedback about both position and
        heading. The visualization uses 2D coordinates in the range [0, 1].

        Args:
            state: BoidsState containing position and velocity arrays. Position should
                have shape (num_boids, 2) with coordinates in [0, 1]. Velocity
                determines the orientation and can have arbitrary magnitude.
            resolution: Size of the output image in pixels for both width and height.
                Higher values produce smoother, more detailed renderings.
            particle_radius: Half-extent of each boid glyph in coordinate space [0, 1]:
                the triangle spans twice this value from base to tip and one radius
                across the base. Larger values make boids more visible but may cause
                overlap.

        Returns:
            RGB image with dtype uint8 and shape (resolution, resolution, 3), where
                boids appear as black triangles on a white background.

        """
        if state.position.shape[-1] != 2:
            raise ValueError("Boids only supports 2D visualization.")

        # Adjust coordinates for rendering
        # - Simulation has y increasing upwards (y=0 bottom, y=1 top).
        # - Image has y increasing downwards (y=0 top, y=1 bottom).
        # - Flip position y: map simulation y to image y with (1 - y).
        # - Negate velocity y: adjust direction to match flipped coordinates.
        position = state.position  # Shape: (num_boids, 2)
        position = position.at[:, 1].set(1 - position[:, 1])  # Flip y-coordinate
        velocity = state.velocity  # Shape: (num_boids, 2)
        velocity = velocity.at[:, 1].set(-velocity[:, 1])  # Negate y-component

        # Compute unit velocity and perpendicular vectors
        v_norm = jnp.linalg.norm(velocity, axis=-1, keepdims=True)
        v_hat = velocity / (v_norm + 1e-8)
        v_perp = jnp.stack([-v_hat[..., 1], v_hat[..., 0]], axis=-1)

        # Define triangle dimensions
        h = 2 * particle_radius  # Height from base to tip
        w = particle_radius  # Base width

        # Compute triangle vertices
        vertex0 = position - (w / 2) * v_perp  # Base left
        vertex1 = position + h * v_hat  # Tip
        vertex2 = position + (w / 2) * v_perp  # Base right
        vertices = jnp.stack(
            [vertex0, vertex1, vertex2], axis=1
        )  # Shape: (num_boids, 3, 2)

        # Rasterize: nearest boid per pixel
        grid = pixel_grid(resolution)  # Shape: (resolution, resolution, 2)
        _, closest_idx = nearest_point(
            grid, position
        )  # Shape: (resolution, resolution)

        # Get vertices of the closest boid
        closest_vertices = vertices[
            closest_idx, :, :
        ]  # Shape: (resolution, resolution, 3, 2)

        # Extract vertices
        a = closest_vertices[..., 0, :]  # vertex0
        b = closest_vertices[..., 1, :]  # vertex1
        c = closest_vertices[..., 2, :]  # vertex2

        # Compute edge vectors
        edge0 = b - a
        edge1 = c - b
        edge2 = a - c

        # Compute cross products for point-in-triangle test
        q = grid
        cross0 = edge0[..., 0] * (q - a)[..., 1] - edge0[..., 1] * (q - a)[..., 0]
        cross1 = edge1[..., 0] * (q - b)[..., 1] - edge1[..., 1] * (q - b)[..., 0]
        cross2 = edge2[..., 0] * (q - c)[..., 1] - edge2[..., 1] * (q - c)[..., 0]

        # Determine if pixel is inside the triangle
        inside = (
            (cross0 > 0) & (cross1 > 0) & (cross2 > 0)
        )  # Shape: (resolution, resolution)

        # Create RGB image
        gray = inside[..., None].astype(
            jnp.float32
        )  # Shape: (resolution, resolution, 1)
        rgb = jnp.repeat(gray, 3, axis=-1)  # Shape: (resolution, resolution, 3)

        return clip_and_uint8(rgb)

__init__(*, dt=0.01, velocity_half_life=jnp.inf, boid_policy)

Initialize Boids.

Parameters:

Name Type Description Default
dt float

Time step of the simulation in arbitrary time units. Smaller values produce smoother motion but require more steps for the same duration.

 0.01
velocity_half_life float

Time constant for velocity decay due to friction. After this time, velocity is halved without steering input. Use jnp.inf for no friction. Smaller values create more damped, sluggish motion.

 inf
boid_policy BoidsPolicy

Policy defining the behavior of the boids.

 required
Source code in src/cax/cs/boids/cs.py 
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 __init__(
    self,
    *,
    dt: float = 0.01,
    velocity_half_life: float = jnp.inf,
    boid_policy: BoidsPolicy,
):
    """Initialize Boids.

    Args:
        dt: Time step of the simulation in arbitrary time units. Smaller values
            produce smoother motion but require more steps for the same duration.
        velocity_half_life: Time constant for velocity decay due to friction. After
            this time, velocity is halved without steering input. Use jnp.inf for no
            friction. Smaller values create more damped, sluggish motion.
        boid_policy: Policy defining the behavior of the boids.

    """
    self.perceive = BoidsPerceive(
        boid_policy=boid_policy,
    )
    self.update = BoidsUpdate(
        dt=dt,
        velocity_half_life=velocity_half_life,
    )

render(state, *, resolution=512, particle_radius=0.01)

Render state to RGB image.

Renders boids as triangular agents pointing in their direction of motion on a white background. Each boid is drawn as a filled triangle with the tip pointing in the velocity direction, providing visual feedback about both position and heading. The visualization uses 2D coordinates in the range [0, 1].

Parameters:

Name Type Description Default
state BoidsState

BoidsState containing position and velocity arrays. Position should have shape (num_boids, 2) with coordinates in [0, 1]. Velocity determines the orientation and can have arbitrary magnitude.

 required
resolution int

Size of the output image in pixels for both width and height. Higher values produce smoother, more detailed renderings.

 512
particle_radius float

Half-extent of each boid glyph in coordinate space [0, 1]: the triangle spans twice this value from base to tip and one radius across the base. Larger values make boids more visible but may cause overlap.

 0.01

Returns:

Type Description
Array

RGB image with dtype uint8 and shape (resolution, resolution, 3), where boids appear as black triangles on a white background.
Source code in src/cax/cs/boids/cs.py 
 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
 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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
@nnx.jit(static_argnames=("resolution", "particle_radius"))
@override
def render(
    self,
    state: BoidsState,
    *,
    resolution: int = 512,
    particle_radius: float = 0.01,
) -> Array:
    """Render state to RGB image.

    Renders boids as triangular agents pointing in their direction of motion on a
    white background. Each boid is drawn as a filled triangle with the tip pointing
    in the velocity direction, providing visual feedback about both position and
    heading. The visualization uses 2D coordinates in the range [0, 1].

    Args:
        state: BoidsState containing position and velocity arrays. Position should
            have shape (num_boids, 2) with coordinates in [0, 1]. Velocity
            determines the orientation and can have arbitrary magnitude.
        resolution: Size of the output image in pixels for both width and height.
            Higher values produce smoother, more detailed renderings.
        particle_radius: Half-extent of each boid glyph in coordinate space [0, 1]:
            the triangle spans twice this value from base to tip and one radius
            across the base. Larger values make boids more visible but may cause
            overlap.

    Returns:
        RGB image with dtype uint8 and shape (resolution, resolution, 3), where
            boids appear as black triangles on a white background.

    """
    if state.position.shape[-1] != 2:
        raise ValueError("Boids only supports 2D visualization.")

    # Adjust coordinates for rendering
    # - Simulation has y increasing upwards (y=0 bottom, y=1 top).
    # - Image has y increasing downwards (y=0 top, y=1 bottom).
    # - Flip position y: map simulation y to image y with (1 - y).
    # - Negate velocity y: adjust direction to match flipped coordinates.
    position = state.position  # Shape: (num_boids, 2)
    position = position.at[:, 1].set(1 - position[:, 1])  # Flip y-coordinate
    velocity = state.velocity  # Shape: (num_boids, 2)
    velocity = velocity.at[:, 1].set(-velocity[:, 1])  # Negate y-component

    # Compute unit velocity and perpendicular vectors
    v_norm = jnp.linalg.norm(velocity, axis=-1, keepdims=True)
    v_hat = velocity / (v_norm + 1e-8)
    v_perp = jnp.stack([-v_hat[..., 1], v_hat[..., 0]], axis=-1)

    # Define triangle dimensions
    h = 2 * particle_radius  # Height from base to tip
    w = particle_radius  # Base width

    # Compute triangle vertices
    vertex0 = position - (w / 2) * v_perp  # Base left
    vertex1 = position + h * v_hat  # Tip
    vertex2 = position + (w / 2) * v_perp  # Base right
    vertices = jnp.stack(
        [vertex0, vertex1, vertex2], axis=1
    )  # Shape: (num_boids, 3, 2)

    # Rasterize: nearest boid per pixel
    grid = pixel_grid(resolution)  # Shape: (resolution, resolution, 2)
    _, closest_idx = nearest_point(
        grid, position
    )  # Shape: (resolution, resolution)

    # Get vertices of the closest boid
    closest_vertices = vertices[
        closest_idx, :, :
    ]  # Shape: (resolution, resolution, 3, 2)

    # Extract vertices
    a = closest_vertices[..., 0, :]  # vertex0
    b = closest_vertices[..., 1, :]  # vertex1
    c = closest_vertices[..., 2, :]  # vertex2

    # Compute edge vectors
    edge0 = b - a
    edge1 = c - b
    edge2 = a - c

    # Compute cross products for point-in-triangle test
    q = grid
    cross0 = edge0[..., 0] * (q - a)[..., 1] - edge0[..., 1] * (q - a)[..., 0]
    cross1 = edge1[..., 0] * (q - b)[..., 1] - edge1[..., 1] * (q - b)[..., 0]
    cross2 = edge2[..., 0] * (q - c)[..., 1] - edge2[..., 1] * (q - c)[..., 0]

    # Determine if pixel is inside the triangle
    inside = (
        (cross0 > 0) & (cross1 > 0) & (cross2 > 0)
    )  # Shape: (resolution, resolution)

    # Create RGB image
    gray = inside[..., None].astype(
        jnp.float32
    )  # Shape: (resolution, resolution, 1)
    rgb = jnp.repeat(gray, 3, axis=-1)  # Shape: (resolution, resolution, 3)

    return clip_and_uint8(rgb)

__call__(state, input=None, *, num_steps=1, input_in_axis=None, return_states=False)

Step the system for multiple time steps.

This method wraps _step inside a JAX scan for efficiency and JIT-compiles the loop. If input is time-varying, set input_in_axis to the axis containing the time dimension so that each step receives the corresponding slice of input.

Under return_states=True, the per-step states are also returned as the scan's stacked outputs, mirroring the (carry, ys) convention of jax.lax.scan. The trajectory holds the state after each step, stacked along a new leading axis of size num_steps — its first element is the state after one step, its last equals the final state, and the initial state is not included.

When remat is enabled, the scan body is wrapped with nnx.remat to reduce memory usage during backpropagation at the cost of recomputing intermediates.

Note that num_steps, input_in_axis, and return_states are static: each distinct combination compiles once, so sweeps over horizons should batch their step counts.

Parameters:

Name Type Description Default
state State

Current state.

 required
input Input | None

Optional input.

 None
num_steps int

Number of steps.

 1
input_in_axis int | None

Axis for input if provided for each step.

 None
return_states bool

Whether to also return the stacked per-step states.

 False

Returns:

Type Description
State | tuple[State, State]

Final state after num_steps applications of _step, or a (final_state, states) tuple under return_states=True, where states stacks the per-step states along a new leading axis of size num_steps.
Source code in src/cax/core/cs.py 
 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
 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
@nnx.jit(static_argnames=("num_steps", "input_in_axis", "return_states"))
def __call__(
    self,
    state: State,
    input: Input | None = None,
    *,
    num_steps: int = 1,
    input_in_axis: int | None = None,
    return_states: bool = False,
) -> State | tuple[State, State]:
    """Step the system for multiple time steps.

    This method wraps `_step` inside a JAX scan for efficiency and JIT-compiles the
    loop. If `input` is time-varying, set `input_in_axis` to the axis containing the
    time dimension so that each step receives the corresponding slice of input.

    Under `return_states=True`, the per-step states are also returned as the scan's
    stacked outputs, mirroring the `(carry, ys)` convention of `jax.lax.scan`. The
    trajectory holds the state *after* each step, stacked along a new leading axis
    of size `num_steps` — its first element is the state after one step, its last
    equals the final state, and the initial state is not included.

    When `remat` is enabled, the scan body is wrapped with `nnx.remat` to reduce
    memory usage during backpropagation at the cost of recomputing intermediates.

    Note that `num_steps`, `input_in_axis`, and `return_states` are static: each
    distinct combination compiles once, so sweeps over horizons should batch their
    step counts.

    Args:
        state: Current state.
        input: Optional input.
        num_steps: Number of steps.
        input_in_axis: Axis for input if provided for each step.
        return_states: Whether to also return the stacked per-step states.

    Returns:
        Final state after `num_steps` applications of `_step`, or a
            `(final_state, states)` tuple under `return_states=True`, where `states`
            stacks the per-step states along a new leading axis of size `num_steps`.

    """

    def step_fn(
        cs: ComplexSystem[State, Input], state: State, input: Input | None
    ) -> tuple[State, State | None]:
        next_state = cs._step(state, input)
        return next_state, (next_state if return_states else None)

    if self.remat:
        step_fn = nnx.remat(step_fn)

    state, states = nnx.scan(
        step_fn,
        in_axes=(nnx.StateAxes({...: nnx.Carry}), nnx.Carry, input_in_axis),
        out_axes=(nnx.Carry, 0),
        length=num_steps,
    )(self, state, input)

    return (state, states) if return_states else state