Added smagorinsky factors and heights

Author: muellch

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

@@ -149,6 +149,16 @@ def __init__(
         hdiff_efdt_ratio: float = 36.0,
         hdiff_w_efdt_ratio: float = 15.0,
         smagorinski_scaling_factor: float = 0.015,
+        smagorinski_scaling_factor2: float = 2e-6
+        * (1600.0 + 25000.0 + math.sqrt(1600.0 * (1600 + 50000.0))),
+        smagorinski_scaling_factor3: float = 0,
+        smagorinski_scaling_factor4: float = 1,
+        smagorinski_scaling_height: float = 32500.0,
+        smagorinski_scaling_height2: float = 1600.0
+        + 50000.0
+        + math.sqrt(1600.0 * (1600 + 50000.0)),
+        smagorinski_scaling_height3: float = 50000.0,
+        smagorinski_scaling_height4: float = 90000.0,
         n_substeps: int = 5,
         zdiffu_t: bool = True,
         velocity_boundary_diffusion_denom: float = 200.0,
@@ -199,10 +209,41 @@ def __init__(
         #: Called 'hdiff_w_efdt_ratio' in mo_diffusion_nml.f90.
         self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio
 
-        #: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below
+        #: Scaling factor for Smagorinsky diffusion at height smagorinski_scaling_height and below
         #: Called 'hdiff_smag_fac' in mo_diffusion_nml.f90
         self.smagorinski_scaling_factor: float = smagorinski_scaling_factor
 
+        #: Scaling factor for Smagorinsky diffusion at height smagorinski_scaling_height to smagorinski_scaling_height2
+        #: This is the linearly interpolated part
+        #: Called 'hdiff_smag_fac2' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_factor2: float = smagorinski_scaling_factor2
+
+        #: Scaling factor for Smagorinsky diffusion at height smagorinski_scaling_height2 to smagorinski_scaling_height3
+        #: This is the quadratically interpolated part
+        #: Called 'hdiff_smag_fac3' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_factor3: float = smagorinski_scaling_factor3
+
+        #: Scaling factor for Smagorinsky diffusion at height smagorinski_scaling_height3 to smagorinski_scaling_height4
+        #: This is the quadratically interpolated part
+        #: Called 'hdiff_smag_fac4' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_factor4: float = smagorinski_scaling_factor4
+
+        #: Height for Smagorinsky diffusion expressing the border between constant and linear interpolation
+        #: Called 'hdiff_smag_z' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_height: float = smagorinski_scaling_height
+
+        #: Height for Smagorinsky diffusion expressing the border between linear and quadratic interpolation
+        #: Called 'hdiff_smag_z2' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_height2: float = smagorinski_scaling_height2
+
+        #: Height for Smagorinsky diffusion expressing the middle of the quadratic interpolation region
+        #: Called 'hdiff_smag_z3' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_height3: float = smagorinski_scaling_height3
+
+        #: Height for Smagorinsky diffusion expressing the top of the quadratic interpolation region
+        #: Called 'hdiff_smag_z4' in mo_diffusion_nml.f90
+        self.smagorinski_scaling_height4: float = smagorinski_scaling_height4
+
         #: 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
@@ -329,53 +370,26 @@ def __post_init__(self, config):
             (1.0 / (config.hdiff_w_efdt_ratio * 36.0) if config.hdiff_w_efdt_ratio > 0 else 0.0),
         )
 
-        (
-            smagorinski_factor,
-            smagorinski_height,
-        ) = self._determine_smagorinski_factor(config)
-        object.__setattr__(self, "smagorinski_factor", smagorinski_factor)
-        object.__setattr__(self, "smagorinski_height", smagorinski_height)
-
-    def _determine_smagorinski_factor(self, config: DiffusionConfig):
-        """Enhanced Smagorinsky diffusion factor.
-
-        Smagorinsky diffusion factor is defined as a profile in height
-        above sea level with 4 height sections.
-
-        It is calculated/used only in the case of diffusion_type 3 or 5
-        """
-        match config.diffusion_type:
-            case 5:
-                (
-                    smagorinski_factor,
-                    smagorinski_height,
-                ) = diffusion_type_5_smagorinski_factor(config)
-            case 4:
-                # according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only
-                # used for diffusion_type (3,5) but the defaults are only defined for iequations=3
-                smagorinski_factor = (
-                    config.smagorinski_scaling_factor
-                    if config.smagorinski_scaling_factor
-                    else 0.15,
-                )
-                smagorinski_height = None
-            case _:
-                raise NotImplementedError("Only implemented for diffusion type 4 and 5")
-        return smagorinski_factor, smagorinski_height
-
-
-def diffusion_type_5_smagorinski_factor(config: DiffusionConfig):
-    """
-    Initialize Smagorinski factors used in diffusion type 5.
-
-    The calculation and magic numbers are taken from mo_diffusion_nml.f90
-    """
-    magic_sqrt = math.sqrt(1600.0 * (1600 + 50000.0))
-    magic_fac2_value = 2e-6 * (1600.0 + 25000.0 + magic_sqrt)
-    magic_z2 = 1600.0 + 50000.0 + magic_sqrt
-    factor = (config.smagorinski_scaling_factor, magic_fac2_value, 0.0, 1.0)
-    heights = (32500.0, magic_z2, 50000.0, 90000.0)
-    return factor, heights
+        object.__setattr__(
+            self,
+            "smagorinski_factor",
+            (
+                config.smagorinski_scaling_factor,
+                config.smagorinski_scaling_factor2,
+                config.smagorinski_scaling_factor3,
+                config.smagorinski_scaling_factor4,
+            ),
+        )
+        object.__setattr__(
+            self,
+            "smagorinski_height",
+            (
+                config.smagorinski_scaling_height,
+                config.smagorinski_scaling_height2,
+                config.smagorinski_scaling_height3,
+                config.smagorinski_scaling_height4,
+            ),
+        )
 
 
 class Diffusion:

model/atmosphere/diffusion/tests/diffusion/integration_tests/test_diffusion.py

@@ -120,16 +120,6 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(experiment):
     assert params.K4W == 0.0
 
 
-def test_smagorinski_factor_for_diffusion_type_4(experiment):
-    config = definitions.construct_diffusion_config(experiment, ndyn_substeps=5)
-    config.smagorinski_scaling_factor = 0.15
-    config.diffusion_type = 4
-
-    params = diffusion.DiffusionParams(config)
-    assert len(params.smagorinski_factor) == 1
-    assert params.smagorinski_factor[0] == pytest.approx(0.15, abs=1e-16)
-    assert params.smagorinski_height is None
-
 
 def test_smagorinski_heights_diffusion_type_5_are_consistent(
     experiment,

tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py

@@ -84,6 +84,13 @@ def diffusion_init(
     hdiff_efdt_ratio: gtx.float64,
     hdiff_w_efdt_ratio: gtx.float64,
     smagorinski_scaling_factor: gtx.float64,
+    smagorinski_scaling_factor2: gtx.float64,
+    smagorinski_scaling_factor3: gtx.float64,
+    smagorinski_scaling_factor4: gtx.float64,
+    smagorinski_scaling_height: gtx.float64,
+    smagorinski_scaling_height2: gtx.float64,
+    smagorinski_scaling_height3: gtx.float64,
+    smagorinski_scaling_height4: gtx.float64,
     hdiff_temp: bool,
     denom_diffu_v: float,
     nudge_max_coeff: float,  # note: this is the scaled ICON value, i.e. not the namelist value
@@ -117,6 +124,13 @@ def diffusion_init(
         hdiff_efdt_ratio=hdiff_efdt_ratio,
         hdiff_w_efdt_ratio=hdiff_w_efdt_ratio,
         smagorinski_scaling_factor=smagorinski_scaling_factor,
+        smagorinski_scaling_factor2=smagorinski_scaling_factor2,
+        smagorinski_scaling_factor3=smagorinski_scaling_factor3,
+        smagorinski_scaling_factor4=smagorinski_scaling_factor4,
+        smagorinski_scaling_height=smagorinski_scaling_height,
+        smagorinski_scaling_height2=smagorinski_scaling_height2,
+        smagorinski_scaling_height3=smagorinski_scaling_height3,
+        smagorinski_scaling_height4=smagorinski_scaling_height4,
         hdiff_temp=hdiff_temp,
         n_substeps=ndyn_substeps,
         velocity_boundary_diffusion_denom=denom_diffu_v,

tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90

@@ -152,6 +152,13 @@ function diffusion_init_wrapper(theta_ref_mc, &
                                       hdiff_efdt_ratio, &
                                       hdiff_w_efdt_ratio, &
                                       smagorinski_scaling_factor, &
+                                      smagorinski_scaling_factor2, &
+                                      smagorinski_scaling_factor3, &
+                                      smagorinski_scaling_factor4, &
+                                      smagorinski_scaling_height, &
+                                      smagorinski_scaling_height2, &
+                                      smagorinski_scaling_height3, &
+                                      smagorinski_scaling_height4, &
                                       hdiff_temp, &
                                       denom_diffu_v, &
                                       nudge_max_coeff, &
@@ -260,6 +267,20 @@ function diffusion_init_wrapper(theta_ref_mc, &
 
          real(c_double), value, target :: smagorinski_scaling_factor
 
+         real(c_double), value, target :: smagorinski_scaling_factor2
+
+         real(c_double), value, target :: smagorinski_scaling_factor3
+
+         real(c_double), value, target :: smagorinski_scaling_factor4
+
+         real(c_double), value, target :: smagorinski_scaling_height
+
+         real(c_double), value, target :: smagorinski_scaling_height2
+
+         real(c_double), value, target :: smagorinski_scaling_height3
+
+         real(c_double), value, target :: smagorinski_scaling_height4
+
          logical(c_int), value, target :: hdiff_temp
 
          real(c_double), value, target :: denom_diffu_v
@@ -493,6 +514,13 @@ subroutine diffusion_init(theta_ref_mc, &
                              hdiff_efdt_ratio, &
                              hdiff_w_efdt_ratio, &
                              smagorinski_scaling_factor, &
+                             smagorinski_scaling_factor2, &
+                             smagorinski_scaling_factor3, &
+                             smagorinski_scaling_factor4, &
+                             smagorinski_scaling_height, &
+                             smagorinski_scaling_height2, &
+                             smagorinski_scaling_height3, &
+                             smagorinski_scaling_height4, &
                              hdiff_temp, &
                              denom_diffu_v, &
                              nudge_max_coeff, &
@@ -550,6 +578,20 @@ subroutine diffusion_init(theta_ref_mc, &
 
       real(c_double), value, target :: smagorinski_scaling_factor
 
+      real(c_double), value, target :: smagorinski_scaling_factor2
+
+      real(c_double), value, target :: smagorinski_scaling_factor3
+
+      real(c_double), value, target :: smagorinski_scaling_factor4
+
+      real(c_double), value, target :: smagorinski_scaling_height
+
+      real(c_double), value, target :: smagorinski_scaling_height2
+
+      real(c_double), value, target :: smagorinski_scaling_height3
+
+      real(c_double), value, target :: smagorinski_scaling_height4
+
       logical(c_int), value, target :: hdiff_temp
 
       real(c_double), value, target :: denom_diffu_v
@@ -752,6 +794,13 @@ subroutine diffusion_init(theta_ref_mc, &
                                   hdiff_efdt_ratio=hdiff_efdt_ratio, &
                                   hdiff_w_efdt_ratio=hdiff_w_efdt_ratio, &
                                   smagorinski_scaling_factor=smagorinski_scaling_factor, &
+                                  smagorinski_scaling_factor2=smagorinski_scaling_factor2, &
+                                  smagorinski_scaling_factor3=smagorinski_scaling_factor3, &
+                                  smagorinski_scaling_factor4=smagorinski_scaling_factor4, &
+                                  smagorinski_scaling_height=smagorinski_scaling_height, &
+                                  smagorinski_scaling_height2=smagorinski_scaling_height2, &
+                                  smagorinski_scaling_height3=smagorinski_scaling_height3, &
+                                  smagorinski_scaling_height4=smagorinski_scaling_height4, &
                                   hdiff_temp=hdiff_temp, &
                                   denom_diffu_v=denom_diffu_v, &
                                   nudge_max_coeff=nudge_max_coeff, &

tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h

@@ -23,6 +23,9 @@ extern int diffusion_init_wrapper(
     int ndyn_substeps, int diffusion_type, int hdiff_w, int hdiff_vn,
     int hdiff_smag_w, int zdiffu_t, int type_t_diffu, int type_vn_diffu,
     double hdiff_efdt_ratio, double hdiff_w_efdt_ratio,
-    double smagorinski_scaling_factor, int hdiff_temp, double denom_diffu_v,
-    double nudge_max_coeff, int itype_sher, int ltkeshs, int backend,
-    int on_gpu);
+    double smagorinski_scaling_factor, double smagorinski_scaling_factor2,
+    double smagorinski_scaling_factor3, double smagorinski_scaling_factor4,
+    double smagorinski_scaling_height, double smagorinski_scaling_height2,
+    double smagorinski_scaling_height3, double smagorinski_scaling_height4,
+    int hdiff_temp, double denom_diffu_v, double nudge_max_coeff,
+    int itype_sher, int ltkeshs, int backend, int on_gpu);

tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py

@@ -391,6 +391,13 @@ def diffusion_init_wrapper(
     hdiff_efdt_ratio,
     hdiff_w_efdt_ratio,
     smagorinski_scaling_factor,
+    smagorinski_scaling_factor2,
+    smagorinski_scaling_factor3,
+    smagorinski_scaling_factor4,
+    smagorinski_scaling_height,
+    smagorinski_scaling_height2,
+    smagorinski_scaling_height3,
+    smagorinski_scaling_height4,
     hdiff_temp,
     denom_diffu_v,
     nudge_max_coeff,
@@ -566,6 +573,13 @@ def diffusion_init_wrapper(
                 hdiff_efdt_ratio=hdiff_efdt_ratio,
                 hdiff_w_efdt_ratio=hdiff_w_efdt_ratio,
                 smagorinski_scaling_factor=smagorinski_scaling_factor,
+                smagorinski_scaling_factor2=smagorinski_scaling_factor2,
+                smagorinski_scaling_factor3=smagorinski_scaling_factor3,
+                smagorinski_scaling_factor4=smagorinski_scaling_factor4,
+                smagorinski_scaling_height=smagorinski_scaling_height,
+                smagorinski_scaling_height2=smagorinski_scaling_height2,
+                smagorinski_scaling_height3=smagorinski_scaling_height3,
+                smagorinski_scaling_height4=smagorinski_scaling_height4,
                 hdiff_temp=hdiff_temp,
                 denom_diffu_v=denom_diffu_v,
                 nudge_max_coeff=nudge_max_coeff,

tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py

@@ -6,6 +6,7 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import math
 from unittest import mock
 
 import cffi
@@ -60,6 +61,13 @@ def test_diffusion_wrapper_granule_inputs(
     hdiff_efdt_ratio = 24.0
     hdiff_w_efdt_ratio = 15.0
     smagorinski_scaling_factor = 0.025
+    smagorinski_scaling_factor2 = 2e-6 * (1600.0 + 25000.0 + math.sqrt(1600.0 * (1600 + 50000.0)))
+    smagorinski_scaling_factor3 = 0.0
+    smagorinski_scaling_factor4 = 1.0
+    smagorinski_scaling_height = 32500.0
+    smagorinski_scaling_height2 = 1600.0 + 50000.0 + math.sqrt(1600.0 * (1600 + 50000.0))
+    smagorinski_scaling_height3 = 50000.0
+    smagorinski_scaling_height4 = 90000.0
     zdiffu_t = True
     denom_diffu_v = 150.0
     max_nudging_coefficient = 0.375
@@ -183,6 +191,13 @@ def test_diffusion_wrapper_granule_inputs(
             hdiff_efdt_ratio=hdiff_efdt_ratio,
             hdiff_w_efdt_ratio=hdiff_w_efdt_ratio,
             smagorinski_scaling_factor=smagorinski_scaling_factor,
+            smagorinski_scaling_factor2=smagorinski_scaling_factor2,
+            smagorinski_scaling_factor3=smagorinski_scaling_factor3,
+            smagorinski_scaling_factor4=smagorinski_scaling_factor4,
+            smagorinski_scaling_height=smagorinski_scaling_height,
+            smagorinski_scaling_height2=smagorinski_scaling_height2,
+            smagorinski_scaling_height3=smagorinski_scaling_height3,
+            smagorinski_scaling_height4=smagorinski_scaling_height4,
             hdiff_temp=hdiff_temp,
             denom_diffu_v=denom_diffu_v,
             nudge_max_coeff=max_nudging_coefficient,
@@ -297,6 +312,13 @@ def test_diffusion_wrapper_single_step(
     hdiff_efdt_ratio = 24.0
     hdiff_w_efdt_ratio = 15.0
     smagorinski_scaling_factor = 0.025
+    smagorinski_scaling_factor2 = 2e-6 * (1600.0 + 25000.0 + math.sqrt(1600.0 * (1600 + 50000.0)))
+    smagorinski_scaling_factor3 = 0.0
+    smagorinski_scaling_factor4 = 1.0
+    smagorinski_scaling_height = 32500.0
+    smagorinski_scaling_height2 = 1600.0 + 50000.0 + math.sqrt(1600.0 * (1600 + 50000.0))
+    smagorinski_scaling_height3 = 50000.0
+    smagorinski_scaling_height4 = 90000.0
     zdiffu_t = True
     denom_diffu_v = 150.0
     max_nudging_coefficient = 0.375
@@ -387,6 +409,13 @@ def test_diffusion_wrapper_single_step(
         hdiff_efdt_ratio=hdiff_efdt_ratio,
         hdiff_w_efdt_ratio=hdiff_w_efdt_ratio,
         smagorinski_scaling_factor=smagorinski_scaling_factor,
+        smagorinski_scaling_factor2=smagorinski_scaling_factor2,
+        smagorinski_scaling_factor3=smagorinski_scaling_factor3,
+        smagorinski_scaling_factor4=smagorinski_scaling_factor4,
+        smagorinski_scaling_height=smagorinski_scaling_height,
+        smagorinski_scaling_height2=smagorinski_scaling_height2,
+        smagorinski_scaling_height3=smagorinski_scaling_height3,
+        smagorinski_scaling_height4=smagorinski_scaling_height4,
         hdiff_temp=hdiff_temp,
         denom_diffu_v=denom_diffu_v,
         nudge_max_coeff=max_nudging_coefficient,