add nudgecoeffs_e and wgtfac_c

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

@@ -19,6 +19,47 @@
 from icon4py.model.common.dimension import CellDim, ECDim, ECVDim, EdgeDim
 
 
+NUM_GHOST_ROWS: Final[int] = 2
+# values from mo_impl_constants.f90
+_ICON_INDEX_OFFSET_CELLS: Final[int] = 8
+_GRF_BOUNDARY_WIDTH_CELL: Final[int] = 4
+_MIN_RL_CELL_INT: Final[int] = -4
+_MIN_RL_CELL: Final[int] = _MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS
+_MAX_RL_CELL: Final[int] = 5
+
+_ICON_INDEX_OFFSET_VERTEX: Final[int] = 7
+_MIN_RL_VERTEX_INT: Final[int] = _MIN_RL_CELL_INT
+_MIN_RL_VERTEX: Final[int] = _MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1)
+_MAX_RL_VERTEX: Final[int] = _MAX_RL_CELL
+
+_ICON_INDEX_OFFSET_EDGES: Final[int] = 13
+_GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9
+_MIN_RL_EDGE_INT: Final[int] = 2 * _MIN_RL_CELL_INT
+_MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1)
+_MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL
+
+_LATERAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES
+_INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES
+_NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES
+_HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES
+_LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES
+_END_EDGES: Final[int] = 0
+
+_LATERAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS
+_INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS
+_NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS
+_HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS
+_LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS
+_END_CELLS: Final[int] = 0
+
+_LATERAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX
+_INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX
+_NUDGING_VERTICES: Final[int] = 0
+_HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX
+_LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX
+_END_VERTICES: Final[int] = 0
+
+
 class HorizontalMarkerIndex:
     """
     Handles constants indexing into the start_index and end_index fields.
@@ -39,46 +80,6 @@ class HorizontalMarkerIndex:
 
     """
 
-    NUM_GHOST_ROWS: Final[int] = 2
-    # values from mo_impl_constants.f90
-    _ICON_INDEX_OFFSET_CELLS: Final[int] = 8
-    _GRF_BOUNDARY_WIDTH_CELL: Final[int] = 4
-    _MIN_RL_CELL_INT: Final[int] = -4
-    _MIN_RL_CELL: Final[int] = _MIN_RL_CELL_INT - 2 * NUM_GHOST_ROWS
-    _MAX_RL_CELL: Final[int] = 5
-
-    _ICON_INDEX_OFFSET_VERTEX: Final[int] = 7
-    _MIN_RL_VERTEX_INT: Final[int] = _MIN_RL_CELL_INT
-    _MIN_RL_VERTEX: Final[int] = _MIN_RL_VERTEX_INT - (NUM_GHOST_ROWS + 1)
-    _MAX_RL_VERTEX: Final[int] = _MAX_RL_CELL
-
-    _ICON_INDEX_OFFSET_EDGES: Final[int] = 13
-    _GRF_BOUNDARY_WIDTH_EDGES: Final[int] = 9
-    _MIN_RL_EDGE_INT: Final[int] = 2 * _MIN_RL_CELL_INT
-    _MIN_RL_EDGE: Final[int] = _MIN_RL_EDGE_INT - (2 * NUM_GHOST_ROWS + 1)
-    _MAX_RL_EDGE: Final[int] = 2 * _MAX_RL_CELL
-
-    _LATERAL_BOUNDARY_EDGES: Final[int] = 1 + _ICON_INDEX_OFFSET_EDGES
-    _INTERIOR_EDGES: Final[int] = _ICON_INDEX_OFFSET_EDGES
-    _NUDGING_EDGES: Final[int] = _GRF_BOUNDARY_WIDTH_EDGES + _ICON_INDEX_OFFSET_EDGES
-    _HALO_EDGES: Final[int] = _MIN_RL_EDGE_INT - 1 + _ICON_INDEX_OFFSET_EDGES
-    _LOCAL_EDGES: Final[int] = _MIN_RL_EDGE_INT + _ICON_INDEX_OFFSET_EDGES
-    _END_EDGES: Final[int] = 0
-
-    _LATERAL_BOUNDARY_CELLS: Final[int] = 1 + _ICON_INDEX_OFFSET_CELLS
-    _INTERIOR_CELLS: Final[int] = _ICON_INDEX_OFFSET_CELLS
-    _NUDGING_CELLS: Final[int] = _GRF_BOUNDARY_WIDTH_CELL + 1 + _ICON_INDEX_OFFSET_CELLS
-    _HALO_CELLS: Final[int] = _MIN_RL_CELL_INT - 1 + _ICON_INDEX_OFFSET_CELLS
-    _LOCAL_CELLS: Final[int] = _MIN_RL_CELL_INT + _ICON_INDEX_OFFSET_CELLS
-    _END_CELLS: Final[int] = 0
-
-    _LATERAL_BOUNDARY_VERTICES = 1 + _ICON_INDEX_OFFSET_VERTEX
-    _INTERIOR_VERTICES: Final[int] = _ICON_INDEX_OFFSET_VERTEX
-    _NUDGING_VERTICES: Final[int] = 0
-    _HALO_VERTICES: Final[int] = _MIN_RL_VERTEX_INT - 1 + _ICON_INDEX_OFFSET_VERTEX
-    _LOCAL_VERTICES: Final[int] = _MIN_RL_VERTEX_INT + _ICON_INDEX_OFFSET_VERTEX
-    _END_VERTICES: Final[int] = 0
-
     _lateral_boundary = {
         dimension.CellDim: _LATERAL_BOUNDARY_CELLS,
         dimension.EdgeDim: _LATERAL_BOUNDARY_EDGES,
@@ -134,6 +135,11 @@ def nudging(cls, dim: Dimension) -> int:
         """Indicate the nudging zone."""
         return cls._nudging[dim]
 
+    @classmethod
+    def nudging_2nd_level(cls, dim: Dimension) -> int:
+        """Indicate the nudging zone for 2nd level."""
+        return cls.nudging(dim) + 1
+
     @classmethod
     def interior(cls, dim: Dimension) -> int:
         """Indicate interior i.e. unordered prognostic cells in ICON."""
@@ -289,3 +295,18 @@ class CellParams:
 
     defined int ICON in mo_model_domain.f90:t_grid_cells%area
     """
+
+
+class RefinCtrlLevel:
+    _boundary_nudging_start = {
+        EdgeDim: _GRF_BOUNDARY_WIDTH_EDGES + 1,
+        CellDim: _GRF_BOUNDARY_WIDTH_CELL + 1,
+    }
+
+    @classmethod
+    def boundary_nudging_start(cls, dim: Dimension) -> int:
+        """Start refin_ctrl levels for boundary nudging (as seen from the child domain)."""
+        try:
+            return cls._boundary_nudging_start[dim]
+        except KeyError:
+            raise ValueError(f"nudging start level only exists for {CellDim} and {EdgeDim}")

model/common/src/icon4py/model/common/metrics/stencils/calc_nudgecoeffs.py

@@ -0,0 +1,59 @@
+# 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, astype, exp, int32, where
+
+from icon4py.model.common.dimension import EdgeDim
+from icon4py.model.common.type_alias import wpfloat
+
+
+@field_operator
+def _calc_nudgecoeffs(
+    refin_ctrl: Field[[EdgeDim], int32],
+    grf_nudge_start_e: int32,
+    nudge_max_coeffs: wpfloat,
+    nudge_efold_width: wpfloat,
+    nudge_zone_width: int32,
+) -> Field[[EdgeDim], wpfloat]:
+
+    return where(
+        ((refin_ctrl > 0) & (refin_ctrl <= (2 * nudge_zone_width + (grf_nudge_start_e - 3)))),
+        nudge_max_coeffs
+        * exp((-(astype(refin_ctrl - grf_nudge_start_e, wpfloat))) / (2.0 * nudge_efold_width)),
+        0.0,
+    )
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def calc_nudgecoeffs(
+    nudgecoeffs_e: Field[[EdgeDim], wpfloat],
+    refin_ctrl: Field[[EdgeDim], int32],
+    grf_nudge_start_e: int32,
+    nudge_max_coeffs: wpfloat,
+    nudge_efold_width: wpfloat,
+    nudge_zone_width: int32,
+    horizontal_start: int32,
+    horizontal_end: int32,
+):
+    """Calculate nudging coefficient for edges based the grid refinement level of an edge."""
+    _calc_nudgecoeffs(
+        refin_ctrl,
+        grf_nudge_start_e,
+        nudge_max_coeffs,
+        nudge_efold_width,
+        nudge_zone_width,
+        out=nudgecoeffs_e,
+        domain={EdgeDim: (horizontal_start, horizontal_end)},
+    )

model/common/src/icon4py/model/common/metrics/stencils/calc_wgtfac_c.py

@@ -0,0 +1,76 @@
+# 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, where
+from gt4py.next.program_processors.runners.gtfn import run_gtfn
+
+from icon4py.model.common.dimension import CellDim, KDim, Koff
+from icon4py.model.common.type_alias import wpfloat
+
+
+@field_operator
+def _calc_wgtfac_c_nlevp1(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+) -> Field[[CellDim, KDim], wpfloat]:
+
+    z_wgtfac_c = (z_ifc(Koff[-1]) - z_ifc) / (z_ifc(Koff[-2]) - z_ifc)
+    return z_wgtfac_c
+
+
+@field_operator
+def _calc_wgtfac_c_0(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+) -> Field[[CellDim, KDim], wpfloat]:
+
+    z_wgtfac_c = (z_ifc(Koff[+1]) - z_ifc) / (z_ifc(Koff[+2]) - z_ifc)
+    return z_wgtfac_c
+
+
+@field_operator
+def _calc_wgtfac_c_inner(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+) -> Field[[CellDim, KDim], wpfloat]:
+
+    z_wgtfac_c = (z_ifc(Koff[-1]) - z_ifc) / (z_ifc(Koff[-1]) - z_ifc(Koff[+1]))
+    return z_wgtfac_c
+
+
+@field_operator
+def _calc_wgtfac_c(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    k_field: Field[[KDim], int32],
+    nlevp1: int32,
+) -> Field[[CellDim, KDim], wpfloat]:
+
+    wgt_fac_c = where((k_field > int32(0)) & (k_field < nlevp1), _calc_wgtfac_c_inner(z_ifc), z_ifc)
+    wgt_fac_c = where(k_field == int32(0), _calc_wgtfac_c_0(z_ifc=z_ifc), wgt_fac_c)
+    wgt_fac_c = where(k_field == nlevp1, _calc_wgtfac_c_nlevp1(z_ifc=z_ifc), wgt_fac_c)
+
+    return wgt_fac_c
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=run_gtfn)
+def calc_wgtfac_c(
+    wgtfac_c: Field[[CellDim, KDim], wpfloat],
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    k_field: Field[[KDim], int32],
+    nlevp1: int32,
+):
+    _calc_wgtfac_c(
+        z_ifc,
+        k_field,
+        nlevp1,
+        out=wgtfac_c,
+    )

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

@@ -563,6 +563,9 @@ def ddxn_z_full(self):
     def ddxt_z_full(self):
         return self._get_field("ddxt_z_full", EdgeDim, KDim)
 
+    def z_ifc(self):
+        return self._get_field("z_ifc", CellDim, KDim)
+
     @optionally_registered
     def mask_hdiff(self):
         return self._get_field("mask_hdiff", CellDim, KDim, dtype=bool)

model/common/tests/metric_tests/test_calc_nudgecoeffs.py

@@ -0,0 +1,80 @@
+# 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
+# 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
+
+import numpy as np
+import pytest
+from gt4py.next.ffront.fbuiltins import int32
+
+from icon4py.model.common.dimension import EdgeDim
+from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex, RefinCtrlLevel
+from icon4py.model.common.metrics.stencils.calc_nudgecoeffs import calc_nudgecoeffs
+from icon4py.model.common.test_utils.datatest_fixtures import (  # noqa: F401  # import fixtures from test_utils package
+    data_provider,
+    datapath,
+    download_ser_data,
+    experiment,
+    grid_savepoint,
+    icon_grid,
+    interpolation_savepoint,
+    processor_props,
+    ranked_data_path,
+)
+from icon4py.model.common.test_utils.helpers import zero_field
+from icon4py.model.common.type_alias import wpfloat
+
+
+@pytest.mark.datatest
+def test_calc_nudgecoeffs_e(
+    grid_savepoint, interpolation_savepoint, icon_grid  # noqa: F811  # fixture
+):
+    nudgecoeff_e = zero_field(icon_grid, EdgeDim, dtype=wpfloat)
+    nudgecoeff_e_ref = interpolation_savepoint.nudgecoeff_e()
+    refin_ctrl = grid_savepoint.refin_ctrl(EdgeDim)
+    grf_nudge_start_e = RefinCtrlLevel.boundary_nudging_start(EdgeDim)
+    nudge_max_coeff = wpfloat(0.375)
+    nudge_efold_width = wpfloat(2.0)
+    nudge_zone_width = int32(10)
+
+    horizontal_start = icon_grid.get_start_index(
+        EdgeDim, HorizontalMarkerIndex.nudging_2nd_level(EdgeDim)
+    )
+    horizontal_end = icon_grid.get_end_index(
+        EdgeDim,
+        HorizontalMarkerIndex.local(EdgeDim),
+    )
+
+    calc_nudgecoeffs(
+        nudgecoeff_e,
+        refin_ctrl,
+        grf_nudge_start_e,
+        nudge_max_coeff,
+        nudge_efold_width,
+        nudge_zone_width,
+        horizontal_start,
+        horizontal_end,
+        offset_provider={},
+    )
+
+    assert np.allclose(nudgecoeff_e.asnumpy(), nudgecoeff_e_ref.asnumpy())