Abelian Sandpile ¶

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.sandpile import Sandpile

import jax import jax.numpy as jnp import mediapy from flax import nnx from cax.cs.sandpile import Sandpile

Configuration¶

In [4]:

seed = 0

num_steps = 3072
spatial_dims = (64, 64)
initial_chips = 4096

seed = 0 num_steps = 3072 spatial_dims = (64, 64) initial_chips = 4096

Instantiate system¶

In [5]:

cs = Sandpile()

cs = Sandpile()

Sample initial state¶

In [6]:

def sample_state():
    """Sample a state with all chips stacked on the center cell."""
    state = jnp.zeros((*spatial_dims, 1))
    mid_x, mid_y = spatial_dims[0] // 2, spatial_dims[1] // 2
    state = state.at[mid_x, mid_y, 0].set(initial_chips)
    return state

def sample_state(): """Sample a state with all chips stacked on the center cell.""" state = jnp.zeros((*spatial_dims, 1)) mid_x, mid_y = spatial_dims[0] // 2, spatial_dims[1] // 2 state = state.at[mid_x, mid_y, 0].set(initial_chips) return state

Run¶

In [7]:

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

# Every toppling sweep is computed; every fourth is shown
states = jnp.concatenate([state_init[None], states])[::4]
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))

# Every toppling sweep is computed; every fourth is shown states = jnp.concatenate([state_init[None], states])[::4] 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))

Self-Organized Criticality¶

Instead of starting with a large pile and letting it relax, we can drive the system by dropping one grain of sand at a random location. The system self-organizes to a critical state where avalanche sizes follow a power law.

The classic protocol separates timescales: a grain falls, and the lattice relaxes fully before the next one. We first drive the empty lattice to the critical state without recording, then drop one grain per relaxation window, so that each avalanche sweeps across a quiet lattice.

In [9]:

num_steps_driven = 4096
num_drops_recorded = 64
relax_steps = 32
spatial_dims_driven = (32, 32)

num_steps_driven = 4096 num_drops_recorded = 64 relax_steps = 32 spatial_dims_driven = (32, 32)

In [10]:

key = jax.random.key(seed)
key1, key2 = jax.random.split(key)

drop_x = jax.random.randint(key1, (num_steps_driven,), 0, spatial_dims_driven[0])
drop_y = jax.random.randint(key2, (num_steps_driven,), 0, spatial_dims_driven[1])

inputs = jnp.zeros((num_steps_driven, *spatial_dims_driven, 1))
inputs = inputs.at[jnp.arange(num_steps_driven), drop_x, drop_y, 0].set(1.0)

key = jax.random.key(seed) key1, key2 = jax.random.split(key) drop_x = jax.random.randint(key1, (num_steps_driven,), 0, spatial_dims_driven[0]) drop_y = jax.random.randint(key2, (num_steps_driven,), 0, spatial_dims_driven[1]) inputs = jnp.zeros((num_steps_driven, *spatial_dims_driven, 1)) inputs = inputs.at[jnp.arange(num_steps_driven), drop_x, drop_y, 0].set(1.0)

In [11]:

cs_driven = Sandpile(padding="ZERO")

# Drive the empty lattice to the critical state
state_init_driven = jnp.zeros((*spatial_dims_driven, 1))
state_critical = cs_driven(
    state_init_driven, input=inputs, num_steps=num_steps_driven, input_in_axis=0
)

# Drop one grain at a time, letting each avalanche play out before the next
key, subkey = jax.random.split(key)
drops = jax.random.randint(
    subkey, (num_drops_recorded, 2), 0, jnp.array(spatial_dims_driven)
)

state = state_critical
trajectory = []
for x, y in drops:
    state = state.at[x, y, 0].add(1.0)
    state, states = cs_driven(state, num_steps=relax_steps, return_states=True)
    trajectory.append(states)
states = jnp.concatenate(trajectory)

cs_driven = Sandpile(padding="ZERO") # Drive the empty lattice to the critical state state_init_driven = jnp.zeros((*spatial_dims_driven, 1)) state_critical = cs_driven( state_init_driven, input=inputs, num_steps=num_steps_driven, input_in_axis=0 ) # Drop one grain at a time, letting each avalanche play out before the next key, subkey = jax.random.split(key) drops = jax.random.randint( subkey, (num_drops_recorded, 2), 0, jnp.array(spatial_dims_driven) ) state = state_critical trajectory = [] for x, y in drops: state = state.at[x, y, 0].add(1.0) state, states = cs_driven(state, num_steps=relax_steps, return_states=True) trajectory.append(states) states = jnp.concatenate(trajectory)

In [12]:

states = jnp.concatenate([state_critical[None], states])
frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs_driven, states)

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

states = jnp.concatenate([state_critical[None], states]) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs_driven, states) mediapy.show_video(frames.repeat(8, axis=-3).repeat(8, axis=-2))