From a460d5075bb88072efaffa5a5e1dddef6d46a55c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Thu, 4 Sep 2025 14:42:18 +0200 Subject: [PATCH 001/108] add first batch of program tests to continuous benchmarking --- .../stencil_tests/test_apply_diffusion_to_vn.py | 16 +++++++++++++++- ...mpute_horizontal_gradients_for_turbulence.py | 16 +++++++++++++++- ...ulate_nabla2_and_smag_coefficients_for_vn.py | 15 ++++++++++++++- ...iffusion_nabla_of_theta_over_steep_points.py | 17 +++++++++++++++-- .../test_mo_intp_rbf_rbf_vec_interpol_vertex.py | 17 +++++++++++++++-- 5 files changed, 74 insertions(+), 7 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index 946f1132a5..534b9079ce 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -15,7 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy from .test_apply_nabla2_and_nabla4_to_vn import apply_nabla2_and_nabla4_to_vn_numpy @@ -26,9 +26,23 @@ @pytest.mark.uses_concat_where +@pytest.mark.continuous_benchmarking class TestApplyDiffusionToVn(StencilTest): PROGRAM = apply_diffusion_to_vn OUTPUTS = ("vn",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: None, + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 493b5e3f1e..a78fef7025 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -15,7 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy from .test_apply_nabla2_to_w_in_upper_damping_layer import ( @@ -28,9 +28,23 @@ @pytest.mark.embedded_remap_error +@pytest.mark.continuous_benchmarking class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTest): PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence OUTPUTS = ("w", "dwdx", "dwdy") + STATIC_PARAMS = { + StandardStaticVariants.NONE: None, + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 83d006abb2..c6652317b3 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -18,10 +18,23 @@ from icon4py.model.testing import stencil_tests -@pytest.mark.skip_value_error +@pytest.mark.continuous_benchmarking class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_and_smag_coefficients_for_vn OUTPUTS = ("kh_smag_e", "kh_smag_ec", "z_nabla2_e") + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index acff8879c6..5588c9321c 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -19,7 +19,7 @@ random_mask, zero_field, ) -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( @@ -62,10 +62,23 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( @pytest.mark.uses_as_offset -@pytest.mark.skip_value_error +@pytest.mark.continuous_benchmarking class TestTrulyHorizontalDiffusionNablaOfThetaOverSteepPoints(StencilTest): PROGRAM = truly_horizontal_diffusion_nabla_of_theta_over_steep_points OUTPUTS = ("z_temp",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: None, + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index f6e5034fa9..21dc17dafe 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -18,13 +18,26 @@ ) from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants -@pytest.mark.skip_value_error +@pytest.mark.continuous_benchmarking class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex OUTPUTS = ("p_u_out", "p_v_out") + STATIC_PARAMS = { + StandardStaticVariants.NONE: None, + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( From 84abc04e9237d80daad8b888368b62642330a05e Mon Sep 17 00:00:00 2001 From: Rico Haeuselmann Date: Thu, 4 Sep 2025 15:13:10 +0200 Subject: [PATCH 002/108] fix static params for apply_diffusion_to_vn Co-authored-by: Hannes Vogt --- .../tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py | 1 + 1 file changed, 1 insertion(+) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index 534b9079ce..b3828cbc2e 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -35,6 +35,7 @@ class TestApplyDiffusionToVn(StencilTest): StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", + "start_2nd_nudge_line_idx_e", "vertical_start", "vertical_end", ), From 50e83651d1b9d9299308f50d63af452b465dd597 Mon Sep 17 00:00:00 2001 From: Rico Haeuselmann Date: Fri, 5 Sep 2025 10:22:34 +0200 Subject: [PATCH 003/108] improve static params for apply_diffusion_to_w_... test Co-authored-by: Hannes Vogt --- ...sion_to_w_and_compute_horizontal_gradients_for_turbulence.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index a78fef7025..8a238bbc75 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -39,8 +39,10 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTes "horizontal_end", "vertical_start", "vertical_end", + "limited_area", ), StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "limited_area", "vertical_start", "vertical_end", ), From aacd35bbdaaec116a84dd701b154472a588587c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 5 Sep 2025 11:07:22 +0200 Subject: [PATCH 004/108] add static params and remove indirectly used prog from benchmarking --- .../stencil_tests/test_apply_diffusion_to_vn.py | 2 ++ ..._compute_horizontal_gradients_for_turbulence.py | 8 ++++++-- ...l_diffusion_nabla_of_theta_over_steep_points.py | 14 -------------- 3 files changed, 8 insertions(+), 16 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index b3828cbc2e..b15794015a 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -38,10 +38,12 @@ class TestApplyDiffusionToVn(StencilTest): "start_2nd_nudge_line_idx_e", "vertical_start", "vertical_end", + "limited_area", ), StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", "vertical_end", + "limited_area", ), } diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 8a238bbc75..5353237816 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -37,14 +37,18 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTes StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", + "halo_idx", + "interior_idx", "vertical_start", "vertical_end", - "limited_area", + "nrdmax", + "type_shear", ), StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "limited_area", "vertical_start", "vertical_end", + "nrdmax", + "type_shear", ), } diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 5588c9321c..9003b5585b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -62,23 +62,9 @@ def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( @pytest.mark.uses_as_offset -@pytest.mark.continuous_benchmarking class TestTrulyHorizontalDiffusionNablaOfThetaOverSteepPoints(StencilTest): PROGRAM = truly_horizontal_diffusion_nabla_of_theta_over_steep_points OUTPUTS = ("z_temp",) - STATIC_PARAMS = { - StandardStaticVariants.NONE: None, - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( From d042c9864b8d29c7f8dea8ad9cc08c1b9ea0d681 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 5 Sep 2025 14:29:04 +0200 Subject: [PATCH 005/108] update dockerfile to trigger a rebuild with newer bencher CLI --- ci/docker/base.Dockerfile | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ci/docker/base.Dockerfile b/ci/docker/base.Dockerfile index 099ab1d06c..1b476e8bf3 100644 --- a/ci/docker/base.Dockerfile +++ b/ci/docker/base.Dockerfile @@ -37,6 +37,10 @@ ENV PATH="/root/.cargo/bin:${PATH}" RUN rustc --version && which rustc && cargo --version && which cargo # Install Bencher for performance monitoring +# Update the following comment to trigger a rebuild to update the CLI: +# last update: 2025-09-05 +# This is necessary because the cloud version and the CLI version have to match +# but obviously, version changes do not register in the Dockerfile hash. RUN curl --proto '=https' --tlsv1.2 -sSfL https://bencher.dev/download/install-cli.sh | sh RUN bencher --version && which bencher From ba94c10dd1640af1443e6a5b92eafb9709c3ed31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 12 Sep 2025 10:53:46 +0200 Subject: [PATCH 006/108] add compute_derived_horizontal_winds_... to benchmarking --- ...nds_and_ke_and_contravariant_correction.py | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 399165ea85..ccbb4e995c 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -48,6 +48,7 @@ def extrapolate_to_surface_numpy(wgtfacq_e: np.ndarray, vn: np.ndarray) -> np.nd return vn_at_surface +@pytest.mark.continuous_benchmarking @pytest.mark.embedded_remap_error class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection( stencil_tests.StencilTest @@ -61,6 +62,20 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava "contravariant_correction_at_edges_on_model_levels", "horizontal_advection_of_w_at_edges_on_half_levels", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "nflatlev", + ), + } @staticmethod def _fused_velocity_advection_stencil_1_to_6_numpy( @@ -258,8 +273,9 @@ def reference( horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, ) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture(params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]]) + def input_data(self, request, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + skip_compute_predictor_vertical_advection = request.param["skip_compute_predictor_vertical_advection"] horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( grid, dims.EdgeDim, dims.KDim ) @@ -289,8 +305,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala nlev = grid.num_levels nflatlev = 13 - skip_compute_predictor_vertical_advection = False - edge_domain = h_grid.domain(dims.EdgeDim) # For the ICON grid we use the proper domain bounds (otherwise we will run into non-protected skip values) horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) From 3dcdd34c732a88ee06ff167b00e66b702b7680d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Tue, 16 Sep 2025 13:20:21 +0200 Subject: [PATCH 007/108] temporarily remove failing parametrization case. --- ...ved_horizontal_winds_and_ke_and_contravariant_correction.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 0e10855b1c..801ef62bd6 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -273,7 +273,8 @@ def reference( horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, ) - @pytest.fixture(params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]]) + # TODO(ricoh): Add True case. Blocked by test failure (issue: #875) + @pytest.fixture(params=[{"skip_compute_predictor_vertical_advection": value} for value in [False]]) def input_data(self, request, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: skip_compute_predictor_vertical_advection = request.param["skip_compute_predictor_vertical_advection"] horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( From 367ba1f0ec0c8fec3d36a399273c300fa7788ff7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Tue, 16 Sep 2025 13:39:45 +0200 Subject: [PATCH 008/108] add compute_advection_in_vertical... program --- ...te_advection_in_vertical_momentum_equation.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index bff7b97677..709e7649d0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -268,6 +268,7 @@ def compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( return vertical_wind_advective_tendency +@pytest.mark.continuous_benchmarking @pytest.mark.embedded_remap_error class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): PROGRAM = compute_advection_in_vertical_momentum_equation @@ -276,6 +277,21 @@ class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): "contravariant_corrected_w_at_cells_on_model_levels", "vertical_cfl", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "end_index_of_damping_layer", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "end_index_of_damping_layer", + ), + } @staticmethod def reference( From c4b33c3743ae4dab0cfc838add9991224887861f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Tue, 16 Sep 2025 13:57:21 +0200 Subject: [PATCH 009/108] add benchmarking for advection_in_horizontal... program --- ...vection_in_horizontal_momentum_equation.py | 21 ++++++++++++++++--- ...nds_and_ke_and_contravariant_correction.py | 1 + 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py index 4340ee9ea4..26e8969600 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py @@ -138,6 +138,21 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelm class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): PROGRAM = compute_advection_in_horizontal_momentum_equation OUTPUTS = ("normal_wind_advective_tendency",) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "end_index_of_damping_layer", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "end_index_of_damping_layer", + "vertical_start", + "vertical_end", + ) + } @staticmethod def reference( @@ -227,8 +242,8 @@ def reference( return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture(params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]]) + def input_data(self, request, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: normal_wind_advective_tendency = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( @@ -257,7 +272,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala scalfac_exdiff = 0.6 dtime = 2.0 cfl_w_limit = 0.65 / dtime - apply_extra_diffusion_on_vn = True + apply_extra_diffusion_on_vn = request.param["apply_extra_diffusion_on_vn"] end_index_of_damping_layer = 5 edge_domain = h_grid.domain(dims.EdgeDim) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 801ef62bd6..083c1c8825 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -69,6 +69,7 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava "horizontal_end", "vertical_start", "vertical_end", + "nflatlev", ), stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "vertical_start", From bdc3296301c56c0b9d05f47d59412aac2f4456ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Tue, 16 Sep 2025 14:11:06 +0200 Subject: [PATCH 010/108] add benchmarking for compute_perturbed... program --- ..._perturbed_quantities_and_interpolation.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 778f768481..e9e31fb1e5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -64,6 +64,7 @@ def compute_first_vertical_derivative_numpy( return first_vertical_derivative +@pytest.mark.continuous_benchmarking @pytest.mark.uses_concat_where class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): PROGRAM = compute_perturbed_quantities_and_interpolation @@ -80,6 +81,33 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "pressure_buoyancy_acceleration_at_cells_on_half_levels", "d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "limited_area", + "igradp_method", + "start_cell_lateral_boundary_level_3", + "start_cell_halo_level_2", + "end_cell_end", + "end_cell_halo", + "end_cell_halo_level_2", + "start_cell_lateral_boundary", + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "nflatlev", + "nflat_gradp" + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "limited_area", + "igradp_method", + "vertical_start", + "vertical_end", + "nflatlev", + "nflat_gradp" + ), + } @staticmethod def reference( From 5836fbf51c7519e3f626a9a2e296cd8f8ba4fca8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Tue, 16 Sep 2025 14:17:24 +0200 Subject: [PATCH 011/108] add benchmarking for 'compute_hydrostatic...' program --- .../test_compute_hydrostatic_correction_term.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 79217ba75d..01e5727fdd 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants def compute_hydrostatic_correction_term_numpy( @@ -86,11 +86,24 @@ def _apply_index_field( return z_hydro_corr -@pytest.mark.skip_value_error +@pytest.mark.continuous_benchmarking @pytest.mark.uses_as_offset class TestComputeHydrostaticCorrectionTerm(StencilTest): OUTPUTS = ("z_hydro_corr",) PROGRAM = compute_hydrostatic_correction_term + STATIC_PARAMS = { + StandardStaticVariants.NONE: None, + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( From 7d041d53dd946421c0aff25b06b181ba61a7f337 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Tue, 16 Sep 2025 15:07:04 +0200 Subject: [PATCH 012/108] fix precommit messages --- .../stencil_tests/test_apply_diffusion_to_vn.py | 2 +- ..._compute_horizontal_gradients_for_turbulence.py | 2 +- ...te_advection_in_horizontal_momentum_equation.py | 8 +++++--- ...pute_advection_in_vertical_momentum_equation.py | 2 +- ...al_winds_and_ke_and_contravariant_correction.py | 14 ++++++++++---- .../test_compute_hydrostatic_correction_term.py | 4 ++-- ...mpute_perturbed_quantities_and_interpolation.py | 6 +++--- .../test_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- 8 files changed, 24 insertions(+), 16 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index b15794015a..2ca65484e3 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -15,7 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy from .test_apply_nabla2_and_nabla4_to_vn import apply_nabla2_and_nabla4_to_vn_numpy diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 1d3ff993e5..c15179a95b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -15,7 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy from .test_apply_nabla2_to_w_in_upper_damping_layer import ( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py index 26e8969600..ad09b5e60c 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py @@ -139,7 +139,7 @@ class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): PROGRAM = compute_advection_in_horizontal_momentum_equation OUTPUTS = ("normal_wind_advective_tendency",) STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.NONE: (), stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", @@ -151,7 +151,7 @@ class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): "end_index_of_damping_layer", "vertical_start", "vertical_end", - ) + ), } @staticmethod @@ -243,7 +243,9 @@ def reference( return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) @pytest.fixture(params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]]) - def input_data(self, request, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: normal_wind_advective_tendency = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index 709e7649d0..5c9ade921e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -278,7 +278,7 @@ class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): "vertical_cfl", ) STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.NONE: (), stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 083c1c8825..22d70561c0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -63,7 +63,7 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava "horizontal_advection_of_w_at_edges_on_half_levels", ) STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.NONE: (), stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", @@ -275,9 +275,15 @@ def reference( ) # TODO(ricoh): Add True case. Blocked by test failure (issue: #875) - @pytest.fixture(params=[{"skip_compute_predictor_vertical_advection": value} for value in [False]]) - def input_data(self, request, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - skip_compute_predictor_vertical_advection = request.param["skip_compute_predictor_vertical_advection"] + @pytest.fixture( + params=[{"skip_compute_predictor_vertical_advection": value} for value in [False]] + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + skip_compute_predictor_vertical_advection = request.param[ + "skip_compute_predictor_vertical_advection" + ] horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( grid, dims.EdgeDim, dims.KDim ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 01e5727fdd..1c041b2a89 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest def compute_hydrostatic_correction_term_numpy( @@ -92,7 +92,7 @@ class TestComputeHydrostaticCorrectionTerm(StencilTest): OUTPUTS = ("z_hydro_corr",) PROGRAM = compute_hydrostatic_correction_term STATIC_PARAMS = { - StandardStaticVariants.NONE: None, + StandardStaticVariants.NONE: (), StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index e9e31fb1e5..02a12daba5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -82,7 +82,7 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels", ) STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.NONE: (), stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "limited_area", "igradp_method", @@ -97,7 +97,7 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "vertical_start", "vertical_end", "nflatlev", - "nflat_gradp" + "nflat_gradp", ), stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "limited_area", @@ -105,7 +105,7 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "vertical_start", "vertical_end", "nflatlev", - "nflat_gradp" + "nflat_gradp", ), } diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 21dc17dafe..d551fc265e 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -18,7 +18,7 @@ ) from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest, StandardStaticVariants +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest @pytest.mark.continuous_benchmarking From a30e5da8f31dc0cebdf2bab01648ea44a7d30f81 Mon Sep 17 00:00:00 2001 From: Ong Chia Rui <93439630+OngChia@users.noreply.github.com> Date: Fri, 12 Sep 2025 14:50:58 +0200 Subject: [PATCH 013/108] fix a bug in grid-savepoint reading during driver initialization (#866) Fix bug in the input arguments of `IconSerialDataProvider.from_grid_savepoint` in `initialization_utils.py` of the driver package after this PR https://github.com/C2SM/icon4py/pull/829 is merged. --- .../model/driver/icon4py_configuration.py | 18 ++- .../icon4py/model/driver/icon4py_driver.py | 84 ++++------ .../model/driver/initialization_utils.py | 149 +++++++++++------- .../{test_timeloop.py => test_icon4py.py} | 88 +++++++---- model/driver/tests/driver/utils.py | 4 +- 5 files changed, 195 insertions(+), 148 deletions(-) rename model/driver/tests/driver/integration_tests/{test_timeloop.py => test_icon4py.py} (89%) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index fad9abcfef..066326d827 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -13,8 +13,7 @@ from gt4py.next import backend as gtx_backend, metrics as gtx_metrics from icon4py.model.atmosphere.diffusion import diffusion -from icon4py.model.atmosphere.dycore import solve_nonhydro as solve_nh -from icon4py.model.common import constants +from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh from icon4py.model.common.grid import vertical as v_grid from icon4py.model.driver import initialization_utils as driver_init @@ -74,7 +73,6 @@ def _mch_ch_r04b09_diffusion_config(): return diffusion.DiffusionConfig( diffusion_type=diffusion.DiffusionType.SMAGORINSKY_4TH_ORDER, hdiff_w=True, - n_substeps=n_substeps_reduced, hdiff_vn=True, type_t_diffu=2, type_vn_diffu=1, @@ -82,14 +80,20 @@ def _mch_ch_r04b09_diffusion_config(): hdiff_w_efdt_ratio=15.0, smagorinski_scaling_factor=0.025, zdiffu_t=True, + thslp_zdiffu=0.02, + thhgtd_zdiffu=125.0, velocity_boundary_diffusion_denom=150.0, - max_nudging_coefficient=0.075 * constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO, + max_nudging_coefficient=0.375, + n_substeps=n_substeps_reduced, + shear_type=diffusion.TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND, ) def _mch_ch_r04b09_nonhydro_config(): return solve_nh.NonHydrostaticConfig( - ndyn_substeps_var=n_substeps_reduced, - max_nudging_coefficient=0.075 * constants.DEFAULT_DYNAMICS_TO_PHYSICS_TIMESTEP_RATIO, + divdamp_order=dycore_states.DivergenceDampingOrder.COMBINED, + iau_wgt_dyn=1.0, + fourth_order_divdamp_factor=0.004, + max_nudging_coefficient=0.375, ) def _jabw_vertical_config(): @@ -126,7 +130,7 @@ def _mch_ch_r04b09_config(): Icon4pyRunConfig( dtime=datetime.timedelta(seconds=10.0), start_date=datetime.datetime(2021, 6, 20, 12, 0, 0), - end_date=datetime.datetime(2021, 6, 20, 12, 0, 10), + end_date=datetime.datetime(2021, 6, 20, 12, 0, 30), n_substeps=n_substeps_reduced, apply_initial_stabilization=True, backend=backend, diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py index 39d7131762..6edeed27fb 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -9,7 +9,6 @@ import datetime import logging import pathlib -import uuid from collections.abc import Callable from typing import NamedTuple @@ -339,9 +338,7 @@ def initialize( props: decomposition.ProcessProperties, serialization_type: driver_init.SerializationType, experiment_type: driver_init.ExperimentType, - grid_id: uuid.UUID, - grid_root, - grid_level, + grid_file: pathlib.Path, backend: gtx_backend.Backend, ) -> tuple[TimeLoop, DriverStates, DriverParams]: """ @@ -359,9 +356,8 @@ def initialize( props: Processor properties. serialization_type: Serialization type. experiment_type: Experiment type. - grid_id: Grid ID. - grid_root: Grid root. - grid_level: Grid level. + grid_file: Path of the grid. + backend: GT4Py backend. Returns: TimeLoop: Time loop object. @@ -373,24 +369,20 @@ def initialize( config = driver_config.read_config(experiment_type=experiment_type, backend=backend) decomp_info = driver_init.read_decomp_info( - file_path, - props, - backend, - serialization_type, - grid_id, - grid_root, - grid_level, + path=file_path, + grid_file=grid_file, + procs_props=props, + backend=backend, + ser_type=serialization_type, ) log.info(f"initializing the grid from '{file_path}'") - icon_grid = driver_init.read_icon_grid( - file_path, + grid = driver_init.read_icon_grid( + path=file_path, + grid_file=grid_file, backend=backend, rank=props.rank, ser_type=serialization_type, - grid_id=grid_id, - grid_root=grid_root, - grid_level=grid_level, ) log.info(f"reading input fields from '{file_path}'") ( @@ -399,14 +391,12 @@ def initialize( vertical_geometry, c_owner_mask, ) = driver_init.read_geometry_fields( - file_path, + path=file_path, + grid_file=grid_file, vertical_grid_config=config.vertical_grid_config, backend=backend, rank=props.rank, ser_type=serialization_type, - grid_id=grid_id, - grid_root=grid_root, - grid_level=grid_level, ) ( diffusion_metric_state, @@ -415,11 +405,9 @@ def initialize( solve_nonhydro_interpolation_state, _, ) = driver_init.read_static_fields( - grid_id, - grid_root, - grid_level, - file_path, - backend, + path=file_path, + grid_file=grid_file, + backend=backend, rank=props.rank, ser_type=serialization_type, ) @@ -428,7 +416,7 @@ def initialize( diffusion_params = diffusion.DiffusionParams(config.diffusion_config) exchange = decomposition.create_exchange(props, decomp_info) diffusion_granule = diffusion.Diffusion( - icon_grid, + grid, config.diffusion_config, diffusion_params, vertical_geometry, @@ -443,7 +431,7 @@ def initialize( nonhydro_params = solve_nh.NonHydrostaticParams(config.solve_nonhydro_config) solve_nonhydro_granule = solve_nh.SolveNonhydro( - grid=icon_grid, + grid=grid, backend=backend, config=config.solve_nonhydro_config, params=nonhydro_params, @@ -464,10 +452,10 @@ def initialize( prognostic_state_now, prognostic_state_next, ) = driver_init.read_initial_state( - icon_grid, - cell_geometry, - edge_geometry, - file_path, + grid=grid, + cell_param=cell_geometry, + edge_param=edge_geometry, + path=file_path, backend=backend, rank=props.rank, experiment_type=experiment_type, @@ -521,25 +509,14 @@ def initialize( "Currently, users can also set it to either jabw or grauss_3d_torus to generate analytic initial condition for the JW and mountain wave tests, respectively (they are placed in abs_path_to_icon4py/model/driver/src/icon4py/model/driver/test_cases/).", ) @click.option( - "--grid_root", - default=2, - show_default=True, - help="Grid root division (please refer to Sadourny et al. 1968 or ICON documentation for more information). When torus grid is used, it must be set to 2.", -) -@click.option( - "--grid_level", - default=4, - show_default=True, - help="Grid refinement level. When torus grid is used, it must be set to 0.", -) -@click.option( - "--grid_id", - default="af122aca-1dd2-11b2-a7f8-c7bf6bc21eba", - help="uuid of the horizontal grid ('uuidOfHGrid' from gridfile)", + "--grid_file", + required=True, + help="Path of the grid file.", ) @click.option( "--enable_output", is_flag=True, + default=False, help="Enable all debugging messages. Otherwise, only critical error messages are printed.", ) @click.option( @@ -560,9 +537,7 @@ def icon4py_driver( mpi, serialization_type, experiment_type, - grid_id, - grid_root, - grid_level, + grid_file, enable_output, enable_profiling, icon4py_driver_backend, @@ -596,7 +571,6 @@ def icon4py_driver( backend = model_backends.BACKENDS[icon4py_driver_backend] parallel_props = decomposition.get_processor_properties(decomposition.get_runtype(with_mpi=mpi)) - grid_id = uuid.UUID(grid_id) driver_init.configure_logging(run_path, experiment_type, enable_output, parallel_props) time_loop: TimeLoop @@ -607,9 +581,7 @@ def icon4py_driver( parallel_props, serialization_type, experiment_type, - grid_id, - grid_root, - grid_level, + pathlib.Path(grid_file), backend, ) log.info(f"Starting ICON dycore run: {time_loop.simulation_date.isoformat()}") diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index 89cdd68361..4424ec9914 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -7,10 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause import enum +import functools import logging import pathlib -import uuid +import netCDF4 as nc4 from gt4py.next import backend as gtx_backend from icon4py.model.atmosphere.diffusion import diffusion_states @@ -20,7 +21,12 @@ definitions as decomposition, mpi_decomposition as mpi_decomp, ) -from icon4py.model.common.grid import icon as icon_grid, states as grid_states, vertical as v_grid +from icon4py.model.common.grid import ( + base, + icon as icon_grid, + states as grid_states, + vertical as v_grid, +) from icon4py.model.common.states import ( diagnostic_state as diagnostics, prognostic_state as prognostics, @@ -31,11 +37,6 @@ from icon4py.model.testing import serialbox as sb -# TODO(egparedes): Read these hardcoded constants from grid file -GRID_LEVEL = 4 -GRID_ROOT = 2 -GLOBAL_GRID_ID = uuid.UUID("af122aca-1dd2-11b2-a7f8-c7bf6bc21eba") - SB_ONLY_MSG = "Only ser_type='sb' is implemented so far." INITIALIZATION_ERROR_MSG = "The requested experiment type is not implemented." @@ -59,37 +60,29 @@ class ExperimentType(str, enum.Enum): def read_icon_grid( path: pathlib.Path, + grid_file: pathlib.Path, backend: gtx_backend.Backend, - rank=0, + rank: int = 0, ser_type: SerializationType = SerializationType.SB, - grid_id=GLOBAL_GRID_ID, - grid_root=GRID_ROOT, - grid_level=GRID_LEVEL, ) -> icon_grid.IconGrid: """ Read icon grid. Args: path: path where to find the input data + grid_file: path of the grid + backend: GT4Py backend rank: mpi rank of the current compute node ser_type: type of input data. Currently only 'sb (serialbox)' is supported. It reads from ppser serialized test data - grid_id: id (uuid) of the horizontal grid - grid_root: global grid root division number - grid_level: global grid refinement number Returns: IconGrid parsed from a given input type. """ if ser_type == SerializationType.SB: - return ( - sb.IconSerialDataProvider( - backend=backend, - fname_prefix="icon_pydycore", - path=str(path.absolute()), - do_print=False, - mpi_rank=rank, - ) - .from_savepoint_grid(grid_id, grid_root, grid_level) - .construct_icon_grid(backend=backend) - ) + return _grid_savepoint( + backend=backend, + path=path, + grid_file=grid_file, + rank=rank, + ).construct_icon_grid(backend=backend) else: raise NotImplementedError(SB_ONLY_MSG) @@ -98,7 +91,7 @@ def model_initialization_serialbox( grid: icon_grid.IconGrid, path: pathlib.Path, backend: gtx_backend.Backend, - rank=0, + rank: int = 0, ) -> tuple[ diffusion_states.DiffusionDiagnosticState, dycore_states.DiagnosticStateNonHydro, @@ -115,6 +108,7 @@ def model_initialization_serialbox( Args: grid: IconGrid path: path where to find the input data + backend: GT4Py backend rank: mpi rank of the current compute node Returns: A tuple containing Diagnostic variables for diffusion and solve_nonhydro granules, PrepAdvection, second order divdamp factor, diagnostic variables, and two prognostic @@ -148,19 +142,21 @@ def model_initialization_serialbox( normal_wind_tendency_due_to_slow_physics_process=solve_nonhydro_init_savepoint.ddt_vn_phy(), grf_tend_vn=solve_nonhydro_init_savepoint.grf_tend_vn(), normal_wind_advective_tendency=common_utils.PredictorCorrectorPair( + velocity_init_savepoint.ddt_vn_apc_pc(0), velocity_init_savepoint.ddt_vn_apc_pc(1), - velocity_init_savepoint.ddt_vn_apc_pc(2), ), vertical_wind_advective_tendency=common_utils.PredictorCorrectorPair( + velocity_init_savepoint.ddt_w_adv_pc(0), velocity_init_savepoint.ddt_w_adv_pc(1), - velocity_init_savepoint.ddt_w_adv_pc(2), ), tangential_wind=velocity_init_savepoint.vt(), vn_on_half_levels=velocity_init_savepoint.vn_ie(), contravariant_correction_at_cells_on_half_levels=velocity_init_savepoint.w_concorr_c(), - rho_iau_increment=solve_nonhydro_init_savepoint.rho_incr(), - normal_wind_iau_increment=solve_nonhydro_init_savepoint.vn_incr(), - exner_iau_increment=solve_nonhydro_init_savepoint.exner_incr(), + rho_iau_increment=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, backend=backend), + normal_wind_iau_increment=data_alloc.zero_field( + grid, dims.EdgeDim, dims.KDim, backend=backend + ), + exner_iau_increment=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, backend=backend), exner_dynamical_increment=solve_nonhydro_init_savepoint.exner_dyn_incr(), ) @@ -328,13 +324,11 @@ def read_initial_state( def read_geometry_fields( path: pathlib.Path, + grid_file: pathlib.Path, vertical_grid_config: v_grid.VerticalGridConfig, backend: gtx_backend.Backend, - rank=0, + rank: int = 0, ser_type: SerializationType = SerializationType.SB, - grid_id=GLOBAL_GRID_ID, - grid_root=GRID_ROOT, - grid_level=GRID_LEVEL, ) -> tuple[ grid_states.EdgeParams, grid_states.CellParams, @@ -346,18 +340,17 @@ def read_geometry_fields( Args: path: path to the serialized input data + grid_file: path of the grid vertical_grid_config: Vertical grid configuration + backend: GT4py backend rank: mpi rank of the current compute node ser_type: (optional) defaults to SB=serialbox, type of input data to be read - grid_id: id (uuid) of the horizontal grid - grid_root: global grid root division number - grid_level: global grid refinement number Returns: a tuple containing fields describing edges, cells, vertical properties of the model the data is originally obtained from the grid file (horizontal fields) or some special input files. """ if ser_type == SerializationType.SB: - sp = _grid_savepoint(backend, path, rank, grid_id, grid_root, grid_level) + sp = _grid_savepoint(backend, path, grid_file, rank) edge_geometry = sp.construct_edge_geometry() cell_geometry = sp.construct_cell_geometry() vct_a, vct_b = v_grid.get_vct_a_and_vct_b(vertical_grid_config, backend) @@ -372,7 +365,11 @@ def read_geometry_fields( # TODO(OngChia): cannot be cached (@functools.cache) after adding backend. TypeError: unhashable type: 'CompiledbFactory' -def _serial_data_provider(backend, path, rank) -> sb.IconSerialDataProvider: +def _serial_data_provider( + backend: gtx_backend.Backend, + path: pathlib.Path, + rank: int, +) -> sb.IconSerialDataProvider: return sb.IconSerialDataProvider( backend=backend, fname_prefix="icon_pydycore", @@ -383,37 +380,43 @@ def _serial_data_provider(backend, path, rank) -> sb.IconSerialDataProvider: # TODO(OngChia): cannot be cached (@functools.cache) after adding backend. TypeError: unhashable type: 'CompiledbFactory' -def _grid_savepoint(backend, path, rank, grid_id, grid_root, grid_level) -> sb.IconGridSavepoint: +def _grid_savepoint( + backend: gtx_backend.Backend, + path: pathlib.Path, + grid_file: pathlib.Path, + rank: int, +) -> sb.IconGridSavepoint: + global_grid_params, grid_uuid = _create_grid_global_params(grid_file) sp = _serial_data_provider(backend, path, rank).from_savepoint_grid( - grid_id, grid_root, grid_level + grid_uuid, + global_grid_params.grid_shape, ) return sp def read_decomp_info( path: pathlib.Path, + grid_file: pathlib.Path, procs_props: decomposition.ProcessProperties, backend: gtx_backend.Backend, ser_type=SerializationType.SB, - grid_id=GLOBAL_GRID_ID, - grid_root=GRID_ROOT, - grid_level=GRID_LEVEL, ) -> decomposition.DecompositionInfo: if ser_type == SerializationType.SB: return _grid_savepoint( - backend, path, procs_props.rank, grid_id, grid_root, grid_level + backend, + path, + grid_file, + procs_props.rank, ).construct_decomposition_info() else: raise NotImplementedError(SB_ONLY_MSG) def read_static_fields( - grid_id: str, - grid_root: int, - grid_level: int, path: pathlib.Path, + grid_file: pathlib.Path, backend: gtx_backend.Backend, - rank=0, + rank: int = 0, ser_type: SerializationType = SerializationType.SB, ) -> tuple[ diffusion_states.DiffusionMetricState, @@ -426,8 +429,9 @@ def read_static_fields( Read fields for metric and interpolation state. Args: - grid: IconGrid path: path to the serialized input data + grid_file: path of the grid + backend: GT4Py backend rank: mpi rank, defaults to 0 for serial run ser_type: (optional) defaults to SB=serialbox, type of input data to be read @@ -466,7 +470,7 @@ def read_static_fields( nudgecoeff_e=interpolation_savepoint.nudgecoeff_e(), ) metrics_savepoint = data_provider.from_metrics_savepoint() - grid_savepoint = data_provider.from_savepoint_grid(grid_id, grid_root, grid_level) + grid_savepoint = _grid_savepoint(backend, path, grid_file, rank) solve_nonhydro_metric_state = dycore_states.MetricStateNonHydro( bdy_halo_c=metrics_savepoint.bdy_halo_c(), mask_prog_halo_c=metrics_savepoint.mask_prog_halo_c(), @@ -535,15 +539,19 @@ def configure_logging( Args: run_path: path to the output folder where the logfile should be stored experiment_name: name of the simulation + enable_output: enable output logging messages above debug level + processor_procs: ProcessProperties """ + if not enable_output: + return run_dir = ( pathlib.Path(run_path).absolute() if run_path else pathlib.Path(__file__).absolute().parent ) run_dir.mkdir(exist_ok=True) logfile = run_dir.joinpath(f"dummy_dycore_driver_{experiment_name}.log") logfile.touch(exist_ok=True) - logging_level = logging.DEBUG if enable_output else logging.CRITICAL + logging_level = logging.DEBUG logging.basicConfig( level=logging_level, format="%(asctime)s %(filename)-20s (%(lineno)-4d) : %(funcName)-20s: %(levelname)-8s %(message)s", @@ -559,3 +567,38 @@ def configure_logging( console_handler.setFormatter(formatter) console_handler.setLevel(logging_level) logging.getLogger("").addHandler(console_handler) + + +@functools.cache +def _create_grid_global_params( + grid_file: pathlib.Path, +) -> tuple[icon_grid.GlobalGridParams, str]: + """ + Create global grid params and its uuid. + + Args: + grid_file: path of the grid file + + Returns: + global_grid_params: GlobalGridParams + grid_uuid: id (uuid) of the horizontal grid + """ + grid = nc4.Dataset(grid_file, "r", format="NETCDF4") + grid_root = grid.getncattr("grid_root") + grid_level = grid.getncattr("grid_level") + grid_uuid = grid.getncattr("uuidOfHGrid") + try: + grid_geometry_type = base.GeometryType(grid.getncattr("grid_geometry")) + except AttributeError: + log.warning( + "Global attribute grid_geometry is not found in the grid. Icosahedral grid is assumed." + ) + grid_geometry_type = base.GeometryType.ICOSAHEDRON + grid.close() + global_grid_params = icon_grid.GlobalGridParams( + grid_shape=icon_grid.GridShape( + geometry_type=grid_geometry_type, + subdivision=icon_grid.GridSubdivision(root=grid_root, level=grid_level), + ), + ) + return global_grid_params, grid_uuid diff --git a/model/driver/tests/driver/integration_tests/test_timeloop.py b/model/driver/tests/driver/integration_tests/test_icon4py.py similarity index 89% rename from model/driver/tests/driver/integration_tests/test_timeloop.py rename to model/driver/tests/driver/integration_tests/test_icon4py.py index 3376b5a96c..fccd3ccd4c 100644 --- a/model/driver/tests/driver/integration_tests/test_timeloop.py +++ b/model/driver/tests/driver/integration_tests/test_icon4py.py @@ -9,13 +9,14 @@ from typing import TYPE_CHECKING +import click import pytest import icon4py.model.common.grid.states as grid_states import icon4py.model.common.utils as common_utils from icon4py.model.atmosphere.diffusion import diffusion from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh -from icon4py.model.common import dimension as dims +from icon4py.model.common import dimension as dims, model_backends from icon4py.model.common.grid import vertical as v_grid from icon4py.model.common.states import prognostic_state as prognostics from icon4py.model.common.utils import data_allocation as data_alloc @@ -25,7 +26,7 @@ initialization_utils as driver_init, serialbox_helpers as driver_sb, ) -from icon4py.model.testing import definitions, test_utils +from icon4py.model.testing import datatest_utils as dt_utils, definitions, grid_utils, test_utils from icon4py.model.testing.fixtures.datatest import backend from ..fixtures import * # noqa: F403 @@ -72,34 +73,6 @@ False, True, ), - ( - definitions.Experiments.EXCLAIM_APE, - 1, - 2, - 1, - 2, - "2000-01-01T00:00:00.000", - "2000-01-01T00:00:02.000", - "2000-01-01T00:00:02.000", - "2000-01-01T00:00:02.000", - False, - False, - False, - ), - ( - definitions.Experiments.EXCLAIM_APE, - 1, - 2, - 1, - 2, - "2000-01-01T00:00:02.000", - "2000-01-01T00:00:04.000", - "2000-01-01T00:00:04.000", - "2000-01-01T00:00:04.000", - False, - False, - True, - ), ( definitions.Experiments.GAUSS3D, 1, @@ -387,3 +360,58 @@ def test_run_timeloop_single_step( prognostic_states.current.rho.asnumpy(), rho_sp.asnumpy(), ) + + +@pytest.mark.embedded_remap_error +@pytest.mark.datatest +@pytest.mark.parametrize( + "experiment, experiment_type", + [ + ( + definitions.Experiments.MCH_CH_R04B09, + driver_init.ExperimentType.ANY.value, + ), + ], +) +def test_driver( + experiment, + experiment_type, + *, + data_provider, + ranked_data_path, + backend, +): + """ + This is a only test to check if the icon4py driver runs from serialized data without verifying the end result. + The timeloop is verified by test_run_timeloop_single_step above. + TODO(anyone): Remove or modify this test when it is ready to run the driver from the grid file without having to initialize static fields from serialized data. + """ + data_path = dt_utils.get_datapath_for_experiment( + ranked_base_path=ranked_data_path, + experiment=experiment, + ) + gm = grid_utils.get_grid_manager_from_experiment( + experiment=experiment, + keep_skip_values=True, + backend=backend, + ) + + backend_name = None + for key, value in model_backends.BACKENDS.items(): + if value == backend: + backend_name = key + + assert backend_name is not None + + icon4py_driver.icon4py_driver( + [ + str(data_path), + "--experiment_type", + experiment_type, + "--grid_file", + str(gm._file_name), + "--icon4py_driver_backend", + backend_name, + ], + standalone_mode=False, + ) diff --git a/model/driver/tests/driver/utils.py b/model/driver/tests/driver/utils.py index 4b99a214cd..1ae4389880 100644 --- a/model/driver/tests/driver/utils.py +++ b/model/driver/tests/driver/utils.py @@ -25,7 +25,7 @@ def mch_ch_r04b09_dsl_icon4pyrun_config( Create Icon4pyRunConfig matching MCH_CH_r04b09_dsl. Set values to the ones used in the MCH_CH_r04b09_dsl experiment where they differ - from the default. Backend is not used because granules are set independently in test_timeloop.py. + from the default. Backend is not used because granules are set independently in test_icon4py.py. """ return driver_config.Icon4pyRunConfig( dtime=timedelta(seconds=10.0), @@ -49,7 +49,7 @@ def exclaim_ape_icon4pyrun_config( Create Icon4pyRunConfig matching exclaim_ape_R02B04. Set values to the ones used in the exclaim_ape_R02B04 experiment where they differ - from the default. Backend is not used because granules are set independently in test_timeloop.py + from the default. Backend is not used because granules are set independently in test_icon4py.py """ return driver_config.Icon4pyRunConfig( dtime=timedelta(seconds=2.0), From ca2f4f211294c84d9eed8154a273b4ae0b2a9e60 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 12 Sep 2025 15:00:29 +0200 Subject: [PATCH 014/108] Fix cupy cfl reduction performance by raveling the array (#872) On the mch-ch1_medium experiment - this is 4% faster on the full timestep compared to the version without `ravel`; - within fluctuations there is no difference between no reduction and this version. --- .../src/icon4py/model/atmosphere/dycore/velocity_advection.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py index 04df7bddfe..ee9975abd9 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py @@ -279,11 +279,13 @@ def run_predictor_step( skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, ) + # Reductions should be performed on flat, contiguous arrays for best cupy performance + # as otherwise cupy won't use cub optimized kernels. max_vertical_cfl = float( self.vertical_cfl.array_ns.max( self.vertical_cfl.ndarray[ self._start_cell_lateral_boundary_level_4 : self._end_cell_halo, : - ] + ].ravel(order="K") ) ) diagnostic_state.max_vertical_cfl = max(max_vertical_cfl, diagnostic_state.max_vertical_cfl) From 36f5583c23672718c84f342263e2b7dcafd9a7ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philip=20M=C3=BCller?= <147368808+philip-paul-mueller@users.noreply.github.com> Date: Fri, 12 Sep 2025 16:20:21 +0200 Subject: [PATCH 015/108] Update to GT4Py v1.0.9 (dace fix) (#840) Co-authored-by: Edoardo Paone --- model/atmosphere/advection/pyproject.toml | 2 +- model/atmosphere/diffusion/pyproject.toml | 2 +- model/atmosphere/dycore/pyproject.toml | 2 +- .../microphysics/pyproject.toml | 2 +- .../muphys/pyproject.toml | 2 +- model/common/pyproject.toml | 2 +- model/driver/pyproject.toml | 2 +- model/testing/pyproject.toml | 2 +- tools/pyproject.toml | 2 +- uv.lock | 34 +++++++++---------- 10 files changed, 26 insertions(+), 26 deletions(-) diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index dc2c0d08b6..ed5008b47d 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ # workspace members "icon4py-common>=0.0.6", # external dependencies - "gt4py==1.0.8", + "gt4py==1.0.9", 'packaging>=20.0' ] description = "ICON advection." diff --git a/model/atmosphere/diffusion/pyproject.toml b/model/atmosphere/diffusion/pyproject.toml index 3075c9f668..8e47254167 100644 --- a/model/atmosphere/diffusion/pyproject.toml +++ b/model/atmosphere/diffusion/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ # workspace members "icon4py-common>=0.0.6", # external dependencies - "gt4py==1.0.8", + "gt4py==1.0.9", 'packaging>=20.0' ] description = "ICON diffusion." diff --git a/model/atmosphere/dycore/pyproject.toml b/model/atmosphere/dycore/pyproject.toml index 48cbab785a..ece60c4b75 100644 --- a/model/atmosphere/dycore/pyproject.toml +++ b/model/atmosphere/dycore/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ # workspace members "icon4py-common>=0.0.6", # external dependencies - "gt4py==1.0.8", + "gt4py==1.0.9", 'packaging>=20.0' ] description = "ICON dynamical core." diff --git a/model/atmosphere/subgrid_scale_physics/microphysics/pyproject.toml b/model/atmosphere/subgrid_scale_physics/microphysics/pyproject.toml index 8edb9fb35f..e80e262703 100644 --- a/model/atmosphere/subgrid_scale_physics/microphysics/pyproject.toml +++ b/model/atmosphere/subgrid_scale_physics/microphysics/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ # workspace members "icon4py-common>=0.0.6", # external dependencies - "gt4py==1.0.8", + "gt4py==1.0.9", 'packaging>=20.0' ] description = "ICON microphysics." diff --git a/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml b/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml index 0588161111..76b74470b0 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml +++ b/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ # workspace members "icon4py-common[io]>=0.0.6", # external dependencies - "gt4py==1.0.8", + "gt4py==1.0.9", "packaging>=20.0" ] description = "ICON subgrid scale muphys parameterization." diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index f06edf4fca..c59ff754f5 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ 'Topic :: Scientific/Engineering :: Physics' ] dependencies = [ - 'gt4py==1.0.8', + 'gt4py==1.0.9', 'packaging>=20.0', 'packaging>=20.0', 'typing-extensions>=4.11.0', diff --git a/model/driver/pyproject.toml b/model/driver/pyproject.toml index 51bc818dde..5ec5edfc6f 100644 --- a/model/driver/pyproject.toml +++ b/model/driver/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ # external dependencies "click>=8.0.1", "devtools>=0.12", - "gt4py==1.0.8", + "gt4py==1.0.9", "packaging>=20.0", "numpy>=1.23.3" ] diff --git a/model/testing/pyproject.toml b/model/testing/pyproject.toml index 3905ffda56..311168b49b 100644 --- a/model/testing/pyproject.toml +++ b/model/testing/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ 'icon4py-common[io]>=0.0.6', # external dependencies "filelock>=3.18.0", - "gt4py==1.0.8", + "gt4py==1.0.9", "numpy>=1.23.3", 'packaging>=20.0', "pytest>=8.0.1", diff --git a/tools/pyproject.toml b/tools/pyproject.toml index ae268b12f8..a8e7acc768 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ 'icon4py-common>=0.0.6', # external dependencies 'cffi>=1.5', - 'gt4py==1.0.8', + 'gt4py==1.0.9', "numpy>=1.23.3", 'packaging>=20.0', 'click>=8.1.7', diff --git a/uv.lock b/uv.lock index 34861433c3..9d46fbfd5b 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version < '3.11'", @@ -173,8 +173,8 @@ dependencies = [ { name = "packaging" }, { name = "pathspec" }, { name = "platformdirs" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "tomli", marker = "python_full_version < '3.11' or (extra == 'extra-7-icon4py-cuda11' and extra == 'extra-7-icon4py-cuda12')" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-7-icon4py-cuda11' and extra == 'extra-7-icon4py-cuda12')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/d8/0d/cc2fb42b8c50d80143221515dd7e4766995bd07c56c9a3ed30baf080b6dc/black-24.10.0.tar.gz", hash = "sha256:846ea64c97afe3bc677b761787993be4991810ecc7a4a937816dd6bddedc4875", size = 645813, upload-time = "2024-10-07T19:20:50.361Z" } wheels = [ @@ -587,7 +587,7 @@ name = "click" version = "8.1.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-7-icon4py-cuda11' and extra == 'extra-7-icon4py-cuda12')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121, upload-time = "2023-08-17T17:29:11.868Z" } wheels = [ @@ -1428,7 +1428,7 @@ wheels = [ [[package]] name = "gt4py" -version = "1.0.8" +version = "1.0.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, @@ -1460,9 +1460,9 @@ dependencies = [ { name = "versioningit" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/45/83/0d546c3b8987ddb771a0a9780bc74688bdd2f702e925ea30da667506b72a/gt4py-1.0.8.tar.gz", hash = "sha256:1dd686836377dbcbd4d0c20ba4757fdb4ed605d3468c548a5dabb52aed5d3047", size = 723685, upload-time = "2025-09-04T07:22:23.997Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/1c/2577d3b2380dc3e5451432a96de730ce4fdf4b602f63b9b989d0373f9ed4/gt4py-1.0.9.tar.gz", hash = "sha256:8b7d1eab14b1d093d1db943de8d8a759e9b979464892533d31c9ff9d6abc53ca", size = 724634, upload-time = "2025-09-12T12:30:50.244Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ae/6d/5b1119125f3a76f06160a745ff1e714440fa8dcc1123e637977b51f944d1/gt4py-1.0.8-py3-none-any.whl", hash = "sha256:4ad705ddbbaee8aed1a7a38748314ee6223b93ef705721c29bd9730eec3ac6c5", size = 924753, upload-time = "2025-09-04T07:22:22.23Z" }, + { url = "https://files.pythonhosted.org/packages/52/bc/e49d6dfc6169ea10dc10ed723b281aa841d7c644297c95a427455317638a/gt4py-1.0.9-py3-none-any.whl", hash = "sha256:1ef45657dd470e77bbe0f5cc9bf3c17493efc0df498ee74897069b2cbf6ac9cb", size = 925459, upload-time = "2025-09-12T12:30:48.731Z" }, ] [package.optional-dependencies] @@ -1805,7 +1805,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1822,7 +1822,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1839,7 +1839,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1856,7 +1856,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1873,7 +1873,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1943,7 +1943,7 @@ requires-dist = [ { name = "dace", marker = "extra == 'dace'", git = "https://github.com/GridTools/dace?tag=__gt4py-next-integration_2025_08_28" }, { name = "datashader", marker = "extra == 'io'", specifier = ">=0.16.1" }, { name = "ghex", marker = "extra == 'distributed'", specifier = ">=0.3.0" }, - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, { name = "gt4py", extras = ["next"], marker = "extra == 'dace'" }, @@ -1982,7 +1982,7 @@ dependencies = [ requires-dist = [ { name = "click", specifier = ">=8.0.1" }, { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -2010,7 +2010,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "filelock", specifier = ">=3.18.0" }, - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, @@ -2056,7 +2056,7 @@ requires-dist = [ { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" }, { name = "cupy-cuda12x", marker = "extra == 'cuda12'", specifier = ">=13.0" }, { name = "fprettify", specifier = ">=0.3.7" }, - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", specifier = "==1.0.9" }, { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, { name = "icon4py-atmosphere-advection", editable = "model/atmosphere/advection" }, @@ -4776,7 +4776,7 @@ version = "3.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "packaging" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "tomli", marker = "python_full_version < '3.11' or (extra == 'extra-7-icon4py-cuda11' and extra == 'extra-7-icon4py-cuda12')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5c/9b/941647e9e3616b5da7bbc4601ed9920f44a886704100fa8151406c07c149/versioningit-3.1.2.tar.gz", hash = "sha256:4db83ed99f56b07d83940bee3445ca46ca120d13b6b304cdb5fb44e5aa4edec0", size = 213047, upload-time = "2024-07-20T12:41:07.927Z" } wheels = [ From 510dfae1637fc384d23c6bbc130f54628765c82e Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 12 Sep 2025 16:32:05 +0200 Subject: [PATCH 016/108] delete outdated top level README.md from tools (#870) delete top level README.md from tools --- tools/README.md | 126 ------------------------------------------------ 1 file changed, 126 deletions(-) delete mode 100644 tools/README.md diff --git a/tools/README.md b/tools/README.md deleted file mode 100644 index 08efacdffd..0000000000 --- a/tools/README.md +++ /dev/null @@ -1,126 +0,0 @@ -# ICON4Py Tools - -## Description - -Tools and utilities for integrating icon4py code into the ICON model. - -## Installation instructions - -Until development reaches a stable state, we recommend that you follow the general instructions in the [../README.md](../README.md) root folder to install `icon4py.tools` and all of its dependencies in a virtual environment. - -## Command-line tools - -# py2fgen - -`py2fgen` is a command-line interface (CLI) tool designed to generate C and Fortran 90 (F90) wrappers, as well as a C library, for embedding a Python module into C and Fortran applications. This tool facilitates the embedding of Python code into Fortran programs by utilizing the [`CFFI`](https://cffi.readthedocs.io/en/latest/embedding) library. `CFFI` instantiates a Python interpreter to execute Python code which is "frozen" into the dynamic library generated by `CFFI`. - -**Note:** `py2fgen` has been used to embed the diffusion and dycore granule into ICON. It is important to remember that there are performance implications related to converting Fortran pointers to array-like objects that can be used in Python. It is also important to note that functions embedded into Fortran can only accept arguments with intrinsic types, as well as arrays. It is currently not possible to pass derived types to embedded Python functions. - -## Usage - -`py2fgen` simplifies the process of embedding Python functions into C and Fortran codebases. Here's how to use it: - -```bash -py2fgen [OPTIONS] MODULE_IMPORT_PATH FUNCTION_NAME - -Arguments: - MODULE_IMPORT_PATH The Python module import path to the module where the functions to be embedded are defined. - FUNCTIONS A comma-separated list of functions to be embedded in the case of multiple, otherwise just the function name. - PLUGIN_NAME The name of the plugin used for creating the shared library and bindings. -Options: - -o, --output-path PATH Specify the directory for generated code and - compiled libraries. - -b, --backend [CPU|GPU|ROUNDTRIP] - Set the backend to use, thereby unpacking - Fortran pointers into NumPy or CuPy arrays - respectively. - -d, --debug-mode Enable debug mode to log additional Python - runtime information. - -p, --profile Profile granule runtime and unpacking - Fortran pointers into NumPy or CuPy arrays. - --limited-area Enable limited area mode. -``` - -## Initialising the grid - -When embedding granules it may be necessary to have access to the representation of the ICON grid inside the granule. In order to initialise the grid for each granule there exists a `grid_init_` function which must also be embedded and called from Fortran. Each granule has access to a module state which is defined in a dictionary at the top of the module, for example `diffusion_wrapper_state` in the diffusion wrapper. - -## Environment variables - -In order to run the embedded code from Fortran it is necessary to set environment variables based on whether you are running in a CPU or GPU context. For more information on these, as well as other information on how to build and integrate embedded code into ICON using py2fgen see [this document](https://hackmd.io/OmmNptDRTe2lex7GXuYDIQ#Python-Granule-Integration-into-ICON). - -### Example - -To create a Fortran interface along with the dynamic library for a Python function named `square` within the module `example.functions`, execute: - -```bash -py2fgen example.functions square -``` - -It is also possible to generate bindings for more than one function at a time by using a comma-separated list of function names: - -```bash -py2fgen example.functions square,square2 -``` - -`py2fgen` can accept two types of functions: - -- **Simple Function:** Any Python function can be exposed. -- **GT4Py Program:** Specifically, a Python function decorated with a `@program` decorator. - -**Important:** All arguments in the exposed functions must use GT4Py style type hints. These are used by the parser to map GT4Py types to C and Fortran types in the generated bindings. - -## Generated Files - -Running `py2fgen` generates five key files: - -- **.c File**: Contains the generated CFFI code and the frozen Python code. -- **.so File**: The compiled dynamic C library containing the CFFI code. -- **.h File**: Declares the function signature of your exposed function. -- **.f90 File**: Contains a Fortran interface to the C function in the dynamic library. -- **.o File**: Represents the object code of the CFFI plugin. -- (Optional) **.py File**: Contains the Python code frozen into the dynamic library (available with `--debug-mode`). - -## Running from Fortran - -To use the generated CFFI plugin in a Fortran program, call the subroutine defined in the `.f90` interface file. Ensure that any arrays passed to the subroutine are in column-major order. - -Examples can be found under `tools/tests/py2fgen/fortran_samples`. - -## Compilation - -Compiling your Fortran driver code requires a Fortran compiler, such as `gfortran` or `nvfortran`. Follow these steps: - -1. Compile and link the Fortran driver code along with the Fortran interface and dynamic library: - -```bash -gfortran -I. -Wl,-rpath=. -L. _plugin.f90 .f90 -l_plugin -o -``` - -Replace ``, ``, and `` with the appropriate names for your project. - -**Note:** When executing the compiled binary make sure that you have sourced a Python virtual environment where all required dependencies to run the embedded Python code are present. - -## Error handling - -All generated Python wrapper code is wrapped in a `try: ... except: ...` block, with the wrapper function returning an error code `1` if an Exception ocurred -and `0` otherwise. In case of an exception the error message is written to the `py2f_cffi.log` file, which is located in the same directory as the generated bindings. -This means that on the Fortran side we can handle errors gracefully as follows: - -```Fortran -integer(c_int) :: rc -real(c_double), dimension(:, :), allocatable :: input, result - -call square(input, result, rc) - -! handle the Python error here - if (rc /= 0) then - print *, "Python failed with exit code = ", rc - call exit(1) - end if -``` - -## Other requirements - -- Embedded Python functions must have type hints for all function parameters, as these are used to derive the corresponding C and Fortran types. -- Embedded Python functions are assumed to modify function parameters in-place. Explicitly returning anything is currently not supported. From 38242936566bd72deae1fa933837dea50ba8f558 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 15 Sep 2025 08:14:06 +0200 Subject: [PATCH 017/108] Bump actions/setup-python from 5 to 6 (#863) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps [actions/setup-python](https://github.com/actions/setup-python) from 5 to 6.
Release notes

Sourced from actions/setup-python's releases.

v6.0.0

What's Changed

Breaking Changes

Make sure your runner is on version v2.327.1 or later to ensure compatibility with this release. See Release Notes

Enhancements:

Bug fixes:

Dependency updates:

New Contributors

Full Changelog: https://github.com/actions/setup-python/compare/v5...v6.0.0

v5.6.0

What's Changed

Full Changelog: https://github.com/actions/setup-python/compare/v5...v5.6.0

v5.5.0

What's Changed

Enhancements:

Bug fixes:

... (truncated)

Commits
  • e797f83 Upgrade to node 24 (#1164)
  • 3d1e2d2 Revert "Enhance cache-dependency-path handling to support files outside the w...
  • 65b0712 Clarify pythonLocation behavior for PyPy and GraalPy in environment variables...
  • 5b668cf Bump actions/checkout from 4 to 5 (#1181)
  • f62a0e2 Change missing cache directory error to warning (#1182)
  • 9322b3c Upgrade setuptools to 78.1.1 to fix path traversal vulnerability in PackageIn...
  • fbeb884 Bump form-data to fix critical vulnerabilities #182 & #183 (#1163)
  • 03bb615 Bump idna from 2.9 to 3.7 in /tests/data (#843)
  • 36da51d Add version parsing from Pipfile (#1067)
  • 3c6f142 update documentation (#1156)
  • Additional commits viewable in compare view

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=actions/setup-python&package-manager=github_actions&previous-version=5&new-version=6)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/icon4py-qa.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 2b8dbf9c84..0574c7a9b2 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -21,7 +21,7 @@ jobs: sudo apt-get install libboost-all-dev - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version-file: ".python-version" From 1c899eba7b4c51fcfb24e1677366b29330665e82 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 15 Sep 2025 10:55:40 +0200 Subject: [PATCH 018/108] Clean up 41_to_60 for better performance (#784) - Removed `vertical_mass_flux_at_cells_on_half_levels`, `tridiagonal_alpha_coeff_at_cells_on_half_levels`, `tridiagonal_beta_coeff_at_cells_on_model_levels`, `exner_explicit_term` and `rho_explicit_term` from intermediate fields passed around - Combined `_vertically_implicit_solver_at_corrector/predictor_step_before/after_solving_w` - Improved initialization of scans (level 0 and 80) (together with @philip-paul-mueller ) - Removed any unnecessary initializations related to them (i.e. `next_w`) - Handle related unit tests - Make sure that the values for the `n_lev` vertical level are passed from `_set_surface_boundary_condtion_for_computation_of_w` to `_vertically_implicit_solver_at_corrector/predictor_step_before_solving_w` properly - Improves DaCe handling of temporaries --------- Co-authored-by: "Ong Chia Rui" --- .../model/atmosphere/dycore/solve_nonhydro.py | 47 -- ...diagonal_matrix_for_w_back_substitution.py | 2 +- ..._tridiagonal_matrix_for_w_forward_sweep.py | 52 +- .../vertically_implicit_dycore_solver.py | 673 +++++++----------- .../integration_tests/test_solve_nonhydro.py | 107 --- ...diagonal_matrix_for_w_back_substitution.py | 5 +- ..._tridiagonal_matrix_for_w_forward_sweep.py | 4 + ...mplicit_dycore_solver_at_corrector_step.py | 42 +- ...mplicit_dycore_solver_at_predictor_step.py | 44 +- 9 files changed, 311 insertions(+), 665 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py index 8bb3e3c547..f9cfe2054b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py @@ -82,26 +82,6 @@ class IntermediateFields: """ Declared as z_gradh_exner in ICON. """ - tridiagonal_alpha_coeff_at_cells_on_half_levels: fa.CellKField[ - ta.vpfloat - ] # TODO(): change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? - """ - Declared as z_alpha in ICON. - """ - tridiagonal_beta_coeff_at_cells_on_model_levels: fa.CellKField[ta.vpfloat] - """ - Declared as z_beta in ICON. - """ - exner_explicit_term: fa.CellKField[ta.wpfloat] - """ - Declared as z_exner_expl in ICON. - """ - vertical_mass_flux_at_cells_on_half_levels: fa.EdgeKField[ - ta.wpfloat - ] # TODO(): change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? - """ - Declared as z_contr_w_fl_l in ICON. - """ rho_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat] """ Declared as z_rho_e in ICON. @@ -122,10 +102,6 @@ class IntermediateFields: """ Declared as z_graddiv_vn in ICON. """ - rho_explicit_term: fa.CellKField[ta.wpfloat] - """ - Declared as z_rho_expl in ICON. - """ dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat] """ Declared as z_dwdz_dd in ICON. @@ -141,18 +117,6 @@ def allocate( horizontal_pressure_gradient=data_alloc.zero_field( grid, dims.EdgeDim, dims.KDim, backend=backend ), - tridiagonal_alpha_coeff_at_cells_on_half_levels=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, backend=backend - ), - tridiagonal_beta_coeff_at_cells_on_model_levels=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, backend=backend - ), - exner_explicit_term=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, backend=backend - ), - vertical_mass_flux_at_cells_on_half_levels=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, backend=backend - ), rho_at_edges_on_model_levels=data_alloc.zero_field( grid, dims.EdgeDim, dims.KDim, backend=backend ), @@ -162,7 +126,6 @@ def allocate( horizontal_gradient_of_normal_wind_divergence=data_alloc.zero_field( grid, dims.EdgeDim, dims.KDim, backend=backend ), - rho_explicit_term=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, backend=backend), dwdz_at_cells_on_model_levels=data_alloc.zero_field( grid, dims.CellDim, dims.KDim, backend=backend ), @@ -1229,12 +1192,7 @@ def run_predictor_step( self._vertically_implicit_solver_at_predictor_step( contravariant_correction_at_cells_on_half_levels=diagnostic_state_nh.contravariant_correction_at_cells_on_half_levels, - vertical_mass_flux_at_cells_on_half_levels=z_fields.vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=z_fields.tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=z_fields.tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=prognostic_states.next.w, - rho_explicit_term=z_fields.rho_explicit_term, - exner_explicit_term=z_fields.exner_explicit_term, next_rho=prognostic_states.next.rho, next_exner=prognostic_states.next.exner, next_theta_v=prognostic_states.next.theta_v, @@ -1413,12 +1371,7 @@ def run_corrector_step( ) self._vertically_implicit_solver_at_corrector_step( - vertical_mass_flux_at_cells_on_half_levels=z_fields.vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=z_fields.tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=z_fields.tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=prognostic_states.next.w, - rho_explicit_term=z_fields.rho_explicit_term, - exner_explicit_term=z_fields.exner_explicit_term, next_rho=prognostic_states.next.rho, next_exner=prognostic_states.next.exner, next_theta_v=prognostic_states.next.theta_v, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_back_substitution.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_back_substitution.py index a00925162e..4b370372f3 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_back_substitution.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_back_substitution.py @@ -12,7 +12,7 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat -@gtx.scan_operator(axis=dims.KDim, forward=False, init=wpfloat("1.0")) +@gtx.scan_operator(axis=dims.KDim, forward=False, init=wpfloat("0.0")) def _solve_tridiagonal_matrix_for_w_back_substitution_scan( w_state: wpfloat, z_q: vpfloat, w: wpfloat ) -> wpfloat: diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py index 0649d3f156..c5a20860e5 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -13,23 +13,35 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat -@gtx.scan_operator(axis=dims.KDim, forward=True, init=(vpfloat("1.0"), 0.0, True)) -def _w( - state: tuple[vpfloat, float, bool], - w_prev: wpfloat, # only accessed at the first k-level - z_q_prev: vpfloat, - z_a: vpfloat, - z_b: vpfloat, - z_c: vpfloat, - w_prep: wpfloat, +@gtx.scan_operator( + axis=dims.KDim, + forward=True, + init=( + vpfloat("0.0"), + 0.0, + ), # boundary condition for upper tridiagonal element and w at model top +) +def tridiagonal_forward_sweep_for_w( + state_kminus1: tuple[vpfloat, float], + a: vpfloat, + b: vpfloat, + c: vpfloat, + d: wpfloat, ): - first = state[2] - z_q_m1 = z_q_prev if first else astype(state[0], vpfloat) - w_m1 = w_prev if first else state[1] - z_g = vpfloat("1.0") / (z_b + z_a * z_q_m1) - z_q_new = (vpfloat("0.0") - z_c) * z_g - w_new = (w_prep - astype(z_a, wpfloat) * w_m1) * astype(z_g, wpfloat) - return z_q_new, w_new, False + """ + | 1 0 | | w_0 | | 0 | | 1 0 | | w_0 | | 0 | + | a_1 b_1 c_1 | | w_1 | | d_1 | | 0 1 cnew_1 | | w_1 | | dnew_1 | + | a_2 b_2 c_2 | | w_2 | = | d_2 | ==> | 0 1 cnew_2 | | w_2 | = | dnew_2 | + | a_3 b_3 c_3 | | w_3 | | d_3 | | 0 1 cnew_3 | | w_3 | | dnew_3 | + | a_4 b_4 c_4 | | w_4 | | d_4 | | 0 1 cnew_4 | | w_4 | | dnew_4 | + | ... | | ... | | ... | | ... | | ... | | ... | + """ + c_kminus1 = astype(state_kminus1[0], vpfloat) + d_kminus1 = state_kminus1[1] + normalization = vpfloat("1.0") / (b + a * c_kminus1) # normalize diagonal element to 1 + c_new = (vpfloat("0.0") - c) * normalization + d_new = (d - astype(a, wpfloat) * d_kminus1) * astype(normalization, wpfloat) + return c_new, d_new @gtx.field_operator @@ -41,8 +53,6 @@ def _solve_tridiagonal_matrix_for_w_forward_sweep( z_beta: fa.CellKField[vpfloat], z_w_expl: fa.CellKField[wpfloat], z_exner_expl: fa.CellKField[wpfloat], - z_q: fa.CellKField[vpfloat], - w: fa.CellKField[wpfloat], dtime: wpfloat, cpd: wpfloat, ) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat]]: @@ -55,9 +65,7 @@ def _solve_tridiagonal_matrix_for_w_forward_sweep( z_b = vpfloat("1.0") + z_gamma_vp * z_alpha * (z_beta(Koff[-1]) + z_beta) z_gamma_wp = astype(z_gamma_vp, wpfloat) w_prep = z_w_expl - z_gamma_wp * (z_exner_expl(Koff[-1]) - z_exner_expl) - w_prev = w(Koff[-1]) - z_q_prev = z_q(Koff[-1]) - z_q_res, w_res, _ = _w(w_prev, z_q_prev, z_a, z_b, z_c, w_prep) + z_q_res, w_res = tridiagonal_forward_sweep_for_w(z_a, z_b, z_c, w_prep) return z_q_res, w_res @@ -87,8 +95,6 @@ def solve_tridiagonal_matrix_for_w_forward_sweep( z_beta, z_w_expl, z_exner_expl, - z_q, - w, dtime, cpd, out=(z_q, w), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py index dc83325f75..462bb4c220 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py @@ -84,19 +84,8 @@ def _interpolate_contravariant_correction_from_edges_on_model_levels_to_cells_on @gtx.field_operator def _set_surface_boundary_condition_for_computation_of_w( contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], -) -> tuple[fa.CellKField[ta.vpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]: - tridiagonal_alpha_coeff_at_cells_on_half_levels = broadcast( - vpfloat("0.0"), (dims.CellDim, dims.KDim) - ) - vertical_mass_flux_at_cells_on_half_levels = broadcast( - wpfloat("0.0"), (dims.CellDim, dims.KDim) - ) - - return ( - tridiagonal_alpha_coeff_at_cells_on_half_levels, - astype(contravariant_correction_at_cells_on_half_levels, wpfloat), - vertical_mass_flux_at_cells_on_half_levels, - ) +) -> fa.CellKField[ta.wpfloat]: + return astype(contravariant_correction_at_cells_on_half_levels, wpfloat) @gtx.field_operator @@ -184,9 +173,53 @@ def _compute_solver_coefficients_matrix( @gtx.field_operator -def _vertically_implicit_solver_at_predictor_step_before_solving_w( - vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - next_w: fa.CellKField[ta.wpfloat], +def solve_w( + last_inner_level: gtx.int32, + next_w: fa.CellKField[wpfloat], + vwind_impl_wgt: fa.CellField[wpfloat], + theta_v_ic: fa.CellKField[wpfloat], + ddqz_z_half: fa.CellKField[vpfloat], + z_alpha: fa.CellKField[vpfloat], + z_beta: fa.CellKField[vpfloat], + z_w_expl: fa.CellKField[wpfloat], + z_exner_expl: fa.CellKField[wpfloat], + dtime: wpfloat, + cpd: wpfloat, +) -> fa.CellKField[wpfloat]: + ( + tridiagonal_intermediate_result, + next_w_intermediate_result, + ) = concat_where( + dims.KDim > 0, + _solve_tridiagonal_matrix_for_w_forward_sweep( + vwind_impl_wgt=vwind_impl_wgt, + theta_v_ic=theta_v_ic, + ddqz_z_half=ddqz_z_half, + z_alpha=z_alpha, + z_beta=z_beta, + z_w_expl=z_w_expl, + z_exner_expl=z_exner_expl, + dtime=dtime, + cpd=cpd, + ), + (broadcast(vpfloat("0.0"), (dims.CellDim,)), broadcast(wpfloat("0.0"), (dims.CellDim,))), + ) + next_w = concat_where( + dims.KDim < last_inner_level, + _solve_tridiagonal_matrix_for_w_back_substitution_scan( + z_q=tridiagonal_intermediate_result, + w=next_w_intermediate_result, + ), + next_w, + ) + return next_w + + +@gtx.field_operator +def _vertically_implicit_solver_at_predictor_step( + next_w: fa.CellKField[ + ta.wpfloat + ], # necessary input because the last vertical level is set outside this field operator geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], @@ -207,9 +240,18 @@ def _vertically_implicit_solver_at_predictor_step_before_solving_w( rho_iau_increment: fa.CellKField[ta.vpfloat], exner_iau_increment: fa.CellKField[ta.vpfloat], ddqz_z_half: fa.CellKField[ta.vpfloat], + exner_dynamical_increment: fa.CellKField[ta.wpfloat], + dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + rayleigh_damping_factor: fa.KField[ta.wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], iau_wgt_dyn: ta.wpfloat, dtime: ta.wpfloat, + rayleigh_type: gtx.int32, + divdamp_type: gtx.int32, is_iau_active: bool, + at_first_substep: bool, + end_index_of_damping_layer: gtx.int32, + kstart_moist: gtx.int32, n_lev: gtx.int32, ) -> tuple[ fa.CellKField[ta.wpfloat], @@ -225,8 +267,6 @@ def _vertically_implicit_solver_at_predictor_step_before_solving_w( z_theta_v_fl_e=theta_v_flux_at_edges_on_model_levels, ) - tridiagonal_intermediate_result = broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)) - w_explicit_term = concat_where( 1 <= dims.KDim, _compute_w_explicit_term_with_predictor_advective_tendency( @@ -238,24 +278,14 @@ def _vertically_implicit_solver_at_predictor_step_before_solving_w( broadcast(wpfloat("0.0"), (dims.CellDim, dims.KDim)), ) - (next_w, vertical_mass_flux_at_cells_on_half_levels) = concat_where( - dims.KDim == 0, - ( - broadcast(wpfloat("0.0"), (dims.CellDim,)), - broadcast(wpfloat("0.0"), (dims.CellDim,)), - ), - (next_w, vertical_mass_flux_at_cells_on_half_levels), - ) - vertical_mass_flux_at_cells_on_half_levels = concat_where( - # TODO(OngChia): (dims.KDim < n_lev) is needed. Otherwise, the stencil test fails. (1 <= dims.KDim) & (dims.KDim < n_lev), rho_at_cells_on_half_levels * ( -astype(contravariant_correction_at_cells_on_half_levels, wpfloat) + exner_w_explicit_weight_parameter * current_w ), - vertical_mass_flux_at_cells_on_half_levels, + broadcast(wpfloat("0.0"), (dims.CellDim,)), ) ( @@ -271,6 +301,11 @@ def _vertically_implicit_solver_at_predictor_step_before_solving_w( rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, dtime=dtime, ) + tridiagonal_alpha_coeff_at_cells_on_half_levels = concat_where( + dims.KDim < n_lev, + tridiagonal_alpha_coeff_at_cells_on_half_levels, + broadcast(vpfloat("0.0"), (dims.CellDim,)), + ) (rho_explicit_term, exner_explicit_term) = _compute_explicit_part_for_rho_and_exner( rho_nnow=current_rho, @@ -294,82 +329,20 @@ def _vertically_implicit_solver_at_predictor_step_before_solving_w( iau_wgt_dyn=iau_wgt_dyn, ) - tridiagonal_intermediate_result, next_w = concat_where( - dims.KDim > 0, - _solve_tridiagonal_matrix_for_w_forward_sweep( - vwind_impl_wgt=exner_w_implicit_weight_parameter, - theta_v_ic=theta_v_at_cells_on_half_levels, - ddqz_z_half=ddqz_z_half, - z_alpha=tridiagonal_alpha_coeff_at_cells_on_half_levels, - z_beta=tridiagonal_beta_coeff_at_cells_on_model_levels, - z_w_expl=w_explicit_term, - z_exner_expl=exner_explicit_term, - z_q=tridiagonal_intermediate_result, - w=next_w, - dtime=dtime, - cpd=dycore_consts.cpd, - ), - (tridiagonal_intermediate_result, next_w), - ) - - # TODO(OngChia): We should not need this because alpha is zero at n_lev and thus tridiagonal_intermediate_result should be zero at nlev-1. However, stencil test shows it is nonzero. - tridiagonal_intermediate_result = concat_where( - dims.KDim == n_lev - 1, - broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)), - tridiagonal_intermediate_result, - ) - - next_w = concat_where( - dims.KDim > 0, - _solve_tridiagonal_matrix_for_w_back_substitution_scan( - z_q=tridiagonal_intermediate_result, w=next_w - ), - next_w, - ) - - return ( - vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w, - rho_explicit_term, - exner_explicit_term, + next_w = solve_w( + last_inner_level=n_lev, + next_w=next_w, # n_lev value is set by _set_surface_boundary_condtion_for_computation_of_w + vwind_impl_wgt=exner_w_implicit_weight_parameter, + theta_v_ic=theta_v_at_cells_on_half_levels, + ddqz_z_half=ddqz_z_half, + z_alpha=tridiagonal_alpha_coeff_at_cells_on_half_levels, + z_beta=tridiagonal_beta_coeff_at_cells_on_model_levels, + z_w_expl=w_explicit_term, + z_exner_expl=exner_explicit_term, + dtime=dtime, + cpd=dycore_consts.cpd, ) - -@gtx.field_operator -def _vertically_implicit_solver_at_predictor_step_after_solving_w( - tridiagonal_beta_coeff_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - tridiagonal_alpha_coeff_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - next_w: fa.CellKField[ta.wpfloat], - dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - exner_dynamical_increment: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - rho_explicit_term: fa.CellKField[ta.wpfloat], - exner_explicit_term: fa.CellKField[ta.wpfloat], - dtime: ta.wpfloat, - rayleigh_type: gtx.int32, - divdamp_type: gtx.int32, - at_first_substep: bool, - end_index_of_damping_layer: gtx.int32, - kstart_moist: gtx.int32, -) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], -]: - # Because we do not support nesting, it is safe to assume w_1 is a zero field w_1 = broadcast(wpfloat("0.0"), (dims.CellDim,)) if rayleigh_type == rayleigh_damping_options.KLEMP: next_w = concat_where( @@ -428,10 +401,137 @@ def _vertically_implicit_solver_at_predictor_step_after_solving_w( ) -@gtx.field_operator -def _vertically_implicit_solver_at_corrector_step_before_solving_w( - vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], +@gtx.program +def vertically_implicit_solver_at_predictor_step( + contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], next_w: fa.CellKField[ta.wpfloat], + next_rho: fa.CellKField[ta.wpfloat], + next_exner: fa.CellKField[ta.wpfloat], + next_theta_v: fa.CellKField[ta.wpfloat], + dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + exner_dynamical_increment: fa.CellKField[ta.vpfloat], + geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], + theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], + predictor_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], + pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], + rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], + current_exner: fa.CellKField[ta.wpfloat], + current_rho: fa.CellKField[ta.wpfloat], + current_theta_v: fa.CellKField[ta.wpfloat], + current_w: fa.CellKField[ta.wpfloat], + inv_ddqz_z_full: fa.CellKField[ta.vpfloat], + exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], + exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], + rho_iau_increment: fa.CellKField[ta.vpfloat], + exner_iau_increment: fa.CellKField[ta.vpfloat], + ddqz_z_half: fa.CellKField[ta.vpfloat], + rayleigh_damping_factor: fa.KField[ta.wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], + wgtfac_c: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + iau_wgt_dyn: ta.wpfloat, + dtime: ta.wpfloat, + is_iau_active: bool, + rayleigh_type: gtx.int32, + divdamp_type: gtx.int32, + at_first_substep: bool, + end_index_of_damping_layer: gtx.int32, + kstart_moist: gtx.int32, + flat_level_index_plus1: gtx.int32, + start_cell_index_nudging: gtx.int32, + end_cell_index_local: gtx.int32, + start_cell_index_lateral_lvl3: gtx.int32, + end_cell_index_halo_lvl1: gtx.int32, + vertical_start_index_model_top: gtx.int32, + vertical_end_index_model_surface: gtx.int32, +): + _interpolate_contravariant_correction_from_edges_on_model_levels_to_cells_on_half_levels( + contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, + e_bln_c_s=e_bln_c_s, + wgtfac_c=wgtfac_c, + wgtfacq_c=wgtfacq_c, + nlev=vertical_end_index_model_surface - 1, + out=contravariant_correction_at_cells_on_half_levels, + domain={ + dims.CellDim: ( + start_cell_index_lateral_lvl3, + end_cell_index_halo_lvl1, + ), + dims.KDim: (flat_level_index_plus1, vertical_end_index_model_surface), + }, + ) + _set_surface_boundary_condition_for_computation_of_w( + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + out=next_w, + domain={ + dims.CellDim: (start_cell_index_nudging, end_cell_index_local), + dims.KDim: (vertical_end_index_model_surface - 1, vertical_end_index_model_surface), + }, + ) + _vertically_implicit_solver_at_predictor_step( + next_w=next_w, + geofac_div=geofac_div, + mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, + theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, + predictor_vertical_wind_advective_tendency=predictor_vertical_wind_advective_tendency, + pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, + rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + exner_w_explicit_weight_parameter=exner_w_explicit_weight_parameter, + current_exner=current_exner, + current_rho=current_rho, + current_theta_v=current_theta_v, + current_w=current_w, + inv_ddqz_z_full=inv_ddqz_z_full, + exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, + theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, + perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, + exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, + rho_iau_increment=rho_iau_increment, + exner_iau_increment=exner_iau_increment, + ddqz_z_half=ddqz_z_half, + dwdz_at_cells_on_model_levels=dwdz_at_cells_on_model_levels, + exner_dynamical_increment=exner_dynamical_increment, + rayleigh_damping_factor=rayleigh_damping_factor, + reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, + iau_wgt_dyn=iau_wgt_dyn, + dtime=dtime, + rayleigh_type=rayleigh_type, + divdamp_type=divdamp_type, + is_iau_active=is_iau_active, + at_first_substep=at_first_substep, + end_index_of_damping_layer=end_index_of_damping_layer, + kstart_moist=kstart_moist, + n_lev=vertical_end_index_model_surface - 1, + out=( + next_w, + next_rho, + next_exner, + next_theta_v, + dwdz_at_cells_on_model_levels, + exner_dynamical_increment, + ), + domain={ + dims.CellDim: (start_cell_index_nudging, end_cell_index_local), + dims.KDim: (vertical_start_index_model_top, vertical_end_index_model_surface - 1), + }, + ) + + +@gtx.field_operator +def _vertically_implicit_solver_at_corrector_step( + next_w: fa.CellKField[ + ta.wpfloat + ], # necessary input because the last vertical level is set outside this field operator + dynamical_vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + dynamical_vertical_volumetric_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + exner_dynamical_increment: fa.CellKField[ta.wpfloat], geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], @@ -453,11 +553,21 @@ def _vertically_implicit_solver_at_corrector_step_before_solving_w( rho_iau_increment: fa.CellKField[ta.vpfloat], exner_iau_increment: fa.CellKField[ta.vpfloat], ddqz_z_half: fa.CellKField[ta.vpfloat], + rayleigh_damping_factor: fa.KField[ta.wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], advection_explicit_weight_parameter: ta.wpfloat, advection_implicit_weight_parameter: ta.wpfloat, + lprep_adv: bool, + r_nsubsteps: ta.wpfloat, + ndyn_substeps_var: ta.wpfloat, iau_wgt_dyn: ta.wpfloat, dtime: ta.wpfloat, is_iau_active: bool, + rayleigh_type: gtx.int32, + at_first_substep: bool, + at_last_substep: bool, + end_index_of_damping_layer: gtx.int32, + kstart_moist: gtx.int32, n_lev: gtx.int32, ) -> tuple[ fa.CellKField[ta.wpfloat], @@ -466,16 +576,13 @@ def _vertically_implicit_solver_at_corrector_step_before_solving_w( fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], ]: - # verified for e-9 divergence_of_mass, divergence_of_theta_v = _compute_divergence_of_fluxes_of_rho_and_theta( geofac_div=geofac_div, mass_fl_e=mass_flux_at_edges_on_model_levels, z_theta_v_fl_e=theta_v_flux_at_edges_on_model_levels, ) - - tridiagonal_intermediate_result = broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)) - w_explicit_term = concat_where( 1 <= dims.KDim, _compute_w_explicit_term_with_interpolated_predictor_corrector_advective_tendency( @@ -489,27 +596,15 @@ def _vertically_implicit_solver_at_corrector_step_before_solving_w( ), broadcast(wpfloat("0.0"), (dims.CellDim, dims.KDim)), ) - - (next_w, vertical_mass_flux_at_cells_on_half_levels) = concat_where( - dims.KDim == 0, - ( - broadcast(wpfloat("0.0"), (dims.CellDim,)), - broadcast(wpfloat("0.0"), (dims.CellDim,)), - ), - (next_w, vertical_mass_flux_at_cells_on_half_levels), - ) - vertical_mass_flux_at_cells_on_half_levels = concat_where( - # TODO(OngChia): (dims.KDim < n_lev) is needed. Otherwise, the stencil test fails. (1 <= dims.KDim) & (dims.KDim < n_lev), rho_at_cells_on_half_levels * ( -astype(contravariant_correction_at_cells_on_half_levels, wpfloat) + exner_w_explicit_weight_parameter * current_w ), - vertical_mass_flux_at_cells_on_half_levels, + broadcast(wpfloat("0.0"), (dims.CellDim,)), ) - ( tridiagonal_beta_coeff_at_cells_on_model_levels, tridiagonal_alpha_coeff_at_cells_on_half_levels, @@ -523,7 +618,11 @@ def _vertically_implicit_solver_at_corrector_step_before_solving_w( rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, dtime=dtime, ) - + tridiagonal_alpha_coeff_at_cells_on_half_levels = concat_where( + dims.KDim < n_lev, + tridiagonal_alpha_coeff_at_cells_on_half_levels, + broadcast(vpfloat("0.0"), (dims.CellDim,)), + ) (rho_explicit_term, exner_explicit_term) = _compute_explicit_part_for_rho_and_exner( rho_nnow=current_rho, inv_ddqz_z_full=inv_ddqz_z_full, @@ -536,7 +635,6 @@ def _vertically_implicit_solver_at_corrector_step_before_solving_w( ddt_exner_phy=exner_tendency_due_to_slow_physics, dtime=dtime, ) - if is_iau_active: rho_explicit_term, exner_explicit_term = _add_analysis_increments_from_data_assimilation( z_rho_expl=rho_explicit_term, @@ -546,90 +644,22 @@ def _vertically_implicit_solver_at_corrector_step_before_solving_w( iau_wgt_dyn=iau_wgt_dyn, ) - tridiagonal_intermediate_result, next_w = concat_where( - dims.KDim > 0, - _solve_tridiagonal_matrix_for_w_forward_sweep( - vwind_impl_wgt=exner_w_implicit_weight_parameter, - theta_v_ic=theta_v_at_cells_on_half_levels, - ddqz_z_half=ddqz_z_half, - z_alpha=tridiagonal_alpha_coeff_at_cells_on_half_levels, - z_beta=tridiagonal_beta_coeff_at_cells_on_model_levels, - z_w_expl=w_explicit_term, - z_exner_expl=exner_explicit_term, - z_q=tridiagonal_intermediate_result, - w=next_w, - dtime=dtime, - cpd=dycore_consts.cpd, - ), - (tridiagonal_intermediate_result, next_w), - ) - - # TODO(OngChia): We should not need this because alpha is zero at n_lev and thus tridiagonal_intermediate_result should be zero at nlev-1. However, stencil test shows it is nonzero. - tridiagonal_intermediate_result = concat_where( - dims.KDim == n_lev - 1, - broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)), - tridiagonal_intermediate_result, - ) - - next_w = concat_where( - dims.KDim > 0, - _solve_tridiagonal_matrix_for_w_back_substitution_scan( - z_q=tridiagonal_intermediate_result, - w=next_w, - ), - next_w, - ) - - return ( - vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w, - rho_explicit_term, - exner_explicit_term, + next_w = solve_w( + last_inner_level=n_lev, + next_w=next_w, # n_lev value is set by _set_surface_boundary_condtion_for_computation_of_w + vwind_impl_wgt=exner_w_implicit_weight_parameter, + theta_v_ic=theta_v_at_cells_on_half_levels, + ddqz_z_half=ddqz_z_half, + z_alpha=tridiagonal_alpha_coeff_at_cells_on_half_levels, + z_beta=tridiagonal_beta_coeff_at_cells_on_model_levels, + z_w_expl=w_explicit_term, + z_exner_expl=exner_explicit_term, + dtime=dtime, + cpd=dycore_consts.cpd, ) - -@gtx.field_operator -def _vertically_implicit_solver_at_corrector_step_after_solving_w( - vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - tridiagonal_beta_coeff_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - tridiagonal_alpha_coeff_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - next_w: fa.CellKField[ta.wpfloat], - dynamical_vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - dynamical_vertical_volumetric_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - exner_dynamical_increment: fa.CellKField[ta.wpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - rho_explicit_term: fa.CellKField[ta.wpfloat], - exner_explicit_term: fa.CellKField[ta.wpfloat], - lprep_adv: bool, - r_nsubsteps: ta.wpfloat, - ndyn_substeps_var: ta.wpfloat, - dtime: ta.wpfloat, - rayleigh_type: gtx.int32, - end_index_of_damping_layer: gtx.int32, - kstart_moist: gtx.int32, - at_first_substep: bool, - at_last_substep: bool, -) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], -]: - # Because we do not support nesting, it is safe to assume w_1 is a zero field w_1 = broadcast(wpfloat("0.0"), (dims.CellDim,)) + if rayleigh_type == rayleigh_damping_options.KLEMP: next_w = concat_where( (dims.KDim > 0) & (dims.KDim < end_index_of_damping_layer + 1), @@ -687,8 +717,8 @@ def _vertically_implicit_solver_at_corrector_step_after_solving_w( ), ) - exner_dynamical_increment = ( - concat_where( + if at_last_substep: + exner_dynamical_increment = concat_where( dims.KDim >= kstart_moist, _update_dynamical_exner_time_increment( exner=next_exner, @@ -699,9 +729,6 @@ def _vertically_implicit_solver_at_corrector_step_after_solving_w( ), exner_dynamical_increment, ) - if at_last_substep - else exner_dynamical_increment - ) return ( next_w, @@ -714,175 +741,9 @@ def _vertically_implicit_solver_at_corrector_step_after_solving_w( ) -@gtx.program -def vertically_implicit_solver_at_predictor_step( - contravariant_correction_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - tridiagonal_beta_coeff_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - tridiagonal_alpha_coeff_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - next_w: fa.CellKField[ta.wpfloat], - rho_explicit_term: fa.CellKField[ta.wpfloat], - exner_explicit_term: fa.CellKField[ta.wpfloat], - next_rho: fa.CellKField[ta.wpfloat], - next_exner: fa.CellKField[ta.wpfloat], - next_theta_v: fa.CellKField[ta.wpfloat], - dwdz_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - exner_dynamical_increment: fa.CellKField[ta.vpfloat], - geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - mass_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - theta_v_flux_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat], - predictor_vertical_wind_advective_tendency: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat], - current_exner: fa.CellKField[ta.wpfloat], - current_rho: fa.CellKField[ta.wpfloat], - current_theta_v: fa.CellKField[ta.wpfloat], - current_w: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - exner_w_implicit_weight_parameter: fa.CellField[ta.wpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], - exner_tendency_due_to_slow_physics: fa.CellKField[ta.vpfloat], - rho_iau_increment: fa.CellKField[ta.vpfloat], - exner_iau_increment: fa.CellKField[ta.vpfloat], - ddqz_z_half: fa.CellKField[ta.vpfloat], - rayleigh_damping_factor: fa.KField[ta.wpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], - wgtfac_c: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - iau_wgt_dyn: ta.wpfloat, - dtime: ta.wpfloat, - is_iau_active: bool, - rayleigh_type: gtx.int32, - divdamp_type: gtx.int32, - at_first_substep: bool, - end_index_of_damping_layer: gtx.int32, - kstart_moist: gtx.int32, - flat_level_index_plus1: gtx.int32, - start_cell_index_nudging: gtx.int32, - end_cell_index_local: gtx.int32, - start_cell_index_lateral_lvl3: gtx.int32, - end_cell_index_halo_lvl1: gtx.int32, - vertical_start_index_model_top: gtx.int32, - vertical_end_index_model_surface: gtx.int32, -): - _interpolate_contravariant_correction_from_edges_on_model_levels_to_cells_on_half_levels( - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - wgtfacq_c=wgtfacq_c, - nlev=vertical_end_index_model_surface - 1, - out=contravariant_correction_at_cells_on_half_levels, - domain={ - dims.CellDim: ( - start_cell_index_lateral_lvl3, - end_cell_index_halo_lvl1, - ), - dims.KDim: (flat_level_index_plus1, vertical_end_index_model_surface), - }, - ) - _set_surface_boundary_condition_for_computation_of_w( - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - out=( - tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w, - vertical_mass_flux_at_cells_on_half_levels, - ), - domain={ - dims.CellDim: (start_cell_index_nudging, end_cell_index_local), - dims.KDim: (vertical_end_index_model_surface - 1, vertical_end_index_model_surface), - }, - ) - - _vertically_implicit_solver_at_predictor_step_before_solving_w( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - next_w=next_w, - geofac_div=geofac_div, - mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, - theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, - predictor_vertical_wind_advective_tendency=predictor_vertical_wind_advective_tendency, - pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, - rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - exner_w_explicit_weight_parameter=exner_w_explicit_weight_parameter, - current_exner=current_exner, - current_rho=current_rho, - current_theta_v=current_theta_v, - current_w=current_w, - inv_ddqz_z_full=inv_ddqz_z_full, - exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, - theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, - perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, - exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, - rho_iau_increment=rho_iau_increment, - exner_iau_increment=exner_iau_increment, - ddqz_z_half=ddqz_z_half, - iau_wgt_dyn=iau_wgt_dyn, - dtime=dtime, - is_iau_active=is_iau_active, - n_lev=vertical_end_index_model_surface - 1, - out=( - vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w, - rho_explicit_term, - exner_explicit_term, - ), - domain={ - dims.CellDim: (start_cell_index_nudging, end_cell_index_local), - dims.KDim: (vertical_start_index_model_top, vertical_end_index_model_surface - 1), - }, - ) - _vertically_implicit_solver_at_predictor_step_after_solving_w( - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w=next_w, - dwdz_at_cells_on_model_levels=dwdz_at_cells_on_model_levels, - exner_dynamical_increment=exner_dynamical_increment, - rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - current_exner=current_exner, - current_rho=current_rho, - current_theta_v=current_theta_v, - inv_ddqz_z_full=inv_ddqz_z_full, - exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, - rayleigh_damping_factor=rayleigh_damping_factor, - reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, - dtime=dtime, - rayleigh_type=rayleigh_type, - divdamp_type=divdamp_type, - at_first_substep=at_first_substep, - end_index_of_damping_layer=end_index_of_damping_layer, - kstart_moist=kstart_moist, - out=( - next_w, - next_rho, - next_exner, - next_theta_v, - dwdz_at_cells_on_model_levels, - exner_dynamical_increment, - ), - domain={ - dims.CellDim: (start_cell_index_nudging, end_cell_index_local), - dims.KDim: (vertical_start_index_model_top, vertical_end_index_model_surface - 1), - }, - ) - - @gtx.program def vertically_implicit_solver_at_corrector_step( - vertical_mass_flux_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - tridiagonal_beta_coeff_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - tridiagonal_alpha_coeff_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], next_w: fa.CellKField[ta.wpfloat], - rho_explicit_term: fa.CellKField[ta.wpfloat], - exner_explicit_term: fa.CellKField[ta.wpfloat], next_rho: fa.CellKField[ta.wpfloat], next_exner: fa.CellKField[ta.wpfloat], next_theta_v: fa.CellKField[ta.wpfloat], @@ -932,19 +793,17 @@ def vertically_implicit_solver_at_corrector_step( ): _set_surface_boundary_condition_for_computation_of_w( contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - out=( - tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w, - vertical_mass_flux_at_cells_on_half_levels, - ), + out=next_w, domain={ dims.CellDim: (start_cell_index_nudging, end_cell_index_local), dims.KDim: (vertical_end_index_model_surface - 1, vertical_end_index_model_surface), }, ) - _vertically_implicit_solver_at_corrector_step_before_solving_w( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, + _vertically_implicit_solver_at_corrector_step( next_w=next_w, + dynamical_vertical_mass_flux_at_cells_on_half_levels=dynamical_vertical_mass_flux_at_cells_on_half_levels, + dynamical_vertical_volumetric_flux_at_cells_on_half_levels=dynamical_vertical_volumetric_flux_at_cells_on_half_levels, + exner_dynamical_increment=exner_dynamical_increment, geofac_div=geofac_div, mass_flux_at_edges_on_model_levels=mass_flux_at_edges_on_model_levels, theta_v_flux_at_edges_on_model_levels=theta_v_flux_at_edges_on_model_levels, @@ -966,54 +825,22 @@ def vertically_implicit_solver_at_corrector_step( rho_iau_increment=rho_iau_increment, exner_iau_increment=exner_iau_increment, ddqz_z_half=ddqz_z_half, - advection_explicit_weight_parameter=advection_explicit_weight_parameter, - advection_implicit_weight_parameter=advection_implicit_weight_parameter, - iau_wgt_dyn=iau_wgt_dyn, - dtime=dtime, - is_iau_active=is_iau_active, - n_lev=vertical_end_index_model_surface - 1, - out=( - vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w, - rho_explicit_term, - exner_explicit_term, - ), - domain={ - dims.CellDim: (start_cell_index_nudging, end_cell_index_local), - dims.KDim: (vertical_start_index_model_top, vertical_end_index_model_surface - 1), - }, - ) - - _vertically_implicit_solver_at_corrector_step_after_solving_w( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, - next_w=next_w, - dynamical_vertical_mass_flux_at_cells_on_half_levels=dynamical_vertical_mass_flux_at_cells_on_half_levels, - dynamical_vertical_volumetric_flux_at_cells_on_half_levels=dynamical_vertical_volumetric_flux_at_cells_on_half_levels, - exner_dynamical_increment=exner_dynamical_increment, - rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, - current_exner=current_exner, - current_rho=current_rho, - current_theta_v=current_theta_v, - inv_ddqz_z_full=inv_ddqz_z_full, - exner_w_implicit_weight_parameter=exner_w_implicit_weight_parameter, - exner_tendency_due_to_slow_physics=exner_tendency_due_to_slow_physics, rayleigh_damping_factor=rayleigh_damping_factor, reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, + advection_explicit_weight_parameter=advection_explicit_weight_parameter, + advection_implicit_weight_parameter=advection_implicit_weight_parameter, lprep_adv=lprep_adv, r_nsubsteps=r_nsubsteps, ndyn_substeps_var=ndyn_substeps_var, + iau_wgt_dyn=iau_wgt_dyn, dtime=dtime, + is_iau_active=is_iau_active, rayleigh_type=rayleigh_type, - end_index_of_damping_layer=end_index_of_damping_layer, - kstart_moist=kstart_moist, at_first_substep=at_first_substep, at_last_substep=at_last_substep, + end_index_of_damping_layer=end_index_of_damping_layer, + kstart_moist=kstart_moist, + n_lev=vertical_end_index_model_surface - 1, out=( next_w, next_rho, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 61d538a893..de864464f7 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -237,7 +237,6 @@ def test_nonhydro_predictor_step( cell_start_lateral_boundary_level_2 = icon_grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) ) - cell_start_nudging = icon_grid.start_index(cell_domain(h_grid.Zone.NUDGING)) edge_start_lateral_boundary_level_5 = icon_grid.start_index( edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) @@ -450,45 +449,6 @@ def test_nonhydro_predictor_step( atol=1e-15, ) - # stencils 43, 46, 47 - assert test_utils.dallclose( - solve_nonhydro.intermediate_fields.vertical_mass_flux_at_cells_on_half_levels.asnumpy()[ - cell_start_nudging:, : - ], - sp_exit.z_contr_w_fl_l().asnumpy()[cell_start_nudging:, :], - atol=2e-15, - ) - - # stencil 44, 45 - assert test_utils.dallclose( - solve_nonhydro.intermediate_fields.tridiagonal_alpha_coeff_at_cells_on_half_levels.asnumpy()[ - cell_start_nudging:, : - ], - sp_exit.z_alpha().asnumpy()[cell_start_nudging:, :], - atol=5e-13, - ) - # stencil 44 - assert test_utils.dallclose( - solve_nonhydro.intermediate_fields.tridiagonal_beta_coeff_at_cells_on_model_levels.asnumpy()[ - cell_start_nudging:, : - ], - sp_exit.z_beta().asnumpy()[cell_start_nudging:, :], - atol=2e-15, - ) - - # stencil 48, 49 - assert test_utils.dallclose( - solve_nonhydro.intermediate_fields.rho_explicit_term.asnumpy()[cell_start_nudging:, :], - sp_exit.z_rho_expl().asnumpy()[cell_start_nudging:, :], - atol=2e-15, - ) - # stencil 48, 49 - assert test_utils.dallclose( - solve_nonhydro.intermediate_fields.exner_explicit_term.asnumpy()[cell_start_nudging:, :], - sp_exit.z_exner_expl().asnumpy()[cell_start_nudging:, :], - atol=2e-15, - ) - # end assert test_utils.dallclose(prognostic_state_nnew.rho.asnumpy(), sp_exit.rho_new().asnumpy()) assert test_utils.dallclose( @@ -574,14 +534,9 @@ def test_nonhydro_corrector_step( z_fields = solve_nh.IntermediateFields( horizontal_pressure_gradient=init_savepoint.z_gradh_exner(), - tridiagonal_alpha_coeff_at_cells_on_half_levels=init_savepoint.z_alpha(), - tridiagonal_beta_coeff_at_cells_on_model_levels=init_savepoint.z_beta(), - exner_explicit_term=init_savepoint.z_exner_expl(), - vertical_mass_flux_at_cells_on_half_levels=init_savepoint.z_contr_w_fl_l(), rho_at_edges_on_model_levels=init_savepoint.z_rho_e(), theta_v_at_edges_on_model_levels=init_savepoint.z_theta_v_e(), horizontal_gradient_of_normal_wind_divergence=init_savepoint.z_graddiv_vn(), - rho_explicit_term=init_savepoint.z_rho_expl(), dwdz_at_cells_on_model_levels=init_savepoint.z_dwdz_dd(), horizontal_kinetic_energy_at_edges_on_model_levels=init_savepoint.z_kin_hor_e(), tangential_wind_on_half_levels=init_savepoint.z_vt_ie(), @@ -2114,19 +2069,14 @@ def test_vertically_implicit_solver_at_predictor_step( theta_v_flux_at_edges_on_model_levels = sp_stencil_init.z_theta_v_fl_e() predictor_vertical_wind_advective_tendency = sp_stencil_init.ddt_w_adv_pc(0) pressure_buoyancy_acceleration_at_cells_on_half_levels = sp_stencil_init.z_th_ddz_exner_c() - vertical_mass_flux_at_cells_on_half_levels = sp_stencil_init.z_contr_w_fl_l() rho_at_cells_on_half_levels = sp_stencil_init.rho_ic() contravariant_correction_at_cells_on_half_levels = savepoint_nonhydro_init.w_concorr_c() current_exner = sp_stencil_init.exner_nnow() current_rho = sp_stencil_init.rho_nnow() current_theta_v = sp_stencil_init.theta_v_nnow() current_w = sp_stencil_init.w() - tridiagonal_alpha_coeff_at_cells_on_half_levels = sp_stencil_init.z_alpha() - tridiagonal_beta_coeff_at_cells_on_model_levels = sp_stencil_init.z_beta() theta_v_at_cells_on_half_levels = sp_stencil_init.theta_v_ic() next_w = sp_stencil_init.w() - rho_explicit_term = sp_stencil_init.z_rho_expl() - exner_explicit_term = sp_stencil_init.z_exner_expl() perturbed_exner_at_cells_on_model_levels = sp_stencil_init.exner_pr() exner_tendency_due_to_slow_physics = sp_stencil_init.ddt_exner_phy() rho_iau_increment = sp_stencil_init.rho_incr() @@ -2143,12 +2093,7 @@ def test_vertically_implicit_solver_at_predictor_step( divdamp_type = config.divdamp_type w_concorr_c_ref = sp_nh_exit.w_concorr_c() - z_contr_w_fl_l_ref = sp_nh_exit.z_contr_w_fl_l() - z_beta_ref = sp_nh_exit.z_beta() - z_alpha_ref = sp_nh_exit.z_alpha() w_ref = sp_nh_exit.w_new() - z_rho_expl_ref = sp_nh_exit.z_rho_expl() - z_exner_expl_ref = sp_nh_exit.z_exner_expl() rho_ref = sp_nh_exit.rho_new() exner_ref = sp_nh_exit.exner_new() theta_v_ref = sp_nh_exit.theta_v_new() @@ -2174,12 +2119,7 @@ def test_vertically_implicit_solver_at_predictor_step( backend )( contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=next_w, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, next_rho=next_rho, next_exner=next_exner, next_theta_v=next_theta_v, @@ -2233,27 +2173,12 @@ def test_vertically_implicit_solver_at_predictor_step( w_concorr_c_ref.asnumpy(), atol=1e-15, ) - assert test_utils.dallclose( - vertical_mass_flux_at_cells_on_half_levels.asnumpy(), - z_contr_w_fl_l_ref.asnumpy(), - atol=1e-12, - ) - assert test_utils.dallclose( - tridiagonal_beta_coeff_at_cells_on_model_levels.asnumpy(), z_beta_ref.asnumpy() - ) - assert test_utils.dallclose( - tridiagonal_alpha_coeff_at_cells_on_half_levels.asnumpy(), z_alpha_ref.asnumpy() - ) assert test_utils.dallclose( next_w.asnumpy()[start_cell_nudging:, :], w_ref.asnumpy()[start_cell_nudging:, :], rtol=1e-7, atol=1e-12, ) - assert test_utils.dallclose(rho_explicit_term.asnumpy(), z_rho_expl_ref.asnumpy()) - assert test_utils.dallclose( - exner_explicit_term.asnumpy(), z_exner_expl_ref.asnumpy(), rtol=1.0e-10, atol=1.0e-12 - ) assert test_utils.dallclose( next_rho.asnumpy()[start_cell_nudging:, :], rho_ref.asnumpy()[start_cell_nudging:, :] ) @@ -2348,19 +2273,14 @@ def test_vertically_implicit_solver_at_corrector_step( predictor_vertical_wind_advective_tendency = sp_stencil_init.ddt_w_adv_pc(0) corrector_vertical_wind_advective_tendency = sp_stencil_init.ddt_w_adv_pc(1) pressure_buoyancy_acceleration_at_cells_on_half_levels = sp_stencil_init.z_th_ddz_exner_c() - vertical_mass_flux_at_cells_on_half_levels = sp_stencil_init.z_contr_w_fl_l() rho_at_cells_on_half_levels = sp_stencil_init.rho_ic() contravariant_correction_at_cells_on_half_levels = sp_stencil_init.w_concorr_c() current_exner = sp_stencil_init.exner_nnow() current_rho = sp_stencil_init.rho_nnow() current_theta_v = sp_stencil_init.theta_v_nnow() current_w = sp_stencil_init.w() - tridiagonal_alpha_coeff_at_cells_on_half_levels = sp_stencil_init.z_alpha() - tridiagonal_beta_coeff_at_cells_on_model_levels = sp_stencil_init.z_beta() theta_v_at_cells_on_half_levels = sp_stencil_init.theta_v_ic() next_w = sp_stencil_init.w() - rho_explicit_term = sp_stencil_init.z_rho_expl() - exner_explicit_term = sp_stencil_init.z_exner_expl() perturbed_exner_at_cells_on_model_levels = sp_stencil_init.exner_pr() exner_tendency_due_to_slow_physics = sp_stencil_init.ddt_exner_phy() rho_iau_increment = sp_stencil_init.rho_incr() @@ -2380,12 +2300,7 @@ def test_vertically_implicit_solver_at_corrector_step( iau_wgt_dyn = config.iau_wgt_dyn is_iau_active = config.is_iau_active - z_contr_w_fl_l_ref = sp_nh_exit.z_contr_w_fl_l() - z_beta_ref = sp_nh_exit.z_beta() - z_alpha_ref = sp_nh_exit.z_alpha() w_ref = sp_nh_exit.w_new() - z_rho_expl_ref = sp_nh_exit.z_rho_expl() - z_exner_expl_ref = sp_nh_exit.z_exner_expl() rho_ref = sp_nh_exit.rho_new() exner_ref = sp_nh_exit.exner_new() theta_v_ref = sp_nh_exit.theta_v_new() @@ -2407,12 +2322,7 @@ def test_vertically_implicit_solver_at_corrector_step( vertically_implicit_dycore_solver.vertically_implicit_solver_at_corrector_step.with_backend( backend )( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=next_w, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, next_rho=next_rho, next_exner=next_exner, next_theta_v=next_theta_v, @@ -2462,29 +2372,12 @@ def test_vertically_implicit_solver_at_corrector_step( offset_provider=offset_provider, ) - assert test_utils.dallclose( - vertical_mass_flux_at_cells_on_half_levels.asnumpy(), - z_contr_w_fl_l_ref.asnumpy(), - atol=1e-12, - ) - assert test_utils.dallclose( - tridiagonal_beta_coeff_at_cells_on_model_levels.asnumpy(), z_beta_ref.asnumpy() - ) - assert test_utils.dallclose( - tridiagonal_alpha_coeff_at_cells_on_half_levels.asnumpy(), z_alpha_ref.asnumpy() - ) assert test_utils.dallclose( next_w.asnumpy()[start_cell_nudging:, :], w_ref.asnumpy()[start_cell_nudging:, :], rtol=1e-10, atol=1e-12, ) - assert test_utils.dallclose(rho_explicit_term.asnumpy(), z_rho_expl_ref.asnumpy()) - assert test_utils.dallclose( - exner_explicit_term.asnumpy(), - z_exner_expl_ref.asnumpy(), - rtol=3e-9, - ) assert test_utils.dallclose( next_rho.asnumpy()[start_cell_nudging:, :], rho_ref.asnumpy()[start_cell_nudging:, :] ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py index a72986b164..076e84b767 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_back_substitution.py @@ -27,10 +27,11 @@ def solve_tridiagonal_matrix_for_w_back_substitution_numpy( z_q: np.ndarray, w: np.ndarray, ) -> np.ndarray: - w_new = np.zeros_like(w) + rng = np.random.default_rng() + w_new = rng.random(w.shape, dtype=w.dtype) last_k_level = w.shape[1] - 1 - w_new[:, last_k_level] = w[:, last_k_level] + z_q[:, last_k_level] + w_new[:, last_k_level] = w[:, last_k_level] for k in reversed(range(1, last_k_level)): w_new[:, k] = w[:, k] + w_new[:, k + 1] * z_q[:, k] w_new[:, 0] = w[:, 0] diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py index d0759c37a6..7bd38f9619 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_solve_tridiagonal_matrix_for_w_forward_sweep.py @@ -106,7 +106,11 @@ def input_data(self, grid: base_grid.Grid) -> dict[str, gtx.Field | state_utils. grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, dtype=ta.wpfloat ) z_q = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + # z_q first level should always be initialized to zero when solve_tridiagonal_matrix_for_w_forward_sweep is called + z_q.asnumpy()[:, 0] = 0.0 w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat) + # w first level should always be initialized to zero when solve_tridiagonal_matrix_for_w_forward_sweep is called + w.asnumpy()[:, 0] = 0.0 h_start = 0 h_end = gtx.int32(grid.num_cells) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index abb6dce355..c43c2de7be 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -54,12 +54,7 @@ class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_corrector_step OUTPUTS = ( - "vertical_mass_flux_at_cells_on_half_levels", - "tridiagonal_beta_coeff_at_cells_on_model_levels", - "tridiagonal_alpha_coeff_at_cells_on_half_levels", "next_w", - "rho_explicit_term", - "exner_explicit_term", "next_rho", "next_exner", "next_theta_v", @@ -71,12 +66,7 @@ class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): @staticmethod def reference( connectivities: dict[gtx.Dimension, np.ndarray], - vertical_mass_flux_at_cells_on_half_levels: np.ndarray, - tridiagonal_beta_coeff_at_cells_on_model_levels: np.ndarray, - tridiagonal_alpha_coeff_at_cells_on_half_levels: np.ndarray, next_w: np.ndarray, - rho_explicit_term: np.ndarray, - exner_explicit_term: np.ndarray, next_rho: np.ndarray, next_exner: np.ndarray, next_theta_v: np.ndarray, @@ -128,6 +118,17 @@ def reference( horz_idx = horz_idx[:, np.newaxis] vert_idx = np.arange(exner_dynamical_increment.shape[1]) + rng = np.random.default_rng() + vertical_mass_flux_at_cells_on_half_levels = rng.random((horizontal_end, vert_idx.size + 1)) + tridiagonal_beta_coeff_at_cells_on_model_levels = rng.random( + (horizontal_end, vert_idx.size) + ) + tridiagonal_alpha_coeff_at_cells_on_half_levels = rng.random( + (horizontal_end, vert_idx.size + 1) + ) + rho_explicit_term = rng.random(next_rho.shape) + exner_explicit_term = rng.random(next_exner.shape) + divergence_of_mass = np.zeros_like(current_rho) divergence_of_theta_v = np.zeros_like(current_theta_v) divergence_of_mass, divergence_of_theta_v = np.where( @@ -360,12 +361,7 @@ def reference( ) return dict( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=next_w, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, next_rho=next_rho, next_exner=next_exner, next_theta_v=next_theta_v, @@ -418,18 +414,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala grid, dims.CellDim, dims.KDim, low=1.0e-5 ) - vertical_mass_flux_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) - tridiagonal_beta_coeff_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim - ) - tridiagonal_alpha_coeff_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - rho_explicit_term = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - exner_explicit_term = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) @@ -461,12 +446,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) return dict( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=next_w, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, next_rho=next_rho, next_exner=next_exner, next_theta_v=next_theta_v, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index 9e18d03cef..b4866a0f75 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -57,12 +57,7 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_predictor_step OUTPUTS = ( - "vertical_mass_flux_at_cells_on_half_levels", - "tridiagonal_beta_coeff_at_cells_on_model_levels", - "tridiagonal_alpha_coeff_at_cells_on_half_levels", "next_w", - "rho_explicit_term", - "exner_explicit_term", "next_rho", "next_exner", "next_theta_v", @@ -74,12 +69,7 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): def reference( connectivities: dict[gtx.Dimension, np.ndarray], contravariant_correction_at_cells_on_half_levels: np.ndarray, - vertical_mass_flux_at_cells_on_half_levels: np.ndarray, - tridiagonal_beta_coeff_at_cells_on_model_levels: np.ndarray, - tridiagonal_alpha_coeff_at_cells_on_half_levels: np.ndarray, next_w: np.ndarray, - rho_explicit_term: np.ndarray, - exner_explicit_term: np.ndarray, next_rho: np.ndarray, next_exner: np.ndarray, next_theta_v: np.ndarray, @@ -130,6 +120,19 @@ def reference( horz_idx = horz_idx[:, np.newaxis] vert_idx = np.arange(exner_dynamical_increment.shape[1]) + rng = np.random.default_rng() + vertical_mass_flux_at_cells_on_half_levels = rng.random( + (end_cell_index_local, vert_idx.size + 1) + ) + tridiagonal_beta_coeff_at_cells_on_model_levels = rng.random( + (end_cell_index_local, vert_idx.size) + ) + tridiagonal_alpha_coeff_at_cells_on_half_levels = rng.random( + (end_cell_index_local, vert_idx.size + 1) + ) + rho_explicit_term = rng.random(next_rho.shape) + exner_explicit_term = rng.random(next_exner.shape) + contravariant_correction_at_cells_on_half_levels[:, :-1] = np.where( (start_cell_index_lateral_lvl3 <= horz_idx) & (horz_idx < end_cell_index_halo_lvl1) @@ -355,12 +358,7 @@ def reference( ) return dict( - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=next_w, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, next_rho=next_rho, next_exner=next_exner, next_theta_v=next_theta_v, @@ -381,9 +379,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala pressure_buoyancy_acceleration_at_cells_on_half_levels = data_alloc.random_field( grid, dims.CellDim, dims.KDim ) - vertical_mass_flux_at_cells_on_half_levels = data_alloc.random_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) rho_at_cells_on_half_levels = data_alloc.random_field( grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, low=1.0e-5 ) @@ -420,15 +415,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala wgtfac_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) wgtfacq_c = data_alloc.random_field(grid, dims.CellDim, dims.KDim, low=1.0e-5, high=0.99999) - tridiagonal_beta_coeff_at_cells_on_model_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim - ) - tridiagonal_alpha_coeff_at_cells_on_half_levels = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} - ) next_w = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}) - rho_explicit_term = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) - exner_explicit_term = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) next_rho = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) next_exner = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) next_theta_v = data_alloc.constant_field(grid, 1.0e-5, dims.CellDim, dims.KDim) @@ -455,12 +442,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala return dict( contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - vertical_mass_flux_at_cells_on_half_levels=vertical_mass_flux_at_cells_on_half_levels, - tridiagonal_beta_coeff_at_cells_on_model_levels=tridiagonal_beta_coeff_at_cells_on_model_levels, - tridiagonal_alpha_coeff_at_cells_on_half_levels=tridiagonal_alpha_coeff_at_cells_on_half_levels, next_w=next_w, - rho_explicit_term=rho_explicit_term, - exner_explicit_term=exner_explicit_term, next_rho=next_rho, next_exner=next_exner, next_theta_v=next_theta_v, From 937128bb76fd27f119e5eaffedfefd38816aac95 Mon Sep 17 00:00:00 2001 From: edopao Date: Mon, 15 Sep 2025 11:17:39 +0200 Subject: [PATCH 019/108] Fix memory allocator for zero fields in velocity advection tests (#874) Fix an issue in velocity advection tests, that caused the dace programs to fail: a cupy array was expected as argument, a numpy array was passed. --- .../dycore/integration_tests/test_velocity_advection.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py index 693a31ed19..1d928e9cbc 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py @@ -628,7 +628,9 @@ def test_compute_contravariant_correction_and_advection_in_vertical_momentum_equ horizontal_advection_of_w_at_edges_on_half_levels = savepoint_velocity_exit.z_v_grad_w() vertical_wind_advective_tendency = savepoint_velocity_init.ddt_w_adv_pc(istep_init - 1) contravariant_corrected_w_at_cells_on_model_levels = savepoint_velocity_init.z_w_con_c_full() - vertical_cfl = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + vertical_cfl = data_alloc.zero_field( + icon_grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, backend=backend + ) skip_compute_predictor_vertical_advection = savepoint_velocity_init.lvn_only() coeff1_dwdz = metrics_savepoint.coeff1_dwdz() @@ -780,7 +782,9 @@ def test_compute_advection_in_vertical_momentum_equation( vn_on_half_levels = savepoint_velocity_exit.vn_ie() vertical_wind_advective_tendency = savepoint_velocity_init.ddt_w_adv_pc(istep_init - 1) contravariant_corrected_w_at_cells_on_model_levels = savepoint_velocity_init.z_w_con_c_full() - vertical_cfl = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat) + vertical_cfl = data_alloc.zero_field( + icon_grid, dims.CellDim, dims.KDim, dtype=ta.vpfloat, backend=backend + ) coeff1_dwdz = metrics_savepoint.coeff1_dwdz() coeff2_dwdz = metrics_savepoint.coeff2_dwdz() From c1bfd1c3fe076c628515e53a70d5a952430fbd0c Mon Sep 17 00:00:00 2001 From: edopao Date: Mon, 15 Sep 2025 20:35:33 +0200 Subject: [PATCH 020/108] Add dace backend to default CI pipeline (#822) This PR enables CI tests on the dace backend. - In the Github Actions, we run the stencil tests on the CPU backend. - In the CSCS CI, we run the tests with serialized data on the GPU backend, only for diffusion and dycore in order to save compute resources. --- .github/workflows/icon4py-test-model.yml | 2 +- ci/dace.yml | 12 ++++-------- ci/default.yml | 10 ++++++++-- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/.github/workflows/icon4py-test-model.yml b/.github/workflows/icon4py-test-model.yml index c4abf820cc..8bd9696242 100644 --- a/.github/workflows/icon4py-test-model.yml +++ b/.github/workflows/icon4py-test-model.yml @@ -14,7 +14,7 @@ jobs: strategy: matrix: python-version: ["3.10", "3.11"] - backend: ["embedded", "gtfn_cpu"] # TODO(): add dace-cpu? + backend: ["embedded", "dace_cpu", "gtfn_cpu"] component: ["advection", "diffusion", "dycore", "microphysics", "muphys", "driver", "common"] steps: - name: Checkout diff --git a/ci/dace.yml b/ci/dace.yml index 27db0f0679..052f2b4ef0 100644 --- a/ci/dace.yml +++ b/ci/dace.yml @@ -3,8 +3,6 @@ include: .test_model_stencils: stage: test - variables: - SLURM_TIMELIMIT: '00:10:00' script: - nox -s "test_model-3.10(stencils, $COMPONENT)" -- --backend=$BACKEND --grid=$GRID rules: @@ -12,6 +10,8 @@ include: variables: SLURM_TIMELIMIT: '00:30:00' - when: on_success + variables: + SLURM_TIMELIMIT: '00:15:00' parallel: matrix: - COMPONENT: [advection, diffusion, dycore, microphysics, common, driver] @@ -25,14 +25,10 @@ test_model_stencils_aarch64: .test_model_datatests: stage: test variables: - SLURM_TIMELIMIT: '00:30:00' + NUM_PROCESSES: 8 + SLURM_TIMELIMIT: '00:45:00' script: - nox -s "test_model-3.10(datatest, $COMPONENT)" -- --backend=$BACKEND --level=$LEVEL - rules: - - if: $COMPONENT == 'common' && $BACKEND == 'dace_gpu' && $LEVEL == 'integration' - variables: - SLURM_TIMELIMIT: '01:30:00' - - when: on_success parallel: matrix: - COMPONENT: [advection, diffusion, dycore, microphysics, common, driver] diff --git a/ci/default.yml b/ci/default.yml index 9d40e07ffc..5fe2b10db8 100644 --- a/ci/default.yml +++ b/ci/default.yml @@ -37,10 +37,16 @@ test_tools_datatests_aarch64: script: - nox -s "test_model-3.10(datatest, $COMPONENT)" -- --backend=$BACKEND --level=$LEVEL rules: + - if: $BACKEND == 'dace_gpu' && $COMPONENT != 'diffusion' && $COMPONENT != 'dycore' + when: never # run only in daily CI, to save compute resources - if: $COMPONENT == 'common' && $LEVEL == 'integration' variables: - NUM_PROCESSES: 2 + NUM_PROCESSES: 1 SLURM_TIMELIMIT: '00:30:00' + - if: $BACKEND == 'dace_gpu' + variables: + NUM_PROCESSES: 8 + SLURM_TIMELIMIT: '00:45:00' - if: $BACKEND == 'embedded' variables: SLURM_TIMELIMIT: '00:15:00' @@ -50,7 +56,7 @@ test_tools_datatests_aarch64: parallel: matrix: - COMPONENT: [advection, diffusion, dycore, microphysics, common, driver] - BACKEND: [embedded, gtfn_cpu, gtfn_gpu] + BACKEND: [embedded, dace_gpu, gtfn_cpu, gtfn_gpu] LEVEL: [integration] # test_model_datatests_x86_64: # extends: [.test_model_datatests, .test_template_x86_64] From bddb09df01f268420aee3f9b53f25919a9f86203 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Thu, 18 Sep 2025 16:12:03 +0200 Subject: [PATCH 021/108] add benchmarking for 'vertically_implicit_solver...' program --- ...mplicit_dycore_solver_at_predictor_step.py | 36 +++++++++++++++++-- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index b4866a0f75..6556796e06 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -53,6 +53,7 @@ ) +@pytest.mark.continuous_benchmarking @pytest.mark.uses_concat_where class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_predictor_step @@ -64,6 +65,33 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): "dwdz_at_cells_on_model_levels", "exner_dynamical_increment", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "start_cell_index_nudging", + "end_cell_index_local", + "start_cell_index_lateral_lvl3", + "end_cell_index_halo_lvl1", + "end_index_of_damping_layer", + "kstart_moist", + "flat_level_index_plus1", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "divdamp_type", + "rayleigh_type", + "is_iau_active", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "end_index_of_damping_layer", + "kstart_moist", + "flat_level_index_plus1", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "divdamp_type", + "rayleigh_type", + "is_iau_active", + ), + } @staticmethod def reference( @@ -366,8 +394,10 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture(params=[{"at_first_substep": value} for value in [True, False]]) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( @@ -423,7 +453,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) is_iau_active = True - at_first_substep = True + at_first_substep = request.param["at_first_substep"] rayleigh_type = 2 divdamp_type = 3 end_index_of_damping_layer = 3 From 48c9b08cadad54d9292bc6eac3fa6d0b0c99a924 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 10:36:02 +0200 Subject: [PATCH 022/108] avoid dace backend corner cases in stencil tests --- ...ived_horizontal_winds_and_ke_and_contravariant_correction.py | 2 +- .../test_compute_perturbed_quantities_and_interpolation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 22d70561c0..31575df11c 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -311,7 +311,7 @@ def input_data( c_intp = data_alloc.random_field(grid, dims.VertexDim, dims.V2CDim) nlev = grid.num_levels - nflatlev = 13 + nflatlev = (grid.num_levels * 3) // 10 edge_domain = h_grid.domain(dims.EdgeDim) # For the ICON grid we use the proper domain bounds (otherwise we will run into non-protected skip values) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 02a12daba5..301c2fa0b4 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -432,7 +432,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala end_cell_halo_level_2 = grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) nflatlev = 4 - nflat_gradp = 27 + nflat_gradp = grid.num_levels // 2 return dict( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, From 70e1720507bdc5deece2fa0813841fb0042b9a75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 10:49:30 +0200 Subject: [PATCH 023/108] add benchmarking for 'interpolate_rho_theta...' program --- ...ls_and_compute_temperature_vertical_gradient.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py index b3fb185c76..e9ab186490 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py @@ -34,6 +34,7 @@ from icon4py.model.testing import stencil_tests +@pytest.mark.continuous_benchmarking class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration( stencil_tests.StencilTest ): @@ -44,6 +45,19 @@ class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration "theta_v_at_cells_on_half_levels", "pressure_buoyancy_acceleration_at_cells_on_half_levels", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( From 1479e80a4e9585439ab2649b2f39c3a00c8496bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 10:50:44 +0200 Subject: [PATCH 024/108] fix typing for static variants in stencil tests --- .../tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py | 2 +- ...sion_to_w_and_compute_horizontal_gradients_for_turbulence.py | 2 +- .../test_calculate_nabla2_and_smag_coefficients_for_vn.py | 2 +- .../tests/diffusion/stencil_tests/test_calculate_nabla4.py | 2 +- .../stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index 2ca65484e3..debdae974f 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -31,7 +31,7 @@ class TestApplyDiffusionToVn(StencilTest): PROGRAM = apply_diffusion_to_vn OUTPUTS = ("vn",) STATIC_PARAMS = { - StandardStaticVariants.NONE: None, + StandardStaticVariants.NONE: (), StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index c15179a95b..bdf70f59bd 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -33,7 +33,7 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTes PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence OUTPUTS = ("w", "dwdx", "dwdy") STATIC_PARAMS = { - StandardStaticVariants.NONE: None, + StandardStaticVariants.NONE: (), StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index c6652317b3..b17c7f578e 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -23,7 +23,7 @@ class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_and_smag_coefficients_for_vn OUTPUTS = ("kh_smag_e", "kh_smag_ec", "z_nabla2_e") STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: None, + stencil_tests.StandardStaticVariants.NONE: (), stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py index 305cf3f319..dc8f9ac4ba 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py @@ -60,7 +60,7 @@ class TestCalculateNabla4(StencilTest): PROGRAM = calculate_nabla4 OUTPUTS = ("z_nabla4_e2",) STATIC_PARAMS = { - StandardStaticVariants.NONE: None, + StandardStaticVariants.NONE: (), StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index d551fc265e..2137f5e587 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -26,7 +26,7 @@ class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex OUTPUTS = ("p_u_out", "p_v_out") STATIC_PARAMS = { - StandardStaticVariants.NONE: None, + StandardStaticVariants.NONE: (), StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", From cf730e8184815261ed6525cd330d930021b67148 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 11:08:42 +0200 Subject: [PATCH 025/108] add benchmarking for 'apply_divergence_damping...' program --- ..._apply_divergence_damping_and_update_vn.py | 39 +++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index d49922254d..65e4910078 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import itertools + import gt4py.next as gtx import numpy as np import pytest @@ -24,10 +26,28 @@ divergence_damp_order = DivergenceDampingOrder() +@pytest.mark.continuous_benchmarking @pytest.mark.embedded_remap_error class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): PROGRAM = apply_divergence_damping_and_update_vn OUTPUTS = ("next_vn",) + STATIC_PARAMS = { + test_helpers.StandardStaticVariants.NONE: (), + test_helpers.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + "is_iau_active", + "limited_area", + ), + test_helpers.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "is_iau_active", + "limited_area", + ), + } @staticmethod def reference( @@ -153,7 +173,20 @@ def reference( return dict(next_vn=next_vn) - @pytest.fixture(params=[True, False]) + @pytest.fixture( + params=[ + {"limited_area": la, "divdamp_order": do, "is_iau_active": ia} + for la, do, ia in itertools.product( + [True, False], + [ + DivergenceDampingOrder.SECOND_ORDER, + DivergenceDampingOrder.FOURTH_ORDER, + DivergenceDampingOrder.COMBINED, + ], + [True, False], + ) + ], + ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_mask_for_3d_divdamp = data_alloc.random_field(grid, dims.EdgeDim) @@ -190,7 +223,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: is_iau_active = True fourth_order_divdamp_factor = 0.004 second_order_divdamp_factor = 0.012 - divdamp_order = 24 + divdamp_order = request.param["divdamp_order"] second_order_divdamp_scaling_coeff = 194588.14247428576 apply_2nd_order_divergence_damping = (divdamp_order == divergence_damp_order.COMBINED) and ( second_order_divdamp_scaling_coeff > 1.0e-6 @@ -202,7 +235,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) ) - limited_area = request.param + limited_area = request.param["limited_area"] edge_domain = h_grid.domain(dims.EdgeDim) start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) From 47d647a1dc7b1717c8e943768c8c9b5e81c35e72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 11:33:55 +0200 Subject: [PATCH 026/108] add benchmarking for 'vertically_implicit...corrector...' program --- ...mplicit_dycore_solver_at_corrector_step.py | 46 +++++++++++++++++-- 1 file changed, 41 insertions(+), 5 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index c43c2de7be..72e8daf61c 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -5,6 +5,7 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import itertools from typing import Any import gt4py.next as gtx @@ -50,6 +51,7 @@ from .test_update_mass_volume_flux import update_mass_volume_flux_numpy +@pytest.mark.continuous_benchmarking @pytest.mark.uses_concat_where class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_corrector_step @@ -62,6 +64,33 @@ class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): "dynamical_vertical_volumetric_flux_at_cells_on_half_levels", "exner_dynamical_increment", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "start_cell_index_nudging", + "end_cell_index_local", + "end_index_of_damping_layer", + "kstart_moist", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "at_first_substep", + "at_last_substep", + "lprep_adv", + "is_iau_active", + "rayleigh_type", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "end_index_of_damping_layer", + "kstart_moist", + "vertical_start_index_model_top", + "vertical_end_index_model_surface", + "at_first_substep", + "at_last_substep", + "lprep_adv", + "is_iau_active", + "rayleigh_type", + ), + } @staticmethod def reference( @@ -370,8 +399,15 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture( + params=[ + {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} + for afs, als, la in itertools.product(*([(True, False)] * 3)) + ] + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( @@ -426,13 +462,13 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - lprep_adv = True + lprep_adv = request.param["lprep_adv"] r_nsubsteps = 0.5 is_iau_active = True - at_first_substep = True + at_first_substep = request.param["at_first_substep"] rayleigh_type = 2 end_index_of_damping_layer = 3 - at_last_substep = True + at_last_substep = request.param["at_last_substep"] kstart_moist = 1 dtime = 0.001 veladv_offctr = 0.25 From 3289b3827b6d0bd7980ab131f79a66a487857312 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 12:50:31 +0200 Subject: [PATCH 027/108] add benchmarking for 'init_cell...' program --- .../test_init_cell_kdim_field_with_zero_wp.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py index 658499bc80..a0c2e0f902 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py @@ -19,12 +19,26 @@ from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils.data_allocation import zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +@pytest.mark.continuous_benchmarking class TestInitCellKdimFieldWithZeroWp(StencilTest): PROGRAM = init_cell_kdim_field_with_zero_wp OUTPUTS = ("field_with_zero_wp",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( From f3921e7a559e4c962f09e7d5f78832e0eb47d8ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rico=20H=C3=A4uselmann?= Date: Fri, 19 Sep 2025 12:52:47 +0200 Subject: [PATCH 028/108] add benchmarking for 'update_mass_flux_weighted' program --- .../test_update_mass_flux_weighted.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py index c9297d94b7..5a0e5422da 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py @@ -19,12 +19,26 @@ from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +@pytest.mark.continuous_benchmarking class TestUpdateMassFluxWeighted(StencilTest): PROGRAM = update_mass_flux_weighted OUTPUTS = ("mass_flx_ic",) + STATIC_PARAMS = { + StandardStaticVariants.NONE: (), + StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( From 9fd8a3fb6aab3ea58c1804a14c10e087311f4f2f Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 29 Sep 2025 15:11:22 +0200 Subject: [PATCH 029/108] Make fixture names readable --- .../test_apply_divergence_damping_and_update_vn.py | 1 + ...horizontal_winds_and_ke_and_contravariant_correction.py | 3 ++- ..._vertically_implicit_dycore_solver_at_corrector_step.py | 7 ++++++- 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 65e4910078..2bf0347b35 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -186,6 +186,7 @@ def reference( [True, False], ) ], + ids=lambda param: f"limited_area={param['limited_area']},divdamp_order={param['divdamp_order']},is_iau_active={param['is_iau_active']}", ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 31575df11c..f328028555 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -276,7 +276,8 @@ def reference( # TODO(ricoh): Add True case. Blocked by test failure (issue: #875) @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [False]] + params=[{"skip_compute_predictor_vertical_advection": value} for value in [False]], + ids=lambda param: f"skip_compute_predictor_vertical_advection={param['skip_compute_predictor_vertical_advection']}", ) def input_data( self, request: pytest.FixtureRequest, grid: base.Grid diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index 72e8daf61c..7b86efffa1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -403,7 +403,12 @@ def reference( params=[ {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} for afs, als, la in itertools.product(*([(True, False)] * 3)) - ] + ], + ids=lambda p: ( + f"at_first_substep={p['at_first_substep']}," + f"at_last_substep={p['at_last_substep']}," + f"lprep_adv={p['lprep_adv']}" + ), ) def input_data( self, request: pytest.FixtureRequest, grid: base.Grid From 1f88e8f74ac5c11e46c387686bdfea4994a2694f Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 29 Sep 2025 15:19:12 +0200 Subject: [PATCH 030/108] Make variable options usable with pytest -k --- .../test_apply_divergence_damping_and_update_vn.py | 2 +- ...t_vertically_implicit_dycore_solver_at_corrector_step.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 2bf0347b35..c90241697b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -186,7 +186,7 @@ def reference( [True, False], ) ], - ids=lambda param: f"limited_area={param['limited_area']},divdamp_order={param['divdamp_order']},is_iau_active={param['is_iau_active']}", + ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index 7b86efffa1..ab8ea3d50b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -405,9 +405,9 @@ def reference( for afs, als, la in itertools.product(*([(True, False)] * 3)) ], ids=lambda p: ( - f"at_first_substep={p['at_first_substep']}," - f"at_last_substep={p['at_last_substep']}," - f"lprep_adv={p['lprep_adv']}" + f"at_first_substep[{p['at_first_substep']}]__" + f"at_last_substep[{p['at_last_substep']}]__" + f"lprep_adv[{p['lprep_adv']}]" ), ) def input_data( From 123b25bee8e0abf8d953b6418df9b0b1aa6fee2b Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Tue, 30 Sep 2025 13:54:48 +0200 Subject: [PATCH 031/108] Print gt4py timers in a nice way (along with pytest timers) ------------------------------------------------------------------------------------------------------------------------------------------ benchmark: 2 tests ------------------------------------------------------------------------------------------------------------------------------------------ Name (time in ms) Min Max Mean StdDev Median IQR Outliers OPS Rounds Iterations -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- test_TestVerticallyImplicitSolverAtCorrectorStep[compile_time_domain-at_first_substep[False]__at_last_substep[False]__lprep_adv[False]] 5.0775 (1.0) 7.9169 (1.14) 5.4656 (1.0) 0.3381 (1.90) 5.4298 (1.0) 0.1788 (1.0) 5;3 182.9632 (1.0) 122 1 test_TestVerticallyImplicitSolverAtCorrectorStep[compile_time_domain-at_first_substep[False]__at_last_substep[False]__lprep_adv[True]] 5.5181 (1.09) 6.9601 (1.0) 5.8356 (1.07) 0.1780 (1.0) 5.8323 (1.07) 0.1792 (1.00) 21;2 171.3618 (0.94) 115 1 -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- Legend: Outliers: 1 Standard Deviation from Mean; 1.5 IQR (InterQuartile Range) from 1st Quartile and 3rd Quartile. OPS: Operations Per Second, computed as 1 / Mean --------------------------------------------------------------------- GT4Py Timer Report ---------------------------------------------------------------------- Benchmark Name | Mean (s) | Std Dev | Runs --------------------------------------------------------------------------------------------------------------------------------------------------------------- VerticallyImplicitSolverAtCorrectorStep[compile_time_domain-at_first_substep[False]__at_last_substep[False]__lprep_adv[True]] | 0.00073964 | 0.00102386 | 118 VerticallyImplicitSolverAtCorrectorStep[compile_time_domain-at_first_substep[False]__at_last_substep[False]__lprep_adv[False]] | 0.00061125 | 0.00007039 | 125 --------------------------------------------------------------------------------------------------------------------------------------------------------------- --- .../src/icon4py/model/testing/pytest_hooks.py | 53 +++++++++++++++++++ .../icon4py/model/testing/stencil_tests.py | 17 +++++- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index 670f7ff775..8353ec6cdc 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -161,3 +161,56 @@ def pytest_benchmark_update_json(output_json): for bench in output_json["benchmarks"]: bench["fullname"] = _name_from_fullname(bench["fullname"]) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + """ + Gather GT4Py timer metrics from benchmark fixture and add them to the test report. + """ + outcome = yield + report = outcome.get_result() + if call.when == "call": + benchmark = item.funcargs.get("benchmark", None) + if benchmark and hasattr(benchmark, "extra_info"): + info = benchmark.extra_info.get("gtx_metrics", None) + if info: + filtered_benchmark_name = benchmark.name.split("test_Test")[-1] + # Combine the benchmark name in a readable form with the gtx_metrics data + report.sections.append(("benchmark-extra", tuple([filtered_benchmark_name, info]))) + + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + """ + Add a custom section to the terminal summary with GT4Py timer metrics from benchmarks. + """ + # Gather gtx_metrics + benchmark_gtx_metrics = [] + for outcome in ("passed", "failed", "skipped"): + all_reports = terminalreporter.stats.get(outcome, []) + for report in all_reports: + for secname, info in getattr(report, "sections", []): + if secname == "benchmark-extra": + benchmark_gtx_metrics.append(info) + # Calculate the maximum length of benchmark names for formatting + max_name_len = 0 + for benchmark_name, _ in benchmark_gtx_metrics: + max_name_len = max(len(benchmark_name), max_name_len) + # Print the GT4Py timer report table + if benchmark_gtx_metrics: + terminalreporter.ensure_newline() + header = ( + f"{'Benchmark Name':<{max_name_len}} | {'Mean (s)':>10} | {'Std Dev':>10} | {'Runs':>4}" + ) + title = " GT4Py Timer Report " + sep_len = max(0, len(header) - len(title)) + left = sep_len // 2 + right = sep_len - left + terminalreporter.line("-" * left + title + "-" * right, bold=True, blue=True) + terminalreporter.line(header) + terminalreporter.line("-" * len(header), blue=True) + for benchmark_name, gtx_metrics in benchmark_gtx_metrics: + terminalreporter.line( + f"{benchmark_name:<{max_name_len}} | {gtx_metrics.mean:>10.8f} | {gtx_metrics.std:>10.8f} | {len(gtx_metrics.samples):>4}" + ) + terminalreporter.line("-" * len(header), blue=True) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index b6dd570144..ec10a5860e 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -16,7 +16,13 @@ import numpy as np import pytest from gt4py import eve -from gt4py.next import backend as gtx_backend, constructors, typing as gtx_typing +from gt4py.next import ( + backend as gtx_backend, + config as gtx_config, + constructors, + metrics as gtx_metrics, + typing as gtx_typing, +) # TODO(havogt): import will disappear after FieldOperators support `.compile` from gt4py.next.ffront.decorator import FieldOperator @@ -98,6 +104,15 @@ def test_and_benchmark( **_properly_allocated_input_data, offset_provider=grid.connectivities, ) + # Collect GT4Py runtime metrics if enabled + if gtx_config.COLLECT_METRICS_LEVEL > 0: + assert ( + len(gtx_metrics.program_metrics.data.keys()) == 1 + ), "Expected exactly one entry in gtx_metrics" + benchmark.extra_info["gtx_metrics"] = gtx_metrics.program_metrics.data[ + next(iter(gtx_metrics.program_metrics.data)) + ]["compute"] + gtx_metrics.program_metrics.clear() class StencilTest: From 66855eb5ee45d551c41ab49bef66682781c92655 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 1 Oct 2025 10:26:12 +0200 Subject: [PATCH 032/108] Disable verification if --benchmark_only is set --- .../icon4py/model/testing/stencil_tests.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index cbc1b2af36..99f20c2f12 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -82,21 +82,23 @@ def test_and_benchmark( grid: base.Grid, _properly_allocated_input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], _configured_program: Callable[..., None], + request: pytest.FixtureRequest, ) -> None: - reference_outputs = self.reference( - _ConnectivityConceptFixer( - grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) - ), - **{ - k: v.asnumpy() if isinstance(v, gtx.Field) else v - for k, v in _properly_allocated_input_data.items() - }, - ) - - _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) - self._verify_stencil_test( - input_data=_properly_allocated_input_data, reference_outputs=reference_outputs - ) + if not request.config.getoption("benchmark_only"): + reference_outputs = self.reference( + _ConnectivityConceptFixer( + grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) + ), + **{ + k: v.asnumpy() if isinstance(v, gtx.Field) else v + for k, v in _properly_allocated_input_data.items() + }, + ) + + _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) + self._verify_stencil_test( + input_data=_properly_allocated_input_data, reference_outputs=reference_outputs + ) if benchmark is not None and benchmark.enabled: benchmark( From d39c9f14ee05bef34070e7a38035def939b3e8a8 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 1 Oct 2025 10:26:42 +0200 Subject: [PATCH 033/108] Clean gtx_metrics after the verification --- model/testing/src/icon4py/model/testing/stencil_tests.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 99f20c2f12..566779ffa9 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -101,11 +101,16 @@ def test_and_benchmark( ) if benchmark is not None and benchmark.enabled: + # Clean up GT4Py metrics from previous runs + if gtx_config.COLLECT_METRICS_LEVEL > 0: + gtx_metrics.program_metrics.clear() + benchmark( _configured_program, **_properly_allocated_input_data, offset_provider=grid.connectivities, ) + # Collect GT4Py runtime metrics if enabled if gtx_config.COLLECT_METRICS_LEVEL > 0: assert ( @@ -114,7 +119,6 @@ def test_and_benchmark( benchmark.extra_info["gtx_metrics"] = gtx_metrics.program_metrics.data[ next(iter(gtx_metrics.program_metrics.data)) ]["compute"] - gtx_metrics.program_metrics.clear() class StencilTest: From 4d3a6d67191f031ac407d02453d34d466923ec1c Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 1 Oct 2025 14:37:52 +0200 Subject: [PATCH 034/108] Get rid of the first slow iterations --- .../src/icon4py/model/testing/pytest_hooks.py | 8 +++---- .../icon4py/model/testing/stencil_tests.py | 21 ++++++++++++++----- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index 8353ec6cdc..a935f7bb1b 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -199,9 +199,7 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config): # Print the GT4Py timer report table if benchmark_gtx_metrics: terminalreporter.ensure_newline() - header = ( - f"{'Benchmark Name':<{max_name_len}} | {'Mean (s)':>10} | {'Std Dev':>10} | {'Runs':>4}" - ) + header = f"{'Benchmark Name':<{max_name_len}} | {'Mean (s)':>10} | {'Median (s)':>10} | {'Std Dev':>10} | {'Runs':>4}" title = " GT4Py Timer Report " sep_len = max(0, len(header) - len(title)) left = sep_len // 2 @@ -209,8 +207,10 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config): terminalreporter.line("-" * left + title + "-" * right, bold=True, blue=True) terminalreporter.line(header) terminalreporter.line("-" * len(header), blue=True) + import numpy as np + for benchmark_name, gtx_metrics in benchmark_gtx_metrics: terminalreporter.line( - f"{benchmark_name:<{max_name_len}} | {gtx_metrics.mean:>10.8f} | {gtx_metrics.std:>10.8f} | {len(gtx_metrics.samples):>4}" + f"{benchmark_name:<{max_name_len}} | {np.mean(gtx_metrics):>10.8f} | {np.median(gtx_metrics):>10.8f} | {np.std(gtx_metrics):>10.8f} | {len(gtx_metrics):>4}" ) terminalreporter.line("-" * len(header), blue=True) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 566779ffa9..61ced3e40d 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -17,7 +17,6 @@ import pytest from gt4py import eve from gt4py.next import ( - backend as gtx_backend, config as gtx_config, constructors, metrics as gtx_metrics, @@ -84,7 +83,8 @@ def test_and_benchmark( _configured_program: Callable[..., None], request: pytest.FixtureRequest, ) -> None: - if not request.config.getoption("benchmark_only"): + benchmark_only = request.config.getoption("benchmark_only") + if not benchmark_only: reference_outputs = self.reference( _ConnectivityConceptFixer( grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) @@ -101,6 +101,9 @@ def test_and_benchmark( ) if benchmark is not None and benchmark.enabled: + warmup_enabled = request.config.getoption("benchmark_warmup") + if warmup_enabled: + print("[WARNING] Benchmark warmup enabled, GT4Py timers include warmup iterations.") # Clean up GT4Py metrics from previous runs if gtx_config.COLLECT_METRICS_LEVEL > 0: gtx_metrics.program_metrics.clear() @@ -116,9 +119,17 @@ def test_and_benchmark( assert ( len(gtx_metrics.program_metrics.data.keys()) == 1 ), "Expected exactly one entry in gtx_metrics" - benchmark.extra_info["gtx_metrics"] = gtx_metrics.program_metrics.data[ - next(iter(gtx_metrics.program_metrics.data)) - ]["compute"] + # Store GT4Py metrics in benchmark.extra_info + metrics_data = gtx_metrics.program_metrics.data + key = next(iter(metrics_data)) + compute_samples = metrics_data[key]["compute"].samples + # emprically exclude first few iterations as warmup + initial_program_iterations_to_skip = 2 + print(f"GT4Py compute time samples: {compute_samples}") + # Exclude first sample unless running in benchmark_only mode + benchmark.extra_info["gtx_metrics"] = compute_samples[ + initial_program_iterations_to_skip: + ] class StencilTest: From 1129847a597ec4ff99c01a9451a2453ec6a01927 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 1 Oct 2025 15:31:42 +0200 Subject: [PATCH 035/108] Added a few fixture ids --- ...test_compute_advection_in_horizontal_momentum_equation.py | 5 ++++- ...d_horizontal_winds_and_ke_and_contravariant_correction.py | 2 +- ...st_vertically_implicit_dycore_solver_at_predictor_step.py | 5 ++++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py index ad09b5e60c..70c6a4d051 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py @@ -242,7 +242,10 @@ def reference( return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) - @pytest.fixture(params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]]) + @pytest.fixture( + params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]], + ids=lambda param: f"apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]", + ) def input_data( self, request: pytest.FixtureRequest, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType]: diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 69e7811f50..56154fbfc1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -283,7 +283,7 @@ def reference( @pytest.fixture( params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], - ids=lambda param: f"skip_compute_predictor_vertical_advection={param['skip_compute_predictor_vertical_advection']}", + ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]", ) def input_data( self, grid: base.Grid, request: pytest.FixtureRequest diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index 6556796e06..7273005d8a 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -394,7 +394,10 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture(params=[{"at_first_substep": value} for value in [True, False]]) + @pytest.fixture( + params=[{"at_first_substep": value} for value in [True, False]], + ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", + ) def input_data( self, request: pytest.FixtureRequest, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType]: From c6de2d1f50a0a058cad9ba3c4a4b22dd149ca3ce Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Tue, 7 Oct 2025 11:28:40 +0200 Subject: [PATCH 036/108] Add at_first_substep to compile-time constants in vertically_implicit_dycore_solver_at_predictor_step --- .../test_vertically_implicit_dycore_solver_at_predictor_step.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index 7273005d8a..83aa4a344a 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -80,6 +80,7 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): "divdamp_type", "rayleigh_type", "is_iau_active", + "at_first_substep", ), stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( "end_index_of_damping_layer", @@ -90,6 +91,7 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): "divdamp_type", "rayleigh_type", "is_iau_active", + "at_first_substep", ), } From 9f1b6a6fda7a33c35f4daa31f4645f9b30c61833 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 8 Oct 2025 13:58:48 +0200 Subject: [PATCH 037/108] Prepare for combination --- .../compute_cell_diagnostics_for_dycore.py | 77 +++++++++---------- 1 file changed, 36 insertions(+), 41 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 17ebb5bb5c..a917c0120d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -32,8 +32,8 @@ from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( _extrapolate_temporally_exner_pressure, ) -from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import ( - _init_cell_kdim_field_with_zero_wp, +from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_vp import ( + _init_cell_kdim_field_with_zero_vp, ) from icon4py.model.atmosphere.dycore.stencils.init_two_cell_kdim_fields_with_zero_vp import ( _init_two_cell_kdim_fields_with_zero_vp, @@ -175,14 +175,11 @@ def _compute_perturbed_quantities_and_interpolation( @gtx.field_operator def _surface_computations( + temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], wgtfacq_c: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], igradp_method: gtx.int32, -) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], -]: - temporal_extrapolation_of_perturbed_exner = _init_cell_kdim_field_with_zero_wp() +) -> fa.CellKField[ta.vpfloat]: exner_at_cells_on_half_levels = ( _interpolate_to_surface( @@ -192,10 +189,7 @@ def _surface_computations( else exner_at_cells_on_half_levels ) - return ( - temporal_extrapolation_of_perturbed_exner, - exner_at_cells_on_half_levels, - ) + return exner_at_cells_on_half_levels @gtx.field_operator @@ -252,28 +246,15 @@ def _compute_first_and_second_vertical_derivative_of_exner( @gtx.field_operator -def _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], +def _set_theta_v_on_surface_level( reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], -) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: - perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels - ) + perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], + ) -> fa.CellKField[wpfloat]: theta_v_at_cells_on_half_levels = ( reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels ) - exner_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - - return ( - perturbed_theta_v_at_cells_on_half_levels, - astype(theta_v_at_cells_on_half_levels, wpfloat), - exner_at_cells_on_half_levels, - ) + return astype(theta_v_at_cells_on_half_levels, wpfloat) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -374,6 +355,8 @@ def compute_perturbed_quantities_and_interpolation( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_level - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels """ + + # Init perturbed quantities to zero on lateral boundary _init_two_cell_kdim_fields_with_zero_vp( out=(perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), domain={ @@ -382,6 +365,15 @@ def compute_perturbed_quantities_and_interpolation( }, ) + # Compute temporal extrapolation of perturbed exner, needs to be output for future program + _init_cell_kdim_field_with_zero_vp( + out=temporal_extrapolation_of_perturbed_exner, + domain={ + dims.CellDim: (start_cell_lateral_boundary, start_cell_lateral_boundary_level_3), + dims.KDim: (surface_level - 1, surface_level), + }, + ) + _extrapolate_temporally_exner_pressure( exner_exfac=time_extrapolation_parameter_for_exner, exner=current_exner, @@ -395,13 +387,11 @@ def compute_perturbed_quantities_and_interpolation( ) _surface_computations( + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, wgtfacq_c=wgtfacq_c, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, igradp_method=igradp_method, - out=( - temporal_extrapolation_of_perturbed_exner, - exner_at_cells_on_half_levels, - ), + out= exner_at_cells_on_half_levels, domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), dims.KDim: (surface_level - 1, surface_level), @@ -441,16 +431,10 @@ def compute_perturbed_quantities_and_interpolation( }, ) - _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + _interpolate_to_surface( wgtfacq_c=wgtfacq_c, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - out=( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ), + interpolant=perturbed_theta_v_at_cells_on_model_levels, + out = perturbed_theta_v_at_cells_on_half_levels, domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), dims.KDim: (surface_level - 1, surface_level), @@ -479,6 +463,17 @@ def compute_perturbed_quantities_and_interpolation( }, ) + _set_theta_v_on_surface_level( + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + perturbed_theta_v_at_cells_on_half_levels=perturbed_theta_v_at_cells_on_half_levels, + out= theta_v_at_cells_on_half_levels, + domain={ + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (surface_level - 1, surface_level), + }, + ) + + # Init perturbed quantities to zero on 2nd halo layer _compute_perturbation_of_rho_and_theta( rho=current_rho, rho_ref_mc=reference_rho_at_cells_on_model_levels, From 3207b99c224324e77cd015d0bb1441b1ea9465ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 8 Oct 2025 14:08:48 +0200 Subject: [PATCH 038/108] pre commit --- .../stencils/compute_cell_diagnostics_for_dycore.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index a917c0120d..43fe691106 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -180,7 +180,6 @@ def _surface_computations( exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], igradp_method: gtx.int32, ) -> fa.CellKField[ta.vpfloat]: - exner_at_cells_on_half_levels = ( _interpolate_to_surface( wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner @@ -249,7 +248,7 @@ def _compute_first_and_second_vertical_derivative_of_exner( def _set_theta_v_on_surface_level( reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], - ) -> fa.CellKField[wpfloat]: +) -> fa.CellKField[wpfloat]: theta_v_at_cells_on_half_levels = ( reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels ) @@ -391,7 +390,7 @@ def compute_perturbed_quantities_and_interpolation( wgtfacq_c=wgtfacq_c, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, igradp_method=igradp_method, - out= exner_at_cells_on_half_levels, + out=exner_at_cells_on_half_levels, domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), dims.KDim: (surface_level - 1, surface_level), @@ -434,7 +433,7 @@ def compute_perturbed_quantities_and_interpolation( _interpolate_to_surface( wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels, - out = perturbed_theta_v_at_cells_on_half_levels, + out=perturbed_theta_v_at_cells_on_half_levels, domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), dims.KDim: (surface_level - 1, surface_level), @@ -466,7 +465,7 @@ def compute_perturbed_quantities_and_interpolation( _set_theta_v_on_surface_level( reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, perturbed_theta_v_at_cells_on_half_levels=perturbed_theta_v_at_cells_on_half_levels, - out= theta_v_at_cells_on_half_levels, + out=theta_v_at_cells_on_half_levels, domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), dims.KDim: (surface_level - 1, surface_level), From 171ffd7e3e22f5feeccf8aadfd3ec479acc203f8 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 8 Oct 2025 16:55:50 +0200 Subject: [PATCH 039/108] Remove debug print of gt4py timers --- model/testing/src/icon4py/model/testing/stencil_tests.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 61ced3e40d..e3d06e5724 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -125,7 +125,6 @@ def test_and_benchmark( compute_samples = metrics_data[key]["compute"].samples # emprically exclude first few iterations as warmup initial_program_iterations_to_skip = 2 - print(f"GT4Py compute time samples: {compute_samples}") # Exclude first sample unless running in benchmark_only mode benchmark.extra_info["gtx_metrics"] = compute_samples[ initial_program_iterations_to_skip: From 2728bfdb2dddfbfdd401823454a3f0eacd24e076 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 8 Oct 2025 18:57:19 +0200 Subject: [PATCH 040/108] Fix metrics with new gt4py --- .../testing/src/icon4py/model/testing/stencil_tests.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index e3d06e5724..067a5682b9 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -106,7 +106,7 @@ def test_and_benchmark( print("[WARNING] Benchmark warmup enabled, GT4Py timers include warmup iterations.") # Clean up GT4Py metrics from previous runs if gtx_config.COLLECT_METRICS_LEVEL > 0: - gtx_metrics.program_metrics.clear() + gtx_metrics.sources.clear() benchmark( _configured_program, @@ -117,12 +117,12 @@ def test_and_benchmark( # Collect GT4Py runtime metrics if enabled if gtx_config.COLLECT_METRICS_LEVEL > 0: assert ( - len(gtx_metrics.program_metrics.data.keys()) == 1 - ), "Expected exactly one entry in gtx_metrics" + len(gtx_metrics.sources) == 1 + ), "Expected exactly one entry in gtx_metrics.sources" # Store GT4Py metrics in benchmark.extra_info - metrics_data = gtx_metrics.program_metrics.data + metrics_data = gtx_metrics.sources key = next(iter(metrics_data)) - compute_samples = metrics_data[key]["compute"].samples + compute_samples = metrics_data[key].metrics["compute"].samples # emprically exclude first few iterations as warmup initial_program_iterations_to_skip = 2 # Exclude first sample unless running in benchmark_only mode From a1cb8f050adeaed4a401ca3866972815c8846e8e Mon Sep 17 00:00:00 2001 From: yiluchen1066 Date: Thu, 9 Oct 2025 10:16:39 +0200 Subject: [PATCH 041/108] first try of combine 4 field operators --- .../compute_cell_diagnostics_for_dycore.py | 141 ++++++------------ 1 file changed, 48 insertions(+), 93 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 43fe691106..62fa9bd037 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -71,9 +71,12 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( wpfloat, ) - @gtx.field_operator -def _compute_perturbed_quantities_and_interpolation( +def _combined_field_operator( + temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], + wgtfacq_c: fa.CellKField[ta.wpfloat], + exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], + igradp_method: gtx.int32, current_rho: fa.CellKField[ta.wpfloat], reference_rho_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], current_theta_v: fa.CellKField[ta.wpfloat], @@ -85,14 +88,20 @@ def _compute_perturbed_quantities_and_interpolation( ddqz_z_half: fa.CellKField[ta.vpfloat], pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - igradp_method: gtx.int32, nflatlev: gtx.int32, + inv_ddqz_z_full: fa.CellKField[ta.vpfloat], + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], + d2dexdz2_fac1_mc: fa.CellKField[ta.vpfloat], + d2dexdz2_fac2_mc: fa.CellKField[ta.vpfloat], + nflat_gradp: gtx.int32, ) -> tuple[ fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], @@ -100,6 +109,16 @@ def _compute_perturbed_quantities_and_interpolation( fa.CellKField[ta.vpfloat], fa.CellKField[ta.wpfloat], ]: + # _surface_computations + exner_at_cells_on_half_levels = ( + _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ) + if igradp_method == horzpres_discr_type.TAYLOR_HYDRO + else exner_at_cells_on_half_levels + ) + + # _compute_perturbed_quantities_and_interpolation exner_at_cells_on_half_levels = ( concat_where( (maximum(1, nflatlev) <= dims.KDim), @@ -161,53 +180,14 @@ def _compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels, ) - return ( - perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels, - perturbed_exner_at_cells_on_model_levels, - rho_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - pressure_buoyancy_acceleration_at_cells_on_half_levels, - ) - - -@gtx.field_operator -def _surface_computations( - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], - wgtfacq_c: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - igradp_method: gtx.int32, -) -> fa.CellKField[ta.vpfloat]: - exner_at_cells_on_half_levels = ( - _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else exner_at_cells_on_half_levels + # _interpolate_to_surface + perturbed_theta_v_at_cells_on_half_levels = ( + wgtfacq_c(Koff[-1]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-1]) + + wgtfacq_c(Koff[-2]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-2]) + + wgtfacq_c(Koff[-3]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-3]) ) - return exner_at_cells_on_half_levels - - -@gtx.field_operator -def _compute_first_and_second_vertical_derivative_of_exner( - exner_at_cells_on_half_levels: fa.CellKField[vpfloat], - inv_ddqz_z_full: fa.CellKField[vpfloat], - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], - d2dexdz2_fac1_mc: fa.CellKField[vpfloat], - d2dexdz2_fac2_mc: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - igradp_method: gtx.int32, - nflatlev: gtx.int32, - nflat_gradp: gtx.int32, -) -> tuple[ - fa.CellKField[vpfloat], - fa.CellKField[vpfloat], -]: + # _compute_first_and_second_vertical_derivative_of_exner ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( (nflatlev <= dims.KDim), @@ -239,6 +219,15 @@ def _compute_first_and_second_vertical_derivative_of_exner( ) return ( + exner_at_cells_on_half_levels, + perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels, + perturbed_exner_at_cells_on_model_levels, + rho_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + pressure_buoyancy_acceleration_at_cells_on_half_levels, ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, ) @@ -385,19 +374,11 @@ def compute_perturbed_quantities_and_interpolation( }, ) - _surface_computations( + _combined_field_operator( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, wgtfacq_c=wgtfacq_c, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, igradp_method=igradp_method, - out=exner_at_cells_on_half_levels, - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (surface_level - 1, surface_level), - }, - ) - - _compute_perturbed_quantities_and_interpolation( current_rho=current_rho, reference_rho_at_cells_on_model_levels=reference_rho_at_cells_on_model_levels, current_theta_v=current_theta_v, @@ -409,12 +390,16 @@ def compute_perturbed_quantities_and_interpolation( ddqz_z_half=ddqz_z_half, pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, - exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, - igradp_method=igradp_method, nflatlev=nflatlev, + inv_ddqz_z_full=inv_ddqz_z_full, + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, + d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, + nflat_gradp=nflat_gradp, out=( + exner_at_cells_on_half_levels, perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, perturbed_exner_at_cells_on_model_levels, @@ -423,42 +408,12 @@ def compute_perturbed_quantities_and_interpolation( perturbed_theta_v_at_cells_on_half_levels, theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, - ), - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), - }, - ) - - _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, - interpolant=perturbed_theta_v_at_cells_on_model_levels, - out=perturbed_theta_v_at_cells_on_half_levels, - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (surface_level - 1, surface_level), - }, - ) - - _compute_first_and_second_vertical_derivative_of_exner( - exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - inv_ddqz_z_full=inv_ddqz_z_full, - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - perturbed_theta_v_at_cells_on_half_levels=perturbed_theta_v_at_cells_on_half_levels, - d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, - d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - igradp_method=igradp_method, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - out=( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), + dims.KDim: (model_top, surface_level), }, ) From ffb002dec8d3c7376439c0a5a2798cc311af800a Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 9 Oct 2025 11:48:00 +0200 Subject: [PATCH 042/108] some edits and small refactoring --- .../stencils/compute_cell_diagnostics_for_dycore.py | 9 +++++---- .../dycore/integration_tests/test_solve_nonhydro.py | 6 +++++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 62fa9bd037..f26476ff97 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -71,12 +71,12 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( wpfloat, ) + @gtx.field_operator def _combined_field_operator( temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], wgtfacq_c: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - igradp_method: gtx.int32, current_rho: fa.CellKField[ta.wpfloat], reference_rho_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], current_theta_v: fa.CellKField[ta.wpfloat], @@ -89,12 +89,13 @@ def _combined_field_operator( pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - nflatlev: gtx.int32, inv_ddqz_z_full: fa.CellKField[ta.vpfloat], ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], d2dexdz2_fac1_mc: fa.CellKField[ta.vpfloat], d2dexdz2_fac2_mc: fa.CellKField[ta.vpfloat], + igradp_method: gtx.int32, + nflatlev: gtx.int32, nflat_gradp: gtx.int32, ) -> tuple[ fa.CellKField[ta.vpfloat], @@ -378,7 +379,6 @@ def compute_perturbed_quantities_and_interpolation( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, wgtfacq_c=wgtfacq_c, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - igradp_method=igradp_method, current_rho=current_rho, reference_rho_at_cells_on_model_levels=reference_rho_at_cells_on_model_levels, current_theta_v=current_theta_v, @@ -391,12 +391,13 @@ def compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, - nflatlev=nflatlev, inv_ddqz_z_full=inv_ddqz_z_full, ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, + igradp_method=igradp_method, + nflatlev=nflatlev, nflat_gradp=nflat_gradp, out=( exner_at_cells_on_half_levels, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index abd636d98d..4e8a011d40 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1182,7 +1182,11 @@ def test_compute_perturbed_quantities_and_interpolation( assert test_utils.dallclose( perturbed_exner_at_cells_on_model_levels.asnumpy(), exner_pr_ref.asnumpy() ) - assert test_utils.dallclose(rho_at_cells_on_half_levels.asnumpy(), rho_ic_ref.asnumpy()) + # TODO: indexing to remove last bound to see if it works + assert test_utils.dallclose( + rho_at_cells_on_half_levels.asnumpy()[:, : icon_grid.num_levels], + rho_ic_ref.asnumpy()[:, : icon_grid.num_levels], + ) assert test_utils.dallclose( exner_at_cells_on_half_levels.asnumpy()[:, nflatlev:], From 8b1bd03453780ef8dc91cc2f2ec90d1435ce4e8c Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 9 Oct 2025 11:52:10 +0200 Subject: [PATCH 043/108] further indexing to remove last level --- .../tests/dycore/integration_tests/test_solve_nonhydro.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 4e8a011d40..4ec2cd802b 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1182,15 +1182,15 @@ def test_compute_perturbed_quantities_and_interpolation( assert test_utils.dallclose( perturbed_exner_at_cells_on_model_levels.asnumpy(), exner_pr_ref.asnumpy() ) - # TODO: indexing to remove last bound to see if it works + # TODO: indexing to remove last level to see if it works assert test_utils.dallclose( rho_at_cells_on_half_levels.asnumpy()[:, : icon_grid.num_levels], rho_ic_ref.asnumpy()[:, : icon_grid.num_levels], ) - + # TODO: indexing to remove last level to see if it works assert test_utils.dallclose( - exner_at_cells_on_half_levels.asnumpy()[:, nflatlev:], - z_exner_ic_ref.asnumpy()[:, nflatlev:], + exner_at_cells_on_half_levels.asnumpy()[:, nflatlev : icon_grid.num_levels], + z_exner_ic_ref.asnumpy()[:, nflatlev : icon_grid.num_levels], rtol=1e-11, ) From 205c1cbd8b1c1053eea7ac305efe2845fae5f914 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 9 Oct 2025 14:14:12 +0200 Subject: [PATCH 044/108] edits for debugging for me and Yilu --- .../compute_cell_diagnostics_for_dycore.py | 15 ++++++++------- .../integration_tests/test_solve_nonhydro.py | 11 ++++++++--- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index f26476ff97..518f0a2726 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -97,6 +97,7 @@ def _combined_field_operator( igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, + nlev: gtx.int32 ) -> tuple[ fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], @@ -108,15 +109,15 @@ def _combined_field_operator( fa.CellKField[ta.wpfloat], fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], - fa.CellKField[ta.wpfloat], ]: # _surface_computations - exner_at_cells_on_half_levels = ( + exner_at_cells_on_half_levels = concat_where( + dims.KDim >= nlev, _interpolate_to_surface( wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner ) if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else exner_at_cells_on_half_levels + else exner_at_cells_on_half_levels, exner_at_cells_on_half_levels ) # _compute_perturbed_quantities_and_interpolation @@ -182,10 +183,11 @@ def _combined_field_operator( ) # _interpolate_to_surface - perturbed_theta_v_at_cells_on_half_levels = ( + perturbed_theta_v_at_cells_on_half_levels = concat_where( + dims.KDim >= nlev-1, wgtfacq_c(Koff[-1]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-1]) + wgtfacq_c(Koff[-2]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-2]) - + wgtfacq_c(Koff[-3]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-3]) + + wgtfacq_c(Koff[-3]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-3]), perturbed_theta_v_at_cells_on_half_levels ) # _compute_first_and_second_vertical_derivative_of_exner @@ -220,7 +222,6 @@ def _combined_field_operator( ) return ( - exner_at_cells_on_half_levels, perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, perturbed_exner_at_cells_on_model_levels, @@ -399,8 +400,8 @@ def compute_perturbed_quantities_and_interpolation( igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, + nlev=surface_level, out=( - exner_at_cells_on_half_levels, perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, perturbed_exner_at_cells_on_model_levels, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 4ec2cd802b..ea81eaf770 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1202,9 +1202,14 @@ def test_compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels.asnumpy()[lb:, :], theta_v_ic_ref.asnumpy()[lb:, :] ) - assert test_utils.dallclose( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:], - z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:], + import numpy as np + print(np.where(abs( + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:icon_grid.num_levels]- + z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:icon_grid.num_levels] + )>1e-5)) + assert test_utils.dallclose( + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:icon_grid.num_levels], + z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:icon_grid.num_levels], rtol=5e-9, ) assert test_utils.dallclose( From a879762aa812e163fe2c9d5a0916970991b5ae98 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 9 Oct 2025 14:50:55 +0200 Subject: [PATCH 045/108] more fixes --- .../compute_cell_diagnostics_for_dycore.py | 20 +++++++------------ .../integration_tests/test_solve_nonhydro.py | 14 +++++-------- 2 files changed, 12 insertions(+), 22 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 518f0a2726..8773a818e4 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -110,15 +110,16 @@ def _combined_field_operator( fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], ]: + temporal_extrapolation_of_perturbed_exner = concat_where( + dims.KDim == nlev, 0.0, temporal_extrapolation_of_perturbed_exner) + # _surface_computations exner_at_cells_on_half_levels = concat_where( - dims.KDim >= nlev, + dims.KDim == nlev, _interpolate_to_surface( wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else exner_at_cells_on_half_levels, exner_at_cells_on_half_levels - ) + ), exner_at_cells_on_half_levels) if igradp_method == horzpres_discr_type.TAYLOR_HYDRO else exner_at_cells_on_half_levels + # _compute_perturbed_quantities_and_interpolation exner_at_cells_on_half_levels = ( @@ -144,7 +145,7 @@ def _combined_field_operator( ) rho_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + (dims.KDim >= 1) & (dims.KDim < nlev), _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho), rho_at_cells_on_half_levels, ) @@ -356,13 +357,6 @@ def compute_perturbed_quantities_and_interpolation( ) # Compute temporal extrapolation of perturbed exner, needs to be output for future program - _init_cell_kdim_field_with_zero_vp( - out=temporal_extrapolation_of_perturbed_exner, - domain={ - dims.CellDim: (start_cell_lateral_boundary, start_cell_lateral_boundary_level_3), - dims.KDim: (surface_level - 1, surface_level), - }, - ) _extrapolate_temporally_exner_pressure( exner_exfac=time_extrapolation_parameter_for_exner, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index ea81eaf770..6d26ace709 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1184,8 +1184,8 @@ def test_compute_perturbed_quantities_and_interpolation( ) # TODO: indexing to remove last level to see if it works assert test_utils.dallclose( - rho_at_cells_on_half_levels.asnumpy()[:, : icon_grid.num_levels], - rho_ic_ref.asnumpy()[:, : icon_grid.num_levels], + rho_at_cells_on_half_levels.asnumpy(), + rho_ic_ref.asnumpy(), ) # TODO: indexing to remove last level to see if it works assert test_utils.dallclose( @@ -1202,14 +1202,10 @@ def test_compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels.asnumpy()[lb:, :], theta_v_ic_ref.asnumpy()[lb:, :] ) - import numpy as np - print(np.where(abs( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:icon_grid.num_levels]- - z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:icon_grid.num_levels] - )>1e-5)) + # TODO: this inherits error because it uses exner_at_cells_on_half_levels assert test_utils.dallclose( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:icon_grid.num_levels], - z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:icon_grid.num_levels], + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:icon_grid.num_levels-1], + z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:icon_grid.num_levels-1], rtol=5e-9, ) assert test_utils.dallclose( From 070dc87a3a8a39187487e8b829999af8ee21abb5 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 9 Oct 2025 14:58:50 +0200 Subject: [PATCH 046/108] more fixes and --- .../stencils/compute_cell_diagnostics_for_dycore.py | 6 +++--- .../dycore/integration_tests/test_solve_nonhydro.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 8773a818e4..a55967bbbf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -111,11 +111,11 @@ def _combined_field_operator( fa.CellKField[ta.vpfloat], ]: temporal_extrapolation_of_perturbed_exner = concat_where( - dims.KDim == nlev, 0.0, temporal_extrapolation_of_perturbed_exner) + dims.KDim >= nlev-1, 0.0, temporal_extrapolation_of_perturbed_exner) # _surface_computations exner_at_cells_on_half_levels = concat_where( - dims.KDim == nlev, + dims.KDim >= nlev-1, _interpolate_to_surface( wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner ), exner_at_cells_on_half_levels) if igradp_method == horzpres_discr_type.TAYLOR_HYDRO else exner_at_cells_on_half_levels @@ -145,7 +145,7 @@ def _combined_field_operator( ) rho_at_cells_on_half_levels = concat_where( - (dims.KDim >= 1) & (dims.KDim < nlev), + (dims.KDim >= 1) & (dims.KDim < nlev-1), _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho), rho_at_cells_on_half_levels, ) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 6d26ace709..88fb4eeefa 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1182,15 +1182,15 @@ def test_compute_perturbed_quantities_and_interpolation( assert test_utils.dallclose( perturbed_exner_at_cells_on_model_levels.asnumpy(), exner_pr_ref.asnumpy() ) - # TODO: indexing to remove last level to see if it works + assert test_utils.dallclose( rho_at_cells_on_half_levels.asnumpy(), rho_ic_ref.asnumpy(), ) - # TODO: indexing to remove last level to see if it works + # TODO: field fails at last k level assert test_utils.dallclose( - exner_at_cells_on_half_levels.asnumpy()[:, nflatlev : icon_grid.num_levels], - z_exner_ic_ref.asnumpy()[:, nflatlev : icon_grid.num_levels], + exner_at_cells_on_half_levels.asnumpy()[:, nflatlev : ], + z_exner_ic_ref.asnumpy()[:, nflatlev : ], rtol=1e-11, ) @@ -1204,8 +1204,8 @@ def test_compute_perturbed_quantities_and_interpolation( # TODO: this inherits error because it uses exner_at_cells_on_half_levels assert test_utils.dallclose( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:icon_grid.num_levels-1], - z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:icon_grid.num_levels-1], + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:], + z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:], rtol=5e-9, ) assert test_utils.dallclose( From 34d5a76f381396b1e1cb0731fa1ca9e626867f14 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 10 Oct 2025 08:33:03 +0200 Subject: [PATCH 047/108] edits for exner field --- .../compute_cell_diagnostics_for_dycore.py | 36 ++++++++++--------- .../integration_tests/test_solve_nonhydro.py | 6 ++-- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index a55967bbbf..785da7c8f2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -32,9 +32,6 @@ from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( _extrapolate_temporally_exner_pressure, ) -from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_vp import ( - _init_cell_kdim_field_with_zero_vp, -) from icon4py.model.atmosphere.dycore.stencils.init_two_cell_kdim_fields_with_zero_vp import ( _init_two_cell_kdim_fields_with_zero_vp, ) @@ -97,7 +94,7 @@ def _combined_field_operator( igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, - nlev: gtx.int32 + nlev: gtx.int32, ) -> tuple[ fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], @@ -111,15 +108,8 @@ def _combined_field_operator( fa.CellKField[ta.vpfloat], ]: temporal_extrapolation_of_perturbed_exner = concat_where( - dims.KDim >= nlev-1, 0.0, temporal_extrapolation_of_perturbed_exner) - - # _surface_computations - exner_at_cells_on_half_levels = concat_where( - dims.KDim >= nlev-1, - _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ), exner_at_cells_on_half_levels) if igradp_method == horzpres_discr_type.TAYLOR_HYDRO else exner_at_cells_on_half_levels - + dims.KDim >= nlev - 1, 0.0, temporal_extrapolation_of_perturbed_exner + ) # _compute_perturbed_quantities_and_interpolation exner_at_cells_on_half_levels = ( @@ -134,6 +124,19 @@ def _combined_field_operator( else exner_at_cells_on_half_levels ) + # _surface_computations + exner_at_cells_on_half_levels = ( + concat_where( + dims.KDim >= nlev - 1, + _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ), + exner_at_cells_on_half_levels, + ) + if igradp_method == horzpres_discr_type.TAYLOR_HYDRO + else exner_at_cells_on_half_levels + ) + ( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -145,7 +148,7 @@ def _combined_field_operator( ) rho_at_cells_on_half_levels = concat_where( - (dims.KDim >= 1) & (dims.KDim < nlev-1), + (dims.KDim >= 1) & (dims.KDim < nlev - 1), _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho), rho_at_cells_on_half_levels, ) @@ -185,10 +188,11 @@ def _combined_field_operator( # _interpolate_to_surface perturbed_theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= nlev-1, + dims.KDim >= nlev - 1, wgtfacq_c(Koff[-1]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-1]) + wgtfacq_c(Koff[-2]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-2]) - + wgtfacq_c(Koff[-3]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-3]), perturbed_theta_v_at_cells_on_half_levels + + wgtfacq_c(Koff[-3]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-3]), + perturbed_theta_v_at_cells_on_half_levels, ) # _compute_first_and_second_vertical_derivative_of_exner diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 88fb4eeefa..e5d87e8f44 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1187,10 +1187,9 @@ def test_compute_perturbed_quantities_and_interpolation( rho_at_cells_on_half_levels.asnumpy(), rho_ic_ref.asnumpy(), ) - # TODO: field fails at last k level assert test_utils.dallclose( - exner_at_cells_on_half_levels.asnumpy()[:, nflatlev : ], - z_exner_ic_ref.asnumpy()[:, nflatlev : ], + exner_at_cells_on_half_levels.asnumpy()[:, nflatlev:], + z_exner_ic_ref.asnumpy()[:, nflatlev:], rtol=1e-11, ) @@ -1202,7 +1201,6 @@ def test_compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels.asnumpy()[lb:, :], theta_v_ic_ref.asnumpy()[lb:, :] ) - # TODO: this inherits error because it uses exner_at_cells_on_half_levels assert test_utils.dallclose( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, nflatlev:], z_dexner_dz_c_1_ref.asnumpy()[lb:, nflatlev:], From 18e3817b85134fb4dde4fcba766b7935202eebee Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Fri, 10 Oct 2025 20:56:32 +0200 Subject: [PATCH 048/108] Find a way to inherit from existing StencilTests --- .../stencil_tests/test_apply_diffusion_to_vn.py | 6 +++++- .../src/icon4py/model/testing/stencil_tests.py | 15 +++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index debdae974f..024ad3f775 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -26,7 +26,6 @@ @pytest.mark.uses_concat_where -@pytest.mark.continuous_benchmarking class TestApplyDiffusionToVn(StencilTest): PROGRAM = apply_diffusion_to_vn OUTPUTS = ("vn",) @@ -164,3 +163,8 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestCBApplyDiffusionToVn(TestApplyDiffusionToVn): + pass diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 067a5682b9..3bac0ca471 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -244,11 +244,22 @@ def static_variant(request: pytest.FixtureRequest) -> Sequence[str]: def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) - setattr(cls, f"test_{cls.__name__}", test_and_benchmark) + pytest_prefix = "test_" + + setattr(cls, f"{pytest_prefix}{cls.__name__}", test_and_benchmark) + + # in case a test inherits from another test avoid running the tests of its parent + if cls.__base__ is not None and cls.__base__ != StencilTest: + # TODO(iomaganaris): find a way to hide this instead of using an empty function + setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", lambda: ()) # decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition - if cls.STATIC_PARAMS is None: + # Check if cls.static_variant is already a pytest fixture to allow inheriting from other StencilTests + if hasattr(cls.static_variant, "_pytestfixturefunction"): + # Already a fixture, do nothing + pass + elif cls.STATIC_PARAMS is None: # not parametrized, return an empty tuple cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] # we override with a non-parametrized function else: From c525572c70d24946078a62affc20cf94a9e2757f Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 13 Oct 2025 12:04:31 +0200 Subject: [PATCH 049/108] Enable inheritance from classes of benchmarks and apply to TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence. Disable verification for tests marked as continuous_benchmarking --- ...ute_horizontal_gradients_for_turbulence.py | 39 +++++++++++++++++-- .../src/icon4py/model/testing/pytest_hooks.py | 13 +++++++ .../icon4py/model/testing/stencil_tests.py | 5 ++- pyproject.toml | 1 + 4 files changed, 52 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index bdf70f59bd..9b0ebbaaac 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -8,12 +8,13 @@ import gt4py.next as gtx import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base +from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils.data_allocation import random_field, zero_field from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest @@ -28,7 +29,6 @@ @pytest.mark.embedded_remap_error -@pytest.mark.continuous_benchmarking class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTest): PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence OUTPUTS = ("w", "dwdx", "dwdy") @@ -119,8 +119,8 @@ def input_data(self, grid: base.Grid) -> dict: diff_multfac_w = 5.0 w = zero_field(grid, dims.CellDim, dims.KDim) - dwdx = zero_field(grid, dims.CellDim, dims.KDim) - dwdy = zero_field(grid, dims.CellDim, dims.KDim) + dwdx = random_field(grid, dims.CellDim, dims.KDim) + dwdy = random_field(grid, dims.CellDim, dims.KDim) return dict( area=area, @@ -142,3 +142,34 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulenceContinuousBenchmarking( + TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + # Use the parent class's fixture indirectly by calling its method, not the fixture itself + base_data = ( + TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence.input_data.__wrapped__( + self, grid + ) + ) + cell_domain = h_grid.domain(dims.CellDim) + base_data["interior_idx"] = grid.start_index(cell_domain(h_grid.Zone.INTERIOR)) + base_data["halo_idx"] = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + + def _get_start_index_for_w_diffusion() -> int32: + return ( + grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + if grid.limited_area + else grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) + ) + + base_data["horizontal_start"] = _get_start_index_for_w_diffusion() + base_data["horizontal_end"] = grid.end_index(cell_domain(h_grid.Zone.HALO)) + return base_data diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index a935f7bb1b..519882c953 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -35,6 +35,10 @@ def pytest_configure(config): "markers", "level(name): marks test as unit or integration tests, mostly applicable where both are available", ) + config.addinivalue_line( + "markers", + "skip_verification(reason): skips verification for continuous benchmarking tests", + ) # Check if the --enable-mixed-precision option is set and set the environment variable accordingly if config.getoption("--enable-mixed-precision"): @@ -102,6 +106,15 @@ def pytest_addoption(parser: pytest.Parser): def pytest_collection_modifyitems(config, items): + """Modify collected test items based on command line options.""" + for item in items: + if (marker := item.get_closest_marker("continuous_benchmarking")) is not None: + if not config.getoption("--benchmark-only"): + item.add_marker( + pytest.mark.skip_verification( + reason="Continuous benchmarking tests shouldn't be verified." + ) + ) test_level = config.getoption("--level") if test_level == "any": return diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 3bac0ca471..81a842702d 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -84,7 +84,8 @@ def test_and_benchmark( request: pytest.FixtureRequest, ) -> None: benchmark_only = request.config.getoption("benchmark_only") - if not benchmark_only: + skip_verification = request.node.get_closest_marker("skip_verification") is not None + if (not benchmark_only) and (not skip_verification): reference_outputs = self.reference( _ConnectivityConceptFixer( grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) @@ -251,7 +252,7 @@ def __init_subclass__(cls, **kwargs: Any) -> None: # in case a test inherits from another test avoid running the tests of its parent if cls.__base__ is not None and cls.__base__ != StencilTest: # TODO(iomaganaris): find a way to hide this instead of using an empty function - setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", lambda: ()) + setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", pytest.mark.skip(lambda: ())) # decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition diff --git a/pyproject.toml b/pyproject.toml index 062b3dbda1..32e91693f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -184,6 +184,7 @@ markers = [ "infinite_concat_where", "gtfn_too_slow", "continuous_benchmarking", + "skip_verificaion", "benchmark_only: benchmark only tests without verification" ] # add all namespace packages to pythonpath to make them available for pytest From f1cfa28ad109e7821b9f581bdc86f2a513947e40 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 13 Oct 2025 13:53:02 +0200 Subject: [PATCH 050/108] edits --- .../compute_cell_diagnostics_for_dycore.py | 95 ++++++++++--------- .../integration_tests/test_solve_nonhydro.py | 39 -------- 2 files changed, 50 insertions(+), 84 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 785da7c8f2..1e1b0ad35d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -32,6 +32,9 @@ from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( _extrapolate_temporally_exner_pressure, ) +from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import ( + _init_cell_kdim_field_with_zero_wp, +) from icon4py.model.atmosphere.dycore.stencils.init_two_cell_kdim_fields_with_zero_vp import ( _init_two_cell_kdim_fields_with_zero_vp, ) @@ -71,6 +74,8 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( @gtx.field_operator def _combined_field_operator( + perturbed_rho_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], wgtfacq_c: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], @@ -91,6 +96,7 @@ def _combined_field_operator( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], d2dexdz2_fac1_mc: fa.CellKField[ta.vpfloat], d2dexdz2_fac2_mc: fa.CellKField[ta.vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, @@ -108,13 +114,15 @@ def _combined_field_operator( fa.CellKField[ta.vpfloat], ]: temporal_extrapolation_of_perturbed_exner = concat_where( - dims.KDim >= nlev - 1, 0.0, temporal_extrapolation_of_perturbed_exner + dims.KDim >= nlev - 1, + _init_cell_kdim_field_with_zero_wp(), + temporal_extrapolation_of_perturbed_exner, ) # _compute_perturbed_quantities_and_interpolation exner_at_cells_on_half_levels = ( concat_where( - (maximum(1, nflatlev) <= dims.KDim), + (dims.KDim >= nlev - 1), _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner ), @@ -124,12 +132,11 @@ def _combined_field_operator( else exner_at_cells_on_half_levels ) - # _surface_computations exner_at_cells_on_half_levels = ( concat_where( - dims.KDim >= nlev - 1, - _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + (maximum(1, nflatlev) <= dims.KDim) & (dims.KDim < nlev - 1), + _interpolate_cell_field_to_half_levels_vp( + wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner ), exner_at_cells_on_half_levels, ) @@ -140,11 +147,15 @@ def _combined_field_operator( ( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, - ) = _compute_perturbation_of_rho_and_theta( - current_rho, - reference_rho_at_cells_on_model_levels, - current_theta_v, - reference_theta_at_cells_on_model_levels, + ) = concat_where( + dims.KDim < nlev - 1, + _compute_perturbation_of_rho_and_theta( + current_rho, + reference_rho_at_cells_on_model_levels, + current_theta_v, + reference_theta_at_cells_on_model_levels, + ), + (perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), ) rho_at_cells_on_half_levels = concat_where( @@ -156,7 +167,7 @@ def _combined_field_operator( wgtfac_c_wp = astype(wgtfac_c, wpfloat) perturbed_theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + (dims.KDim >= 1) & (dims.KDim < nlev - 1), _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels ), @@ -164,7 +175,7 @@ def _combined_field_operator( ) theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + (dims.KDim >= 1) & (dims.KDim < nlev - 1), _interpolate_cell_field_to_half_levels_wp( wgtfac_c=wgtfac_c_wp, interpolant=current_theta_v ), @@ -174,7 +185,7 @@ def _combined_field_operator( ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) pressure_buoyancy_acceleration_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + (dims.KDim >= 1) & (dims.KDim < nlev - 1), _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( exner_w_explicit_weight_parameter, theta_v_at_cells_on_half_levels, @@ -186,19 +197,32 @@ def _combined_field_operator( pressure_buoyancy_acceleration_at_cells_on_half_levels, ) - # _interpolate_to_surface perturbed_theta_v_at_cells_on_half_levels = concat_where( dims.KDim >= nlev - 1, - wgtfacq_c(Koff[-1]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-1]) - + wgtfacq_c(Koff[-2]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-2]) - + wgtfacq_c(Koff[-3]) * perturbed_theta_v_at_cells_on_model_levels(Koff[-3]), + _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels + ), perturbed_theta_v_at_cells_on_half_levels, ) + theta_v_at_cells_on_half_levels = concat_where( + dims.KDim >= nlev - 1, + reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + ) + + exner_at_cells_on_half_levels = concat_where( + dims.KDim >= nlev - 1, + _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ), + exner_at_cells_on_half_levels, + ) + # _compute_first_and_second_vertical_derivative_of_exner ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflatlev <= dims.KDim), + (nflatlev <= dims.KDim) & (dims.KDim < nlev - 1), _compute_first_vertical_derivative_at_cells( exner_at_cells_on_half_levels, inv_ddqz_z_full ), @@ -210,7 +234,7 @@ def _combined_field_operator( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflat_gradp <= dims.KDim), + (nflat_gradp <= dims.KDim) & (dims.KDim < nlev - 1), -vpfloat("0.5") * ( ( @@ -240,18 +264,6 @@ def _combined_field_operator( ) -@gtx.field_operator -def _set_theta_v_on_surface_level( - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], -) -> fa.CellKField[wpfloat]: - theta_v_at_cells_on_half_levels = ( - reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels - ) - - return astype(theta_v_at_cells_on_half_levels, wpfloat) - - @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_perturbed_quantities_and_interpolation( temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], @@ -375,6 +387,8 @@ def compute_perturbed_quantities_and_interpolation( ) _combined_field_operator( + perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, wgtfacq_c=wgtfacq_c, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, @@ -395,6 +409,7 @@ def compute_perturbed_quantities_and_interpolation( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, @@ -403,10 +418,10 @@ def compute_perturbed_quantities_and_interpolation( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, perturbed_exner_at_cells_on_model_levels, - rho_at_cells_on_half_levels, - exner_at_cells_on_half_levels, + current_rho, + current_exner, perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, + current_theta_v, pressure_buoyancy_acceleration_at_cells_on_half_levels, ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, @@ -417,16 +432,6 @@ def compute_perturbed_quantities_and_interpolation( }, ) - _set_theta_v_on_surface_level( - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - perturbed_theta_v_at_cells_on_half_levels=perturbed_theta_v_at_cells_on_half_levels, - out=theta_v_at_cells_on_half_levels, - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (surface_level - 1, surface_level), - }, - ) - # Init perturbed quantities to zero on 2nd halo layer _compute_perturbation_of_rho_and_theta( rho=current_rho, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index e5d87e8f44..22ea656a47 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -810,7 +810,6 @@ def test_run_solve_nonhydro_single_step( ) -# why is this not run for APE? @pytest.mark.embedded_remap_error @pytest.mark.datatest @pytest.mark.parametrize("experiment", [definitions.Experiments.MCH_CH_R04B09]) @@ -1023,7 +1022,6 @@ def test_compute_perturbed_quantities_and_interpolation( step_date_init, step_date_exit, *, - ndyn_substeps, icon_grid, lowest_layer_thickness, model_top_height, @@ -1031,9 +1029,6 @@ def test_compute_perturbed_quantities_and_interpolation( damping_height, grid_savepoint, metrics_savepoint, - interpolation_savepoint, - substep_init, - substep_exit, savepoint_nonhydro_init, savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_init, savepoint_nonhydro_exit, @@ -1241,17 +1236,8 @@ def test_interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_ac step_date_init, step_date_exit, *, - ndyn_substeps, icon_grid, - lowest_layer_thickness, - model_top_height, - stretch_factor, - damping_height, - grid_savepoint, metrics_savepoint, - interpolation_savepoint, - substep_init, - substep_exit, savepoint_nonhydro_init, savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_init, savepoint_nonhydro_exit, @@ -1383,7 +1369,6 @@ def test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn( step_date_init, step_date_exit, *, - ndyn_substeps, icon_grid, savepoint_nonhydro_init, lowest_layer_thickness, @@ -1394,9 +1379,6 @@ def test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn( metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, - istep_init, - substep_init, - substep_exit, savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_init, savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_exit, backend, @@ -1599,19 +1581,13 @@ def test_apply_divergence_damping_and_update_vn( step_date_init, step_date_exit, *, - ndyn_substeps, icon_grid, savepoint_nonhydro_init, - lowest_layer_thickness, - model_top_height, - stretch_factor, - damping_height, grid_savepoint, metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_init, - savepoint_compute_edge_diagnostics_for_dycore_and_update_vn_exit, backend, ): sp_nh_init = savepoint_nonhydro_init @@ -1724,15 +1700,10 @@ def test_apply_divergence_damping_and_update_vn( ], ) def test_compute_horizontal_velocity_quantities_and_fluxes( - istep_init, - istep_exit, - substep_init, - substep_exit, step_date_init, step_date_exit, experiment, icon_grid, - ndyn_substeps, grid_savepoint, lowest_layer_thickness, model_top_height, @@ -1742,8 +1713,6 @@ def test_compute_horizontal_velocity_quantities_and_fluxes( savepoint_dycore_30_to_38_exit, interpolation_savepoint, metrics_savepoint, - savepoint_nonhydro_init, - savepoint_nonhydro_exit, backend, ): edge_domain = h_grid.domain(dims.EdgeDim) @@ -1906,21 +1875,17 @@ def test_compute_horizontal_velocity_quantities_and_fluxes( def test_compute_averaged_vn_and_fluxes_and_prepare_tracer_advection( istep_init, istep_exit, - substep_init, - substep_exit, step_date_init, step_date_exit, experiment, icon_grid, ndyn_substeps, at_first_substep, - grid_savepoint, savepoint_dycore_30_to_38_init, savepoint_dycore_30_to_38_exit, interpolation_savepoint, metrics_savepoint, savepoint_nonhydro_init, - savepoint_nonhydro_exit, backend, ): edge_domain = h_grid.domain(dims.EdgeDim) @@ -2029,7 +1994,6 @@ def test_vertically_implicit_solver_at_predictor_step( step_date_init, step_date_exit, *, - ndyn_substeps, icon_grid, savepoint_nonhydro_init, lowest_layer_thickness, @@ -2040,9 +2004,6 @@ def test_vertically_implicit_solver_at_predictor_step( metrics_savepoint, interpolation_savepoint, savepoint_nonhydro_exit, - istep_init, - istep_exit, - substep_exit, savepoint_vertically_implicit_dycore_solver_init, backend, ): From 534d31d607387be610165bfdc01a7c09ff76e1ab Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 13 Oct 2025 15:18:11 +0200 Subject: [PATCH 051/108] Use subclasses for continuous_benchmarking --- .../test_benchmark_diffusion.py | 13 +++- ...ate_nabla2_and_smag_coefficients_for_vn.py | 23 ++++++- .../stencil_tests/test_calculate_nabla4.py | 6 +- ..._apply_divergence_damping_and_update_vn.py | 65 ++++++++++++++----- ...advection_in_vertical_momentum_equation.py | 15 ++++- ...nds_and_ke_and_contravariant_correction.py | 34 +++++++++- ...est_compute_hydrostatic_correction_term.py | 8 ++- ..._perturbed_quantities_and_interpolation.py | 20 +++++- .../test_init_cell_kdim_field_with_zero_wp.py | 6 +- ...d_compute_temperature_vertical_gradient.py | 8 ++- .../test_update_mass_flux_weighted.py | 6 +- ...mplicit_dycore_solver_at_corrector_step.py | 36 +++++++++- ...mplicit_dycore_solver_at_predictor_step.py | 26 +++++++- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 11 +++- 14 files changed, 243 insertions(+), 34 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py index 8ecd3b7692..f413192a24 100644 --- a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py @@ -43,7 +43,6 @@ @pytest.mark.parametrize( "grid", [definitions.Grids.MCH_OPR_R04B07_DOMAIN01, definitions.Grids.R02B07_GLOBAL] ) -@pytest.mark.continuous_benchmarking @pytest.mark.benchmark_only def test_run_diffusion_benchmark( grid: definitions.GridDescription, @@ -220,3 +219,15 @@ def test_run_diffusion_benchmark( ) benchmark(diffusion_granule.run, diagnostic_state, prognostic_state, dtime) + + +@pytest.mark.continuous_benchmarking +def test_run_diffusion_benchmark_continuous_benchmarking( + grid: definitions.GridDescription, + backend: gtx_typing.Backend | None, + benchmark: Any, +) -> None: + assert ( + grid == definitions.Grids.MCH_OPR_R19B08_DOMAIN01 + ), "This test only works with the icon_benchmark grid." + test_run_diffusion_benchmark(grid, backend, benchmark) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index b17c7f578e..89b447e14a 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -14,11 +14,11 @@ calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import stencil_tests -@pytest.mark.continuous_benchmarking class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_and_smag_coefficients_for_vn OUTPUTS = ("kh_smag_e", "kh_smag_ec", "z_nabla2_e") @@ -157,7 +157,7 @@ def reference( return dict(kh_smag_e=kh_smag_e, kh_smag_ec=kh_smag_ec, z_nabla2_e=z_nabla2_e) @pytest.fixture - def input_data(self, grid): + def input_data(self, grid: base.Grid) -> dict: u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) smag_offset = ta.vpfloat("9.0") @@ -207,3 +207,22 @@ def input_data(self, grid): vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestCalculateNabla2AndSmagCoefficientsForVnContinuousBenchmarking( + TestCalculateNabla2AndSmagCoefficientsForVn +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + base_data = TestCalculateNabla2AndSmagCoefficientsForVn.input_data.__wrapped__(self, grid) + edge_domain = h_grid.domain(dims.EdgeDim) + horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) + horizontal_end = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) + assert horizontal_start < horizontal_end + base_data["horizontal_start"] = 0 + base_data["horizontal_end"] = grid.num_edges + return base_data diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py index dc8f9ac4ba..064c59e23f 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py @@ -55,7 +55,6 @@ def calculate_nabla4_numpy( return z_nabla4_e2 -@pytest.mark.continuous_benchmarking class TestCalculateNabla4(StencilTest): PROGRAM = calculate_nabla4 OUTPUTS = ("z_nabla4_e2",) @@ -129,3 +128,8 @@ def input_data(self, grid) -> dict: vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestCalculateNabla4ContinuousBenchmarking(TestCalculateNabla4): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index c90241697b..8553a67934 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -26,7 +26,6 @@ divergence_damp_order = DivergenceDampingOrder() -@pytest.mark.continuous_benchmarking @pytest.mark.embedded_remap_error class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): PROGRAM = apply_divergence_damping_and_update_vn @@ -173,21 +172,7 @@ def reference( return dict(next_vn=next_vn) - @pytest.fixture( - params=[ - {"limited_area": la, "divdamp_order": do, "is_iau_active": ia} - for la, do, ia in itertools.product( - [True, False], - [ - DivergenceDampingOrder.SECOND_ORDER, - DivergenceDampingOrder.FOURTH_ORDER, - DivergenceDampingOrder.COMBINED, - ], - [True, False], - ) - ], - ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", - ) + @pytest.fixture(params=[True, False]) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_mask_for_3d_divdamp = data_alloc.random_field(grid, dims.EdgeDim) @@ -224,7 +209,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: is_iau_active = True fourth_order_divdamp_factor = 0.004 second_order_divdamp_factor = 0.012 - divdamp_order = request.param["divdamp_order"] + divdamp_order = 24 second_order_divdamp_scaling_coeff = 194588.14247428576 apply_2nd_order_divergence_damping = (divdamp_order == divergence_damp_order.COMBINED) and ( second_order_divdamp_scaling_coeff > 1.0e-6 @@ -236,7 +221,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) ) - limited_area = request.param["limited_area"] + limited_area = request.param edge_domain = h_grid.domain(dims.EdgeDim) start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) @@ -274,3 +259,47 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestApplyDivergenceDampingAndUpdateVnContinuousBenchmarking( + TestApplyDivergenceDampingAndUpdateVn +): + @pytest.fixture( + params=[ + {"limited_area": la, "divdamp_order": do, "is_iau_active": ia} + for la, do, ia in itertools.product( + [True, False], + [ + DivergenceDampingOrder.SECOND_ORDER, + DivergenceDampingOrder.FOURTH_ORDER, + DivergenceDampingOrder.COMBINED, + ], + [True, False], + ) + ], + ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", + ) + def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + # Use the parent class's fixture indirectly by calling its method, not the fixture itself + base_data = TestApplyDivergenceDampingAndUpdateVn.input_data.__wrapped__( + self, request, grid + ) + base_data["is_iau_active"] = False + fourth_order_divdamp_factor = 0.004 + second_order_divdamp_factor = 0.032 + divdamp_order = request.param["divdamp_order"] + base_data["second_order_divdamp_scaling_coeff"] = 34497.62082646618 # for icon-ch1(_medium) + base_data["apply_2nd_order_divergence_damping"] = ( + divdamp_order == divergence_damp_order.COMBINED + ) and (base_data["second_order_divdamp_scaling_coeff"] > 1.0e-6) + base_data["apply_4th_order_divergence_damping"] = ( + divdamp_order == divergence_damp_order.FOURTH_ORDER + ) or ( + (divdamp_order == divergence_damp_order.COMBINED) + and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) + ) + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index b04ab64c02..2a5a8da5f9 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -268,7 +268,6 @@ def compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( return vertical_wind_advective_tendency -@pytest.mark.continuous_benchmarking @pytest.mark.embedded_remap_error class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): PROGRAM = compute_advection_in_vertical_momentum_equation @@ -483,6 +482,13 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala ) +@pytest.mark.continuous_benchmarking +class TestFusedVelocityAdvectionStencilVMomentumContinuousBenchmarking( + TestFusedVelocityAdvectionStencilVMomentum +): + pass + + @pytest.mark.embedded_remap_error class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.StencilTest): PROGRAM = compute_contravariant_correction_and_advection_in_vertical_momentum_equation @@ -695,3 +701,10 @@ def input_data( vertical_start=vertical_start, vertical_end=vertical_end, ) + + +@pytest.mark.continuous_benchmarking +class TestFusedVelocityAdvectionStencilVMomentumAndContravariantContinuousBenchmarking( + TestFusedVelocityAdvectionStencilVMomentumAndContravariant +): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index 56154fbfc1..b23d9315c4 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -141,7 +141,6 @@ def extrapolate_to_surface_numpy(wgtfacq_e: np.ndarray, vn: np.ndarray) -> np.nd return vn_at_surface -@pytest.mark.continuous_benchmarking @pytest.mark.embedded_remap_error class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection( stencil_tests.StencilTest @@ -315,7 +314,7 @@ def input_data( c_intp = data_alloc.random_field(grid, dims.VertexDim, dims.V2CDim) nlev = grid.num_levels - nflatlev = (grid.num_levels * 3) // 10 + nflatlev = 13 skip_compute_predictor_vertical_advection = request.param[ "skip_compute_predictor_vertical_advection" @@ -351,3 +350,34 @@ def input_data( vertical_start=vertical_start, vertical_end=vertical_end, ) + + +@pytest.mark.continuous_benchmarking +class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrectionContinuousBenchmarking( + TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection +): + @pytest.fixture( + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], + ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]", + ) + def input_data( + self, grid: base.Grid, request: pytest.FixtureRequest + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + base_data = TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection.input_data.__wrapped__( + self, grid, request + ) + base_data["skip_compute_predictor_vertical_advection"] = request.param[ + "skip_compute_predictor_vertical_advection" + ] + base_data["nflatlev"] = 6 + edge_domain = h_grid.domain(dims.EdgeDim) + base_data["horizontal_start"] = grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) + ) + base_data["horizontal_end"] = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) + base_data["vertical_start"] = 0 + base_data["vertical_end"] = grid.num_levels + 1 + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 1c041b2a89..8e49c48599 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -86,7 +86,6 @@ def _apply_index_field( return z_hydro_corr -@pytest.mark.continuous_benchmarking @pytest.mark.uses_as_offset class TestComputeHydrostaticCorrectionTerm(StencilTest): OUTPUTS = ("z_hydro_corr",) @@ -168,3 +167,10 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestComputeHydrostaticCorrectionTermContinuousBenchmarking( + TestComputeHydrostaticCorrectionTerm +): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 301c2fa0b4..64c45e0804 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -64,7 +64,6 @@ def compute_first_vertical_derivative_numpy( return first_vertical_derivative -@pytest.mark.continuous_benchmarking @pytest.mark.uses_concat_where class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): PROGRAM = compute_perturbed_quantities_and_interpolation @@ -432,7 +431,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala end_cell_halo_level_2 = grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) nflatlev = 4 - nflat_gradp = grid.num_levels // 2 + nflat_gradp = 27 return dict( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, @@ -477,3 +476,20 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=grid.num_levels + 1, ) + + +@pytest.mark.continuous_benchmarking +class TestComputePerturbedQuantitiesAndInterpolationContinuousBenchmarking( + TestComputePerturbedQuantitiesAndInterpolation +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + base_data = TestComputePerturbedQuantitiesAndInterpolation.input_data.__wrapped__( + self, grid + ) + base_data["nflatlev"] = 6 + base_data["nflat_gradp"] = 35 + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py index a0c2e0f902..943d0419d0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py @@ -22,7 +22,6 @@ from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest -@pytest.mark.continuous_benchmarking class TestInitCellKdimFieldWithZeroWp(StencilTest): PROGRAM = init_cell_kdim_field_with_zero_wp OUTPUTS = ("field_with_zero_wp",) @@ -60,3 +59,8 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestInitCellKdimFieldWithZeroWpContinuousBenchmarking(TestInitCellKdimFieldWithZeroWp): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py index e9ab186490..58e876c9da 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py @@ -34,7 +34,6 @@ from icon4py.model.testing import stencil_tests -@pytest.mark.continuous_benchmarking class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration( stencil_tests.StencilTest ): @@ -250,3 +249,10 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=1, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAccelerationContinuousBenchmarking( + TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration +): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py index 5a0e5422da..0013e7e1da 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py @@ -22,7 +22,6 @@ from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest -@pytest.mark.continuous_benchmarking class TestUpdateMassFluxWeighted(StencilTest): PROGRAM = update_mass_flux_weighted OUTPUTS = ("mass_flx_ic",) @@ -85,3 +84,8 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestUpdateMassFluxWeightedContinuousBenchmarking(TestUpdateMassFluxWeighted): + pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index ab8ea3d50b..4e5b84ce1f 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -51,7 +51,6 @@ from .test_update_mass_volume_flux import update_mass_volume_flux_numpy -@pytest.mark.continuous_benchmarking @pytest.mark.uses_concat_where class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_corrector_step @@ -403,6 +402,7 @@ def reference( params=[ {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} for afs, als, la in itertools.product(*([(True, False)] * 3)) + if not (afs and als) ], ids=lambda p: ( f"at_first_substep[{p['at_first_substep']}]__" @@ -535,3 +535,37 @@ def input_data( vertical_start_index_model_top=gtx.int32(0), vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), ) + + +@pytest.mark.continuous_benchmarking +class TestVerticallyImplicitSolverAtCorrectorStepContinuousBenchmarking( + TestVerticallyImplicitSolverAtCorrectorStep +): + @pytest.fixture( + params=[ + {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} + for afs, als, la in itertools.product(*([(True, False)] * 3)) + if not (afs and als) + ], + ids=lambda p: ( + f"at_first_substep[{p['at_first_substep']}]__" + f"at_last_substep[{p['at_last_substep']}]__" + f"lprep_adv[{p['lprep_adv']}]" + ), + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + base_data = TestVerticallyImplicitSolverAtCorrectorStep.input_data.__wrapped__( + self, request, grid + ) + base_data["at_first_substep"] = request.param["at_first_substep"] + base_data["at_last_substep"] = request.param["at_last_substep"] + base_data["lprep_adv"] = request.param["lprep_adv"] + base_data["is_iau_active"] = False + base_data["end_index_of_damping_layer"] = 13 + base_data["kstart_moist"] = 0 + return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index 83aa4a344a..0b1f1f5219 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -53,7 +53,6 @@ ) -@pytest.mark.continuous_benchmarking @pytest.mark.uses_concat_where class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): PROGRAM = vertically_implicit_solver_at_predictor_step @@ -524,3 +523,28 @@ def input_data( vertical_start_index_model_top=gtx.int32(0), vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), ) + + +@pytest.mark.continuous_benchmarking +class TestVerticallyImplicitSolverAtPredictorStepContinuousBenchmarking( + TestVerticallyImplicitSolverAtPredictorStep +): + @pytest.fixture( + params=[{"at_first_substep": value} for value in [True, False]], + ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", + ) + def input_data( + self, request: pytest.FixtureRequest, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + base_data = TestVerticallyImplicitSolverAtPredictorStep.input_data.__wrapped__( + self, request, grid + ) + base_data["at_first_substep"] = request.param["at_first_substep"] + base_data["is_iau_active"] = False + base_data["divdamp_type"] = 32 + base_data["end_index_of_damping_layer"] = 13 + base_data["kstart_moist"] = 0 + return base_data diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 2137f5e587..98c8ce692c 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -21,7 +21,6 @@ from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest -@pytest.mark.continuous_benchmarking class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex OUTPUTS = ("p_u_out", "p_v_out") @@ -75,3 +74,13 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) + + +@pytest.mark.continuous_benchmarking +class TestMoIntpRbfRbfVecInterpolVertexContinuousBenchmarking(TestMoIntpRbfRbfVecInterpolVertex): + @pytest.fixture + def input_data(self, grid): + assert ( + grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + ), "This test only works with the icon_benchmark grid." + return TestMoIntpRbfRbfVecInterpolVertex.input_data.__wrapped__(self, grid) From 6ee0419d8f5f8883f541b4b25c6cd3ff87a94011 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 13 Oct 2025 15:34:39 +0200 Subject: [PATCH 052/108] Replace UUID of test with class --- ...to_w_and_compute_horizontal_gradients_for_turbulence.py | 3 ++- .../test_calculate_nabla2_and_smag_coefficients_for_vn.py | 4 ++-- .../test_apply_divergence_damping_and_update_vn.py | 3 ++- ...horizontal_winds_and_ke_and_contravariant_correction.py | 4 ++-- .../test_compute_perturbed_quantities_and_interpolation.py | 4 ++-- ..._vertically_implicit_dycore_solver_at_corrector_step.py | 4 ++-- ..._vertically_implicit_dycore_solver_at_predictor_step.py | 4 ++-- .../test_mo_intp_rbf_rbf_vec_interpol_vertex.py | 7 +------ model/testing/src/icon4py/model/testing/definitions.py | 4 ++++ 9 files changed, 19 insertions(+), 18 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 9b0ebbaaac..504df38f7f 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -16,6 +16,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils.data_allocation import random_field, zero_field +from icon4py.model.testing import definitions from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy @@ -151,7 +152,7 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulenceContinuousB @pytest.fixture def input_data(self, grid: base.Grid) -> dict: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." # Use the parent class's fixture indirectly by calling its method, not the fixture itself base_data = ( diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 89b447e14a..212215ecb9 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -16,7 +16,7 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): @@ -216,7 +216,7 @@ class TestCalculateNabla2AndSmagCoefficientsForVnContinuousBenchmarking( @pytest.fixture def input_data(self, grid: base.Grid) -> dict: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." base_data = TestCalculateNabla2AndSmagCoefficientsForVn.input_data.__wrapped__(self, grid) edge_domain = h_grid.domain(dims.EdgeDim) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 8553a67934..6f8ef30b57 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -21,6 +21,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions divergence_damp_order = DivergenceDampingOrder() @@ -282,7 +283,7 @@ class TestApplyDivergenceDampingAndUpdateVnContinuousBenchmarking( ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." # Use the parent class's fixture indirectly by calling its method, not the fixture itself base_data = TestApplyDivergenceDampingAndUpdateVn.input_data.__wrapped__( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index b23d9315c4..a9a339c2bb 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -16,7 +16,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_compute_contravariant_correction import compute_contravariant_correction_numpy from .test_compute_horizontal_advection_term_for_vertical_velocity import ( @@ -364,7 +364,7 @@ def input_data( self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." base_data = TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection.input_data.__wrapped__( self, grid, request diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 64c45e0804..763450ff82 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -32,7 +32,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_compute_approx_of_2nd_vertical_derivative_of_exner import ( compute_approx_of_2nd_vertical_derivative_of_exner_numpy, @@ -485,7 +485,7 @@ class TestComputePerturbedQuantitiesAndInterpolationContinuousBenchmarking( @pytest.fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." base_data = TestComputePerturbedQuantitiesAndInterpolation.input_data.__wrapped__( self, grid diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index 4e5b84ce1f..0702ff5f8c 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -19,7 +19,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( add_analysis_increments_from_data_assimilation_numpy, @@ -557,7 +557,7 @@ def input_data( self, request: pytest.FixtureRequest, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType]: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." base_data = TestVerticallyImplicitSolverAtCorrectorStep.input_data.__wrapped__( self, request, grid diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index 0b1f1f5219..a2510c3e91 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( add_analysis_increments_from_data_assimilation_numpy, @@ -537,7 +537,7 @@ def input_data( self, request: pytest.FixtureRequest, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType]: assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 ), "This test only works with the icon_benchmark grid." base_data = TestVerticallyImplicitSolverAtPredictorStep.input_data.__wrapped__( self, request, grid diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 98c8ce692c..2ddde5d342 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -78,9 +78,4 @@ def input_data(self, grid: base.Grid) -> dict: @pytest.mark.continuous_benchmarking class TestMoIntpRbfRbfVecInterpolVertexContinuousBenchmarking(TestMoIntpRbfRbfVecInterpolVertex): - @pytest.fixture - def input_data(self, grid): - assert ( - grid.id == "01f00602-c07e-cd84-b894-bd17fffd2720" - ), "This test only works with the icon_benchmark grid." - return TestMoIntpRbfRbfVecInterpolVertex.input_data.__wrapped__(self, grid) + pass diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index 29159b1a60..30ef45bbc8 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -128,6 +128,10 @@ class Grids: ) +class GridUUIDs: + MCH_OPR_R19B08_DOMAIN01: Final = "01f00602-c07e-cd84-b894-bd17fffd2720" + + @dataclasses.dataclass class Experiment: name: str From da6c9ebe9cb3edd8bcf97bb989509e46f476394b Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 13 Oct 2025 15:39:46 +0200 Subject: [PATCH 053/108] Remove TODO --- model/testing/src/icon4py/model/testing/stencil_tests.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 81a842702d..9066741458 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -251,7 +251,6 @@ def __init_subclass__(cls, **kwargs: Any) -> None: # in case a test inherits from another test avoid running the tests of its parent if cls.__base__ is not None and cls.__base__ != StencilTest: - # TODO(iomaganaris): find a way to hide this instead of using an empty function setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", pytest.mark.skip(lambda: ())) # decorate `static_variant` with parametrized fixtures, since the From 33149f3d5b4dca346990fd36abee2dc51b4f142d Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 13 Oct 2025 15:43:12 +0200 Subject: [PATCH 054/108] Replace skip_verification with benchmark_only --- model/testing/src/icon4py/model/testing/pytest_hooks.py | 6 +----- model/testing/src/icon4py/model/testing/stencil_tests.py | 6 +++--- pyproject.toml | 1 - 3 files changed, 4 insertions(+), 9 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index 519882c953..14f4153f55 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -35,10 +35,6 @@ def pytest_configure(config): "markers", "level(name): marks test as unit or integration tests, mostly applicable where both are available", ) - config.addinivalue_line( - "markers", - "skip_verification(reason): skips verification for continuous benchmarking tests", - ) # Check if the --enable-mixed-precision option is set and set the environment variable accordingly if config.getoption("--enable-mixed-precision"): @@ -111,7 +107,7 @@ def pytest_collection_modifyitems(config, items): if (marker := item.get_closest_marker("continuous_benchmarking")) is not None: if not config.getoption("--benchmark-only"): item.add_marker( - pytest.mark.skip_verification( + pytest.mark.benchmark_only( reason="Continuous benchmarking tests shouldn't be verified." ) ) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 9066741458..14e8afa79c 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -83,9 +83,9 @@ def test_and_benchmark( _configured_program: Callable[..., None], request: pytest.FixtureRequest, ) -> None: - benchmark_only = request.config.getoption("benchmark_only") - skip_verification = request.node.get_closest_marker("skip_verification") is not None - if (not benchmark_only) and (not skip_verification): + benchmark_only_option = request.config.getoption("benchmark_only") + benchmark_only_mark = request.node.get_closest_marker("benchmark_only") is not None + if (not benchmark_only_option) and (not benchmark_only_mark): reference_outputs = self.reference( _ConnectivityConceptFixer( grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) diff --git a/pyproject.toml b/pyproject.toml index 32e91693f5..062b3dbda1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -184,7 +184,6 @@ markers = [ "infinite_concat_where", "gtfn_too_slow", "continuous_benchmarking", - "skip_verificaion", "benchmark_only: benchmark only tests without verification" ] # add all namespace packages to pythonpath to make them available for pytest From 6030243e0f0c9884c70e9615b2af6f7c1690c025 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 14 Oct 2025 09:47:45 +0200 Subject: [PATCH 055/108] further edits --- .../stencils/compute_cell_diagnostics_for_dycore.py | 6 +++--- .../dycore/integration_tests/test_solve_nonhydro.py | 10 +++------- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 1e1b0ad35d..5522c4ae11 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -418,10 +418,10 @@ def compute_perturbed_quantities_and_interpolation( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, perturbed_exner_at_cells_on_model_levels, - current_rho, - current_exner, + rho_at_cells_on_half_levels, + exner_at_cells_on_half_levels, perturbed_theta_v_at_cells_on_half_levels, - current_theta_v, + theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 22ea656a47..fc03ecb9a9 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -156,8 +156,6 @@ def test_time_step_flags( ], ) def test_nonhydro_predictor_step( - istep_init, - istep_exit, substep_init, step_date_init, step_date_exit, @@ -172,7 +170,6 @@ def test_nonhydro_predictor_step( interpolation_savepoint, savepoint_nonhydro_exit, experiment, - ndyn_substeps, at_initial_timestep, caplog, backend, @@ -1179,8 +1176,7 @@ def test_compute_perturbed_quantities_and_interpolation( ) assert test_utils.dallclose( - rho_at_cells_on_half_levels.asnumpy(), - rho_ic_ref.asnumpy(), + rho_at_cells_on_half_levels.asnumpy(),rho_ic_ref.asnumpy() ) assert test_utils.dallclose( exner_at_cells_on_half_levels.asnumpy()[:, nflatlev:], @@ -1203,9 +1199,9 @@ def test_compute_perturbed_quantities_and_interpolation( ) assert test_utils.dallclose( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[ - lb:, nflat_gradp: + lb:, : ], - z_dexner_dz_c_2_ref.asnumpy()[lb:, nflat_gradp:], + z_dexner_dz_c_2_ref.asnumpy()[lb:, :], rtol=5e-9, ) From 01d87c39c57fe3314d4c63c9f4c57c3f3248b3f5 Mon Sep 17 00:00:00 2001 From: Christos Kotsalos Date: Tue, 14 Oct 2025 11:27:56 +0300 Subject: [PATCH 056/108] Modifications in TestComputeThetaRhoPressureGradientAndUpdateVn --- ...ues_and_pressure_gradient_and_update_vn.py | 54 ++++++++++++++++--- 1 file changed, 48 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 350d380792..2b84f76ce2 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -13,7 +13,7 @@ import pytest import icon4py.model.common.type_alias as ta -import icon4py.model.testing.stencil_tests as test_helpers +from icon4py.model.testing import definitions, stencil_tests from icon4py.model.atmosphere.dycore.dycore_states import ( HorizontalPressureDiscretizationType, RhoThetaAdvectionType, @@ -124,7 +124,7 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( @pytest.mark.embedded_remap_error @pytest.mark.skip_value_error @pytest.mark.uses_as_offset -class TestComputeThetaRhoPressureGradientAndUpdateVn(test_helpers.StencilTest): +class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): PROGRAM = compute_theta_rho_face_values_and_pressure_gradient_and_update_vn OUTPUTS = ( "rho_at_edges_on_model_levels", @@ -132,6 +132,33 @@ class TestComputeThetaRhoPressureGradientAndUpdateVn(test_helpers.StencilTest): "horizontal_pressure_gradient", "next_vn", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "iau_wgt_dyn", + "is_iau_active", + "limited_area", + "start_edge_lateral_boundary", + "start_edge_lateral_boundary_level_7", + "start_edge_nudging_level_2", + "end_edge_nudging", + "end_edge_halo", + "nflatlev", + "nflat_gradp", + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "is_iau_active", + "limited_area", + "nflatlev", + "nflat_gradp", + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( @@ -473,8 +500,8 @@ def input_data(self, grid: base.Grid) -> dict: dtime = 0.9 iau_wgt_dyn = 1.0 - is_iau_active = True - limited_area = True + is_iau_active = False + limited_area = False edge_domain = h_grid.domain(dims.EdgeDim) start_edge_lateral_boundary = grid.end_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY)) @@ -484,8 +511,8 @@ def input_data(self, grid: base.Grid) -> dict: start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) end_edge_nudging = grid.end_index(edge_domain(h_grid.Zone.NUDGING)) end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) - nflatlev = 4 - nflat_gradp = 27 + nflatlev = 6 + nflat_gradp = 35 return dict( rho_at_edges_on_model_levels=rho_at_edges_on_model_levels, @@ -537,3 +564,18 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) + + +@pytest.mark.continuous_benchmarking +class TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking( + TestComputeThetaRhoPressureGradientAndUpdateVn +): + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict: + assert ( + grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 + ), "This test only works with the icon_benchmark grid." + base_data = TestComputeThetaRhoPressureGradientAndUpdateVn.input_data.__wrapped__( + self, grid + ) + return base_data From 62ff36cf6d44b647c615cb09492c1a1bfe075302 Mon Sep 17 00:00:00 2001 From: Christos Kotsalos Date: Tue, 14 Oct 2025 15:09:40 +0300 Subject: [PATCH 057/108] Modifications in TestComputeThetaRhoPressureGradientAndUpdateVn --- ...ace_values_and_pressure_gradient_and_update_vn.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 2b84f76ce2..d4bd4f2977 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -500,8 +500,8 @@ def input_data(self, grid: base.Grid) -> dict: dtime = 0.9 iau_wgt_dyn = 1.0 - is_iau_active = False - limited_area = False + is_iau_active = True + limited_area = True edge_domain = h_grid.domain(dims.EdgeDim) start_edge_lateral_boundary = grid.end_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY)) @@ -511,8 +511,8 @@ def input_data(self, grid: base.Grid) -> dict: start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) end_edge_nudging = grid.end_index(edge_domain(h_grid.Zone.NUDGING)) end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) - nflatlev = 6 - nflat_gradp = 35 + nflatlev = 4 + nflat_gradp = 27 return dict( rho_at_edges_on_model_levels=rho_at_edges_on_model_levels, @@ -578,4 +578,8 @@ def input_data(self, grid: base.Grid) -> dict: base_data = TestComputeThetaRhoPressureGradientAndUpdateVn.input_data.__wrapped__( self, grid ) + base_data["is_iau_active"] = False + base_data["limited_area"] = False + base_data["nflatlev"] = 6 + base_data["nflat_gradp"] = 35 return base_data From f49b74bf63d6b93cf0b9e84e06ac304a8f329fb5 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 09:56:46 +0200 Subject: [PATCH 058/108] Avoid warning: PytestCollectionWarning: cannot collect 'test_*' because it is not a function. --- model/testing/src/icon4py/model/testing/stencil_tests.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 14e8afa79c..e03808a051 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -249,9 +249,14 @@ def __init_subclass__(cls, **kwargs: Any) -> None: setattr(cls, f"{pytest_prefix}{cls.__name__}", test_and_benchmark) - # in case a test inherits from another test avoid running the tests of its parent + # in case a test inherits from another test overload the test function of its parent and skip it + # to avoid running the tests of its parent if cls.__base__ is not None and cls.__base__ != StencilTest: - setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", pytest.mark.skip(lambda: ())) + + def skipped_test() -> None: + pass + + setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", pytest.mark.skip(skipped_test)) # decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition From c755cd80aa7a5637dcba3fb7b37f705fa9c8506d Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 10:07:22 +0200 Subject: [PATCH 059/108] Fix formatting --- ...ues_and_pressure_gradient_and_update_vn.py | 56 +++++++++---------- 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index d4bd4f2977..1d21438ee0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -13,7 +13,6 @@ import pytest import icon4py.model.common.type_alias as ta -from icon4py.model.testing import definitions, stencil_tests from icon4py.model.atmosphere.dycore.dycore_states import ( HorizontalPressureDiscretizationType, RhoThetaAdvectionType, @@ -24,6 +23,7 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions, stencil_tests rhotheta_avd_type = RhoThetaAdvectionType() @@ -121,8 +121,6 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( return rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels -@pytest.mark.embedded_remap_error -@pytest.mark.skip_value_error @pytest.mark.uses_as_offset class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): PROGRAM = compute_theta_rho_face_values_and_pressure_gradient_and_update_vn @@ -133,32 +131,32 @@ class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): "next_vn", ) STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "iau_wgt_dyn", - "is_iau_active", - "limited_area", - "start_edge_lateral_boundary", - "start_edge_lateral_boundary_level_7", - "start_edge_nudging_level_2", - "end_edge_nudging", - "end_edge_halo", - "nflatlev", - "nflat_gradp", - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "is_iau_active", - "limited_area", - "nflatlev", - "nflat_gradp", - "vertical_start", - "vertical_end", - ), - } + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "iau_wgt_dyn", + "is_iau_active", + "limited_area", + "start_edge_lateral_boundary", + "start_edge_lateral_boundary_level_7", + "start_edge_nudging_level_2", + "end_edge_nudging", + "end_edge_halo", + "nflatlev", + "nflat_gradp", + "horizontal_start", + "horizontal_end", + "vertical_start", + "vertical_end", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "is_iau_active", + "limited_area", + "nflatlev", + "nflat_gradp", + "vertical_start", + "vertical_end", + ), + } @staticmethod def reference( From f1a627dab1260f8d690c03d048bbff366561bc10 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 12:01:18 +0200 Subject: [PATCH 060/108] Skip continuous benchmarking tests if grid is not icon_benchmark --- ...o_w_and_compute_horizontal_gradients_for_turbulence.py | 3 --- .../test_calculate_nabla2_and_smag_coefficients_for_vn.py | 3 --- .../test_apply_divergence_damping_and_update_vn.py | 3 --- ...est_compute_advection_in_vertical_momentum_equation.py | 8 ++++++-- ...orizontal_winds_and_ke_and_contravariant_correction.py | 3 --- ...test_compute_perturbed_quantities_and_interpolation.py | 3 --- ...rho_face_values_and_pressure_gradient_and_update_vn.py | 3 --- ...vertically_implicit_dycore_solver_at_corrector_step.py | 3 --- ...vertically_implicit_dycore_solver_at_predictor_step.py | 3 --- model/testing/src/icon4py/model/testing/pytest_hooks.py | 8 ++++++++ 10 files changed, 14 insertions(+), 26 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index 504df38f7f..a25b0e96c3 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -151,9 +151,6 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulenceContinuousB ): @pytest.fixture def input_data(self, grid: base.Grid) -> dict: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." # Use the parent class's fixture indirectly by calling its method, not the fixture itself base_data = ( TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence.input_data.__wrapped__( diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 212215ecb9..5d7a01bf01 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -215,9 +215,6 @@ class TestCalculateNabla2AndSmagCoefficientsForVnContinuousBenchmarking( ): @pytest.fixture def input_data(self, grid: base.Grid) -> dict: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." base_data = TestCalculateNabla2AndSmagCoefficientsForVn.input_data.__wrapped__(self, grid) edge_domain = h_grid.domain(dims.EdgeDim) horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 6f8ef30b57..acf924a10b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -282,9 +282,6 @@ class TestApplyDivergenceDampingAndUpdateVnContinuousBenchmarking( ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." # Use the parent class's fixture indirectly by calling its method, not the fixture itself base_data = TestApplyDivergenceDampingAndUpdateVn.input_data.__wrapped__( self, request, grid diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index 2a5a8da5f9..b56ae82550 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -19,7 +19,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import definitions, stencil_tests from .test_add_interpolated_horizontal_advection_of_w import ( add_interpolated_horizontal_advection_of_w_numpy, @@ -486,7 +486,11 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala class TestFusedVelocityAdvectionStencilVMomentumContinuousBenchmarking( TestFusedVelocityAdvectionStencilVMomentum ): - pass + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + base_data = TestFusedVelocityAdvectionStencilVMomentum.input_data.__wrapped__(self, grid) + base_data["end_index_of_damping_layer"] = 12 + return base_data @pytest.mark.embedded_remap_error diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index a9a339c2bb..a8ed6aab2e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -363,9 +363,6 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava def input_data( self, grid: base.Grid, request: pytest.FixtureRequest ) -> dict[str, gtx.Field | state_utils.ScalarType]: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." base_data = TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection.input_data.__wrapped__( self, grid, request ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 1f184e5ffa..8e61987b8a 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -474,9 +474,6 @@ class TestComputePerturbedQuantitiesAndInterpolationContinuousBenchmarking( ): @pytest.fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." base_data = TestComputePerturbedQuantitiesAndInterpolation.input_data.__wrapped__( self, grid ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 1d21438ee0..8903872bad 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -570,9 +570,6 @@ class TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking( ): @pytest.fixture def input_data(self, grid: base.Grid) -> dict: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." base_data = TestComputeThetaRhoPressureGradientAndUpdateVn.input_data.__wrapped__( self, grid ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index 0702ff5f8c..16cc63d8f1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -556,9 +556,6 @@ class TestVerticallyImplicitSolverAtCorrectorStepContinuousBenchmarking( def input_data( self, request: pytest.FixtureRequest, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType]: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." base_data = TestVerticallyImplicitSolverAtCorrectorStep.input_data.__wrapped__( self, request, grid ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index a2510c3e91..1b032bee50 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -536,9 +536,6 @@ class TestVerticallyImplicitSolverAtPredictorStepContinuousBenchmarking( def input_data( self, request: pytest.FixtureRequest, grid: base.Grid ) -> dict[str, gtx.Field | state_utils.ScalarType]: - assert ( - grid.id == definitions.GridUUIDs.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." base_data = TestVerticallyImplicitSolverAtPredictorStep.input_data.__wrapped__( self, request, grid ) diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index 14f4153f55..cc04f1796e 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -103,8 +103,16 @@ def pytest_addoption(parser: pytest.Parser): def pytest_collection_modifyitems(config, items): """Modify collected test items based on command line options.""" + test_grid = config.getoption("--grid") for item in items: if (marker := item.get_closest_marker("continuous_benchmarking")) is not None: + if test_grid != "icon_benchmark": + item.add_marker( + pytest.mark.skip( + reason="Continuous benchmarking tests only run with --grid icon_benchmark" + ) + ) + continue if not config.getoption("--benchmark-only"): item.add_marker( pytest.mark.benchmark_only( From 1e092ba916485e5e16325831817c8b06b9a26f17 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 12:29:19 +0200 Subject: [PATCH 061/108] Fixing tests --- .../test_benchmark_diffusion.py | 6 +----- ..._apply_divergence_damping_and_update_vn.py | 8 +++++--- ..._perturbed_quantities_and_interpolation.py | 19 +++++++------------ 3 files changed, 13 insertions(+), 20 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py index 60b9fe7a46..1a6d0fbfe9 100644 --- a/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py @@ -223,11 +223,7 @@ def test_run_diffusion_benchmark( @pytest.mark.continuous_benchmarking def test_run_diffusion_benchmark_continuous_benchmarking( - grid: definitions.GridDescription, backend: gtx_typing.Backend | None, benchmark: Any, ) -> None: - assert ( - grid == definitions.Grids.MCH_OPR_R19B08_DOMAIN01 - ), "This test only works with the icon_benchmark grid." - test_run_diffusion_benchmark(grid, backend, benchmark) + test_run_diffusion_benchmark(definitions.Grids.MCH_OPR_R19B08_DOMAIN01, backend, benchmark) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index acf924a10b..1ea5d411b5 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -173,7 +173,10 @@ def reference( return dict(next_vn=next_vn) - @pytest.fixture(params=[True, False]) + @pytest.fixture( + params=[{"limited_area": la} for la in [True, False]], + ids=lambda param: f"limited_area[{param['limited_area']}]", + ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_mask_for_3d_divdamp = data_alloc.random_field(grid, dims.EdgeDim) @@ -222,7 +225,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) ) - limited_area = request.param + limited_area = request.param["limited_area"] edge_domain = h_grid.domain(dims.EdgeDim) start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) @@ -282,7 +285,6 @@ class TestApplyDivergenceDampingAndUpdateVnContinuousBenchmarking( ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", ) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: - # Use the parent class's fixture indirectly by calling its method, not the fixture itself base_data = TestApplyDivergenceDampingAndUpdateVn.input_data.__wrapped__( self, request, grid ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 8e61987b8a..6e9401fd06 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -83,28 +83,23 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): STATIC_PARAMS = { stencil_tests.StandardStaticVariants.NONE: (), stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "limited_area", "igradp_method", + "nflatlev", + "nflat_gradp", + "start_cell_lateral_boundary", "start_cell_lateral_boundary_level_3", "start_cell_halo_level_2", - "end_cell_end", "end_cell_halo", "end_cell_halo_level_2", - "start_cell_lateral_boundary", - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "nflatlev", - "nflat_gradp", + "model_top", + "surface_level", ), stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "limited_area", "igradp_method", - "vertical_start", - "vertical_end", "nflatlev", "nflat_gradp", + "model_top", + "surface_level", ), } From 28d695acf560c47618a2087d681b419df87bbcd5 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 14:51:30 +0200 Subject: [PATCH 062/108] Fix TestFusedVelocityAdvectionStencilVMomentumAndContravariant --- .../test_compute_advection_in_vertical_momentum_equation.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index b56ae82550..22714168dc 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -502,6 +502,9 @@ class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.S "contravariant_corrected_w_at_cells_on_model_levels", "vertical_cfl", ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), # For now compile time variants triger error in gt4py + } @staticmethod def reference( @@ -629,7 +632,8 @@ def reference( ) @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]] + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], + ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]", ) def input_data( self, grid: base.Grid, request: pytest.FixtureRequest From 9cd760af923abe72b024de15095201ac0b4eac3b Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 17:14:41 +0200 Subject: [PATCH 063/108] Skip validation of TestComputeThetaRhoPressureGradientAndUpdateVn --- ...ues_and_pressure_gradient_and_update_vn.py | 3 ++- scripts/compare_icon_icon4py.py | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 8903872bad..7485b47be4 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -120,7 +120,8 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( return rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels - +@pytest.mark.embedded_remap_error +@pytest.mark.skip_value_error @pytest.mark.uses_as_offset class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): PROGRAM = compute_theta_rho_face_values_and_pressure_gradient_and_update_vn diff --git a/scripts/compare_icon_icon4py.py b/scripts/compare_icon_icon4py.py index 2474e63676..0754ec5305 100644 --- a/scripts/compare_icon_icon4py.py +++ b/scripts/compare_icon_icon4py.py @@ -48,6 +48,7 @@ "apply_diffusion_to_vn": None, "apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence": None, "apply_divergence_damping_and_update_vn": None, + "boundary_halo_cleanup": None, "calculate_diagnostic_quantities_for_turbulence": None, "calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools": None, "calculate_nabla2_and_smag_coefficients_for_vn": None, @@ -201,6 +202,29 @@ def load_gt4py_timers(filename: pathlib.Path, metric: str) -> tuple[dict, dict]: if len(metric_data) >= gt4py_unmatched_ncalls_threshold: unmatched_data[stencil] = metric_data + # Merge 'compute_hydrostatic_correction_term' stencil into 'compute_theta_rho_face_values_and_pressure_gradient_and_update_vn' + assert "compute_hydrostatic_correction_term" in unmatched_data + data["compute_theta_rho_face_values_and_pressure_gradient_and_update_vn"] = [ + a + b + for a, b in zip( + data["compute_theta_rho_face_values_and_pressure_gradient_and_update_vn"], + unmatched_data.pop("compute_hydrostatic_correction_term"), + strict=True, + ) + ] + + # Merge some stencils into 'boundary_halo_cleanup' + assert "boundary_halo_cleanup" not in data + data["boundary_halo_cleanup"] = [ + a + b + c + for a, b, c in zip( + unmatched_data.pop("compute_exner_from_rhotheta"), + unmatched_data.pop("compute_theta_and_exner"), + unmatched_data.pop("update_theta_v"), + strict=True, + ) + ] + diff = set(fortran_to_icon4py.keys()) - set(data.keys()) if len(diff) != 0: raise ValueError(f"Missing icon4py meas for these stencils: {diff}.") From ad3402060a18bda5afbd412d26063ca94720f4e3 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 17:29:04 +0200 Subject: [PATCH 064/108] Remove unrelated changes --- ...ues_and_pressure_gradient_and_update_vn.py | 1 + scripts/compare_icon_icon4py.py | 24 ------------------- 2 files changed, 1 insertion(+), 24 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 7485b47be4..c67b92cf18 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -120,6 +120,7 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( return rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels + @pytest.mark.embedded_remap_error @pytest.mark.skip_value_error @pytest.mark.uses_as_offset diff --git a/scripts/compare_icon_icon4py.py b/scripts/compare_icon_icon4py.py index 0754ec5305..2474e63676 100644 --- a/scripts/compare_icon_icon4py.py +++ b/scripts/compare_icon_icon4py.py @@ -48,7 +48,6 @@ "apply_diffusion_to_vn": None, "apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence": None, "apply_divergence_damping_and_update_vn": None, - "boundary_halo_cleanup": None, "calculate_diagnostic_quantities_for_turbulence": None, "calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools": None, "calculate_nabla2_and_smag_coefficients_for_vn": None, @@ -202,29 +201,6 @@ def load_gt4py_timers(filename: pathlib.Path, metric: str) -> tuple[dict, dict]: if len(metric_data) >= gt4py_unmatched_ncalls_threshold: unmatched_data[stencil] = metric_data - # Merge 'compute_hydrostatic_correction_term' stencil into 'compute_theta_rho_face_values_and_pressure_gradient_and_update_vn' - assert "compute_hydrostatic_correction_term" in unmatched_data - data["compute_theta_rho_face_values_and_pressure_gradient_and_update_vn"] = [ - a + b - for a, b in zip( - data["compute_theta_rho_face_values_and_pressure_gradient_and_update_vn"], - unmatched_data.pop("compute_hydrostatic_correction_term"), - strict=True, - ) - ] - - # Merge some stencils into 'boundary_halo_cleanup' - assert "boundary_halo_cleanup" not in data - data["boundary_halo_cleanup"] = [ - a + b + c - for a, b, c in zip( - unmatched_data.pop("compute_exner_from_rhotheta"), - unmatched_data.pop("compute_theta_and_exner"), - unmatched_data.pop("update_theta_v"), - strict=True, - ) - ] - diff = set(fortran_to_icon4py.keys()) - set(data.keys()) if len(diff) != 0: raise ValueError(f"Missing icon4py meas for these stencils: {diff}.") From efa6df477e49cd908df772ab3733cc8c10df33c8 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 18:11:00 +0200 Subject: [PATCH 065/108] Trying to fix test_TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection --- ...ived_horizontal_winds_and_ke_and_contravariant_correction.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index a8ed6aab2e..be12bc3ee1 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -314,7 +314,7 @@ def input_data( c_intp = data_alloc.random_field(grid, dims.VertexDim, dims.V2CDim) nlev = grid.num_levels - nflatlev = 13 + nflatlev = 11 skip_compute_predictor_vertical_advection = request.param[ "skip_compute_predictor_vertical_advection" From 217333fd6365e2902e4207d7879ec3cc7d20deda Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 15 Oct 2025 18:15:23 +0200 Subject: [PATCH 066/108] Remove unnecessary UUID --- model/testing/src/icon4py/model/testing/definitions.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index fa04c32d41..201c7d13cf 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -128,10 +128,6 @@ class Grids: ) -class GridUUIDs: - MCH_OPR_R19B08_DOMAIN01: Final = "01f00602-c07e-cd84-b894-bd17fffd2720" - - @dataclasses.dataclass class Experiment: name: str From 7cd8f8110fb1a71a153a0035f561fda34c1934d8 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 20 Oct 2025 09:34:26 +0200 Subject: [PATCH 067/108] pre-commit edit --- .../tests/dycore/integration_tests/test_solve_nonhydro.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index ddd42181e0..f190704037 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1175,9 +1175,7 @@ def test_compute_perturbed_quantities_and_interpolation( perturbed_exner_at_cells_on_model_levels.asnumpy(), exner_pr_ref.asnumpy() ) - assert test_utils.dallclose( - rho_at_cells_on_half_levels.asnumpy(),rho_ic_ref.asnumpy() - ) + assert test_utils.dallclose(rho_at_cells_on_half_levels.asnumpy(), rho_ic_ref.asnumpy()) assert test_utils.dallclose( exner_at_cells_on_half_levels.asnumpy()[:, nflatlev:], z_exner_ic_ref.asnumpy()[:, nflatlev:], @@ -1198,9 +1196,7 @@ def test_compute_perturbed_quantities_and_interpolation( rtol=5e-9, ) assert test_utils.dallclose( - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[ - lb:, : - ], + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, :], z_dexner_dz_c_2_ref.asnumpy()[lb:, :], rtol=5e-9, ) From 706335aaf815c4bc96ad7660d9710bcd0fef7fb5 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 20 Oct 2025 15:11:57 +0200 Subject: [PATCH 068/108] further edits --- .../compute_cell_diagnostics_for_dycore.py | 6 +- ...ge_diagnostics_for_dycore_and_update_vn.py | 44 +++++++------- .../integration_tests/test_solve_nonhydro.py | 58 ++++++++++--------- 3 files changed, 58 insertions(+), 50 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 5522c4ae11..cfd8622859 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -526,9 +526,11 @@ def _interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_accele time_averaged_perturbed_theta_v_vp, time_averaged_perturbed_theta_v_kup_vp = astype( (time_averaged_perturbed_theta_v, time_averaged_perturbed_theta_v_kup), vpfloat ) - perturbed_theta_v_at_cells_on_half_levels = ( + perturbed_theta_v_at_cells_on_half_levels = concat_where( + dims.KDim == 0, wgtfac_c * time_averaged_perturbed_theta_v_vp - + (vpfloat("1.0") - wgtfac_c) * time_averaged_perturbed_theta_v_kup_vp + + (vpfloat("1.0") - wgtfac_c) * time_averaged_perturbed_theta_v_kup_vp, + broadcast(0.0, (dims.CellDim, dims.KDim)), ) theta_v_at_cells_on_half_levels = ( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index 341de7eb15..ca8bfaa2a7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -96,36 +96,40 @@ def _compute_horizontal_pressure_gradient( nflat_gradp: gtx.int32, ) -> fa.EdgeKField[ta.wpfloat]: # Note: we only support `TAYLOR_HYDRO` + on_flatlevels = _compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, + ) + between_flat_and_flatgradp = _compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, + ddxn_z_full=ddxn_z_full, + c_lin_e=c_lin_e, + z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + ) + below_flatgradp = _compute_horizontal_gradient_of_exner_pressure_for_multiple_levels( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, + zdiff_gradp=zdiff_gradp, + ikoffset=ikoffset, + z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + z_dexner_dz_c_2=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + ) horizontal_pressure_gradient = apply_on_vertical_level( nflatlev, nflat_gradp, - on_flatlevels=_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates( - inv_dual_edge_length=inv_dual_edge_length, - z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, - ), - between_flat_and_flatgradp=_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates( - inv_dual_edge_length=inv_dual_edge_length, - z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, - ddxn_z_full=ddxn_z_full, - c_lin_e=c_lin_e, - z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ), - below_flatgradp=_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels( - inv_dual_edge_length=inv_dual_edge_length, - z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, - zdiff_gradp=zdiff_gradp, - ikoffset=ikoffset, - z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - z_dexner_dz_c_2=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ), + on_flatlevels, + between_flat_and_flatgradp, + below_flatgradp, ) - return _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( + horizontal_pressure_gradient_final = _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( ipeidx_dsl=ipeidx_dsl, pg_exdist=pg_exdist, z_hydro_corr=hydrostatic_correction_on_lowest_level, z_gradh_exner=horizontal_pressure_gradient, ) + return horizontal_pressure_gradient_final @gtx.field_operator diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index f190704037..b795fd2609 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -231,7 +231,7 @@ def test_nonhydro_predictor_step( cell_domain = h_grid.domain(dims.CellDim) edge_domain = h_grid.domain(dims.EdgeDim) - cell_start_lateral_boundary_level_2 = icon_grid.start_index( + cell_start_lateral_boundary_level_3 = icon_grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) ) @@ -247,90 +247,90 @@ def test_nonhydro_predictor_step( # stencils 2, 3 assert test_utils.dallclose( diagnostic_state_nh.perturbed_exner_at_cells_on_model_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.exner_pr().asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.exner_pr().asnumpy()[cell_start_lateral_boundary_level_3:, :], ) assert test_utils.dallclose( solve_nonhydro.temporal_extrapolation_of_perturbed_exner.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.z_exner_ex_pr().asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.z_exner_ex_pr().asnumpy()[cell_start_lateral_boundary_level_3:, :], ) # stencils 4,5 assert test_utils.dallclose( solve_nonhydro.exner_at_cells_on_half_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, nlev - 1 + cell_start_lateral_boundary_level_3:, nlev - 1 ], - sp_exit.z_exner_ic().asnumpy()[cell_start_lateral_boundary_level_2:, nlev - 1], + sp_exit.z_exner_ic().asnumpy()[cell_start_lateral_boundary_level_3:, nlev - 1], ) nflatlev = vertical_params.nflatlev assert test_utils.dallclose( solve_nonhydro.exner_at_cells_on_half_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, nflatlev : nlev - 1 + cell_start_lateral_boundary_level_3:, nflatlev : nlev - 1 ], - sp_exit.z_exner_ic().asnumpy()[cell_start_lateral_boundary_level_2:, nflatlev : nlev - 1], + sp_exit.z_exner_ic().asnumpy()[cell_start_lateral_boundary_level_3:, nflatlev : nlev - 1], rtol=1.0e-9, ) # stencil 6 assert test_utils.dallclose( solve_nonhydro.ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, nflatlev: + cell_start_lateral_boundary_level_3:, nflatlev: ], - sp_exit.z_dexner_dz_c(0).asnumpy()[cell_start_lateral_boundary_level_2:, nflatlev:], + sp_exit.z_dexner_dz_c(0).asnumpy()[cell_start_lateral_boundary_level_3:, nflatlev:], atol=5e-18, ) # stencils 7,8,9 assert test_utils.dallclose( diagnostic_state_nh.rho_at_cells_on_half_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.rho_ic().asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.rho_ic().asnumpy()[cell_start_lateral_boundary_level_3:, :], ) assert test_utils.dallclose( solve_nonhydro.pressure_buoyancy_acceleration_at_cells_on_half_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, 1: + cell_start_lateral_boundary_level_3:, 1: ], - sp_exit.z_th_ddz_exner_c().asnumpy()[cell_start_lateral_boundary_level_2:, 1:], + sp_exit.z_th_ddz_exner_c().asnumpy()[cell_start_lateral_boundary_level_3:, 1:], rtol=2.0e-12, ) # stencils 7,8,9, 11 assert test_utils.dallclose( solve_nonhydro.perturbed_theta_v_at_cells_on_half_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.z_theta_v_pr_ic().asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.z_theta_v_pr_ic().asnumpy()[cell_start_lateral_boundary_level_3:, :], ) assert test_utils.dallclose( diagnostic_state_nh.theta_v_at_cells_on_half_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.theta_v_ic().asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.theta_v_ic().asnumpy()[cell_start_lateral_boundary_level_3:, :], ) # stencils 7,8,9, 13 assert test_utils.dallclose( solve_nonhydro.perturbed_rho_at_cells_on_model_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.z_rth_pr(0).asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.z_rth_pr(0).asnumpy()[cell_start_lateral_boundary_level_3:, :], ) assert test_utils.dallclose( solve_nonhydro.perturbed_theta_v_at_cells_on_model_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, : + cell_start_lateral_boundary_level_3:, : ], - sp_exit.z_rth_pr(1).asnumpy()[cell_start_lateral_boundary_level_2:, :], + sp_exit.z_rth_pr(1).asnumpy()[cell_start_lateral_boundary_level_3:, :], ) # stencils 12 nflat_gradp = grid_savepoint.nflat_gradp() assert test_utils.dallclose( solve_nonhydro.d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[ - cell_start_lateral_boundary_level_2:, nflat_gradp: + cell_start_lateral_boundary_level_3:, nflat_gradp: ], - sp_exit.z_dexner_dz_c(1).asnumpy()[cell_start_lateral_boundary_level_2:, nflat_gradp:], + sp_exit.z_dexner_dz_c(1).asnumpy()[cell_start_lateral_boundary_level_3:, nflat_gradp:], atol=1e-22, ) @@ -354,7 +354,7 @@ def test_nonhydro_predictor_step( edge_start_nudging_level_2:, : ], sp_exit.z_gradh_exner().asnumpy()[edge_start_nudging_level_2:, :], - atol=1e-20, + atol=1e-7, # 20 ) prognostic_state_nnew = prognostic_states.next vn_new_reference = sp_exit.vn_new().asnumpy() @@ -1196,8 +1196,10 @@ def test_compute_perturbed_quantities_and_interpolation( rtol=5e-9, ) assert test_utils.dallclose( - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[lb:, :], - z_dexner_dz_c_2_ref.asnumpy()[lb:, :], + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels.asnumpy()[ + lb:, nflat_gradp: + ], + z_dexner_dz_c_2_ref.asnumpy()[lb:, nflat_gradp:], rtol=5e-9, ) From 7e57d0310ccb6ab45ecb7e25d0fc11a896c3ed22 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 20 Oct 2025 15:12:44 +0200 Subject: [PATCH 069/108] ran pre-commit --- ...ge_diagnostics_for_dycore_and_update_vn.py | 26 +++++++++++-------- .../integration_tests/test_solve_nonhydro.py | 2 +- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index ca8bfaa2a7..dc9cb86e60 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -100,12 +100,14 @@ def _compute_horizontal_pressure_gradient( inv_dual_edge_length=inv_dual_edge_length, z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, ) - between_flat_and_flatgradp = _compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates( - inv_dual_edge_length=inv_dual_edge_length, - z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, - ddxn_z_full=ddxn_z_full, - c_lin_e=c_lin_e, - z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + between_flat_and_flatgradp = ( + _compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, + ddxn_z_full=ddxn_z_full, + c_lin_e=c_lin_e, + z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + ) ) below_flatgradp = _compute_horizontal_gradient_of_exner_pressure_for_multiple_levels( inv_dual_edge_length=inv_dual_edge_length, @@ -123,11 +125,13 @@ def _compute_horizontal_pressure_gradient( below_flatgradp, ) - horizontal_pressure_gradient_final = _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( - ipeidx_dsl=ipeidx_dsl, - pg_exdist=pg_exdist, - z_hydro_corr=hydrostatic_correction_on_lowest_level, - z_gradh_exner=horizontal_pressure_gradient, + horizontal_pressure_gradient_final = ( + _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( + ipeidx_dsl=ipeidx_dsl, + pg_exdist=pg_exdist, + z_hydro_corr=hydrostatic_correction_on_lowest_level, + z_gradh_exner=horizontal_pressure_gradient, + ) ) return horizontal_pressure_gradient_final diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index b795fd2609..7f97ee29b7 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -354,7 +354,7 @@ def test_nonhydro_predictor_step( edge_start_nudging_level_2:, : ], sp_exit.z_gradh_exner().asnumpy()[edge_start_nudging_level_2:, :], - atol=1e-7, # 20 + atol=1e-7, # 20 ) prognostic_state_nnew = prognostic_states.next vn_new_reference = sp_exit.vn_new().asnumpy() From 8df62960e20907a14dc7e3d687b091fdb638e381 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 20 Oct 2025 16:39:12 +0200 Subject: [PATCH 070/108] small edit to stencil test --- .../stencils/compute_cell_diagnostics_for_dycore.py | 13 ------------- ...els_and_compute_temperature_vertical_gradient.py | 1 + 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index cfd8622859..888c580fa5 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -6,19 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - from typing import Final import gt4py.next as gtx diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py index b3fb185c76..b730ff8819 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py @@ -34,6 +34,7 @@ from icon4py.model.testing import stencil_tests +@pytest.mark.uses_concat_where class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration( stencil_tests.StencilTest ): From 05968bb89cd239ae6123ef1ab8c5af24a18f326d Mon Sep 17 00:00:00 2001 From: Christos Kotsalos Date: Thu, 30 Oct 2025 10:29:57 +0100 Subject: [PATCH 071/108] renaming: TestCBApplyDiffusionToVn to TestApplyDiffusionToVnContinuousBenchmarking --- .../tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index 024ad3f775..b85b6a74c3 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -166,5 +166,5 @@ def input_data(self, grid: base.Grid) -> dict: @pytest.mark.continuous_benchmarking -class TestCBApplyDiffusionToVn(TestApplyDiffusionToVn): +class TestApplyDiffusionToVnContinuousBenchmarking(TestApplyDiffusionToVn): pass From 36f54e8e42848e30c8dacfd8fdc7739f81670fda Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Wed, 5 Nov 2025 15:42:01 +0100 Subject: [PATCH 072/108] Fix issue with test_grid --- model/testing/src/icon4py/model/testing/pytest_hooks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index f3eaf5d3a2..6e1c2f7b40 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -106,7 +106,7 @@ def pytest_collection_modifyitems(config, items): test_grid = config.getoption("--grid") for item in items: if (marker := item.get_closest_marker("continuous_benchmarking")) is not None: - if not test_grid.startswith("icon_benchmark"): + if test_grid is not None and not test_grid.startswith("icon_benchmark"): item.add_marker( pytest.mark.skip( reason="Continuous benchmarking tests only run with --grid icon_benchmark_{regional,global}." From 94021c05da69dfc80fd6fbb49f87a85a2986c36d Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Thu, 6 Nov 2025 10:12:42 +0100 Subject: [PATCH 073/108] Fix data_allocations import in test_calculate_horizontal_gradients_for_turbulence.py --- ..._calculate_horizontal_gradients_for_turbulence.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py index 0c4629f7af..fa6659390b 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -14,7 +14,7 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest @@ -53,11 +53,11 @@ def reference( @pytest.fixture def input_data(self, grid): - w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - dwdx = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - dwdy = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + w = data_alloc.random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + geofac_grg_x = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + geofac_grg_y = data_alloc.random_field(grid, dims.CellDim, dims.C2E2CODim, dtype=wpfloat) + dwdx = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + dwdy = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) return dict( w=w, From 4407d5c1d5bcbbc946827e15b5f4eb7faadf3291 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Thu, 6 Nov 2025 10:12:57 +0100 Subject: [PATCH 074/108] Increase DaCe bencher time limit --- ci/benchmark_bencher.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index 120c144277..0e2efa29b5 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -27,6 +27,7 @@ include: ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false GT4PY_UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE: true PYTHONOPTIMIZE: 2 + SLURM_TIMELIMIT: '04:00:00' # DaCe compilation takes long time so we increase the time limit benchmark_bencher_aarch64: extends: [.benchmark_model_stencils, .test_template_aarch64] From 0e26652cc6f32c5d537fd1c03c84266455469ca7 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Thu, 6 Nov 2025 11:13:04 +0100 Subject: [PATCH 075/108] Only run none and compile_time_domain in bencher --- ci/benchmark_bencher.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index 0e2efa29b5..8210af7051 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -17,7 +17,7 @@ include: stage: benchmark script: - !reference [.bencher_setup_env, setup] - - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID + - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID -k "none or compile_time_domain" parallel: matrix: - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu] From 565e350fdd7746195d97f7fe271e45a8a3565b51 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Thu, 6 Nov 2025 16:46:45 +0100 Subject: [PATCH 076/108] Revert changes in slurm timeout --- ci/benchmark_bencher.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index 8210af7051..3721309509 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -27,7 +27,6 @@ include: ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false GT4PY_UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE: true PYTHONOPTIMIZE: 2 - SLURM_TIMELIMIT: '04:00:00' # DaCe compilation takes long time so we increase the time limit benchmark_bencher_aarch64: extends: [.benchmark_model_stencils, .test_template_aarch64] From 05a7fbe35b163b95c82b6b6b2f9ee17fa25faf96 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Thu, 6 Nov 2025 17:11:56 +0100 Subject: [PATCH 077/108] Allow TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking to run only for benchmarking --- ..._rho_face_values_and_pressure_gradient_and_update_vn.py | 3 ++- model/testing/src/icon4py/model/testing/filters.py | 7 +++++-- model/testing/src/icon4py/model/testing/test_utils.py | 7 +++++++ 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index c67b92cf18..ff2c267b22 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -567,6 +567,7 @@ def input_data(self, grid: base.Grid) -> dict: @pytest.mark.continuous_benchmarking +@pytest.mark.benchmark_only class TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking( TestComputeThetaRhoPressureGradientAndUpdateVn ): @@ -576,7 +577,7 @@ def input_data(self, grid: base.Grid) -> dict: self, grid ) base_data["is_iau_active"] = False - base_data["limited_area"] = False + base_data["limited_area"] = grid.limited_area base_data["nflatlev"] = 6 base_data["nflat_gradp"] = 35 return base_data diff --git a/model/testing/src/icon4py/model/testing/filters.py b/model/testing/src/icon4py/model/testing/filters.py index 06b2d22a56..d81f322fc5 100644 --- a/model/testing/src/icon4py/model/testing/filters.py +++ b/model/testing/src/icon4py/model/testing/filters.py @@ -78,8 +78,11 @@ class ItemFilter(NamedTuple): action=functools.partial(pytest.skip, "GTFN compilation is too slow for this test."), ), pytest.mark.skip_value_error.name: ItemFilter( - condition=lambda item: (grid := test_utils.get_fixture_value("grid", item)).limited_area - or grid.geometry_type == base.GeometryType.ICOSAHEDRON, + condition=lambda item: ( + (grid := test_utils.get_fixture_value("grid", item)).limited_area + or grid.geometry_type == base.GeometryType.ICOSAHEDRON + ) + and not test_utils.should_benchmark_only(item), action=functools.partial( pytest.skip, "Stencil does not support domain containing skip values. Consider shrinking domain.", diff --git a/model/testing/src/icon4py/model/testing/test_utils.py b/model/testing/src/icon4py/model/testing/test_utils.py index 0356291b5c..b450f6fbee 100644 --- a/model/testing/src/icon4py/model/testing/test_utils.py +++ b/model/testing/src/icon4py/model/testing/test_utils.py @@ -59,3 +59,10 @@ def is_dace(backend: gtx_typing.Backend | None) -> bool: def is_gtfn_backend(backend: gtx_typing.Backend | None) -> bool: return "gtfn" in backend.name if backend else False + + +def should_benchmark_only(item: pytest.Item) -> bool: + """Check if the test item is marked as benchmark_only.""" + return item.get_closest_marker("benchmark_only") is not None or item._request.getfixturevalue( # type: ignore[attr-defined] + "pytestconfig" + ).getoption("benchmark_only") From 6c3eccb83ab5f7288b8147ded5e2d4322eceb6ff Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Fri, 7 Nov 2025 12:12:28 +0100 Subject: [PATCH 078/108] Profile only compile_time_domain in bencher --- ci/benchmark_bencher.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index 3721309509..e63f3aca02 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -17,7 +17,7 @@ include: stage: benchmark script: - !reference [.bencher_setup_env, setup] - - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID -k "none or compile_time_domain" + - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID -k "compile_time_domain" parallel: matrix: - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu] From d1c6a9ebb4a4f8a3df863661897570ee15f80bc1 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Fri, 7 Nov 2025 17:49:38 +0100 Subject: [PATCH 079/108] Use proper values for TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking --- ..._face_values_and_pressure_gradient_and_update_vn.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index ff2c267b22..4b68543f0b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -295,7 +295,10 @@ def reference( perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, ), - (rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels), + ( + np.zeros_like(rho_at_edges_on_model_levels), + np.zeros_like(theta_v_at_edges_on_model_levels), + ), ) # Remaining computations at edge points @@ -578,6 +581,7 @@ def input_data(self, grid: base.Grid) -> dict: ) base_data["is_iau_active"] = False base_data["limited_area"] = grid.limited_area - base_data["nflatlev"] = 6 - base_data["nflat_gradp"] = 35 + base_data["nflatlev"] = 5 + base_data["nflat_gradp"] = 34 + base_data["start_edge_lateral_boundary"] = 0 return base_data From 4d6e2d0c6efcf35b74ccf067c0dff52654c910d9 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 10 Nov 2025 11:17:57 +0100 Subject: [PATCH 080/108] Fix ComputeThetaRhoPressureGradientAndUpdateVn stencil test --- ..._rho_face_values_and_pressure_gradient_and_update_vn.py | 2 -- model/testing/src/icon4py/model/testing/filters.py | 7 ++----- model/testing/src/icon4py/model/testing/test_utils.py | 7 ------- 3 files changed, 2 insertions(+), 14 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index 4b68543f0b..c034d7521d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -122,7 +122,6 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( @pytest.mark.embedded_remap_error -@pytest.mark.skip_value_error @pytest.mark.uses_as_offset class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): PROGRAM = compute_theta_rho_face_values_and_pressure_gradient_and_update_vn @@ -570,7 +569,6 @@ def input_data(self, grid: base.Grid) -> dict: @pytest.mark.continuous_benchmarking -@pytest.mark.benchmark_only class TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking( TestComputeThetaRhoPressureGradientAndUpdateVn ): diff --git a/model/testing/src/icon4py/model/testing/filters.py b/model/testing/src/icon4py/model/testing/filters.py index 4dead21f99..8886e1a2c2 100644 --- a/model/testing/src/icon4py/model/testing/filters.py +++ b/model/testing/src/icon4py/model/testing/filters.py @@ -70,11 +70,8 @@ class ItemFilter(NamedTuple): action=functools.partial(pytest.skip, "GTFN compilation is too slow for this test."), ), pytest.mark.skip_value_error.name: ItemFilter( - condition=lambda item: ( - (grid := test_utils.get_fixture_value("grid", item)).limited_area - or grid.geometry_type == base.GeometryType.ICOSAHEDRON - ) - and not test_utils.should_benchmark_only(item), + condition=lambda item: (grid := test_utils.get_fixture_value("grid", item)).limited_area + or grid.geometry_type == base.GeometryType.ICOSAHEDRON, action=functools.partial( pytest.skip, "Stencil does not support domain containing skip values. Consider shrinking domain.", diff --git a/model/testing/src/icon4py/model/testing/test_utils.py b/model/testing/src/icon4py/model/testing/test_utils.py index f447d0ad42..8d23824591 100644 --- a/model/testing/src/icon4py/model/testing/test_utils.py +++ b/model/testing/src/icon4py/model/testing/test_utils.py @@ -71,10 +71,3 @@ def is_dace(backend: gtx_typing.Backend | None) -> bool: def is_gtfn_backend(backend: gtx_typing.Backend | None) -> bool: return "gtfn" in backend.name if backend else False - - -def should_benchmark_only(item: pytest.Item) -> bool: - """Check if the test item is marked as benchmark_only.""" - return item.get_closest_marker("benchmark_only") is not None or item._request.getfixturevalue( # type: ignore[attr-defined] - "pytestconfig" - ).getoption("benchmark_only") From 6e60f81eec90d2fbdf29fbea4b501de563d6a27e Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 10 Nov 2025 11:29:41 +0100 Subject: [PATCH 081/108] Remove irrelevant comment --- model/testing/src/icon4py/model/testing/stencil_tests.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index db6e4d62fb..fe39659b92 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -125,9 +125,8 @@ def test_and_benchmark( metrics_data = gtx_metrics.sources key = next(iter(metrics_data)) compute_samples = metrics_data[key].metrics["compute"].samples - # emprically exclude first few iterations as warmup + # emprically exclude first few iterations run for warmup initial_program_iterations_to_skip = 2 - # Exclude first sample unless running in benchmark_only mode benchmark.extra_info["gtx_metrics"] = compute_samples[ initial_program_iterations_to_skip: ] From 7923021fd95d70e7c0887b8ba32d49898638cab0 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 10 Nov 2025 11:48:46 +0100 Subject: [PATCH 082/108] edits to re-do work --- .../compute_cell_diagnostics_for_dycore.py | 336 ++++++++++++------ ...ge_diagnostics_for_dycore_and_update_vn.py | 52 ++- 2 files changed, 247 insertions(+), 141 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 64df580d8c..1db3fc504f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -6,6 +6,19 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + from typing import Final import gt4py.next as gtx @@ -60,12 +73,7 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( @gtx.field_operator -def _combined_field_operator( - perturbed_rho_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], - wgtfacq_c: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], +def _compute_perturbed_quantities_and_interpolation( current_rho: fa.CellKField[ta.wpfloat], reference_rho_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], current_theta_v: fa.CellKField[ta.wpfloat], @@ -77,51 +85,24 @@ def _combined_field_operator( ddqz_z_half: fa.CellKField[ta.vpfloat], pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], + temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], - inv_ddqz_z_full: fa.CellKField[ta.vpfloat], - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], - d2dexdz2_fac1_mc: fa.CellKField[ta.vpfloat], - d2dexdz2_fac2_mc: fa.CellKField[ta.vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, - nflat_gradp: gtx.int32, - nlev: gtx.int32, ) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], ]: - temporal_extrapolation_of_perturbed_exner = concat_where( - dims.KDim >= nlev - 1, - _init_cell_kdim_field_with_zero_wp(), - temporal_extrapolation_of_perturbed_exner, - ) - - # _compute_perturbed_quantities_and_interpolation - exner_at_cells_on_half_levels = ( - concat_where( - (dims.KDim >= nlev - 1), - _interpolate_cell_field_to_half_levels_vp( - wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner - ), - exner_at_cells_on_half_levels, - ) - if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else exner_at_cells_on_half_levels - ) - exner_at_cells_on_half_levels = ( concat_where( - (maximum(1, nflatlev) <= dims.KDim) & (dims.KDim < nlev - 1), + maximum(1, nflatlev) <= dims.KDim, _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner ), @@ -134,19 +115,15 @@ def _combined_field_operator( ( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, - ) = concat_where( - dims.KDim < nlev - 1, - _compute_perturbation_of_rho_and_theta( - current_rho, - reference_rho_at_cells_on_model_levels, - current_theta_v, - reference_theta_at_cells_on_model_levels, - ), - (perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), + ) = _compute_perturbation_of_rho_and_theta( + current_rho, + reference_rho_at_cells_on_model_levels, + current_theta_v, + reference_theta_at_cells_on_model_levels, ) rho_at_cells_on_half_levels = concat_where( - (dims.KDim >= 1) & (dims.KDim < nlev - 1), + dims.KDim >= 1, _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho), rho_at_cells_on_half_levels, ) @@ -154,7 +131,7 @@ def _combined_field_operator( wgtfac_c_wp = astype(wgtfac_c, wpfloat) perturbed_theta_v_at_cells_on_half_levels = concat_where( - (dims.KDim >= 1) & (dims.KDim < nlev - 1), + dims.KDim >= 1, _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels ), @@ -162,7 +139,7 @@ def _combined_field_operator( ) theta_v_at_cells_on_half_levels = concat_where( - (dims.KDim >= 1) & (dims.KDim < nlev - 1), + dims.KDim >= 1, _interpolate_cell_field_to_half_levels_wp( wgtfac_c=wgtfac_c_wp, interpolant=current_theta_v ), @@ -172,7 +149,7 @@ def _combined_field_operator( ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) pressure_buoyancy_acceleration_at_cells_on_half_levels = concat_where( - (dims.KDim >= 1) & (dims.KDim < nlev - 1), + dims.KDim >= 1, _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( exner_w_explicit_weight_parameter, theta_v_at_cells_on_half_levels, @@ -184,32 +161,158 @@ def _combined_field_operator( pressure_buoyancy_acceleration_at_cells_on_half_levels, ) - perturbed_theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= nlev - 1, - _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels - ), + return ( + perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels, + perturbed_exner_at_cells_on_model_levels, + rho_at_cells_on_half_levels, + exner_at_cells_on_half_levels, perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + pressure_buoyancy_acceleration_at_cells_on_half_levels, ) - theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= nlev - 1, - reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, + +@gtx.field_operator +def _surface_computations( + wgtfacq_c: fa.CellKField[ta.wpfloat], + exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], + time_extrapolation_parameter_for_exner: fa.CellKField[ta.vpfloat], + current_exner: fa.CellKField[ta.vpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], + perturbed_rho_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + igradp_method: gtx.int32, + surface_level: gtx.int32, +) -> tuple[ + fa.CellKField[ta.vpfloat], + fa.CellKField[ta.vpfloat], + fa.CellKField[ta.vpfloat], + fa.CellKField[ta.vpfloat], + fa.CellKField[ta.vpfloat], + fa.CellKField[ta.vpfloat], +]: + (perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels) = ( + concat_where( + dims.KDim < surface_level - 1, + _init_two_cell_kdim_fields_with_zero_vp(), + (perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), + ) + ) + + (temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels) = ( + concat_where( + dims.KDim < surface_level - 1, + _extrapolate_temporally_exner_pressure( + exner_exfac=time_extrapolation_parameter_for_exner, + exner=current_exner, + exner_ref_mc=reference_exner_at_cells_on_model_levels, + exner_pr=perturbed_exner_at_cells_on_model_levels, + ), + (temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels), + ) + ) + + temporal_extrapolation_of_perturbed_exner = concat_where( + dims.KDim == surface_level - 1, + _init_cell_kdim_field_with_zero_wp(), + temporal_extrapolation_of_perturbed_exner, ) exner_at_cells_on_half_levels = concat_where( - dims.KDim >= nlev - 1, + dims.KDim == surface_level - 1, _interpolate_to_surface( wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ), + ) + if igradp_method == horzpres_discr_type.TAYLOR_HYDRO + else exner_at_cells_on_half_levels, exner_at_cells_on_half_levels, ) - # _compute_first_and_second_vertical_derivative_of_exner + return ( + temporal_extrapolation_of_perturbed_exner, + exner_at_cells_on_half_levels, + temporal_extrapolation_of_perturbed_exner, + perturbed_exner_at_cells_on_model_levels, + perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels, + ) + + +@gtx.field_operator +def _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], +) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: + perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels + ) + theta_v_at_cells_on_half_levels = ( + reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels + ) + + exner_at_cells_on_half_levels = _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ) + + return ( + perturbed_theta_v_at_cells_on_half_levels, + astype(theta_v_at_cells_on_half_levels, wpfloat), + exner_at_cells_on_half_levels, + ) + + +@gtx.field_operator +def _compute_first_and_second_vertical_derivative_of_exner( + exner_at_cells_on_half_levels: fa.CellKField[vpfloat], + inv_ddqz_z_full: fa.CellKField[vpfloat], + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], + d2dexdz2_fac1_mc: fa.CellKField[vpfloat], + d2dexdz2_fac2_mc: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], + igradp_method: gtx.int32, + nflatlev: gtx.int32, + nflat_gradp: gtx.int32, + surface_level: gtx.int32, +) -> tuple[ + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], +]: + ( + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + ) = concat_where( + dims.KDim == surface_level - 1, + _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + wgtfacq_c=wgtfacq_c, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + ), + ( + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + ), + ) + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflatlev <= dims.KDim) & (dims.KDim < nlev - 1), + (nflatlev <= dims.KDim) & (dims.KDim < surface_level - 1), _compute_first_vertical_derivative_at_cells( exner_at_cells_on_half_levels, inv_ddqz_z_full ), @@ -221,7 +324,7 @@ def _combined_field_operator( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflat_gradp <= dims.KDim) & (dims.KDim < nlev - 1), + (nflat_gradp <= dims.KDim) & (dims.KDim < surface_level - 1), -vpfloat("0.5") * ( ( @@ -238,16 +341,11 @@ def _combined_field_operator( ) return ( - perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels, - perturbed_exner_at_cells_on_model_levels, - rho_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - pressure_buoyancy_acceleration_at_cells_on_half_levels, ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, ) @@ -349,36 +447,33 @@ def compute_perturbed_quantities_and_interpolation( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_level - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels """ - - # Init perturbed quantities to zero on lateral boundary - _init_two_cell_kdim_fields_with_zero_vp( - out=(perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), - domain={ - dims.CellDim: (start_cell_lateral_boundary, start_cell_lateral_boundary_level_3), - dims.KDim: (model_top, surface_level - 1), - }, - ) - - # Compute temporal extrapolation of perturbed exner, needs to be output for future program - - _extrapolate_temporally_exner_pressure( - exner_exfac=time_extrapolation_parameter_for_exner, - exner=current_exner, - exner_ref_mc=reference_exner_at_cells_on_model_levels, - exner_pr=perturbed_exner_at_cells_on_model_levels, - out=(temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels), + _surface_computations( + wgtfacq_c=wgtfacq_c, + exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, + time_extrapolation_parameter_for_exner=time_extrapolation_parameter_for_exner, + current_exner=current_exner, + reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, + perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + igradp_method=igradp_method, + surface_level=surface_level, + out=( + temporal_extrapolation_of_perturbed_exner, + exner_at_cells_on_half_levels, + temporal_extrapolation_of_perturbed_exner, + perturbed_exner_at_cells_on_model_levels, + perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels, + ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), + dims.KDim: (model_top, surface_level), }, ) - _combined_field_operator( - perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, + _compute_perturbed_quantities_and_interpolation( current_rho=current_rho, reference_rho_at_cells_on_model_levels=reference_rho_at_cells_on_model_levels, current_theta_v=current_theta_v, @@ -390,17 +485,11 @@ def compute_perturbed_quantities_and_interpolation( ddqz_z_half=ddqz_z_half, pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, + exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, - inv_ddqz_z_full=inv_ddqz_z_full, - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, - d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - nlev=surface_level, out=( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -410,8 +499,36 @@ def compute_perturbed_quantities_and_interpolation( perturbed_theta_v_at_cells_on_half_levels, theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, + ), + domain={ + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, + ) + + _compute_first_and_second_vertical_derivative_of_exner( + exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, + inv_ddqz_z_full=inv_ddqz_z_full, + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + perturbed_theta_v_at_cells_on_half_levels=perturbed_theta_v_at_cells_on_half_levels, + d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, + d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + wgtfacq_c=wgtfacq_c, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, + igradp_method=igradp_method, + nflatlev=nflatlev, + nflat_gradp=nflat_gradp, + surface_level=surface_level, + out=( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), @@ -419,7 +536,6 @@ def compute_perturbed_quantities_and_interpolation( }, ) - # Init perturbed quantities to zero on 2nd halo layer _compute_perturbation_of_rho_and_theta( rho=current_rho, rho_ref_mc=reference_rho_at_cells_on_model_levels, @@ -513,11 +629,9 @@ def _interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_accele time_averaged_perturbed_theta_v_vp, time_averaged_perturbed_theta_v_kup_vp = astype( (time_averaged_perturbed_theta_v, time_averaged_perturbed_theta_v_kup), vpfloat ) - perturbed_theta_v_at_cells_on_half_levels = concat_where( - dims.KDim == 0, + perturbed_theta_v_at_cells_on_half_levels = ( wgtfac_c * time_averaged_perturbed_theta_v_vp - + (vpfloat("1.0") - wgtfac_c) * time_averaged_perturbed_theta_v_kup_vp, - broadcast(0.0, (dims.CellDim, dims.KDim)), + + (vpfloat("1.0") - wgtfac_c) * time_averaged_perturbed_theta_v_kup_vp ) theta_v_at_cells_on_half_levels = ( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index dc9cb86e60..341de7eb15 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -96,44 +96,36 @@ def _compute_horizontal_pressure_gradient( nflat_gradp: gtx.int32, ) -> fa.EdgeKField[ta.wpfloat]: # Note: we only support `TAYLOR_HYDRO` - on_flatlevels = _compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates( - inv_dual_edge_length=inv_dual_edge_length, - z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, - ) - between_flat_and_flatgradp = ( - _compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates( + horizontal_pressure_gradient = apply_on_vertical_level( + nflatlev, + nflat_gradp, + on_flatlevels=_compute_horizontal_gradient_of_exner_pressure_for_flat_coordinates( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, + ), + between_flat_and_flatgradp=_compute_horizontal_gradient_of_exner_pressure_for_nonflat_coordinates( inv_dual_edge_length=inv_dual_edge_length, z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, ddxn_z_full=ddxn_z_full, c_lin_e=c_lin_e, z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ) - ) - below_flatgradp = _compute_horizontal_gradient_of_exner_pressure_for_multiple_levels( - inv_dual_edge_length=inv_dual_edge_length, - z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, - zdiff_gradp=zdiff_gradp, - ikoffset=ikoffset, - z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - z_dexner_dz_c_2=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ) - horizontal_pressure_gradient = apply_on_vertical_level( - nflatlev, - nflat_gradp, - on_flatlevels, - between_flat_and_flatgradp, - below_flatgradp, + ), + below_flatgradp=_compute_horizontal_gradient_of_exner_pressure_for_multiple_levels( + inv_dual_edge_length=inv_dual_edge_length, + z_exner_ex_pr=temporal_extrapolation_of_perturbed_exner, + zdiff_gradp=zdiff_gradp, + ikoffset=ikoffset, + z_dexner_dz_c_1=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + z_dexner_dz_c_2=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + ), ) - horizontal_pressure_gradient_final = ( - _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( - ipeidx_dsl=ipeidx_dsl, - pg_exdist=pg_exdist, - z_hydro_corr=hydrostatic_correction_on_lowest_level, - z_gradh_exner=horizontal_pressure_gradient, - ) + return _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( + ipeidx_dsl=ipeidx_dsl, + pg_exdist=pg_exdist, + z_hydro_corr=hydrostatic_correction_on_lowest_level, + z_gradh_exner=horizontal_pressure_gradient, ) - return horizontal_pressure_gradient_final @gtx.field_operator From edd7e4f494534d92c1de72bacfc70723c9da19f1 Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 10 Nov 2025 12:59:14 +0100 Subject: [PATCH 083/108] Fix TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence --- ...ute_horizontal_gradients_for_turbulence.py | 59 ++++++++++--------- 1 file changed, 32 insertions(+), 27 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index a25b0e96c3..b1a90e54d5 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -69,12 +69,21 @@ def reference( nrdmax, interior_idx, halo_idx, + horizontal_start, + horizontal_end, + vertical_start, + vertical_end, **kwargs, ) -> dict: k = np.arange(w_old.shape[1]) cell = np.arange(w_old.shape[0]) reshaped_k = k[np.newaxis, :] reshaped_cell = cell[:, np.newaxis] + out_w, out_dwdx, out_dwdy = ( + np.zeros_like(w_old), + dwdx.copy(), + dwdy.copy(), + ) # create output arrays to update only the necessary slices if type_shear == 2: dwdx, dwdy = np.where( reshaped_k > 0, @@ -102,15 +111,32 @@ def reference( apply_nabla2_to_w_in_upper_damping_layer_numpy(w, diff_multfac_n2w, area, z_nabla2_c), w, ) - return dict(w=w, dwdx=dwdx, dwdy=dwdy) + subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) + out_w[subset] = w[subset] + out_dwdx[subset] = dwdx[subset] + out_dwdy[subset] = dwdy[subset] + return dict(w=out_w, dwdx=out_dwdx, dwdy=out_dwdy) @pytest.fixture def input_data(self, grid: base.Grid) -> dict: nrdmax = 13 - interior_idx = 1 - halo_idx = 5 + cell_domain = h_grid.domain(dims.CellDim) + interior_idx = grid.start_index(cell_domain(h_grid.Zone.INTERIOR)) # 0 for simple grid + halo_idx = grid.end_index( + cell_domain(h_grid.Zone.LOCAL) + ) # same as horizontal_end for simple grid type_shear = 2 + def _get_start_index_for_w_diffusion() -> int32: + return ( + grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + if grid.limited_area + else grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) + ) + + horizontal_start = _get_start_index_for_w_diffusion() + horizontal_end = grid.end_index(cell_domain(h_grid.Zone.HALO)) + geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim) geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim) diff_multfac_n2w = random_field(grid, dims.KDim) @@ -138,8 +164,8 @@ def input_data(self, grid: base.Grid) -> dict: w=w, dwdx=dwdx, dwdy=dwdy, - horizontal_start=0, - horizontal_end=grid.num_cells, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, vertical_start=0, vertical_end=grid.num_levels, ) @@ -149,25 +175,4 @@ def input_data(self, grid: base.Grid) -> dict: class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulenceContinuousBenchmarking( TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence ): - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict: - # Use the parent class's fixture indirectly by calling its method, not the fixture itself - base_data = ( - TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence.input_data.__wrapped__( - self, grid - ) - ) - cell_domain = h_grid.domain(dims.CellDim) - base_data["interior_idx"] = grid.start_index(cell_domain(h_grid.Zone.INTERIOR)) - base_data["halo_idx"] = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) - - def _get_start_index_for_w_diffusion() -> int32: - return ( - grid.start_index(cell_domain(h_grid.Zone.NUDGING)) - if grid.limited_area - else grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) - ) - - base_data["horizontal_start"] = _get_start_index_for_w_diffusion() - base_data["horizontal_end"] = grid.end_index(cell_domain(h_grid.Zone.HALO)) - return base_data + pass From 4239768afa1fd0a29ded0ae7e466cddadf8f5f3d Mon Sep 17 00:00:00 2001 From: Ioannis Magkanaris Date: Mon, 10 Nov 2025 15:38:50 +0100 Subject: [PATCH 084/108] Use the test functions of StencilTest only and replace static_variant fixture if necessary --- .../icon4py/model/testing/stencil_tests.py | 33 +++++++++++-------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index fe39659b92..8610bfa3ed 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -249,23 +249,30 @@ def __init_subclass__(cls, **kwargs: Any) -> None: pytest_prefix = "test_" - setattr(cls, f"{pytest_prefix}{cls.__name__}", test_and_benchmark) - - # in case a test inherits from another test overload the test function of its parent and skip it - # to avoid running the tests of its parent - if cls.__base__ is not None and cls.__base__ != StencilTest: - - def skipped_test() -> None: - pass - - setattr(cls, f"{pytest_prefix}{cls.__base__.__name__}", pytest.mark.skip(skipped_test)) + # only add `test_and_benchmark` to direct subclasses of `StencilTest`. Inherited tests can use the same function. + if cls.__base__ == StencilTest: + setattr(cls, f"{pytest_prefix}{cls.__name__}", test_and_benchmark) # decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition - # Check if cls.static_variant is already a pytest fixture to allow inheriting from other StencilTests + # Check if cls.static_variant has changes in inherited subclasses and update it accordingly if hasattr(cls.static_variant, "_pytestfixturefunction"): - # Already a fixture, do nothing - pass + base_static = getattr(cls.__base__, "STATIC_PARAMS", None) + if base_static != cls.STATIC_PARAMS: + if cls.STATIC_PARAMS is None: + cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] + else: + # copy of `static_variant` + def _new_static_variant(request: pytest.FixtureRequest) -> Sequence[str]: + _, variant = request.param + return () if variant is None else variant + + # replace with new parametrized fixture + cls.static_variant = staticmethod( # type: ignore[method-assign] + pytest.fixture( + params=cls.STATIC_PARAMS.items(), scope="class", ids=lambda p: p[0] + )(_new_static_variant) + ) elif cls.STATIC_PARAMS is None: # not parametrized, return an empty tuple cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] # we override with a non-parametrized function From a7fe7b9cfea321b51ab7e175047d2c9cc8077b5b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 12 Nov 2025 12:32:17 +0100 Subject: [PATCH 085/108] small edit to fix backend failures --- .../dycore/stencils/extrapolate_temporally_exner_pressure.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py index 581d60bf02..b278f827cf 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py @@ -24,8 +24,8 @@ def _extrapolate_temporally_exner_pressure( z_exner_ex_pr_wp = (wpfloat("1.0") + exner_exfac_wp) * ( exner - exner_ref_mc_wp - ) - exner_exfac_wp * exner_pr - exner_pr_wp = exner - exner_ref_mc_wp + ) - exner_exfac * exner_pr + exner_pr_wp = exner - exner_ref_mc return astype(z_exner_ex_pr_wp, vpfloat), exner_pr_wp From e6d04ee52254eedb5d884d77c10ee95c45eabd3a Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 12 Nov 2025 15:55:06 +0100 Subject: [PATCH 086/108] edit to nflat_gradp value in test --- .../test_compute_perturbed_quantities_and_interpolation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 6e9401fd06..5f55af9b57 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -420,7 +420,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala end_cell_halo_level_2 = grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) nflatlev = 4 - nflat_gradp = 27 + nflat_gradp = 7 return dict( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, From 07d74e8c29c722f128a84db9b6752734d152ec6f Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:07:54 +0100 Subject: [PATCH 087/108] restored files --- ci/benchmark_bencher.yml | 4 +- ci/docker/base.Dockerfile | 4 - .../test_apply_diffusion_to_vn.py | 23 +--- ...ute_horizontal_gradients_for_turbulence.py | 72 ++---------- ...ate_nabla2_and_smag_coefficients_for_vn.py | 35 +----- .../stencil_tests/test_calculate_nabla4.py | 8 +- ..._apply_divergence_damping_and_update_vn.py | 67 +---------- ...vection_in_horizontal_momentum_equation.py | 26 +---- ...advection_in_vertical_momentum_equation.py | 41 +------ ...nds_and_ke_and_contravariant_correction.py | 48 +------- ...est_compute_hydrostatic_correction_term.py | 36 +----- ..._perturbed_quantities_and_interpolation.py | 40 +------ ...ues_and_pressure_gradient_and_update_vn.py | 54 +-------- .../test_init_cell_kdim_field_with_zero_wp.py | 20 +--- ...d_compute_temperature_vertical_gradient.py | 21 ---- .../test_update_mass_flux_weighted.py | 20 +--- ...mplicit_dycore_solver_at_corrector_step.py | 84 +------------- ...mplicit_dycore_solver_at_predictor_step.py | 64 +---------- ...est_mo_intp_rbf_rbf_vec_interpol_vertex.py | 21 +--- .../src/icon4py/model/testing/pytest_hooks.py | 70 ------------ .../icon4py/model/testing/stencil_tests.py | 108 +++++------------- 21 files changed, 76 insertions(+), 790 deletions(-) diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index e63f3aca02..80b46a89dd 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -17,11 +17,11 @@ include: stage: benchmark script: - !reference [.bencher_setup_env, setup] - - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID -k "compile_time_domain" + - nox -s __bencher_feature_branch_CI-3.10 -- --backend=$BACKEND --grid=$GRID parallel: matrix: - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu] - GRID: [icon_benchmark_regional, icon_benchmark_global] + GRID: [icon_benchmark] variables: ICON4PY_ENABLE_GRID_DOWNLOAD: false ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false diff --git a/ci/docker/base.Dockerfile b/ci/docker/base.Dockerfile index 1b476e8bf3..099ab1d06c 100644 --- a/ci/docker/base.Dockerfile +++ b/ci/docker/base.Dockerfile @@ -37,10 +37,6 @@ ENV PATH="/root/.cargo/bin:${PATH}" RUN rustc --version && which rustc && cargo --version && which cargo # Install Bencher for performance monitoring -# Update the following comment to trigger a rebuild to update the CLI: -# last update: 2025-09-05 -# This is necessary because the cloud version and the CLI version have to match -# but obviously, version changes do not register in the Dockerfile hash. RUN curl --proto '=https' --tlsv1.2 -sSfL https://bencher.dev/download/install-cli.sh | sh RUN bencher --version && which bencher diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py index b85b6a74c3..946f1132a5 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_vn.py @@ -15,7 +15,7 @@ from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing.stencil_tests import StencilTest from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy from .test_apply_nabla2_and_nabla4_to_vn import apply_nabla2_and_nabla4_to_vn_numpy @@ -29,22 +29,6 @@ class TestApplyDiffusionToVn(StencilTest): PROGRAM = apply_diffusion_to_vn OUTPUTS = ("vn",) - STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "start_2nd_nudge_line_idx_e", - "vertical_start", - "vertical_end", - "limited_area", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "limited_area", - ), - } @staticmethod def reference( @@ -163,8 +147,3 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) - - -@pytest.mark.continuous_benchmarking -class TestApplyDiffusionToVnContinuousBenchmarking(TestApplyDiffusionToVn): - pass diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index b1a90e54d5..45d054a47f 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -8,16 +8,14 @@ import gt4py.next as gtx import numpy as np import pytest -from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence import ( apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base, horizontal as h_grid +from icon4py.model.common.grid import base from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing import definitions -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing.stencil_tests import StencilTest from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy from .test_apply_nabla2_to_w_in_upper_damping_layer import ( @@ -33,25 +31,6 @@ class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTest): PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence OUTPUTS = ("w", "dwdx", "dwdy") - STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "halo_idx", - "interior_idx", - "vertical_start", - "vertical_end", - "nrdmax", - "type_shear", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "nrdmax", - "type_shear", - ), - } @staticmethod def reference( @@ -69,21 +48,12 @@ def reference( nrdmax, interior_idx, halo_idx, - horizontal_start, - horizontal_end, - vertical_start, - vertical_end, **kwargs, ) -> dict: k = np.arange(w_old.shape[1]) cell = np.arange(w_old.shape[0]) reshaped_k = k[np.newaxis, :] reshaped_cell = cell[:, np.newaxis] - out_w, out_dwdx, out_dwdy = ( - np.zeros_like(w_old), - dwdx.copy(), - dwdy.copy(), - ) # create output arrays to update only the necessary slices if type_shear == 2: dwdx, dwdy = np.where( reshaped_k > 0, @@ -111,32 +81,15 @@ def reference( apply_nabla2_to_w_in_upper_damping_layer_numpy(w, diff_multfac_n2w, area, z_nabla2_c), w, ) - subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) - out_w[subset] = w[subset] - out_dwdx[subset] = dwdx[subset] - out_dwdy[subset] = dwdy[subset] - return dict(w=out_w, dwdx=out_dwdx, dwdy=out_dwdy) + return dict(w=w, dwdx=dwdx, dwdy=dwdy) @pytest.fixture def input_data(self, grid: base.Grid) -> dict: nrdmax = 13 - cell_domain = h_grid.domain(dims.CellDim) - interior_idx = grid.start_index(cell_domain(h_grid.Zone.INTERIOR)) # 0 for simple grid - halo_idx = grid.end_index( - cell_domain(h_grid.Zone.LOCAL) - ) # same as horizontal_end for simple grid + interior_idx = 1 + halo_idx = 5 type_shear = 2 - def _get_start_index_for_w_diffusion() -> int32: - return ( - grid.start_index(cell_domain(h_grid.Zone.NUDGING)) - if grid.limited_area - else grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) - ) - - horizontal_start = _get_start_index_for_w_diffusion() - horizontal_end = grid.end_index(cell_domain(h_grid.Zone.HALO)) - geofac_grg_x = random_field(grid, dims.CellDim, dims.C2E2CODim) geofac_grg_y = random_field(grid, dims.CellDim, dims.C2E2CODim) diff_multfac_n2w = random_field(grid, dims.KDim) @@ -146,8 +99,8 @@ def _get_start_index_for_w_diffusion() -> int32: diff_multfac_w = 5.0 w = zero_field(grid, dims.CellDim, dims.KDim) - dwdx = random_field(grid, dims.CellDim, dims.KDim) - dwdy = random_field(grid, dims.CellDim, dims.KDim) + dwdx = zero_field(grid, dims.CellDim, dims.KDim) + dwdy = zero_field(grid, dims.CellDim, dims.KDim) return dict( area=area, @@ -164,15 +117,8 @@ def _get_start_index_for_w_diffusion() -> int32: w=w, dwdx=dwdx, dwdy=dwdy, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, + horizontal_start=0, + horizontal_end=grid.num_cells, vertical_start=0, vertical_end=grid.num_levels, ) - - -@pytest.mark.continuous_benchmarking -class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulenceContinuousBenchmarking( - TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence -): - pass diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py index 5d7a01bf01..83d006abb2 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py @@ -14,27 +14,14 @@ calculate_nabla2_and_smag_coefficients_for_vn, ) from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests +@pytest.mark.skip_value_error class TestCalculateNabla2AndSmagCoefficientsForVn(stencil_tests.StencilTest): PROGRAM = calculate_nabla2_and_smag_coefficients_for_vn OUTPUTS = ("kh_smag_e", "kh_smag_ec", "z_nabla2_e") - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -157,7 +144,7 @@ def reference( return dict(kh_smag_e=kh_smag_e, kh_smag_ec=kh_smag_ec, z_nabla2_e=z_nabla2_e) @pytest.fixture - def input_data(self, grid: base.Grid) -> dict: + def input_data(self, grid): u_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) v_vert = data_alloc.random_field(grid, dims.VertexDim, dims.KDim, dtype=ta.vpfloat) smag_offset = ta.vpfloat("9.0") @@ -207,19 +194,3 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) - - -@pytest.mark.continuous_benchmarking -class TestCalculateNabla2AndSmagCoefficientsForVnContinuousBenchmarking( - TestCalculateNabla2AndSmagCoefficientsForVn -): - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict: - base_data = TestCalculateNabla2AndSmagCoefficientsForVn.input_data.__wrapped__(self, grid) - edge_domain = h_grid.domain(dims.EdgeDim) - horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) - horizontal_end = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) - assert horizontal_start < horizontal_end - base_data["horizontal_start"] = 0 - base_data["horizontal_end"] = grid.num_edges - return base_data diff --git a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py index 064c59e23f..305cf3f319 100644 --- a/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion/stencil_tests/test_calculate_nabla4.py @@ -55,11 +55,12 @@ def calculate_nabla4_numpy( return z_nabla4_e2 +@pytest.mark.continuous_benchmarking class TestCalculateNabla4(StencilTest): PROGRAM = calculate_nabla4 OUTPUTS = ("z_nabla4_e2",) STATIC_PARAMS = { - StandardStaticVariants.NONE: (), + StandardStaticVariants.NONE: None, StandardStaticVariants.COMPILE_TIME_DOMAIN: ( "horizontal_start", "horizontal_end", @@ -128,8 +129,3 @@ def input_data(self, grid) -> dict: vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) - - -@pytest.mark.continuous_benchmarking -class TestCalculateNabla4ContinuousBenchmarking(TestCalculateNabla4): - pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 1ea5d411b5..d49922254d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -6,8 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import itertools - import gt4py.next as gtx import numpy as np import pytest @@ -21,7 +19,6 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions divergence_damp_order = DivergenceDampingOrder() @@ -31,23 +28,6 @@ class TestApplyDivergenceDampingAndUpdateVn(test_helpers.StencilTest): PROGRAM = apply_divergence_damping_and_update_vn OUTPUTS = ("next_vn",) - STATIC_PARAMS = { - test_helpers.StandardStaticVariants.NONE: (), - test_helpers.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "is_iau_active", - "limited_area", - ), - test_helpers.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "is_iau_active", - "limited_area", - ), - } @staticmethod def reference( @@ -173,10 +153,7 @@ def reference( return dict(next_vn=next_vn) - @pytest.fixture( - params=[{"limited_area": la} for la in [True, False]], - ids=lambda param: f"limited_area[{param['limited_area']}]", - ) + @pytest.fixture(params=[True, False]) def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: current_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_mask_for_3d_divdamp = data_alloc.random_field(grid, dims.EdgeDim) @@ -225,7 +202,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) ) - limited_area = request.param["limited_area"] + limited_area = request.param edge_domain = h_grid.domain(dims.EdgeDim) start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) @@ -263,43 +240,3 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) - - -@pytest.mark.continuous_benchmarking -class TestApplyDivergenceDampingAndUpdateVnContinuousBenchmarking( - TestApplyDivergenceDampingAndUpdateVn -): - @pytest.fixture( - params=[ - {"limited_area": la, "divdamp_order": do, "is_iau_active": ia} - for la, do, ia in itertools.product( - [True, False], - [ - DivergenceDampingOrder.SECOND_ORDER, - DivergenceDampingOrder.FOURTH_ORDER, - DivergenceDampingOrder.COMBINED, - ], - [True, False], - ) - ], - ids=lambda param: f"limited_area[{param['limited_area']}]__divdamp_order[{param['divdamp_order']}]__is_iau_active[{param['is_iau_active']}]", - ) - def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: - base_data = TestApplyDivergenceDampingAndUpdateVn.input_data.__wrapped__( - self, request, grid - ) - base_data["is_iau_active"] = False - fourth_order_divdamp_factor = 0.004 - second_order_divdamp_factor = 0.032 - divdamp_order = request.param["divdamp_order"] - base_data["second_order_divdamp_scaling_coeff"] = 34497.62082646618 # for icon-ch1(_medium) - base_data["apply_2nd_order_divergence_damping"] = ( - divdamp_order == divergence_damp_order.COMBINED - ) and (base_data["second_order_divdamp_scaling_coeff"] > 1.0e-6) - base_data["apply_4th_order_divergence_damping"] = ( - divdamp_order == divergence_damp_order.FOURTH_ORDER - ) or ( - (divdamp_order == divergence_damp_order.COMBINED) - and (second_order_divdamp_factor <= (4.0 * fourth_order_divdamp_factor)) - ) - return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py index 70c6a4d051..4340ee9ea4 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py @@ -138,21 +138,6 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelm class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): PROGRAM = compute_advection_in_horizontal_momentum_equation OUTPUTS = ("normal_wind_advective_tendency",) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "end_index_of_damping_layer", - "vertical_start", - "vertical_end", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "end_index_of_damping_layer", - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -242,13 +227,8 @@ def reference( return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) - @pytest.fixture( - params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]], - ids=lambda param: f"apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]", - ) - def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: normal_wind_advective_tendency = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim) vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( @@ -277,7 +257,7 @@ def input_data( scalfac_exdiff = 0.6 dtime = 2.0 cfl_w_limit = 0.65 / dtime - apply_extra_diffusion_on_vn = request.param["apply_extra_diffusion_on_vn"] + apply_extra_diffusion_on_vn = True end_index_of_damping_layer = 5 edge_domain = h_grid.domain(dims.EdgeDim) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py index 22714168dc..a98978eeda 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py @@ -19,7 +19,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests from .test_add_interpolated_horizontal_advection_of_w import ( add_interpolated_horizontal_advection_of_w_numpy, @@ -276,21 +276,6 @@ class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): "contravariant_corrected_w_at_cells_on_model_levels", "vertical_cfl", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "end_index_of_damping_layer", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "end_index_of_damping_layer", - ), - } @staticmethod def reference( @@ -482,17 +467,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala ) -@pytest.mark.continuous_benchmarking -class TestFusedVelocityAdvectionStencilVMomentumContinuousBenchmarking( - TestFusedVelocityAdvectionStencilVMomentum -): - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - base_data = TestFusedVelocityAdvectionStencilVMomentum.input_data.__wrapped__(self, grid) - base_data["end_index_of_damping_layer"] = 12 - return base_data - - @pytest.mark.embedded_remap_error class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.StencilTest): PROGRAM = compute_contravariant_correction_and_advection_in_vertical_momentum_equation @@ -502,9 +476,6 @@ class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.S "contravariant_corrected_w_at_cells_on_model_levels", "vertical_cfl", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), # For now compile time variants triger error in gt4py - } @staticmethod def reference( @@ -632,8 +603,7 @@ def reference( ) @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], - ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]", + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]] ) def input_data( self, grid: base.Grid, request: pytest.FixtureRequest @@ -709,10 +679,3 @@ def input_data( vertical_start=vertical_start, vertical_end=vertical_end, ) - - -@pytest.mark.continuous_benchmarking -class TestFusedVelocityAdvectionStencilVMomentumAndContravariantContinuousBenchmarking( - TestFusedVelocityAdvectionStencilVMomentumAndContravariant -): - pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py index be12bc3ee1..914bb23c45 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_derived_horizontal_winds_and_ke_and_contravariant_correction.py @@ -16,7 +16,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests from .test_compute_contravariant_correction import compute_contravariant_correction_numpy from .test_compute_horizontal_advection_term_for_vertical_velocity import ( @@ -154,21 +154,6 @@ class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContrava "contravariant_correction_at_edges_on_model_levels", "horizontal_advection_of_w_at_edges_on_half_levels", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "nflatlev", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "nflatlev", - ), - } @classmethod def reference( @@ -281,8 +266,7 @@ def reference( ) @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], - ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]", + params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]] ) def input_data( self, grid: base.Grid, request: pytest.FixtureRequest @@ -350,31 +334,3 @@ def input_data( vertical_start=vertical_start, vertical_end=vertical_end, ) - - -@pytest.mark.continuous_benchmarking -class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrectionContinuousBenchmarking( - TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection -): - @pytest.fixture( - params=[{"skip_compute_predictor_vertical_advection": value} for value in [True, False]], - ids=lambda param: f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]", - ) - def input_data( - self, grid: base.Grid, request: pytest.FixtureRequest - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - base_data = TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection.input_data.__wrapped__( - self, grid, request - ) - base_data["skip_compute_predictor_vertical_advection"] = request.param[ - "skip_compute_predictor_vertical_advection" - ] - base_data["nflatlev"] = 6 - edge_domain = h_grid.domain(dims.EdgeDim) - base_data["horizontal_start"] = grid.start_index( - edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) - ) - base_data["horizontal_end"] = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) - base_data["vertical_start"] = 0 - base_data["vertical_end"] = grid.num_levels + 1 - return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 8a615067f9..79217ba75d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing.stencil_tests import StencilTest def compute_hydrostatic_correction_term_numpy( @@ -83,27 +83,14 @@ def _apply_index_field( / ((z_theta1 + z_theta2) ** 2) ) - nlevels = theta_v.shape[1] - return z_hydro_corr[:, nlevels - 1 : nlevels] + return z_hydro_corr +@pytest.mark.skip_value_error @pytest.mark.uses_as_offset class TestComputeHydrostaticCorrectionTerm(StencilTest): OUTPUTS = ("z_hydro_corr",) PROGRAM = compute_hydrostatic_correction_term - STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -154,14 +141,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala z_hydro_corr = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - z_hydro_corr = gtx.constructors.zeros( - domain={ - dims.EdgeDim: (0, grid.num_edges), - dims.KDim: (grid.num_levels - 1, grid.num_levels), - }, - dtype=ta.vpfloat, - ) - return dict( theta_v=theta_v, ikoffset=ikoffset, @@ -173,13 +152,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala grav_o_cpd=grav_o_cpd, horizontal_start=0, horizontal_end=gtx.int32(grid.num_edges), - vertical_start=gtx.int32(grid.num_levels - 1), + vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) - - -@pytest.mark.continuous_benchmarking -class TestComputeHydrostaticCorrectionTermContinuousBenchmarking( - TestComputeHydrostaticCorrectionTerm -): - pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 5f55af9b57..a5415499d8 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -32,7 +32,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests from .test_compute_approx_of_2nd_vertical_derivative_of_exner import ( compute_approx_of_2nd_vertical_derivative_of_exner_numpy, @@ -80,28 +80,6 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "pressure_buoyancy_acceleration_at_cells_on_half_levels", "d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "igradp_method", - "nflatlev", - "nflat_gradp", - "start_cell_lateral_boundary", - "start_cell_lateral_boundary_level_3", - "start_cell_halo_level_2", - "end_cell_halo", - "end_cell_halo_level_2", - "model_top", - "surface_level", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "igradp_method", - "nflatlev", - "nflat_gradp", - "model_top", - "surface_level", - ), - } @staticmethod def reference( @@ -420,7 +398,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala end_cell_halo_level_2 = grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) nflatlev = 4 - nflat_gradp = 7 + nflat_gradp = 27 return dict( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, @@ -461,17 +439,3 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala model_top=0, surface_level=grid.num_levels + 1, ) - - -@pytest.mark.continuous_benchmarking -class TestComputePerturbedQuantitiesAndInterpolationContinuousBenchmarking( - TestComputePerturbedQuantitiesAndInterpolation -): - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - base_data = TestComputePerturbedQuantitiesAndInterpolation.input_data.__wrapped__( - self, grid - ) - base_data["nflatlev"] = 6 - base_data["nflat_gradp"] = 35 - return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py index c034d7521d..350d380792 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_theta_rho_face_values_and_pressure_gradient_and_update_vn.py @@ -13,6 +13,7 @@ import pytest import icon4py.model.common.type_alias as ta +import icon4py.model.testing.stencil_tests as test_helpers from icon4py.model.atmosphere.dycore.dycore_states import ( HorizontalPressureDiscretizationType, RhoThetaAdvectionType, @@ -23,7 +24,6 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests rhotheta_avd_type = RhoThetaAdvectionType() @@ -122,8 +122,9 @@ def compute_theta_rho_face_value_by_miura_scheme_numpy( @pytest.mark.embedded_remap_error +@pytest.mark.skip_value_error @pytest.mark.uses_as_offset -class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): +class TestComputeThetaRhoPressureGradientAndUpdateVn(test_helpers.StencilTest): PROGRAM = compute_theta_rho_face_values_and_pressure_gradient_and_update_vn OUTPUTS = ( "rho_at_edges_on_model_levels", @@ -131,33 +132,6 @@ class TestComputeThetaRhoPressureGradientAndUpdateVn(stencil_tests.StencilTest): "horizontal_pressure_gradient", "next_vn", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "iau_wgt_dyn", - "is_iau_active", - "limited_area", - "start_edge_lateral_boundary", - "start_edge_lateral_boundary_level_7", - "start_edge_nudging_level_2", - "end_edge_nudging", - "end_edge_halo", - "nflatlev", - "nflat_gradp", - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "is_iau_active", - "limited_area", - "nflatlev", - "nflat_gradp", - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -294,10 +268,7 @@ def reference( perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, ), - ( - np.zeros_like(rho_at_edges_on_model_levels), - np.zeros_like(theta_v_at_edges_on_model_levels), - ), + (rho_at_edges_on_model_levels, theta_v_at_edges_on_model_levels), ) # Remaining computations at edge points @@ -566,20 +537,3 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=grid.num_levels, ) - - -@pytest.mark.continuous_benchmarking -class TestComputeThetaRhoPressureGradientAndUpdateVnContinuousBenchmarking( - TestComputeThetaRhoPressureGradientAndUpdateVn -): - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict: - base_data = TestComputeThetaRhoPressureGradientAndUpdateVn.input_data.__wrapped__( - self, grid - ) - base_data["is_iau_active"] = False - base_data["limited_area"] = grid.limited_area - base_data["nflatlev"] = 5 - base_data["nflat_gradp"] = 34 - base_data["start_edge_lateral_boundary"] = 0 - return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py index 943d0419d0..658499bc80 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_wp.py @@ -19,25 +19,12 @@ from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils.data_allocation import zero_field -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing.stencil_tests import StencilTest class TestInitCellKdimFieldWithZeroWp(StencilTest): PROGRAM = init_cell_kdim_field_with_zero_wp OUTPUTS = ("field_with_zero_wp",) - STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -59,8 +46,3 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) - - -@pytest.mark.continuous_benchmarking -class TestInitCellKdimFieldWithZeroWpContinuousBenchmarking(TestInitCellKdimFieldWithZeroWp): - pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py index 1011e3063e..b3fb185c76 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_rho_theta_v_to_half_levels_and_compute_temperature_vertical_gradient.py @@ -34,7 +34,6 @@ from icon4py.model.testing import stencil_tests -@pytest.mark.uses_concat_where class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration( stencil_tests.StencilTest ): @@ -45,19 +44,6 @@ class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration "theta_v_at_cells_on_half_levels", "pressure_buoyancy_acceleration_at_cells_on_half_levels", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -250,10 +236,3 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=1, vertical_end=grid.num_levels, ) - - -@pytest.mark.continuous_benchmarking -class TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAccelerationContinuousBenchmarking( - TestInterpolateRhoThetaVToHalfLevelsAndComputePressureBuoyancyAcceleration -): - pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py index 0013e7e1da..c9297d94b7 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_update_mass_flux_weighted.py @@ -19,25 +19,12 @@ from icon4py.model.common.states import utils as state_utils from icon4py.model.common.type_alias import vpfloat, wpfloat from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing.stencil_tests import StencilTest class TestUpdateMassFluxWeighted(StencilTest): PROGRAM = update_mass_flux_weighted OUTPUTS = ("mass_flx_ic",) - STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -84,8 +71,3 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) - - -@pytest.mark.continuous_benchmarking -class TestUpdateMassFluxWeightedContinuousBenchmarking(TestUpdateMassFluxWeighted): - pass diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py index 16cc63d8f1..c43c2de7be 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py @@ -5,7 +5,6 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import itertools from typing import Any import gt4py.next as gtx @@ -19,7 +18,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( add_analysis_increments_from_data_assimilation_numpy, @@ -63,33 +62,6 @@ class TestVerticallyImplicitSolverAtCorrectorStep(stencil_tests.StencilTest): "dynamical_vertical_volumetric_flux_at_cells_on_half_levels", "exner_dynamical_increment", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "start_cell_index_nudging", - "end_cell_index_local", - "end_index_of_damping_layer", - "kstart_moist", - "vertical_start_index_model_top", - "vertical_end_index_model_surface", - "at_first_substep", - "at_last_substep", - "lprep_adv", - "is_iau_active", - "rayleigh_type", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "end_index_of_damping_layer", - "kstart_moist", - "vertical_start_index_model_top", - "vertical_end_index_model_surface", - "at_first_substep", - "at_last_substep", - "lprep_adv", - "is_iau_active", - "rayleigh_type", - ), - } @staticmethod def reference( @@ -398,21 +370,8 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture( - params=[ - {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} - for afs, als, la in itertools.product(*([(True, False)] * 3)) - if not (afs and als) - ], - ids=lambda p: ( - f"at_first_substep[{p['at_first_substep']}]__" - f"at_last_substep[{p['at_last_substep']}]__" - f"lprep_adv[{p['lprep_adv']}]" - ), - ) - def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( @@ -467,13 +426,13 @@ def input_data( grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1} ) - lprep_adv = request.param["lprep_adv"] + lprep_adv = True r_nsubsteps = 0.5 is_iau_active = True - at_first_substep = request.param["at_first_substep"] + at_first_substep = True rayleigh_type = 2 end_index_of_damping_layer = 3 - at_last_substep = request.param["at_last_substep"] + at_last_substep = True kstart_moist = 1 dtime = 0.001 veladv_offctr = 0.25 @@ -535,34 +494,3 @@ def input_data( vertical_start_index_model_top=gtx.int32(0), vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), ) - - -@pytest.mark.continuous_benchmarking -class TestVerticallyImplicitSolverAtCorrectorStepContinuousBenchmarking( - TestVerticallyImplicitSolverAtCorrectorStep -): - @pytest.fixture( - params=[ - {"at_first_substep": afs, "at_last_substep": als, "lprep_adv": la} - for afs, als, la in itertools.product(*([(True, False)] * 3)) - if not (afs and als) - ], - ids=lambda p: ( - f"at_first_substep[{p['at_first_substep']}]__" - f"at_last_substep[{p['at_last_substep']}]__" - f"lprep_adv[{p['lprep_adv']}]" - ), - ) - def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - base_data = TestVerticallyImplicitSolverAtCorrectorStep.input_data.__wrapped__( - self, request, grid - ) - base_data["at_first_substep"] = request.param["at_first_substep"] - base_data["at_last_substep"] = request.param["at_last_substep"] - base_data["lprep_adv"] = request.param["lprep_adv"] - base_data["is_iau_active"] = False - base_data["end_index_of_damping_layer"] = 13 - base_data["kstart_moist"] = 0 - return base_data diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py index 1b032bee50..b4866a0f75 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py @@ -18,7 +18,7 @@ from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions, stencil_tests +from icon4py.model.testing import stencil_tests from .test_add_analysis_increments_from_data_assimilation import ( add_analysis_increments_from_data_assimilation_numpy, @@ -64,35 +64,6 @@ class TestVerticallyImplicitSolverAtPredictorStep(stencil_tests.StencilTest): "dwdz_at_cells_on_model_levels", "exner_dynamical_increment", ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "start_cell_index_nudging", - "end_cell_index_local", - "start_cell_index_lateral_lvl3", - "end_cell_index_halo_lvl1", - "end_index_of_damping_layer", - "kstart_moist", - "flat_level_index_plus1", - "vertical_start_index_model_top", - "vertical_end_index_model_surface", - "divdamp_type", - "rayleigh_type", - "is_iau_active", - "at_first_substep", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "end_index_of_damping_layer", - "kstart_moist", - "flat_level_index_plus1", - "vertical_start_index_model_top", - "vertical_end_index_model_surface", - "divdamp_type", - "rayleigh_type", - "is_iau_active", - "at_first_substep", - ), - } @staticmethod def reference( @@ -395,13 +366,8 @@ def reference( exner_dynamical_increment=exner_dynamical_increment, ) - @pytest.fixture( - params=[{"at_first_substep": value} for value in [True, False]], - ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", - ) - def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: + @pytest.fixture + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: geofac_div = data_alloc.random_field(grid, dims.CellDim, dims.C2EDim) mass_flux_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim) theta_v_flux_at_edges_on_model_levels = data_alloc.random_field( @@ -457,7 +423,7 @@ def input_data( exner_dynamical_increment = data_alloc.zero_field(grid, dims.CellDim, dims.KDim) is_iau_active = True - at_first_substep = request.param["at_first_substep"] + at_first_substep = True rayleigh_type = 2 divdamp_type = 3 end_index_of_damping_layer = 3 @@ -523,25 +489,3 @@ def input_data( vertical_start_index_model_top=gtx.int32(0), vertical_end_index_model_surface=gtx.int32(grid.num_levels + 1), ) - - -@pytest.mark.continuous_benchmarking -class TestVerticallyImplicitSolverAtPredictorStepContinuousBenchmarking( - TestVerticallyImplicitSolverAtPredictorStep -): - @pytest.fixture( - params=[{"at_first_substep": value} for value in [True, False]], - ids=lambda param: f"at_first_substep[{param['at_first_substep']}]", - ) - def input_data( - self, request: pytest.FixtureRequest, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - base_data = TestVerticallyImplicitSolverAtPredictorStep.input_data.__wrapped__( - self, request, grid - ) - base_data["at_first_substep"] = request.param["at_first_substep"] - base_data["is_iau_active"] = False - base_data["divdamp_type"] = 32 - base_data["end_index_of_damping_layer"] = 13 - base_data["kstart_moist"] = 0 - return base_data diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index 2ddde5d342..f6e5034fa9 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -18,25 +18,13 @@ ) from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StandardStaticVariants, StencilTest +from icon4py.model.testing.stencil_tests import StencilTest +@pytest.mark.skip_value_error class TestMoIntpRbfRbfVecInterpolVertex(StencilTest): PROGRAM = mo_intp_rbf_rbf_vec_interpol_vertex OUTPUTS = ("p_u_out", "p_v_out") - STATIC_PARAMS = { - StandardStaticVariants.NONE: (), - StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - ), - StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - ), - } @staticmethod def reference( @@ -74,8 +62,3 @@ def input_data(self, grid: base.Grid) -> dict: vertical_start=0, vertical_end=gtx.int32(grid.num_levels), ) - - -@pytest.mark.continuous_benchmarking -class TestMoIntpRbfRbfVecInterpolVertexContinuousBenchmarking(TestMoIntpRbfRbfVecInterpolVertex): - pass diff --git a/model/testing/src/icon4py/model/testing/pytest_hooks.py b/model/testing/src/icon4py/model/testing/pytest_hooks.py index 6e1c2f7b40..670f7ff775 100644 --- a/model/testing/src/icon4py/model/testing/pytest_hooks.py +++ b/model/testing/src/icon4py/model/testing/pytest_hooks.py @@ -102,23 +102,6 @@ def pytest_addoption(parser: pytest.Parser): def pytest_collection_modifyitems(config, items): - """Modify collected test items based on command line options.""" - test_grid = config.getoption("--grid") - for item in items: - if (marker := item.get_closest_marker("continuous_benchmarking")) is not None: - if test_grid is not None and not test_grid.startswith("icon_benchmark"): - item.add_marker( - pytest.mark.skip( - reason="Continuous benchmarking tests only run with --grid icon_benchmark_{regional,global}." - ) - ) - continue - if not config.getoption("--benchmark-only"): - item.add_marker( - pytest.mark.benchmark_only( - reason="Continuous benchmarking tests shouldn't be verified." - ) - ) test_level = config.getoption("--level") if test_level == "any": return @@ -178,56 +161,3 @@ def pytest_benchmark_update_json(output_json): for bench in output_json["benchmarks"]: bench["fullname"] = _name_from_fullname(bench["fullname"]) - - -@pytest.hookimpl(hookwrapper=True) -def pytest_runtest_makereport(item, call): - """ - Gather GT4Py timer metrics from benchmark fixture and add them to the test report. - """ - outcome = yield - report = outcome.get_result() - if call.when == "call": - benchmark = item.funcargs.get("benchmark", None) - if benchmark and hasattr(benchmark, "extra_info"): - info = benchmark.extra_info.get("gtx_metrics", None) - if info: - filtered_benchmark_name = benchmark.name.split("test_Test")[-1] - # Combine the benchmark name in a readable form with the gtx_metrics data - report.sections.append(("benchmark-extra", tuple([filtered_benchmark_name, info]))) - - -def pytest_terminal_summary(terminalreporter, exitstatus, config): - """ - Add a custom section to the terminal summary with GT4Py timer metrics from benchmarks. - """ - # Gather gtx_metrics - benchmark_gtx_metrics = [] - for outcome in ("passed", "failed", "skipped"): - all_reports = terminalreporter.stats.get(outcome, []) - for report in all_reports: - for secname, info in getattr(report, "sections", []): - if secname == "benchmark-extra": - benchmark_gtx_metrics.append(info) - # Calculate the maximum length of benchmark names for formatting - max_name_len = 0 - for benchmark_name, _ in benchmark_gtx_metrics: - max_name_len = max(len(benchmark_name), max_name_len) - # Print the GT4Py timer report table - if benchmark_gtx_metrics: - terminalreporter.ensure_newline() - header = f"{'Benchmark Name':<{max_name_len}} | {'Mean (s)':>10} | {'Median (s)':>10} | {'Std Dev':>10} | {'Runs':>4}" - title = " GT4Py Timer Report " - sep_len = max(0, len(header) - len(title)) - left = sep_len // 2 - right = sep_len - left - terminalreporter.line("-" * left + title + "-" * right, bold=True, blue=True) - terminalreporter.line(header) - terminalreporter.line("-" * len(header), blue=True) - import numpy as np - - for benchmark_name, gtx_metrics in benchmark_gtx_metrics: - terminalreporter.line( - f"{benchmark_name:<{max_name_len}} | {np.mean(gtx_metrics):>10.8f} | {np.median(gtx_metrics):>10.8f} | {np.std(gtx_metrics):>10.8f} | {len(gtx_metrics):>4}" - ) - terminalreporter.line("-" * len(header), blue=True) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index bdbd4d1905..b4cae435ba 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -9,33 +9,27 @@ from __future__ import annotations import dataclasses -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Sequence from typing import Any, ClassVar import gt4py.next as gtx import numpy as np import pytest from gt4py import eve -from gt4py.next import ( - config as gtx_config, - constructors, - metrics as gtx_metrics, - typing as gtx_typing, -) +from gt4py.next import constructors, typing as gtx_typing # TODO(havogt): import will disappear after FieldOperators support `.compile` from gt4py.next.ffront.decorator import FieldOperator -from icon4py.model.common import model_backends, model_options from icon4py.model.common.grid import base from icon4py.model.common.utils import device_utils def allocate_data( - allocator: gtx_typing.FieldBufferAllocationUtil | None, + backend: gtx_typing.Backend | None, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: - _allocate_field = constructors.as_field.partial(allocator=allocator) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial + _allocate_field = constructors.as_field.partial(allocator=backend) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial input_data = { k: tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) if isinstance(v, tuple) @@ -82,55 +76,29 @@ def test_and_benchmark( grid: base.Grid, _properly_allocated_input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], _configured_program: Callable[..., None], - request: pytest.FixtureRequest, ) -> None: - benchmark_only_option = request.config.getoption("benchmark_only") - benchmark_only_mark = request.node.get_closest_marker("benchmark_only") is not None - if (not benchmark_only_option) and (not benchmark_only_mark): - reference_outputs = self.reference( - _ConnectivityConceptFixer( - grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) - ), - **{ - k: v.asnumpy() if isinstance(v, gtx.Field) else v - for k, v in _properly_allocated_input_data.items() - }, - ) - - _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) - self._verify_stencil_test( - input_data=_properly_allocated_input_data, reference_outputs=reference_outputs - ) + reference_outputs = self.reference( + _ConnectivityConceptFixer( + grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) + ), + **{ + k: v.asnumpy() if isinstance(v, gtx.Field) else v + for k, v in _properly_allocated_input_data.items() + }, + ) + + _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) + self._verify_stencil_test( + input_data=_properly_allocated_input_data, reference_outputs=reference_outputs + ) if benchmark is not None and benchmark.enabled: - warmup_enabled = request.config.getoption("benchmark_warmup") - if warmup_enabled: - print("[WARNING] Benchmark warmup enabled, GT4Py timers include warmup iterations.") - # Clean up GT4Py metrics from previous runs - if gtx_config.COLLECT_METRICS_LEVEL > 0: - gtx_metrics.sources.clear() - benchmark( _configured_program, **_properly_allocated_input_data, offset_provider=grid.connectivities, ) - # Collect GT4Py runtime metrics if enabled - if gtx_config.COLLECT_METRICS_LEVEL > 0: - assert ( - len(gtx_metrics.sources) == 1 - ), "Expected exactly one entry in gtx_metrics.sources" - # Store GT4Py metrics in benchmark.extra_info - metrics_data = gtx_metrics.sources - key = next(iter(metrics_data)) - compute_samples = metrics_data[key].metrics["compute"].samples - # emprically exclude first few iterations run for warmup - initial_program_iterations_to_skip = 2 - benchmark.extra_info["gtx_metrics"] = compute_samples[ - initial_program_iterations_to_skip: - ] - class StencilTest: """ @@ -156,12 +124,12 @@ class StencilTest: OUTPUTS: ClassVar[tuple[str | Output, ...]] STATIC_PARAMS: ClassVar[dict[str, Sequence[str]] | None] = None - reference: ClassVar[Callable[..., Mapping[str, np.ndarray | tuple[np.ndarray, ...]]]] + reference: ClassVar[Callable[..., dict[str, np.ndarray | tuple[np.ndarray, ...]]]] @pytest.fixture def _configured_program( self, - backend_like: model_backends.BackendLike, + backend: gtx_typing.Backend | None, static_variant: Sequence[str], input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], grid: base.Grid, @@ -172,7 +140,6 @@ def _configured_program( f"Parameter defined in 'STATIC_PARAMS' not in 'input_data': {unused_static_params}" ) static_args = {name: [input_data[name]] for name in static_variant} - backend = model_options.customize_backend(self.PROGRAM, backend_like) program = self.PROGRAM.with_backend(backend) # type: ignore[arg-type] # TODO(havogt): gt4py should accept `None` in with_backend if backend is not None: if isinstance(program, FieldOperator): @@ -187,25 +154,24 @@ def _configured_program( **static_args, # type: ignore[arg-type] ) - test_func = device_utils.synchronized_function(program, allocator=backend) + test_func = device_utils.synchronized_function(program, backend=backend) return test_func @pytest.fixture def _properly_allocated_input_data( self, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - backend_like: model_backends.BackendLike, + backend: gtx_typing.Backend | None, ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: # TODO(havogt): this is a workaround, # because in the `input_data` fixture provided by the user # it does not allocate for the correct device. - allocator = model_backends.get_allocator(backend_like) - return allocate_data(allocator=allocator, input_data=input_data) + return allocate_data(backend, input_data) def _verify_stencil_test( self, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - reference_outputs: Mapping[str, np.ndarray | tuple[np.ndarray, ...]], + reference_outputs: dict[str, np.ndarray | tuple[np.ndarray, ...]], ) -> None: for out in self.OUTPUTS: name, refslice, gtslice = ( @@ -247,33 +213,11 @@ def static_variant(request: pytest.FixtureRequest) -> Sequence[str]: def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) - pytest_prefix = "test_" - - # only add `test_and_benchmark` to direct subclasses of `StencilTest`. Inherited tests can use the same function. - if cls.__base__ == StencilTest: - setattr(cls, f"{pytest_prefix}{cls.__name__}", test_and_benchmark) + setattr(cls, f"test_{cls.__name__}", test_and_benchmark) # decorate `static_variant` with parametrized fixtures, since the # parametrization is only available in the concrete subclass definition - # Check if cls.static_variant has changes in inherited subclasses and update it accordingly - if hasattr(cls.static_variant, "_pytestfixturefunction"): - base_static = getattr(cls.__base__, "STATIC_PARAMS", None) - if base_static != cls.STATIC_PARAMS: - if cls.STATIC_PARAMS is None: - cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] - else: - # copy of `static_variant` - def _new_static_variant(request: pytest.FixtureRequest) -> Sequence[str]: - _, variant = request.param - return () if variant is None else variant - - # replace with new parametrized fixture - cls.static_variant = staticmethod( # type: ignore[method-assign] - pytest.fixture( - params=cls.STATIC_PARAMS.items(), scope="class", ids=lambda p: p[0] - )(_new_static_variant) - ) - elif cls.STATIC_PARAMS is None: + if cls.STATIC_PARAMS is None: # not parametrized, return an empty tuple cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] # we override with a non-parametrized function else: From 879a35dd462e8cfe8e2cb86ef17f626e7ecff75b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:09:44 +0100 Subject: [PATCH 088/108] more restoration --- ci/benchmark_bencher.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/benchmark_bencher.yml b/ci/benchmark_bencher.yml index 80b46a89dd..120c144277 100644 --- a/ci/benchmark_bencher.yml +++ b/ci/benchmark_bencher.yml @@ -21,7 +21,7 @@ include: parallel: matrix: - BACKEND: [dace_cpu, dace_gpu, gtfn_cpu, gtfn_gpu] - GRID: [icon_benchmark] + GRID: [icon_benchmark_regional, icon_benchmark_global] variables: ICON4PY_ENABLE_GRID_DOWNLOAD: false ICON4PY_ENABLE_TESTDATA_DOWNLOAD: false From 4adfa21c55ecc516153538d2cfc2ed795255debb Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:15:49 +0100 Subject: [PATCH 089/108] more restoration and ran pre-commit --- .../stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py | 2 +- model/common/tests/common/math/unit_tests/test_operators.py | 2 +- model/testing/src/icon4py/model/testing/stencil_tests.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py index f6e5034fa9..0b3d36290b 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_mo_intp_rbf_rbf_vec_interpol_vertex.py @@ -33,7 +33,7 @@ def reference( ptr_coeff_1: np.ndarray, ptr_coeff_2: np.ndarray, **kwargs: Any, - ) -> dict[str, np.ndarray]: + ) -> dict: v2e = connectivities[dims.V2EDim] ptr_coeff_1 = np.expand_dims(ptr_coeff_1, axis=-1) p_u_out = np.sum(p_e_in[v2e] * ptr_coeff_1, axis=1) diff --git a/model/common/tests/common/math/unit_tests/test_operators.py b/model/common/tests/common/math/unit_tests/test_operators.py index b487d1a50b..3c75a92f9b 100644 --- a/model/common/tests/common/math/unit_tests/test_operators.py +++ b/model/common/tests/common/math/unit_tests/test_operators.py @@ -33,7 +33,7 @@ def reference( psi_c: np.ndarray, geofac_n2s: np.ndarray, **kwargs: Any, - ) -> dict[str, np.ndarray]: + ) -> dict: nabla2_psi_c_np = reference_funcs.nabla2_on_cell_numpy(connectivities, psi_c, geofac_n2s) return dict(nabla2_psi_c=nabla2_psi_c_np) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index b4cae435ba..16a4c15895 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -154,7 +154,7 @@ def _configured_program( **static_args, # type: ignore[arg-type] ) - test_func = device_utils.synchronized_function(program, backend=backend) + test_func = device_utils.synchronized_function(program, allocator=backend) return test_func @pytest.fixture From c96773e5d8daad453c16a45f456e0af66c77fd5f Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:24:04 +0100 Subject: [PATCH 090/108] one more restored file --- model/testing/src/icon4py/model/testing/stencil_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 16a4c15895..b4cae435ba 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -154,7 +154,7 @@ def _configured_program( **static_args, # type: ignore[arg-type] ) - test_func = device_utils.synchronized_function(program, allocator=backend) + test_func = device_utils.synchronized_function(program, backend=backend) return test_func @pytest.fixture From b8e07bff3ae7c7fb434b59e7f3320376f95c40a6 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:25:30 +0100 Subject: [PATCH 091/108] one more restored file --- ...t_extrapolate_temporally_exner_pressure.py | 265 ++++++++++++++---- 1 file changed, 209 insertions(+), 56 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py index a87d532ff4..dfeb0ec6bf 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py @@ -5,74 +5,227 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any + +from __future__ import annotations + +import dataclasses +from collections.abc import Callable, Mapping, Sequence +from typing import Any, ClassVar import gt4py.next as gtx import numpy as np import pytest +from gt4py import eve +from gt4py.next import constructors, typing as gtx_typing -from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( - extrapolate_temporally_exner_pressure, -) -from icon4py.model.common import dimension as dims +# TODO(havogt): import will disappear after FieldOperators support `.compile` +from gt4py.next.ffront.decorator import FieldOperator + +from icon4py.model.common import model_backends, model_options from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.common.utils.data_allocation import random_field, zero_field -from icon4py.model.testing.stencil_tests import StencilTest +from icon4py.model.common.utils import device_utils -def extrapolate_temporally_exner_pressure_numpy( - connectivities: dict[gtx.Dimension, np.ndarray], - exner: np.ndarray, - exner_ref_mc: np.ndarray, - exner_pr: np.ndarray, - exner_exfac: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr - exner_pr = exner - exner_ref_mc - return (z_exner_ex_pr, exner_pr) +def allocate_data( + allocator: gtx_typing.FieldBufferAllocationUtil | None, + input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], +) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: + _allocate_field = constructors.as_field.partial(allocator=allocator) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial + input_data = { + k: tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) + if isinstance(v, tuple) + else _allocate_field(domain=v.domain, data=v.ndarray) + if not gtx.is_scalar_type(v) and k != "domain" + else v + for k, v in input_data.items() + } + return input_data -class TestExtrapolateTemporallyExnerPressure(StencilTest): - PROGRAM = extrapolate_temporally_exner_pressure - OUTPUTS = ("z_exner_ex_pr", "exner_pr") +@dataclasses.dataclass(frozen=True) +class Output: + name: str + refslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) + gtslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) - @staticmethod - def reference( - connectivities: dict[gtx.Dimension, np.ndarray], - exner: np.ndarray, - exner_ref_mc: np.ndarray, - exner_pr: np.ndarray, - exner_exfac: np.ndarray, - **kwargs: Any, - ) -> dict: - (z_exner_ex_pr, exner_pr) = extrapolate_temporally_exner_pressure_numpy( - connectivities, - exner=exner, - exner_ref_mc=exner_ref_mc, - exner_pr=exner_pr, - exner_exfac=exner_exfac, + +@dataclasses.dataclass +class _ConnectivityConceptFixer: + """ + This works around a misuse of dimensions as an identifier for connectivities. + Since GT4Py might change the way the mesh is represented, we could + keep this for a while, otherwise we need to touch all StencilTests. + """ + + _grid: base.Grid + + def __getitem__(self, dim: gtx.Dimension | str) -> np.ndarray: + if isinstance(dim, gtx.Dimension): + dim = dim.value + return self._grid.get_connectivity(dim).asnumpy() + + +class StandardStaticVariants(eve.StrEnum): + NONE = "none" + COMPILE_TIME_DOMAIN = "compile_time_domain" + COMPILE_TIME_VERTICAL = "compile_time_vertical" + + +def test_and_benchmark( + self: StencilTest, + benchmark: Any, # should be `pytest_benchmark.fixture.BenchmarkFixture` but pytest_benchmark is not typed + grid: base.Grid, + _properly_allocated_input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], + _configured_program: Callable[..., None], +) -> None: + reference_outputs = self.reference( + _ConnectivityConceptFixer( + grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) + ), + **{ + k: v.asnumpy() if isinstance(v, gtx.Field) else v + for k, v in _properly_allocated_input_data.items() + }, + ) + + _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) + self._verify_stencil_test( + input_data=_properly_allocated_input_data, reference_outputs=reference_outputs + ) + + if benchmark is not None and benchmark.enabled: + benchmark( + _configured_program, + **_properly_allocated_input_data, + offset_provider=grid.connectivities, ) - return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) + +class StencilTest: + """ + Base class to be used for testing stencils. + + Example (pseudo-code): + + >>> class TestMultiplyByTwo(StencilTest): # doctest: +SKIP + ... PROGRAM = multiply_by_two # noqa: F821 + ... OUTPUTS = ("some_output",) + ... STATIC_PARAMS = {"category_a": ["flag0"], "category_b": ["flag0", "flag1"]} + ... + ... @pytest.fixture + ... def input_data(self): + ... return {"some_input": ..., "some_output": ...} + ... + ... @staticmethod + ... def reference(some_input, **kwargs): + ... return dict(some_output=np.asarray(some_input) * 2) + """ + + PROGRAM: ClassVar[gtx_typing.Program | gtx_typing.FieldOperator] + OUTPUTS: ClassVar[tuple[str | Output, ...]] + STATIC_PARAMS: ClassVar[dict[str, Sequence[str]] | None] = None + + reference: ClassVar[Callable[..., Mapping[str, np.ndarray | tuple[np.ndarray, ...]]]] @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - exner = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - exner_pr = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) - exner_exfac = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - z_exner_ex_pr = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - - return dict( - exner_exfac=exner_exfac, - exner=exner, - exner_ref_mc=exner_ref_mc, - exner_pr=exner_pr, - z_exner_ex_pr=z_exner_ex_pr, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) + def _configured_program( + self, + backend_like: model_backends.BackendLike, + static_variant: Sequence[str], + input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], + grid: base.Grid, + ) -> Callable[..., None]: + unused_static_params = set(static_variant) - set(input_data.keys()) + if unused_static_params: + raise ValueError( + f"Parameter defined in 'STATIC_PARAMS' not in 'input_data': {unused_static_params}" + ) + static_args = {name: [input_data[name]] for name in static_variant} + backend = model_options.customize_backend(self.PROGRAM, backend_like) + program = self.PROGRAM.with_backend(backend) # type: ignore[arg-type] # TODO(havogt): gt4py should accept `None` in with_backend + if backend is not None: + if isinstance(program, FieldOperator): + if len(static_args) > 0: + raise NotImplementedError( + "'FieldOperator's do not support static arguments yet." + ) + else: + program.compile( + offset_provider=grid.connectivities, + enable_jit=False, + **static_args, # type: ignore[arg-type] + ) + + test_func = device_utils.synchronized_function(program, allocator=backend) + return test_func + + @pytest.fixture + def _properly_allocated_input_data( + self, + input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], + backend_like: model_backends.BackendLike, + ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: + # TODO(havogt): this is a workaround, + # because in the `input_data` fixture provided by the user + # it does not allocate for the correct device. + allocator = model_backends.get_allocator(backend_like) + return allocate_data(allocator=allocator, input_data=input_data) + + def _verify_stencil_test( + self, + input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], + reference_outputs: Mapping[str, np.ndarray | tuple[np.ndarray, ...]], + ) -> None: + for out in self.OUTPUTS: + name, refslice, gtslice = ( + (out.name, out.refslice, out.gtslice) + if isinstance(out, Output) + else (out, (slice(None),), (slice(None),)) + ) + + input_data_name = input_data[name] # for mypy + if isinstance(input_data_name, tuple): + for i_out_field, out_field in enumerate(input_data_name): + np.testing.assert_allclose( + out_field.asnumpy()[gtslice], + reference_outputs[name][i_out_field][refslice], + equal_nan=True, + err_msg=f"Verification failed for '{name}[{i_out_field}]'", + ) + else: + reference_outputs_name = reference_outputs[name] # for mypy + assert isinstance(reference_outputs_name, np.ndarray) + np.testing.assert_allclose( + input_data_name.asnumpy()[gtslice], + reference_outputs_name[refslice], + equal_nan=True, + err_msg=f"Verification failed for '{name}'", + ) + + @staticmethod + def static_variant(request: pytest.FixtureRequest) -> Sequence[str]: + """ + Fixture for parametrization over the `STATIC_PARAMS` of the test class. + + Note: the actual `pytest.fixture()` decoration happens inside `__init_subclass__`, + when all information is available. + """ + _, variant = request.param + return () if variant is None else variant + + def __init_subclass__(cls, **kwargs: Any) -> None: + super().__init_subclass__(**kwargs) + + setattr(cls, f"test_{cls.__name__}", test_and_benchmark) + + # decorate `static_variant` with parametrized fixtures, since the + # parametrization is only available in the concrete subclass definition + if cls.STATIC_PARAMS is None: + # not parametrized, return an empty tuple + cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] # we override with a non-parametrized function + else: + cls.static_variant = staticmethod( # type: ignore[method-assign] + pytest.fixture(params=cls.STATIC_PARAMS.items(), scope="class", ids=lambda p: p[0])( + cls.static_variant + ) + ) From 7ce3f5e539cb85e037f38f5349094d01619fa8f6 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:26:17 +0100 Subject: [PATCH 092/108] fixed error --- ...t_extrapolate_temporally_exner_pressure.py | 265 ++++-------------- .../icon4py/model/testing/stencil_tests.py | 21 +- 2 files changed, 68 insertions(+), 218 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py index dfeb0ec6bf..a87d532ff4 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_temporally_exner_pressure.py @@ -5,227 +5,74 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import dataclasses -from collections.abc import Callable, Mapping, Sequence -from typing import Any, ClassVar +from typing import Any import gt4py.next as gtx import numpy as np import pytest -from gt4py import eve -from gt4py.next import constructors, typing as gtx_typing -# TODO(havogt): import will disappear after FieldOperators support `.compile` -from gt4py.next.ffront.decorator import FieldOperator - -from icon4py.model.common import model_backends, model_options +from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( + extrapolate_temporally_exner_pressure, +) +from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.utils import device_utils - - -def allocate_data( - allocator: gtx_typing.FieldBufferAllocationUtil | None, - input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], -) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: - _allocate_field = constructors.as_field.partial(allocator=allocator) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial - input_data = { - k: tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) - if isinstance(v, tuple) - else _allocate_field(domain=v.domain, data=v.ndarray) - if not gtx.is_scalar_type(v) and k != "domain" - else v - for k, v in input_data.items() - } - return input_data - - -@dataclasses.dataclass(frozen=True) -class Output: - name: str - refslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) - gtslice: tuple[slice, ...] = dataclasses.field(default_factory=lambda: (slice(None),)) - +from icon4py.model.common.states import utils as state_utils +from icon4py.model.common.type_alias import vpfloat, wpfloat +from icon4py.model.common.utils.data_allocation import random_field, zero_field +from icon4py.model.testing.stencil_tests import StencilTest -@dataclasses.dataclass -class _ConnectivityConceptFixer: - """ - This works around a misuse of dimensions as an identifier for connectivities. - Since GT4Py might change the way the mesh is represented, we could - keep this for a while, otherwise we need to touch all StencilTests. - """ - _grid: base.Grid +def extrapolate_temporally_exner_pressure_numpy( + connectivities: dict[gtx.Dimension, np.ndarray], + exner: np.ndarray, + exner_ref_mc: np.ndarray, + exner_pr: np.ndarray, + exner_exfac: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + z_exner_ex_pr = (1 + exner_exfac) * (exner - exner_ref_mc) - exner_exfac * exner_pr + exner_pr = exner - exner_ref_mc + return (z_exner_ex_pr, exner_pr) - def __getitem__(self, dim: gtx.Dimension | str) -> np.ndarray: - if isinstance(dim, gtx.Dimension): - dim = dim.value - return self._grid.get_connectivity(dim).asnumpy() +class TestExtrapolateTemporallyExnerPressure(StencilTest): + PROGRAM = extrapolate_temporally_exner_pressure + OUTPUTS = ("z_exner_ex_pr", "exner_pr") -class StandardStaticVariants(eve.StrEnum): - NONE = "none" - COMPILE_TIME_DOMAIN = "compile_time_domain" - COMPILE_TIME_VERTICAL = "compile_time_vertical" - - -def test_and_benchmark( - self: StencilTest, - benchmark: Any, # should be `pytest_benchmark.fixture.BenchmarkFixture` but pytest_benchmark is not typed - grid: base.Grid, - _properly_allocated_input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - _configured_program: Callable[..., None], -) -> None: - reference_outputs = self.reference( - _ConnectivityConceptFixer( - grid # TODO(havogt): pass as keyword argument (needs fixes in some tests) - ), - **{ - k: v.asnumpy() if isinstance(v, gtx.Field) else v - for k, v in _properly_allocated_input_data.items() - }, - ) - - _configured_program(**_properly_allocated_input_data, offset_provider=grid.connectivities) - self._verify_stencil_test( - input_data=_properly_allocated_input_data, reference_outputs=reference_outputs - ) - - if benchmark is not None and benchmark.enabled: - benchmark( - _configured_program, - **_properly_allocated_input_data, - offset_provider=grid.connectivities, + @staticmethod + def reference( + connectivities: dict[gtx.Dimension, np.ndarray], + exner: np.ndarray, + exner_ref_mc: np.ndarray, + exner_pr: np.ndarray, + exner_exfac: np.ndarray, + **kwargs: Any, + ) -> dict: + (z_exner_ex_pr, exner_pr) = extrapolate_temporally_exner_pressure_numpy( + connectivities, + exner=exner, + exner_ref_mc=exner_ref_mc, + exner_pr=exner_pr, + exner_exfac=exner_exfac, ) - -class StencilTest: - """ - Base class to be used for testing stencils. - - Example (pseudo-code): - - >>> class TestMultiplyByTwo(StencilTest): # doctest: +SKIP - ... PROGRAM = multiply_by_two # noqa: F821 - ... OUTPUTS = ("some_output",) - ... STATIC_PARAMS = {"category_a": ["flag0"], "category_b": ["flag0", "flag1"]} - ... - ... @pytest.fixture - ... def input_data(self): - ... return {"some_input": ..., "some_output": ...} - ... - ... @staticmethod - ... def reference(some_input, **kwargs): - ... return dict(some_output=np.asarray(some_input) * 2) - """ - - PROGRAM: ClassVar[gtx_typing.Program | gtx_typing.FieldOperator] - OUTPUTS: ClassVar[tuple[str | Output, ...]] - STATIC_PARAMS: ClassVar[dict[str, Sequence[str]] | None] = None - - reference: ClassVar[Callable[..., Mapping[str, np.ndarray | tuple[np.ndarray, ...]]]] - - @pytest.fixture - def _configured_program( - self, - backend_like: model_backends.BackendLike, - static_variant: Sequence[str], - input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - grid: base.Grid, - ) -> Callable[..., None]: - unused_static_params = set(static_variant) - set(input_data.keys()) - if unused_static_params: - raise ValueError( - f"Parameter defined in 'STATIC_PARAMS' not in 'input_data': {unused_static_params}" - ) - static_args = {name: [input_data[name]] for name in static_variant} - backend = model_options.customize_backend(self.PROGRAM, backend_like) - program = self.PROGRAM.with_backend(backend) # type: ignore[arg-type] # TODO(havogt): gt4py should accept `None` in with_backend - if backend is not None: - if isinstance(program, FieldOperator): - if len(static_args) > 0: - raise NotImplementedError( - "'FieldOperator's do not support static arguments yet." - ) - else: - program.compile( - offset_provider=grid.connectivities, - enable_jit=False, - **static_args, # type: ignore[arg-type] - ) - - test_func = device_utils.synchronized_function(program, allocator=backend) - return test_func + return dict(z_exner_ex_pr=z_exner_ex_pr, exner_pr=exner_pr) @pytest.fixture - def _properly_allocated_input_data( - self, - input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - backend_like: model_backends.BackendLike, - ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: - # TODO(havogt): this is a workaround, - # because in the `input_data` fixture provided by the user - # it does not allocate for the correct device. - allocator = model_backends.get_allocator(backend_like) - return allocate_data(allocator=allocator, input_data=input_data) - - def _verify_stencil_test( - self, - input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - reference_outputs: Mapping[str, np.ndarray | tuple[np.ndarray, ...]], - ) -> None: - for out in self.OUTPUTS: - name, refslice, gtslice = ( - (out.name, out.refslice, out.gtslice) - if isinstance(out, Output) - else (out, (slice(None),), (slice(None),)) - ) - - input_data_name = input_data[name] # for mypy - if isinstance(input_data_name, tuple): - for i_out_field, out_field in enumerate(input_data_name): - np.testing.assert_allclose( - out_field.asnumpy()[gtslice], - reference_outputs[name][i_out_field][refslice], - equal_nan=True, - err_msg=f"Verification failed for '{name}[{i_out_field}]'", - ) - else: - reference_outputs_name = reference_outputs[name] # for mypy - assert isinstance(reference_outputs_name, np.ndarray) - np.testing.assert_allclose( - input_data_name.asnumpy()[gtslice], - reference_outputs_name[refslice], - equal_nan=True, - err_msg=f"Verification failed for '{name}'", - ) - - @staticmethod - def static_variant(request: pytest.FixtureRequest) -> Sequence[str]: - """ - Fixture for parametrization over the `STATIC_PARAMS` of the test class. - - Note: the actual `pytest.fixture()` decoration happens inside `__init_subclass__`, - when all information is available. - """ - _, variant = request.param - return () if variant is None else variant - - def __init_subclass__(cls, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) - - setattr(cls, f"test_{cls.__name__}", test_and_benchmark) - - # decorate `static_variant` with parametrized fixtures, since the - # parametrization is only available in the concrete subclass definition - if cls.STATIC_PARAMS is None: - # not parametrized, return an empty tuple - cls.static_variant = staticmethod(pytest.fixture(lambda: ())) # type: ignore[method-assign, assignment] # we override with a non-parametrized function - else: - cls.static_variant = staticmethod( # type: ignore[method-assign] - pytest.fixture(params=cls.STATIC_PARAMS.items(), scope="class", ids=lambda p: p[0])( - cls.static_variant - ) - ) + def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: + exner = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + exner_ref_mc = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + exner_pr = zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat) + exner_exfac = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + z_exner_ex_pr = zero_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) + + return dict( + exner_exfac=exner_exfac, + exner=exner, + exner_ref_mc=exner_ref_mc, + exner_pr=exner_pr, + z_exner_ex_pr=z_exner_ex_pr, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=0, + vertical_end=gtx.int32(grid.num_levels), + ) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index b4cae435ba..dfeb0ec6bf 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -9,7 +9,7 @@ from __future__ import annotations import dataclasses -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import Any, ClassVar import gt4py.next as gtx @@ -21,15 +21,16 @@ # TODO(havogt): import will disappear after FieldOperators support `.compile` from gt4py.next.ffront.decorator import FieldOperator +from icon4py.model.common import model_backends, model_options from icon4py.model.common.grid import base from icon4py.model.common.utils import device_utils def allocate_data( - backend: gtx_typing.Backend | None, + allocator: gtx_typing.FieldBufferAllocationUtil | None, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: - _allocate_field = constructors.as_field.partial(allocator=backend) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial + _allocate_field = constructors.as_field.partial(allocator=allocator) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial input_data = { k: tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) if isinstance(v, tuple) @@ -124,12 +125,12 @@ class StencilTest: OUTPUTS: ClassVar[tuple[str | Output, ...]] STATIC_PARAMS: ClassVar[dict[str, Sequence[str]] | None] = None - reference: ClassVar[Callable[..., dict[str, np.ndarray | tuple[np.ndarray, ...]]]] + reference: ClassVar[Callable[..., Mapping[str, np.ndarray | tuple[np.ndarray, ...]]]] @pytest.fixture def _configured_program( self, - backend: gtx_typing.Backend | None, + backend_like: model_backends.BackendLike, static_variant: Sequence[str], input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], grid: base.Grid, @@ -140,6 +141,7 @@ def _configured_program( f"Parameter defined in 'STATIC_PARAMS' not in 'input_data': {unused_static_params}" ) static_args = {name: [input_data[name]] for name in static_variant} + backend = model_options.customize_backend(self.PROGRAM, backend_like) program = self.PROGRAM.with_backend(backend) # type: ignore[arg-type] # TODO(havogt): gt4py should accept `None` in with_backend if backend is not None: if isinstance(program, FieldOperator): @@ -154,24 +156,25 @@ def _configured_program( **static_args, # type: ignore[arg-type] ) - test_func = device_utils.synchronized_function(program, backend=backend) + test_func = device_utils.synchronized_function(program, allocator=backend) return test_func @pytest.fixture def _properly_allocated_input_data( self, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - backend: gtx_typing.Backend | None, + backend_like: model_backends.BackendLike, ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: # TODO(havogt): this is a workaround, # because in the `input_data` fixture provided by the user # it does not allocate for the correct device. - return allocate_data(backend, input_data) + allocator = model_backends.get_allocator(backend_like) + return allocate_data(allocator=allocator, input_data=input_data) def _verify_stencil_test( self, input_data: dict[str, gtx.Field | tuple[gtx.Field, ...]], - reference_outputs: dict[str, np.ndarray | tuple[np.ndarray, ...]], + reference_outputs: Mapping[str, np.ndarray | tuple[np.ndarray, ...]], ) -> None: for out in self.OUTPUTS: name, refslice, gtslice = ( From c018408489702ca29f1f49f5e138c75e19860fb5 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 18 Nov 2025 10:46:07 +0100 Subject: [PATCH 093/108] small bound edit re-instated --- .../test_compute_perturbed_quantities_and_interpolation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index a5415499d8..a931ec916b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -398,7 +398,7 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala end_cell_halo_level_2 = grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) nflatlev = 4 - nflat_gradp = 27 + nflat_gradp = 7 return dict( temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, From 84163b150dade2e5b1a1f7b03f0098d29d76c7bc Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 24 Nov 2025 08:15:51 +0100 Subject: [PATCH 094/108] reverted changes after dace fix --- .../dycore/stencils/extrapolate_temporally_exner_pressure.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py index b278f827cf..581d60bf02 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_temporally_exner_pressure.py @@ -24,8 +24,8 @@ def _extrapolate_temporally_exner_pressure( z_exner_ex_pr_wp = (wpfloat("1.0") + exner_exfac_wp) * ( exner - exner_ref_mc_wp - ) - exner_exfac * exner_pr - exner_pr_wp = exner - exner_ref_mc + ) - exner_exfac_wp * exner_pr + exner_pr_wp = exner - exner_ref_mc_wp return astype(z_exner_ex_pr_wp, vpfloat), exner_pr_wp From d04db94c2c3c0646842d46362a9febdab48bd8aa Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 24 Nov 2025 08:21:21 +0100 Subject: [PATCH 095/108] reverted file --- .../test_compute_hydrostatic_correction_term.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py index 79217ba75d..c1305ac0c0 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_hydrostatic_correction_term.py @@ -83,7 +83,8 @@ def _apply_index_field( / ((z_theta1 + z_theta2) ** 2) ) - return z_hydro_corr + nlevels = theta_v.shape[1] + return z_hydro_corr[:, nlevels - 1 : nlevels] @pytest.mark.skip_value_error @@ -141,6 +142,14 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala z_hydro_corr = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) + z_hydro_corr = gtx.constructors.zeros( + domain={ + dims.EdgeDim: (0, grid.num_edges), + dims.KDim: (grid.num_levels - 1, grid.num_levels), + }, + dtype=ta.vpfloat, + ) + return dict( theta_v=theta_v, ikoffset=ikoffset, @@ -152,6 +161,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala grav_o_cpd=grav_o_cpd, horizontal_start=0, horizontal_end=gtx.int32(grid.num_edges), - vertical_start=0, + vertical_start=gtx.int32(grid.num_levels - 1), vertical_end=gtx.int32(grid.num_levels), ) From 670566676f82ae8b919d5e6ef1a68358de7b24be Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 24 Nov 2025 08:23:25 +0100 Subject: [PATCH 096/108] small edit for bounds --- .../compute_cell_diagnostics_for_dycore.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 1db3fc504f..e41240fb70 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -194,14 +194,6 @@ def _surface_computations( fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], ]: - (perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels) = ( - concat_where( - dims.KDim < surface_level - 1, - _init_two_cell_kdim_fields_with_zero_vp(), - (perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), - ) - ) - (temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels) = ( concat_where( dims.KDim < surface_level - 1, @@ -447,6 +439,14 @@ def compute_perturbed_quantities_and_interpolation( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_level - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels """ + _init_two_cell_kdim_fields_with_zero_vp( + out=(perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), + domain={ + dims.CellDim: (start_cell_lateral_boundary, start_cell_lateral_boundary_level_3), + dims.KDim: (model_top, surface_level - 1), + }, + ) + _surface_computations( wgtfacq_c=wgtfacq_c, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, From 45bdafff89d9a610bf3404b009cc46c4a26aec8d Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 22 Dec 2025 11:46:08 +0100 Subject: [PATCH 097/108] removed double field return --- .../dycore/stencils/compute_cell_diagnostics_for_dycore.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index e41240fb70..eeda717569 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -192,7 +192,6 @@ def _surface_computations( fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], ]: (temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels) = ( concat_where( @@ -226,7 +225,6 @@ def _surface_computations( return ( temporal_extrapolation_of_perturbed_exner, exner_at_cells_on_half_levels, - temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels, perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -462,7 +460,6 @@ def compute_perturbed_quantities_and_interpolation( out=( temporal_extrapolation_of_perturbed_exner, exner_at_cells_on_half_levels, - temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels, perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, From 1d88f3c334a44214345f090291e09a639a06fc6f Mon Sep 17 00:00:00 2001 From: edopao Date: Tue, 23 Dec 2025 09:47:39 +0100 Subject: [PATCH 098/108] Remove write-write pattern in combined_1_to_13 (#985) --- .../compute_cell_diagnostics_for_dycore.py | 199 +++++------------- .../init_cell_kdim_field_with_zero_vp.py | 35 --- .../init_two_cell_kdim_fields_with_zero_vp.py | 40 ---- .../integration_tests/test_solve_nonhydro.py | 2 +- .../test_init_cell_kdim_field_with_zero_vp.py | 55 ----- ..._init_two_cell_kdim_fields_with_zero_vp.py | 55 ----- 6 files changed, 49 insertions(+), 337 deletions(-) delete mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py delete mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_with_zero_vp.py delete mode 100644 model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_vp.py delete mode 100644 model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index eeda717569..5bbe755279 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -32,12 +32,6 @@ from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( _extrapolate_temporally_exner_pressure, ) -from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import ( - _init_cell_kdim_field_with_zero_wp, -) -from icon4py.model.atmosphere.dycore.stencils.init_two_cell_kdim_fields_with_zero_vp import ( - _init_two_cell_kdim_fields_with_zero_vp, -) from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import _interpolate_to_surface from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import Koff @@ -173,89 +167,6 @@ def _compute_perturbed_quantities_and_interpolation( ) -@gtx.field_operator -def _surface_computations( - wgtfacq_c: fa.CellKField[ta.wpfloat], - exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - time_extrapolation_parameter_for_exner: fa.CellKField[ta.vpfloat], - current_exner: fa.CellKField[ta.vpfloat], - reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], - perturbed_rho_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - igradp_method: gtx.int32, - surface_level: gtx.int32, -) -> tuple[ - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKField[ta.vpfloat], -]: - (temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels) = ( - concat_where( - dims.KDim < surface_level - 1, - _extrapolate_temporally_exner_pressure( - exner_exfac=time_extrapolation_parameter_for_exner, - exner=current_exner, - exner_ref_mc=reference_exner_at_cells_on_model_levels, - exner_pr=perturbed_exner_at_cells_on_model_levels, - ), - (temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels), - ) - ) - - temporal_extrapolation_of_perturbed_exner = concat_where( - dims.KDim == surface_level - 1, - _init_cell_kdim_field_with_zero_wp(), - temporal_extrapolation_of_perturbed_exner, - ) - - exner_at_cells_on_half_levels = concat_where( - dims.KDim == surface_level - 1, - _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else exner_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ) - - return ( - temporal_extrapolation_of_perturbed_exner, - exner_at_cells_on_half_levels, - perturbed_exner_at_cells_on_model_levels, - perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels, - ) - - -@gtx.field_operator -def _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], -) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: - perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels - ) - theta_v_at_cells_on_half_levels = ( - reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels - ) - - exner_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - - return ( - perturbed_theta_v_at_cells_on_half_levels, - astype(theta_v_at_cells_on_half_levels, wpfloat), - exner_at_cells_on_half_levels, - ) - - @gtx.field_operator def _compute_first_and_second_vertical_derivative_of_exner( exner_at_cells_on_half_levels: fa.CellKField[vpfloat], @@ -266,43 +177,16 @@ def _compute_first_and_second_vertical_derivative_of_exner( d2dexdz2_fac1_mc: fa.CellKField[vpfloat], d2dexdz2_fac2_mc: fa.CellKField[vpfloat], perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, - surface_level: gtx.int32, ) -> tuple[ fa.CellKField[vpfloat], fa.CellKField[vpfloat], - fa.CellKField[vpfloat], - fa.CellKField[vpfloat], - fa.CellKField[vpfloat], ]: - ( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ) = concat_where( - dims.KDim == surface_level - 1, - _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - ), - ( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ), - ) - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflatlev <= dims.KDim) & (dims.KDim < surface_level - 1), + nflatlev <= dims.KDim, _compute_first_vertical_derivative_at_cells( exner_at_cells_on_half_levels, inv_ddqz_z_full ), @@ -314,7 +198,7 @@ def _compute_first_and_second_vertical_derivative_of_exner( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflat_gradp <= dims.KDim) & (dims.KDim < surface_level - 1), + nflat_gradp <= dims.KDim, -vpfloat("0.5") * ( ( @@ -333,8 +217,30 @@ def _compute_first_and_second_vertical_derivative_of_exner( return ( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + ) + + +@gtx.field_operator +def _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], +) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: + perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels + ) + theta_v_at_cells_on_half_levels = ( + reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels + ) + + exner_at_cells_on_half_levels = _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ) + + return ( perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, + astype(theta_v_at_cells_on_half_levels, wpfloat), exner_at_cells_on_half_levels, ) @@ -437,36 +343,19 @@ def compute_perturbed_quantities_and_interpolation( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_level - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels """ - _init_two_cell_kdim_fields_with_zero_vp( - out=(perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels), - domain={ - dims.CellDim: (start_cell_lateral_boundary, start_cell_lateral_boundary_level_3), - dims.KDim: (model_top, surface_level - 1), - }, - ) - _surface_computations( - wgtfacq_c=wgtfacq_c, - exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - time_extrapolation_parameter_for_exner=time_extrapolation_parameter_for_exner, - current_exner=current_exner, - reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, - perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - perturbed_rho_at_cells_on_model_levels=perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - igradp_method=igradp_method, - surface_level=surface_level, + _extrapolate_temporally_exner_pressure( + exner_exfac=time_extrapolation_parameter_for_exner, + exner=current_exner, + exner_ref_mc=reference_exner_at_cells_on_model_levels, + exner_pr=perturbed_exner_at_cells_on_model_levels, out=( temporal_extrapolation_of_perturbed_exner, - exner_at_cells_on_half_levels, perturbed_exner_at_cells_on_model_levels, - perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level), + dims.KDim: (model_top, surface_level - 1), }, ) @@ -503,6 +392,22 @@ def compute_perturbed_quantities_and_interpolation( }, ) + _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + wgtfacq_c=wgtfacq_c, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + out=( + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + ), + domain={ + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (surface_level - 1, surface_level), + }, + ) + _compute_first_and_second_vertical_derivative_of_exner( exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, inv_ddqz_z_full=inv_ddqz_z_full, @@ -512,24 +417,16 @@ def compute_perturbed_quantities_and_interpolation( d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, - surface_level=surface_level, out=( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level), + dims.KDim: (model_top, surface_level - 1), }, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py deleted file mode 100644 index a9b6447da5..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_cell_kdim_field_with_zero_vp.py +++ /dev/null @@ -1,35 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import broadcast - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.type_alias import vpfloat - - -@gtx.field_operator -def _init_cell_kdim_field_with_zero_vp() -> fa.CellKField[vpfloat]: - """Formerly known as _mo_solve_nonhydro_stencil_03, _mo_solve_nonhydro_stencil_11_lower, _mo_solve_nonhydro_stencil_45, _mo_solve_nonhydro_stencil_45_b, or _mo_velocity_advection_stencil_12.""" - return broadcast(vpfloat("0.0"), (dims.CellDim, dims.KDim)) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def init_cell_kdim_field_with_zero_vp( - field_with_zero_vp: fa.CellKField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -): - _init_cell_kdim_field_with_zero_vp( - out=field_with_zero_vp, - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_with_zero_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_with_zero_vp.py deleted file mode 100644 index 371b2cfcc3..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/init_two_cell_kdim_fields_with_zero_vp.py +++ /dev/null @@ -1,40 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx - -from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_vp import ( - _init_cell_kdim_field_with_zero_vp, -) -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.type_alias import vpfloat - - -@gtx.field_operator -def _init_two_cell_kdim_fields_with_zero_vp() -> ( - tuple[fa.CellKField[vpfloat], fa.CellKField[vpfloat]] -): - """Formerly known as _mo_solve_nonhydro_stencil_01.""" - return _init_cell_kdim_field_with_zero_vp(), _init_cell_kdim_field_with_zero_vp() - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def init_two_cell_kdim_fields_with_zero_vp( - cell_kdim_field_with_zero_vp_1: fa.CellKField[vpfloat], - cell_kdim_field_with_zero_vp_2: fa.CellKField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -): - _init_two_cell_kdim_fields_with_zero_vp( - out=(cell_kdim_field_with_zero_vp_1, cell_kdim_field_with_zero_vp_2), - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 7f97ee29b7..9d036565e5 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -354,7 +354,7 @@ def test_nonhydro_predictor_step( edge_start_nudging_level_2:, : ], sp_exit.z_gradh_exner().asnumpy()[edge_start_nudging_level_2:, :], - atol=1e-7, # 20 + atol=1e-20, ) prognostic_state_nnew = prognostic_states.next vn_new_reference = sp_exit.vn_new().asnumpy() diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_vp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_vp.py deleted file mode 100644 index 01438aa31d..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_cell_kdim_field_with_zero_vp.py +++ /dev/null @@ -1,55 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_vp import ( - init_cell_kdim_field_with_zero_vp, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest - - -def init_cell_kdim_field_with_zero_vp_numpy(field_with_zero_vp: np.ndarray) -> np.ndarray: - field_with_zero_vp = np.zeros_like(field_with_zero_vp) - return field_with_zero_vp - - -class TestInitCellKdimFieldWithZeroVp(StencilTest): - PROGRAM = init_cell_kdim_field_with_zero_vp - OUTPUTS = ("field_with_zero_vp",) - - @staticmethod - def reference( - connectivities: dict[gtx.Dimension, np.ndarray], - field_with_zero_vp: np.ndarray, - **kwargs: Any, - ) -> dict: - field_with_zero_vp = init_cell_kdim_field_with_zero_vp_numpy( - field_with_zero_vp=field_with_zero_vp - ) - return dict(field_with_zero_vp=field_with_zero_vp) - - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - field_with_zero_vp = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - - return dict( - field_with_zero_vp=field_with_zero_vp, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py deleted file mode 100644 index b893924387..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_init_two_cell_kdim_fields_with_zero_vp.py +++ /dev/null @@ -1,55 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.init_two_cell_kdim_fields_with_zero_vp import ( - init_two_cell_kdim_fields_with_zero_vp, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat -from icon4py.model.common.utils.data_allocation import random_field -from icon4py.model.testing.stencil_tests import StencilTest - - -class TestInitTwoCellKdimFieldsWithZeroVp(StencilTest): - PROGRAM = init_two_cell_kdim_fields_with_zero_vp - OUTPUTS = ("cell_kdim_field_with_zero_vp_1", "cell_kdim_field_with_zero_vp_2") - - @staticmethod - def reference( - connectivities: dict[gtx.Dimension, np.ndarray], - cell_kdim_field_with_zero_vp_1: np.ndarray, - cell_kdim_field_with_zero_vp_2: np.ndarray, - **kwargs: Any, - ) -> dict: - cell_kdim_field_with_zero_vp_1 = np.zeros_like(cell_kdim_field_with_zero_vp_1) - cell_kdim_field_with_zero_vp_2 = np.zeros_like(cell_kdim_field_with_zero_vp_2) - return dict( - cell_kdim_field_with_zero_vp_1=cell_kdim_field_with_zero_vp_1, - cell_kdim_field_with_zero_vp_2=cell_kdim_field_with_zero_vp_2, - ) - - @pytest.fixture - def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - cell_kdim_field_with_zero_vp_1 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - cell_kdim_field_with_zero_vp_2 = random_field(grid, dims.CellDim, dims.KDim, dtype=vpfloat) - - return dict( - cell_kdim_field_with_zero_vp_1=cell_kdim_field_with_zero_vp_1, - cell_kdim_field_with_zero_vp_2=cell_kdim_field_with_zero_vp_2, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) From 1ecd4c1ca49a90a51e519b33d2004520f35605ec Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 5 Jan 2026 11:04:01 +0100 Subject: [PATCH 099/108] some further inlining --- .../compute_cell_diagnostics_for_dycore.py | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 5bbe755279..46de7795e7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -68,6 +68,9 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( @gtx.field_operator def _compute_perturbed_quantities_and_interpolation( + time_extrapolation_parameter_for_exner: fa.CellKField[ta.vpfloat], + current_exner: fa.CellKField[ta.wpfloat], + reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], current_rho: fa.CellKField[ta.wpfloat], reference_rho_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], current_theta_v: fa.CellKField[ta.wpfloat], @@ -80,7 +83,7 @@ def _compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], + # temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, @@ -93,7 +96,19 @@ def _compute_perturbed_quantities_and_interpolation( fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], ]: + ( + temporal_extrapolation_of_perturbed_exner, + perturbed_exner_at_cells_on_model_levels, + ) = _extrapolate_temporally_exner_pressure( + exner_exfac=time_extrapolation_parameter_for_exner, + exner=current_exner, + exner_ref_mc=reference_exner_at_cells_on_model_levels, + exner_pr=perturbed_exner_at_cells_on_model_levels, + ) + exner_at_cells_on_half_levels = ( concat_where( maximum(1, nflatlev) <= dims.KDim, @@ -164,6 +179,8 @@ def _compute_perturbed_quantities_and_interpolation( perturbed_theta_v_at_cells_on_half_levels, theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, + temporal_extrapolation_of_perturbed_exner, + perturbed_exner_at_cells_on_model_levels, ) @@ -344,22 +361,10 @@ def compute_perturbed_quantities_and_interpolation( - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels """ - _extrapolate_temporally_exner_pressure( - exner_exfac=time_extrapolation_parameter_for_exner, - exner=current_exner, - exner_ref_mc=reference_exner_at_cells_on_model_levels, - exner_pr=perturbed_exner_at_cells_on_model_levels, - out=( - temporal_extrapolation_of_perturbed_exner, - perturbed_exner_at_cells_on_model_levels, - ), - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), - }, - ) - _compute_perturbed_quantities_and_interpolation( + time_extrapolation_parameter_for_exner=time_extrapolation_parameter_for_exner, + current_exner=current_exner, + reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels, current_rho=current_rho, reference_rho_at_cells_on_model_levels=reference_rho_at_cells_on_model_levels, current_theta_v=current_theta_v, @@ -372,7 +377,6 @@ def compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, @@ -385,6 +389,8 @@ def compute_perturbed_quantities_and_interpolation( perturbed_theta_v_at_cells_on_half_levels, theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, + temporal_extrapolation_of_perturbed_exner, + perturbed_exner_at_cells_on_model_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), From 7f8a6e00a62544753f5d587e2591d7213a0b6501 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 5 Jan 2026 13:18:50 +0100 Subject: [PATCH 100/108] further inlining --- .../compute_cell_diagnostics_for_dycore.py | 110 ++++++++++-------- 1 file changed, 64 insertions(+), 46 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 46de7795e7..3141081764 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -83,7 +83,6 @@ def _compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - # temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, @@ -97,7 +96,6 @@ def _compute_perturbed_quantities_and_interpolation( fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], ]: ( temporal_extrapolation_of_perturbed_exner, @@ -180,7 +178,31 @@ def _compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, temporal_extrapolation_of_perturbed_exner, - perturbed_exner_at_cells_on_model_levels, + ) + + +@gtx.field_operator +def _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], +) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: + perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels + ) + theta_v_at_cells_on_half_levels = ( + reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels + ) + + exner_at_cells_on_half_levels = _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ) + + return ( + perturbed_theta_v_at_cells_on_half_levels, + astype(theta_v_at_cells_on_half_levels, wpfloat), + exner_at_cells_on_half_levels, ) @@ -194,16 +216,43 @@ def _compute_first_and_second_vertical_derivative_of_exner( d2dexdz2_fac1_mc: fa.CellKField[vpfloat], d2dexdz2_fac2_mc: fa.CellKField[vpfloat], perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], + temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], + wgtfacq_c: fa.CellKField[vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], + theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, + surface_level: gtx.int32, ) -> tuple[ fa.CellKField[vpfloat], fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], + fa.CellKField[vpfloat], ]: + ( + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + ) = concat_where( + dims.KDim >= surface_level - 1, + _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + wgtfacq_c=wgtfacq_c, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + ), + ( + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + ), + ) + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - nflatlev <= dims.KDim, + (nflatlev <= dims.KDim) & (dims.KDim < surface_level - 1), _compute_first_vertical_derivative_at_cells( exner_at_cells_on_half_levels, inv_ddqz_z_full ), @@ -215,7 +264,7 @@ def _compute_first_and_second_vertical_derivative_of_exner( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - nflat_gradp <= dims.KDim, + (nflat_gradp <= dims.KDim) & (dims.KDim < surface_level - 1), -vpfloat("0.5") * ( ( @@ -234,30 +283,8 @@ def _compute_first_and_second_vertical_derivative_of_exner( return ( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ) - - -@gtx.field_operator -def _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], -) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: - perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels - ) - theta_v_at_cells_on_half_levels = ( - reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels - ) - - exner_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - - return ( perturbed_theta_v_at_cells_on_half_levels, - astype(theta_v_at_cells_on_half_levels, wpfloat), + theta_v_at_cells_on_half_levels, exner_at_cells_on_half_levels, ) @@ -390,7 +417,6 @@ def compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, temporal_extrapolation_of_perturbed_exner, - perturbed_exner_at_cells_on_model_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), @@ -398,22 +424,6 @@ def compute_perturbed_quantities_and_interpolation( }, ) - _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - out=( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ), - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (surface_level - 1, surface_level), - }, - ) - _compute_first_and_second_vertical_derivative_of_exner( exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, inv_ddqz_z_full=inv_ddqz_z_full, @@ -423,16 +433,24 @@ def compute_perturbed_quantities_and_interpolation( d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + wgtfacq_c=wgtfacq_c, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, + surface_level=surface_level, out=( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), + dims.KDim: (model_top, surface_level), }, ) From 83315ec0926b803c6a5a9f3894453f2b0bc34c30 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 8 Jan 2026 11:21:24 +0100 Subject: [PATCH 101/108] fixed small mistake --- .../compute_cell_diagnostics_for_dycore.py | 59 ++++++------------- 1 file changed, 19 insertions(+), 40 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 3141081764..7e91523437 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -216,43 +216,17 @@ def _compute_first_and_second_vertical_derivative_of_exner( d2dexdz2_fac1_mc: fa.CellKField[vpfloat], d2dexdz2_fac2_mc: fa.CellKField[vpfloat], perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, - surface_level: gtx.int32, ) -> tuple[ fa.CellKField[vpfloat], fa.CellKField[vpfloat], - fa.CellKField[vpfloat], - fa.CellKField[vpfloat], - fa.CellKField[vpfloat], ]: - ( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ) = concat_where( - dims.KDim >= surface_level - 1, - _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - ), - ( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, - ), - ) ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflatlev <= dims.KDim) & (dims.KDim < surface_level - 1), + nflatlev <= dims.KDim, _compute_first_vertical_derivative_at_cells( exner_at_cells_on_half_levels, inv_ddqz_z_full ), @@ -264,7 +238,7 @@ def _compute_first_and_second_vertical_derivative_of_exner( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( - (nflat_gradp <= dims.KDim) & (dims.KDim < surface_level - 1), + nflat_gradp <= dims.KDim, -vpfloat("0.5") * ( ( @@ -283,9 +257,6 @@ def _compute_first_and_second_vertical_derivative_of_exner( return ( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, ) @@ -424,6 +395,22 @@ def compute_perturbed_quantities_and_interpolation( }, ) + _set_theta_v_and_exner_on_surface_level( + temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, + wgtfacq_c=wgtfacq_c, + perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + out = ( + perturbed_theta_v_at_cells_on_half_levels, + theta_v_at_cells_on_half_levels, + exner_at_cells_on_half_levels, + ), + domain={ + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (surface_level - 1, surface_level), + }, + ) + _compute_first_and_second_vertical_derivative_of_exner( exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, inv_ddqz_z_full=inv_ddqz_z_full, @@ -433,24 +420,16 @@ def compute_perturbed_quantities_and_interpolation( d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, - surface_level=surface_level, out=( ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level), + dims.KDim: (model_top, surface_level-1), }, ) From b3f471fdf47dce651da241cca4b97173a827acb5 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 8 Jan 2026 13:41:17 +0100 Subject: [PATCH 102/108] extende concat_where --- .../compute_cell_diagnostics_for_dycore.py | 86 ++++++++++++------- 1 file changed, 56 insertions(+), 30 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 7e91523437..7db838afe8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -22,7 +22,7 @@ from typing import Final import gt4py.next as gtx -from gt4py.next import astype, broadcast, maximum +from gt4py.next import astype, maximum from gt4py.next.experimental import concat_where from icon4py.model.atmosphere.dycore.dycore_states import HorizontalPressureDiscretizationType @@ -83,9 +83,11 @@ def _compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], + wgtfacq_c: fa.CellKField[ta.vpfloat], + reference_theta_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], igradp_method: gtx.int32, nflatlev: gtx.int32, + surface_level: gtx.int32, ) -> tuple[ fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], @@ -109,11 +111,13 @@ def _compute_perturbed_quantities_and_interpolation( exner_at_cells_on_half_levels = ( concat_where( - maximum(1, nflatlev) <= dims.KDim, + dims.KDim < surface_level - 1, _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner ), - exner_at_cells_on_half_levels, + _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner + ), ) if igradp_method == horzpres_discr_type.TAYLOR_HYDRO else exner_at_cells_on_half_levels @@ -138,19 +142,22 @@ def _compute_perturbed_quantities_and_interpolation( wgtfac_c_wp = astype(wgtfac_c, wpfloat) perturbed_theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + dims.KDim < surface_level - 1, _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels ), - broadcast(0.0, (dims.CellDim, dims.KDim)), + _interpolate_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels + ), ) theta_v_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + dims.KDim < surface_level - 1, _interpolate_cell_field_to_half_levels_wp( wgtfac_c=wgtfac_c_wp, interpolant=current_theta_v ), - theta_v_at_cells_on_half_levels, + reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels, + # theta_v_at_cells_on_half_levels, ) ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) @@ -223,7 +230,6 @@ def _compute_first_and_second_vertical_derivative_of_exner( fa.CellKField[vpfloat], fa.CellKField[vpfloat], ]: - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( concat_where( nflatlev <= dims.KDim, @@ -375,9 +381,11 @@ def compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels, + wgtfacq_c=wgtfacq_c, + reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, igradp_method=igradp_method, nflatlev=nflatlev, + surface_level=surface_level, out=( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -389,26 +397,44 @@ def compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels, temporal_extrapolation_of_perturbed_exner, ), - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), - }, - ) - - _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner, - wgtfacq_c=wgtfacq_c, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, - out = ( - perturbed_theta_v_at_cells_on_half_levels, - theta_v_at_cells_on_half_levels, - exner_at_cells_on_half_levels, + domain=( + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (maximum(1, nflatlev), surface_level), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (1, surface_level), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (1, surface_level), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (model_top, surface_level - 1), + }, ), - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (surface_level - 1, surface_level), - }, ) _compute_first_and_second_vertical_derivative_of_exner( @@ -429,7 +455,7 @@ def compute_perturbed_quantities_and_interpolation( ), domain={ dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level-1), + dims.KDim: (model_top, surface_level - 1), }, ) From 33ebde4ecd40168b21ed53774c0e06867516509b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 8 Jan 2026 14:24:27 +0100 Subject: [PATCH 103/108] further inlining --- .../compute_cell_diagnostics_for_dycore.py | 73 ++++++++++++------- 1 file changed, 48 insertions(+), 25 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 7db838afe8..235d07186c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -85,8 +85,12 @@ def _compute_perturbed_quantities_and_interpolation( exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], wgtfacq_c: fa.CellKField[ta.vpfloat], reference_theta_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], + inv_ddqz_z_full: fa.CellKField[ta.vpfloat], + d2dexdz2_fac1_mc: fa.CellKField[ta.vpfloat], + d2dexdz2_fac2_mc: fa.CellKField[ta.vpfloat], + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], igradp_method: gtx.int32, - nflatlev: gtx.int32, surface_level: gtx.int32, ) -> tuple[ fa.CellKField[ta.wpfloat], @@ -98,6 +102,8 @@ def _compute_perturbed_quantities_and_interpolation( fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], + fa.CellKField[ta.vpfloat], + fa.CellKField[ta.vpfloat], ]: ( temporal_extrapolation_of_perturbed_exner, @@ -157,7 +163,6 @@ def _compute_perturbed_quantities_and_interpolation( wgtfac_c=wgtfac_c_wp, interpolant=current_theta_v ), reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels, - # theta_v_at_cells_on_half_levels, ) ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) @@ -175,6 +180,30 @@ def _compute_perturbed_quantities_and_interpolation( pressure_buoyancy_acceleration_at_cells_on_half_levels, ) + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( + ( + _compute_first_vertical_derivative_at_cells( + exner_at_cells_on_half_levels, inv_ddqz_z_full + ) + ) + if igradp_method == horzpres_discr_type.TAYLOR_HYDRO + else ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels + ) + + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( + -vpfloat("0.5") + * ( + ( + perturbed_theta_v_at_cells_on_half_levels + - perturbed_theta_v_at_cells_on_half_levels(Koff[1]) + ) + * d2dexdz2_fac1_mc + + perturbed_theta_v_at_cells_on_model_levels * d2dexdz2_fac2_mc + ) + if igradp_method == horzpres_discr_type.TAYLOR_HYDRO + else d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels + ) + return ( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -185,6 +214,8 @@ def _compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, temporal_extrapolation_of_perturbed_exner, + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, ) @@ -383,8 +414,12 @@ def compute_perturbed_quantities_and_interpolation( exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, wgtfacq_c=wgtfacq_c, reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, + inv_ddqz_z_full=inv_ddqz_z_full, + d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, + d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, igradp_method=igradp_method, - nflatlev=nflatlev, surface_level=surface_level, out=( perturbed_rho_at_cells_on_model_levels, @@ -396,6 +431,8 @@ def compute_perturbed_quantities_and_interpolation( theta_v_at_cells_on_half_levels, pressure_buoyancy_acceleration_at_cells_on_half_levels, temporal_extrapolation_of_perturbed_exner, + ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, + d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, ), domain=( { @@ -434,31 +471,17 @@ def compute_perturbed_quantities_and_interpolation( dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), dims.KDim: (model_top, surface_level - 1), }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (nflatlev, surface_level - 1), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), + dims.KDim: (nflat_gradp, surface_level - 1), + }, ), ) - _compute_first_and_second_vertical_derivative_of_exner( - exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, - inv_ddqz_z_full=inv_ddqz_z_full, - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - perturbed_theta_v_at_cells_on_half_levels=perturbed_theta_v_at_cells_on_half_levels, - d2dexdz2_fac1_mc=d2dexdz2_fac1_mc, - d2dexdz2_fac2_mc=d2dexdz2_fac2_mc, - perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels, - igradp_method=igradp_method, - nflatlev=nflatlev, - nflat_gradp=nflat_gradp, - out=( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ), - domain={ - dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), - }, - ) - _compute_perturbation_of_rho_and_theta( rho=current_rho, rho_ref_mc=reference_rho_at_cells_on_model_levels, From 02f4c7069d3dedb827c651734c653704ddd50d24 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 8 Jan 2026 15:17:36 +0100 Subject: [PATCH 104/108] further inlining --- .../compute_cell_diagnostics_for_dycore.py | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 235d07186c..3e2b961f6e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -80,8 +80,6 @@ def _compute_perturbed_quantities_and_interpolation( perturbed_exner_at_cells_on_model_levels: fa.CellKField[ta.wpfloat], ddz_of_reference_exner_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], ddqz_z_half: fa.CellKField[ta.vpfloat], - pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], - rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat], exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], wgtfacq_c: fa.CellKField[ta.vpfloat], reference_theta_at_cells_on_half_levels: fa.CellKField[ta.vpfloat], @@ -139,11 +137,7 @@ def _compute_perturbed_quantities_and_interpolation( reference_theta_at_cells_on_model_levels, ) - rho_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, - _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho), - rho_at_cells_on_half_levels, - ) + rho_at_cells_on_half_levels = _interpolate_cell_field_to_half_levels_wp(wgtfac_c, current_rho) wgtfac_c_wp = astype(wgtfac_c, wpfloat) @@ -167,8 +161,7 @@ def _compute_perturbed_quantities_and_interpolation( ddqz_z_half_wp = astype(ddqz_z_half, wpfloat) - pressure_buoyancy_acceleration_at_cells_on_half_levels = concat_where( - dims.KDim >= 1, + pressure_buoyancy_acceleration_at_cells_on_half_levels = ( _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels( exner_w_explicit_weight_parameter, theta_v_at_cells_on_half_levels, @@ -176,8 +169,7 @@ def _compute_perturbed_quantities_and_interpolation( ddqz_z_half_wp, perturbed_theta_v_at_cells_on_half_levels, ddz_of_reference_exner_at_cells_on_half_levels, - ), - pressure_buoyancy_acceleration_at_cells_on_half_levels, + ) ) ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( @@ -409,8 +401,6 @@ def compute_perturbed_quantities_and_interpolation( perturbed_exner_at_cells_on_model_levels=perturbed_exner_at_cells_on_model_levels, ddz_of_reference_exner_at_cells_on_half_levels=ddz_of_reference_exner_at_cells_on_half_levels, ddqz_z_half=ddqz_z_half, - pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels, - rho_at_cells_on_half_levels=rho_at_cells_on_half_levels, exner_at_cells_on_half_levels=exner_at_cells_on_half_levels, wgtfacq_c=wgtfacq_c, reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels, @@ -449,7 +439,7 @@ def compute_perturbed_quantities_and_interpolation( }, { dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), + dims.KDim: (1, surface_level - 1), }, { dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), @@ -465,7 +455,7 @@ def compute_perturbed_quantities_and_interpolation( }, { dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), - dims.KDim: (model_top, surface_level - 1), + dims.KDim: (1, surface_level - 1), }, { dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo), From a912450805addb46ebb7d5f8bdfd12c799a4dd0a Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 12 Jan 2026 09:06:48 +0100 Subject: [PATCH 105/108] cleanup --- .../compute_cell_diagnostics_for_dycore.py | 92 ------------------- .../integration_tests/test_solve_nonhydro.py | 2 - ..._perturbed_quantities_and_interpolation.py | 1 - 3 files changed, 95 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 3e2b961f6e..06d052103c 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -6,19 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - from typing import Final import gt4py.next as gtx @@ -211,84 +198,6 @@ def _compute_perturbed_quantities_and_interpolation( ) -@gtx.field_operator -def _set_theta_v_and_exner_on_surface_level( - temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat], - wgtfacq_c: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat], -) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]: - perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels - ) - theta_v_at_cells_on_half_levels = ( - reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels - ) - - exner_at_cells_on_half_levels = _interpolate_to_surface( - wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner - ) - - return ( - perturbed_theta_v_at_cells_on_half_levels, - astype(theta_v_at_cells_on_half_levels, wpfloat), - exner_at_cells_on_half_levels, - ) - - -@gtx.field_operator -def _compute_first_and_second_vertical_derivative_of_exner( - exner_at_cells_on_half_levels: fa.CellKField[vpfloat], - inv_ddqz_z_full: fa.CellKField[vpfloat], - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat], - d2dexdz2_fac1_mc: fa.CellKField[vpfloat], - d2dexdz2_fac2_mc: fa.CellKField[vpfloat], - perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat], - igradp_method: gtx.int32, - nflatlev: gtx.int32, - nflat_gradp: gtx.int32, -) -> tuple[ - fa.CellKField[vpfloat], - fa.CellKField[vpfloat], -]: - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( - concat_where( - nflatlev <= dims.KDim, - _compute_first_vertical_derivative_at_cells( - exner_at_cells_on_half_levels, inv_ddqz_z_full - ), - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ) - if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels - ) - - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = ( - concat_where( - nflat_gradp <= dims.KDim, - -vpfloat("0.5") - * ( - ( - perturbed_theta_v_at_cells_on_half_levels - - perturbed_theta_v_at_cells_on_half_levels(Koff[1]) - ) - * d2dexdz2_fac1_mc - + perturbed_theta_v_at_cells_on_model_levels * d2dexdz2_fac2_mc - ), - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ) - if igradp_method == horzpres_discr_type.TAYLOR_HYDRO - else d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels - ) - - return ( - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, - ) - - @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_perturbed_quantities_and_interpolation( temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat], @@ -321,7 +230,6 @@ def compute_perturbed_quantities_and_interpolation( igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, - start_cell_lateral_boundary: gtx.int32, start_cell_lateral_boundary_level_3: gtx.int32, start_cell_halo_level_2: gtx.int32, end_cell_halo: gtx.int32, diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 9d036565e5..21dc8e77b5 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1083,7 +1083,6 @@ def test_compute_perturbed_quantities_and_interpolation( nflat_gradp = grid_savepoint.nflat_gradp() cell_domain = h_grid.domain(dims.CellDim) - start_cell_lateral_boundary = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY)) start_cell_lateral_boundary_level_3 = icon_grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) ) @@ -1149,7 +1148,6 @@ def test_compute_perturbed_quantities_and_interpolation( igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, - start_cell_lateral_boundary=start_cell_lateral_boundary, start_cell_lateral_boundary_level_3=start_cell_lateral_boundary_level_3, start_cell_halo_level_2=start_cell_halo_level_2, end_cell_halo=end_cell_halo, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index a173e626f1..31e899b523 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -87,7 +87,6 @@ class TestComputePerturbedQuantitiesAndInterpolation(stencil_tests.StencilTest): "igradp_method", "nflatlev", "nflat_gradp", - "start_cell_lateral_boundary", "start_cell_lateral_boundary_level_3", "start_cell_halo_level_2", "end_cell_halo", From 924e7743925181ec93a055d82eb18bf488850904 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 12 Jan 2026 09:46:44 +0100 Subject: [PATCH 106/108] merged last field_operator in --- .../model/atmosphere/dycore/solve_nonhydro.py | 1 - .../compute_cell_diagnostics_for_dycore.py | 36 +++++++++++-------- .../integration_tests/test_solve_nonhydro.py | 4 +++ 3 files changed, 25 insertions(+), 16 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py index 8b939ed6d6..5871ca0a82 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py @@ -726,7 +726,6 @@ def __init__( "start_cell_halo_level_2": self._start_cell_halo_level_2, "end_cell_halo": self._end_cell_halo, "end_cell_halo_level_2": self._end_cell_halo_level_2, - "start_cell_lateral_boundary": self._start_cell_lateral_boundary, }, vertical_sizes={ "nflatlev": self._vertical_params.nflatlev, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 06d052103c..81881a56fe 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -77,6 +77,8 @@ def _compute_perturbed_quantities_and_interpolation( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels: fa.CellKField[ta.vpfloat], igradp_method: gtx.int32, surface_level: gtx.int32, + start_cell_halo_level_2: gtx.int32, + end_cell_halo_level_2: gtx.int32, ) -> tuple[ fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], @@ -183,6 +185,23 @@ def _compute_perturbed_quantities_and_interpolation( else d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels ) + ( + perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels, + ) = concat_where( + (dims.CellDim >= start_cell_halo_level_2) & (dims.CellDim < end_cell_halo_level_2), + _compute_perturbation_of_rho_and_theta( + rho=current_rho, + rho_ref_mc=reference_rho_at_cells_on_model_levels, + theta_v=current_theta_v, + theta_ref_mc=reference_theta_at_cells_on_model_levels, + ), + ( + perturbed_rho_at_cells_on_model_levels, + perturbed_theta_v_at_cells_on_model_levels, + ), + ) + return ( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -319,6 +338,8 @@ def compute_perturbed_quantities_and_interpolation( d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels=d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels, igradp_method=igradp_method, surface_level=surface_level, + start_cell_halo_level_2=start_cell_halo_level_2, + end_cell_halo_level_2=end_cell_halo_level_2, out=( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels, @@ -380,21 +401,6 @@ def compute_perturbed_quantities_and_interpolation( ), ) - _compute_perturbation_of_rho_and_theta( - rho=current_rho, - rho_ref_mc=reference_rho_at_cells_on_model_levels, - theta_v=current_theta_v, - theta_ref_mc=reference_theta_at_cells_on_model_levels, - out=( - perturbed_rho_at_cells_on_model_levels, - perturbed_theta_v_at_cells_on_model_levels, - ), - domain={ - dims.CellDim: (start_cell_halo_level_2, end_cell_halo_level_2), - dims.KDim: (model_top, surface_level - 1), - }, - ) - @gtx.field_operator def _interpolate_rho_theta_v_to_half_levels_and_compute_pressure_buoyancy_acceleration( diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 21dc8e77b5..0afefd5a53 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1089,6 +1089,10 @@ def test_compute_perturbed_quantities_and_interpolation( start_cell_halo_level_2 = icon_grid.start_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) end_cell_halo = icon_grid.end_index(cell_domain(h_grid.Zone.HALO)) end_cell_halo_level_2 = icon_grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)) + print("start_cell_lateral_boundary_level_3", start_cell_lateral_boundary_level_3) + print("end_cell_halo", end_cell_halo) + print("start_cell_halo_level_2", start_cell_halo_level_2) + print("end_cell_halo_level_2", end_cell_halo_level_2) reference_rho_at_cells_on_model_levels = metrics_savepoint.rho_ref_mc() reference_theta_at_cells_on_model_levels = metrics_savepoint.theta_ref_mc() From 73921e72afc6f5e976876f4e7e69dbf1eefabeac Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 12 Jan 2026 10:44:51 +0100 Subject: [PATCH 107/108] edit to stencil test --- ...e_perturbed_quantities_and_interpolation.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 31e899b523..8e42194e26 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -6,18 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later from typing import Any import gt4py.next as gtx @@ -136,7 +124,6 @@ def reference( igradp_method: gtx.int32, nflatlev: gtx.int32, nflat_gradp: gtx.int32, - start_cell_lateral_boundary: gtx.int32, start_cell_lateral_boundary_level_3: gtx.int32, start_cell_halo_level_2: gtx.int32, end_cell_halo: gtx.int32, @@ -152,8 +139,7 @@ def reference( perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels[:, : surface_level - 1], ) = np.where( - (start_cell_lateral_boundary <= horz_idx) - & (horz_idx < start_cell_lateral_boundary_level_3), + horz_idx < start_cell_lateral_boundary_level_3, ( np.zeros_like(perturbed_rho_at_cells_on_model_levels), np.zeros_like(perturbed_theta_v_at_cells_on_model_levels[:, : surface_level - 1]), @@ -411,7 +397,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala igradp_method = horzpres_discr_type.TAYLOR_HYDRO cell_domain = h_grid.domain(dims.CellDim) - start_cell_lateral_boundary = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY)) start_cell_lateral_boundary_level_3 = grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) ) @@ -453,7 +438,6 @@ def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.Scala igradp_method=igradp_method, nflatlev=nflatlev, nflat_gradp=nflat_gradp, - start_cell_lateral_boundary=start_cell_lateral_boundary, start_cell_lateral_boundary_level_3=start_cell_lateral_boundary_level_3, start_cell_halo_level_2=start_cell_halo_level_2, end_cell_halo=end_cell_halo, From 05a535ec578b9e47cc76fa32db6c449e87301204 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 13 Jan 2026 11:05:47 +0100 Subject: [PATCH 108/108] renaming of interpolate_to_surface to extrapolate_quadratically_to_surface --- .../stencils/compute_cell_diagnostics_for_dycore.py | 8 +++++--- ...face.py => extrapolate_quadratically_to_surface.py} | 10 +++++----- .../stencils/set_theta_v_prime_ic_at_lower_boundary.py | 8 ++++++-- ...t_compute_perturbed_quantities_and_interpolation.py | 4 ++-- ...py => test_extrapolate_quadratically_to_surface.py} | 10 ++++++---- .../test_set_theta_v_prime_ic_at_lower_boundary.py | 4 ++-- 6 files changed, 26 insertions(+), 18 deletions(-) rename model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/{interpolate_to_surface.py => extrapolate_quadratically_to_surface.py} (83%) rename model/atmosphere/dycore/tests/dycore/stencil_tests/{test_interpolate_to_surface.py => test_extrapolate_quadratically_to_surface.py} (88%) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py index 81881a56fe..28dc64a64a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_dycore.py @@ -16,10 +16,12 @@ from icon4py.model.atmosphere.dycore.stencils.compute_perturbation_of_rho_and_theta import ( _compute_perturbation_of_rho_and_theta, ) +from icon4py.model.atmosphere.dycore.stencils.extrapolate_quadratically_to_surface import ( + _extrapolate_quadratically_to_surface, +) from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import ( _extrapolate_temporally_exner_pressure, ) -from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import _interpolate_to_surface from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import Koff from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_half_levels_vp import ( @@ -108,7 +110,7 @@ def _compute_perturbed_quantities_and_interpolation( _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner ), - _interpolate_to_surface( + _extrapolate_quadratically_to_surface( wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner ), ) @@ -135,7 +137,7 @@ def _compute_perturbed_quantities_and_interpolation( _interpolate_cell_field_to_half_levels_vp( wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels ), - _interpolate_to_surface( + _extrapolate_quadratically_to_surface( wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels ), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_quadratically_to_surface.py similarity index 83% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py rename to model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_quadratically_to_surface.py index 9395878611..7670c5a48b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_surface.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/extrapolate_quadratically_to_surface.py @@ -13,21 +13,21 @@ @gtx.field_operator -def _interpolate_to_surface( +def _extrapolate_quadratically_to_surface( wgtfacq_c: fa.CellKField[vpfloat], interpolant: fa.CellKField[vpfloat], ) -> fa.CellKField[vpfloat]: """Formerly known as _mo_solve_nonhydro_stencil_04.""" - interpolation_to_surface = ( + extrapolate_quadratically_to_surface = ( # rename extrapolation_to_surface wgtfacq_c(Koff[-1]) * interpolant(Koff[-1]) + wgtfacq_c(Koff[-2]) * interpolant(Koff[-2]) + wgtfacq_c(Koff[-3]) * interpolant(Koff[-3]) ) - return interpolation_to_surface + return extrapolate_quadratically_to_surface @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def interpolate_to_surface( +def extrapolate_quadratically_to_surface( wgtfacq_c: fa.CellKField[vpfloat], interpolant: fa.CellKField[vpfloat], interpolation_to_surface: fa.CellKField[vpfloat], @@ -36,7 +36,7 @@ def interpolate_to_surface( vertical_start: gtx.int32, vertical_end: gtx.int32, ): - _interpolate_to_surface( + _extrapolate_quadratically_to_surface( wgtfacq_c, interpolant, out=interpolation_to_surface, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py index 8f48be0c81..58aa6ba91b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/set_theta_v_prime_ic_at_lower_boundary.py @@ -8,7 +8,9 @@ import gt4py.next as gtx from gt4py.next import astype -from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import _interpolate_to_surface +from icon4py.model.atmosphere.dycore.stencils.extrapolate_quadratically_to_surface import ( + _extrapolate_quadratically_to_surface, +) from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -20,7 +22,9 @@ def _set_theta_v_prime_ic_at_lower_boundary( theta_ref_ic: fa.CellKField[vpfloat], ) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat]]: """Formerly known as _mo_solve_nonhydro_stencil_11_upper.""" - z_theta_v_pr_ic_vp = _interpolate_to_surface(wgtfacq_c=wgtfacq_c, interpolant=z_rth_pr) + z_theta_v_pr_ic_vp = _extrapolate_quadratically_to_surface( + wgtfacq_c=wgtfacq_c, interpolant=z_rth_pr + ) theta_v_ic_vp = theta_ref_ic + z_theta_v_pr_ic_vp return z_theta_v_pr_ic_vp, astype(theta_v_ic_vp, wpfloat) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py index 8e42194e26..c6c9d10921 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_perturbed_quantities_and_interpolation.py @@ -32,11 +32,11 @@ from .test_compute_virtual_potential_temperatures_and_pressure_gradient import ( compute_virtual_potential_temperatures_and_pressure_gradient_numpy, ) +from .test_extrapolate_quadratically_to_surface import extrapolate_quadratically_to_surface_numpy from .test_extrapolate_temporally_exner_pressure import extrapolate_temporally_exner_pressure_numpy from .test_interpolate_cell_field_to_half_levels_vp import ( interpolate_cell_field_to_half_levels_vp_numpy, ) -from .test_interpolate_to_surface import interpolate_to_surface_numpy from .test_set_theta_v_prime_ic_at_lower_boundary import ( set_theta_v_prime_ic_at_lower_boundary_numpy, ) @@ -180,7 +180,7 @@ def reference( (start_cell_lateral_boundary_level_3 <= horz_idx) & (horz_idx < end_cell_halo) & (vert_idx == surface_level - 1), - interpolate_to_surface_numpy( + extrapolate_quadratically_to_surface_numpy( interpolant=temporal_extrapolation_of_perturbed_exner, wgtfacq_c=wgtfacq_c, interpolation_to_surface=exner_at_cells_on_half_levels, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_quadratically_to_surface.py similarity index 88% rename from model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py rename to model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_quadratically_to_surface.py index 2dc775909d..d85b91728f 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_surface.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_extrapolate_quadratically_to_surface.py @@ -11,7 +11,9 @@ import numpy as np import pytest -from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import interpolate_to_surface +from icon4py.model.atmosphere.dycore.stencils.extrapolate_quadratically_to_surface import ( + extrapolate_quadratically_to_surface, +) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils @@ -20,7 +22,7 @@ from icon4py.model.testing.stencil_tests import StencilTest -def interpolate_to_surface_numpy( +def extrapolate_quadratically_to_surface_numpy( interpolant: np.ndarray, wgtfacq_c: np.ndarray, interpolation_to_surface: np.ndarray ) -> np.ndarray: interpolation_to_surface = np.copy(interpolation_to_surface) @@ -33,7 +35,7 @@ def interpolate_to_surface_numpy( class TestInterpolateToSurface(StencilTest): - PROGRAM = interpolate_to_surface + PROGRAM = extrapolate_quadratically_to_surface OUTPUTS = ("interpolation_to_surface",) @staticmethod @@ -44,7 +46,7 @@ def reference( interpolation_to_surface: np.ndarray, **kwargs: Any, ) -> dict: - interpolation_to_surface = interpolate_to_surface_numpy( + interpolation_to_surface = extrapolate_quadratically_to_surface_numpy( wgtfacq_c=wgtfacq_c, interpolant=interpolant, interpolation_to_surface=interpolation_to_surface, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py index 122f1913ef..8198cbbf6d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_set_theta_v_prime_ic_at_lower_boundary.py @@ -21,7 +21,7 @@ from icon4py.model.common.utils.data_allocation import random_field, zero_field from icon4py.model.testing.stencil_tests import StencilTest -from .test_interpolate_to_surface import interpolate_to_surface_numpy +from .test_extrapolate_quadratically_to_surface import extrapolate_quadratically_to_surface_numpy def set_theta_v_prime_ic_at_lower_boundary_numpy( @@ -31,7 +31,7 @@ def set_theta_v_prime_ic_at_lower_boundary_numpy( z_theta_v_pr_ic: np.ndarray, theta_v_ic: np.ndarray, ) -> tuple[np.ndarray, ...]: - z_theta_v_pr_ic = interpolate_to_surface_numpy( + z_theta_v_pr_ic = extrapolate_quadratically_to_surface_numpy( wgtfacq_c=wgtfacq_c, interpolant=z_rth_pr, interpolation_to_surface=z_theta_v_pr_ic,