Port z ifc

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

@@ -34,6 +34,7 @@ def calculate_nabla2_for_w(
     vertical_start: gtx.int32,
     vertical_end: gtx.int32,
 ):
+    # TODO: replace this by common/math/stencils/compute_nabla2_on_cell_k
     _calculate_nabla2_for_w(
         w,
         geofac_n2s,

model/common/src/icon4py/model/common/external_parameters.py

@@ -0,0 +1,30 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import dataclasses
+
+import icon4py.model.common.external_parameters
+from icon4py.model.common import field_type_aliases as fa
+from icon4py.model.common.test_utils import serialbox_utils as sb
+
+
+@dataclasses.dataclass
+class ExternalParameters:
+    """Dataclass containing external parameters."""
+
+    topo_c: fa.CellField[float]
+    topo_smt_c: fa.CellField[float]
+
+
+def construct_external_parameters_state(
+    savepoint: sb.ExternalParametersSavepoint, num_k_lev
+) -> icon4py.model.common.external_parameters.ExternalParameters:
+    return icon4py.model.common.external_parameters.ExternalParameters(
+        topo_c=savepoint.topo_c(),
+        topo_smt_c=savepoint.topo_smt_c(),
+    )

model/common/src/icon4py/model/common/grid/topography.py

@@ -0,0 +1,87 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import gt4py.next as gtx
+
+from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common.grid import icon as icon_grid
+from icon4py.model.common.math.stencils.compute_nabla2_on_cell import compute_nabla2_on_cell
+from icon4py.model.common.settings import xp
+
+
+@gtx.field_operator
+def _update_topo_smooth(
+    topography_smoothed: fa.CellField[ta.wpfloat],
+    nabla2_topo: fa.CellField[ta.wpfloat],
+    cell_areas: fa.CellField[ta.wpfloat],
+) -> fa.CellField[ta.wpfloat]:
+    """
+    Updates the smoothed topography field inside the loop.
+    """
+    return topography_smoothed + 0.125 * nabla2_topo * cell_areas
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def update_topo_smooth(
+    topography_smoothed: fa.CellField[ta.wpfloat],
+    nabla2_topo: fa.CellField[ta.wpfloat],
+    cell_areas: fa.CellField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+):
+    _update_topo_smooth(
+        topography_smoothed=topography_smoothed,
+        nabla2_topo=nabla2_topo,
+        cell_areas=cell_areas,
+        out=topography_smoothed,
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+        },
+    )
+
+
+def compute_smooth_topo(
+    topography: fa.CellField[ta.wpfloat],
+    grid: icon_grid.IconGrid,
+    cell_areas: fa.CellField[ta.wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat],
+    backend,
+    num_iterations: int = 25,
+) -> fa.CellField[ta.wpfloat]:
+    """
+    Computes the smoothed (laplacian-filtered) topography needed by the SLEVE
+    coordinate.
+    """
+
+    topography_smoothed = gtx.as_field((dims.CellDim,), topography.ndarray)
+
+    nabla2_topo_np = xp.zeros((grid.num_cells,), dtype=ta.wpfloat)
+    nabla2_topo = gtx.as_field((dims.CellDim,), nabla2_topo_np)
+
+    for _ in range(num_iterations):
+        compute_nabla2_on_cell.with_backend(backend)(
+            psi_c=topography_smoothed,
+            geofac_n2s=geofac_n2s,
+            nabla2_psi_c=nabla2_topo,
+            horizontal_start=0,
+            horizontal_end=grid.num_cells,
+            offset_provider={
+                "C2E2CO": grid.get_offset_provider("C2E2CO"),
+            },
+        )
+
+        update_topo_smooth.with_backend(backend)(
+            topography_smoothed=topography_smoothed,
+            nabla2_topo=nabla2_topo,
+            cell_areas=cell_areas,
+            horizontal_start=0,
+            horizontal_end=grid.num_cells,
+            offset_provider={},
+        )
+
+    return topography_smoothed

model/common/src/icon4py/model/common/grid/vertical.py

@@ -16,7 +16,13 @@
 import gt4py.next as gtx
 
 import icon4py.model.common.states.metadata as data
-from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa
+from icon4py.model.common import (
+    dimension as dims,
+    exceptions,
+    field_type_aliases as fa,
+    type_alias as ta,
+)
+from icon4py.model.common.grid import icon as icon_grid
 from icon4py.model.common.settings import xp
 
 
@@ -97,6 +103,23 @@ class VerticalGridConfig:
     #: file name containing vct_a and vct_b table
     file_path: pathlib.Path = None
 
