fix(solver): apply ns-grid forcing in every PhiFlow RK4 stage

.github/workflows/benchmark.yml

@@ -528,6 +528,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Decide status
+        env:
+          GH_TOKEN: ${{ github.token }}
         run: |
           echo "plan: ${{ needs.plan.result }}"
           echo "build: ${{ needs.build.result }}"
@@ -536,6 +538,22 @@ jobs:
           echo "label: ${{ needs.plan.outputs.label }}"
           echo "mosaic_touched: ${{ needs.plan.outputs.mosaic_touched }}"
 
+          # A pull request opened with a label fires both `opened` and
+          # `labeled`; the concurrency group cancels the first run. When a
+          # newer run exists for the same commit, that run decides, so this
+          # one passes instead of leaving a failed check next to it. A run
+          # cancelled by hand, with no newer run, still fails below.
+          if [[ "${{ needs.plan.result }}" == "cancelled" \
+             && "${{ github.event_name }}" == "pull_request" ]]; then
+            newer=$(gh api \
+              "repos/${{ github.repository }}/actions/workflows/benchmark.yml/runs?head_sha=${{ github.event.pull_request.head.sha }}" \
+              --jq '[.workflow_runs[] | select(.id > ${{ github.run_id }})] | length')
+            if [[ "$newer" -gt 0 ]]; then
+              echo "Superseded by a newer run for this commit, which decides the status."
+              exit 0
+            fi
+          fi
+
           # plan must always succeed
           if [[ "${{ needs.plan.result }}" != "success" ]]; then
             echo "::error::Plan failed."

docs/generate_results.py

@@ -411,13 +411,14 @@ def _rank_gradient(problem: str) -> tuple[list[str], list[tuple]] | None:
     ]
 
 
-def _cost_at_max_n(results: list[dict]) -> dict[str, tuple[float, float]]:
-    """For each solver, return ``(N, mean_s)`` at the largest N it completed.
+def _cost_at_max_n(results: list[dict]) -> dict[str, tuple[float, float, float]]:
+    """Return ``(N, mean_s, first_call_s)`` per solver at the largest N it completed.
 
     Picking the time at the max successful size avoids flattering a solver that
     OOMs early. ``results`` is a flat cost ``results`` list (sweep_value = N).
+    ``first_call_s`` is NaN for results recorded before the harness kept it.
     """
-    out: dict[str, tuple[float, float]] = {}
+    out: dict[str, tuple[float, float, float]] = {}
     for r in results:
         m = r.get("metrics") or {}
         if m.get("status") == "failed":
@@ -427,9 +428,11 @@ def _cost_at_max_n(results: list[dict]) -> dict[str, tuple[float, float]]:
         if not isinstance(mean_s, int | float) or not math.isfinite(mean_s):
             continue
         nf = float(n) if isinstance(n, int | float | str) and str(n) else float("nan")
+        first_s = m.get("first_call_s")
+        first_s = float(first_s) if isinstance(first_s, int | float) else float("nan")
         prev = out.get(r["solver"])
         if prev is None or nf > prev[0]:
-            out[r["solver"]] = (nf, float(mean_s))
+            out[r["solver"]] = (nf, float(mean_s), first_s)
     return out
 
 
@@ -439,25 +442,30 @@ def _rank_cost(problem: str) -> tuple[list[str], list[tuple]] | None:
     Forward time comes from the ``spatial_cost`` sweep and VJP time from
     ``vjp_cost/by_N``; each is taken at the largest N that solver completed
     (shown in the "at N" column). Forward-only solvers have no VJP entry, which
-    shows as "—". Ranked by forward time (faster wins).
+    shows as "—". Ranked by forward time (faster wins). The forward pass's
+    first call at that N, which includes compilation and setup, is shown next
+    to it and does not affect the ranking.
     """
     fwd = _cost_at_max_n(_load_suite_results(problem, "cost", "spatial_cost"))
     vjp = _cost_at_max_n(_load_suite_results(problem, "cost", "vjp_cost/by_N"))
     if len(fwd) < 2:
         return None
 
-    def _time(v: tuple[float, float] | None) -> str:
+    def _time(v: tuple[float, float, float] | None) -> str:
         # Annotate each time with its own N: forward and VJP can top out at
         # different sizes (a solver may OOM on the heavier backward pass first).
         if not v:
             return "—"
-        n, t = v
+        n, t, _ = v
         n_str = f" @ N={int(n)}" if math.isfinite(n) else ""
         return f"{t:.3g} s{n_str}"
 
+    def _first(v: tuple[float, float, float]) -> str:
+        return f"{v[2]:.3g} s" if math.isfinite(v[2]) else "—"
+
     rows = sorted(fwd.items(), key=lambda t: t[1][1])
-    header = ["Solver", "Forward time", "VJP time"]
-    return header, [(s, _time(fv), _time(vjp.get(s))) for s, fv in rows]
+    header = ["Solver", "Forward time", "Forward first call", "VJP time"]
+    return header, [(s, _time(fv), _first(fv), _time(vjp.get(s))) for s, fv in rows]
 
 
 # final-objective metric name by domain optimization experiment, in priority

mosaic/benchmarks/core/harness.py

@@ -109,6 +109,8 @@ class TimedResult:
             ``capture_value=True``; otherwise ``None``.
         failure: ``None`` on success; otherwise
             ``{"failure_type", "exc_type", "exc_msg", "traceback"}``.
+        warmup_elapsed: Elapsed seconds for the warmup call, or ``None`` if
+            it raised.
     """
 
     times: list[float]
@@ -117,6 +119,7 @@ class TimedResult:
     first_elapsed: float | None
     last_value: Any | None
     failure: dict | None
+    warmup_elapsed: float | None = None
 
     @property
     def succeeded(self) -> bool:
@@ -126,21 +129,32 @@ def succeeded(self) -> bool:
     def as_record(self, **extra: object) -> dict:
         """Return a dict suitable for inserting into a suite result.
 
-        On success: ``{"mean", "std", "trials_s", **mem, **extra}``.
-        On failure: ``{"status": "failed", "trials_s", **failure, **mem, **extra}``.
+        On success: ``{"mean", "std", "trials_s", "first_call_s", **mem, **extra}``.
+        On failure: ``{"status": "failed", "trials_s", **failure, **mem, **extra}``,
+        plus ``first_call_s`` if the warmup completed.
 
         ``trials_s`` is the raw list of per-trial elapsed seconds (in trial
         order). It's emitted unconditionally so downstream tooling can
         recompute distributional statistics (min, median, percentiles, …)
         without re-running the benchmark. On warmup-time failure the list
         is empty; on mid-trial failure it contains every successful trial
         before the exception.
+
+        ``first_call_s`` is the warmup call's elapsed time. It includes the
+        compilation and setup that the timed trials reuse, so it is reported
+        next to ``mean`` but is not part of it.
         """
         trials = [float(t) for t in self.times]
+        first_call = (
+            {"first_call_s": float(self.warmup_elapsed)}
+            if self.warmup_elapsed is not None
+            else {}
+        )
         if self.failure is not None:
             return {
                 "status": "failed",
                 "trials_s": trials,
+                **first_call,
                 **self.failure,
                 **self.mem,
                 **extra,
@@ -151,6 +165,7 @@ def as_record(self, **extra: object) -> dict:
             if len(self.times) > 1
             else 0.0,
             "trials_s": trials,
+            **first_call,
             **self.mem,
             **extra,
         }
@@ -171,9 +186,10 @@ def run_timed_trials(
     Behaviour:
       * ``setup()``, when given, runs once before the warmup. It is neither
         timed nor returned, but a failure in it is recorded like one in ``fn``.
-      * One unreported warmup ``fn()`` call is performed first to absorb
-        per-solver JIT compilation, first-touch CUDA kernel caching, and
-        scan-unroll tracing.
+      * One warmup ``fn()`` call is performed first to absorb per-solver JIT
+        compilation, first-touch CUDA kernel caching, and scan-unroll tracing.
+        Its elapsed time is kept as ``warmup_elapsed`` and is not one of the
+        trials.
       * ``n_trials`` timed calls follow, each bracketed by ``perf_counter``.
       * If the first timed trial exceeds ``wall_limit_s``, the loop stops
         immediately and ``wall_limit_hit`` is set on the result.
@@ -190,12 +206,15 @@ def run_timed_trials(
     last_value: Any = None
     wall_limit_hit = False
     first_elapsed: float | None = None
+    warmup_elapsed: float | None = None
     failure: dict | None = None
     try:
         with sampler:
             if setup is not None:
                 setup()
+            t0 = time.perf_counter()
             fn()  # warmup; return value intentionally discarded
+            warmup_elapsed = time.perf_counter() - t0
             for i in range(n_trials):
                 t0 = time.perf_counter()
                 v = fn()
@@ -223,4 +242,5 @@ def run_timed_trials(
         first_elapsed=first_elapsed,
         last_value=last_value,
         failure=failure,
+        warmup_elapsed=warmup_elapsed,
     )

mosaic/benchmarks/core/status.py

@@ -306,12 +306,19 @@ def _median(values: list[float]) -> float:
     return vs[mid] if len(vs) % 2 else 0.5 * (vs[mid - 1] + vs[mid])
 
 
-def _refine_cost(data: dict, cells: dict[str, Cell], checks: list) -> None:
+def _refine_cost(
+    data: dict,
+    cells: dict[str, Cell],
+    checks: list,
+    device_class: dict[str, str] | None = None,
+) -> None:
     """Walk ``checks`` against per-solver :class:`CostSummary` instances.
 
