Skip to content

Release v0.2.0 #15

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
maxencefaldor merged 2 commits into from
Apr 21, 2025
Merged

Release v0.2.0 #15

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
276b309
Generalize Elementary CA and Life
maxencefaldor Apr 21, 2025
14847c0
Format and lint code
maxencefaldor Apr 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,15 +22,15 @@ Designed for speed and flexibility, CAX allows you to easily experiment with sel

## Why CAX? 💡

CAX supports discrete and continuous models, including neural cellular automata, across any number of dimensions. Beyond traditional cellular automata, it also handles particle systems and more, all unified under a single, intuitive API.
CAX supports discrete and continuous systems, including neural cellular automata, across any number of dimensions. Beyond traditional cellular automata, it also handles particle systems and more, all unified under a single, intuitive API.

### Rich 🎨

CAX provides a comprehensive collection of 15+ ready-to-use systems. From simulating one-dimensional [elementary cellular automata](examples/10_elementary_ca.ipynb) to training three-dimensional [self-autoencoding neural cellular automata](examples/45_self_autoencoding_mnist.ipynb), or even creating beautiful [Lenia](examples/20_lenia.ipynb) simulations, CAX provides a versatile platform for exploring the rich world of self-organizing systems.

### Flexible 🧩

CAX makes it easy to extend existing models or build custom ones from scratch for endless experimentation and discovery. Design your own experiments to probe the boundaries of artificial open-ended evolution and emergent complexity.
CAX makes it easy to extend existing systems or build custom ones from scratch for endless experimentation and discovery. Design your own experiments to probe the boundaries of artificial open-ended evolution and emergent complexity.

### Fast 🚀

Expand All @@ -40,7 +40,7 @@ CAX is built on top of the JAX/Flax ecosystem for speed and scalability. The lib

The library is thoroughly tested and documented with numerous examples to get you started! Our comprehensive guides walk you through everything from basic cellular automata to advanced neural implementations.

## Implemented Models 🦎
## Implemented Systems 🦎

