diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py
index 1900843ae9..5ce6847f76 100644
--- a/model/common/src/icon4py/model/common/metrics/metric_fields.py
+++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py
@@ -24,7 +24,9 @@
int32,
maximum,
minimum,
+ neighbor_sum,
program,
+ scan_operator,
sin,
tanh,
where,
@@ -32,7 +34,10 @@
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
diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py
new file mode 100644
index 0000000000..63e78ad739
--- /dev/null
+++ b/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 .
+#
+# 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
diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py
index 70461ea639..02461a7d1c 100644
--- a/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py
+++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_zdiff_gradp_dsl.py
@@ -29,19 +29,48 @@ def compute_zdiff_gradp_dsl(
zdiff_gradp[horizontal_start:, :, :] = (
np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :]
)
+ """
+ First part for loop implementation with gt4py code
+
+ >>> z_ifc_off_koff = zero_field(icon_grid, EdgeDim, KDim, extend={KDim: 1})
+ >>> z_ifc_off = as_field((EdgeDim, KDim,), z_ifc[e2c[:, 0], :])
+ >>> _compute_z_ifc_off_koff(
+ >>> z_ifc_off=z_ifc_off,
+ >>> domain={EdgeDim: (horizontal_start, nedges), KDim: (0, nlev)},
+ >>> out=z_ifc_off_koff,
+ >>> offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}
+ >>> )
+ """
for je in range(horizontal_start, nedges):
- jk_start = int(flat_idx[je])
for jk in range(int(flat_idx[je]) + 1, nlev):
- for jk1 in range(jk_start, nlev):
+ """
+ Second part for loop implementation with gt4py code
+ >>> param_2 = as_field((KDim,), np.asarray([False] * nlev))
+ >>> param_3 = as_field((KDim,), np.arange(nlev))
+ >>> z_ifc_off_e = as_field((KDim,), z_ifc[e2c[je, 0], :])
+ >>> _compute_param.with_backend(backend)(
+ >>> z_me_jk=z_me[je, jk],
+ >>> z_ifc_off=z_ifc_off_e,
+ >>> z_ifc_off_koff=as_field((KDim,), z_ifc_off_koff.asnumpy()[je, :]),
+ >>> lower=int(flat_idx[je]),
+ >>> nlev=nlev - 1,
+ >>> out=(param_3, param_2),
+ >>> offset_provider={}
+ >>> )
+ >>> zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], np.where(param_2.asnumpy())[0][0]]
+ """
+
+ param = [False] * nlev
+ for jk1 in range(int(flat_idx[je]), nlev):
if (
jk1 == nlev - 1
or z_me[je, jk] <= z_ifc[e2c[je, 0], jk1]
and z_me[je, jk] >= z_ifc[e2c[je, 0], jk1 + 1]
):
- zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], jk1]
- jk_start = jk1
- break
+ param[jk1] = True
+
+ zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], np.where(param)[0][0]]
jk_start = int(flat_idx[je])
for jk in range(int(flat_idx[je]) + 1, nlev):
diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
index e7946b2c54..b147f66995 100644
--- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
+++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
@@ -497,6 +497,9 @@ def construct_cell_geometry(self) -> horizontal.CellParams:
class InterpolationSavepoint(IconSavepoint):
+ def c_bln_avg(self):
+ return self._get_field("c_bln_avg", CellDim, C2E2CODim)
+
def c_intp(self):
return self._get_field("c_intp", VertexDim, V2CDim)
@@ -705,6 +708,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)
diff --git a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py
index a89d86c722..a79cc83ed6 100644
--- a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py
+++ b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py
@@ -10,18 +10,6 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
-# 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 .
-#
-# SPDX-License-Identifier: GPL-3.0-or-later
import pytest
diff --git a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py
new file mode 100644
index 0000000000..bf7fd920a4
--- /dev/null
+++ b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py
@@ -0,0 +1,150 @@
+# 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 .
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+
+import pytest
+from gt4py.next import as_field
+
+from icon4py.model.common.dimension import (
+ C2E2CDim,
+ CECDim,
+ CellDim,
+ KDim,
+)
+from icon4py.model.common.grid.horizontal import (
+ HorizontalMarkerIndex,
+)
+from icon4py.model.common.metrics.metric_fields import (
+ compute_max_nbhgt,
+ compute_maxslp_maxhgtd,
+ compute_weighted_cell_neighbor_sum,
+ compute_z_mc,
+)
+from icon4py.model.common.metrics.stencils.compute_diffusion_metrics import (
+ compute_diffusion_metrics,
+)
+from icon4py.model.common.test_utils.helpers import (
+ constant_field,
+ dallclose,
+ flatten_first_two_dims,
+ is_roundtrip,
+ zero_field,
+)
+
+
+@pytest.mark.datatest
+def test_compute_diffusion_metrics(
+ metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend
+):
+ if is_roundtrip(backend):
+ pytest.skip("skipping: slow backend")
+ mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy()
+ zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy()
+ z_vintcoeff = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy()
+ zd_intcoef_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy()
+ z_maxslp_avg = zero_field(icon_grid, CellDim, KDim)
+ z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim)
+ zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy()
+ maxslp = zero_field(icon_grid, CellDim, KDim)
+ maxhgtd = zero_field(icon_grid, CellDim, KDim)
+ max_nbhgt = zero_field(icon_grid, CellDim)
+
+ c2e2c = icon_grid.connectivities[C2E2CDim]
+ nbidx = constant_field(icon_grid, 1, CellDim, C2E2CDim, KDim, dtype=int).asnumpy()
+ c_bln_avg = interpolation_savepoint.c_bln_avg()
+ thslp_zdiffu = 0.02
+ thhgtd_zdiffu = 125
+ cell_nudging = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim))
+ cell_lateral = icon_grid.get_start_index(
+ CellDim,
+ HorizontalMarkerIndex.lateral_boundary(CellDim) + 1,
+ )
+ nlev = icon_grid.num_levels
+
+ compute_maxslp_maxhgtd.with_backend(backend)(
+ ddxn_z_full=metrics_savepoint.ddxn_z_full(),
+ dual_edge_length=grid_savepoint.dual_edge_length(),
+ z_maxslp=maxslp,
+ z_maxhgtd=maxhgtd,
+ horizontal_start=cell_lateral,
+ horizontal_end=icon_grid.num_cells,
+ vertical_start=0,
+ vertical_end=nlev,
+ offset_provider={"C2E": icon_grid.get_offset_provider("C2E")},
+ )
+
+ z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1})
+ compute_z_mc.with_backend(backend)(
+ metrics_savepoint.z_ifc(),
+ z_mc,
+ horizontal_start=0,
+ horizontal_end=icon_grid.num_cells,
+ vertical_start=0,
+ vertical_end=nlev,
+ offset_provider={"Koff": icon_grid.get_offset_provider("Koff")},
+ )
+
+ compute_weighted_cell_neighbor_sum.with_backend(backend)(
+ maxslp=maxslp,
+ maxhgtd=maxhgtd,
+ c_bln_avg=c_bln_avg,
+ z_maxslp_avg=z_maxslp_avg,
+ z_maxhgtd_avg=z_maxhgtd_avg,
+ horizontal_start=cell_lateral,
+ horizontal_end=icon_grid.num_cells,
+ vertical_start=0,
+ vertical_end=nlev,
+ offset_provider={
+ "C2E2C": icon_grid.get_offset_provider("C2E2C"),
+ "C2E2CO": icon_grid.get_offset_provider("C2E2CO"),
+ },
+ )
+
+ compute_max_nbhgt.with_backend(backend)(
+ z_mc_nlev=as_field((CellDim,), z_mc.asnumpy()[:, nlev - 1]),
+ max_nbhgt=max_nbhgt,
+ horizontal_start=cell_nudging,
+ horizontal_end=icon_grid.num_cells,
+ offset_provider={"C2E2C": icon_grid.get_offset_provider("C2E2C")},
+ )
+
+ mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_metrics(
+ z_mc=z_mc.asnumpy(),
+ z_mc_off=z_mc.asnumpy()[c2e2c],
+ max_nbhgt=max_nbhgt.asnumpy(),
+ c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(),
+ nbidx=nbidx,
+ z_vintcoeff=z_vintcoeff,
+ z_maxslp_avg=z_maxslp_avg.asnumpy(),
+ z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(),
+ mask_hdiff=mask_hdiff,
+ zd_diffcoef_dsl=zd_diffcoef_dsl,
+ zd_intcoef_dsl=zd_intcoef_dsl,
+ zd_vertoffset_dsl=zd_vertoffset_dsl,
+ thslp_zdiffu=thslp_zdiffu,
+ thhgtd_zdiffu=thhgtd_zdiffu,
+ cell_nudging=cell_nudging,
+ n_cells=icon_grid.num_cells,
+ nlev=nlev,
+ )
+ zd_intcoef_dsl = flatten_first_two_dims(
+ CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_intcoef_dsl)
+ )
+ zd_vertoffset_dsl = flatten_first_two_dims(
+ CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_vertoffset_dsl)
+ )
+
+ assert dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy())
+ assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11)
+ assert dallclose(zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy())
+ assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef().asnumpy())
diff --git a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py
index 6227f2d00c..0fe2f4cd1e 100644
--- a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py
+++ b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py
@@ -10,18 +10,6 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
-# 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 .
-#
-# SPDX-License-Identifier: GPL-3.0-or-later
import numpy as np
import pytest
diff --git a/model/common/tests/metric_tests/test_compute_wgtfac_c.py b/model/common/tests/metric_tests/test_compute_wgtfac_c.py
index 008d3c272c..ebd41a4df7 100644
--- a/model/common/tests/metric_tests/test_compute_wgtfac_c.py
+++ b/model/common/tests/metric_tests/test_compute_wgtfac_c.py
@@ -10,18 +10,6 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
-# 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 .
-#
-# SPDX-License-Identifier: GPL-3.0-or-later
import pytest
diff --git a/model/common/tests/metric_tests/test_compute_wgtfacq.py b/model/common/tests/metric_tests/test_compute_wgtfacq.py
index c38afa6262..b65d25d207 100644
--- a/model/common/tests/metric_tests/test_compute_wgtfacq.py
+++ b/model/common/tests/metric_tests/test_compute_wgtfacq.py
@@ -10,18 +10,6 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
-# 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 .
-#
-# SPDX-License-Identifier: GPL-3.0-or-later
import pytest
diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py
index 31462a2193..fd36dcfca1 100644
--- a/model/common/tests/metric_tests/test_metric_fields.py
+++ b/model/common/tests/metric_tests/test_metric_fields.py
@@ -386,8 +386,6 @@ def test_compute_ddqz_z_full_e(
vertical_start = 0
vertical_end = icon_grid.num_levels
ddqz_z_full_e = zero_field(icon_grid, EdgeDim, KDim)
- # TODO: perhaps write a program with ddqz_z_full_e name and call fieldop _cells2edges... from there
- # TODO: This way it's clear where this field is computed and we cna more easily avoid duplicates
cell_2_edge_interpolation.with_backend(backend)(
in_field=ddqz_z_full,
coeff=c_lin_e,