Tracer advection stencil porting ml -> tracer advection stencils (#98)

* adding hflx_limiter_mo_stencil_02

* changes order of args and return a tuple

* Revert "changes order of args and return a tuple"

This reverts commit 99638df.

* porting step_advection_stencil_01

* porting step_advection_stencil_02

* porting upwind_hflux_miura_stencil_02

* hflx_limiter_mo_stencil_04

* fix step_advection_stencil_01

* hflx_limiter_mo_stencil_03

* clean up: move upwind_hflux_miura_stencil_02.py

* add test for conditional branch in hflx_limiter_mo_stencil_04

* switch to using min_over in hflx_limiter_mo_stencil_03

* refactor hflx_limiter_mo_stencil_04, fix tests for Koffset

* add upwind_vflux_ppm_stencil_01

* upwind_hflux_miura_stencil_02 - fix sparse field issue, split into separate sums instead of using neighbor_sum

* remove commented out code

* fix: upwind_vflux_ppm_stencil_01

* fix: hflx_limiter_mo_stencil_01: return tuple, use int32

* hflx_limiter_pd_stencil_02.py: use int fields

advection/src/icon4py/advection/hflx_limiter_mo_stencil_02.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.common import Field
+from functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import int32, maximum, minimum, where
+
+from icon4py.common.dimension import CellDim, KDim
+
+
+@field_operator
+def _hflx_limiter_mo_stencil_02(
+    refin_ctrl: Field[[CellDim], int32],
+    p_cc: Field[[CellDim, KDim], float],
+    z_tracer_new_low: Field[[CellDim, KDim], float],
+    lo_bound: int32,
+    hi_bound: int32,
+) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
+    condition = (refin_ctrl == lo_bound) | (refin_ctrl == hi_bound)
+    z_tracer_new_tmp = where(
+        condition,
+        minimum(1.1 * p_cc, maximum(0.9 * p_cc, z_tracer_new_low)),
+        z_tracer_new_low,
+    )
+    return where(
+        condition,
+        (maximum(p_cc, z_tracer_new_tmp), minimum(p_cc, z_tracer_new_tmp)),
+        (z_tracer_new_low, z_tracer_new_low),
+    )
+
+
+@program
+def hflx_limiter_mo_stencil_02(
+    refin_ctrl: Field[[CellDim], int32],
+    p_cc: Field[[CellDim, KDim], float],
+    z_tracer_new_low: Field[[CellDim, KDim], float],
+    lo_bound: int32,
+    hi_bound: int32,
+    z_tracer_max: Field[[CellDim, KDim], float],
+    z_tracer_min: Field[[CellDim, KDim], float],
+):
+    _hflx_limiter_mo_stencil_02(
+        refin_ctrl,
+        p_cc,
+        z_tracer_new_low,
+        lo_bound,
+        hi_bound,
+        out=(z_tracer_max, z_tracer_min),
+    )

advection/src/icon4py/advection/hflx_limiter_mo_stencil_03.py

@@ -0,0 +1,92 @@
+# 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.common import Field
+from functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import max_over, maximum, min_over, minimum
+
+from icon4py.common.dimension import C2E2C, C2E2CDim, CellDim, KDim
+
+
+@field_operator
+def _hflx_limiter_mo_stencil_03_min_max(
+    z_tracer_max: Field[[CellDim, KDim], float],
+    z_tracer_min: Field[[CellDim, KDim], float],
+    beta_fct: float,
+    r_beta_fct: float,
+) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
+    z_max = beta_fct * maximum(
+        max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max
+    )
+    z_min = r_beta_fct * minimum(
+        min_over(z_tracer_min(C2E2C), axis=C2E2CDim), z_tracer_min
+    )
+    return z_max, z_min
+
+
+@program
+def hflx_limiter_mo_stencil_03_min_max(
+    z_tracer_max: Field[[CellDim, KDim], float],
+    z_tracer_min: Field[[CellDim, KDim], float],
+    beta_fct: float,
+    r_beta_fct: float,
+    z_max: Field[[CellDim, KDim], float],
+    z_min: Field[[CellDim, KDim], float],
+):
+    _hflx_limiter_mo_stencil_03_min_max(
+        z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_max, z_min)
+    )
+
+
+@field_operator
+def _hflx_limiter_mo_stencil_03(
+    z_mflx_anti_in: Field[[CellDim, KDim], float],
+    z_mflx_anti_out: Field[[CellDim, KDim], float],
+    z_tracer_new_low: Field[[CellDim, KDim], float],
+    z_max: Field[[CellDim, KDim], float],
+    z_min: Field[[CellDim, KDim], float],
+    dbl_eps: float,
+) -> tuple[[Field[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
+    r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps)
+    r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps)
+    return r_p, r_m
+
+
+@program
+def hflx_limiter_mo_stencil_03(
+    z_tracer_max: Field[[CellDim, KDim], float],
+    z_tracer_min: Field[[CellDim, KDim], float],
+    beta_fct: float,
+    r_beta_fct: float,
+    z_max: Field[[CellDim, KDim], float],
+    z_min: Field[[CellDim, KDim], float],
+    z_mflx_anti_in: Field[[CellDim, KDim], float],
+    z_mflx_anti_out: Field[[CellDim, KDim], float],
+    z_tracer_new_low: Field[[CellDim, KDim], float],
+    dbl_eps: float,
+    r_p: Field[[CellDim, KDim], float],
+    r_m: Field[[CellDim, KDim], float],
+):
+
+    _hflx_limiter_mo_stencil_03_min_max(
+        z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_min, z_max)
+    )
+    _hflx_limiter_mo_stencil_03(
+        z_mflx_anti_in,
+        z_mflx_anti_out,
+        z_tracer_new_low,
+        z_max,
+        z_min,
+        dbl_eps,
+        out=(r_p, r_m),
+    )

