Metrics reference atmosphere fields

model/common/src/icon4py/model/common/constants.py

@@ -13,30 +13,50 @@
 import sys
 from typing import Final
 
+from icon4py.model.common.type_alias import wpfloat
+
 
 #: Gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90),
 #: see https://glossary.ametsoc.org/wiki/Gas_constant.
-GAS_CONSTANT_DRY_AIR: Final[float] = 287.04
-RD: Final[float] = GAS_CONSTANT_DRY_AIR
+GAS_CONSTANT_DRY_AIR: Final[wpfloat] = 287.04
+RD: Final[wpfloat] = GAS_CONSTANT_DRY_AIR
 
 #: Specific heat at constant pressure [J/K/kg]
-CPD: Final[float] = 1004.64
+SPECIFIC_HEAT_CONSTANT_PRESSURE: Final[wpfloat] = 1004.64
+CPD = SPECIFIC_HEAT_CONSTANT_PRESSURE
 
 #: [J/K/kg] specific heat at constant volume
-CVD: Final[float] = CPD - RD
-CVD_O_RD: Final[float] = CVD / RD
+SPECIFIC_HEAT_CONSTANT_VOLUME: Final[wpfloat] = CPD - RD
+CVD: Final[wpfloat] = SPECIFIC_HEAT_CONSTANT_VOLUME
+CVD_O_RD: Final[wpfloat] = CVD / RD
 
 #: Gas constant for water vapor [J/K/kg], rv in ICON.
-GAS_CONSTANT_WATER_VAPOR: Final[float] = 461.51
-RV: Final[float] = GAS_CONSTANT_WATER_VAPOR
+GAS_CONSTANT_WATER_VAPOR: Final[wpfloat] = 461.51
+RV: Final[wpfloat] = GAS_CONSTANT_WATER_VAPOR
 
 #: Av. gravitational acceleration [m/s^2]
-GRAVITATIONAL_ACCELERATION: Final[float] = 9.80665
-GRAV: Final[float] = GRAVITATIONAL_ACCELERATION
+GRAVITATIONAL_ACCELERATION: Final[wpfloat] = 9.80665
+GRAV: Final[wpfloat] = GRAVITATIONAL_ACCELERATION
 
 #: reference pressure for Exner function [Pa]
-P0REF: Final[float] = 100000.0
+REFERENCE_PRESSURE: Final[wpfloat] = 100000.0
+P0REF: Final[wpfloat] = REFERENCE_PRESSURE
+
+#: sea level pressure [Pa]
+SEAL_LEVEL_PRESSURE: Final[wpfloat] = 101325.0
+P0SL_BG: Final[wpfloat] = SEAL_LEVEL_PRESSURE
+
+#: sea level temperature for reference atmosphere [K]
+SEA_LEVEL_TEMPERATURE: Final[wpfloat] = 288.15
+T0SL_BG: Final[wpfloat] = SEA_LEVEL_TEMPERATURE
+
+#: difference between sea level temperature and asymptotic stratospheric temperature
+DELTA_TEMPERATURE: Final[wpfloat] = 75.0
+DEL_T_BG: Final[wpfloat] = DELTA_TEMPERATURE
 
+#: height scale for reference atmosphere [m], defined  in mo_vertical_grid
+#: scale height [m]
+_H_SCAL_BG: Final[wpfloat] = 10000.0
 
 # Math constants
 dbl_eps = sys.float_info.epsilon  # EPSILON(1._wp)

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

@@ -13,10 +13,12 @@
 from dataclasses import dataclass
 from typing import Final
 
-from gt4py.next.common import Dimension, Field
+from gt4py.next import Dimension, Field, GridType, field_operator, neighbor_sum, program
+from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.common import dimension
-from icon4py.model.common.dimension import CellDim, ECDim, ECVDim, EdgeDim
+from icon4py.model.common.dimension import E2C, CellDim, E2CDim, ECDim, ECVDim, EdgeDim, KDim
+from icon4py.model.common.type_alias import wpfloat
 
 
 NUM_GHOST_ROWS: Final[int] = 2
