doc: Add in data assimilation tutorial and refactor example gallery

.github/workflows/run_tests.yml

@@ -271,8 +271,50 @@ jobs:
           files: coverage*.xml
           fail_ci_if_error: true
 
+  tests-demos:
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        # test with oldest supported Python version only (for slow tests)
+        python-version: ["3.10"]
+
+        demo:
+          - data-assimilation-4dvar
+
+      fail-fast: false
+
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Set up Git repository
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Restore UV environment
+        run: cp production.uv.lock uv.lock
+
+      - name: Install dev requirements
+        run: |
+          uv sync --extra dev --frozen
+
+      - name: Run demos
+        working-directory: demo/${{matrix.demo}}
+        run: |
+          uv pip install jupyter
+          uv run --no-sync jupyter nbconvert --to notebook --execute demo.ipynb
+
+
   all-ok:
-    needs: [tests-base, tests-e2e]
+    needs: [tests-base, tests-e2e, tests-demos]
 
     runs-on: ubuntu-latest
 
@@ -289,4 +331,9 @@ jobs:
             exit 1
           fi
 
+          if [ "${{ needs.tests-demos.result }}" != "success" ]; then
+            echo "Demo tests failed"
+            exit 1
+          fi
+
           echo "All tests passed!"

demo/data-assimilation-4dvar/lorenz_tesseract/tesseract_api.py

@@ -0,0 +1,191 @@
+# Copyright 2025 Pasteur Labs. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Tesseract API module for lorenz
+# Generated by tesseract 0.8.5 on 2025-04-30T09:44:04.234338
+
+from typing import Any
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from pydantic import BaseModel, Field
+
+from tesseract_core.runtime import Array, Differentiable, Float32
+from tesseract_core.runtime.tree_transforms import filter_func, flatten_with_paths
+
+
+class InputSchema(BaseModel):
+    """Input schema for forecasting of Lorenz 96 system."""
+
+    state: Differentiable[Array[(None,), Float32]] = Field(
+        description="A state vector for the Lorenz 96 system"
+    )
+    F: float = Field(description="Forcing parameter for Lorenz 96", default=8.0)
+    dt: float = Field(description="Time step for integration", default=0.05)
+    n_steps: int = Field(description="Number of integration steps", default=1)
+
+
+class OutputSchema(BaseModel):
+    """Output schema for forecasting of Lorenz 96 system."""
+
+    result: Differentiable[Array[(None, None), Float32]] = Field(
+        description="A trajectorie of predictions after integration"
+    )
+
+
+def lorenz96_step(state: jnp.ndarray, F: float, dt: float) -> jnp.ndarray:
+    """Perform one step of RK4 integration for the Lorenz 96 system."""
+
+    def lorenz96_derivatives(x: jnp.ndarray) -> jnp.ndarray:
+        """Compute the derivatives for Lorenz 96 system."""
+        N = x.shape[0]
+        # Create arrays for indices with wraparound
+        ip1 = (jnp.arange(N) + 1) % N  # i+1 with wraparound
+        im1 = (jnp.arange(N) - 1) % N  # i-1 with wraparound
+        im2 = (jnp.arange(N) - 2) % N  # i-2 with wraparound
+        # Compute derivatives: dx_i/dt = (x_{i+1} - x_{i-2}) * x_{i-1} - x_i + F
+        d = (x[ip1] - x[im2]) * x[im1] - x + F
+        return d
+
+    # RK4 integration
+    k1 = lorenz96_derivatives(state)
+    k2 = lorenz96_derivatives(state + dt * k1 / 2)
+    k3 = lorenz96_derivatives(state + dt * k2 / 2)
+    k4 = lorenz96_derivatives(state + dt * k3)
+    return state + dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6
+
+
+def lorenz96_multi_step(
+    state: jnp.ndarray, F: float, dt: float, n_steps: int
+) -> jnp.ndarray:
+    """Perform multiple steps of Lorenz 96 integration using scan."""
+
+    def step_fn(state: jnp.ndarray, _: Any) -> tuple:
+        return lorenz96_step(state, F, dt), state
+
+    _, trajectory = jax.lax.scan(step_fn, state, None, length=n_steps)
+    return trajectory
+
+
+@eqx.filter_jit
+def apply_jit(inputs: dict) -> dict:
+    """The apply_jit function for the Lorenz 96 tesseract."""
+    trajectory = lorenz96_multi_step(**inputs)
+    return dict(result=trajectory)
+
+
+def apply(inputs: InputSchema) -> OutputSchema:
+    """The apply function for the Lorenz 96 tesseract."""
+    return apply_jit(inputs.model_dump())
+
+
+def abstract_eval(abstract_inputs: Any) -> Any:
+    """Calculate output shape of apply from the shape of its inputs."""
+    is_shapedtype_dict = lambda x: type(x) is dict and (x.keys() == {"shape", "dtype"})
+    is_shapedtype_struct = lambda x: isinstance(x, jax.ShapeDtypeStruct)
+
+    jaxified_inputs = jax.tree.map(
+        lambda x: jax.ShapeDtypeStruct(**x) if is_shapedtype_dict(x) else x,
+        abstract_inputs.model_dump(),
+        is_leaf=is_shapedtype_dict,
+    )
+    dynamic_inputs, static_inputs = eqx.partition(
+        jaxified_inputs, filter_spec=is_shapedtype_struct
+    )
+
+    def wrapped_apply(dynamic_inputs: Any) -> Any:
+        inputs = eqx.combine(static_inputs, dynamic_inputs)
+        return apply_jit(inputs)
+
+    jax_shapes = jax.eval_shape(wrapped_apply, dynamic_inputs)
+    return jax.tree.map(
+        lambda x: (
+            {"shape": x.shape, "dtype": str(x.dtype)} if is_shapedtype_struct(x) else x
+        ),
+        jax_shapes,
+        is_leaf=is_shapedtype_struct,
+    )
+
+
+@eqx.filter_jit
+def jac_jit(
+    inputs: dict,
+    jac_inputs: tuple[str],
+    jac_outputs: tuple[str],
+) -> dict:
+    """The jac_jit function for the Lorenz 96 tesseract."""
+    filtered_apply = filter_func(apply_jit, inputs, jac_outputs)
+    return jax.jacrev(filtered_apply)(
+        flatten_with_paths(inputs, include_paths=jac_inputs)
+    )
+
+
+def jacobian(
+    inputs: InputSchema,
+    jac_inputs: set[str],
+    jac_outputs: set[str],
+) -> Any:
+    """The jacobian function for the Lorenz 96 tesseract."""
+    return jac_jit(inputs.model_dump(), tuple(jac_inputs), tuple(jac_outputs))
+
+
+@eqx.filter_jit
+def jvp_jit(
+    inputs: dict,
+    jvp_inputs: tuple[str],
+    jvp_outputs: tuple[str],
+    tangent_vector: dict,
+) -> Any:
+    """The jvp_jit function for the Lorenz 96 tesseract."""
+    filtered_apply = filter_func(apply_jit, inputs, jvp_outputs)
+    return jax.jvp(
+        filtered_apply,
+        [flatten_with_paths(inputs, include_paths=jvp_inputs)],
+        [tangent_vector],
+    )[1]
+
+
+def jacobian_vector_product(
+    inputs: InputSchema,
+    jvp_inputs: set[str],
+    jvp_outputs: set[str],
+    tangent_vector: dict[str, Any],
+) -> Any:
+    """The jacobian vector product function for the Lorenz 96 tesseract."""
+    return jvp_jit(
+        inputs.model_dump(),
+        tuple(jvp_inputs),
+        tuple(jvp_outputs),
+        tangent_vector,
+    )
+
+
+@eqx.filter_jit
+def vjp_jit(
+    inputs: dict,
+    vjp_inputs: tuple[str],
+    vjp_outputs: tuple[str],
+    cotangent_vector: dict,
+) -> Any:
+    """The vjp_jit function for the Lorenz 96 tesseract."""
+    filtered_apply = filter_func(apply_jit, inputs, vjp_outputs)
+    _, vjp_func = jax.vjp(
+        filtered_apply, flatten_with_paths(inputs, include_paths=vjp_inputs)
+    )
+    return vjp_func(cotangent_vector)[0]
+
+
+def vector_jacobian_product(
+    inputs: InputSchema,
+    vjp_inputs: set[str],
+    vjp_outputs: set[str],
+    cotangent_vector: dict[str, Any],
+) -> Any:
+    """The vector jacobian product function for the Lorenz 96 tesseract."""
+    return vjp_jit(
+        inputs.model_dump(),
+        tuple(vjp_inputs),
+        tuple(vjp_outputs),
+        cotangent_vector,
+    )

demo/data-assimilation-4dvar/lorenz_tesseract/tesseract_config.yaml

@@ -0,0 +1,27 @@
+# Tesseract configuration file
+# Generated by tesseract 0.8.5 on 2025-04-30T09:44:04.234338
+
+name: "lorenz"
+version: "0+unknown"
+description: "Tesseract that runs a scan of forecasting for a single-scale Lorenz 96 system"
+
+build_config:
+  # Base image to use for the container, must be Ubuntu or Debian-based
+  # base_image: "debian:bookworm-slim"
+
+  # Platform to build the container for. In general, images can only be executed
+  # on the platform they were built for.
+  target_platform: "native"
+
+  # Additional packages to install in the container (via apt-get)
+  # extra_packages:
+  #  - package_name
+
+  # Data to copy into the container, relative to the project root
+  # package_data:
+  #  - [path/to/source, path/to/destination]
+
+  # Additional Dockerfile commands to run during the build process
+  # custom_build_steps:
+  #  - |
+  #   RUN echo "Hello, World!"

demo/data-assimilation-4dvar/lorenz_tesseract/tesseract_requirements.txt

@@ -0,0 +1,10 @@
+# Tesseract requirements file
+# Generated by tesseract 0.8.5 on 2025-04-30T09:44:04.234338
+
+# Add Python requirements like this:
+numpy==1.26.0
+jax[cpu]==0.5.2
+equinox
+# This may contain private dependencies via SSH URLs:
+# git+ssh://git@github.com/username/repo.git@branch
+# (use `tesseract build --forward-ssh-agent` to grant the builder access to your SSH keys)

demo/data-assimilation-4dvar/requirements.txt

@@ -0,0 +1,4 @@
+jax[cpu]
+numpy
+scipy
+matplotlib

docs/conf.py

@@ -11,6 +11,8 @@
 
 import os
 import re
+import shutil
+from pathlib import Path
 
 from tesseract_core import __version__
 
@@ -40,7 +42,7 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
 extensions = [
-    "myst_parser",
+    "myst_nb",
     "sphinx.ext.intersphinx",
     "sphinx.ext.autodoc",
     "sphinx.ext.napoleon",
@@ -106,3 +108,15 @@ def setup(app) -> None:
     """Sphinx setup function. Used to register custom stuff."""
     # HACK: We zip the examples folder here so that it can be downloaded
     zip_examples_folder()
+
+
+# -- Handle Jupyter notebooks ------------------------------------------------
+
+# Do not execute notebooks during build (just take existing output)
+nb_execution_mode = "off"
+
+# Copy example notebooks to auto_examples folder on every build
+for example_notebook in Path("../demo").glob("*/demo.ipynb"):
+    # Copy the example notebook to the docs folder
+    dest = (Path("content/demo") / example_notebook.parent.name).with_suffix(".ipynb")
+    shutil.copyfile(example_notebook, dest)

docs/content/demo/.gitignore

@@ -0,0 +1,2 @@
+# Notebooks are copied on demand by conf.py
+*.ipynb

docs/content/demo/demo.md

@@ -0,0 +1,26 @@
+# Demo: 4D-variational data assimilation with Tesseract
+
+```{toctree}
+:caption: Contents
+:maxdepth: 1
+:hidden:
+
+data-assimilation-4dvar.ipynb
+lorenz_tesseract.md
+```
+
+(cards-clickable)=
+
+
+:::{card} 4D-Variational Data Assimilation for a Chaotic Dynamical System
+:link: data-assimilation-4dvar.html
+
+Demo using differentiable Tesseracts to implement a 4D-Var data assimilation scheme
+for a chaotic dynamical system
+:::
+
+:::{card} Lorenz Tesseract
+:link: lorenz_tesseract.html
+
+Detailed implementation of the Lorenz Tesseract used in the 4D-Var demo.
+:::

docs/content/demo/lorenz_tesseract.md

@@ -0,0 +1,23 @@
+# Building the JAX Solver Tesseract for Lorenz-96
+
+This examples demonstrates how the JAX solver Tesseract for the Lorenz-96 model is built for the purposes of the data assimilation demo.
+
+Below is the input and output schema definition for the Lorenz Tesseract.
+
+```{literalinclude} ../../../demo/data-assimilation-4dvar/lorenz_tesseract/tesseract_api.py
+:language: python
+:pyobject: InputSchema
+```
+
+```{literalinclude} ../../../demo/data-assimilation-4dvar/lorenz_tesseract/tesseract_api.py
+:language: python
+:pyobject: OutputSchema
+```
+
+Below is the implementation of the apply function, which takes in an initial condition and returns a trajectory of physical states.
+
+```{literalinclude} ../../../demo/data-assimilation-4dvar/lorenz_tesseract/tesseract_api.py
+:language: python
+:start-at: def lorenz96_step
+:end-before:     return apply_jit(inputs.model_dump())
+```