From 0dbb5ea16fd8edf44646ea062c1683fb2f5cbc21 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 13 Dec 2023 14:39:36 +0100 Subject: [PATCH 01/26] Refactor a few advection stencils --- .../advection/btraj_dreg_stencil_01.py | 2 +- .../advection/btraj_dreg_stencil_02.py | 2 +- .../advection/btraj_dreg_stencil_03.py | 2 +- .../advection/face_val_ppm_stencil_01.py | 2 +- .../advection/face_val_ppm_stencil_05.py | 2 +- .../test_btraj_dreg_stencil_01.py | 65 +++++------ .../test_btraj_dreg_stencil_02.py | 74 ++++++------ .../test_face_val_ppm_stencil_05.py | 107 ++++++++---------- .../model/common/test_utils/helpers.py | 2 +- 9 files changed, 115 insertions(+), 143 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py index 60109d3b6d..9c21bd6693 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py @@ -30,7 +30,7 @@ def _btraj_dreg_stencil_01( return lvn_sys_pos -@program +@program(grid_type="unstructured") def btraj_dreg_stencil_01( lcounterclock: bool, p_vn: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py index a11392c323..0431332a36 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py @@ -34,7 +34,7 @@ def _btraj_dreg_stencil_02( return opt_famask_dsl -@program +@program(grid_type="unstructured") def btraj_dreg_stencil_02( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index 6dce27f72a..e95f95625b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -102,7 +102,7 @@ def _btraj_dreg_stencil_03( ) -@program +@program(grid_type="unstructured") def btraj_dreg_stencil_03( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py index f47d3c8273..d0e9ce14fc 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py @@ -70,7 +70,7 @@ def _face_val_ppm_stencil_01( return z_slope -@program +@program(grid_type="unstructured") def face_val_ppm_stencil_01( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py index 884937b7cc..018e720496 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py @@ -51,7 +51,7 @@ def _face_val_ppm_stencil_05( return p_face -@program +@program(grid_type="unstructured") def face_val_ppm_stencil_05( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py index f87672cdf7..229ac10410 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_01.py @@ -10,54 +10,43 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 from icon4py.model.common.dimension import EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field - - -def btraj_dreg_stencil_01_numpy( - lcounterclock: bool, - p_vn: np.array, - tangent_orientation: np.array, -): - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - - tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) - - lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) - - mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field - lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) - return lvn_sys_pos +class TestBtrajDregStencil01(StencilTest): + PROGRAM = btraj_dreg_stencil_01 + OUTPUTS = ("lvn_sys_pos",) + @staticmethod + def reference( + grid, lcounterclock: bool, p_vn: np.array, tangent_orientation: np.array, **kwargs + ): + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) -def test_btraj_dreg_stencil_01(backend): - grid = SimpleGrid() - lcounterclock = True - p_vn = random_field(grid, EdgeDim, KDim) + tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) - tangent_orientation = random_field(grid, EdgeDim) + lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) - lvn_sys_pos = zero_field(grid, EdgeDim, KDim, dtype=bool) + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) - ref = btraj_dreg_stencil_01_numpy( - lcounterclock, - p_vn.asnumpy(), - tangent_orientation.asnumpy(), - ) + lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) - btraj_dreg_stencil_01.with_backend(backend)( - lcounterclock, - p_vn, - tangent_orientation, - lvn_sys_pos, - offset_provider={}, - ) + return dict(lvn_sys_pos=lvn_sys_pos) - assert np.allclose(ref, lvn_sys_pos.asnumpy()) + @pytest.fixture + def input_data(self, grid): + lcounterclock = True + p_vn = random_field(grid, EdgeDim, KDim) + tangent_orientation = random_field(grid, EdgeDim) + lvn_sys_pos = zero_field(grid, EdgeDim, KDim, dtype=bool) + return dict( + lcounterclock=lcounterclock, + p_vn=p_vn, + tangent_orientation=tangent_orientation, + lvn_sys_pos=lvn_sys_pos, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py index 74fea0327b..d602eae212 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_02.py @@ -12,61 +12,55 @@ # 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.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid from icon4py.model.common.test_utils.helpers import ( + StencilTest, numpy_to_1D_sparse_field, random_field, zero_field, ) -def btraj_dreg_stencil_02_numpy( - p_vn: np.array, - p_vt: np.array, - edge_cell_length: np.array, - p_dt: float, -): - lvn_pos = np.where(p_vn >= 0.0, True, False) +class TestBtrajDregStencil02(StencilTest): + PROGRAM = btraj_dreg_stencil_02 + OUTPUTS = ("opt_famask_dsl",) - traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt + @staticmethod + def reference(grid, p_vn: np.array, p_vt: np.array, p_dt: float, **kwargs): + lvn_pos = np.where(p_vn >= 0.0, True, False) - edge_cell_length = np.expand_dims(edge_cell_length, axis=-1) - e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) + traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt - opt_famask_dsl = np.where( - traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), - int32(1), - int32(0), - ) + edge_cell_length = np.expand_dims( + np.asarray(grid.connectivities[E2CDim], dtype=float), axis=-1 + ) + e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) - return opt_famask_dsl + opt_famask_dsl = np.where( + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), + int32(1), + int32(0), + ) + return dict(opt_famask_dsl=opt_famask_dsl) -def test_btraj_dreg_stencil_02(backend): - grid = SimpleGrid() - p_vn = random_field(grid, EdgeDim, KDim) - p_vt = random_field(grid, EdgeDim, KDim) - edge_cell_length = np.asarray(grid.connectivities[E2CDim], dtype=float) - edge_cell_length_new = numpy_to_1D_sparse_field(edge_cell_length, ECDim) - p_dt = 1.0 - opt_famask_dsl = zero_field(grid, EdgeDim, KDim, dtype=int32) + @pytest.fixture + def input_data(self, grid): + p_vn = random_field(grid, EdgeDim, KDim) + p_vt = random_field(grid, EdgeDim, KDim) + edge_cell_length = np.asarray(grid.connectivities[E2CDim], dtype=float) + edge_cell_length_new = numpy_to_1D_sparse_field(edge_cell_length, ECDim) + p_dt = 1.0 + opt_famask_dsl = zero_field(grid, EdgeDim, KDim, dtype=int32) - ref = btraj_dreg_stencil_02_numpy(p_vn.asnumpy(), p_vt.asnumpy(), edge_cell_length, p_dt) - - btraj_dreg_stencil_02.with_backend(backend)( - p_vn, - p_vt, - edge_cell_length_new, - p_dt, - opt_famask_dsl, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - "E2EC": grid.get_offset_provider("E2EC"), - }, - ) - - assert np.allclose(ref, opt_famask_dsl.asnumpy()) + return dict( + p_vn=p_vn, + p_vt=p_vt, + edge_cell_length=edge_cell_length_new, + p_dt=p_dt, + opt_famask_dsl=opt_famask_dsl, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py index 86bc568e9c..7e5c163fc5 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py @@ -12,75 +12,64 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def face_val_ppm_stencil_05_numpy( - p_cc: np.array, - p_cellhgt_mc_now: np.array, - z_slope: np.array, -): - p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] - p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] - p_cellhgt_mc_now_k_plus_1 = p_cellhgt_mc_now[:, 3:] - p_cellhgt_mc_now = p_cellhgt_mc_now[:, 2:-1] +class TestFaceValPpmStencil05(StencilTest): + PROGRAM = face_val_ppm_stencil_05 + OUTPUTS = ("p_face",) - p_cc_k_minus_1 = p_cc[:, 1:-1] - p_cc = p_cc[:, 2:] - z_slope_k_minus_1 = z_slope[:, 1:-1] - z_slope = z_slope[:, 2:] + @staticmethod + def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.array, **kwargs): + p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] + p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] + p_cellhgt_mc_now_k_plus_1 = p_cellhgt_mc_now[:, 3:] + p_cellhgt_mc_now = p_cellhgt_mc_now[:, 2:-1] - zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) - zgeo2 = 1.0 / ( - p_cellhgt_mc_now_k_minus_2 - + p_cellhgt_mc_now_k_minus_1 - + p_cellhgt_mc_now - + p_cellhgt_mc_now_k_plus_1 - ) - zgeo3 = (p_cellhgt_mc_now_k_minus_2 + p_cellhgt_mc_now_k_minus_1) / ( - 2.0 * p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now - ) - zgeo4 = (p_cellhgt_mc_now_k_plus_1 + p_cellhgt_mc_now) / ( - 2 * p_cellhgt_mc_now + p_cellhgt_mc_now_k_minus_1 - ) + p_cc_k_minus_1 = p_cc[:, 1:-1] + p_cc = p_cc[:, 2:] + z_slope_k_minus_1 = z_slope[:, 1:-1] + z_slope = z_slope[:, 2:] - p_face = ( - p_cc_k_minus_1 - + zgeo1 * (p_cc - p_cc_k_minus_1) - + zgeo2 - * ( - (2 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc_k_minus_1) - - zgeo3 * p_cellhgt_mc_now_k_minus_1 * z_slope - + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 + zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) + zgeo2 = 1.0 / ( + p_cellhgt_mc_now_k_minus_2 + + p_cellhgt_mc_now_k_minus_1 + + p_cellhgt_mc_now + + p_cellhgt_mc_now_k_plus_1 + ) + zgeo3 = (p_cellhgt_mc_now_k_minus_2 + p_cellhgt_mc_now_k_minus_1) / ( + 2.0 * p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now + ) + zgeo4 = (p_cellhgt_mc_now_k_plus_1 + p_cellhgt_mc_now) / ( + 2 * p_cellhgt_mc_now + p_cellhgt_mc_now_k_minus_1 ) - ) - - return p_face - -def test_face_val_ppm_stencil_05(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) - z_slope = random_field(grid, CellDim, KDim) - p_face = zero_field(grid, CellDim, KDim) + p_face = ( + p_cc_k_minus_1 + + zgeo1 * (p_cc - p_cc_k_minus_1) + + zgeo2 + * ( + (2 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc_k_minus_1) + - zgeo3 * p_cellhgt_mc_now_k_minus_1 * z_slope + + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 + ) + ) - ref = face_val_ppm_stencil_05_numpy( - p_cc.asnumpy(), - p_cellhgt_mc_now.asnumpy(), - z_slope.asnumpy(), - ) + # Concatenate the NaN column to the existing array + nan_column = np.full((p_face.shape[0], 2), np.nan) + p_face = np.concatenate((nan_column, p_face), axis=1) + return dict(p_face=p_face) - face_val_ppm_stencil_05.with_backend(backend)( - p_cc, - p_cellhgt_mc_now, - z_slope, - p_face, - offset_provider={"Koff": KDim}, - ) + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) + z_slope = random_field(grid, CellDim, KDim) + p_face = zero_field(grid, CellDim, KDim) - assert np.allclose(ref[:, :], p_face.asnumpy()[:, 2:]) + return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, z_slope=z_slope, p_face=p_face) 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 883a0aa75e..3c5ed303a0 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -163,7 +163,7 @@ def _test_validation(self, grid, backend, input_data): ) for name in self.OUTPUTS: assert np.allclose( - input_data[name].asnumpy(), reference_outputs[name] + input_data[name].asnumpy(), reference_outputs[name], equal_nan=True ), f"Validation failed for '{name}'" From 0fd8ea4bdf5e1c23b343e2845b5c0d9acdf57581 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 14 Dec 2023 10:10:54 +0100 Subject: [PATCH 02/26] Disable C2V offset provider due to missing connectivity --- model/common/src/icon4py/model/common/grid/icon.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 5c0d462545..e4b59cce4c 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -56,7 +56,7 @@ def __init__(self): "E2C2V": (self._get_offset_provider, E2C2VDim, EdgeDim, VertexDim), "V2E": (self._get_offset_provider, V2EDim, VertexDim, EdgeDim), "V2C": (self._get_offset_provider, V2CDim, VertexDim, CellDim), - "C2V": (self._get_offset_provider, C2VDim, CellDim, VertexDim), + # "C2V": (self._get_offset_provider, C2VDim, CellDim, VertexDim), "E2ECV": (self._get_offset_provider_for_sparse_fields, E2C2VDim, EdgeDim, ECVDim), "C2CEC": (self._get_offset_provider_for_sparse_fields, C2E2CDim, CellDim, CECDim), "C2CE": (self._get_offset_provider_for_sparse_fields, C2EDim, CellDim, CEDim), From a8a622f4cbfe87d58179cd4e1f5c931709de0515 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 8 Jan 2024 12:39:39 +0100 Subject: [PATCH 03/26] Allow comparing subset of output field --- .../advection/face_val_ppm_stencil_02b.py | 2 +- .../advection/face_val_ppm_stencil_02c.py | 2 +- .../test_face_val_ppm_stencil_02b.py | 38 +++++++---------- .../test_face_val_ppm_stencil_02c.py | 42 +++++++------------ .../model/common/test_utils/helpers.py | 10 ++++- 5 files changed, 39 insertions(+), 55 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py index d1a4d1242d..413b1a25e4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py @@ -25,7 +25,7 @@ def _face_val_ppm_stencil_02b( return p_face -@program +@program(grid_type="unstructured") def face_val_ppm_stencil_02b( p_cc: Field[[CellDim, KDim], float], p_face: Field[[CellDim, KDim], float], 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 4bb148785e..72a093666e 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 @@ -25,7 +25,7 @@ def _face_val_ppm_stencil_02c( return p_face -@program +@program(grid_type="unstructured") def face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], p_face: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py index 9144f3c62e..1a0a540106 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02b.py @@ -12,34 +12,24 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field -def face_val_ppm_stencil_02b_numpy( - p_cc: np.array, -): - p_face = p_cc.copy() +class TestFaceValPpmStencil02b(StencilTest): + PROGRAM = face_val_ppm_stencil_02b + OUTPUTS = ("p_face",) - return p_face + @staticmethod + def reference(grid, p_cc: np.array, **kwargs): + p_face = p_cc.copy() + return dict(p_face=p_face) - -def test_face_val_ppm_stencil_02b(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim) - - ref = face_val_ppm_stencil_02b_numpy( - p_cc.asnumpy(), - ) - - face_val_ppm_stencil_02b.with_backend(backend)( - p_cc, - p_face, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(ref, p_face.asnumpy()) + @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) 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 ff25183ecf..b83ace8859 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 @@ -10,38 +10,26 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest 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.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field - - -def face_val_ppm_stencil_02c_numpy( - p_cc: np.array, -): - p_face = p_cc.copy() - - p_face[:, 1:] = p_cc[:, :-1] - - return p_face - +from icon4py.model.common.test_utils.helpers import StencilTest, random_field -def test_face_val_ppm_stencil_02c(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim) - ref = face_val_ppm_stencil_02c_numpy( - p_cc.asnumpy(), - ) +class TestFaceValPpmStencil02c(StencilTest): + PROGRAM = face_val_ppm_stencil_02c + OUTPUTS = (("p_face", (slice(None), slice(1, None))),) - face_val_ppm_stencil_02c.with_backend(backend)( - p_cc, - p_face, - offset_provider={"Koff": KDim}, - ) + @staticmethod + def reference(grid, p_cc: np.array, **kwargs): + p_face = p_cc.copy() + p_face[:, 1:] = p_cc[:, :-1] + return dict(p_face=p_face) - assert np.allclose(ref[:, 1:], p_face.asnumpy()[:, 1:]) + @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) 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 85fbb8bc59..968218297c 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -163,9 +163,15 @@ def _test_validation(self, grid, backend, input_data): **input_data, offset_provider=grid.get_all_offset_providers(), ) - for name in self.OUTPUTS: + for out in self.OUTPUTS: + if isinstance(out, tuple): + name, subset = out + else: + name = out + subset = slice(None) + assert np.allclose( - input_data[name].asnumpy(), reference_outputs[name], equal_nan=True + input_data[name].asnumpy()[subset], reference_outputs[name][subset], equal_nan=True ), f"Validation failed for '{name}'" From c983e7cc5f152933c9f06478dcfb79fd1981da4c Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 8 Jan 2024 14:52:02 +0100 Subject: [PATCH 04/26] port btraj_dreg_03 --- .../test_btraj_dreg_stencil_03.py | 412 +++++++++--------- 1 file changed, 207 insertions(+), 205 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py index 74cc38a1c7..5555a0c665 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py @@ -10,14 +10,14 @@ # distribution for a copy of the license or check . # # 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.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid from icon4py.model.common.test_utils.helpers import ( + StencilTest, as_1D_sparse_field, constant_field, numpy_to_1D_sparse_field, @@ -25,208 +25,210 @@ ) -def btraj_dreg_stencil_03_numpy( - p_vn: np.array, - p_vt: np.array, - cell_idx: np.array, - cell_blk: np.array, - edge_verts_1_x: np.array, - edge_verts_2_x: np.array, - edge_verts_1_y: np.array, - edge_verts_2_y: np.array, - pos_on_tplane_e_1_x: np.array, - pos_on_tplane_e_2_x: np.array, - pos_on_tplane_e_1_y: np.array, - pos_on_tplane_e_2_y: np.array, - primal_normal_cell_x: np.array, - primal_normal_cell_y: np.array, - dual_normal_cell_x: np.array, - dual_normal_cell_y: np.array, - lvn_sys_pos: np.array, - p_dt: float, -) -> tuple[np.array]: - lvn_pos = np.where(p_vn >= 0.0, True, False) - cell_idx = np.expand_dims(cell_idx, axis=-1) - cell_blk = np.expand_dims(cell_blk, axis=-1) - primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) - dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) - primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) - dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) - edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) - edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) - edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) - edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) - pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) - pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) - pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) - pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) - - p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) - p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) - - depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt - depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt - depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt - depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt - - pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) - pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) - - pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x - pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y - pos_dreg_vert_c_2_x = np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x - pos_dreg_vert_c_2_y = np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y - pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x - pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y - pos_dreg_vert_c_4_x = np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x - pos_dreg_vert_c_4_y = np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y - - pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) - pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) - dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) - dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) - - p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 - p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 - p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 - p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 - p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 - p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 - p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 - p_coords_dreg_v_4_lat_dsl = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 - - return ( - p_cell_idx, - p_cell_rel_idx_dsl, - p_cell_blk, - p_coords_dreg_v_1_lon_dsl, - p_coords_dreg_v_2_lon_dsl, - p_coords_dreg_v_3_lon_dsl, - p_coords_dreg_v_4_lon_dsl, - p_coords_dreg_v_1_lat_dsl, - p_coords_dreg_v_2_lat_dsl, - p_coords_dreg_v_3_lat_dsl, - p_coords_dreg_v_4_lat_dsl, +class TestBtrajDregStencil03(StencilTest): + PROGRAM = btraj_dreg_stencil_03 + OUTPUTS = ( + "p_cell_idx", + "p_cell_rel_idx_dsl", + "p_cell_blk", + "p_coords_dreg_v_1_lon_dsl", + "p_coords_dreg_v_2_lon_dsl", + "p_coords_dreg_v_3_lon_dsl", + "p_coords_dreg_v_4_lon_dsl", + "p_coords_dreg_v_1_lat_dsl", + "p_coords_dreg_v_2_lat_dsl", + "p_coords_dreg_v_3_lat_dsl", + "p_coords_dreg_v_4_lat_dsl", ) - -def test_btraj_dreg_stencil_03(backend): - grid = SimpleGrid() - - p_vn = random_field(grid, EdgeDim, KDim) - p_vt = random_field(grid, EdgeDim, KDim) - cell_idx = np.asarray(grid.connectivities[E2CDim], dtype=int32) - cell_idx_new = numpy_to_1D_sparse_field(cell_idx, ECDim) - cell_blk = constant_field(grid, 1, EdgeDim, E2CDim, dtype=int32) - cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) - - edge_verts_1_x = random_field(grid, EdgeDim) - edge_verts_2_x = random_field(grid, EdgeDim) - edge_verts_1_y = random_field(grid, EdgeDim) - edge_verts_2_y = random_field(grid, EdgeDim) - pos_on_tplane_e_1_x = random_field(grid, EdgeDim) - pos_on_tplane_e_2_x = random_field(grid, EdgeDim) - pos_on_tplane_e_1_y = random_field(grid, EdgeDim) - pos_on_tplane_e_2_y = random_field(grid, EdgeDim) - primal_normal_cell_x = random_field(grid, EdgeDim, E2CDim) - primal_normal_cell_x_new = as_1D_sparse_field(primal_normal_cell_x, ECDim) - dual_normal_cell_x = random_field(grid, EdgeDim, E2CDim) - dual_normal_cell_x_new = as_1D_sparse_field(dual_normal_cell_x, ECDim) - primal_normal_cell_y = random_field(grid, EdgeDim, E2CDim) - primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) - dual_normal_cell_y = random_field(grid, EdgeDim, E2CDim) - dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) - lvn_sys_pos = constant_field(grid, True, EdgeDim, KDim, dtype=bool) - p_dt = 2.0 - p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) - p_coords_dreg_v_1_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_lon_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_1_lat_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_2_lat_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_3_lat_dsl = random_field(grid, EdgeDim, KDim) - p_coords_dreg_v_4_lat_dsl = random_field(grid, EdgeDim, KDim) - - ( - p_cell_idx_ref, - p_cell_rel_idx_dsl_ref, - p_cell_blk_ref, - p_coords_dreg_v_1_lon_dsl_ref, - p_coords_dreg_v_2_lon_dsl_ref, - p_coords_dreg_v_3_lon_dsl_ref, - p_coords_dreg_v_4_lon_dsl_ref, - p_coords_dreg_v_1_lat_dsl_ref, - p_coords_dreg_v_2_lat_dsl_ref, - p_coords_dreg_v_3_lat_dsl_ref, - p_coords_dreg_v_4_lat_dsl_ref, - ) = btraj_dreg_stencil_03_numpy( - p_vn.asnumpy(), - p_vt.asnumpy(), - cell_idx, - cell_blk.asnumpy(), - edge_verts_1_x.asnumpy(), - edge_verts_2_x.asnumpy(), - edge_verts_1_y.asnumpy(), - edge_verts_2_y.asnumpy(), - pos_on_tplane_e_1_x.asnumpy(), - pos_on_tplane_e_2_x.asnumpy(), - pos_on_tplane_e_1_y.asnumpy(), - pos_on_tplane_e_2_y.asnumpy(), - primal_normal_cell_x.asnumpy(), - primal_normal_cell_y.asnumpy(), - dual_normal_cell_x.asnumpy(), - dual_normal_cell_y.asnumpy(), - lvn_sys_pos.asnumpy(), - p_dt, - ) - - btraj_dreg_stencil_03.with_backend(backend)( - p_vn, - p_vt, - cell_idx_new, - cell_blk_new, - edge_verts_1_x, - edge_verts_2_x, - edge_verts_1_y, - edge_verts_2_y, - pos_on_tplane_e_1_x, - pos_on_tplane_e_2_x, - pos_on_tplane_e_1_y, - pos_on_tplane_e_2_y, - primal_normal_cell_x_new, - primal_normal_cell_y_new, - dual_normal_cell_x_new, - dual_normal_cell_y_new, - lvn_sys_pos, - p_dt, - p_cell_idx, - p_cell_rel_idx_dsl, - p_cell_blk, - p_coords_dreg_v_1_lon_dsl, - p_coords_dreg_v_2_lon_dsl, - p_coords_dreg_v_3_lon_dsl, - p_coords_dreg_v_4_lon_dsl, - p_coords_dreg_v_1_lat_dsl, - p_coords_dreg_v_2_lat_dsl, - p_coords_dreg_v_3_lat_dsl, - p_coords_dreg_v_4_lat_dsl, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - "E2EC": grid.get_offset_provider("E2EC"), - }, - ) - assert np.allclose(p_cell_idx.asnumpy(), p_cell_idx_ref) - assert np.allclose(p_cell_rel_idx_dsl.asnumpy(), p_cell_rel_idx_dsl_ref) - assert np.allclose(p_cell_blk.asnumpy(), p_cell_blk_ref) - assert np.allclose(p_coords_dreg_v_1_lon_dsl.asnumpy(), p_coords_dreg_v_1_lon_dsl_ref) - assert np.allclose(p_coords_dreg_v_2_lon_dsl.asnumpy(), p_coords_dreg_v_2_lon_dsl_ref) - assert np.allclose(p_coords_dreg_v_3_lon_dsl.asnumpy(), p_coords_dreg_v_3_lon_dsl_ref) - assert np.allclose(p_coords_dreg_v_4_lon_dsl.asnumpy(), p_coords_dreg_v_4_lon_dsl_ref) - assert np.allclose(p_coords_dreg_v_1_lat_dsl.asnumpy(), p_coords_dreg_v_1_lat_dsl_ref) - assert np.allclose(p_coords_dreg_v_2_lat_dsl.asnumpy(), p_coords_dreg_v_2_lat_dsl_ref) - assert np.allclose(p_coords_dreg_v_3_lat_dsl.asnumpy(), p_coords_dreg_v_3_lat_dsl_ref) - assert np.allclose(p_coords_dreg_v_4_lat_dsl.asnumpy(), p_coords_dreg_v_4_lat_dsl_ref) + @staticmethod + def reference( + grid, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + edge_verts_1_x: np.array, + edge_verts_2_x: np.array, + edge_verts_1_y: np.array, + edge_verts_2_y: np.array, + pos_on_tplane_e_1_x: np.array, + pos_on_tplane_e_2_x: np.array, + pos_on_tplane_e_1_y: np.array, + pos_on_tplane_e_2_y: np.array, + primal_normal_cell_x: np.array, + primal_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, + lvn_sys_pos: np.array, + p_dt: float, + **kwargs, + ): + def reshape(arr): + return np.reshape(arr, grid.connectivities[E2CDim].shape) + + cell_idx = reshape(cell_idx) + cell_blk = reshape(cell_blk) + primal_normal_cell_x = reshape(primal_normal_cell_x) + dual_normal_cell_x = reshape(dual_normal_cell_x) + primal_normal_cell_y = reshape(primal_normal_cell_y) + dual_normal_cell_y = reshape(dual_normal_cell_y) + + lvn_pos = np.where(p_vn >= 0.0, True, False) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) + dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) + primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) + dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) + edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) + edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) + edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) + edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) + pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) + pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) + pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) + pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) + + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) + + depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt + depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt + depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt + depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt + + pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x + pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = ( + np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_2_y = ( + np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + ) + pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x + pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y + pos_dreg_vert_c_4_x = ( + np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + ) + pos_dreg_vert_c_4_y = ( + np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + ) + + pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) + pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) + dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) + dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) + + p_coords_dreg_v_1_lon_dsl = ( + pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + ) + p_coords_dreg_v_2_lon_dsl = ( + pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + ) + p_coords_dreg_v_3_lon_dsl = ( + pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + ) + p_coords_dreg_v_4_lon_dsl = ( + pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + ) + p_coords_dreg_v_1_lat_dsl = ( + pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + ) + p_coords_dreg_v_2_lat_dsl = ( + pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + ) + p_coords_dreg_v_3_lat_dsl = ( + pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + ) + p_coords_dreg_v_4_lat_dsl = ( + pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + ) + return { + "p_cell_idx": p_cell_idx, + "p_cell_rel_idx_dsl": p_cell_rel_idx_dsl, + "p_cell_blk": p_cell_blk, + "p_coords_dreg_v_1_lon_dsl": p_coords_dreg_v_1_lon_dsl, + "p_coords_dreg_v_2_lon_dsl": p_coords_dreg_v_2_lon_dsl, + "p_coords_dreg_v_3_lon_dsl": p_coords_dreg_v_3_lon_dsl, + "p_coords_dreg_v_4_lon_dsl": p_coords_dreg_v_4_lon_dsl, + "p_coords_dreg_v_1_lat_dsl": p_coords_dreg_v_1_lat_dsl, + "p_coords_dreg_v_2_lat_dsl": p_coords_dreg_v_2_lat_dsl, + "p_coords_dreg_v_3_lat_dsl": p_coords_dreg_v_3_lat_dsl, + "p_coords_dreg_v_4_lat_dsl": p_coords_dreg_v_4_lat_dsl, + } + + @pytest.fixture + def input_data(self, grid): + p_vn = random_field(grid, EdgeDim, KDim) + p_vt = random_field(grid, EdgeDim, KDim) + cell_idx = np.asarray(grid.connectivities[E2CDim], dtype=int32) + cell_idx_new = numpy_to_1D_sparse_field(cell_idx, ECDim) + cell_blk = constant_field(grid, 1, EdgeDim, E2CDim, dtype=int32) + cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) + + edge_verts_1_x = random_field(grid, EdgeDim) + edge_verts_2_x = random_field(grid, EdgeDim) + edge_verts_1_y = random_field(grid, EdgeDim) + edge_verts_2_y = random_field(grid, EdgeDim) + pos_on_tplane_e_1_x = random_field(grid, EdgeDim) + pos_on_tplane_e_2_x = random_field(grid, EdgeDim) + pos_on_tplane_e_1_y = random_field(grid, EdgeDim) + pos_on_tplane_e_2_y = random_field(grid, EdgeDim) + primal_normal_cell_x = random_field(grid, EdgeDim, E2CDim) + primal_normal_cell_x_new = as_1D_sparse_field(primal_normal_cell_x, ECDim) + dual_normal_cell_x = random_field(grid, EdgeDim, E2CDim) + dual_normal_cell_x_new = as_1D_sparse_field(dual_normal_cell_x, ECDim) + primal_normal_cell_y = random_field(grid, EdgeDim, E2CDim) + primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) + dual_normal_cell_y = random_field(grid, EdgeDim, E2CDim) + dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) + lvn_sys_pos = constant_field(grid, True, EdgeDim, KDim, dtype=bool) + p_dt = 2.0 + p_cell_idx = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_cell_blk = constant_field(grid, 0, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_lon_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_1_lat_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_2_lat_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_3_lat_dsl = random_field(grid, EdgeDim, KDim) + p_coords_dreg_v_4_lat_dsl = random_field(grid, EdgeDim, KDim) + + return dict( + p_vn=p_vn, + p_vt=p_vt, + cell_idx=cell_idx_new, + cell_blk=cell_blk_new, + edge_verts_1_x=edge_verts_1_x, + edge_verts_2_x=edge_verts_2_x, + edge_verts_1_y=edge_verts_1_y, + edge_verts_2_y=edge_verts_2_y, + pos_on_tplane_e_1_x=pos_on_tplane_e_1_x, + pos_on_tplane_e_2_x=pos_on_tplane_e_2_x, + pos_on_tplane_e_1_y=pos_on_tplane_e_1_y, + pos_on_tplane_e_2_y=pos_on_tplane_e_2_y, + primal_normal_cell_x=primal_normal_cell_x_new, + primal_normal_cell_y=primal_normal_cell_y_new, + dual_normal_cell_x=dual_normal_cell_x_new, + dual_normal_cell_y=dual_normal_cell_y_new, + lvn_sys_pos=lvn_sys_pos, + p_dt=p_dt, + p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, + p_cell_blk=p_cell_blk, + p_coords_dreg_v_1_lon_dsl=p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl=p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl=p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl=p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl=p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl=p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl=p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl=p_coords_dreg_v_4_lat_dsl, + ) From 7a028fc1daca74ff7f0b682e637e7e8208364511 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 8 Jan 2024 17:15:06 +0100 Subject: [PATCH 05/26] Add face_val_ppm_stencil_02 --- .../advection/face_val_ppm_stencil_02.py | 2 +- .../test_face_val_ppm_stencil_02.py | 109 ++++++++---------- 2 files changed, 51 insertions(+), 60 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py index 72579f6bfd..a033f91457 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py @@ -71,7 +71,7 @@ def _face_val_ppm_stencil_02( return p_face -@program +@program(grid_type="unstructured") def face_val_ppm_stencil_02( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py index 160d986deb..bfa201e5e3 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py @@ -12,74 +12,65 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from gt4py.next import as_field from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import _shape, random_field +from icon4py.model.common.test_utils.helpers import StencilTest, _shape, random_field, zero_field -def face_val_ppm_stencil_02_numpy( - p_cc: np.array, - p_cellhgt_mc_now: np.array, - p_face_in: np.array, - k: np.array, - slev: int32, - elev: int32, - slevp1: int32, - elevp1: int32, -): - p_face_a = p_face_in +class TestFaceValPpmStencil02(StencilTest): + PROGRAM = face_val_ppm_stencil_02 + OUTPUTS = ("p_face",) - p_face_a[:, 1:] = p_cc[:, 1:] * (1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1])) + ( - p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) - ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) + @staticmethod + def reference( + grid, + p_cc: np.array, + p_cellhgt_mc_now: np.array, + p_face_in: np.array, + k: np.array, + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, + **kwargs, + ): + p_face_a = p_face_in + p_face_a[:, 1:] = p_cc[:, 1:] * ( + 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) + ) + (p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:])) * ( + (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] + ) - p_face = np.where((k == slevp1) | (k == elev), p_face_a, p_face_in) - p_face = np.where((k == slev), p_cc, p_face) - p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) + p_face = np.where((k == slevp1) | (k == elev), p_face_a, p_face_in) + p_face = np.where((k == slev), p_cc, p_face) + p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) + return dict(p_face=p_face) - return p_face + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim) + p_face_in = random_field(grid, CellDim, KDim) + p_face = zero_field(grid, CellDim, KDim) + k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) + slev = int32(1) + slevp1 = slev + int32(1) + elev = k[-3] + elevp1 = elev + int32(1) -def test_face_val_ppm_stencil_02(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim) - p_face_in = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim) - - k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) - - slev = int32(1) - slevp1 = slev + int32(1) - elev = k[-3] - elevp1 = elev + int32(1) - - ref = face_val_ppm_stencil_02_numpy( - p_cc.asnumpy(), - p_cellhgt_mc_now.asnumpy(), - p_face_in.asnumpy(), - k.asnumpy(), - slev, - elev, - slevp1, - elevp1, - ) - - face_val_ppm_stencil_02.with_backend(backend)( - p_cc, - p_cellhgt_mc_now, - p_face_in, - k, - slev, - elev, - slevp1, - elevp1, - p_face, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(ref, p_face.asnumpy()) + return dict( + p_cc=p_cc, + p_cellhgt_mc_now=p_cellhgt_mc_now, + p_face_in=p_face_in, + k=k, + slev=slev, + elev=elev, + slevp1=slevp1, + elevp1=elevp1, + p_face=p_face, + ) From 00bc933bec3c2234ec2a8497f9ad114f210de922 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Mon, 8 Jan 2024 17:46:12 +0100 Subject: [PATCH 06/26] Add face_val_ppm_02a --- .../advection/face_val_ppm_stencil_02a.py | 2 +- .../test_face_val_ppm_stencil_02a.py | 61 ++++++++----------- 2 files changed, 25 insertions(+), 38 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py index 3653de1a59..2052867155 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py @@ -29,7 +29,7 @@ def _face_val_ppm_stencil_02a( return p_face -@program +@program(grid_type="unstructured") def face_val_ppm_stencil_02a( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py index 9bb44714a1..4b188192c7 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py @@ -10,44 +10,31 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field - - -def face_val_ppm_stencil_02a_numpy( - p_cc: np.array, - p_cellhgt_mc_now: np.array, -): - p_face = p_cc.copy() - - p_face[:, 1:] = p_cc[:, 1:] * (1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1])) + ( - p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) - ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) - - return p_face - - -def test_face_val_ppm_stencil_02a(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim) - - ref = face_val_ppm_stencil_02a_numpy( - p_cc.asnumpy(), - p_cellhgt_mc_now.asnumpy(), - ) - - face_val_ppm_stencil_02a.with_backend(backend)( - p_cc, - p_cellhgt_mc_now, - p_face, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(ref[:, 1:], p_face.asnumpy()[:, 1:]) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field + + +class TestFaceValPpmStencil02a(StencilTest): + PROGRAM = face_val_ppm_stencil_02a + OUTPUTS = (("p_face", (slice(None), slice(1, None))),) + + @staticmethod + def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs): + p_face = p_cc.copy() + p_face[:, 1:] = p_cc[:, 1:] * ( + 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) + ) + (p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:])) * ( + (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] + ) + return dict(p_face=p_face) + + @pytest.fixture + def input_data(self, grid): + p_face = random_field(grid, CellDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim) + return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, p_face=p_face) From db284feefdd44f7efd7d858c95d6d8ceb0201a8b Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Jan 2024 14:25:00 +0100 Subject: [PATCH 07/26] Add Output helper class --- .../test_face_val_ppm_stencil_01.py | 115 ++++++++---------- .../test_face_val_ppm_stencil_02a.py | 7 +- .../model/common/test_utils/helpers.py | 20 ++- 3 files changed, 71 insertions(+), 71 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py index 8149e7fd2f..b5ceb9b413 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py @@ -10,80 +10,69 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest from gt4py.next import as_field from gt4py.next.ffront.fbuiltins import int32 from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import _shape, random_field - - -def face_val_ppm_stencil_01_numpy( - p_cc: np.array, - p_cellhgt_mc_now: np.array, - k: np.array, - elev: int32, -): - # this is a comment: k = np.broadcast_to(k, p_cc.shape) +from icon4py.model.common.test_utils.helpers import ( + Output, + StencilTest, + _shape, + random_field, + zero_field, +) - # 01a - zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( - p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] - ) - zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1]) - z_slope_a = ( - p_cellhgt_mc_now[:, 1:-1] - / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:]) - ) * ( - (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac - + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 - ) - # 01b - zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( - p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] - ) - zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) - z_slope_b = ( - p_cellhgt_mc_now[:, 1:-1] - / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) - ) * ( - (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac - + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 +class TestFaceValPpmStencil01(StencilTest): + PROGRAM = face_val_ppm_stencil_01 + OUTPUTS = ( + Output( + "z_slope", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(1, -1)) + ), ) - z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) + @staticmethod + def reference( + grid, p_cc: np.array, p_cellhgt_mc_now: np.array, k: np.array, elev: int32, **kwargs + ): + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1]) + z_slope_a = ( + p_cellhgt_mc_now[:, 1:-1] + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:]) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 + ) - return z_slope + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1] + ) + z_slope_b = ( + p_cellhgt_mc_now[:, 1:-1] + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 + ) + z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) + return dict(z_slope=z_slope) -def test_face_val_ppm_stencil_01(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1}) - p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) - - k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32)) - elev = k[-2] - - z_slope = random_field(grid, CellDim, KDim) - - ref = face_val_ppm_stencil_01_numpy( - p_cc.asnumpy(), - p_cellhgt_mc_now.asnumpy(), - k.asnumpy(), - elev, - ) - - face_val_ppm_stencil_01.with_backend(backend)( - p_cc, - p_cellhgt_mc_now, - k, - elev, - z_slope, - offset_provider={"Koff": KDim}, - ) + @pytest.fixture + def input_data(self, grid): + z_slope = zero_field(grid, CellDim, KDim) + p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1}) + p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) + k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32)) + elev = k[-2] - assert np.allclose(ref[:, :-1], z_slope.asnumpy()[:, 1:-1]) + return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, k=k, elev=elev, z_slope=z_slope) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py index 4b188192c7..d1d8d5d900 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02a.py @@ -15,12 +15,15 @@ from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field + + +outslice = (slice(None), slice(1, None)) class TestFaceValPpmStencil02a(StencilTest): PROGRAM = face_val_ppm_stencil_02a - OUTPUTS = (("p_face", (slice(None), slice(1, None))),) + OUTPUTS = (Output("p_face", refslice=outslice, gtslice=outslice),) @staticmethod def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs): 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 968218297c..2bffe5de77 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +from dataclasses import dataclass, field from typing import ClassVar, Optional import numpy as np @@ -148,6 +149,13 @@ def allocate_data(backend, input_data): return input_data +@dataclass(frozen=True) +class Output: + name: str + refslice: tuple[slice, ...] = field(default_factory=lambda: (slice(None),)) + gtslice: tuple[slice, ...] = field(default_factory=lambda: (slice(None),)) + + def _test_validation(self, grid, backend, input_data): reference_outputs = self.reference( grid, @@ -164,14 +172,14 @@ def _test_validation(self, grid, backend, input_data): offset_provider=grid.get_all_offset_providers(), ) for out in self.OUTPUTS: - if isinstance(out, tuple): - name, subset = out - else: - name = out - subset = slice(None) + name, refslice, gtslice = ( + (out.name, out.refslice, out.gtslice) + if isinstance(out, Output) + else (out, (slice(None),), (slice(None),)) + ) assert np.allclose( - input_data[name].asnumpy()[subset], reference_outputs[name][subset], equal_nan=True + input_data[name].asnumpy()[gtslice], reference_outputs[name][refslice], equal_nan=True ), f"Validation failed for '{name}'" From ebcee6a29a3bbd0c837f11a767a2ab768e361122 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Jan 2024 15:47:34 +0100 Subject: [PATCH 08/26] Add more stencils and reshape helper --- .../test_btraj_dreg_stencil_03.py | 17 +- .../test_hflx_limiter_mo_stencil_01a.py | 84 ++++---- .../test_hflx_limiter_mo_stencil_01b.py | 195 +++++++++--------- .../model/common/test_utils/helpers.py | 4 + 4 files changed, 147 insertions(+), 153 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py index 5555a0c665..927bccd654 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_btraj_dreg_stencil_03.py @@ -22,6 +22,7 @@ constant_field, numpy_to_1D_sparse_field, random_field, + reshape, ) @@ -64,15 +65,13 @@ def reference( p_dt: float, **kwargs, ): - def reshape(arr): - return np.reshape(arr, grid.connectivities[E2CDim].shape) - - cell_idx = reshape(cell_idx) - cell_blk = reshape(cell_blk) - primal_normal_cell_x = reshape(primal_normal_cell_x) - dual_normal_cell_x = reshape(dual_normal_cell_x) - primal_normal_cell_y = reshape(primal_normal_cell_y) - dual_normal_cell_y = reshape(dual_normal_cell_y) + e2c_shape = grid.connectivities[E2CDim].shape + cell_idx = reshape(cell_idx, e2c_shape) + cell_blk = reshape(cell_blk, e2c_shape) + primal_normal_cell_x = reshape(primal_normal_cell_x, e2c_shape) + dual_normal_cell_x = reshape(dual_normal_cell_x, e2c_shape) + primal_normal_cell_y = reshape(primal_normal_cell_y, e2c_shape) + dual_normal_cell_y = reshape(dual_normal_cell_y, e2c_shape) lvn_pos = np.where(p_vn >= 0.0, True, False) cell_idx = np.expand_dims(cell_idx, axis=-1) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py index ea68a95943..0dbf30ca22 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01a.py @@ -12,58 +12,50 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( hflx_limiter_mo_stencil_01a, ) from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def hflx_limiter_mo_stencil_01a_numpy( - e2c: np.array, - p_mflx_tracer_h: np.ndarray, - p_mass_flx_e: np.ndarray, - p_cc: np.ndarray, -): - p_cc_e2c = p_cc[e2c] - - z_mflx_low = 0.5 * ( - p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) - - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) - ) - - z_anti = p_mflx_tracer_h - z_mflx_low - - return (z_mflx_low, z_anti) - - -def test_hflx_limiter_mo_stencil_01a(backend): - grid = SimpleGrid() - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - p_mass_flx_e = random_field(grid, EdgeDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - z_mflx_low = zero_field(grid, EdgeDim, KDim) - z_anti = zero_field(grid, EdgeDim, KDim) - - ref_1, ref_2 = hflx_limiter_mo_stencil_01a_numpy( - grid.connectivities[E2CDim], - p_mflx_tracer_h.asnumpy(), - p_mass_flx_e.asnumpy(), - p_cc.asnumpy(), - ) - - hflx_limiter_mo_stencil_01a.with_backend(backend)( - p_mflx_tracer_h, - p_mass_flx_e, - p_cc, - z_mflx_low, - z_anti, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, +class TestHflxLimiterMoStencil01a(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_01a + OUTPUTS = ( + "z_mflx_low", + "z_anti", ) - assert np.allclose(z_mflx_low.asnumpy(), ref_1) - assert np.allclose(z_anti.asnumpy(), ref_2) + @staticmethod + def reference( + grid, p_mflx_tracer_h: np.ndarray, p_mass_flx_e: np.ndarray, p_cc: np.ndarray, **kwargs + ): + e2c = grid.connectivities[E2CDim] + p_cc_e2c = p_cc[e2c] + + z_mflx_low = 0.5 * ( + p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) + - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) + ) + + z_anti = p_mflx_tracer_h - z_mflx_low + + return dict(z_mflx_low=z_mflx_low, z_anti=z_anti) + + @pytest.fixture + def input_data(self, grid): + p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) + p_mass_flx_e = random_field(grid, EdgeDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + z_mflx_low = zero_field(grid, EdgeDim, KDim) + z_anti = zero_field(grid, EdgeDim, KDim) + + return dict( + p_mflx_tracer_h=p_mflx_tracer_h, + p_mass_flx_e=p_mass_flx_e, + p_cc=p_cc, + z_mflx_low=z_mflx_low, + z_anti=z_anti, + ) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py index 0ebb6b94e7..70181f6105 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_01b.py @@ -12,110 +12,109 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import ( hflx_limiter_mo_stencil_01b, ) from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field - - -def hflx_limiter_mo_stencil_01b_numpy( - c2e: np.array, - geofac_div: np.ndarray, - p_rhodz_now: np.ndarray, - p_rhodz_new: np.ndarray, - z_mflx_low: np.ndarray, - z_anti: np.ndarray, - p_cc: np.ndarray, - p_dtime: float, -): - z_anti_c2e = z_anti[c2e] - geofac_div = np.expand_dims(geofac_div, axis=-1) - - zero_array = np.zeros(p_rhodz_now.shape) - - z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] - z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] - z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] - - z_mflx_anti_in = -1.0 * ( - np.minimum(zero_array, z_mflx_anti_1) - + np.minimum(zero_array, z_mflx_anti_2) - + np.minimum(zero_array, z_mflx_anti_3) - ) - - z_mflx_anti_out = ( - np.maximum(zero_array, z_mflx_anti_1) - + np.maximum(zero_array, z_mflx_anti_2) - + np.maximum(zero_array, z_mflx_anti_3) - ) - - z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) - - z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new - z_tracer_max = np.maximum(p_cc, z_tracer_new_low) - z_tracer_min = np.minimum(p_cc, z_tracer_new_low) - - return ( - z_mflx_anti_in, - z_mflx_anti_out, - z_tracer_new_low, - z_tracer_max, - z_tracer_min, - ) - +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) -def test_hflx_limiter_mo_stencil_01b(backend): - grid = SimpleGrid() - - geofac_div = random_field(grid, CellDim, C2EDim) - geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) - p_rhodz_now = random_field(grid, CellDim, KDim) - p_rhodz_new = random_field(grid, CellDim, KDim) - z_mflx_low = random_field(grid, EdgeDim, KDim) - z_anti = random_field(grid, EdgeDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - p_dtime = 5.0 - z_mflx_anti_in = random_field(grid, CellDim, KDim) - z_mflx_anti_out = random_field(grid, CellDim, KDim) - z_tracer_new_low = random_field(grid, CellDim, KDim) - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - - ref_1, ref_2, ref_3, ref_4, ref_5 = hflx_limiter_mo_stencil_01b_numpy( - grid.connectivities[C2EDim], - geofac_div.asnumpy(), - p_rhodz_now.asnumpy(), - p_rhodz_new.asnumpy(), - z_mflx_low.asnumpy(), - z_anti.asnumpy(), - p_cc.asnumpy(), - p_dtime, - ) - hflx_limiter_mo_stencil_01b.with_backend(backend)( - geofac_div_new, - p_rhodz_now, - p_rhodz_new, - z_mflx_low, - z_anti, - p_cc, - p_dtime, - z_mflx_anti_in, - z_mflx_anti_out, - z_tracer_new_low, - z_tracer_max, - z_tracer_min, - offset_provider={ - "C2E": grid.get_offset_provider("C2E"), - "C2CE": grid.get_offset_provider("C2CE"), - }, +class TestHflxLimiterMoStencil01b(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_01b + OUTPUTS = ( + "z_mflx_anti_in", + "z_mflx_anti_out", + "z_tracer_new_low", + "z_tracer_max", + "z_tracer_min", ) - assert np.allclose(z_mflx_anti_in.asnumpy(), ref_1) - assert np.allclose(z_mflx_anti_out.asnumpy(), ref_2) - assert np.allclose(z_tracer_new_low.asnumpy(), ref_3) - assert np.allclose(z_tracer_max.asnumpy(), ref_4) - assert np.allclose(z_tracer_min.asnumpy(), ref_5) + @staticmethod + def reference( + grid, + geofac_div: np.ndarray, + p_rhodz_now: np.ndarray, + p_rhodz_new: np.ndarray, + z_mflx_low: np.ndarray, + z_anti: np.ndarray, + p_cc: np.ndarray, + p_dtime: float, + **kwargs, + ): + c2e = grid.connectivities[C2EDim] + z_anti_c2e = z_anti[c2e] + + geofac_div = reshape(geofac_div, c2e.shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + + zero_array = np.zeros(p_rhodz_now.shape) + + z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] + z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] + z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] + + z_mflx_anti_in = -1.0 * ( + np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) + ) + + z_mflx_anti_out = ( + np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) + ) + + z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) + + z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new + z_tracer_max = np.maximum(p_cc, z_tracer_new_low) + z_tracer_min = np.minimum(p_cc, z_tracer_new_low) + + return dict( + z_mflx_anti_in=z_mflx_anti_in, + z_mflx_anti_out=z_mflx_anti_out, + z_tracer_new_low=z_tracer_new_low, + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + ) + + @pytest.fixture + def input_data(self, grid): + geofac_div = random_field(grid, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_rhodz_now = random_field(grid, CellDim, KDim) + p_rhodz_new = random_field(grid, CellDim, KDim) + z_mflx_low = random_field(grid, EdgeDim, KDim) + z_anti = random_field(grid, EdgeDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + p_dtime = 5.0 + + z_mflx_anti_in = zero_field(grid, CellDim, KDim) + z_mflx_anti_out = zero_field(grid, CellDim, KDim) + z_tracer_new_low = zero_field(grid, CellDim, KDim) + z_tracer_max = zero_field(grid, CellDim, KDim) + z_tracer_min = zero_field(grid, CellDim, KDim) + + return dict( + geofac_div=geofac_div_new, + p_rhodz_now=p_rhodz_now, + p_rhodz_new=p_rhodz_new, + z_mflx_low=z_mflx_low, + z_anti=z_anti, + p_cc=p_cc, + p_dtime=p_dtime, + z_mflx_anti_in=z_mflx_anti_in, + z_mflx_anti_out=z_mflx_anti_out, + z_tracer_new_low=z_tracer_new_low, + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + ) 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 2bffe5de77..160141d37b 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -245,3 +245,7 @@ def uses_icon_grid_with_otf(backend, grid): if isinstance(backend.executor, OTFCompileExecutor): return True return False + + +def reshape(arr: np.array, shape: tuple[int, ...]): + return np.reshape(arr, shape) From d89a007aabac69b058176443f349a589710b0b2f Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 9 Jan 2024 15:57:26 +0100 Subject: [PATCH 09/26] Fix test --- .../test_mo_solve_nonhydro_stencil_11_lower.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py index e40672ae0e..76637bcb4b 100644 --- a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np import pytest from gt4py.next.ffront.fbuiltins import int32 @@ -28,7 +29,7 @@ class TestMoSolveNonhydroStencil11Lower(StencilTest): @staticmethod def reference(grid, **kwargs) -> dict: - z_theta_v_pr_ic = 0 + z_theta_v_pr_ic = np.zeros((grid.num_cells, grid.num_levels)) return dict(z_theta_v_pr_ic=z_theta_v_pr_ic) @pytest.fixture From d0b14dc9ac541f018651e426174e2f1b9074bdb6 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 10 Jan 2024 13:17:40 +0100 Subject: [PATCH 10/26] Add hflx limiter 2 --- .../advection/hflx_limiter_mo_stencil_02.py | 2 +- .../test_hflx_limiter_mo_stencil_02.py | 180 +++++++----------- 2 files changed, 69 insertions(+), 113 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py index 6a6a9f47b4..fef2e48c94 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py @@ -45,7 +45,7 @@ def _hflx_limiter_mo_stencil_02( return (z_tracer_new_out, z_tracer_max_out, z_tracer_min_out) -@program +@program(grid_type="unstructured") def hflx_limiter_mo_stencil_02( refin_ctrl: Field[[CellDim], int32], p_cc: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py index b84200d8ca..4925d83987 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_02.py @@ -10,120 +10,76 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest from numpy import int32 from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import hflx_limiter_mo_stencil_02 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field - - -def hflx_limiter_mo_stencil_02_numpy( - refin_ctrl: np.ndarray, - p_cc: np.ndarray, - z_tracer_new_low: np.ndarray, - z_tracer_max: np.ndarray, - z_tracer_min: np.ndarray, - lo_bound: float, - hi_bound: float, -): - refin_ctrl = np.expand_dims(refin_ctrl, axis=1) - condition = np.logical_or( - np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), - np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), - ) - z_tracer_new_out = np.where( - condition, - np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), - z_tracer_new_low, - ) - z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) - z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) - return z_tracer_new_out, z_tracer_max_out, z_tracer_min_out - - -def test_hflx_limiter_mo_stencil_02_some_matching_condition(backend): - grid = SimpleGrid() - - hi_bound = np.int32(1) - lo_bound = np.int32(5) - - refin_ctrl = constant_field(grid, hi_bound, CellDim, dtype=int32) - - refin_ctrl[0:2] = np.int32(3) - - p_cc = random_field(grid, CellDim, KDim) - z_tracer_new_low_in = random_field(grid, CellDim, KDim) - z_tracer_max_in = random_field(grid, CellDim, KDim) - z_tracer_min_in = random_field(grid, CellDim, KDim) - - z_tracer_new_low = zero_field(grid, CellDim, KDim) - z_tracer_max = zero_field(grid, CellDim, KDim) - z_tracer_min = zero_field(grid, CellDim, KDim) - - ref_new_low, ref_max, ref_min = hflx_limiter_mo_stencil_02_numpy( - refin_ctrl.asnumpy(), - p_cc.asnumpy(), - z_tracer_new_low_in.asnumpy(), - z_tracer_max_in.asnumpy(), - z_tracer_min_in.asnumpy(), - lo_bound, - hi_bound, - ) - - hflx_limiter_mo_stencil_02.with_backend(backend)( - refin_ctrl, - p_cc, - z_tracer_new_low_in, - z_tracer_max_in, - z_tracer_min_in, - lo_bound, - hi_bound, - z_tracer_new_low, - z_tracer_max, - z_tracer_min, - offset_provider={}, - ) - - assert np.allclose(z_tracer_new_low.asnumpy(), ref_new_low) - assert np.allclose(z_tracer_max.asnumpy(), ref_max) - assert np.allclose(z_tracer_min.asnumpy(), ref_min) - - -def test_hflx_limiter_mo_stencil_02_none_matching_condition(backend): - grid = SimpleGrid() - - hi_bound = np.int32(3) - lo_bound = np.int32(1) - - refin_ctrl = constant_field(grid, 2, CellDim, dtype=int32) - - p_cc = random_field(grid, CellDim, KDim) - - z_tracer_new_low_in = random_field(grid, CellDim, KDim) - z_tracer_max_in = random_field(grid, CellDim, KDim) - z_tracer_min_in = random_field(grid, CellDim, KDim) - - z_tracer_new_low = zero_field(grid, CellDim, KDim) - z_tracer_max = zero_field(grid, CellDim, KDim) - z_tracer_min = zero_field(grid, CellDim, KDim) - - hflx_limiter_mo_stencil_02.with_backend(backend)( - refin_ctrl, - p_cc, - z_tracer_new_low_in, - z_tracer_max_in, - z_tracer_min_in, - lo_bound, - hi_bound, - z_tracer_new_low, - z_tracer_max, - z_tracer_min, - offset_provider={}, - ) - - assert np.allclose(z_tracer_new_low_in.asnumpy(), z_tracer_new_low.asnumpy()) - assert np.allclose(z_tracer_min_in.asnumpy(), z_tracer_min.asnumpy()) - assert np.allclose(z_tracer_max_in.asnumpy(), z_tracer_max.asnumpy()) +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + constant_field, + random_field, + zero_field, +) + + +class TestHflxLimiterMoStencil02(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_02 + OUTPUTS = ("z_tracer_new_low", "z_tracer_max", "z_tracer_min") + + @staticmethod + def reference( + grid, + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + lo_bound: float, + hi_bound: float, + **kwargs, + ): + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), + ) + z_tracer_new_out = np.where( + condition, + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) + return dict( + z_tracer_new_low=z_tracer_new_out, + z_tracer_max=z_tracer_max_out, + z_tracer_min=z_tracer_min_out, + ) + + @pytest.fixture() + def input_data(self, grid): + hi_bound, lo_bound = 3, 1 + refin_ctrl = constant_field(grid, 2, CellDim, dtype=int32) + p_cc = random_field(grid, CellDim, KDim) + z_tracer_new_low_in = random_field(grid, CellDim, KDim) + z_tracer_max_in = random_field(grid, CellDim, KDim) + z_tracer_min_in = random_field(grid, CellDim, KDim) + + z_tracer_new_low_out = zero_field(grid, CellDim, KDim) + z_tracer_max_out = zero_field(grid, CellDim, KDim) + z_tracer_min_out = zero_field(grid, CellDim, KDim) + + return dict( + refin_ctrl=refin_ctrl, + p_cc=p_cc, + z_tracer_new_low=z_tracer_new_low_in, + z_tracer_max=z_tracer_max_in, + z_tracer_min=z_tracer_min_in, + lo_bound=lo_bound, + hi_bound=hi_bound, + z_tracer_new_low_out=z_tracer_new_low_out, + z_tracer_max_out=z_tracer_max_out, + z_tracer_min_out=z_tracer_min_out, + ) From afaba3f3ea7fd11a36bc17239a3a9867c76c24ed Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 10 Jan 2024 16:19:29 +0100 Subject: [PATCH 11/26] Add hflx limiter 4 --- .../test_hflx_limiter_mo_stencil_04.py | 67 ++++++++----------- 1 file changed, 29 insertions(+), 38 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py index 8876560c47..d1b5e3244d 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_04.py @@ -12,48 +12,39 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import hflx_limiter_mo_stencil_04 from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def hflx_limiter_mo_stencil_04_numpy( - e2c: np.ndarray, - z_anti: np.ndarray, - r_m: np.ndarray, - r_p: np.ndarray, - z_mflx_low: np.ndarray, -): - r_frac = np.where( - z_anti >= 0, - np.minimum(r_m[e2c[:, 0]], r_p[e2c[:, 1]]), - np.minimum(r_m[e2c[:, 1]], r_p[e2c[:, 0]]), - ) - return z_mflx_low + np.minimum(1.0, r_frac) * z_anti +class TestHflxLimiterMoStencil04(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_04 + OUTPUTS = ("p_mflx_tracer_h",) + @staticmethod + def reference( + grid, z_anti: np.ndarray, r_m: np.ndarray, r_p: np.ndarray, z_mflx_low: np.ndarray, **kwargs + ): + r_frac = np.where( + z_anti >= 0, + np.minimum( + r_m[grid.connectivities[E2CDim][:, 0]], r_p[grid.connectivities[E2CDim][:, 1]] + ), + np.minimum( + r_m[grid.connectivities[E2CDim][:, 1]], r_p[grid.connectivities[E2CDim][:, 0]] + ), + ) + return dict(p_mflx_tracer_h=z_mflx_low + np.minimum(1.0, r_frac) * z_anti) -def test_hflx_limiter_mo_stencil_04(backend): - grid = SimpleGrid() - z_anti = random_field(grid, EdgeDim, KDim, low=-2.0, high=2.0) - r_m = random_field(grid, CellDim, KDim) - r_p = random_field(grid, CellDim, KDim) - z_mflx_low = random_field(grid, EdgeDim, KDim) - p_mflx_tracer_h = zero_field(grid, EdgeDim, KDim) - ref = hflx_limiter_mo_stencil_04_numpy( - grid.connectivities[E2CDim], - z_anti.asnumpy(), - r_m.asnumpy(), - r_p.asnumpy(), - z_mflx_low.asnumpy(), - ) - hflx_limiter_mo_stencil_04.with_backend(backend)( - z_anti, - r_m, - r_p, - z_mflx_low, - p_mflx_tracer_h, - offset_provider={"E2C": grid.get_offset_provider("E2C")}, - ) - assert np.allclose(p_mflx_tracer_h.asnumpy(), ref) + @pytest.fixture + def input_data(self, grid): + z_anti = random_field(grid, EdgeDim, KDim, low=-2.0, high=2.0) + r_m = random_field(grid, CellDim, KDim) + r_p = random_field(grid, CellDim, KDim) + z_mflx_low = random_field(grid, EdgeDim, KDim) + p_mflx_tracer_h = zero_field(grid, EdgeDim, KDim) + return dict( + z_anti=z_anti, r_m=r_m, r_p=r_p, z_mflx_low=z_mflx_low, p_mflx_tracer_h=p_mflx_tracer_h + ) From a810ab72184c891a78c5607399c7f83c5c68c2d0 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 10 Jan 2024 16:29:31 +0100 Subject: [PATCH 12/26] Add hflx limiter pd 1 --- .../test_hflx_limiter_pd_stencil_01.py | 108 +++++++++--------- 1 file changed, 55 insertions(+), 53 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py index f3c4e86f25..9ebfe4692a 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_01.py @@ -10,66 +10,68 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import hflx_limiter_pd_stencil_01 from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field - - -def hflx_limiter_pd_stencil_01_numpy( - c2e: np.array, - geofac_div: np.array, - p_cc: np.array, - p_rhodz_now: np.array, - p_mflx_tracer_h: np.array, - p_dtime, - dbl_eps, -): - geofac_div = np.expand_dims(geofac_div, axis=-1) - p_m_0 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 0]] * geofac_div[:, 0] * p_dtime) - p_m_1 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 1]] * geofac_div[:, 1] * p_dtime) - p_m_2 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 2]] * geofac_div[:, 2] * p_dtime) - - p_m = p_m_0 + p_m_1 + p_m_2 - r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) - - return r_m +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) -def test_hflx_limiter_pd_stencil_01(backend): - grid = SimpleGrid() - geofac_div = random_field(grid, CellDim, C2EDim) - p_cc = random_field(grid, CellDim, KDim) - p_rhodz_now = random_field(grid, CellDim, KDim) - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - r_m = zero_field(grid, CellDim, KDim) - p_dtime = np.float64(5) - dbl_eps = np.float64(1e-9) +class TestHflxLimiterPdStencil01(StencilTest): + PROGRAM = hflx_limiter_pd_stencil_01 + OUTPUTS = ("r_m",) - ref = hflx_limiter_pd_stencil_01_numpy( - grid.connectivities[C2EDim], - geofac_div.asnumpy(), - p_cc.asnumpy(), - p_rhodz_now.asnumpy(), - p_mflx_tracer_h.asnumpy(), + @staticmethod + def reference( + grid, + geofac_div: np.ndarray, + p_cc: np.ndarray, + p_rhodz_now: np.ndarray, + p_mflx_tracer_h: np.ndarray, p_dtime, dbl_eps, - ) + **kwargs, + ): + geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + p_m_0 = np.maximum( + 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 0]] * geofac_div[:, 0] * p_dtime + ) + p_m_1 = np.maximum( + 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 1]] * geofac_div[:, 1] * p_dtime + ) + p_m_2 = np.maximum( + 0.0, p_mflx_tracer_h[grid.connectivities[C2EDim][:, 2]] * geofac_div[:, 2] * p_dtime + ) - hflx_limiter_pd_stencil_01.with_backend(backend)( - as_1D_sparse_field(geofac_div, CEDim), - p_cc, - p_rhodz_now, - p_mflx_tracer_h, - r_m, - p_dtime, - dbl_eps, - offset_provider={ - "C2CE": grid.get_offset_provider("C2CE"), - "C2E": grid.get_offset_provider("C2E"), - }, - ) - assert np.allclose(r_m.asnumpy(), ref) + p_m = p_m_0 + p_m_1 + p_m_2 + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) + + return dict(r_m=r_m) + + @pytest.fixture + def input_data(self, grid): + geofac_div = random_field(grid, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_cc = random_field(grid, CellDim, KDim) + p_rhodz_now = random_field(grid, CellDim, KDim) + p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) + r_m = zero_field(grid, CellDim, KDim) + p_dtime = np.float64(5) + dbl_eps = np.float64(1e-9) + return dict( + geofac_div=geofac_div_new, + p_cc=p_cc, + p_rhodz_now=p_rhodz_now, + p_mflx_tracer_h=p_mflx_tracer_h, + p_dtime=p_dtime, + dbl_eps=dbl_eps, + r_m=r_m, + ) From 62b543f0b5266184e8f7e04a6b438118f4cd55a0 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 10:38:40 +0100 Subject: [PATCH 13/26] add horadv stencil 1 --- .../stencil_tests/test_hor_adv_stencil_01.py | 111 +++++++++--------- 1 file changed, 53 insertions(+), 58 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py index 50ce69a4d6..51fb24ffb9 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hor_adv_stencil_01.py @@ -12,68 +12,63 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + reshape, + zero_field, +) -def hor_adv_stencil_01_numpy( - c2e: np.array, - p_mflx_tracer_h: np.array, - deepatmo_divh: np.array, - tracer_now: np.array, - rhodz_now: np.array, - rhodz_new: np.array, - geofac_div: np.array, - p_dtime, -) -> np.array: - geofac_div = np.expand_dims(geofac_div, axis=-1) +class TestHorAdvStencil01(StencilTest): + PROGRAM = hor_adv_stencil_01 + OUTPUTS = ("tracer_new",) - tracer_new = ( - tracer_now * rhodz_now - - p_dtime * deepatmo_divh * np.sum(p_mflx_tracer_h[c2e] * geofac_div, axis=1) - ) / rhodz_new - - return tracer_new - - -def test_hor_adv_stencil_01(backend): - grid = SimpleGrid() - - p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) - deepatmo_divh = random_field(grid, KDim) - tracer_now = random_field(grid, CellDim, KDim) - rhodz_now = random_field(grid, CellDim, KDim) - rhodz_new = random_field(grid, CellDim, KDim) - geofac_div = random_field(grid, CellDim, C2EDim) - geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) - tracer_new = random_field(grid, CellDim, KDim) - p_dtime = np.float64(5.0) - - ref = hor_adv_stencil_01_numpy( - grid.connectivities[C2EDim], - p_mflx_tracer_h.asnumpy(), - deepatmo_divh.asnumpy(), - tracer_now.asnumpy(), - rhodz_now.asnumpy(), - rhodz_new.asnumpy(), - geofac_div.asnumpy(), - p_dtime, - ) - hor_adv_stencil_01.with_backend(backend)( - p_mflx_tracer_h, - deepatmo_divh, - tracer_now, - rhodz_now, - rhodz_new, - geofac_div_new, - tracer_new, + @staticmethod + def reference( + grid, + p_mflx_tracer_h: np.array, + deepatmo_divh: np.array, + tracer_now: np.array, + rhodz_now: np.array, + rhodz_new: np.array, + geofac_div: np.array, p_dtime, - offset_provider={ - "C2E": grid.get_offset_provider("C2E"), - "C2CE": grid.get_offset_provider("C2CE"), - }, - ) - assert np.allclose(tracer_new.asnumpy(), ref) + **kwargs, + ) -> np.array: + geofac_div = reshape(geofac_div, grid.connectivities[C2EDim].shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + tracer_new = ( + tracer_now * rhodz_now + - p_dtime + * deepatmo_divh + * np.sum(p_mflx_tracer_h[grid.connectivities[C2EDim]] * geofac_div, axis=1) + ) / rhodz_new + return dict(tracer_new=tracer_new) + + @pytest.fixture + def input_data(self, grid): + p_mflx_tracer_h = random_field(grid, EdgeDim, KDim) + deepatmo_divh = random_field(grid, KDim) + tracer_now = random_field(grid, CellDim, KDim) + rhodz_now = random_field(grid, CellDim, KDim) + rhodz_new = random_field(grid, CellDim, KDim) + geofac_div = random_field(grid, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_dtime = np.float64(5.0) + tracer_new = zero_field(grid, CellDim, KDim) + return dict( + p_mflx_tracer_h=p_mflx_tracer_h, + deepatmo_divh=deepatmo_divh, + tracer_now=tracer_now, + rhodz_now=rhodz_now, + rhodz_new=rhodz_new, + geofac_div=geofac_div_new, + p_dtime=p_dtime, + tracer_new=tracer_new, + ) From e786e2103785e311fbf8750edefb53047d58bf6e Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 11:10:27 +0100 Subject: [PATCH 14/26] rbf intp edge stencil 1 --- .../test_rbf_intp_edge_stencil_01.py | 48 +++++++------------ 1 file changed, 16 insertions(+), 32 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py index 9f6c90d2a5..47882226d9 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_rbf_intp_edge_stencil_01.py @@ -12,42 +12,26 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def rbf_intp_edge_stencil_01_numpy( - e2c2e: np.array, - p_vn_in: np.array, - ptr_coeff: np.array, -) -> np.array: - ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) - p_vt_out = np.sum(p_vn_in[e2c2e] * ptr_coeff, axis=1) - return p_vt_out +class TestRbfIntpEdgeStencil01(StencilTest): + PROGRAM = rbf_intp_edge_stencil_01 + OUTPUTS = ("p_vt_out",) + @staticmethod + def reference(grid, p_vn_in: np.array, ptr_coeff: np.array, **kwargs) -> np.array: + ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) + p_vt_out = np.sum(p_vn_in[grid.connectivities[E2C2EDim]] * ptr_coeff, axis=1) + return dict(p_vt_out=p_vt_out) -def test_rbf_intp_edge_stencil_01(backend): - grid = SimpleGrid() - - p_vn_in = random_field(grid, EdgeDim, KDim) - ptr_coeff = random_field(grid, EdgeDim, E2C2EDim) - p_vt_out = zero_field(grid, EdgeDim, KDim) - - ref = rbf_intp_edge_stencil_01_numpy( - grid.connectivities[E2C2EDim], - p_vn_in.asnumpy(), - ptr_coeff.asnumpy(), - ) - - rbf_intp_edge_stencil_01.with_backend(backend)( - p_vn_in, - ptr_coeff, - p_vt_out, - offset_provider={ - "E2C2E": grid.get_offset_provider("E2C2E"), - }, - ) - assert np.allclose(p_vt_out.asnumpy(), ref) + @pytest.fixture + def input_data(self, grid): + p_vn_in = random_field(grid, EdgeDim, KDim) + ptr_coeff = random_field(grid, EdgeDim, E2C2EDim) + p_vt_out = zero_field(grid, EdgeDim, KDim) + return dict(p_vn_in=p_vn_in, ptr_coeff=ptr_coeff, p_vt_out=p_vt_out) From 37b1c4d58dc2c29c0f063ee060e8761cf975f3dd Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 11:55:25 +0100 Subject: [PATCH 15/26] add zero stencils --- .../model/atmosphere/advection/set_zero_c.py | 2 +- .../atmosphere/advection/set_zero_c_k.py | 2 +- .../tests/stencil_tests/test_set_zero_c.py | 21 ++++++++++++------- .../tests/stencil_tests/test_set_zero_c_k.py | 21 ++++++++++++------- 4 files changed, 28 insertions(+), 18 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py index e0860c9fb3..2d2905e2e4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py @@ -23,6 +23,6 @@ def _set_zero_c() -> Field[[CellDim], float]: return broadcast(0.0, (CellDim,)) -@program +@program(grid_type="unstructured") def set_zero_c(field: Field[[CellDim], float]): _set_zero_c(out=field) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py index f9db607eaf..8b352f8508 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py @@ -23,6 +23,6 @@ def _set_zero_c_k() -> Field[[CellDim, KDim], float]: return broadcast(0.0, (CellDim, KDim)) -@program +@program(grid_type="unstructured") def set_zero_c_k(field: Field[[CellDim, KDim], float]): _set_zero_c_k(out=field) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c.py b/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c.py index 4b6f2992c6..f15449b3f4 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c.py @@ -11,17 +11,22 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np +import pytest from icon4py.model.atmosphere.advection.set_zero_c import set_zero_c from icon4py.model.common.dimension import CellDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def test_set_zero_cell_k(backend): - grid = SimpleGrid() - field = random_field(grid, CellDim) +class TestSetZeroCellK(StencilTest): + PROGRAM = set_zero_c + OUTPUTS = ("field",) - set_zero_c.with_backend(backend)(field, offset_provider={}) - assert np.allclose(field.asnumpy(), zero_field(grid, CellDim).asnumpy()) + @staticmethod + def reference(grid, **kwargs): + return dict(field=zero_field(grid, CellDim).asnumpy()) + + @pytest.fixture + def input_data(self, grid): + field = random_field(grid, CellDim) + return dict(field=field) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c_k.py b/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c_k.py index 8811ce3768..f72a85188f 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c_k.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_set_zero_c_k.py @@ -11,17 +11,22 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np +import pytest from icon4py.model.atmosphere.advection.set_zero_c_k import set_zero_c_k from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def test_set_zero_c_k(backend): - grid = SimpleGrid() - field = random_field(grid, CellDim, KDim) +class TestSetZeroCK(StencilTest): + PROGRAM = set_zero_c_k + OUTPUTS = ("field",) - set_zero_c_k.with_backend(backend)(field, offset_provider={}) - assert np.allclose(field.asnumpy(), zero_field(grid, CellDim, KDim).asnumpy()) + @staticmethod + def reference(grid, **kwargs): + return dict(field=zero_field(grid, CellDim, KDim).asnumpy()) + + @pytest.fixture + def input_data(self, grid): + field = random_field(grid, CellDim, KDim) + return dict(field=field) From f150408373010106a377cb22aa96ff5fbf73179a Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 13:31:03 +0100 Subject: [PATCH 16/26] add upwind_vflux_ppm 1 --- .../advection/upwind_vflux_ppm_stencil_01.py | 2 +- .../test_upwind_vflux_ppm_stencil_01.py | 56 +++++++++---------- 2 files changed, 27 insertions(+), 31 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py index 953b507708..4d8c476b2f 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py @@ -29,7 +29,7 @@ def _upwind_vflux_ppm_stencil_01( return z_delta_q, z_a1 -@program +@program(grid_type="unstructured") def upwind_vflux_ppm_stencil_01( z_face_up: Field[[CellDim, KDim], float], z_face_low: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py index 446ee49525..a8f9c3303b 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_upwind_vflux_ppm_stencil_01.py @@ -12,38 +12,34 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.upwind_vflux_ppm_stencil_01 import ( upwind_vflux_ppm_stencil_01, ) from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field - - -def upwind_vflux_ppm_stencil_01_numpy( - z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray -) -> tuple[np.ndarray]: - z_delta_q = 0.5 * (z_face_up - z_face_low) - z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) - return z_delta_q, z_a1 - - -def test_upwind_vflux_ppm_stencil_01(backend): - grid = SimpleGrid() - z_face_up = random_field(grid, CellDim, KDim) - z_face_down = random_field(grid, CellDim, KDim) - p_cc = random_field(grid, CellDim, KDim) - z_delta_q = zero_field(grid, CellDim, KDim) - z_a1 = zero_field(grid, CellDim, KDim) - - ref_z_delta_q, ref_z_a1 = upwind_vflux_ppm_stencil_01_numpy( - z_face_up.asnumpy(), z_face_down.asnumpy(), p_cc.asnumpy() - ) - - upwind_vflux_ppm_stencil_01.with_backend(backend)( - z_face_up, z_face_down, p_cc, z_delta_q, z_a1, offset_provider={} - ) - - assert np.allclose(ref_z_delta_q, z_delta_q.asnumpy()) - assert np.allclose(ref_z_a1, z_a1.asnumpy()) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestUpwindVfluxPpmStencil01(StencilTest): + PROGRAM = upwind_vflux_ppm_stencil_01 + OUTPUTS = ("z_delta_q", "z_a1") + + @staticmethod + def reference( + grid, z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray, **kwargs + ) -> tuple[np.ndarray]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + return dict(z_delta_q=z_delta_q, z_a1=z_a1) + + @pytest.fixture + def input_data(self, grid): + z_face_up = random_field(grid, CellDim, KDim) + z_face_low = random_field(grid, CellDim, KDim) + p_cc = random_field(grid, CellDim, KDim) + z_delta_q = zero_field(grid, CellDim, KDim) + z_a1 = zero_field(grid, CellDim, KDim) + return dict( + z_face_up=z_face_up, z_face_low=z_face_low, p_cc=p_cc, z_delta_q=z_delta_q, z_a1=z_a1 + ) From 7c817b7ac793e8218d5728c27d43000d1bc752ec Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 13:59:31 +0100 Subject: [PATCH 17/26] TestVertAdvStencil01 --- .../advection/vert_adv_stencil_01.py | 2 +- .../stencil_tests/test_vert_adv_stencil_01.py | 94 +++++++++---------- 2 files changed, 44 insertions(+), 52 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index 78736914db..7e236e846b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -35,7 +35,7 @@ def _vert_adv_stencil_01( return tracer_new -@program +@program(grid_type="unstructured") def vert_adv_stencil_01( tracer_now: Field[[CellDim, KDim], float], rhodz_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py index b44ea19888..598d496e3c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py @@ -12,61 +12,53 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def vert_adv_stencil_01_numpy( - tracer_now: np.array, - rhodz_now: np.array, - p_mflx_tracer_v: np.array, - deepatmo_divzl: np.array, - deepatmo_divzu: np.array, - rhodz_new: np.array, - p_dtime, -) -> np.array: - tracer_new = ( - tracer_now * rhodz_now - + p_dtime - * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) - ) / rhodz_new +class TestVertAdvStencil01(StencilTest): + PROGRAM = vert_adv_stencil_01 + OUTPUTS = ("tracer_new",) - return tracer_new - - -def test_vert_adv_stencil_01(backend): - grid = SimpleGrid() - - tracer_now = random_field(grid, CellDim, KDim) - rhodz_now = random_field(grid, CellDim, KDim) - p_mflx_tracer_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) - deepatmo_divzl = random_field(grid, KDim) - deepatmo_divzu = random_field(grid, KDim) - rhodz_new = random_field(grid, CellDim, KDim) - tracer_new = zero_field(grid, CellDim, KDim) - p_dtime = np.float64(5.0) - - ref = vert_adv_stencil_01_numpy( - tracer_now.asnumpy(), - rhodz_now.asnumpy(), - p_mflx_tracer_v.asnumpy(), - deepatmo_divzl.asnumpy(), - deepatmo_divzu.asnumpy(), - rhodz_new.asnumpy(), - p_dtime, - ) - vert_adv_stencil_01.with_backend(backend)( - tracer_now, - rhodz_now, - p_mflx_tracer_v, - deepatmo_divzl, - deepatmo_divzu, - rhodz_new, - tracer_new, + @staticmethod + def reference( + grid, + tracer_now: np.array, + rhodz_now: np.array, + p_mflx_tracer_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + rhodz_new: np.array, p_dtime, - offset_provider={"Koff": KDim}, - ) - assert np.allclose(tracer_new.asnumpy()[:, :-1], ref[:, :-1]) + **kwargs, + ) -> np.array: + tracer_new = ( + tracer_now * rhodz_now + + p_dtime + * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) + ) / rhodz_new + return dict(tracer_new=tracer_new) + + @pytest.fixture + def input_data(self, grid): + tracer_now = random_field(grid, CellDim, KDim) + rhodz_now = random_field(grid, CellDim, KDim) + p_mflx_tracer_v = random_field(grid, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(grid, KDim) + deepatmo_divzu = random_field(grid, KDim) + rhodz_new = random_field(grid, CellDim, KDim) + p_dtime = np.float64(5.0) + tracer_new = zero_field(grid, CellDim, KDim) + return dict( + tracer_now=tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=deepatmo_divzl, + deepatmo_divzu=deepatmo_divzu, + rhodz_new=rhodz_new, + p_dtime=p_dtime, + tracer_new=tracer_new, + ) From b82d9c7f6d8e8fb21af59054bc9ffad2dfa2e75b Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 14:33:51 +0100 Subject: [PATCH 18/26] Add vlimit prbl sm stencils --- .../advection/v_limit_prbl_sm_stencil_01.py | 2 +- .../advection/v_limit_prbl_sm_stencil_02.py | 2 +- .../test_vlimit_prbl_sm_stencil_01.py | 56 +++++------- .../test_vlimit_prbl_sm_stencil_02.py | 87 ++++++++----------- 4 files changed, 61 insertions(+), 86 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py index 4b35c4de09..51ff924fab 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py @@ -39,7 +39,7 @@ def _v_limit_prbl_sm_stencil_01( return l_limit -@program +@program(grid_type="unstructured") def v_limit_prbl_sm_stencil_01( p_face: Field[[CellDim, KDim], float], p_cc: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py index be7f31fd37..7d6f08ff8e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py @@ -43,7 +43,7 @@ def _v_limit_prbl_sm_stencil_02( return q_face_up, q_face_low -@program +@program(grid_type="unstructured") def v_limit_prbl_sm_stencil_02( l_limit: Field[[CellDim, KDim], int32], p_face: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py index b4a5770f7c..f6fc3c2ff8 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_01.py @@ -12,42 +12,28 @@ # 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.v_limit_prbl_sm_stencil_01 import v_limit_prbl_sm_stencil_01 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field - - -def v_limit_prbl_sm_stencil_01_numpy( - p_face: np.array, - p_cc: np.array, -): - z_delta = p_face[:, :-1] - p_face[:, 1:] - z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - - l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) - - return l_limit - - -def test_v_limit_prbl_sm_stencil_01(backend): - grid = SimpleGrid() - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) - l_limit = zero_field(grid, CellDim, KDim, dtype=int32) - - l_limit_ref = v_limit_prbl_sm_stencil_01_numpy( - p_face.asnumpy(), - p_cc.asnumpy(), - ) - - v_limit_prbl_sm_stencil_01.with_backend(backend)( - p_face, - p_cc, - l_limit, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(l_limit_ref, l_limit.asnumpy()) +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + + +class TestVLimitPrblSmStencil01(StencilTest): + PROGRAM = v_limit_prbl_sm_stencil_01 + OUTPUTS = ("l_limit",) + + @staticmethod + def reference(grid, p_face: np.array, p_cc: np.array, **kwargs): + z_delta = p_face[:, :-1] - p_face[:, 1:] + z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) + l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) + return dict(l_limit=l_limit) + + @pytest.fixture + def input_data(self, grid): + p_cc = random_field(grid, CellDim, KDim) + p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) + l_limit = zero_field(grid, CellDim, KDim, dtype=int32) + return dict(p_face=p_face, p_cc=p_cc, l_limit=l_limit) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py index 92a6b7ae6f..c79595b338 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_vlimit_prbl_sm_stencil_02.py @@ -12,58 +12,47 @@ # 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.v_limit_prbl_sm_stencil_02 import v_limit_prbl_sm_stencil_02 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field - - -def v_limit_prbl_sm_stencil_02_numpy( - l_limit: np.array, - p_face: np.array, - p_cc: np.array, -): - q_face_up, q_face_low = np.where( - l_limit != int32(0), - np.where( - (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), - (p_cc, p_cc), +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + random_field, + random_mask, + zero_field, +) + + +class TestVLimitPrblSmStencil02(StencilTest): + PROGRAM = v_limit_prbl_sm_stencil_02 + OUTPUTS = ("p_face_up", "p_face_low") + + @staticmethod + def reference(grid, l_limit: np.array, p_face: np.array, p_cc: np.array, **kwargs): + q_face_up, q_face_low = np.where( + l_limit != int32(0), np.where( - p_face[:, :-1] > p_face[:, 1:], - (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), - (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), + (p_cc, p_cc), + np.where( + p_face[:, :-1] > p_face[:, 1:], + (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), + (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), + ), ), - ), - (p_face[:, :-1], p_face[:, 1:]), - ) - - return q_face_up, q_face_low - - -def test_v_limit_prbl_sm_stencil_02(backend): - grid = SimpleGrid() - l_limit = random_mask(grid, CellDim, KDim, dtype=int32) - p_cc = random_field(grid, CellDim, KDim) - p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) - p_face_up = zero_field(grid, CellDim, KDim) - p_face_low = zero_field(grid, CellDim, KDim) - - p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_02_numpy( - l_limit.asnumpy(), - p_face.asnumpy(), - p_cc.asnumpy(), - ) - - v_limit_prbl_sm_stencil_02.with_backend(backend)( - l_limit, - p_face, - p_cc, - p_face_up, - p_face_low, - offset_provider={"Koff": KDim}, - ) - - assert np.allclose(p_face_up_ref[:, :-1], p_face_up.asnumpy()[:, :-1]) - assert np.allclose(p_face_low_ref[:, :-1], p_face_low.asnumpy()[:, :-1]) + (p_face[:, :-1], p_face[:, 1:]), + ) + return dict(p_face_up=q_face_up, p_face_low=q_face_low) + + @pytest.fixture + def input_data(self, grid): + l_limit = random_mask(grid, CellDim, KDim, dtype=int32) + p_cc = random_field(grid, CellDim, KDim) + p_face = random_field(grid, CellDim, KDim, extend={KDim: 1}) + p_face_up = zero_field(grid, CellDim, KDim) + p_face_low = zero_field(grid, CellDim, KDim) + return dict( + l_limit=l_limit, p_cc=p_cc, p_face=p_face, p_face_up=p_face_up, p_face_low=p_face_low + ) From 90db4d844c3ae46aef1335ee96c41fbc18ab864b Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Thu, 11 Jan 2024 16:13:37 +0100 Subject: [PATCH 19/26] Add missing assert in test --- .../tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py index ce20330f59..7cfc429956 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py @@ -118,5 +118,5 @@ def test_hflx_diffusion_mo_stencil_03(backend): r_m, offset_provider={"C2E2C": grid.get_offset_provider("C2E2C")}, ) - np.allclose(r_p_ref, r_p.asnumpy()) - np.allclose(r_m_ref, r_m.asnumpy()) + assert np.allclose(r_p_ref, r_p.asnumpy()) + assert np.allclose(r_m_ref, r_m.asnumpy()) From 057d9c2e72f967486714b8732d17ccbffe2c88ef Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 12 Jan 2024 17:39:11 +0100 Subject: [PATCH 20/26] Fix error in numpy test_hflx_limiter_mo_stencil_03 --- .../tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py index 7cfc429956..3f461168a3 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py @@ -36,7 +36,7 @@ def hflx_limiter_mo_stencil_03_numpy( c2e2c, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct ) r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) - r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out * dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) return r_p, r_m From b154bc10f94ce5a7a6eded61f1a7c9fea32ae94d Mon Sep 17 00:00:00 2001 From: Nina Burgdorfer Date: Fri, 12 Jan 2024 17:40:25 +0100 Subject: [PATCH 21/26] Fix missing numpy ref. stencils calls --- .../test_hflx_limiter_pd_stencil_02.py | 38 ++++++++----------- 1 file changed, 15 insertions(+), 23 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py index 0c55e19330..84429baa15 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py @@ -67,7 +67,6 @@ def test_hflx_limiter_pd_stencil_02_nowhere_matching_refin_ctl(backend): assert np.allclose(p_mflx_tracer_h_in.asnumpy(), ref) -# TODO (Nina) what is this test for? its trivially true def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(backend): grid = SimpleGrid() bound = np.int32(7) @@ -75,6 +74,14 @@ def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(backend): r_m = random_field(grid, CellDim, KDim) p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) + ref = hflx_limiter_pd_stencil_02_numpy( + grid.connectivities[E2CDim], + refin_ctrl.asnumpy(), + r_m.asnumpy(), + p_mflx_tracer_h_in.asnumpy(), + bound, + ) + hflx_limiter_pd_stencil_02.with_backend(backend)( refin_ctrl, r_m, @@ -84,10 +91,9 @@ def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(backend): "E2C": grid.get_offset_provider("E2C"), }, ) - assert np.allclose(p_mflx_tracer_h_in.asnumpy(), p_mflx_tracer_h_in.asnumpy()) + assert np.allclose(p_mflx_tracer_h_in.asnumpy(), ref) -# TODO (Nina) what is this test for? its trivially true def test_hflx_limiter_pd_stencil_02_partly_matching_refin_ctl(backend): grid = SimpleGrid() bound = np.int32(4) @@ -96,27 +102,13 @@ def test_hflx_limiter_pd_stencil_02_partly_matching_refin_ctl(backend): r_m = random_field(grid, CellDim, KDim) p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) - hflx_limiter_pd_stencil_02.with_backend(backend)( - refin_ctrl, - r_m, - p_mflx_tracer_h_in, + ref = hflx_limiter_pd_stencil_02_numpy( + grid.connectivities[E2CDim], + refin_ctrl.asnumpy(), + r_m.asnumpy(), + p_mflx_tracer_h_in.asnumpy(), bound, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, ) - assert np.allclose(p_mflx_tracer_h_in.asnumpy(), p_mflx_tracer_h_in.asnumpy()) - - -# TODO (Nina) what is this test for? its trivially true -def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl_does_not_change_inout_arg( - backend, -): - grid = SimpleGrid() - bound = np.int32(7) - refin_ctrl = constant_field(grid, bound, EdgeDim, dtype=np.int32) - r_m = random_field(grid, CellDim, KDim) - p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) hflx_limiter_pd_stencil_02.with_backend(backend)( refin_ctrl, @@ -127,4 +119,4 @@ def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl_does_not_chang "E2C": grid.get_offset_provider("E2C"), }, ) - assert np.allclose(p_mflx_tracer_h_in.asnumpy(), p_mflx_tracer_h_in.asnumpy()) + assert np.allclose(p_mflx_tracer_h_in.asnumpy(), ref) From ffb6b79d79941e86da85b864dc775a8506a7f1cc Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 16 Jan 2024 14:13:10 +0100 Subject: [PATCH 22/26] refactor hflx limiter mo stencil 03 --- .../test_hflx_limiter_mo_stencil_03.py | 155 ++++++++---------- 1 file changed, 67 insertions(+), 88 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py index 3f461168a3..f8819ee83e 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_mo_stencil_03.py @@ -11,101 +11,52 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_03 import ( hflx_limiter_mo_stencil_03, hflx_limiter_mo_stencil_03_min_max, ) from icon4py.model.common.dimension import C2E2CDim, CellDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field -def hflx_limiter_mo_stencil_03_numpy( - c2e2c: np.ndarray, - z_tracer_max: np.ndarray, - z_tracer_min: np.ndarray, - beta_fct: float, - r_beta_fct: float, - z_mflx_anti_in: np.ndarray, - z_mflx_anti_out: np.ndarray, - z_tracer_new_low: np.ndarray, - dbl_eps: float, -): - z_max, z_min = hflx_limiter_mo_stencil_03_min_max_numpy( - c2e2c, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct - ) - r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) - r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) - return r_p, r_m +class TestHflxLimiterMoStencil03MinMax(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_03_min_max + OUTPUTS = ("z_max", "z_min") + @staticmethod + def reference(grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, **kwargs): + c2e2c = grid.connectivities[C2E2CDim] + z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) + z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) + return dict(z_max=z_max, z_min=z_min) -def hflx_limiter_mo_stencil_03_min_max_numpy( - c2e2c: np.array, - z_tracer_max: np.ndarray, - z_tracer_min: np.ndarray, - beta_fct: float, - r_beta_fct: float, -) -> tuple[np.ndarray]: - z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) - z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) - return z_max, z_min + @pytest.fixture + def input_data(self, grid): + z_tracer_max = random_field(grid, CellDim, KDim) + z_tracer_min = random_field(grid, CellDim, KDim) + beta_fct = 0.9 + r_beta_fct = 0.3 + z_max = zero_field(grid, CellDim, KDim) + z_min = zero_field(grid, CellDim, KDim) + return dict( + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + beta_fct=beta_fct, + r_beta_fct=r_beta_fct, + z_max=z_max, + z_min=z_min, + ) -def test_hflx_diffusion_mo_stencil_03_min_max(backend): - grid = SimpleGrid() - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - z_max = zero_field(grid, CellDim, KDim) - z_min = zero_field(grid, CellDim, KDim) - beta_fct = 0.9 - r_beta_fct = 0.3 - z_max_ref, z_min_ref = hflx_limiter_mo_stencil_03_min_max_numpy( - grid.connectivities[C2E2CDim], - z_tracer_max.asnumpy(), - z_tracer_min.asnumpy(), - beta_fct, - r_beta_fct, - ) - hflx_limiter_mo_stencil_03_min_max.with_backend(backend)( - z_tracer_max, - z_tracer_min, - beta_fct, - r_beta_fct, - z_max, - z_min, - offset_provider={"C2E2C": grid.get_offset_provider("C2E2C")}, - ) - assert np.allclose(z_max.asnumpy(), z_max_ref) - assert np.allclose(z_min.asnumpy(), z_min_ref) - +class TestHflxLimiterMoStencil03(StencilTest): + PROGRAM = hflx_limiter_mo_stencil_03 + OUTPUTS = ("r_p", "r_m") -def test_hflx_diffusion_mo_stencil_03(backend): - grid = SimpleGrid() - z_tracer_max = random_field(grid, CellDim, KDim) - z_tracer_min = random_field(grid, CellDim, KDim) - beta_fct = 0.4 - r_beta_fct = 0.6 - z_mflx_anti_in = random_field(grid, CellDim, KDim) - z_mflx_anti_out = random_field(grid, CellDim, KDim) - z_tracer_new_low = random_field(grid, CellDim, KDim) - dbl_eps = 1e-5 - r_p = zero_field(grid, CellDim, KDim) - r_m = zero_field(grid, CellDim, KDim) - - r_p_ref, r_m_ref = hflx_limiter_mo_stencil_03_numpy( - grid.connectivities[C2E2CDim], - z_tracer_max.asnumpy(), - z_tracer_min.asnumpy(), - beta_fct, - r_beta_fct, - z_mflx_anti_in.asnumpy(), - z_mflx_anti_out.asnumpy(), - z_tracer_new_low.asnumpy(), - dbl_eps, - ) - - hflx_limiter_mo_stencil_03.with_backend(backend)( + @staticmethod + def reference( + grid, z_tracer_max, z_tracer_min, beta_fct, @@ -114,9 +65,37 @@ def test_hflx_diffusion_mo_stencil_03(backend): z_mflx_anti_out, z_tracer_new_low, dbl_eps, - r_p, - r_m, - offset_provider={"C2E2C": grid.get_offset_provider("C2E2C")}, - ) - assert np.allclose(r_p_ref, r_p.asnumpy()) - assert np.allclose(r_m_ref, r_m.asnumpy()) + **kwargs, + ): + res = TestHflxLimiterMoStencil03MinMax.reference( + grid, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) + z_max, z_min = res["z_max"], res["z_min"] + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) + return dict(r_p=r_p, r_m=r_m) + + @pytest.fixture + def input_data(self, grid): + z_tracer_max = random_field(grid, CellDim, KDim) + z_tracer_min = random_field(grid, CellDim, KDim) + beta_fct = 0.4 + r_beta_fct = 0.6 + z_mflx_anti_in = random_field(grid, CellDim, KDim) + z_mflx_anti_out = random_field(grid, CellDim, KDim) + z_tracer_new_low = random_field(grid, CellDim, KDim) + dbl_eps = 1e-5 + r_p = zero_field(grid, CellDim, KDim) + r_m = zero_field(grid, CellDim, KDim) + return dict( + z_tracer_max=z_tracer_max, + z_tracer_min=z_tracer_min, + beta_fct=beta_fct, + r_beta_fct=r_beta_fct, + z_mflx_anti_in=z_mflx_anti_in, + z_mflx_anti_out=z_mflx_anti_out, + z_tracer_new_low=z_tracer_new_low, + dbl_eps=dbl_eps, + r_p=r_p, + r_m=r_m, + ) From d971c059b708da39073e011d241f7aa526ca4b30 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 16 Jan 2024 15:41:33 +0100 Subject: [PATCH 23/26] hflx limiter pd stencil 02 --- .../advection/hflx_limiter_pd_stencil_02.py | 10 +- .../test_hflx_limiter_pd_stencil_02.py | 152 ++++++------------ 2 files changed, 51 insertions(+), 111 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py index 9d1be46771..2380296e9e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -21,16 +21,16 @@ def _hflx_limiter_pd_stencil_02( refin_ctrl: Field[[EdgeDim], int32], r_m: Field[[CellDim, KDim], float], - p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], bound: int32, ) -> Field[[EdgeDim, KDim], float]: p_mflx_tracer_h_out = where( refin_ctrl == bound, - p_mflx_tracer_h_in, + p_mflx_tracer_h, where( - p_mflx_tracer_h_in >= 0.0, - p_mflx_tracer_h_in * r_m(E2C[0]), - p_mflx_tracer_h_in * r_m(E2C[1]), + p_mflx_tracer_h >= 0.0, + p_mflx_tracer_h * r_m(E2C[0]), + p_mflx_tracer_h * r_m(E2C[1]), ), ) return p_mflx_tracer_h_out diff --git a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py index 84429baa15..212f95664a 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_hflx_limiter_pd_stencil_02.py @@ -10,113 +10,53 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later - import numpy as np +import pytest from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 from icon4py.model.common.dimension import CellDim, E2CDim, EdgeDim, KDim -from icon4py.model.common.grid.simple import SimpleGrid -from icon4py.model.common.test_utils.helpers import constant_field, random_field - - -def hflx_limiter_pd_stencil_02_numpy( - e2c: np.array, - refin_ctrl: np.array, - r_m: np.array, - p_mflx_tracer_h_in: np.array, - bound, -): - r_m_e2c = r_m[e2c] - refin_ctrl = np.expand_dims(refin_ctrl, axis=-1) - p_mflx_tracer_h_out = np.where( - refin_ctrl != bound, - np.where( - p_mflx_tracer_h_in >= 0, - p_mflx_tracer_h_in * r_m_e2c[:, 0], - p_mflx_tracer_h_in * r_m_e2c[:, 1], - ), - p_mflx_tracer_h_in, - ) - return p_mflx_tracer_h_out - - -def test_hflx_limiter_pd_stencil_02_nowhere_matching_refin_ctl(backend): - grid = SimpleGrid() - bound = np.int32(7) - refin_ctrl = constant_field(grid, 4, EdgeDim, dtype=np.int32) - r_m = random_field(grid, CellDim, KDim) - p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) - - ref = hflx_limiter_pd_stencil_02_numpy( - grid.connectivities[E2CDim], - refin_ctrl.asnumpy(), - r_m.asnumpy(), - p_mflx_tracer_h_in.asnumpy(), - bound, - ) - - hflx_limiter_pd_stencil_02.with_backend(backend)( - refin_ctrl, - r_m, - p_mflx_tracer_h_in, - bound, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) - assert np.allclose(p_mflx_tracer_h_in.asnumpy(), ref) - - -def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(backend): - grid = SimpleGrid() - bound = np.int32(7) - refin_ctrl = constant_field(grid, bound, EdgeDim, dtype=np.int32) - r_m = random_field(grid, CellDim, KDim) - p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) - - ref = hflx_limiter_pd_stencil_02_numpy( - grid.connectivities[E2CDim], - refin_ctrl.asnumpy(), - r_m.asnumpy(), - p_mflx_tracer_h_in.asnumpy(), - bound, - ) - - hflx_limiter_pd_stencil_02.with_backend(backend)( - refin_ctrl, - r_m, - p_mflx_tracer_h_in, - bound, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) - assert np.allclose(p_mflx_tracer_h_in.asnumpy(), ref) - - -def test_hflx_limiter_pd_stencil_02_partly_matching_refin_ctl(backend): - grid = SimpleGrid() - bound = np.int32(4) - refin_ctrl = constant_field(grid, 5, EdgeDim, dtype=np.int32) - refin_ctrl[2:6] = bound - r_m = random_field(grid, CellDim, KDim) - p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) - - ref = hflx_limiter_pd_stencil_02_numpy( - grid.connectivities[E2CDim], - refin_ctrl.asnumpy(), - r_m.asnumpy(), - p_mflx_tracer_h_in.asnumpy(), - bound, - ) - - hflx_limiter_pd_stencil_02.with_backend(backend)( - refin_ctrl, - r_m, - p_mflx_tracer_h_in, - bound, - offset_provider={ - "E2C": grid.get_offset_provider("E2C"), - }, - ) - assert np.allclose(p_mflx_tracer_h_in.asnumpy(), ref) +from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, random_field + + +class TestHflxLimiterPdStencil02(StencilTest): + PROGRAM = hflx_limiter_pd_stencil_02 + OUTPUTS = ("p_mflx_tracer_h",) + + @staticmethod + def reference(grid, refin_ctrl, r_m, p_mflx_tracer_h, bound, **kwargs): + e2c = grid.connectivities[E2CDim] + r_m_e2c = r_m[e2c] + refin_ctrl_expanded = np.expand_dims(refin_ctrl, axis=-1) + p_mflx_tracer_h_out = np.where( + refin_ctrl_expanded != bound, + np.where( + p_mflx_tracer_h >= 0, + p_mflx_tracer_h * r_m_e2c[:, 0], + p_mflx_tracer_h * r_m_e2c[:, 1], + ), + p_mflx_tracer_h, + ) + return dict(p_mflx_tracer_h=p_mflx_tracer_h_out) + + @pytest.fixture(params=[("no_match", 4), ("everywhere_match", 7), ("partly_match", 4)]) + def input_data(self, request, grid): + bound = np.int32(7) + scenario, ctrl_value = request.param + + if scenario == "no_match": + refin_ctrl = constant_field(grid, ctrl_value, EdgeDim, dtype=np.int32) + elif scenario == "everywhere_match": + refin_ctrl = constant_field(grid, bound, EdgeDim, dtype=np.int32) + elif scenario == "partly_match": + refin_ctrl = constant_field(grid, 5, EdgeDim, dtype=np.int32) + refin_ctrl[2:6] = bound + + r_m = random_field(grid, CellDim, KDim) + p_mflx_tracer_h_in = random_field(grid, EdgeDim, KDim) + + return dict( + refin_ctrl=refin_ctrl, + r_m=r_m, + p_mflx_tracer_h=p_mflx_tracer_h_in, + bound=bound, + ) From fab872279f9f70e96e8c165cf7ca329c8c26cdf0 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 16 Jan 2024 16:57:53 +0100 Subject: [PATCH 24/26] Fix test --- .../stencil_tests/test_face_val_ppm_stencil_02c.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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 b83ace8859..7bdf4edd17 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 @@ -15,12 +15,18 @@ 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 StencilTest, random_field +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field class TestFaceValPpmStencil02c(StencilTest): PROGRAM = face_val_ppm_stencil_02c - OUTPUTS = (("p_face", (slice(None), slice(1, None))),) + OUTPUTS = ( + Output( + name="p_face", + refslice=(slice(None), slice(1, None)), + gtslice=(slice(None), slice(1, None)), + ), + ) @staticmethod def reference(grid, p_cc: np.array, **kwargs): From c6cb54d1ab609565a16eefd62736009809dcfe1a Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 17 Jan 2024 10:54:18 +0100 Subject: [PATCH 25/26] Use GridType --- .../model/atmosphere/advection/btraj_dreg_stencil_01.py | 4 ++-- .../model/atmosphere/advection/btraj_dreg_stencil_02.py | 4 ++-- .../model/atmosphere/advection/btraj_dreg_stencil_03.py | 4 ++-- .../model/atmosphere/advection/face_val_ppm_stencil_01.py | 4 ++-- .../model/atmosphere/advection/face_val_ppm_stencil_02.py | 4 ++-- .../model/atmosphere/advection/face_val_ppm_stencil_02a.py | 4 ++-- .../model/atmosphere/advection/face_val_ppm_stencil_02b.py | 4 ++-- .../model/atmosphere/advection/face_val_ppm_stencil_02c.py | 4 ++-- .../model/atmosphere/advection/face_val_ppm_stencil_05.py | 4 ++-- .../model/atmosphere/advection/hflx_limiter_mo_stencil_02.py | 4 ++-- .../src/icon4py/model/atmosphere/advection/set_zero_c.py | 4 ++-- .../src/icon4py/model/atmosphere/advection/set_zero_c_k.py | 4 ++-- .../model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py | 4 ++-- .../model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py | 4 ++-- .../model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py | 4 ++-- .../icon4py/model/atmosphere/advection/vert_adv_stencil_01.py | 4 ++-- 16 files changed, 32 insertions(+), 32 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py index 9c21bd6693..d278ba27c2 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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, broadcast, where @@ -30,7 +30,7 @@ def _btraj_dreg_stencil_01( return lvn_sys_pos -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def btraj_dreg_stencil_01( lcounterclock: bool, p_vn: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py index 0431332a36..3537d58045 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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, broadcast, int32, sqrt, where @@ -34,7 +34,7 @@ def _btraj_dreg_stencil_02( return opt_famask_dsl -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def btraj_dreg_stencil_02( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py index e95f95625b..d2c2fd2c72 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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, int32, where @@ -102,7 +102,7 @@ def _btraj_dreg_stencil_03( ) -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def btraj_dreg_stencil_03( p_vn: Field[[EdgeDim, KDim], float], p_vt: Field[[EdgeDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py index d0e9ce14fc..860d7a33f4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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, broadcast, int32, where @@ -70,7 +70,7 @@ def _face_val_ppm_stencil_01( return z_slope -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def face_val_ppm_stencil_01( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py index a033f91457..bb2ae279b7 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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, broadcast, int32, where @@ -71,7 +71,7 @@ def _face_val_ppm_stencil_02( return p_face -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def face_val_ppm_stencil_02( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py index 2052867155..a6b426e243 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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 @@ -29,7 +29,7 @@ def _face_val_ppm_stencil_02a( return p_face -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def face_val_ppm_stencil_02a( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py index 413b1a25e4..09a5188d42 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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 @@ -25,7 +25,7 @@ def _face_val_ppm_stencil_02b( return p_face -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def face_val_ppm_stencil_02b( p_cc: Field[[CellDim, KDim], float], p_face: Field[[CellDim, KDim], float], 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 72a093666e..0bf3ddd9a1 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 @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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 @@ -25,7 +25,7 @@ def _face_val_ppm_stencil_02c( return p_face -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], p_face: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py index 018e720496..b5c3ee3132 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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 @@ -51,7 +51,7 @@ def _face_val_ppm_stencil_05( return p_face -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def face_val_ppm_stencil_05( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py index fef2e48c94..51e34da114 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import int32, maximum, minimum, where @@ -45,7 +45,7 @@ def _hflx_limiter_mo_stencil_02( return (z_tracer_new_out, z_tracer_max_out, z_tracer_min_out) -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def hflx_limiter_mo_stencil_02( refin_ctrl: Field[[CellDim], int32], p_cc: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py index 2d2905e2e4..93e21929e3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast @@ -23,6 +23,6 @@ def _set_zero_c() -> Field[[CellDim], float]: return broadcast(0.0, (CellDim,)) -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def set_zero_c(field: Field[[CellDim], float]): _set_zero_c(out=field) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py index 8b352f8508..a6b50d09ef 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import broadcast @@ -23,6 +23,6 @@ def _set_zero_c_k() -> Field[[CellDim, KDim], float]: return broadcast(0.0, (CellDim, KDim)) -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def set_zero_c_k(field: Field[[CellDim, KDim], float]): _set_zero_c_k(out=field) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py index 4d8c476b2f..b0824647f6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py @@ -11,7 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from gt4py.next.common import Field +from gt4py.next.common import Field, GridType from gt4py.next.ffront.decorator import field_operator, program from icon4py.model.common.dimension import CellDim, KDim @@ -29,7 +29,7 @@ def _upwind_vflux_ppm_stencil_01( return z_delta_q, z_a1 -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def upwind_vflux_ppm_stencil_01( z_face_up: Field[[CellDim, KDim], float], z_face_low: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py index 51ff924fab..ea12f547d4 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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 ( # noqa: A004 # import gt4py builtin Field, @@ -39,7 +39,7 @@ def _v_limit_prbl_sm_stencil_01( return l_limit -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def v_limit_prbl_sm_stencil_01( p_face: Field[[CellDim, KDim], float], p_cc: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py index 7d6f08ff8e..8717d8fd06 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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, FieldOffset, int32, minimum, where @@ -43,7 +43,7 @@ def _v_limit_prbl_sm_stencil_02( return q_face_up, q_face_low -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def v_limit_prbl_sm_stencil_02( l_limit: Field[[CellDim, KDim], int32], p_face: Field[[CellDim, KDim], float], diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py index 7e236e846b..44f933b18f 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -10,7 +10,7 @@ # distribution for a copy of the license or check . # # 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 @@ -35,7 +35,7 @@ def _vert_adv_stencil_01( return tracer_new -@program(grid_type="unstructured") +@program(grid_type=GridType.UNSTRUCTURED) def vert_adv_stencil_01( tracer_now: Field[[CellDim, KDim], float], rhodz_now: Field[[CellDim, KDim], float], From de88c3f4d3bc68e1cd458521a0a42bd80ca8fdf3 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 17 Jan 2024 11:07:43 +0100 Subject: [PATCH 26/26] Slice p_face --- .../tests/stencil_tests/test_face_val_ppm_stencil_05.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py index 7e5c163fc5..6d7115b942 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_05.py @@ -16,12 +16,12 @@ from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 from icon4py.model.common.dimension import CellDim, KDim -from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field +from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field, zero_field class TestFaceValPpmStencil05(StencilTest): PROGRAM = face_val_ppm_stencil_05 - OUTPUTS = ("p_face",) + OUTPUTS = (Output("p_face", gtslice=(slice(None), slice(2, None))),) @staticmethod def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.array, **kwargs): @@ -59,10 +59,6 @@ def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.arra + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 ) ) - - # Concatenate the NaN column to the existing array - nan_column = np.full((p_face.shape[0], 2), np.nan) - p_face = np.concatenate((nan_column, p_face), axis=1) return dict(p_face=p_face) @pytest.fixture