| Cellular Automata | Reference | Example |
| --- | --- | --- |
Expand Down
50 changes: 34 additions & 16 deletions examples/00_getting_started.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -178,7 +178,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"In this section, we'll demonstrate the basic usage of CAX with pre-implemented cellular automata. We'll instantiate Conway's Game of Life and visualize a glider pattern, showing how easily you can get started with existing models in the library."
"In this section, we'll demonstrate the basic usage of CAX with pre-implemented cellular automata. We'll instantiate Conway's Game of Life and visualize a glider pattern, showing how easily you can get started with existing systems in the library."
]
},
{
Expand Down Expand Up @@ -215,7 +215,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"### Model"
"### System"
]
},
{
Expand All @@ -231,11 +231,20 @@
"metadata": {},
"outputs": [],
"source": [
"from cax.models.life import Life\n",
"from cax.systems.life import Life\n",
"\n",
"ca = Life(rngs=rngs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ca.update.update_birth_survival_from_string(\"B3/S23\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand All @@ -247,7 +256,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Then, we define a function to sample an initial state, which is essential for running a model."
"Then, we define a function to sample an initial state, which is essential for running a system."
]
},
{
Expand Down Expand Up @@ -282,7 +291,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Given an initial state and the model, we can simulate for `num_steps`."
"Given an initial state and the system, we can simulate for `num_steps`."
]
},
{
Expand All @@ -308,7 +317,7 @@
"source": [
"Finally, we can visualize the trajectory of states.\n",
"\n",
"All models should include a `render` method to convert a state into an RGB frame. For the Game of Life, a ready-to-use `render` method is provided, allowing you to easily generate a frame with a simple call: `frame = ca.render(state)`.\n",
"All systems should include a `render` method to convert a state into an RGB frame. For the Game of Life, a ready-to-use `render` method is provided, allowing you to easily generate a frame with a simple call: `frame = ca.render(state)`.\n",
"\n",
"Enjoy! 👾"
]
Expand Down Expand Up @@ -349,9 +358,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"While visualizing the trajectory of states can produce captivating simulations, we often also want to track additional metrics to better understand how the model evolves over time.\n",
"While visualizing the trajectory of states can produce captivating simulations, we often also want to track additional metrics to better understand how the system evolves over time.\n",
"\n",
"By default, the metrics function in CAX returns the states encountered during the model’s rollout. However, this behavior is highly customizable. In this section, we’ll explore how to create a tailored metrics function to log custom metrics suited to your needs.\n",
"By default, the metrics function in CAX returns the states encountered during the system’s rollout. However, this behavior is highly customizable. In this section, we’ll explore how to create a tailored metrics function to log custom metrics suited to your needs.\n",
"\n",
"A metrics function must accept the next state, the current state, the perception, and the input as its parameters. For the Game of Life, however, the input parameter is not utilized, and can safely be ignored."
]
Expand Down Expand Up @@ -388,6 +397,15 @@
"ca = Life(rngs=rngs, metrics_fn=custom_metrics_fn)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ca.update.update_birth_survival_from_string(\"B3/S23\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand All @@ -410,7 +428,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's run the model:"
"Let's run the system:"
]
},
{
Expand Down Expand Up @@ -510,7 +528,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"When using `nnx.Module`s with rngs in the state, we need to use `nnx.split_rngs` to properly vectorize over rngs state across parallel operations. For Conway's Game of Life specifically, the model doesn't use randomness during execution, so we could use `nnx.vmap` directly. However, we'll demonstrate the more general approach with `nnx.split_rngs` below, which works for any model that maintains rngs state."
"When using `nnx.Module`s with rngs in the state, we need to use `nnx.split_rngs` to properly vectorize over rngs state across parallel operations. For Conway's Game of Life specifically, the system doesn't use randomness during execution, so we could use `nnx.vmap` directly. However, we'll demonstrate the more general approach with `nnx.split_rngs` below, which works for any systems that maintain rngs state."
]
},
{
Expand Down Expand Up @@ -614,7 +632,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"CAX provides a set of perceive models. In this notebook, we will use a simple convolution perceive module."
"CAX provides a set of perceive modules. In this notebook, we will use a simple convolution perceive module."
]
},
{
Expand Down Expand Up @@ -643,7 +661,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"CAX provides a set of update models. In this notebook, we will use a residual MLP update module."
"CAX provides a set of update modules. In this notebook, we will use a residual MLP update module."
]
},
{
Expand Down Expand Up @@ -685,8 +703,8 @@
"metadata": {},
"outputs": [],
"source": [
"from cax.core.ca import CA\n",
"from cax.utils.render import clip_and_uint8, rgba_to_rgb\n",
"from cax.core import CA\n",
"from cax.utils import clip_and_uint8, rgba_to_rgb\n",
"\n",
"\n",
"class MyCustomCA(CA):\n",
Expand Down Expand Up @@ -856,7 +874,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.8"
"version": "3.13.3"
}
},
"nbformat": 4,
Expand Down
21 changes: 15 additions & 6 deletions examples/10_elementary_ca.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,15 +55,15 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import jax.numpy as jnp\n",
"import mediapy\n",
"from flax import nnx\n",
"\n",
"from cax.models.elementary import ElementaryCA"
"from cax.systems.elementary import ElementaryCA"
]
},
{
Expand Down Expand Up @@ -102,7 +102,16 @@
"metadata": {},
"outputs": [],
"source": [
"ca = ElementaryCA(rngs=rngs, wolfram_code=wolfram_code)"
"ca = ElementaryCA(rngs=rngs)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"ca.update.update_wolfram_code_from_string(wolfram_code)"
]
},
{
Expand All @@ -114,7 +123,7 @@
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -150,7 +159,7 @@
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": 9,
"metadata": {},
"outputs": [
{
Expand Down Expand Up @@ -190,7 +199,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.8"
"version": "3.13.3"
}
},
"nbformat": 4,
Expand Down
Loading