Skip to content

Particle Lenia

cax.cs.particle_lenia.cs.ParticleLenia

Bases: ComplexSystem[ParticleLeniaState, Array]

Particle Lenia class.

Source code in src/cax/cs/particle_lenia/cs.py 
 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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
class ParticleLenia(ComplexSystem[ParticleLeniaState, Array]):
    """Particle Lenia class."""

    def __init__(
        self,
        *,
        num_spatial_dims: int,
        T: float,
        kernel_fn: Callable[[Array, ParticleLeniaKernelParams], Array] = peak_kernel_fn,
        growth_fn: Callable[[Array, ParticleLeniaGrowthParams], Array] = peak_growth_fn,
        rule_params: ParticleLeniaRuleParams,
    ):
        """Initialize Particle Lenia.

        Args:
            num_spatial_dims: Number of spatial dimensions (e.g., 2 for 2D, 3 for 3D).
                Determines the dimensionality of particle positions and field
                computations.
            T: Time resolution controlling the temporal discretization. Higher values
                produce smoother temporal dynamics with smaller update steps.
            kernel_fn: Callable that computes pairwise kernel weights between particles
                based on their distance. Takes rule parameters and returns kernel
                values.
            growth_fn: Callable that maps kernel field values to growth field values.
                Defines how particles respond to their local neighborhood density.
            rule_params: Instance of ParticleLeniaRuleParams containing kernel and
                growth parameters such as radii, peak positions, widths, and heights.

        """
        self.num_spatial_dims = num_spatial_dims
        self.perceive = ParticleLeniaPerceive(
            num_spatial_dims=num_spatial_dims,
            kernel_fn=kernel_fn,
            growth_fn=growth_fn,
            rule_params=rule_params,
        )
        self.update = ParticleLeniaUpdate(
            T=T,
        )

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

        return next_state

    @nnx.jit(static_argnames=("resolution", "extent", "particle_radius", "mode"))
    @override
    def render(
        self,
        state: ParticleLeniaState,
        *,
        resolution: int = 512,
        extent: float = 15.0,
        particle_radius: float = 0.3,
        mode: Literal["particles", "UG", "E"] = "UG",
    ) -> Array:
        """Render state to RGB image.

        Renders Particle Lenia state as particles optionally overlaid on field
        visualizations. Particles appear as blue circles. The background can show kernel
        field (U), growth field (G), or energy field (E) to visualize the underlying
        dynamics driving particle motion. Field visualizations use color mapping to
        represent field intensities across space.

        Args:
            state: ParticleLeniaState containing particle positions in continuous space.
                Currently only 2D visualization is supported.
            resolution: Size of the output image in pixels for both width and height.
                Higher values produce smoother field gradients but increase computation
                cost.
            extent: Half-width of the viewing area in coordinate space. The view spans
                from -extent to +extent in each dimension. Adjust to zoom in or out on
                the particle system.
            particle_radius: Radius of each particle in coordinate space. Particles are
                drawn as smooth circles with anti-aliased edges.
            mode: Visualization mode determining what fields to display:
                "particles": Only show particles on white background.
                "UG": Show particles overlaid on kernel (U) and growth (G) field
                    visualization (default).
                "E": Show particles overlaid on energy field visualization.

        Returns:
            RGB image with dtype uint8 and shape (resolution, resolution, 3), showing
            particles and optionally the underlying field structure that drives their
            motion.

        """
        if self.num_spatial_dims != 2:
            raise ValueError("Particle Lenia only supports 2D visualization.")

        # Create a grid of coordinates
        grid = pixel_grid(
            resolution, low=-extent, high=extent
        )  # (resolution, resolution, 2)

        # Reshape grid for computation
        flat_grid = grid.reshape(-1, 2)

        # Vectorize the field computation over all grid points. nnx transforms reject
        # bound methods, so the unbound method is vmapped with the module as first arg.
        flat_U, flat_G, flat_R = nnx.vmap(
            ParticleLeniaPerceive.compute_fields, in_axes=(None, None, 0)
        )(self.perceive, state, flat_grid)

        # Reshape back to grid; the energy field is repulsion minus growth
        U_field = flat_U.reshape(resolution, resolution)
        G_field = flat_G.reshape(resolution, resolution)
        R_field = flat_R.reshape(resolution, resolution)
        E_field = R_field - G_field

        # Helper functions for colormapping
        def lerp(x: Array, a: Array, b: Array) -> Array:
            return a * (1.0 - x) + b * x

        def cmap_e(e: Array) -> Array:
            stacked = jnp.stack([e, -e], -1).clip(0)
            colors = jnp.array([[0.3, 1.0, 1.0], [1.0, 0.3, 1.0]], dtype=jnp.float32)
            return 1.0 - jnp.matmul(stacked, colors)

        def cmap_ug(u: Array, g: Array) -> Array:
            vis = lerp(
                u[..., None], jnp.array([0.1, 0.1, 0.3]), jnp.array([0.2, 0.7, 1.0])
            )
            return lerp(g[..., None], vis, jnp.array([1.17, 0.91, 0.13]))

        # Calculate particle mask
        distance_sq_min, _ = nearest_point(grid, state.position)
        particle_mask = soft_disk_mask(distance_sq_min, particle_radius)

        # Normalize fields for visualization
        U_norm = (U_field - jnp.min(U_field)) / (
            jnp.max(U_field) - jnp.min(U_field) + 1e-8
        )
        G_norm = (G_field - jnp.min(G_field)) / (
            jnp.max(G_field) - jnp.min(G_field) + 1e-8
        )

        # Create visualizations
        vis_e = cmap_e(E_field)
        vis_ug = cmap_ug(U_norm, G_norm)

        # Apply particle mask
        particle_mask = particle_mask[:, :, None]

        # Create base particle visualization (blue particles on white background)
        vis_particle = jnp.ones((resolution, resolution, 3))
        vis_particle = (
            vis_particle * (1.0 - particle_mask)
            + jnp.array([0.0, 0.0, 1.0]) * particle_mask
        )

        # Choose visualization based on mode
        if mode == "UG":
            # Blend particles with UG field
            rgb = vis_ug * (1.0 - particle_mask * 0.7) + vis_particle * (
                particle_mask * 0.7
            )
        elif mode == "E":
            # Blend particles with E field
            rgb = vis_e * (1.0 - particle_mask * 0.7) + vis_particle * (
                particle_mask * 0.7
            )
        elif mode == "particles":
            rgb = vis_particle
        else:
            raise ValueError(
                f"mode must be one of 'particles', 'UG', 'E', got {mode!r}"
            )

        return clip_and_uint8(rgb)