advection/src/icon4py/advection/hflx_limiter_mo_stencil_04.py

@@ -0,0 +1,44 @@
+# 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.common import Field
+from functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import minimum, where
+
+from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim
+
+
+@field_operator
+def _hflx_limiter_mo_stencil_04(
+    z_anti: Field[[EdgeDim, KDim], float],
+    r_m: Field[[CellDim, KDim], float],
+    r_p: Field[[CellDim, KDim], float],
+    z_mflx_low: Field[[EdgeDim, KDim], float],
+) -> Field[[EdgeDim, KDim], float]:
+    r_frac = where(
+        z_anti >= 0.0,
+        minimum(r_m(E2C[0]), r_p(E2C[1])),
+        minimum(r_m(E2C[1]), r_p(E2C[0])),
+    )
+    return z_mflx_low + minimum(1.0, r_frac) * z_anti
+
+
+@program
+def hflx_limiter_mo_stencil_04(
+    z_anti: Field[[EdgeDim, KDim], float],
+    r_m: Field[[CellDim, KDim], float],
+    r_p: Field[[CellDim, KDim], float],
+    z_mflx_low: Field[[EdgeDim, KDim], float],
+    p_mflx_tracer_h: Field[[EdgeDim, KDim], float],
+):
+    _hflx_limiter_mo_stencil_04(z_anti, r_m, r_p, z_mflx_low, out=p_mflx_tracer_h)

advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py

@@ -12,23 +12,23 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from functional.ffront.decorator import field_operator, program
-from functional.ffront.fbuiltins import Field, where
+from functional.ffront.fbuiltins import Field, int32, where
 
 from icon4py.common.dimension import E2C, CellDim, EdgeDim, KDim
 
 
 @field_operator
 def _hflx_limiter_pd_stencil_02(
-    refin_ctrl: Field[[EdgeDim], float],
+    refin_ctrl: Field[[EdgeDim], int32],
     r_m: Field[[CellDim, KDim], float],
     p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float],
