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), + ) 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), )