@@ -287,14 +289,43 @@ def __init__(
 
 @dataclass(frozen=True)
 class CellParams:
+    #: Area of a cell, defined in ICON in mo_model_domain.f90:t_grid_cells%area
     area: Field[[CellDim], float]
+
     mean_cell_area: float
 
+
+@field_operator
+def _cell_2_edge_interpolation(
+    in_field: Field[[CellDim, KDim], wpfloat], coeff: Field[[EdgeDim, E2CDim], wpfloat]
+) -> Field[[EdgeDim, KDim], wpfloat]:
     """
-    Area of a cell.
+    Interpolate a Cell Field to Edges.
 
-    defined int ICON in mo_model_domain.f90:t_grid_cells%area
+    There is a special handling of lateral boundary edges in `subroutine cells2edges_scalar`
+    where the value is set to the one valid in_field value without multiplication by coeff.
+    This essentially means: the skip value neighbor in the neighbor_sum is skipped and coeff needs to
+    be 1 for this Edge index.
     """
+    return neighbor_sum(in_field(E2C) * coeff, axis=E2CDim)
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def cell_2_edge_interpolation(
+    in_field: Field[[CellDim, KDim], wpfloat],
+    coeff: Field[[EdgeDim, E2CDim], wpfloat],
+    out_field: Field[[EdgeDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _cell_2_edge_interpolation(
+        in_field,
+        coeff,
+        out=out_field,
+        domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
 
 
 class RefinCtrlLevel:

model/common/src/icon4py/model/common/math/helpers.py

@@ -0,0 +1,71 @@
+# 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
+
+from gt4py.next import Field, field_operator
+
+from icon4py.model.common.dimension import CellDim, KDim, Koff
+from icon4py.model.common.type_alias import wpfloat
+
+
+@field_operator
+def average_k_level_up(
+    half_level_field: Field[[CellDim, KDim], wpfloat]
+) -> Field[[CellDim, KDim], wpfloat]:
+    """
+    Calculate the mean value of adjacent interface levels.
+
+    Computes the average of two adjacent interface levels upwards over a cell field for storage
+    in the corresponding full levels.
+    Args:
+        half_level_field: Field[[CellDim, KDim], wpfloat]
+
+    Returns: Field[[CellDim, KDim], wpfloat] full level field
+
+    """
+    return 0.5 * (half_level_field + half_level_field(Koff[1]))
+
+
+@field_operator
+def difference_k_level_down(
+    half_level_field: Field[[CellDim, KDim], wpfloat]
+) -> Field[[CellDim, KDim], wpfloat]:
+    """
+    Calculate the difference value of adjacent interface levels.
+
+    Computes the difference of two adjacent interface levels downwards over a cell field for storage
+    in the corresponding full levels.
+    Args:
+        half_level_field: Field[[CellDim, KDim], wpfloat]
+
+    Returns: Field[[CellDim, KDim], wpfloat] full level field
+
+    """
+    return half_level_field(Koff[-1]) - half_level_field
+
+
+@field_operator
+def difference_k_level_up(
+    half_level_field: Field[[CellDim, KDim], wpfloat]
+) -> Field[[CellDim, KDim], wpfloat]:
+    """
+    Calculate the difference value of adjacent interface levels.
+
+    Computes the difference of two adjacent interface levels upwards over a cell field for storage
+    in the corresponding full levels.
+    Args:
+        half_level_field: Field[[CellDim, KDim], wpfloat]
+
+    Returns: Field[[CellDim, KDim], wpfloat] full level field
+
+    """
+    return half_level_field - half_level_field(Koff[1])

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

@@ -0,0 +1,12 @@
+# 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

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

@@ -0,0 +1,153 @@
+# 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
+
+from gt4py.next import Field, GridType, field_operator, int32, program, where
+
+from icon4py.model.common.dimension import CellDim, KDim, Koff
+from icon4py.model.common.math.helpers import (
+    average_k_level_up,
+    difference_k_level_down,
+    difference_k_level_up,
+)
+from icon4py.model.common.type_alias import wpfloat
+
+
+"""
+Contains metric fields calculations for the vertical grid, ported from mo_vertical_grid.f90.
+"""
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def compute_z_mc(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    z_mc: Field[[CellDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute the geometric height of full levels from the geometric height of half levels (z_ifc).
+
+    This assumes that the input field z_ifc is defined on half levels (KHalfDim) and the
+    returned fields is defined on full levels (KDim)
+
+    Args:
+        z_ifc: Field[[CellDim, KDim], wpfloat] geometric height on half levels
+        z_mc: Field[[CellDim, KDim], wpfloat] output, geometric height defined on full levels
+        horizontal_start:int32 start index of horizontal domain
+        horizontal_end:int32 end index of horizontal domain
+        vertical_start:int32 start index of vertical domain
+        vertical_end:int32 end index of vertical domain
+
+    """
+    average_k_level_up(
+        z_ifc,
+        out=z_mc,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_ddqz_z_half(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    z_mc: Field[[CellDim, KDim], wpfloat],
+    k: Field[[KDim], int32],
+    nlev: int32,
+) -> Field[[CellDim, KDim], wpfloat]:
+    ddqz_z_half = where(
+        (k > int32(0)) & (k < nlev),
+        difference_k_level_down(z_mc),
+        where(k == 0, 2.0 * (z_ifc - z_mc), 2.0 * (z_mc(Koff[-1]) - z_ifc)),
+    )
+    return ddqz_z_half
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=None)
+def compute_ddqz_z_half(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    z_mc: Field[[CellDim, KDim], wpfloat],
+    k: Field[[KDim], int32],
+    ddqz_z_half: Field[[CellDim, KDim], wpfloat],
+    nlev: int32,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute functional determinant of the metrics (is positive) on half levels.
+
+    See mo_vertical_grid.f90
+
+    Args:
+        z_ifc: geometric height on half levels
+        z_mc: geometric height on full levels
+        k: vertical dimension index
+        nlev: total number of levels
+        ddqz_z_half: (output)
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+    """
+    _compute_ddqz_z_half(
+        z_ifc,
+        z_mc,
+        k,
+        nlev,
+        out=ddqz_z_half,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )
+
+
+@field_operator
+def _compute_ddqz_z_full(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
+    ddqz_z_full = difference_k_level_up(z_ifc)
+    inverse_ddqz_z_full = 1.0 / ddqz_z_full
+    return ddqz_z_full, inverse_ddqz_z_full
+
+
+@program(grid_type=GridType.UNSTRUCTURED)
+def compute_ddqz_z_full(
+    z_ifc: Field[[CellDim, KDim], wpfloat],
+    ddqz_z_full: Field[[CellDim, KDim], wpfloat],
+    inv_ddqz_z_full: Field[[CellDim, KDim], wpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    """
+    Compute ddqz_z_full and its inverse inv_ddqz_z_full.
+
+    Functional determinant of the metrics (is positive) on full levels and inverse inverse layer thickness(for runtime optimization).
+    See mo_vertical_grid.f90
+
+    Args:
+        z_ifc: geometric height on half levels
+        ddqz_z_full: (output)
+        inv_ddqz_z_full: (output)
+        horizontal_start: horizontal start index
+        horizontal_end: horizontal end index
+        vertical_start: vertical start index
+        vertical_end: vertical end index
+
+    """
+    _compute_ddqz_z_full(
+        z_ifc,
+        out=(ddqz_z_full, inv_ddqz_z_full),
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )