Metrics fields factory

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

@@ -868,7 +868,6 @@ def time_step(
             f"running timestep: dtime = {dtime}, initial_timestep = {at_initial_timestep}, first_substep = {at_first_substep}, last_substep = {at_last_substep}, prep_adv = {lprep_adv}"
         )
 
-        # # TODO: abishekg7 move this to tests
         if self.p_test_run:
             self._init_test_fields(
                 self.intermediate_fields.z_rho_e,

model/common/src/icon4py/model/common/field_type_aliases.py

@@ -20,6 +20,7 @@
 EdgeField: TypeAlias = Field[Dims[dims.EdgeDim], T]
 VertexField: TypeAlias = Field[Dims[dims.VertexDim], T]
 KField: TypeAlias = Field[Dims[dims.KDim], T]
+KHalfField: TypeAlias = Field[Dims[dims.KHalfDim], T]
 
 CellKField: TypeAlias = Field[Dims[dims.CellDim, dims.KDim], T]
 EdgeKField: TypeAlias = Field[Dims[dims.EdgeDim, dims.KDim], T]

model/common/src/icon4py/model/common/grid/geometry.py

@@ -132,6 +132,7 @@ def __init__(
             {
                 # TODO (@magdalena) rescaled by grid_length_rescale_factor (mo_grid_tools.f90)
                 attrs.EDGE_CELL_DISTANCE: extra_fields[gm.GeometryName.EDGE_CELL_DISTANCE],
+                attrs.EDGE_VERTEX_DISTANCE: extra_fields[gm.GeometryName.EDGE_VERTEX_DISTANCE],
                 attrs.CELL_AREA: extra_fields[gm.GeometryName.CELL_AREA],
                 attrs.DUAL_AREA: extra_fields[gm.GeometryName.DUAL_AREA],
                 attrs.TANGENT_ORIENTATION: extra_fields[gm.GeometryName.TANGENT_ORIENTATION],

model/common/src/icon4py/model/common/grid/geometry_attributes.py

@@ -31,6 +31,7 @@
 EDGE_AREA: Final[str] = "edge_area"
 DUAL_AREA: Final[str] = "dual_area"
 EDGE_CELL_DISTANCE: Final[str] = "edge_midpoint_to_cell_center_distance"
+EDGE_VERTEX_DISTANCE: Final[str] = "edge_midpoint_to_vertex_distance"
 TANGENT_ORIENTATION: Final[str] = "edge_orientation"
 CELL_NORMAL_ORIENTATION: Final[str] = "orientation_of_normal_to_cell_edges"
 VERTEX_EDGE_ORIENTATION: Final[str] = "orientation_of_edges_around_vertex"
@@ -126,6 +127,14 @@
         icon_var_name="t_grid_edges%edge_cell_length",
         dtype=ta.wpfloat,
     ),
+    EDGE_VERTEX_DISTANCE: dict(
+        standard_name=EDGE_VERTEX_DISTANCE,
+        long_name="distances between edge midpoint and adjacent vertices",
+        units="m",
+        dims=(dims.EdgeDim, dims.E2VDim),
+        icon_var_name="t_grid_edges%edge_vert_length",
+        dtype=ta.wpfloat,
+    ),
     DUAL_EDGE_LENGTH: dict(
         standard_name=DUAL_EDGE_LENGTH,
         long_name="length of the dual edge",

model/common/src/icon4py/model/common/grid/grid_manager.py

@@ -193,6 +193,7 @@ class GeometryName(FieldName):
     EDGE_ORIENTATION_ON_VERTEX = "edge_orientation"
     # TODO (@halungge) compute from coordinates
     EDGE_CELL_DISTANCE = "edge_cell_distance"
+    EDGE_VERTEX_DISTANCE = "edge_vert_distance"
 
 
 class CoordinateName(FieldName):
@@ -467,6 +468,10 @@ def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]):
                 self._reader.variable(GeometryName.EDGE_CELL_DISTANCE, transpose=True),
             ),
             # TODO (@halungge) recompute from coordinates? field in gridfile contains NaN on boundary edges
