Refactor interpolation and metrics field

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

@@ -7,6 +7,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import gt4py.next as gtx
+import numpy as np
 from gt4py.next import backend as gtx_backend
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
@@ -46,6 +47,53 @@ def update_smoothed_topography(
     )
 
 
+def compute_nabla2_on_cell_numpy(
+    psi_c: data_alloc.NDArray,
+    geofac_n2s: data_alloc.NDArray,
+    c2e2co: data_alloc.NDArray,
+) -> np.ndarray:
+    """
+    Computes the Laplacian (nabla squared) of a scalar field defined on cell
+    centres. (Numpy version)
+    """
+    nabla2_psi_c = np.sum(psi_c[c2e2co] * geofac_n2s, axis=1)
+    return nabla2_psi_c
+
+
+def update_smoothed_topography_numpy(
+    smoothed_topography: np.ndarray,
+    nabla2_topo: np.ndarray,
+    cell_areas: np.ndarray,
+) -> np.ndarray:
+    """
+    Updates the smoothed topography field inside the loop. (Numpy version)
+    """
+    return smoothed_topography + 0.125 * nabla2_topo * cell_areas
+
+
+def smooth_topography_numpy(
+    topography: data_alloc.NDArray,
+    cell_areas: data_alloc.NDArray,
+    geofac_n2s: data_alloc.NDArray,
+    c2e2co: data_alloc.NDArray,
+    num_iterations: int = 25,
+) -> data_alloc.NDArray:
+    """
+    Computes the smoothed (laplacian-filtered) topography needed by the SLEVE
+    coordinate.
+    """
+
+    smoothed_topography = topography.copy()
+
+    for _ in range(num_iterations):
+        nabla2_topo = compute_nabla2_on_cell_numpy(smoothed_topography, geofac_n2s, c2e2co)
+        smoothed_topography = update_smoothed_topography_numpy(
+            smoothed_topography, nabla2_topo, cell_areas
+        )
+
+    return smoothed_topography
+
+
 def smooth_topography(
     topography: fa.CellField[ta.wpfloat],
     cell_areas: fa.CellField[ta.wpfloat],
@@ -85,5 +133,4 @@ def smooth_topography(
             horizontal_end=grid.num_cells,
             offset_provider={},
         )
-
     return smoothed_topography

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

@@ -19,6 +19,7 @@
 from gt4py.next import backend as gtx_backend
 
 import icon4py.model.common.states.metadata as data
+import icon4py.model.common.type_alias as ta
 from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa
 from icon4py.model.common.grid import base, icon as icon_grid, topography as topo
 from icon4py.model.common.utils import data_allocation as data_alloc
@@ -70,7 +71,7 @@ def _domain(marker: Zone):
     return _domain
 
 
-@dataclasses.dataclass(frozen=True)
+@dataclasses.dataclass(frozen=False)
 class VerticalGridConfig:
     """
     Contains necessary parameter to configure vertical grid.
@@ -99,7 +100,7 @@ class VerticalGridConfig:
     #: Defined in ICON namelist nonhydrostatic_nml. Height [m] above which moist physics and advection of cloud and precipitation variables are turned off.
     htop_moist_proc: Final[float] = 22500.0
     #: file name containing vct_a and vct_b table
-    file_path: pathlib.Path = None
+    file_path: Optional[pathlib.Path] = None
 
     # Parameters for setting up the decay function of the topographic signal for
     # SLEVE. Default values from mo_sleve_nml.
@@ -739,6 +740,183 @@ def _check_flatness_of_flat_level(
         raise exceptions.InvalidComputationError("Level nflatlev is not flat")
 
 
+def _compute_SLEVE_coordinate_from_vcta_and_topography_numpy(
+    vct_a: data_alloc.NDArray,
+    topography: data_alloc.NDArray,
+    cell_areas: data_alloc.NDArray,
+    geofac_n2s: data_alloc.NDArray,
+    c2e2co: data_alloc.NDArray,
+    num_cells: int,
+    num_levels: int,
+    nflatlev: int,
+    model_top_height: ta.wpfloat,
+    SLEVE_decay_scale_1: ta.wpfloat,
+    SLEVE_decay_exponent: ta.wpfloat,
+    SLEVE_decay_scale_2: ta.wpfloat,
+    array_ns: ModuleType = np,
+) -> data_alloc.NDArray:
+    """
+    Compute the 3D vertical coordinate field using the SLEVE coordinate
+    https://doi.org/10.1175/1520-0493(2002)130%3C2459:ANTFVC%3E2.0.CO;2
+
+    This is the same as vct_a in the flat levels (above nflatlev).
+    Below it is vct_a corrected by smothed and decaying topography such that it
+    blends smothly into the surface layer at num_lev + 1 which is the
+    topography.
+    """
+
+    def _decay_func(
+        vct_a: data_alloc.NDArray,
+        model_top_height: float,
+        decay_scale: float,
+        decay_exponent: float,
+    ) -> data_alloc.NDArray:
+        return array_ns.sinh(
+            (model_top_height / decay_scale) ** decay_exponent
+            - (vct_a / decay_scale) ** decay_exponent
+        ) / array_ns.sinh((model_top_height / decay_scale) ** decay_exponent)
+
+    smoothed_topography = topo.smooth_topography_numpy(
+        topography=topography,
+        cell_areas=cell_areas,
+        geofac_n2s=geofac_n2s,
+        c2e2co=c2e2co,
+    )
+
+    vertical_coordinate = array_ns.zeros((num_cells, num_levels + 1), dtype=float)
+    vertical_coordinate[:, num_levels] = topography
+
+    # Small-scale topography (i.e. full topo - smooth topo)
+    small_scale_topography = topography - smoothed_topography
+
+    k = range(nflatlev + 1)
+    vertical_coordinate[:, k] = vct_a[k]
+
+    k = range(nflatlev + 1, num_levels)
+    # Scaling factors for large-scale and small-scale topography
+    z_fac1 = _decay_func(
+        vct_a[k],
+        model_top_height,
+        SLEVE_decay_scale_1,
+        SLEVE_decay_exponent,
+    )
+    z_fac2 = _decay_func(
+        vct_a[k],
+        model_top_height,
+        SLEVE_decay_scale_2,
+        SLEVE_decay_exponent,
+    )
+    vertical_coordinate[:, k] = (
+        vct_a[k][array_ns.newaxis, :]
+        + smoothed_topography[:, array_ns.newaxis] * z_fac1
+        + small_scale_topography[:, array_ns.newaxis] * z_fac2
+    )
+
+    return vertical_coordinate
+
+
+def _check_and_correct_layer_thickness_numpy(
+    vertical_coordinate: data_alloc.NDArray,
+    vct_a: data_alloc.NDArray,
+    num_cells: int,
+    num_levels: int,
+    SLEVE_minimum_layer_thickness_1: ta.wpfloat,
+    SLEVE_minimum_relative_layer_thickness_1: ta.wpfloat,
+    SLEVE_minimum_layer_thickness_2: ta.wpfloat,
+    SLEVE_minimum_relative_layer_thickness_2: ta.wpfloat,
+    lowest_layer_thickness: ta.wpfloat,
+    array_ns: ModuleType = np,
+) -> data_alloc.NDArray:
+    ktop_thicklimit = array_ns.asarray(num_cells * [num_levels], dtype=int)
+    # Ensure that layer thicknesses are not too small; this would potentially
+    # cause instabilities in vertical advection
+    for k in reversed(range(num_levels)):
+        delta_vct_a = vct_a[k] - vct_a[k + 1]
+        if delta_vct_a < SLEVE_minimum_layer_thickness_1:
+            # limit layer thickness to SLEVE_minimum_relative_layer_thickness_1 times its nominal value
+            minimum_layer_thickness = SLEVE_minimum_relative_layer_thickness_1 * delta_vct_a
+        elif delta_vct_a < 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 = (
+                (SLEVE_minimum_layer_thickness_2 - delta_vct_a)
+                / (SLEVE_minimum_layer_thickness_2 - SLEVE_minimum_layer_thickness_1)
+            ) ** 2
+            minimum_layer_thickness = (
+                SLEVE_minimum_relative_layer_thickness_1 * layer_thickness_adjustment_factor
+                + SLEVE_minimum_relative_layer_thickness_2
+                * (1.0 - layer_thickness_adjustment_factor)
+            ) * delta_vct_a
+        else:
+            # limitation factor decreases again
+            minimum_layer_thickness = (
+                SLEVE_minimum_relative_layer_thickness_2
+                * SLEVE_minimum_layer_thickness_2
+                * (delta_vct_a / SLEVE_minimum_layer_thickness_2) ** (1.0 / 3.0)
+            )
+
+        minimum_layer_thickness = max(minimum_layer_thickness, min(50, lowest_layer_thickness))
+
+        # Ensure that the layer thickness is not too small, if so fix it and
+        # save the layer index
+        cell_ids = array_ns.argwhere(
+            vertical_coordinate[:, k + 1] + minimum_layer_thickness > vertical_coordinate[:, k]
+        )
+        vertical_coordinate[cell_ids, k] = (
+            vertical_coordinate[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 = array_ns.argwhere((ktop_thicklimit <= num_levels - 3) & (ktop_thicklimit >= 3))
+    if cell_ids.size > 0:
+        delta_z1 = (
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] + 1]
+            - vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] + 2]
+        )
+        delta_z2 = (
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] - 3]
+            - vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] - 2]
+        )
+        stretching_factor = (delta_z2 / delta_z1) ** 0.25
+        delta_z3 = (
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] - 2]
+            - vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] + 1]
+        ) / (stretching_factor * (1.0 + stretching_factor * (1.0 + stretching_factor)))
+        vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids]] = array_ns.maximum(
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids]],
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] + 1]
+            + delta_z3 * stretching_factor,
+        )
+        vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] - 1] = array_ns.maximum(
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids] - 1],
+            vertical_coordinate[cell_ids, ktop_thicklimit[cell_ids]]
+            + delta_z3 * stretching_factor**2,
+        )
+
+    # Check if ktop_thicklimit is sufficiently far away from the model top
+    if not array_ns.all(ktop_thicklimit > 2):
+        if num_levels > 6:
+            raise exceptions.InvalidConfigError(
+                f"Model top is too low and num_levels, {num_levels}, > 6."
+            )
+        else:
+            log.warning(f"Model top is too low. But num_levels, {num_levels}, <= 6. ")
+
+    return vertical_coordinate
+
+
+def _check_flatness_of_flat_level_numpy(
+    vertical_coordinate: data_alloc.NDArray,
+    vct_a: data_alloc.NDArray,
+    nflatlev: int,
+    array_ns: ModuleType = np,
+) -> None:
+    # Check if level nflatlev is still flat
+    if not array_ns.all(vertical_coordinate[:, nflatlev - 1] == vct_a[nflatlev - 1]):
+        raise exceptions.InvalidComputationError("Level nflatlev is not flat")
+
+
 def compute_vertical_coordinate(
     vct_a: data_alloc.NDArray,
     topography: data_alloc.NDArray,
@@ -797,3 +975,101 @@ def compute_vertical_coordinate(
     )
 
     return vertical_coordinate
+
+
+def compute_vertical_coordinate_numpy(
+    vct_a: data_alloc.NDArray,
+    topography: data_alloc.NDArray,
+    cell_areas: data_alloc.NDArray,
+    geofac_n2s: data_alloc.NDArray,
+    c2e2co: data_alloc.NDArray,
+    num_cells: int,
+    num_levels: int,
+    nflatlev: int,
+    model_top_height: ta.wpfloat,
+    SLEVE_decay_scale_1: ta.wpfloat,
+    SLEVE_decay_exponent: ta.wpfloat,
+    SLEVE_decay_scale_2: ta.wpfloat,
+    SLEVE_minimum_layer_thickness_1: ta.wpfloat,
+    SLEVE_minimum_relative_layer_thickness_1: ta.wpfloat,
+    SLEVE_minimum_layer_thickness_2: ta.wpfloat,
+    SLEVE_minimum_relative_layer_thickness_2: ta.wpfloat,
+    lowest_layer_thickness: ta.wpfloat,
+    array_ns: ModuleType = np,
+) -> data_alloc.NDArray:
+    """
+    Compute the (Cell, K) vertical coordinate field starting from the
+    "flat/uniform/aquaplanet" vertical coordinate.
+
+    Args:
+        vct_a: Vertical coordinate with flat topography.
+        topography: Topography field.
+        cell_areas: Cell areas field.
+        geofac_n2s: Coefficients for nabla2 computation.
+        grid: Grid object.
+        vertical_geometry: Vertical grid object.
+        backend: Backend to use for computations.
+
+    Returns:
+        The (Cell, K) vertical coordinate field.
+
+    Raises:
+        exceptions.InvalidComputationError: If level nflatlev is not flat.
+        exceptions.InvalidConfigError: If model top is too low and num_levels > 6.
+    """
+
+    vertical_coordinate = _compute_SLEVE_coordinate_from_vcta_and_topography_numpy(
+        vct_a=vct_a,
+        topography=topography,
+        cell_areas=cell_areas,
+        geofac_n2s=geofac_n2s,
+        c2e2co=c2e2co,
+        num_cells=num_cells,
+        num_levels=num_levels,
+        nflatlev=nflatlev,
+        model_top_height=model_top_height,
+        SLEVE_decay_scale_1=SLEVE_decay_scale_1,
+        SLEVE_decay_exponent=SLEVE_decay_exponent,
+        SLEVE_decay_scale_2=SLEVE_decay_scale_2,
+        array_ns=array_ns,
+    )
+
+    vertical_coordinate = _check_and_correct_layer_thickness_numpy(
+        vertical_coordinate=vertical_coordinate,
+        vct_a=vct_a,
+        num_cells=num_cells,
+        num_levels=num_levels,
+        SLEVE_minimum_layer_thickness_1=SLEVE_minimum_layer_thickness_1,
+        SLEVE_minimum_relative_layer_thickness_1=SLEVE_minimum_relative_layer_thickness_1,
+        SLEVE_minimum_layer_thickness_2=SLEVE_minimum_layer_thickness_2,
+        SLEVE_minimum_relative_layer_thickness_2=SLEVE_minimum_relative_layer_thickness_2,
+        lowest_layer_thickness=lowest_layer_thickness,
+        array_ns=array_ns,
+    )
+
+    _check_flatness_of_flat_level_numpy(
+        vertical_coordinate=vertical_coordinate,
+        vct_a=vct_a,
+        nflatlev=nflatlev,
+        array_ns=array_ns,
+    )
+
+    return vertical_coordinate
+
+
+def compute_surface_elevation(
+    vertical_coordinate: data_alloc.NDArray,
+    num_levels: int,
+) -> data_alloc.NDArray:
+    """
+    Compute the surface elevation from the vertical coordinate field.
+
+    Args:
+        vertical_coordinate: The (Cell, K) vertical coordinate field.
+        vertical_geometry: Vertical grid object.
+        array_ns: NumPy module for array operations.
+
+    Returns:
+        The surface elevation field.
+    """
+    return vertical_coordinate[:, num_levels]