Fix hflx limiter pd stencil (#91)

* run pre-commit
* fix hflx_limiter_pd_stencil_01

advection/setup.py

@@ -1,3 +1,4 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
 #
 # Copyright (c) 2022, ETH Zurich and MeteoSwiss
 # All rights reserved.

advection/src/icon4py/advection/__init__.py

@@ -10,4 +10,3 @@
 # distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
-

advection/src/icon4py/advection/hflx_limiter_pd_stencil_01.py

@@ -12,45 +12,49 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 from functional.ffront.decorator import field_operator, program
-from functional.ffront.fbuiltins import Field, maximum, minimum, neighbor_sum, broadcast
-from icon4py.common.dimension import CellDim, EdgeDim, KDim, C2EDim, C2E
+from functional.ffront.fbuiltins import Field, broadcast, maximum, minimum
+
+from icon4py.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim
+
 
 @field_operator
 def _hflx_limiter_pd_stencil_01(
-    geofac_div: Field[[CellDim, C2EDim], float],
+    geofac_div: Field[[CEDim], float],
     p_cc: Field[[CellDim, KDim], float],
     p_rhodz_now: Field[[CellDim, KDim], float],
     p_mflx_tracer_h: Field[[EdgeDim, KDim], float],
-    p_dtime:float,
+    p_dtime: float,
     dbl_eps: float,
 ) -> Field[[CellDim, KDim], float]:
-#    p_m: Field[[CellDim, KDim], float]
-#    p_m = maximum(0.0, p_mflx_tracer_h(C2E[0])*geofac_div(0)*p_dtime) 
-#    p_m = maximum(0.0, p_mflx_tracer_h(C2E[0])*geofac_div(0)*p_dtime) + maximum(0.0, p_mflx_tracer_h(C2E[1])*geofac_div(1)*p_dtime) + maximum(0.0, p_mflx_tracer_h(C2E[2])*geofac_div(2)*p_dtime) 
-    #p_m = neighbor_sum(a_a, axis=C2EDim)
-    #r_m = minimum(1.0, (p_cc*p_rhodz_now)/(p_m+dbl_eps))
-#    r_m = minimum(1.0, (p_cc*p_rhodz_now)/(1.0+dbl_eps))
-
-    p_m = neighbor_sum(maximum(broadcast(0.0, (CellDim, KDim)), p_mflx_tracer_h(C2E)*geofac_div*p_dtime), axis=C2EDim)
-    r_m = minimum(broadcast(1.0, (CellDim, KDim)), (p_cc*p_rhodz_now)/(p_m+dbl_eps))
+    zero = broadcast(0.0, (CellDim, KDim))
+
+    pm_0 = maximum(zero, p_mflx_tracer_h(C2E[0]) * geofac_div(C2CE[0]) * p_dtime)
+    pm_1 = maximum(zero, p_mflx_tracer_h(C2E[1]) * geofac_div(C2CE[1]) * p_dtime)
+    pm_2 = maximum(zero, p_mflx_tracer_h(C2E[2]) * geofac_div(C2CE[2]) * p_dtime)
+    p_m = pm_0 + pm_1 + pm_2
+    r_m = minimum(
+        broadcast(1.0, (CellDim, KDim)), (p_cc * p_rhodz_now) / (p_m + dbl_eps)
+    )
+
     return r_m
 
+
 @program
 def hflx_limiter_pd_stencil_01(
-    geofac_div: Field[[CellDim, C2EDim], float],
+    geofac_div: Field[[CEDim], float],
     p_cc: Field[[CellDim, KDim], float],
     p_rhodz_now: Field[[CellDim, KDim], float],
     p_mflx_tracer_h: Field[[EdgeDim, KDim], float],
     r_m: Field[[CellDim, KDim], float],
-    p_dtime:float,
+    p_dtime: float,
     dbl_eps: float,
 ):
     _hflx_limiter_pd_stencil_01(
-    geofac_div,
-    p_cc,
-    p_rhodz_now,
-    p_mflx_tracer_h,
-    p_dtime,
-    dbl_eps,
-    out=r_m,
+        geofac_div,
+        p_cc,
+        p_rhodz_now,
+        p_mflx_tracer_h,
+        p_dtime,
+        dbl_eps,
+        out=r_m,
     )

advection/src/icon4py/advection/hflx_limiter_pd_stencil_02.py

@@ -14,26 +14,28 @@
 from functional.ffront.decorator import field_operator, program
 from functional.ffront.fbuiltins import Field, where
 
-from icon4py.common.dimension import CellDim, EdgeDim, KDim, E2C
+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], float],
     r_m: Field[[CellDim, KDim], float],
     p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float],
     bound: float,
 ) -> Field[[EdgeDim, KDim], float]:
-    p_mflx_tracer_h_out = where (
+    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),
+        where(
+            (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]),
         ),
     )
     return p_mflx_tracer_h_out
 
+
 @program
 def hflx_limiter_pd_stencil_02(
     refin_ctrl: Field[[EdgeDim], float],
@@ -43,9 +45,9 @@ def hflx_limiter_pd_stencil_02(
     bound: float,
 ):
     _hflx_limiter_pd_stencil_02(
-    refin_ctrl,
-    r_m,
-    p_mflx_tracer_h_in,
-    bound,
-    out=p_mflx_tracer_h_out,
-)
+        refin_ctrl,
+        r_m,
+        p_mflx_tracer_h_in,
+        bound,
+        out=p_mflx_tracer_h_out,
+    )

advection/src/icon4py/advection/rbf_intp_edge_stencil_01.py

@@ -19,17 +19,17 @@
 
 @field_operator
 def _rbf_intp_edge_stencil_01(
-    p_vn_in: Field[[EdgeDim, KDim], float], 
-   # ptr_coeff: Field[[EdgeDim, E2CDim, E2C2EDim], float]
+    p_vn_in: Field[[EdgeDim, KDim], float],
     ptr_coeff: Field[[EdgeDim, E2C2EDim], float],
 ) -> Field[[EdgeDim, KDim], float]:
-    p_vt_out = neighbor_sum(p_vn_in(E2C2E) * ptr_coeff, axis=E2C2EDim) 
+    p_vt_out = neighbor_sum(p_vn_in(E2C2E) * ptr_coeff, axis=E2C2EDim)
     return p_vt_out
 
+
 @program
 def rbf_intp_edge_stencil_01(
     p_vn_in: Field[[EdgeDim, KDim], float],
     ptr_coeff: Field[[EdgeDim, E2C2EDim], float],
     p_vt_out: Field[[EdgeDim, KDim], float],
 ):
-    _rbf_intp_edge_stencil_01(p_vn_in, ptr_coeff, out=p_vt_out)  
+    _rbf_intp_edge_stencil_01(p_vn_in, ptr_coeff, out=p_vt_out)

advection/src/icon4py/advection/step_advection_stencil_04.py

@@ -16,20 +16,24 @@
 
 from icon4py.common.dimension import CellDim, KDim
 
+
 @field_operator
 def _step_advection_stencil_04(
     p_tracer_now: Field[[CellDim, KDim], float],
     p_tracer_new: Field[[CellDim, KDim], float],
     p_dtime: float,
 ) -> Field[[CellDim, KDim], float]:
-    opt_ddt_tracer_adv=(p_tracer_new-p_tracer_now)/p_dtime
+    opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime
     return opt_ddt_tracer_adv
 
+
 @program
 def step_advection_stencil_04(
     p_tracer_now: Field[[CellDim, KDim], float],
     p_tracer_new: Field[[CellDim, KDim], float],
     opt_ddt_tracer_adv: Field[[CellDim, KDim], float],
     p_dtime: float,
 ):
-    _step_advection_stencil_04(p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv)
+    _step_advection_stencil_04(
+        p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv
+    )

advection/tests/test_hflx_limiter_pd_stencil_01.py

@@ -11,14 +11,16 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-import numpy as np 
+import numpy as np
+from functional.iterator.embedded import StridedNeighborOffsetProvider
 
 from icon4py.advection.hflx_limiter_pd_stencil_01 import (
     hflx_limiter_pd_stencil_01,
 )
-from icon4py.common.dimension import CellDim, EdgeDim, KDim, C2EDim
+from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim
 from icon4py.testutils.simple_mesh import SimpleMesh
