feat(domain): add 2D solver-in-the-loop corrector benchmark

README.md

@@ -29,12 +29,12 @@ Each solver is packaged as a [Tesseract](https://github.com/pasteurlabs/tesserac
 
 ## 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 |
+| 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) | Drag optimization and neural correction | 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
 

docs/generate_results.py

@@ -171,8 +171,9 @@
         "**Can you optimize through it?** End-to-end optimization benchmarks run a "
         "gradient-based optimizer using each solver's own gradients: recovery of "
         "initial conditions or physical parameters, topology optimization, and drag "
-        "minimization. This is the ultimate test, since a gradient can pass the "
-        "finite-difference check yet still fail to drive a full optimization loop."
+        "minimization, plus recurrent neural-corrector training in 2D flow. This is "
+        "the ultimate test, since a gradient can pass the finite-difference check yet "
+        "still fail to drive a full optimization or learning loop."
     ),
     "cost": (
         "**What does it cost?** Wall-clock scaling of the forward and VJP passes "
@@ -457,7 +458,13 @@ def _time(v: tuple[float, float] | None) -> str:
 
 # final-objective metric name by domain optimization experiment, in priority
 # order (the first key present in a result's metrics is used to rank it).
-_OPT_FINAL_KEYS = ("final_error", "final_drag", "final_compliance", "final_ic_error")
+_OPT_FINAL_KEYS = (
+    "final_error",
+    "final_drag",
+    "final_compliance",
+    "final_ic_error",
+    "final_rollout_error",
+)
 
 # Human-readable column labels for the ranked metric. Falls back to a generic
 # de-underscored title when a key is missing here.
@@ -466,6 +473,7 @@ def _time(v: tuple[float, float] | None) -> str:
     "final_drag": "Final drag",
     "final_compliance": "Final compliance",
     "final_ic_error": "Final IC recovery error",
+    "final_rollout_error": "Final rollout error",
 }
 
 

docs/index.qmd

