Verify global model

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py

@@ -14,7 +14,6 @@
 import logging
 import math
 import sys
-from collections import namedtuple
 from dataclasses import InitVar, dataclass, field
 from enum import Enum
 from typing import Final, Optional
@@ -43,8 +42,8 @@
     zero_field,
 )
 from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn
-from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import (
-    apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance,
+from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence import (
+    apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence,
 )
 from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostic_quantities_for_turbulence import (
     calculate_diagnostic_quantities_for_turbulence,
@@ -81,11 +80,15 @@
 from icon4py.model.common.states.prognostic_state import PrognosticState
 
 
+"""
+Diffusion module ported from ICON mo_nh_diffusion.f90.
+
+Supports only diffusion_type (=hdiff_order) 5 from the diffusion namelist.
+"""
+
 # flake8: noqa
 log = logging.getLogger(__name__)
 
-VectorTuple = namedtuple("VectorTuple", "x y")
-
 cached_backend = run_gtfn_cached
 compiled_backend = run_gtfn
 imperative_backend = run_gtfn_imperative
@@ -95,6 +98,7 @@
 class DiffusionType(int, Enum):
     """
     Order of nabla operator for diffusion.
+
     Note: Called `hdiff_order` in `mo_diffusion_nml.f90`.
     Note: We currently only support type 5.
     """
@@ -106,6 +110,23 @@ class DiffusionType(int, Enum):
     SMAGORINSKY_4TH_ORDER = 5  #: Smagorinsky diffusion with fourth-order background diffusion
 
 
+class TurbulenceShearForcingType(int, Enum):
+    """
+    Type of shear forcing used in turbulance.
+
+    Note: called `itype_sher` in `mo_turbdiff_nml.f90`
+    """
+
+    VERTICAL_OF_HORIZONTAL_WIND = 0  #: only vertical shear of horizontal wind
+    VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND = (
+        1  #: as `VERTICAL_ONLY` plus horizontal shar correction
+    )
+    VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND = (
+        2  #: as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` plus shear form vertical velocity
+    )
+    VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND_LTHESH = 3  #: same as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` but scaling of coarse-grid horizontal shear production term with 1/sqrt(Ri) (if LTKESH = TRUE)
+
+
 class DiffusionConfig:
     """
     Contains necessary parameter to configure a diffusion run.
@@ -132,11 +153,14 @@ def __init__(
         smagorinski_scaling_factor: float = 0.015,
         n_substeps: int = 5,
         zdiffu_t: bool = True,
+        thslp_zdiffu: float = 0.025,
+        thhgtd_zdiffu: float = 200.0,
         hdiff_rcf: bool = True,
         velocity_boundary_diffusion_denom: float = 200.0,
         temperature_boundary_diffusion_denom: float = 135.0,
         max_nudging_coeff: float = 0.02,
         nudging_decay_rate: float = 2.0,
+        shear_type: TurbulenceShearForcingType = TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND,
     ):
         """Set the diffusion configuration parameters with the ICON default values."""
         # parameters from namelist diffusion_nml
@@ -176,13 +200,18 @@ def __init__(
         self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio
 
         #: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below
-        #: Called `hdiff_smag_fac` inmo_diffusion_nml.f90
+        #: Called `hdiff_smag_fac` in mo_diffusion_nml.f90
         self.smagorinski_scaling_factor: float = smagorinski_scaling_factor
 
         #: If True, apply truly horizontal temperature diffusion over steep slopes
         #: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90
         self.apply_zdiffusion_t: bool = zdiffu_t
 
+        #:slope threshold (temperature diffusion): is used to build up an index list for application of truly horizontal diffusion in mo_vertical_grid.f89
+        self.thslp_zdiffu = thslp_zdiffu
+        #: threshold [m] for height difference between adjacent grid points, defaults to 200m (temperature diffusion)
+        self.thhgtd_zdiffu = thhgtd_zdiffu
+
         # from other namelists:
         # from parent namelist mo_nonhydrostatic_nml
 
@@ -219,6 +248,10 @@ def __init__(
         #: Called `nudge_efold_width` in mo_interpol_nml.f90
         self.nudge_efold_width: float = nudging_decay_rate
 
+        #: Type of shear forcing used in turbulence
+        #: Called itype_shear in `mo_turbdiff_nml.f90
+        self.shear_type = shear_type
+
         self._validate()
 
     def _validate(self):
@@ -237,8 +270,15 @@ def _validate(self):
             self.apply_to_temperature = True
             self.apply_to_horizontal_wind = True
 
-        if not self.apply_zdiffusion_t:
-            raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)")
+        if self.shear_type not in (
+            TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND,
+            TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND,
+        ):
+            raise NotImplementedError(
+                f"Turbulence Shear only {TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND} "
+                f"and {TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND} "
+                f"implemented"
+            )
 
     @functools.cached_property
     def substep_as_float(self):
@@ -618,7 +658,7 @@ def _do_diffusion_step(
 
         # 2.  HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
         log.debug("communication rbf extrapolation of vn - start")
-        h = self._exchange.exchange_and_wait(VertexDim, self.u_vert, self.v_vert)
+        self._exchange.exchange_and_wait(VertexDim, self.u_vert, self.v_vert)
         log.debug("communication rbf extrapolation of vn - end")
 
         log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start")
@@ -649,30 +689,38 @@ def _do_diffusion_step(
             },
         )
         log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end")
-        log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start")
-        calculate_diagnostic_quantities_for_turbulence.with_backend(backend)(
-            kh_smag_ec=self.kh_smag_ec,
-            vn=prognostic_state.vn,
-            e_bln_c_s=self.interpolation_state.e_bln_c_s,
-            geofac_div=self.interpolation_state.geofac_div,
-            diff_multfac_smag=self.diff_multfac_smag,
-            wgtfac_c=self.metric_state.wgtfac_c,
-            div_ic=diagnostic_state.div_ic,
-            hdef_ic=diagnostic_state.hdef_ic,
-            horizontal_start=cell_start_nudging,
-            horizontal_end=cell_end_local,
-            vertical_start=1,
-            vertical_end=klevels,
-            offset_provider={
-                "C2E": self.grid.get_offset_provider("C2E"),
-                "C2CE": self.grid.get_offset_provider("C2CE"),
-                "Koff": KDim,
-            },
-        )
-        log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end")
-
-        # HALO EXCHANGE  IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH
+        if (
+            self.config.shear_type
+            >= TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND
+        ):
+            log.debug(
+                "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start"
+            )
+            calculate_diagnostic_quantities_for_turbulence.with_backend(backend)(
+                kh_smag_ec=self.kh_smag_ec,
+                vn=prognostic_state.vn,
+                e_bln_c_s=self.interpolation_state.e_bln_c_s,
+                geofac_div=self.interpolation_state.geofac_div,
+                diff_multfac_smag=self.diff_multfac_smag,
+                wgtfac_c=self.metric_state.wgtfac_c,
+                div_ic=diagnostic_state.div_ic,
+                hdef_ic=diagnostic_state.hdef_ic,
+                horizontal_start=cell_start_nudging,
+                horizontal_end=cell_end_local,
+                vertical_start=1,
+                vertical_end=klevels,
+                offset_provider={
+                    "C2E": self.grid.get_offset_provider("C2E"),
+                    "C2CE": self.grid.get_offset_provider("C2CE"),
+                    "Koff": KDim,
+                },
+            )
+            log.debug(
+                "running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end"
+            )
 
+        # HALO EXCHANGE  IF (discr_vn > 1) THEN CALL sync_patch_array
+        # TODO (magdalena) move this up and do asynchronous exchange
         if self.config.type_vn_diffu > 1:
             log.debug("communication rbf extrapolation of z_nable2_e - start")
             self._exchange.exchange_and_wait(EdgeDim, self.z_nabla2_e)
@@ -731,17 +779,18 @@ def _do_diffusion_step(
         handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn)
 
         log.debug(
-            "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start"
+            "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start"
         )
         # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify?
         copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={})
-        apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)(
+        apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_backend(backend)(
             area=self.cell_params.area,
             geofac_n2s=self.interpolation_state.geofac_n2s,
             geofac_grg_x=self.interpolation_state.geofac_grg_x,
             geofac_grg_y=self.interpolation_state.geofac_grg_y,
             w_old=self.w_tmp,
             w=prognostic_state.w,
+            type_shear=int32(self.config.shear_type.value),
             dwdx=diagnostic_state.dwdx,
             dwdy=diagnostic_state.dwdy,
             diff_multfac_w=self.diff_multfac_w,
@@ -762,7 +811,7 @@ def _do_diffusion_step(
             },
         )
         log.debug(
-            "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end"
+            "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): end"
         )
 
         log.debug(
@@ -807,29 +856,30 @@ def _do_diffusion_step(
         log.debug(
             "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start"
         )
-        truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)(
-            mask=self.metric_state.mask_hdiff,
-            zd_vertoffset=self.metric_state.zd_vertoffset,
-            zd_diffcoef=self.metric_state.zd_diffcoef,
-            geofac_n2s_c=self.interpolation_state.geofac_n2s_c,
-            geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh,
-            vcoef=self.metric_state.zd_intcoef,
-            theta_v=prognostic_state.theta_v,
-            z_temp=self.z_temp,
-            horizontal_start=cell_start_nudging,
-            horizontal_end=cell_end_local,
-            vertical_start=0,
-            vertical_end=klevels,
-            offset_provider={
-                "C2CEC": self.grid.get_offset_provider("C2CEC"),
-                "C2E2C": self.grid.get_offset_provider("C2E2C"),
-                "Koff": KDim,
-            },
-        )
+        if self.config.apply_zdiffusion_t:
+            truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_backend(backend)(
+                mask=self.metric_state.mask_hdiff,
+                zd_vertoffset=self.metric_state.zd_vertoffset,
+                zd_diffcoef=self.metric_state.zd_diffcoef,
+                geofac_n2s_c=self.interpolation_state.geofac_n2s_c,
+                geofac_n2s_nbh=self.interpolation_state.geofac_n2s_nbh,
+                vcoef=self.metric_state.zd_intcoef,
+                theta_v=prognostic_state.theta_v,
+                z_temp=self.z_temp,
+                horizontal_start=cell_start_nudging,
+                horizontal_end=cell_end_local,
+                vertical_start=0,
+                vertical_end=klevels,
+                offset_provider={
+                    "C2CEC": self.grid.get_offset_provider("C2CEC"),
+                    "C2E2C": self.grid.get_offset_provider("C2E2C"),
+                    "Koff": KDim,
+                },
+            )
 
-        log.debug(
-            "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end"
-        )
+            log.debug(
+                "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end"
+            )
         log.debug("running stencil 16 (update_theta_and_exner): start")
         update_theta_and_exner.with_backend(backend)(
             z_temp=self.z_temp,

model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py → model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py

@@ -29,12 +29,13 @@
 
 
 @field_operator
-def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
+def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
     area: Field[[CellDim], wpfloat],
     geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat],
     geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat],
     geofac_grg_y: Field[[CellDim, C2E2CODim], wpfloat],
     w_old: Field[[CellDim, KDim], wpfloat],
+    type_shear: int32,
     dwdx: Field[[CellDim, KDim], vpfloat],
     dwdy: Field[[CellDim, KDim], vpfloat],
     diff_multfac_w: wpfloat,
@@ -50,11 +51,14 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
     Field[[CellDim, KDim], vpfloat],
 ]:
     vert_idx = broadcast(vert_idx, (CellDim, KDim))
-
-    dwdx, dwdy = where(
-        int32(0) < vert_idx,
-        _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y),
-        (dwdx, dwdy),
+    dwdx, dwdy = (
+        where(
+            int32(0) < vert_idx,
+            _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y),
+            (dwdx, dwdy),
+        )
+        if type_shear == int32(2)
+        else (dwdx, dwdy)
     )
 
     z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s)
@@ -78,13 +82,14 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
 
 
 @program
-def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
+def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
     area: Field[[CellDim], wpfloat],
     geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat],
     geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat],
     geofac_grg_y: Field[[CellDim, C2E2CODim], wpfloat],
     w_old: Field[[CellDim, KDim], wpfloat],
     w: Field[[CellDim, KDim], wpfloat],
+    type_shear: int32,
     dwdx: Field[[CellDim, KDim], vpfloat],
     dwdy: Field[[CellDim, KDim], vpfloat],
     diff_multfac_w: wpfloat,
@@ -99,12 +104,13 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
     vertical_start: int32,
     vertical_end: int32,
 ):
-    _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
+    _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
         area,
         geofac_n2s,
         geofac_grg_x,
         geofac_grg_y,
         w_old,
+        type_shear,
         dwdx,
         dwdy,
         diff_multfac_w,