Benchmarking Neural Operators

The neojax.benchmark module provides utilities for running standardized and reproducible evaluations of neural operators.

---

Motivation

Comparing neural operator results across papers is difficult in practice, for a few recurring reasons:

1. Non-standardized normalization: Papers apply different normalization schemes (global vs. per-channel, min-max, non-dimensionalization), and normalization choice alone can shift accuracy metrics by several percentage points.
2. Custom simulation setups: Datasets are often generated with different solvers, grid resolutions, domains, boundary values, or time step sizes, which makes direct comparison across papers unreliable.
3. Rollout discrepancies: Time-dependent problems are evaluated inconsistently — some papers report single-step error (predicting \(t+1\) from ground truth at \(t\)), others report autoregressive rollout error over long horizons.
4. Inconsistent metrics: Definitions of relative \(L^2\)/\(L^p\) error vary across papers, and spectral accuracy or physical-invariant conservation are often not reported at all.

neojax.benchmark addresses this by standardizing the pieces that make results comparable:

• DataBundle gives every dataset the same interface for coordinates, time, discretized fields, physical parameters, and boundary values (see neojax.data.bundles.data_bundle).
• BaseSchema decouples the dataset representation from model-specific input formats, so any JAX model (Equinox, Flax, or custom) can be benchmarked as long as a schema exists that maps DataBundle to its expected input.
• Steady-state and time-dependent problems are evaluated through separate evaluators (see Evaluators) with a consistent metric interface, so single-step and rollout results aren't conflated.
• The dataset registry (see Data Downloading) fetches public datasets directly, removing per-paper simulation setup as a variable.

For an example of a benchmarking setup, see here.

---

Registered Benchmarking Datasets

The registry contains 37 standardized PDE datasets (including 20 PDEgym datasets from Hugging Face, 2 verification datasets from Zenodo, and 15 complete PDEBench datasets from Stuttgart's DaRUS repository).

You can choose any subset to benchmark with and add your own datasets on top if desired.

See the download page for details.

---

Benchmark Runner

The high-level orchestrator that takes the benchmark configuration and executes the full benchmark run.

neojax.benchmark.runner.BenchmarkRunner

Orchestrates the entire benchmarking suite across datasets, pipelines, and models.

Parameters:

Name	Type	Description	Default
config	BenchmarkConfig	The parsed BenchmarkConfig object.	required
custom_model_builders	dict[str, Callable[[Any], Callable]] | None	Optional mapping of model names to custom callables defining how given non-neojax models should be built. Default is None. Each callable should return a "pure" Jax callable, meaning e.g. Flax models should be wrapped in a lambda around flax_model.apply().	None

Attributes

• config (BenchmarkConfig): The benchmark configuration.
• custom_model_builders (dict[str, Callable] | None): Optional mapping of model names to custom callables defining how given non-neojax models should be built. Default is None.

Examples:

from jaxtyping import PRNGKeyArray
from neojax.models.geo_fno import GeoFNO
from neojax.benchmark.config import BenchmarkConfig
from neojax.benchmark.runner import BenchmarkRunner

config_dict = {...}
benchmark_config = BenchmarkConfig(**config_dict)

def build_custom_geo_fno(config: Any, key: PRNGKeyArray, dataset: Any) -> Callable:
    return GeoFNO(
        key=key,
        in_channels=config.hyperparameters["in_channels"],
        out_channels=config.hyperparameters["out_channels"],
        hidden_channels=config.hyperparameters["hidden_channels"],
        n_layers=config.hyperparameters["n_layers"],
        modes=config.hyperparameters["modes"],
        grid_resolution=config.hyperparameters["grid_resolution"],
        use_coord_projection=config.hyperparameters.get("use_coord_projection",
        True),
    )

custom_model_builders = {
    "GeoFNO": build_custom_geo_fno
}

runner = BenchmarkRunner(benchmark_config, custom_model_builders)
results = runner.run()

run() -> dict[str, typing.Any]

Executes the benchmarking suite.

Iterates through all tasks, loads datasets, fits normalizers on the fly if needed, builds models, and evaluates them.

Returns:

Type	Description
dict[str, Any]	A dictionary containing the aggregated benchmark results.

---

Saving Results

We provide a utility save_results to serialize results to JSON or binary NumPy format.

neojax.benchmark.evaluators.base_evaluator.save_results(results: dict[str, typing.Any], filepath: str | pathlib.Path) -> None

Saves evaluation results to a file.

Parameters:

Name	Type	Description	Default
results	dict[str, Any]	The dictionary of results (containing 'summary' and 'details').	required
filepath	str | Path	The path to save the file. Supports .json and .npz formats.	required

Raises:

Type	Description
ValueError	If the given file format is unsupported.