+    # Parameters for setting up the decay function of the topographic signal for
+    # SLEVE. Default values from mo_sleve_nml.
+    #: Decay scale for large-scale topography component
+    SLEVE_decay_scale_1: Final[float] = 4000.0
+    #: Decay scale for small-scale topography component
+    SLEVE_decay_scale_2: Final[float] = 2500.0
+    #: Exponent for decay function
+    SLEVE_decay_exponent: Final[float] = 1.2
+    #: minimum absolute layer thickness 1 for SLEVE coordinates
+    SLEVE_minimum_layer_thickness_1: Final[float] = 100.0
+    #: minimum absolute layer thickness 2 for SLEVE coordinates
+    SLEVE_minimum_layer_thickness_2: Final[float] = 500.0
+    #: minimum relative layer thickness for nominal thicknesses <= SLEVE_minimum_layer_thickness_1
+    SLEVE_minimum_relative_layer_thickness_1: Final[float] = 1.0 / 3.0
+    #: minimum relative layer thickness for a nominal thickness of SLEVE_minimum_layer_thickness_2
+    SLEVE_minimum_relative_layer_thickness_2: Final[float] = 0.5
+
 
 @dataclasses.dataclass(frozen=True)
 class VerticalGrid:
@@ -517,3 +540,154 @@ def get_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KField,
         if vertical_config.file_path
         else _compute_vct_a_and_vct_b(vertical_config)
     )
