CAX: 使用 JAX 加速的细胞自动机库-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00131/article/details/147406108

CAX: 使用 JAX 加速的细胞自动机库

cax Cellular Automata Accelerated in JAX (Oral at ICLR 2025) 项目地址: https://gitcode.com/gh_mirrors/cax/cax

1. 项目介绍

CAX（Cellular Automata Accelerated in JAX）是一个高性能、灵活的开源库，旨在加速人工生命研究。它支持离散和连续系统，包括神经细胞自动机，适用于任意维度。CAX不仅处理传统的细胞自动机，还统一在直观的API下处理粒子系统等。它提供了15+个 ready-to-use 系统，从模拟一维的基本细胞自动机到训练三维的自编码神经细胞自动机，或创建美丽的Lenia模拟，CAX 为探索自组织系统的丰富世界提供了一个多功能的平台。

2. 项目快速启动

在开始之前，请确保您的系统中已安装 Python 3.10 或更高版本，并且已经配置好了 JAX 环境。

安装 CAX 的最简单方式是使用 pip：

pip install cax

或者，如果您在虚拟环境中使用 uv：

uv pip install cax

以下是一个基本的 CAX API 使用示例：

import jax
from flax import nnx
from cax.core.ca import CA
from cax.core.perceive import ConvPerceive
from cax.core.update import NCAUpdate

seed = 0
channel_size = 16
num_kernels = 3
hidden_layer_sizes = (128,)
cell_dropout_rate = 0.5
key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

perceive = ConvPerceive(
    channel_size=channel_size,
    perception_size=num_kernels * channel_size,
    rngs=rngs,
    feature_group_count=channel_size,
)

update = NCAUpdate(
    channel_size=channel_size,
    perception_size=num_kernels * channel_size,
    hidden_layer_sizes=hidden_layer_sizes,
    rngs=rngs,
    cell_dropout_rate=cell_dropout_rate,
    zeros_init=True,
)

ca = CA(perceive, update)

state = jax.random.normal(key, (64, 64, channel_size))
state, metrics = ca(state, num_steps=128)

3. 应用案例和最佳实践

CAX 提供了多种细胞自动机的实现，以下是一些应用案例：

Elementary Cellular Automata：基础的一维细胞自动机，可以用来探索基本的自组织行为。
Conway's Game of Life：康威生命游戏，一个经典的二维细胞自动机。
Lenia：一个美丽的细胞自动机模拟，可以用来生成动态的图案。

以下是如何创建一个简单的一维细胞自动机的代码示例：

from cax.core.ca import CA1D
from cax.core.rules import ElementaryCA

# 创建一个基本的一维细胞自动机
ca1d = CA1D(ElementaryCA(rule=30), size=100)

# 初始化状态
state = jax.random.uniform(jax.random.PRNGKey(0), (ca1d.size,))

# 运行100步
for _ in range(100):
    state = ca1d.update(state)

# 输出最终状态
print(state)