Merge remote-tracking branch 'origin/main' into ci-gh200

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

@@ -14,19 +14,19 @@
 import enum
 from typing import Final, Optional
 
-from icon4py.model.common import field_type_aliases as fa
-import gt4py.next as gtx
 
+import gt4py.next as gtx
 
 from icon4py.model.atmosphere.diffusion import diffusion_utils, diffusion_states, cached
-
-
-from icon4py.model.common import constants
+from icon4py.model.common import field_type_aliases as fa, constants, dimension as dims
 from icon4py.model.common.decomposition import definitions as decomposition
-from icon4py.model.common import dimension as dims
-from icon4py.model.common.grid import horizontal as h_grid, vertical as v_grid, icon as icon_grid
+from icon4py.model.common.grid import (
+    horizontal as h_grid,
+    vertical as v_grid,
+    icon as icon_grid,
+    geometry,
+)
 from icon4py.model.common.utils import gt4py_field_allocation as field_alloc
-
 import icon4py.model.common.states.prognostic_state as prognostics
 from icon4py.model.common.settings import xp
 
@@ -337,9 +337,9 @@ def __init__(
         self.fac_bdydiff_v: Optional[float] = None
         self.bdy_diff: Optional[float] = None
         self.nudgezone_diff: Optional[float] = None
-        self.edge_params: Optional[h_grid.EdgeParams] = None
-        self.cell_params: Optional[h_grid.CellParams] = None
-        self._horizontal_start_index_w_diffusion: gtx.int32 = 0
+        self.edge_params: Optional[geometry.EdgeParams] = None
+        self.cell_params: Optional[geometry.CellParams] = None
+        self._horizontal_start_index_w_diffusion: gtx.int32 = gtx.int32(0)
 
     def init(
         self,
@@ -349,8 +349,8 @@ def init(
         vertical_grid: v_grid.VerticalGrid,
         metric_state: diffusion_states.DiffusionMetricState,
         interpolation_state: diffusion_states.DiffusionInterpolationState,
-        edge_params: h_grid.EdgeParams,
-        cell_params: h_grid.CellParams,
+        edge_params: geometry.EdgeParams,
+        cell_params: geometry.CellParams,
     ):
         """
         Initialize Diffusion granule with configuration.
@@ -379,13 +379,11 @@ def init(
         self._allocate_temporary_fields()
 
         def _get_start_index_for_w_diffusion() -> gtx.int32:
-            return self.grid.get_start_index(
-                dims.CellDim,
-                (
-                    h_grid.HorizontalMarkerIndex.nudging(dims.CellDim)
-                    if self.grid.limited_area
-                    else h_grid.HorizontalMarkerIndex.interior(dims.CellDim)
-                ),
+            cell_domain = h_grid.domain(dims.CellDim)
+            return (
+                self.grid.start_index(cell_domain(h_grid.Zone.NUDGING))
+                if self.grid.limited_area
+                else self.grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4))
             )
 
         self.nudgezone_diff: float = 0.04 / (params.scaled_nudge_max_coeff + sys.float_info.epsilon)
@@ -541,45 +539,30 @@ def _do_diffusion_step(
             smag_offset:
 
         """
-        klevels = self.grid.num_levels
-        cell_start_interior = self.grid.get_start_index(
-            dims.CellDim, h_grid.HorizontalMarkerIndex.interior(dims.CellDim)
-        )
-        cell_start_nudging = self.grid.get_start_index(
-            dims.CellDim, h_grid.HorizontalMarkerIndex.nudging(dims.CellDim)
-        )
-        cell_end_local = self.grid.get_end_index(
-            dims.CellDim, h_grid.HorizontalMarkerIndex.local(dims.CellDim)
-        )
-        cell_end_halo = self.grid.get_end_index(
-            dims.CellDim, h_grid.HorizontalMarkerIndex.halo(dims.CellDim)
-        )
+        num_levels = self.grid.num_levels
+        cell_domain = h_grid.domain(dims.CellDim)
+        cell_start_interior = self.grid.start_index(cell_domain(h_grid.Zone.INTERIOR))
 
-        edge_start_nudging_plus_one = self.grid.get_start_index(
-            dims.EdgeDim, h_grid.HorizontalMarkerIndex.nudging(dims.EdgeDim) + 1
-        )
-        edge_start_nudging = self.grid.get_start_index(
-            dims.EdgeDim, h_grid.HorizontalMarkerIndex.nudging(dims.EdgeDim)
-        )
-        edge_start_lb_plus4 = self.grid.get_start_index(
-            dims.EdgeDim, h_grid.HorizontalMarkerIndex.lateral_boundary(dims.EdgeDim) + 4
-        )
-        edge_end_local = self.grid.get_end_index(
-            dims.EdgeDim, h_grid.HorizontalMarkerIndex.local(dims.EdgeDim)
-        )
-        edge_end_local_minus2 = self.grid.get_end_index(
-            dims.EdgeDim, h_grid.HorizontalMarkerIndex.local(dims.EdgeDim) - 2
-        )
-        edge_end_halo = self.grid.get_end_index(
-            dims.EdgeDim, h_grid.HorizontalMarkerIndex.halo(dims.EdgeDim)
-        )
+        cell_start_nudging = self.grid.start_index(cell_domain(h_grid.Zone.NUDGING))
+        cell_end_local = self.grid.end_index(cell_domain(h_grid.Zone.LOCAL))
+        cell_end_halo = self.grid.end_index(cell_domain(h_grid.Zone.HALO))
 
-        vertex_start_lb_plus1 = self.grid.get_start_index(
-            dims.VertexDim, h_grid.HorizontalMarkerIndex.lateral_boundary(dims.VertexDim) + 1
+        edge_domain = h_grid.domain(dims.EdgeDim)
+        edge_start_nudging_level_2 = self.grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2))
+        edge_start_nudging = self.grid.start_index(edge_domain(h_grid.Zone.NUDGING))
+
+        edge_start_lateral_boundary_level_5 = self.grid.start_index(
+            edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)
         )
