Refactor advection stencils to use StencilTest

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 <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@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],

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

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

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 <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@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],

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
+@program(grid_type=GridType.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

@@ -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
+@program(grid_type=GridType.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

@@ -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
+@program(grid_type=GridType.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

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

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 <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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

@@ -10,7 +10,7 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # 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
+@program(grid_type=GridType.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,
+        )