+            GeometryName.EDGE_VERTEX_DISTANCE.value: gtx.as_field(
+                (dims.EdgeDim, dims.E2VDim),
+                self._reader.variable(GeometryName.EDGE_VERTEX_DISTANCE, transpose=True),
+            ),
             GeometryName.TANGENT_ORIENTATION.value: gtx.as_field(
                 (dims.EdgeDim,),
                 self._reader.variable(GeometryName.TANGENT_ORIENTATION),

model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py

@@ -20,6 +20,7 @@
 GEOFAC_GRDIV: Final[str] = "geometrical_factor_for_gradient_of_divergence"
 GEOFAC_GRG_X: Final[str] = "geometrical_factor_for_green_gauss_gradient_x"
 GEOFAC_GRG_Y: Final[str] = "geometrical_factor_for_green_gauss_gradient_y"
+CELL_AW_VERTS: Final[str] = "geometrical_factor_for_cells_aw_verts"
 
 attrs: dict[str, model.FieldMetaData] = {
     C_LIN_E: dict(
@@ -86,4 +87,12 @@
         icon_var_name="geofac_grg",
         dtype=ta.wpfloat,
     ),
+    CELL_AW_VERTS: dict(
+        standard_name=CELL_AW_VERTS,
+        long_name="geometrical factor for cells_aw_verts",
+        units="",
+        dims=(dims.VertexDim, dims.V2CDim),
+        icon_var_name="cells_aw_verts",
+        dtype=ta.wpfloat,
+    ),
 }

model/common/src/icon4py/model/common/interpolation/interpolation_factory.py

@@ -28,6 +28,7 @@
 
 cell_domain = h_grid.domain(dims.CellDim)
 edge_domain = h_grid.domain(dims.EdgeDim)
+vertex_domain = h_grid.domain(dims.VertexDim)
 
 
 class InterpolationFieldsFactory(factory.FieldSource, factory.GridProvider):
@@ -181,9 +182,31 @@ def _register_computed_fields(self):
                 )
             },
         )
-
         self.register_provider(geofac_grg)
 
+        cells_aw_verts = factory.NumpyFieldsProvider(
+            func=functools.partial(interpolation_fields.compute_cells_aw_verts),
+            fields=(attrs.CELL_AW_VERTS,),
+            domain=(dims.VertexDim, dims.V2CDim),
+            deps={
+                "dual_area": geometry_attrs.DUAL_AREA,
+                "edge_vert_length": geometry_attrs.EDGE_VERTEX_DISTANCE,
+                "edge_cell_length": geometry_attrs.EDGE_CELL_DISTANCE,
+            },
+            connectivities={
+                "v2e": dims.V2EDim,
+                "e2v": dims.E2VDim,
+                "v2c": dims.V2CDim,
+                "e2c": dims.E2CDim,
+            },
+            params={
+                "horizontal_start": self.grid.start_index(
+                    vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)
+                )
+            },
+        )
+        self.register_provider(cells_aw_verts)
+
     @property
     def metadata(self) -> dict[str, model.FieldMetaData]:
         return self._attrs

model/common/src/icon4py/model/common/interpolation/interpolation_fields.py

@@ -832,7 +832,7 @@ def compute_cells_aw_verts(
 
     Args:
         dual_area: numpy array, representing a gtx.Field[gtx.Dims[VertexDim], ta.wpfloat]
-        edge_vert_length: \\ numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat]
+        edge_vert_length: \\ numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2VDim], ta.wpfloat]
         edge_cell_length: //
         owner_mask: numpy array, representing a gtx.Field[gtx.Dims[VertexDim], bool]
         v2e: numpy array, representing a gtx.Field[gtx.Dims[VertexDim, V2EDim], gtx.int32]
