Generalize Elementary CA and Life

Author: maxencefaldor

README.md

@@ -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 🚀
 
@@ -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 |
 | --- | --- | --- |

examples/00_getting_started.ipynb

@@ -157,7 +157,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -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."
    ]
   },
   {
@@ -215,7 +215,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Model"
+    "### System"
    ]
   },
   {
@@ -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": {},
@@ -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."
    ]
   },
   {
@@ -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`."
    ]
   },
   {
@@ -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! 👾"
    ]
@@ -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."
    ]
@@ -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": {},
@@ -410,7 +428,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Let's run the model:"
+    "Let's run the system:"
    ]
   },
   {
@@ -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."
    ]
   },
   {
@@ -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."
    ]
   },
   {
@@ -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."
    ]
   },
   {
@@ -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",
@@ -856,7 +874,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.13.3"
   }
  },
  "nbformat": 4,

examples/10_elementary_ca.ipynb

@@ -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"
    ]
   },
   {
@@ -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)"
    ]
   },
   {
@@ -114,7 +123,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -150,7 +159,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
@@ -190,7 +199,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.13.3"
   }
  },
  "nbformat": 4,