Skip to content

Particle Life

cax.cs.particle_life.cs.ParticleLife

Bases: ComplexSystem[ParticleLifeState, Array]

Particle Life class.

Source code in src/cax/cs/particle_life/cs.py 
 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
161
162
163
164
165
166
167
class ParticleLife(ComplexSystem[ParticleLifeState, Array]):
    """Particle Life class."""

    def __init__(
        self,
        *,
        num_classes: int,
        dt: float = 0.01,
        force_factor: float = 1.0,
        velocity_half_life: float = 0.01,
        r_max: float = 0.15,
        beta: float = 0.3,
        attraction_matrix: Array,
    ):
        """Initialize Particle Life.

        Args:
            num_classes: Number of distinct particle types (classes). Each type can have
                different interactions with other types as specified in the attraction
                matrix.
            dt: Time step of the simulation in arbitrary time units. Smaller values
                produce smoother motion but require more steps for the same duration.
            force_factor: Global scaling factor for all interaction forces. Higher
                values create stronger, more dynamic interactions.
            velocity_half_life: Time constant for velocity decay due to friction. After
                this time, velocity is halved without force input. Smaller values create
                more damped, viscous dynamics.
            r_max: Maximum interaction distance in coordinate space [0, 1]. Particles
                beyond this distance do not interact. Larger values increase computation
                cost.
            beta: Distance threshold parameter controlling the transition from repulsion
                to attraction. Typically in range [0, 1], where smaller values create
                stronger short-range repulsion.
            attraction_matrix: Attraction matrix of shape (num_classes, num_classes)
                where entry (i, j) defines the attraction strength from type i to type
                j. Positive values attract, negative values repel. Values typically
                range from -1 to 1.

        """
        self.num_classes = num_classes

        self.perceive = ParticleLifePerceive(
            force_factor=force_factor,
            r_max=r_max,
            beta=beta,
            attraction_matrix=attraction_matrix,
        )
        self.update = ParticleLifeUpdate(
            dt=dt,
            velocity_half_life=velocity_half_life,
        )

    @override
    def _step(
        self, state: ParticleLifeState, input: Array | None = None
    ) -> ParticleLifeState:
        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: ParticleLifeState,
        *,
        resolution: int = 512,
        particle_radius: float = 0.005,
    ) -> Array:
        """Render state to RGB image.

        Renders particles as colored circles on a black background. Each particle type
        (class) is assigned a distinct hue from the color spectrum, with colors evenly
        distributed across the HSV color space. Particles are drawn with smooth
        anti-aliased edges based on their distance from pixel centers. The visualization
        uses 2D coordinates in the range [0, 1].

        Args:
            state: ParticleLifeState containing class_id, position, and velocity arrays.
                Position should have shape (num_particles, 2) with coordinates in
                [0, 1]. Class array determines the color of each particle.
            resolution: Size of the output image in pixels for both width and height.
                Higher values produce smoother, more detailed renderings.
            particle_radius: Radius of each particle in coordinate space [0, 1].
                Particles appear as smooth circles with this radius. Larger values make
                particles more visible but may cause overlap.

        Returns:
            RGB image with dtype uint8 and shape (resolution, resolution, 3), where
                particles appear as colored circles on a black background, with colors
                determined by particle type.

        """
        if state.position.shape[-1] != 2:
            raise ValueError("Particle Life 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).
        positions = state.position  # Shape: (num_particles, 2)
        positions = positions.at[:, 1].set(1 - positions[:, 1])

        # Rasterize: nearest particle per pixel
        grid = pixel_grid(resolution)  # Shape: (resolution, resolution, 2)
        min_distance_sq, closest_particle_idx = nearest_point(grid, positions)

        # Get class of the closest particle for each pixel
        closest_class = state.class_id[
            closest_particle_idx
        ]  # Shape: (resolution, resolution)

        # Compute smooth mask based on distance to closest particle
        mask = soft_disk_mask(
            min_distance_sq, particle_radius
        )  # Shape: (resolution, resolution)

        # Generate colors for each class using HSV
        hues = jnp.linspace(0, 1, self.num_classes, endpoint=False)
        hsv = jnp.stack([hues, jnp.ones_like(hues), jnp.ones_like(hues)], axis=-1)
        colors = hsv_to_rgb(hsv)  # Shape: (num_classes, 3)

        # Assign colors based on closest particle's class
        particle_colors = colors[closest_class]  # Shape: (resolution, resolution, 3)

        # Create black background
        background = jnp.zeros(
            (resolution, resolution, 3)
        )  # Shape: (resolution, resolution, 3)

        # Blend particle colors with background using the mask
        rgb = (
            background * (1.0 - mask[..., None]) + particle_colors * mask[..., None]
        )  # Shape: (resolution, resolution, 3)

        return clip_and_uint8(rgb)

__init__(*, num_classes, dt=0.01, force_factor=1.0, velocity_half_life=0.01, r_max=0.15, beta=0.3, attraction_matrix)

Initialize Particle Life.

Parameters:

Name Type Description Default
num_classes int

Number of distinct particle types (classes). Each type can have different interactions with other types as specified in the attraction matrix.

 required
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
force_factor float

Global scaling factor for all interaction forces. Higher values create stronger, more dynamic interactions.

 1.0
velocity_half_life float

Time constant for velocity decay due to friction. After this time, velocity is halved without force input. Smaller values create more damped, viscous dynamics.

 0.01
r_max float

Maximum interaction distance in coordinate space [0, 1]. Particles beyond this distance do not interact. Larger values increase computation cost.

 0.15
beta float

Distance threshold parameter controlling the transition from repulsion to attraction. Typically in range [0, 1], where smaller values create stronger short-range repulsion.

 0.3
attraction_matrix Array

Attraction matrix of shape (num_classes, num_classes) where entry (i, j) defines the attraction strength from type i to type j. Positive values attract, negative values repel. Values typically range from -1 to 1.

 required
Source code in src/cax/cs/particle_life/cs.py 
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
def __init__(
    self,
    *,
    num_classes: int,
    dt: float = 0.01,
    force_factor: float = 1.0,
    velocity_half_life: float = 0.01,
    r_max: float = 0.15,
    beta: float = 0.3,
    attraction_matrix: Array,
):
    """Initialize Particle Life.

    Args:
        num_classes: Number of distinct particle types (classes). Each type can have
            different interactions with other types as specified in the attraction
            matrix.
        dt: Time step of the simulation in arbitrary time units. Smaller values
            produce smoother motion but require more steps for the same duration.
        force_factor: Global scaling factor for all interaction forces. Higher
            values create stronger, more dynamic interactions.
        velocity_half_life: Time constant for velocity decay due to friction. After
            this time, velocity is halved without force input. Smaller values create
            more damped, viscous dynamics.
        r_max: Maximum interaction distance in coordinate space [0, 1]. Particles
            beyond this distance do not interact. Larger values increase computation
            cost.
        beta: Distance threshold parameter controlling the transition from repulsion
            to attraction. Typically in range [0, 1], where smaller values create
            stronger short-range repulsion.
        attraction_matrix: Attraction matrix of shape (num_classes, num_classes)
            where entry (i, j) defines the attraction strength from type i to type
            j. Positive values attract, negative values repel. Values typically
            range from -1 to 1.

    """
    self.num_classes = num_classes

    self.perceive = ParticleLifePerceive(
        force_factor=force_factor,
        r_max=r_max,
        beta=beta,
        attraction_matrix=attraction_matrix,
    )
    self.update = ParticleLifeUpdate(
        dt=dt,
        velocity_half_life=velocity_half_life,
    )

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

Render state to RGB image.

Renders particles as colored circles on a black background. Each particle type (class) is assigned a distinct hue from the color spectrum, with colors evenly distributed across the HSV color space. Particles are drawn with smooth anti-aliased edges based on their distance from pixel centers. The visualization uses 2D coordinates in the range [0, 1].

Parameters:

Name Type Description Default
state ParticleLifeState

ParticleLifeState containing class_id, position, and velocity arrays. Position should have shape (num_particles, 2) with coordinates in [0, 1]. Class array determines the color of each particle.

 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

Radius of each particle in coordinate space [0, 1]. Particles appear as smooth circles with this radius. Larger values make particles more visible but may cause overlap.

 0.005

Returns:

Type Description
Array

RGB image with dtype uint8 and shape (resolution, resolution, 3), where particles appear as colored circles on a black background, with colors determined by particle type.
Source code in src/cax/cs/particle_life/cs.py 
 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
161
162
163
164
165
166
167
@nnx.jit(static_argnames=("resolution", "particle_radius"))
@override
def render(
    self,
    state: ParticleLifeState,
    *,
    resolution: int = 512,
    particle_radius: float = 0.005,
) -> Array:
    """Render state to RGB image.

    Renders particles as colored circles on a black background. Each particle type
    (class) is assigned a distinct hue from the color spectrum, with colors evenly
    distributed across the HSV color space. Particles are drawn with smooth
    anti-aliased edges based on their distance from pixel centers. The visualization
    uses 2D coordinates in the range [0, 1].

    Args:
        state: ParticleLifeState containing class_id, position, and velocity arrays.
            Position should have shape (num_particles, 2) with coordinates in
            [0, 1]. Class array determines the color of each particle.
        resolution: Size of the output image in pixels for both width and height.
            Higher values produce smoother, more detailed renderings.
        particle_radius: Radius of each particle in coordinate space [0, 1].
            Particles appear as smooth circles with this radius. Larger values make
            particles more visible but may cause overlap.

    Returns:
        RGB image with dtype uint8 and shape (resolution, resolution, 3), where
            particles appear as colored circles on a black background, with colors
            determined by particle type.

    """
    if state.position.shape[-1] != 2:
        raise ValueError("Particle Life 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).
    positions = state.position  # Shape: (num_particles, 2)
    positions = positions.at[:, 1].set(1 - positions[:, 1])

    # Rasterize: nearest particle per pixel
    grid = pixel_grid(resolution)  # Shape: (resolution, resolution, 2)
    min_distance_sq, closest_particle_idx = nearest_point(grid, positions)

    # Get class of the closest particle for each pixel
    closest_class = state.class_id[
        closest_particle_idx
    ]  # Shape: (resolution, resolution)

    # Compute smooth mask based on distance to closest particle
    mask = soft_disk_mask(
        min_distance_sq, particle_radius
    )  # Shape: (resolution, resolution)

    # Generate colors for each class using HSV
    hues = jnp.linspace(0, 1, self.num_classes, endpoint=False)
    hsv = jnp.stack([hues, jnp.ones_like(hues), jnp.ones_like(hues)], axis=-1)
    colors = hsv_to_rgb(hsv)  # Shape: (num_classes, 3)

    # Assign colors based on closest particle's class
    particle_colors = colors[closest_class]  # Shape: (resolution, resolution, 3)

    # Create black background
    background = jnp.zeros(
        (resolution, resolution, 3)
    )  # Shape: (resolution, resolution, 3)

    # Blend particle colors with background using the mask
    rgb = (
        background * (1.0 - mask[..., None]) + particle_colors * mask[..., None]
    )  # 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