Refactor advection stencils to use StencilTest

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],

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],

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],

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],

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],

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],

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],

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],

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],

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],

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

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)

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)

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],

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],

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],

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],

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

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