Metrics fields second batch (#454)

* further interpolation coefficients

model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_cell.py

@@ -0,0 +1,61 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# 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, int32, neighbor_sum
+
+from icon4py.model.common.dimension import (
+    C2E,
+    C2EDim,
+    CellDim,
+    EdgeDim,
+    KDim,
+)
+from icon4py.model.common.settings import backend
+from icon4py.model.common.type_alias import wpfloat
+
+
+@field_operator
+def _mo_intp_rbf_rbf_vec_interpol_cell(
+    p_vn_in: Field[[EdgeDim, KDim], wpfloat],
+    ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat],
+    ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat],
+) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
+    p_u_out = neighbor_sum(ptr_coeff_1 * p_vn_in(C2E), axis=C2EDim)
+    p_v_out = neighbor_sum(ptr_coeff_2 * p_vn_in(C2E), axis=C2EDim)
+    return p_u_out, p_v_out
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+def mo_intp_rbf_rbf_vec_interpol_cell(
+    p_vn_in: Field[[EdgeDim, KDim], wpfloat],
+    ptr_coeff_1: Field[[CellDim, C2EDim], wpfloat],
+    ptr_coeff_2: Field[[CellDim, C2EDim], wpfloat],
+    p_u_out: Field[[CellDim, KDim], wpfloat],
+    p_v_out: Field[[CellDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _mo_intp_rbf_rbf_vec_interpol_cell(
+        p_vn_in,
+        ptr_coeff_1,
+        ptr_coeff_2,
+        out=(p_u_out, p_v_out),
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/common/src/icon4py/model/common/interpolation/stencils/mo_intp_rbf_rbf_vec_interpol_vertex.py

@@ -15,7 +15,13 @@
 from gt4py.next.ffront.decorator import field_operator, program
 from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum
 
-from icon4py.model.common.dimension import V2E, EdgeDim, KDim, V2EDim, VertexDim
+from icon4py.model.common.dimension import (
+    V2E,
+    EdgeDim,
+    KDim,
+    V2EDim,
+    VertexDim,
+)
 from icon4py.model.common.settings import backend
 from icon4py.model.common.type_alias import wpfloat
 

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

@@ -25,12 +25,17 @@
     where,
 )
 
-from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim
+from icon4py.model.common.dimension import (
+    CellDim,
+    EdgeDim,
+    KDim,
+    Koff,
+    VertexDim,
+)
 from icon4py.model.common.math.helpers import (
     _grad_fd_tang,
     average_cell_kdim_level_up,
     average_edge_kdim_level_up,
-    difference_k_level_down,
     difference_k_level_up,
     grad_fd_norm,
 )
@@ -79,12 +84,11 @@ def _compute_ddqz_z_half(
     z_mc: Field[[CellDim, KDim], wpfloat],
     k: Field[[KDim], int32],
     nlev: int32,
-) -> Field[[CellDim, KDim], wpfloat]:
-    ddqz_z_half = where(
-        (k > int32(0)) & (k < nlev),
-        difference_k_level_down(z_mc),
-        where(k == 0, 2.0 * (z_ifc - z_mc), 2.0 * (z_mc(Koff[-1]) - z_ifc)),
-    )
+):
+    # TODO: change this to concat_where once it's merged
+    ddqz_z_half = where(k == 0, 2.0 * (z_ifc - z_mc), 0.0)
+    ddqz_z_half = where((k > 0) & (k < nlev), z_mc(Koff[-1]) - z_mc, ddqz_z_half)
+    ddqz_z_half = where(k == nlev, 2.0 * (z_mc(Koff[-1]) - z_ifc), ddqz_z_half)
     return ddqz_z_half
 
 
@@ -127,7 +131,7 @@ def compute_ddqz_z_half(
 
 
 @field_operator
-def _compute_ddqz_z_full(
+def _compute_ddqz_z_full_and_inverse(
     z_ifc: Field[[CellDim, KDim], wpfloat],
 ) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
     ddqz_z_full = difference_k_level_up(z_ifc)
@@ -136,7 +140,7 @@ def _compute_ddqz_z_full(
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
-def compute_ddqz_z_full(
+def compute_ddqz_z_full_and_inverse(
     z_ifc: Field[[CellDim, KDim], wpfloat],
     ddqz_z_full: Field[[CellDim, KDim], wpfloat],
     inv_ddqz_z_full: Field[[CellDim, KDim], wpfloat],
@@ -161,7 +165,7 @@ def compute_ddqz_z_full(
         vertical_end: vertical end index
 
     """
-    _compute_ddqz_z_full(
+    _compute_ddqz_z_full_and_inverse(
         z_ifc,
         out=(ddqz_z_full, inv_ddqz_z_full),
         domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
@@ -488,7 +492,39 @@ def compute_ddxt_z_half_e(
 
 
 @program
-def compute_ddxnt_z_full(
+def compute_ddxn_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)
+
+
+@field_operator
+def _compute_vwind_expl_wgt(vwind_impl_wgt: Field[[CellDim], wpfloat]) -> Field[[CellDim], wpfloat]:
+    return 1.0 - vwind_impl_wgt
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def compute_vwind_expl_wgt(
+    vwind_impl_wgt: Field[[CellDim], wpfloat],
+    vwind_expl_wgt: Field[[CellDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+):
+    """
+    Compute vwind_expl_wgt.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        vwind_impl_wgt: offcentering in vertical mass flux
+        vwind_expl_wgt: (output) 1 - of vwind_impl_wgt
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+
+    """
+
+    _compute_vwind_expl_wgt(
+        vwind_impl_wgt=vwind_impl_wgt,
+        out=vwind_expl_wgt,
+        domain={CellDim: (horizontal_start, horizontal_end)},
+    )

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

@@ -200,6 +200,9 @@ def edge_areas(self):
     def inv_dual_edge_length(self):
         return self._get_field("inv_dual_edge_length", EdgeDim)
 
+    def dual_edge_length(self):
+        return self._get_field("dual_edge_length", EdgeDim)
+
     def edge_cell_length(self):
         return self._get_field("edge_cell_length", EdgeDim, E2CDim)