From c0eeef0b214189c2ddb760d442423765768203a3 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 21 Feb 2024 08:23:58 +0100 Subject: [PATCH 1/4] Fix test cases for interpolate_vn_to_ie --- ...late_vn_to_ie_and_compute_ekin_on_edges.py | 45 ++++++++++++------- .../test_interpolate_vt_to_ie.py | 29 +++++++----- 2 files changed, 47 insertions(+), 27 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py index 651b1f1684..93b4b823da 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vn_to_ie_and_compute_ekin_on_edges.py @@ -19,7 +19,7 @@ interpolate_vn_to_ie_and_compute_ekin_on_edges, ) from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -55,16 +55,29 @@ class TestMoVelocityAdvectionStencil02VnIe(StencilTest): PROGRAM = interpolate_vn_to_ie_and_compute_ekin_on_edges OUTPUTS = ("vn_ie", "z_kin_hor_e") - @classmethod - def reference(cls, grid, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs) -> dict: - vn_ie, z_kin_hor_e = interpolate_vn_to_ie_and_compute_ekin_on_edges_numpy( - grid, wgtfac_e, vn, vt + @staticmethod + def reference( + grid, + wgtfac_e: np.array, + vn: np.array, + vt: np.array, + vn_ie: np.array, + z_kin_hor_e: np.array, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, + ) -> dict: + subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) + vn_ie, z_kin_hor_e = vn_ie.copy(), z_kin_hor_e.copy() + vn_ie[subset], z_kin_hor_e[subset] = ( + x[subset] + for x in interpolate_vn_to_ie_and_compute_ekin_on_edges_numpy(grid, wgtfac_e, vn, vt) ) + return dict( - vn_ie=vn_ie[int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels)], - z_kin_hor_e=z_kin_hor_e[ - int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels) - ], + vn_ie=vn_ie, + z_kin_hor_e=z_kin_hor_e, ) @pytest.fixture @@ -73,19 +86,17 @@ def input_data(self, grid): vn = random_field(grid, EdgeDim, KDim, dtype=wpfloat) vt = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - vn_ie = zero_field(grid, EdgeDim, KDim, dtype=vpfloat) - z_kin_hor_e = zero_field(grid, EdgeDim, KDim, dtype=vpfloat) + vn_ie = random_field(grid, EdgeDim, KDim, dtype=vpfloat) + z_kin_hor_e = random_field(grid, EdgeDim, KDim, dtype=vpfloat) return dict( wgtfac_e=wgtfac_e, vn=vn, vt=vt, - vn_ie=vn_ie[int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels)], - z_kin_hor_e=z_kin_hor_e[ - int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels) - ], - horizontal_start=int32(1), - horizontal_end=int32(grid.num_cells), + vn_ie=vn_ie, + z_kin_hor_e=z_kin_hor_e, + horizontal_start=int32(0), + horizontal_end=int32(grid.num_edges), vertical_start=int32(1), vertical_end=int32(grid.num_levels), ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_ie.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_ie.py index 91d9a4410e..75b24fa475 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_ie.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_interpolate_vt_to_ie.py @@ -17,7 +17,7 @@ from icon4py.model.atmosphere.dycore.interpolate_vt_to_ie import interpolate_vt_to_ie from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field from icon4py.model.common.type_alias import vpfloat @@ -33,25 +33,34 @@ class TestMoVelocityAdvectionStencil03(StencilTest): OUTPUTS = ("z_vt_ie",) @staticmethod - def reference(grid, wgtfac_e: np.array, vt: np.array, **kwargs) -> dict: - z_vt_ie = interpolate_vt_to_ie_numpy(grid, wgtfac_e, vt) - return dict( - z_vt_ie=z_vt_ie[int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels)] - ) + def reference( + grid, + wgtfac_e: np.array, + vt: np.array, + z_vt_ie: np.array, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, + ) -> dict: + subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) + z_vt_ie = z_vt_ie.copy() + z_vt_ie[subset] = interpolate_vt_to_ie_numpy(grid, wgtfac_e, vt)[subset] + return dict(z_vt_ie=z_vt_ie) @pytest.fixture def input_data(self, grid): wgtfac_e = random_field(grid, EdgeDim, KDim, dtype=vpfloat) vt = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - z_vt_ie = zero_field(grid, EdgeDim, KDim, dtype=vpfloat) + z_vt_ie = random_field(grid, EdgeDim, KDim, dtype=vpfloat) return dict( wgtfac_e=wgtfac_e, vt=vt, - z_vt_ie=z_vt_ie[int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels)], - horizontal_start=int32(1), - horizontal_end=int32(grid.num_cells), + z_vt_ie=z_vt_ie, + horizontal_start=int32(0), + horizontal_end=int32(grid.num_edges), vertical_start=int32(1), vertical_end=int32(grid.num_levels), ) From 17429d7ed6f973ea2f4ddb7371abbbd62b5a523f Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 21 Feb 2024 08:29:14 +0100 Subject: [PATCH 2/4] Fix face_val_ppm_stencil_02c --- .../advection/face_val_ppm_stencil_02c.py | 10 +++++- .../test_face_val_ppm_stencil_02c.py | 35 ++++++++++++------- 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py index 0bf3ddd9a1..af702d49f5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field +from gt4py.next.ffront.fbuiltins import Field, int32 from icon4py.model.common.dimension import CellDim, KDim, Koff @@ -29,8 +29,16 @@ def _face_val_ppm_stencil_02c( def face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], p_face: Field[[CellDim, KDim], float], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): _face_val_ppm_stencil_02c( p_cc, out=p_face, + domain={ + CellDim: (horizontal_start, horizontal_end), + KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py index 7bdf4edd17..bdc5529eaa 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02c.py @@ -12,30 +12,41 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np import pytest +from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field class TestFaceValPpmStencil02c(StencilTest): PROGRAM = face_val_ppm_stencil_02c - OUTPUTS = ( - Output( - name="p_face", - refslice=(slice(None), slice(1, None)), - gtslice=(slice(None), slice(1, None)), - ), - ) + OUTPUTS = ("p_face",) @staticmethod - def reference(grid, p_cc: np.array, **kwargs): - p_face = p_cc.copy() - p_face[:, 1:] = p_cc[:, :-1] + def reference( + grid, + p_cc: np.array, + p_face: np.array, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, + ): + subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) + p_face = p_face.copy() + p_face[subset] = np.roll(p_cc, shift=1, axis=1)[subset] return dict(p_face=p_face) @pytest.fixture def input_data(self, grid): p_cc = random_field(grid, CellDim, KDim) p_face = random_field(grid, CellDim, KDim) - return dict(p_cc=p_cc, p_face=p_face) + return dict( + p_cc=p_cc, + p_face=p_face, + horizontal_start=int32(0), + horizontal_end=int32(grid.num_cells), + vertical_start=int32(1), + vertical_end=int32(grid.num_levels), + ) From 1ea7e7680eb68261a65258e2d7f1f2cfa699af18 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 21 Feb 2024 09:39:45 +0100 Subject: [PATCH 3/4] Add runtime exception when fields are passed as array slices. --- .../src/icon4py/model/common/test_utils/helpers.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 2858ee09ba..4a9eae9688 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -144,6 +144,18 @@ def dallclose(a, b, rtol=1.0e-12, atol=0.0, equal_nan=False): def allocate_data(backend, input_data): + non_zero_offsets = [ + (param, dim, dim_range) + for param, v in input_data.items() + if not is_scalar_type(v) + for dim, dim_range in zip(v.domain.dims, v.domain.ranges) + if dim_range.start != 0 + ] + if non_zero_offsets: + param, dim, dim_range = non_zero_offsets[0] + raise RuntimeError( + f"Field '{param}' passed as array slice with offset {dim_range.start} on dimension {dim.value}." + ) _allocate_field = constructors.as_field.partial(allocator=backend) input_data = { k: _allocate_field(domain=v.domain, data=v.ndarray) if not is_scalar_type(v) else v From 549316cf5c771318af1c6ccf945113a1cadf759d Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 21 Feb 2024 11:18:03 +0100 Subject: [PATCH 4/4] Remove runtime exception for sliced arrays --- .../src/icon4py/model/common/test_utils/helpers.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 4a9eae9688..2858ee09ba 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -144,18 +144,6 @@ def dallclose(a, b, rtol=1.0e-12, atol=0.0, equal_nan=False): def allocate_data(backend, input_data): - non_zero_offsets = [ - (param, dim, dim_range) - for param, v in input_data.items() - if not is_scalar_type(v) - for dim, dim_range in zip(v.domain.dims, v.domain.ranges) - if dim_range.start != 0 - ] - if non_zero_offsets: - param, dim, dim_range = non_zero_offsets[0] - raise RuntimeError( - f"Field '{param}' passed as array slice with offset {dim_range.start} on dimension {dim.value}." - ) _allocate_field = constructors.as_field.partial(allocator=backend) input_data = { k: _allocate_field(domain=v.domain, data=v.ndarray) if not is_scalar_type(v) else v