Diffusion metrics fields

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

@@ -24,15 +24,20 @@
     int32,
     maximum,
     minimum,
+    neighbor_sum,
     program,
+    scan_operator,
     sin,
     tanh,
     where,
 )
 
 from icon4py.model.common.dimension import (
     C2E,
+    C2E2C,
+    C2E2CO,
     E2C,
+    C2E2CODim,
     CellDim,
     E2CDim,
     EdgeDim,
@@ -564,6 +569,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 +1060,112 @@ def compute_hmask_dd3d(
         out=hmask_dd3d,
         domain={EdgeDim: (horizontal_start, horizontal_end)},
     )
+
+
+@field_operator
+def _compute_weighted_cell_neighbor_sum(
+    field: Field[[CellDim, KDim], wpfloat],
+    c_bln_avg: Field[[CellDim, C2E2CODim], wpfloat],
+) -> Field[[CellDim, KDim], wpfloat]:
+    field_avg = neighbor_sum(field(C2E2CO) * c_bln_avg, axis=C2E2CODim)
+    return field_avg
+
+
+@program
+def compute_weighted_cell_neighbor_sum(
+    maxslp: Field[[CellDim, KDim], wpfloat],
+    maxhgtd: Field[[CellDim, KDim], wpfloat],
+    c_bln_avg: Field[[CellDim, C2E2CODim], 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: Interpolation field
+        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_weighted_cell_neighbor_sum(
+        field=maxslp,
+        c_bln_avg=c_bln_avg,
+        out=z_maxslp_avg,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+    _compute_weighted_cell_neighbor_sum(
+        field=maxhgtd,
+        c_bln_avg=c_bln_avg,
+        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)}
+    )
+
+
+@scan_operator(axis=KDim, forward=True, init=(0, False))
+def _compute_param(
+    param: tuple[int32, bool],
+    z_me_jk: float,
+    z_ifc_off: float,
+    z_ifc_off_koff: float,
+    lower: int32,
+    nlev: int32,
+) -> tuple[int32, bool]:
+    param_0, param_1 = param
+    if param_0 >= lower:
+        if (param_0 == nlev) | (z_me_jk <= z_ifc_off) & (z_me_jk >= z_ifc_off_koff):
+            param_1 = True
+    return param_0 + 1, param_1
+
+
+@field_operator(grid_type=GridType.UNSTRUCTURED)
+def _compute_z_ifc_off_koff(
+    z_ifc_off: Field[[EdgeDim, KDim], wpfloat],
+) -> Field[[EdgeDim, KDim], wpfloat]:
+    n = z_ifc_off(Koff[1])
+    return n

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

@@ -0,0 +1,203 @@
+# 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_ls_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]:
+    indlist = [0] * n_cells
+    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
+            indlist[ji] = jc
+
+            if all((k_start[jc], k_end[jc])) and k_start[jc] > k_end[jc]:
+                listreduce += 1
+            else:
+                ji_ind += 1
+                indlist[ji_ind] = jc
+
+    return indlist, 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,
+    max_nbhgt: np.array,
+    c_owner_mask: np.array,
+    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,
+    cell_nudging: int,
+    n_cells: int,
+    nlev: int,
+) -> tuple[np.array, np.array, np.array, np.array]:
+    k_start, k_end = _compute_k_start_end(
+        z_mc=z_mc,
+        max_nbhgt=max_nbhgt,
+        z_maxslp_avg=z_maxslp_avg,
+        z_maxhgtd_avg=z_maxhgtd_avg,
+        c_owner_mask=c_owner_mask,
+        thslp_zdiffu=thslp_zdiffu,
+        thhgtd_zdiffu=thhgtd_zdiffu,
+        cell_nudging=cell_nudging,
+        n_cells=n_cells,
+        nlev=nlev,
+    )
+
+    indlist, listreduce, ji = _compute_ls_params(
+        k_start=k_start,
+        k_end=k_end,
+        z_maxslp_avg=z_maxslp_avg,
+        z_maxhgtd_avg=z_maxhgtd_avg,
+        c_owner_mask=c_owner_mask,
+        thslp_zdiffu=thslp_zdiffu,
+        thhgtd_zdiffu=thhgtd_zdiffu,
+        cell_nudging=cell_nudging,
+        n_cells=n_cells,
+        nlev=nlev,
+    )
+
+    listdim = ji - listreduce
+
+    for ji in range(listdim):
+        jc = 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