WIP: Greenline metrics coefficients aj (#429)

* further interpolation coefficients

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

@@ -17,7 +17,18 @@
 from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.common import dimension
-from icon4py.model.common.dimension import E2C, CellDim, E2CDim, ECDim, ECVDim, EdgeDim, KDim
+from icon4py.model.common.dimension import (
+    E2C,
+    V2C,
+    CellDim,
+    E2CDim,
+    ECDim,
+    ECVDim,
+    EdgeDim,
+    KDim,
+    V2CDim,
+    VertexDim,
+)
 from icon4py.model.common.type_alias import wpfloat
 
 
@@ -343,3 +354,42 @@ def boundary_nudging_start(cls, dim: Dimension) -> int:
             raise ValueError(
                 f"nudging start level only exists for {CellDim} and {EdgeDim}"
             ) from err
+
+
+@field_operator
+def _compute_cells2edges(
+    p_cell_in: Field[[CellDim, KDim], float],
+    c_int: Field[[EdgeDim, E2CDim], float],
+) -> Field[[EdgeDim, KDim], float]:
+    p_vert_out = neighbor_sum(c_int * p_cell_in(E2C), axis=E2CDim)
+    return p_vert_out
+
+
+@program
+def compute_cells2edges(
+    p_cell_in: Field[[CellDim, KDim], float],
+    c_int: Field[[EdgeDim, E2CDim], float],
+    p_vert_out: Field[[EdgeDim, KDim], float],
+    horizontal_start_edge: int32,
+    horizontal_end_edge: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _compute_cells2edges(
+        p_cell_in,
+        c_int,
+        out=p_vert_out,
+        domain={
+            EdgeDim: (horizontal_start_edge, horizontal_end_edge),
+            KDim: (vertical_start, vertical_end),
+        },
+    )
+
+
+@field_operator
+def _compute_cells2verts(
+    p_cell_in: Field[[CellDim, KDim], float],
+    c_int: Field[[VertexDim, V2CDim], float],
+) -> Field[[VertexDim, KDim], float]:
+    p_vert_out = neighbor_sum(c_int * p_cell_in(V2C), axis=V2CDim)
+    return p_vert_out

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

@@ -49,3 +49,35 @@ def compute_c_lin_e(
     c_lin_e[0:second_boundary_layer_start_index, :] = 0.0
     mask = np.transpose(np.tile(owner_mask, (2, 1)))
     return np.where(mask, c_lin_e, 0.0)
+
+
+def compute_cells_aw_verts(
+    dual_area: np.array,
+    edge_vert_length: np.array,
+    edge_cell_length: np.array,
+    owner_mask: np.array,
+    e2c: np.array,
+    v2c: np.array,
+    v2e: np.array,
+    e2v: np.array,
+    horizontal_start: np.int32,
+    horizontal_end: np.int32,
+) -> np.array:
+    llb = horizontal_start
+    cells_aw_verts = np.zeros([horizontal_end, 6])
+    index = np.repeat(np.arange(horizontal_end, dtype=float), 6).reshape(horizontal_end, 6)
+    idx_ve = np.where(index == e2v[v2e, 0], 0, 1)
+
+    for i in range(2):
+        for j in range(6):
+            for k in range(6):
+                cells_aw_verts[llb:, k] = np.where(
+                    np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]),
+                    cells_aw_verts[llb:, k]
+                    + 0.5
+                    / dual_area[llb:]
+                    * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]]
+                    * edge_cell_length[v2e[llb:, j], i],
+                    cells_aw_verts[llb:, k],
+                )
+    return cells_aw_verts

model/common/src/icon4py/model/common/math/helpers.py

@@ -13,12 +13,12 @@
 
 from gt4py.next import Field, field_operator
 
-from icon4py.model.common.dimension import CellDim, KDim, Koff
+from icon4py.model.common.dimension import E2C, E2V, CellDim, EdgeDim, KDim, Koff, VertexDim
 from icon4py.model.common.type_alias import wpfloat
 
 
 @field_operator
-def average_k_level_up(
+def average_cell_kdim_level_up(
     half_level_field: Field[[CellDim, KDim], wpfloat],
 ) -> Field[[CellDim, KDim], wpfloat]:
     """
@@ -35,6 +35,24 @@ def average_k_level_up(
     return 0.5 * (half_level_field + half_level_field(Koff[1]))
 
 
+@field_operator
+def average_edge_kdim_level_up(
+    half_level_field: Field[[EdgeDim, KDim], wpfloat],
+) -> Field[[EdgeDim, KDim], wpfloat]:
+    """
+    Calculate the mean value of adjacent interface levels.
+
+    Computes the average of two adjacent interface levels upwards over an edge field for storage
+    in the corresponding full levels.
+    Args:
+        half_level_field: Field[[EdgeDim, KDim], wpfloat]
+
+    Returns: Field[[EdgeDim, KDim], wpfloat] full level field
+
+    """
+    return 0.5 * (half_level_field + half_level_field(Koff[1]))
+
+
 @field_operator
 def difference_k_level_down(
     half_level_field: Field[[CellDim, KDim], wpfloat],
@@ -69,3 +87,33 @@ def difference_k_level_up(
 
     """
     return half_level_field - half_level_field(Koff[1])
+
+
+@field_operator
+def grad_fd_norm(
+    psi_c: Field[[CellDim, KDim], float],
+    inv_dual_edge_length: Field[[EdgeDim], float],
+) -> Field[[EdgeDim, KDim], float]:
+    """
+    Calculate the gradient value of adjacent interface levels.
+
+    Computes the difference of two offseted values multiplied by a field of the offseted dimension
+    Args:
+        psi_c: Field[[CellDim, KDim], float],
+        inv_dual_edge_length: Field[[EdgeDim], float],
+
+    Returns: Field[[EdgeDim, KDim], float]
+
+    """
+    grad_norm_psi_e = (psi_c(E2C[1]) - psi_c(E2C[0])) * inv_dual_edge_length
+    return grad_norm_psi_e
+
+
+@field_operator
+def _grad_fd_tang(
+    psi_v: Field[[VertexDim, KDim], float],
+    inv_primal_edge_length: Field[[EdgeDim], float],
+    tangent_orientation: Field[[EdgeDim], float],
+) -> Field[[EdgeDim, KDim], float]:
+    grad_tang_psi_e = tangent_orientation * (psi_v(E2V[1]) - psi_v(E2V[0])) * inv_primal_edge_length
+    return grad_tang_psi_e

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

@@ -25,11 +25,14 @@
     where,
 )
 
