{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Self-autoencoding MNIST Digits [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/maxencefaldor/cax/blob/main/examples/45_self_autoencoding_mnist.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 jax\n",
    "import jax.numpy as jnp\n",
    "import mediapy\n",
    "import optax\n",
    "import torchvision\n",
    "from flax import nnx\n",
    "from jax import Array\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from cax.core import ComplexSystem\n",
    "from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel\n",
    "from cax.core.update.nca_update import NCAUpdate\n",
    "from cax.nn.pool import Pool\n",
    "from cax.utils import clip_and_uint8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Self-autoencoding MNIST Digits](https://raw.githubusercontent.com/maxencefaldor/cax/main/docs/assets/self_autoencoding_mnist.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 0\n",
    "\n",
    "channel_size = 16\n",
    "spatial_dims = (28, 28, 42)\n",
    "num_kernels = 4\n",
    "hidden_size = 256\n",
    "cell_dropout_rate = 0.5\n",
    "\n",
    "num_steps = 96\n",
    "pool_size = 1_024\n",
    "batch_size = 8\n",
    "learning_rate = 1e-3\n",
    "\n",
    "key = jax.random.key(seed)\n",
    "rngs = nnx.Rngs(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load MNIST dataset\n",
    "ds_train = torchvision.datasets.MNIST(root=\"./data\", train=True, download=True)\n",
    "ds_test = torchvision.datasets.MNIST(root=\"./data\", train=False, download=True)\n",
    "\n",
    "# Convert to jax.Array\n",
    "x_train = jnp.array([x.resize(spatial_dims[:2]) for x, _ in ds_train])[..., None] / 255\n",
    "x_test = jnp.array([x.resize(spatial_dims[:2]) for x, _ in ds_test])[..., None] / 255\n",
    "\n",
    "# Visualize\n",
    "mediapy.show_images(x_train[:8], width=128, height=128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SelfAutoencodingNCA(ComplexSystem):\n",
    "\t\"\"\"Self-Autoencoding Neural Cellular Automata class.\"\"\"\n",
    "\n",
    "\tdef __init__(self, *, rngs: nnx.Rngs):\n",
    "\t\t\"\"\"Initialize Self-Autoencoding NCA.\"\"\"\n",
    "\t\tself.perceive = ConvPerceive(\n",
    "\t\t\tchannel_size=channel_size,\n",
    "\t\t\tperception_size=num_kernels * channel_size,\n",
    "\t\t\tkernel_size=(3, 3, 3),\n",
    "\t\t\tfeature_group_count=channel_size,\n",
    "\t\t\trngs=rngs,\n",
    "\t\t)\n",
    "\t\tself.update = NCAUpdate(\n",
    "\t\t\tchannel_size=channel_size,\n",
    "\t\t\tperception_size=num_kernels * channel_size,\n",
    "\t\t\thidden_layer_sizes=(hidden_size,),\n",
    "\t\t\tcell_dropout_rate=cell_dropout_rate,\n",
    "\t\t\tkernel_size=(3, 3, 3),\n",
    "\t\t\tzeros_init=True,\n",
    "\t\t\trngs=rngs,\n",
    "\t\t)\n",
    "\n",
    "\t\t# Initialize kernel with sobel filters\n",
    "\t\tkernel = jnp.concatenate([identity_kernel(num_dims=3), grad_kernel(num_dims=3)], axis=-1)\n",
    "\t\tkernel = jnp.expand_dims(jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2)\n",
    "\t\tself.perceive.conv.kernel[...] = kernel\n",
    "\n",
    "\tdef _step(self, state: Array, input: Array | None = None, *, sow: bool = False) -> Array:\n",
    "\t\t\"\"\"Perform a single step.\"\"\"\n",
    "\t\t# Extract x\n",
    "\t\tx = state[..., 0, -1:]\n",
    "\n",
    "\t\t# Step\n",
    "\t\tperception = self.perceive(state)\n",
    "\t\tnext_state = self.update(state, perception, input)\n",
    "\n",
    "\t\t# Mask\n",
    "\t\tmid = tuple(size // 2 for size in spatial_dims)\n",
    "\t\tcenter = next_state[..., *mid, :]\n",
    "\t\tnext_state = next_state.at[..., mid[-1], :].set(0.0)  # Mask\n",
    "\t\tnext_state = next_state.at[..., *mid, :].set(center)  # Except center cell\n",
    "\n",
    "\t\t# Override\n",
    "\t\tnext_state = next_state.at[..., 0, -1:].set(x)\n",
    "\n",
    "\t\tif sow:\n",
    "\t\t\tself.sow(nnx.Intermediate, \"state\", next_state)\n",
    "\n",
    "\t\treturn next_state\n",
    "\n",
    "\t@nnx.jit\n",
    "\tdef render(self, state):\n",
    "\t\t\"\"\"Render state to RGB.\"\"\"\n",
    "\t\tgray = state[..., -1, -1:]\n",
    "\t\trgb = jnp.repeat(gray, 3, axis=-1)\n",
    "\n",
    "\t\t# Clip values to valid range and convert to uint8\n",
    "\t\treturn clip_and_uint8(rgb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cs = SelfAutoencodingNCA(rngs=rngs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = nnx.state(cs, nnx.Param)\n",
    "print(\"Number of params:\", sum(x.size for x in jax.tree.leaves(params)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample initial state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_state(key):\n",
    "\t\"\"\"Sample a state with a random image.\"\"\"\n",
    "\t# Init state\n",
    "\tstate = jnp.zeros(spatial_dims + (channel_size,))\n",
    "\n",
    "\t# Sample random image\n",
    "\tx_idx = jax.random.choice(key, x_train.shape[0])\n",
    "\tx = x_train[x_idx]\n",
    "\n",
    "\t# Set image in state\n",
    "\tstate = state.at[..., 0, -1:].set(x)\n",
    "\treturn state, x_idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key, subkey = jax.random.split(key)\n",
    "\n",
    "keys = jax.random.split(subkey, pool_size)\n",
    "state, x_idx = jax.vmap(sample_state)(keys)\n",
    "\n",
    "pool = Pool.create({\"state\": state, \"x_idx\": x_idx})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_sched = optax.linear_schedule(\n",
    "\tinit_value=learning_rate, end_value=0.5 * learning_rate, transition_steps=2_000\n",
    ")\n",
    "\n",
    "optimizer = optax.chain(\n",
    "\toptax.clip_by_global_norm(1.0),\n",
    "\toptax.adam(learning_rate=lr_sched),\n",
    ")\n",
    "\n",
    "update_params = nnx.All(\n",
    "\tnnx.Param,\n",
    "\t# nnx.PathContains(\"update\"),\n",
    ")\n",
    "optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mse(state, x):\n",
    "\t\"\"\"Mean Squared Error.\"\"\"\n",
    "\treturn jnp.mean(jnp.square(state[..., :, -1:] - x[..., None, :]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nnx.jit\n",
    "def loss_fn(cs, state, x):\n",
    "\t\"\"\"Loss function.\"\"\"\n",
    "\tstate_axes = nnx.StateAxes({nnx.RngState: 0, nnx.Intermediate: 0, ...: None})\n",
    "\tnnx.split_rngs(splits=batch_size)(\n",
    "\t\tnnx.vmap(\n",
    "\t\t\tlambda cs, state: cs(state, num_steps=num_steps, sow=True),\n",
    "\t\t\tin_axes=(state_axes, 0),\n",
    "\t\t)\n",
    "\t)(cs, state)\n",
    "\n",
    "\t# Get intermediate states\n",
    "\tintermediates = nnx.pop(cs, nnx.Intermediate)\n",
    "\tstate = intermediates.state[0]\n",
    "\n",
    "\t# Sample a random step\n",
    "\tidx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps)\n",
    "\tstate = state[jnp.arange(batch_size), idx]\n",
    "\n",
    "\tloss = mse(state, x)\n",
    "\treturn loss, state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@nnx.jit\n",
    "def train_step(cs, optimizer, pool, key):\n",
    "\t\"\"\"Train step.\"\"\"\n",
    "\tsample_key, sample_state_key = jax.random.split(key)\n",
    "\n",
    "\t# Sample from pool\n",
    "\tpool_idx, batch = pool.sample(sample_key, batch_size=batch_size)\n",
    "\tcurrent_state = batch[\"state\"]\n",
    "\tcurrent_x_idx = batch[\"x_idx\"]\n",
    "\tcurrent_x = x_train[current_x_idx]\n",
    "\n",
    "\t# Sort by descending loss\n",
    "\tsort_idx = jnp.argsort(jax.vmap(mse)(current_state, current_x), descending=True)\n",
    "\tpool_idx = pool_idx[sort_idx]\n",
    "\tcurrent_state = current_state[sort_idx]\n",
    "\tcurrent_x_idx = current_x_idx[sort_idx]\n",
    "\n",
    "\t# Sample a new state to replace the worst\n",
    "\tnew_state, new_x_idx = sample_state(sample_state_key)\n",
    "\tcurrent_state = current_state.at[0].set(new_state)\n",
    "\tcurrent_x_idx = current_x_idx.at[0].set(new_x_idx)\n",
    "\tcurrent_x = x_train[current_x_idx]\n",
    "\n",
    "\t(loss, current_state), grad = nnx.value_and_grad(\n",
    "\t\tloss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params)\n",
    "\t)(cs, current_state, current_x)\n",
    "\toptimizer.update(cs, grad)\n",
    "\n",
    "\tpool = pool.update(pool_idx, {\"state\": current_state, \"x_idx\": current_x_idx})\n",
    "\treturn loss, pool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_train_steps = 2 * 8_192\n",
    "print_interval = 128\n",
    "\n",
    "pbar = tqdm(range(num_train_steps), desc=\"Training\", unit=\"train_step\")\n",
    "losses = []\n",
    "for i in pbar:\n",
    "\tkey, subkey = jax.random.split(key)\n",
    "\tloss, pool = train_step(cs, optimizer, pool, subkey)\n",
    "\tlosses.append(loss)\n",
    "\n",
    "\tif i % print_interval == 0 or i == num_train_steps - 1:\n",
    "\t\tavg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:])\n",
    "\t\tpbar.set_postfix({\"Average Loss\": f\"{avg_loss:.3e}\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_examples = 8\n",
    "\n",
    "key, subkey = jax.random.split(key)\n",
    "keys = jax.random.split(subkey, num_examples)\n",
    "state_init, x_idx = jax.vmap(sample_state)(keys)\n",
    "\n",
    "state_axes = nnx.StateAxes({nnx.RngState: 0, nnx.Intermediate: 0, ...: None})\n",
    "state_final = nnx.split_rngs(splits=num_examples)(\n",
    "\tnnx.vmap(\n",
    "\t\tlambda cs, state: cs(state, num_steps=2 * num_steps, sow=True),\n",
    "\t\tin_axes=(state_axes, 0),\n",
    "\t)\n",
    ")(cs, state_init)"
   ]
  },
  {
   "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], axis=1)\n",
    "frames = nnx.vmap(\n",
    "\tlambda cs, state: cs.render(state),\n",
    "\tin_axes=(None, 0),\n",
    ")(cs, states)\n",
    "\n",
    "mediapy.show_images(x_train[x_idx], width=128, height=128)\n",
    "mediapy.show_videos(frames, width=128, height=128, codec=\"gif\")"
   ]
  }
 ],
 "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}