Tracer advection stencils - Part 1

advection/src/icon4py/advection/face_val_ppm_stencil_05.py

@@ -0,0 +1,67 @@
+# 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 functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import Field
+
+from icon4py.common.dimension import CellDim, KDim, Koff
+
+
+@field_operator
+def _face_val_ppm_stencil_05(
+    p_cc: Field[[CellDim, KDim], float],
+    p_cellhgt_mc_now: Field[[CellDim, KDim], float],
+    z_slope: Field[[CellDim, KDim], float],
+) -> Field[[CellDim, KDim], float]:
+
+    zgeo1 = p_cellhgt_mc_now(Koff[-1]) / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now)
+    zgeo2 = 1.0 / (
+        p_cellhgt_mc_now(Koff[-2])
+        + p_cellhgt_mc_now(Koff[-1])
+        + p_cellhgt_mc_now
+        + p_cellhgt_mc_now(Koff[1])
+    )
+    zgeo3 = (p_cellhgt_mc_now(Koff[-2]) + p_cellhgt_mc_now(Koff[-1])) / (
+        2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now
+    )
+    zgeo4 = (p_cellhgt_mc_now(Koff[1]) + p_cellhgt_mc_now) / (
+        2.0 * p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])
+    )
+
+    p_face = (
+        p_cc(Koff[-1])
+        + zgeo1 * (p_cc - p_cc(Koff[-1]))
+        + zgeo2
+        * (
+            (2.0 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc(Koff[-1]))
+            - zgeo3 * p_cellhgt_mc_now(Koff[-1]) * z_slope
+            + zgeo4 * p_cellhgt_mc_now * z_slope(Koff[-1])
+        )
+    )
+
+    return p_face
+
+
+@program
+def face_val_ppm_stencil_05(
+    p_cc: Field[[CellDim, KDim], float],
+    p_cellhgt_mc_now: Field[[CellDim, KDim], float],
+    z_slope: Field[[CellDim, KDim], float],
+    p_face: Field[[CellDim, KDim], float],
+):
+    _face_val_ppm_stencil_05(
+        p_cc,
+        p_cellhgt_mc_now,
+        z_slope,
+        out=p_face,
+    )

advection/src/icon4py/advection/hor_adv_stencil_01.py

@@ -0,0 +1,68 @@
+# 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 functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import Field, neighbor_sum
+
+from icon4py.common.dimension import (
+    C2CE,
+    C2E,
+    C2EDim,
+    CEDim,
+    CellDim,
+    EdgeDim,
+    KDim,
+)
+
+
+@field_operator
+def _hor_adv_stencil_01(
+    p_mflx_tracer_h: Field[[EdgeDim, KDim], float],
+    deepatmo_divh: Field[[KDim], float],
+    tracer_now: Field[[CellDim, KDim], float],
+    rhodz_now: Field[[CellDim, KDim], float],
+    rhodz_new: Field[[CellDim, KDim], float],
+    geofac_div: Field[[CEDim], float],
+    p_dtime: float,
+) -> Field[[CellDim, KDim], float]:
+    tracer_new = (
+        tracer_now * rhodz_now
+        - p_dtime
+        * deepatmo_divh
+        * neighbor_sum(p_mflx_tracer_h(C2E) * geofac_div(C2CE), axis=C2EDim)
+    ) / rhodz_new
+
+    return tracer_new
+
+
+@program
+def hor_adv_stencil_01(
+    p_mflx_tracer_h: Field[[EdgeDim, KDim], float],
+    deepatmo_divh: Field[[KDim], float],
+    tracer_now: Field[[CellDim, KDim], float],
+    rhodz_now: Field[[CellDim, KDim], float],
+    rhodz_new: Field[[CellDim, KDim], float],
+    geofac_div: Field[[CEDim], float],
+    tracer_new: Field[[CellDim, KDim], float],
+    p_dtime: float,
+):
+    _hor_adv_stencil_01(
+        p_mflx_tracer_h,
+        deepatmo_divh,
+        tracer_now,
+        rhodz_now,
+        rhodz_new,
+        geofac_div,
+        p_dtime,
+        out=tracer_new,
+    )

advection/src/icon4py/advection/step_advection_stencil_03.py

@@ -0,0 +1,39 @@
+# 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 functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import Field, maximum
+
+from icon4py.common.dimension import CellDim, KDim
+
+
+@field_operator
+def _step_advection_stencil_03(
+    p_tracer_now: Field[[CellDim, KDim], float],
+    p_grf_tend_tracer: Field[[CellDim, KDim], float],
+    p_dtime: float,
+) -> Field[[CellDim, KDim], float]:
+    p_tracer_new = maximum(0.0, p_tracer_now + p_dtime * p_grf_tend_tracer)
+    return p_tracer_new
+
+
+@program
+def step_advection_stencil_03(
+    p_tracer_now: Field[[CellDim, KDim], float],
+    p_grf_tend_tracer: Field[[CellDim, KDim], float],
+    p_tracer_new: Field[[CellDim, KDim], float],
+    p_dtime: float,
+):
+    _step_advection_stencil_03(
+        p_tracer_now, p_grf_tend_tracer, p_dtime, out=p_tracer_new
+    )

