Porting of metric fields

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

@@ -248,6 +248,8 @@ def __init__(
         is_iau_active: bool = False,
         iau_wgt_dyn: float = 0.0,
         divdamp_type: int = 3,
+        divdamp_trans_start: float = 12500.0,
+        divdamp_trans_end: float = 17500.0,
         l_vert_nested: bool = False,
         rhotheta_offctr: float = -0.1,
         veladv_offctr: float = 0.25,
@@ -285,6 +287,8 @@ def __init__(
 
         #: type of divergence damping
         self.divdamp_type: int = divdamp_type
+        self.divdamp_trans_start: float = divdamp_trans_start
+        self.divdamp_trans_end: float = divdamp_trans_end
 
         #: off-centering for density and potential temperature at interface levels.
         #: Specifying a negative value here reduces the amount of vertical

model/common/src/icon4py/model/common/constants.py

@@ -69,3 +69,6 @@
 #: Klemp (2008) type Rayleigh damping
 # TODO (magdalena) not a constant, move somewhere else, convert to enum
 RAYLEIGH_KLEMP: Final[int] = 2
+RAYLEIGH_CLASSIC: Final[
+    int
+] = 1  # classical Rayleigh damping, which makes use of a reference state.

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

@@ -11,15 +11,26 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from gt4py.next import Field, GridType, field_operator, int32, program, where
+from gt4py.next import (
+    Field,
+    GridType,
+    broadcast,
+    field_operator,
+    int32,
+    maximum,
+    program,
+    sin,
+    tanh,
+    where,
+)
 
 from icon4py.model.common.dimension import CellDim, KDim, Koff
 from icon4py.model.common.math.helpers import (
     average_k_level_up,
     difference_k_level_down,
     difference_k_level_up,
 )
-from icon4py.model.common.type_alias import wpfloat
+from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
 """
@@ -151,3 +162,139 @@ def compute_ddqz_z_full(
         out=(ddqz_z_full, inv_ddqz_z_full),
         domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
     )
+
+
+@field_operator
+def _compute_scalfac_dd3d(
+    vct_a: Field[[KDim], wpfloat],
+    divdamp_trans_start: wpfloat,
+    divdamp_trans_end: wpfloat,
+    divdamp_type: int32,
+) -> Field[[KDim], wpfloat]:
+    scalfac_dd3d = broadcast(1.0, (KDim,))
+    if divdamp_type == 32:
+        zf = 0.5 * (vct_a + vct_a(Koff[1]))
+        scalfac_dd3d = where(zf >= divdamp_trans_end, 0.0, scalfac_dd3d)
+        scalfac_dd3d = where(
+            zf >= divdamp_trans_start,
+            (divdamp_trans_end - zf) / (divdamp_trans_end - divdamp_trans_start),
+            scalfac_dd3d,
+        )
+    return scalfac_dd3d
+
+
+@program
+def compute_scalfac_dd3d(
+    vct_a: Field[[KDim], wpfloat],
+    scalfac_dd3d: Field[[KDim], wpfloat],
+    divdamp_trans_start: wpfloat,
+    divdamp_trans_end: wpfloat,
+    divdamp_type: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute scalfac_dd3d.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        vct_a: Field[[KDim], float],
+        scalfac_dd3d: (output)
+        divdamp_trans_start: lower bound of transition zone between 2D and 3D div damping in case of divdamp_type = 32
+        divdamp_trans_end: upper bound of transition zone between 2D and 3D div damping in case of divdamp_type = 32
+        divdamp_type: type of divergence damping (2D or 3D divergence)
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+    _compute_scalfac_dd3d(
+        vct_a,
+        divdamp_trans_start,
+        divdamp_trans_end,
+        divdamp_type,
+        out=scalfac_dd3d,
+        domain={KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_rayleigh_w(
+    vct_a: Field[[KDim], wpfloat],
+    vct_a_1: Field[[], wpfloat],
+    damping_height: wpfloat,
+    rayleigh_type: int32,
+    rayleigh_classic: int32,
+    rayleigh_klemp: int32,
+    rayleigh_coeff: wpfloat,
+    pi_const: wpfloat,
+) -> Field[[KDim], wpfloat]:
+    rayleigh_w = broadcast(0.0, (KDim,))
+    z_sin_diff = maximum(0.0, vct_a - damping_height)
+    z_tanh_diff = vct_a_1 - vct_a  # vct_a(1) - vct_a
+    if rayleigh_type == rayleigh_classic:
+        rayleigh_w = (
+            rayleigh_coeff
+            * sin(pi_const / 2.0 * z_sin_diff / maximum(0.001, vct_a_1 - damping_height)) ** 2
+        )
+    elif rayleigh_type == rayleigh_klemp:
+        rayleigh_w = rayleigh_coeff * (
+            1.0 - tanh(3.8 * z_tanh_diff / maximum(0.000001, vct_a_1 - damping_height))
+        )
+    return rayleigh_w
+
+
+@program
+def compute_rayleigh_w(
+    rayleigh_w: Field[[KDim], wpfloat],
+    vct_a: Field[[KDim], wpfloat],
+    vct_a_1: Field[[], wpfloat],
+    damping_height: wpfloat,
+    rayleigh_type: int32,
+    rayleigh_classic: int32,
+    rayleigh_klemp: int32,
+    rayleigh_coeff: wpfloat,
+    pi_const: wpfloat,
+    vertical_start: int32,  # 1, nrdmax(jg)
+    vertical_end: int32,
+):
+    _compute_rayleigh_w(
+        vct_a,
+        vct_a_1,
+        damping_height,
+        rayleigh_type,
+        rayleigh_classic,
+        rayleigh_klemp,
+        rayleigh_coeff,
+        pi_const,
+        out=rayleigh_w,
+        domain={KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_coeff_dwdz(
+    ddqz_z_full: Field[[CellDim, KDim], float], z_ifc: Field[[CellDim, KDim], wpfloat]
+) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]:
+    coeff1_dwdz = ddqz_z_full / ddqz_z_full(Koff[-1]) / (z_ifc(Koff[-1]) - z_ifc(Koff[1]))
+    coeff2_dwdz = ddqz_z_full(Koff[-1]) / ddqz_z_full / (z_ifc(Koff[-1]) - z_ifc(Koff[1]))
+
+    return coeff1_dwdz, coeff2_dwdz
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def compute_coeff_dwdz(
+    ddqz_z_full: Field[[CellDim, KDim], float],
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    coeff1_dwdz: Field[[CellDim, KDim], vpfloat],
+    coeff2_dwdz: Field[[CellDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _compute_coeff_dwdz(
+        ddqz_z_full,
+        z_ifc,
+        out=(coeff1_dwdz, coeff2_dwdz),
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )

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

@@ -0,0 +1,17 @@
+# 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
+
+
+def compute_kstart_dd3d(scalfac_dd3d: np.array) -> int:
+    return np.min(scalfac_dd3d[np.where(scalfac_dd3d > 0)])

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

@@ -518,6 +518,9 @@ def rho_ref_me(self):
     def scalfac_dd3d(self):
         return self._get_field("scalfac_dd3d", KDim)
 
+    def kstart_dd3d_ref(self):
+        return self.serializer.read("kstart_dd3d_ref", self.savepoint)[0]
+
     def theta_ref_ic(self):
         return self._get_field("theta_ref_ic", CellDim, KDim)
 

model/common/tests/metric_tests/test_metric_fields.py

@@ -16,12 +16,17 @@
 from gt4py.next import as_field
 from gt4py.next.ffront.fbuiltins import int32
 
+from icon4py.model.common import constants
 from icon4py.model.common.dimension import CellDim, KDim
 from icon4py.model.common.metrics.metric_fields import (
+    compute_coeff_dwdz,
     compute_ddqz_z_full,
     compute_ddqz_z_half,
+    compute_rayleigh_w,
+    compute_scalfac_dd3d,
     compute_z_mc,
 )
+from icon4py.model.common.metrics.metric_scalars import compute_kstart_dd3d
 from icon4py.model.common.test_utils.helpers import (
     StencilTest,
     dallclose,
@@ -120,3 +125,83 @@ def test_compute_ddqz_z_full(icon_grid, metrics_savepoint, backend):
     )
 
     assert dallclose(inv_ddqz_z_full.asnumpy(), inv_ddqz_full_ref.asnumpy())
+
+
+def test_compute_scalfac_dd3d(icon_grid, metrics_savepoint, grid_savepoint, backend):
+    scalfac_dd3d_ref = metrics_savepoint.scalfac_dd3d()
+    scalfac_dd3d_full = zero_field(icon_grid, KDim)
+    divdamp_trans_start = 12500.0
+    divdamp_trans_end = 17500.0
+    divdamp_type = 3
+    kstart_dd3d_ref = 1.0
+
+    compute_scalfac_dd3d.with_backend(backend=backend)(
+        vct_a=grid_savepoint.vct_a(),
+        scalfac_dd3d=scalfac_dd3d_full,
+        divdamp_trans_start=divdamp_trans_start,
+        divdamp_trans_end=divdamp_trans_end,
+        divdamp_type=divdamp_type,
+        vertical_start=int32(0),
+        vertical_end=icon_grid.num_levels,
+        offset_provider={"Koff": icon_grid.get_offset_provider("Koff")},
+    )
+
+    kstart_dd3d_full = compute_kstart_dd3d(
+        scalfac_dd3d=scalfac_dd3d_full.asnumpy(),
+    )
+
+    assert dallclose(scalfac_dd3d_ref.asnumpy(), scalfac_dd3d_full.asnumpy())
+    assert dallclose(kstart_dd3d_ref, kstart_dd3d_full)
+
+
+def test_compute_rayleigh_w(icon_grid, metrics_savepoint, grid_savepoint, backend):
+    import math
+
+    rayleigh_w_ref = metrics_savepoint.rayleigh_w()
+    vct_a_1 = as_field((), grid_savepoint.vct_a().asnumpy()[0])
+    rayleigh_w_full = zero_field(icon_grid, KDim, extend={KDim: 1})
+    rayleigh_type = 2
+    rayleigh_coeff = 5.0
+    damping_height = 12500.0
+    compute_rayleigh_w.with_backend(backend=backend)(
+        rayleigh_w=rayleigh_w_full,
+        vct_a=grid_savepoint.vct_a(),
+        vct_a_1=vct_a_1,
+        damping_height=damping_height,
+        rayleigh_type=rayleigh_type,
+        rayleigh_classic=constants.RAYLEIGH_CLASSIC,
+        rayleigh_klemp=constants.RAYLEIGH_KLEMP,
+        rayleigh_coeff=rayleigh_coeff,
+        pi_const=math.pi,
+        vertical_start=int32(0),
+        vertical_end=grid_savepoint.nrdmax().item() + 1,
+        offset_provider={},
+    )
+
+    assert dallclose(rayleigh_w_full.asnumpy(), rayleigh_w_ref.asnumpy())
+
+
+def test_compute_coeff_dwdz(icon_grid, metrics_savepoint, grid_savepoint, backend):
+    if is_roundtrip(backend):
+        pytest.skip("skipping: slow backend")
+    coeff1_dwdz_ref = metrics_savepoint.coeff1_dwdz()
+    coeff2_dwdz_ref = metrics_savepoint.coeff2_dwdz()
+
+    coeff1_dwdz_full = zero_field(icon_grid, CellDim, KDim)
+    coeff2_dwdz_full = zero_field(icon_grid, CellDim, KDim)
+    ddqz_z_full = as_field((CellDim, KDim), 1 / metrics_savepoint.inv_ddqz_z_full().asnumpy())
+
+    compute_coeff_dwdz.with_backend(backend=backend)(
+        ddqz_z_full=ddqz_z_full,
+        z_ifc=metrics_savepoint.z_ifc(),
+        coeff1_dwdz=coeff1_dwdz_full,
+        coeff2_dwdz=coeff2_dwdz_full,
+        horizontal_start=int32(0),
+        horizontal_end=icon_grid.num_cells,
+        vertical_start=int32(1),
+        vertical_end=int32(icon_grid.num_levels),
+        offset_provider={"Koff": icon_grid.get_offset_provider("Koff")},
+    )
+
+    assert dallclose(coeff1_dwdz_full.asnumpy(), coeff1_dwdz_ref.asnumpy())
+    assert dallclose(coeff2_dwdz_full.asnumpy(), coeff2_dwdz_ref.asnumpy())