Reaction-Diffusion (Gray-Scott) ¶

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.reaction_diffusion import ReactionDiffusion

import jax import jax.numpy as jnp import mediapy from flax import nnx from cax.cs.reaction_diffusion import ReactionDiffusion

Configuration¶

In [4]:

seed = 0

num_steps = 10000
spatial_dims = (128, 128)
diffusion_rate_u = 0.16
diffusion_rate_v = 0.08
feed_rate = 0.06
kill_rate = 0.062
dt = 1.0

rngs = nnx.Rngs(seed)

seed = 0 num_steps = 10000 spatial_dims = (128, 128) diffusion_rate_u = 0.16 diffusion_rate_v = 0.08 feed_rate = 0.06 kill_rate = 0.062 dt = 1.0 rngs = nnx.Rngs(seed)

Instantiate system¶

In [5]:

cs = ReactionDiffusion(
    diffusion_rate_u=diffusion_rate_u,
    diffusion_rate_v=diffusion_rate_v,
    feed_rate=feed_rate,
    kill_rate=kill_rate,
    dt=dt,
)

cs = ReactionDiffusion( diffusion_rate_u=diffusion_rate_u, diffusion_rate_v=diffusion_rate_v, feed_rate=feed_rate, kill_rate=kill_rate, dt=dt, )

Sample initial state¶

In [6]:

def sample_state():
    """Sample a state with U≈1 everywhere and a small square seed of V in the center."""
    key = jax.random.key(seed)
    key1, key2 = jax.random.split(key)

    random_influence = 0.2
    u = (1.0 - random_influence) * jnp.ones(
        spatial_dims
    ) + random_influence * jax.random.uniform(key1, spatial_dims)
    v = random_influence * jax.random.uniform(key2, spatial_dims)

    mid_x, mid_y = spatial_dims[0] // 2, spatial_dims[1] // 2
    radius = spatial_dims[0] // 10
    u = u.at[mid_x - radius : mid_x + radius, mid_y - radius : mid_y + radius].set(0.5)
    v = v.at[mid_x - radius : mid_x + radius, mid_y - radius : mid_y + radius].set(0.25)

    state = jnp.stack([u, v], axis=-1)
    return state

def sample_state(): """Sample a state with U≈1 everywhere and a small square seed of V in the center.""" key = jax.random.key(seed) key1, key2 = jax.random.split(key) random_influence = 0.2 u = (1.0 - random_influence) * jnp.ones( spatial_dims ) + random_influence * jax.random.uniform(key1, spatial_dims) v = random_influence * jax.random.uniform(key2, spatial_dims) mid_x, mid_y = spatial_dims[0] // 2, spatial_dims[1] // 2 radius = spatial_dims[0] // 10 u = u.at[mid_x - radius : mid_x + radius, mid_y - radius : mid_y + radius].set(0.5) v = v.at[mid_x - radius : mid_x + radius, mid_y - radius : mid_y + radius].set(0.25) state = jnp.stack([u, v], axis=-1) 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]:

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

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

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

Pattern Gallery¶

You can experiment with other Gray-Scott patterns by changing the feed rate and kill rate. Different parameter regimes produce diverse Turing patterns.

Bacteria (f=0.035, k=0.065)¶

In [9]:

cs = ReactionDiffusion(
    diffusion_rate_u=0.14,
    diffusion_rate_v=0.06,
    feed_rate=0.035,
    kill_rate=0.065,
)

cs = ReactionDiffusion( diffusion_rate_u=0.14, diffusion_rate_v=0.06, feed_rate=0.035, kill_rate=0.065, )

Coral (f=0.062, k=0.063)¶

In [10]:

cs = ReactionDiffusion(
    diffusion_rate_u=0.16,
    diffusion_rate_v=0.08,
    feed_rate=0.062,
    kill_rate=0.063,
)

cs = ReactionDiffusion( diffusion_rate_u=0.16, diffusion_rate_v=0.08, feed_rate=0.062, kill_rate=0.063, )

Run¶

In [11]:

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

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

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

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