Elementary Cellular Automata ¶

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.numpy as jnp
import mediapy
from flax import nnx

from cax.cs.elementary import Elementary

import jax.numpy as jnp import mediapy from flax import nnx from cax.cs.elementary import Elementary

Configuration¶

In [4]:

seed = 0

num_steps = 512
spatial_dims = (1_024,)
wolfram_code_int = 110  # Rule 110

rngs = nnx.Rngs(seed)

seed = 0 num_steps = 512 spatial_dims = (1_024,) wolfram_code_int = 110 # Rule 110 rngs = nnx.Rngs(seed)

Instantiate system¶

In [5]:

wolfram_code = Elementary.wolfram_code_from_rule_number(wolfram_code_int)
wolfram_code

wolfram_code = Elementary.wolfram_code_from_rule_number(wolfram_code_int) wolfram_code

Out[5]:

Array([0., 1., 1., 0., 1., 1., 1., 0.], dtype=float32)

In [6]:

cs = Elementary(wolfram_code=wolfram_code)

cs = Elementary(wolfram_code=wolfram_code)

Sample initial state¶

In [7]:

def sample_state():
    """Sample a state with a single active cell."""
    state = jnp.zeros((*spatial_dims, 1))
    return state.at[spatial_dims[0] // 2].set(1.0)

def sample_state(): """Sample a state with a single active cell.""" state = jnp.zeros((*spatial_dims, 1)) return state.at[spatial_dims[0] // 2].set(1.0)

Run¶

In [8]:

state_init = sample_state()
state_final, states = cs(state_init, num_steps=num_steps, return_states=True)

state_init = sample_state() state_final, states = cs(state_init, num_steps=num_steps, return_states=True)

Visualize¶

In [9]:

states = jnp.concatenate([state_init[None], states])
frame = cs.render(states)

mediapy.show_image(frame)

states = jnp.concatenate([state_init[None], states]) frame = cs.render(states) mediapy.show_image(frame)