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,9 @@ def __init__(
 
         #: type of divergence damping
         self.divdamp_type: int = divdamp_type
+        #: Lower and upper bound of transition zone between 2D and 3D divergence damping in case of divdamp_type = 32 [m]
+        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
@@ -1257,7 +1262,7 @@ def run_predictor_step(
             offset_provider={},
         )
 
-        if self.config.rayleigh_type == constants.RAYLEIGH_KLEMP:
+        if self.config.rayleigh_type == constants.RayleighType.RAYLEIGH_KLEMP:
             apply_rayleigh_damping_mechanism(
                 z_raylfac=self.z_raylfac,
                 w_1=prognostic_state[nnew].w_1,
@@ -1807,7 +1812,7 @@ def run_corrector_step(
             offset_provider={},
         )
 
-        if self.config.rayleigh_type == constants.RAYLEIGH_KLEMP:
+        if self.config.rayleigh_type == constants.RayleighType.RAYLEIGH_KLEMP:
             log.debug(f"corrector start stencil 54")
             apply_rayleigh_damping_mechanism(
                 z_raylfac=self.z_raylfac,

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

@@ -11,6 +11,7 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 import sys
+from enum import IntEnum
 from typing import Final
 
 from icon4py.model.common.type_alias import wpfloat
@@ -66,6 +67,9 @@
 # TODO (magdalena) not a constant, this is a default config parameter
 DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0
 
-#: Klemp (2008) type Rayleigh damping
-# TODO (magdalena) not a constant, move somewhere else, convert to enum
-RAYLEIGH_KLEMP: Final[int] = 2
+
+class RayleighType(IntEnum):
+    RAYLEIGH_CLASSIC: Final[
+        int
+    ] = 1  # classical Rayleigh damping, which makes use of a reference state.
+    RAYLEIGH_KLEMP: Final[int] = 2  # Klemp (2008) type Rayleigh damping

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

@@ -11,15 +11,27 @@
 #
 # 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,
+    exp,
+    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
 
 
 """
@@ -95,7 +107,7 @@ def compute_ddqz_z_half(
         z_mc: geometric height on full levels
         k: vertical dimension index
         nlev: total number of levels
-        ddqz_z_half: (output)
+        ddqz_z_half: (output) functional determinant of the metrics (is positive), half levels
         horizontal_start: horizontal start index
         horizontal_end: horizontal end index
         vertical_start: vertical start index
@@ -138,8 +150,8 @@ def compute_ddqz_z_full(
 
     Args:
         z_ifc: geometric height on half levels
-        ddqz_z_full: (output)
-        inv_ddqz_z_full: (output)
+        ddqz_z_full: (output) functional determinant of the metrics (is positive), full levels
+        inv_ddqz_z_full: (output) inverse layer thickness (for runtime optimization)
         horizontal_start: horizontal start index
         horizontal_end: horizontal end index
         vertical_start: vertical start index
@@ -151,3 +163,278 @@ 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]))  # depends on nshift_total, assumed to be always 0
+        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 scaling factor for 3D divergence damping terms.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        vct_a: Field[[KDim], float],
+        scalfac_dd3d: (output) scaling factor for 3D divergence damping terms, and start level from which they are > 0
+        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,
+    vertical_end: int32,
+):
+    """
+    Compute rayleigh_w factor.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        rayleigh_w: (output) Rayleigh damping
+        vct_a: Field[[KDim], float]
+        vct_a_1: 1D of vct_a
+        damping_height: height at which w-damping and sponge layer start
+        rayleigh_type: type of Rayleigh damping (1: CLASSIC, 2: Klemp (2008))
+        rayleigh_classic: classical Rayleigh damping, which makes use of a reference state.
+        rayleigh_klemp: Klemp (2008) type Rayleigh damping
+        rayleigh_coeff: Rayleigh damping coefficient in w-equation
+        pi_const: pi constant
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+    _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 coeff1_dwdz and coeff2_dwdz factors.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        ddqz_z_full: functional determinant of the metrics (is positive), full levels
+        z_ifc: geometric height of half levels
+        coeff1_dwdz: coefficient for second-order acurate dw/dz term
+        coeff2_dwdz: coefficient for second-order acurate dw/dz term
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+
+    _compute_coeff_dwdz(
+        ddqz_z_full,
+        z_ifc,
+        out=(coeff1_dwdz, coeff2_dwdz),
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_d2dexdz2_fac1_mc(
+    theta_ref_mc: Field[[CellDim, KDim], vpfloat],
+    inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat],
+    d2dexdz2_fac1_mc: Field[[CellDim, KDim], vpfloat],
+    cpd: float,
+    grav: wpfloat,
+    igradp_method: int32,
+) -> Field[[CellDim, KDim], vpfloat]:
+    if igradp_method <= int32(3):
+        d2dexdz2_fac1_mc = -grav / (cpd * theta_ref_mc**2) * inv_ddqz_z_full
+
+    return d2dexdz2_fac1_mc
+
+
+@field_operator
+def _compute_d2dexdz2_fac2_mc(
+    theta_ref_mc: Field[[CellDim, KDim], vpfloat],
+    exner_ref_mc: Field[[CellDim, KDim], vpfloat],
+    z_mc: Field[[CellDim, KDim], wpfloat],
+    d2dexdz2_fac2_mc: Field[[CellDim, KDim], vpfloat],
+    cpd: float,
+    grav: wpfloat,
+    del_t_bg: wpfloat,
+    h_scal_bg: wpfloat,
+    igradp_method: int32,
+) -> Field[[CellDim, KDim], vpfloat]:
+    if igradp_method <= int32(3):
+        d2dexdz2_fac2_mc = (
+            2.0
+            * grav
+            / (cpd * theta_ref_mc**3)
+            * (grav / cpd - del_t_bg / h_scal_bg * exp(-z_mc / h_scal_bg))
+            / exner_ref_mc
+        )
+    return d2dexdz2_fac2_mc
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def compute_d2dexdz2_fac_mc(
+    theta_ref_mc: Field[[CellDim, KDim], vpfloat],
+    inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat],
+    exner_ref_mc: Field[[CellDim, KDim], vpfloat],
+    z_mc: Field[[CellDim, KDim], wpfloat],
+    d2dexdz2_fac1_mc: Field[[CellDim, KDim], vpfloat],
+    d2dexdz2_fac2_mc: Field[[CellDim, KDim], vpfloat],
+    cpd: float,
+    grav: wpfloat,
+    del_t_bg: wpfloat,
+    h_scal_bg: wpfloat,
+    igradp_method: int32,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute d2dexdz2_fac1_mc and d2dexdz2_fac2_mc factors.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        theta_ref_mc: reference Potential temperature, full level mass points
+        inv_ddqz_z_full: inverse layer thickness (for runtime optimization)
+        exner_ref_mc: reference Exner pressure, full level mass points
+        z_mc: geometric height defined on full levels
+        d2dexdz2_fac1_mc: (output) first vertical derivative of reference Exner pressure, full level mass points, divided by theta_ref
+        d2dexdz2_fac2_mc: (output) vertical derivative of d_exner_dz/theta_ref, full level mass points
+        cpd: Specific heat at constant pressure [J/K/kg]
+        grav: avergae gravitational acceleratio
+        del_t_bg: difference between sea level temperature and asymptotic stratospheric temperature
+        h_scal_bg: height scale for reference atmosphere [m]
+        igradp_method: method for computing the horizontal presure gradient
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+
+    _compute_d2dexdz2_fac1_mc(
+        theta_ref_mc,
+        inv_ddqz_z_full,
+        d2dexdz2_fac1_mc,
+        cpd,
+        grav,
+        igradp_method,
+        out=d2dexdz2_fac1_mc,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+    _compute_d2dexdz2_fac2_mc(
+        theta_ref_mc,
+        exner_ref_mc,
+        z_mc,
+        d2dexdz2_fac2_mc,
+        cpd,
+        grav,
+        del_t_bg,
+        h_scal_bg,
+        igradp_method,
+        out=d2dexdz2_fac2_mc,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )