{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Lenia [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/maxencefaldor/cax/blob/main/examples/20_lenia.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install -U \"jax[cuda]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then, install CAX from PyPi:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install -U \"cax[examples]\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import pickle\n",
        "from importlib import resources\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import mediapy\n",
        "from flax import nnx\n",
        "\n",
        "from cax.cs.lenia import (\n",
        "\tFreeKernelParams,\n",
        "\tGrowthParams,\n",
        "\tKernelParams,\n",
        "\tLenia,\n",
        "\tLeniaRuleParams,\n",
        "\texponential_growth_fn,\n",
        "\tfree_kernel_fn,\n",
        "\tgaussian_kernel_fn,\n",
        "\tpolynomial_kernel_fn,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "seed = 0\n",
        "\n",
        "num_steps = 256\n",
        "spatial_dims = (128, 128)\n",
        "channel_size = 3\n",
        "R = 12\n",
        "T = 1 / 3\n",
        "state_scale = 2\n",
        "\n",
        "key = jax.random.key(seed)\n",
        "rngs = nnx.Rngs(seed)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Instantiate system"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This section demonstrates how to visualize well-known Lenia creatures by loading rule parameters and patterns from `cax/systems/lenia/patterns`.\n",
        "\n",
        "You can also experiment with combining rule parameters from one soliton with the pattern of another to observe novel emergent behaviors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "patterns = [\"5N7KKM\", \"5N7KKM_gyrating\", \"5N7KKM_twin\", \"VT049W\"]\n",
        "pattern = patterns[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Rule Params\n",
        "with (\n",
        "\tresources.files(f\"cax.cs.lenia.patterns.{pattern}\")\n",
        "\t.joinpath(\"rule_params.pickle\")\n",
        "\t.open(\"rb\") as f\n",
        "):\n",
        "\trule_params = pickle.load(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Pattern\n",
        "with resources.files(f\"cax.cs.lenia.patterns.{pattern}\").joinpath(\"pattern.pickle\").open(\"rb\") as f:\n",
        "\tpattern = pickle.load(f)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cs = Lenia(\n",
        "\tspatial_dims=spatial_dims,\n",
        "\tchannel_size=channel_size,\n",
        "\tR=R,\n",
        "\tT=T,\n",
        "\tstate_scale=state_scale,\n",
        "\trule_params=rule_params,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sample initial state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def sample_state(pattern):\n",
        "\t\"\"\"Sample a state with a pattern at the center.\"\"\"\n",
        "\t# Calculate the center of the state for each dimension\n",
        "\tmid = tuple(dim // 2 for dim in spatial_dims)\n",
        "\n",
        "\t# Scale pattern\n",
        "\tpattern_scaled = pattern\n",
        "\tfor axis in range(len(spatial_dims)):\n",
        "\t\tpattern_scaled = pattern_scaled.repeat(state_scale, axis=axis)\n",
        "\n",
        "\t# Get the shape of the scaled cells\n",
        "\tpattern_spatial_dims = pattern_scaled.shape[:-1]  # Exclude the channel dimension\n",
        "\n",
        "\t# Create empty state with the shape defined by spatial_dims and channel_size\n",
        "\tstate = jnp.zeros((*spatial_dims, channel_size))\n",
        "\n",
        "\t# Calculate the slice indices for each dimension\n",
        "\tslices = tuple(slice(m - c // 2, m + c - c // 2) for m, c in zip(mid, pattern_spatial_dims))\n",
        "\n",
        "\t# Place the scaled cells in the center of the state\n",
        "\tstate = state.at[slices].set(pattern_scaled)\n",
        "\n",
        "\treturn state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "state_init = sample_state(pattern)\n",
        "state_final = cs(state_init, num_steps=num_steps, sow=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualize"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "intermediates = nnx.pop(cs, nnx.Intermediate)\n",
        "states = intermediates.state[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "states = jnp.concatenate([state_init[None], states])\n",
        "frames = nnx.vmap(\n",
        "\tlambda cs, state: cs.render(state),\n",
        "\tin_axes=(None, 0),\n",
        ")(cs, states)\n",
        "\n",
        "mediapy.show_video(frames, width=256, height=256)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Rule Parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In this section, we will deep dive into what are the rule parameters in Lenia and explore the `RuleParams` object provided by CAX, how it works, and how you can define your own rule parameters to experiment with new Lenia worlds.\n",
        "\n",
        "The Lenia implementation in CAX rigorously follows:\n",
        "- [1] [Lenia - Biology of Artificial Life, Chan (2018)](https://arxiv.org/abs/1812.05433)\n",
        "- [2] [Lenia and Expanded Universe, Chan (2020)](https://arxiv.org/abs/2005.03742)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In Lenia, the state is a $d$-dimensional lattice $A^{t}$ with $c$ channels that evolves from $A^t$ to $A^{t+dt}$ based on a set of rules. Each rule consists of a source channel, a target channel, a weight, a kernel and a growth mapping function, see [1, 2] for more details.\n",
        "\n",
        "In CAX, the parameters of a rule are defined through the `RuleParams` object. For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "rule_params = LeniaRuleParams(\n",
        "\tchannel_source=0,\n",
        "\tchannel_target=1,\n",
        "\tweight=1.0,\n",
        "\tkernel_params=KernelParams(\n",
        "\t\tr=0.82,\n",
        "\t\tbeta=jnp.array([1 / 6, 1.0, 0.2]),\n",
        "\t),\n",
        "\tgrowth_params=GrowthParams(\n",
        "\t\tmean=0.0,\n",
        "\t\tstd=0.05,\n",
        "\t),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will now explore how to define custom kernels and growth mapping functions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Kernels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This section explores the kernel functions available in Lenia through CAX. We\"ll examine the built-in kernel functions and demonstrate how to create custom kernels for your own experiments.\n",
        "\n",
        "The original kernel functions are defined in [1, 2] and the free kernel function is defined in Flow-Lenia."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Original kernel function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "kernel_params = KernelParams(\n",
        "\tr=0.82,\n",
        "\tbeta=jnp.array([0.2, 1.0, 0.4]),\n",
        ")\n",
        "\n",
        "# Plot the original kernel function\n",
        "radius = jnp.linspace(0.0, 1.0, 1000)\n",
        "kernel_values = gaussian_kernel_fn(radius, kernel_params)\n",
        "\n",
        "plt.figure(figsize=(5, 3))\n",
        "plt.plot(radius, kernel_values)\n",
        "\n",
        "# Add horizontal lines for beta values\n",
        "for b in kernel_params.beta:\n",
        "\tplt.axhline(y=b, color=\"gray\", linestyle=\"--\", alpha=0.5)\n",
        "\n",
        "# Add vertical lines at peaks (segment boundaries)\n",
        "number_of_segments = jnp.count_nonzero(~jnp.isnan(kernel_params.beta), axis=-1)\n",
        "for i in range(number_of_segments):\n",
        "\tpeak_x = (i + 0.5) * kernel_params.r / number_of_segments\n",
        "\tplt.axvline(x=peak_x, color=\"gray\", linestyle=\"--\", alpha=0.5)\n",
        "\n",
        "plt.xlabel(\"Radius\")\n",
        "plt.ylabel(\"Kernel Value\")\n",
        "plt.title(\"Original Kernel Function\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Create a 2D grid for the kernel\n",
        "x = jnp.mgrid[[slice(-dim // 2, dim // 2) for dim in spatial_dims]] / (state_scale * R)\n",
        "d = jnp.linalg.norm(x, axis=0)\n",
        "\n",
        "# Apply the kernel function to the distances\n",
        "kernel_2d = gaussian_kernel_fn(jnp.array(d), kernel_params)\n",
        "\n",
        "# Plot the 2D kernel\n",
        "plt.figure(figsize=(5, 3))\n",
        "plt.imshow(kernel_2d, origin=\"lower\")\n",
        "\n",
        "plt.title(\"2D Kernel\")\n",
        "plt.colorbar(label=\"Kernel Value\")\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Free kernel function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "kernel_params = FreeKernelParams(\n",
        "\tr=0.91,\n",
        "\tb=jnp.array([0.1, 0.6, 0.7]),\n",
        "\ta=jnp.array([0.1, 0.4, 0.8]),\n",
        "\tw=jnp.array([0.05, 0.05, 0.05]),\n",
        ")\n",
        "\n",
        "# Plot the free kernel function\n",
        "radius = jnp.linspace(0.0, 1.0, 1000)\n",
        "kernel_values = jax.vmap(free_kernel_fn, in_axes=(0, None))(radius, kernel_params)\n",
        "\n",
        "plt.figure(figsize=(5, 3))\n",
        "plt.plot(radius, kernel_values)\n",
        "\n",
        "# Add horizontal lines for b values\n",
        "for b in kernel_params.b:\n",
        "\tplt.axhline(y=b, color=\"gray\", linestyle=\"--\", alpha=0.5)\n",
        "\n",
        "# Add vertical lines at peak positions\n",
        "for a in kernel_params.a:\n",
        "\tplt.axvline(x=a * kernel_params.r, color=\"gray\", linestyle=\"--\", alpha=0.5)\n",
        "\n",
        "plt.xlabel(\"Radius\")\n",
        "plt.ylabel(\"Kernel Value\")\n",
        "plt.title(\"Free Kernel Function\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Create a 2D grid for the kernel\n",
        "x = jnp.mgrid[[slice(-dim // 2, dim // 2) for dim in spatial_dims]] / (state_scale * R)\n",
        "d = jnp.linalg.norm(x, axis=0)\n",
        "\n",
        "# Apply the kernel function to the distances\n",
        "kernel_2d = free_kernel_fn(jnp.array(d), kernel_params)\n",
        "\n",
        "# Plot the 2D kernel\n",
        "plt.figure(figsize=(5, 3))\n",
        "plt.imshow(kernel_2d, origin=\"lower\")\n",
        "\n",
        "plt.title(\"2D Kernel\")\n",
        "plt.colorbar(label=\"Kernel Value\")\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Growth Mapping Function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "growth_params = GrowthParams(\n",
        "\tmean=0.22,\n",
        "\tstd=0.06,\n",
        ")\n",
        "\n",
        "# Plot the growth mapping function\n",
        "u_values = jnp.linspace(0, 1, 1000)\n",
        "growth_values = exponential_growth_fn(u_values, growth_params)\n",
        "\n",
        "plt.figure(figsize=(5, 3))\n",
        "plt.plot(u_values, growth_values)\n",
        "\n",
        "# Add vertical line at the mean\n",
        "plt.axvline(x=growth_params.mean, color=\"gray\", linestyle=\"--\", alpha=0.5)\n",
        "\n",
        "plt.xlabel(\"Perception (u)\")\n",
        "plt.ylabel(\"Growth Rate\")\n",
        "\n",
        "plt.title(\"Growth Mapping Function\")\n",
        "plt.grid(True, alpha=0.3)\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Orbium - Putting it Together"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "spatial_dims = (128, 128)\n",
        "channel_size = 1\n",
        "R = 13\n",
        "T = 10\n",
        "state_scale = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "kernel_params = KernelParams(\n",
        "\tr=jnp.array(1.0, jnp.float32),\n",
        "\tbeta=jnp.array([1.0]),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Growth Mapping Function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "growth_params = GrowthParams(\n",
        "\tmean=jnp.array(0.15, jnp.float32),\n",
        "\tstd=jnp.array(0.015, jnp.float32),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Rule Params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "rule_params = LeniaRuleParams(\n",
        "\tchannel_source=jnp.array(0, jnp.int32),\n",
        "\tchannel_target=jnp.array(0, jnp.int32),\n",
        "\tweight=jnp.array(1.0, jnp.float32),\n",
        "\tkernel_params=kernel_params,\n",
        "\tgrowth_params=growth_params,\n",
        ")\n",
        "rule_params = jax.tree.map(lambda x: x[None], rule_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Instantiate system"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cs = Lenia(\n",
        "\tspatial_dims=spatial_dims,\n",
        "\tchannel_size=channel_size,\n",
        "\tR=R,\n",
        "\tT=T,\n",
        "\tstate_scale=state_scale,\n",
        "\trule_params=rule_params,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Sample initial state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def sample_state(pattern):\n",
        "\t\"\"\"Sample a state with a pattern at the center.\"\"\"\n",
        "\t# Calculate the center of the state for each dimension\n",
        "\tmid = tuple(dim // 2 for dim in spatial_dims)\n",
        "\n",
        "\t# Scale pattern\n",
        "\tpattern_scaled = pattern\n",
        "\tfor axis in range(len(spatial_dims)):\n",
        "\t\tpattern_scaled = pattern_scaled.repeat(state_scale, axis=axis)\n",
        "\n",
        "\t# Get the shape of the scaled cells\n",
        "\tpattern_spatial_dims = pattern_scaled.shape[:-1]  # Exclude the channel dimension\n",
        "\n",
        "\t# Create empty state with the shape defined by spatial_dims and channel_size\n",
        "\tstate = jnp.zeros((*spatial_dims, channel_size))\n",
        "\n",
        "\t# Calculate the slice indices for each dimension\n",
        "\tslices = tuple(slice(m - c // 2, m + c - c // 2) for m, c in zip(mid, pattern_spatial_dims))\n",
        "\n",
        "\t# Place the scaled cells in the center of the state\n",
        "\tstate = state.at[slices].set(pattern_scaled)\n",
        "\n",
        "\treturn state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# ruff: noqa: E501\n",
        "# fmt: off\n",
        "orbium = jnp.array(\n",
        "\t[\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.1, 0.14, 0.1, 0.0, 0.0, 0.03, 0.03, 0.0, 0.0, 0.3, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.08, 0.24, 0.3, 0.3, 0.18, 0.14, 0.15, 0.16, 0.15, 0.09, 0.2, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.15, 0.34, 0.44, 0.46, 0.38, 0.18, 0.14, 0.11, 0.13, 0.19, 0.18, 0.45, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.06, 0.13, 0.39, 0.5, 0.5, 0.37, 0.06, 0.0, 0.0, 0.0, 0.02, 0.16, 0.68, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.11, 0.17, 0.17, 0.33, 0.4, 0.38, 0.28, 0.14, 0.0, 0.0, 0.0, 0.0, 0.0, 0.18, 0.42, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.09, 0.18, 0.13, 0.06, 0.08, 0.26, 0.32, 0.32, 0.27, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.82, 0.0, 0.0],\n",
        "\t\t[0.27, 0.0, 0.16, 0.12, 0.0, 0.0, 0.0, 0.25, 0.38, 0.44, 0.45, 0.34, 0.0, 0.0, 0.0, 0.0, 0.0, 0.22, 0.17, 0.0],\n",
        "\t\t[0.0, 0.07, 0.2, 0.02, 0.0, 0.0, 0.0, 0.31, 0.48, 0.57, 0.6, 0.57, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.49, 0.0],\n",
        "\t\t[0.0, 0.59, 0.19, 0.0, 0.0, 0.0, 0.0, 0.2, 0.57, 0.69, 0.76, 0.76, 0.49, 0.0, 0.0, 0.0, 0.0, 0.0, 0.36, 0.0],\n",
        "\t\t[0.0, 0.58, 0.19, 0.0, 0.0, 0.0, 0.0, 0.0, 0.67, 0.83, 0.9, 0.92, 0.87, 0.12, 0.0, 0.0, 0.0, 0.0, 0.22, 0.07],\n",
        "\t\t[0.0, 0.0, 0.46, 0.0, 0.0, 0.0, 0.0, 0.0, 0.7, 0.93, 1.0, 1.0, 1.0, 0.61, 0.0, 0.0, 0.0, 0.0, 0.18, 0.11],\n",
        "\t\t[0.0, 0.0, 0.82, 0.0, 0.0, 0.0, 0.0, 0.0, 0.47, 1.0, 1.0, 0.98, 1.0, 0.96, 0.27, 0.0, 0.0, 0.0, 0.19, 0.1],\n",
        "\t\t[0.0, 0.0, 0.46, 0.0, 0.0, 0.0, 0.0, 0.0, 0.25, 1.0, 1.0, 0.84, 0.92, 0.97, 0.54, 0.14, 0.04, 0.1, 0.21, 0.05],\n",
        "\t\t[0.0, 0.0, 0.0, 0.4, 0.0, 0.0, 0.0, 0.0, 0.09, 0.8, 1.0, 0.82, 0.8, 0.85, 0.63, 0.31, 0.18, 0.19, 0.2, 0.01],\n",
        "\t\t[0.0, 0.0, 0.0, 0.36, 0.1, 0.0, 0.0, 0.0, 0.05, 0.54, 0.86, 0.79, 0.74, 0.72, 0.6, 0.39, 0.28, 0.24, 0.13, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.01, 0.3, 0.07, 0.0, 0.0, 0.08, 0.36, 0.64, 0.7, 0.64, 0.6, 0.51, 0.39, 0.29, 0.19, 0.04, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.1, 0.24, 0.14, 0.1, 0.15, 0.29, 0.45, 0.53, 0.52, 0.46, 0.4, 0.31, 0.21, 0.08, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.08, 0.21, 0.21, 0.22, 0.29, 0.36, 0.39, 0.37, 0.33, 0.26, 0.18, 0.09, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.03, 0.13, 0.19, 0.22, 0.24, 0.24, 0.23, 0.18, 0.13, 0.05, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.02, 0.06, 0.08, 0.09, 0.07, 0.05, 0.01, 0.0, 0.0, 0.0, 0.0, 0.0]\n",
        "\t]\n",
        ")[..., None]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "state_init = sample_state(orbium)\n",
        "state_final = cs(state_init, num_steps=num_steps, sow=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Visualize"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "intermediates = nnx.pop(cs, nnx.Intermediate)\n",
        "states = intermediates.state[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "states = jnp.concatenate([state_init[None], states])\n",
        "frames = nnx.vmap(\n",
        "\tlambda cs, state: cs.render(state),\n",
        "\tin_axes=(None, 0),\n",
        ")(cs, states)\n",
        "\n",
        "mediapy.show_video(frames, width=256, height=256)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Rotator - Using a different type of kernel function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In this section, we demonstrate how to visualize a different pattern, called rotator.\n",
        "\n",
        "We also show the effect of using a larger state size and state scale"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "spatial_dims = (128, 128)\n",
        "channel_size = 1\n",
        "R = 13\n",
        "T = 10\n",
        "state_scale = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Rule Params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "rule_params = LeniaRuleParams(\n",
        "\tchannel_source=jnp.array(0, jnp.int32),\n",
        "\tchannel_target=jnp.array(0, jnp.int32),\n",
        "\tweight=jnp.array(1.0, jnp.float32),\n",
        "\tkernel_params=KernelParams(r=jnp.array(1.0), beta=jnp.array([1.0])),\n",
        "\tgrowth_params=GrowthParams(mean=jnp.array(0.156), std=jnp.array(0.0224)),\n",
        ")\n",
        "rule_params = jax.tree.map(lambda x: x[None], rule_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Instantiate system"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In the following, we will use a polynomial kernel function instead of the Gaussian one."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cs = Lenia(\n",
        "\tspatial_dims=spatial_dims,\n",
        "\tchannel_size=channel_size,\n",
        "\tR=R,\n",
        "\tT=T,\n",
        "\tstate_scale=state_scale,\n",
        "\tkernel_fn=polynomial_kernel_fn,\n",
        "\trule_params=rule_params,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Sample initial state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def sample_state(pattern):\n",
        "\t\"\"\"Sample a state with a pattern at the center.\"\"\"\n",
        "\t# Calculate the center of the state for each dimension\n",
        "\tmid = tuple(dim // 2 for dim in spatial_dims)\n",
        "\n",
        "\t# Scale pattern\n",
        "\tpattern_scaled = pattern\n",
        "\tfor axis in range(len(spatial_dims)):\n",
        "\t\tpattern_scaled = pattern_scaled.repeat(state_scale, axis=axis)\n",
        "\n",
        "\t# Get the shape of the scaled cells\n",
        "\tpattern_spatial_dims = pattern_scaled.shape[:-1]  # Exclude the channel dimension\n",
        "\n",
        "\t# Create empty state with the shape defined by spatial_dims and channel_size\n",
        "\tstate = jnp.zeros((*spatial_dims, channel_size))\n",
        "\n",
        "\t# Calculate the slice indices for each dimension\n",
        "\tslices = tuple(slice(m - c // 2, m + c - c // 2) for m, c in zip(mid, pattern_spatial_dims))\n",
        "\n",
        "\t# Place the scaled cells in the center of the state\n",
        "\tstate = state.at[slices].set(pattern_scaled)\n",
        "\n",
        "\treturn state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# fmt: off\n",
        "rotator = jnp.array(\n",
        "\t[\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.003978, 0.016492, 0.004714, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.045386, 0.351517, 0.417829, 0.367137, 0.37766, 0.426948, 0.431058, 0.282864, 0.081247, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.325473, 0.450995, 0.121737, 0.0, 0.0, 0.0, 0.003113, 0.224278, 0.47101, 0.456459, 0.247231, 0.071609, 0.013126, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.386337, 0.454077, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.27848, 0.524466, 0.464281, 0.242651, 0.096721, 0.038476, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.258817, 0.583802, 0.150994, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.226639, 0.548329, 0.550422, 0.334764, 0.153108, 0.087049, 0.042872, 0.0],\n",
        "\t\t[0.0, 0.008021, 0.502406, 0.524042, 0.059531, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.033946, 0.378866, 0.615467, 0.577527, 0.357306, 0.152872, 0.090425, 0.058275, 0.023345],\n",
        "\t\t[0.0, 0.179756, 0.596317, 0.533619, 0.162612, 0.0, 0.0, 0.0, 0.0, 0.015021, 0.107673, 0.325125, 0.594765, 0.682434, 0.594688, 0.381172, 0.152078, 0.073544, 0.054424, 0.030592],\n",
        "\t\t[0.0, 0.266078, 0.614339, 0.605474, 0.379255, 0.195176, 0.16516, 0.179148, 0.204498, 0.299535, 0.760743, 1.0, 1.0, 1.0, 1.0, 0.490799, 0.237826, 0.069989, 0.043549, 0.022165],\n",
        "\t\t[0.0, 0.333031, 0.64057, 0.686886, 0.60698, 0.509866, 0.450525, 0.389552, 0.434978, 0.859115, 0.94097, 1.0, 1.0, 1.0, 1.0, 1.0, 0.747866, 0.118317, 0.037712, 0.006271],\n",
        "\t\t[0.0, 0.417887, 0.6856, 0.805342, 0.824229, 0.771553, 0.69251, 0.614328, 0.651704, 0.843665, 0.910114, 1.0, 1.0, 0.81765, 0.703404, 0.858469, 1.0, 0.613961, 0.035691, 0.0],\n",
        "\t\t[0.04674, 0.526827, 0.787644, 0.895984, 0.734214, 0.661746, 0.670024, 0.646184, 0.69904, 0.723163, 0.682438, 0.618645, 0.589858, 0.374017, 0.30658, 0.404027, 0.746403, 0.852551, 0.031459, 0.0],\n",
        "\t\t[0.130727, 0.658494, 0.899652, 0.508352, 0.065875, 0.009245, 0.232702, 0.419661, 0.461988, 0.470213, 0.390198, 0.007773, 0.0, 0.010182, 0.080666, 0.17231, 0.44588, 0.819878, 0.034815, 0.0],\n",
        "\t\t[0.198532, 0.810417, 0.63725, 0.031385, 0.0, 0.0, 0.0, 0.0, 0.315842, 0.319248, 0.321024, 0.0, 0.0, 0.0, 0.0, 0.021482, 0.27315, 0.747039, 0.0, 0.0],\n",
        "\t\t[0.217619, 0.968727, 0.104843, 0.0, 0.0, 0.0, 0.0, 0.0, 0.152033, 0.158413, 0.114036, 0.0, 0.0, 0.0, 0.0, 0.0, 0.224751, 0.647423, 0.0, 0.0],\n",
        "\t\t[0.138866, 1.0, 0.093672, 0.0, 0.0, 0.0, 0.0, 0.0, 0.000052, 0.015966, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.281471, 0.455713, 0.0, 0.0],\n",
        "\t\t[0.0, 1.0, 0.145606, 0.005319, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.016878, 0.381439, 0.173336, 0.0, 0.0],\n",
        "\t\t[0.0, 0.97421, 0.262735, 0.096478, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.013827, 0.217967, 0.287352, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.593133, 0.2981, 0.251901, 0.167326, 0.088798, 0.041468, 0.013086, 0.002207, 0.009404, 0.032743, 0.061718, 0.102995, 0.1595, 0.24721, 0.233961, 0.002389, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.610166, 0.15545, 0.200204, 0.228209, 0.241863, 0.243451, 0.270572, 0.446258, 0.376504, 0.174319, 0.154149, 0.12061, 0.074709, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.354313, 0.32245, 0.0, 0.0, 0.0, 0.151173, 0.479517, 0.650744, 0.392183, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.329339, 0.328926, 0.176186, 0.198788, 0.335721, 0.534118, 0.549606, 0.361315, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.090407, 0.217992, 0.190592, 0.174636, 0.222482, 0.375871, 0.265924, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.050256, 0.235176, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.180145, 0.132616, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n",
        "\t\t[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.092581, 0.188519, 0.118256, 0.0, 0.0, 0.0, 0.0]\n",
        "\t]\n",
        ")[..., None]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "state_init = sample_state(rotator)\n",
        "state_final = cs(state_init, num_steps=num_steps, sow=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Visualize"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "intermediates = nnx.pop(cs, nnx.Intermediate)\n",
        "states = intermediates.state[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "states = jnp.concatenate([state_init[None], states])\n",
        "frames = nnx.vmap(\n",
        "\tlambda cs, state: cs.render(state),\n",
        "\tin_axes=(None, 0),\n",
        ")(cs, states)\n",
        "\n",
        "mediapy.show_video(frames, width=256, height=256)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sampling Rules"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "spatial_dims = (128, 128)\n",
        "channel_size = 3\n",
        "R = 12\n",
        "T = 10\n",
        "state_scale = 2\n",
        "\n",
        "# Parameter ranges\n",
        "num_self_kernels = 3\n",
        "num_cross_kernels = 1\n",
        "kernel_rank = 3\n",
        "\n",
        "kernel_r_range = (0.2, 1.0)\n",
        "kernel_beta_range = (0.0, 1.0)\n",
        "growth_mean_range = (0.05, 0.35)\n",
        "growth_std_range = (0.005, 0.1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Define sampling distribution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def sample_kernel_params(key: jax.Array):\n",
        "\t\"\"\"Sample kernel parameters.\"\"\"\n",
        "\tkey_r, key_beta, key_rank, key_max_idx = jax.random.split(key, 4)\n",
        "\n",
        "\t# Spatial extent\n",
        "\tr = jax.random.uniform(key_r, minval=kernel_r_range[0], maxval=kernel_r_range[1])\n",
        "\n",
        "\t# Rank: uniform over {1, ..., kernel_rank}\n",
        "\trank = jax.random.randint(key_rank, (), minval=1, maxval=kernel_rank + 1)\n",
        "\n",
        "\t# Pick which active ring will have the dominant peak height\n",
        "\tmax_idx = jax.random.randint(key_max_idx, (), minval=0, maxval=rank)\n",
        "\n",
        "\t# Sample peak heights uniformly from kernel_beta_range\n",
        "\tbeta = jax.random.uniform(\n",
        "\t\tkey_beta, (kernel_rank,), minval=kernel_beta_range[0], maxval=kernel_beta_range[1]\n",
        "\t)\n",
        "\n",
        "\t# Set the chosen peak to kernel_beta_range[1]\n",
        "\tbeta = beta.at[max_idx].set(kernel_beta_range[1])\n",
        "\n",
        "\t# Mask: active rings keep their peak height, inactive become nan\n",
        "\tactive_mask = jnp.arange(kernel_rank) < rank\n",
        "\tbeta = jnp.where(active_mask, beta, jnp.nan)\n",
        "\n",
        "\treturn KernelParams(r=r, beta=beta)\n",
        "\n",
        "\n",
        "def sample_growth_params(key: jax.Array):\n",
        "\t\"\"\"Sample growth parameters.\"\"\"\n",
        "\tkey_mean, key_std = jax.random.split(key)\n",
        "\n",
        "\tmean = jax.random.uniform(key_mean, minval=growth_mean_range[0], maxval=growth_mean_range[1])\n",
        "\tstd = jax.random.uniform(key_std, minval=growth_std_range[0], maxval=growth_std_range[1])\n",
        "\treturn GrowthParams(mean=mean, std=std)\n",
        "\n",
        "\n",
        "def sample_rule_params(key: jax.Array):\n",
        "\t\"\"\"Sample rule parameters with structured channel connectivity and Dirichlet weights.\"\"\"\n",
        "\tkey_weight, key_kernel_params, key_growth_params = jax.random.split(key, 3)\n",
        "\n",
        "\t# Self-kernels: each channel maps to itself\n",
        "\tself_sources = jnp.repeat(jnp.arange(channel_size), num_self_kernels)\n",
        "\tself_targets = self_sources\n",
        "\n",
        "\t# Cross-kernels: all ordered pairs (i, j) where i != j\n",
        "\tcross_sources = jnp.array(\n",
        "\t\t[i for i in range(channel_size) for j in range(channel_size) if i != j]\n",
        "\t)\n",
        "\tcross_targets = jnp.array(\n",
        "\t\t[j for i in range(channel_size) for j in range(channel_size) if i != j]\n",
        "\t)\n",
        "\tcross_sources = jnp.repeat(cross_sources, num_cross_kernels)\n",
        "\tcross_targets = jnp.repeat(cross_targets, num_cross_kernels)\n",
        "\n",
        "\t# Build structured channel connectivity (following Chan 2020)\n",
        "\tchannel_source = jnp.concatenate([self_sources, cross_sources])\n",
        "\tchannel_target = jnp.concatenate([self_targets, cross_targets])\n",
        "\tnum_rules = channel_source.shape[0]\n",
        "\n",
        "\t# Sample weights jointly via Dirichlet (sum to 1 by construction)\n",
        "\tweight = jax.random.dirichlet(key_weight, alpha=jnp.ones(num_rules))\n",
        "\n",
        "\t# Sample kernel and growth parameters\n",
        "\tkeys_kernel = jax.random.split(key_kernel_params, num_rules)\n",
        "\tkernel_params = jax.vmap(sample_kernel_params)(keys_kernel)\n",
        "\n",
        "\tkeys_growth = jax.random.split(key_growth_params, num_rules)\n",
        "\tgrowth_params = jax.vmap(sample_growth_params)(keys_growth)\n",
        "\n",
        "\treturn LeniaRuleParams(\n",
        "\t\tchannel_source=channel_source,\n",
        "\t\tchannel_target=channel_target,\n",
        "\t\tweight=weight,\n",
        "\t\tkernel_params=kernel_params,\n",
        "\t\tgrowth_params=growth_params,\n",
        "\t)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Sample rules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "seed = 0\n",
        "\n",
        "key = jax.random.key(seed)\n",
        "rule_params = sample_rule_params(key)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Instantiate Lenia with sampled rules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "cs = Lenia(\n",
        "\tspatial_dims=spatial_dims,\n",
        "\tchannel_size=channel_size,\n",
        "\tR=R,\n",
        "\tT=T,\n",
        "\tstate_scale=state_scale,\n",
        "\trule_params=rule_params,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "state_init = sample_state(rotator)\n",
        "state_final = cs(state_init, num_steps=num_steps, sow=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Visualize"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Uniformly sampled Lenia rule parameters often result in patterns that either vanish or explode. For a more stable approach, see [examples/21_flow_lenia.ipynb](../examples/21_flow_lenia.ipynb)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "intermediates = nnx.pop(cs, nnx.Intermediate)\n",
        "states = intermediates.state[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "states = jnp.concatenate([state_init[None], states])\n",
        "frames = nnx.vmap(\n",
        "\tlambda cs, state: cs.render(state),\n",
        "\tin_axes=(None, 0),\n",
        ")(cs, states)\n",
        "\n",
        "mediapy.show_video(frames, width=256, height=256)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.14.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}