@@ -849,14 +849,18 @@ def compute_cells_aw_verts(
         cells_aw_verts[jv, :] = 0.0
         for je in range(v2e.shape[1]):
             # INVALID_INDEX
-            if je > gm.GridFile.INVALID_INDEX and (je > 0 and v2e[jv, je] == v2e[jv, je - 1]):
+            if v2e[jv, je] == gm.GridFile.INVALID_INDEX or (
+                je > 0 and v2e[jv, je] == v2e[jv, je - 1]
+            ):
                 continue
             ile = v2e[jv, je]
             idx_ve = 0 if e2v[ile, 0] == jv else 1
             cell_offset_idx_0 = e2c[ile, 0]
             cell_offset_idx_1 = e2c[ile, 1]
             for jc in range(v2e.shape[1]):
-                if jc > gm.GridFile.INVALID_INDEX and (jc > 0 and v2c[jv, jc] == v2c[jv, jc - 1]):
+                if v2c[jv, jc] == gm.GridFile.INVALID_INDEX or (
+                    jc > 0 and v2c[jv, jc] == v2c[jv, jc - 1]
+                ):
                     continue
                 if cell_offset_idx_0 == v2c[jv, jc]:
                     cells_aw_verts[jv, jc] = (

model/common/src/icon4py/model/common/metrics/compute_coeff_gradekin.py

@@ -13,11 +13,11 @@
 
 
 def compute_coeff_gradekin(
-    edge_cell_length: np.array,
-    inv_dual_edge_length: np.array,
-    horizontal_start: float,
-    horizontal_end: float,
-) -> np.array:
+    edge_cell_length: data_alloc.NDArray,
+    inv_dual_edge_length: data_alloc.NDArray,
+    horizontal_start: int,
+    horizontal_end: int,
+):
     """
     Compute coefficients for improved calculation of kinetic energy gradient
 

model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py

@@ -11,6 +11,15 @@
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
+def compute_max_nbhgt_np(
+    c2e2c: data_alloc.NDArray, z_mc: data_alloc.NDArray, nlev: int
+) -> data_alloc.NDArray:
+    z_mc_nlev = z_mc[:, nlev - 1]
+    max_nbhgt_0_1 = np.maximum(z_mc_nlev[c2e2c[:, 0]], z_mc_nlev[c2e2c[:, 1]])
+    max_nbhgt = np.maximum(max_nbhgt_0_1, z_mc_nlev[c2e2c[:, 2]])
+    return max_nbhgt
+
+
 def _compute_nbidx(
     k_range: range,
     z_mc: data_alloc.NDArray,
@@ -62,8 +71,8 @@ def _compute_z_vintcoeff(
 def _compute_ls_params(
     k_start: list,
     k_end: list,
-    z_maxslp_avg: data_alloc.NDArray,
-    z_maxhgtd_avg: data_alloc.NDArray,
+    maxslp_avg: data_alloc.NDArray,
+    maxhgtd_avg: data_alloc.NDArray,
     c_owner_mask: data_alloc.NDArray,
     thslp_zdiffu: float,
     thhgtd_zdiffu: float,
@@ -78,8 +87,7 @@ def _compute_ls_params(
 
     for jc in range(cell_nudging, n_cells):
         if (
-            z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu
-            or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu
+            maxslp_avg[jc, nlev - 1] >= thslp_zdiffu or maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu
         ) and c_owner_mask[jc]:
             ji += 1
             indlist[ji] = jc
@@ -96,8 +104,8 @@ def _compute_ls_params(
 def _compute_k_start_end(
     z_mc: data_alloc.NDArray,
     max_nbhgt: data_alloc.NDArray,
-    z_maxslp_avg: data_alloc.NDArray,
-    z_maxhgtd_avg: data_alloc.NDArray,
+    maxslp_avg: data_alloc.NDArray,
+    maxhgtd_avg: data_alloc.NDArray,
     c_owner_mask: data_alloc.NDArray,
     thslp_zdiffu: float,
     thhgtd_zdiffu: float,
@@ -109,16 +117,15 @@ def _compute_k_start_end(
     k_end = [None] * n_cells
     for jc in range(cell_nudging, n_cells):
         if (
-            z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu
-            or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu
+            maxslp_avg[jc, nlev - 1] >= thslp_zdiffu or maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu
         ) and c_owner_mask[jc]:
             for jk in reversed(range(nlev)):
                 if z_mc[jc, jk] >= max_nbhgt[jc]:
                     k_end[jc] = jk + 1
                     break
 
             for jk in range(nlev):
-                if z_maxslp_avg[jc, jk] >= thslp_zdiffu or z_maxhgtd_avg[jc, jk] >= thhgtd_zdiffu:
+                if maxslp_avg[jc, jk] >= thslp_zdiffu or maxhgtd_avg[jc, jk] >= thhgtd_zdiffu:
                     k_start[jc] = jk
                     break
 
@@ -129,29 +136,31 @@ def _compute_k_start_end(
 
 
 def compute_diffusion_metrics(
+    c2e2c: data_alloc.NDArray,
     z_mc: data_alloc.NDArray,
-    z_mc_off: data_alloc.NDArray,
     max_nbhgt: data_alloc.NDArray,
     c_owner_mask: data_alloc.NDArray,
-    nbidx: data_alloc.NDArray,
-    z_vintcoeff: data_alloc.NDArray,
-    z_maxslp_avg: data_alloc.NDArray,
-    z_maxhgtd_avg: data_alloc.NDArray,
-    mask_hdiff: data_alloc.NDArray,
-    zd_diffcoef_dsl: data_alloc.NDArray,
-    zd_intcoef_dsl: data_alloc.NDArray,
-    zd_vertoffset_dsl: data_alloc.NDArray,
+    maxslp_avg: data_alloc.NDArray,
+    maxhgtd_avg: data_alloc.NDArray,
     thslp_zdiffu: float,
     thhgtd_zdiffu: float,
+    n_c2e2c: int,
     cell_nudging: int,
     n_cells: int,
     nlev: int,
 ) -> tuple[data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray]:
+    z_mc_off = z_mc[c2e2c]
+    nbidx = np.ones(shape=(n_cells, n_c2e2c, nlev), dtype=int)
+    z_vintcoeff = np.zeros(shape=(n_cells, n_c2e2c, nlev))
+    mask_hdiff = np.zeros(shape=(n_cells, nlev), dtype=bool)
+    zd_vertoffset_dsl = np.zeros(shape=(n_cells, n_c2e2c, nlev))
+    zd_intcoef_dsl = np.zeros(shape=(n_cells, n_c2e2c, nlev))
+    zd_diffcoef_dsl = np.zeros(shape=(n_cells, nlev))
     k_start, k_end = _compute_k_start_end(
         z_mc=z_mc,
         max_nbhgt=max_nbhgt,
-        z_maxslp_avg=z_maxslp_avg,
-        z_maxhgtd_avg=z_maxhgtd_avg,
+        maxslp_avg=maxslp_avg,
+        maxhgtd_avg=maxhgtd_avg,
         c_owner_mask=c_owner_mask,
         thslp_zdiffu=thslp_zdiffu,
         thhgtd_zdiffu=thhgtd_zdiffu,
@@ -163,8 +172,8 @@ def compute_diffusion_metrics(
     indlist, listreduce, ji = _compute_ls_params(
         k_start=k_start,
         k_end=k_end,
-        z_maxslp_avg=z_maxslp_avg,
-        z_maxhgtd_avg=z_maxhgtd_avg,
+        maxslp_avg=maxslp_avg,
+        maxhgtd_avg=maxhgtd_avg,
         c_owner_mask=c_owner_mask,
         thslp_zdiffu=thslp_zdiffu,
         thhgtd_zdiffu=thhgtd_zdiffu,
@@ -191,10 +200,18 @@ def compute_diffusion_metrics(
             zd_diffcoef_dsl_var = np.maximum(
                 0.0,
                 np.maximum(
-                    np.sqrt(np.maximum(0.0, z_maxslp_avg[jc, k_range] - thslp_zdiffu)) / 250.0,
-                    2.0e-4 * np.sqrt(np.maximum(0.0, z_maxhgtd_avg[jc, k_range] - thhgtd_zdiffu)),
+                    np.sqrt(np.maximum(0.0, maxslp_avg[jc, k_range] - thslp_zdiffu)) / 250.0,
+                    2.0e-4 * np.sqrt(np.maximum(0.0, maxhgtd_avg[jc, k_range] - thhgtd_zdiffu)),
                 ),
             )
             zd_diffcoef_dsl[jc, k_range] = np.minimum(0.002, zd_diffcoef_dsl_var)
 
+    # flatten first two dims:
+    zd_intcoef_dsl = zd_intcoef_dsl.reshape(
+        (zd_intcoef_dsl.shape[0] * zd_intcoef_dsl.shape[1],) + zd_intcoef_dsl.shape[2:]
+    )
+    zd_vertoffset_dsl = zd_vertoffset_dsl.reshape(
+        (zd_vertoffset_dsl.shape[0] * zd_vertoffset_dsl.shape[1],) + zd_vertoffset_dsl.shape[2:]
+    )
+
     return mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl

model/common/src/icon4py/model/common/metrics/compute_flat_idx_max.py

@@ -0,0 +1,39 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import numpy as np
+
+from icon4py.model.common.utils import data_allocation as data_alloc
+
+
+def compute_flat_idx_max(
+    e2c: data_alloc.NDArray,
+    z_mc: data_alloc.NDArray,
+    c_lin_e: data_alloc.NDArray,
+    z_ifc: data_alloc.NDArray,
+    k_lev: data_alloc.NDArray,
+    horizontal_lower: int,
+    horizontal_upper: int,
+) -> data_alloc.NDArray:
+    z_me = np.sum(z_mc[e2c] * np.expand_dims(c_lin_e, axis=-1), axis=1)
+    z_ifc_e_0 = z_ifc[e2c[:, 0]]
+    z_ifc_e_k_0 = np.roll(z_ifc_e_0, -1, axis=1)
+    z_ifc_e_1 = z_ifc[e2c[:, 1]]
+    z_ifc_e_k_1 = np.roll(z_ifc_e_1, -1, axis=1)
+    flat_idx = np.zeros_like(z_me)
+    for je in range(horizontal_lower, horizontal_upper):
+        for jk in range(k_lev.shape[0] - 1):
+            if (
+                (z_me[je, jk] <= z_ifc_e_0[je, jk])
+                and (z_me[je, jk] >= z_ifc_e_k_0[je, jk])
+                and (z_me[je, jk] <= z_ifc_e_1[je, jk])
+                and (z_me[je, jk] >= z_ifc_e_k_1[je, jk])
+            ):
+                flat_idx[je, jk] = k_lev[jk]
+    flat_idx_max = np.amax(flat_idx, axis=1)
+    return flat_idx_max.astype(np.int32)