Langton's Ant ¶

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.langton_ant import LangtonAnt, LangtonAntState

import jax import jax.numpy as jnp import mediapy from flax import nnx from cax.cs.langton_ant import LangtonAnt, LangtonAntState

Configuration¶

In [4]:

seed = 0

num_steps = 11_000
spatial_dims = (64, 64)
rule_string = "RL"  # Classic Langton's Ant

rngs = nnx.Rngs(seed)

seed = 0 num_steps = 11_000 spatial_dims = (64, 64) rule_string = "RL" # Classic Langton's Ant rngs = nnx.Rngs(seed)

Instantiate system¶

In [5]:

turns = LangtonAnt.turns_from_rule_string(rule_string)
turns

turns = LangtonAnt.turns_from_rule_string(rule_string) turns

Out[5]:

Array([1, 3], dtype=int32)

In [6]:

cs = LangtonAnt(turns=turns)

cs = LangtonAnt(turns=turns)

Sample initial state¶

In [7]:

def sample_state():
    """Sample a state with the ant at the center of an empty grid."""
    grid = jnp.zeros((*spatial_dims, 1))
    position = jnp.array([spatial_dims[0] // 2, spatial_dims[1] // 2], dtype=jnp.int32)
    direction = jnp.array(0, dtype=jnp.int32)  # North
    return LangtonAntState(grid=grid, position=position, direction=direction)

def sample_state(): """Sample a state with the ant at the center of an empty grid.""" grid = jnp.zeros((*spatial_dims, 1)) position = jnp.array([spatial_dims[0] // 2, spatial_dims[1] // 2], dtype=jnp.int32) direction = jnp.array(0, dtype=jnp.int32) # North return LangtonAntState(grid=grid, position=position, direction=direction)

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 = jax.tree.map(
    lambda state_init, states: jnp.concatenate([state_init[None], states]),
    state_init,
    states,
)
frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, states)

mediapy.show_video(frames.repeat(4, axis=-3).repeat(4, axis=-2))

states = jax.tree.map( lambda state_init, states: jnp.concatenate([state_init[None], states]), state_init, states, ) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_video(frames.repeat(4, axis=-3).repeat(4, axis=-2))

Langton's Ant Family¶

You can experiment with other generalized Langton's Ants by changing the rule string.

LLRR¶

In [10]:

turns = LangtonAnt.turns_from_rule_string("LLRR")
cs = LangtonAnt(turns=turns)

turns = LangtonAnt.turns_from_rule_string("LLRR") cs = LangtonAnt(turns=turns)

LRRRRRLLR¶

In [11]:

turns = LangtonAnt.turns_from_rule_string("LRRRRRLLR")
cs = LangtonAnt(turns=turns)

turns = LangtonAnt.turns_from_rule_string("LRRRRRLLR") cs = LangtonAnt(turns=turns)

Run¶

In [12]:

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 [13]:

states = jax.tree.map(
    lambda state_init, states: jnp.concatenate([state_init[None], states]),
    state_init,
    states,
)
frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, states)

mediapy.show_video(frames.repeat(4, axis=-3).repeat(4, axis=-2))

states = jax.tree.map( lambda state_init, states: jnp.concatenate([state_init[None], states]), state_init, states, ) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_video(frames.repeat(4, axis=-3).repeat(4, axis=-2))