From 6167047289b1e5e7998fe59a2e4eff1bceec7189 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felix=20K=C3=B6hler?= <27728103+Ceyron@users.noreply.github.com> Date: Wed, 23 Oct 2024 11:16:27 +0200 Subject: [PATCH] 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 | 39 +++++++++-- apebench/scenarios/difficulty/_convection.py | 4 +- apebench/scenarios/difficulty/_linear.py | 8 +-- apebench/scenarios/difficulty/_nonlinear.py | 2 +- apebench/scenarios/normalized/_convection.py | 6 +- apebench/scenarios/normalized/_linear.py | 8 +-- apebench/scenarios/normalized/_nonlinear.py | 6 +- apebench/scenarios/physical/_convection.py | 6 +- apebench/scenarios/physical/_gray_scott.py | 4 +- apebench/scenarios/physical/_linear.py | 8 +-- apebench/scenarios/physical/_nonlinear.py | 6 +- apebench/scenarios/physical/_polynomial.py | 6 +- .../scenarios/physical/_swift_hohenberg.py | 4 +- pyproject.toml | 2 +- tests/test_builtin_scenarios.py | 65 +++++++++++++++---- 15 files changed, 122 insertions(+), 52 deletions(-) diff --git a/apebench/_base_scenario.py b/apebench/_base_scenario.py index dc35f82..4096b11 100644 --- a/apebench/_base_scenario.py +++ b/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, diff --git a/apebench/scenarios/difficulty/_convection.py b/apebench/scenarios/difficulty/_convection.py index fd6f93d..f80e55b 100644 --- a/apebench/scenarios/difficulty/_convection.py +++ b/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, diff --git a/apebench/scenarios/difficulty/_linear.py b/apebench/scenarios/difficulty/_linear.py index f82a600..d38176e 100644 --- a/apebench/scenarios/difficulty/_linear.py +++ b/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: diff --git a/apebench/scenarios/difficulty/_nonlinear.py b/apebench/scenarios/difficulty/_nonlinear.py index 5ca7dbf..0fdad62 100644 --- a/apebench/scenarios/difficulty/_nonlinear.py +++ b/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, diff --git a/apebench/scenarios/normalized/_convection.py b/apebench/scenarios/normalized/_convection.py index fe5b5f3..2dd7d36 100644 --- a/apebench/scenarios/normalized/_convection.py +++ b/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, diff --git a/apebench/scenarios/normalized/_linear.py b/apebench/scenarios/normalized/_linear.py index 5717bd3..84d170f 100644 --- a/apebench/scenarios/normalized/_linear.py +++ b/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 ), ) diff --git a/apebench/scenarios/normalized/_nonlinear.py b/apebench/scenarios/normalized/_nonlinear.py index 5f0f3c9..c6d35c2 100644 --- a/apebench/scenarios/normalized/_nonlinear.py +++ b/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, diff --git a/apebench/scenarios/physical/_convection.py b/apebench/scenarios/physical/_convection.py index daa6aba..2031f27 100644 --- a/apebench/scenarios/physical/_convection.py +++ b/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, diff --git a/apebench/scenarios/physical/_gray_scott.py b/apebench/scenarios/physical/_gray_scott.py index 8a31c2f..3ea82ef 100644 --- a/apebench/scenarios/physical/_gray_scott.py +++ b/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, diff --git a/apebench/scenarios/physical/_linear.py b/apebench/scenarios/physical/_linear.py index 982abbb..5f56e90 100644 --- a/apebench/scenarios/physical/_linear.py +++ b/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: diff --git a/apebench/scenarios/physical/_nonlinear.py b/apebench/scenarios/physical/_nonlinear.py index 117b209..53b49b0 100644 --- a/apebench/scenarios/physical/_nonlinear.py +++ b/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, diff --git a/apebench/scenarios/physical/_polynomial.py b/apebench/scenarios/physical/_polynomial.py index 734740e..c61578e 100644 --- a/apebench/scenarios/physical/_polynomial.py +++ b/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, diff --git a/apebench/scenarios/physical/_swift_hohenberg.py b/apebench/scenarios/physical/_swift_hohenberg.py index b48024a..1d61a63 100644 --- a/apebench/scenarios/physical/_swift_hohenberg.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 10c41b5..b695e91 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/tests/test_builtin_scenarios.py b/tests/test_builtin_scenarios.py index 0c91762..56f4832 100644 --- a/tests/test_builtin_scenarios.py +++ b/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