-    Solver-medians and peer-median are computed once across OK solvers (need
-    ≥2 to form a meaningful peer comparison), then each check function is
-    given a CostSummary per solver. First anomaly wins.
+    Each solver's peers are the OK solvers of its device class (``device_class``
+    maps solver name to e.g. ``"gpu"`` or ``"cpu"``; without it, all solvers
+    form one class), as the results pages compare cost. A CPU solver being
+    slower than GPU solvers is the hardware, not an anomaly. A class needs ≥2
+    solvers to form a meaningful peer comparison. First anomaly wins.
     """
     from .status_checks import CostSummary
 
@@ -333,16 +340,21 @@ def _refine_cost(data: dict, cells: dict[str, Cell], checks: list) -> None:
         ]
         if times:
             solver_medians[solver] = _median(times)
-    if len(solver_medians) < 2:
-        return
-    peer_median = _median(list(solver_medians.values()))
-    if peer_median <= 0:
-        return
+    classes: dict[str, dict[str, float]] = {}
     for solver, med in solver_medians.items():
-        summary = CostSummary(solver_median_time=med, peer_median_time=peer_median)
-        verdict = _run_checks(checks, summary)
-        if verdict:
-            cells[solver] = Cell(*verdict)
+        cls = (device_class or {}).get(solver, "")
+        classes.setdefault(cls, {})[solver] = med
+    for medians in classes.values():
+        if len(medians) < 2:
+            continue
+        peer_median = _median(list(medians.values()))
+        if peer_median <= 0:
+            continue
+        for solver, med in medians.items():
+            summary = CostSummary(solver_median_time=med, peer_median_time=peer_median)
+            verdict = _run_checks(checks, summary)
+            if verdict:
+                cells[solver] = Cell(*verdict)
 
 
 def _refine_forward(data: dict, cells: dict[str, Cell], checks: list) -> None:
@@ -864,7 +876,12 @@ def _is_analytic_grad_check(suite: str, exp_label: str) -> bool:
 
 
 def _refine_for_suite(
-    suite: str, exp_label: str, data: dict, cells: dict[str, Cell], checks: list
+    suite: str,
+    exp_label: str,
+    data: dict,
+    cells: dict[str, Cell],
+    checks: list,
+    device_class: dict[str, str] | None = None,
 ) -> None:
     """Dispatch suite-specific anomaly refinements (no-op without thresholds)."""
     view = _v1_to_legacy_view(data) if data.get("schema_version") == 1 else data
@@ -874,7 +891,7 @@ def _refine_for_suite(
     if suite == "forward" or _is_analytic_grad_check(suite, exp_label):
         _refine_forward(view, cells, checks)
     elif suite == "cost":
-        _refine_cost(view, cells, checks)
+        _refine_cost(view, cells, checks, device_class)
     elif suite == "gradient" and exp_label.split("/")[0] in (
         "fd_check",
         "source_fd_check",
@@ -1246,7 +1263,8 @@ def _build_row(
     # Capture numeric scalars *before* anomaly refinement downgrades any OK
     # cell — the scalar is worth diffing whether or not it trips a threshold.
     _collect_metrics_for_suite(suite, exp_label, data, row.cells)
-    _refine_for_suite(suite, exp_label, data, row.cells, checks)
+    device_class = {s.name: "gpu" if s.uses_gpu else "cpu" for s in cfg.solvers}
+    _refine_for_suite(suite, exp_label, data, row.cells, checks, device_class)
     _apply_staleness(
         cfg, data, row.cells, solvers, tesseract_hash_cache, harness_hash_cache
     )

mosaic/benchmarks/core/status_checks.py

@@ -173,7 +173,10 @@ def _check(s: ForwardSummary) -> CheckOutcome:
 
 
 def max_peer_k(k: float) -> Callable[[CostSummary], CheckOutcome]:
-    """Cost suite: anomaly if median wall-clock time > ``k × peer median``."""
+    """Cost suite: anomaly if median wall-clock time > ``k × peer median``.
+
+    The peers are the solvers on the same device class (GPU or CPU).
+    """
 
     def _check(s: CostSummary) -> CheckOutcome:
         if s.solver_median_time is None or s.peer_median_time is None:

mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py

@@ -6,8 +6,8 @@
 Two layers:
 
 * Reusable :class:`Exclusion` constants for failure modes that recur across
-  multiple experiments (e.g. OpenFOAM's per-step cost explains both
-  cost/spatial_cost and cost/temporal_cost).
+  multiple experiments (e.g. OpenFOAM's missing VJP excludes it from every
+  gradient and VJP-cost experiment).
 * :func:`register` — applies every ``problem.exclude(...)`` call, so
   ``config.py`` keeps a single ``_register_exclusions(problem)`` line and the
   long reason strings live here instead of cluttering the assembly.
@@ -65,15 +65,6 @@
     "order drops. More correctors remove the floor at a multiple of the wall "
     "time and reverse-mode memory",
 )
-OPENFOAM_SERIAL_IMPLICIT_COST = Exclusion(
-    ExclusionCategory.ANOMALY_EXPLAINED,
-    "OpenFOAM runs icoFoam on one CPU core, while most peers run on the GPU. "
-    "Each PISO step solves the momentum equation once and the pressure "
-    "equation twice with PCG/DIC, the setup of the icoFoam tutorials. The "
-    "pressure solve's iteration count grows in proportion to N, so the cost "
-    "per step grows roughly as N³",
-)
-
 _OBSTACLE_GATE = {
     "jax_cfd": JAX_CFD_NO_OBSTACLE,
     "ins_jl": INS_JL_NO_OBSTACLE,
@@ -87,8 +78,6 @@ def register(problem: Problem) -> None:
     problem.exclude("forward/cylinder", _OBSTACLE_GATE)
 
     # Cost
-    problem.exclude("cost/spatial_cost", {"openfoam": OPENFOAM_SERIAL_IMPLICIT_COST})
-    problem.exclude("cost/temporal_cost", {"openfoam": OPENFOAM_SERIAL_IMPLICIT_COST})
     problem.exclude("cost/vjp_cost", {"openfoam": OPENFOAM_NO_VJP})
 
     # Gradient — suite-level, covers every gradient/* below

mosaic/benchmarks/problems/shared/cost.py

@@ -19,9 +19,12 @@
     for a millisecond presentation.
   * ``trials_s`` carries the raw per-trial elapsed times so distributional
     statistics can be recomputed without rerunning.
-  * Each (solver, sweep-point) pair does one unreported warmup call before
-    the trial loop — absorbs per-solver JIT compilation, first-touch CUDA
-    kernel caching, and scan-unroll tracing.
+  * Each (solver, sweep-point) pair does one warmup call before the trial
+    loop — absorbs per-solver JIT compilation, first-touch CUDA kernel
+    caching, and scan-unroll tracing. Its time is recorded as
+    ``first_call_s`` and is not part of ``mean``. Sweep points share a
+    container, so work a solver does once for all sizes lands in the first
+    point's ``first_call_s`` only.
   * Each entry also records ``vram_peak_mib`` and ``ram_peak_mib`` (peak GPU
     VRAM and container RAM during the trial batch).
   * A per-trial wall-clock limit of ``_SPATIAL_WALL_S`` seconds is enforced.
@@ -210,9 +213,9 @@ def _vjp_cost_aggregate(
 def spatial_cost(t: Any, ctx: KernelContext) -> dict:
     """Forward-pass wall-clock timing at one (solver, N) point.
 
