Diffusion metrics fields

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

@@ -32,6 +32,7 @@
 
 from icon4py.model.common.dimension import (
     C2E,
+    C2E2C,
     E2C,
     CellDim,
     E2CDim,
@@ -564,6 +565,40 @@ def _compute_maxslp_maxhgtd(
     return z_maxslp, z_maxhgtd
 
 
+@program
+def compute_maxslp_maxhgtd(
+    ddxn_z_full: Field[[EdgeDim, KDim], wpfloat],
+    dual_edge_length: Field[[EdgeDim], wpfloat],
+    z_maxslp: Field[[CellDim, KDim], wpfloat],
+    z_maxhgtd: Field[[CellDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute z_maxslp and z_maxhgtd.
+
+    See mo_vertical_grid.f90.
+
+    Args:
+        ddxn_z_full: dual_edge_length
+        dual_edge_length: dual_edge_length
+        z_maxslp: output
+        z_maxhgtd: output
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+    _compute_maxslp_maxhgtd(
+        ddxn_z_full=ddxn_z_full,
+        dual_edge_length=dual_edge_length,
+        out=(z_maxslp, z_maxhgtd),
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+
 @field_operator
 def _compute_exner_exfac(
     ddxn_z_full: Field[[EdgeDim, KDim], wpfloat],
@@ -1021,3 +1056,156 @@ def compute_hmask_dd3d(
         out=hmask_dd3d,
         domain={EdgeDim: (horizontal_start, horizontal_end)},
     )
+
+
+@field_operator
+def _compute_mask_hdiff() -> Field[[CellDim, KDim], bool]:
+    return broadcast(True, (CellDim, KDim))
+
+
+@program
+def compute_mask_hdiff(
+    mask_hdiff: Field[[CellDim, KDim], bool],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute mask_hdiff.
+
+    See mo_vertical_grid.f90.
+
+    Args:
+        mask_hdiff: output
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+    _compute_mask_hdiff(
+        out=mask_hdiff,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_z_maxslp_avg(
+    maxslp: Field[[CellDim, KDim], wpfloat],
+    c_bln_avg_0: Field[[CellDim], wpfloat],
+    c_bln_avg_1: Field[[CellDim], wpfloat],
+    c_bln_avg_2: Field[[CellDim], wpfloat],
+    c_bln_avg_3: Field[[CellDim], wpfloat],
+) -> Field[[CellDim, KDim], wpfloat]:
+    z_maxslp_avg = (
+        maxslp * c_bln_avg_0
+        + maxslp(C2E2C[0]) * c_bln_avg_1
+        + maxslp(C2E2C[1]) * c_bln_avg_2
+        + maxslp(C2E2C[2]) * c_bln_avg_3
+    )
+    return z_maxslp_avg
+
+
+@field_operator
+def _compute_z_maxhgtd_avg(
+    maxhgtd: Field[[CellDim, KDim], wpfloat],
+    c_bln_avg_0: Field[[CellDim], wpfloat],
+    c_bln_avg_1: Field[[CellDim], wpfloat],
+    c_bln_avg_2: Field[[CellDim], wpfloat],
+    c_bln_avg_3: Field[[CellDim], wpfloat],
+) -> Field[[CellDim, KDim], wpfloat]:
+    z_maxhgtd_avg = (
+        maxhgtd * c_bln_avg_0
+        + maxhgtd(C2E2C[0]) * c_bln_avg_1
+        + maxhgtd(C2E2C[1]) * c_bln_avg_2
+        + maxhgtd(C2E2C[2]) * c_bln_avg_3
+    )
+    return z_maxhgtd_avg
+
+
+@program
+def compute_z_maxslp_avg_z_maxhgtd_avg(
+    maxslp: Field[[CellDim, KDim], wpfloat],
+    maxhgtd: Field[[CellDim, KDim], wpfloat],
+    c_bln_avg_0: Field[[CellDim], wpfloat],
+    c_bln_avg_1: Field[[CellDim], wpfloat],
+    c_bln_avg_2: Field[[CellDim], wpfloat],
+    c_bln_avg_3: Field[[CellDim], wpfloat],
+    z_maxslp_avg: Field[[CellDim, KDim], wpfloat],
+    z_maxhgtd_avg: Field[[CellDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute z_maxslp_avg and z_maxhgtd_avg.
+
+    See mo_vertical_grid.f90.
+
+    Args:
+        maxslp: Max field over ddxn_z_full offset
+        maxhgtd: Max field over ddxn_z_full offset*dual_edge_length offset
+        c_bln_avg_0: Interpolation field 0th K element
+        c_bln_avg_1: Interpolation field 1st K element
+        c_bln_avg_2: Interpolation field 2nd K element
+        c_bln_avg_3: Interpolation field 3rd K element
+        z_maxslp_avg: output
+        z_maxhgtd_avg: output
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+
+    _compute_z_maxslp_avg(
+        maxslp=maxslp,
+        c_bln_avg_0=c_bln_avg_0,
+        c_bln_avg_1=c_bln_avg_1,
+        c_bln_avg_2=c_bln_avg_2,
+        c_bln_avg_3=c_bln_avg_3,
+        out=z_maxslp_avg,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+    _compute_z_maxhgtd_avg(
+        maxhgtd=maxhgtd,
+        c_bln_avg_0=c_bln_avg_0,
+        c_bln_avg_1=c_bln_avg_1,
+        c_bln_avg_2=c_bln_avg_2,
+        c_bln_avg_3=c_bln_avg_3,
+        out=z_maxhgtd_avg,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_max_nbhgt(
+    z_mc_nlev: Field[[CellDim], wpfloat],
+) -> Field[[CellDim], wpfloat]:
+    max_nbhgt_0_1 = maximum(z_mc_nlev(C2E2C[0]), z_mc_nlev(C2E2C[1]))
+    max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2]))
+    return max_nbhgt
+
+
+@program
+def compute_max_nbhgt(
+    z_mc_nlev: Field[[CellDim], wpfloat],
+    max_nbhgt: Field[[CellDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+):
+    """
+    Compute max_nbhgt.
+
+    See mo_vertical_grid.f90.
+
+    Args:
+        z_mc_nlev: Last K level of z_mc
+        max_nbhgt: output
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+    """
+    _compute_max_nbhgt(
+        z_mc_nlev=z_mc_nlev, out=max_nbhgt, domain={CellDim: (horizontal_start, horizontal_end)}
+    )

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

@@ -0,0 +1,175 @@
+# 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_nbidx(
+    k_range: range,
+    z_mc: np.array,
+    z_mc_off: np.array,
+    nbidx: np.array,
+    jc: int,
+    nlev: int,
+) -> np.array:
+    for ind in range(3):
+        jk_start = nlev - 1
+        for jk in reversed(k_range):
+            for jk1 in reversed(range(jk_start)):
+                if (
+                    z_mc[jc, jk] <= z_mc_off[jc, ind, jk1]
+                    and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1]
+                ):
+                    nbidx[jc, ind, jk] = jk1
+                    jk_start = jk1 + 1
+                    break
+
+    return nbidx[jc, :, :]
+
+
+def _compute_z_vintcoeff(
+    k_range: range,
+    z_mc: np.array,
+    z_mc_off: np.array,
+    z_vintcoeff: np.array,
+    jc: int,
+    nlev: int,
+) -> np.array:
+    for ind in range(3):
+        jk_start = nlev - 1
+        for jk in reversed(k_range):
+            for jk1 in reversed(range(jk_start)):
+                if (
+                    z_mc[jc, jk] <= z_mc_off[jc, ind, jk1]
+                    and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1]
+                ):
+                    z_vintcoeff[jc, ind, jk] = (z_mc[jc, jk] - z_mc_off[jc, ind, jk1 + 1]) / (
+                        z_mc_off[jc, ind, jk1] - z_mc_off[jc, ind, jk1 + 1]
+                    )
+                    jk_start = jk1 + 1
+                    break
+
+    return z_vintcoeff[jc, :, :]
+
+
+def _compute_i_params(
+    k_start: list,
+    k_end: list,
+    z_maxslp_avg: np.array,
+    z_maxhgtd_avg: np.array,
+    c_owner_mask: np.array,
+    thslp_zdiffu: float,
+    thhgtd_zdiffu: float,
+    cell_nudging: int,
+    n_cells: int,
+    nlev: int,
+) -> tuple[list, int, int]:
+    i_indlist = [0] * n_cells
+    i_listreduce = 0
+    ji = -1
+    ji_ind = -1
+
+    for jc in range(cell_nudging, n_cells):
+        if (
+            z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu
+            or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu
+        ) and c_owner_mask[jc]:
+            ji += 1
+            i_indlist[ji] = jc
+
+            if all((k_start[jc], k_end[jc])) and k_start[jc] > k_end[jc]:
+                i_listreduce += 1
+            else:
+                ji_ind += 1
+                i_indlist[ji_ind] = jc
+
+    return i_indlist, i_listreduce, ji
+
+
+def _compute_k_start_end(
+    z_mc: np.array,
+    max_nbhgt: np.array,
+    z_maxslp_avg: np.array,
+    z_maxhgtd_avg: np.array,
+    c_owner_mask: np.array,
+    thslp_zdiffu: float,
+    thhgtd_zdiffu: float,
+    cell_nudging: int,
+    n_cells: int,
+    nlev: int,
+) -> tuple[list, list]:
+    k_start = [None] * n_cells
+    k_end = [None] * n_cells
+    for jc in range(cell_nudging, n_cells):
+        if (
+            z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu
+            or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu
+        ) and c_owner_mask[jc]:
+            for jk in reversed(range(nlev)):
+                if z_mc[jc, jk] >= max_nbhgt[jc]:
+                    k_end[jc] = jk + 1
+                    break
+
+            for jk in range(nlev):
+                if z_maxslp_avg[jc, jk] >= thslp_zdiffu or z_maxhgtd_avg[jc, jk] >= thhgtd_zdiffu:
+                    k_start[jc] = jk
+                    break
+
+            if all((k_start[jc], k_end[jc])) and k_start[jc] > k_end[jc]:
+                k_start[jc] = nlev - 1
+
+    return k_start, k_end
+
+
+def compute_diffusion_metrics(
+    z_mc: np.array,
+    z_mc_off: np.array,
+    k_start: list,
+    k_end: list,
+    i_indlist: list,
+    i_listdim: int,
+    nbidx: np.array,
+    z_vintcoeff: np.array,
+    z_maxslp_avg: np.array,
+    z_maxhgtd_avg: np.array,
+    mask_hdiff: np.array,
+    zd_diffcoef_dsl: np.array,
+    zd_intcoef_dsl: np.array,
+    zd_vertoffset_dsl: np.array,
+    thslp_zdiffu: float,
+    thhgtd_zdiffu: float,
+    nlev: int,
+) -> np.array:
+    for ji in range(i_listdim):
+        jc = i_indlist[ji]
+        k_range = range(k_start[jc], k_end[jc])
+        if all((k_range)):
+            nbidx[jc, :, :] = _compute_nbidx(k_range, z_mc, z_mc_off, nbidx, jc, nlev)
+            z_vintcoeff[jc, :, :] = _compute_z_vintcoeff(
+                k_range, z_mc, z_mc_off, z_vintcoeff, jc, nlev
+            )
+
+            zd_intcoef_dsl[jc, :, k_range] = z_vintcoeff[jc, :, k_range]
+            zd_vertoffset_dsl[jc, :, k_range] = nbidx[jc, :, k_range] - np.transpose([k_range] * 3)
+            mask_hdiff[jc, k_range] = True
+
+            zd_diffcoef_dsl_var = np.maximum(
+                0.0,
+                np.maximum(
+                    np.sqrt(np.maximum(0.0, z_maxslp_avg[jc, k_range] - thslp_zdiffu)) / 250.0,
+                    2.0e-4 * np.sqrt(np.maximum(0.0, z_maxhgtd_avg[jc, k_range] - thhgtd_zdiffu)),
+                ),
+            )
+            zd_diffcoef_dsl[jc, k_range] = np.minimum(0.002, zd_diffcoef_dsl_var)
+
+    return mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl

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

@@ -484,6 +484,9 @@ def construct_cell_geometry(self) -> CellParams:
 
 
 class InterpolationSavepoint(IconSavepoint):
+    def c_bln_avg(self):
+        return self._get_field("c_bln_avg", CellDim, KDim)
+
     def c_intp(self):
         return self._get_field("c_intp", VertexDim, V2CDim)
 
@@ -692,6 +695,7 @@ def geopot(self):
 
     def _read_and_reorder_sparse_field(self, name: str, sparse_size=3):
         ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[:, :, :]
+        ser_input = self._reduce_to_dim_size(ser_input, (CellDim, C2E2CDim, KDim))
         if ser_input.shape[1] != sparse_size:
             ser_input = np.moveaxis(ser_input, 1, -1)