Remove pre_levelmask from vel. advection and fix bugs (#258)

Author: muellch

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py

@@ -22,13 +22,9 @@
 def _mo_velocity_advection_stencil_14(
     ddqz_z_half: Field[[CellDim, KDim], float],
     z_w_con_c: Field[[CellDim, KDim], float],
-    cfl_clipping: Field[[CellDim, KDim], bool],
-    pre_levelmask: Field[[CellDim, KDim], bool],
-    vcfl: Field[[CellDim, KDim], float],
     cfl_w_limit: float,
     dtime: float,
 ) -> tuple[
-    Field[[CellDim, KDim], bool],
     Field[[CellDim, KDim], bool],
     Field[[CellDim, KDim], float],
     Field[[CellDim, KDim], float],
@@ -39,34 +35,30 @@ def _mo_velocity_advection_stencil_14(
         False,
     )
 
-    pre_levelmask = where(cfl_clipping, broadcast(True, (CellDim, KDim)), False)
-
+    # TOOO: As soon as reductions of arbitrar dimensions are possible in gt4py, this stencil
+    # should reduce the vertical cfl to a scalar and the levmask to a per level boolean field (see Fortran dycore).
     vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0)
 
     z_w_con_c = where((cfl_clipping) & (vcfl < -0.85), -0.85 * ddqz_z_half / dtime, z_w_con_c)
 
     z_w_con_c = where((cfl_clipping) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c)
 
-    return cfl_clipping, pre_levelmask, vcfl, z_w_con_c
+    return cfl_clipping, vcfl, z_w_con_c
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
 def mo_velocity_advection_stencil_14(
     ddqz_z_half: Field[[CellDim, KDim], float],
     z_w_con_c: Field[[CellDim, KDim], float],
     cfl_clipping: Field[[CellDim, KDim], bool],
-    pre_levelmask: Field[[CellDim, KDim], bool],
     vcfl: Field[[CellDim, KDim], float],
     cfl_w_limit: float,
     dtime: float,
 ):
     _mo_velocity_advection_stencil_14(
         ddqz_z_half,
         z_w_con_c,
-        cfl_clipping,
-        pre_levelmask,
-        vcfl,
         cfl_w_limit,
         dtime,
-        out=(cfl_clipping, pre_levelmask, vcfl, z_w_con_c),
+        out=(cfl_clipping, vcfl, z_w_con_c),
     )

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py

@@ -20,7 +20,7 @@
 
 @field_operator
 def _mo_velocity_advection_stencil_18(
-    levelmask: Field[[KDim], bool],
+    levmask: Field[[KDim], bool],
     cfl_clipping: Field[[CellDim, KDim], bool],
     owner_mask: Field[[CellDim], bool],
     z_w_con_c: Field[[CellDim, KDim], float],
@@ -34,7 +34,7 @@ def _mo_velocity_advection_stencil_18(
     dtime: float,
 ) -> Field[[CellDim, KDim], float]:
     difcoef = where(
-        levelmask & cfl_clipping & owner_mask,
+        levmask & cfl_clipping & owner_mask,
         scalfac_exdiff
         * minimum(
             0.85 - cfl_w_limit * dtime,
@@ -44,7 +44,7 @@ def _mo_velocity_advection_stencil_18(
     )
 
     ddt_w_adv = where(
-        levelmask & cfl_clipping & owner_mask,
+        levmask & cfl_clipping & owner_mask,
         ddt_w_adv + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=C2E2CODim),
         ddt_w_adv,
     )
@@ -54,7 +54,7 @@ def _mo_velocity_advection_stencil_18(
 
 @program(grid_type=GridType.UNSTRUCTURED)
 def mo_velocity_advection_stencil_18(
-    levelmask: Field[[KDim], bool],
+    levmask: Field[[KDim], bool],
     cfl_clipping: Field[[CellDim, KDim], bool],
     owner_mask: Field[[CellDim], bool],
     z_w_con_c: Field[[CellDim, KDim], float],
@@ -68,7 +68,7 @@ def mo_velocity_advection_stencil_18(
     dtime: float,
 ):
     _mo_velocity_advection_stencil_18(
-        levelmask,
+        levmask,
         cfl_clipping,
         owner_mask,
         z_w_con_c,

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py

@@ -30,7 +30,6 @@
     CellDim,
     E2C2EODim,
     E2CDim,
-    E2VDim,
     EdgeDim,
     KDim,
     Koff,
@@ -75,8 +74,12 @@ def _mo_velocity_advection_stencil_20(
     ddt_vn_adv = where(
         (levelmask | levelmask(Koff[1])) & (abs(w_con_e) > cfl_w_limit * ddqz_z_full_e),
         ddt_vn_adv
-        + difcoef * area_edge * neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=E2C2EODim)
-        + tangent_orientation * inv_primal_edge_length * neighbor_sum(zeta(E2V), axis=E2VDim),
+        + difcoef
+        * area_edge
+        * (
+            neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=E2C2EODim)
+            + tangent_orientation * inv_primal_edge_length * (zeta(E2V[1]) - zeta(E2V[0]))
+        ),
         ddt_vn_adv,
     )
     return ddt_vn_adv
@@ -114,5 +117,5 @@ def mo_velocity_advection_stencil_20(
         cfl_w_limit,
         scalfac_exdiff,
         d_time,
-        out=ddt_vn_adv[:, :-1],
+        out=(ddt_vn_adv),
     )

model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py

@@ -28,7 +28,7 @@
 
 class TestMoVelocityAdvectionStencil14(StencilTest):
     PROGRAM = mo_velocity_advection_stencil_14
-    OUTPUTS = ("cfl_clipping", "pre_levelmask", "vcfl", "z_w_con_c")
+    OUTPUTS = ("cfl_clipping", "vcfl", "z_w_con_c")
 
     @staticmethod
     def reference(
@@ -41,11 +41,6 @@ def reference(
             np.zeros_like(z_w_con_c),
         )
         num_rows, num_cols = cfl_clipping.shape
-        pre_levelmask = np.where(
-            cfl_clipping == 1.0,
-            np.ones([num_rows, num_cols]),
-            np.zeros_like(cfl_clipping),
-        )
         vcfl = np.where(cfl_clipping == 1.0, z_w_con_c * dtime / ddqz_z_half, 0.0)
         z_w_con_c = np.where(
             (cfl_clipping == 1.0) & (vcfl < -0.85),
@@ -58,7 +53,6 @@ def reference(
 
         return dict(
             cfl_clipping=cfl_clipping,
-            pre_levelmask=pre_levelmask,
             vcfl=vcfl,
             z_w_con_c=z_w_con_c,
         )
@@ -68,9 +62,6 @@ def input_data(self, mesh):
         ddqz_z_half = random_field(mesh, CellDim, KDim)
         z_w_con_c = random_field(mesh, CellDim, KDim)
         cfl_clipping = random_mask(mesh, CellDim, KDim, dtype=bool)
-        pre_levelmask = random_mask(
-            mesh, CellDim, KDim, dtype=bool
-        )  # TODO should be just a K field
         vcfl = zero_field(mesh, CellDim, KDim)
         cfl_w_limit = 5.0
         dtime = 9.0
@@ -79,7 +70,6 @@ def input_data(self, mesh):
             ddqz_z_half=ddqz_z_half,
             z_w_con_c=z_w_con_c,
             cfl_clipping=cfl_clipping,
-            pre_levelmask=pre_levelmask,
             vcfl=vcfl,
             cfl_w_limit=cfl_w_limit,
             dtime=dtime,

model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py

@@ -32,7 +32,7 @@ class TestMoVelocityAdvectionStencil18(StencilTest):
     @staticmethod
     def reference(
         mesh,
-        levelmask: np.array,
+        levmask: np.array,
         cfl_clipping: np.array,
         owner_mask: np.array,
         z_w_con_c: np.array,
@@ -46,13 +46,13 @@ def reference(
         dtime: float,
         **kwargs,
     ):
-        levelmask = np.expand_dims(levelmask, axis=0)
+        levmask = np.expand_dims(levmask, axis=0)
         owner_mask = np.expand_dims(owner_mask, axis=-1)
         area = np.expand_dims(area, axis=-1)
         geofac_n2s = np.expand_dims(geofac_n2s, axis=-1)
 
         difcoef = np.where(
-            (levelmask == 1) & (cfl_clipping == 1) & (owner_mask == 1),
+            (levmask == 1) & (cfl_clipping == 1) & (owner_mask == 1),
             scalfac_exdiff
             * np.minimum(
                 0.85 - cfl_w_limit * dtime,
@@ -62,7 +62,7 @@ def reference(
         )
 
         ddt_w_adv = np.where(
-            (levelmask == 1) & (cfl_clipping == 1) & (owner_mask == 1),
+            (levmask == 1) & (cfl_clipping == 1) & (owner_mask == 1),
             ddt_w_adv + difcoef * area * np.sum(w[mesh.c2e2cO] * geofac_n2s, axis=1),
             ddt_w_adv,
         )
@@ -71,7 +71,7 @@ def reference(
 
     @pytest.fixture
     def input_data(self, mesh):
-        levelmask = random_mask(mesh, KDim)
+        levmask = random_mask(mesh, KDim)
         cfl_clipping = random_mask(mesh, CellDim, KDim)
         owner_mask = random_mask(mesh, CellDim)
         z_w_con_c = random_field(mesh, CellDim, KDim)
@@ -85,7 +85,7 @@ def input_data(self, mesh):
         dtime = 2.0
 
         return dict(
-            levelmask=levelmask,
+            levmask=levmask,
             cfl_clipping=cfl_clipping,
             owner_mask=owner_mask,
             z_w_con_c=z_w_con_c,

model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py

@@ -50,7 +50,8 @@ def mo_velocity_advection_stencil_20_numpy(
     w_con_e = np.zeros_like(vn)
     difcoef = np.zeros_like(vn)
 
-    levelmask_offset_1 = np.roll(levelmask, shift=-1, axis=0)
+    levelmask_offset_0 = levelmask[:-1]
+    levelmask_offset_1 = levelmask[1:]
 
     c_lin_e = np.expand_dims(c_lin_e, axis=-1)
     geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1)
@@ -59,12 +60,12 @@ def mo_velocity_advection_stencil_20_numpy(
     inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1)
 
     w_con_e = np.where(
-        (levelmask == 1) | (levelmask_offset_1 == 1),
+        (levelmask_offset_0) | (levelmask_offset_1),
         np.sum(c_lin_e * z_w_con_c_full[e2c], axis=1),
         w_con_e,
     )
     difcoef = np.where(
-        ((levelmask == 1) | (levelmask_offset_1 == 1))
+        ((levelmask_offset_0) | (levelmask_offset_1))
         & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e),
         scalfac_exdiff
         * np.minimum(
@@ -74,11 +75,15 @@ def mo_velocity_advection_stencil_20_numpy(
         difcoef,
     )
     ddt_vn_adv = np.where(
-        ((levelmask == 1) | (levelmask_offset_1 == 1))
+        ((levelmask_offset_0) | (levelmask_offset_1))
         & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e),
         ddt_vn_adv
-        + difcoef * area_edge * np.sum(geofac_grdiv * vn[e2c2eO], axis=1)
-        + tangent_orientation * inv_primal_edge_length * np.sum(zeta[e2v], axis=1),
+        + difcoef
+        * area_edge
+        * (
+            np.sum(geofac_grdiv * vn[e2c2eO], axis=1)
+            + tangent_orientation * inv_primal_edge_length * (zeta[e2v][:, 1] - zeta[e2v][:, 0])
+        ),
         ddt_vn_adv,
     )
     return ddt_vn_adv
@@ -87,7 +92,7 @@ def mo_velocity_advection_stencil_20_numpy(
 def test_mo_velocity_advection_stencil_20():
     mesh = SimpleMesh()
 
-    levelmask = random_mask(mesh, KDim)
+    levelmask = random_mask(mesh, KDim, extend={KDim: 1})
     c_lin_e = random_field(mesh, EdgeDim, E2CDim)
     z_w_con_c_full = random_field(mesh, CellDim, KDim)
     ddqz_z_full_e = random_field(mesh, EdgeDim, KDim)
@@ -145,4 +150,4 @@ def test_mo_velocity_advection_stencil_20():
         },
     )
 
-    assert np.allclose(ddt_vn_adv[:, :-1], ddt_vn_adv_ref[:, :-1])
+    assert np.allclose(ddt_vn_adv, ddt_vn_adv_ref)