-from icon4py.model.common.dimension import CellDim, KDim, Koff
+from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim
 from icon4py.model.common.math.helpers import (
-    average_k_level_up,
+    _grad_fd_tang,
+    average_cell_kdim_level_up,
+    average_edge_kdim_level_up,
     difference_k_level_down,
     difference_k_level_up,
+    grad_fd_norm,
 )
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
@@ -63,7 +66,7 @@ def compute_z_mc(
         vertical_end:int32 end index of vertical domain
 
     """
-    average_k_level_up(
+    average_cell_kdim_level_up(
         z_ifc,
         out=z_mc,
         domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
@@ -438,3 +441,54 @@ def compute_d2dexdz2_fac_mc(
         out=d2dexdz2_fac2_mc,
         domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
     )
+
+
+@program
+def compute_ddxn_z_half_e(
+    z_ifc: Field[[CellDim, KDim], float],
+    inv_dual_edge_length: Field[[EdgeDim], float],
+    ddxn_z_half_e: Field[[EdgeDim, KDim], float],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    grad_fd_norm(
+        z_ifc,
+        inv_dual_edge_length,
+        out=ddxn_z_half_e,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )
+
+
+@program
+def compute_ddxt_z_half_e(
+    z_ifv: Field[[VertexDim, KDim], float],
+    inv_primal_edge_length: Field[[EdgeDim], float],
+    tangent_orientation: Field[[EdgeDim], float],
+    ddxt_z_half_e: Field[[EdgeDim, KDim], float],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _grad_fd_tang(
+        z_ifv,
+        inv_primal_edge_length,
+        tangent_orientation,
+        out=ddxt_z_half_e,
+        domain={
+            EdgeDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )
+
+
+@program
+def compute_ddxnt_z_full(
+    z_ddxnt_z_half_e: Field[[EdgeDim, KDim], float], ddxn_z_full: Field[[EdgeDim, KDim], float]
+):
+    average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full)

model/common/src/icon4py/model/common/test_utils/serialbox_utils.py

@@ -142,6 +142,12 @@ def _read(self, name: str, offset=0, dtype=int):
 
 
 class IconGridSavepoint(IconSavepoint):
+    def v_dual_area(self):
+        return self._get_field("v_dual_area", VertexDim)
+
+    def edge_vert_length(self):
+        return self._get_field("edge_vert_length", EdgeDim, E2C2VDim)
+
     def vct_a(self):
         return self._get_field("vct_a", KDim)
 
@@ -219,6 +225,9 @@ def edge_end_index(self):
         # one off accounts for being exclusive [from:to)
         return self.serializer.read("e_end_index", self.savepoint)
 
+    def v_owner_mask(self):
+        return self._get_field("v_owner_mask", VertexDim, dtype=bool)
+
     def c_owner_mask(self):
         return self._get_field("c_owner_mask", CellDim, dtype=bool)
 

model/common/tests/interpolation_tests/test_interpolation_fields.py

@@ -26,9 +26,29 @@
 import numpy as np
 import pytest
 
-from icon4py.model.common.dimension import EdgeDim
-from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex
-from icon4py.model.common.interpolation.interpolation_fields import compute_c_lin_e
+from icon4py.model.common.dimension import (
+    E2CDim,
+    E2VDim,
+    EdgeDim,
+    V2CDim,
+    V2EDim,
+    VertexDim,
+)
+from icon4py.model.common.grid.horizontal import (
+    HorizontalMarkerIndex,
+)
+from icon4py.model.common.interpolation.interpolation_fields import (
+    compute_c_lin_e,
+    compute_cells_aw_verts,
+)
+from icon4py.model.common.test_utils.datatest_fixtures import (  # noqa: F401  # import fixtures from test_utils package
+    data_provider,
+    datapath,
+    download_ser_data,
+    experiment,
+    processor_props,
+    ranked_data_path,
+)
 
 
 @pytest.mark.datatest
@@ -49,3 +69,40 @@ def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid):
     )
 
     assert np.allclose(c_lin_e, c_lin_e_ref.asnumpy())
+
+
+@pytest.mark.datatest
+def test_compute_cells_aw_verts(
+    grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint
+):
+    dual_area = grid_savepoint.v_dual_area().asnumpy()
+    edge_vert_length = grid_savepoint.edge_vert_length().asnumpy()
+    edge_cell_length = grid_savepoint.edge_cell_length().asnumpy()
+    owner_mask = grid_savepoint.v_owner_mask()
+    e2c = icon_grid.connectivities[E2CDim]
+    v2c = icon_grid.connectivities[V2CDim]
+    v2e = icon_grid.connectivities[V2EDim]
+    e2v = icon_grid.connectivities[E2VDim]
+    horizontal_start_vertex = icon_grid.get_start_index(
+        VertexDim,
+        HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1,
+    )
+    horizontal_end_vertex = icon_grid.get_end_index(
+        VertexDim,
+        HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1,
+    )
+
+    cells_aw_verts = compute_cells_aw_verts(
+        dual_area,
+        edge_vert_length,
+        edge_cell_length,
+        owner_mask,
+        e2c,
+        v2c,
+        v2e,
+        e2v,
+        horizontal_start_vertex,
+        horizontal_end_vertex,
+    )
+    cells_aw_verts_ref = interpolation_savepoint.c_intp().asnumpy()
+    assert np.allclose(cells_aw_verts, cells_aw_verts_ref)