+
+
+def init_vert_coord(
+    vct_a: fa.KField[ta.wpfloat],
+    topography: fa.CellField[ta.wpfloat],
+    topography_smoothed: fa.CellField[ta.wpfloat],
+    grid: icon_grid.IconGrid,
+    vertical_config: VerticalGridConfig,
+    vertical_geometry: VerticalGrid,
+    nshift: int = 0,  # shift parameter used for vertical nesting
+) -> fa.CellKField[ta.wpfloat]:
+    """
+    Computes the 3D vertical coordinate fields.
+    """
+
+    def _decay_func(
+        vct_a: xp.ndarray, model_top_height: float, decay_scale: float, decay_exponent: float
+    ) -> xp.ndarray:
+        return xp.sinh(
+            (model_top_height / decay_scale) ** decay_exponent
+            - (vct_a / decay_scale) ** decay_exponent
+        ) / xp.sinh((model_top_height / decay_scale) ** decay_exponent)
+
+    z3d_i = xp.zeros((grid.num_cells, grid.num_levels + 1), dtype=ta.wpfloat)
+
+    topography = topography.ndarray
+    topography_smoothed = topography_smoothed.ndarray
+    vct_a = vct_a.ndarray
+
+    z3d_i[:, grid.num_levels] = topography
+    ktop_thicklimit = xp.asarray(grid.num_cells * [grid.num_levels], dtype=ta.wpfloat)
+
+    # Small-scale topography (i.e. full topo - smooth topo)
+    topo_deviation = topography - topography_smoothed
+
+    k = range(vertical_geometry.nflatlev + 1)
+    k1 = range(nshift, vertical_geometry.nflatlev + 1 + nshift)
+    z3d_i[:, k] = vct_a[k1]
+
+    k = range(vertical_geometry.nflatlev + 1, grid.num_levels)
+    k1 = range(vertical_geometry.nflatlev + 1 + nshift, grid.num_levels + nshift)
+    # Scaling factors for large-scale and small-scale topography
+    z_fac1 = _decay_func(
+        vct_a[k1],
+        vertical_config.model_top_height,
+        vertical_config.SLEVE_decay_scale_1,
+        vertical_config.SLEVE_decay_exponent,
+    )
+    z_fac2 = _decay_func(
+        vct_a[k1],
+        vertical_config.model_top_height,
+        vertical_config.SLEVE_decay_scale_2,
+        vertical_config.SLEVE_decay_exponent,
+    )
+    z3d_i[:, k] = (
+        vct_a[k1][xp.newaxis, :]
+        + topography_smoothed[:, xp.newaxis] * z_fac1
+        + topo_deviation[:, xp.newaxis] * z_fac2
+    )
+
+    # Ensure that layer thicknesses are not too small; this would potentially
+    # cause instabilities in vertical advection
+    for k in reversed(range(grid.num_levels)):
+        k1 = k + nshift
+        delta_vct_a = vct_a[k1] - vct_a[k1 + 1]
+        if delta_vct_a < vertical_config.SLEVE_minimum_layer_thickness_1:
+            # limit layer thickness to SLEVE_minimum_relative_layer_thickness_1 times its nominal value
+            minimum_layer_thickness = (
+                vertical_config.SLEVE_minimum_relative_layer_thickness_1 * delta_vct_a
+            )
+        elif delta_vct_a < vertical_config.SLEVE_minimum_layer_thickness_2:
+            # limitation factor changes from SLEVE_minimum_relative_layer_thickness_1 to SLEVE_minimum_relative_layer_thickness_2
+            layer_thickness_adjustment_factor = (
+                (vertical_config.SLEVE_minimum_layer_thickness_2 - delta_vct_a)
+                / (
+                    vertical_config.SLEVE_minimum_layer_thickness_2
+                    - vertical_config.SLEVE_minimum_layer_thickness_1
+                )
+            ) ** 2
+            minimum_layer_thickness = (
+                vertical_config.SLEVE_minimum_relative_layer_thickness_1
+                * layer_thickness_adjustment_factor
+                + vertical_config.SLEVE_minimum_relative_layer_thickness_2
+                * (1.0 - layer_thickness_adjustment_factor)
+            ) * delta_vct_a
+        else:
+            # limitation factor decreases again
+            minimum_layer_thickness = (
+                vertical_config.SLEVE_minimum_relative_layer_thickness_2
+                * vertical_config.SLEVE_minimum_layer_thickness_2
+                * (delta_vct_a / vertical_config.SLEVE_minimum_layer_thickness_2) ** (1.0 / 3.0)
+            )
+
+        minimum_layer_thickness = max(
+            minimum_layer_thickness, min(50, vertical_config.lowest_layer_thickness)
+        )
+
+        # Ensure that the layer thickness is not too small, if so fix it and
+        # save the layer index
+        cell_ids = xp.argwhere(z3d_i[:, k + 1] + minimum_layer_thickness > z3d_i[:, k])
+        z3d_i[cell_ids, k] = z3d_i[cell_ids, k + 1] + minimum_layer_thickness
+        ktop_thicklimit[cell_ids] = k
+
+    # Smooth layer thickness ratios in the transition layer of columns where the
+    # thickness limiter has been active (exclude lowest and highest layers)
+    cell_ids = xp.argwhere((ktop_thicklimit <= grid.num_levels - 3) & (ktop_thicklimit >= 3))
+    if cell_ids.size > 0:
+        delta_z1 = (
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids] + 1]
+            - z3d_i[cell_ids, ktop_thicklimit[cell_ids] + 2]
+        )
+        delta_z2 = (
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids] - 3]
+            - z3d_i[cell_ids, ktop_thicklimit[cell_ids] - 2]
+        )
+        stretching_factor = (delta_z2 / delta_z1) ** 0.25
+        delta_z3 = (
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids] - 2]
+            - z3d_i[cell_ids, ktop_thicklimit[cell_ids] + 1]
+        ) / (stretching_factor * (1.0 + stretching_factor * (1.0 + stretching_factor)))
+        z3d_i[cell_ids, ktop_thicklimit[cell_ids]] = xp.maximum(
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids]],
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids] + 1] + delta_z3 * stretching_factor,
+        )
+        z3d_i[cell_ids, ktop_thicklimit[cell_ids] - 1] = xp.maximum(
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids] - 1],
+            z3d_i[cell_ids, ktop_thicklimit[cell_ids]] + delta_z3 * stretching_factor**2,
+        )
+
+    # Check if level nflatlev is still flat
+    try:
+        assert xp.all(
+            z3d_i[:, vertical_geometry.nflatlev - 1]
+            == vct_a[vertical_geometry.nflatlev - 1 + nshift]
+        )
+    except AssertionError:
+        log.error("Level nflatlev is not flat")
+        exit(1)
+    # Check if ktop_thicklimit is sufficiently far away from the model top
+    try:
+        assert xp.all(ktop_thicklimit > 2)
+    except AssertionError:
+        if vertical_config.num_levels > 6:
+            log.error(f"Model top is too low and num_levels, {vertical_config.num_levels}, > 6.")
+            exit(1)
+        else:
+            log.warning(
+                f"Model top is too low. But num_levels, {vertical_config.num_levels}, <= 6. "
+            )
+
+    return gtx.as_field((dims.CellDim, dims.KDim), z3d_i)

model/common/src/icon4py/model/common/math/stencils/compute_nabla2_on_cell.py

@@ -0,0 +1,30 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import gt4py.next as gtx
+
+from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common.math.stencils.operators import _compute_nabla2_on_cell
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def compute_nabla2_on_cell(
+    psi_c: fa.CellField[ta.wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat],
+    nabla2_psi_c: fa.CellField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+):
+    _compute_nabla2_on_cell(
+        psi_c,
+        geofac_n2s,
+        out=nabla2_psi_c,
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+        },
+    )

model/common/src/icon4py/model/common/math/stencils/compute_nabla2_on_cell_k.py

@@ -0,0 +1,33 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import gt4py.next as gtx
+
+from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common.math.stencils.operators import _compute_nabla2_on_cell_k
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def compute_nabla2_on_cell_k(
+    psi_c: fa.CellKField[ta.wpfloat],
+    geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat],
+    nabla2_psi_c: fa.CellKField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
+):
+    _compute_nabla2_on_cell_k(
+        psi_c,
+        geofac_n2s,
+        out=nabla2_psi_c,
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+            dims.KDim: (vertical_start, vertical_end),
+        },
+    )