Repository: tum-pbs/apebench

Issue:
Using multiple `fourier_nRMSE` report metrics have all of them being calculated incorrectly
Maybe an issue with how multiple metrics are parsed?

Discussion:

Solution patch:
diff --git a/apebench/__init__.py b/apebench/__init__.py
index 8dba510..3b4d79b 100644
--- a/apebench/__init__.py
+++ b/apebench/__init__.py
@@ -3,7 +3,7 @@
 import trainax
 
 from . import _scraper as scraper
-from . import scenarios
+from . import components, scenarios
 from ._base_scenario import BaseScenario
 from ._extensions import arch_extensions
 from ._run import (
@@ -51,4 +51,5 @@
     "aggregate_gmean",
     "relative_by_config",
     "split_train",
+    "components",
 ]
diff --git a/apebench/_base_scenario.py b/apebench/_base_scenario.py
index 4096b11..fc9b035 100644
--- a/apebench/_base_scenario.py
+++ b/apebench/_base_scenario.py
@@ -15,6 +15,7 @@
 
 from ._corrected_stepper import CorrectedStepper
 from ._extensions import arch_extensions
+from .components import metric_dict
 
 
 class BaseScenario(eqx.Module, ABC):
@@ -818,59 +819,20 @@ def perform_tests(
 
         results = {}
 
-        for metric in metrics:
-            if metric == "mean_MSE":
-                results["mean_MSE"] = lambda pred, ref: ex.metrics.mean_metric(
-                    ex.metrics.MSE,
-                    pred,
-                    ref,
-                )
-            elif metric == "mean_nMSE":
-                results["mean_nMSE"] = lambda pred, ref: ex.metrics.mean_metric(
-                    ex.metrics.nMSE,
-                    pred,
-                    ref,
-                )
-            elif metric == "mean_RMSE":
-                results["mean_RMSE"] = lambda pred, ref: ex.metrics.mean_metric(
-                    ex.metrics.RMSE,
-                    pred,
-                    ref,
-                )
-            elif metric == "mean_nRMSE":
-                results["mean_nRMSE"] = lambda pred, ref: ex.metrics.mean_metric(
-                    ex.metrics.nRMSE,
-                    pred,
-                    ref,
-                )
-            elif metric == "mean_correlation":
-                results["mean_correlation"] = lambda pred, ref: ex.metrics.mean_metric(
-                    ex.metrics.correlation,
-                    pred,
-                    ref,
-                )
-            else:
-                metric_args = metric.split(";")
-                if metric_args[0] == "mean_fourier_nRMSE":
-                    low = int(metric_args[1])
-                    high = int(metric_args[2])
-                    results[metric] = lambda pred, ref: ex.metrics.mean_metric(
-                        ex.metrics.fourier_nRMSE,
-                        pred,
-                        ref,
-                        low=low,
-                        high=high,
-                    )
-                else:
-                    raise ValueError(f"Unknown metric '{metric}'")
+        for metric_config in metrics:
+            metric_args = metric_config.split(";")
+            metric_name = metric_args[0]
+            metric_constructor = metric_dict[metric_name]
+            metric_fn = metric_constructor(metric_config)
+            results[metric_config] = metric_fn
 
-        for metric, func in results.items():
+        for metric_config, func in results.items():
             exec_func = lambda model: self.perform_test_rollout(model, func)
             if remove_singleton_axis:
                 # add singleton axis for compatibility
-                results[metric] = exec_func(neural_stepper)[None]
+                results[metric_config] = exec_func(neural_stepper)[None]
             else:
-                results[metric] = eqx.filter_vmap(exec_func)(neural_stepper)
+                results[metric_config] = eqx.filter_vmap(exec_func)(neural_stepper)
 
         return results
 
diff --git a/apebench/components/__init__.py b/apebench/components/__init__.py
new file mode 100644
index 0000000..f5afbce
--- /dev/null
+++ b/apebench/components/__init__.py
@@ -0,0 +1,3 @@
+from ._metrics import metric_dict
+
+__all__ = ["metric_dict"]
diff --git a/apebench/components/_metrics.py b/apebench/components/_metrics.py
new file mode 100644
index 0000000..934c183
--- /dev/null
+++ b/apebench/components/_metrics.py
@@ -0,0 +1,149 @@
+from typing import Callable, Dict
+
+import exponax as ex
+from jaxtyping import Array, Float
+
+metric_dict: Dict[
+    str,
+    Callable[
+        [
+            str,  # metric_config
+        ],
+        Callable[
+            [
+                Float[Array, "batch channel ... N"],  # batched pred
+                Float[Array, "batch channel ... N"],  # batched target
+            ],
+            float,
+        ],
+    ],
+] = {
+    "mean_MAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.MAE,
+        pred,
+        ref,
+    ),
+    "mean_nMAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.nMAE,
+        pred,
+        ref,
+    ),
+    "mean_sMAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.sMAE,
+        pred,
+        ref,
+    ),
+    "mean_MSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.MSE,
+        pred,
+        ref,
+    ),
+    "mean_nMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.nMSE,
+        pred,
+        ref,
+    ),
+    "mean_sMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.sMSE,
+        pred,
+        ref,
+    ),
+    "mean_RMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.RMSE,
+        pred,
+        ref,
+    ),
+    "mean_nRMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.nRMSE,
+        pred,
+        ref,
+    ),
+    "mean_sRMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.sRMSE,
+        pred,
+        ref,
+    ),
+    "mean_fourier_MAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.fourier_MAE,
+        pred,
+        ref,
+        low=int(metric_config.split(";")[1]),
+        high=int(metric_config.split(";")[2]),
+        derivative_order=float(metric_config.split(";")[3]),
+    ),
+    "mean_fourier_nMAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.fourier_nMAE,
+        pred,
+        ref,
+        low=int(metric_config.split(";")[1]),
+        high=int(metric_config.split(";")[2]),
+        derivative_order=float(metric_config.split(";")[3]),
+    ),
+    "mean_fourier_MSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.fourier_MSE,
+        pred,
+        ref,
+        low=int(metric_config.split(";")[1]),
+        high=int(metric_config.split(";")[2]),
+        derivative_order=float(metric_config.split(";")[3]),
+    ),
+    "mean_fourier_nMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.fourier_nMSE,
+        pred,
+        ref,
+        low=int(metric_config.split(";")[1]),
+        high=int(metric_config.split(";")[2]),
+        derivative_order=float(metric_config.split(";")[3]),
+    ),
+    "mean_fourier_RMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.fourier_RMSE,
+        pred,
+        ref,
+        low=int(metric_config.split(";")[1]),
+        high=int(metric_config.split(";")[2]),
+        derivative_order=float(metric_config.split(";")[3]),
+    ),
+    "mean_fourier_nRMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.fourier_nRMSE,
+        pred,
+        ref,
+        low=int(metric_config.split(";")[1]),
+        high=int(metric_config.split(";")[2]),
+        derivative_order=float(metric_config.split(";")[3]),
+    ),
+    "mean_H1_MAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.H1_MAE,
+        pred,
+        ref,
+    ),
+    "mean_H1_nMAE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.H1_nMAE,
+        pred,
+        ref,
+    ),
+    "mean_H1_MSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.H1_MSE,
+        pred,
+        ref,
+    ),
+    "mean_H1_nMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.H1_nMSE,
+        pred,
+        ref,
+    ),
+    "mean_H1_RMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.H1_RMSE,
+        pred,
+        ref,
+    ),
+    "mean_H1_nRMSE": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.H1_nRMSE,
+        pred,
+        ref,
+    ),
+    "mean_correlation": lambda metric_config: lambda pred, ref: ex.metrics.mean_metric(
+        ex.metrics.correlation,
+        pred,
+        ref,
+    ),
+}


