https://github.com/pasteurlabs/mosaic

A benchmark suite and re-usable collection of differentiable physics solvers.
https://github.com/pasteurlabs/mosaic

automatic-differentiation benchmark differentiable-physics differentiable-programming scientific-computing scientific-machine-learning tesseract

Last synced: 2 days ago
JSON representation

A benchmark suite and re-usable collection of differentiable physics solvers.

Host: GitHub
URL: https://github.com/pasteurlabs/mosaic
Owner: pasteurlabs
License: apache-2.0
Created: 2026-05-04T19:07:03.000Z (3 months ago)
Default Branch: main
Last Pushed: 2026-07-06T21:49:10.000Z (18 days ago)
Last Synced: 2026-07-06T22:17:46.057Z (18 days ago)
Topics: automatic-differentiation, benchmark, differentiable-physics, differentiable-programming, scientific-computing, scientific-machine-learning, tesseract
Language: Python
Homepage: https://docs.pasteurlabs.ai/projects/mosaic/stable
Size: 73.9 MB
Stars: 11
Watchers: 2
Forks: 0
Open Issues: 15
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md

Awesome Lists containing this project

README

# 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.

![Overview of Mosaic: diverse solver backends are wrapped behind a uniform containerized interface (Tesseract), enabling cross-solver comparison on shared benchmark tasks across different physical domains.](docs/visual_abstract.png)

## 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](https://github.com/pasteurlabs/tesseract-core) 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

## 📊 Results

**[Browse the benchmark results →](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results_ns_grid.html)** — no setup required.

Per-domain pages with every plot, solver rankings, and the full evaluation protocol, refreshed on each release:
[Navier–Stokes 2D](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results_ns_grid.html) ·
[Navier–Stokes 3D](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results_ns_3d_grid.html) ·
[Structural mechanics](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results_structural_mesh.html) ·
[Heat transfer](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results_thermal_mesh.html)

## So, which solver is best?

It depends, and there likely isn't even a "best" solver for a given task. These solvers were built for different jobs, in different languages, with different differentiation strategies, and very different user experiences. Mosaic runs them side by side to make those differences visible, but also marginalizes over a lot of nuance.

What we can offer are a few patterns that recur across domains and are unlikely to shift with more configuration tuning, because they follow from fundamental constraints rather than from how exactly a solver is set up.

> [!NOTE]
> These takeaways are the high-level version and are glossing over a lot of detail. For the full analysis see our [paper](https://arxiv.org/abs/2606.27895).

**On well-conditioned problems, compute performance and integration effort matter most.** Heat transfer and structural mechanics are our controlled cases, where the physics is linear and the objective is smooth. Every differentiable solver reproduces the finite-difference gradient direction almost exactly (cosine similarity above 0.999), and although pointwise accuracy still spans a few orders of magnitude, that spread does not decide the outcome, since all backends agree on the forward solution and converge to effectively the same optimum. What is left to distinguish them is how fast the forward-plus-VJP pass runs and how much work it took to wire up. For a large class of inverse and design problems, you can simply pick the solver that is fastest and/or easiest to integrate.

**Rollout length is a fundamental bottleneck.** Pushing the 3D Navier–Stokes gradient out to longer and longer horizons separates a solver's success into two separate properties. The first is _gradient conditioning_, set by the numerical discretization: as the flow turns chaotic, projection-based schemes tend to accumulate spurious non-physical modes that corrupt the gradient (return non-finite values within a couple hundred steps), whereas lattice-Boltzmann and finite-volume schemes keep the gradient bounded far longer. The second is _memory_, set by the differentiation strategy: tape-based reverse-mode AD stores every step and eventually runs out, while methods that avoid unrolling the full trace stay within budget. So at long horizons, prefer a well-conditioned discretization with an AD strategy that does not grow with the rollout, and always check the gradient magnitude rather than trusting that a run which completes has produced a usable gradient.

**Whether a solver can even run the task is a major constraint.** For example, on the 2D drag-minimization problem, only three of seven fluid solvers can participate at all. Some assume periodic boundaries and cannot represent the obstacle, others rely on a spectral projection incompatible with the geometry. Before comparing gradient quality, check whether a solver's assumptions match your problem.

Note that today's tasks are deliberately on the simpler end — laminar flow, linear elasticity, modest horizons — and configurations are mostly out-of-the-box rather than expert-tuned. Harder regimes may well surface cases where raw gradient quality becomes the binding constraint.

> [!TIP]
> Browse the [per-domain pages](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/results.html) for the numbers behind each of these claims.

## 📖 Documentation

Two versions are published. **You most likely want to use [stable](https://docs.pasteurlabs.ai/projects/mosaic/stable/) — it tracks the latest release and is the most reliable (all solvers benchmarked in the same run).** [Latest](https://docs.pasteurlabs.ai/projects/mosaic/latest/) tracks the `main` branch and may aggregate results from different runs.

Start here:
[Getting Started](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/getting-started.html) ·
[Use Solvers Elsewhere](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/standalone.html) ·
[Solver Reference](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/solvers.html) ·
[How it works](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/internals.html) ·
[Add a Backend](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/tutorial.html)

> [!TIP]
> **Reproducing our paper?** See the [`v0.1+paper-repro`](https://github.com/pasteurlabs/mosaic/tree/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](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html).

> [!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.

```bash
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):

```bash
$ 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

```bash
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:

```bash
# 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 ` 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.

```bash
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](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/getting-started.html).

## Use Tesseracts in your own code

Every solver is a standalone [Tesseract](https://github.com/pasteurlabs/tesseract-core) you can call from your own research code — no benchmark harness required.

```bash
# 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:

```python
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](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/standalone.html) (GPU, mesh-based solvers, gotchas) · [Solver Reference](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/solvers.html) (per-solver catalog with image names)

Programmatic API — run evaluations without the CLI

```python
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](CONTRIBUTING.md#tuning-an-existing-solver)
- **Add a solver** to an existing domain — three files under `mosaic/tesseracts///`. → [Add a Solver tutorial](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/tutorial.html#part-a--add-a-solver-to-an-existing-domain)
- **Add a benchmark domain** — scaffold with `mosaic new-domain --from-template `. → [Add a Domain tutorial](https://docs.pasteurlabs.ai/projects/mosaic/stable/docs/tutorial.html#part-b--add-a-new-benchmark-domain)

[CONTRIBUTING.md](CONTRIBUTING.md) covers code style, the PR workflow, and building the docs locally. For questions, visit the [Tesseract Forum](https://si-tesseract.discourse.group/).

## 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
```

## Citation

If you use Mosaic in your research, please cite the paper ([arXiv:2606.27895](https://arxiv.org/abs/2606.27895)):

```bibtex
@misc{rehmann2026mosaic,
title = {Mosaic: A Benchmark Suite for Differentiable Physics Solvers},
author = {Rehmann, Andrin and Zimmermann, Heiko and H{\"a}fner, Dion},
year = {2026},
eprint = {2606.27895},
archivePrefix = {arXiv},
primaryClass = {physics.comp-ph},
url = {https://arxiv.org/abs/2606.27895}
}
```

## License

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

ecosyste.ms

Data

Tools

Indexes

Applications

Experiments

Awesome

https://github.com/pasteurlabs/mosaic

Awesome Lists containing this project

README