Evaluators

Benchmark Evaluators

To perform standardized and robust evaluation of neural operators, neojax provides two main benchmark evaluators:

1. Steady State Evaluator

Used for stationary or steady-state PDE problems where time is not a dynamical coordinate that requires sequential temporal propagation.

neojax.benchmark.evaluators.steady_state_eval.SteadyStateEvaluator

Evaluator for steady-state PDE problems.

Transforms data, runs the model, reconstructs the physical outputs, and evaluates all configured metrics for a single sample.

Parameters:

Name	Type	Description	Default
metrics	dict[str, BaseMetric]	A dictionary mapping metric names to BaseMetric instances.	required
in_schema	BaseSchema	Input schema for the model.	required
out_schema	BaseSchema	Output schema for the model.	required
processor	BundleProcessor | None	Optional BundleProcessor. Default is None.	None

evaluate_sample(model: Module, bundle: DataBundle) -> dict[str, Float[Array, '']]

Evaluates a single steady-state sample.

Processes the input sample through normalizers and structural schemas, executes the neural operator, and computes the error metrics.

Parameters:

Name	Type	Description	Default
model	Module	The neural operator model to evaluate.	required
bundle	DataBundle	A single unbatched DataBundle.	required

Returns:

Type	Description
dict[str, Float[Array, '']]	Dictionary of computed metric values.

2. Time Dependent Evaluator

Used for time-dependent PDE problems where future states are predicted autoregressively from history. The evaluator uses a stateful scan rollout loop to evaluate accuracy over long time horizons.

neojax.benchmark.evaluators.time_dependent_eval.TimeDependentEvaluator

Evaluator for time-dependent PDE problems using autoregressive rollouts.

Iterates over temporal trajectories by feeding predicted future steps back as inputs via jax.lax.scan.

Parameters:

Name	Type	Description	Default
metrics	dict[str, BaseMetric]	A dictionary mapping metric names to BaseMetric instances.	required
in_schema	BaseSchema	Input schema for the model.	required
out_schema	BaseSchema	Output schema for the model.	required
history_steps	int	Number of lookback/history steps required for the model input.	required
processor	BundleProcessor | None	Optional BundleProcessor. Default is None.	None

evaluate_sample(model: Module, bundle: DataBundle) -> dict[str, Float[Array, '']]

Evaluates a single time-dependent sample using autoregressive rollout.

Parameters:

Name	Type	Description	Default
model	Module	The neural operator model to evaluate.	required
bundle	DataBundle	A single unbatched DataBundle.	required

Returns:

Type	Description
dict[str, Float[Array, '']]	Dictionary of computed metric values.