advection/src/icon4py/advection/upwind_hflux_miura_stencil_01.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 functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import Field, where
+
+from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim
+
+
+@field_operator
+def _upwind_hflux_miura_stencil_01(
+    p_vn: Field[[EdgeDim, KDim], float],
+    p_cc: Field[[CellDim, KDim], float],
+    distv_bary_1: Field[[EdgeDim, KDim], float],
+    distv_bary_2: Field[[EdgeDim, KDim], float],
+    z_grad_1: Field[[CellDim, KDim], float],
+    z_grad_2: Field[[CellDim, KDim], float],
+    p_mass_flx_e: Field[[EdgeDim, KDim], float],
+) -> Field[[EdgeDim, KDim], float]:
+
+    p_out_e = where(
+        p_vn > 0.0,
+        (
+            p_cc(E2C[0])
+            + distv_bary_1 * z_grad_1(E2C[0])
+            + distv_bary_2 * z_grad_2(E2C[0])
+        )
+        * p_mass_flx_e,
+        (
+            p_cc(E2C[1])
+            + distv_bary_1 * z_grad_1(E2C[1])
+            + distv_bary_2 * z_grad_2(E2C[1])
+        )
+        * p_mass_flx_e,
+    )
+
+    return p_out_e
+
+
+@program
+def upwind_hflux_miura_stencil_01(
+    p_vn: Field[[EdgeDim, KDim], float],
+    p_cc: Field[[CellDim, KDim], float],
+    distv_bary_1: Field[[EdgeDim, KDim], float],
+    distv_bary_2: Field[[EdgeDim, KDim], float],
+    z_grad_1: Field[[CellDim, KDim], float],
+    z_grad_2: Field[[CellDim, KDim], float],
+    p_mass_flx_e: Field[[EdgeDim, KDim], float],
+    p_out_e: Field[[EdgeDim, KDim], float],
+):
+    _upwind_hflux_miura_stencil_01(
+        p_vn,
+        p_cc,
+        distv_bary_1,
+        distv_bary_2,
+        z_grad_1,
+        z_grad_2,
+        p_mass_flx_e,
+        out=p_out_e,
+    )

advection/src/icon4py/advection/v_limit_prbl_sm_stencil_01.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 functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import (
+    Field,
+    FieldOffset,
+    abs,
+    broadcast,
+    minimum,
+    where,
+)
+
+from icon4py.common.dimension import CellDim, KDim
+
+
+Koff = FieldOffset("Koff", source=KDim, target=(KDim,))
+
+
+@field_operator
+def _v_limit_prbl_sm_stencil_01(
+    p_face: Field[[CellDim, KDim], float],
+    p_cc: Field[[CellDim, KDim], float],
+) -> Field[[CellDim, KDim], float]:
+
+    z_delta = p_face - p_face(Koff[1])
+    z_a6i = -6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1])))
+
+    q_face_up = broadcast(1.0, (CellDim, KDim))
+    q_face_low = broadcast(1.0, (CellDim, KDim))
+
+    q_face_up, q_face_low = where(
+        abs(z_delta) < z_a6i,
+        where(
+            (p_cc < minimum(p_face, p_face(Koff[1]))),
+            (p_cc, p_cc),
+            where(
+                p_face > p_face(Koff[1]),
+                (3.0 * p_cc - 2.0 * p_face(Koff[1]), p_face(Koff[1])),
+                (p_face, 3.0 * p_cc - 2.0 * p_face),
+            ),
+        ),
+        (p_face, p_face(Koff[1])),
+    )
+
+    return q_face_up, q_face_low
+
+
+@program
+def v_limit_prbl_sm_stencil_01(
+    p_face: Field[[CellDim, KDim], float],
+    p_cc: Field[[CellDim, KDim], float],
+    p_face_up: Field[[CellDim, KDim], float],
+    p_face_low: Field[[CellDim, KDim], float],
+):
+    _v_limit_prbl_sm_stencil_01(
+        p_face,
+        p_cc,
+        out=(p_face_up, p_face_low),
+    )

advection/src/icon4py/advection/vert_adv_stencil_01.py

@@ -0,0 +1,59 @@
+# 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 functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import Field
+
+from icon4py.common.dimension import CellDim, KDim, Koff
+
+
+@field_operator
+def _vert_adv_stencil_01(
+    tracer_now: Field[[CellDim, KDim], float],
+    rhodz_now: Field[[CellDim, KDim], float],
+    p_mflx_tracer_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    rhodz_new: Field[[CellDim, KDim], float],
+    p_dtime: float,
+) -> Field[[CellDim, KDim], float]:
+    tracer_new = (
+        tracer_now * rhodz_now
+        + p_dtime
+        * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu)
+    ) / rhodz_new
+
+    return tracer_new
+
+
+@program
+def vert_adv_stencil_01(
+    tracer_now: Field[[CellDim, KDim], float],
+    rhodz_now: Field[[CellDim, KDim], float],
+    p_mflx_tracer_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    rhodz_new: Field[[CellDim, KDim], float],
+    tracer_new: Field[[CellDim, KDim], float],
+    p_dtime: float,
+):
+    _vert_adv_stencil_01(
+        tracer_now,
+        rhodz_now,
+        p_mflx_tracer_v,
+        deepatmo_divzl,
+        deepatmo_divzu,
+        rhodz_new,
+        p_dtime,
+        out=tracer_new,
+    )