Benchmark Config

Benchmark Configuration¤

The benchmark configuration (a nested dictionary) defines every model, dataset and component used for all benchmarking tasks. neojax then uses pydantic to validate all the schemas and instantiate all benchmark components.

A benchmarking configuration has to adhere to the following structure:

config_dict = {
    # Benchmark config as outermost schema
    "benchmark_name": "<name>",
    # Define global settings (GlobalSettingsConfig)
    "global_settings": {
        "batch_size": <batch_size>,
        "output_dir": "<output_path>",
        "seed": <rng_seed>,
    },
    # Then define each benchmark task (TaskConfig)
    "tasks": [
        {
            "name": "<task_name>",
            # Define datasets used for this task (DatasetConfig)
            "datasets": [
                {
                    "name": "<dataset_name>",
                    "type": "RawDataset",
                    "source": "<source>",
                    "path": "<dataset_path>",
                    "field_mapping": {
                        "fields": "<dataset_fieldname>",
                        ...
                    }
                    # and more optional settings
                    ...
                },
                ...
            ],
            # Define evaluator to use (EvaluatorConfig)
            "evaluator": {
                "type": "<evaluator_class>",
                # Define metrics to use in this task (ComponentConfig)
                "metrics": [
                    {
                        "type": "<metric_class>", "kwargs": {...}
                    },
                    ...
                ]
            },
            # Define models to benchmark in this task (ModelConfig)
            "models": [
                {
                    "name": "<model_name>",
                    "architecture": "<model_class>",
                    "framework": "custom",
                    "hyperparameters": {
                        ...
                    },
                    # Define a pipeline for each model (PipelineConfig)
                    "pipeline": {
                        # Define each schema for the pipeline (ComponentConfig)
                        "in_schema": {
                            "type": "<schema_class>",
                            "kwargs": {...}
                        },
                        "out_schema": {
                            "type": "<schema_class>",
                            "kwargs": {...}
                        }
                    }
                },
                ...
            ]
        },
        # And so on for all tasks to run
        ...
    ],
    ...
}

If you'd like to instantiate each config separately, instantiate from the innermost nested level outwards.

Configuring Benchmark¤

`neojax.benchmark.config.BenchmarkConfig` `pydantic-model` ¤

Root configuration schema for a Benchmark Runner suite.

Attributes:

Name	Type	Description
`benchmark_name`	`str`	The name of the overall benchmark suite.
`global_settings`	`GlobalSettingsConfig`	Global configuration settings.
`tasks`	`list[TaskConfig]`	A list of benchmarking tasks to execute.

Configuring Global Settings¤

`neojax.benchmark.config.GlobalSettingsConfig` `pydantic-model` ¤

Global settings for the benchmark run.

Attributes:

Name	Type	Description
`batch_size`	`int`	The default batch size for evaluation.
`output_dir`	`str \| Path`	The directory where results should be saved.
`seed`	`int`	Random seed for reproducibility.

Configuring Tasks¤

`neojax.benchmark.config.TaskConfig` `pydantic-model` ¤

Configuration for a benchmark task.

A task applies the same evaluator across a collection of datasets and models.

Attributes:

Name	Type	Description
`name`	`str`	A unique identifier for the task.
`evaluator`	`EvaluatorConfig`	The evaluator configuration.
`datasets`	`list[DatasetConfig]`	A list of dataset configurations to evaluate.
`models`	`list[ModelConfig]`	A list of model configurations to benchmark.

Configuring Datasets¤

`neojax.benchmark.config.DatasetConfig` `pydantic-model` ¤

Configuration for a dataset used in benchmarking.

Attributes:

Name	Type	Description
`name`	`str`	A unique identifier for the dataset in this benchmark.
`type`	`str`	Dataset class name (e.g., 'BundleDataset', 'RawDataset').
`source`	`str`	Source type (e.g., 'pdebench', 'pdegym').
`path`	`str \| Path`	Path to the dataset directory or file.
`field_mapping`	`FieldMappingConfig`	Dictionary mapping for DataBundle keys.
`train_split`	`tuple[float, float] \| None`	Optional tuple of [start, end] fractions (e.g., (0.0, 0.8)) for computing normalizer statistics on the fly.
`test_split`	`tuple[float, float]`	Tuple of [start, end] fractions (e.g., (0.8, 1.0)) for evaluation.

Configuring Dataset Field Mappings¤

`neojax.benchmark.config.FieldMappingConfig` `pydantic-model` ¤

Configuration for mapping raw datasets to DataBundle keys.

Attributes:

Name	Type	Description
`fields`	`str \| list[str]`	Name or list of names for fields.
`coords`	`str \| list[str]`	Name or list of names for coordinates.
`parameters`	`str \| list[str] \| None`	Optional name or list of names for parameters.
`bc_masks`	`str \| list[str] \| None`	Optional name or list of names for boundary condition masks.
`bc_values`	`str \| list[str] \| None`	Optional name or list of names for boundary condition values.
`edge_indices`	`str \| list[str] \| None`	Optional name or list of names for edge indices.

Configuring Evaluators¤

`neojax.benchmark.config.EvaluatorConfig` `pydantic-model` ¤

Configuration for the evaluator used in a task.

Attributes:

Name	Type	Description
`type`	`str`	The class name of the evaluator ('SteadyStateEvaluator' or 'TimeDependentEvaluator').
`metrics`	`list[str \| ComponentConfig]`	A list of metric names or ComponentConfigs.
`autoregressive_steps`	`int \| None`	Number of rollout steps (for time_dependent evaluator).

Configuring Models¤

`neojax.benchmark.config.ModelConfig` `pydantic-model` ¤

Configuration for a neural operator model to be benchmarked.

Attributes:

Name	Type	Description
`name`	`str`	A unique identifier for the model in the benchmark.
`architecture`	`str`	The class name of the model (e.g., 'FNO', 'UNet').
`framework`	`str`	The framework the model is implemented in ('neojax', 'custom'). If 'custom', ensure to pass `custom_model_builders` to `BenchmarkRunner`.
`checkpoint_path`	`str \| None`	Path to the Orbax checkpoint directory.
`hyperparameters`	`dict[str, Any]`	Dictionary of hyperparameters to pass to the model constructor.
`pipeline`	`PipelineConfig`	The data processing pipeline required by this model.

Configuring Pipelines¤

`neojax.benchmark.config.PipelineConfig` `pydantic-model` ¤

Configuration for the data processing pipeline.

Attributes:

Name	Type	Description
`in_schema`	`ComponentConfig \| list[ComponentConfig]`	Input schema component(s). Can be a single config or a list.
`out_schema`	`ComponentConfig`	ComponentConfig for the output.
`normalizers`	`dict[str, NormalizerConfig]`	A dictionary mapping DataBundle keys to NormalizerConfigs.

Configuring Normalizers¤

`neojax.benchmark.config.NormalizerConfig` `pydantic-model` ¤

Configuration for normalizers within the pipeline.

Attributes:

Name	Type	Description
`type`	`str`	The class name of the normalizer (e.g., 'UnitGaussianNormalizer').
`compute_on_fly`	`bool`	Whether to fit this normalizer on the `train_split` of the dataset.
`kwargs`	`dict[str, Any]`	Additional arguments for the normalizer constructor.

Configuring Metrics & Schemas¤

`neojax.benchmark.config.NormalizerConfig` `pydantic-model` ¤

Configuration for normalizers within the pipeline.

Attributes:

Name	Type	Description
`type`	`str`	The class name of the normalizer (e.g., 'UnitGaussianNormalizer').
`compute_on_fly`	`bool`	Whether to fit this normalizer on the `train_split` of the dataset.
`kwargs`	`dict[str, Any]`	Additional arguments for the normalizer constructor.

Benchmark Config