From 358e2312241992ce945af286ada5f1972a36501a Mon Sep 17 00:00:00 2001 From: andrinr Date: Fri, 24 Jul 2026 16:45:22 +0200 Subject: [PATCH 01/15] feat(domain): add 2D solver-in-the-loop corrector benchmark --- README.md | 12 +- docs/generate_results.py | 14 +- docs/index.qmd | 4 +- mosaic/benchmarks/core/utils.py | 20 + .../problems/navier_stokes_grid/config.py | 59 ++- .../problems/navier_stokes_grid/corrector.py | 295 ++++++++++++ .../problems/navier_stokes_grid/ics.py | 16 +- .../problems/navier_stokes_grid/plots.py | 118 ++++- .../navier_stokes_grid/solver_in_loop.py | 443 ++++++++++++++++++ .../tesseract_api.py | 92 ++++ tests/test_dummy_integration.py | 41 ++ tests/test_solver_in_loop.py | 182 +++++++ 12 files changed, 1279 insertions(+), 17 deletions(-) create mode 100644 mosaic/benchmarks/problems/navier_stokes_grid/corrector.py create mode 100644 mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py create mode 100644 tests/dummy_tesseracts/navier_stokes_grid_identity/tesseract_api.py create mode 100644 tests/test_solver_in_loop.py diff --git a/README.md b/README.md index b304fd2a..42e488b1 100644 --- a/README.md +++ b/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 diff --git a/docs/generate_results.py b/docs/generate_results.py index 7d792a3f..e9c6ef8b 100644 --- a/docs/generate_results.py +++ b/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", } diff --git a/docs/index.qmd b/docs/index.qmd index 2f17f300..be6914f9 100644 --- a/docs/index.qmd +++ b/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 diff --git a/mosaic/benchmarks/core/utils.py b/mosaic/benchmarks/core/utils.py index 88a66f57..824e0ea0 100644 --- a/mosaic/benchmarks/core/utils.py +++ b/mosaic/benchmarks/core/utils.py @@ -131,6 +131,26 @@ 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 + dataset = run.get("dataset", {}) + for key in ("train_seeds", "test_seeds"): + if key in dataset: + dataset[key] = list(dataset[key])[:1] + 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) cost = run.get("cost", {}) for key in ("N_values", "steps_values"): if key in cost: diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index 2a1e93cf..e5725431 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -10,6 +10,7 @@ re-exported from :mod:`mosaic.benchmarks.problems.shared.diagnostics`. - :mod:`.optimization` β€” drag-minimisation runner. +- :mod:`.solver_in_loop` β€” neural corrector training through each solver. - :mod:`.plots` β€” per-experiment plot fns wired in below. - :mod:`.exclusions` β€” per-(solver, experiment) opt-outs. - :mod:`.extras` β€” cross-experiment aggregator plots. @@ -67,7 +68,8 @@ from .ics import _flat_inflow, _multimode, _tgv, _tgv_analytic, _uniform_flow from .optimization import drag_opt from .physics import DIAGNOSTICS, make_inputs -from .plots import plot_drag_opt +from .plots import plot_drag_opt, plot_solver_in_loop +from .solver_in_loop import solver_in_loop _TESSERACT_SLUG = "navier-stokes-grid" @@ -442,6 +444,61 @@ }, plot=plot_drag_opt, ) +problem.add_experiment( + "optimization/solver_in_loop", + solver_in_loop, + description=( + "Train an identical residual convolutional corrector through each " + "differentiable solver and evaluate its held-out recurrent rollout." + ), + plot_description=( + "Corrector training loss, held-out rollout error, stable horizon, and " + "wall-clock cost for each differentiable solver." + ), + # Keep dataset seed lists inside an explicit run payload: they are data, + # not benchmark sweep coordinates. + runs=[ + { + "ic": {"name": "multimode", "seed": 0}, + "physics": { + # ``steps`` is the number of native PDE steps between neural + # corrections. Recurrent training feeds the corrected + # canonical velocity back through the Tesseract as ``v0``. + "N": 32, + "nu": 0.001, + "dt": 0.02, + "steps": 4, + }, + "dataset": { + "reference_factor": 2, + "reference_substeps": 2, + "train_seeds": [0, 1, 2, 3], + "test_seeds": [100, 101], + "train_frames": 16, + "k0": 6.0, + "sigma_k": 1.0, + "amplitude": 0.3, + }, + "training": { + "max_updates": 100, + "unroll": 4, + "lr": 1e-4, + "clip_norm": 1.0, + "hidden_channels": 32, + "kernel_size": 5, + "seed": 2026, + "model_seed": 0, + "check_grad": True, + "fd_epsilon": 1e-2, + }, + "evaluation": { + "rollout_frames": 24, + "stable_error_threshold": 1.0, + }, + } + ], + plot=plot_solver_in_loop, +) # Bonus plot (not paired with an experiment). problem.add_extra_plot( "_extra/gradient/jacobian_svd_comparison", diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py new file mode 100644 index 00000000..88def793 --- /dev/null +++ b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py @@ -0,0 +1,295 @@ +# Copyright 2026 Pasteur Labs. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark-owned neural corrector and periodic reference dynamics. + +The corrector deliberately lives in the benchmark harness instead of any +solver Tesseract. Every candidate solver therefore sees the same model, +initial weights, projection, optimiser, and reference trajectories; only the +``v0 -> result`` transition and its VJP vary between benchmark cells. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np + + +def init_corrector( + key: jax.Array, + *, + hidden_channels: int = 32, + kernel_size: int = 5, +) -> tuple[dict[str, jax.Array], ...]: + """Initialise a small three-layer periodic residual CNN. + + The final layer starts at a much smaller scale than the feature layers, so + the initial corrected rollout remains close to the underlying solver while + gradients can still reach every layer on the first update. + """ + if kernel_size % 2 != 1: + raise ValueError("kernel_size must be odd") + channels = (2, hidden_channels, hidden_channels * 2, 2) + keys = jax.random.split(key, len(channels) - 1) + layers: list[dict[str, jax.Array]] = [] + for idx, (cin, cout, layer_key) in enumerate( + zip(channels[:-1], channels[1:], keys, strict=True) + ): + fan_in = kernel_size * kernel_size * cin + scale = np.sqrt(2.0 / fan_in) + if idx == len(channels) - 2: + scale *= 1e-2 + kernel = ( + jax.random.normal( + layer_key, + (kernel_size, kernel_size, cin, cout), + dtype=jnp.float32, + ) + * scale + ) + layers.append( + { + "kernel": kernel, + "bias": jnp.zeros((cout,), dtype=jnp.float32), + } + ) + return tuple(layers) + + +def _periodic_conv(x: jax.Array, layer: dict[str, jax.Array]) -> jax.Array: + """Apply one stride-one NHWC convolution with periodic padding.""" + kernel = layer["kernel"] + pad = kernel.shape[0] // 2 + padded = jnp.pad(x, ((pad, pad), (pad, pad), (0, 0)), mode="wrap") + y = jax.lax.conv_general_dilated( + padded[jnp.newaxis, ...], + kernel, + window_strides=(1, 1), + padding="VALID", + dimension_numbers=("NHWC", "HWIO", "NHWC"), + )[0] + return y + layer["bias"] + + +def apply_corrector( + params: Sequence[dict[str, jax.Array]], + velocity: jax.Array, + *, + velocity_scale: float | jax.Array, +) -> jax.Array: + """Predict an additive velocity correction on ``(N, N, 1, 2)`` fields.""" + x = velocity[:, :, 0, :] / jnp.asarray(velocity_scale, dtype=velocity.dtype) + for layer in params[:-1]: + x = jax.nn.gelu(_periodic_conv(x, layer)) + delta = _periodic_conv(x, params[-1]) + return delta[:, :, jnp.newaxis, :] * jnp.asarray( + velocity_scale, dtype=velocity.dtype + ) + + +def project_periodic_correction(delta: jax.Array, domain_extent: float) -> jax.Array: + """Project a velocity correction onto zero-mean divergence-free fields.""" + field = delta[:, :, 0, :] + nx, ny = field.shape[:2] + dx, dy = domain_extent / nx, domain_extent / ny + kx = 2.0 * jnp.pi * jnp.fft.fftfreq(nx, d=dx) + ky = 2.0 * jnp.pi * jnp.fft.fftfreq(ny, d=dy) + kx_grid, ky_grid = jnp.meshgrid(kx, ky, indexing="ij") + k2 = kx_grid**2 + ky_grid**2 + + u_hat = jnp.fft.fft2(field[..., 0]) + v_hat = jnp.fft.fft2(field[..., 1]) + k_dot_u = kx_grid * u_hat + ky_grid * v_hat + safe_k2 = jnp.where(k2 == 0, 1.0, k2) + u_hat = u_hat - kx_grid * k_dot_u / safe_k2 + v_hat = v_hat - ky_grid * k_dot_u / safe_k2 + # The signed Nyquist wave number has no conjugate partner on an even + # real-valued grid. Removing those two lines keeps the projected spectrum + # Hermitian, rather than silently discarding an imaginary component below. + if nx % 2 == 0: + u_hat = u_hat.at[nx // 2, :].set(0.0) + v_hat = v_hat.at[nx // 2, :].set(0.0) + if ny % 2 == 0: + u_hat = u_hat.at[:, ny // 2].set(0.0) + v_hat = v_hat.at[:, ny // 2].set(0.0) + # A corrector should not change the conserved mean flow. + u_hat = u_hat.at[0, 0].set(0.0) + v_hat = v_hat.at[0, 0].set(0.0) + + projected = jnp.stack( + [jnp.fft.ifft2(u_hat).real, jnp.fft.ifft2(v_hat).real], + axis=-1, + ) + return projected[:, :, jnp.newaxis, :].astype(delta.dtype) + + +def corrected_velocity( + params: Sequence[dict[str, jax.Array]], + provisional: jax.Array, + *, + velocity_scale: float | jax.Array, + domain_extent: float, +) -> jax.Array: + """Apply the learned residual and enforce a physical correction subspace.""" + delta = apply_corrector(params, provisional, velocity_scale=velocity_scale) + return provisional + project_periodic_correction(delta, domain_extent) + + +def relative_l2(predicted: Any, target: Any) -> float: + """Relative L2 error with a stable zero-target denominator.""" + p = np.asarray(predicted, dtype=np.float64) + q = np.asarray(target, dtype=np.float64) + return float(np.linalg.norm(p - q) / (np.linalg.norm(q) + 1e-12)) + + +def divergence_rms(velocity: Any, domain_extent: float) -> float: + """Spectral RMS divergence for one canonical 2D velocity field.""" + field = np.asarray(velocity) + if field.ndim == 4: + field = field[:, :, 0, :] + n = field.shape[0] + dx = domain_extent / n + wave = 2.0 * np.pi * np.fft.fftfreq(n, d=dx) + kx, ky = np.meshgrid(wave, wave, indexing="ij") + div_hat = 1j * (kx * np.fft.fft2(field[..., 0]) + ky * np.fft.fft2(field[..., 1])) + return float(np.sqrt(np.mean(np.fft.ifft2(div_hat).real ** 2))) + + +def spectral_restrict(velocity: Any, target_n: int) -> np.ndarray: + """Fourier-restrict a periodic velocity field to ``target_n`` cells.""" + field = np.asarray(velocity) + had_z_axis = field.ndim == 4 + if had_z_axis: + field = field[:, :, 0, :] + source_n = field.shape[0] + if source_n == target_n: + result = field.astype(np.float32, copy=True) + else: + if source_n < target_n or source_n % target_n != 0: + raise ValueError("source grid must be an integer multiple of target_n") + spectrum = np.fft.fftshift(np.fft.fft2(field, axes=(0, 1)), axes=(0, 1)) + start = (source_n - target_n) // 2 + stop = start + target_n + cropped = spectrum[start:stop, start:stop, :] + result = np.fft.ifft2(np.fft.ifftshift(cropped, axes=(0, 1)), axes=(0, 1)).real + result *= (target_n / source_n) ** 2 + result = result.astype(np.float32) + return result[:, :, np.newaxis, :] if had_z_axis else result + + +def _velocity_to_vorticity_hat( + velocity: np.ndarray, domain_extent: float +) -> np.ndarray: + """Return the spectral scalar vorticity of a 2D velocity field.""" + field = velocity[:, :, 0, :] if velocity.ndim == 4 else velocity + n = field.shape[0] + dx = domain_extent / n + wave = 2.0 * np.pi * np.fft.fftfreq(n, d=dx) + kx, ky = np.meshgrid(wave, wave, indexing="ij") + u_hat = np.fft.fft2(field[..., 0]) + v_hat = np.fft.fft2(field[..., 1]) + return 1j * kx * v_hat - 1j * ky * u_hat + + +def _vorticity_hat_to_velocity( + omega_hat: np.ndarray, domain_extent: float +) -> np.ndarray: + """Recover canonical velocity from spectral vorticity.""" + n = omega_hat.shape[0] + dx = domain_extent / n + wave = 2.0 * np.pi * np.fft.fftfreq(n, d=dx) + kx, ky = np.meshgrid(wave, wave, indexing="ij") + k2 = kx**2 + ky**2 + safe_k2 = np.where(k2 == 0, 1.0, k2) + psi_hat = omega_hat / safe_k2 + psi_hat[0, 0] = 0.0 + u = np.fft.ifft2(1j * ky * psi_hat).real + v = np.fft.ifft2(-1j * kx * psi_hat).real + return np.stack([u, v], axis=-1)[:, :, np.newaxis, :].astype(np.float32) + + +def _vorticity_rhs( + omega_hat: np.ndarray, + *, + viscosity: float, + domain_extent: float, +) -> np.ndarray: + """Dealiased pseudo-spectral RHS for 2D vorticity dynamics.""" + n = omega_hat.shape[0] + dx = domain_extent / n + wave = 2.0 * np.pi * np.fft.fftfreq(n, d=dx) + mode = np.fft.fftfreq(n) * n + kx, ky = np.meshgrid(wave, wave, indexing="ij") + mx, my = np.meshgrid(mode, mode, indexing="ij") + k2 = kx**2 + ky**2 + safe_k2 = np.where(k2 == 0, 1.0, k2) + psi_hat = omega_hat / safe_k2 + psi_hat[0, 0] = 0.0 + + u = np.fft.ifft2(1j * ky * psi_hat).real + v = np.fft.ifft2(-1j * kx * psi_hat).real + grad_x = np.fft.ifft2(1j * kx * omega_hat).real + grad_y = np.fft.ifft2(1j * ky * omega_hat).real + advection_hat = np.fft.fft2(u * grad_x + v * grad_y) + cutoff = n // 3 + dealias = (np.abs(mx) <= cutoff) & (np.abs(my) <= cutoff) + advection_hat *= dealias + return -advection_hat - viscosity * k2 * omega_hat + + +def reference_trajectory( + initial_velocity: Any, + *, + viscosity: float, + dt: float, + frame_steps: int, + n_frames: int, + substeps: int, + domain_extent: float, +) -> np.ndarray: + """Integrate a high-resolution periodic reference with spectral RK4.""" + velocity = np.asarray(initial_velocity, dtype=np.float64) + omega_hat = _velocity_to_vorticity_hat(velocity, domain_extent) + step_dt = dt / substeps + trajectory = [velocity.astype(np.float32)] + for _ in range(n_frames): + for _ in range(frame_steps * substeps): + k1 = _vorticity_rhs( + omega_hat, + viscosity=viscosity, + domain_extent=domain_extent, + ) + k2 = _vorticity_rhs( + omega_hat + 0.5 * step_dt * k1, + viscosity=viscosity, + domain_extent=domain_extent, + ) + k3 = _vorticity_rhs( + omega_hat + 0.5 * step_dt * k2, + viscosity=viscosity, + domain_extent=domain_extent, + ) + k4 = _vorticity_rhs( + omega_hat + step_dt * k3, + viscosity=viscosity, + domain_extent=domain_extent, + ) + omega_hat = omega_hat + step_dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6 + trajectory.append(_vorticity_hat_to_velocity(omega_hat, domain_extent)) + return np.stack(trajectory) + + +__all__ = [ + "apply_corrector", + "corrected_velocity", + "divergence_rms", + "init_corrector", + "project_periodic_correction", + "reference_trajectory", + "relative_l2", + "spectral_restrict", +] diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/ics.py b/mosaic/benchmarks/problems/navier_stokes_grid/ics.py index 4a2fb473..b7c10450 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/ics.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/ics.py @@ -11,13 +11,21 @@ import jax.numpy as jnp -def _multimode(N: int, L: float = 2 * jnp.pi, seed: int = 42, **_: Any) -> jax.Array: - """Energy ring at k=2, Οƒ=0.5, max speed normalised to 0.3.""" +def _multimode( + N: int, + L: float = 2 * jnp.pi, + seed: int = 42, + k0: float = 2.0, + sigma_k: float = 0.5, + amplitude: float = 0.3, + **_: Any, +) -> jax.Array: + """Random divergence-free field with energy concentrated in a Fourier ring.""" key = jax.random.PRNGKey(seed) kn = jnp.fft.fftfreq(N, d=1.0 / N) KX, KY = jnp.meshgrid(kn, kn, indexing="ij") K_abs = jnp.sqrt(KX**2 + KY**2) - envelope = jnp.exp(-0.5 * ((K_abs - 2.0) / 0.5) ** 2) + envelope = jnp.exp(-0.5 * ((K_abs - k0) / sigma_k) ** 2) phases = jax.random.uniform(key, (N, N), minval=0.0, maxval=2.0 * jnp.pi) psi_hat = envelope * jnp.exp(1j * phases) psi_hat = 0.5 * (psi_hat + jnp.conj(psi_hat[::-1, ::-1])) @@ -26,7 +34,7 @@ def _multimode(N: int, L: float = 2 * jnp.pi, seed: int = 42, **_: Any) -> jax.A vx = jnp.real(jnp.fft.ifft2(1j * KY * kfac * psi_hat)) vy = jnp.real(jnp.fft.ifft2(-1j * KX * kfac * psi_hat)) u_max = float(jnp.sqrt(vx**2 + vy**2).max()) or 1.0 - vx, vy = vx / u_max * 0.3, vy / u_max * 0.3 + vx, vy = vx / u_max * amplitude, vy / u_max * amplitude return jnp.stack([vx, vy], axis=-1)[:, :, None, :].astype(jnp.float32) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 6da936e8..b79ccd45 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -1,7 +1,7 @@ # Copyright 2026 Pasteur Labs. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 -"""Per-problem plots for the navier-stokes-grid drag-optimisation experiments.""" +"""Per-problem plots for 2D Navier--Stokes optimisation experiments.""" from __future__ import annotations @@ -44,6 +44,122 @@ _DRAG_OPT_SOLVER_ORDER = ["xlb", "phiflow", "pict"] +def plot_solver_in_loop( + cfg: Problem, + *, + save: bool = True, + suffix: str = "", + **_kw: Any, +) -> list: + """Plot corrector trainability, rollout quality, and time-to-quality.""" + out_dir = experiment_dir( + results_dir(), + cfg.name, + "optimization", + f"solver_in_loop{suffix}", + ) + result_path = out_dir / "result.json" + fields_path = out_dir / "corrector_fields.npz" + if not result_path.exists() or not fields_path.exists(): + print(f"[solver_in_loop] missing results in {out_dir} β€” skipping") + return [] + + data = v1_to_legacy(load_json(result_path)) + arrays = try_load_npz(fields_path) + names = [str(v) for v in arrays.get("solver_names", np.array([])).tolist()] + if not names: + return [] + + times = np.asarray(arrays.get("evaluation_times", np.array([]))) + fig, axes = paper_row(3, squeeze=False) + ax_loss, ax_roll, ax_cost = np.atleast_1d(axes).ravel() + present: list[str] = [] + + for idx, name in enumerate(names): + metrics = data.get("by_solver", {}).get(name, {}) + label, color, linestyle, marker = solver_props(name) + loss = np.asarray(arrays.get(f"loss_{idx}", np.array([]))) + corrected = np.asarray(arrays.get(f"error_corrected_{idx}", np.array([]))) + uncorrected = np.asarray(arrays.get(f"error_uncorrected_{idx}", np.array([]))) + if loss.size: + ax_loss.plot( + np.arange(1, loss.size + 1), + loss, + color=color, + linestyle=linestyle, + label=label, + ) + if corrected.size: + x = times[: corrected.size] if times.size else np.arange(corrected.size) + ax_roll.plot( + x[1:], + corrected[1:], + color=color, + linestyle=linestyle, + label=label, + ) + if uncorrected.size: + x = times[: uncorrected.size] if times.size else np.arange(uncorrected.size) + ax_roll.plot( + x[1:], + uncorrected[1:], + color=color, + linestyle=":", + alpha=0.45, + ) + + wall = metrics.get("training_wall_time_s") + error = metrics.get("final_rollout_error") + if wall is not None and error is not None: + ax_cost.scatter( + [wall], + [error], + color=color, + marker=marker, + s=30, + label=label, + ) + present.append(name) + + if any( + np.any(np.asarray(arrays.get(f"loss_{idx}", np.array([]))) > 0) + for idx in range(len(names)) + ): + ax_loss.set_yscale("log") + ax_loss.set_xlabel("Optimizer update") + ax_loss.set_ylabel("Training loss") + ax_loss.set_title("Corrector training") + + ax_roll.set_yscale("log") + ax_roll.set_xlabel("Physical time") + ax_roll.set_ylabel("Relative $L^2$ error") + ax_roll.set_title("Held-out rollout") + ax_roll.text( + 0.98, + 0.96, + "solid: corrected\n dotted: uncorrected", + transform=ax_roll.transAxes, + ha="right", + va="top", + fontsize=6.5, + ) + + ax_cost.set_xscale("log") + ax_cost.set_yscale("log") + ax_cost.set_xlabel("Training wall time [s]") + ax_cost.set_ylabel("Final rollout error") + ax_cost.set_title("Time to quality") + + aliases = [ + alias for name in present if (alias := resolve_solver_alias(name)) is not None + ] + solver_legend(fig, aliases, order=NS_ORDER) + fig.suptitle("Solver-in-the-loop neural correction") + if save: + save_fig(fig, "solver_in_loop", out_dir) + return [fig] + + def plot_drag_opt( cfg: Problem, *, diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py new file mode 100644 index 00000000..92415f31 --- /dev/null +++ b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py @@ -0,0 +1,443 @@ +# Copyright 2026 Pasteur Labs. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Solver-in-the-loop training benchmark for periodic 2D Navier--Stokes. + +Each candidate Tesseract supplies the recurrent coarse-grid transition and +its VJP. The data, neural corrector, optimiser, rollout loss, and evaluation +are benchmark-owned so differences in training outcomes can be attributed to +the solver interface under test. +""" + +from __future__ import annotations + +import hashlib +import json +import threading +import time +from functools import lru_cache, partial +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +import optax +from tesseract_jax import apply_tesseract + +from mosaic.benchmarks.core.experiment import KernelContext, kernel +from mosaic.benchmarks.core.utils import active_differentiable_solvers + +from .corrector import ( + corrected_velocity, + divergence_rms, + init_corrector, + reference_trajectory, + relative_l2, + spectral_restrict, +) +from .ics import _multimode + +_DATASET_LOCK = threading.Lock() + + +def _dataset_digest(config: dict[str, Any]) -> str: + """Stable short identifier for generated reference data.""" + payload = json.dumps(config, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest()[:16] + + +@lru_cache(maxsize=8) +def _cached_reference_dataset( + *, + n: int, + reference_factor: int, + viscosity: float, + dt: float, + frame_steps: int, + n_frames: int, + reference_substeps: int, + domain_extent: float, + seeds: tuple[int, ...], + k0: float, + sigma_k: float, + amplitude: float, +) -> np.ndarray: + """Generate and cache coarse views of deterministic fine-grid trajectories.""" + trajectories: list[np.ndarray] = [] + reference_n = n * reference_factor + for seed in seeds: + fine_initial = _multimode( + reference_n, + L=domain_extent, + seed=seed, + k0=k0, + sigma_k=sigma_k, + amplitude=amplitude, + ) + fine = reference_trajectory( + fine_initial, + viscosity=viscosity, + dt=dt, + frame_steps=frame_steps, + n_frames=n_frames, + substeps=reference_substeps, + domain_extent=domain_extent, + ) + trajectories.append(np.stack([spectral_restrict(frame, n) for frame in fine])) + return np.stack(trajectories).astype(np.float32) + + +def make_reference_dataset( + *, + physics: dict[str, Any], + dataset: dict[str, Any], + evaluation: dict[str, Any], + training: dict[str, Any], + domain_extent: float, +) -> tuple[np.ndarray, np.ndarray, str]: + """Build the train/test trajectory tensors used by every solver.""" + n = int(physics["N"]) + train_seeds = tuple(int(v) for v in dataset.get("train_seeds", [0, 1, 2, 3])) + test_seeds = tuple(int(v) for v in dataset.get("test_seeds", [100, 101])) + train_frames = int(dataset.get("train_frames", 16)) + eval_frames = int(evaluation.get("rollout_frames", 32)) + unroll = int(training.get("unroll", 4)) + n_frames = max(train_frames, eval_frames, unroll) + config = { + "n": n, + "reference_factor": int(dataset.get("reference_factor", 2)), + "viscosity": float(physics["nu"]), + "dt": float(physics["dt"]), + "frame_steps": int(physics["steps"]), + "n_frames": n_frames, + "reference_substeps": int(dataset.get("reference_substeps", 2)), + "domain_extent": float(domain_extent), + "seeds": train_seeds + test_seeds, + "k0": float(dataset.get("k0", 6.0)), + "sigma_k": float(dataset.get("sigma_k", 1.0)), + "amplitude": float(dataset.get("amplitude", 0.3)), + } + # functools.lru_cache is thread-safe for its dictionary, but concurrent + # cache misses may still execute the wrapped function more than once. + # Reference generation is substantial enough to serialize that first miss. + with _DATASET_LOCK: + all_trajectories = _cached_reference_dataset(**config) + n_train = len(train_seeds) + train = all_trajectories[:n_train, : train_frames + 1] + test = all_trajectories[n_train:, : eval_frames + 1] + return train, test, _dataset_digest(config) + + +def _solver_advance( + t: Any, + ctx: KernelContext, + velocity: jax.Array, + *, + frame_steps: int, +) -> jax.Array: + """Advance one canonical correction interval through a Tesseract.""" + physics = {**ctx.phys, "steps": frame_steps} + inputs = ctx.make_inputs(ctx.name, velocity, **physics) + outputs = apply_tesseract(t, inputs) + if ctx.output_key not in outputs: + raise RuntimeError( + f"Solver '{ctx.name}' did not return output {ctx.output_key!r}" + ) + return outputs[ctx.output_key] + + +def _window_loss( + params: Any, + targets: jax.Array, + *, + t: Any, + ctx: KernelContext, + frame_steps: int, + velocity_scale: float, +) -> jax.Array: + """Mean normalized state error over one recurrent training window.""" + state = targets[0] + losses: list[jax.Array] = [] + for target in targets[1:]: + provisional = _solver_advance(t, ctx, state, frame_steps=frame_steps) + state = corrected_velocity( + params, + provisional, + velocity_scale=velocity_scale, + domain_extent=ctx.domain_extent, + ) + losses.append(jnp.sum((state - target) ** 2) / (jnp.sum(target**2) + 1e-12)) + return jnp.mean(jnp.stack(losses)) + + +def _directional_fd( + loss_fn: Any, + params: Any, + grads: Any, + key: jax.Array, + *, + epsilon: float, +) -> float: + """Relative error of one end-to-end directional finite difference.""" + direction = jax.tree_util.tree_map( + lambda p, k: jax.random.normal(k, p.shape, dtype=p.dtype), + params, + jax.tree_util.tree_unflatten( + jax.tree_util.tree_structure(params), + jax.random.split(key, len(jax.tree_util.tree_leaves(params))), + ), + ) + direction_norm = optax.tree.norm(direction) + direction = jax.tree_util.tree_map( + lambda value: value / (direction_norm + 1e-12), direction + ) + plus = jax.tree_util.tree_map(lambda p, d: p + epsilon * d, params, direction) + minus = jax.tree_util.tree_map(lambda p, d: p - epsilon * d, params, direction) + fd = (loss_fn(plus) - loss_fn(minus)) / (2.0 * epsilon) + ad = sum( + jnp.vdot(g, d) + for g, d in zip( + jax.tree_util.tree_leaves(grads), + jax.tree_util.tree_leaves(direction), + strict=True, + ) + ) + return float(jnp.abs(fd - ad) / (jnp.abs(fd) + jnp.abs(ad) + 1e-12)) + + +def _train_corrector( + t: Any, + ctx: KernelContext, + train: np.ndarray, + *, + frame_steps: int, + training: dict[str, Any], + velocity_scale: float, +) -> tuple[Any, list[float], list[float], float | None, bool]: + """Train one solver-specific corrector with a fixed stochastic schedule.""" + max_updates = int(training.get("max_updates", 100)) + unroll = int(training.get("unroll", 4)) + lr = float(training.get("lr", 1e-4)) + clip_norm = float(training.get("clip_norm", 1.0)) + seed = int(training.get("seed", 2026)) + model_seed = int(training.get("model_seed", 0)) + hidden_channels = int(training.get("hidden_channels", 32)) + kernel_size = int(training.get("kernel_size", 5)) + fd_epsilon = float(training.get("fd_epsilon", 1e-2)) + + params = init_corrector( + jax.random.PRNGKey(model_seed), + hidden_channels=hidden_channels, + kernel_size=kernel_size, + ) + optimiser = optax.chain( + optax.clip_by_global_norm(clip_norm), + optax.adam(lr), + ) + opt_state = optimiser.init(params) + rng = np.random.RandomState(seed) + losses: list[float] = [] + grad_norms: list[float] = [] + fd_error: float | None = None + completed = True + + for update in range(max_updates): + trajectory_idx = int(rng.randint(train.shape[0])) + max_start = train.shape[1] - unroll - 1 + start = int(rng.randint(max_start + 1)) if max_start > 0 else 0 + targets = jnp.asarray(train[trajectory_idx, start : start + unroll + 1]) + + loss_fn = partial( + _window_loss, + targets=targets, + t=t, + ctx=ctx, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + ) + + loss, grads = jax.value_and_grad(loss_fn)(params) + loss_value = float(loss) + grad_norm = float(optax.tree.norm(grads)) + if not np.isfinite(loss_value) or not np.isfinite(grad_norm): + completed = False + break + if update == 0 and bool(training.get("check_grad", True)): + fd_error = _directional_fd( + loss_fn, + params, + grads, + jax.random.PRNGKey(model_seed + 1), + epsilon=fd_epsilon, + ) + updates, opt_state = optimiser.update(grads, opt_state, params) + params = optax.apply_updates(params, updates) + losses.append(loss_value) + grad_norms.append(grad_norm) + return params, losses, grad_norms, fd_error, completed + + +def _evaluate_rollout( + t: Any, + ctx: KernelContext, + params: Any, + reference: np.ndarray, + *, + frame_steps: int, + velocity_scale: float, + corrected: bool, +) -> tuple[np.ndarray, list[float]]: + """Roll out one held-out sequence with or without neural corrections.""" + state = jnp.asarray(reference[0]) + states = [np.asarray(state)] + errors = [0.0] + for target in reference[1:]: + state = _solver_advance(t, ctx, state, frame_steps=frame_steps) + if corrected: + state = corrected_velocity( + params, + state, + velocity_scale=velocity_scale, + domain_extent=ctx.domain_extent, + ) + state_np = np.asarray(state) + states.append(state_np) + errors.append(relative_l2(state_np, target)) + return np.stack(states), errors + + +def _first_unstable(errors: list[float], threshold: float) -> int: + """Return first bad frame, or the full completed horizon.""" + for idx, error in enumerate(errors): + if not np.isfinite(error) or error > threshold: + return idx + return len(errors) - 1 + + +@kernel( + sweep_mode="none", + selector_fn=active_differentiable_solvers, + snapshot_filename="corrector_fields.npz", + snapshot_prefixes=( + "loss", + "grad_norm", + "error_corrected", + "error_uncorrected", + "rollout_corrected", + "rollout_uncorrected", + ), +) +def solver_in_loop(t: Any, ctx: KernelContext) -> dict: + """Train and evaluate one neural corrector through one candidate solver.""" + training = ctx.run.get("training", {}) + dataset_cfg = ctx.run.get("dataset", {}) + evaluation = ctx.run.get("evaluation", {}) + frame_steps = int(ctx.phys["steps"]) + + train, test, dataset_hash = make_reference_dataset( + physics=ctx.phys, + dataset=dataset_cfg, + evaluation=evaluation, + training=training, + domain_extent=ctx.domain_extent, + ) + velocity_scale = float(np.sqrt(np.mean(train**2)) + 1e-8) + + started = time.perf_counter() + params, losses, grad_norms, fd_error, completed = _train_corrector( + t, + ctx, + train, + frame_steps=frame_steps, + training=training, + velocity_scale=velocity_scale, + ) + training_wall = time.perf_counter() - started + + corrected_errors: list[list[float]] = [] + uncorrected_errors: list[list[float]] = [] + first_corrected: np.ndarray | None = None + first_uncorrected: np.ndarray | None = None + for reference in test: + corrected_rollout, corr_err = _evaluate_rollout( + t, + ctx, + params, + reference, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + uncorrected_rollout, raw_err = _evaluate_rollout( + t, + ctx, + params, + reference, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=False, + ) + corrected_errors.append(corr_err) + uncorrected_errors.append(raw_err) + if first_corrected is None: + first_corrected = corrected_rollout + first_uncorrected = uncorrected_rollout + + corrected_error = np.mean(np.asarray(corrected_errors), axis=0) + uncorrected_error = np.mean(np.asarray(uncorrected_errors), axis=0) + final_corrected = float(corrected_error[-1]) + final_uncorrected = float(uncorrected_error[-1]) + mean_corrected = float(np.mean(corrected_error[1:])) + mean_uncorrected = float(np.mean(uncorrected_error[1:])) + threshold = float(evaluation.get("stable_error_threshold", 1.0)) + interval_time = float(ctx.phys["dt"]) * frame_steps + + metrics = { + "final_rollout_error": final_corrected, + "uncorrected_rollout_error": final_uncorrected, + "mean_rollout_error": mean_corrected, + "uncorrected_mean_rollout_error": mean_uncorrected, + "improvement_pct": 100.0 + * (final_uncorrected - final_corrected) + / (final_uncorrected + 1e-12), + "stable_horizon": _first_unstable(corrected_error.tolist(), threshold) + * interval_time, + "uncorrected_stable_horizon": _first_unstable( + uncorrected_error.tolist(), threshold + ) + * interval_time, + "final_train_loss": losses[-1] if losses else None, + "best_train_loss": min(losses) if losses else None, + "n_updates": len(losses), + "training_wall_time_s": training_wall, + "seconds_per_update": training_wall / max(len(losses), 1), + "final_grad_norm": grad_norms[-1] if grad_norms else None, + "end_to_end_fd_rel_error": fd_error, + "final_divergence_rms": divergence_rms(first_corrected[-1], ctx.domain_extent) + if first_corrected is not None + else None, + "completed": completed, + "dataset_hash": dataset_hash, + "state_restart": True, + } + snapshots = { + "loss": np.asarray(losses, dtype=np.float32), + "grad_norm": np.asarray(grad_norms, dtype=np.float32), + "error_corrected": np.asarray(corrected_error, dtype=np.float32), + "error_uncorrected": np.asarray(uncorrected_error, dtype=np.float32), + } + if first_corrected is not None and first_uncorrected is not None: + snapshots["rollout_corrected"] = first_corrected + snapshots["rollout_uncorrected"] = first_uncorrected + return { + "metrics": metrics, + "snapshots": snapshots, + "shared": { + "reference_rollout": test[0], + "evaluation_times": np.arange(test.shape[1], dtype=np.float32) + * interval_time, + }, + } + + +__all__ = ["make_reference_dataset", "solver_in_loop"] diff --git a/tests/dummy_tesseracts/navier_stokes_grid_identity/tesseract_api.py b/tests/dummy_tesseracts/navier_stokes_grid_identity/tesseract_api.py new file mode 100644 index 00000000..3770e5ad --- /dev/null +++ b/tests/dummy_tesseracts/navier_stokes_grid_identity/tesseract_api.py @@ -0,0 +1,92 @@ +# Copyright 2026 Pasteur Labs. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: E402 β€” sys.path bootstrap must precede workspace imports +"""Identity Navier--Stokes Tesseract for recurrent corrector smoke tests. + +Unlike the corpus-wide constant-output dummy, this fixture exposes a nonzero, +exact VJP so a solver-in-the-loop test can verify gradient flow through more +than one recurrent ``apply_tesseract`` call. +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from typing import Any + +_TESSERACTS_DIR = Path(__file__).resolve().parents[3] / "mosaic" / "tesseracts" +if str(_TESSERACTS_DIR) not in sys.path: + sys.path.insert(0, str(_TESSERACTS_DIR)) + +import numpy as np +from mosaic_shared.problems.navier_stokes_grid import ( + InputSchema as _CanonicalInputSchema, +) +from mosaic_shared.problems.navier_stokes_grid import ( + OutputSchema as _CanonicalOutputSchema, +) +from mosaic_shared.schema_types import make_differentiable +from pydantic import ConfigDict +from tesseract_core.runtime import ShapeDType + + +class InputSchema( + make_differentiable( + _CanonicalInputSchema, ["v0", "viscosity", "dt", "inflow_profile"] + ) +): + model_config = ConfigDict(extra="ignore") + + +class OutputSchema(make_differentiable(_CanonicalOutputSchema, ["result", "drag"])): + pass + + +def apply(inputs: InputSchema) -> OutputSchema: + """Return the input velocity unchanged.""" + return OutputSchema( + result=np.asarray(inputs.v0, dtype=np.float32), + drag=np.asarray([0.0], dtype=np.float32), + ) + + +def vector_jacobian_product( + inputs: InputSchema, + vjp_inputs: set[str], + vjp_outputs: set[str], + cotangent_vector: dict[str, Any], +) -> dict[str, np.ndarray]: + """Return the exact identity VJP for ``result = v0``.""" + del vjp_outputs + out: dict[str, np.ndarray] = {} + if "v0" in vjp_inputs: + out["v0"] = np.asarray( + cotangent_vector.get("result", np.zeros_like(inputs.v0)), + dtype=np.float32, + ) + if "viscosity" in vjp_inputs: + out["viscosity"] = np.zeros((1,), dtype=np.float32) + if "dt" in vjp_inputs: + out["dt"] = np.zeros((1,), dtype=np.float32) + if "inflow_profile" in vjp_inputs and inputs.inflow_profile is not None: + out["inflow_profile"] = np.zeros_like( + np.asarray(inputs.inflow_profile), dtype=np.float32 + ) + return out + + +def abstract_eval(abstract_inputs: InputSchema) -> dict[str, ShapeDType]: + """Declare that ``result`` has the same shape and dtype as ``v0``.""" + v0 = abstract_inputs.model_dump()["v0"] + if isinstance(v0, dict) and "shape" in v0: + shape = tuple(v0["shape"]) + dtype = v0.get("dtype", "float32") + else: + array = np.asarray(v0) + shape = array.shape + dtype = str(array.dtype) + return { + "result": ShapeDType(shape=shape, dtype=dtype), + "drag": ShapeDType(shape=(1,), dtype="float32"), + } diff --git a/tests/test_dummy_integration.py b/tests/test_dummy_integration.py index cb6f0eda..1c1cc77f 100644 --- a/tests/test_dummy_integration.py +++ b/tests/test_dummy_integration.py @@ -279,6 +279,47 @@ def _maybe_shrink(cfg, problem: str, exp_key: str) -> None: fresh per-call ``get_config(problem)``, so this never leaks across tests). No-op for experiments not in :data:`_SHRINK_PHYSICS`. """ + if (problem, exp_key) == ("ns-grid", "optimization/solver_in_loop"): + from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( + solver_in_loop, + ) + + cfg.add_experiment( + exp_key, + solver_in_loop, + runs=[ + { + "ic": {"name": "multimode", "seed": 0}, + "physics": { + "N": 8, + "nu": 0.001, + "dt": 0.02, + "steps": 1, + }, + "dataset": { + "reference_factor": 2, + "reference_substeps": 1, + "train_seeds": [0], + "test_seeds": [100], + "train_frames": 2, + "k0": 2.0, + }, + "training": { + "max_updates": 1, + "unroll": 2, + "hidden_channels": 4, + "kernel_size": 3, + "check_grad": False, + }, + "evaluation": { + "rollout_frames": 2, + "stable_error_threshold": 1.0, + }, + } + ], + ) + return + physics = _SHRINK_PHYSICS.get((problem, exp_key)) if physics is None: return diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py new file mode 100644 index 00000000..66c83358 --- /dev/null +++ b/tests/test_solver_in_loop.py @@ -0,0 +1,182 @@ +# Copyright 2026 Pasteur Labs. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Numerical and configuration tests for the 2D solver-in-the-loop task.""" + +from __future__ import annotations + +import dataclasses +from pathlib import Path + +import jax +import jax.numpy as jnp +import numpy as np + +from mosaic.benchmarks.core.utils import _debug_run +from mosaic.benchmarks.problems.navier_stokes_grid.corrector import ( + apply_corrector, + divergence_rms, + init_corrector, + project_periodic_correction, + reference_trajectory, + relative_l2, + spectral_restrict, +) +from mosaic.benchmarks.problems.navier_stokes_grid.ics import _tgv, _tgv_analytic +from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( + solver_in_loop, +) + +_IDENTITY_DUMMY = ( + Path(__file__).parent + / "dummy_tesseracts" + / "navier_stokes_grid_identity" + / "tesseract_api.py" +).resolve() + + +def test_periodic_corrector_is_translation_equivariant(): + params = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) + velocity = jax.random.normal(jax.random.PRNGKey(1), (8, 8, 1, 2)) + shifted = jnp.roll(velocity, shift=(2, -1), axis=(0, 1)) + + expected = jnp.roll( + apply_corrector(params, velocity, velocity_scale=1.0), + shift=(2, -1), + axis=(0, 1), + ) + actual = apply_corrector(params, shifted, velocity_scale=1.0) + + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_periodic_projection_is_divergence_free_and_zero_mean(): + delta = jax.random.normal(jax.random.PRNGKey(2), (16, 16, 1, 2)) + projected = project_periodic_correction(delta, 2.0 * np.pi) + + assert divergence_rms(projected, 2.0 * np.pi) < 1e-5 + np.testing.assert_allclose( + np.asarray(projected).mean(axis=(0, 1, 2)), + np.zeros(2), + atol=1e-6, + ) + + +def test_spectral_restriction_preserves_low_mode_tgv(): + fine = _tgv(32) + coarse = spectral_restrict(fine, 16) + + np.testing.assert_allclose(coarse, _tgv(16), rtol=1e-5, atol=1e-5) + + +def test_spectral_reference_matches_tgv_decay(): + viscosity = 0.05 + dt = 0.01 + steps = 10 + initial = _tgv(16) + trajectory = reference_trajectory( + initial, + viscosity=viscosity, + dt=dt, + frame_steps=steps, + n_frames=1, + substeps=1, + domain_extent=2.0 * np.pi, + ) + expected = _tgv_analytic( + initial, + nu=viscosity, + t=dt * steps, + L=2.0 * np.pi, + ) + + assert relative_l2(trajectory[-1], expected) < 1e-5 + + +def test_debug_run_caps_corrector_training(): + run = { + "physics": {"N": 32, "steps": 4}, + "training": { + "max_updates": 100, + "unroll": 8, + "hidden_channels": 32, + "kernel_size": 5, + "check_grad": True, + }, + "dataset": { + "train_seeds": [0, 1, 2], + "test_seeds": [100, 101], + "train_frames": 16, + }, + "evaluation": {"rollout_frames": 24}, + } + + _debug_run(run) + + assert run["physics"]["N"] == 16 + assert run["physics"]["steps"] == 4 + assert run["training"] == { + "max_updates": 2, + "unroll": 2, + "hidden_channels": 8, + "kernel_size": 3, + "check_grad": False, + } + assert run["dataset"]["train_seeds"] == [0] + assert run["dataset"]["test_seeds"] == [100] + assert run["dataset"]["train_frames"] == 3 + assert run["evaluation"]["rollout_frames"] == 3 + + +def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): + """One update crosses the apply/VJP boundary and writes canonical artifacts.""" + from mosaic.benchmarks.problems import get_config + + monkeypatch.setenv("MOSAIC_RESULTS_DIR", str(tmp_path)) + base = get_config("ns-grid") + jax_cfd = next(spec for spec in base.solvers if spec.key == "jax_cfd") + cfg = dataclasses.replace(base, solvers=[jax_cfd]) + cfg.add_experiment( + "optimization/solver_in_loop_smoke", + solver_in_loop, + runs=[ + { + "ic": {"name": "multimode", "seed": 0}, + "physics": {"N": 8, "nu": 0.001, "dt": 0.02, "steps": 1}, + "dataset": { + "reference_factor": 2, + "reference_substeps": 1, + "train_seeds": [0], + "test_seeds": [100], + "train_frames": 2, + "k0": 2.0, + }, + "training": { + "max_updates": 1, + "unroll": 2, + "hidden_channels": 4, + "kernel_size": 3, + "check_grad": True, + "fd_epsilon": 1e-3, + }, + "evaluation": { + "rollout_frames": 2, + "stable_error_threshold": 1.0, + }, + } + ], + ) + result = cfg.experiments["optimization/solver_in_loop_smoke"].fn( + cfg, + {jax_cfd.name: f"inprocess:{_IDENTITY_DUMMY}"}, + ) + + assert len(result["results"]) == 1 + metrics = result["results"][0]["metrics"] + assert metrics["n_updates"] == 1 + assert metrics["completed"] is True + assert metrics["final_grad_norm"] > 0 + assert metrics["end_to_end_fd_rel_error"] < 5e-2 + out_dir = tmp_path / "ns-grid" / "optimization" / "solver_in_loop_smoke" + assert (out_dir / "result.json").exists() + assert (out_dir / "corrector_fields.npz").exists() From 91c98fdac070358b304db222ec4a3113b5ff3852 Mon Sep 17 00:00:00 2001 From: andrinr Date: Fri, 24 Jul 2026 22:08:45 +0200 Subject: [PATCH 02/15] feat(domain): render solver correction fields --- .../problems/navier_stokes_grid/plots.py | 120 +++++++++++++++++- tests/test_solver_in_loop.py | 44 +++++++ 2 files changed, 163 insertions(+), 1 deletion(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index b79ccd45..720f37e1 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -155,9 +155,127 @@ def plot_solver_in_loop( ] solver_legend(fig, aliases, order=NS_ORDER) fig.suptitle("Solver-in-the-loop neural correction") + figs = [fig] if save: save_fig(fig, "solver_in_loop", out_dir) - return [fig] + fields_fig = _plot_solver_in_loop_fields(arrays, names, out_dir, save=save) + if fields_fig is not None: + figs.append(fields_fig) + return figs + + +def _periodic_vorticity_2d( + velocity: np.ndarray, + *, + domain_extent: float = 2.0 * np.pi, +) -> np.ndarray: + """Return centered-difference vorticity for a periodic 2-D velocity field.""" + ux, uy = _vel_components_2d(np.asarray(velocity)) + dx = domain_extent / ux.shape[0] + dy = domain_extent / ux.shape[1] + d_uy_dx = (np.roll(uy, -1, axis=0) - np.roll(uy, 1, axis=0)) / (2.0 * dx) + d_ux_dy = (np.roll(ux, -1, axis=1) - np.roll(ux, 1, axis=1)) / (2.0 * dy) + return d_uy_dx - d_ux_dy + + +def _plot_solver_in_loop_fields( + arrays: dict[str, np.ndarray], + names: list[str], + out_dir: Path, + *, + save: bool, +) -> plt.Figure | None: + """Render final held-out reference, raw-solver, and corrected vorticity. + + Each solver gets one row. The reference column is deliberately repeated so + every raw/corrected pair can be read horizontally without looking across + rows. All panels share one robust symmetric colour scale. + """ + reference = np.asarray(arrays.get("reference_rollout", np.array([]))) + if reference.size == 0 or reference.shape[0] == 0: + return None + + rows: list[tuple[str, np.ndarray, np.ndarray]] = [] + for idx, name in enumerate(names): + raw = np.asarray(arrays.get(f"rollout_uncorrected_{idx}", np.array([]))) + corrected = np.asarray(arrays.get(f"rollout_corrected_{idx}", np.array([]))) + if raw.size and corrected.size and raw.shape[0] and corrected.shape[0]: + rows.append((name, raw[-1], corrected[-1])) + if not rows: + return None + solver_order = {alias: idx for idx, alias in enumerate(NS_ORDER)} + rows.sort( + key=lambda row: solver_order.get( + resolve_solver_alias(row[0]) or row[0], + len(solver_order), + ) + ) + + reference_final = reference[-1] + reference_vorticity = _periodic_vorticity_2d(reference_final) + rendered_rows: list[tuple[str, np.ndarray, np.ndarray]] = [] + scale_fields = [reference_vorticity] + for name, raw, corrected in rows: + raw_vorticity = _periodic_vorticity_2d(raw) + corrected_vorticity = _periodic_vorticity_2d(corrected) + rendered_rows.append((name, raw_vorticity, corrected_vorticity)) + scale_fields.extend((raw_vorticity, corrected_vorticity)) + + magnitudes = np.concatenate([np.abs(field).ravel() for field in scale_fields]) + vmax = float(np.percentile(magnitudes, 99.0)) or 1.0 + + plt.rcParams.update(RCPARAMS) + fig, axes = plt.subplots( + len(rendered_rows), + 3, + figsize=(TEXTWIDTH, max(1.8, 0.95 * len(rendered_rows))), + squeeze=False, + layout="constrained", + ) + column_titles = ("Reference", "Solver only", "Solver + corrector") + image = None + for row_idx, (name, raw_vorticity, corrected_vorticity) in enumerate(rendered_rows): + label, _color, _linestyle, _marker = solver_props(name) + for col_idx, field in enumerate( + (reference_vorticity, raw_vorticity, corrected_vorticity) + ): + ax = axes[row_idx, col_idx] + image = ax.imshow( + field.T, + origin="lower", + cmap="RdBu_r", + vmin=-vmax, + vmax=vmax, + interpolation="nearest", + ) + ax.set_xticks([]) + ax.set_yticks([]) + if row_idx == 0: + ax.set_title(column_titles[col_idx]) + axes[row_idx, 0].set_ylabel( + label, + rotation=0, + ha="right", + va="center", + labelpad=8, + ) + + if image is not None: + colorbar = fig.colorbar( + image, + ax=axes.ravel().tolist(), + location="right", + shrink=0.82, + pad=0.02, + ) + colorbar.set_label(r"Vorticity $\omega$") + times = np.asarray(arrays.get("evaluation_times", np.array([]))) + time_suffix = f" at $t={float(times[-1]):g}$" if times.size else "" + fig.suptitle(f"Held-out final-frame flow{time_suffix}") + + if save: + save_fig(fig, "solver_in_loop_fields", out_dir) + return fig def plot_drag_opt( diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index 66c83358..b88e5bd2 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -23,6 +23,10 @@ spectral_restrict, ) from mosaic.benchmarks.problems.navier_stokes_grid.ics import _tgv, _tgv_analytic +from mosaic.benchmarks.problems.navier_stokes_grid.plots import ( + _periodic_vorticity_2d, + _plot_solver_in_loop_fields, +) from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( solver_in_loop, ) @@ -128,6 +132,46 @@ def test_debug_run_caps_corrector_training(): assert run["evaluation"]["rollout_frames"] == 3 +def test_solver_in_loop_fields_render_reference_raw_and_corrected(tmp_path): + n = 16 + coordinates = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + x, y = np.meshgrid(coordinates, coordinates, indexing="ij") + reference = np.zeros((n, n, 1, 2), dtype=np.float32) + reference[:, :, 0, 0] = np.sin(y) + reference[:, :, 0, 1] = -np.sin(x) + + expected_vorticity = -np.cos(x) - np.cos(y) + np.testing.assert_allclose( + _periodic_vorticity_2d(reference), + expected_vorticity, + rtol=3e-2, + atol=3e-2, + ) + + arrays = { + "reference_rollout": np.stack((reference, reference)), + "evaluation_times": np.asarray([0.0, 0.1]), + "rollout_uncorrected_0": np.stack((reference, 0.8 * reference)), + "rollout_corrected_0": np.stack((reference, 0.95 * reference)), + } + fig = _plot_solver_in_loop_fields( + arrays, + ["jax-cfd"], + tmp_path, + save=True, + ) + + assert fig is not None + assert [axis.get_title() for axis in fig.axes[:3]] == [ + "Reference", + "Solver only", + "Solver + corrector", + ] + rendered = tmp_path / "solver_in_loop_fields.png" + assert rendered.exists() + assert rendered.stat().st_size > 0 + + def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): """One update crosses the apply/VJP boundary and writes canonical artifacts.""" from mosaic.benchmarks.problems import get_config From 1ea41df29ec566cbfc9201ec646ad19ab1615318 Mon Sep 17 00:00:00 2001 From: andrinr Date: Fri, 24 Jul 2026 22:09:10 +0200 Subject: [PATCH 03/15] fix(runtime): support daemon-free GPU solver runs --- mosaic/benchmarks/core/runner.py | 12 + .../jax-cfd/tesseract_requirements.txt | 2 +- .../phiflow/tesseract_requirements.txt | 2 +- .../pict/tesseract_config.yaml | 9 +- .../pict/tesseract_requirements.txt | 2 +- .../warp-ns/tesseract_config.yaml | 4 + .../xlb/tesseract_requirements.txt | 2 +- production.uv.lock | 347 +++++++++++++++++- pyproject.toml | 3 + tests/core/test_runner.py | 28 ++ 10 files changed, 403 insertions(+), 8 deletions(-) diff --git a/mosaic/benchmarks/core/runner.py b/mosaic/benchmarks/core/runner.py index 5b65e615..f9df501f 100644 --- a/mosaic/benchmarks/core/runner.py +++ b/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, diff --git a/mosaic/tesseracts/navier-stokes-grid/jax-cfd/tesseract_requirements.txt b/mosaic/tesseracts/navier-stokes-grid/jax-cfd/tesseract_requirements.txt index 53f5234a..6d19694c 100644 --- a/mosaic/tesseracts/navier-stokes-grid/jax-cfd/tesseract_requirements.txt +++ b/mosaic/tesseracts/navier-stokes-grid/jax-cfd/tesseract_requirements.txt @@ -1,4 +1,4 @@ jax-cfd==0.2.1 -jax[cuda12-local]==0.10.0 +jax[cuda12]==0.10.0 equinox==0.13.8 ../../../mosaic_shared diff --git a/mosaic/tesseracts/navier-stokes-grid/phiflow/tesseract_requirements.txt b/mosaic/tesseracts/navier-stokes-grid/phiflow/tesseract_requirements.txt index 3e7acb6f..54fbe1c1 100644 --- a/mosaic/tesseracts/navier-stokes-grid/phiflow/tesseract_requirements.txt +++ b/mosaic/tesseracts/navier-stokes-grid/phiflow/tesseract_requirements.txt @@ -1,4 +1,4 @@ phiflow==3.4.0 -jax[cuda12-local]==0.10.0 +jax[cuda12]==0.10.0 equinox==0.13.8 ../../../mosaic_shared diff --git a/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_config.yaml b/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_config.yaml index b6dbfbbc..aaa7c8de 100644 --- a/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_config.yaml +++ b/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_config.yaml @@ -17,7 +17,8 @@ build_config: # runtime, nvidia-cuda-nvcc-cu12 provides nvcc for building PISOtorch's # CUDA extensions, and the GPU driver is injected at runtime by # nvidia-container-toolkit on the host. - # torch==2.5.1+cu121 is pinned in requirements to avoid cu128/cu130 mismatch. + # torch==2.7.1+cu128 is the oldest stable wheel with NVIDIA Blackwell + # support and stays aligned with the CUDA 12.8 extension toolchain below. # No extra_packages β€” build-only deps (ninja, gcc, nvcc) are installed in # custom_build_steps so they only appear in run_stage, avoiding duplication # across the two-stage Dockerfile and saving ~200 MB of disk during build. @@ -32,7 +33,7 @@ build_config: && dpkg -i /tmp/cuda-keyring.deb && rm /tmp/cuda-keyring.deb \ && apt-get update \ && apt-get install -y --no-install-recommends \ - cuda-nvcc-12-3 cuda-cudart-dev-12-3 libcublas-dev-12-3 libcusparse-dev-12-3 \ + cuda-nvcc-12-8 cuda-cudart-dev-12-8 libcublas-dev-12-8 libcusparse-dev-12-8 \ && rm -rf /var/lib/apt/lists/* - | RUN TORCH_LIB=$(python -c "import torch; print(torch.__path__[0])")/lib \ @@ -57,7 +58,9 @@ build_config: env: CUDA_HOME: "/usr/local/cuda" PATH: "/usr/local/cuda/bin:${PATH}" - TORCH_CUDA_ARCH_LIST: "7.0;7.5;8.0;8.6;9.0" + TORCH_CUDA_ARCH_LIST: "7.0;7.5;8.0;8.6;9.0;10.0;12.0" + NVIDIA_VISIBLE_DEVICES: "all" + NVIDIA_DRIVER_CAPABILITIES: "compute,utility" metadata: mosaic: diff --git a/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_requirements.txt b/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_requirements.txt index 19af4360..6a20d916 100644 --- a/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_requirements.txt +++ b/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_requirements.txt @@ -1,4 +1,4 @@ -https://download.pytorch.org/whl/cu121/torch-2.5.1%2Bcu121-cp311-cp311-linux_x86_64.whl +https://download.pytorch.org/whl/cu128/torch-2.7.1%2Bcu128-cp311-cp311-manylinux_2_28_x86_64.whl numpy==2.4.4 scipy==1.17.1 imageio==2.37.3 diff --git a/mosaic/tesseracts/navier-stokes-grid/warp-ns/tesseract_config.yaml b/mosaic/tesseracts/navier-stokes-grid/warp-ns/tesseract_config.yaml index 1dca1a34..9cc82840 100644 --- a/mosaic/tesseracts/navier-stokes-grid/warp-ns/tesseract_config.yaml +++ b/mosaic/tesseracts/navier-stokes-grid/warp-ns/tesseract_config.yaml @@ -9,6 +9,10 @@ build_config: custom_build_steps: - "RUN /python-env/bin/pip install /tmp/mosaic_shared" +env: + NVIDIA_VISIBLE_DEVICES: "all" + NVIDIA_DRIVER_CAPABILITIES: "compute,utility" + metadata: mosaic: name: "Warp-NS" diff --git a/mosaic/tesseracts/navier-stokes-grid/xlb/tesseract_requirements.txt b/mosaic/tesseracts/navier-stokes-grid/xlb/tesseract_requirements.txt index 55d94e21..f809a252 100644 --- a/mosaic/tesseracts/navier-stokes-grid/xlb/tesseract_requirements.txt +++ b/mosaic/tesseracts/navier-stokes-grid/xlb/tesseract_requirements.txt @@ -1,4 +1,4 @@ -jax[cuda12-local]==0.10.0 +jax[cuda12]==0.10.0 equinox==0.13.8 xlb==0.3.1 ../../../mosaic_shared diff --git a/production.uv.lock b/production.uv.lock index 5e15f609..9eb2596b 100644 --- a/production.uv.lock +++ b/production.uv.lock @@ -1275,6 +1275,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/a8/97ef0cbb7a17143ace2643d600a7b80d6705b2266fc31078229e406bdef2/jax-0.6.2-py3-none-any.whl", hash = "sha256:bb24a82dc60ccf704dcaf6dbd07d04957f68a6c686db19630dd75260d1fb788c", size = 2722396, upload-time = "2025-06-17T23:10:25.293Z" }, ] +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, extra = ["with-cuda"], marker = "python_full_version < '3.11'" }, + { name = "jaxlib", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, +] + [[package]] name = "jax" version = "0.10.2" @@ -1296,6 +1302,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/89/82/5ab5211079a151b6f661529369c0c8e98ec64cabf5c0cf22a0a05af124d8/jax-0.10.2-py3-none-any.whl", hash = "sha256:724d73c4678d8b06f6a6ab4db1b8a2fea8cd4f1e2c2564f99601634ec7b8d1c6", size = 3219515, upload-time = "2026-06-17T23:42:41.259Z" }, ] +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", version = "0.10.2", source = { registry = "https://pypi.org/simple" }, extra = ["with-cuda"], marker = "python_full_version == '3.11.*'" }, + { name = "jaxlib", version = "0.10.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, +] + [[package]] name = "jax" version = "0.11.0" @@ -1323,6 +1335,185 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d4/5a/26a1305f8978bc41f889a4172198c7d8976e8e394fc4d03597d9edfd6b7e/jax-0.11.0-py3-none-any.whl", hash = "sha256:b31ed66c4321d5c9e48b861f51a72ed5f7960c93514296202d2041cbb9ac45bf", size = 3255281, upload-time = "2026-07-16T18:58:25.02Z" }, ] +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", version = "0.11.0", source = { registry = "https://pypi.org/simple" }, extra = ["with-cuda"], marker = "python_full_version >= '3.12'" }, + { name = "jaxlib", version = "0.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.6.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/8b/096a12d91b1bc76a13cd8e63f2d8eae5ca9b5693b5ee2687e46be3b5d779/jax_cuda12_pjrt-0.6.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:22faf020d2e8f7ca1e2915633241f7df7678b73c7078f5f0b2f113248337f7de", size = 111228681, upload-time = "2025-06-17T23:11:55.179Z" }, + { url = "https://files.pythonhosted.org/packages/9c/8e/21c21b4335fce1c022c339da5e6b6249c246ad062e924d28fb0eda4bcef0/jax_cuda12_pjrt-0.6.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8cd9ead7948ea2c778a508fef5d1159e8b7abf4fccc7037c3fe1dbfcd95012dc", size = 125263999, upload-time = "2025-06-17T23:11:59.986Z" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.10.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/bf/b0/e1d392ee24ecd53caa202194746cb12058befde5882d1f3307922829783f/jax_cuda12_pjrt-0.10.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b6f2d66548ee1ee910a836b5fe3d0ebbe1da04a62d0f468ba4822189036a32b3", size = 169847949, upload-time = "2026-06-17T23:42:44.729Z" }, + { url = "https://files.pythonhosted.org/packages/77/00/326facd192b7dc0731e43b30f5ada5ef235e2a3325a151d94ed3db041536/jax_cuda12_pjrt-0.10.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:806d1fd29038b6acf5a2b289dd62192abea977ee26aef60ea295b1d28a23acf8", size = 174364763, upload-time = "2026-06-17T23:42:49.931Z" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/f5/d3c0db11be08d5180d4634327e996597ef7847fcfa0c5e53282dd9256485/jax_cuda12_pjrt-0.11.0-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:67ec5915e7e494775d5dc0d73d1e33cb74a47b6a6bb1089eef7b517e4e873e33", size = 170618370, upload-time = "2026-07-16T18:58:28.763Z" }, + { url = "https://files.pythonhosted.org/packages/fe/73/dfe1781a2f752dbb25bd7abcc0f5a7ed74f7df3c78c4d53db32a591b01cd/jax_cuda12_pjrt-0.11.0-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:242df99c90827a2937df1c444983d4076173f9836dc0bc1d20df86960144f35a", size = 175803444, upload-time = "2026-07-16T18:58:35.744Z" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.6.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "jax-cuda12-pjrt", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/29/4b8822ca459da39bda9be7454908ae4e29d88cfb99b480b641cbb063af7a/jax_cuda12_plugin-0.6.2-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:bc5c3a75d05519b4d326e4669d0f7ad0fe0f0acf875f9313d913748ccca5a9ea", size = 15873729, upload-time = "2025-06-17T23:12:05.046Z" }, + { url = "https://files.pythonhosted.org/packages/4d/3d/f543bab6ef7eebb9840d618fb2272bdcc0e990e60aa012f14a3564532823/jax_cuda12_plugin-0.6.2-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:1751f88989269b3cdb0dfe4f7b072a6442149818c9bc98c3a395c8acaf910a79", size = 15879965, upload-time = "2025-06-17T23:12:07.323Z" }, + { url = "https://files.pythonhosted.org/packages/9e/99/90f81c660bf662698be17e8b959d8302682c5cd5ce0729c0bfc883d8affe/jax_cuda12_plugin-0.6.2-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:2cd8e279a59a38ba0c978a831e13adeb6ee9e4572fba387c7975ba3ad535dd38", size = 15873970, upload-time = "2025-06-17T23:12:09.492Z" }, + { url = "https://files.pythonhosted.org/packages/ad/00/e733c87a2fb7265c96f48c991a896552c873de949217823d519288724d91/jax_cuda12_plugin-0.6.2-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:0896cbb308d95291e205cd89d254029dee3a1df43d66e9831331a9afd2d27870", size = 15879563, upload-time = "2025-06-17T23:12:11.439Z" }, + { url = "https://files.pythonhosted.org/packages/1d/53/6ea0db7230ac3dcb26b732452f4f2e2c6868b75d3603be19cee388fef279/jax_cuda12_plugin-0.6.2-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:6c9b002d13b1fcb9403713eedd3876a227ad1ffbdfb3811b1f9f89af4c25a5f7", size = 15867462, upload-time = "2025-06-17T23:12:14.093Z" }, + { url = "https://files.pythonhosted.org/packages/a5/db/e6643143caf573273eedb991cb1af2bea964b84594a8887802eb0b6ba64a/jax_cuda12_plugin-0.6.2-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:febd099f970d350eb8fa5a2c9a2fb4b0ea7b3d6a89df1496663edfa7afe590e5", size = 15876401, upload-time = "2025-06-17T23:12:15.964Z" }, + { url = "https://files.pythonhosted.org/packages/b4/10/74fbae1c1bb9d11b113c62e03c445bdf0aaaa4981bc13d2de8e3f4f503ca/jax_cuda12_plugin-0.6.2-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:773efa8b55a837406c561f0ef02144dda9019181193760ec5419eec9dd2b9aac", size = 15868561, upload-time = "2025-06-17T23:12:18.189Z" }, + { url = "https://files.pythonhosted.org/packages/2b/96/53928ad62ecddbf76f4c413025fdeab5a90adf7fbd970d800162399e504a/jax_cuda12_plugin-0.6.2-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:db4c6103c912d8cd1adf94c34d313bb4760ca7f01c897ca7cd62e65f27994199", size = 15876276, upload-time = "2025-06-17T23:12:20.361Z" }, + { url = "https://files.pythonhosted.org/packages/22/29/f1c1790d77ccad3587e6c97045b567006a8050fec151581d4cf883779e25/jax_cuda12_plugin-0.6.2-cp313-cp313t-manylinux2014_aarch64.whl", hash = "sha256:ed5316ca1818db7ef53230ee0a41398d3a60942e361dfb857a952eb4d92fc8d7", size = 15964355, upload-time = "2025-06-17T23:12:22.176Z" }, + { url = "https://files.pythonhosted.org/packages/07/af/c3224aafbc1d2d7654359a5410319bf15361067635bd3616cf5ad3e16af2/jax_cuda12_plugin-0.6.2-cp313-cp313t-manylinux2014_x86_64.whl", hash = "sha256:83345f52f610cdb8e90044566d8e120864150b8090968c8ab6dd8e0bfb9a6a9f", size = 16037754, upload-time = "2025-06-17T23:12:24.066Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cudnn-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cufft-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cusolver-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-cusparse-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-nccl-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version < '3.11'" }, + { name = "nvidia-nvshmem-cu12", marker = "python_full_version < '3.11'" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.10.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +dependencies = [ + { name = "jax-cuda12-pjrt", version = "0.10.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/25/ac0e8a02fd46e34d4d0e67214fc022688701d3c26ea93de15d1b2c74b0d6/jax_cuda12_plugin-0.10.2-cp311-cp311-manylinux_2_27_aarch64.whl", hash = "sha256:9cd0937070bbd28e44b91435b52c7225c8ed6cb37ddb6a90e9c246873a03f9cd", size = 8098522, upload-time = "2026-06-17T23:42:53.439Z" }, + { url = "https://files.pythonhosted.org/packages/fa/55/627b9ed487e82047c6f7d1dc8174a31df7c350093292a3c664183908faf1/jax_cuda12_plugin-0.10.2-cp311-cp311-manylinux_2_27_x86_64.whl", hash = "sha256:76bbfece2f6effb4e8908eb8a5fbf1abbc1125de667683ec2ffba51c63dc57af", size = 8144162, upload-time = "2026-06-17T23:42:55.513Z" }, + { url = "https://files.pythonhosted.org/packages/94/4a/382ec14e57d774148b079c3d251fc9b9e1ed7f92dd2327823f35e0a968a7/jax_cuda12_plugin-0.10.2-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:767a1482dbf652688403c4e22ff01470e1add13c48f9d817169fd8f7440afce5", size = 8093630, upload-time = "2026-06-17T23:42:57.019Z" }, + { url = "https://files.pythonhosted.org/packages/ca/25/19ee8936b00ca74a12bfaebd4703695cee2407f954a6e4e67cdd7f82b7bf/jax_cuda12_plugin-0.10.2-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:4eb6e8e0992fd9897db2468da33f276ae414ec84dc436804124404430502eeac", size = 8141521, upload-time = "2026-06-17T23:42:58.483Z" }, + { url = "https://files.pythonhosted.org/packages/94/33/a9c4ea148d31843764383b26fd66b38cfedad9a6dbf08db634a23d9312c3/jax_cuda12_plugin-0.10.2-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:69776ac849112f4fc2d6fa616f8772106daff38fcdb825c4e39f47e878f27610", size = 8093605, upload-time = "2026-06-17T23:43:00.071Z" }, + { url = "https://files.pythonhosted.org/packages/92/07/3098ae537e68d76bbc98ed61d53e9856f63f4f1bd64f8aeb4d432dad1290/jax_cuda12_plugin-0.10.2-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:948a988927cb10843b501a215181ccb1daed3fa70531b2f1ae0842fe1c08b05e", size = 8141236, upload-time = "2026-06-17T23:43:01.55Z" }, + { url = "https://files.pythonhosted.org/packages/72/0c/ff1e1975ef108cb2e45deacf5dea812822efaf3210af075f464abd40399c/jax_cuda12_plugin-0.10.2-cp313-cp313t-manylinux_2_27_aarch64.whl", hash = "sha256:febd9d3f488d182090758936042ad647a60470923715b48a385bc312cb943c63", size = 8108408, upload-time = "2026-06-17T23:43:03.398Z" }, + { url = "https://files.pythonhosted.org/packages/09/c2/04c70eb73beb4df23f49b72d159ca74ee31cccb1be95cdab7491a5ebd259/jax_cuda12_plugin-0.10.2-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:1f1e20fba0aaa3f28ab0a3273754c8c098fa40b74bf4dac5b76c9f2a70425487", size = 8151184, upload-time = "2026-06-17T23:43:05.106Z" }, + { url = "https://files.pythonhosted.org/packages/95/34/5d5b2993f6b7bbfe781920671de61036faffa5a8bf3740b4ca23d790b6f9/jax_cuda12_plugin-0.10.2-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:4946ec57ddb795e2da8288ceb1ceddd58c2a204a0f7765af89eae486a1d3ace7", size = 8094107, upload-time = "2026-06-17T23:43:06.546Z" }, + { url = "https://files.pythonhosted.org/packages/02/ac/85bc07d7a0766c7999f9bdfcf7f7cb2487fd0eb5d6bbe2bd41a3a54cd379/jax_cuda12_plugin-0.10.2-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:0c7a0204155585cc1c4fcb30c66b7f881b38e57e1f0d8e97d48e1654bf0d27f8", size = 8141889, upload-time = "2026-06-17T23:43:08.018Z" }, + { url = "https://files.pythonhosted.org/packages/f3/7d/9fa25e24b96ab06217e1e3dd311a80f3534076bb1b19d46cbc8cf9aa55b9/jax_cuda12_plugin-0.10.2-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:c26a7dd5e71c9edf203b95480f02ec4ad60c9c7e11e38b7310b1fd1dd03ba116", size = 8108639, upload-time = "2026-06-17T23:43:09.651Z" }, + { url = "https://files.pythonhosted.org/packages/20/c2/98fa02ba204f5dc88680128129fbc7cd833d5b5ddae2734af93ef85db853/jax_cuda12_plugin-0.10.2-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:c539b052056e8181082fa2635f17e098b3f625d87da53ba35470a2c13480801d", size = 8151505, upload-time = "2026-06-17T23:43:11.023Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "python_full_version == '3.11.*' and sys_platform == 'linux'" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +dependencies = [ + { name = "jax-cuda12-pjrt", version = "0.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/12/5d35d4e3ac48e3976a07616dcb82403ec05d91fb764da4306c83cd980169/jax_cuda12_plugin-0.11.0-cp312-cp312-manylinux_2_27_aarch64.whl", hash = "sha256:3e95cb9deac2988e15226c6d8927e72bffc841505d236231e4843a8b16679f25", size = 8209123, upload-time = "2026-07-16T18:58:39.248Z" }, + { url = "https://files.pythonhosted.org/packages/b9/ba/c6a2168852154cb32aa7021074f8b849407a99bb5b7f88fe35d34cfa43c8/jax_cuda12_plugin-0.11.0-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:31ffe245d0e2fd1cf4ee57254ffc4c88505bb8e7f5eb7f6b3a37d691d736650a", size = 8229603, upload-time = "2026-07-16T18:58:41.37Z" }, + { url = "https://files.pythonhosted.org/packages/f6/0e/fff3d0eaf710b8dd3b76b3f177b0d26ff8cfed493b4d3bedd9ccdb8b687b/jax_cuda12_plugin-0.11.0-cp313-cp313-manylinux_2_27_aarch64.whl", hash = "sha256:ef103b35364795708c4a4380376d8d73fc3d487356a7037fcf20cab735fb59a6", size = 8208891, upload-time = "2026-07-16T18:58:42.877Z" }, + { url = "https://files.pythonhosted.org/packages/8a/36/fb9714a8c71cf33c114ab7ca22426c615516cf1532744b7224019e4a45d3/jax_cuda12_plugin-0.11.0-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:b6f706d6c10fe9bfafe0334c5d52e40e6e99f343c24d8e74f045ac67b7f7ace2", size = 8230117, upload-time = "2026-07-16T18:58:44.61Z" }, + { url = "https://files.pythonhosted.org/packages/38/bf/935faeaa1834a31ef4324a55288bc3cf547dc90c4f24e578361e269950f5/jax_cuda12_plugin-0.11.0-cp314-cp314-manylinux_2_27_aarch64.whl", hash = "sha256:2f8dcc6ff34833158251b52b711a3b5c73785ec364771294bc3827aaaea75450", size = 8209555, upload-time = "2026-07-16T18:58:46.313Z" }, + { url = "https://files.pythonhosted.org/packages/c6/2a/af0001e6b8e59792ace3f0c1d11e1fb81a278f57cf2d0aeb4b2e8c7a15ca/jax_cuda12_plugin-0.11.0-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:4b6b060fe73ffde206e62c2637079618c34c2b4b74be5cb0371b22a43a907c4a", size = 8230380, upload-time = "2026-07-16T18:58:48.081Z" }, + { url = "https://files.pythonhosted.org/packages/fa/46/e21fc196562672d45f0d0166d8d34f73001badd48a29c9a2be883270739d/jax_cuda12_plugin-0.11.0-cp314-cp314t-manylinux_2_27_aarch64.whl", hash = "sha256:1dd78abc57481d542876ff7d1a2c460bee7d8b191059617a0cb2c5fb4b191630", size = 8223073, upload-time = "2026-07-16T18:58:49.605Z" }, + { url = "https://files.pythonhosted.org/packages/2f/90/bade91ca305f3bc7b9e2cc67715978cc3a8a4d428d170fc1656815575c53/jax_cuda12_plugin-0.11.0-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:fbd22d6bbfa636f22bf27e4a3d6ca3e9424b7fcb7f5ccae698634a3ea5b6a242", size = 8239448, upload-time = "2026-07-16T18:58:51.71Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvcc-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "python_full_version >= '3.12' and sys_platform == 'linux'" }, +] + [[package]] name = "jaxlib" version = "0.6.2" @@ -1950,12 +2141,18 @@ dev = [ { name = "pyyaml" }, { name = "ruff" }, ] +gpu = [ + { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, extra = ["cuda12"], marker = "python_full_version < '3.11'" }, + { name = "jax", version = "0.10.2", source = { registry = "https://pypi.org/simple" }, extra = ["cuda12"], marker = "python_full_version == '3.11.*'" }, + { name = "jax", version = "0.11.0", source = { registry = "https://pypi.org/simple" }, extra = ["cuda12"], marker = "python_full_version >= '3.12'" }, +] [package.metadata] requires-dist = [ { name = "docker", specifier = ">=6" }, { name = "filelock", specifier = ">=3" }, { name = "jax", extras = ["cpu"] }, + { name = "jax", extras = ["cuda12"], marker = "extra == 'gpu'" }, { name = "matplotlib" }, { name = "mosaic-shared", marker = "extra == 'dev'", editable = "mosaic/mosaic_shared" }, { name = "numpy" }, @@ -1974,7 +2171,7 @@ requires-dist = [ { name = "tesseract-jax" }, { name = "typer", specifier = ">=0.12" }, ] -provides-extras = ["dev"] +provides-extras = ["dev", "gpu"] [[package]] name = "mosaic-shared" @@ -2296,6 +2493,123 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/ce/e5ec180bc41812edcd8daeb8639d205622c0e8c02259d8ab25a0201b3c2a/numpy-2.4.6-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:2803abfebfc990042cd494d8ce2d5f82e9d847af6d35ec486923aa19dbad5e73", size = 12504263, upload-time = "2026-05-18T23:37:09.715Z" }, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.9.2.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/a2/c96163a0fff1839c0c9548bbdeae7b853b867009e33b9b9264adc238b1cf/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:5572131a59c3eebeeb1c4c8144f772d49372c20124916e072a0e3fc30df421d5", size = 575012079, upload-time = "2026-04-08T18:51:47.303Z" }, + { url = "https://files.pythonhosted.org/packages/cb/c0/0a517bfe63ccd3b92eb254d264e28fca3c7cab75d07daea315250fb1bf73/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:e4f53a8ca8c5d6e8c492d0d0a3d565ecb59a751b19cfdaa4f6da0ab2104c1702", size = 581240110, upload-time = "2026-04-08T18:52:31.532Z" }, + { url = "https://files.pythonhosted.org/packages/20/e2/fc9a0e985249d873150276d5afb02e39a66817fedbf1a385724393e505ed/nvidia_cublas_cu12-12.9.2.10-py3-none-win_amd64.whl", hash = "sha256:623f43027d40d44ceadf0043f002bd25cf353e8f13ce90b9a87057019f560661", size = 553162896, upload-time = "2026-04-08T18:53:10.035Z" }, +] + +[[package]] +name = "nvidia-cuda-cccl-cu12" +version = "12.9.27" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/7e/82e49956b046bdc506c789235c587d9b3ef58b8bc1782258c1e247229647/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7898b38aa68beaa234d48f0868273702342a196d6e2e9d0ef058dca2390ebea", size = 3152245, upload-time = "2025-05-01T19:32:04.802Z" }, + { url = "https://files.pythonhosted.org/packages/18/2a/d4cd8506d2044e082f8cd921be57392e6a9b5ccd3ffdf050362430a3d5d5/nvidia_cuda_cccl_cu12-12.9.27-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:37869e17ce2e1ecec6eddf1927cca0f8c34e64fd848d40453df559091e2d7117", size = 3152243, upload-time = "2025-05-01T19:32:13.955Z" }, + { url = "https://files.pythonhosted.org/packages/ce/9b/1daf405620c7ac371b76b823c6336dd742673d41a150d9a04eec2c690379/nvidia_cuda_cccl_cu12-12.9.27-py3-none-win_amd64.whl", hash = "sha256:72106f95a9bb3be18472806b4f663ebf0f9248a86d14b4ae3305725b855d9d92", size = 3152175, upload-time = "2025-05-01T19:45:11.372Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/78/351b5c8cdbd9a6b4fb0d6ee73fb176dcdc1b6b6ad47c2ffff5ae8ca4a1f7/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:791853b030602c6a11d08b5578edfb957cadea06e9d3b26adbf8d036135a4afe", size = 10077166, upload-time = "2025-06-05T20:01:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:096bcf334f13e1984ba36685ad4c1d6347db214de03dbb6eebb237b41d9d934f", size = 10814997, upload-time = "2025-06-05T20:01:10.168Z" }, + { url = "https://files.pythonhosted.org/packages/3b/b4/298983ab1a83de500f77d0add86d16d63b19d1a82c59f8eaf04f90445703/nvidia_cuda_cupti_cu12-12.9.79-py3-none-win_amd64.whl", hash = "sha256:1848a9380067560d5bee10ed240eecc22991713e672c0515f9c3d9396adf93c8", size = 7730496, upload-time = "2025-06-05T20:11:26.444Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/48/b54a06168a2190572a312bfe4ce443687773eb61367ced31e064953dd2f7/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5d6a0d32fdc7ea39917c20065614ae93add6f577d840233237ff08e9a38f58f0", size = 40546229, upload-time = "2025-06-05T20:01:53.357Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" }, + { url = "https://files.pythonhosted.org/packages/d2/9e/c71c53655a65d7531c89421c282359e2f626838762f1ce6180ea0bbebd29/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:8ed7f0b17dea662755395be029376db3b94fed5cbb17c2d35cc866c5b1b84099", size = 34669845, upload-time = "2025-06-05T20:11:56.308Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/85/e4af82cc9202023862090bfca4ea827d533329e925c758f0cde964cb54b7/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:210cf05005a447e29214e9ce50851e83fc5f4358df8b453155d5e1918094dcb4", size = 89568129, upload-time = "2025-06-05T20:02:41.973Z" }, + { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" }, + { url = "https://files.pythonhosted.org/packages/52/de/823919be3b9d0ccbf1f784035423c5f18f4267fb0123558d58b813c6ec86/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:72972ebdcf504d69462d3bcd67e7b81edd25d0fb85a2c46d3ea3517666636349", size = 76408187, upload-time = "2025-06-05T20:12:27.819Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e0/0279bd94539fda525e0c8538db29b72a5a8495b0c12173113471d28bce78/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83469a846206f2a733db0c42e223589ab62fd2fabac4432d2f8802de4bded0a4", size = 3515012, upload-time = "2025-06-05T20:00:35.519Z" }, + { url = "https://files.pythonhosted.org/packages/bc/46/a92db19b8309581092a3add7e6fceb4c301a3fd233969856a8cbf042cd3c/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:25bba2dfb01d48a9b59ca474a1ac43c6ebf7011f1b0b8cc44f54eb6ac48a96c3", size = 3493179, upload-time = "2025-06-05T20:00:53.735Z" }, + { url = "https://files.pythonhosted.org/packages/59/df/e7c3a360be4f7b93cee39271b792669baeb3846c58a4df6dfcf187a7ffab/nvidia_cuda_runtime_cu12-12.9.79-py3-none-win_amd64.whl", hash = "sha256:8e018af8fa02363876860388bd10ccb89eb9ab8fb0aa749aaf58430a9f7c4891", size = 3591604, upload-time = "2025-06-05T20:11:17.036Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.24.0.43" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/f1/cd42563325fa827f54ff30da05686c747652bdbd4cb5654cea54d7d0ad4f/nvidia_cudnn_cu12-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a42996943f0cd78ddfd61c8bf59361672a19b63e0491aa22a53d6fe63a3f854a", size = 856490582, upload-time = "2026-07-02T16:21:30.924Z" }, + { url = "https://files.pythonhosted.org/packages/10/13/b8887c869cf2471339a24b60d3c28e761facbb534935f572b61423371abb/nvidia_cudnn_cu12-9.24.0.43-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:f424192dd85e7d29f44be18df2dae4c80d32c67a29c0d42f5c283c40cfdf871c", size = 799083985, upload-time = "2026-07-02T16:25:37.467Z" }, + { url = "https://files.pythonhosted.org/packages/29/28/2c9a2a97a8b3fedcf74a14f38fd5edfae12274380a829fdc6b16ce29be4c/nvidia_cudnn_cu12-9.24.0.43-py3-none-win_amd64.whl", hash = "sha256:cbd41a0ab084422c936dc9fb2fc89be5ea9a85bc421c6f23d0243bdfc945fbef", size = 737103728, upload-time = "2026-07-02T16:30:10.901Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.4.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/2b/76445b0af890da61b501fde30650a1a4bd910607261b209cccb5235d3daa/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1a28c9b12260a1aa7a8fd12f5ebd82d027963d635ba82ff39a1acfa7c4c0fbcf", size = 200822453, upload-time = "2025-06-05T20:05:27.889Z" }, + { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, + { url = "https://files.pythonhosted.org/packages/20/ee/29955203338515b940bd4f60ffdbc073428f25ef9bfbce44c9a066aedc5c/nvidia_cufft_cu12-11.4.1.4-py3-none-win_amd64.whl", hash = "sha256:8e5bfaac795e93f80611f807d42844e8e27e340e0cde270dcb6c65386d795b80", size = 200067309, upload-time = "2025-06-05T20:13:59.762Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.5.82" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cusparse-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" }, + { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" }, + { url = "https://files.pythonhosted.org/packages/32/5d/feb7f86b809f89b14193beffebe24cf2e4bf7af08372ab8cdd34d19a65a0/nvidia_cusolver_cu12-11.7.5.82-py3-none-win_amd64.whl", hash = "sha256:77666337237716783c6269a658dea310195cddbd80a5b2919b1ba8735cec8efd", size = 326215953, upload-time = "2025-06-05T20:14:41.76Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.10.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" }, + { url = "https://files.pythonhosted.org/packages/12/46/b0fd4b04f86577921feb97d8e2cf028afe04f614d17fb5013de9282c9216/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:73060ce019ac064a057267c585bf1fd5a353734151f87472ff02b2c5c9984e78", size = 366465088, upload-time = "2025-06-05T20:08:20.413Z" }, + { url = "https://files.pythonhosted.org/packages/73/ef/063500c25670fbd1cbb0cd3eb7c8a061585b53adb4dd8bf3492bb49b0df3/nvidia_cusparse_cu12-12.5.10.65-py3-none-win_amd64.whl", hash = "sha256:9e487468a22a1eaf1fbd1d2035936a905feb79c4ce5c2f67626764ee4f90227c", size = 362504719, upload-time = "2025-06-05T20:15:17.947Z" }, +] + [[package]] name = "nvidia-ml-py" version = "13.610.43" @@ -2305,6 +2619,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl", hash = "sha256:f13c72698edef492f985cc225f14faafe68ae065a2e407f45bdf6f4b9b43fde8", size = 53163, upload-time = "2026-06-01T18:54:07.704Z" }, ] +[[package]] +name = "nvidia-nccl-cu12" +version = "2.30.7" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/8c/554bb020501d6c04ad8127d83f728137f8f9123f991666efbdcf9095a221/nvidia_nccl_cu12-2.30.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:03ecd776fd1d58fd2c9a0a687dcf8db9ecd0057382dba646fa3d65786d4a9ea1", size = 303277471, upload-time = "2026-06-09T03:24:16.327Z" }, + { url = "https://files.pythonhosted.org/packages/50/32/e7ffa9c324ae260e5dbb4af2cd557bf7a8d155c8ac7b79a785fe1796fb92/nvidia_nccl_cu12-2.30.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:8ce1b8213f61f2bfac132e6df890af6450b77cbd140c6ce4e98cb0c2d8e678c9", size = 303361239, upload-time = "2026-06-09T03:24:53.816Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, + { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, + { url = "https://files.pythonhosted.org/packages/dd/7e/2eecb277d8a98184d881fb98a738363fd4f14577a4d2d7f8264266e82623/nvidia_nvjitlink_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:cc6fcec260ca843c10e34c936921a1c426b351753587fdd638e8cff7b16bb9db", size = 35584936, upload-time = "2025-06-05T20:16:08.525Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.7.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-cccl-cu12", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/a1/9ca9cc3338217311cee21f7ba795dad4f38b7a279a6da7d3d549258f60b1/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:98495de30994f1401a12cc8158ffacbeada44579f54185d75665381468056faa", size = 229963056, upload-time = "2026-07-17T19:23:44.067Z" }, + { url = "https://files.pythonhosted.org/packages/63/03/cde130e3ff706784f9efd47204f299b15113ea08203e98a556603dac245d/nvidia_nvshmem_cu12-3.7.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:02b6d3482d90ea8ba214bc400362297f573257763990a98c9846ec5c786a7227", size = 230171132, upload-time = "2026-07-17T19:25:07.092Z" }, +] + [[package]] name = "opentelemetry-api" version = "1.44.0" diff --git a/pyproject.toml b/pyproject.toml index 127f830d..b1545828 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,9 @@ dependencies = [ ] [project.optional-dependencies] +gpu = [ + "jax[cuda12]", +] dev = [ "pytest", "pytest-cov", diff --git a/tests/core/test_runner.py b/tests/core/test_runner.py index 7fe97041..4426126a 100644 --- a/tests/core/test_runner.py +++ b/tests/core/test_runner.py @@ -57,6 +57,34 @@ def test_timeout_tuple_matches_components(): ) +# ── Tesseract launch modes ─────────────────────────────────────────────────── + + +def test_tracked_tesseract_connects_to_url_without_container_lifecycle(monkeypatch): + """``url:`` tags connect to a caller-owned service and do not enter it.""" + expected = object() + calls: list[tuple[str, tuple[float, float]]] = [] + + class FakeTesseract: + @staticmethod + def from_url(url, *, timeout): + calls.append((url, timeout)) + return expected + + @staticmethod + def from_image(*_args, **_kwargs): # pragma: no cover - must not run + raise AssertionError("url tags must not launch a container") + + monkeypatch.setattr(runner, "Tesseract", FakeTesseract) + live_before = set(runner._live_containers) + + with runner._tracked_tesseract("url:http://127.0.0.1:8123", [], {}) as actual: + assert actual is expected + + assert calls == [("http://127.0.0.1:8123", runner.MOSAIC_TESSERACT_TIMEOUT_TUPLE)] + assert runner._live_containers == live_before + + # ── safe_apply_with_extras: success paths ───────────────────────────────────── From 8c3389ab0fc8c505c0c003205b1ce6b1a8cfefda Mon Sep 17 00:00:00 2001 From: andrinr Date: Fri, 24 Jul 2026 23:29:50 +0200 Subject: [PATCH 04/15] feat(domain): benchmark solver-loop trainability --- mosaic/benchmarks/core/utils.py | 2 + .../problems/navier_stokes_grid/config.py | 75 +- .../problems/navier_stokes_grid/corrector.py | 171 +++-- .../problems/navier_stokes_grid/plots.py | 699 +++++++++++++++++- .../navier_stokes_grid/solver_in_loop.py | 617 ++++++++++++++-- production.uv.lock | 70 ++ pyproject.toml | 1 + requirements.txt | 14 + tests/test_dummy_integration.py | 16 +- tests/test_solver_in_loop.py | 117 ++- 10 files changed, 1611 insertions(+), 171 deletions(-) diff --git a/mosaic/benchmarks/core/utils.py b/mosaic/benchmarks/core/utils.py index 824e0ea0..97556880 100644 --- a/mosaic/benchmarks/core/utils.py +++ b/mosaic/benchmarks/core/utils.py @@ -142,6 +142,8 @@ def _debug_run(run: dict) -> None: 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: diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index e5725431..ae2fb418 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -68,7 +68,7 @@ from .ics import _flat_inflow, _multimode, _tgv, _tgv_analytic, _uniform_flow from .optimization import drag_opt from .physics import DIAGNOSTICS, make_inputs -from .plots import plot_drag_opt, plot_solver_in_loop +from .plots import plot_drag_opt, plot_solver_in_loop, plot_solver_in_loop_tgv from .solver_in_loop import solver_in_loop _TESSERACT_SLUG = "navier-stokes-grid" @@ -448,12 +448,15 @@ "optimization/solver_in_loop", solver_in_loop, description=( - "Train an identical residual convolutional corrector through each " - "differentiable solver and evaluate its held-out recurrent rollout." + "Train an identical Equinox periodic residual corrector through each " + "differentiable solver. Paired stop-gradient training and native versus " + "restarted rollouts separate raw accuracy, coupling cost, correctability, " + "and the incremental benefit of the solver VJP." ), plot_description=( - "Corrector training loss, held-out rollout error, stable horizon, and " - "wall-clock cost for each differentiable solver." + "Corrector training, held-out trajectories, physical diagnostics, " + "state-restart penalty, correction gain, solver-VJP lift, and wall-clock " + "cost for each differentiable solver." ), # Keep dataset seed lists inside an explicit run payload: they are data, # not benchmark sweep coordinates. @@ -470,6 +473,7 @@ "steps": 4, }, "dataset": { + "reference_kind": "pseudo_spectral_multimode", "reference_factor": 2, "reference_substeps": 2, "train_seeds": [0, 1, 2, 3], @@ -477,28 +481,83 @@ "train_frames": 16, "k0": 6.0, "sigma_k": 1.0, - "amplitude": 0.3, + # The stronger field and longer held-out horizon produce + # substantial nonlinear evolution (rather than testing a + # nearly frozen decaying texture). + "amplitude": 0.5, }, "training": { "max_updates": 100, "unroll": 4, "lr": 1e-4, "clip_norm": 1.0, + "architecture": "periodic_residual_cnn", "hidden_channels": 32, "kernel_size": 5, "seed": 2026, - "model_seed": 0, + "model_seeds": [0, 1, 2], "check_grad": True, "fd_epsilon": 1e-2, }, "evaluation": { - "rollout_frames": 24, + "rollout_frames": 36, "stable_error_threshold": 1.0, }, } ], plot=plot_solver_in_loop, ) +problem.add_experiment( + "optimization/solver_in_loop_tgv", + solver_in_loop, + description=( + "Repeat solver-in-the-loop training against an analytic, translated " + "Taylor--Green vortex. This low-complexity regime audits physical time, " + "viscous decay, and canonical state-restart effects separately from the " + "nonlinear multimode comparison." + ), + plot_description=( + "Native and restarted solver error, neural-correction gain, solver-VJP " + "lift, and held-out trajectories against the analytic TGV solution." + ), + runs=[ + { + "ic": {"name": "tgv", "seed": 0}, + "physics": { + "N": 32, + "nu": 0.05, + "dt": 0.02, + "steps": 4, + }, + "dataset": { + "reference_kind": "analytic_tgv", + "reference_factor": 2, + "reference_substeps": 1, + "train_seeds": [0, 1, 2, 3], + "test_seeds": [100, 101], + "train_frames": 16, + }, + "training": { + "max_updates": 100, + "unroll": 4, + "lr": 1e-4, + "clip_norm": 1.0, + "architecture": "periodic_residual_cnn", + "hidden_channels": 32, + "kernel_size": 5, + "seed": 2026, + "model_seeds": [0, 1, 2], + "check_grad": True, + "fd_epsilon": 1e-2, + }, + "evaluation": { + "rollout_frames": 24, + "stable_error_threshold": 1.0, + }, + } + ], + plot=plot_solver_in_loop_tgv, +) # Bonus plot (not paired with an experiment). problem.add_extra_plot( "_extra/gradient/jacobian_svd_comparison", diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py index 88def793..4f0d20ef 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py @@ -11,81 +11,106 @@ from __future__ import annotations -from collections.abc import Sequence from typing import Any +import equinox as eqx import jax import jax.numpy as jnp import numpy as np +class PeriodicResidualCNN(eqx.Module): + """Small translation-equivariant corrector built from periodic convolutions.""" + + layers: tuple[eqx.nn.Conv2d, ...] + architecture: str = eqx.field(static=True) + hidden_channels: int = eqx.field(static=True) + kernel_size: int = eqx.field(static=True) + + def __init__( + self, + key: jax.Array, + *, + hidden_channels: int, + kernel_size: int, + ) -> None: + if kernel_size % 2 != 1: + raise ValueError("kernel_size must be odd") + channels = (2, hidden_channels, hidden_channels * 2, 2) + keys = jax.random.split(key, len(channels) - 1) + layers: list[eqx.nn.Conv2d] = [] + for idx, (cin, cout, layer_key) in enumerate( + zip(channels[:-1], channels[1:], keys, strict=True) + ): + layer = eqx.nn.Conv2d( + cin, + cout, + kernel_size, + padding=kernel_size // 2, + padding_mode="CIRCULAR", + key=layer_key, + ) + fan_in = kernel_size * kernel_size * cin + scale = np.sqrt(2.0 / fan_in) + if idx == len(channels) - 2: + scale *= 1e-2 + weight = ( + jax.random.normal(layer_key, layer.weight.shape, dtype=jnp.float32) + * scale + ) + layer = eqx.tree_at( + lambda value: (value.weight, value.bias), + layer, + ( + weight, + jnp.zeros((cout, 1, 1), dtype=jnp.float32), + ), + ) + layers.append(layer) + self.layers = tuple(layers) + self.architecture = "periodic_residual_cnn" + self.hidden_channels = hidden_channels + self.kernel_size = kernel_size + + def __call__(self, velocity: jax.Array) -> jax.Array: + """Map a channel-first velocity field to an additive correction.""" + x = velocity + for layer in self.layers[:-1]: + x = jax.nn.gelu(layer(x)) + return self.layers[-1](x) + + def init_corrector( key: jax.Array, *, hidden_channels: int = 32, kernel_size: int = 5, -) -> tuple[dict[str, jax.Array], ...]: - """Initialise a small three-layer periodic residual CNN. + architecture: str = "periodic_residual_cnn", +) -> PeriodicResidualCNN: + """Initialise the benchmark-owned Equinox corrector. The final layer starts at a much smaller scale than the feature layers, so the initial corrected rollout remains close to the underlying solver while gradients can still reach every layer on the first update. """ - if kernel_size % 2 != 1: - raise ValueError("kernel_size must be odd") - channels = (2, hidden_channels, hidden_channels * 2, 2) - keys = jax.random.split(key, len(channels) - 1) - layers: list[dict[str, jax.Array]] = [] - for idx, (cin, cout, layer_key) in enumerate( - zip(channels[:-1], channels[1:], keys, strict=True) - ): - fan_in = kernel_size * kernel_size * cin - scale = np.sqrt(2.0 / fan_in) - if idx == len(channels) - 2: - scale *= 1e-2 - kernel = ( - jax.random.normal( - layer_key, - (kernel_size, kernel_size, cin, cout), - dtype=jnp.float32, - ) - * scale - ) - layers.append( - { - "kernel": kernel, - "bias": jnp.zeros((cout,), dtype=jnp.float32), - } - ) - return tuple(layers) - - -def _periodic_conv(x: jax.Array, layer: dict[str, jax.Array]) -> jax.Array: - """Apply one stride-one NHWC convolution with periodic padding.""" - kernel = layer["kernel"] - pad = kernel.shape[0] // 2 - padded = jnp.pad(x, ((pad, pad), (pad, pad), (0, 0)), mode="wrap") - y = jax.lax.conv_general_dilated( - padded[jnp.newaxis, ...], - kernel, - window_strides=(1, 1), - padding="VALID", - dimension_numbers=("NHWC", "HWIO", "NHWC"), - )[0] - return y + layer["bias"] + if architecture != "periodic_residual_cnn": + raise ValueError(f"unknown corrector architecture: {architecture!r}") + return PeriodicResidualCNN( + key, + hidden_channels=hidden_channels, + kernel_size=kernel_size, + ) def apply_corrector( - params: Sequence[dict[str, jax.Array]], + model: PeriodicResidualCNN, velocity: jax.Array, *, velocity_scale: float | jax.Array, ) -> jax.Array: """Predict an additive velocity correction on ``(N, N, 1, 2)`` fields.""" x = velocity[:, :, 0, :] / jnp.asarray(velocity_scale, dtype=velocity.dtype) - for layer in params[:-1]: - x = jax.nn.gelu(_periodic_conv(x, layer)) - delta = _periodic_conv(x, params[-1]) + delta = jnp.moveaxis(model(jnp.moveaxis(x, -1, 0)), 0, -1) return delta[:, :, jnp.newaxis, :] * jnp.asarray( velocity_scale, dtype=velocity.dtype ) @@ -128,14 +153,14 @@ def project_periodic_correction(delta: jax.Array, domain_extent: float) -> jax.A def corrected_velocity( - params: Sequence[dict[str, jax.Array]], + model: PeriodicResidualCNN, provisional: jax.Array, *, velocity_scale: float | jax.Array, domain_extent: float, ) -> jax.Array: """Apply the learned residual and enforce a physical correction subspace.""" - delta = apply_corrector(params, provisional, velocity_scale=velocity_scale) + delta = apply_corrector(model, provisional, velocity_scale=velocity_scale) return provisional + project_periodic_correction(delta, domain_extent) @@ -159,6 +184,46 @@ def divergence_rms(velocity: Any, domain_extent: float) -> float: return float(np.sqrt(np.mean(np.fft.ifft2(div_hat).real ** 2))) +def centered_divergence_rms(velocity: Any, domain_extent: float) -> float: + """Centered-difference RMS divergence for one canonical velocity field.""" + field = np.asarray(velocity) + if field.ndim == 4: + field = field[:, :, 0, :] + dx = domain_extent / field.shape[0] + dy = domain_extent / field.shape[1] + du_dx = (np.roll(field[..., 0], -1, axis=0) - np.roll(field[..., 0], 1, axis=0)) / ( + 2.0 * dx + ) + dv_dy = (np.roll(field[..., 1], -1, axis=1) - np.roll(field[..., 1], 1, axis=1)) / ( + 2.0 * dy + ) + return float(np.sqrt(np.mean((du_dx + dv_dy) ** 2))) + + +def kinetic_energy(velocity: Any) -> float: + """Mean kinetic energy density for one canonical velocity field.""" + field = np.asarray(velocity) + if field.ndim == 4: + field = field[:, :, 0, :] + return float(0.5 * np.mean(np.sum(field**2, axis=-1))) + + +def enstrophy(velocity: Any, domain_extent: float) -> float: + """Mean enstrophy density using centered periodic differences.""" + field = np.asarray(velocity) + if field.ndim == 4: + field = field[:, :, 0, :] + dx = domain_extent / field.shape[0] + dy = domain_extent / field.shape[1] + dv_dx = (np.roll(field[..., 1], -1, axis=0) - np.roll(field[..., 1], 1, axis=0)) / ( + 2.0 * dx + ) + du_dy = (np.roll(field[..., 0], -1, axis=1) - np.roll(field[..., 0], 1, axis=1)) / ( + 2.0 * dy + ) + return float(0.5 * np.mean((dv_dx - du_dy) ** 2)) + + def spectral_restrict(velocity: Any, target_n: int) -> np.ndarray: """Fourier-restrict a periodic velocity field to ``target_n`` cells.""" field = np.asarray(velocity) @@ -284,10 +349,14 @@ def reference_trajectory( __all__ = [ + "PeriodicResidualCNN", "apply_corrector", + "centered_divergence_rms", "corrected_velocity", "divergence_rms", + "enstrophy", "init_corrector", + "kinetic_energy", "project_periodic_correction", "reference_trajectory", "relative_l2", diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 720f37e1..bbc35fd7 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -12,6 +12,7 @@ import matplotlib.gridspec as gridspec import matplotlib.lines as mlines import matplotlib.pyplot as plt +import matplotlib.ticker as mticker import numpy as np from mosaic.benchmarks.core.config import Problem @@ -44,11 +45,50 @@ _DRAG_OPT_SOLVER_ORDER = ["xlb", "phiflow", "pict"] +def _solver_loop_legend( + fig: plt.Figure, + names: list[str], + *, + extra_handles: list[Any] | None = None, +) -> None: + """Identify solvers by colour/marker while reserving line style for semantics.""" + by_alias = { + alias: name + for name in names + if (alias := resolve_solver_alias(name)) is not None + } + handles: list[Any] = [] + for alias in NS_ORDER: + name = by_alias.get(alias) + if name is None: + continue + label, color, _linestyle, marker = solver_props(name) + handles.append( + mlines.Line2D( + [], + [], + color=color, + marker=marker, + linestyle="-", + label=label, + ) + ) + handles.extend(extra_handles or []) + if handles: + fig.legend( + handles=handles, + loc="outside lower center", + ncol=min(len(handles), 6), + handlelength=2.0, + ) + + def plot_solver_in_loop( cfg: Problem, *, save: bool = True, suffix: str = "", + exp_key: str = "solver_in_loop", **_kw: Any, ) -> list: """Plot corrector trainability, rollout quality, and time-to-quality.""" @@ -56,7 +96,7 @@ def plot_solver_in_loop( results_dir(), cfg.name, "optimization", - f"solver_in_loop{suffix}", + f"{exp_key}{suffix}", ) result_path = out_dir / "result.json" fields_path = out_dir / "corrector_fields.npz" @@ -74,30 +114,74 @@ def plot_solver_in_loop( fig, axes = paper_row(3, squeeze=False) ax_loss, ax_roll, ax_cost = np.atleast_1d(axes).ravel() present: list[str] = [] + has_stopped_rollout = False for idx, name in enumerate(names): metrics = data.get("by_solver", {}).get(name, {}) - label, color, linestyle, marker = solver_props(name) + label, color, _linestyle, marker = solver_props(name) loss = np.asarray(arrays.get(f"loss_{idx}", np.array([]))) + loss_std = np.asarray(arrays.get(f"loss_seed_std_{idx}", np.array([]))) + stopped_loss = np.asarray(arrays.get(f"loss_stop_gradient_{idx}", np.array([]))) corrected = np.asarray(arrays.get(f"error_corrected_{idx}", np.array([]))) + corrected_std = np.asarray( + arrays.get(f"error_corrected_seed_std_{idx}", np.array([])) + ) + stopped = np.asarray(arrays.get(f"error_stop_gradient_{idx}", np.array([]))) uncorrected = np.asarray(arrays.get(f"error_uncorrected_{idx}", np.array([]))) if loss.size: + updates = np.arange(1, loss.size + 1) ax_loss.plot( - np.arange(1, loss.size + 1), + updates, loss, color=color, - linestyle=linestyle, + linestyle="-", label=label, ) + if loss_std.shape == loss.shape: + ax_loss.fill_between( + updates, + np.maximum(loss - loss_std, 0.2 * loss), + loss + loss_std, + color=color, + alpha=0.12, + linewidth=0, + ) + if stopped_loss.size: + ax_loss.plot( + np.arange(1, stopped_loss.size + 1), + stopped_loss, + color=color, + linestyle="--", + alpha=0.55, + ) if corrected.size: x = times[: corrected.size] if times.size else np.arange(corrected.size) ax_roll.plot( x[1:], corrected[1:], color=color, - linestyle=linestyle, + linestyle="-", label=label, ) + if corrected_std.shape == corrected.shape: + ax_roll.fill_between( + x[1:], + np.maximum(corrected[1:] - corrected_std[1:], 1e-12), + corrected[1:] + corrected_std[1:], + color=color, + alpha=0.12, + linewidth=0, + ) + if stopped.size: + has_stopped_rollout = True + x = times[: stopped.size] if times.size else np.arange(stopped.size) + ax_roll.plot( + x[1:], + stopped[1:], + color=color, + linestyle="--", + alpha=0.55, + ) if uncorrected.size: x = times[: uncorrected.size] if times.size else np.arange(uncorrected.size) ax_roll.plot( @@ -136,24 +220,24 @@ def plot_solver_in_loop( ax_roll.set_title("Held-out rollout") ax_roll.text( 0.98, - 0.96, - "solid: corrected\n dotted: uncorrected", + 0.04, + ( + "solid: full VJP\n dashed: stop-gradient\n dotted: solver only" + if has_stopped_rollout + else "solid: solver + corrector\n dotted: solver only" + ), transform=ax_roll.transAxes, ha="right", - va="top", + va="bottom", fontsize=6.5, + bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.75}, ) - ax_cost.set_xscale("log") - ax_cost.set_yscale("log") ax_cost.set_xlabel("Training wall time [s]") ax_cost.set_ylabel("Final rollout error") ax_cost.set_title("Time to quality") - aliases = [ - alias for name in present if (alias := resolve_solver_alias(name)) is not None - ] - solver_legend(fig, aliases, order=NS_ORDER) + _solver_loop_legend(fig, present) fig.suptitle("Solver-in-the-loop neural correction") figs = [fig] if save: @@ -161,9 +245,34 @@ def plot_solver_in_loop( fields_fig = _plot_solver_in_loop_fields(arrays, names, out_dir, save=save) if fields_fig is not None: figs.append(fields_fig) + fairness_fig = _plot_solver_in_loop_fairness(data, names, out_dir, save=save) + if fairness_fig is not None: + figs.append(fairness_fig) + physics_fig = _plot_solver_in_loop_physics(arrays, names, out_dir, save=save) + if physics_fig is not None: + figs.append(physics_fig) + if save: + _save_solver_in_loop_animation(arrays, names, out_dir) return figs +def plot_solver_in_loop_tgv( + cfg: Problem, + *, + save: bool = True, + suffix: str = "", + **kwargs: Any, +) -> list: + """Plot the analytic Taylor--Green solver-in-the-loop sanity regime.""" + return plot_solver_in_loop( + cfg, + save=save, + suffix=suffix, + exp_key="solver_in_loop_tgv", + **kwargs, + ) + + def _periodic_vorticity_2d( velocity: np.ndarray, *, @@ -178,6 +287,78 @@ def _periodic_vorticity_2d( return d_uy_dx - d_ux_dy +def _periodic_divergence_2d( + velocity: np.ndarray, + *, + domain_extent: float = 2.0 * np.pi, + spectral: bool, +) -> np.ndarray: + """Return a common spectral or centered periodic divergence diagnostic.""" + ux, uy = _vel_components_2d(np.asarray(velocity)) + dx = domain_extent / ux.shape[0] + dy = domain_extent / ux.shape[1] + if spectral: + kx = 2.0 * np.pi * np.fft.fftfreq(ux.shape[0], d=dx) + ky = 2.0 * np.pi * np.fft.fftfreq(ux.shape[1], d=dy) + kx_grid, ky_grid = np.meshgrid(kx, ky, indexing="ij") + divergence_hat = 1j * (kx_grid * np.fft.fft2(ux) + ky_grid * np.fft.fft2(uy)) + return np.fft.ifft2(divergence_hat).real + du_dx = (np.roll(ux, -1, axis=0) - np.roll(ux, 1, axis=0)) / (2.0 * dx) + dv_dy = (np.roll(uy, -1, axis=1) - np.roll(uy, 1, axis=1)) / (2.0 * dy) + return du_dx + dv_dy + + +def _trajectory_diagnostics( + trajectory: np.ndarray, + *, + domain_extent: float = 2.0 * np.pi, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return energy, enstrophy, and two divergence curves for a trajectory.""" + energy: list[float] = [] + enstrophy_values: list[float] = [] + spectral_divergence: list[float] = [] + centered_divergence: list[float] = [] + for frame in np.asarray(trajectory): + ux, uy = _vel_components_2d(frame) + vorticity = _periodic_vorticity_2d(frame, domain_extent=domain_extent) + energy.append(float(0.5 * np.mean(ux**2 + uy**2))) + enstrophy_values.append(float(0.5 * np.mean(vorticity**2))) + spectral_divergence.append( + float( + np.sqrt( + np.mean( + _periodic_divergence_2d( + frame, + domain_extent=domain_extent, + spectral=True, + ) + ** 2 + ) + ) + ) + ) + centered_divergence.append( + float( + np.sqrt( + np.mean( + _periodic_divergence_2d( + frame, + domain_extent=domain_extent, + spectral=False, + ) + ** 2 + ) + ) + ) + ) + return ( + np.asarray(energy), + np.asarray(enstrophy_values), + np.asarray(spectral_divergence), + np.asarray(centered_divergence), + ) + + def _plot_solver_in_loop_fields( arrays: dict[str, np.ndarray], names: list[str], @@ -195,29 +376,17 @@ def _plot_solver_in_loop_fields( if reference.size == 0 or reference.shape[0] == 0: return None - rows: list[tuple[str, np.ndarray, np.ndarray]] = [] - for idx, name in enumerate(names): - raw = np.asarray(arrays.get(f"rollout_uncorrected_{idx}", np.array([]))) - corrected = np.asarray(arrays.get(f"rollout_corrected_{idx}", np.array([]))) - if raw.size and corrected.size and raw.shape[0] and corrected.shape[0]: - rows.append((name, raw[-1], corrected[-1])) + rows = _ordered_solver_rollouts(arrays, names) if not rows: return None - solver_order = {alias: idx for idx, alias in enumerate(NS_ORDER)} - rows.sort( - key=lambda row: solver_order.get( - resolve_solver_alias(row[0]) or row[0], - len(solver_order), - ) - ) reference_final = reference[-1] reference_vorticity = _periodic_vorticity_2d(reference_final) rendered_rows: list[tuple[str, np.ndarray, np.ndarray]] = [] scale_fields = [reference_vorticity] for name, raw, corrected in rows: - raw_vorticity = _periodic_vorticity_2d(raw) - corrected_vorticity = _periodic_vorticity_2d(corrected) + raw_vorticity = _periodic_vorticity_2d(raw[-1]) + corrected_vorticity = _periodic_vorticity_2d(corrected[-1]) rendered_rows.append((name, raw_vorticity, corrected_vorticity)) scale_fields.extend((raw_vorticity, corrected_vorticity)) @@ -278,6 +447,478 @@ def _plot_solver_in_loop_fields( return fig +def _ordered_solver_rollouts( + arrays: dict[str, np.ndarray], + names: list[str], +) -> list[tuple[str, np.ndarray, np.ndarray]]: + """Return complete raw/corrected rollouts in canonical solver order.""" + rows: list[tuple[str, np.ndarray, np.ndarray]] = [] + for idx, name in enumerate(names): + raw = np.asarray(arrays.get(f"rollout_uncorrected_{idx}", np.array([]))) + corrected = np.asarray(arrays.get(f"rollout_corrected_{idx}", np.array([]))) + if raw.size and corrected.size and raw.shape[0] and corrected.shape[0]: + rows.append((name, raw, corrected)) + solver_order = {alias: idx for idx, alias in enumerate(NS_ORDER)} + rows.sort( + key=lambda row: solver_order.get( + resolve_solver_alias(row[0]) or row[0], + len(solver_order), + ) + ) + return rows + + +def _plot_solver_in_loop_physics( + arrays: dict[str, np.ndarray], + names: list[str], + out_dir: Path, + *, + save: bool, +) -> plt.Figure | None: + """Compare energy, enstrophy, and divergence along held-out rollouts.""" + reference = np.asarray(arrays.get("reference_rollout", np.array([]))) + rows = _ordered_solver_rollouts(arrays, names) + if reference.size == 0 or not rows: + return None + + n_frames = min( + reference.shape[0], + *(min(raw.shape[0], corrected.shape[0]) for _, raw, corrected in rows), + ) + times = np.asarray(arrays.get("evaluation_times", np.array([])))[:n_frames] + if times.size != n_frames: + times = np.arange(n_frames, dtype=float) + reference_diagnostics = _trajectory_diagnostics(reference[:n_frames]) + energy_scale = reference_diagnostics[0][0] + 1e-12 + enstrophy_scale = reference_diagnostics[1][0] + 1e-12 + + plt.rcParams.update(RCPARAMS) + fig, axes = plt.subplots( + 2, + 2, + figsize=(TEXTWIDTH, 4.5), + squeeze=False, + layout="constrained", + ) + ax_energy, ax_enstrophy, ax_spectral, ax_centered = axes.ravel() + diagnostic_axes = (ax_energy, ax_enstrophy, ax_spectral, ax_centered) + reference_curves = ( + reference_diagnostics[0] / energy_scale, + reference_diagnostics[1] / enstrophy_scale, + reference_diagnostics[2], + reference_diagnostics[3], + ) + for ax, curve in zip(diagnostic_axes, reference_curves, strict=True): + ax.plot(times, np.maximum(curve, 1e-12), color="black", linewidth=1.25) + + present: list[str] = [] + for name, raw, corrected in rows: + _label, color, _linestyle, _marker = solver_props(name) + raw_diagnostics = _trajectory_diagnostics(raw[:n_frames]) + corrected_diagnostics = _trajectory_diagnostics(corrected[:n_frames]) + raw_curves = ( + raw_diagnostics[0] / energy_scale, + raw_diagnostics[1] / enstrophy_scale, + raw_diagnostics[2], + raw_diagnostics[3], + ) + corrected_curves = ( + corrected_diagnostics[0] / energy_scale, + corrected_diagnostics[1] / enstrophy_scale, + corrected_diagnostics[2], + corrected_diagnostics[3], + ) + for ax, raw_curve, corrected_curve in zip( + diagnostic_axes, + raw_curves, + corrected_curves, + strict=True, + ): + ax.plot(times, np.maximum(raw_curve, 1e-12), color=color, linestyle=":") + ax.plot( + times, + np.maximum(corrected_curve, 1e-12), + color=color, + linestyle="-", + ) + present.append(name) + + ax_energy.set_ylabel("Energy / initial reference") + ax_energy.set_title("Kinetic energy") + ax_enstrophy.set_ylabel("Enstrophy / initial reference") + ax_enstrophy.set_title("Resolved enstrophy") + ax_spectral.set_ylabel("RMS divergence") + ax_spectral.set_title("Common spectral divergence") + ax_centered.set_ylabel("RMS divergence") + ax_centered.set_title("Common centered-grid divergence") + for ax in diagnostic_axes: + ax.set_xlabel("Physical time") + ax_spectral.set_yscale("log") + ax_centered.set_yscale("log") + + _solver_loop_legend( + fig, + present, + extra_handles=[ + mlines.Line2D([], [], color="black", linewidth=1.25, label="reference"), + mlines.Line2D([], [], color="0.4", linestyle=":", label="solver only"), + mlines.Line2D([], [], color="0.4", linestyle="-", label="full corrector"), + ], + ) + fig.suptitle("Held-out trajectory physics (divergence is operator-dependent)") + if save: + save_fig(fig, "solver_in_loop_physics", out_dir) + return fig + + +def _plot_solver_in_loop_fairness( + data: dict[str, Any], + names: list[str], + out_dir: Path, + *, + save: bool, +) -> plt.Figure | None: + """Separate raw quality, correctability, and benefit from the solver VJP.""" + by_solver = data.get("by_solver", {}) + rows: list[tuple[str, dict[str, Any]]] = [] + solver_order = {alias: idx for idx, alias in enumerate(NS_ORDER)} + for name in names: + metrics = by_solver.get(name, {}) + if ( + metrics.get("uncorrected_mean_rollout_error") is not None + and metrics.get("mean_rollout_error") is not None + ): + rows.append((name, metrics)) + rows.sort( + key=lambda row: solver_order.get( + resolve_solver_alias(row[0]) or row[0], + len(solver_order), + ) + ) + if not rows: + return None + + plt.rcParams.update(RCPARAMS) + fig, axes = plt.subplots( + 2, + 2, + figsize=(TEXTWIDTH, 4.5), + squeeze=False, + layout="constrained", + ) + ax_restart, ax_quality, ax_gain, ax_vjp = axes.ravel() + positions = np.arange(len(rows), dtype=float) + labels: list[str] = [] + + all_errors: list[float] = [] + full_gain: list[float] = [] + stopped_gain: list[float] = [] + vjp_lift: list[float] = [] + full_gain_log_std: list[float] = [] + stopped_gain_log_std: list[float] = [] + vjp_lift_log_std: list[float] = [] + has_counterfactual = True + has_native = True + has_seed_uncertainty = False + for name, metrics in rows: + label, color, _linestyle, marker = solver_props(name) + labels.append(label) + raw = float(metrics["uncorrected_mean_rollout_error"]) + corrected = float(metrics["mean_rollout_error"]) + all_errors.extend((raw, corrected)) + ax_quality.scatter(raw, corrected, color=color, marker=marker, s=30) + native = metrics.get("native_final_rollout_error") + restarted = metrics.get("uncorrected_rollout_error") + if native is None or restarted is None: + has_native = False + else: + ax_restart.scatter( + float(native), + float(restarted), + color=color, + marker=marker, + s=30, + ) + + full = metrics.get("geometric_error_reduction") + if full is None: + full = raw / max(corrected, 1e-12) + full_gain.append(float(full)) + full_gain_log_std.append(float(metrics.get("rollout_log_gain_seed_std", 0.0))) + has_seed_uncertainty |= "rollout_log_gain_seed_std" in metrics + + stopped = metrics.get("stop_gradient_geometric_error_reduction") + lift = metrics.get("solver_vjp_geometric_lift") + if stopped is None or lift is None: + has_counterfactual = False + stopped_gain.append(float(stopped) if stopped is not None else np.nan) + vjp_lift.append(float(lift) if lift is not None else np.nan) + stopped_gain_log_std.append( + float(metrics.get("stop_gradient_rollout_log_gain_seed_std", 0.0)) + ) + vjp_lift_log_std.append(float(metrics.get("solver_vjp_log_lift_seed_std", 0.0))) + + error_min = max(min(all_errors) * 0.8, 1e-4) + error_max = max(all_errors) * 1.25 + ax_quality.plot( + [error_min, error_max], + [error_min, error_max], + color="0.45", + linestyle=":", + linewidth=1.0, + ) + if error_max / error_min >= 10.0: + ax_quality.set_xscale("log") + ax_quality.set_yscale("log") + ax_quality.xaxis.set_minor_formatter(mticker.NullFormatter()) + ax_quality.yaxis.set_minor_formatter(mticker.NullFormatter()) + ax_quality.set_xlim(error_min, error_max) + ax_quality.set_ylim(error_min, error_max) + ax_quality.set_xlabel("Solver-only mean error") + ax_quality.set_ylabel("Corrected mean error") + ax_quality.set_title("Absolute quality") + + if has_native: + native_values = [ + float(metrics["native_final_rollout_error"]) for _name, metrics in rows + ] + restarted_values = [ + float(metrics["uncorrected_rollout_error"]) for _name, metrics in rows + ] + restart_min = max(min(native_values + restarted_values) * 0.8, 1e-4) + restart_max = max(native_values + restarted_values) * 1.25 + ax_restart.plot( + [restart_min, restart_max], + [restart_min, restart_max], + color="0.45", + linestyle=":", + linewidth=1.0, + ) + if restart_max / restart_min >= 10.0: + ax_restart.set_xscale("log") + ax_restart.set_yscale("log") + ax_restart.xaxis.set_minor_formatter(mticker.NullFormatter()) + ax_restart.yaxis.set_minor_formatter(mticker.NullFormatter()) + ax_restart.set_xlim(restart_min, restart_max) + ax_restart.set_ylim(restart_min, restart_max) + ax_restart.set_xlabel("Native single-call final error") + ax_restart.set_ylabel("Restarted final error") + ax_restart.set_title("Coupling / restart penalty") + else: + ax_restart.axis("off") + ax_restart.text( + 0.5, + 0.5, + "Native single-call baseline\nnot available in this run", + ha="center", + va="center", + transform=ax_restart.transAxes, + ) + + width = 0.36 if has_counterfactual else 0.62 + colors = [solver_props(name)[1] for name, _metrics in rows] + full_error = _log_scale_errorbars(full_gain, full_gain_log_std) + ax_gain.bar( + positions - (width / 2 if has_counterfactual else 0.0), + full_gain, + width=width, + color=colors, + alpha=0.9, + label="full solver VJP", + yerr=full_error, + capsize=2, + ) + if has_counterfactual: + ax_gain.bar( + positions + width / 2, + stopped_gain, + width=width, + color=colors, + alpha=0.35, + hatch="//", + label="stop-gradient", + yerr=_log_scale_errorbars(stopped_gain, stopped_gain_log_std), + capsize=2, + ) + ax_gain.legend(loc="best", fontsize=6.5) + ax_gain.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_gain.set_xticks(positions, labels, rotation=35, ha="right") + ax_gain.set_ylabel("Geometric error reduction [Γ—]") + ax_gain.set_title("Correctability") + + if has_counterfactual: + vjp_lift_pct = [100.0 * (value - 1.0) for value in vjp_lift] + vjp_error = _log_scale_errorbars(vjp_lift, vjp_lift_log_std) + if vjp_error is not None: + vjp_error = 100.0 * vjp_error + ax_vjp.bar( + positions, + vjp_lift_pct, + color=colors, + alpha=0.9, + yerr=vjp_error, + capsize=2, + ) + ax_vjp.axhline(0.0, color="0.45", linestyle=":", linewidth=1.0) + ax_vjp.set_xticks(positions, labels, rotation=35, ha="right") + ax_vjp.set_ylabel("Solver-VJP lift [%]") + ax_vjp.set_title("Benefit from solver VJP") + else: + ax_vjp.axis("off") + ax_vjp.text( + 0.5, + 0.5, + "VJP counterfactual\nnot available in this run", + ha="center", + va="center", + transform=ax_vjp.transAxes, + ) + + _solver_loop_legend(fig, [name for name, _metrics in rows]) + uncertainty_suffix = " (error bars: Β±1 seed SD)" if has_seed_uncertainty else "" + fig.suptitle(f"Fair solver-in-the-loop decomposition{uncertainty_suffix}") + if save: + save_fig(fig, "solver_in_loop_fairness", out_dir) + return fig + + +def _log_scale_errorbars( + values: list[float], + log_standard_deviations: list[float], +) -> np.ndarray | None: + """Convert symmetric log-space uncertainty to asymmetric linear bars.""" + centers = np.asarray(values, dtype=float) + deviations = np.asarray(log_standard_deviations, dtype=float) + if not np.any(deviations > 0): + return None + lower = centers - centers * np.exp(-deviations) + upper = centers * np.exp(deviations) - centers + return np.stack((lower, upper)) + + +def _save_solver_in_loop_animation( + arrays: dict[str, np.ndarray], + names: list[str], + out_dir: Path, +) -> None: + """Animate reference, solver-only, and fully corrected held-out rollouts.""" + reference = np.asarray(arrays.get("reference_rollout", np.array([]))) + rows = _ordered_solver_rollouts(arrays, names) + if reference.size == 0 or not rows: + return + + available_frames = min( + reference.shape[0], + *(min(raw.shape[0], corrected.shape[0]) for _, raw, corrected in rows), + ) + # Keep the full simulated horizon in the animation while bounding the GIF + # payload embedded in PRs. Evenly spaced samples retain both endpoints. + frame_indices = np.unique( + np.linspace( + 0, + available_frames - 1, + num=min(available_frames, 24), + dtype=int, + ) + ) + n_frames = len(frame_indices) + reference_vorticity = [ + _periodic_vorticity_2d(reference[index]) for index in frame_indices + ] + rollout_vorticity = [ + ( + name, + [_periodic_vorticity_2d(raw[index]) for index in frame_indices], + [_periodic_vorticity_2d(corrected[index]) for index in frame_indices], + ) + for name, raw, corrected in rows + ] + scale_fields = list(reference_vorticity) + for _name, raw_fields, corrected_fields in rollout_vorticity: + scale_fields.extend(raw_fields) + scale_fields.extend(corrected_fields) + magnitudes = np.concatenate([np.abs(field).ravel() for field in scale_fields]) + vmax = float(np.percentile(magnitudes, 99.0)) or 1.0 + + plt.rcParams.update(RCPARAMS) + fig, axes = plt.subplots( + len(rows), + 3, + figsize=(TEXTWIDTH, max(1.8, 0.95 * len(rows))), + dpi=100, + squeeze=False, + layout="constrained", + ) + images: list[tuple[Any, int, int]] = [] + for row_idx, (name, raw_fields, corrected_fields) in enumerate(rollout_vorticity): + label, _color, _linestyle, _marker = solver_props(name) + for col_idx, fields in enumerate( + (reference_vorticity, raw_fields, corrected_fields) + ): + ax = axes[row_idx, col_idx] + image = ax.imshow( + fields[0].T, + origin="lower", + cmap="RdBu_r", + vmin=-vmax, + vmax=vmax, + interpolation="nearest", + animated=True, + ) + images.append((image, row_idx, col_idx)) + ax.set_xticks([]) + ax.set_yticks([]) + if row_idx == 0: + ax.set_title( + ("Reference", "Solver only", "Solver + corrector")[col_idx] + ) + axes[row_idx, 0].set_ylabel( + label, + rotation=0, + ha="right", + va="center", + labelpad=8, + ) + + colorbar = fig.colorbar( + images[-1][0], + ax=axes.ravel().tolist(), + location="right", + shrink=0.82, + pad=0.02, + ) + colorbar.set_label(r"Vorticity $\omega$") + all_times = np.asarray(arrays.get("evaluation_times", np.array([]))) + times = ( + all_times[frame_indices] if all_times.size >= available_frames else np.array([]) + ) + title = fig.suptitle("Held-out trajectory at $t=0$") + + def _update(frame: int) -> tuple[Any, ...]: + for image, row_idx, col_idx in images: + if col_idx == 0: + field = reference_vorticity[frame] + elif col_idx == 1: + field = rollout_vorticity[row_idx][1][frame] + else: + field = rollout_vorticity[row_idx][2][frame] + image.set_data(field.T) + time_value = float(times[frame]) if frame < times.size else float(frame) + title.set_text(f"Held-out trajectory at $t={time_value:g}$") + artists = [image for image, _row, _col in images] + artists.append(title) + return tuple(artists) + + animation = manimation.FuncAnimation( + fig, + _update, + frames=n_frames, + interval=180, + blit=False, + ) + _save_animation(animation, "solver_in_loop_trajectory", out_dir, fps=6) + + def plot_drag_opt( cfg: Problem, *, diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py index 92415f31..90b0489d 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py @@ -18,6 +18,7 @@ from functools import lru_cache, partial from typing import Any +import equinox as eqx import jax import jax.numpy as jnp import numpy as np @@ -28,14 +29,17 @@ from mosaic.benchmarks.core.utils import active_differentiable_solvers from .corrector import ( + centered_divergence_rms, corrected_velocity, divergence_rms, + enstrophy, init_corrector, + kinetic_energy, reference_trajectory, relative_l2, spectral_restrict, ) -from .ics import _multimode +from .ics import _multimode, _tgv _DATASET_LOCK = threading.Lock() @@ -49,6 +53,7 @@ def _dataset_digest(config: dict[str, Any]) -> str: @lru_cache(maxsize=8) def _cached_reference_dataset( *, + reference_kind: str, n: int, reference_factor: int, viscosity: float, @@ -66,23 +71,46 @@ def _cached_reference_dataset( trajectories: list[np.ndarray] = [] reference_n = n * reference_factor for seed in seeds: - fine_initial = _multimode( - reference_n, - L=domain_extent, - seed=seed, - k0=k0, - sigma_k=sigma_k, - amplitude=amplitude, - ) - fine = reference_trajectory( - fine_initial, - viscosity=viscosity, - dt=dt, - frame_steps=frame_steps, - n_frames=n_frames, - substeps=reference_substeps, - domain_extent=domain_extent, - ) + if reference_kind == "pseudo_spectral_multimode": + fine_initial = _multimode( + reference_n, + L=domain_extent, + seed=seed, + k0=k0, + sigma_k=sigma_k, + amplitude=amplitude, + ) + fine = reference_trajectory( + fine_initial, + viscosity=viscosity, + dt=dt, + frame_steps=frame_steps, + n_frames=n_frames, + substeps=reference_substeps, + domain_extent=domain_extent, + ) + elif reference_kind == "analytic_tgv": + fine_initial = np.asarray( + _tgv(reference_n, L=domain_extent), + dtype=np.float32, + ) + # Integer translations preserve the exact periodic TGV solution + # while giving train/test trajectories distinct phases. + shift = ( + int(seed) % reference_n, + (7 * int(seed) + 3) % reference_n, + ) + fine_initial = np.roll(fine_initial, shift=shift, axis=(0, 1)) + frame_dt = dt * frame_steps + decay_rate = 2.0 * viscosity * (2.0 * np.pi / domain_extent) ** 2 + fine = np.stack( + [ + fine_initial * np.exp(-decay_rate * frame_dt * frame) + for frame in range(n_frames + 1) + ] + ).astype(np.float32) + else: + raise ValueError(f"unknown reference_kind: {reference_kind!r}") trajectories.append(np.stack([spectral_restrict(frame, n) for frame in fine])) return np.stack(trajectories).astype(np.float32) @@ -104,6 +132,9 @@ def make_reference_dataset( unroll = int(training.get("unroll", 4)) n_frames = max(train_frames, eval_frames, unroll) config = { + "reference_kind": str( + dataset.get("reference_kind", "pseudo_spectral_multimode") + ), "n": n, "reference_factor": int(dataset.get("reference_factor", 2)), "viscosity": float(physics["nu"]), @@ -147,21 +178,27 @@ def _solver_advance( def _window_loss( - params: Any, + model: Any, targets: jax.Array, *, t: Any, ctx: KernelContext, frame_steps: int, velocity_scale: float, + differentiate_solver: bool, ) -> jax.Array: """Mean normalized state error over one recurrent training window.""" state = targets[0] losses: list[jax.Array] = [] for target in targets[1:]: provisional = _solver_advance(t, ctx, state, frame_steps=frame_steps) + if not differentiate_solver: + # Keep the identical recurrent forward trajectory but cut the + # backward graph at every solver transition. This counterfactual + # measures what the same corrector can learn without the solver VJP. + provisional = jax.lax.stop_gradient(provisional) state = corrected_velocity( - params, + model, provisional, velocity_scale=velocity_scale, domain_extent=ctx.domain_extent, @@ -172,27 +209,46 @@ def _window_loss( def _directional_fd( loss_fn: Any, - params: Any, + model: Any, grads: Any, key: jax.Array, *, epsilon: float, ) -> float: """Relative error of one end-to-end directional finite difference.""" - direction = jax.tree_util.tree_map( - lambda p, k: jax.random.normal(k, p.shape, dtype=p.dtype), - params, - jax.tree_util.tree_unflatten( - jax.tree_util.tree_structure(params), - jax.random.split(key, len(jax.tree_util.tree_leaves(params))), - ), + dynamic, static = eqx.partition(model, eqx.is_inexact_array) + leaves, structure = jax.tree_util.tree_flatten(dynamic) + direction = jax.tree_util.tree_unflatten( + structure, + [ + jax.random.normal(part_key, leaf.shape, dtype=leaf.dtype) + for leaf, part_key in zip( + leaves, + jax.random.split(key, len(leaves)), + strict=True, + ) + ], ) direction_norm = optax.tree.norm(direction) direction = jax.tree_util.tree_map( lambda value: value / (direction_norm + 1e-12), direction ) - plus = jax.tree_util.tree_map(lambda p, d: p + epsilon * d, params, direction) - minus = jax.tree_util.tree_map(lambda p, d: p - epsilon * d, params, direction) + plus = eqx.combine( + jax.tree_util.tree_map( + lambda parameter, value: parameter + epsilon * value, + dynamic, + direction, + ), + static, + ) + minus = eqx.combine( + jax.tree_util.tree_map( + lambda parameter, value: parameter - epsilon * value, + dynamic, + direction, + ), + static, + ) fd = (loss_fn(plus) - loss_fn(minus)) / (2.0 * epsilon) ad = sum( jnp.vdot(g, d) @@ -213,35 +269,40 @@ def _train_corrector( frame_steps: int, training: dict[str, Any], velocity_scale: float, -) -> tuple[Any, list[float], list[float], float | None, bool]: + differentiate_solver: bool, + model_seed: int, +) -> tuple[Any, list[float], list[float], list[float], float | None, bool]: """Train one solver-specific corrector with a fixed stochastic schedule.""" max_updates = int(training.get("max_updates", 100)) unroll = int(training.get("unroll", 4)) lr = float(training.get("lr", 1e-4)) clip_norm = float(training.get("clip_norm", 1.0)) seed = int(training.get("seed", 2026)) - model_seed = int(training.get("model_seed", 0)) hidden_channels = int(training.get("hidden_channels", 32)) kernel_size = int(training.get("kernel_size", 5)) + architecture = str(training.get("architecture", "periodic_residual_cnn")) fd_epsilon = float(training.get("fd_epsilon", 1e-2)) - params = init_corrector( + model = init_corrector( jax.random.PRNGKey(model_seed), hidden_channels=hidden_channels, kernel_size=kernel_size, + architecture=architecture, ) optimiser = optax.chain( optax.clip_by_global_norm(clip_norm), optax.adam(lr), ) - opt_state = optimiser.init(params) + opt_state = optimiser.init(eqx.filter(model, eqx.is_inexact_array)) rng = np.random.RandomState(seed) losses: list[float] = [] grad_norms: list[float] = [] + update_times: list[float] = [] fd_error: float | None = None completed = True for update in range(max_updates): + update_started = time.perf_counter() trajectory_idx = int(rng.randint(train.shape[0])) max_start = train.shape[1] - unroll - 1 start = int(rng.randint(max_start + 1)) if max_start > 0 else 0 @@ -254,33 +315,43 @@ def _train_corrector( ctx=ctx, frame_steps=frame_steps, velocity_scale=velocity_scale, + differentiate_solver=differentiate_solver, ) - loss, grads = jax.value_and_grad(loss_fn)(params) + loss, grads = eqx.filter_value_and_grad(loss_fn)(model) loss_value = float(loss) grad_norm = float(optax.tree.norm(grads)) if not np.isfinite(loss_value) or not np.isfinite(grad_norm): completed = False break - if update == 0 and bool(training.get("check_grad", True)): + if ( + update == 0 + and differentiate_solver + and bool(training.get("check_grad", True)) + ): fd_error = _directional_fd( loss_fn, - params, + model, grads, jax.random.PRNGKey(model_seed + 1), epsilon=fd_epsilon, ) - updates, opt_state = optimiser.update(grads, opt_state, params) - params = optax.apply_updates(params, updates) + updates, opt_state = optimiser.update( + grads, + opt_state, + eqx.filter(model, eqx.is_inexact_array), + ) + model = eqx.apply_updates(model, updates) losses.append(loss_value) grad_norms.append(grad_norm) - return params, losses, grad_norms, fd_error, completed + update_times.append(time.perf_counter() - update_started) + return model, losses, grad_norms, update_times, fd_error, completed def _evaluate_rollout( t: Any, ctx: KernelContext, - params: Any, + model: Any, reference: np.ndarray, *, frame_steps: int, @@ -295,7 +366,7 @@ def _evaluate_rollout( state = _solver_advance(t, ctx, state, frame_steps=frame_steps) if corrected: state = corrected_velocity( - params, + model, state, velocity_scale=velocity_scale, domain_extent=ctx.domain_extent, @@ -314,16 +385,68 @@ def _first_unstable(errors: list[float], threshold: float) -> int: return len(errors) - 1 +def _rollout_log_gain( + baseline_errors: np.ndarray | list[float], + corrected_errors: np.ndarray | list[float], +) -> float: + """Return rollout-wide geometric error reduction on non-initial frames.""" + baseline = np.asarray(baseline_errors, dtype=np.float64)[1:] + corrected = np.asarray(corrected_errors, dtype=np.float64)[1:] + if baseline.shape != corrected.shape or baseline.size == 0: + raise ValueError("rollout errors must have matching non-empty shapes") + return float(np.mean(np.log((baseline + 1e-12) / (corrected + 1e-12)))) + + +def _median_steady_update_time(update_times: list[float]) -> float | None: + """Median update time after the compile/FD-heavy first optimizer update.""" + steady = update_times[1:] if len(update_times) > 1 else update_times + return float(np.median(steady)) if steady else None + + +def _mean_curve(curves: list[list[float]]) -> np.ndarray: + """Average completed prefixes of repeated optimization or rollout curves.""" + if not curves: + return np.asarray([], dtype=np.float32) + common_length = min(len(curve) for curve in curves) + if common_length == 0: + return np.asarray([], dtype=np.float32) + return np.mean( + np.stack([np.asarray(curve[:common_length]) for curve in curves]), + axis=0, + ).astype(np.float32) + + +def _mean_optional(values: list[float | None]) -> float | None: + """Average present scalar measurements from repeated model seeds.""" + present = [float(value) for value in values if value is not None] + return float(np.mean(present)) if present else None + + +def _std(values: list[float]) -> float: + """Population standard deviation, including the well-defined singleton case.""" + return float(np.std(np.asarray(values, dtype=np.float64))) + + @kernel( sweep_mode="none", selector_fn=active_differentiable_solvers, snapshot_filename="corrector_fields.npz", snapshot_prefixes=( "loss", + "loss_seed_std", + "loss_stop_gradient", + "loss_stop_gradient_seed_std", "grad_norm", + "grad_norm_stop_gradient", + "update_time", + "update_time_stop_gradient", "error_corrected", + "error_corrected_seed_std", + "error_stop_gradient", + "error_stop_gradient_seed_std", "error_uncorrected", "rollout_corrected", + "rollout_stop_gradient", "rollout_uncorrected", ), ) @@ -343,91 +466,431 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: ) velocity_scale = float(np.sqrt(np.mean(train**2)) + 1e-8) - started = time.perf_counter() - params, losses, grad_norms, fd_error, completed = _train_corrector( - t, - ctx, - train, - frame_steps=frame_steps, - training=training, - velocity_scale=velocity_scale, - ) - training_wall = time.perf_counter() - started - - corrected_errors: list[list[float]] = [] + configured_seeds = training.get("model_seeds") + if configured_seeds is None: + model_seeds = (int(training.get("model_seed", 0)),) + else: + model_seeds = tuple(int(seed) for seed in configured_seeds) + if not model_seeds: + raise ValueError("training.model_seeds must contain at least one seed") + + # Evaluate the solver itself once: these trajectories do not depend on a + # neural-model seed. A long native call and a sequence of short canonical + # restarts expose coupling penalties separately from integration accuracy. uncorrected_errors: list[list[float]] = [] - first_corrected: np.ndarray | None = None + native_final_errors: list[float] = [] first_uncorrected: np.ndarray | None = None for reference in test: - corrected_rollout, corr_err = _evaluate_rollout( + uncorrected_rollout, raw_err = _evaluate_rollout( t, ctx, - params, + None, reference, frame_steps=frame_steps, velocity_scale=velocity_scale, - corrected=True, + corrected=False, ) - uncorrected_rollout, raw_err = _evaluate_rollout( + native_final = _solver_advance( t, ctx, - params, - reference, - frame_steps=frame_steps, - velocity_scale=velocity_scale, - corrected=False, + jnp.asarray(reference[0]), + frame_steps=frame_steps * (reference.shape[0] - 1), ) - corrected_errors.append(corr_err) uncorrected_errors.append(raw_err) - if first_corrected is None: - first_corrected = corrected_rollout + native_final_errors.append(relative_l2(np.asarray(native_final), reference[-1])) + if first_uncorrected is None: first_uncorrected = uncorrected_rollout - corrected_error = np.mean(np.asarray(corrected_errors), axis=0) uncorrected_error = np.mean(np.asarray(uncorrected_errors), axis=0) + + models: list[Any] = [] + losses_by_seed: list[list[float]] = [] + grad_norms_by_seed: list[list[float]] = [] + update_times_by_seed: list[list[float]] = [] + fd_errors: list[float | None] = [] + completed_by_seed: list[bool] = [] + training_walls: list[float] = [] + stop_gradient_losses_by_seed: list[list[float]] = [] + stop_gradient_grad_norms_by_seed: list[list[float]] = [] + stop_gradient_update_times_by_seed: list[list[float]] = [] + stop_gradient_completed_by_seed: list[bool] = [] + stop_gradient_training_walls: list[float] = [] + corrected_errors_by_seed: list[np.ndarray] = [] + stop_gradient_errors_by_seed: list[np.ndarray] = [] + first_corrected: np.ndarray | None = None + first_stop_gradient: np.ndarray | None = None + + for seed_idx, model_seed in enumerate(model_seeds): + seed_training = { + **training, + # One end-to-end FD check is enough for a shared architecture and + # solver VJP; repeating it for initialisation replicates adds cost + # without strengthening the paired training comparison. + "check_grad": bool(training.get("check_grad", True)) and seed_idx == 0, + } + started = time.perf_counter() + ( + model, + losses, + grad_norms, + update_times, + fd_error, + completed, + ) = _train_corrector( + t, + ctx, + train, + frame_steps=frame_steps, + training=seed_training, + velocity_scale=velocity_scale, + differentiate_solver=True, + model_seed=model_seed, + ) + training_walls.append(time.perf_counter() - started) + + stop_gradient_started = time.perf_counter() + ( + stop_gradient_model, + stop_gradient_losses, + stop_gradient_grad_norms, + stop_gradient_update_times, + _stop_gradient_fd_error, + stop_gradient_completed, + ) = _train_corrector( + t, + ctx, + train, + frame_steps=frame_steps, + training=seed_training, + velocity_scale=velocity_scale, + differentiate_solver=False, + model_seed=model_seed, + ) + stop_gradient_training_walls.append(time.perf_counter() - stop_gradient_started) + + seed_corrected_errors: list[list[float]] = [] + seed_stop_gradient_errors: list[list[float]] = [] + for reference_idx, reference in enumerate(test): + corrected_rollout, corr_err = _evaluate_rollout( + t, + ctx, + model, + reference, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + stop_gradient_rollout, stop_gradient_err = _evaluate_rollout( + t, + ctx, + stop_gradient_model, + reference, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + seed_corrected_errors.append(corr_err) + seed_stop_gradient_errors.append(stop_gradient_err) + if seed_idx == 0 and reference_idx == 0: + first_corrected = corrected_rollout + first_stop_gradient = stop_gradient_rollout + + models.append(model) + losses_by_seed.append(losses) + grad_norms_by_seed.append(grad_norms) + update_times_by_seed.append(update_times) + fd_errors.append(fd_error) + completed_by_seed.append(completed) + stop_gradient_losses_by_seed.append(stop_gradient_losses) + stop_gradient_grad_norms_by_seed.append(stop_gradient_grad_norms) + stop_gradient_update_times_by_seed.append(stop_gradient_update_times) + stop_gradient_completed_by_seed.append(stop_gradient_completed) + corrected_errors_by_seed.append( + np.mean(np.asarray(seed_corrected_errors), axis=0) + ) + stop_gradient_errors_by_seed.append( + np.mean(np.asarray(seed_stop_gradient_errors), axis=0) + ) + + corrected_errors_array = np.stack(corrected_errors_by_seed) + stop_gradient_errors_array = np.stack(stop_gradient_errors_by_seed) + corrected_error = np.mean(corrected_errors_array, axis=0) + corrected_error_std = np.std(corrected_errors_array, axis=0) + stop_gradient_error = np.mean(stop_gradient_errors_array, axis=0) + stop_gradient_error_std = np.std(stop_gradient_errors_array, axis=0) + losses = _mean_curve(losses_by_seed) + stop_gradient_losses = _mean_curve(stop_gradient_losses_by_seed) + grad_norms = _mean_curve(grad_norms_by_seed) + stop_gradient_grad_norms = _mean_curve(stop_gradient_grad_norms_by_seed) + update_times = _mean_curve(update_times_by_seed) + stop_gradient_update_times = _mean_curve(stop_gradient_update_times_by_seed) + loss_seed_std = ( + np.std(np.stack([curve[: len(losses)] for curve in losses_by_seed]), axis=0) + if losses.size + else np.asarray([], dtype=np.float32) + ) + stop_gradient_loss_seed_std = ( + np.std( + np.stack( + [ + curve[: len(stop_gradient_losses)] + for curve in stop_gradient_losses_by_seed + ] + ), + axis=0, + ) + if stop_gradient_losses.size + else np.asarray([], dtype=np.float32) + ) + training_wall = float(np.sum(training_walls)) + stop_gradient_training_wall = float(np.sum(stop_gradient_training_walls)) + total_updates = sum(len(curve) for curve in losses_by_seed) + stop_gradient_total_updates = sum( + len(curve) for curve in stop_gradient_losses_by_seed + ) + steady_update_times = [ + value + for seed_times in update_times_by_seed + for value in (seed_times[1:] if len(seed_times) > 1 else seed_times) + ] + stop_gradient_steady_update_times = [ + value + for seed_times in stop_gradient_update_times_by_seed + for value in (seed_times[1:] if len(seed_times) > 1 else seed_times) + ] + final_corrected = float(corrected_error[-1]) + final_stop_gradient = float(stop_gradient_error[-1]) final_uncorrected = float(uncorrected_error[-1]) + native_final_error = float(np.mean(native_final_errors)) mean_corrected = float(np.mean(corrected_error[1:])) + mean_stop_gradient = float(np.mean(stop_gradient_error[1:])) mean_uncorrected = float(np.mean(uncorrected_error[1:])) + rollout_log_gains = [ + _rollout_log_gain(uncorrected_error, seed_error) + for seed_error in corrected_errors_by_seed + ] + stop_gradient_log_gains = [ + _rollout_log_gain(uncorrected_error, seed_error) + for seed_error in stop_gradient_errors_by_seed + ] + solver_vjp_log_lifts = [ + _rollout_log_gain(stopped, full) + for stopped, full in zip( + stop_gradient_errors_by_seed, + corrected_errors_by_seed, + strict=True, + ) + ] + rollout_log_gain = float(np.mean(rollout_log_gains)) + stop_gradient_log_gain = float(np.mean(stop_gradient_log_gains)) + solver_vjp_log_lift = float(np.mean(solver_vjp_log_lifts)) threshold = float(evaluation.get("stable_error_threshold", 1.0)) interval_time = float(ctx.phys["dt"]) * frame_steps + model = models[0] + parameter_count = sum( + int(leaf.size) + for leaf in jax.tree_util.tree_leaves(eqx.filter(model, eqx.is_inexact_array)) + ) + median_update_time = ( + float(np.median(steady_update_times)) if steady_update_times else None + ) + stop_gradient_median_update_time = ( + float(np.median(stop_gradient_steady_update_times)) + if stop_gradient_steady_update_times + else None + ) metrics = { "final_rollout_error": final_corrected, + "final_rollout_error_seed_std": float(corrected_error_std[-1]), + "stop_gradient_final_rollout_error": final_stop_gradient, + "stop_gradient_final_rollout_error_seed_std": float( + stop_gradient_error_std[-1] + ), "uncorrected_rollout_error": final_uncorrected, + "native_final_rollout_error": native_final_error, + "state_restart_error_ratio": final_uncorrected / (native_final_error + 1e-12), "mean_rollout_error": mean_corrected, + "stop_gradient_mean_rollout_error": mean_stop_gradient, "uncorrected_mean_rollout_error": mean_uncorrected, "improvement_pct": 100.0 * (final_uncorrected - final_corrected) / (final_uncorrected + 1e-12), + "mean_improvement_pct": 100.0 + * (mean_uncorrected - mean_corrected) + / (mean_uncorrected + 1e-12), + "rollout_log_gain": rollout_log_gain, + "rollout_log_gain_seed_std": _std(rollout_log_gains), + "geometric_error_reduction": float(np.exp(rollout_log_gain)), + "stop_gradient_rollout_log_gain": stop_gradient_log_gain, + "stop_gradient_rollout_log_gain_seed_std": _std(stop_gradient_log_gains), + "stop_gradient_geometric_error_reduction": float( + np.exp(stop_gradient_log_gain) + ), + "solver_vjp_log_lift": solver_vjp_log_lift, + "solver_vjp_log_lift_seed_std": _std(solver_vjp_log_lifts), + "solver_vjp_geometric_lift": float(np.exp(solver_vjp_log_lift)), "stable_horizon": _first_unstable(corrected_error.tolist(), threshold) * interval_time, + "stop_gradient_stable_horizon": _first_unstable( + stop_gradient_error.tolist(), threshold + ) + * interval_time, "uncorrected_stable_horizon": _first_unstable( uncorrected_error.tolist(), threshold ) * interval_time, - "final_train_loss": losses[-1] if losses else None, - "best_train_loss": min(losses) if losses else None, + "initial_train_loss": float(losses[0]) if losses.size else None, + "final_train_loss": float(losses[-1]) if losses.size else None, + "best_train_loss": float(np.min(losses)) if losses.size else None, + "train_loss_log_gain": float( + np.log((losses[0] + 1e-12) / (np.min(losses) + 1e-12)) + ) + if losses.size + else None, "n_updates": len(losses), + "total_optimizer_updates": total_updates, "training_wall_time_s": training_wall, - "seconds_per_update": training_wall / max(len(losses), 1), - "final_grad_norm": grad_norms[-1] if grad_norms else None, - "end_to_end_fd_rel_error": fd_error, + "seconds_per_update": training_wall / max(total_updates, 1), + "median_update_time_s": median_update_time, + "stop_gradient_initial_train_loss": float(stop_gradient_losses[0]) + if stop_gradient_losses.size + else None, + "stop_gradient_final_train_loss": float(stop_gradient_losses[-1]) + if stop_gradient_losses.size + else None, + "stop_gradient_best_train_loss": float(np.min(stop_gradient_losses)) + if stop_gradient_losses.size + else None, + "stop_gradient_n_updates": len(stop_gradient_losses), + "stop_gradient_total_optimizer_updates": stop_gradient_total_updates, + "stop_gradient_training_wall_time_s": stop_gradient_training_wall, + "stop_gradient_seconds_per_update": stop_gradient_training_wall + / max(stop_gradient_total_updates, 1), + "stop_gradient_median_update_time_s": stop_gradient_median_update_time, + "solver_vjp_update_overhead_ratio": median_update_time + / (stop_gradient_median_update_time + 1e-12) + if median_update_time is not None + and stop_gradient_median_update_time is not None + else None, + "final_grad_norm": float(grad_norms[-1]) if grad_norms.size else None, + "stop_gradient_final_grad_norm": float(stop_gradient_grad_norms[-1]) + if stop_gradient_grad_norms.size + else None, + "end_to_end_fd_rel_error": _mean_optional(fd_errors), "final_divergence_rms": divergence_rms(first_corrected[-1], ctx.domain_extent) if first_corrected is not None else None, - "completed": completed, + "stop_gradient_final_divergence_rms": divergence_rms( + first_stop_gradient[-1], ctx.domain_extent + ) + if first_stop_gradient is not None + else None, + "uncorrected_final_divergence_rms": divergence_rms( + first_uncorrected[-1], ctx.domain_extent + ) + if first_uncorrected is not None + else None, + "final_centered_divergence_rms": centered_divergence_rms( + first_corrected[-1], ctx.domain_extent + ) + if first_corrected is not None + else None, + "stop_gradient_final_centered_divergence_rms": centered_divergence_rms( + first_stop_gradient[-1], ctx.domain_extent + ) + if first_stop_gradient is not None + else None, + "uncorrected_final_centered_divergence_rms": centered_divergence_rms( + first_uncorrected[-1], ctx.domain_extent + ) + if first_uncorrected is not None + else None, + "final_energy_ratio_to_reference": kinetic_energy(first_corrected[-1]) + / (kinetic_energy(test[0, -1]) + 1e-12) + if first_corrected is not None + else None, + "stop_gradient_final_energy_ratio_to_reference": kinetic_energy( + first_stop_gradient[-1] + ) + / (kinetic_energy(test[0, -1]) + 1e-12) + if first_stop_gradient is not None + else None, + "uncorrected_final_energy_ratio_to_reference": kinetic_energy( + first_uncorrected[-1] + ) + / (kinetic_energy(test[0, -1]) + 1e-12) + if first_uncorrected is not None + else None, + "final_enstrophy_ratio_to_reference": enstrophy( + first_corrected[-1], ctx.domain_extent + ) + / (enstrophy(test[0, -1], ctx.domain_extent) + 1e-12) + if first_corrected is not None + else None, + "uncorrected_final_enstrophy_ratio_to_reference": enstrophy( + first_uncorrected[-1], ctx.domain_extent + ) + / (enstrophy(test[0, -1], ctx.domain_extent) + 1e-12) + if first_uncorrected is not None + else None, + "corrector_architecture": model.architecture, + "corrector_parameter_count": parameter_count, + "model_seeds": list(model_seeds), + "n_model_seeds": len(model_seeds), + "n_test_trajectories": int(test.shape[0]), + "visualization_model_seed": model_seeds[0], + "completed": all(completed_by_seed) and all(stop_gradient_completed_by_seed), "dataset_hash": dataset_hash, + "reference_kind": str( + dataset_cfg.get("reference_kind", "pseudo_spectral_multimode") + ), + "correction_intervals": int(test.shape[1] - 1), + "native_steps": frame_steps * int(test.shape[1] - 1), + "rollout_final_time": interval_time * int(test.shape[1] - 1), "state_restart": True, } snapshots = { - "loss": np.asarray(losses, dtype=np.float32), - "grad_norm": np.asarray(grad_norms, dtype=np.float32), + "loss": losses, + "loss_seed_std": np.asarray(loss_seed_std, dtype=np.float32), + "loss_stop_gradient": stop_gradient_losses, + "loss_stop_gradient_seed_std": np.asarray( + stop_gradient_loss_seed_std, + dtype=np.float32, + ), + "grad_norm": grad_norms, + "grad_norm_stop_gradient": np.asarray( + stop_gradient_grad_norms, + dtype=np.float32, + ), + "update_time": update_times, + "update_time_stop_gradient": np.asarray( + stop_gradient_update_times, + dtype=np.float32, + ), "error_corrected": np.asarray(corrected_error, dtype=np.float32), + "error_corrected_seed_std": np.asarray( + corrected_error_std, + dtype=np.float32, + ), + "error_stop_gradient": np.asarray( + stop_gradient_error, + dtype=np.float32, + ), + "error_stop_gradient_seed_std": np.asarray( + stop_gradient_error_std, + dtype=np.float32, + ), "error_uncorrected": np.asarray(uncorrected_error, dtype=np.float32), } - if first_corrected is not None and first_uncorrected is not None: + if ( + first_corrected is not None + and first_stop_gradient is not None + and first_uncorrected is not None + ): snapshots["rollout_corrected"] = first_corrected + snapshots["rollout_stop_gradient"] = first_stop_gradient snapshots["rollout_uncorrected"] = first_uncorrected return { "metrics": metrics, diff --git a/production.uv.lock b/production.uv.lock index 9eb2596b..e3d5e1ce 100644 --- a/production.uv.lock +++ b/production.uv.lock @@ -939,6 +939,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/75/23/529140fe1aab80fc6992f93a706deec709140a6397439139a054e1515c45/docker-7.2.0-py3-none-any.whl", hash = "sha256:a3f45fdeb9165e2d25d9a1d02ddf3bc70fb572cf5ebbf9b58558c22caf29b71f", size = 148775, upload-time = "2026-07-09T14:53:45.224Z" }, ] +[[package]] +name = "equinox" +version = "0.13.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "jax", version = "0.10.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "jax", version = "0.11.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "jaxtyping", version = "0.3.7", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "jaxtyping", version = "0.3.11", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "typing-extensions" }, + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/ff/522336d2f8264f2ad97119710b76e2cddf66145d03a1e89899175d26b192/equinox-0.13.8.tar.gz", hash = "sha256:dd075050018e2dd02e252e9d29d3060f7e67f085622d8d27a8e89e24bb8523db", size = 145257, upload-time = "2026-05-05T10:03:43.258Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/d6/69a76c8ccdef14af687c497040292a46e59fc7a0ab24724b60e50ca61030/equinox-0.13.8-py3-none-any.whl", hash = "sha256:ca004348533cc30a63ebe8823d7dd4bb626dce17743d40bbddb89b402ef2a240", size = 185813, upload-time = "2026-05-05T10:03:41.673Z" }, +] + [[package]] name = "exceptiongroup" version = "1.3.1" @@ -1624,6 +1642,47 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ac/ff/2edc74b4de40ac4222ee48b3805a6ea5e40f34abe4cd10b65bc0fa0864db/jaxlib-0.11.0-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:005ae3a663193d3a89eafadc073880c67d6b829e0ad989690275d02b5adc810a", size = 87367515, upload-time = "2026-07-16T19:00:10.661Z" }, ] +[[package]] +name = "jaxtyping" +version = "0.3.7" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11'", +] +dependencies = [ + { name = "wadler-lindig", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/40/a2ea3ce0e3e5f540eb970de7792c90fa58fef1b27d34c83f9fa94fea4729/jaxtyping-0.3.7.tar.gz", hash = "sha256:3bd7d9beb7d3cb01a89f93f90581c6f4fff3e5c5dc3c9307e8f8687a040d10c4", size = 45721, upload-time = "2026-01-30T14:18:47.409Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/42/caf65e9a0576a3abadc537e2f831701ba9081f21317fb3be87d64451587a/jaxtyping-0.3.7-py3-none-any.whl", hash = "sha256:303ab8599edf412eeb40bf06c863e3168fa186cf0e7334703fa741ddd7046e66", size = 56101, upload-time = "2026-01-30T14:18:45.954Z" }, +] + +[[package]] +name = "jaxtyping" +version = "0.3.11" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and sys_platform == 'win32'", + "python_full_version >= '3.14' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", +] +dependencies = [ + { name = "wadler-lindig", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/c1/091b8852bd7cbf50bd655543c8506033cf4029300c67f8c176c1286879a9/jaxtyping-0.3.11.tar.gz", hash = "sha256:b09c14acf6686feb9e0df5b0d8c6e7c5b6f8d36bf059ee54cd522a186c2ef050", size = 46489, upload-time = "2026-06-13T18:35:23.167Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/38/c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0/jaxtyping-0.3.11-py3-none-any.whl", hash = "sha256:8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3", size = 56593, upload-time = "2026-06-13T18:35:22.01Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -2111,6 +2170,7 @@ name = "mosaic-bench" source = { editable = "." } dependencies = [ { name = "docker" }, + { name = "equinox" }, { name = "filelock" }, { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "jax", version = "0.10.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, @@ -2150,6 +2210,7 @@ gpu = [ [package.metadata] requires-dist = [ { name = "docker", specifier = ">=6" }, + { name = "equinox", specifier = "==0.13.8" }, { name = "filelock", specifier = ">=3" }, { name = "jax", extras = ["cpu"] }, { name = "jax", extras = ["cuda12"], marker = "extra == 'gpu'" }, @@ -4316,6 +4377,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c1/7c/4e7225d46d634a0d8d534dd8a6ce0c319d09b4d0cf0337eb314ca4789d8c/virtualenv-21.6.1-py3-none-any.whl", hash = "sha256:afe991df855715a2b2f60edfcc0107ef95a79fdfd8cb4cdaa71603d1c12e463b", size = 5506392, upload-time = "2026-07-10T19:33:51.629Z" }, ] +[[package]] +name = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + [[package]] name = "wrapt" version = "2.2.2" diff --git a/pyproject.toml b/pyproject.toml index b1545828..5a5c6f17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "pandas", "matplotlib", "jax[cpu]", + "equinox==0.13.8", "optax", "tesseract-core[runtime]>=1.10.0", "tesseract-jax", diff --git a/requirements.txt b/requirements.txt index b2cb6f8e..fe9a93a5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -61,6 +61,8 @@ debugpy==1.8.21 # via tesseract-core docker==7.2.0 # via mosaic-bench +equinox==0.13.8 + # via mosaic-bench exceptiongroup==1.3.1 ; python_full_version < '3.11' # via anyio fastapi==0.137.0 @@ -96,16 +98,19 @@ importlib-metadata==9.0.0 # via mlflow-skinny jax==0.6.2 ; python_full_version < '3.11' # via + # equinox # mosaic-bench # optax # tesseract-jax jax==0.10.2 ; python_full_version == '3.11.*' # via + # equinox # mosaic-bench # optax # tesseract-jax jax==0.11.0 ; python_full_version >= '3.12' # via + # equinox # mosaic-bench # optax # tesseract-jax @@ -121,6 +126,10 @@ jaxlib==0.11.0 ; python_full_version >= '3.12' # via # jax # optax +jaxtyping==0.3.7 ; python_full_version < '3.11' + # via equinox +jaxtyping==0.3.11 ; python_full_version >= '3.11' + # via equinox jinja2==3.1.6 # via tesseract-core jmespath==1.1.0 @@ -309,6 +318,7 @@ typing-extensions==4.16.0 # aiosignal # anyio # cryptography + # equinox # exceptiongroup # fastapi # mlflow-skinny @@ -337,6 +347,10 @@ uvicorn==0.49.0 # via # mlflow-skinny # tesseract-core +wadler-lindig==0.1.7 + # via + # equinox + # jaxtyping wrapt==2.2.2 # via aiobotocore yarl==1.24.2 diff --git a/tests/test_dummy_integration.py b/tests/test_dummy_integration.py index 1c1cc77f..e1428b88 100644 --- a/tests/test_dummy_integration.py +++ b/tests/test_dummy_integration.py @@ -279,24 +279,34 @@ def _maybe_shrink(cfg, problem: str, exp_key: str) -> None: fresh per-call ``get_config(problem)``, so this never leaks across tests). No-op for experiments not in :data:`_SHRINK_PHYSICS`. """ - if (problem, exp_key) == ("ns-grid", "optimization/solver_in_loop"): + if problem == "ns-grid" and exp_key in { + "optimization/solver_in_loop", + "optimization/solver_in_loop_tgv", + }: from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( solver_in_loop, ) + analytic = exp_key.endswith("_tgv") cfg.add_experiment( exp_key, solver_in_loop, runs=[ { - "ic": {"name": "multimode", "seed": 0}, + "ic": { + "name": "tgv" if analytic else "multimode", + "seed": 0, + }, "physics": { "N": 8, - "nu": 0.001, + "nu": 0.05 if analytic else 0.001, "dt": 0.02, "steps": 1, }, "dataset": { + "reference_kind": ( + "analytic_tgv" if analytic else "pseudo_spectral_multimode" + ), "reference_factor": 2, "reference_substeps": 1, "train_seeds": [0], diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index b88e5bd2..b392770e 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -8,13 +8,16 @@ import dataclasses from pathlib import Path +import equinox as eqx import jax import jax.numpy as jnp import numpy as np from mosaic.benchmarks.core.utils import _debug_run from mosaic.benchmarks.problems.navier_stokes_grid.corrector import ( + PeriodicResidualCNN, apply_corrector, + centered_divergence_rms, divergence_rms, init_corrector, project_periodic_correction, @@ -25,9 +28,14 @@ from mosaic.benchmarks.problems.navier_stokes_grid.ics import _tgv, _tgv_analytic from mosaic.benchmarks.problems.navier_stokes_grid.plots import ( _periodic_vorticity_2d, + _plot_solver_in_loop_fairness, _plot_solver_in_loop_fields, + _plot_solver_in_loop_physics, + _save_solver_in_loop_animation, ) from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( + _rollout_log_gain, + make_reference_dataset, solver_in_loop, ) @@ -40,16 +48,19 @@ def test_periodic_corrector_is_translation_equivariant(): - params = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) + model = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) velocity = jax.random.normal(jax.random.PRNGKey(1), (8, 8, 1, 2)) shifted = jnp.roll(velocity, shift=(2, -1), axis=(0, 1)) + assert isinstance(model, PeriodicResidualCNN) + assert isinstance(model, eqx.Module) + assert model.architecture == "periodic_residual_cnn" expected = jnp.roll( - apply_corrector(params, velocity, velocity_scale=1.0), + apply_corrector(model, velocity, velocity_scale=1.0), shift=(2, -1), axis=(0, 1), ) - actual = apply_corrector(params, shifted, velocity_scale=1.0) + actual = apply_corrector(model, shifted, velocity_scale=1.0) np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) @@ -64,6 +75,7 @@ def test_periodic_projection_is_divergence_free_and_zero_mean(): np.zeros(2), atol=1e-6, ) + assert centered_divergence_rms(_tgv(16), 2.0 * np.pi) < 1e-6 def test_spectral_restriction_preserves_low_mode_tgv(): @@ -97,6 +109,40 @@ def test_spectral_reference_matches_tgv_decay(): assert relative_l2(trajectory[-1], expected) < 1e-5 +def test_analytic_tgv_reference_has_exact_decay_and_distinct_phases(): + viscosity = 0.05 + dt = 0.02 + frame_steps = 4 + train, test, dataset_hash = make_reference_dataset( + physics={"N": 16, "nu": viscosity, "dt": dt, "steps": frame_steps}, + dataset={ + "reference_kind": "analytic_tgv", + "reference_factor": 2, + "train_seeds": [0], + "test_seeds": [100], + "train_frames": 2, + }, + evaluation={"rollout_frames": 2}, + training={"unroll": 1}, + domain_extent=2.0 * np.pi, + ) + + expected_decay = np.exp(-2.0 * viscosity * dt * frame_steps * 2) + actual_decay = np.linalg.norm(test[0, -1]) / np.linalg.norm(test[0, 0]) + np.testing.assert_allclose(actual_decay, expected_decay, rtol=1e-6) + assert not np.allclose(train[0, 0], test[0, 0]) + assert len(dataset_hash) == 16 + + +def test_rollout_log_gain_is_geometric_and_ignores_initial_frame(): + baseline = np.asarray([0.0, 4.0, 8.0]) + corrected = np.asarray([0.0, 2.0, 2.0]) + + gain = _rollout_log_gain(baseline, corrected) + + np.testing.assert_allclose(np.exp(gain), np.sqrt(8.0)) + + def test_debug_run_caps_corrector_training(): run = { "physics": {"N": 32, "steps": 4}, @@ -105,6 +151,7 @@ def test_debug_run_caps_corrector_training(): "unroll": 8, "hidden_channels": 32, "kernel_size": 5, + "model_seeds": [0, 1, 2], "check_grad": True, }, "dataset": { @@ -124,6 +171,7 @@ def test_debug_run_caps_corrector_training(): "unroll": 2, "hidden_channels": 8, "kernel_size": 3, + "model_seeds": [0], "check_grad": False, } assert run["dataset"]["train_seeds"] == [0] @@ -172,6 +220,61 @@ def test_solver_in_loop_fields_render_reference_raw_and_corrected(tmp_path): assert rendered.stat().st_size > 0 +def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): + n = 8 + reference = np.asarray(_tgv(n)) + reference_rollout = np.stack((reference, 0.98 * reference, 0.96 * reference)) + arrays = { + "reference_rollout": reference_rollout, + "evaluation_times": np.asarray([0.0, 0.1, 0.2]), + "rollout_uncorrected_0": np.stack( + (reference, 0.85 * reference, 0.7 * reference) + ), + "rollout_corrected_0": np.stack((reference, 0.95 * reference, 0.9 * reference)), + } + data = { + "by_solver": { + "jax-cfd": { + "native_final_rollout_error": 0.1, + "uncorrected_rollout_error": 0.3, + "uncorrected_mean_rollout_error": 0.25, + "mean_rollout_error": 0.12, + "geometric_error_reduction": 2.0, + "rollout_log_gain_seed_std": 0.1, + "stop_gradient_geometric_error_reduction": 1.4, + "stop_gradient_rollout_log_gain_seed_std": 0.08, + "solver_vjp_geometric_lift": 2.0 / 1.4, + "solver_vjp_log_lift_seed_std": 0.06, + } + } + } + + fairness = _plot_solver_in_loop_fairness( + data, + ["jax-cfd"], + tmp_path, + save=True, + ) + physics = _plot_solver_in_loop_physics( + arrays, + ["jax-cfd"], + tmp_path, + save=True, + ) + _save_solver_in_loop_animation(arrays, ["jax-cfd"], tmp_path) + + assert fairness is not None + assert physics is not None + for filename in ( + "solver_in_loop_fairness.png", + "solver_in_loop_physics.png", + "solver_in_loop_trajectory.gif", + ): + rendered = tmp_path / filename + assert rendered.exists() + assert rendered.stat().st_size > 0 + + def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): """One update crosses the apply/VJP boundary and writes canonical artifacts.""" from mosaic.benchmarks.problems import get_config @@ -200,6 +303,7 @@ def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): "unroll": 2, "hidden_channels": 4, "kernel_size": 3, + "model_seeds": [0, 1], "check_grad": True, "fd_epsilon": 1e-3, }, @@ -218,9 +322,16 @@ def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): assert len(result["results"]) == 1 metrics = result["results"][0]["metrics"] assert metrics["n_updates"] == 1 + assert metrics["total_optimizer_updates"] == 2 + assert metrics["stop_gradient_total_optimizer_updates"] == 2 + assert metrics["n_model_seeds"] == 2 assert metrics["completed"] is True assert metrics["final_grad_norm"] > 0 assert metrics["end_to_end_fd_rel_error"] < 5e-2 + assert metrics["native_final_rollout_error"] >= 0 + assert metrics["solver_vjp_geometric_lift"] > 0 + assert metrics["solver_vjp_update_overhead_ratio"] > 0 + assert metrics["corrector_architecture"] == "periodic_residual_cnn" out_dir = tmp_path / "ns-grid" / "optimization" / "solver_in_loop_smoke" assert (out_dir / "result.json").exists() assert (out_dir / "corrector_fields.npz").exists() From 985584a0303916cf52045d64820f215bd213ea0f Mon Sep 17 00:00:00 2001 From: andrinr Date: Sat, 25 Jul 2026 10:47:50 +0200 Subject: [PATCH 05/15] feat(domain): sharpen solver-loop comparison --- .../problems/navier_stokes_grid/config.py | 39 +- .../problems/navier_stokes_grid/plots.py | 325 +++++++++++--- .../navier_stokes_grid/solver_in_loop.py | 425 ++++++++++++++---- tests/test_solver_in_loop.py | 37 +- 4 files changed, 680 insertions(+), 146 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index ae2fb418..f21a7716 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -449,9 +449,10 @@ solver_in_loop, description=( "Train an identical Equinox periodic residual corrector through each " - "differentiable solver. Paired stop-gradient training and native versus " - "restarted rollouts separate raw accuracy, coupling cost, correctability, " - "and the incremental benefit of the solver VJP." + "differentiable solver. A solver-terminal recurrent auxiliary, paired " + "stop-gradient controls, native versus restarted rollouts, and seen versus " + "held-out IC horizons separate integration, coupling, trainability, and " + "the incremental benefit of the solver VJP." ), plot_description=( "Corrector training, held-out trajectories, physical diagnostics, " @@ -476,9 +477,9 @@ "reference_kind": "pseudo_spectral_multimode", "reference_factor": 2, "reference_substeps": 2, - "train_seeds": [0, 1, 2, 3], - "test_seeds": [100, 101], - "train_frames": 16, + "train_seeds": list(range(16)), + "test_seeds": list(range(100, 108)), + "train_frames": 24, "k0": 6.0, "sigma_k": 1.0, # The stronger field and longer held-out horizon produce @@ -487,10 +488,18 @@ "amplitude": 0.5, }, "training": { - "max_updates": 100, - "unroll": 4, + "max_updates": 200, + "unroll": 8, + # A solver-only terminal auxiliary measures temporal credit + # without letting it destabilize the locally supervised rollout. + "loss_mode": "solver_terminal", + "solver_loss_weight": 0.1, + # Scale each solver's loss by its own uncorrected recurrent + # baseline so clipping/optimisation do not inherit raw accuracy. + "loss_normalization": "solver_baseline", + "loss_scale_floor": 1e-6, "lr": 1e-4, - "clip_norm": 1.0, + "clip_norm": 5.0, "architecture": "periodic_residual_cnn", "hidden_channels": 32, "kernel_size": 5, @@ -501,6 +510,7 @@ }, "evaluation": { "rollout_frames": 36, + "seen_ic_trajectories": 8, "stable_error_threshold": 1.0, }, } @@ -538,10 +548,14 @@ "train_frames": 16, }, "training": { - "max_updates": 100, - "unroll": 4, + "max_updates": 150, + "unroll": 8, + "loss_mode": "solver_terminal", + "solver_loss_weight": 0.1, + "loss_normalization": "solver_baseline", + "loss_scale_floor": 1e-6, "lr": 1e-4, - "clip_norm": 1.0, + "clip_norm": 5.0, "architecture": "periodic_residual_cnn", "hidden_channels": 32, "kernel_size": 5, @@ -552,6 +566,7 @@ }, "evaluation": { "rollout_frames": 24, + "seen_ic_trajectories": 2, "stable_error_threshold": 1.0, }, } diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index bbc35fd7..e72f04d9 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -126,6 +126,9 @@ def plot_solver_in_loop( corrected_std = np.asarray( arrays.get(f"error_corrected_seed_std_{idx}", np.array([])) ) + corrected_ic_std = np.asarray( + arrays.get(f"error_corrected_ic_std_{idx}", np.array([])) + ) stopped = np.asarray(arrays.get(f"error_stop_gradient_{idx}", np.array([]))) uncorrected = np.asarray(arrays.get(f"error_uncorrected_{idx}", np.array([]))) if loss.size: @@ -164,10 +167,16 @@ def plot_solver_in_loop( label=label, ) if corrected_std.shape == corrected.shape: + rollout_uncertainty = corrected_std + if corrected_ic_std.shape == corrected.shape: + rollout_uncertainty = np.hypot( + corrected_std, + corrected_ic_std, + ) ax_roll.fill_between( x[1:], - np.maximum(corrected[1:] - corrected_std[1:], 1e-12), - corrected[1:] + corrected_std[1:], + np.maximum(corrected[1:] - rollout_uncertainty[1:], 1e-12), + corrected[1:] + rollout_uncertainty[1:], color=color, alpha=0.12, linewidth=0, @@ -211,7 +220,7 @@ def plot_solver_in_loop( ): ax_loss.set_yscale("log") ax_loss.set_xlabel("Optimizer update") - ax_loss.set_ylabel("Training loss") + ax_loss.set_ylabel("Normalized training loss") ax_loss.set_title("Corrector training") ax_roll.set_yscale("log") @@ -251,6 +260,14 @@ def plot_solver_in_loop( physics_fig = _plot_solver_in_loop_physics(arrays, names, out_dir, save=save) if physics_fig is not None: figs.append(physics_fig) + diagnostics_fig = _plot_solver_in_loop_diagnostics( + data, + names, + out_dir, + save=save, + ) + if diagnostics_fig is not None: + figs.append(diagnostics_fig) if save: _save_solver_in_loop_animation(arrays, names, out_dir) return figs @@ -606,7 +623,7 @@ def _plot_solver_in_loop_fairness( squeeze=False, layout="constrained", ) - ax_restart, ax_quality, ax_gain, ax_vjp = axes.ravel() + ax_quality, ax_gain, ax_vjp, ax_efficiency = axes.ravel() positions = np.arange(len(rows), dtype=float) labels: list[str] = [] @@ -617,9 +634,11 @@ def _plot_solver_in_loop_fairness( full_gain_log_std: list[float] = [] stopped_gain_log_std: list[float] = [] vjp_lift_log_std: list[float] = [] + update_times: list[float] = [] has_counterfactual = True - has_native = True + has_cost = True has_seed_uncertainty = False + has_ic_uncertainty = False for name, metrics in rows: label, color, _linestyle, marker = solver_props(name) labels.append(label) @@ -627,25 +646,24 @@ def _plot_solver_in_loop_fairness( corrected = float(metrics["mean_rollout_error"]) all_errors.extend((raw, corrected)) ax_quality.scatter(raw, corrected, color=color, marker=marker, s=30) - native = metrics.get("native_final_rollout_error") - restarted = metrics.get("uncorrected_rollout_error") - if native is None or restarted is None: - has_native = False - else: - ax_restart.scatter( - float(native), - float(restarted), - color=color, - marker=marker, - s=30, - ) + update_time = metrics.get("median_update_time_s") + has_cost &= update_time is not None + update_times.append(float(update_time) if update_time is not None else np.nan) full = metrics.get("geometric_error_reduction") if full is None: full = raw / max(corrected, 1e-12) full_gain.append(float(full)) - full_gain_log_std.append(float(metrics.get("rollout_log_gain_seed_std", 0.0))) + full_gain_log_std.append( + float( + np.hypot( + float(metrics.get("rollout_log_gain_seed_std", 0.0)), + float(metrics.get("rollout_log_gain_ic_std", 0.0)), + ) + ) + ) has_seed_uncertainty |= "rollout_log_gain_seed_std" in metrics + has_ic_uncertainty |= "rollout_log_gain_ic_std" in metrics stopped = metrics.get("stop_gradient_geometric_error_reduction") lift = metrics.get("solver_vjp_geometric_lift") @@ -654,9 +672,27 @@ def _plot_solver_in_loop_fairness( stopped_gain.append(float(stopped) if stopped is not None else np.nan) vjp_lift.append(float(lift) if lift is not None else np.nan) stopped_gain_log_std.append( - float(metrics.get("stop_gradient_rollout_log_gain_seed_std", 0.0)) + float( + np.hypot( + float( + metrics.get( + "stop_gradient_rollout_log_gain_seed_std", + 0.0, + ) + ), + float( + metrics.get( + "stop_gradient_rollout_log_gain_ic_std", + 0.0, + ) + ), + ) + ) ) - vjp_lift_log_std.append(float(metrics.get("solver_vjp_log_lift_seed_std", 0.0))) + seed_lift_std = float(metrics.get("solver_vjp_log_lift_seed_std", 0.0)) + ic_lift_std = float(metrics.get("solver_vjp_log_lift_ic_std", 0.0)) + vjp_lift_log_std.append(float(np.hypot(seed_lift_std, ic_lift_std))) + has_ic_uncertainty |= "solver_vjp_log_lift_ic_std" in metrics error_min = max(min(all_errors) * 0.8, 1e-4) error_max = max(all_errors) * 1.25 @@ -678,43 +714,6 @@ def _plot_solver_in_loop_fairness( ax_quality.set_ylabel("Corrected mean error") ax_quality.set_title("Absolute quality") - if has_native: - native_values = [ - float(metrics["native_final_rollout_error"]) for _name, metrics in rows - ] - restarted_values = [ - float(metrics["uncorrected_rollout_error"]) for _name, metrics in rows - ] - restart_min = max(min(native_values + restarted_values) * 0.8, 1e-4) - restart_max = max(native_values + restarted_values) * 1.25 - ax_restart.plot( - [restart_min, restart_max], - [restart_min, restart_max], - color="0.45", - linestyle=":", - linewidth=1.0, - ) - if restart_max / restart_min >= 10.0: - ax_restart.set_xscale("log") - ax_restart.set_yscale("log") - ax_restart.xaxis.set_minor_formatter(mticker.NullFormatter()) - ax_restart.yaxis.set_minor_formatter(mticker.NullFormatter()) - ax_restart.set_xlim(restart_min, restart_max) - ax_restart.set_ylim(restart_min, restart_max) - ax_restart.set_xlabel("Native single-call final error") - ax_restart.set_ylabel("Restarted final error") - ax_restart.set_title("Coupling / restart penalty") - else: - ax_restart.axis("off") - ax_restart.text( - 0.5, - 0.5, - "Native single-call baseline\nnot available in this run", - ha="center", - va="center", - transform=ax_restart.transAxes, - ) - width = 0.36 if has_counterfactual else 0.62 colors = [solver_props(name)[1] for name, _metrics in rows] full_error = _log_scale_errorbars(full_gain, full_gain_log_std) @@ -763,6 +762,33 @@ def _plot_solver_in_loop_fairness( ax_vjp.set_xticks(positions, labels, rotation=35, ha="right") ax_vjp.set_ylabel("Solver-VJP lift [%]") ax_vjp.set_title("Benefit from solver VJP") + if has_cost: + for idx, (name, _metrics) in enumerate(rows): + _label, color, _linestyle, marker = solver_props(name) + error = None if vjp_error is None else vjp_error[:, idx].reshape(2, 1) + ax_efficiency.errorbar( + update_times[idx], + vjp_lift_pct[idx], + yerr=error, + color=color, + marker=marker, + linestyle="none", + capsize=2, + ) + positive_times = [value for value in update_times if value > 0] + if positive_times and max(positive_times) / min(positive_times) >= 10.0: + ax_efficiency.set_xscale("log") + ax_efficiency.axhline( + 0.0, + color="0.45", + linestyle=":", + linewidth=1.0, + ) + ax_efficiency.set_xlabel("Full-VJP update time [s]") + ax_efficiency.set_ylabel("Solver-VJP lift [%]") + ax_efficiency.set_title("VJP benefit versus cost") + else: + ax_efficiency.axis("off") else: ax_vjp.axis("off") ax_vjp.text( @@ -773,15 +799,198 @@ def _plot_solver_in_loop_fairness( va="center", transform=ax_vjp.transAxes, ) + ax_efficiency.axis("off") _solver_loop_legend(fig, [name for name, _metrics in rows]) - uncertainty_suffix = " (error bars: Β±1 seed SD)" if has_seed_uncertainty else "" + uncertainty_suffix = "" + if has_seed_uncertainty and has_ic_uncertainty: + uncertainty_suffix = " (error bars: combined seed/IC SD)" + elif has_seed_uncertainty: + uncertainty_suffix = " (error bars: Β±1 seed SD)" fig.suptitle(f"Fair solver-in-the-loop decomposition{uncertainty_suffix}") if save: save_fig(fig, "solver_in_loop_fairness", out_dir) return fig +def _plot_solver_in_loop_diagnostics( + data: dict[str, Any], + names: list[str], + out_dir: Path, + *, + save: bool, +) -> plt.Figure | None: + """Separate solver-interface effects, IC generalization, and extrapolation.""" + by_solver = data.get("by_solver", {}) + required = ( + "first_interval_rollout_error", + "native_final_rollout_error", + "uncorrected_rollout_error", + "state_restart_error_ratio", + "seen_ic_matched_horizon_error", + "heldout_ic_matched_horizon_error", + "seen_ic_long_horizon_error", + "heldout_ic_long_horizon_error", + ) + solver_order = {alias: idx for idx, alias in enumerate(NS_ORDER)} + rows = [ + (name, by_solver[name]) + for name in names + if name in by_solver + and all(by_solver[name].get(key) is not None for key in required) + ] + rows.sort( + key=lambda row: solver_order.get( + resolve_solver_alias(row[0]) or row[0], + len(solver_order), + ) + ) + if not rows: + return None + + labels = [solver_props(name)[0] for name, _metrics in rows] + colors = [solver_props(name)[1] for name, _metrics in rows] + markers = [solver_props(name)[3] for name, _metrics in rows] + positions = np.arange(len(rows), dtype=float) + width = 0.36 + + def values(key: str) -> list[float]: + return [float(metrics.get(key, 0.0)) for _name, metrics in rows] + + plt.rcParams.update(RCPARAMS) + fig, axes = plt.subplots( + 2, + 3, + figsize=(TEXTWIDTH, 5.0), + squeeze=False, + layout="constrained", + ) + ax_first, ax_restart, ax_restart_ratio, ax_matched, ax_long, ax_gaps = axes.ravel() + + ax_first.bar( + positions, + values("first_interval_rollout_error"), + color=colors, + yerr=values("first_interval_rollout_error_ic_std"), + capsize=2, + ) + ax_first.set_ylabel("Relative $L^2$ error") + ax_first.set_title("First canonical interval") + + native = values("native_final_rollout_error") + restarted = values("uncorrected_rollout_error") + lower = max(min(native + restarted) * 0.8, 1e-5) + upper = max(native + restarted) * 1.25 + ax_restart.plot( + [lower, upper], + [lower, upper], + color="0.45", + linestyle=":", + linewidth=1.0, + ) + for x_value, y_value, color, marker in zip( + native, + restarted, + colors, + markers, + strict=True, + ): + ax_restart.scatter(x_value, y_value, color=color, marker=marker, s=30) + if upper / lower >= 10.0: + ax_restart.set_xscale("log") + ax_restart.set_yscale("log") + ax_restart.set(xlim=(lower, upper), ylim=(lower, upper)) + ax_restart.set_xlabel("Native final error") + ax_restart.set_ylabel("Repeated-call final error") + ax_restart.set_title("State-coupling effect") + + ax_restart_ratio.bar( + positions, + values("state_restart_error_ratio"), + color=colors, + ) + ax_restart_ratio.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_restart_ratio.set_ylabel("Repeated / native error [Γ—]") + ax_restart_ratio.set_title("Restart penalty") + + for ax, horizon, title in ( + (ax_matched, "matched_horizon", "Matched training horizon"), + (ax_long, "long_horizon", "Long rollout horizon"), + ): + seen = values(f"seen_ic_{horizon}_error") + heldout = values(f"heldout_ic_{horizon}_error") + ax.bar( + positions - width / 2, + seen, + width, + color=colors, + alpha=0.9, + label="seen IC", + yerr=values(f"seen_ic_{horizon}_error_ic_std"), + capsize=2, + ) + ax.bar( + positions + width / 2, + heldout, + width, + color=colors, + alpha=0.35, + hatch="//", + label="held-out IC", + yerr=values(f"heldout_ic_{horizon}_error_ic_std"), + capsize=2, + ) + positive = [value for value in (*seen, *heldout) if value > 0] + if positive and max(positive) / min(positive) >= 10.0: + ax.set_yscale("log") + ax.set_xticks(positions, labels, rotation=35, ha="right") + ax.set_ylabel("Corrected relative $L^2$ error") + ax.set_title(title) + ax_matched.legend(loc="best", fontsize=6.5) + + gap_width = 0.24 + for offset, key, label, alpha in ( + ( + -gap_width, + "ic_generalization_ratio_at_matched_horizon", + "held-out / seen", + 0.9, + ), + (0.0, "seen_ic_temporal_extrapolation_ratio", "seen long / matched", 0.6), + ( + gap_width, + "heldout_ic_temporal_extrapolation_ratio", + "held-out long / matched", + 0.3, + ), + ): + ax_gaps.bar( + positions + offset, + values(key), + gap_width, + color=colors, + alpha=alpha, + label=label, + ) + ax_gaps.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_gaps.set_ylabel("Error ratio [Γ—]") + ax_gaps.set_title("Generalization gaps") + ax_gaps.legend(loc="lower left", fontsize=6.0) + + for ax in (ax_first, ax_restart_ratio, ax_matched, ax_long, ax_gaps): + ax.set_xticks(positions, labels, rotation=35, ha="right") + + matched_time = rows[0][1].get("matched_horizon_time") + long_time = rows[0][1].get("rollout_final_time") + suffix = "" + if matched_time is not None and long_time is not None: + suffix = f" ($t_\\mathrm{{match}}={matched_time:g}$, $t_\\mathrm{{long}}={long_time:g}$)" + fig.suptitle(f"Solver interface and corrector generalization{suffix}") + if save: + save_fig(fig, "solver_in_loop_diagnostics", out_dir) + return fig + + def _log_scale_errorbars( values: list[float], log_standard_deviations: list[float], diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py index 90b0489d..c66b45ae 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py @@ -115,15 +115,15 @@ def _cached_reference_dataset( return np.stack(trajectories).astype(np.float32) -def make_reference_dataset( +def _make_reference_datasets( *, physics: dict[str, Any], dataset: dict[str, Any], evaluation: dict[str, Any], training: dict[str, Any], domain_extent: float, -) -> tuple[np.ndarray, np.ndarray, str]: - """Build the train/test trajectory tensors used by every solver.""" +) -> tuple[np.ndarray, np.ndarray, np.ndarray, str]: + """Build train windows plus seen- and held-out-IC evaluation trajectories.""" n = int(physics["N"]) train_seeds = tuple(int(v) for v in dataset.get("train_seeds", [0, 1, 2, 3])) test_seeds = tuple(int(v) for v in dataset.get("test_seeds", [100, 101])) @@ -155,8 +155,28 @@ def make_reference_dataset( all_trajectories = _cached_reference_dataset(**config) n_train = len(train_seeds) train = all_trajectories[:n_train, : train_frames + 1] + train_rollouts = all_trajectories[:n_train, : eval_frames + 1] test = all_trajectories[n_train:, : eval_frames + 1] - return train, test, _dataset_digest(config) + return train, train_rollouts, test, _dataset_digest(config) + + +def make_reference_dataset( + *, + physics: dict[str, Any], + dataset: dict[str, Any], + evaluation: dict[str, Any], + training: dict[str, Any], + domain_extent: float, +) -> tuple[np.ndarray, np.ndarray, str]: + """Return training windows and held-out trajectories for public callers.""" + train, _train_rollouts, test, dataset_hash = _make_reference_datasets( + physics=physics, + dataset=dataset, + evaluation=evaluation, + training=training, + domain_extent=domain_extent, + ) + return train, test, dataset_hash def _solver_advance( @@ -186,17 +206,25 @@ def _window_loss( frame_steps: int, velocity_scale: float, differentiate_solver: bool, + loss_mode: str, + solver_loss_weight: float, + loss_scale: float, ) -> jax.Array: - """Mean normalized state error over one recurrent training window.""" + """Normalized recurrent loss with configurable temporal credit assignment.""" state = targets[0] losses: list[jax.Array] = [] - for target in targets[1:]: + solver_terminal_loss: jax.Array | None = None + for step, target in enumerate(targets[1:]): provisional = _solver_advance(t, ctx, state, frame_steps=frame_steps) if not differentiate_solver: # Keep the identical recurrent forward trajectory but cut the # backward graph at every solver transition. This counterfactual # measures what the same corrector can learn without the solver VJP. provisional = jax.lax.stop_gradient(provisional) + if step == targets.shape[0] - 2: + solver_terminal_loss = jnp.sum((provisional - target) ** 2) / ( + jnp.sum(target**2) + 1e-12 + ) state = corrected_velocity( model, provisional, @@ -204,7 +232,17 @@ def _window_loss( domain_extent=ctx.domain_extent, ) losses.append(jnp.sum((state - target) ** 2) / (jnp.sum(target**2) + 1e-12)) - return jnp.mean(jnp.stack(losses)) + if loss_mode == "mean": + loss = jnp.mean(jnp.stack(losses)) + elif loss_mode == "terminal": + loss = losses[-1] + elif loss_mode == "solver_terminal": + if solver_terminal_loss is None: + raise RuntimeError("solver-terminal loss requires a non-empty rollout") + loss = solver_loss_weight * solver_terminal_loss + jnp.mean(jnp.stack(losses)) + else: + raise ValueError(f"unknown training.loss_mode: {loss_mode!r}") + return loss / jnp.asarray(loss_scale, dtype=loss.dtype) def _directional_fd( @@ -269,6 +307,7 @@ def _train_corrector( frame_steps: int, training: dict[str, Any], velocity_scale: float, + loss_scale: float, differentiate_solver: bool, model_seed: int, ) -> tuple[Any, list[float], list[float], list[float], float | None, bool]: @@ -281,7 +320,17 @@ def _train_corrector( hidden_channels = int(training.get("hidden_channels", 32)) kernel_size = int(training.get("kernel_size", 5)) architecture = str(training.get("architecture", "periodic_residual_cnn")) + loss_mode = str(training.get("loss_mode", "mean")) + solver_loss_weight = float(training.get("solver_loss_weight", 0.1)) fd_epsilon = float(training.get("fd_epsilon", 1e-2)) + if unroll < 1 or train.shape[1] < unroll + 1: + raise ValueError( + "training.unroll must be positive and fit inside dataset.train_frames" + ) + if loss_mode not in {"mean", "terminal", "solver_terminal"}: + raise ValueError(f"unknown training.loss_mode: {loss_mode!r}") + if solver_loss_weight < 0: + raise ValueError("training.solver_loss_weight must be non-negative") model = init_corrector( jax.random.PRNGKey(model_seed), @@ -316,6 +365,9 @@ def _train_corrector( frame_steps=frame_steps, velocity_scale=velocity_scale, differentiate_solver=differentiate_solver, + loss_mode=loss_mode, + solver_loss_weight=solver_loss_weight, + loss_scale=loss_scale, ) loss, grads = eqx.filter_value_and_grad(loss_fn)(model) @@ -377,6 +429,35 @@ def _evaluate_rollout( return np.stack(states), errors +def _evaluate_reference_set( + t: Any, + ctx: KernelContext, + model: Any, + references: np.ndarray, + *, + frame_steps: int, + velocity_scale: float, + corrected: bool, +) -> tuple[np.ndarray | None, np.ndarray]: + """Evaluate one model on several ICs while retaining only the first rollout.""" + errors: list[list[float]] = [] + first_rollout: np.ndarray | None = None + for reference in references: + rollout, reference_errors = _evaluate_rollout( + t, + ctx, + model, + reference, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=corrected, + ) + errors.append(reference_errors) + if first_rollout is None: + first_rollout = rollout + return first_rollout, np.asarray(errors, dtype=np.float64) + + def _first_unstable(errors: list[float], threshold: float) -> int: """Return first bad frame, or the full completed horizon.""" for idx, error in enumerate(errors): @@ -397,6 +478,34 @@ def _rollout_log_gain( return float(np.mean(np.log((baseline + 1e-12) / (corrected + 1e-12)))) +def _rollout_log_gain_samples( + baseline_errors: np.ndarray, + corrected_errors: np.ndarray, +) -> np.ndarray: + """Return paired rollout gains with leading model-seed and IC axes.""" + baseline = np.asarray(baseline_errors, dtype=np.float64) + corrected = np.asarray(corrected_errors, dtype=np.float64) + if baseline.ndim == 2: + baseline = baseline[None, ...] + if baseline.shape[-2:] != corrected.shape[-2:]: + raise ValueError("rollout ensembles must have matching IC/time axes") + return np.mean( + np.log((baseline[..., 1:] + 1e-12) / (corrected[..., 1:] + 1e-12)), + axis=-1, + ) + + +def _ensemble_curve_stats( + errors: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Mean curve plus independent model-seed and initial-condition spread.""" + return ( + np.mean(errors, axis=(0, 1)), + np.std(np.mean(errors, axis=1), axis=0), + np.std(np.mean(errors, axis=0), axis=0), + ) + + def _median_steady_update_time(update_times: list[float]) -> float | None: """Median update time after the compile/FD-heavy first optimizer update.""" steady = update_times[1:] if len(update_times) > 1 else update_times @@ -442,9 +551,18 @@ def _std(values: list[float]) -> float: "update_time_stop_gradient", "error_corrected", "error_corrected_seed_std", + "error_corrected_ic_std", "error_stop_gradient", "error_stop_gradient_seed_std", + "error_stop_gradient_ic_std", "error_uncorrected", + "error_uncorrected_ic_std", + "error_seen_corrected", + "error_seen_stop_gradient", + "error_seen_uncorrected", + "rollout_log_gain_samples", + "stop_gradient_log_gain_samples", + "solver_vjp_log_lift_samples", "rollout_corrected", "rollout_stop_gradient", "rollout_uncorrected", @@ -457,7 +575,7 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: evaluation = ctx.run.get("evaluation", {}) frame_steps = int(ctx.phys["steps"]) - train, test, dataset_hash = make_reference_dataset( + train, train_rollouts, test, dataset_hash = _make_reference_datasets( physics=ctx.phys, dataset=dataset_cfg, evaluation=evaluation, @@ -477,31 +595,56 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: # Evaluate the solver itself once: these trajectories do not depend on a # neural-model seed. A long native call and a sequence of short canonical # restarts expose coupling penalties separately from integration accuracy. - uncorrected_errors: list[list[float]] = [] + seen_trajectory_count = min( + int(evaluation.get("seen_ic_trajectories", test.shape[0])), + train_rollouts.shape[0], + ) + if seen_trajectory_count < 1: + raise ValueError("evaluation.seen_ic_trajectories must select at least one IC") + seen_references = train_rollouts[:seen_trajectory_count] + first_uncorrected, uncorrected_errors_array = _evaluate_reference_set( + t, + ctx, + None, + test, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=False, + ) + _first_seen_uncorrected, seen_uncorrected_errors = _evaluate_reference_set( + t, + ctx, + None, + seen_references, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=False, + ) native_final_errors: list[float] = [] - first_uncorrected: np.ndarray | None = None for reference in test: - uncorrected_rollout, raw_err = _evaluate_rollout( - t, - ctx, - None, - reference, - frame_steps=frame_steps, - velocity_scale=velocity_scale, - corrected=False, - ) native_final = _solver_advance( t, ctx, jnp.asarray(reference[0]), frame_steps=frame_steps * (reference.shape[0] - 1), ) - uncorrected_errors.append(raw_err) native_final_errors.append(relative_l2(np.asarray(native_final), reference[-1])) - if first_uncorrected is None: - first_uncorrected = uncorrected_rollout - uncorrected_error = np.mean(np.asarray(uncorrected_errors), axis=0) + uncorrected_error = np.mean(uncorrected_errors_array, axis=0) + uncorrected_error_ic_std = np.std(uncorrected_errors_array, axis=0) + seen_uncorrected_error = np.mean(seen_uncorrected_errors, axis=0) + training_horizon_frame = min(train.shape[1] - 1, seen_uncorrected_error.size - 1) + loss_normalization = str(training.get("loss_normalization", "target_energy")) + if loss_normalization == "target_energy": + training_loss_scale = 1.0 + elif loss_normalization == "solver_baseline": + loss_scale_floor = float(training.get("loss_scale_floor", 1e-6)) + training_loss_scale = max( + float(np.mean(seen_uncorrected_error[1 : training_horizon_frame + 1] ** 2)), + loss_scale_floor, + ) + else: + raise ValueError(f"unknown training.loss_normalization: {loss_normalization!r}") models: list[Any] = [] losses_by_seed: list[list[float]] = [] @@ -517,6 +660,8 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: stop_gradient_training_walls: list[float] = [] corrected_errors_by_seed: list[np.ndarray] = [] stop_gradient_errors_by_seed: list[np.ndarray] = [] + seen_corrected_errors_by_seed: list[np.ndarray] = [] + seen_stop_gradient_errors_by_seed: list[np.ndarray] = [] first_corrected: np.ndarray | None = None first_stop_gradient: np.ndarray | None = None @@ -543,6 +688,7 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: frame_steps=frame_steps, training=seed_training, velocity_scale=velocity_scale, + loss_scale=training_loss_scale, differentiate_solver=True, model_seed=model_seed, ) @@ -563,37 +709,53 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: frame_steps=frame_steps, training=seed_training, velocity_scale=velocity_scale, + loss_scale=training_loss_scale, differentiate_solver=False, model_seed=model_seed, ) stop_gradient_training_walls.append(time.perf_counter() - stop_gradient_started) - seed_corrected_errors: list[list[float]] = [] - seed_stop_gradient_errors: list[list[float]] = [] - for reference_idx, reference in enumerate(test): - corrected_rollout, corr_err = _evaluate_rollout( - t, - ctx, - model, - reference, - frame_steps=frame_steps, - velocity_scale=velocity_scale, - corrected=True, - ) - stop_gradient_rollout, stop_gradient_err = _evaluate_rollout( + corrected_rollout, seed_corrected_errors = _evaluate_reference_set( + t, + ctx, + model, + test, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + stop_gradient_rollout, seed_stop_gradient_errors = _evaluate_reference_set( + t, + ctx, + stop_gradient_model, + test, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + _seen_corrected_rollout, seen_seed_corrected_errors = _evaluate_reference_set( + t, + ctx, + model, + seen_references, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + _seen_stop_gradient_rollout, seen_seed_stop_gradient_errors = ( + _evaluate_reference_set( t, ctx, stop_gradient_model, - reference, + seen_references, frame_steps=frame_steps, velocity_scale=velocity_scale, corrected=True, ) - seed_corrected_errors.append(corr_err) - seed_stop_gradient_errors.append(stop_gradient_err) - if seed_idx == 0 and reference_idx == 0: - first_corrected = corrected_rollout - first_stop_gradient = stop_gradient_rollout + ) + if seed_idx == 0: + first_corrected = corrected_rollout + first_stop_gradient = stop_gradient_rollout models.append(model) losses_by_seed.append(losses) @@ -605,19 +767,33 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: stop_gradient_grad_norms_by_seed.append(stop_gradient_grad_norms) stop_gradient_update_times_by_seed.append(stop_gradient_update_times) stop_gradient_completed_by_seed.append(stop_gradient_completed) - corrected_errors_by_seed.append( - np.mean(np.asarray(seed_corrected_errors), axis=0) - ) - stop_gradient_errors_by_seed.append( - np.mean(np.asarray(seed_stop_gradient_errors), axis=0) - ) + corrected_errors_by_seed.append(seed_corrected_errors) + stop_gradient_errors_by_seed.append(seed_stop_gradient_errors) + seen_corrected_errors_by_seed.append(seen_seed_corrected_errors) + seen_stop_gradient_errors_by_seed.append(seen_seed_stop_gradient_errors) corrected_errors_array = np.stack(corrected_errors_by_seed) stop_gradient_errors_array = np.stack(stop_gradient_errors_by_seed) - corrected_error = np.mean(corrected_errors_array, axis=0) - corrected_error_std = np.std(corrected_errors_array, axis=0) - stop_gradient_error = np.mean(stop_gradient_errors_array, axis=0) - stop_gradient_error_std = np.std(stop_gradient_errors_array, axis=0) + seen_corrected_errors_array = np.stack(seen_corrected_errors_by_seed) + seen_stop_gradient_errors_array = np.stack(seen_stop_gradient_errors_by_seed) + corrected_error, corrected_error_seed_std, corrected_error_ic_std = ( + _ensemble_curve_stats(corrected_errors_array) + ) + ( + stop_gradient_error, + stop_gradient_error_seed_std, + stop_gradient_error_ic_std, + ) = _ensemble_curve_stats(stop_gradient_errors_array) + ( + seen_corrected_error, + _seen_corrected_error_seed_std, + seen_corrected_error_ic_std, + ) = _ensemble_curve_stats(seen_corrected_errors_array) + ( + seen_stop_gradient_error, + _seen_stop_gradient_error_seed_std, + _seen_stop_gradient_error_ic_std, + ) = _ensemble_curve_stats(seen_stop_gradient_errors_array) losses = _mean_curve(losses_by_seed) stop_gradient_losses = _mean_curve(stop_gradient_losses_by_seed) grad_norms = _mean_curve(grad_norms_by_seed) @@ -666,27 +842,38 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: mean_corrected = float(np.mean(corrected_error[1:])) mean_stop_gradient = float(np.mean(stop_gradient_error[1:])) mean_uncorrected = float(np.mean(uncorrected_error[1:])) - rollout_log_gains = [ - _rollout_log_gain(uncorrected_error, seed_error) - for seed_error in corrected_errors_by_seed - ] - stop_gradient_log_gains = [ - _rollout_log_gain(uncorrected_error, seed_error) - for seed_error in stop_gradient_errors_by_seed - ] - solver_vjp_log_lifts = [ - _rollout_log_gain(stopped, full) - for stopped, full in zip( - stop_gradient_errors_by_seed, - corrected_errors_by_seed, - strict=True, - ) - ] - rollout_log_gain = float(np.mean(rollout_log_gains)) - stop_gradient_log_gain = float(np.mean(stop_gradient_log_gains)) - solver_vjp_log_lift = float(np.mean(solver_vjp_log_lifts)) + rollout_log_gain_samples = _rollout_log_gain_samples( + uncorrected_errors_array, + corrected_errors_array, + ) + stop_gradient_log_gain_samples = _rollout_log_gain_samples( + uncorrected_errors_array, + stop_gradient_errors_array, + ) + solver_vjp_log_lift_samples = _rollout_log_gain_samples( + stop_gradient_errors_array, + corrected_errors_array, + ) + rollout_log_gains = np.mean(rollout_log_gain_samples, axis=1) + stop_gradient_log_gains = np.mean( + stop_gradient_log_gain_samples, + axis=1, + ) + solver_vjp_log_lifts = np.mean( + solver_vjp_log_lift_samples, + axis=1, + ) + rollout_log_gain = float(np.mean(rollout_log_gain_samples)) + stop_gradient_log_gain = float(np.mean(stop_gradient_log_gain_samples)) + solver_vjp_log_lift = float(np.mean(solver_vjp_log_lift_samples)) threshold = float(evaluation.get("stable_error_threshold", 1.0)) interval_time = float(ctx.phys["dt"]) * frame_steps + matched_horizon_frame = min(train.shape[1] - 1, test.shape[1] - 1) + long_horizon_frame = test.shape[1] - 1 + seen_matched_error = float(seen_corrected_error[matched_horizon_frame]) + heldout_matched_error = float(corrected_error[matched_horizon_frame]) + seen_long_error = float(seen_corrected_error[long_horizon_frame]) + heldout_long_error = float(corrected_error[long_horizon_frame]) model = models[0] parameter_count = sum( int(leaf.size) @@ -703,12 +890,19 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: metrics = { "final_rollout_error": final_corrected, - "final_rollout_error_seed_std": float(corrected_error_std[-1]), + "final_rollout_error_seed_std": float(corrected_error_seed_std[-1]), + "final_rollout_error_ic_std": float(corrected_error_ic_std[-1]), "stop_gradient_final_rollout_error": final_stop_gradient, "stop_gradient_final_rollout_error_seed_std": float( - stop_gradient_error_std[-1] + stop_gradient_error_seed_std[-1] + ), + "stop_gradient_final_rollout_error_ic_std": float( + stop_gradient_error_ic_std[-1] ), "uncorrected_rollout_error": final_uncorrected, + "uncorrected_rollout_error_ic_std": float(uncorrected_error_ic_std[-1]), + "first_interval_rollout_error": float(uncorrected_error[1]), + "first_interval_rollout_error_ic_std": float(uncorrected_error_ic_std[1]), "native_final_rollout_error": native_final_error, "state_restart_error_ratio": final_uncorrected / (native_final_error + 1e-12), "mean_rollout_error": mean_corrected, @@ -722,15 +916,58 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: / (mean_uncorrected + 1e-12), "rollout_log_gain": rollout_log_gain, "rollout_log_gain_seed_std": _std(rollout_log_gains), + "rollout_log_gain_ic_std": _std( + np.mean(rollout_log_gain_samples, axis=0).tolist() + ), "geometric_error_reduction": float(np.exp(rollout_log_gain)), "stop_gradient_rollout_log_gain": stop_gradient_log_gain, "stop_gradient_rollout_log_gain_seed_std": _std(stop_gradient_log_gains), + "stop_gradient_rollout_log_gain_ic_std": _std( + np.mean(stop_gradient_log_gain_samples, axis=0).tolist() + ), "stop_gradient_geometric_error_reduction": float( np.exp(stop_gradient_log_gain) ), "solver_vjp_log_lift": solver_vjp_log_lift, "solver_vjp_log_lift_seed_std": _std(solver_vjp_log_lifts), + "solver_vjp_log_lift_ic_std": _std( + np.mean(solver_vjp_log_lift_samples, axis=0).tolist() + ), "solver_vjp_geometric_lift": float(np.exp(solver_vjp_log_lift)), + "seen_ic_matched_horizon_error": seen_matched_error, + "seen_ic_matched_horizon_error_ic_std": float( + seen_corrected_error_ic_std[matched_horizon_frame] + ), + "heldout_ic_matched_horizon_error": heldout_matched_error, + "heldout_ic_matched_horizon_error_ic_std": float( + corrected_error_ic_std[matched_horizon_frame] + ), + "seen_ic_long_horizon_error": seen_long_error, + "seen_ic_long_horizon_error_ic_std": float( + seen_corrected_error_ic_std[long_horizon_frame] + ), + "heldout_ic_long_horizon_error": heldout_long_error, + "heldout_ic_long_horizon_error_ic_std": float( + corrected_error_ic_std[long_horizon_frame] + ), + "ic_generalization_ratio_at_matched_horizon": heldout_matched_error + / (seen_matched_error + 1e-12), + "seen_ic_temporal_extrapolation_ratio": seen_long_error + / (seen_matched_error + 1e-12), + "heldout_ic_temporal_extrapolation_ratio": heldout_long_error + / (heldout_matched_error + 1e-12), + "stop_gradient_seen_ic_matched_horizon_error": float( + seen_stop_gradient_error[matched_horizon_frame] + ), + "stop_gradient_heldout_ic_matched_horizon_error": float( + stop_gradient_error[matched_horizon_frame] + ), + "uncorrected_seen_ic_matched_horizon_error": float( + seen_uncorrected_error[matched_horizon_frame] + ), + "uncorrected_heldout_ic_matched_horizon_error": float( + uncorrected_error[matched_horizon_frame] + ), "stable_horizon": _first_unstable(corrected_error.tolist(), threshold) * interval_time, "stop_gradient_stable_horizon": _first_unstable( @@ -839,7 +1076,16 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: "corrector_parameter_count": parameter_count, "model_seeds": list(model_seeds), "n_model_seeds": len(model_seeds), + "n_train_trajectories": int(train.shape[0]), "n_test_trajectories": int(test.shape[0]), + "n_seen_ic_evaluation_trajectories": seen_trajectory_count, + "training_frames": int(train.shape[1] - 1), + "training_unroll": int(training.get("unroll", 4)), + "training_loss_mode": str(training.get("loss_mode", "mean")), + "training_solver_loss_weight": float(training.get("solver_loss_weight", 0.1)), + "training_loss_normalization": loss_normalization, + "training_loss_scale": training_loss_scale, + "matched_horizon_time": matched_horizon_frame * interval_time, "visualization_model_seed": model_seeds[0], "completed": all(completed_by_seed) and all(stop_gradient_completed_by_seed), "dataset_hash": dataset_hash, @@ -871,7 +1117,11 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: ), "error_corrected": np.asarray(corrected_error, dtype=np.float32), "error_corrected_seed_std": np.asarray( - corrected_error_std, + corrected_error_seed_std, + dtype=np.float32, + ), + "error_corrected_ic_std": np.asarray( + corrected_error_ic_std, dtype=np.float32, ), "error_stop_gradient": np.asarray( @@ -879,10 +1129,35 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: dtype=np.float32, ), "error_stop_gradient_seed_std": np.asarray( - stop_gradient_error_std, + stop_gradient_error_seed_std, + dtype=np.float32, + ), + "error_stop_gradient_ic_std": np.asarray( + stop_gradient_error_ic_std, dtype=np.float32, ), "error_uncorrected": np.asarray(uncorrected_error, dtype=np.float32), + "error_uncorrected_ic_std": np.asarray( + uncorrected_error_ic_std, + dtype=np.float32, + ), + "error_seen_corrected": np.asarray( + seen_corrected_error, + dtype=np.float32, + ), + "error_seen_stop_gradient": np.asarray( + seen_stop_gradient_error, + dtype=np.float32, + ), + "error_seen_uncorrected": np.asarray( + seen_uncorrected_error, + dtype=np.float32, + ), + "rollout_log_gain_samples": rollout_log_gain_samples.astype(np.float32), + "stop_gradient_log_gain_samples": stop_gradient_log_gain_samples.astype( + np.float32 + ), + "solver_vjp_log_lift_samples": solver_vjp_log_lift_samples.astype(np.float32), } if ( first_corrected is not None diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index b392770e..32ab1a02 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -28,6 +28,7 @@ from mosaic.benchmarks.problems.navier_stokes_grid.ics import _tgv, _tgv_analytic from mosaic.benchmarks.problems.navier_stokes_grid.plots import ( _periodic_vorticity_2d, + _plot_solver_in_loop_diagnostics, _plot_solver_in_loop_fairness, _plot_solver_in_loop_fields, _plot_solver_in_loop_physics, @@ -236,7 +237,10 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): "by_solver": { "jax-cfd": { "native_final_rollout_error": 0.1, + "first_interval_rollout_error": 0.08, + "first_interval_rollout_error_ic_std": 0.01, "uncorrected_rollout_error": 0.3, + "state_restart_error_ratio": 3.0, "uncorrected_mean_rollout_error": 0.25, "mean_rollout_error": 0.12, "geometric_error_reduction": 2.0, @@ -245,6 +249,16 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): "stop_gradient_rollout_log_gain_seed_std": 0.08, "solver_vjp_geometric_lift": 2.0 / 1.4, "solver_vjp_log_lift_seed_std": 0.06, + "solver_vjp_log_lift_ic_std": 0.03, + "seen_ic_matched_horizon_error": 0.1, + "heldout_ic_matched_horizon_error": 0.12, + "seen_ic_long_horizon_error": 0.15, + "heldout_ic_long_horizon_error": 0.2, + "ic_generalization_ratio_at_matched_horizon": 1.2, + "seen_ic_temporal_extrapolation_ratio": 1.5, + "heldout_ic_temporal_extrapolation_ratio": 5.0 / 3.0, + "matched_horizon_time": 0.1, + "rollout_final_time": 0.2, } } } @@ -261,11 +275,19 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): tmp_path, save=True, ) + diagnostics = _plot_solver_in_loop_diagnostics( + data, + ["jax-cfd"], + tmp_path, + save=True, + ) _save_solver_in_loop_animation(arrays, ["jax-cfd"], tmp_path) assert fairness is not None assert physics is not None + assert diagnostics is not None for filename in ( + "solver_in_loop_diagnostics.png", "solver_in_loop_fairness.png", "solver_in_loop_physics.png", "solver_in_loop_trajectory.gif", @@ -301,6 +323,9 @@ def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): "training": { "max_updates": 1, "unroll": 2, + "loss_mode": "solver_terminal", + "solver_loss_weight": 0.1, + "loss_normalization": "solver_baseline", "hidden_channels": 4, "kernel_size": 3, "model_seeds": [0, 1], @@ -332,6 +357,16 @@ def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): assert metrics["solver_vjp_geometric_lift"] > 0 assert metrics["solver_vjp_update_overhead_ratio"] > 0 assert metrics["corrector_architecture"] == "periodic_residual_cnn" + assert metrics["training_loss_mode"] == "solver_terminal" + assert metrics["training_solver_loss_weight"] == 0.1 + assert metrics["training_loss_normalization"] == "solver_baseline" + assert metrics["training_loss_scale"] > 0 + assert metrics["n_train_trajectories"] == 1 + assert metrics["n_test_trajectories"] == 1 + assert metrics["seen_ic_matched_horizon_error"] >= 0 + assert metrics["heldout_ic_long_horizon_error"] >= 0 + assert metrics["final_rollout_error_ic_std"] == 0 out_dir = tmp_path / "ns-grid" / "optimization" / "solver_in_loop_smoke" assert (out_dir / "result.json").exists() - assert (out_dir / "corrector_fields.npz").exists() + with np.load(out_dir / "corrector_fields.npz") as snapshots: + assert snapshots["solver_vjp_log_lift_samples_0"].shape == (2, 1) From 4cfa5d1273eb69fd8648b36d907602191bc801e8 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sat, 25 Jul 2026 13:15:27 +0200 Subject: [PATCH 06/15] feat(ns-grid): add recurrence-gated self-reference corrector --- mosaic/benchmarks/core/utils.py | 16 + .../problems/navier_stokes_grid/config.py | 72 +++- .../problems/navier_stokes_grid/exclusions.py | 21 + .../problems/navier_stokes_grid/plots.py | 247 +++++++++--- .../navier_stokes_grid/solver_in_loop.py | 363 +++++++++++++++++- tests/test_dummy_integration.py | 13 +- tests/test_solver_in_loop.py | 108 +++++- 7 files changed, 756 insertions(+), 84 deletions(-) diff --git a/mosaic/benchmarks/core/utils.py b/mosaic/benchmarks/core/utils.py index 97556880..73a04ac9 100644 --- a/mosaic/benchmarks/core/utils.py +++ b/mosaic/benchmarks/core/utils.py @@ -148,11 +148,27 @@ def _debug_run(run: dict) -> None: 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: diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index f21a7716..780c19cf 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -68,7 +68,12 @@ from .ics import _flat_inflow, _multimode, _tgv, _tgv_analytic, _uniform_flow from .optimization import drag_opt from .physics import DIAGNOSTICS, make_inputs -from .plots import plot_drag_opt, plot_solver_in_loop, plot_solver_in_loop_tgv +from .plots import ( + plot_drag_opt, + plot_solver_in_loop, + plot_solver_in_loop_self_reference, + plot_solver_in_loop_tgv, +) from .solver_in_loop import solver_in_loop _TESSERACT_SLUG = "navier-stokes-grid" @@ -573,6 +578,71 @@ ], plot=plot_solver_in_loop_tgv, ) +problem.add_experiment( + "optimization/solver_in_loop_self_reference", + solver_in_loop, + description=( + "Train the paired full-VJP and stop-gradient Equinox correctors against " + "each admitted solver's own spatially and temporally refined trajectory. " + "A direct semigroup audit gates training so the comparison measures " + "solver-in-the-loop credit assignment rather than state-reset artifacts." + ), + plot_description=( + "Within-solver refined-reference correction, recurrence admission, paired " + "solver-VJP lift, held-out trajectories, and training cost. Absolute errors " + "are not compared across solver-specific references." + ), + runs=[ + { + "ic": {"name": "multimode", "seed": 0}, + "physics": { + "N": 32, + "nu": 0.001, + "dt": 0.02, + "steps": 4, + }, + "dataset": { + "reference_kind": "solver_self_refined", + "reference_factor": 2, + "reference_temporal_factor": 2, + "train_seeds": list(range(16)), + "test_seeds": list(range(100, 108)), + "train_frames": 24, + "k0": 6.0, + "sigma_k": 1.0, + "amplitude": 0.5, + "prefix_audit_seeds": [0, 1, 100, 101], + "prefix_audit_frames": [1, 8, 24, 36], + "closure_relative_tolerance": 0.01, + "closure_to_signal_tolerance": 0.1, + "minimum_refinement_signal": 1e-4, + }, + "training": { + "max_updates": 200, + "unroll": 8, + "loss_mode": "solver_terminal", + "solver_loss_weight": 0.1, + "loss_normalization": "solver_baseline", + "loss_scale_floor": 1e-6, + "lr": 1e-4, + "clip_norm": 5.0, + "architecture": "periodic_residual_cnn", + "hidden_channels": 32, + "kernel_size": 5, + "seed": 2026, + "model_seeds": [0, 1, 2], + "check_grad": True, + "fd_epsilon": 1e-2, + }, + "evaluation": { + "rollout_frames": 36, + "seen_ic_trajectories": 8, + "stable_error_threshold": 1.0, + }, + } + ], + plot=plot_solver_in_loop_self_reference, +) # Bonus plot (not paired with an experiment). problem.add_extra_plot( "_extra/gradient/jacobian_svd_comparison", diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py index e39109b6..6f9bac05 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py @@ -88,6 +88,18 @@ "remaining floor is O(dxΒ²) LBM spatial discretization at N=64, not reducible " "by further sub-stepping (tested k=9..27); valid=True", ) +STAGGERED_STATE_NOT_CLOSED = Exclusion( + ExclusionCategory.INFEASIBLE, + "the canonical velocity call boundary repeatedly converts between collocated " + "and staggered states; the resulting transition fails the semigroup admission " + "test and cannot support an intrinsic solver-VJP ranking", +) +XLB_STATE_NOT_CLOSED = Exclusion( + ExclusionCategory.INFEASIBLE, + "the canonical velocity-only call boundary rebuilds equilibrium populations " + "and discards density and non-equilibrium moments; native recurrent state must " + "be preserved before an intrinsic solver-VJP ranking", +) _OBSTACLE_GATE = { "jax_cfd": JAX_CFD_NO_OBSTACLE, @@ -124,3 +136,12 @@ def register(problem: Problem) -> None: problem.exclude("optimization", {"openfoam": OPENFOAM_NON_DIFFERENTIABLE_OPT}) problem.exclude("optimization/drag_opt", _OBSTACLE_GATE) problem.exclude("optimization/drag_opt_bfgs", _OBSTACLE_GATE) + problem.exclude( + "optimization/solver_in_loop_self_reference", + { + "jax_cfd": STAGGERED_STATE_NOT_CLOSED, + "ins_jl": STAGGERED_STATE_NOT_CLOSED, + "phiflow": STAGGERED_STATE_NOT_CLOSED, + "xlb": XLB_STATE_NOT_CLOSED, + }, + ) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index e72f04d9..64c95570 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -89,6 +89,8 @@ def plot_solver_in_loop( save: bool = True, suffix: str = "", exp_key: str = "solver_in_loop", + title: str = "Solver-in-the-loop neural correction", + solver_specific_reference: bool = False, **_kw: Any, ) -> list: """Plot corrector trainability, rollout quality, and time-to-quality.""" @@ -202,11 +204,16 @@ def plot_solver_in_loop( ) wall = metrics.get("training_wall_time_s") - error = metrics.get("final_rollout_error") - if wall is not None and error is not None: + quality = ( + 100.0 * (float(metrics["solver_vjp_geometric_lift"]) - 1.0) + if solver_specific_reference + and metrics.get("solver_vjp_geometric_lift") is not None + else metrics.get("final_rollout_error") + ) + if wall is not None and quality is not None: ax_cost.scatter( [wall], - [error], + [quality], color=color, marker=marker, s=30, @@ -225,7 +232,11 @@ def plot_solver_in_loop( ax_roll.set_yscale("log") ax_roll.set_xlabel("Physical time") - ax_roll.set_ylabel("Relative $L^2$ error") + ax_roll.set_ylabel( + "Error to solver-specific reference" + if solver_specific_reference + else "Relative $L^2$ error" + ) ax_roll.set_title("Held-out rollout") ax_roll.text( 0.98, @@ -243,18 +254,30 @@ def plot_solver_in_loop( ) ax_cost.set_xlabel("Training wall time [s]") - ax_cost.set_ylabel("Final rollout error") - ax_cost.set_title("Time to quality") + ax_cost.set_ylabel( + "Solver-VJP lift [%]" if solver_specific_reference else "Final rollout error" + ) + ax_cost.set_title( + "VJP benefit versus cost" if solver_specific_reference else "Time to quality" + ) + if solver_specific_reference: + ax_cost.axhline(0.0, color="0.45", linestyle=":", linewidth=1.0) _solver_loop_legend(fig, present) - fig.suptitle("Solver-in-the-loop neural correction") + fig.suptitle(title) figs = [fig] if save: save_fig(fig, "solver_in_loop", out_dir) fields_fig = _plot_solver_in_loop_fields(arrays, names, out_dir, save=save) if fields_fig is not None: figs.append(fields_fig) - fairness_fig = _plot_solver_in_loop_fairness(data, names, out_dir, save=save) + fairness_fig = _plot_solver_in_loop_fairness( + data, + names, + out_dir, + save=save, + solver_specific_reference=solver_specific_reference, + ) if fairness_fig is not None: figs.append(fairness_fig) physics_fig = _plot_solver_in_loop_physics(arrays, names, out_dir, save=save) @@ -290,6 +313,38 @@ def plot_solver_in_loop_tgv( ) +def plot_solver_in_loop_self_reference( + cfg: Problem, + *, + save: bool = True, + suffix: str = "", + **kwargs: Any, +) -> list: + """Plot the recurrence-admitted, solver-specific refined-reference task.""" + return plot_solver_in_loop( + cfg, + save=save, + suffix=suffix, + exp_key="solver_in_loop_self_reference", + title="Solver-specific refined-reference neural correction", + solver_specific_reference=True, + **kwargs, + ) + + +def _solver_reference_rollout( + arrays: dict[str, np.ndarray], + solver_index: int, +) -> np.ndarray: + """Return a solver-specific target when present, else the shared target.""" + return np.asarray( + arrays.get( + f"reference_rollout_{solver_index}", + arrays.get("reference_rollout", np.array([])), + ) + ) + + def _periodic_vorticity_2d( velocity: np.ndarray, *, @@ -389,23 +444,25 @@ def _plot_solver_in_loop_fields( every raw/corrected pair can be read horizontally without looking across rows. All panels share one robust symmetric colour scale. """ - reference = np.asarray(arrays.get("reference_rollout", np.array([]))) - if reference.size == 0 or reference.shape[0] == 0: - return None - rows = _ordered_solver_rollouts(arrays, names) if not rows: return None - reference_final = reference[-1] - reference_vorticity = _periodic_vorticity_2d(reference_final) - rendered_rows: list[tuple[str, np.ndarray, np.ndarray]] = [] - scale_fields = [reference_vorticity] + rendered_rows: list[tuple[str, np.ndarray, np.ndarray, np.ndarray]] = [] + scale_fields: list[np.ndarray] = [] for name, raw, corrected in rows: + reference = _solver_reference_rollout(arrays, names.index(name)) + if reference.size == 0 or reference.shape[0] == 0: + continue + reference_vorticity = _periodic_vorticity_2d(reference[-1]) raw_vorticity = _periodic_vorticity_2d(raw[-1]) corrected_vorticity = _periodic_vorticity_2d(corrected[-1]) - rendered_rows.append((name, raw_vorticity, corrected_vorticity)) - scale_fields.extend((raw_vorticity, corrected_vorticity)) + rendered_rows.append( + (name, reference_vorticity, raw_vorticity, corrected_vorticity) + ) + scale_fields.extend((reference_vorticity, raw_vorticity, corrected_vorticity)) + if not rendered_rows: + return None magnitudes = np.concatenate([np.abs(field).ravel() for field in scale_fields]) vmax = float(np.percentile(magnitudes, 99.0)) or 1.0 @@ -420,7 +477,12 @@ def _plot_solver_in_loop_fields( ) column_titles = ("Reference", "Solver only", "Solver + corrector") image = None - for row_idx, (name, raw_vorticity, corrected_vorticity) in enumerate(rendered_rows): + for row_idx, ( + name, + reference_vorticity, + raw_vorticity, + corrected_vorticity, + ) in enumerate(rendered_rows): label, _color, _linestyle, _marker = solver_props(name) for col_idx, field in enumerate( (reference_vorticity, raw_vorticity, corrected_vorticity) @@ -493,21 +555,28 @@ def _plot_solver_in_loop_physics( save: bool, ) -> plt.Figure | None: """Compare energy, enstrophy, and divergence along held-out rollouts.""" - reference = np.asarray(arrays.get("reference_rollout", np.array([]))) rows = _ordered_solver_rollouts(arrays, names) - if reference.size == 0 or not rows: + solver_references = { + name: _solver_reference_rollout(arrays, names.index(name)) + for name, _raw, _corrected in rows + } + rows = [row for row in rows if solver_references[row[0]].size] + if not rows: return None n_frames = min( - reference.shape[0], - *(min(raw.shape[0], corrected.shape[0]) for _, raw, corrected in rows), + [ + min( + solver_references[name].shape[0], + raw.shape[0], + corrected.shape[0], + ) + for name, raw, corrected in rows + ] ) times = np.asarray(arrays.get("evaluation_times", np.array([])))[:n_frames] if times.size != n_frames: times = np.arange(n_frames, dtype=float) - reference_diagnostics = _trajectory_diagnostics(reference[:n_frames]) - energy_scale = reference_diagnostics[0][0] + 1e-12 - enstrophy_scale = reference_diagnostics[1][0] + 1e-12 plt.rcParams.update(RCPARAMS) fig, axes = plt.subplots( @@ -519,18 +588,21 @@ def _plot_solver_in_loop_physics( ) ax_energy, ax_enstrophy, ax_spectral, ax_centered = axes.ravel() diagnostic_axes = (ax_energy, ax_enstrophy, ax_spectral, ax_centered) - reference_curves = ( - reference_diagnostics[0] / energy_scale, - reference_diagnostics[1] / enstrophy_scale, - reference_diagnostics[2], - reference_diagnostics[3], - ) - for ax, curve in zip(diagnostic_axes, reference_curves, strict=True): - ax.plot(times, np.maximum(curve, 1e-12), color="black", linewidth=1.25) present: list[str] = [] for name, raw, corrected in rows: _label, color, _linestyle, _marker = solver_props(name) + reference_diagnostics = _trajectory_diagnostics( + solver_references[name][:n_frames] + ) + energy_scale = reference_diagnostics[0][0] + 1e-12 + enstrophy_scale = reference_diagnostics[1][0] + 1e-12 + reference_curves = ( + reference_diagnostics[0] / energy_scale, + reference_diagnostics[1] / enstrophy_scale, + reference_diagnostics[2], + reference_diagnostics[3], + ) raw_diagnostics = _trajectory_diagnostics(raw[:n_frames]) corrected_diagnostics = _trajectory_diagnostics(corrected[:n_frames]) raw_curves = ( @@ -545,12 +617,20 @@ def _plot_solver_in_loop_physics( corrected_diagnostics[2], corrected_diagnostics[3], ) - for ax, raw_curve, corrected_curve in zip( + for ax, reference_curve, raw_curve, corrected_curve in zip( diagnostic_axes, + reference_curves, raw_curves, corrected_curves, strict=True, ): + ax.plot( + times, + np.maximum(reference_curve, 1e-12), + color=color, + linestyle="--", + alpha=0.7, + ) ax.plot(times, np.maximum(raw_curve, 1e-12), color=color, linestyle=":") ax.plot( times, @@ -577,7 +657,9 @@ def _plot_solver_in_loop_physics( fig, present, extra_handles=[ - mlines.Line2D([], [], color="black", linewidth=1.25, label="reference"), + mlines.Line2D( + [], [], color="0.4", linestyle="--", label="solver-specific reference" + ), mlines.Line2D([], [], color="0.4", linestyle=":", label="solver only"), mlines.Line2D([], [], color="0.4", linestyle="-", label="full corrector"), ], @@ -594,6 +676,7 @@ def _plot_solver_in_loop_fairness( out_dir: Path, *, save: bool, + solver_specific_reference: bool = False, ) -> plt.Figure | None: """Separate raw quality, correctability, and benefit from the solver VJP.""" by_solver = data.get("by_solver", {}) @@ -712,7 +795,11 @@ def _plot_solver_in_loop_fairness( ax_quality.set_ylim(error_min, error_max) ax_quality.set_xlabel("Solver-only mean error") ax_quality.set_ylabel("Corrected mean error") - ax_quality.set_title("Absolute quality") + ax_quality.set_title( + "Within-solver target error" + if solver_specific_reference + else "Absolute quality" + ) width = 0.36 if has_counterfactual else 0.62 colors = [solver_props(name)[1] for name, _metrics in rows] @@ -807,7 +894,12 @@ def _plot_solver_in_loop_fairness( uncertainty_suffix = " (error bars: combined seed/IC SD)" elif has_seed_uncertainty: uncertainty_suffix = " (error bars: Β±1 seed SD)" - fig.suptitle(f"Fair solver-in-the-loop decomposition{uncertainty_suffix}") + title = ( + "Solver-specific self-consistency decomposition" + if solver_specific_reference + else "Fair solver-in-the-loop decomposition" + ) + fig.suptitle(f"{title}{uncertainty_suffix}") if save: save_fig(fig, "solver_in_loop_fairness", out_dir) return fig @@ -904,14 +996,31 @@ def values(key: str) -> list[float]: ax_restart.set_ylabel("Repeated-call final error") ax_restart.set_title("State-coupling effect") - ax_restart_ratio.bar( - positions, - values("state_restart_error_ratio"), - color=colors, - ) - ax_restart_ratio.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) - ax_restart_ratio.set_ylabel("Repeated / native error [Γ—]") - ax_restart_ratio.set_title("Restart penalty") + if all(metrics.get("semigroup_error_p95") is not None for _name, metrics in rows): + closure_pct = 100.0 * np.asarray(values("semigroup_error_p95")) + tolerance_pct = 100.0 * max(values("semigroup_p95_tolerance")) + ax_restart_ratio.bar(positions, closure_pct, color=colors) + ax_restart_ratio.axhline( + tolerance_pct, + color="0.45", + linestyle=":", + linewidth=1.0, + label="admission limit", + ) + positive = closure_pct[closure_pct > 0] + if positive.size and max(positive) / min(positive) >= 10.0: + ax_restart_ratio.set_yscale("log") + ax_restart_ratio.set_ylabel("95th percentile residual [%]") + ax_restart_ratio.set_title("Semigroup admission") + else: + ax_restart_ratio.bar( + positions, + values("state_restart_error_ratio"), + color=colors, + ) + ax_restart_ratio.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_restart_ratio.set_ylabel("Repeated / native error [Γ—]") + ax_restart_ratio.set_title("Restart penalty") for ax, horizon, title in ( (ax_matched, "matched_horizon", "Matched training horizon"), @@ -1011,14 +1120,24 @@ def _save_solver_in_loop_animation( out_dir: Path, ) -> None: """Animate reference, solver-only, and fully corrected held-out rollouts.""" - reference = np.asarray(arrays.get("reference_rollout", np.array([]))) rows = _ordered_solver_rollouts(arrays, names) - if reference.size == 0 or not rows: + solver_references = { + name: _solver_reference_rollout(arrays, names.index(name)) + for name, _raw, _corrected in rows + } + rows = [row for row in rows if solver_references[row[0]].size] + if not rows: return available_frames = min( - reference.shape[0], - *(min(raw.shape[0], corrected.shape[0]) for _, raw, corrected in rows), + [ + min( + solver_references[name].shape[0], + raw.shape[0], + corrected.shape[0], + ) + for name, raw, corrected in rows + ] ) # Keep the full simulated horizon in the animation while bounding the GIF # payload embedded in PRs. Evenly spaced samples retain both endpoints. @@ -1031,19 +1150,21 @@ def _save_solver_in_loop_animation( ) ) n_frames = len(frame_indices) - reference_vorticity = [ - _periodic_vorticity_2d(reference[index]) for index in frame_indices - ] rollout_vorticity = [ ( name, + [ + _periodic_vorticity_2d(solver_references[name][index]) + for index in frame_indices + ], [_periodic_vorticity_2d(raw[index]) for index in frame_indices], [_periodic_vorticity_2d(corrected[index]) for index in frame_indices], ) for name, raw, corrected in rows ] - scale_fields = list(reference_vorticity) - for _name, raw_fields, corrected_fields in rollout_vorticity: + scale_fields: list[np.ndarray] = [] + for _name, reference_fields, raw_fields, corrected_fields in rollout_vorticity: + scale_fields.extend(reference_fields) scale_fields.extend(raw_fields) scale_fields.extend(corrected_fields) magnitudes = np.concatenate([np.abs(field).ravel() for field in scale_fields]) @@ -1059,10 +1180,15 @@ def _save_solver_in_loop_animation( layout="constrained", ) images: list[tuple[Any, int, int]] = [] - for row_idx, (name, raw_fields, corrected_fields) in enumerate(rollout_vorticity): + for row_idx, ( + name, + reference_fields, + raw_fields, + corrected_fields, + ) in enumerate(rollout_vorticity): label, _color, _linestyle, _marker = solver_props(name) for col_idx, fields in enumerate( - (reference_vorticity, raw_fields, corrected_fields) + (reference_fields, raw_fields, corrected_fields) ): ax = axes[row_idx, col_idx] image = ax.imshow( @@ -1105,12 +1231,7 @@ def _save_solver_in_loop_animation( def _update(frame: int) -> tuple[Any, ...]: for image, row_idx, col_idx in images: - if col_idx == 0: - field = reference_vorticity[frame] - elif col_idx == 1: - field = rollout_vorticity[row_idx][1][frame] - else: - field = rollout_vorticity[row_idx][2][frame] + field = rollout_vorticity[row_idx][col_idx + 1][frame] image.set_data(field.T) time_value = float(times[frame]) if frame < times.size else float(frame) title.set_text(f"Held-out trajectory at $t={time_value:g}$") diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py index c66b45ae..c6646478 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py @@ -187,7 +187,30 @@ def _solver_advance( frame_steps: int, ) -> jax.Array: """Advance one canonical correction interval through a Tesseract.""" - physics = {**ctx.phys, "steps": frame_steps} + return _solver_advance_with_physics( + t, + ctx, + velocity, + dt=float(ctx.phys["dt"]), + steps=frame_steps, + ) + + +def _solver_advance_with_physics( + t: Any, + ctx: KernelContext, + velocity: jax.Array | np.ndarray, + *, + dt: float, + steps: int, +) -> jax.Array: + """Advance a state with an explicit grid and temporal compute budget.""" + physics = { + **ctx.phys, + "N": int(velocity.shape[0]), + "dt": float(dt), + "steps": int(steps), + } inputs = ctx.make_inputs(ctx.name, velocity, **physics) outputs = apply_tesseract(t, inputs) if ctx.output_key not in outputs: @@ -197,6 +220,228 @@ def _solver_advance( return outputs[ctx.output_key] +def _make_solver_self_reference_datasets( + t: Any, + ctx: KernelContext, + *, + dataset: dict[str, Any], + evaluation: dict[str, Any], + training: dict[str, Any], +) -> tuple[np.ndarray, np.ndarray, np.ndarray, str, dict[str, Any]]: + """Build and audit targets from the same solver at a refined compute budget.""" + started = time.perf_counter() + n = int(ctx.phys["N"]) + reference_factor = int(dataset.get("reference_factor", 2)) + temporal_factor = int(dataset.get("reference_temporal_factor", 2)) + if reference_factor < 1 or temporal_factor < 1: + raise ValueError("reference refinement factors must be positive") + + train_seeds = tuple(int(v) for v in dataset.get("train_seeds", [0, 1, 2, 3])) + test_seeds = tuple(int(v) for v in dataset.get("test_seeds", [100, 101])) + all_seeds = train_seeds + test_seeds + train_frames = int(dataset.get("train_frames", 16)) + eval_frames = int(evaluation.get("rollout_frames", 32)) + unroll = int(training.get("unroll", 4)) + n_frames = max(train_frames, eval_frames, unroll) + + fine_n = n * reference_factor + coarse_dt = float(ctx.phys["dt"]) + coarse_steps = int(ctx.phys["steps"]) + fine_dt = coarse_dt / temporal_factor + fine_steps = coarse_steps * temporal_factor + k0 = float(dataset.get("k0", 6.0)) + sigma_k = float(dataset.get("sigma_k", 1.0)) + amplitude = float(dataset.get("amplitude", 0.3)) + + apply_count = 0 + + def advance( + velocity: np.ndarray, + *, + dt: float, + steps: int, + ) -> np.ndarray: + nonlocal apply_count + apply_count += 1 + return np.asarray( + _solver_advance_with_physics( + t, + ctx, + velocity, + dt=dt, + steps=steps, + ) + ) + + trajectories: dict[int, np.ndarray] = {} + fine_initials: dict[int, np.ndarray] = {} + for seed in all_seeds: + fine_state = np.asarray( + _multimode( + fine_n, + L=ctx.domain_extent, + seed=seed, + k0=k0, + sigma_k=sigma_k, + amplitude=amplitude, + ), + dtype=np.float32, + ) + fine_initials[seed] = fine_state + frames = [spectral_restrict(fine_state, n)] + for _frame in range(n_frames): + fine_state = advance(fine_state, dt=fine_dt, steps=fine_steps) + frames.append(spectral_restrict(fine_state, n)) + trajectories[seed] = np.asarray(frames, dtype=np.float32) + + requested_audit_seeds = tuple( + int(v) + for v in dataset.get( + "prefix_audit_seeds", + [train_seeds[0], test_seeds[0]], + ) + ) + missing_audit_seeds = set(requested_audit_seeds) - set(all_seeds) + if missing_audit_seeds: + raise ValueError( + "dataset.prefix_audit_seeds must belong to the train/test split: " + f"{sorted(missing_audit_seeds)}" + ) + audit_frames = tuple( + sorted( + { + int(frame) + for frame in dataset.get( + "prefix_audit_frames", + [1, min(train_frames, n_frames), n_frames], + ) + if 0 < int(frame) <= n_frames + } + ) + ) + if not audit_frames: + raise ValueError("dataset.prefix_audit_frames must select a positive frame") + + coarse_closure_errors: list[float] = [] + fine_closure_errors: list[float] = [] + refinement_signals: list[float] = [] + finite = True + audit_frame_set = set(audit_frames) + max_audit_frame = max(audit_frames) + for seed in requested_audit_seeds: + target = trajectories[seed] + coarse_state = np.asarray(target[0]) + coarse_repeated: dict[int, np.ndarray] = {} + for frame in range(1, max_audit_frame + 1): + coarse_state = advance( + coarse_state, + dt=coarse_dt, + steps=coarse_steps, + ) + if frame in audit_frame_set: + coarse_repeated[frame] = coarse_state + + for frame in audit_frames: + coarse_native = advance( + np.asarray(target[0]), + dt=coarse_dt, + steps=coarse_steps * frame, + ) + fine_native = advance( + fine_initials[seed], + dt=fine_dt, + steps=fine_steps * frame, + ) + fine_native_coarse = spectral_restrict(fine_native, n) + coarse_closure_errors.append( + relative_l2(coarse_repeated[frame], coarse_native) + ) + fine_closure_errors.append(relative_l2(target[frame], fine_native_coarse)) + refinement_signals.append( + relative_l2(coarse_repeated[frame], target[frame]) + ) + finite = finite and bool( + np.all(np.isfinite(coarse_repeated[frame])) + and np.all(np.isfinite(coarse_native)) + and np.all(np.isfinite(target[frame])) + and np.all(np.isfinite(fine_native_coarse)) + ) + + closure_tolerance = float(dataset.get("closure_relative_tolerance", 0.01)) + closure_to_signal_tolerance = float(dataset.get("closure_to_signal_tolerance", 0.1)) + minimum_signal = float(dataset.get("minimum_refinement_signal", 1e-4)) + max_coarse_closure = float(max(coarse_closure_errors)) + max_fine_closure = float(max(fine_closure_errors)) + mean_signal = float(np.mean(refinement_signals)) + closure_to_signal = [ + closure / (signal + 1e-12) + for closure, signal in zip( + fine_closure_errors, + refinement_signals, + strict=True, + ) + ] + eligible = bool( + finite + and max_coarse_closure <= closure_tolerance + and max_fine_closure <= closure_tolerance + and max(closure_to_signal) <= closure_to_signal_tolerance + and mean_signal > minimum_signal + ) + + config = { + "reference_kind": "solver_self_refined", + "solver": ctx.name, + "n": n, + "reference_factor": reference_factor, + "reference_temporal_factor": temporal_factor, + "viscosity": float(ctx.phys["nu"]), + "coarse_dt": coarse_dt, + "coarse_steps": coarse_steps, + "fine_dt": fine_dt, + "fine_steps": fine_steps, + "n_frames": n_frames, + "domain_extent": float(ctx.domain_extent), + "train_seeds": train_seeds, + "test_seeds": test_seeds, + "k0": k0, + "sigma_k": sigma_k, + "amplitude": amplitude, + "prefix_audit_seeds": requested_audit_seeds, + "prefix_audit_frames": audit_frames, + "closure_relative_tolerance": closure_tolerance, + "closure_to_signal_tolerance": closure_to_signal_tolerance, + "minimum_refinement_signal": minimum_signal, + } + all_trajectories = np.stack([trajectories[seed] for seed in all_seeds]) + n_train = len(train_seeds) + train = all_trajectories[:n_train, : train_frames + 1] + train_rollouts = all_trajectories[:n_train, : eval_frames + 1] + test = all_trajectories[n_train:, : eval_frames + 1] + audit = { + "eligible_for_corrector_training": eligible, + "reference_generation_wall_time_s": time.perf_counter() - started, + "reference_generation_apply_count": apply_count, + "reference_grid_size": fine_n, + "reference_dt": fine_dt, + "reference_steps_per_interval": fine_steps, + "prefix_audit_seeds": list(requested_audit_seeds), + "prefix_audit_frames": list(audit_frames), + "coarse_closure_errors": coarse_closure_errors, + "fine_closure_errors": fine_closure_errors, + "refinement_signals": refinement_signals, + "max_coarse_closure_error": max_coarse_closure, + "max_fine_closure_error": max_fine_closure, + "max_fine_closure_to_signal_ratio": float(max(closure_to_signal)), + "mean_refinement_signal": mean_signal, + "closure_relative_tolerance": closure_tolerance, + "closure_to_signal_tolerance": closure_to_signal_tolerance, + "minimum_refinement_signal": minimum_signal, + "reference_states_finite": finite, + } + return train, train_rollouts, test, _dataset_digest(config), audit + + def _window_loss( model: Any, targets: jax.Array, @@ -566,6 +811,7 @@ def _std(values: list[float]) -> float: "rollout_corrected", "rollout_stop_gradient", "rollout_uncorrected", + "reference_rollout", ), ) def solver_in_loop(t: Any, ctx: KernelContext) -> dict: @@ -574,14 +820,49 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: dataset_cfg = ctx.run.get("dataset", {}) evaluation = ctx.run.get("evaluation", {}) frame_steps = int(ctx.phys["steps"]) - - train, train_rollouts, test, dataset_hash = _make_reference_datasets( - physics=ctx.phys, - dataset=dataset_cfg, - evaluation=evaluation, - training=training, - domain_extent=ctx.domain_extent, - ) + reference_kind = str(dataset_cfg.get("reference_kind", "pseudo_spectral_multimode")) + reference_audit: dict[str, Any] = {} + if reference_kind == "solver_self_refined": + ( + train, + train_rollouts, + test, + dataset_hash, + reference_audit, + ) = _make_solver_self_reference_datasets( + t, + ctx, + dataset=dataset_cfg, + evaluation=evaluation, + training=training, + ) + else: + train, train_rollouts, test, dataset_hash = _make_reference_datasets( + physics=ctx.phys, + dataset=dataset_cfg, + evaluation=evaluation, + training=training, + domain_extent=ctx.domain_extent, + ) + interval_time = float(ctx.phys["dt"]) * frame_steps + if not reference_audit.get("eligible_for_corrector_training", True): + return { + "metrics": { + **reference_audit, + "completed": True, + "n_updates": 0, + "stop_gradient_n_updates": 0, + "dataset_hash": dataset_hash, + "reference_kind": reference_kind, + "correction_intervals": int(test.shape[1] - 1), + "rollout_final_time": interval_time * int(test.shape[1] - 1), + }, + "snapshots": {"reference_rollout": test[0]}, + "shared": { + "evaluation_times": np.arange(test.shape[1], dtype=np.float32) + * interval_time + }, + } velocity_scale = float(np.sqrt(np.mean(train**2)) + 1e-8) configured_seeds = training.get("model_seeds") @@ -602,6 +883,39 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: if seen_trajectory_count < 1: raise ValueError("evaluation.seen_ic_trajectories must select at least one IC") seen_references = train_rollouts[:seen_trajectory_count] + semigroup_errors: list[float] = [] + for reference in test: + one_interval = _solver_advance( + t, + ctx, + jnp.asarray(reference[0]), + frame_steps=frame_steps, + ) + repeated = _solver_advance( + t, + ctx, + one_interval, + frame_steps=frame_steps, + ) + uninterrupted = _solver_advance( + t, + ctx, + jnp.asarray(reference[0]), + frame_steps=2 * frame_steps, + ) + semigroup_errors.append( + relative_l2(np.asarray(repeated), np.asarray(uninterrupted)) + ) + semigroup_median = float(np.median(semigroup_errors)) + semigroup_p95 = float(np.percentile(semigroup_errors, 95.0)) + semigroup_median_tolerance = float( + dataset_cfg.get("semigroup_median_tolerance", 0.005) + ) + semigroup_p95_tolerance = float(dataset_cfg.get("semigroup_p95_tolerance", 0.01)) + valid_for_vjp_ranking = bool( + semigroup_median <= semigroup_median_tolerance + and semigroup_p95 <= semigroup_p95_tolerance + ) first_uncorrected, uncorrected_errors_array = _evaluate_reference_set( t, ctx, @@ -867,7 +1181,6 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: stop_gradient_log_gain = float(np.mean(stop_gradient_log_gain_samples)) solver_vjp_log_lift = float(np.mean(solver_vjp_log_lift_samples)) threshold = float(evaluation.get("stable_error_threshold", 1.0)) - interval_time = float(ctx.phys["dt"]) * frame_steps matched_horizon_frame = min(train.shape[1] - 1, test.shape[1] - 1) long_horizon_frame = test.shape[1] - 1 seen_matched_error = float(seen_corrected_error[matched_horizon_frame]) @@ -889,6 +1202,14 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: ) metrics = { + **reference_audit, + "eligible_for_corrector_training": True, + "valid_for_vjp_ranking": valid_for_vjp_ranking, + "semigroup_error_median": semigroup_median, + "semigroup_error_p95": semigroup_p95, + "semigroup_error_samples": semigroup_errors, + "semigroup_median_tolerance": semigroup_median_tolerance, + "semigroup_p95_tolerance": semigroup_p95_tolerance, "final_rollout_error": final_corrected, "final_rollout_error_seed_std": float(corrected_error_seed_std[-1]), "final_rollout_error_ic_std": float(corrected_error_ic_std[-1]), @@ -1089,9 +1410,7 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: "visualization_model_seed": model_seeds[0], "completed": all(completed_by_seed) and all(stop_gradient_completed_by_seed), "dataset_hash": dataset_hash, - "reference_kind": str( - dataset_cfg.get("reference_kind", "pseudo_spectral_multimode") - ), + "reference_kind": reference_kind, "correction_intervals": int(test.shape[1] - 1), "native_steps": frame_steps * int(test.shape[1] - 1), "rollout_final_time": interval_time * int(test.shape[1] - 1), @@ -1167,14 +1486,22 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: snapshots["rollout_corrected"] = first_corrected snapshots["rollout_stop_gradient"] = first_stop_gradient snapshots["rollout_uncorrected"] = first_uncorrected - return { - "metrics": metrics, - "snapshots": snapshots, - "shared": { + if reference_kind == "solver_self_refined": + snapshots["reference_rollout"] = test[0] + shared = { + "evaluation_times": np.arange(test.shape[1], dtype=np.float32) + * interval_time + } + else: + shared = { "reference_rollout": test[0], "evaluation_times": np.arange(test.shape[1], dtype=np.float32) * interval_time, - }, + } + return { + "metrics": metrics, + "snapshots": snapshots, + "shared": shared, } diff --git a/tests/test_dummy_integration.py b/tests/test_dummy_integration.py index e1428b88..c7b8db95 100644 --- a/tests/test_dummy_integration.py +++ b/tests/test_dummy_integration.py @@ -281,6 +281,7 @@ def _maybe_shrink(cfg, problem: str, exp_key: str) -> None: """ if problem == "ns-grid" and exp_key in { "optimization/solver_in_loop", + "optimization/solver_in_loop_self_reference", "optimization/solver_in_loop_tgv", }: from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( @@ -288,6 +289,7 @@ def _maybe_shrink(cfg, problem: str, exp_key: str) -> None: ) analytic = exp_key.endswith("_tgv") + self_reference = exp_key.endswith("_self_reference") cfg.add_experiment( exp_key, solver_in_loop, @@ -305,14 +307,23 @@ def _maybe_shrink(cfg, problem: str, exp_key: str) -> None: }, "dataset": { "reference_kind": ( - "analytic_tgv" if analytic else "pseudo_spectral_multimode" + "analytic_tgv" + if analytic + else ( + "solver_self_refined" + if self_reference + else "pseudo_spectral_multimode" + ) ), "reference_factor": 2, + "reference_temporal_factor": 2, "reference_substeps": 1, "train_seeds": [0], "test_seeds": [100], "train_frames": 2, "k0": 2.0, + "prefix_audit_seeds": [0, 100], + "prefix_audit_frames": [1, 2], }, "training": { "max_updates": 1, diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index 32ab1a02..c94b4494 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -7,6 +7,7 @@ import dataclasses from pathlib import Path +from types import SimpleNamespace import equinox as eqx import jax @@ -35,6 +36,7 @@ _save_solver_in_loop_animation, ) from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( + _make_solver_self_reference_datasets, _rollout_log_gain, make_reference_dataset, solver_in_loop, @@ -135,6 +137,59 @@ def test_analytic_tgv_reference_has_exact_decay_and_distinct_phases(): assert len(dataset_hash) == 16 +def test_solver_self_reference_matches_physical_time_and_passes_closure( + monkeypatch, +): + calls: list[tuple[int, float, int]] = [] + + def _closed_refined_step(_t, _ctx, velocity, *, dt, steps): + calls.append((velocity.shape[0], dt, steps)) + per_step_decay = 1.0 - dt / velocity.shape[0] + return jnp.asarray(velocity) * per_step_decay**steps + + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop." + "_solver_advance_with_physics", + _closed_refined_step, + ) + ctx = SimpleNamespace( + name="closed-dummy", + phys={"N": 8, "nu": 0.001, "dt": 0.02, "steps": 1}, + domain_extent=2.0 * np.pi, + ) + train, train_rollouts, test, dataset_hash, audit = ( + _make_solver_self_reference_datasets( + None, + ctx, + dataset={ + "reference_factor": 2, + "reference_temporal_factor": 2, + "train_seeds": [0, 1], + "test_seeds": [100], + "train_frames": 2, + "prefix_audit_seeds": [0, 100], + "prefix_audit_frames": [1, 2], + "k0": 2.0, + "minimum_refinement_signal": 1e-5, + }, + evaluation={"rollout_frames": 3}, + training={"unroll": 2}, + ) + ) + + assert train.shape == (2, 3, 8, 8, 1, 2) + assert train_rollouts.shape == (2, 4, 8, 8, 1, 2) + assert test.shape == (1, 4, 8, 8, 1, 2) + assert len(dataset_hash) == 16 + assert audit["eligible_for_corrector_training"] is True + assert audit["max_coarse_closure_error"] < 1e-5 + assert audit["max_fine_closure_error"] < 1e-5 + assert audit["mean_refinement_signal"] > 1e-5 + assert (16, 0.01, 2) in calls + assert (8, 0.02, 1) in calls + assert 0.01 * 2 == 0.02 * 1 + + def test_rollout_log_gain_is_geometric_and_ignores_initial_frame(): baseline = np.asarray([0.0, 4.0, 8.0]) corrected = np.asarray([0.0, 2.0, 2.0]) @@ -198,7 +253,7 @@ def test_solver_in_loop_fields_render_reference_raw_and_corrected(tmp_path): ) arrays = { - "reference_rollout": np.stack((reference, reference)), + "reference_rollout_0": np.stack((reference, reference)), "evaluation_times": np.asarray([0.0, 0.1]), "rollout_uncorrected_0": np.stack((reference, 0.8 * reference)), "rollout_corrected_0": np.stack((reference, 0.95 * reference)), @@ -370,3 +425,54 @@ def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): assert (out_dir / "result.json").exists() with np.load(out_dir / "corrector_fields.npz") as snapshots: assert snapshots["solver_vjp_log_lift_samples_0"].shape == (2, 1) + + +def test_solver_self_reference_skips_training_below_refinement_floor( + tmp_path, + monkeypatch, +): + """An admitted interface still skips training when refinement has no signal.""" + from mosaic.benchmarks.problems import get_config + + monkeypatch.setenv("MOSAIC_RESULTS_DIR", str(tmp_path)) + base = get_config("ns-grid") + pict = next(spec for spec in base.solvers if spec.key == "pict") + cfg = dataclasses.replace(base, solvers=[pict]) + cfg.add_experiment( + "optimization/solver_self_reference_ineligible_smoke", + solver_in_loop, + runs=[ + { + "ic": {"name": "multimode", "seed": 0}, + "physics": {"N": 8, "nu": 0.001, "dt": 0.02, "steps": 1}, + "dataset": { + "reference_kind": "solver_self_refined", + "reference_factor": 2, + "reference_temporal_factor": 2, + "train_seeds": [0], + "test_seeds": [100], + "train_frames": 1, + "prefix_audit_seeds": [0, 100], + "prefix_audit_frames": [1], + "k0": 2.0, + }, + "training": { + "max_updates": 1, + "unroll": 1, + "hidden_channels": 4, + "kernel_size": 3, + }, + "evaluation": {"rollout_frames": 1}, + } + ], + ) + result = cfg.experiments["optimization/solver_self_reference_ineligible_smoke"].fn( + cfg, + {pict.name: f"inprocess:{_IDENTITY_DUMMY}"}, + ) + + metrics = result["results"][0]["metrics"] + assert metrics["completed"] is True + assert metrics["eligible_for_corrector_training"] is False + assert metrics["mean_refinement_signal"] == 0 + assert metrics["n_updates"] == 0 From 9f15c6c63e44b26755c705e05baaacc754d02af7 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sat, 25 Jul 2026 13:34:17 +0200 Subject: [PATCH 07/15] fix(ns-grid): keep self-reference comparisons normalized --- .../problems/navier_stokes_grid/plots.py | 149 +++++++++++++----- 1 file changed, 106 insertions(+), 43 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 64c95570..6b9b5437 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -133,6 +133,20 @@ def plot_solver_in_loop( ) stopped = np.asarray(arrays.get(f"error_stop_gradient_{idx}", np.array([]))) uncorrected = np.asarray(arrays.get(f"error_uncorrected_{idx}", np.array([]))) + if solver_specific_reference and uncorrected.size: + # Each self-reference cell has a different target. Divide by that + # cell's raw trajectory so shared axes show within-solver correction + # gain rather than incomparable absolute target errors. + baseline = np.maximum(uncorrected, 1e-12) + if corrected.size: + corrected = corrected / baseline + if corrected_std.size: + corrected_std = corrected_std / baseline + if corrected_ic_std.size: + corrected_ic_std = corrected_ic_std / baseline + if stopped.size: + stopped = stopped / baseline + uncorrected = np.ones_like(uncorrected) if loss.size: updates = np.arange(1, loss.size + 1) ax_loss.plot( @@ -233,7 +247,7 @@ def plot_solver_in_loop( ax_roll.set_yscale("log") ax_roll.set_xlabel("Physical time") ax_roll.set_ylabel( - "Error to solver-specific reference" + "Error / solver-only error" if solver_specific_reference else "Relative $L^2$ error" ) @@ -280,17 +294,24 @@ def plot_solver_in_loop( ) if fairness_fig is not None: figs.append(fairness_fig) - physics_fig = _plot_solver_in_loop_physics(arrays, names, out_dir, save=save) - if physics_fig is not None: - figs.append(physics_fig) - diagnostics_fig = _plot_solver_in_loop_diagnostics( - data, + physics_fig = _plot_solver_in_loop_physics( + arrays, names, out_dir, save=save, + solver_specific_reference=solver_specific_reference, ) - if diagnostics_fig is not None: - figs.append(diagnostics_fig) + if physics_fig is not None: + figs.append(physics_fig) + if not solver_specific_reference: + diagnostics_fig = _plot_solver_in_loop_diagnostics( + data, + names, + out_dir, + save=save, + ) + if diagnostics_fig is not None: + figs.append(diagnostics_fig) if save: _save_solver_in_loop_animation(arrays, names, out_dir) return figs @@ -553,6 +574,7 @@ def _plot_solver_in_loop_physics( out_dir: Path, *, save: bool, + solver_specific_reference: bool = False, ) -> plt.Figure | None: """Compare energy, enstrophy, and divergence along held-out rollouts.""" rows = _ordered_solver_rollouts(arrays, names) @@ -590,7 +612,7 @@ def _plot_solver_in_loop_physics( diagnostic_axes = (ax_energy, ax_enstrophy, ax_spectral, ax_centered) present: list[str] = [] - for name, raw, corrected in rows: + for row_index, (name, raw, corrected) in enumerate(rows): _label, color, _linestyle, _marker = solver_props(name) reference_diagnostics = _trajectory_diagnostics( solver_references[name][:n_frames] @@ -624,13 +646,14 @@ def _plot_solver_in_loop_physics( corrected_curves, strict=True, ): - ax.plot( - times, - np.maximum(reference_curve, 1e-12), - color=color, - linestyle="--", - alpha=0.7, - ) + if solver_specific_reference or row_index == 0: + ax.plot( + times, + np.maximum(reference_curve, 1e-12), + color=color if solver_specific_reference else "0.25", + linestyle="--", + alpha=0.7, + ) ax.plot(times, np.maximum(raw_curve, 1e-12), color=color, linestyle=":") ax.plot( times, @@ -658,7 +681,15 @@ def _plot_solver_in_loop_physics( present, extra_handles=[ mlines.Line2D( - [], [], color="0.4", linestyle="--", label="solver-specific reference" + [], + [], + color="0.4", + linestyle="--", + label=( + "solver-specific reference" + if solver_specific_reference + else "shared reference" + ), ), mlines.Line2D([], [], color="0.4", linestyle=":", label="solver only"), mlines.Line2D([], [], color="0.4", linestyle="-", label="full corrector"), @@ -711,6 +742,8 @@ def _plot_solver_in_loop_fairness( labels: list[str] = [] all_errors: list[float] = [] + normalized_full_errors: list[float] = [] + normalized_stopped_errors: list[float] = [] full_gain: list[float] = [] stopped_gain: list[float] = [] vjp_lift: list[float] = [] @@ -727,8 +760,17 @@ def _plot_solver_in_loop_fairness( labels.append(label) raw = float(metrics["uncorrected_mean_rollout_error"]) corrected = float(metrics["mean_rollout_error"]) - all_errors.extend((raw, corrected)) - ax_quality.scatter(raw, corrected, color=color, marker=marker, s=30) + normalized_full_errors.append(corrected / max(raw, 1e-12)) + stopped_error = metrics.get("stop_gradient_mean_rollout_error") + has_counterfactual &= stopped_error is not None + normalized_stopped_errors.append( + float(stopped_error) / max(raw, 1e-12) + if stopped_error is not None + else np.nan + ) + if not solver_specific_reference: + all_errors.extend((raw, corrected)) + ax_quality.scatter(raw, corrected, color=color, marker=marker, s=30) update_time = metrics.get("median_update_time_s") has_cost &= update_time is not None update_times.append(float(update_time) if update_time is not None else np.nan) @@ -777,32 +819,53 @@ def _plot_solver_in_loop_fairness( vjp_lift_log_std.append(float(np.hypot(seed_lift_std, ic_lift_std))) has_ic_uncertainty |= "solver_vjp_log_lift_ic_std" in metrics - error_min = max(min(all_errors) * 0.8, 1e-4) - error_max = max(all_errors) * 1.25 - ax_quality.plot( - [error_min, error_max], - [error_min, error_max], - color="0.45", - linestyle=":", - linewidth=1.0, - ) - if error_max / error_min >= 10.0: - ax_quality.set_xscale("log") - ax_quality.set_yscale("log") - ax_quality.xaxis.set_minor_formatter(mticker.NullFormatter()) - ax_quality.yaxis.set_minor_formatter(mticker.NullFormatter()) - ax_quality.set_xlim(error_min, error_max) - ax_quality.set_ylim(error_min, error_max) - ax_quality.set_xlabel("Solver-only mean error") - ax_quality.set_ylabel("Corrected mean error") - ax_quality.set_title( - "Within-solver target error" - if solver_specific_reference - else "Absolute quality" - ) - width = 0.36 if has_counterfactual else 0.62 colors = [solver_props(name)[1] for name, _metrics in rows] + if solver_specific_reference: + ax_quality.bar( + positions - (width / 2 if has_counterfactual else 0.0), + normalized_full_errors, + width=width, + color=colors, + alpha=0.9, + label="full solver VJP", + ) + if has_counterfactual: + ax_quality.bar( + positions + width / 2, + normalized_stopped_errors, + width=width, + color=colors, + alpha=0.35, + hatch="//", + label="stop-gradient", + ) + ax_quality.legend(loc="best", fontsize=6.5) + ax_quality.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_quality.set_xticks(positions, labels, rotation=35, ha="right") + ax_quality.set_ylabel("Mean error / solver-only error") + ax_quality.set_title("Target-normalized quality") + else: + error_min = max(min(all_errors) * 0.8, 1e-4) + error_max = max(all_errors) * 1.25 + ax_quality.plot( + [error_min, error_max], + [error_min, error_max], + color="0.45", + linestyle=":", + linewidth=1.0, + ) + if error_max / error_min >= 10.0: + ax_quality.set_xscale("log") + ax_quality.set_yscale("log") + ax_quality.xaxis.set_minor_formatter(mticker.NullFormatter()) + ax_quality.yaxis.set_minor_formatter(mticker.NullFormatter()) + ax_quality.set_xlim(error_min, error_max) + ax_quality.set_ylim(error_min, error_max) + ax_quality.set_xlabel("Solver-only mean error") + ax_quality.set_ylabel("Corrected mean error") + ax_quality.set_title("Absolute quality") + full_error = _log_scale_errorbars(full_gain, full_gain_log_std) ax_gain.bar( positions - (width / 2 if has_counterfactual else 0.0), From 96546b0d83ebc973e067195df15f954899bc2588 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sat, 25 Jul 2026 13:36:18 +0200 Subject: [PATCH 08/15] test(ns-grid): guard solver-loop plot fairness --- .../problems/navier_stokes_grid/plots.py | 3 +- tests/test_solver_in_loop.py | 40 ++++++++++++++++++- 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 6b9b5437..597779b0 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -762,7 +762,8 @@ def _plot_solver_in_loop_fairness( corrected = float(metrics["mean_rollout_error"]) normalized_full_errors.append(corrected / max(raw, 1e-12)) stopped_error = metrics.get("stop_gradient_mean_rollout_error") - has_counterfactual &= stopped_error is not None + if solver_specific_reference: + has_counterfactual &= stopped_error is not None normalized_stopped_errors.append( float(stopped_error) / max(raw, 1e-12) if stopped_error is not None diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index c94b4494..9f885dee 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -287,6 +287,12 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): (reference, 0.85 * reference, 0.7 * reference) ), "rollout_corrected_0": np.stack((reference, 0.95 * reference, 0.9 * reference)), + "rollout_uncorrected_1": np.stack( + (reference, 0.9 * reference, 0.8 * reference) + ), + "rollout_corrected_1": np.stack( + (reference, 0.97 * reference, 0.94 * reference) + ), } data = { "by_solver": { @@ -326,7 +332,7 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): ) physics = _plot_solver_in_loop_physics( arrays, - ["jax-cfd"], + ["jax-cfd", "warp-ns"], tmp_path, save=True, ) @@ -341,6 +347,8 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): assert fairness is not None assert physics is not None assert diagnostics is not None + # One shared reference plus raw/corrected curves for each solver. + assert all(len(axis.lines) == 5 for axis in physics.axes[:4]) for filename in ( "solver_in_loop_diagnostics.png", "solver_in_loop_fairness.png", @@ -352,6 +360,36 @@ def test_solver_in_loop_fairness_physics_and_animation_render(tmp_path): assert rendered.stat().st_size > 0 +def test_self_reference_fairness_normalizes_solver_specific_targets(tmp_path): + data = { + "by_solver": { + "pict": { + "uncorrected_mean_rollout_error": 0.4, + "mean_rollout_error": 0.2, + "stop_gradient_mean_rollout_error": 0.3, + "geometric_error_reduction": 2.0, + "stop_gradient_geometric_error_reduction": 4.0 / 3.0, + "solver_vjp_geometric_lift": 1.5, + } + } + } + + figure = _plot_solver_in_loop_fairness( + data, + ["pict"], + tmp_path, + save=False, + solver_specific_reference=True, + ) + + assert figure is not None + np.testing.assert_allclose( + [patch.get_height() for patch in figure.axes[0].patches], + [0.5, 0.75], + ) + assert figure.axes[0].get_title() == "Target-normalized quality" + + def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): """One update crosses the apply/VJP boundary and writes canonical artifacts.""" from mosaic.benchmarks.problems import get_config From eb7492e8950dd7937269bff916450883e78dddb0 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 16:14:03 +0200 Subject: [PATCH 09/15] fix(ns-grid): align solver-loop forward control --- .../problems/navier_stokes_grid/config.py | 26 +++++++++---------- .../problems/navier_stokes_grid/exclusions.py | 14 +++++----- .../navier-stokes-grid/pict/tesseract_api.py | 8 +++--- tests/test_solver_in_loop.py | 18 +++++++++++++ 4 files changed, 42 insertions(+), 24 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index 780c19cf..cf6b3dd9 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -526,10 +526,10 @@ "optimization/solver_in_loop_tgv", solver_in_loop, description=( - "Repeat solver-in-the-loop training against an analytic, translated " - "Taylor--Green vortex. This low-complexity regime audits physical time, " - "viscous decay, and canonical state-restart effects separately from the " - "nonlinear multimode comparison." + "Run a matched-forward-accuracy solver-in-the-loop control using the " + "same Taylor--Green setup as forward/agreement. This low-complexity " + "regime separates raw solver accuracy from corrector trainability; the " + "nonlinear multimode experiment remains the stress task." ), plot_description=( "Native and restarted solver error, neural-correction gain, solver-VJP " @@ -539,22 +539,22 @@ { "ic": {"name": "tgv", "seed": 0}, "physics": { - "N": 32, - "nu": 0.05, - "dt": 0.02, - "steps": 4, + "N": 64, + "nu": 0.005, + "dt": 0.05, + "steps": 20, }, "dataset": { "reference_kind": "analytic_tgv", - "reference_factor": 2, + "reference_factor": 1, "reference_substeps": 1, "train_seeds": [0, 1, 2, 3], "test_seeds": [100, 101], - "train_frames": 16, + "train_frames": 8, }, "training": { - "max_updates": 150, - "unroll": 8, + "max_updates": 50, + "unroll": 4, "loss_mode": "solver_terminal", "solver_loss_weight": 0.1, "loss_normalization": "solver_baseline", @@ -570,7 +570,7 @@ "fd_epsilon": 1e-2, }, "evaluation": { - "rollout_frames": 24, + "rollout_frames": 12, "seen_ic_trajectories": 2, "stable_error_threshold": 1.0, }, diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py index 6f9bac05..c4e0e34f 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py @@ -106,6 +106,12 @@ "ins_jl": INS_JL_NO_OBSTACLE, "warp_ns": WARP_NS_NO_OBSTACLE, } +_RECURRENT_STATE_GATE = { + "jax_cfd": STAGGERED_STATE_NOT_CLOSED, + "ins_jl": STAGGERED_STATE_NOT_CLOSED, + "phiflow": STAGGERED_STATE_NOT_CLOSED, + "xlb": XLB_STATE_NOT_CLOSED, +} def register(problem: Problem) -> None: @@ -136,12 +142,8 @@ def register(problem: Problem) -> None: problem.exclude("optimization", {"openfoam": OPENFOAM_NON_DIFFERENTIABLE_OPT}) problem.exclude("optimization/drag_opt", _OBSTACLE_GATE) problem.exclude("optimization/drag_opt_bfgs", _OBSTACLE_GATE) + problem.exclude("optimization/solver_in_loop_tgv", _RECURRENT_STATE_GATE) problem.exclude( "optimization/solver_in_loop_self_reference", - { - "jax_cfd": STAGGERED_STATE_NOT_CLOSED, - "ins_jl": STAGGERED_STATE_NOT_CLOSED, - "phiflow": STAGGERED_STATE_NOT_CLOSED, - "xlb": XLB_STATE_NOT_CLOSED, - }, + _RECURRENT_STATE_GATE, ) diff --git a/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_api.py b/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_api.py index 487cba13..11ef2b43 100644 --- a/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_api.py +++ b/mosaic/tesseracts/navier-stokes-grid/pict/tesseract_api.py @@ -1137,11 +1137,9 @@ def _post_step(domain: PISOtorch.Domain, time_step: Any, **_kw: Any) -> None: domain, substeps=1, time_step=float(dt_val), - # PISO is conventionally run with 2 pressure correctors (Issa 1986); - # 4 was a conservative default and roughly doubles per-step wall. - # Re=20 cylinder & decaying-TGV regression checks remained within - # the existing status_checks bounds, so we run with 2 by default. - corrector_steps=2, + # Four pressure correctors preserve the forward-agreement accuracy + # contract on current CUDA runtimes; two under-converge the projection. + corrector_steps=4, non_orthogonal=False, differentiable=differentiable, log_dir=None, diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index 9f885dee..df9751bf 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -6,6 +6,7 @@ from __future__ import annotations import dataclasses +import inspect from pathlib import Path from types import SimpleNamespace @@ -50,6 +51,23 @@ ).resolve() +def test_tgv_control_uses_a_forward_agreement_cell(): + from mosaic.benchmarks.problems import get_config + + cfg = get_config("ns-grid") + forward = cfg.experiments["forward/agreement/tgv"] + control = cfg.experiments["optimization/solver_in_loop_tgv"] + forward_run = inspect.signature(forward.fn).parameters["_kw"].default["runs"][0] + control_run = inspect.signature(control.fn).parameters["_kw"].default["runs"][0] + forward_physics = forward_run["physics"] + control_physics = control_run["physics"] + + for key in ("N", "dt", "steps"): + assert control_physics[key] == forward_physics[key] + assert forward_run["sweep"]["key"] == "nu" + assert control_physics["nu"] in forward_run["sweep"]["values"] + + def test_periodic_corrector_is_translation_equivariant(): model = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) velocity = jax.random.normal(jax.random.PRNGKey(1), (8, 8, 1, 2)) From 875718fbfa69a16d2bd301ad11b0760743ea5aa2 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 16:26:29 +0200 Subject: [PATCH 10/15] test: isolate dummy experiment configs --- tests/test_dummy_integration.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_dummy_integration.py b/tests/test_dummy_integration.py index c7b8db95..55c6707e 100644 --- a/tests/test_dummy_integration.py +++ b/tests/test_dummy_integration.py @@ -25,6 +25,7 @@ from __future__ import annotations import contextlib +import copy import json import warnings from pathlib import Path @@ -275,8 +276,8 @@ def test_thermal_mesh_forward_runs_with_dummy(thermal_mesh_tags, tmp_path, monke def _maybe_shrink(cfg, problem: str, exp_key: str) -> None: """Re-register ``exp_key`` at a smaller N if it's a known heavy case. - Mutates ``cfg.experiments[exp_key]`` in place (the config object is a - fresh per-call ``get_config(problem)``, so this never leaks across tests). + Mutates ``cfg.experiments[exp_key]`` in place. Callers must pass an + experiment-local copy so the smaller smoke configuration cannot leak. No-op for experiments not in :data:`_SHRINK_PHYSICS`. """ if problem == "ns-grid" and exp_key in { @@ -449,7 +450,7 @@ def dummy_corpus(tmp_path_factory): try: with _suppress_dummy_warnings(): for problem, exp_key in _all_experiments(): - cfg = get_config(problem) + cfg = copy.deepcopy(get_config(problem)) _maybe_shrink(cfg, problem, exp_key) tag = f"inprocess:{_DUMMY_FOR[problem]}" tags = dict.fromkeys(cfg.solver_names, tag) From 65b30e15fbb64afa67a5392a6a8f6c1dd9c75708 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 17:31:37 +0200 Subject: [PATCH 11/15] fix(ns-grid): gate solver VJP comparison plots --- .../problems/navier_stokes_grid/plots.py | 43 ++++++++++++++----- tests/test_solver_in_loop.py | 30 +++++++++++++ 2 files changed, 62 insertions(+), 11 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 597779b0..46e6b2b6 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -897,36 +897,57 @@ def _plot_solver_in_loop_fairness( ax_gain.set_title("Correctability") if has_counterfactual: - vjp_lift_pct = [100.0 * (value - 1.0) for value in vjp_lift] + admitted = [ + idx + for idx, (_name, metrics) in enumerate(rows) + if metrics.get("valid_for_vjp_ranking", True) + ] + admitted_positions = np.arange(len(admitted), dtype=float) + admitted_labels = [labels[idx] for idx in admitted] + admitted_colors = [colors[idx] for idx in admitted] + vjp_lift_pct = [100.0 * (vjp_lift[idx] - 1.0) for idx in admitted] vjp_error = _log_scale_errorbars(vjp_lift, vjp_lift_log_std) if vjp_error is not None: - vjp_error = 100.0 * vjp_error + vjp_error = 100.0 * vjp_error[:, admitted] ax_vjp.bar( - positions, + admitted_positions, vjp_lift_pct, - color=colors, + color=admitted_colors, alpha=0.9, yerr=vjp_error, capsize=2, ) ax_vjp.axhline(0.0, color="0.45", linestyle=":", linewidth=1.0) - ax_vjp.set_xticks(positions, labels, rotation=35, ha="right") + ax_vjp.set_xticks( + admitted_positions, + admitted_labels, + rotation=35, + ha="right", + ) ax_vjp.set_ylabel("Solver-VJP lift [%]") - ax_vjp.set_title("Benefit from solver VJP") + title_suffix = " (admitted only)" if len(admitted) < len(rows) else "" + ax_vjp.set_title(f"Benefit from solver VJP{title_suffix}") if has_cost: - for idx, (name, _metrics) in enumerate(rows): + for ranked_idx, idx in enumerate(admitted): + name, _metrics = rows[idx] _label, color, _linestyle, marker = solver_props(name) - error = None if vjp_error is None else vjp_error[:, idx].reshape(2, 1) + error = ( + None + if vjp_error is None + else vjp_error[:, ranked_idx].reshape(2, 1) + ) ax_efficiency.errorbar( update_times[idx], - vjp_lift_pct[idx], + vjp_lift_pct[ranked_idx], yerr=error, color=color, marker=marker, linestyle="none", capsize=2, ) - positive_times = [value for value in update_times if value > 0] + positive_times = [ + update_times[idx] for idx in admitted if update_times[idx] > 0 + ] if positive_times and max(positive_times) / min(positive_times) >= 10.0: ax_efficiency.set_xscale("log") ax_efficiency.axhline( @@ -937,7 +958,7 @@ def _plot_solver_in_loop_fairness( ) ax_efficiency.set_xlabel("Full-VJP update time [s]") ax_efficiency.set_ylabel("Solver-VJP lift [%]") - ax_efficiency.set_title("VJP benefit versus cost") + ax_efficiency.set_title(f"VJP benefit versus cost{title_suffix}") else: ax_efficiency.axis("off") else: diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index df9751bf..1ad8e85d 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -408,6 +408,36 @@ def test_self_reference_fairness_normalizes_solver_specific_targets(tmp_path): assert figure.axes[0].get_title() == "Target-normalized quality" +def test_solver_vjp_panels_only_show_recurrence_admitted_cells(tmp_path): + common = { + "uncorrected_mean_rollout_error": 0.4, + "mean_rollout_error": 0.2, + "stop_gradient_mean_rollout_error": 0.22, + "geometric_error_reduction": 2.0, + "stop_gradient_geometric_error_reduction": 1.8, + "solver_vjp_geometric_lift": 2.0 / 1.8, + "median_update_time_s": 1.0, + } + data = { + "by_solver": { + "jax-cfd": {**common, "valid_for_vjp_ranking": False}, + "pict": {**common, "valid_for_vjp_ranking": True}, + } + } + + figure = _plot_solver_in_loop_fairness( + data, + ["jax-cfd", "pict"], + tmp_path, + save=False, + ) + + assert figure is not None + assert len(figure.axes[2].patches) == 1 + assert [tick.get_text() for tick in figure.axes[2].get_xticklabels()] == ["PICT"] + assert figure.axes[2].get_title() == "Benefit from solver VJP (admitted only)" + + def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): """One update crosses the apply/VJP boundary and writes canonical artifacts.""" from mosaic.benchmarks.problems import get_config From 9a347793acb612df92427a547fe1460421bf201a Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 17:54:24 +0200 Subject: [PATCH 12/15] fix(ns-grid): select the correct forward reference --- .../problems/navier_stokes_grid/config.py | 34 ++++++++++++++----- mosaic/benchmarks/problems/shared/forward.py | 11 ++++-- tests/test_solver_in_loop.py | 10 ++++++ 3 files changed, 44 insertions(+), 11 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index cf6b3dd9..6be27e4b 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -212,16 +212,32 @@ "forward/agreement", agreement, plot_description=( - "Relative error vs viscosity $\\nu$ for each IC, with vorticity field snapshots" - " compared against the analytic TGV reference." + "Relative error vs viscosity $\\nu$ for each IC, using the analytic " + "Taylor--Green reference for TGV and cross-solver consensus for multimode." ), - ic=[{"name": "tgv", "seed": 42}, {"name": "multimode", "seed": 42}], - physics={ - "N": 64, - "dt": 0.05, - "steps": 20, - "nu": [0.001, 0.005, 0.01, 0.02, 0.05], - }, + runs=[ + { + "name": "tgv", + "ic": {"name": "tgv", "seed": 42}, + "physics": { + "N": 64, + "dt": 0.05, + "steps": 20, + "nu": [0.001, 0.005, 0.01, 0.02, 0.05], + }, + }, + { + "name": "multimode", + "ic": {"name": "multimode", "seed": 42}, + "physics": { + "N": 64, + "dt": 0.05, + "steps": 20, + "nu": [0.001, 0.005, 0.01, 0.02, 0.05], + }, + "reference": "consensus", + }, + ], plot=plot_agreement, ) problem.add_experiment( diff --git a/mosaic/benchmarks/problems/shared/forward.py b/mosaic/benchmarks/problems/shared/forward.py index 3f362a11..4d8ca0e1 100644 --- a/mosaic/benchmarks/problems/shared/forward.py +++ b/mosaic/benchmarks/problems/shared/forward.py @@ -85,6 +85,14 @@ def _analytic_reference( # ── agreement ──────────────────────────────────────────────────────────────── +def _agreement_analytic_reference(run: dict, fallback: Any) -> Any: + """Resolve a per-run analytic reference, allowing explicit consensus.""" + reference = run.get("reference") + if reference == "consensus": + return None + return reference if callable(reference) else fallback + + def _agreement_aggregate( by_solver: Any, *, @@ -114,8 +122,7 @@ def _agreement_aggregate( seed = ic_cfg.get("seed", 0) phys = run.get("physics", {}) - run_reference = run.get("reference") - analytic_fn = run_reference if callable(run_reference) else cfg.reference + analytic_fn = _agreement_analytic_reference(run, cfg.reference) analytic_params = ( set(inspect.signature(analytic_fn).parameters) if analytic_fn else set() ) diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index 1ad8e85d..551d85cc 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -68,6 +68,16 @@ def test_tgv_control_uses_a_forward_agreement_cell(): assert control_physics["nu"] in forward_run["sweep"]["values"] +def test_multimode_forward_agreement_does_not_use_tgv_reference(): + from mosaic.benchmarks.problems import get_config + + cfg = get_config("ns-grid") + experiment = cfg.experiments["forward/agreement/multimode"] + run = inspect.signature(experiment.fn).parameters["_kw"].default["runs"][0] + + assert run["reference"] == "consensus" + + def test_periodic_corrector_is_translation_equivariant(): model = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) velocity = jax.random.normal(jax.random.PRNGKey(1), (8, 8, 1, 2)) From dba83f8801fb8354696bfdf719f063b131d51b1e Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 18:27:52 +0200 Subject: [PATCH 13/15] fix(ns-grid): require all-solver recurrent admission --- .../problems/navier_stokes_grid/config.py | 27 +++-- .../problems/navier_stokes_grid/corrector.py | 23 ++-- .../problems/navier_stokes_grid/exclusions.py | 1 - .../problems/navier_stokes_grid/plots.py | 22 +++- .../navier_stokes_grid/solver_in_loop.py | 100 ++++++++++++++---- tests/test_solver_in_loop.py | 42 ++++++-- 6 files changed, 166 insertions(+), 49 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index 6be27e4b..19f449d4 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -542,10 +542,11 @@ "optimization/solver_in_loop_tgv", solver_in_loop, description=( - "Run a matched-forward-accuracy solver-in-the-loop control using the " - "same Taylor--Green setup as forward/agreement. This low-complexity " - "regime separates raw solver accuracy from corrector trainability; the " - "nonlinear multimode experiment remains the stress task." + "Run an all-solver solver-in-the-loop control using the N=64 cell from " + "forward/baseline, where every differentiable solver has less than 0.75% " + "one-step error. Repeated calls expose recurrent-state defects instead of " + "hiding them behind solver exclusions; the nonlinear multimode experiment " + "remains the learning stress task." ), plot_description=( "Native and restarted solver error, neural-correction gain, solver-VJP " @@ -556,9 +557,10 @@ "ic": {"name": "tgv", "seed": 0}, "physics": { "N": 64, - "nu": 0.005, - "dt": 0.05, - "steps": 20, + "nu": 0.05, + "dt": 0.01, + "steps": 1, + "lbm_N_base": 64, }, "dataset": { "reference_kind": "analytic_tgv", @@ -566,7 +568,8 @@ "reference_substeps": 1, "train_seeds": [0, 1, 2, 3], "test_seeds": [100, 101], - "train_frames": 8, + "train_frames": 24, + "long_closure_tolerance": 0.01, }, "training": { "max_updates": 50, @@ -575,7 +578,9 @@ "solver_loss_weight": 0.1, "loss_normalization": "solver_baseline", "loss_scale_floor": 1e-6, - "lr": 1e-4, + # Calibrated on training loss only. A nonlinear-task-size Adam + # step overwhelms this near-floor, per-native-step residual. + "lr": 1e-9, "clip_norm": 5.0, "architecture": "periodic_residual_cnn", "hidden_channels": 32, @@ -586,9 +591,11 @@ "fd_epsilon": 1e-2, }, "evaluation": { - "rollout_frames": 12, + "rollout_frames": 100, "seen_ic_trajectories": 2, "stable_error_threshold": 1.0, + "first_interval_error_tolerance": 0.01, + "native_long_error_tolerance": 0.01, }, } ], diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py index 4f0d20ef..1a6a1a4b 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py @@ -50,14 +50,17 @@ def __init__( padding_mode="CIRCULAR", key=layer_key, ) - fan_in = kernel_size * kernel_size * cin - scale = np.sqrt(2.0 / fan_in) if idx == len(channels) - 2: - scale *= 1e-2 - weight = ( - jax.random.normal(layer_key, layer.weight.shape, dtype=jnp.float32) - * scale - ) + # Start from the underlying solver exactly. This matters when + # its forward residual is already near numerical precision. + weight = jnp.zeros_like(layer.weight) + else: + fan_in = kernel_size * kernel_size * cin + scale = np.sqrt(2.0 / fan_in) + weight = ( + jax.random.normal(layer_key, layer.weight.shape, dtype=jnp.float32) + * scale + ) layer = eqx.tree_at( lambda value: (value.weight, value.bias), layer, @@ -89,9 +92,9 @@ def init_corrector( ) -> PeriodicResidualCNN: """Initialise the benchmark-owned Equinox corrector. - The final layer starts at a much smaller scale than the feature layers, so - the initial corrected rollout remains close to the underlying solver while - gradients can still reach every layer on the first update. + The final layer starts at zero, so the initial corrected rollout is exactly + the underlying solver. The first update trains the readout; subsequent + updates propagate through the feature layers. """ if architecture != "periodic_residual_cnn": raise ValueError(f"unknown corrector architecture: {architecture!r}") diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py index c4e0e34f..77d38539 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py @@ -142,7 +142,6 @@ def register(problem: Problem) -> None: problem.exclude("optimization", {"openfoam": OPENFOAM_NON_DIFFERENTIABLE_OPT}) problem.exclude("optimization/drag_opt", _OBSTACLE_GATE) problem.exclude("optimization/drag_opt_bfgs", _OBSTACLE_GATE) - problem.exclude("optimization/solver_in_loop_tgv", _RECURRENT_STATE_GATE) problem.exclude( "optimization/solver_in_loop_self_reference", _RECURRENT_STATE_GATE, diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 46e6b2b6..9d55385a 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -1081,7 +1081,25 @@ def values(key: str) -> list[float]: ax_restart.set_ylabel("Repeated-call final error") ax_restart.set_title("State-coupling effect") - if all(metrics.get("semigroup_error_p95") is not None for _name, metrics in rows): + if all( + metrics.get("long_closure_error_p95") is not None for _name, metrics in rows + ): + closure_pct = 100.0 * np.asarray(values("long_closure_error_p95")) + tolerance_pct = 100.0 * max(values("long_closure_tolerance")) + ax_restart_ratio.bar(positions, closure_pct, color=colors) + ax_restart_ratio.axhline( + tolerance_pct, + color="0.45", + linestyle=":", + linewidth=1.0, + label="admission limit", + ) + positive = closure_pct[closure_pct > 0] + if positive.size and max(positive) / min(positive) >= 10.0: + ax_restart_ratio.set_yscale("log") + ax_restart_ratio.set_ylabel("95th percentile residual [%]") + ax_restart_ratio.set_title("Long-horizon closure admission") + elif all(metrics.get("semigroup_error_p95") is not None for _name, metrics in rows): closure_pct = 100.0 * np.asarray(values("semigroup_error_p95")) tolerance_pct = 100.0 * max(values("semigroup_p95_tolerance")) ax_restart_ratio.bar(positions, closure_pct, color=colors) @@ -1096,7 +1114,7 @@ def values(key: str) -> list[float]: if positive.size and max(positive) / min(positive) >= 10.0: ax_restart_ratio.set_yscale("log") ax_restart_ratio.set_ylabel("95th percentile residual [%]") - ax_restart_ratio.set_title("Semigroup admission") + ax_restart_ratio.set_title("Two-step closure admission") else: ax_restart_ratio.bar( positions, diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py index c6646478..788488b1 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py @@ -16,7 +16,7 @@ import threading import time from functools import lru_cache, partial -from typing import Any +from typing import Any, NamedTuple import equinox as eqx import jax @@ -674,6 +674,12 @@ def _evaluate_rollout( return np.stack(states), errors +class _ReferenceEvaluation(NamedTuple): + first_rollout: np.ndarray | None + errors: np.ndarray + final_states: np.ndarray + + def _evaluate_reference_set( t: Any, ctx: KernelContext, @@ -683,9 +689,10 @@ def _evaluate_reference_set( frame_steps: int, velocity_scale: float, corrected: bool, -) -> tuple[np.ndarray | None, np.ndarray]: - """Evaluate one model on several ICs while retaining only the first rollout.""" +) -> _ReferenceEvaluation: + """Evaluate several ICs, retaining the first rollout and every final state.""" errors: list[list[float]] = [] + final_states: list[np.ndarray] = [] first_rollout: np.ndarray | None = None for reference in references: rollout, reference_errors = _evaluate_rollout( @@ -698,9 +705,14 @@ def _evaluate_reference_set( corrected=corrected, ) errors.append(reference_errors) + final_states.append(rollout[-1]) if first_rollout is None: first_rollout = rollout - return first_rollout, np.asarray(errors, dtype=np.float64) + return _ReferenceEvaluation( + first_rollout=first_rollout, + errors=np.asarray(errors, dtype=np.float64), + final_states=np.asarray(final_states), + ) def _first_unstable(errors: list[float], threshold: float) -> int: @@ -916,7 +928,7 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: semigroup_median <= semigroup_median_tolerance and semigroup_p95 <= semigroup_p95_tolerance ) - first_uncorrected, uncorrected_errors_array = _evaluate_reference_set( + uncorrected_eval = _evaluate_reference_set( t, ctx, None, @@ -925,7 +937,10 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: velocity_scale=velocity_scale, corrected=False, ) - _first_seen_uncorrected, seen_uncorrected_errors = _evaluate_reference_set( + first_uncorrected = uncorrected_eval.first_rollout + uncorrected_errors_array = uncorrected_eval.errors + restarted_final_states = uncorrected_eval.final_states + seen_uncorrected_eval = _evaluate_reference_set( t, ctx, None, @@ -934,8 +949,10 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: velocity_scale=velocity_scale, corrected=False, ) + seen_uncorrected_errors = seen_uncorrected_eval.errors native_final_errors: list[float] = [] - for reference in test: + long_closure_errors: list[float] = [] + for reference, restarted_final in zip(test, restarted_final_states, strict=True): native_final = _solver_advance( t, ctx, @@ -943,6 +960,37 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: frame_steps=frame_steps * (reference.shape[0] - 1), ) native_final_errors.append(relative_l2(np.asarray(native_final), reference[-1])) + long_closure_errors.append( + relative_l2(np.asarray(restarted_final), np.asarray(native_final)) + ) + + first_interval_error_p95 = float( + np.percentile(uncorrected_errors_array[:, 1], 95.0) + ) + long_closure_median = float(np.median(long_closure_errors)) + long_closure_p95 = float(np.percentile(long_closure_errors, 95.0)) + native_final_error_p95 = float(np.percentile(native_final_errors, 95.0)) + long_closure_tolerance = float( + dataset_cfg.get("long_closure_tolerance", semigroup_p95_tolerance) + ) + first_interval_error_tolerance = float( + evaluation.get( + "first_interval_error_tolerance", + evaluation.get("stable_error_threshold", 1.0), + ) + ) + native_long_error_tolerance = float( + evaluation.get( + "native_long_error_tolerance", + evaluation.get("stable_error_threshold", 1.0), + ) + ) + valid_for_vjp_ranking = bool( + valid_for_vjp_ranking + and first_interval_error_p95 <= first_interval_error_tolerance + and long_closure_p95 <= long_closure_tolerance + and native_final_error_p95 <= native_long_error_tolerance + ) uncorrected_error = np.mean(uncorrected_errors_array, axis=0) uncorrected_error_ic_std = np.std(uncorrected_errors_array, axis=0) @@ -1029,7 +1077,7 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: ) stop_gradient_training_walls.append(time.perf_counter() - stop_gradient_started) - corrected_rollout, seed_corrected_errors = _evaluate_reference_set( + corrected_eval = _evaluate_reference_set( t, ctx, model, @@ -1038,7 +1086,9 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: velocity_scale=velocity_scale, corrected=True, ) - stop_gradient_rollout, seed_stop_gradient_errors = _evaluate_reference_set( + corrected_rollout = corrected_eval.first_rollout + seed_corrected_errors = corrected_eval.errors + stop_gradient_eval = _evaluate_reference_set( t, ctx, stop_gradient_model, @@ -1047,7 +1097,9 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: velocity_scale=velocity_scale, corrected=True, ) - _seen_corrected_rollout, seen_seed_corrected_errors = _evaluate_reference_set( + stop_gradient_rollout = stop_gradient_eval.first_rollout + seed_stop_gradient_errors = stop_gradient_eval.errors + seen_corrected_eval = _evaluate_reference_set( t, ctx, model, @@ -1056,17 +1108,17 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: velocity_scale=velocity_scale, corrected=True, ) - _seen_stop_gradient_rollout, seen_seed_stop_gradient_errors = ( - _evaluate_reference_set( - t, - ctx, - stop_gradient_model, - seen_references, - frame_steps=frame_steps, - velocity_scale=velocity_scale, - corrected=True, - ) + seen_seed_corrected_errors = seen_corrected_eval.errors + seen_stop_gradient_eval = _evaluate_reference_set( + t, + ctx, + stop_gradient_model, + seen_references, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, ) + seen_seed_stop_gradient_errors = seen_stop_gradient_eval.errors if seed_idx == 0: first_corrected = corrected_rollout first_stop_gradient = stop_gradient_rollout @@ -1210,6 +1262,14 @@ def solver_in_loop(t: Any, ctx: KernelContext) -> dict: "semigroup_error_samples": semigroup_errors, "semigroup_median_tolerance": semigroup_median_tolerance, "semigroup_p95_tolerance": semigroup_p95_tolerance, + "first_interval_error_p95": first_interval_error_p95, + "first_interval_error_tolerance": first_interval_error_tolerance, + "long_closure_error_median": long_closure_median, + "long_closure_error_p95": long_closure_p95, + "long_closure_error_samples": long_closure_errors, + "long_closure_tolerance": long_closure_tolerance, + "native_final_rollout_error_p95": native_final_error_p95, + "native_long_error_tolerance": native_long_error_tolerance, "final_rollout_error": final_corrected, "final_rollout_error_seed_std": float(corrected_error_seed_std[-1]), "final_rollout_error_ic_std": float(corrected_error_ic_std[-1]), diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index 551d85cc..b9d0f78b 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -15,7 +15,7 @@ import jax.numpy as jnp import numpy as np -from mosaic.benchmarks.core.utils import _debug_run +from mosaic.benchmarks.core.utils import _debug_run, active_solvers from mosaic.benchmarks.problems.navier_stokes_grid.corrector import ( PeriodicResidualCNN, apply_corrector, @@ -51,21 +51,33 @@ ).resolve() -def test_tgv_control_uses_a_forward_agreement_cell(): +def test_tgv_control_uses_all_solver_forward_baseline_cell(): from mosaic.benchmarks.problems import get_config cfg = get_config("ns-grid") - forward = cfg.experiments["forward/agreement/tgv"] + forward = cfg.experiments["forward/baseline"] control = cfg.experiments["optimization/solver_in_loop_tgv"] forward_run = inspect.signature(forward.fn).parameters["_kw"].default["runs"][0] control_run = inspect.signature(control.fn).parameters["_kw"].default["runs"][0] forward_physics = forward_run["physics"] control_physics = control_run["physics"] + control_dataset = control_run["dataset"] + control_evaluation = control_run["evaluation"] - for key in ("N", "dt", "steps"): + for key in ("nu", "dt", "steps", "lbm_N_base"): assert control_physics[key] == forward_physics[key] - assert forward_run["sweep"]["key"] == "nu" - assert control_physics["nu"] in forward_run["sweep"]["values"] + assert forward_run["sweep"]["key"] == "N" + assert control_physics["N"] in forward_run["sweep"]["values"] + differentiable = {"jax_cfd", "ins_jl", "phiflow", "pict", "warp_ns", "xlb"} + active_forward = { + cfg.solver(name).key for name in active_solvers(cfg, "forward", "baseline") + } + assert differentiable <= active_forward + for solver in differentiable: + assert "optimization/solver_in_loop_tgv" not in cfg.exclusions.get(solver, {}) + assert control_evaluation["rollout_frames"] == 100 + assert control_dataset["long_closure_tolerance"] == 0.01 + assert control_evaluation["native_long_error_tolerance"] == 0.01 def test_multimode_forward_agreement_does_not_use_tgv_reference(): @@ -80,6 +92,11 @@ def test_multimode_forward_agreement_does_not_use_tgv_reference(): def test_periodic_corrector_is_translation_equivariant(): model = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) + model = eqx.tree_at( + lambda value: value.layers[-1].weight, + model, + jax.random.normal(jax.random.PRNGKey(2), model.layers[-1].weight.shape), + ) velocity = jax.random.normal(jax.random.PRNGKey(1), (8, 8, 1, 2)) shifted = jnp.roll(velocity, shift=(2, -1), axis=(0, 1)) @@ -96,6 +113,16 @@ def test_periodic_corrector_is_translation_equivariant(): np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) +def test_corrector_starts_from_the_uncorrected_solver(): + model = init_corrector(jax.random.PRNGKey(0), hidden_channels=4, kernel_size=3) + velocity = jax.random.normal(jax.random.PRNGKey(1), (8, 8, 1, 2)) + + np.testing.assert_array_equal( + apply_corrector(model, velocity, velocity_scale=1.0), + jnp.zeros_like(velocity), + ) + + def test_periodic_projection_is_divergence_free_and_zero_mean(): delta = jax.random.normal(jax.random.PRNGKey(2), (16, 16, 1, 2)) projected = project_periodic_correction(delta, 2.0 * np.pi) @@ -505,6 +532,9 @@ def test_solver_in_loop_runs_recurrently_through_dummy(tmp_path, monkeypatch): assert metrics["final_grad_norm"] > 0 assert metrics["end_to_end_fd_rel_error"] < 5e-2 assert metrics["native_final_rollout_error"] >= 0 + assert metrics["native_final_rollout_error_p95"] >= 0 + assert metrics["long_closure_error_p95"] < 1e-6 + assert metrics["first_interval_error_p95"] >= 0 assert metrics["solver_vjp_geometric_lift"] > 0 assert metrics["solver_vjp_update_overhead_ratio"] > 0 assert metrics["corrector_architecture"] == "periodic_residual_cnn" From feab2f8404e2e22a872e5f952add7616395b8e0d Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 18:35:47 +0200 Subject: [PATCH 14/15] fix(ns-grid): keep closure diagnostics legible --- mosaic/benchmarks/problems/navier_stokes_grid/plots.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py index 9d55385a..26089613 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/plots.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/plots.py @@ -1079,7 +1079,7 @@ def values(key: str) -> list[float]: ax_restart.set(xlim=(lower, upper), ylim=(lower, upper)) ax_restart.set_xlabel("Native final error") ax_restart.set_ylabel("Repeated-call final error") - ax_restart.set_title("State-coupling effect") + ax_restart.set_title("Restart vs native") if all( metrics.get("long_closure_error_p95") is not None for _name, metrics in rows @@ -1098,7 +1098,7 @@ def values(key: str) -> list[float]: if positive.size and max(positive) / min(positive) >= 10.0: ax_restart_ratio.set_yscale("log") ax_restart_ratio.set_ylabel("95th percentile residual [%]") - ax_restart_ratio.set_title("Long-horizon closure admission") + ax_restart_ratio.set_title("K-step closure gate") elif all(metrics.get("semigroup_error_p95") is not None for _name, metrics in rows): closure_pct = 100.0 * np.asarray(values("semigroup_error_p95")) tolerance_pct = 100.0 * max(values("semigroup_p95_tolerance")) From 7d6a683a11e8b60fb8da90b3a5a890abc5e75a64 Mon Sep 17 00:00:00 2001 From: andrinr Date: Sun, 26 Jul 2026 19:08:48 +0200 Subject: [PATCH 15/15] fix(ns-grid): calibrate all-solver forward control --- .../problems/navier_stokes_grid/config.py | 19 +++++++++------ tests/test_solver_in_loop.py | 24 +++++++++++++++++-- 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index 19f449d4..04fca0a6 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/config.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/config.py @@ -542,11 +542,12 @@ "optimization/solver_in_loop_tgv", solver_in_loop, description=( - "Run an all-solver solver-in-the-loop control using the N=64 cell from " - "forward/baseline, where every differentiable solver has less than 0.75% " - "one-step error. Repeated calls expose recurrent-state defects instead of " - "hiding them behind solver exclusions; the nonlinear multimode experiment " - "remains the learning stress task." + "Run an all-solver solver-in-the-loop control using the analytic N=64 " + "forward/baseline case, with an explicit Mach-safe XLB substep budget so " + "every differentiable solver passes both one-interval and uninterrupted " + "forward gates. Repeated calls then expose recurrent-state defects instead " + "of hiding them behind solver exclusions; the nonlinear multimode " + "experiment remains the learning stress task." ), plot_description=( "Native and restarted solver error, neural-correction gain, solver-VJP " @@ -560,7 +561,9 @@ "nu": 0.05, "dt": 0.01, "steps": 1, - "lbm_N_base": 64, + # Four XLB substeps are the smallest tested budget that keeps + # both one-interval and uninterrupted t=1 error below 1%. + "lbm_N_base": 16, }, "dataset": { "reference_kind": "analytic_tgv", @@ -588,7 +591,9 @@ "seed": 2026, "model_seeds": [0, 1, 2], "check_grad": True, - "fd_epsilon": 1e-2, + # The zero-initialized residual needs a smaller perturbation + # than the nonlinear stress task for its local VJP check. + "fd_epsilon": 3e-4, }, "evaluation": { "rollout_frames": 100, diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py index b9d0f78b..7211a472 100644 --- a/tests/test_solver_in_loop.py +++ b/tests/test_solver_in_loop.py @@ -51,7 +51,7 @@ ).resolve() -def test_tgv_control_uses_all_solver_forward_baseline_cell(): +def test_tgv_control_uses_forward_baseline_with_mach_safe_xlb_budget(): from mosaic.benchmarks.problems import get_config cfg = get_config("ns-grid") @@ -64,10 +64,30 @@ def test_tgv_control_uses_all_solver_forward_baseline_cell(): control_dataset = control_run["dataset"] control_evaluation = control_run["evaluation"] - for key in ("nu", "dt", "steps", "lbm_N_base"): + for key in ("nu", "dt", "steps"): assert control_physics[key] == forward_physics[key] + assert forward_physics["lbm_N_base"] == 64 + assert control_physics["lbm_N_base"] == 16 assert forward_run["sweep"]["key"] == "N" assert control_physics["N"] in forward_run["sweep"]["values"] + initial = _tgv(control_physics["N"]) + by_key = {solver.key: solver for solver in cfg.solvers} + xlb_inputs = cfg.make_inputs( + by_key["xlb"].name, + initial, + domain_extent=cfg.domain_extent, + **control_physics, + ) + pict_inputs = cfg.make_inputs( + by_key["pict"].name, + initial, + domain_extent=cfg.domain_extent, + **control_physics, + ) + assert xlb_inputs["steps"] == 4 + assert np.isclose(float(xlb_inputs["dt"][0]), 0.0025) + assert pict_inputs["steps"] == 1 + assert np.isclose(float(pict_inputs["dt"][0]), 0.01) differentiable = {"jax_cfd", "ins_jl", "phiflow", "pict", "warp_ns", "xlb"} active_forward = { cfg.solver(name).key for name in active_solvers(cfg, "forward", "baseline")