Flow Lenia

cax.cs.flow_lenia.cs.FlowLenia

Bases: Lenia

Flow Lenia class.

Subclasses Lenia: perception and rendering are Lenia's, while the update adds flow-based advection with mass conservation. Like Lenia, Flow Lenia works in any number of spatial dimensions.

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

class FlowLenia(Lenia):
    """Flow Lenia class.

    Subclasses Lenia: perception and rendering are Lenia's, while the update adds
    flow-based advection with mass conservation. Like Lenia, Flow Lenia works in any
    number of spatial dimensions.
    """

    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,
        # Flow Lenia parameters
        theta_A: float | None = None,
        n: int = 2,
        dd: int = 5,
        sigma: float = 0.65,
    ):
        """Initialize Flow Lenia.

        Args:
            spatial_dims: Spatial dimensions (e.g., (64, 64) for 2D or (32, 32, 32) for
                3D). Note that reintegration tracking considers
                (2 * dd + 1) ** num_spatial_dims displacements per cell, which sets the
                memory and compute cost.
            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.
            theta_A: Threshold value for computing the flow activation alpha. Higher
                values make flow less sensitive to local density. Defaults to
                ``channel_size``, matching the official implementation.
            n: Exponent controlling the nonlinearity of flow activation. Higher values
                create sharper transitions between flow and no-flow regions.
            dd: Maximum displacement distance in pixels that flow can induce per time
                step. Controls the strength of advective transport.
            sigma: Spread parameter for the displacement kernel. Smaller values create
                more localized flow, larger values produce smoother displacement fields.

        """
        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 = FlowLeniaUpdate(
            channel_size=channel_size,
            T=T,
            growth_fn=growth_fn,
            rule_params=rule_params,
            theta_A=theta_A,
            n=n,
            dd=dd,
            sigma=sigma,
        )

__init__(*, spatial_dims, channel_size, R, T, state_scale=1.0, kernel_fn=gaussian_kernel_fn, growth_fn=exponential_growth_fn, rule_params, theta_A=None, n=2, dd=5, sigma=0.65)

Initialize Flow Lenia.

Parameters:

Name	Type	Description	Default
spatial_dims	tuple[int, ...]	Spatial dimensions (e.g., (64, 64) for 2D or (32, 32, 32) for 3D). Note that reintegration tracking considers (2 * dd + 1) ** num_spatial_dims displacements per cell, which sets the memory and compute cost.	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
theta_A	float | None	Threshold value for computing the flow activation alpha. Higher values make flow less sensitive to local density. Defaults to channel_size, matching the official implementation.	None
n	int	Exponent controlling the nonlinearity of flow activation. Higher values create sharper transitions between flow and no-flow regions.	2
dd	int	Maximum displacement distance in pixels that flow can induce per time step. Controls the strength of advective transport.	5
sigma	float	Spread parameter for the displacement kernel. Smaller values create more localized flow, larger values produce smoother displacement fields.	0.65

Source code in src/cax/cs/flow_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,
    # Flow Lenia parameters
    theta_A: float | None = None,
    n: int = 2,
    dd: int = 5,
    sigma: float = 0.65,
):
    """Initialize Flow Lenia.

    Args:
        spatial_dims: Spatial dimensions (e.g., (64, 64) for 2D or (32, 32, 32) for
            3D). Note that reintegration tracking considers
            (2 * dd + 1) ** num_spatial_dims displacements per cell, which sets the
            memory and compute cost.
        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.
        theta_A: Threshold value for computing the flow activation alpha. Higher
            values make flow less sensitive to local density. Defaults to
            ``channel_size``, matching the official implementation.
        n: Exponent controlling the nonlinearity of flow activation. Higher values
            create sharper transitions between flow and no-flow regions.
        dd: Maximum displacement distance in pixels that flow can induce per time
            step. Controls the strength of advective transport.
        sigma: Spread parameter for the displacement kernel. Smaller values create
            more localized flow, larger values produce smoother displacement fields.

    """
    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 = FlowLeniaUpdate(
        channel_size=channel_size,
        T=T,
        growth_fn=growth_fn,
        rule_params=rule_params,
        theta_A=theta_A,
        n=n,
        dd=dd,
        sigma=sigma,
    )

__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

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)