Pass domain sizes to all programs

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py

@@ -87,6 +87,10 @@ def mo_advection_traj_btraj_compute_o1_dsl(
     p_distv_bary_1: Field[[EdgeDim, KDim], vpfloat],
     p_distv_bary_2: Field[[EdgeDim, KDim], vpfloat],
     p_dthalf: wpfloat,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _mo_advection_traj_btraj_compute_o1_dsl(
         p_vn,
@@ -101,4 +105,8 @@ def mo_advection_traj_btraj_compute_o1_dsl(
         dual_normal_cell_2,
         p_dthalf,
         out=(p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2),
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/advection/tests/stencil_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py

@@ -47,7 +47,7 @@ def reference(
         dual_normal_cell_2: np.array,
         p_dthalf: float,
         **kwargs,
-    ) -> np.array:
+    ) -> dict:
         e2c = grid.connectivities[E2CDim]
         cell_idx = cell_idx.reshape(e2c.shape)
         cell_blk = cell_blk.reshape(e2c.shape)
@@ -143,4 +143,8 @@ def input_data(self, grid):
             p_distv_bary_1=p_distv_bary_1,
             p_distv_bary_2=p_distv_bary_2,
             p_dthalf=p_dthalf,
+            horizontal_start=int32(0),
+            horizontal_end=int32(grid.num_edges),
+            vertical_start=int32(0),
+            vertical_end=int32(grid.num_levels),
         )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_global_to_vn.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype
+from gt4py.next.ffront.fbuiltins import Field, astype, int32
 
 from icon4py.model.common.dimension import EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -44,6 +44,10 @@ def apply_nabla2_and_nabla4_global_to_vn(
     z_nabla4_e2: Field[[EdgeDim, KDim], vpfloat],
     diff_multfac_vn: Field[[KDim], wpfloat],
     vn: Field[[EdgeDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _apply_nabla2_and_nabla4_global_to_vn(
         area_edge,
@@ -53,4 +57,8 @@ def apply_nabla2_and_nabla4_global_to_vn(
         diff_multfac_vn,
         vn,
         out=vn,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_and_nabla4_to_vn.py

@@ -12,7 +12,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype, maximum
+from gt4py.next.ffront.fbuiltins import Field, astype, int32, maximum
 
 from icon4py.model.common.dimension import EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -50,6 +50,10 @@ def apply_nabla2_and_nabla4_to_vn(
     nudgecoeff_e: Field[[EdgeDim], wpfloat],
     vn: Field[[EdgeDim, KDim], wpfloat],
     nudgezone_diff: vpfloat,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _apply_nabla2_and_nabla4_to_vn(
         area_edge,
@@ -61,4 +65,8 @@ def apply_nabla2_and_nabla4_to_vn(
         vn,
         nudgezone_diff,
         out=vn,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_nabla2_to_vn_in_lateral_boundary.py

@@ -12,7 +12,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field
+from gt4py.next.ffront.fbuiltins import Field, int32
 
 from icon4py.model.common.dimension import EdgeDim, KDim
 from icon4py.model.common.type_alias import wpfloat
@@ -35,5 +35,19 @@ def apply_nabla2_to_vn_in_lateral_boundary(
     area_edge: Field[[EdgeDim], wpfloat],
     vn: Field[[EdgeDim, KDim], wpfloat],
     fac_bdydiff_v: wpfloat,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _apply_nabla2_to_vn_in_lateral_boundary(z_nabla2_e, area_edge, vn, fac_bdydiff_v, out=vn)
+    _apply_nabla2_to_vn_in_lateral_boundary(
+        z_nabla2_e,
+        area_edge,
+        vn,
+        fac_bdydiff_v,
+        out=vn,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_for_z.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype
+from gt4py.next.ffront.fbuiltins import Field, astype, int32
 
 from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -37,10 +37,18 @@ def calculate_nabla2_for_z(
     inv_dual_edge_length: Field[[EdgeDim], wpfloat],
     theta_v: Field[[CellDim, KDim], wpfloat],
     z_nabla2_e: Field[[EdgeDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _calculate_nabla2_for_z(
         kh_smag_e,
         inv_dual_edge_length,
         theta_v,
         out=z_nabla2_e,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla2_of_theta.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype, neighbor_sum
+from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum
 
 from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -33,5 +33,17 @@ def calculate_nabla2_of_theta(
     z_nabla2_e: Field[[EdgeDim, KDim], wpfloat],
     geofac_div: Field[[CEDim], wpfloat],
     z_temp: Field[[CellDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _calculate_nabla2_of_theta(z_nabla2_e, geofac_div, out=z_temp)
+    _calculate_nabla2_of_theta(
+        z_nabla2_e,
+        geofac_div,
+        out=z_temp,
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/calculate_nabla4.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype
+from gt4py.next.ffront.fbuiltins import Field, astype, int32
 
 from icon4py.model.common.dimension import E2C2V, E2ECV, ECVDim, EdgeDim, KDim, VertexDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -72,6 +72,10 @@ def calculate_nabla4(
     inv_vert_vert_length: Field[[EdgeDim], wpfloat],
     inv_primal_edge_length: Field[[EdgeDim], wpfloat],
     z_nabla4_e2: Field[[EdgeDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _calculate_nabla4(
         u_vert,
@@ -82,4 +86,8 @@ def calculate_nabla4(
         inv_vert_vert_length,
         inv_primal_edge_length,
         out=z_nabla4_e2,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/enhance_diffusion_coefficient_for_grid_point_cold_pools.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, max_over, maximum
+from gt4py.next.ffront.fbuiltins import Field, int32, max_over, maximum
 
 from icon4py.model.common.dimension import E2C, CellDim, E2CDim, EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat
@@ -32,5 +32,17 @@ def _enhance_diffusion_coefficient_for_grid_point_cold_pools(
 def enhance_diffusion_coefficient_for_grid_point_cold_pools(
     kh_smag_e: Field[[EdgeDim, KDim], vpfloat],
     enh_diffu_3d: Field[[CellDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _enhance_diffusion_coefficient_for_grid_point_cold_pools(kh_smag_e, enh_diffu_3d, out=kh_smag_e)
+    _enhance_diffusion_coefficient_for_grid_point_cold_pools(
+        kh_smag_e,
+        enh_diffu_3d,
+        out=kh_smag_e,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_field_for_grid_point_cold_pools_enhancement.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype, neighbor_sum, where
+from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum, where
 
 from icon4py.model.common.dimension import C2E2C, C2E2CDim, CellDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -48,7 +48,19 @@ def temporary_field_for_grid_point_cold_pools_enhancement(
     enh_diffu_3d: Field[[CellDim, KDim], vpfloat],
     thresh_tdiff: wpfloat,
     smallest_vpfloat: vpfloat,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _temporary_field_for_grid_point_cold_pools_enhancement(
-        theta_v, theta_ref_mc, thresh_tdiff, smallest_vpfloat, out=enh_diffu_3d
+        theta_v,
+        theta_ref_mc,
+        thresh_tdiff,
+        smallest_vpfloat,
+        out=enh_diffu_3d,
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/temporary_fields_for_turbulence_diagnostics.py

@@ -13,7 +13,7 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype, neighbor_sum
+from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum
 
 from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
@@ -43,7 +43,20 @@ def temporary_fields_for_turbulence_diagnostics(
     diff_multfac_smag: Field[[KDim], vpfloat],
     kh_c: Field[[CellDim, KDim], vpfloat],
     div: Field[[CellDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
     _temporary_fields_for_turbulence_diagnostics(
-        kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, out=(kh_c, div)
+        kh_smag_ec,
+        vn,
+        e_bln_c_s,
+        geofac_div,
+        diff_multfac_smag,
+        out=(kh_c, div),
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
     )

model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import pytest
+from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import (
     apply_nabla2_and_nabla4_global_to_vn,
@@ -42,11 +43,19 @@ def input_data(self, grid):
             z_nabla4_e2=z_nabla4_e2,
             diff_multfac_vn=diff_multfac_vn,
             vn=vn,
+            horizontal_start=int32(0),
+            horizontal_end=int32(grid.num_edges),
+            vertical_start=int32(0),
+            vertical_end=int32(grid.num_levels),
         )
 
     @staticmethod
-    def reference(grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn):
+    def reference(
+        grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn, **kwargs
+    ):
         area_edge = np.expand_dims(area_edge, axis=-1)
         diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0)
         vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge)
-        return dict(vn=vn)
+        return dict(
+            vn=vn,
+        )

model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import pytest
+from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import (
     apply_nabla2_and_nabla4_global_to_vn,
@@ -49,6 +50,10 @@ def input_data(self, grid):
             nudgecoeff_e=nudgecoeff_e,
             vn=vn,
             nudgezone_diff=nudgezone_diff,
+            horizontal_start=int32(0),
+            horizontal_end=int32(grid.num_edges),
+            vertical_start=int32(0),
+            vertical_end=int32(grid.num_levels),
         )
 
     @staticmethod
@@ -62,6 +67,7 @@ def reference(
         nudgecoeff_e,
         vn,
         nudgezone_diff,
+        **kwargs,
     ):
         area_edge = np.expand_dims(area_edge, axis=-1)
         diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0)
@@ -94,11 +100,19 @@ def input_data(self, grid):
             z_nabla4_e2=z_nabla4_e2,
             diff_multfac_vn=diff_multfac_vn,
             vn=vn,
+            horizontal_start=int32(0),
+            horizontal_end=int32(grid.num_edges),
+            vertical_start=int32(0),
+            vertical_end=int32(grid.num_levels),
         )
 
     @staticmethod
-    def reference(grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn):
+    def reference(
+        grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn, **kwargs
+    ):
         area_edge = np.expand_dims(area_edge, axis=-1)
         diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0)
         vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge)
-        return dict(vn=vn)
+        return dict(
+            vn=vn,
+        )