-        vertex_end_local = self.grid.get_end_index(
-            dims.VertexDim, h_grid.HorizontalMarkerIndex.local(dims.VertexDim)
+        edge_end_local = self.grid.end_index(edge_domain(h_grid.Zone.LOCAL))
+        edge_end_halo_level_2 = self.grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2))
+        edge_end_halo = self.grid.end_index(edge_domain(h_grid.Zone.HALO))
+
+        vertex_domain = h_grid.domain(dims.VertexDim)
+        vertex_start_lateral_boundary_level_2 = self.grid.start_index(
+            vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)
         )
+        vertex_end_local = self.grid.end_index(vertex_domain(h_grid.Zone.LOCAL))
 
         # dtime dependent: enh_smag_factor,
         cached.scale_k(self.enh_smag_fac, dtime, self.diff_multfac_smag, offset_provider={})
@@ -591,10 +574,10 @@ def _do_diffusion_step(
             ptr_coeff_2=self.interpolation_state.rbf_coeff_2,
             p_u_out=self.u_vert,
             p_v_out=self.v_vert,
-            horizontal_start=vertex_start_lb_plus1,
+            horizontal_start=vertex_start_lateral_boundary_level_2,
             horizontal_end=vertex_end_local,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug("rbf interpolation 1: end")
@@ -622,10 +605,10 @@ def _do_diffusion_step(
             kh_smag_ec=self.kh_smag_ec,
             z_nabla2_e=self.z_nabla2_e,
             smag_offset=smag_offset,
-            horizontal_start=edge_start_lb_plus4,
-            horizontal_end=edge_end_local_minus2,
+            horizontal_start=edge_start_lateral_boundary_level_5,
+            horizontal_end=edge_end_halo_level_2,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end")
@@ -648,7 +631,7 @@ def _do_diffusion_step(
                 horizontal_start=cell_start_nudging,
                 horizontal_end=cell_end_local,
                 vertical_start=1,
-                vertical_end=klevels,
+                vertical_end=num_levels,
                 offset_provider=self.grid.offset_providers,
             )
             log.debug(
@@ -669,10 +652,10 @@ def _do_diffusion_step(
             ptr_coeff_2=self.interpolation_state.rbf_coeff_2,
             p_u_out=self.u_vert,
             p_v_out=self.v_vert,
-            horizontal_start=vertex_start_lb_plus1,
+            horizontal_start=vertex_start_lateral_boundary_level_2,
             horizontal_end=vertex_end_local,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug("2nd rbf interpolation: end")
@@ -699,12 +682,12 @@ def _do_diffusion_step(
             edge=self.horizontal_edge_index,
             nudgezone_diff=self.nudgezone_diff,
             fac_bdydiff_v=self.fac_bdydiff_v,
-            start_2nd_nudge_line_idx_e=gtx.int32(edge_start_nudging_plus_one),
+            start_2nd_nudge_line_idx_e=gtx.int32(edge_start_nudging_level_2),
             limited_area=self.grid.limited_area,
-            horizontal_start=edge_start_lb_plus4,
+            horizontal_start=edge_start_lateral_boundary_level_5,
             horizontal_end=edge_end_local,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug("running stencils 04 05 06 (apply_diffusion_to_vn): end")
@@ -739,7 +722,7 @@ def _do_diffusion_step(
             horizontal_start=self._horizontal_start_index_w_diffusion,
             horizontal_end=cell_end_halo,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug(
@@ -758,8 +741,8 @@ def _do_diffusion_step(
             kh_smag_e=self.kh_smag_e,
             horizontal_start=edge_start_nudging,
             horizontal_end=edge_end_halo,
-            vertical_start=(klevels - 2),
-            vertical_end=klevels,
+            vertical_start=(num_levels - 2),
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug(
@@ -776,7 +759,7 @@ def _do_diffusion_step(
             horizontal_start=cell_start_nudging,
             horizontal_end=cell_end_local,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider=self.grid.offset_providers,
         )
         log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end")
@@ -796,7 +779,7 @@ def _do_diffusion_step(
                 horizontal_start=cell_start_nudging,
                 horizontal_end=cell_end_local,
                 vertical_start=0,
-                vertical_end=klevels,
+                vertical_end=num_levels,
                 offset_provider=self.grid.offset_providers,
             )
 
@@ -813,7 +796,7 @@ def _do_diffusion_step(
             horizontal_start=cell_start_nudging,
             horizontal_end=cell_end_local,
             vertical_start=0,
-            vertical_end=klevels,
+            vertical_end=num_levels,
             offset_provider={},
         )
         log.debug("running stencil 16 (update_theta_and_exner): end")

model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py

@@ -38,10 +38,12 @@ def test_parallel_diffusion(
     model_top_height,
     stretch_factor,
     damping_height,
+    caplog,
 ):
+    caplog.set_level("INFO")
     parallel_helpers.check_comm_size(processor_props)
     print(
-        f"rank={processor_props.rank}/{processor_props.comm_size}: inializing diffusion for experiment '{datatest_utils.REGIONAL_EXPERIMENT}'"
+        f"rank={processor_props.rank}/{processor_props.comm_size}: initializing diffusion for experiment '{datatest_utils.REGIONAL_EXPERIMENT}'"
     )
     print(
         f"rank={processor_props.rank}/{processor_props.comm_size}: decomposition info : klevels = {decomposition_info.klevels}, "

model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py

@@ -11,7 +11,7 @@
 
 from icon4py.model.atmosphere.diffusion import diffusion, diffusion_utils
 from icon4py.model.common.grid import vertical as v_grid
-from icon4py.model.common.grid.horizontal import CellParams, EdgeParams
+from icon4py.model.common.grid.geometry import CellParams, EdgeParams
 from icon4py.model.common.settings import backend
 from icon4py.model.common.test_utils import (
     datatest_utils as dt_utils,