Forward Exponax Release (#10)

* Requires updated Exponax

* Reactivate instantiation tests

* Fix to usage of

* Change to new Exponax interface

* Add test that runs a simple training

* Only test on one scenario per mode

* Fix usage of metrics

* Also test in higher dimensions

* Hardcode Exponax version for reproducibility

* Change to conservative mode by default to ensure consistency with previous versions

* Reduce resolution to have test run easier on GitHub CI

apebench/_base_scenario.py

@@ -777,7 +777,11 @@ def full_loss(
     def perform_test_rollout(
         self,
         neural_stepper: eqx.Module,
-        mean_error_fn: Callable = ex.metrics.mean_nRMSE,
+        mean_error_fn: Callable = lambda pred, ref: ex.metrics.mean_metric(
+            ex.metrics.nRMSE,
+            pred,
+            ref,
+        ),
     ) -> Float[Array, "test_temporal_horizon"]:
         """
         Rollout the neural stepper starting from the test initial condition and
@@ -816,21 +820,42 @@ def perform_tests(
 
         for metric in metrics:
             if metric == "mean_MSE":
-                results["mean_MSE"] = ex.metrics.mean_MSE
+                results["mean_MSE"] = lambda pred, ref: ex.metrics.mean_metric(
+                    ex.metrics.MSE,
+                    pred,
+                    ref,
+                )
             elif metric == "mean_nMSE":
-                results["mean_nMSE"] = ex.metrics.mean_nMSE
+                results["mean_nMSE"] = lambda pred, ref: ex.metrics.mean_metric(
+                    ex.metrics.nMSE,
+                    pred,
+                    ref,
+                )
             elif metric == "mean_RMSE":
-                results["mean_RMSE"] = ex.metrics.mean_RMSE
+                results["mean_RMSE"] = lambda pred, ref: ex.metrics.mean_metric(
+                    ex.metrics.RMSE,
+                    pred,
+                    ref,
+                )
             elif metric == "mean_nRMSE":
-                results["mean_nRMSE"] = ex.metrics.mean_nRMSE
+                results["mean_nRMSE"] = lambda pred, ref: ex.metrics.mean_metric(
+                    ex.metrics.nRMSE,
+                    pred,
+                    ref,
+                )
             elif metric == "mean_correlation":
-                results["mean_correlation"] = ex.metrics.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_fourier_nRMSE(
+                    results[metric] = lambda pred, ref: ex.metrics.mean_metric(
+                        ex.metrics.fourier_nRMSE,
                         pred,
                         ref,
                         low=low,

apebench/scenarios/difficulty/_convection.py

@@ -10,6 +10,7 @@
 class Convection(BaseScenario):
     gammas: tuple[float, ...] = (0.0, 0.0, 1.5, 0.0, 0.0)
     convection_delta: float = -1.5
+    conservative: bool = True
 
     num_substeps: int = 1
 
@@ -27,12 +28,13 @@ def _build_stepper(self, gammas, delta):
         substepped_gammas = tuple(g / self.num_substeps for g in gammas)
         substepped_delta = delta / self.num_substeps
 
-        substepped_stepper = ex.normalized.DifficultyConvectionStepper(
+        substepped_stepper = ex.stepper.generic.DifficultyConvectionStepper(
             self.num_spatial_dims,
             self.num_points,
             linear_difficulties=substepped_gammas,
             # Need minus to move the convection to the right hand side
             convection_difficulty=-substepped_delta,
+            conservative=self.conservative,
             order=self.order,
             dealiasing_fraction=self.dealiasing_fraction,
             num_circle_points=self.num_circle_points,

apebench/scenarios/difficulty/_linear.py

@@ -8,17 +8,17 @@ class Linear(BaseScenario):
     coarse_proportion: float = 0.5
 
     def get_ref_stepper(self):
-        return ex.normalized.DifficultyLinearStepper(
+        return ex.stepper.generic.DifficultyLinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             num_points=self.num_points,
-            difficulties=self.gammas,
+            linear_difficulties=self.gammas,
         )
 
     def get_coarse_stepper(self) -> ex.BaseStepper:
-        return ex.normalized.DifficultyLinearStepper(
+        return ex.stepper.generic.DifficultyLinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             num_points=self.num_points,
-            difficulties=tuple(f * self.coarse_proportion for f in self.gammas),
+            linear_difficulties=tuple(f * self.coarse_proportion for f in self.gammas),
         )
 
     def get_scenario_name(self) -> str:

apebench/scenarios/difficulty/_nonlinear.py

@@ -34,7 +34,7 @@ def _build_stepper(self, gammas, deltas):
         substepped_gammas = tuple(g / self.num_substeps for g in gammas)
         substepped_deltas = tuple(d / self.num_substeps for d in deltas)
 
-        substepped_stepper = ex.normalized.DifficultyGeneralNonlinearStepper(
+        substepped_stepper = ex.stepper.generic.DifficultyNonlinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             num_points=self.num_points,
             linear_difficulties=substepped_gammas,

apebench/scenarios/normalized/_convection.py

@@ -6,6 +6,7 @@
 class Convection(BaseScenario):
     alphas: tuple[float, ...] = (0.0, 0.0, 3.0e-5, 0.0, 0.0)
     convection_beta: float = -1.25e-2
+    conservative: bool = True
 
     num_substeps: int = 1
 
@@ -23,12 +24,13 @@ def _build_stepper(self, convection, alphas):
         substepped_convection = convection / self.num_substeps
         substepped_alphas = tuple(a / self.num_substeps for a in alphas)
 
-        substepped_stepper = ex.normalized.NormalizedConvectionStepper(
+        substepped_stepper = ex.stepper.generic.NormalizedConvectionStepper(
             self.num_spatial_dims,
             self.num_points,
-            normalized_coefficients=substepped_alphas,
+            normalized_linear_coefficients=substepped_alphas,
             # Need minus to move the convection to the right hand side
             normalized_convection_scale=-substepped_convection,
+            conservative=self.conservative,
             order=self.order,
             dealiasing_fraction=self.dealiasing_fraction,
             num_circle_points=self.num_circle_points,

apebench/scenarios/normalized/_linear.py

@@ -8,17 +8,17 @@ class Linear(BaseScenario):
     coarse_proportion: float = 0.5
 
     def get_ref_stepper(self):
-        return ex.normalized.NormalizedLinearStepper(
+        return ex.stepper.generic.NormalizedLinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             num_points=self.num_points,
-            normalized_coefficients=self.alphas,
+            normalized_linear_coefficients=self.alphas,
         )
 
     def get_coarse_stepper(self) -> ex.BaseStepper:
-        return ex.normalized.NormalizedLinearStepper(
+        return ex.stepper.generic.NormalizedLinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             num_points=self.num_points,
-            normalized_coefficients=tuple(
+            normalized_linear_coefficients=tuple(
                 f * self.coarse_proportion for f in self.alphas
             ),
         )

apebench/scenarios/normalized/_nonlinear.py

@@ -28,11 +28,11 @@ def _build_stepper(self, alphas, betas):
         substepped_alphas = tuple(a / self.num_substeps for a in alphas)
         substepped_betas = tuple(b / self.num_substeps for b in betas)
 
-        substepped_stepper = ex.normalized.NormlizedGeneralNonlinearStepper(
+        substepped_stepper = ex.stepper.generic.NormalizedNonlinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             num_points=self.num_points,
-            normalized_coefficients_linear=substepped_alphas,
-            normalized_coefficients_nonlinear=substepped_betas,
+            normalized_linear_coefficients=substepped_alphas,
+            normalized_nonlinear_coefficients=substepped_betas,
             order=self.order,
             dealiasing_fraction=self.dealiasing_fraction,
             num_circle_points=self.num_circle_points,

apebench/scenarios/physical/_convection.py

@@ -9,6 +9,7 @@ class Convection(BaseScenario):
 
     a_coefs: tuple[float, ...] = (0.0, 0.0, 0.0003, 0.0, 0.0)
     convection_coef: float = -0.125
+    conservative: bool = True
 
     num_substeps: int = 1
 
@@ -23,14 +24,15 @@ def __post_init__(self):
         self.num_channels = self.num_spatial_dims  # Overwrite
 
     def _build_stepper(self, dt):
-        substepped_stepper = ex.stepper.GeneralConvectionStepper(
+        substepped_stepper = ex.stepper.generic.GeneralConvectionStepper(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
             dt=dt / self.num_substeps,
-            coefficients=self.a_coefs,
+            linear_coefficients=self.a_coefs,
             # Need minus to move the convection to the right hand side
             convection_scale=-self.convection_coef,
+            conservative=self.conservative,
             order=self.order,
             dealiasing_fraction=self.dealiasing_fraction,
             num_circle_points=self.num_circle_points,

apebench/scenarios/physical/_gray_scott.py

@@ -37,7 +37,7 @@ def get_ic_generator(self) -> BaseRandomICGenerator:
         )
 
     def _build_stepper(self, dt):
-        substepped_stepper = ex.reaction.GrayScott(
+        substepped_stepper = ex.stepper.reaction.GrayScott(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
@@ -177,7 +177,7 @@ def get_ic_generator(self) -> BaseRandomICGenerator:
 
     def _build_stepper(self, dt):
         feed_rate, kill_rate = self.get_feed_and_kill_rate(self.pattern_type)
-        substepped_stepper = ex.reaction.GrayScott(
+        substepped_stepper = ex.stepper.reaction.GrayScott(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,

apebench/scenarios/physical/_linear.py

@@ -11,21 +11,21 @@ class Linear(BaseScenario):
     coarse_proportion: float = 0.5
 
     def get_ref_stepper(self):
-        return ex.stepper.GeneralLinearStepper(
+        return ex.stepper.generic.GeneralLinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
             dt=self.dt,
-            coefficients=self.a_coefs,
+            linear_coefficients=self.a_coefs,
         )
 
     def get_coarse_stepper(self) -> ex.BaseStepper:
-        return ex.stepper.GeneralLinearStepper(
+        return ex.stepper.generic.GeneralLinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
             dt=self.dt * self.coarse_proportion,
-            coefficients=self.a_coefs,
+            linear_coefficients=self.a_coefs,
         )
 
     def get_scenario_name(self) -> str:

apebench/scenarios/physical/_nonlinear.py

@@ -28,13 +28,13 @@ def __post_init__(self):
         pass
 
     def _build_stepper(self, dt):
-        substepped_stepper = ex.stepper.GeneralNonlinearStepper(
+        substepped_stepper = ex.stepper.generic.GeneralNonlinearStepper(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
             dt=dt / self.num_substeps,
-            coefficients_linear=self.a_coefs,
-            coefficients_nonlinear=self.b_coefs,
+            linear_coefficients=self.a_coefs,
+            nonlinear_coefficients=self.b_coefs,
             order=self.order,
             dealiasing_fraction=self.dealiasing_fraction,
             num_circle_points=self.num_circle_points,

apebench/scenarios/physical/_polynomial.py

@@ -23,13 +23,13 @@ def __post_init__(self):
         pass
 
     def _build_stepper(self, dt):
-        substepped_stepper = ex.stepper.GeneralPolynomialStepper(
+        substepped_stepper = ex.stepper.generic.GeneralPolynomialStepper(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
             dt=dt / self.num_substeps,
-            coefficients=self.a_coefs,
-            polynomial_scales=self.poly_coefs,
+            linear_coefficients=self.a_coefs,
+            polynomial_coefficients=self.poly_coefs,
             order=self.order,
             dealiasing_fraction=self.dealiasing_fraction,
             num_circle_points=self.num_circle_points,

apebench/scenarios/physical/_swift_hohenberg.py

@@ -22,14 +22,14 @@ def __post_init__(self):
             raise ValueError("Swift-Hohenberg is only supported for 2D and 3D")
 
     def _build_stepper(self, dt):
-        substepped_stepper = ex.reaction.SwiftHohenberg(
+        substepped_stepper = ex.stepper.reaction.SwiftHohenberg(
             num_spatial_dims=self.num_spatial_dims,
             domain_extent=self.domain_extent,
             num_points=self.num_points,
             dt=dt / self.num_substeps,
             reactivity=self.reactivity,
             critical_number=self.critical_number,
-            polynomial_coefficients=self.polynomial_coefficients,
+            polynomial_linear_coefficients=self.polynomial_coefficients,
         )
 
         if self.num_substeps == 1:

pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
   "matplotlib>=3.8.1",
   "pandas>=2.2.0",
   "seaborn>=0.13.0",
-  "exponax==0.0.1",
+  "exponax==0.1.0",
   "pdequinox==0.1.2",
   "trainax==0.0.2",
 ]

tests/test_builtin_scenarios.py

@@ -1,4 +1,4 @@
-# import pytest
+import pytest
 
 import apebench
 
@@ -7,17 +7,56 @@ def test_simple():
     apebench.scenarios.difficulty.Advection()
 
 
-# @pytest.mark.parametrize(
-#     "name",
-#     list(apebench.scenarios.scenario_dict.keys()),
-# )
-# def test_builtin_scenarios(name: str):
-#     # Some scenarios might not work in 1d, (which is the default number of spatial dims)
-#     try:
-#         scene = apebench.scenarios.scenario_dict[name]()
-#     except ValueError:
-#         return
+@pytest.mark.parametrize(
+    "name",
+    list(apebench.scenarios.scenario_dict.keys()),
+)
+def test_builtin_scenarios(name: str):
+    # Some scenarios might not work in 1d, (which is the default number of spatial dims)
+    try:
+        scene = apebench.scenarios.scenario_dict[name]()
+    except ValueError:
+        return
 
-#     ref = scene.get_ref_sample_data()
+    ref = scene.get_ref_sample_data()
 
-#     del ref
+    del ref
+
+
+@pytest.mark.parametrize(
+    "name,num_spatial_dims",
+    [
+        (name, num_spatial_dims)
+        for name in [
+            "phy_adv",
+            "norm_adv",
+            "diff_adv",
+        ]
+        for num_spatial_dims in [1, 2, 3]
+    ],
+)
+def test_simple_training(name: str, num_spatial_dims: int):
+    NUM_TRAIN_SAMPLES = 5
+    NUM_TEST_SAMPLES = 5
+    NUM_POINTS = 15
+    OPTIM_CONFIG = "adam;10;constant;1e-4"
+
+    # Some scenarios might not work in 1d, (which is the default number of spatial dims)
+    try:
+        scene = apebench.scenarios.scenario_dict[name](
+            num_spatial_dims=num_spatial_dims,
+            num_train_samples=NUM_TRAIN_SAMPLES,
+            num_test_samples=NUM_TEST_SAMPLES,
+            num_points=NUM_POINTS,
+            optim_config=OPTIM_CONFIG,
+        )
+    except ValueError:
+        return
+
+    NETWORK_CONFIG = "Conv;10;2;relu"
+
+    data, trained_net = scene(
+        network_config=NETWORK_CONFIG,
+    )
+
+    del data, trained_net