-from icon4py.testutils.utils import random_field, zero_field
+from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field
+
 
 def hflx_limiter_pd_stencil_01_numpy(
     c2e: np.array,
@@ -29,16 +31,17 @@ def hflx_limiter_pd_stencil_01_numpy(
     p_dtime,
     dbl_eps,
 ):
-#    p_m: np.array
-    p_mflx_tracer_h_c2e = p_mflx_tracer_h[c2e] 
-    #p_m = max(0.0, geofac_div[:, 0]*p_mflx_tracer_h_c2e[:, 0]*p_dtime) + max(0.0, geofac_div[:, 1]*p_mflx_tracer_h_c2e[:, 1]*p_dtime) + max(0.0, geofac_div[:, 2]*p_mflx_tracer_h_c2e[:, 2]*p_dtime)
-    geofac_div=np.expand_dims(geofac_div,axis=-1)
-#    p_m = max(0.0, geofac_div[:, 0]*p_mflx_tracer_h_c2e[:, 0]*p_dtime) 
-    p_m = np.sum(np.maximum(np.zeros(p_mflx_tracer_h_c2e.shape,dtype=float), p_mflx_tracer_h_c2e*geofac_div*p_dtime))
-    r_m = np.minimum(np.ones(p_cc.shape,dtype=float),(p_cc*p_rhodz_now)/(p_m+dbl_eps))
+    geofac_div = np.expand_dims(geofac_div, axis=-1)
+    p_m_0 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 0]] * geofac_div[:, 0] * p_dtime)
+    p_m_1 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 1]] * geofac_div[:, 1] * p_dtime)
+    p_m_2 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 2]] * geofac_div[:, 2] * p_dtime)
+
+    p_m = p_m_0 + p_m_1 + p_m_2
+    r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps))
 
     return r_m
 
+
 def test_hflx_limiter_pd_stencil_01():
     mesh = SimpleMesh()
     geofac_div = random_field(mesh, CellDim, C2EDim)
@@ -60,14 +63,15 @@ def test_hflx_limiter_pd_stencil_01():
     )
 
     hflx_limiter_pd_stencil_01(
-        geofac_div,
+        as_1D_sparse_field(geofac_div, CEDim),
         p_cc,
         p_rhodz_now,
         p_mflx_tracer_h,
         r_m,
         p_dtime,
         dbl_eps,
         offset_provider={
+            "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e),
             "C2E": mesh.get_c2e_offset_provider(),
         },
     )

advection/tests/test_hflx_limiter_pd_stencil_02.py

@@ -11,43 +11,45 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-import numpy as np 
+import numpy as np
 
 from icon4py.advection.hflx_limiter_pd_stencil_02 import (
     hflx_limiter_pd_stencil_02,
 )
 from icon4py.common.dimension import CellDim, EdgeDim, KDim
 from icon4py.testutils.simple_mesh import SimpleMesh
-from icon4py.testutils.utils import random_field, zero_field
+from icon4py.testutils.utils import random_field
+
 
 def hflx_limiter_pd_stencil_02_numpy(
     e2c: np.array,
-    refin_ctrl: np.array, 
+    refin_ctrl: np.array,
     r_m: np.array,
     p_mflx_tracer_h_in: np.array,
     bound,
 ):
     r_m_e2c = r_m[e2c]
     refin_ctrl = np.expand_dims(refin_ctrl, axis=-1)
