jax-fractal-renderer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# JAX Fractal Renderer\n",
    "\n",
    "[https://github.com/tripplyons/jax-fractal-renderer](https://github.com/tripplyons/jax-fractal-renderer)\n",
    "\n",
    "Hint: If you are using Colab, select \"Runtime\" > \"Change runtime type\" > \"Hardware accelerator\" > \"GPU\" to use GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install jax tqdm optax matplotlib pillow numpy psgd-jax flax\n",
    "!mkdir -p output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import optax\n",
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "from IPython.display import display\n",
    "import numpy as np\n",
    "# the optimizer that I found to work the best\n",
    "from psgd_jax.affine import scale_by_affine\n",
    "\n",
    "index_dtype = jnp.int32\n",
    "dtype = jnp.float32\n",
    "\n",
    "def apply_transform(x, y, settings):\n",
    "    a, b, c, d, e, f, power = settings\n",
    "\n",
    "    # apply a matrix multiplication and vector addition using the constants\n",
    "    x_prime = x * a + y * b + c\n",
    "    y_prime = x * d + y * e + f\n",
    "\n",
    "    # use a non-linearity that helps keep the points centered\n",
    "    scale_squared = x_prime ** 2 + y_prime ** 2\n",
    "    # divide by 2 for square root\n",
    "    x_prime /= scale_squared ** (power / 2)\n",
    "    y_prime /= scale_squared ** (power / 2)\n",
    "\n",
    "    return x_prime, y_prime\n",
    "\n",
    "def update(x, y, settings, rng):\n",
    "    rate, settings_1, settings_2 = settings\n",
    "\n",
    "    # apply the first transform to the points\n",
    "    x_prime_1, y_prime_1 = apply_transform(x, y, settings_1)\n",
    "    # apply the second transform to the points\n",
    "    x_prime_2, y_prime_2 = apply_transform(x, y, settings_2)\n",
    "\n",
    "    # randomly choose one of the transforms (chaos game)\n",
    "    mask = jax.random.uniform(rng, shape=(x.shape[0],), minval=0, maxval=1, dtype=dtype)\n",
    "    mask = 1 * (mask < rate)\n",
    "\n",
    "    # choose the first transform where the mask is 1\n",
    "    # and the second transform where the mask is 0\n",
    "    x = x_prime_1 * mask + x_prime_2 * (1 - mask)\n",
    "    y = y_prime_1 * mask + y_prime_2 * (1 - mask)\n",
    "\n",
    "    return x, y\n",
    "\n",
    "# density is the vector of the density values at each point\n",
    "# it is initialized to zeros\n",
    "# it is updated where the current points are\n",
    "def optimize(density, batch_size, x_res, y_res, settings, iterations, bounds, rng):\n",
    "    min_x, max_x, min_y, max_y = bounds\n",
    "\n",
    "    # use a warmup schedule to start the optimization at a lower learning rate\n",
    "    schedule = optax.schedules.warmup_constant_schedule(0, 1, iterations // 10)\n",
    "    optimizer = optax.chain(\n",
    "        scale_by_affine(), # optimizer\n",
    "        optax.scale_by_schedule(schedule) # learning rate schedule\n",
    "    )\n",
    "    opt_state = optimizer.init(density)\n",
    "\n",
    "    # start at the origin\n",
    "    x = jnp.zeros(batch_size, dtype=dtype)\n",
    "    y = jnp.zeros(batch_size, dtype=dtype)\n",
    "\n",
    "    # calculate how large each pixel is\n",
    "    dx = (max_x - min_x) / (x_res - 1)\n",
    "    dy = (max_y - min_y) / (y_res - 1)\n",
    "\n",
    "    @jax.jit\n",
    "    def step(density, x, y, opt_state, rng):\n",
    "        update_rng, rng = jax.random.split(rng, 2)\n",
    "\n",
    "        # update the points\n",
    "        x, y = update(x, y, settings, update_rng)\n",
    "\n",
    "        # find the nearest pixel to the current points\n",
    "        x_index = jnp.round((x - min_x) / dx).astype(index_dtype)\n",
    "        y_index = jnp.round((y - min_y) / dy).astype(index_dtype)\n",
    "\n",
    "        # go to an invalid index after the end of our vector if the coordinate is negative or too large\n",
    "        x_index = x_index + (-x_index + x_res) * ((x_index < 0) + (x_index >= x_res))\n",
    "        y_index = y_index + (-y_index + y_res * x_res) * ((y_index < 0) + (y_index >= y_res))\n",
    "\n",
    "        # set the gradient to 0 by default\n",
    "        grad = jnp.zeros((x_res * y_res,), dtype=dtype)\n",
    "        # set the gradient to 1 at the current locations\n",
    "        grad = grad.at[x_index * y_res + y_index].set(1)\n",
    "        # normalize the gradient to help the optimizer converge faster\n",
    "        grad = jax.nn.standardize(grad, axis=0)\n",
    "\n",
    "        updates, opt_state = optimizer.update(grad, opt_state, density)\n",
    "        density = optax.apply_updates(density, updates)\n",
    "\n",
    "        return density, x, y, opt_state, rng\n",
    "\n",
    "    # run the optimization for the specified number of iterations\n",
    "    for _ in tqdm(range(iterations)):\n",
    "        density, x, y, opt_state, rng = step(density, x, y, opt_state, rng)\n",
    "\n",
    "    return density\n",
    "\n",
    "def main(rng, res):\n",
    "    x_res = res\n",
    "    y_res = res\n",
    "    batch_size = 1024\n",
    "    iterations = 10000\n",
    "    scale = 3\n",
    "    bounds = (-scale, scale, -scale, scale)\n",
    "\n",
    "    density = jnp.zeros((x_res * y_res), dtype=dtype)\n",
    "\n",
    "    params_key, rng = jax.random.split(rng, 2)\n",
    "\n",
    "    min_scale = 1\n",
    "    max_scale = 1.5\n",
    "\n",
    "    min_bias = 0.25\n",
    "    max_bias = 0.75\n",
    "\n",
    "    min_power = 0.2\n",
    "    max_power = 0.3\n",
    "\n",
    "    min_rate = 0.4\n",
    "    max_rate = 0.5\n",
    "\n",
    "    def make_settings(rng):\n",
    "        angle1_key, angle2_key, angle3_key, scale1_key, scale2_key, bias_scale_key, power_key = jax.random.split(rng, 7)\n",
    "\n",
    "        angle1 = jax.random.uniform(angle1_key, shape=(), minval=0, maxval=2 * jnp.pi, dtype=dtype)\n",
    "        angle2 = jax.random.uniform(angle2_key, shape=(), minval=0, maxval=2 * jnp.pi, dtype=dtype)\n",
    "        angle3 = jax.random.uniform(angle3_key, shape=(), minval=0, maxval=2 * jnp.pi, dtype=dtype)\n",
    "        scale1 = jax.random.uniform(scale1_key, shape=(), minval=min_scale, maxval=max_scale, dtype=dtype)\n",
    "        scale2 = jax.random.uniform(scale2_key, shape=(), minval=min_scale, maxval=max_scale, dtype=dtype)\n",
    "        bias_scale = jax.random.uniform(bias_scale_key, shape=(), minval=min_bias, maxval=max_bias, dtype=dtype)\n",
    "        power = jax.random.uniform(power_key, shape=(), minval=min_power, maxval=max_power, dtype=dtype)\n",
    "\n",
    "        print(scale1, scale2, bias_scale, power)\n",
    "\n",
    "        settings = (\n",
    "            jnp.cos(angle1) * scale1,\n",
    "            jnp.cos(angle2) * scale2,\n",
    "            jnp.cos(angle3) * bias_scale,\n",
    "            jnp.sin(angle1) * scale1,\n",
    "            jnp.sin(angle2) * scale2,\n",
    "            jnp.sin(angle3) * bias_scale,\n",
    "            power\n",
    "        )\n",
    "\n",
    "        return settings\n",
    "\n",
    "    settings1_key, settings2_key, rate_key = jax.random.split(params_key, 3)\n",
    "    settings1 = make_settings(settings1_key)\n",
    "    settings2 = make_settings(settings2_key)\n",
    "    rate = jax.random.uniform(rate_key, shape=(), minval=min_rate, maxval=max_rate, dtype=dtype)\n",
    "    print(rate)\n",
    "\n",
    "    settings = (rate, settings1, settings2)\n",
    "\n",
    "    density = optimize(density, batch_size, x_res, y_res, settings, iterations, bounds, rng)\n",
    "\n",
    "    return density.reshape((x_res, y_res))\n",
    "\n",
    "def plot(name, density):\n",
    "    density -= density.min()\n",
    "    density /= density.max()\n",
    "\n",
    "    density = (255 * density).astype(jnp.uint8)\n",
    "    density = np.array(density)\n",
    "\n",
    "    img = Image.fromarray(density)\n",
    "    display(img)\n",
    "    img.save(f'output/{name}.png')\n",
    "\n",
    "    return density\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    res = 1024\n",
    "    num_rows = 5\n",
    "    num_cols = 5\n",
    "    num_images = num_rows * num_cols\n",
    "    images = []\n",
    "\n",
    "    for seed in range(num_images):\n",
    "        rng = jax.random.PRNGKey(seed)\n",
    "        density = main(rng, res)\n",
    "        images.append(plot(seed, density))\n",
    "\n",
    "    grid = np.array(images).reshape((num_rows, num_cols, res, res))\n",
    "    grid = np.transpose(grid, (0, 2, 1, 3)).reshape(res * num_rows, res * num_cols)\n",
    "\n",
    "    grid_img = Image.fromarray(grid).resize((res, res))\n",
    "    display(grid_img)\n",
    "    grid_img.save('output/grid.png')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fractal",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}