heading structure and blurb

Author: jacanchaplais

examples/cfd/demo.ipynb

@@ -12,7 +12,19 @@
     "All examples are expected to run from the `examples/<example_name>` directory of the [Tesseract-JAX repository](https://github.com/pasteurlabs/tesseract-jax).\n",
     "</div>\n",
     "\n",
-    "`cfd-tesseract` is a differentiable Navier-Stokes solver based on [JAX-CFD](https://github.com/google/jax-cfd) that is wrapped in a Tesseract. "
+    "`cfd-tesseract` is a differentiable Navier-Stokes solver based on [JAX-CFD](https://github.com/google/jax-cfd) that is wrapped in a Tesseract. \n",
+    "\n",
+    "In this demo, you will learn how to:\n",
+    "1. Define a fluid velocity field with periodic boundary conditions, using JAX-CFD\n",
+    "1. Pass that velocity field to a Tesseract which applies a differentiable Navier-Stokes solver to evolve it\n",
+    "1. Use gradient-based optimizations to find an initial field configuration which spontaneously evolves into the Pasteur Labs Logo!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 1: Build `cfd-tesseract` + install dependencies"
    ]
   },
   {
@@ -64,16 +76,30 @@
     "%pip install -r requirements.txt -q"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 2: Test the forward evaluation with Tesseract-JAX"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Let's set up the Tesseract with Tesseract-JAX and test a simple forward evaluation."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "First, we'll define an initial guess for the velocity field over a grid. The resulting `vx` and `vy` give our horizontal and vertical velocity fields, respectively."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -97,8 +123,15 @@
     "v0 = cfd.initial_conditions.filtered_velocity_field(\n",
     "    jax.random.PRNGKey(seed), grid, max_velocity\n",
     ")\n",
-    "vx, vy = v0\n",
-    "\n",
+    "vx, vy = v0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "cfd_tesseract = Tesseract.from_image(\"jax-cfd\")\n",
     "cfd_tesseract.serve()\n",
     "\n",
@@ -115,17 +148,22 @@
     "\n",
     "v0 = np.stack([np.array(vx.array.data), np.array(vy.array.data)], axis=-1)\n",
     "\n",
+    "\n",
     "def cfd_tesseract_fn(v0):\n",
     "    return apply_tesseract(cfd_tesseract, inputs=dict(v0=v0, **params))\n",
     "\n",
+    "\n",
     "outputs = cfd_tesseract_fn(v0)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Lets look at the results of the forward pass."
+    "## Step 3: Visualise the outputs from our Tesseract\n",
+    "\n",
+    "Using the results of the forward pass, we can set up a basic approach for visualising our velocity field.\n",
+    "We'll use `matplotlib` to show the $x$ and $y$ components of the velocity as heatmaps."
    ]
   },
   {
@@ -156,9 +194,16 @@
     "plt.show()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can take this further, and view the vorticity field of the fluid; first we define periodic boundary conditions."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -167,34 +212,15 @@
     "        (cfd.boundaries.BCType.PERIODIC, cfd.boundaries.BCType.PERIODIC),\n",
     "        (cfd.boundaries.BCType.PERIODIC, cfd.boundaries.BCType.PERIODIC),\n",
     "    )\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "HomogeneousBoundaryConditions(types=(('periodic', 'periodic'), ('periodic', 'periodic')), bc_values=((0.0, 0.0), (0.0, 0.0)))"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "bc"
+    ")\n",
+    "print(bc)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "And look at its vorticity field."
+    "Next we define a vorticity function (recalling that vorticity is the curl of the flow velocity, *ie.* $\\omega = \\nabla \\times \\mathbf{v}$."
    ]
   },
   {
@@ -240,9 +266,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Optimization\n",
+    "## Step 4: Optimizing the fluid to evolve into the Pasteur Labs logo\n",
+    "\n",
+    "Now we want to perform an actual optimization.\n",
+    "Our goal is to find the initial state, such that the final state resembles logo of Pasteur Labs.\n",
     "\n",
-    "Now we want to perform an actual optimization. The target is to find the initial state, such that the final state looks a bit like the logo of Pasteur Labs. Herefore we first load the logo.\n"
+    "Let's start by loading in our logo!"
    ]
   },
   {
@@ -334,7 +363,7 @@
     "    vort = vorticity(vxn, vyn)\n",
     "\n",
     "    # decrase difference of vorticity and image and ensure the field is divergence free\n",
-    "    return mse(vort, img) + 0.05 * mse(div, 0.0)\n"
+    "    return mse(vort, img) + 0.05 * mse(div, 0.0)"
    ]
   },
   {
@@ -376,11 +405,11 @@
     "from tqdm import tqdm\n",
     "\n",
     "\n",
-    "def loss_fn_capt(v0_flat, img=img, xlen: int= grid.shape[0]):\n",
+    "def loss_fn_capt(v0_flat, img=img, xlen: int = grid.shape[0]):\n",
     "    total_len = len(v0_flat)\n",
     "    ylen = (total_len // 2) // xlen\n",
-    "    v0x = v0_flat[:total_len//2].reshape(xlen, ylen)\n",
-    "    v0y = v0_flat[total_len//2:].reshape(xlen, ylen)\n",
+    "    v0x = v0_flat[: total_len // 2].reshape(xlen, ylen)\n",
+    "    v0y = v0_flat[total_len // 2 :].reshape(xlen, ylen)\n",
     "\n",
     "    div = divergence(v0x, v0y)\n",
     "\n",
@@ -395,7 +424,6 @@
     "    return mse(vort, img) + 0.05 * mse(div, 0.0)\n",
     "\n",
     "\n",
-    "\n",
     "v0_field = cfd.initial_conditions.filtered_velocity_field(\n",
     "    jax.random.PRNGKey(221), grid, max_velocity\n",
     ")\n",
@@ -406,12 +434,13 @@
     "max_iter = 400\n",
     "with tqdm(total=max_iter) as pbar:\n",
     "    i = 0\n",
+    "\n",
     "    def callback(intermediate_result):\n",
     "        global i\n",
     "        i += 1\n",
     "        pbar.set_postfix(loss=f\"{intermediate_result.fun:.4f}\")\n",
     "        pbar.update(1)\n",
-    "    \n",
+    "\n",
     "    opt = minimize(\n",
     "        grad_fn,\n",
     "        v0_flat,\n",
@@ -423,13 +452,13 @@
     "\n",
     "if i < max_iter:\n",
     "    print(\"Optimisation converged before reaching max_iter!\")\n",
-    "    \n",
+    "\n",
     "# Reshape to generate gif in next cell\n",
     "v0_flat = opt.x\n",
     "xlen = grid.shape[0]\n",
     "ylen = grid.shape[1]\n",
-    "v0x = v0_flat[:xlen*ylen].reshape(xlen, ylen)\n",
-    "v0y = v0_flat[xlen*ylen:].reshape(xlen, ylen)\n",
+    "v0x = v0_flat[: xlen * ylen].reshape(xlen, ylen)\n",
+    "v0y = v0_flat[xlen * ylen :].reshape(xlen, ylen)\n",
     "v0 = jnp.stack([v0x, v0y], axis=-1)"
    ]
   },
@@ -11791,10 +11820,8 @@
     }
    ],
    "source": [
-    "from IPython.display import HTML\n",
-    "\n",
     "import matplotlib.animation as animation\n",
-    "\n",
+    "from IPython.display import HTML\n",
     "\n",
     "trajectory = []\n",
     "\n",