Initial condition for jw test

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

@@ -30,6 +30,8 @@
 SPECIFIC_HEAT_CONSTANT_VOLUME: Final[wpfloat] = CPD - RD
 CVD: Final[wpfloat] = SPECIFIC_HEAT_CONSTANT_VOLUME
 CVD_O_RD: Final[wpfloat] = CVD / RD
+RD_O_CPD: Final[wpfloat] = RD / CPD
+CPD_O_RD: Final[wpfloat] = CPD / RD
 
 #: Gas constant for water vapor [J/K/kg], rv in ICON.
 GAS_CONSTANT_WATER_VAPOR: Final[wpfloat] = 461.51
@@ -38,6 +40,7 @@
 #: Av. gravitational acceleration [m/s^2]
 GRAVITATIONAL_ACCELERATION: Final[wpfloat] = 9.80665
 GRAV: Final[wpfloat] = GRAVITATIONAL_ACCELERATION
+GRAV_O_RD: Final[wpfloat] = GRAV / RD
 
 #: reference pressure for Exner function [Pa]
 REFERENCE_PRESSURE: Final[wpfloat] = 100000.0
@@ -50,6 +53,9 @@
 # average earth radius in [m]
 EARTH_RADIUS: Final[float] = 6.371229e6
 
+#: Earth angular velocity [rad/s]
+EARTH_ANGULAR_VELOCITY: Final[wpfloat] = 7.29212e-5
+
 #: sea level temperature for reference atmosphere [K]
 SEA_LEVEL_TEMPERATURE: Final[wpfloat] = 288.15
 T0SL_BG: Final[wpfloat] = SEA_LEVEL_TEMPERATURE

model/common/src/icon4py/model/common/diagnostic_calculations/stencils/__init__.py

@@ -0,0 +1,13 @@
+# 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/diagnostic_calculations/stencils/diagnose_pressure.py