__init__(*, num_spatial_dims, T, kernel_fn=peak_kernel_fn, growth_fn=peak_growth_fn, rule_params)

Initialize Particle Lenia.

Parameters:

Name Type Description Default
num_spatial_dims int

Number of spatial dimensions (e.g., 2 for 2D, 3 for 3D). Determines the dimensionality of particle positions and field computations.

 required
T float

Time resolution controlling the temporal discretization. Higher values produce smoother temporal dynamics with smaller update steps.

 required
kernel_fn Callable[[Array, ParticleLeniaKernelParams], Array]

Callable that computes pairwise kernel weights between particles based on their distance. Takes rule parameters and returns kernel values.

 peak_kernel_fn
growth_fn Callable[[Array, ParticleLeniaGrowthParams], Array]

Callable that maps kernel field values to growth field values. Defines how particles respond to their local neighborhood density.

 peak_growth_fn
rule_params ParticleLeniaRuleParams

Instance of ParticleLeniaRuleParams containing kernel and growth parameters such as radii, peak positions, widths, and heights.

 required
Source code in src/cax/cs/particle_lenia/cs.py 
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
def __init__(
    self,
    *,
    num_spatial_dims: int,
    T: float,
    kernel_fn: Callable[[Array, ParticleLeniaKernelParams], Array] = peak_kernel_fn,
    growth_fn: Callable[[Array, ParticleLeniaGrowthParams], Array] = peak_growth_fn,
    rule_params: ParticleLeniaRuleParams,
):
    """Initialize Particle Lenia.

    Args:
        num_spatial_dims: Number of spatial dimensions (e.g., 2 for 2D, 3 for 3D).
            Determines the dimensionality of particle positions and field
            computations.
        T: Time resolution controlling the temporal discretization. Higher values
            produce smoother temporal dynamics with smaller update steps.
        kernel_fn: Callable that computes pairwise kernel weights between particles
            based on their distance. Takes rule parameters and returns kernel
            values.
        growth_fn: Callable that maps kernel field values to growth field values.
            Defines how particles respond to their local neighborhood density.
        rule_params: Instance of ParticleLeniaRuleParams containing kernel and
            growth parameters such as radii, peak positions, widths, and heights.

    """
    self.num_spatial_dims = num_spatial_dims
    self.perceive = ParticleLeniaPerceive(
        num_spatial_dims=num_spatial_dims,
        kernel_fn=kernel_fn,
        growth_fn=growth_fn,
        rule_params=rule_params,
    )
    self.update = ParticleLeniaUpdate(
        T=T,
    )

