mosaic-bench 0.1.1

A benchmark suite for differentiable physics solvers

Mosaic

A benchmark suite and reusable collection of differentiable physics solvers.
Think OpenAI Gym, but for differentiable physics: a growing catalog of tasks across physical domains, with a standardized interface and evaluation protocol for every solver and their gradients.

What Mosaic measures

If you optimize or train through a physics simulation, the solver must return two correct things: the forward prediction and its gradient (the vector–Jacobian product, VJP). Most benchmarks check only the forward pass. Mosaic checks both, and scores every solver on three axes:

• Gradient accuracy — does the VJP match a finite-difference ground truth?
• Computational cost — wall-clock time (forward + VJP) and peak memory.
• Setup compatibility — does the solver even run on the task, or do structural constraints rule it out?

Each solver is packaged as a Tesseract container exposing a uniform apply / vjp interface. A single harness can therefore compare solvers across languages and AD backends (JAX, PyTorch, Julia, hand-written C++ adjoints) by talking only to that common interface.

Domains & solvers

ID	Domain	Optimization task	Solvers
H	Heat transfer	Conductivity inversion	deal.II, FEniCS, Firedrake, JAX-FEM, torch-fem
S	Structural mechanics	Compliance minimization (SIMP)	deal.II, FEniCS, Firedrake, JAX-FEM, TopOpt.jl
F2	Incompressible fluids (2D)	Inflow optimization (drag)	JAX-CFD, PhiFlow, INS.jl, XLB, PICT, Warp-NS, OpenFOAM
F3	3D Navier–Stokes	Initial condition recovery	PhiFlow, XLB, PICT, Warp-NS, Exponax, INS.jl, OpenFOAM

📊 Results

Browse the benchmark results → — no setup required.

Per-domain pages with every plot, solver rankings, and the full evaluation protocol, refreshed on each release: Navier–Stokes 2D · Navier–Stokes 3D · Structural mechanics · Heat transfer

📖 Documentation

Two versions are published. You most likely want to use stable — it tracks the latest release and is the most reliable (all solvers benchmarked in the same run). Latest tracks the main branch and may aggregate results from different runs.

Start here: Getting Started · Use Solvers Elsewhere · Solver Reference · How it works · Add a Backend

[!TIP] Reproducing our paper? See the v0.1+paper-repro tag for figure-generation code, pinned dependencies, and step-by-step instructions.

---

Run the benchmarks

Requires Python ≥ 3.10, Docker, and — for GPU solvers — the NVIDIA Container Toolkit.

[!WARNING] We strongly recommend Linux with Docker Engine. Docker Desktop on macOS/Windows runs containers in a VM, adding significant overhead and ARM compatibility issues on Apple Silicon. On macOS/Windows, prefer a Linux VM or WSL 2 with Docker Engine installed natively.

git clone https://github.com/pasteurlabs/mosaic && cd mosaic
uv sync          # or: pip install -e .
mosaic run       # builds containers, runs experiments, generates plots

Verify your setup with a single-problem --debug run (reduced grid sizes, finishes in minutes):

$ mosaic run -p thermal-mesh --suites forward --debug
──────────────────────────── problem: thermal-mesh ─────────────────────────────
──────────────────────────────────── build ─────────────────────────────────────
  deal.II          → dealii_heat_thermal_mesh:latest     (3.6s)
  FEniCS           → fenics_heat_thermal_mesh:latest     (3.2s)
  Firedrake        → firedrake_heat_thermal_mesh:latest  (2.4s)
  JAX-FEM          → jax_fem_thermal_mesh:latest         (5.1s)
  torch-fem        → torch_fem_thermal_mesh:latest       (4.8s)
─────────────────────────────────── summary ────────────────────────────────────
┏━━━━━━━━━━━━━━┳━━━━━━━━━┓
┃ problem      ┃ forward ┃
┡━━━━━━━━━━━━━━╇━━━━━━━━━┩
│ thermal-mesh │   ok    │
└──────────────┴─────────┘

Common workflows — inspect results, pick solvers, re-run a subset

Inspect results

mosaic status                        # per-experiment completion table
mosaic status -p ns-grid -f          # single domain with failure reasons
mosaic status --format md > report.md
mosaic status --format json > snap.json

Pick which solvers run

-s / --solvers takes either a flat CSV (union across every problem) or a per-problem map:

# Flat CSV — each problem keeps only the listed solvers that exist there.
mosaic run -s OpenFOAM,XLB,deal.II,JAX-FEM

