patching with @muellch vel adv bug fixes (#280)

Author: abishekg7

atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_14.py

@@ -21,13 +21,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],
@@ -38,8 +34,6 @@ def _mo_velocity_advection_stencil_14(
         False,
     )
 
-    pre_levelmask = where(cfl_clipping, broadcast(True, (CellDim, KDim)), False)
-
     vcfl = where(cfl_clipping, z_w_con_c * dtime / ddqz_z_half, 0.0)
 
     z_w_con_c = where(
@@ -50,26 +44,22 @@ def _mo_velocity_advection_stencil_14(
         (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
 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),
     )

atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_18.py

@@ -19,7 +19,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],
@@ -33,7 +33,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,
@@ -43,7 +43,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
 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,

atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_velocity_advection_stencil_20.py

@@ -29,7 +29,6 @@
     CellDim,
     E2C2EODim,
     E2CDim,
-    E2VDim,
     EdgeDim,
     KDim,
     Koff,
@@ -74,10 +73,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

atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_14.py

@@ -33,10 +33,7 @@ def mo_velocity_advection_stencil_14_numpy(
         np.ones([num_rows, num_cols]),
         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), -0.85 * ddqz_z_half / dtime, z_w_con_c
@@ -45,7 +42,7 @@ def mo_velocity_advection_stencil_14_numpy(
         (cfl_clipping == 1.0) & (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
 
 
 def test_mo_velocity_advection_stencil_14():
@@ -54,16 +51,12 @@ def test_mo_velocity_advection_stencil_14():
     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
 
     (
         cfl_clipping_ref,
-        pre_levelmask_ref,
         vcfl_ref,
         z_w_con_c_ref,
     ) = mo_velocity_advection_stencil_14_numpy(
@@ -77,13 +70,11 @@ def test_mo_velocity_advection_stencil_14():
         ddqz_z_half,
         z_w_con_c,
         cfl_clipping,
-        pre_levelmask,
         vcfl,
         cfl_w_limit,
         dtime,
         offset_provider={},
     )
     assert np.allclose(cfl_clipping, cfl_clipping_ref)
-    assert np.allclose(pre_levelmask, pre_levelmask_ref)
     assert np.allclose(vcfl, vcfl_ref)
     assert np.allclose(z_w_con_c, z_w_con_c_ref)

atm_dyn_iconam/tests/test_mo_velocity_advection_stencil_18.py

@@ -23,7 +23,7 @@
 
 def mo_velocity_advection_stencil_18_numpy(
     c2e2c0: np.array,
-    levelmask: np.array,
+    levmask: np.array,
     cfl_clipping: np.array,
     owner_mask: np.array,
     z_w_con_c: np.array,
@@ -36,13 +36,13 @@ def mo_velocity_advection_stencil_18_numpy(
     cfl_w_limit: float,
     dtime: float,
 ):
-    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,
@@ -52,7 +52,7 @@ def mo_velocity_advection_stencil_18_numpy(
     )
 
     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[c2e2c0] * geofac_n2s, axis=1),
         ddt_w_adv,
     )
@@ -63,7 +63,7 @@ def mo_velocity_advection_stencil_18_numpy(
 def test_mo_velocity_advection_stencil_18():
     mesh = SimpleMesh()
 
-    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)
@@ -78,7 +78,7 @@ def test_mo_velocity_advection_stencil_18():
 
     ref = mo_velocity_advection_stencil_18_numpy(
         mesh.c2e2cO,
-        np.asarray(levelmask),
+        np.asarray(levmask),
         np.asarray(cfl_clipping),
         np.asarray(owner_mask),
         np.asarray(z_w_con_c),
@@ -93,7 +93,7 @@ def test_mo_velocity_advection_stencil_18():
     )
 
     mo_velocity_advection_stencil_18(
-        levelmask,
+        levmask,
         cfl_clipping,
         owner_mask,
         z_w_con_c,

atm_dyn_iconam/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,11 @@ 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 +88,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)
@@ -144,5 +145,5 @@ def test_mo_velocity_advection_stencil_20():
             "E2V": mesh.get_e2v_offset_provider(),
         },
     )
-
+    
     assert np.allclose(ddt_vn_adv[:, :-1], ddt_vn_adv_ref[:, :-1])