Particle Lenia ¶

Installation¶

You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:

In [1]:

%pip install -U "jax[cuda]"

%pip install -U "jax[cuda]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Then, install CAX from PyPi:

In [2]:

%pip install -U "cax[examples]"

%pip install -U "cax[examples]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Import¶

In [3]:

import jax
import jax.numpy as jnp
import mediapy
from flax import nnx

from cax.cs.particle_lenia import (
    ParticleLenia,
    ParticleLeniaGrowthParams,
    ParticleLeniaKernelParams,
    ParticleLeniaRuleParams,
    ParticleLeniaState,
    bell,
)
from cax.utils import render_states

import jax import jax.numpy as jnp import mediapy from flax import nnx from cax.cs.particle_lenia import ( ParticleLenia, ParticleLeniaGrowthParams, ParticleLeniaKernelParams, ParticleLeniaRuleParams, ParticleLeniaState, bell, ) from cax.utils import render_states

Configuration¶

In [4]:

seed = 0

num_steps = 8_192
num_spatial_dims = 2
num_particles = 200
T = 10

key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

seed = 0 num_steps = 8_192 num_spatial_dims = 2 num_particles = 200 T = 10 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Instantiate system¶

Rule parameters¶

In [5]:

mean = 4.0
std = 1.0


def compute_weight(mean, std, num_spatial_dims):
    """Compute weight for the kernel."""
    r = jnp.linspace(max(mean - 4 * std, 0.0), mean + 4 * std, 51)
    y = bell(r, mean, std) * r ** (num_spatial_dims - 1)
    s = jnp.trapezoid(y, r) * {2: 2, 3: 4}[num_spatial_dims] * jnp.pi
    return 1 / s


weight = compute_weight(mean, std, num_spatial_dims)

mean = 4.0 std = 1.0 def compute_weight(mean, std, num_spatial_dims): """Compute weight for the kernel.""" r = jnp.linspace(max(mean - 4 * std, 0.0), mean + 4 * std, 51) y = bell(r, mean, std) * r ** (num_spatial_dims - 1) s = jnp.trapezoid(y, r) * {2: 2, 3: 4}[num_spatial_dims] * jnp.pi return 1 / s weight = compute_weight(mean, std, num_spatial_dims)

In [6]:

kernel_params = ParticleLeniaKernelParams(
    weight=weight,
    mean=mean,
    std=std,
)

growth_params = ParticleLeniaGrowthParams(
    mean=0.6,
    std=0.15,
)

rule_params = ParticleLeniaRuleParams(
    c_rep=1.0,
    kernel_params=kernel_params,
    growth_params=growth_params,
)

kernel_params = ParticleLeniaKernelParams( weight=weight, mean=mean, std=std, ) growth_params = ParticleLeniaGrowthParams( mean=0.6, std=0.15, ) rule_params = ParticleLeniaRuleParams( c_rep=1.0, kernel_params=kernel_params, growth_params=growth_params, )

In [7]:

cs = ParticleLenia(
    num_spatial_dims=num_spatial_dims,
    T=T,
    rule_params=rule_params,
)

cs = ParticleLenia( num_spatial_dims=num_spatial_dims, T=T, rule_params=rule_params, )

Sample initial state¶

In [8]:

def sample_state(key):
    """Sample a state with random particle positions."""
    position = 12.0 * (jax.random.uniform(key, (num_particles, num_spatial_dims)) - 0.5)
    return ParticleLeniaState(position=position)

def sample_state(key): """Sample a state with random particle positions.""" position = 12.0 * (jax.random.uniform(key, (num_particles, num_spatial_dims)) - 0.5) return ParticleLeniaState(position=position)

Run¶

In [9]:

key, subkey = jax.random.split(key)
state_init = sample_state(subkey)
state_final, states = cs(state_init, num_steps=num_steps, return_states=True)

key, subkey = jax.random.split(key) state_init = sample_state(subkey) state_final, states = cs(state_init, num_steps=num_steps, return_states=True)

Visualize¶

In [10]:

states = jax.tree.map(
    lambda first, rest: jnp.concatenate([first[None], rest]), state_init, states
)
frames = render_states(cs, states, resolution=512, particle_radius=0.3)

mediapy.show_video(frames, width=256, height=256, fps=600)

states = jax.tree.map( lambda first, rest: jnp.concatenate([first[None], rest]), state_init, states ) frames = render_states(cs, states, resolution=512, particle_radius=0.3) mediapy.show_video(frames, width=256, height=256, fps=600)