Add Output helper class

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 <https://www.gnu.org/licenses/>.
 #
 # 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)

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):

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}'"