render(state, *, resolution=512, extent=15.0, particle_radius=0.3, mode='UG')

Render state to RGB image.

Renders Particle Lenia state as particles optionally overlaid on field visualizations. Particles appear as blue circles. The background can show kernel field (U), growth field (G), or energy field (E) to visualize the underlying dynamics driving particle motion. Field visualizations use color mapping to represent field intensities across space.

Parameters:

Name Type Description Default
state ParticleLeniaState

ParticleLeniaState containing particle positions in continuous space. Currently only 2D visualization is supported.

 required
resolution int

Size of the output image in pixels for both width and height. Higher values produce smoother field gradients but increase computation cost.

 512
extent float

Half-width of the viewing area in coordinate space. The view spans from -extent to +extent in each dimension. Adjust to zoom in or out on the particle system.

 15.0
particle_radius float

Radius of each particle in coordinate space. Particles are drawn as smooth circles with anti-aliased edges.

 0.3
mode Literal['particles', 'UG', 'E']

Visualization mode determining what fields to display: "particles": Only show particles on white background. "UG": Show particles overlaid on kernel (U) and growth (G) field visualization (default). "E": Show particles overlaid on energy field visualization.

 'UG'

Returns:

Type Description
Array

RGB image with dtype uint8 and shape (resolution, resolution, 3), showing
Array

particles and optionally the underlying field structure that drives their
Array

motion.
Source code in src/cax/cs/particle_lenia/cs.py 
 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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
@nnx.jit(static_argnames=("resolution", "extent", "particle_radius", "mode"))
@override
def render(
    self,
    state: ParticleLeniaState,
    *,
    resolution: int = 512,
    extent: float = 15.0,
    particle_radius: float = 0.3,
    mode: Literal["particles", "UG", "E"] = "UG",
) -> Array:
    """Render state to RGB image.

    Renders Particle Lenia state as particles optionally overlaid on field
    visualizations. Particles appear as blue circles. The background can show kernel
    field (U), growth field (G), or energy field (E) to visualize the underlying
    dynamics driving particle motion. Field visualizations use color mapping to
    represent field intensities across space.

    Args:
        state: ParticleLeniaState containing particle positions in continuous space.
            Currently only 2D visualization is supported.
        resolution: Size of the output image in pixels for both width and height.
            Higher values produce smoother field gradients but increase computation
            cost.
        extent: Half-width of the viewing area in coordinate space. The view spans
            from -extent to +extent in each dimension. Adjust to zoom in or out on
            the particle system.
        particle_radius: Radius of each particle in coordinate space. Particles are
            drawn as smooth circles with anti-aliased edges.
        mode: Visualization mode determining what fields to display:
            "particles": Only show particles on white background.
            "UG": Show particles overlaid on kernel (U) and growth (G) field
                visualization (default).
            "E": Show particles overlaid on energy field visualization.

    Returns:
        RGB image with dtype uint8 and shape (resolution, resolution, 3), showing
        particles and optionally the underlying field structure that drives their
        motion.

    """
    if self.num_spatial_dims != 2:
        raise ValueError("Particle Lenia only supports 2D visualization.")

    # Create a grid of coordinates
    grid = pixel_grid(
        resolution, low=-extent, high=extent
    )  # (resolution, resolution, 2)

    # Reshape grid for computation
    flat_grid = grid.reshape(-1, 2)

    # Vectorize the field computation over all grid points. nnx transforms reject
    # bound methods, so the unbound method is vmapped with the module as first arg.
    flat_U, flat_G, flat_R = nnx.vmap(
        ParticleLeniaPerceive.compute_fields, in_axes=(None, None, 0)
    )(self.perceive, state, flat_grid)

    # Reshape back to grid; the energy field is repulsion minus growth
    U_field = flat_U.reshape(resolution, resolution)
    G_field = flat_G.reshape(resolution, resolution)
    R_field = flat_R.reshape(resolution, resolution)
    E_field = R_field - G_field

    # Helper functions for colormapping
    def lerp(x: Array, a: Array, b: Array) -> Array:
        return a * (1.0 - x) + b * x

    def cmap_e(e: Array) -> Array:
        stacked = jnp.stack([e, -e], -1).clip(0)
        colors = jnp.array([[0.3, 1.0, 1.0], [1.0, 0.3, 1.0]], dtype=jnp.float32)
        return 1.0 - jnp.matmul(stacked, colors)

    def cmap_ug(u: Array, g: Array) -> Array:
        vis = lerp(
            u[..., None], jnp.array([0.1, 0.1, 0.3]), jnp.array([0.2, 0.7, 1.0])
        )
        return lerp(g[..., None], vis, jnp.array([1.17, 0.91, 0.13]))

    # Calculate particle mask
    distance_sq_min, _ = nearest_point(grid, state.position)
    particle_mask = soft_disk_mask(distance_sq_min, particle_radius)

    # Normalize fields for visualization
    U_norm = (U_field - jnp.min(U_field)) / (
        jnp.max(U_field) - jnp.min(U_field) + 1e-8
    )
    G_norm = (G_field - jnp.min(G_field)) / (
        jnp.max(G_field) - jnp.min(G_field) + 1e-8
    )

    # Create visualizations
    vis_e = cmap_e(E_field)
    vis_ug = cmap_ug(U_norm, G_norm)

    # Apply particle mask
    particle_mask = particle_mask[:, :, None]

    # Create base particle visualization (blue particles on white background)
    vis_particle = jnp.ones((resolution, resolution, 3))
    vis_particle = (
        vis_particle * (1.0 - particle_mask)
        + jnp.array([0.0, 0.0, 1.0]) * particle_mask
    )

    # Choose visualization based on mode
    if mode == "UG":
        # Blend particles with UG field
        rgb = vis_ug * (1.0 - particle_mask * 0.7) + vis_particle * (
            particle_mask * 0.7
        )
    elif mode == "E":
        # Blend particles with E field
        rgb = vis_e * (1.0 - particle_mask * 0.7) + vis_particle * (
            particle_mask * 0.7
        )
    elif mode == "particles":
        rgb = vis_particle
    else:
        raise ValueError(
            f"mode must be one of 'particles', 'UG', 'E', got {mode!r}"
        )

    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