Merge branch 'main' into diffusion_backend

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py

@@ -5,7 +5,6 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
 import dataclasses
 import functools
 
@@ -38,9 +37,9 @@ class DiffusionMetricState:
         float
     ]  # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c)
     mask_hdiff: fa.CellKField[bool]
-    zd_vertoffset: gtx.Field[[dims.CECDim, dims.KDim], gtx.int32]
+    zd_vertoffset: gtx.Field[gtx.Dims[dims.CECDim, dims.KDim], gtx.int32]
     zd_diffcoef: fa.CellKField[float]
-    zd_intcoef: gtx.Field[[dims.CECDim, dims.KDim], float]
+    zd_intcoef: gtx.Field[gtx.Dims[dims.CECDim, dims.KDim], float]
 
 
 @dataclasses.dataclass(frozen=True)
@@ -57,12 +56,14 @@ class DiffusionInterpolationState:
         [dims.VertexDim, dims.V2EDim], float
     ]  # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v)
 
-    geofac_div: gtx.Field[[dims.CEDim], float]  # factor for divergence (nproma,cell_type,nblks_c)
+    geofac_div: gtx.Field[
+        gtx.Dims[dims.CEDim], float
+    ]  # factor for divergence (nproma,cell_type,nblks_c)
 
     geofac_n2s: gtx.Field[
         [dims.CellDim, dims.C2E2CODim], float
     ]  # factor for nabla2-scalar (nproma,cell_type+1,nblks_c)
-    geofac_grg_x: gtx.Field[[dims.CellDim, dims.C2E2CODim], float]
+    geofac_grg_x: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], float]
     geofac_grg_y: gtx.Field[
         [dims.CellDim, dims.C2E2CODim], float
     ]  # factors for green gauss gradient (nproma,4,nblks_c,2)
@@ -73,7 +74,7 @@ def geofac_n2s_c(self) -> fa.CellField[float]:
         return gtx.as_field((dims.CellDim,), data=self.geofac_n2s.ndarray[:, 0])
 
     @functools.cached_property