-    Returns ``{"metrics": {mean, std, trials_s, vram_peak_mib, ram_peak_mib},
-    "stop_sweep": bool}``; on failure / wall-limit-hit, the metrics dict
-    carries the failure shape and ``stop_sweep=True``.
+    Returns ``{"metrics": {mean, std, trials_s, first_call_s, vram_peak_mib,
+    ram_peak_mib}, "stop_sweep": bool}``; on failure / wall-limit-hit, the
+    metrics dict carries the failure shape and ``stop_sweep=True``.
     """
     inputs = ctx.make_inputs(ctx.name, ctx.ic, **ctx.phys)
     return _timed_kernel(t, ctx, timed_call=lambda: t.apply(inputs))

mosaic/mosaic_shared/problems/navier_stokes_grid/schemas.py

@@ -126,6 +126,20 @@ class InputSchema(BaseModel):
             "Requires x_lo to be 'dirichlet' with no value."
         ),
     )
+    forcing: GridVectorField | None = Field(
+        default=None,
+        description=(
+            "Body force per unit mass (an acceleration) at the cell centers, "
+            "constant in time, same shape as v0. 2-D periodic boxes only."
+        ),
+    )
+    linear_drag: Array[(1,), Float32] | None = Field(
+        default=None,
+        description=(
+            "Coefficient α of a linear drag -αu on the right-hand side of the "
+            "momentum equation. 2-D periodic boxes only."
+        ),
+    )
 
     @property
     def extent(self) -> tuple[float, ...]:
@@ -179,6 +193,18 @@ def _check_bcs(self) -> "InputSchema":
                     f"inflow_profile has shape {tuple(self.inflow_profile.shape)}, "
                     f"expected ({ny},) to match v0"
                 )
+        if self.forcing is not None or self.linear_drag is not None:
+            if not bc.is_fully_periodic or self.obstacle is not None:
+                raise ValueError(
+                    "forcing and linear_drag require a periodic box without an obstacle"
+                )
+            if self.v0.shape[-1] != 2:
+                raise ValueError("forcing and linear_drag are only defined in 2-D")
+        if self.forcing is not None and self.forcing.shape != self.v0.shape:
+            raise ValueError(
+                f"forcing has shape {tuple(self.forcing.shape)}, expected "
+                f"{tuple(self.v0.shape)} to match v0"
+            )
         return self