doc: Add new landing page for Tesseract ecosystem (#562)

#### Relevant issue or PR

n/a

#### Description of changes

- Adds a landing page for the Tesseract ecosystem as docs index.
- Pulls in demos from Tesseract-JAX so we unify all ecosystem demos in
one place.

#### Testing done

:eyes:

Author: dionhaefner

.github/workflows/run_tests.yml

@@ -281,9 +281,23 @@ jobs:
         run: |
           uv sync --extra dev --frozen
 
+      - name: Install system dependencies
+        run: |
+          sudo apt-get update -qq
+          sudo apt-get install -y -qq xvfb libgl1 libosmesa6
+
+      - name: Install demo requirements
+        working-directory: demo/${{matrix.demo}}
+        run: |
+          uv pip install -r requirements.txt
+
       - name: Run demos
         working-directory: demo/${{matrix.demo}}
+        env:
+          DISPLAY: ":99"
+          PYVISTA_OFF_SCREEN: "true"
         run: |
+          Xvfb :99 -screen 0 1024x768x24 &
           uv pip install jupyter
           uv run --no-sync jupyter nbconvert --to notebook --execute demo.ipynb
 

README.md

@@ -5,7 +5,7 @@
 
 ### Tesseract Core
 
-Universal, autodiff-native software components for [Simulation Intelligence](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/content/misc/faq.html#what-is-simulation-intelligence) 📦
+Universal components for differentiable scientific computing 📦
 
 [Read the docs](https://docs.pasteurlabs.ai/projects/tesseract-core/latest/) |
 [Showcases & tutorials](https://si-tesseract.discourse.group/c/showcase/11) |

demo/cfd-optimization/cfd-tesseract/tesseract_api.py

@@ -0,0 +1,222 @@
+# Copyright 2025 Pasteur Labs. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax_cfd.base as cfd
+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):
+    v0: Differentiable[
+        Array[
+            (
+                None,
+                None,
+                None,
+            ),
+            Float32,
+        ]
+    ] = Field(description="3D Array defining the initial velocity field")
+    density: float = Field(description="Density of the fluid")
+    viscosity: float = Field(description="Viscosity of the fluid")
+    inner_steps: int = Field(
+        description="Number of solver steps for each timestep", default=25
+    )
+    outer_steps: int = Field(description="Number of timesteps steps", default=10)
+    max_velocity: float = Field(description="Maximum velocity", default=2.0)
+    cfl_safety_factor: float = Field(description="CFL safety factor", default=0.5)
+    domain_size_x: float = Field(description="Domain size x", default=1.0)
+    domain_size_y: float = Field(description="Domain size y", default=1.0)
+
+
+class OutputSchema(BaseModel):
+    result: Differentiable[Array[(None, None, None), Float32]] = Field(
+        description="3D Array defining the final velocity field"
+    )
+
+
+def cfd_fwd(
+    v0: jnp.ndarray,
+    density: float,
+    viscosity: float,
+    inner_steps: int,
+    outer_steps: int,
+    max_velocity: float,
+    cfl_safety_factor: float,
+    domain_size_x: float,
+    domain_size_y: float,
+) -> tuple[jax.Array, jax.Array]:
+    """Compute the final velocity field using the semi-implicit Navier-Stokes equations.
+
+    Args:
+        v0: Initial velocity field.
+        density: Density of the fluid.
+        viscosity: Viscosity of the fluid.
+        inner_steps: Number of solver steps for each timestep.
+        outer_steps: Number of timesteps steps.
+        max_velocity: Maximum velocity.
+        cfl_safety_factor: CFL safety factor.
+        domain_size_x: Domain size in x direction.
+        domain_size_y: Domain size in y direction.
+
+    Returns:
+        Final velocity field.
+    """
+    vx0 = v0[..., 0]
+    vy0 = v0[..., 1]
+    bc = cfd.boundaries.HomogeneousBoundaryConditions(
+        (
+            (cfd.boundaries.BCType.PERIODIC, cfd.boundaries.BCType.PERIODIC),
+            (cfd.boundaries.BCType.PERIODIC, cfd.boundaries.BCType.PERIODIC),
+        )
+    )
+
+    # reconstruct grid from input
+    grid = cfd.grids.Grid(
+        vx0.shape, domain=((0.0, domain_size_x), (0.0, domain_size_y))
+    )
+
+    vx0 = cfd.grids.GridArray(vx0, grid=grid, offset=(1.0, 0.5))
+    vy0 = cfd.grids.GridArray(vy0, grid=grid, offset=(0.5, 1.0))
+
+    # reconstruct GridVariable from input
+    vx0 = cfd.grids.GridVariable(vx0, bc)
+    vy0 = cfd.grids.GridVariable(vy0, bc)
+    v0 = (vx0, vy0)
+
+    # Choose a time step.
+    dt = cfd.equations.stable_time_step(
+        max_velocity, cfl_safety_factor, viscosity, grid
+    )
+
+    # Define a step function and use it to compute a trajectory.
+    step_fn = cfd.funcutils.repeated(
+        cfd.equations.semi_implicit_navier_stokes(
+            density=density, viscosity=viscosity, dt=dt, grid=grid
+        ),
+        steps=inner_steps,
+    )
+    rollout_fn = cfd.funcutils.trajectory(step_fn, outer_steps)
+    _, trajectory = jax.device_get(rollout_fn(v0))
+    vxn = trajectory[0].array.data[-1]
+    vyn = trajectory[1].array.data[-1]
+    return jnp.stack([vxn, vyn], axis=-1)
+
+
+@eqx.filter_jit
+def apply_jit(inputs: dict) -> dict:
+    vn = cfd_fwd(**inputs)
+    return dict(result=vn)
+
+
+def apply(inputs: InputSchema) -> OutputSchema:
+    return apply_jit(inputs.model_dump())
+
+
+def jacobian(
+    inputs: InputSchema,
+    jac_inputs: set[str],
+    jac_outputs: set[str],
+):
+    return jac_jit(inputs.model_dump(), tuple(jac_inputs), tuple(jac_outputs))
+
+
+def jacobian_vector_product(
+    inputs: InputSchema,
+    jvp_inputs: set[str],
+    jvp_outputs: set[str],
+    tangent_vector: dict[str, Any],
+):
+    return jvp_jit(
+        inputs.model_dump(),
+        tuple(jvp_inputs),
+        tuple(jvp_outputs),
+        tangent_vector,
+    )
+
+
+def vector_jacobian_product(
+    inputs: InputSchema,
+    vjp_inputs: set[str],
+    vjp_outputs: set[str],
+    cotangent_vector: dict[str, Any],
+):
+    return vjp_jit(
+        inputs.model_dump(),
+        tuple(vjp_inputs),
+        tuple(vjp_outputs),
+        cotangent_vector,
+    )
+
+
+def abstract_eval(abstract_inputs):
+    """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):
+        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],
+):
+    filtered_apply = filter_func(apply_jit, inputs, jac_outputs)
+    return jax.jacrev(filtered_apply)(
+        flatten_with_paths(inputs, include_paths=jac_inputs)
+    )
+
+
+@eqx.filter_jit
+def jvp_jit(
+    inputs: dict, jvp_inputs: tuple[str], jvp_outputs: tuple[str], tangent_vector: dict
+):
+    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]
+
+
+@eqx.filter_jit
+def vjp_jit(
+    inputs: dict,
+    vjp_inputs: tuple[str],
+    vjp_outputs: tuple[str],
+    cotangent_vector: dict,
+):
+    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]

demo/cfd-optimization/cfd-tesseract/tesseract_config.yaml

@@ -0,0 +1,8 @@
+name: jax-cfd
+version: "0.1.0"
+description: |
+  Tesseract that runs a differentiable 2D Navier Stokes simulation on a 2D grid.
+
+build_config:
+  package_data: []
+  custom_build_steps: []

demo/cfd-optimization/cfd-tesseract/tesseract_requirements.txt

@@ -0,0 +1,4 @@
+numpy==1.26.4
+jax-cfd==0.2.1
+jax[cpu]==0.6.0
+equinox

demo/cfd-optimization/demo.ipynb

@@ -0,0 +1,6 @@
+jax_cfd
+matplotlib
+scipy
+tqdm
+pillow
+tesseract-jax

demo/cfd-optimization/pl.png

@@ -512,15 +512,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "But wait! Our model is a Tesseract, which is not natively compatible with JAX's automatic differentiation. So we need to wrap the Tesseract in a JAX-compatible function that can be differentiated, by registering a JAX custom derivative for the Tesseract's `apply` method, via [`jax.pure_callback`](https://docs.jax.dev/en/latest/_autosummary/jax.pure_callback.html) and [`jax.custom_vjp`](https://docs.jax.dev/en/latest/_autosummary/jax.custom_vjp.html).\n",
-    "\n",
-    "<div class=admonition note alert alert-warning>\n",
-    "<p class=admonition-title>Note</p>\n",
-    "\n",
-    "While this works in this example, registering Tesseracts by hand is not the recommended way to use Tesseracts in JAX. Instead, you should check out [Tesseract-JAX](https://github.com/pasteurlabs/tesseract-jax), which automatically registers Tesseracts as JAX-compatible functions, allowing you to use them seamlessly in JAX code without needing to manually define callbacks or register custom derivatives.\n",
-    "</div>"
-   ]
+   "source": "But wait! Our model is a Tesseract, which is not natively compatible with JAX's automatic differentiation. So we need to wrap the Tesseract in a JAX-compatible function that can be differentiated, by registering a JAX custom derivative for the Tesseract's `apply` method, via [`jax.pure_callback`](https://docs.jax.dev/en/latest/_autosummary/jax.pure_callback.html) and [`jax.custom_vjp`](https://docs.jax.dev/en/latest/_autosummary/jax.custom_vjp.html).\n\nConcretely, we define three functions:\n- **`M(x)`** — calls the Tesseract's `apply` endpoint via `jax.pure_callback`, so JAX treats it as an opaque function.\n- **`M_fwd(x)`** — the forward pass, which runs `M(x)` and saves the input `x` as a residual for the backward pass.\n- **`M_bwd(residuals, cotangents)`** — the backward pass, which calls the Tesseract's `vector_jacobian_product` endpoint (again via `pure_callback`) to compute the gradient.\n\nThis is the manual version of what [Tesseract-JAX](https://github.com/pasteurlabs/tesseract-jax) does automatically. The other two demos in this series (CFD optimization and FEM shape optimization) use Tesseract-JAX instead, which handles all of this registration behind the scenes. We show the manual approach here to illustrate what happens under the hood.\n\n<div class=admonition note alert alert-warning>\n<p class=admonition-title>Note</p>\n\nFor your own projects, we recommend using [Tesseract-JAX](https://github.com/pasteurlabs/tesseract-jax) rather than registering custom derivatives by hand. It automatically registers Tesseracts as JAX-compatible functions, allowing you to use them seamlessly in JAX code.\n</div>"
   },
   {
    "cell_type": "code",