-    def geofac_n2s_nbh(self) -> gtx.Field[[dims.CECDim], float]:
+    def geofac_n2s_nbh(self) -> gtx.Field[gtx.Dims[dims.CECDim], float]:
         geofac_nbh_ar = self.geofac_n2s.ndarray[:, 1:]
         old_shape = geofac_nbh_ar.shape
         return gtx.as_field(

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py

@@ -5,7 +5,6 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
 from typing import Tuple
 
 import gt4py.next as gtx

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

@@ -5,10 +5,9 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, int32
 
 from icon4py.model.atmosphere.diffusion.stencils.calculate_nabla2_for_z import (
     _calculate_nabla2_for_z,
@@ -32,13 +31,13 @@ def _apply_diffusion_to_theta_and_exner(
     kh_smag_e: fa.EdgeKField[vpfloat],
     inv_dual_edge_length: fa.EdgeField[wpfloat],
     theta_v_in: fa.CellKField[wpfloat],
-    geofac_div: Field[[dims.CEDim], wpfloat],
+    geofac_div: gtx.Field[gtx.Dims[dims.CEDim], wpfloat],
     mask: fa.CellKField[bool],
-    zd_vertoffset: Field[[dims.CECDim, dims.KDim], int32],
+    zd_vertoffset: gtx.Field[gtx.Dims[dims.CECDim, dims.KDim], gtx.int32],
     zd_diffcoef: fa.CellKField[wpfloat],
     geofac_n2s_c: fa.CellField[wpfloat],
-    geofac_n2s_nbh: Field[[dims.CECDim], wpfloat],
-    vcoef: Field[[dims.CECDim, dims.KDim], wpfloat],
+    geofac_n2s_nbh: gtx.Field[gtx.Dims[dims.CECDim], wpfloat],
+    vcoef: gtx.Field[gtx.Dims[dims.CECDim, dims.KDim], wpfloat],
     area: fa.CellField[wpfloat],
     exner: fa.CellKField[wpfloat],
     rd_o_cvd: vpfloat,
@@ -64,21 +63,21 @@ def apply_diffusion_to_theta_and_exner(
     kh_smag_e: fa.EdgeKField[vpfloat],
     inv_dual_edge_length: fa.EdgeField[wpfloat],
     theta_v_in: fa.CellKField[wpfloat],
-    geofac_div: Field[[dims.CEDim], wpfloat],
+    geofac_div: gtx.Field[gtx.Dims[dims.CEDim], wpfloat],
     mask: fa.CellKField[bool],
-    zd_vertoffset: Field[[dims.CECDim, dims.KDim], int32],
+    zd_vertoffset: gtx.Field[gtx.Dims[dims.CECDim, dims.KDim], gtx.int32],
     zd_diffcoef: fa.CellKField[wpfloat],
     geofac_n2s_c: fa.CellField[wpfloat],
-    geofac_n2s_nbh: Field[[dims.CECDim], wpfloat],
-    vcoef: Field[[dims.CECDim, dims.KDim], wpfloat],
+    geofac_n2s_nbh: gtx.Field[gtx.Dims[dims.CECDim], wpfloat],
+    vcoef: gtx.Field[gtx.Dims[dims.CECDim, dims.KDim], wpfloat],
     area: fa.CellField[wpfloat],
     theta_v: fa.CellKField[wpfloat],
     exner: fa.CellKField[wpfloat],
     rd_o_cvd: vpfloat,
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_diffusion_to_theta_and_exner(
         kh_smag_e,

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

@@ -5,10 +5,10 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, int32, where
+from gt4py.next.ffront.fbuiltins import where
 
 from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import (
     _apply_nabla2_and_nabla4_global_to_vn,
@@ -29,8 +29,8 @@
 def _apply_diffusion_to_vn(
     u_vert: fa.VertexKField[vpfloat],
     v_vert: fa.VertexKField[vpfloat],
-    primal_normal_vert_v1: Field[[dims.ECVDim], wpfloat],
-    primal_normal_vert_v2: Field[[dims.ECVDim], wpfloat],
+    primal_normal_vert_v1: gtx.Field[gtx.Dims[dims.ECVDim], wpfloat],
+    primal_normal_vert_v2: gtx.Field[gtx.Dims[dims.ECVDim], wpfloat],
     z_nabla2_e: fa.EdgeKField[wpfloat],
     inv_vert_vert_length: fa.EdgeField[wpfloat],
     inv_primal_edge_length: fa.EdgeField[wpfloat],
@@ -39,10 +39,10 @@ def _apply_diffusion_to_vn(
     diff_multfac_vn: fa.KField[wpfloat],
     nudgecoeff_e: fa.EdgeField[wpfloat],
     vn: fa.EdgeKField[wpfloat],
-    edge: fa.EdgeField[int32],
+    edge: fa.EdgeField[gtx.int32],
     nudgezone_diff: vpfloat,
     fac_bdydiff_v: wpfloat,
-    start_2nd_nudge_line_idx_e: int32,
+    start_2nd_nudge_line_idx_e: gtx.int32,
     limited_area: bool,
 ) -> fa.EdgeKField[wpfloat]:
     z_nabla4_e2 = _calculate_nabla4(
@@ -93,8 +93,8 @@ def _apply_diffusion_to_vn(
 def apply_diffusion_to_vn(
     u_vert: fa.VertexKField[vpfloat],
     v_vert: fa.VertexKField[vpfloat],
-    primal_normal_vert_v1: Field[[dims.ECVDim], wpfloat],
-    primal_normal_vert_v2: Field[[dims.ECVDim], wpfloat],
+    primal_normal_vert_v1: gtx.Field[gtx.Dims[dims.ECVDim], wpfloat],
+    primal_normal_vert_v2: gtx.Field[gtx.Dims[dims.ECVDim], wpfloat],
     z_nabla2_e: fa.EdgeKField[wpfloat],
     inv_vert_vert_length: fa.EdgeField[wpfloat],
     inv_primal_edge_length: fa.EdgeField[wpfloat],
@@ -103,15 +103,15 @@ def apply_diffusion_to_vn(
     diff_multfac_vn: fa.KField[wpfloat],
     nudgecoeff_e: fa.EdgeField[wpfloat],
     vn: fa.EdgeKField[wpfloat],
-    edge: fa.EdgeField[int32],
+    edge: fa.EdgeField[gtx.int32],
     nudgezone_diff: vpfloat,
     fac_bdydiff_v: wpfloat,
-    start_2nd_nudge_line_idx_e: int32,
+    start_2nd_nudge_line_idx_e: gtx.int32,
     limited_area: bool,
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_diffusion_to_vn(
         u_vert,

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

@@ -5,10 +5,10 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where
+from gt4py.next.ffront.fbuiltins import broadcast, where
 
 from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w import _apply_nabla2_to_w
 from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_to_w_in_upper_damping_layer import (
@@ -29,20 +29,20 @@
 @field_operator
 def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
     area: fa.CellField[wpfloat],
-    geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat],
-    geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat],
-    geofac_grg_y: Field[[CellDim, C2E2CODim], wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[CellDim, C2E2CODim], wpfloat],
+    geofac_grg_x: gtx.Field[gtx.Dims[CellDim, C2E2CODim], wpfloat],
+    geofac_grg_y: gtx.Field[gtx.Dims[CellDim, C2E2CODim], wpfloat],
     w_old: fa.CellKField[wpfloat],
-    type_shear: int32,
+    type_shear: gtx.int32,
     dwdx: fa.CellKField[vpfloat],
     dwdy: fa.CellKField[vpfloat],
     diff_multfac_w: wpfloat,
     diff_multfac_n2w: fa.KField[wpfloat],
-    k: fa.KField[int32],
-    cell: fa.CellField[int32],
-    nrdmax: int32,
-    interior_idx: int32,
-    halo_idx: int32,
+    k: fa.KField[gtx.int32],
+    cell: fa.CellField[gtx.int32],
+    nrdmax: gtx.int32,
+    interior_idx: gtx.int32,
+    halo_idx: gtx.int32,
 ) -> tuple[
     fa.CellKField[wpfloat],
     fa.CellKField[vpfloat],
@@ -79,25 +79,25 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
 @program(grid_type=GridType.UNSTRUCTURED)
 def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
     area: fa.CellField[wpfloat],
-    geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat],
-    geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat],
-    geofac_grg_y: Field[[CellDim, C2E2CODim], wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[CellDim, C2E2CODim], wpfloat],
+    geofac_grg_x: gtx.Field[gtx.Dims[CellDim, C2E2CODim], wpfloat],
+    geofac_grg_y: gtx.Field[gtx.Dims[CellDim, C2E2CODim], wpfloat],
     w_old: fa.CellKField[wpfloat],
     w: fa.CellKField[wpfloat],
-    type_shear: int32,
+    type_shear: gtx.int32,
     dwdx: fa.CellKField[vpfloat],
     dwdy: fa.CellKField[vpfloat],
     diff_multfac_w: wpfloat,
     diff_multfac_n2w: fa.KField[wpfloat],
-    k: fa.KField[int32],
-    cell: fa.CellField[int32],
-    nrdmax: int32,
-    interior_idx: int32,
-    halo_idx: int32,
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    k: fa.KField[gtx.int32],
+    cell: fa.CellField[gtx.int32],
+    nrdmax: gtx.int32,
+    interior_idx: gtx.int32,
+    halo_idx: gtx.int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
         area,

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

@@ -5,10 +5,10 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import astype, int32
+from gt4py.next.ffront.fbuiltins import astype
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.settings import backend
@@ -40,10 +40,10 @@ def apply_nabla2_and_nabla4_global_to_vn(
     z_nabla4_e2: fa.EdgeKField[vpfloat],
     diff_multfac_vn: fa.KField[wpfloat],
     vn: fa.EdgeKField[wpfloat],
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_nabla2_and_nabla4_global_to_vn(
         area_edge,

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

@@ -5,10 +5,10 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import astype, broadcast, int32, maximum
+from gt4py.next.ffront.fbuiltins import astype, broadcast, maximum
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.settings import backend
@@ -48,10 +48,10 @@ def apply_nabla2_and_nabla4_to_vn(
     nudgecoeff_e: fa.EdgeField[wpfloat],
     vn: fa.EdgeKField[wpfloat],
     nudgezone_diff: vpfloat,
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_nabla2_and_nabla4_to_vn(
         area_edge,

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

@@ -5,10 +5,9 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.settings import backend
@@ -32,10 +31,10 @@ def apply_nabla2_to_vn_in_lateral_boundary(
     area_edge: fa.EdgeField[wpfloat],
     vn: fa.EdgeKField[wpfloat],
     fac_bdydiff_v: wpfloat,
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_nabla2_to_vn_in_lateral_boundary(
         z_nabla2_e,

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

@@ -5,10 +5,10 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum
+from gt4py.next.ffront.fbuiltins import astype, neighbor_sum
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa
 from icon4py.model.common.dimension import C2E2CO, C2E2CODim
@@ -20,7 +20,7 @@
 def _apply_nabla2_to_w(
     area: fa.CellField[wpfloat],
     z_nabla2_c: fa.CellKField[vpfloat],
-    geofac_n2s: Field[[dims.CellDim, C2E2CODim], wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat],
     w: fa.CellKField[wpfloat],
     diff_multfac_w: wpfloat,
 ) -> fa.CellKField[wpfloat]:
@@ -36,13 +36,13 @@ def _apply_nabla2_to_w(
 def apply_nabla2_to_w(
     area: fa.CellField[wpfloat],
     z_nabla2_c: fa.CellKField[vpfloat],
-    geofac_n2s: Field[[dims.CellDim, C2E2CODim], wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat],
     w: fa.CellKField[wpfloat],
     diff_multfac_w: wpfloat,
-    horizontal_start: int32,
-    horizontal_end: int32,
-    vertical_start: int32,
-    vertical_end: int32,
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
 ):
     _apply_nabla2_to_w(
         area,