@@ -22,7 +22,7 @@ Built on [tesseract-core](https://github.com/pasteurlabs/tesseract-core), which
 |:---|:-------|:------------------|:-------------|:---------|
 | **H** | Heat transfer | Conductivity inversion | 128 | [deal.II](solvers.qmd#thermal-mesh-dealii-heat), [FEniCS](solvers.qmd#thermal-mesh-fenics-heat), [Firedrake](solvers.qmd#thermal-mesh-firedrake-heat), [JAX-FEM](solvers.qmd#thermal-mesh-jax-fem), [torch-fem](solvers.qmd#thermal-mesh-torch-fem-thermal) |
 | **S** | Structural mechanics | Compliance minimization (SIMP) | 2048 | [deal.II](solvers.qmd#structural-mesh-dealii-structural), [FEniCS](solvers.qmd#structural-mesh-fenics-structural), [Firedrake](solvers.qmd#structural-mesh-firedrake-structural), [JAX-FEM](solvers.qmd#structural-mesh-jax-fem), [TopOpt.jl](solvers.qmd#structural-mesh-topopt-jl) |
-| **F2** | Incompressible fluids (2D) | Inflow optimization for drag min. | 32 | [JAX-CFD](solvers.qmd#navier-stokes-grid-jax-cfd), [PhiFlow](solvers.qmd#navier-stokes-grid-phiflow), [INS.jl](solvers.qmd#navier-stokes-grid-ins-jl), [XLB](solvers.qmd#navier-stokes-grid-xlb), [PICT](solvers.qmd#navier-stokes-grid-pict), [Warp-NS](solvers.qmd#navier-stokes-grid-warp-ns), [OpenFOAM](solvers.qmd#navier-stokes-grid-openfoam) |
+| **F2** | Incompressible fluids (2D) | Drag optimization and solver-in-the-loop neural correction | 32 | [JAX-CFD](solvers.qmd#navier-stokes-grid-jax-cfd), [PhiFlow](solvers.qmd#navier-stokes-grid-phiflow), [INS.jl](solvers.qmd#navier-stokes-grid-ins-jl), [XLB](solvers.qmd#navier-stokes-grid-xlb), [PICT](solvers.qmd#navier-stokes-grid-pict), [Warp-NS](solvers.qmd#navier-stokes-grid-warp-ns), [OpenFOAM](solvers.qmd#navier-stokes-grid-openfoam) |
 | **F3** | 3D Navier-Stokes | Initial condition recovery | 12288 | [PhiFlow](solvers.qmd#navier-stokes-grid-phiflow), [XLB](solvers.qmd#navier-stokes-grid-xlb), [PICT](solvers.qmd#navier-stokes-grid-pict), [Warp-NS](solvers.qmd#navier-stokes-grid-warp-ns), [Exponax](solvers.qmd#navier-stokes-grid-exponax), [INS.jl](solvers.qmd#navier-stokes-grid-ins-jl), [OpenFOAM](solvers.qmd#navier-stokes-grid-openfoam) |
 
 The full catalog — per-solver numerical scheme, AD strategy, image name, schema fields, and known limitations — is on the [Solver Reference](solvers.qmd) page.
@@ -54,7 +54,7 @@ The FD step size is swept over $\varepsilon \in \{10^{-6}, \dots, 10^{-1}\}$ and
 
 **Forward accuracy.** Resolution sweep against a reference solver (OpenFOAM for fluids, deal.II for structures/thermal). Where an analytical solution exists (e.g. the Taylor–Green vortex), precision is also measured against it. Physical-law adherence (e.g. a divergence-free velocity field, $\nabla \cdot \mathbf{u} = 0$) is checked where applicable.
 
-**Optimization convergence.** The ultimate test: can you actually optimize *through* the solver? We run with L-BFGS on each benchmark task for a fixed iteration budget (500 iterations for H, F2, F3; 2500 for S). A solver **succeeds** if it reaches a final objective within 1% of the best solution any solver achieved within the budget. This catches solvers whose gradients pass the FD check pointwise but still fail to drive a full optimization loop.
+**Optimization convergence.** The ultimate test: can you actually optimize *through* the solver? Alongside inverse and design problems, the 2D fluid benchmark trains an identical neural corrector recurrently through every differentiable solver. Fixed update budgets expose solvers whose gradients pass a pointwise FD check but still fail to drive a full optimization or learning loop.
 
 ## Benchmark suites
 

mosaic/benchmarks/core/runner.py

@@ -167,6 +167,8 @@ def _tracked_tesseract(tag: str, gpus: object, docker_args: object):
     :meth:`Tesseract.from_tesseract_api` (no Docker, just imports a
     ``tesseract_api.py``) — meant for end-to-end framework tests
     with the dummy tesseracts in ``tests/dummy_tesseracts/``.
+    Tags prefixed with ``url:`` connect to an already-running Tesseract
+    service, which supports daemon-free launchers such as Slurm + Pyxis.
     Every other tag is a Docker image and goes through
     :meth:`Tesseract.from_image`.
     """
@@ -175,6 +177,16 @@ def _tracked_tesseract(tag: str, gpus: object, docker_args: object):
         yield _inprocess_tesseract(tag[len("inprocess:") :])
         return
 
+    if tag.startswith("url:"):
+        # The remote service is owned by the caller (for example an enclosing
+        # Slurm allocation), so Mosaic must neither enter a serve context nor
+        # attempt container cleanup.
+        yield Tesseract.from_url(
+            tag[len("url:") :],
+            timeout=MOSAIC_TESSERACT_TIMEOUT_TUPLE,
+        )
+        return
+
     t = Tesseract.from_image(
         tag,
         gpus=gpus,

mosaic/benchmarks/core/utils.py

@@ -131,6 +131,44 @@ def _debug_run(run: dict) -> None:
     for key, cap in [("max_iters", 50), ("patience", 10)]:
         if key in optim:
             optim[key] = min(optim[key], cap)
+    training = run.get("training", {})
+    for key, cap in [
+        ("max_updates", 2),
+        ("unroll", 2),
+        ("hidden_channels", 8),
+        ("kernel_size", 3),
+    ]:
+        if key in training:
+            training[key] = min(training[key], cap)
+    if training:
+        training["check_grad"] = False
+    if "model_seeds" in training:
+        training["model_seeds"] = list(training["model_seeds"])[:1]
+    dataset = run.get("dataset", {})
+    for key in ("train_seeds", "test_seeds"):
+        if key in dataset:
+            dataset[key] = list(dataset[key])[:1]
+    if "prefix_audit_seeds" in dataset:
+        retained_seeds = set(dataset.get("train_seeds", [])) | set(
+            dataset.get("test_seeds", [])
+        )
+        dataset["prefix_audit_seeds"] = [
+            seed for seed in dataset["prefix_audit_seeds"] if seed in retained_seeds
+        ]
+    if "train_frames" in dataset:
+        dataset["train_frames"] = min(dataset["train_frames"], 3)
+    evaluation = run.get("evaluation", {})
+    if "rollout_frames" in evaluation:
+        evaluation["rollout_frames"] = min(evaluation["rollout_frames"], 3)
+    if "prefix_audit_frames" in dataset:
+        max_frame = max(
+            int(dataset.get("train_frames", 0)),
+            int(evaluation.get("rollout_frames", 0)),
+            int(training.get("unroll", 0)),
+        )
+        dataset["prefix_audit_frames"] = [
+            frame for frame in dataset["prefix_audit_frames"] if frame <= max_frame
+        ]
     cost = run.get("cost", {})
     for key in ("N_values", "steps_values"):
         if key in cost: