Lenia

cax.cs.lenia.cs.Lenia

Bases: ComplexSystem[Array, Array]

Lenia class.

Source code in src/cax/cs/lenia/cs.py

class Lenia(ComplexSystem[Array, Array]):
    """Lenia class."""

    def __init__(
        self,
        *,
        spatial_dims: tuple[int, ...],
        channel_size: int,
        R: int,
        T: float,
        state_scale: float = 1.0,
        kernel_fn: Callable[[Array, LeniaKernelParams], Array] = gaussian_kernel_fn,
        growth_fn: Callable[[Array, LeniaGrowthParams], Array] = exponential_growth_fn,
        rule_params: LeniaRuleParams,
    ):
        """Initialize Lenia.

        Args:
            spatial_dims: Spatial dimensions (e.g., (64, 64) for 2D or (32, 32, 32) for
                3D).
            channel_size: Number of channels.
            R: Space resolution defining the kernel radius. Larger values create wider
                neighborhoods and smoother patterns.
            T: Time resolution controlling the temporal discretization. Higher values
                produce smoother temporal dynamics with smaller update steps.
            state_scale: Scaling factor applied to state values.
            kernel_fn: Callable that generates convolution kernels. Takes rule
                parameters and returns kernel weights.
            growth_fn: Callable that maps neighborhood potential to growth values.
                Defines how cells respond to their local environment.
            rule_params: Instance of LeniaRuleParams containing kernel and growth
                parameters for each channel.

        """
        self.perceive = LeniaPerceive(
            spatial_dims=spatial_dims,
            channel_size=channel_size,
            R=R,
            state_scale=state_scale,
            kernel_fn=kernel_fn,
            rule_params=rule_params,
        )
        self.update = LeniaUpdate(
            channel_size=channel_size,
            T=T,
            growth_fn=growth_fn,
            rule_params=rule_params,
        )

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

        return next_state

    @nnx.jit
    @override
    def render(self, state: Array) -> Array:
        """Render state to RGB image.

        Converts the multi-channel Lenia state to an RGB visualization. Channels are
        mapped to color channels (Red, Green, Blue) for visualization. If there are
        more than 3 channels, the last three are displayed. If there are fewer than
        3 channels, the missing channels are filled with zeros.

        Args:
            state: Array with shape (*spatial_dims, channel_size) representing the Lenia
                state, where each cell contains continuous values typically in [0, 1].

        Returns:
            RGB image with dtype uint8 and shape (*spatial_dims, 3), where state
                values are mapped to colors in the range [0, 255].

        """
        rgb = render_array_with_channels_to_rgb(state)

        return clip_and_uint8(rgb)

__init__(*, spatial_dims, channel_size, R, T, state_scale=1.0, kernel_fn=gaussian_kernel_fn, growth_fn=exponential_growth_fn, rule_params)

Initialize Lenia.

Parameters:

Name	Type	Description	Default
spatial_dims	tuple[int, ...]	Spatial dimensions (e.g., (64, 64) for 2D or (32, 32, 32) for 3D).	required
channel_size	int	Number of channels.	required
R	int	Space resolution defining the kernel radius. Larger values create wider neighborhoods and smoother patterns.	required
T	float	Time resolution controlling the temporal discretization. Higher values produce smoother temporal dynamics with smaller update steps.	required
state_scale	float	Scaling factor applied to state values.	1.0
kernel_fn	Callable[[Array, LeniaKernelParams], Array]	Callable that generates convolution kernels. Takes rule parameters and returns kernel weights.	gaussian_kernel_fn
growth_fn	Callable[[Array, LeniaGrowthParams], Array]	Callable that maps neighborhood potential to growth values. Defines how cells respond to their local environment.	exponential_growth_fn
rule_params	LeniaRuleParams	Instance of LeniaRuleParams containing kernel and growth parameters for each channel.	required

Source code in src/cax/cs/lenia/cs.py

def __init__(
    self,
    *,
    spatial_dims: tuple[int, ...],
    channel_size: int,
    R: int,
    T: float,
    state_scale: float = 1.0,
    kernel_fn: Callable[[Array, LeniaKernelParams], Array] = gaussian_kernel_fn,
    growth_fn: Callable[[Array, LeniaGrowthParams], Array] = exponential_growth_fn,
    rule_params: LeniaRuleParams,
):
    """Initialize Lenia.

    Args:
        spatial_dims: Spatial dimensions (e.g., (64, 64) for 2D or (32, 32, 32) for
            3D).
        channel_size: Number of channels.
        R: Space resolution defining the kernel radius. Larger values create wider
            neighborhoods and smoother patterns.
        T: Time resolution controlling the temporal discretization. Higher values
            produce smoother temporal dynamics with smaller update steps.
        state_scale: Scaling factor applied to state values.
        kernel_fn: Callable that generates convolution kernels. Takes rule
            parameters and returns kernel weights.
        growth_fn: Callable that maps neighborhood potential to growth values.
            Defines how cells respond to their local environment.
        rule_params: Instance of LeniaRuleParams containing kernel and growth
            parameters for each channel.

    """
    self.perceive = LeniaPerceive(
        spatial_dims=spatial_dims,
        channel_size=channel_size,
        R=R,
        state_scale=state_scale,
        kernel_fn=kernel_fn,
        rule_params=rule_params,
    )
    self.update = LeniaUpdate(
        channel_size=channel_size,
        T=T,
        growth_fn=growth_fn,
        rule_params=rule_params,
    )

render(state)

Render state to RGB image.

Converts the multi-channel Lenia state to an RGB visualization. Channels are mapped to color channels (Red, Green, Blue) for visualization. If there are more than 3 channels, the last three are displayed. If there are fewer than 3 channels, the missing channels are filled with zeros.

Parameters:

Name	Type	Description	Default
state	Array	Array with shape (*spatial_dims, channel_size) representing the Lenia state, where each cell contains continuous values typically in [0, 1].	required

Returns:

Type	Description
Array	RGB image with dtype uint8 and shape (*spatial_dims, 3), where state values are mapped to colors in the range [0, 255].

Source code in src/cax/cs/lenia/cs.py

@nnx.jit
@override
def render(self, state: Array) -> Array:
    """Render state to RGB image.

    Converts the multi-channel Lenia state to an RGB visualization. Channels are
    mapped to color channels (Red, Green, Blue) for visualization. If there are
    more than 3 channels, the last three are displayed. If there are fewer than
    3 channels, the missing channels are filled with zeros.

    Args:
        state: Array with shape (*spatial_dims, channel_size) representing the Lenia
            state, where each cell contains continuous values typically in [0, 1].

    Returns:
        RGB image with dtype uint8 and shape (*spatial_dims, 3), where state
            values are mapped to colors in the range [0, 255].

    """
    rgb = render_array_with_channels_to_rgb(state)

    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

@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