# Per-problem map — explicit picks per domain.
mosaic run -s "ns-grid=XLB,jax-cfd;structural-mesh=Firedrake,JAX-FEM"

Re-run a subset

mosaic run --only <state[,…]> re-executes only cells in the given state, leaving fresh-ok cells alone — handy for iterating on one solver or recovering from a partial failure.

mosaic run --only failed              # re-run only failed cells
mosaic run --only failed,stale        # plus anything invalidated by the harness/source
mosaic run --only missing             # first-time runs only
mosaic run -s PhiFlow --only excluded # re-check after dropping an exclusion

States: failed, anom, missing, stale, excluded. Combine with -p / --suites / -e / -s for finer scoping.

The full CLI reference and smoke-test workflow live in Getting Started.

Use Tesseracts in your own code

Every solver is a standalone Tesseract you can call from your own research code — no benchmark harness required.

# Shared schemas (deps: pydantic + tesseract-core only)
pip install -e mosaic/mosaic_shared

# For containerised usage (recommended): also install tesseract-jax
pip install tesseract-core tesseract-jax jax

Via container (works for every solver regardless of language). Build the image once, then call it from JAX with full grad support:

import jax
import jax.numpy as jnp
from tesseract_core import Tesseract
from tesseract_jax import apply_tesseract
from mosaic_shared.problems.navier_stokes_grid.schemas import make_vortex_ic

ic = make_vortex_ic(N=64, seed=42)
inputs = {"v0": ic, "viscosity": jnp.array([0.01]), "steps": 50}

with Tesseract.from_image("exponax_navier_stokes_grid:latest") as t:
    outputs = apply_tesseract(t, inputs)
    grad_v0 = jax.grad(lambda v0: jnp.mean(
        apply_tesseract(t, {**inputs, "v0": v0})["result"] ** 2
    ))(inputs["v0"])

A local (no Docker) path is also available for Python-only solvers — see the full guide below.

📖 Standalone Usage (GPU, mesh-based solvers, gotchas) · Solver Reference (per-solver catalog with image names)

Programmatic API — run evaluations without the CLI

from mosaic import get_config, PROBLEMS

cfg = get_config("ns-grid")           # Problem for 2-D Navier-Stokes
print(cfg.solver_names)               # available solver backends

# Each (suite, experiment) is registered on the Problem as an Experiment
# closure. Invoke one directly with a {solver_name: image_tag} mapping:
tags = {s.name: s.image_tag for s in cfg.solvers}
results = cfg.experiments["gradient/fd_check"].fn(cfg, tags)

Top-level imports: PROBLEMS, get_config, Problem, SolverSpec, IcSpec, and the shared suite-kernel modules forward, gradient, cost, optimization (from mosaic.benchmarks.problems.shared).

Contribute

Mosaic is designed to grow with the community. Three ways in, roughly by scope:

• Tune an existing solver — improve an out-of-the-box config. Snapshot mosaic status --format json before/after and include the diff. → CONTRIBUTING.md
• Add a solver to an existing domain — three files under mosaic/tesseracts/<domain>/<solver-name>/. → Add a Solver tutorial
• Add a benchmark domain — scaffold with mosaic new-domain <name> --from-template <template>. → Add a Domain tutorial

CONTRIBUTING.md covers code style, the PR workflow, and building the docs locally. For questions, visit the Tesseract Forum.

Project structure

mosaic/
  benchmarks/             # evaluation harness (Python package: mosaic.benchmarks)
    cli.py                # command-line interface
    core/                 # runner, config, hardware detection, solver auto-discovery
    problems/             # per-domain packages (ns-grid, ns-3d-grid, structural-mesh, thermal-mesh)
      shared/             # cross-domain suite kernels (forward, gradient, cost, optimization) + plots
    plots/                # plotting infrastructure
  templates/              # task templates for scaffolding new domains
  tesseracts/             # solver backends (each is a Tesseract container)
    mosaic_shared/        # shared Tesseract interface schemas (also pip-installable)
    navier-stokes-grid/   # JAX-CFD, PhiFlow, XLB, PICT, Warp-NS, etc.
    structural-mesh/      # deal.II, FEniCS, Firedrake, JAX-FEM, TopOpt.jl
    thermal-mesh/         # deal.II, FEniCS, Firedrake, JAX-FEM, torch-fem
  tests/                  # unit tests (run with pytest)
docs/                     # Quarto documentation site

License

Apache 2.0. Individual solver backends retain their upstream licenses, documented per solver in the repository.