-    p_mflx_tracer_h_out = np.where (
+    p_mflx_tracer_h_out = np.where(
         refin_ctrl != bound,
-        np.where (
-            p_mflx_tracer_h_in >=0,
+        np.where(
+            p_mflx_tracer_h_in >= 0,
             p_mflx_tracer_h_in * r_m_e2c[:, 0],
             p_mflx_tracer_h_in * r_m_e2c[:, 1],
         ),
         p_mflx_tracer_h_in,
     )
     return p_mflx_tracer_h_out
 
+
 def test_hflx_limiter_pd_stencil_02():
     mesh = SimpleMesh()
 
     refin_ctrl = random_field(mesh, EdgeDim)
     r_m = random_field(mesh, CellDim, KDim)
     p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim)
     p_mflx_tracer_h_out = random_field(mesh, EdgeDim, KDim)
-    bound=np.float64(7.0)
+    bound = np.float64(7.0)
 
     ref = hflx_limiter_pd_stencil_02_numpy(
         mesh.e2c,
@@ -58,7 +60,7 @@ def test_hflx_limiter_pd_stencil_02():
     )
 
     hflx_limiter_pd_stencil_02(
-        refin_ctrl, 
+        refin_ctrl,
         r_m,
         p_mflx_tracer_h_in,
         p_mflx_tracer_h_out,

advection/tests/test_rbf_intp_edge_stencil_01.py

@@ -11,34 +11,31 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-import numpy as np 
+import numpy as np
 
-from icon4py.advection.rbf_intp_edge_stencil_01 import (
-    rbf_intp_edge_stencil_01,
-)
-from icon4py.common.dimension import EdgeDim, E2C2EDim, KDim
+from icon4py.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01
+from icon4py.common.dimension import E2C2EDim, EdgeDim, KDim
 from icon4py.testutils.simple_mesh import SimpleMesh
 from icon4py.testutils.utils import random_field, zero_field
 
+
 def rbf_intp_edge_stencil_01_numpy(
     e2c2e: np.array,
-    p_vn_in: np.array, 
+    p_vn_in: np.array,
     ptr_coeff: np.array,
 ) -> np.array:
 
     ptr_coeff = np.expand_dims(ptr_coeff, axis=-1)
     p_vt_out = np.sum(p_vn_in[e2c2e] * ptr_coeff, axis=1)
-#    p_vn_ptr = p_vn_in[e2c2e] * ptr_coeff
-#    p_vt_out = np.sum(p_vn_ptr, axis=1) 
-
     return p_vt_out
 
+
 def test_rbf_intp_edge_stencil_01():
-    mesh= SimpleMesh()
+    mesh = SimpleMesh()
 
     p_vn_in = random_field(mesh, EdgeDim, KDim)
     ptr_coeff = random_field(mesh, EdgeDim, E2C2EDim)
-    p_vt_out = zero_field(mesh, EdgeDim, KDim) 
+    p_vt_out = zero_field(mesh, EdgeDim, KDim)
 
     ref = rbf_intp_edge_stencil_01_numpy(
         mesh.e2c2e,

advection/tests/test_step_advection_stencil_04.py

@@ -13,40 +13,39 @@
 
 import numpy as np
 
-from icon4py.advection.step_advection_stencil_04 import (
-    step_advection_stencil_04,
-)
-
+from icon4py.advection.step_advection_stencil_04 import step_advection_stencil_04
 from icon4py.common.dimension import CellDim, KDim
 from icon4py.testutils.simple_mesh import SimpleMesh
 from icon4py.testutils.utils import random_field, zero_field
 
+
 def step_advection_stencil_04_numpy(
-    p_tracer_now: np.array, 
-    p_tracer_new: np.array, 
+    p_tracer_now: np.array,
+    p_tracer_new: np.array,
     p_dtime,
 ) -> np.array:
-    opt_ddt_tracer_adv = (p_tracer_new-p_tracer_now)/p_dtime
+    opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime
     return opt_ddt_tracer_adv
 
+
 def test_step_advection_stencil_04():
     mesh = SimpleMesh()
 
     p_tracer_now = random_field(mesh, CellDim, KDim)
     p_tracer_new = random_field(mesh, CellDim, KDim)
-    opt_ddt_tracer_adv = zero_field(mesh, CellDim, KDim)    
-    p_dtime = np.float64(5.0)  
+    opt_ddt_tracer_adv = zero_field(mesh, CellDim, KDim)
+    p_dtime = np.float64(5.0)
 
     ref = step_advection_stencil_04_numpy(
         np.asarray(p_tracer_now),
-        np.asarray(p_tracer_new), 
-        p_dtime, 
+        np.asarray(p_tracer_new),
+        p_dtime,
     )
     step_advection_stencil_04(
         p_tracer_now,
         p_tracer_new,
         opt_ddt_tracer_adv,
         p_dtime,
         offset_provider={},
-)
+    )
     assert np.allclose(opt_ddt_tracer_adv, ref)