-    bound: float,
+    bound: int32,
 ) -> Field[[EdgeDim, KDim], float]:
     p_mflx_tracer_h_out = where(
         refin_ctrl == bound,
         p_mflx_tracer_h_in,
         where(
-            (p_mflx_tracer_h_in > 0.0) | (p_mflx_tracer_h_in == 0.0),
+            p_mflx_tracer_h_in >= 0.0,
             p_mflx_tracer_h_in * r_m(E2C[0]),
             p_mflx_tracer_h_in * r_m(E2C[1]),
         ),
@@ -38,10 +38,10 @@ def _hflx_limiter_pd_stencil_02(
 
 @program
 def hflx_limiter_pd_stencil_02(
-    refin_ctrl: Field[[EdgeDim], float],
+    refin_ctrl: Field[[EdgeDim], int32],
     r_m: Field[[CellDim, KDim], float],
     p_mflx_tracer_h: Field[[EdgeDim, KDim], float],
-    bound: float,
+    bound: int32,
 ):
     _hflx_limiter_pd_stencil_02(
         refin_ctrl,

advection/src/icon4py/advection/step_advection_stencil_01.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 functional.common import Field
+from functional.ffront.decorator import field_operator, program
+
+from icon4py.common.dimension import CellDim, KDim, Koff
+
+
+@field_operator
+def _step_advection_stencil_01(
+    rhodz_ast: Field[[CellDim, KDim], float],
+    p_mflx_contra_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    p_dtime: float,
+) -> Field[[CellDim, KDim], float]:
+    k_offset_up_low = p_dtime * (
+        p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu
+    )
+    return rhodz_ast + k_offset_up_low
+
+
+@program
+def step_advection_stencil_01(
+    rhodz_ast: Field[[CellDim, KDim], float],
+    p_mflx_contra_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    p_dtime: float,
+    rhodz_ast2: Field[[CellDim, KDim], float],
+):
+    _step_advection_stencil_01(
+        rhodz_ast,
+        p_mflx_contra_v,
+        deepatmo_divzl,
+        deepatmo_divzu,
+        p_dtime,
+        out=rhodz_ast2,
+    )

advection/src/icon4py/advection/step_advection_stencil_02.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 functional.common import Field
+from functional.ffront.decorator import field_operator, program
+from functional.ffront.fbuiltins import maximum
+
+from icon4py.common.dimension import CellDim, KDim, Koff
+
+
+@field_operator
+def _step_advection_stencil_02(
+    p_rhodz_new: Field[[CellDim, KDim], float],
+    p_mflx_contra_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    p_dtime: float,
+) -> Field[[CellDim, KDim], float]:
+    return maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * (
+        p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu
+    )
+
+
+@program
+def step_advection_stencil_02(
+    p_rhodz_new: Field[[CellDim, KDim], float],
+    p_mflx_contra_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    p_dtime: float,
+    rhodz_ast2: Field[[CellDim, KDim], float],
+):
+    _step_advection_stencil_02(
+        p_rhodz_new,
+        p_mflx_contra_v,
+        deepatmo_divzl,
+        deepatmo_divzu,
+        p_dtime,
+        out=rhodz_ast2,
+    )

advection/src/icon4py/advection/upwind_hflux_miura_stencil_02.py

@@ -0,0 +1,49 @@
+# 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.common import Field
+from functional.ffront.decorator import field_operator, program
+
+from icon4py.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim
+
+
+@field_operator
+def _upwind_hflux_miura_stencil_02(
+    p_cc: Field[[CellDim, KDim], float],
+    lsq_pseudoinv_1: Field[[CECDim], float],
+    lsq_pseudoinv_2: Field[[CECDim], float],
+) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]:
+    p_coeff_1 = (
+        lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc)
+        + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc)
+        + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc)
+    )
+    p_coeff_2 = (
+        lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc)
+        + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc)
+        + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc)
+    )
+    return p_coeff_1, p_coeff_2
+
+
+@program
+def upwind_hflux_miura_stencil_02(
+    p_cc: Field[[CellDim, KDim], float],
+    lsq_pseudoinv_1: Field[[CECDim], float],
+    lsq_pseudoinv_2: Field[[CECDim], float],
+    p_coeff_1: Field[[CellDim, KDim], float],
+    p_coeff_2: Field[[CellDim, KDim], float],
+):
+    _upwind_hflux_miura_stencil_02(
+        p_cc, lsq_pseudoinv_1, lsq_pseudoinv_2, out=(p_coeff_1, p_coeff_2)
+    )