@@ -0,0 +1,73 @@
+# 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.common import GridType
+from gt4py.next.ffront.decorator import field_operator, program, scan_operator
+from gt4py.next.ffront.fbuiltins import Field, exp, int32, sqrt
+
+from icon4py.model.common.dimension import CellDim, KDim
+from icon4py.model.common.settings import backend
+from icon4py.model.common.type_alias import vpfloat, wpfloat
+
+
+@scan_operator(axis=KDim, forward=False, init=(0.0, 0.0, True))
+def _scan_pressure(
+    state: tuple[vpfloat, vpfloat, bool],
+    ddqz_z_full: vpfloat,
+    temperature: vpfloat,
+    pressure_sfc: vpfloat,
+    grav_o_rd: wpfloat,
+):
+    pressure_interface = (
+        pressure_sfc * exp(-grav_o_rd * ddqz_z_full / temperature)
+        if state[2]
+        else state[1] * exp(-grav_o_rd * ddqz_z_full / temperature)
+    )
+    pressure = (
+        sqrt(pressure_sfc * pressure_interface) if state[2] else sqrt(state[1] * pressure_interface)
+    )
+    return pressure, pressure_interface, False
+
+
+@field_operator
+def _diagnose_pressure(
+    ddqz_z_full: Field[[CellDim, KDim], wpfloat],
+    temperature: Field[[CellDim, KDim], vpfloat],
+    pressure_sfc: Field[[CellDim], vpfloat],
+    grav_o_rd: wpfloat,
+) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]:
+    pressure, pressure_ifc, _ = _scan_pressure(ddqz_z_full, temperature, pressure_sfc, grav_o_rd)
+    return pressure, pressure_ifc
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+def diagnose_pressure(
+    ddqz_z_full: Field[[CellDim, KDim], wpfloat],
+    temperature: Field[[CellDim, KDim], vpfloat],
+    pressure_sfc: Field[[CellDim], vpfloat],
+    pressure: Field[[CellDim, KDim], vpfloat],
+    pressure_ifc: Field[[CellDim, KDim], vpfloat],
+    grav_o_rd: wpfloat,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _diagnose_pressure(
+        ddqz_z_full,
+        temperature,
+        pressure_sfc,
+        grav_o_rd,
+        out=(pressure, pressure_ifc),
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )

model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py

@@ -0,0 +1,70 @@
+# 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.common import GridType
+from gt4py.next.ffront.decorator import field_operator, program
+from gt4py.next.ffront.fbuiltins import Field, exp, int32, log
+
+from icon4py.model.common.dimension import CellDim, KDim, Koff
+from icon4py.model.common.settings import backend
+from icon4py.model.common.type_alias import vpfloat, wpfloat
+
+
+@field_operator
+def _diagnose_surface_pressure(
+    exner: Field[[CellDim, KDim], vpfloat],
+    temperature: Field[[CellDim, KDim], vpfloat],
+    ddqz_z_full: Field[[CellDim, KDim], wpfloat],
+    cpd_o_rd: wpfloat,
+    p0ref: wpfloat,
+    grav_o_rd: wpfloat,
+) -> Field[[CellDim, KDim], vpfloat]:
+    pressure_sfc = p0ref * exp(
+        cpd_o_rd * log(exner(Koff[-3]))
+        + grav_o_rd
+        * (
+            ddqz_z_full(Koff[-1]) / temperature(Koff[-1])
+            + ddqz_z_full(Koff[-2]) / temperature(Koff[-2])
+            + 0.5 * ddqz_z_full(Koff[-3]) / temperature(Koff[-3])
+        )
+    )
+    return pressure_sfc
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+def diagnose_surface_pressure(
+    exner: Field[[CellDim, KDim], vpfloat],
+    temperature: Field[[CellDim, KDim], vpfloat],
+    ddqz_z_full: Field[[CellDim, KDim], wpfloat],
+    pressure_sfc: Field[[CellDim, KDim], vpfloat],
+    cpd_o_rd: wpfloat,
+    p0ref: wpfloat,
+    grav_o_rd: wpfloat,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _diagnose_surface_pressure(
+        exner,
+        temperature,
+        ddqz_z_full,
+        cpd_o_rd,
+        p0ref,
+        grav_o_rd,
+        out=pressure_sfc,
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py

@@ -0,0 +1,50 @@
+# 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.common import GridType
+from gt4py.next.ffront.decorator import field_operator, program
+from gt4py.next.ffront.fbuiltins import Field, int32
+
+from icon4py.model.common.dimension import CellDim, KDim
+from icon4py.model.common.settings import backend
+from icon4py.model.common.type_alias import vpfloat
+
+
+@field_operator
+def _diagnose_temperature(
+    theta_v: Field[[CellDim, KDim], vpfloat],
+    exner: Field[[CellDim, KDim], vpfloat],
+) -> Field[[CellDim, KDim], vpfloat]:
+    temperature = theta_v * exner
+    return temperature
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+def diagnose_temperature(
+    theta_v: Field[[CellDim, KDim], vpfloat],
+    exner: Field[[CellDim, KDim], vpfloat],
+    temperature: Field[[CellDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _diagnose_temperature(
+        theta_v,
+        exner,
+        out=(temperature),
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/common/src/icon4py/model/common/dimension.py

@@ -24,7 +24,6 @@
 CECDim = Dimension("CEC")
 ECDim = Dimension("EC")
 ECVDim = Dimension("ECV")
-CECDim = Dimension("CEC")
 CECECDim = Dimension("CECEC")
 E2CDim = Dimension("E2C", DimensionKind.LOCAL)
 E2VDim = Dimension("E2V", DimensionKind.LOCAL)
@@ -38,6 +37,7 @@
 E2C2EODim = Dimension("E2C2EO", DimensionKind.LOCAL)
 E2C2EDim = Dimension("E2C2E", DimensionKind.LOCAL)
 C2E2CDim = Dimension("C2E2C", DimensionKind.LOCAL)
+C2E2C2EDim = Dimension("C2E2C2E", DimensionKind.LOCAL)
 C2E2C2E2CDim = Dimension("C2E2C2E2C", DimensionKind.LOCAL)
 E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim, E2CDim))
 C2E = FieldOffset("C2E", source=EdgeDim, target=(CellDim, C2EDim))
@@ -54,8 +54,8 @@
 E2C2EO = FieldOffset("E2C2EO", source=EdgeDim, target=(EdgeDim, E2C2EODim))
 E2C2E = FieldOffset("E2C2E", source=EdgeDim, target=(EdgeDim, E2C2EDim))
 C2E2C = FieldOffset("C2E2C", source=CellDim, target=(CellDim, C2E2CDim))
+C2E2C2E = FieldOffset("C2E2C2E", source=EdgeDim, target=(CellDim, C2E2C2EDim))
 C2E2C2E2C = FieldOffset("C2E2C2E2C", source=CellDim, target=(CellDim, C2E2C2E2CDim))
-C2CEC = FieldOffset("C2CEC", source=CECDim, target=(CellDim, C2E2CDim))
 C2CECEC = FieldOffset("C2CECEC", source=CECECDim, target=(CellDim, C2E2C2E2CDim))
 
 Koff = FieldOffset("Koff", source=KDim, target=(KDim,))