From f929579d1c43369415dbabd276f0ed6369ce59a9 Mon Sep 17 00:00:00 2001 From: muellch <60387010+muellch@users.noreply.github.com> Date: Wed, 30 Aug 2023 15:15:11 +0200 Subject: [PATCH] Remove pre_levelmask from vel. advection and fix bugs (#258) --- .../mo_velocity_advection_stencil_14.py | 16 ++++---------- .../mo_velocity_advection_stencil_18.py | 10 ++++----- .../mo_velocity_advection_stencil_20.py | 11 ++++++---- .../test_mo_velocity_advection_stencil_14.py | 12 +---------- .../test_mo_velocity_advection_stencil_18.py | 12 +++++------ .../test_mo_velocity_advection_stencil_20.py | 21 ++++++++++++------- 6 files changed, 36 insertions(+), 46 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py index 3bce7d7ef2..b61a42a54a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_14.py +++ b/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,15 +35,15 @@ 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) @@ -55,7 +51,6 @@ 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, @@ -63,10 +58,7 @@ def mo_velocity_advection_stencil_14( _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), ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py index c4e8d32d2b..7dbbe1797a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_18.py +++ b/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, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py index 173bd2e9b1..5817ee1718 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_20.py +++ b/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), ) diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py index f691e32ad2..335cc2fe31 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_14.py +++ b/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, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py index 56d33d5454..8408b8baeb 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_18.py +++ b/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, diff --git a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py index 5a37883c2c..50253026c3 100644 --- a/model/atmosphere/dycore/tests/test_mo_velocity_advection_stencil_20.py +++ b/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)