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/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/benchmarks/core/utils.py b/mosaic/benchmarks/core/utils.py index 88a66f57..73a04ac9 100644 --- a/mosaic/benchmarks/core/utils.py +++ b/mosaic/benchmarks/core/utils.py @@ -131,6 +131,44 @@ def _debug_run(run: dict) -> None: for key, cap in [("max_iters", 50), ("patience", 10)]: if key in optim: optim[key] = min(optim[key], cap) + training = run.get("training", {}) + for key, cap in [ + ("max_updates", 2), + ("unroll", 2), + ("hidden_channels", 8), + ("kernel_size", 3), + ]: + if key in training: + training[key] = min(training[key], cap) + if training: + training["check_grad"] = False + if "model_seeds" in training: + training["model_seeds"] = list(training["model_seeds"])[:1] + dataset = run.get("dataset", {}) + for key in ("train_seeds", "test_seeds"): + if key in dataset: + dataset[key] = list(dataset[key])[:1] + if "prefix_audit_seeds" in dataset: + retained_seeds = set(dataset.get("train_seeds", [])) | set( + dataset.get("test_seeds", []) + ) + dataset["prefix_audit_seeds"] = [ + seed for seed in dataset["prefix_audit_seeds"] if seed in retained_seeds + ] + if "train_frames" in dataset: + dataset["train_frames"] = min(dataset["train_frames"], 3) + evaluation = run.get("evaluation", {}) + if "rollout_frames" in evaluation: + evaluation["rollout_frames"] = min(evaluation["rollout_frames"], 3) + if "prefix_audit_frames" in dataset: + max_frame = max( + int(dataset.get("train_frames", 0)), + int(evaluation.get("rollout_frames", 0)), + int(training.get("unroll", 0)), + ) + dataset["prefix_audit_frames"] = [ + frame for frame in dataset["prefix_audit_frames"] if frame <= max_frame + ] cost = run.get("cost", {}) for key in ("N_values", "steps_values"): if key in cost: diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/config.py b/mosaic/benchmarks/problems/navier_stokes_grid/config.py index 2a1e93cf..d1188d08 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,14 @@ 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, + plot_solver_in_loop_reference_sensitivity, + plot_solver_in_loop_self_reference, + plot_solver_in_loop_tgv, +) +from .solver_in_loop import solver_in_loop _TESSERACT_SLUG = "navier-stokes-grid" @@ -205,16 +213,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( @@ -442,6 +466,297 @@ }, plot=plot_drag_opt, ) +problem.add_experiment( + "optimization/solver_in_loop", + solver_in_loop, + description=( + "Train an identical Equinox periodic residual corrector through each " + "differentiable solver. A solver-terminal recurrent auxiliary, paired " + "stop-gradient controls, native versus repeated-call 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, " + "recurrent-call closure, 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. + runs=[ + { + "ic": {"name": "multimode", "seed": 0}, + "physics": { + # ``steps`` is the number of native PDE steps between neural + # corrections. Recurrent training feeds back the corrected + # canonical ``v0`` while preserving any solver-native checkpoint. + "N": 32, + "nu": 0.001, + "dt": 0.02, + "steps": 4, + }, + "dataset": { + "reference_kind": "pseudo_spectral_multimode", + "reference_factor": 2, + "reference_substeps": 2, + "train_seeds": list(range(16)), + "test_seeds": list(range(100, 108)), + "train_frames": 24, + # Begin from the well-resolved spectrum used by the forward + # agreement benchmark. The stronger amplitude and long horizon + # still allow nonlinear transfer to populate smaller scales. + "k0": 2.0, + "sigma_k": 0.5, + "amplitude": 0.5, + }, + "training": { + "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": 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, + # Predeclare forward-admission bounds for the nonlinear task. + # Long chaotic trajectories may decorrelate, but a candidate + # must start close to the common reference and remain finite + # enough for a meaningful relative-gain comparison. + "first_interval_error_tolerance": 0.05, + "native_long_error_tolerance": 0.5, + }, + } + ], + plot=plot_solver_in_loop, +) + + +def _reference_sensitivity_run(name: str, reference_kind: str) -> dict: + """Return one common-target solver-loop reference variant.""" + return { + "name": name, + "ic": {"name": "multimode", "seed": 0}, + "physics": { + "N": 32, + "nu": 0.001, + "dt": 0.02, + "steps": 4, + }, + "dataset": { + "reference_kind": reference_kind, + # Production targets use 128Β² and dt/4. A disjoint convergence + # audit compares selected frames against 256Β² and dt/8 before + # corrector training is admitted. + "reference_factor": 4, + "reference_substeps": 4, + "reference_audit_factor": 8, + "reference_audit_substeps": 8, + "reference_audit_seeds": [0, 100], + "reference_audit_frames": [1, 8, 24, 36], + "reference_convergence_tolerance": 0.005, + "train_seeds": list(range(16)), + "test_seeds": list(range(100, 108)), + "train_frames": 24, + "k0": 2.0, + "sigma_k": 0.5, + "amplitude": 0.5, + }, + "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, + "first_interval_error_tolerance": 0.05, + "native_long_error_tolerance": 0.5, + }, + } + + +problem.add_experiment( + "optimization/solver_in_loop_reference_sensitivity", + solver_in_loop, + description=( + "Repeat the common all-solver corrector comparison against independently " + "discretized, space-time-converged pseudo-spectral and conservative " + "finite-volume references. Paired VJP conclusions are reported separately " + "from reference-dependent absolute error." + ), + plot_description=( + "Reference disagreement, convergence audits, absolute error, correction " + "gain, and paired solver-VJP lift under spectral and finite-volume targets." + ), + runs=[ + _reference_sensitivity_run("spectral", "pseudo_spectral_multimode"), + _reference_sensitivity_run("finite_volume", "finite_volume_multimode"), + ], + plot=plot_solver_in_loop_reference_sensitivity, +) + +problem.add_experiment( + "optimization/solver_in_loop_tgv", + solver_in_loop, + description=( + "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 validate each opt-in recurrent-state " + "adapter under a common closure gate; the nonlinear multimode experiment " + "remains the learning stress task." + ), + plot_description=( + "Native and repeated-call 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": 64, + "nu": 0.05, + "dt": 0.01, + "steps": 1, + # 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", + "reference_factor": 1, + "reference_substeps": 1, + "train_seeds": [0, 1, 2, 3], + "test_seeds": [100, 101], + "train_frames": 24, + "long_closure_tolerance": 0.01, + }, + "training": { + "max_updates": 50, + "unroll": 4, + "loss_mode": "solver_terminal", + "solver_loss_weight": 0.1, + "loss_normalization": "solver_baseline", + "loss_scale_floor": 1e-6, + # 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, + "kernel_size": 5, + "seed": 2026, + "model_seeds": [0, 1, 2], + "check_grad": True, + # 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, + "seen_ic_trajectories": 2, + "stable_error_threshold": 1.0, + "first_interval_error_tolerance": 0.01, + "native_long_error_tolerance": 0.01, + }, + } + ], + 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": 2.0, + "sigma_k": 0.5, + "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/corrector.py b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py new file mode 100644 index 00000000..5c912c18 --- /dev/null +++ b/mosaic/benchmarks/problems/navier_stokes_grid/corrector.py @@ -0,0 +1,574 @@ +# 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 +``velocity -> result`` transition and its VJP vary between benchmark cells. +""" + +from __future__ import annotations + +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, + ) + if idx == len(channels) - 2: + # 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, + ( + 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, + architecture: str = "periodic_residual_cnn", +) -> PeriodicResidualCNN: + """Initialise the benchmark-owned Equinox corrector. + + 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}") + return PeriodicResidualCNN( + key, + hidden_channels=hidden_channels, + kernel_size=kernel_size, + ) + + +def apply_corrector( + 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) + delta = jnp.moveaxis(model(jnp.moveaxis(x, -1, 0)), 0, -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( + 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(model, 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 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) + 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 spectral_prolong(velocity: Any, target_n: int) -> np.ndarray: + """Fourier-prolong 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 target_n % source_n != 0: + raise ValueError("target grid must be an integer multiple of source_n") + spectrum = np.fft.fftshift(np.fft.fft2(field, axes=(0, 1)), axes=(0, 1)) + padded = np.zeros((target_n, target_n, field.shape[-1]), dtype=np.complex128) + start = (target_n - source_n) // 2 + stop = start + source_n + padded[start:stop, start:stop, :] = spectrum + result = np.fft.ifft2( + np.fft.ifftshift(padded, 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, + operators: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None = None, +) -> np.ndarray: + """Dealiased pseudo-spectral RHS for 2D vorticity dynamics.""" + n = omega_hat.shape[0] + if operators is None: + 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 + cutoff = n // 3 + dealias = (np.abs(mx) <= cutoff) & (np.abs(my) <= cutoff) + else: + kx, ky, k2, dealias = operators + 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) + advection_hat *= dealias + return -advection_hat - viscosity * k2 * omega_hat + + +def _minmod3(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> np.ndarray: + """Three-argument minmod limiter for periodic finite-volume slopes.""" + same_sign = (np.signbit(a) == np.signbit(b)) & (np.signbit(b) == np.signbit(c)) + magnitude = np.minimum(np.minimum(np.abs(a), np.abs(b)), np.abs(c)) + return np.where(same_sign, np.copysign(magnitude, a), 0.0) + + +def _finite_volume_vorticity_to_velocity( + omega: np.ndarray, + *, + dx: float, + safe_k2: np.ndarray, +) -> np.ndarray: + """Recover cell-centred velocity with the finite-volume discrete operators.""" + psi_hat = np.fft.fft2(omega) / safe_k2 + psi_hat[0, 0] = 0.0 + psi = np.fft.ifft2(psi_hat).real + u = (np.roll(psi, -1, axis=1) - np.roll(psi, 1, axis=1)) / (2.0 * dx) + v = -(np.roll(psi, -1, axis=0) - np.roll(psi, 1, axis=0)) / (2.0 * dx) + return np.stack([u, v], axis=-1)[:, :, np.newaxis, :] + + +def _finite_volume_vorticity_rhs( + omega: np.ndarray, + *, + viscosity: float, + domain_extent: float, + operators: tuple[float, np.ndarray] | None = None, +) -> np.ndarray: + """Conservative MUSCL finite-volume RHS for periodic 2-D vorticity. + + The advective update uses monotonized-central reconstruction and a local + Rusanov flux. Velocity is recovered from a discrete streamfunction whose + Fourier symbol exactly inverts the cell-centred five-point Laplacian. The + evolution is therefore independently discretized from + :func:`_vorticity_rhs`: only the periodic elliptic inversion uses FFTs, + matching the pressure solve used by several candidate grid solvers. + """ + n = omega.shape[0] + if operators is None: + dx = domain_extent / n + mode = np.fft.fftfreq(n) * n + mx, my = np.meshgrid(mode, mode, indexing="ij") + discrete_k2 = ( + 4.0 * (np.sin(np.pi * mx / n) ** 2 + np.sin(np.pi * my / n) ** 2) / dx**2 + ) + safe_k2 = np.where(discrete_k2 == 0.0, 1.0, discrete_k2) + else: + dx, safe_k2 = operators + + velocity = _finite_volume_vorticity_to_velocity( + omega, + dx=dx, + safe_k2=safe_k2, + ) + u = velocity[:, :, 0, 0] + v = velocity[:, :, 0, 1] + u_face = 0.5 * (u + np.roll(u, -1, axis=0)) + v_face = 0.5 * (v + np.roll(v, -1, axis=1)) + + backward_x = omega - np.roll(omega, 1, axis=0) + forward_x = np.roll(omega, -1, axis=0) - omega + centered_x = 0.5 * (np.roll(omega, -1, axis=0) - np.roll(omega, 1, axis=0)) + slope_x = _minmod3(2.0 * backward_x, centered_x, 2.0 * forward_x) + omega_x_left = omega + 0.5 * slope_x + omega_x_right = np.roll(omega, -1, axis=0) - 0.5 * np.roll(slope_x, -1, axis=0) + flux_x = 0.5 * u_face * (omega_x_left + omega_x_right) - 0.5 * np.abs(u_face) * ( + omega_x_right - omega_x_left + ) + + backward_y = omega - np.roll(omega, 1, axis=1) + forward_y = np.roll(omega, -1, axis=1) - omega + centered_y = 0.5 * (np.roll(omega, -1, axis=1) - np.roll(omega, 1, axis=1)) + slope_y = _minmod3(2.0 * backward_y, centered_y, 2.0 * forward_y) + omega_y_left = omega + 0.5 * slope_y + omega_y_right = np.roll(omega, -1, axis=1) - 0.5 * np.roll(slope_y, -1, axis=1) + flux_y = 0.5 * v_face * (omega_y_left + omega_y_right) - 0.5 * np.abs(v_face) * ( + omega_y_right - omega_y_left + ) + + advection = ( + -((flux_x - np.roll(flux_x, 1, axis=0)) + (flux_y - np.roll(flux_y, 1, axis=1))) + / dx + ) + laplacian = ( + np.roll(omega, -1, axis=0) + + np.roll(omega, 1, axis=0) + + np.roll(omega, -1, axis=1) + + np.roll(omega, 1, axis=1) + - 4.0 * omega + ) / dx**2 + return advection + viscosity * laplacian + + +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) + 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_operator = kx**2 + ky**2 + cutoff = n // 3 + spectral_operators = ( + kx, + ky, + k2_operator, + (np.abs(mx) <= cutoff) & (np.abs(my) <= cutoff), + ) + 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, + operators=spectral_operators, + ) + k2 = _vorticity_rhs( + omega_hat + 0.5 * step_dt * k1, + viscosity=viscosity, + domain_extent=domain_extent, + operators=spectral_operators, + ) + k3 = _vorticity_rhs( + omega_hat + 0.5 * step_dt * k2, + viscosity=viscosity, + domain_extent=domain_extent, + operators=spectral_operators, + ) + k4 = _vorticity_rhs( + omega_hat + step_dt * k3, + viscosity=viscosity, + domain_extent=domain_extent, + operators=spectral_operators, + ) + 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) + + +def finite_volume_reference_trajectory( + initial_velocity: Any, + *, + viscosity: float, + dt: float, + frame_steps: int, + n_frames: int, + substeps: int, + domain_extent: float, +) -> np.ndarray: + """Integrate an independent conservative finite-volume reference with RK4.""" + velocity = np.asarray(initial_velocity, dtype=np.float64) + omega = np.fft.ifft2(_velocity_to_vorticity_hat(velocity, domain_extent)).real + n = omega.shape[0] + dx = domain_extent / n + mode = np.fft.fftfreq(n) * n + mx, my = np.meshgrid(mode, mode, indexing="ij") + discrete_k2 = ( + 4.0 * (np.sin(np.pi * mx / n) ** 2 + np.sin(np.pi * my / n) ** 2) / dx**2 + ) + finite_volume_operators = ( + dx, + np.where(discrete_k2 == 0.0, 1.0, discrete_k2), + ) + step_dt = dt / substeps + trajectory = [velocity.astype(np.float32)] + for _ in range(n_frames): + for _ in range(frame_steps * substeps): + k1 = _finite_volume_vorticity_rhs( + omega, + viscosity=viscosity, + domain_extent=domain_extent, + operators=finite_volume_operators, + ) + k2 = _finite_volume_vorticity_rhs( + omega + 0.5 * step_dt * k1, + viscosity=viscosity, + domain_extent=domain_extent, + operators=finite_volume_operators, + ) + k3 = _finite_volume_vorticity_rhs( + omega + 0.5 * step_dt * k2, + viscosity=viscosity, + domain_extent=domain_extent, + operators=finite_volume_operators, + ) + k4 = _finite_volume_vorticity_rhs( + omega + step_dt * k3, + viscosity=viscosity, + domain_extent=domain_extent, + operators=finite_volume_operators, + ) + omega = omega + step_dt * (k1 + 2 * k2 + 2 * k3 + k4) / 6 + trajectory.append( + _finite_volume_vorticity_to_velocity( + omega, + dx=dx, + safe_k2=finite_volume_operators[1], + ).astype(np.float32) + ) + return np.stack(trajectory) + + +__all__ = [ + "PeriodicResidualCNN", + "apply_corrector", + "centered_divergence_rms", + "corrected_velocity", + "divergence_rms", + "enstrophy", + "finite_volume_reference_trajectory", + "init_corrector", + "kinetic_energy", + "project_periodic_correction", + "reference_trajectory", + "relative_l2", + "spectral_prolong", + "spectral_restrict", +] diff --git a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py index e39109b6..2dec2148 100644 --- a/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py +++ b/mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py @@ -75,12 +75,6 @@ "at all nu values (0.0001–0.05); 0.0309 at nu=0.05 is 12.0Γ— peer median; " "not reducible by further sub-stepping (tested k=9..27); valid=True", ) -PHIFLOW_TGV_DAMPING = Exclusion( - ExclusionCategory.ANOMALY_EXPLAINED, - "phiflow's double CenteredGrid↔StaggeredGrid resampling gives 4.18% amplitude " - "damping (ratio=0.9582); cosine=0.9999924 (pattern correct); arithmetic-average " - "output conversion fix worsened error 9Γ—; upstream library change required", -) XLB_TGV_LBM_FLOOR = Exclusion( ExclusionCategory.ANOMALY_EXPLAINED, "automatic k=9 sub-steps reduce Ma 0.88β†’0.098 (81Γ— MaΒ² reduction); " @@ -88,7 +82,6 @@ "remaining floor is O(dxΒ²) LBM spatial discretization at N=64, not reducible " "by further sub-stepping (tested k=9..27); valid=True", ) - _OBSTACLE_GATE = { "jax_cfd": JAX_CFD_NO_OBSTACLE, "ins_jl": INS_JL_NO_OBSTACLE, @@ -109,7 +102,7 @@ def register(problem: Problem) -> None: ) problem.exclude( "forward/agreement/tgv", - {"phiflow": PHIFLOW_TGV_DAMPING, "xlb": XLB_TGV_LBM_FLOOR}, + {"xlb": XLB_TGV_LBM_FLOOR}, ) problem.exclude("forward/tgv_nu_sweep", {"xlb": XLB_DX2_FLOOR}) problem.exclude("forward/cylinder", _OBSTACLE_GATE) 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..13c25ca6 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 @@ -9,10 +9,11 @@ from typing import Any import matplotlib.animation as manimation -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 matplotlib import gridspec from mosaic.benchmarks.core.config import Problem from mosaic.benchmarks.core.io import ( @@ -44,6 +45,1661 @@ _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", + 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.""" + out_dir = experiment_dir( + results_dir(), + cfg.name, + "optimization", + f"{exp_key}{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] = [] + has_stopped_rollout = False + + 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([]))) + 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([])) + ) + 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 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( + updates, + loss, + color=color, + 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="-", + 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:] - rollout_uncertainty[1:], 1e-12), + corrected[1:] + rollout_uncertainty[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( + x[1:], + uncorrected[1:], + color=color, + linestyle=":", + alpha=0.45, + ) + + wall = metrics.get("training_wall_time_s") + admitted = bool(metrics.get("valid_for_vjp_ranking", False)) + quality = ( + 100.0 * (float(metrics["solver_vjp_geometric_lift"]) - 1.0) + if solver_specific_reference + and admitted + and metrics.get("solver_vjp_geometric_lift") is not None + else ( + None + if solver_specific_reference + else metrics.get("final_rollout_error") + ) + ) + if wall is not None and quality is not None: + ax_cost.scatter( + [wall], + [quality], + 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("Normalized training loss") + ax_loss.set_title("Corrector training") + + ax_roll.set_yscale("log") + ax_roll.set_xlabel("Physical time") + ax_roll.set_ylabel( + "Error / solver-only error" + if solver_specific_reference + else "Relative $L^2$ error" + ) + ax_roll.set_title("Held-out rollout") + ax_roll.text( + 0.98, + 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="bottom", + fontsize=6.5, + bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.75}, + ) + + ax_cost.set_xlabel("Training wall time [s]") + ax_cost.set_ylabel( + "Solver-VJP lift [%]" if solver_specific_reference else "Final rollout error" + ) + ax_cost.set_title( + "VJP benefit versus cost (admitted only)" + 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(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, + 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, + solver_specific_reference=solver_specific_reference, + ) + 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 + + +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 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 plot_solver_in_loop_reference_sensitivity( + cfg: Problem, + *, + sub_keys: list[str], + save: bool = True, + suffix: str = "", + **_kwargs: Any, +) -> list: + """Render both reference variants and their cross-reference sensitivity.""" + payloads: dict[str, tuple[dict, dict[str, np.ndarray], list[str], Path]] = {} + figs: list[plt.Figure] = [] + for sub_key in sub_keys: + suite, _, exp_key = sub_key.partition("/") + variant = exp_key.rsplit("/", 1)[-1] + out_dir = experiment_dir( + results_dir(), + cfg.name, + suite, + f"{exp_key}{suffix}", + ) + result_path = out_dir / "result.json" + fields_path = out_dir / "corrector_fields.npz" + if not result_path.exists() or not fields_path.exists(): + continue + 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: + continue + payloads[variant] = (data, arrays, names, out_dir) + figs.extend( + plot_solver_in_loop( + cfg, + save=save, + suffix=suffix, + exp_key=exp_key, + title=f"Solver-in-the-loop correction: {variant.replace('_', ' ')} reference", + ) + ) + + if {"spectral", "finite_volume"} - set(payloads): + return figs + + reference_labels = { + "spectral": "Pseudo-spectral", + "finite_volume": "Finite-volume", + } + reference_colors = { + "spectral": "#6A51A3", + "finite_volume": "#238B45", + } + aliases: list[str] = [] + names_by_variant: dict[str, dict[str, str]] = {} + for variant, (_data, _arrays, names, _out_dir) in payloads.items(): + names_by_variant[variant] = {} + for name in names: + alias = resolve_solver_alias(name) or name + names_by_variant[variant][alias] = name + if alias not in aliases: + aliases.append(alias) + aliases.sort( + key=lambda alias: ( + NS_ORDER.index(alias) if alias in NS_ORDER else len(NS_ORDER), + alias, + ) + ) + + parent_dir = experiment_dir( + results_dir(), + cfg.name, + "optimization", + f"solver_in_loop_reference_sensitivity{suffix}", + ) + parent_dir.mkdir(parents=True, exist_ok=True) + fig, axes = plt.subplots( + 2, + 3, + figsize=(TEXTWIDTH, 0.72 * TEXTWIDTH), + layout="constrained", + ) + ( + ax_disagreement, + ax_convergence, + ax_solver_only, + ax_corrected, + ax_gain, + ax_vjp, + ) = axes.ravel() + + spectral_arrays = payloads["spectral"][1] + finite_volume_arrays = payloads["finite_volume"][1] + spectral_reference = np.asarray( + spectral_arrays.get("reference_rollout", np.array([])) + ) + finite_volume_reference = np.asarray( + finite_volume_arrays.get("reference_rollout", np.array([])) + ) + times = np.asarray( + spectral_arrays.get( + "evaluation_times", + np.arange(min(len(spectral_reference), len(finite_volume_reference))), + ) + ) + n_frames = min(len(spectral_reference), len(finite_volume_reference), len(times)) + if n_frames: + disagreement = np.asarray( + [ + np.linalg.norm( + spectral_reference[index] - finite_volume_reference[index] + ) + / (np.linalg.norm(spectral_reference[index]) + 1e-12) + for index in range(n_frames) + ] + ) + ax_disagreement.plot( + times[:n_frames], + disagreement, + color="#333333", + marker="o", + markersize=2.5, + ) + ax_disagreement.set_xlabel("Physical time") + ax_disagreement.set_ylabel("Reference disagreement") + ax_disagreement.set_title("Target sensitivity") + + convergence_values: list[float] = [] + for variant in ("spectral", "finite_volume"): + data = payloads[variant][0] + values = [ + float(metrics["reference_convergence_error_p95"]) + for metrics in data.get("by_solver", {}).values() + if metrics.get("reference_convergence_error_p95") is not None + ] + convergence_values.append(float(np.median(values)) if values else np.nan) + ax_convergence.bar( + np.arange(2), + convergence_values, + color=[reference_colors["spectral"], reference_colors["finite_volume"]], + ) + tolerance = 0.05 + for data, _arrays, _names, _out_dir in payloads.values(): + for metrics in data.get("by_solver", {}).values(): + if metrics.get("reference_convergence_tolerance") is not None: + tolerance = float(metrics["reference_convergence_tolerance"]) + break + ax_convergence.axhline(tolerance, color="0.45", linestyle=":", linewidth=1.0) + ax_convergence.set_xticks( + np.arange(2), + [reference_labels["spectral"], reference_labels["finite_volume"]], + rotation=20, + ha="right", + ) + ax_convergence.set_ylabel("128Β² vs 256Β² p95") + ax_convergence.set_title("Reference convergence") + + x = np.arange(len(aliases), dtype=float) + width = 0.36 + + def metric_values(variant: str, key: str) -> np.ndarray: + data = payloads[variant][0] + by_solver = data.get("by_solver", {}) + return np.asarray( + [ + float( + by_solver.get(names_by_variant[variant].get(alias, ""), {}).get( + key, np.nan + ) + ) + for alias in aliases + ] + ) + + for variant_index, variant in enumerate(("spectral", "finite_volume")): + offset = (variant_index - 0.5) * width + common = { + "width": width, + "color": reference_colors[variant], + "alpha": 0.88, + } + ax_solver_only.bar( + x + offset, + metric_values(variant, "uncorrected_rollout_error"), + **common, + ) + ax_corrected.bar( + x + offset, + metric_values(variant, "final_rollout_error"), + **common, + ) + ax_gain.bar( + x + offset, + metric_values(variant, "geometric_error_reduction"), + **common, + ) + lift = 100.0 * (metric_values(variant, "solver_vjp_geometric_lift") - 1.0) + data = payloads[variant][0] + by_solver = data.get("by_solver", {}) + log_std = np.asarray( + [ + np.hypot( + float( + by_solver.get(names_by_variant[variant].get(alias, ""), {}).get( + "solver_vjp_log_lift_seed_std", 0.0 + ) + ), + float( + by_solver.get(names_by_variant[variant].get(alias, ""), {}).get( + "solver_vjp_log_lift_ic_std", 0.0 + ) + ), + ) + for alias in aliases + ] + ) + lift_scale = metric_values(variant, "solver_vjp_geometric_lift") + ax_vjp.bar( + x + offset, + lift, + yerr=100.0 * lift_scale * log_std, + capsize=2, + **common, + ) + + for axis, ylabel, title in ( + (ax_solver_only, "Final relative $L^2$ error", "Solver-only accuracy"), + (ax_corrected, "Final relative $L^2$ error", "Corrected accuracy"), + (ax_gain, "Geometric error reduction [Γ—]", "Correctability"), + (ax_vjp, "Solver-VJP lift [%]", "VJP conclusion"), + ): + axis.set_xticks( + x, + [solver_props(alias)[0] for alias in aliases], + rotation=50, + ha="right", + rotation_mode="anchor", + fontsize=6.5, + ) + axis.set_ylabel(ylabel) + axis.set_title(title) + ax_gain.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_vjp.axhline(0.0, color="0.45", linestyle=":", linewidth=1.0) + reference_handles = [ + mlines.Line2D( + [], + [], + color=reference_colors[variant], + linewidth=6.0, + label=reference_labels[variant], + ) + for variant in ("spectral", "finite_volume") + ] + fig.legend( + handles=reference_handles, + loc="outside lower center", + ncol=2, + frameon=False, + ) + fig.suptitle("Solver-in-the-loop reference sensitivity") + if save: + save_fig(fig, "solver_in_loop_reference_sensitivity", parent_dir) + figs.append(fig) + + if n_frames: + field_indices = sorted({0, max(1, n_frames * 2 // 3), n_frames - 1}) + fields_fig, fields_axes = plt.subplots( + len(field_indices), + 3, + figsize=(0.70 * TEXTWIDTH, 0.62 * TEXTWIDTH), + squeeze=False, + ) + spectral_vorticity = [ + _periodic_vorticity_2d(spectral_reference[index]) for index in field_indices + ] + finite_volume_vorticity = [ + _periodic_vorticity_2d(finite_volume_reference[index]) + for index in field_indices + ] + vmax = float( + np.percentile( + np.abs( + np.concatenate( + [ + *(value.ravel() for value in spectral_vorticity), + *(value.ravel() for value in finite_volume_vorticity), + ] + ) + ), + 99, + ) + ) + differences = [ + finite - spectral + for spectral, finite in zip( + spectral_vorticity, finite_volume_vorticity, strict=True + ) + ] + diff_vmax = float( + np.percentile( + np.abs(np.concatenate([value.ravel() for value in differences])), + 99, + ) + ) + for row, index in enumerate(field_indices): + for col, (value, scale) in enumerate( + ( + (spectral_vorticity[row], vmax), + (finite_volume_vorticity[row], vmax), + (differences[row], diff_vmax), + ) + ): + image = fields_axes[row, col].imshow( + value.T, + origin="lower", + cmap="RdBu_r", + vmin=-max(scale, 1e-12), + vmax=max(scale, 1e-12), + ) + fields_axes[row, col].set_xticks([]) + fields_axes[row, col].set_yticks([]) + if col == 2: + fields_fig.colorbar(image, ax=fields_axes[row, col], shrink=0.75) + fields_axes[row, 0].set_ylabel(f"$t={times[index]:.2f}$") + for axis, title in zip( + fields_axes[0], + ("Pseudo-spectral", "Finite-volume", "FV βˆ’ spectral"), + strict=True, + ): + axis.set_title(title) + fields_fig.suptitle("Reference vorticity fields") + fields_fig.subplots_adjust(hspace=0.16, wspace=0.16) + if save: + save_fig( + fields_fig, + "solver_in_loop_reference_fields", + parent_dir, + ) + figs.append(fields_fig) + return figs + + +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, + *, + 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 _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], + 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. + """ + rows = _ordered_solver_rollouts(arrays, names) + if not rows: + return None + + 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, 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 + + 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, + 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) + ): + 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 _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, + solver_specific_reference: bool = False, +) -> plt.Figure | None: + """Compare energy, enstrophy, and divergence along held-out rollouts.""" + rows = _ordered_solver_rollouts(arrays, names) + 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( + [ + 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) + + 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) + + present: list[str] = [] + 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] + ) + 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 = ( + 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, reference_curve, raw_curve, corrected_curve in zip( + diagnostic_axes, + reference_curves, + raw_curves, + corrected_curves, + strict=True, + ): + 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, + 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="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"), + ], + ) + 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, + solver_specific_reference: bool = False, +) -> 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_quality, ax_gain, ax_vjp, ax_efficiency = axes.ravel() + positions = np.arange(len(rows), dtype=float) + 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] = [] + 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_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) + raw = float(metrics["uncorrected_mean_rollout_error"]) + corrected = float(metrics["mean_rollout_error"]) + normalized_full_errors.append(corrected / max(raw, 1e-12)) + stopped_error = metrics.get("stop_gradient_mean_rollout_error") + 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 + 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) + + 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( + 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") + 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( + 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, + ) + ), + ) + ) + ) + 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 + + 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), + 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: + admitted = [ + idx + for idx, (_name, metrics) in enumerate(rows) + if metrics.get("valid_for_vjp_ranking", False) + ] + 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[:, admitted] + ax_vjp.bar( + admitted_positions, + vjp_lift_pct, + 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( + admitted_positions, + admitted_labels, + rotation=35, + ha="right", + ) + ax_vjp.set_ylabel("Solver-VJP lift [%]") + 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 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[:, ranked_idx].reshape(2, 1) + ) + ax_efficiency.errorbar( + update_times[idx], + vjp_lift_pct[ranked_idx], + yerr=error, + color=color, + marker=marker, + linestyle="none", + capsize=2, + ) + 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( + 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(f"VJP benefit versus cost{title_suffix}") + else: + ax_efficiency.axis("off") + 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, + ) + ax_efficiency.axis("off") + + _solver_loop_legend(fig, [name for name, _metrics in rows]) + 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)" + 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 + + +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", + "recurrent_to_native_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_recurrent, ax_closure, 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") + recurrent = values("uncorrected_rollout_error") + lower = max(min(native + recurrent) * 0.8, 1e-5) + upper = max(native + recurrent) * 1.25 + ax_recurrent.plot( + [lower, upper], + [lower, upper], + color="0.45", + linestyle=":", + linewidth=1.0, + ) + for x_value, y_value, color, marker in zip( + native, + recurrent, + colors, + markers, + strict=True, + ): + ax_recurrent.scatter(x_value, y_value, color=color, marker=marker, s=30) + if upper / lower >= 10.0: + ax_recurrent.set_xscale("log") + ax_recurrent.set_yscale("log") + ax_recurrent.set(xlim=(lower, upper), ylim=(lower, upper)) + ax_recurrent.set_xlabel("Native final error") + ax_recurrent.set_ylabel("Repeated-call final error") + ax_recurrent.set_title("Repeated calls vs native") + + 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_closure.bar(positions, closure_pct, color=colors) + ax_closure.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_closure.set_yscale("log") + ax_closure.set_ylabel("95th percentile residual [%]") + ax_closure.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")) + ax_closure.bar(positions, closure_pct, color=colors) + ax_closure.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_closure.set_yscale("log") + ax_closure.set_ylabel("95th percentile residual [%]") + ax_closure.set_title("Two-step closure admission") + else: + ax_closure.bar( + positions, + values("recurrent_to_native_error_ratio"), + color=colors, + ) + ax_closure.axhline(1.0, color="0.45", linestyle=":", linewidth=1.0) + ax_closure.set_ylabel("Repeated / native error [Γ—]") + ax_closure.set_title("Repeated-call 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_closure, 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], +) -> 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.""" + rows = _ordered_solver_rollouts(arrays, names) + 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( + [ + 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. + frame_indices = np.unique( + np.linspace( + 0, + available_frames - 1, + num=min(available_frames, 24), + dtype=int, + ) + ) + n_frames = len(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[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]) + 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, + reference_fields, + raw_fields, + corrected_fields, + ) in enumerate(rollout_vorticity): + label, _color, _linestyle, _marker = solver_props(name) + for col_idx, fields in enumerate( + (reference_fields, 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: + 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}$") + 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 new file mode 100644 index 00000000..ff1381a2 --- /dev/null +++ b/mosaic/benchmarks/problems/navier_stokes_grid/solver_in_loop.py @@ -0,0 +1,1843 @@ +# 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 +import weakref +from functools import lru_cache, partial +from typing import Any, NamedTuple + +import equinox as eqx +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 ( + centered_divergence_rms, + corrected_velocity, + divergence_rms, + enstrophy, + finite_volume_reference_trajectory, + init_corrector, + kinetic_energy, + reference_trajectory, + relative_l2, + spectral_prolong, + spectral_restrict, +) +from .ics import _multimode, _tgv + +_DATASET_LOCK = threading.Lock() +_NATIVE_STATE_SUPPORT_LOCK = threading.Lock() +_NATIVE_STATE_SUPPORT: weakref.WeakKeyDictionary[Any, bool] = ( + weakref.WeakKeyDictionary() +) + + +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( + *, + reference_kind: str, + 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: + if reference_kind in { + "pseudo_spectral_multimode", + "finite_volume_multimode", + }: + fine_initial = _multimode( + reference_n, + L=domain_extent, + seed=seed, + k0=k0, + sigma_k=sigma_k, + amplitude=amplitude, + ) + trajectory_fn = ( + reference_trajectory + if reference_kind == "pseudo_spectral_multimode" + else finite_volume_reference_trajectory + ) + fine = trajectory_fn( + 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) + + +def _reference_convergence_audit( + *, + production_trajectories: np.ndarray, + production_seeds: tuple[int, ...], + physics: dict[str, Any], + dataset: dict[str, Any], + evaluation: dict[str, Any], + domain_extent: float, +) -> dict[str, Any]: + """Compare the production reference with a finer space-time realization.""" + audit_factor = dataset.get("reference_audit_factor") + if audit_factor is None: + return {} + + n = int(physics["N"]) + reference_kind = str(dataset["reference_kind"]) + audit_seeds = tuple( + int(value) + for value in dataset.get( + "reference_audit_seeds", + [production_seeds[0], production_seeds[-1]], + ) + ) + missing = sorted(set(audit_seeds) - set(production_seeds)) + if missing: + raise ValueError( + f"reference audit seeds are absent from the dataset: {missing}" + ) + n_frames = int(evaluation.get("rollout_frames", 32)) + production_factor = int(dataset.get("reference_factor", 2)) + production_n = n * production_factor + audit_n = n * int(audit_factor) + audit_substeps = int(dataset.get("reference_audit_substeps", audit_factor)) + if audit_n < production_n or audit_n % production_n != 0: + raise ValueError( + "reference_audit_factor must be an integer refinement of reference_factor" + ) + if reference_kind not in { + "pseudo_spectral_multimode", + "finite_volume_multimode", + }: + raise ValueError( + f"reference convergence audit is unsupported for {reference_kind!r}" + ) + trajectory_fn = ( + reference_trajectory + if reference_kind == "pseudo_spectral_multimode" + else finite_volume_reference_trajectory + ) + audit_trajectories = [] + for seed in audit_seeds: + production_initial = _multimode( + production_n, + L=domain_extent, + seed=seed, + k0=float(dataset.get("k0", 6.0)), + sigma_k=float(dataset.get("sigma_k", 1.0)), + amplitude=float(dataset.get("amplitude", 0.3)), + ) + audit_initial = spectral_prolong(production_initial, audit_n) + audit_fine = trajectory_fn( + audit_initial, + viscosity=float(physics["nu"]), + dt=float(physics["dt"]), + frame_steps=int(physics["steps"]), + n_frames=n_frames, + substeps=audit_substeps, + domain_extent=float(domain_extent), + ) + audit_trajectories.append( + np.stack([spectral_restrict(frame, n) for frame in audit_fine]) + ) + audit_trajectories = np.stack(audit_trajectories) + + frames = tuple( + sorted( + { + int(frame) + for frame in dataset.get( + "reference_audit_frames", + [1, max(1, n_frames // 3), n_frames], + ) + if 0 < int(frame) <= n_frames + } + ) + ) + if not frames: + raise ValueError("reference_audit_frames must contain a positive valid frame") + production_by_seed = { + seed: production_trajectories[index] + for index, seed in enumerate(production_seeds) + } + errors = np.asarray( + [ + relative_l2( + production_by_seed[seed][frame], + audit_trajectories[seed_index, frame], + ) + for seed_index, seed in enumerate(audit_seeds) + for frame in frames + ], + dtype=np.float64, + ) + tolerance = float(dataset.get("reference_convergence_tolerance", 0.05)) + p95 = float(np.percentile(errors, 95)) + return { + "reference_convergence_audit_applied": True, + "reference_convergence_scheme": reference_kind, + "reference_production_grid_size": production_n, + "reference_audit_grid_size": audit_n, + "reference_production_dt": float( + physics["dt"] / int(dataset.get("reference_substeps", 2)) + ), + "reference_audit_dt": float(physics["dt"] / audit_substeps), + "reference_convergence_frames": list(frames), + "reference_convergence_seeds": list(audit_seeds), + "reference_convergence_errors": errors.tolist(), + "reference_convergence_error_median": float(np.median(errors)), + "reference_convergence_error_p95": p95, + "reference_convergence_error_max": float(np.max(errors)), + "reference_convergence_tolerance": tolerance, + "reference_convergence_passed": bool(p95 <= tolerance), + "eligible_for_corrector_training": bool(p95 <= tolerance), + } + + +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, np.ndarray, str, dict[str, Any]]: + """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])) + 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 = { + "reference_kind": str( + dataset.get("reference_kind", "pseudo_spectral_multimode") + ), + "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] + train_rollouts = all_trajectories[:n_train, : eval_frames + 1] + test = all_trajectories[n_train:, : eval_frames + 1] + audit = _reference_convergence_audit( + production_trajectories=all_trajectories, + production_seeds=train_seeds + test_seeds, + physics=physics, + dataset=dataset, + evaluation=evaluation, + domain_extent=domain_extent, + ) + return train, train_rollouts, test, _dataset_digest(config), audit + + +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, _audit = _make_reference_datasets( + physics=physics, + dataset=dataset, + evaluation=evaluation, + training=training, + domain_extent=domain_extent, + ) + return train, test, dataset_hash + + +def _solver_advance( + t: Any, + ctx: KernelContext, + velocity: jax.Array, + *, + frame_steps: int, + native_state: Any | None = None, +) -> tuple[jax.Array, Any | None]: + """Advance one canonical correction interval through a Tesseract.""" + return _solver_advance_with_physics( + t, + ctx, + velocity, + dt=float(ctx.phys["dt"]), + steps=frame_steps, + native_state=native_state, + ) + + +def _supports_native_state(t: Any) -> bool: + """Return whether apply advertises a differentiable recurrent checkpoint.""" + with _NATIVE_STATE_SUPPORT_LOCK: + try: + return _NATIVE_STATE_SUPPORT[t] + except (KeyError, TypeError): + pass + try: + schemas = t.openapi_schema["components"]["schemas"] + properties = schemas["Apply_InputSchema"]["properties"] + differentiable_inputs = schemas["ApplyInputSchema"]["differentiable_arrays"] + differentiable_outputs = schemas["ApplyOutputSchema"]["differentiable_arrays"] + except (AttributeError, KeyError, TypeError): + return False + supports_native_state = ( + "return_state" in properties + and "state" in properties + and "state" in differentiable_inputs + and "state" in differentiable_outputs + ) + with _NATIVE_STATE_SUPPORT_LOCK: + try: + _NATIVE_STATE_SUPPORT[t] = supports_native_state + except TypeError: + pass + return supports_native_state + + +def _solver_advance_with_physics( + t: Any, + ctx: KernelContext, + velocity: jax.Array | np.ndarray, + *, + dt: float, + steps: int, + native_state: Any | None = None, +) -> tuple[jax.Array, Any | None]: + """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) + supports_native_state = _supports_native_state(t) + if supports_native_state: + inputs = {**inputs, "return_state": True} + if native_state is not None: + if not supports_native_state: + raise RuntimeError( + f"Solver '{ctx.name}' received native state without advertising " + "'state' and 'return_state' apply inputs" + ) + inputs = {**inputs, "state": native_state} + 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}" + ) + next_native_state = outputs.get("state") if supports_native_state else None + if supports_native_state and next_native_state is None: + raise RuntimeError( + f"Solver '{ctx.name}' advertised recurrent state but did not return 'state'" + ) + return outputs[ctx.output_key], next_native_state + + +def _stop_recurrent_gradient( + velocity: jax.Array, + native_state: Any | None, +) -> tuple[jax.Array, Any | None]: + """Cut solver-VJP paths through canonical and solver-native recurrent state.""" + return ( + jax.lax.stop_gradient(velocity), + jax.tree_util.tree_map(jax.lax.stop_gradient, native_state), + ) + + +def _passes_reference_accuracy_gate( + reference_kind: str, + *, + first_interval_error: float, + first_interval_tolerance: float, + native_long_error: float, + native_long_tolerance: float, +) -> bool: + """Gate only references whose absolute accuracy is shared across solvers.""" + if reference_kind == "solver_self_refined": + return True + return bool( + first_interval_error <= first_interval_tolerance + and native_long_error <= native_long_tolerance + ) + + +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, + native_state: Any | None = None, + ) -> tuple[np.ndarray, Any | None]: + nonlocal apply_count + apply_count += 1 + next_velocity, next_native_state = _solver_advance_with_physics( + t, + ctx, + velocity, + dt=dt, + steps=steps, + native_state=native_state, + ) + return np.asarray(next_velocity), next_native_state + + 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)] + fine_native_state = None + for _frame in range(n_frames): + fine_state, fine_native_state = advance( + fine_state, + dt=fine_dt, + steps=fine_steps, + native_state=fine_native_state, + ) + 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_native_state = None + coarse_repeated: dict[int, np.ndarray] = {} + for frame in range(1, max_audit_frame + 1): + coarse_state, coarse_native_state = advance( + coarse_state, + dt=coarse_dt, + steps=coarse_steps, + native_state=coarse_native_state, + ) + 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)) + coarse_closure_to_signal = [ + closure / (signal + 1e-12) + for closure, signal in zip( + coarse_closure_errors, + refinement_signals, + strict=True, + ) + ] + fine_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(coarse_closure_to_signal) <= closure_to_signal_tolerance + and max(fine_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_coarse_closure_to_signal_ratio": float(max(coarse_closure_to_signal)), + "max_fine_closure_to_signal_ratio": float(max(fine_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, + *, + t: Any, + ctx: KernelContext, + frame_steps: int, + velocity_scale: float, + differentiate_solver: bool, + loss_mode: str, + solver_loss_weight: float, + loss_scale: float, +) -> jax.Array: + """Normalized recurrent loss with configurable temporal credit assignment.""" + state = targets[0] + native_state = None + losses: list[jax.Array] = [] + solver_terminal_loss: jax.Array | None = None + for step, target in enumerate(targets[1:]): + provisional, native_state = _solver_advance( + t, + ctx, + state, + frame_steps=frame_steps, + native_state=native_state, + ) + if not differentiate_solver: + # Keep the identical recurrent forward trajectory but cut the + # backward graph through both canonical and solver-native recurrent + # state. This counterfactual measures what the same corrector can + # learn without the solver VJP. + provisional, native_state = _stop_recurrent_gradient( + provisional, + native_state, + ) + 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, + velocity_scale=velocity_scale, + domain_extent=ctx.domain_extent, + ) + losses.append(jnp.sum((state - target) ** 2) / (jnp.sum(target**2) + 1e-12)) + 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( + loss_fn: Any, + model: Any, + grads: Any, + key: jax.Array, + *, + epsilon: float, +) -> float: + """Relative error of one end-to-end directional finite difference.""" + 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 = 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) + 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, + loss_scale: float, + 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)) + 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), + 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(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 + 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, + 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) + 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 differentiate_solver + and bool(training.get("check_grad", True)) + ): + fd_error = _directional_fd( + loss_fn, + model, + grads, + jax.random.PRNGKey(model_seed + 1), + epsilon=fd_epsilon, + ) + 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) + 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, + model: 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]) + native_state = None + states = [np.asarray(state)] + errors = [0.0] + for target in reference[1:]: + state, native_state = _solver_advance( + t, + ctx, + state, + frame_steps=frame_steps, + native_state=native_state, + ) + if corrected: + state = corrected_velocity( + model, + 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 + + +class _ReferenceEvaluation(NamedTuple): + first_rollout: np.ndarray | None + errors: np.ndarray + final_states: np.ndarray + + +def _evaluate_reference_set( + t: Any, + ctx: KernelContext, + model: Any, + references: np.ndarray, + *, + frame_steps: int, + velocity_scale: float, + corrected: bool, +) -> _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( + t, + ctx, + model, + reference, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=corrected, + ) + errors.append(reference_errors) + final_states.append(rollout[-1]) + if first_rollout is None: + first_rollout = rollout + 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: + """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 + + +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 _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 + 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_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", + "reference_rollout", + ), +) +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"]) + 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, + reference_audit, + ) = _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), + "native_state_threading": _supports_native_state(t), + }, + "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") + 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 + # intervals expose recurrent-call closure separately from integration accuracy. + 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] + semigroup_errors: list[float] = [] + for reference in test: + one_interval, one_interval_native_state = _solver_advance( + t, + ctx, + jnp.asarray(reference[0]), + frame_steps=frame_steps, + ) + repeated, _ = _solver_advance( + t, + ctx, + one_interval, + frame_steps=frame_steps, + native_state=one_interval_native_state, + ) + 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 + ) + uncorrected_eval = _evaluate_reference_set( + t, + ctx, + None, + test, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=False, + ) + first_uncorrected = uncorrected_eval.first_rollout + uncorrected_errors_array = uncorrected_eval.errors + recurrent_final_states = uncorrected_eval.final_states + seen_uncorrected_eval = _evaluate_reference_set( + t, + ctx, + None, + seen_references, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=False, + ) + seen_uncorrected_errors = seen_uncorrected_eval.errors + native_final_errors: list[float] = [] + long_closure_errors: list[float] = [] + for reference, recurrent_final in zip(test, recurrent_final_states, strict=True): + native_final, _ = _solver_advance( + t, + ctx, + jnp.asarray(reference[0]), + 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(recurrent_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 long_closure_p95 <= long_closure_tolerance + and _passes_reference_accuracy_gate( + reference_kind, + first_interval_error=first_interval_error_p95, + first_interval_tolerance=first_interval_error_tolerance, + native_long_error=native_final_error_p95, + native_long_tolerance=native_long_error_tolerance, + ) + ) + + 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]] = [] + 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] = [] + 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 + + 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, + loss_scale=training_loss_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, + loss_scale=training_loss_scale, + differentiate_solver=False, + model_seed=model_seed, + ) + stop_gradient_training_walls.append(time.perf_counter() - stop_gradient_started) + + corrected_eval = _evaluate_reference_set( + t, + ctx, + model, + test, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + corrected_rollout = corrected_eval.first_rollout + seed_corrected_errors = corrected_eval.errors + stop_gradient_eval = _evaluate_reference_set( + t, + ctx, + stop_gradient_model, + test, + frame_steps=frame_steps, + velocity_scale=velocity_scale, + corrected=True, + ) + 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, + 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 + + 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(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) + 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) + 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_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)) + 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) + 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 = { + **reference_audit, + "eligible_for_corrector_training": True, + "valid_for_vjp_ranking": valid_for_vjp_ranking, + "reference_accuracy_gate_applied": reference_kind != "solver_self_refined", + "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, + "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]), + "stop_gradient_final_rollout_error": final_stop_gradient, + "stop_gradient_final_rollout_error_seed_std": float( + 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, + "recurrent_to_native_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), + "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( + stop_gradient_error.tolist(), threshold + ) + * interval_time, + "uncorrected_stable_horizon": _first_unstable( + uncorrected_error.tolist(), threshold + ) + * interval_time, + "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(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, + "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_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, + "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), + "native_state_threading": _supports_native_state(t), + } + snapshots = { + "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_seed_std, + dtype=np.float32, + ), + "error_corrected_ic_std": np.asarray( + corrected_error_ic_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_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 + 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 + 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, + } + + +__all__ = ["make_reference_dataset", "solver_in_loop"] 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/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_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/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..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" @@ -1275,6 +1293,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 +1320,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 +1353,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" @@ -1433,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" @@ -1920,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.*'" }, @@ -1950,12 +2201,19 @@ 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 = "equinox", specifier = "==0.13.8" }, { 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 +2232,7 @@ requires-dist = [ { name = "tesseract-jax" }, { name = "typer", specifier = ">=0.12" }, ] -provides-extras = ["dev"] +provides-extras = ["dev", "gpu"] [[package]] name = "mosaic-shared" @@ -2296,6 +2554,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 +2680,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" @@ -3971,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 127f830d..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", @@ -30,6 +31,9 @@ dependencies = [ ] [project.optional-dependencies] +gpu = [ + "jax[cuda12]", +] dev = [ "pytest", "pytest-cov", 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/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 ───────────────────────────────────── 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..94043f07 --- /dev/null +++ b/tests/dummy_tesseracts/navier_stokes_grid_identity/tesseract_api.py @@ -0,0 +1,94 @@ +# 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"] + ) +): + # The smoke test presents this fixture as different real solvers, whose + # solver-specific tuning inputs must be accepted and ignored. + 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 abc89a74..00069093 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 @@ -276,10 +277,84 @@ 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`. """ + solver_loop_sensitivity = exp_key.startswith( + "optimization/solver_in_loop_reference_sensitivity/" + ) + if problem == "ns-grid" and ( + exp_key + in { + "optimization/solver_in_loop", + "optimization/solver_in_loop_self_reference", + "optimization/solver_in_loop_tgv", + } + or solver_loop_sensitivity + ): + from mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop import ( + solver_in_loop, + ) + + analytic = exp_key.endswith("_tgv") + self_reference = exp_key.endswith("_self_reference") + finite_volume_reference = exp_key.endswith("/finite_volume") + cfg.add_experiment( + exp_key, + solver_in_loop, + runs=[ + { + "ic": { + "name": "tgv" if analytic else "multimode", + "seed": 0, + }, + "physics": { + "N": 8, + "nu": 0.05 if analytic else 0.001, + "dt": 0.02, + "steps": 1, + }, + "dataset": { + "reference_kind": ( + "analytic_tgv" + if analytic + else ( + "solver_self_refined" + if self_reference + else ( + "finite_volume_multimode" + if finite_volume_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, + "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 @@ -388,7 +463,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) diff --git a/tests/test_solver_in_loop.py b/tests/test_solver_in_loop.py new file mode 100644 index 00000000..9b402e51 --- /dev/null +++ b/tests/test_solver_in_loop.py @@ -0,0 +1,1072 @@ +# 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 +import inspect +from pathlib import Path +from types import SimpleNamespace + +import equinox as eqx +import jax +import jax.numpy as jnp +import numpy as np + +from mosaic.benchmarks.core.utils import _debug_run, active_solvers +from mosaic.benchmarks.problems.navier_stokes_grid.corrector import ( + PeriodicResidualCNN, + apply_corrector, + centered_divergence_rms, + divergence_rms, + finite_volume_reference_trajectory, + init_corrector, + project_periodic_correction, + reference_trajectory, + relative_l2, + spectral_prolong, + 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_diagnostics, + _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 ( + _evaluate_rollout, + _make_reference_datasets, + _make_solver_self_reference_datasets, + _passes_reference_accuracy_gate, + _rollout_log_gain, + _solver_advance, + _stop_recurrent_gradient, + _window_loss, + make_reference_dataset, + solver_in_loop, +) + +_IDENTITY_DUMMY = ( + Path(__file__).parent + / "dummy_tesseracts" + / "navier_stokes_grid_identity" + / "tesseract_api.py" +).resolve() + + +def test_tgv_control_uses_forward_baseline_with_mach_safe_xlb_budget(): + from mosaic.benchmarks.problems import get_config + + cfg = get_config("ns-grid") + 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 ("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") + } + 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(): + 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_shared_solver_loop_ranking_declares_forward_admission_bounds(): + from mosaic.benchmarks.problems import get_config + + cfg = get_config("ns-grid") + experiment = cfg.experiments["optimization/solver_in_loop"] + run = inspect.signature(experiment.fn).parameters["_kw"].default["runs"][0] + assert run["dataset"]["k0"] == 2.0 + assert run["dataset"]["sigma_k"] == 0.5 + assert run["evaluation"]["first_interval_error_tolerance"] == 0.05 + assert run["evaluation"]["native_long_error_tolerance"] == 0.5 + self_reference = cfg.experiments["optimization/solver_in_loop_self_reference"] + self_reference_run = ( + inspect.signature(self_reference.fn).parameters["_kw"].default["runs"][0] + ) + assert self_reference_run["dataset"]["k0"] == run["dataset"]["k0"] + assert self_reference_run["dataset"]["sigma_k"] == run["dataset"]["sigma_k"] + + +def test_reference_sensitivity_declares_independent_converged_targets(): + from mosaic.benchmarks.problems import get_config + + cfg = get_config("ns-grid") + prefix = "optimization/solver_in_loop_reference_sensitivity/" + variants = { + key.removeprefix(prefix): inspect.signature(experiment.fn) + .parameters["_kw"] + .default["runs"][0] + for key, experiment in cfg.experiments.items() + if key.startswith(prefix) + } + + assert set(variants) == {"spectral", "finite_volume"} + assert variants["spectral"]["dataset"]["reference_kind"] == ( + "pseudo_spectral_multimode" + ) + assert variants["finite_volume"]["dataset"]["reference_kind"] == ( + "finite_volume_multimode" + ) + for run in variants.values(): + dataset = run["dataset"] + assert dataset["reference_factor"] == 4 + assert dataset["reference_substeps"] == 4 + assert dataset["reference_audit_factor"] == 8 + assert dataset["reference_audit_substeps"] == 8 + assert dataset["reference_convergence_tolerance"] == 0.005 + + +def test_self_reference_does_not_gate_the_learnable_refinement_signal(): + assert _passes_reference_accuracy_gate( + "solver_self_refined", + first_interval_error=0.6, + first_interval_tolerance=0.15, + native_long_error=0.7, + native_long_tolerance=0.5, + ) + assert not _passes_reference_accuracy_gate( + "pseudo_spectral_multimode", + first_interval_error=0.6, + first_interval_tolerance=0.15, + native_long_error=0.7, + native_long_tolerance=0.5, + ) + + +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)) + + assert isinstance(model, PeriodicResidualCNN) + assert isinstance(model, eqx.Module) + assert model.architecture == "periodic_residual_cnn" + expected = jnp.roll( + apply_corrector(model, velocity, velocity_scale=1.0), + shift=(2, -1), + axis=(0, 1), + ) + actual = apply_corrector(model, shifted, velocity_scale=1.0) + + 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) + + 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, + ) + assert centered_divergence_rms(_tgv(16), 2.0 * np.pi) < 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_prolongation_preserves_low_mode_tgv(): + coarse = _tgv(16) + fine = spectral_prolong(coarse, 32) + + np.testing.assert_allclose(fine, _tgv(32), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose( + spectral_restrict(fine, 16), + coarse, + 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_finite_volume_reference_converges_to_tgv_decay(): + viscosity = 0.05 + errors = [] + for n in (16, 32): + initial = _tgv(n) + trajectory = finite_volume_reference_trajectory( + initial, + viscosity=viscosity, + dt=0.01, + frame_steps=4, + n_frames=1, + substeps=1, + domain_extent=2.0 * np.pi, + ) + expected = _tgv_analytic( + initial, + nu=viscosity, + t=0.04, + L=2.0 * np.pi, + ) + errors.append(relative_l2(trajectory[-1], expected)) + assert np.isfinite(trajectory).all() + + assert errors[1] < errors[0] + assert errors[1] < 0.01 + + +def test_finite_volume_reference_passes_space_time_convergence_audit(): + train, train_rollouts, test, dataset_hash, audit = _make_reference_datasets( + physics={"N": 8, "nu": 0.01, "dt": 0.01, "steps": 1}, + dataset={ + "reference_kind": "finite_volume_multimode", + "reference_factor": 2, + "reference_substeps": 2, + "reference_audit_factor": 4, + "reference_audit_substeps": 4, + "reference_audit_seeds": [0, 100], + "reference_audit_frames": [1, 2], + "reference_convergence_tolerance": 0.2, + "train_seeds": [0], + "test_seeds": [100], + "train_frames": 2, + "k0": 1.0, + "sigma_k": 0.5, + "amplitude": 0.1, + }, + evaluation={"rollout_frames": 2}, + training={"unroll": 1}, + domain_extent=2.0 * np.pi, + ) + + assert train.shape == (1, 3, 8, 8, 1, 2) + assert train_rollouts.shape == train.shape + assert test.shape == train.shape + assert len(dataset_hash) == 16 + assert audit["reference_convergence_passed"] is True + assert audit["eligible_for_corrector_training"] is True + assert audit["reference_production_grid_size"] == 16 + assert audit["reference_audit_grid_size"] == 32 + assert audit["reference_convergence_error_p95"] < 0.2 + + +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_solver_self_reference_matches_physical_time_and_passes_closure( + monkeypatch, +): + calls: list[tuple[int, float, int, object | None]] = [] + + def _closed_refined_step( + _t, + _ctx, + velocity, + *, + dt, + steps, + native_state=None, + ): + calls.append((velocity.shape[0], dt, steps, native_state)) + per_step_decay = 1.0 - dt / velocity.shape[0] + next_native_state = ( + jnp.asarray([1.0]) + if native_state is None + else jnp.asarray(native_state) + 1.0 + ) + return jnp.asarray(velocity) * per_step_decay**steps, next_native_state + + 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["max_coarse_closure_to_signal_ratio"] < 1e-4 + assert audit["max_fine_closure_to_signal_ratio"] < 1e-4 + assert audit["mean_refinement_signal"] > 1e-5 + assert any(call[:3] == (16, 0.01, 2) for call in calls) + assert any(call[:3] == (8, 0.02, 1) for call in calls) + assert any(call[3] is not None for call in calls) + assert 0.01 * 2 == 0.02 * 1 + + +def test_solver_self_reference_rejects_coarse_closure_larger_than_signal( + monkeypatch, +): + def _coarse_call_biased_step( + _t, + _ctx, + velocity, + *, + dt, + steps, + native_state=None, + ): + per_step_decay = 1.0 - dt / velocity.shape[0] + call_bias = 5e-4 if velocity.shape[0] == 8 else 0.0 + next_native_state = ( + jnp.asarray([1.0]) + if native_state is None + else jnp.asarray(native_state) + 1.0 + ) + return ( + jnp.asarray(velocity) * per_step_decay**steps + call_bias, + next_native_state, + ) + + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop." + "_solver_advance_with_physics", + _coarse_call_biased_step, + ) + ctx = SimpleNamespace( + name="coarse-call-biased-dummy", + phys={"N": 8, "nu": 0.001, "dt": 0.02, "steps": 1}, + domain_extent=2.0 * np.pi, + ) + + *_datasets, audit = _make_solver_self_reference_datasets( + None, + ctx, + dataset={ + "reference_factor": 2, + "reference_temporal_factor": 2, + "train_seeds": [0], + "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": 2}, + training={"unroll": 2}, + ) + + assert audit["max_coarse_closure_error"] < audit["closure_relative_tolerance"] + assert ( + audit["max_coarse_closure_to_signal_ratio"] + > audit["closure_to_signal_tolerance"] + ) + assert audit["max_fine_closure_to_signal_ratio"] < 1e-5 + assert audit["eligible_for_corrector_training"] is False + + +def test_solver_advance_threads_optional_native_state(monkeypatch): + calls: list[dict] = [] + + def _apply(_t, inputs): + calls.append(inputs) + velocity = jnp.asarray(inputs["v0"]) + native_state = inputs.get("state") + if native_state is None: + native_state = jnp.zeros_like(velocity) + return { + "result": velocity + native_state + 1.0, + "state": native_state + 2.0, + } + + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop.apply_tesseract", + _apply, + ) + ctx = SimpleNamespace( + name="stateful-dummy", + phys={"N": 2, "nu": 0.001, "dt": 0.02, "steps": 1}, + output_key="result", + make_inputs=lambda _name, velocity, **_physics: {"v0": velocity}, + ) + t = SimpleNamespace( + openapi_schema={ + "components": { + "schemas": { + "Apply_InputSchema": { + "properties": { + "v0": {}, + "state": {}, + "return_state": {}, + } + }, + "ApplyInputSchema": { + "differentiable_arrays": {"v0": {}, "state": {}} + }, + "ApplyOutputSchema": { + "differentiable_arrays": {"result": {}, "state": {}} + }, + } + } + } + ) + initial = jnp.ones((2, 2, 1, 2)) + + first_velocity, first_native_state = _solver_advance( + t, + ctx, + initial, + frame_steps=1, + ) + second_velocity, second_native_state = _solver_advance( + t, + ctx, + first_velocity, + frame_steps=1, + native_state=first_native_state, + ) + + assert "state" not in calls[0] + assert calls[0]["return_state"] is True + assert calls[1]["return_state"] is True + np.testing.assert_array_equal(calls[1]["state"], first_native_state) + np.testing.assert_array_equal(first_velocity, 2.0 * jnp.ones_like(initial)) + np.testing.assert_array_equal(first_native_state, 2.0 * jnp.ones_like(initial)) + np.testing.assert_array_equal(second_velocity, 5.0 * jnp.ones_like(initial)) + np.testing.assert_array_equal(second_native_state, 4.0 * jnp.ones_like(initial)) + + +def test_solver_advance_leaves_stateless_schema_unchanged(monkeypatch): + calls: list[dict] = [] + + def _apply(_t, inputs): + calls.append(inputs) + return {"result": jnp.asarray(inputs["v0"]) + 1.0} + + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop.apply_tesseract", + _apply, + ) + t = SimpleNamespace( + openapi_schema={ + "components": { + "schemas": { + "Apply_InputSchema": { + "properties": { + "v0": {}, + "state": {}, + "return_state": {}, + } + }, + "ApplyInputSchema": {"differentiable_arrays": {"v0": {}}}, + "ApplyOutputSchema": {"differentiable_arrays": {"result": {}}}, + } + } + } + ) + ctx = SimpleNamespace( + name="stateless-dummy", + phys={"N": 2, "nu": 0.001, "dt": 0.02, "steps": 1}, + output_key="result", + make_inputs=lambda _name, velocity, **_physics: {"v0": velocity}, + ) + + velocity, native_state = _solver_advance( + t, + ctx, + jnp.zeros((2, 2, 1, 2)), + frame_steps=1, + ) + + assert set(calls[0]) == {"v0"} + assert native_state is None + np.testing.assert_array_equal(velocity, jnp.ones_like(velocity)) + + +def test_corrected_rollout_threads_velocity_and_native_state(monkeypatch): + calls: list[tuple[np.ndarray, int | None]] = [] + + def _advance(_t, _ctx, velocity, *, frame_steps, native_state=None): + del frame_steps + calls.append( + ( + np.asarray(velocity), + None if native_state is None else int(native_state), + ) + ) + next_native_state = 1 if native_state is None else native_state + 1 + return jnp.asarray(velocity) + next_native_state, next_native_state + + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop._solver_advance", + _advance, + ) + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop." + "corrected_velocity", + lambda _model, velocity, **_kwargs: velocity + 10.0, + ) + reference = np.zeros((3, 2, 2, 1, 2), dtype=np.float32) + ctx = SimpleNamespace(domain_extent=2.0 * np.pi) + + rollout, _ = _evaluate_rollout( + None, + ctx, + None, + reference, + frame_steps=1, + velocity_scale=1.0, + corrected=True, + ) + + assert calls[0][1] is None + assert calls[1][1] == 1 + np.testing.assert_array_equal(calls[1][0], 11.0 * np.ones_like(reference[0])) + np.testing.assert_array_equal(rollout[-1], 23.0 * np.ones_like(reference[0])) + + +def test_training_window_threads_corrected_velocity_and_native_state(monkeypatch): + calls: list[tuple[np.ndarray, int | None]] = [] + + def _advance(_t, _ctx, velocity, *, frame_steps, native_state=None): + del frame_steps + calls.append( + ( + np.asarray(velocity), + None if native_state is None else int(native_state), + ) + ) + next_native_state = 1 if native_state is None else native_state + 1 + return jnp.asarray(velocity) + next_native_state, next_native_state + + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop._solver_advance", + _advance, + ) + monkeypatch.setattr( + "mosaic.benchmarks.problems.navier_stokes_grid.solver_in_loop." + "corrected_velocity", + lambda model, velocity, **_kwargs: velocity + model, + ) + targets = jnp.ones((3, 2, 2, 1, 2)) + ctx = SimpleNamespace(domain_extent=2.0 * np.pi) + + loss = _window_loss( + jnp.asarray(10.0), + targets, + t=None, + ctx=ctx, + frame_steps=1, + velocity_scale=1.0, + differentiate_solver=True, + loss_mode="mean", + solver_loss_weight=0.1, + loss_scale=1.0, + ) + + assert np.isfinite(float(loss)) + assert calls[0][1] is None + assert calls[1][1] == 1 + np.testing.assert_array_equal(calls[1][0], 12.0 * np.ones_like(targets[0])) + + +def test_stop_gradient_cuts_velocity_and_native_state(): + def stopped(value): + velocity, native_state = _stop_recurrent_gradient( + value, + {"memory": 2.0 * value}, + ) + return velocity + native_state["memory"] + + primal, tangent = jax.jvp( + stopped, + (jnp.asarray(3.0),), + (jnp.asarray(1.0),), + ) + + assert float(primal) == 9.0 + assert float(tangent) == 0.0 + + +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}, + "training": { + "max_updates": 100, + "unroll": 8, + "hidden_channels": 32, + "kernel_size": 5, + "model_seeds": [0, 1, 2], + "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, + "model_seeds": [0], + "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_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_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)), + } + 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_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)), + "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": { + "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, + "recurrent_to_native_error_ratio": 3.0, + "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, + "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, + } + } + } + + fairness = _plot_solver_in_loop_fairness( + data, + ["jax-cfd"], + tmp_path, + save=True, + ) + physics = _plot_solver_in_loop_physics( + arrays, + ["jax-cfd", "warp-ns"], + 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 + # 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", + "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_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_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 + + 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, + "loss_mode": "solver_terminal", + "solver_loss_weight": 0.1, + "loss_normalization": "solver_baseline", + "hidden_channels": 4, + "kernel_size": 3, + "model_seeds": [0, 1], + "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["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["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" + 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() + 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