Test patch:
diff --git a/tests/test_metric_system.py b/tests/test_metric_system.py
new file mode 100644
index 0000000..a07d34d
--- /dev/null
+++ b/tests/test_metric_system.py
@@ -0,0 +1,43 @@
+import jax
+import pytest
+
+import apebench
+
+
+@pytest.mark.parametrize(
+    "metric_config",
+    [
+        "mean_MAE",
+        "mean_nMAE",
+        "mean_sMAE",
+        "mean_MSE",
+        "mean_nMSE",
+        "mean_sMSE",
+        "mean_RMSE",
+        "mean_nRMSE",
+        "mean_sRMSE",
+        "mean_fourier_MAE;0;4;0",
+        "mean_fourier_nMAE;0;4;0",
+        "mean_fourier_MSE;0;4;0",
+        "mean_fourier_nMSE;0;4;0",
+        "mean_fourier_RMSE;0;4;0",
+        "mean_fourier_nRMSE;0;4;0",
+        "mean_H1_MAE",
+        "mean_H1_nMAE",
+        "mean_H1_MSE",
+        "mean_H1_nMSE",
+        "mean_H1_RMSE",
+        "mean_H1_nRMSE",
+        "mean_correlation",
+    ],
+)
+def test_metric_execution(metric_config: str):
+    pred = jax.random.uniform(jax.random.PRNGKey(0), (20, 3, 10))
+    ref = jax.random.uniform(jax.random.PRNGKey(1), (20, 3, 10))
+
+    metric_name = metric_config.split(";")[0]
+    metric_fn = apebench.components.metric_dict[metric_name](metric_config)
+    metric_result = metric_fn(pred, ref)
+
+    # Must be a float type in form of a shape-less JAX array
+    assert metric_result.shape == ()