diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py index 142dcfd83..7d678b48e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/fused_velocity_advection_stencil_19_to_20.py @@ -85,8 +85,8 @@ def _fused_velocity_advection_stencil_19_to_20( ), ddt_vn_apc, ) - if extra_diffu - else ddt_vn_apc + # if extra_diffu + # else ddt_vn_apc ) return ddt_vn_apc diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py b/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py index 53b8d2c05..7dff25206 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py @@ -8,6 +8,8 @@ import pytest from icon4py.model.atmosphere.dycore import dycore_states, velocity_advection as advection +from icon4py.model.atmosphere.dycore.stencils import fused_velocity_advection_stencil_19_to_20, \ + fused_velocity_advection_stencil_15_to_18, fused_velocity_advection_stencil_8_to_13 from icon4py.model.common import dimension as dims, utils as common_utils from icon4py.model.common.grid import ( horizontal as h_grid, @@ -16,8 +18,15 @@ ) from icon4py.model.common.states import prognostic_state as prognostics from icon4py.model.testing import datatest_utils as dt_utils, helpers - +from icon4py.model.common.utils.data_allocation import ( + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) from . import utils +import gt4py.next as gtx +import numpy as np def create_vertical_params(vertical_config, grid_savepoint): @@ -450,3 +459,242 @@ def test_velocity_corrector_step( icon_result_ddt_vn_apc_pc, atol=5.0e-16, ) + +@pytest.mark.datatest +@pytest.mark.parametrize( + "experiment, step_date_init, step_date_exit", + [ + (dt_utils.REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + #(dt_utils.GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), + ], +) +def test_velocity_fused_19_20( + step_date_init, + step_date_exit, + icon_grid, + grid_savepoint, + savepoint_velocity_init, + savepoint_velocity_exit, + interpolation_savepoint, + metrics_savepoint, + backend, +): + vn = savepoint_velocity_init.vn_something() # TODO + geofac_rot = interpolation_savepoint.geofac_rot() + z_kin_hor_e = savepoint_velocity_init.init_z_kin_hor_e_19_20() # TODO + coeff_gradekin = metrics_savepoint.coeff_gradekin() + z_ekinh = savepoint_velocity_init.init_z_ekinh_19_20() # TODO + vt = savepoint_velocity_init.init_vt_19_20() # TODO + f_e = grid_savepoint.f_e() + c_lin_e = interpolation_savepoint.c_lin_e() + z_w_con_c_full = savepoint_velocity_init.init_z_w_con_c_full_19_20() # TODO + vn_ie = savepoint_velocity_init.init_vn_ie_19_20() # TODO + ddqz_z_full_e = metrics_savepoint.ddqz_z_full_e() + levelmask = savepoint_velocity_init.init_levelmask_19_20() # TODO + area_edge = grid_savepoint.edge_areas() + tangent_orientation = grid_savepoint.tangent_orientation() + inv_primal_edge_length = grid_savepoint.inv_primal_edge_length() + geofac_grdiv = interpolation_savepoint.geofac_grdiv() + ddt_vn_apc = savepoint_velocity_init.init_ddt_vn_apc_19_20() # TODO + k = gtx.as_field((dims.KDim,), np.arange(icon_grid.num_levels, dtype=gtx.int32)) + + d_time = 2.0 + extra_diffu = True + nrdmax = grid_savepoint.nrdmax() + + ddt_vn_apc_ref = savepoint_velocity_exit.x_ddt_vn_apc_19_20() # TODO + edge_domain = h_grid.domain(dims.EdgeDim) + horizontal_start = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) + + scalfac_exdiff = savepoint_velocity_init.scalfac_exdiff() + cfl_w_limit = savepoint_velocity_init.cfl_w_limit() + + fused_velocity_advection_stencil_19_to_20.fused_velocity_advection_stencil_19_to_20.with_backend(backend)( + vn=vn, + geofac_rot=geofac_rot, + z_kin_hor_e=z_kin_hor_e, + coeff_gradekin=coeff_gradekin, + z_ekinh=z_ekinh, + vt=vt, + f_e=f_e, + c_lin_e=c_lin_e, + z_w_con_c_full=z_w_con_c_full, + vn_ie=vn_ie, + ddqz_z_full_e=ddqz_z_full_e, + levelmask=levelmask, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + geofac_grdiv=geofac_grdiv, + ddt_vn_apc=ddt_vn_apc, + k=k, + cfl_w_limit=cfl_w_limit, + scalfac_exdiff=scalfac_exdiff, + d_time=d_time, + extra_diffu=extra_diffu, + nlev=icon_grid.num_levels, + nrdmax=nrdmax, + horizontal_start=horizontal_start, + horizontal_end=icon_grid.num_edges, + vertical_start=0, + vertical_end=icon_grid.num_levels, + offset_provider={"V2E": icon_grid.get_offset_provider("V2E"), + "E2EC": icon_grid.get_offset_provider("E2EC"), + "E2V": icon_grid.get_offset_provider("E2V"), + "E2C": icon_grid.get_offset_provider("E2C"), + "E2C2EO": icon_grid.get_offset_provider("E2C2EO"), + "Koff": dims.KDim} + ) + + assert helpers.dallclose(ddt_vn_apc_ref.asnumpy(), ddt_vn_apc.asnumpy()) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "experiment, step_date_init, step_date_exit", + [ + (dt_utils.REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + #(dt_utils.GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), + ], +) +def test_velocity_fused_15_18( + step_date_init, + step_date_exit, + icon_grid, + grid_savepoint, + savepoint_velocity_init, + savepoint_velocity_exit, + interpolation_savepoint, + metrics_savepoint, + savepoint_nonhydro_exit, + backend, +): + z_w_con_c = savepoint_velocity_init.init_z_w_con_c_15_18() # TODO + w = savepoint_velocity_init.init_w_15_18() # TODO + coeff1_dwdz = metrics_savepoint.coeff1_dwdz() + coeff2_dwdz = metrics_savepoint.coeff2_dwdz() + ddt_w_adv = savepoint_velocity_init.init_ddt_w_adv_15_18() # TODO + e_bln_c_s = interpolation_savepoint.e_bln_c_s() + z_v_grad_w = savepoint_velocity_init.init_z_v_grad_w_15_18() # TODO + levelmask = savepoint_velocity_init.init_levelmask_15_18() # TODO + cfl_clipping = savepoint_nonhydro_exit.cfl_clipping() + owner_mask = grid_savepoint.c_owner_mask() + ddqz_z_half = metrics_savepoint.ddqz_z_half() + area = grid_savepoint.cell_areas() + geofac_n2s = interpolation_savepoint.geofac_n2s() + z_w_con_c_full = savepoint_velocity_init.init_z_w_con_c_full_15_18() # TODO + z_w_con_c_full_ref = savepoint_velocity_exit.x_z_w_con_c_full_15_18() # TODO + ddt_w_adv_ref = savepoint_velocity_exit.x_ddt_w_adv_15_18() # TODO + + k = gtx.as_field((dims.KDim,), np.arange(icon_grid.num_levels, dtype=gtx.int32)) + cell = gtx.as_field((dims.CellDim,), np.arange(icon_grid.num_cells, dtype=gtx.int32)) + + nrdmax = grid_savepoint.nrdmax() + extra_diffu = True + + cell_lower_bound = 2 # TODO + cell_upper_bound = 4 # TODO + + lvn_only = False + + scalfac_exdiff = savepoint_velocity_init.scalfac_exdiff() + cfl_w_limit = savepoint_velocity_init.cfl_w_limit() + dtime = 2.0 + + fused_velocity_advection_stencil_15_to_18.fused_velocity_advection_stencil_15_to_18.with_backend(backend)( + z_w_con_c=z_w_con_c, + w=w, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + ddt_w_adv=ddt_w_adv, + e_bln_c_s=e_bln_c_s, + z_v_grad_w=z_v_grad_w, + levelmask=levelmask, + cfl_clipping=cfl_clipping, + owner_mask=owner_mask, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + z_w_con_c_full=z_w_con_c_full, + cell=cell, + k=k, + scalfac_exdiff=scalfac_exdiff, + cfl_w_limit=cfl_w_limit, + dtime=dtime, + cell_lower_bound=cell_lower_bound, + cell_upper_bound=cell_upper_bound, + nlev=icon_grid.num_levels, + nrdmax=nrdmax, + lvn_only=lvn_only, + extra_diffu=extra_diffu, + offset_provider={"C2E": icon_grid.get_offset_provider("C2E"), + "C2CE": icon_grid.get_offset_provider("C2CE"), + "C2E2CO": icon_grid.get_offset_provider("C2E2CO"), + "Koff": dims.KDim} + ) + + assert helpers.dallclose(z_w_con_c_full_ref.asnumpy(), z_w_con_c_full.asnumpy()) + assert helpers.dallclose(ddt_w_adv_ref.asnumpy(), ddt_w_adv.asnumpy()) + +@pytest.mark.datatest +@pytest.mark.parametrize( + "experiment, step_date_init, step_date_exit", + [ + (dt_utils.REGIONAL_EXPERIMENT, "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"), + #(dt_utils.GLOBAL_EXPERIMENT, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), + ], +) +def test_velocity_fused_8_13( + step_date_init, + step_date_exit, + icon_grid, + grid_savepoint, + savepoint_velocity_init, + savepoint_velocity_exit, + metrics_savepoint, + istep, + backend, +): + z_ekinh_ref = savepoint_velocity_exit.x_z_ekinh_8_13() # TODO + w_concorr_c_ref = savepoint_velocity_exit.x_w_concorr_c_8_13() # TODO + z_w_con_c_ref = savepoint_velocity_exit.x_z_w_con_c_8_13() # TODO + + w = savepoint_velocity_init.init_w_8_13() # TODO + e_bln_c_s = savepoint_velocity_init.init_e_bln_c_s_8_13() + z_kin_hor_e = savepoint_velocity_init.init_z_kin_hor_e_8_13() # TODO + z_w_concorr_me = savepoint_velocity_init.init_z_w_concorr_me_8_13() # TODO + wgtfac_c = metrics_savepoint.wgtfac_c() + z_w_concorr_mc = savepoint_velocity_init.init_z_w_concorr_mc_8_13() # TODO + w_concorr_c = savepoint_velocity_init.init_w_concorr_c_8_13() # TODO + z_ekinh = savepoint_velocity_init.init_z_ekinh_8_13() # TODO + + k = gtx.as_field((dims.KDim,), np.arange(icon_grid.num_levels, dtype=gtx.int32)) + z_w_con_c = savepoint_velocity_exit.init_z_w_con_c_8_13() + + nflatlev = grid_savepoint.nflatlev() + + fused_velocity_advection_stencil_8_to_13.fused_velocity_advection_stencil_8_to_13.with_backend(backend)( + z_kin_hor_e=z_kin_hor_e, + e_bln_c_s=e_bln_c_s, + z_w_concorr_me=z_w_concorr_me, + wgtfac_c=wgtfac_c, + w=w, + z_w_concorr_mc=z_w_concorr_mc, + w_concorr_c=w_concorr_c, + z_ekinh=z_ekinh, + z_w_con_c=z_w_con_c, + k=k, + istep=istep, + nlev=icon_grid.num_levels, + nflatlev=nflatlev, + horizontal_start=0, + horizontal_end=icon_grid.num_cells, + vertical_start=0, + vertical_end=icon_grid.num_levels+1, + offset_provider={"C2E": icon_grid.get_offset_provider("C2E"), + "Koff": dims.KDim} + ) + assert helpers.dallclose(z_ekinh_ref.asnumpy(), z_ekinh.asnumpy()) + assert helpers.dallclose(w_concorr_c_ref.asnumpy(), w_concorr_c.asnumpy()) + assert helpers.dallclose(z_w_con_c_ref.asnumpy(), z_w_con_c.asnumpy()) + diff --git a/model/testing/src/icon4py/model/testing/serialbox.py b/model/testing/src/icon4py/model/testing/serialbox.py index 1c33e4599..22bbe53be 100644 --- a/model/testing/src/icon4py/model/testing/serialbox.py +++ b/model/testing/src/icon4py/model/testing/serialbox.py @@ -1352,6 +1352,15 @@ def z_dwdz_dd(self): def exner_dyn_incr(self): return self._get_field("x_exner_dyn_incr", dims.CellDim, dims.KDim) + def x_ddt_vn_apc_19_20(self): + return self._get_field("x_ddt_vn_apc_19_20", dims.EdgeDim, dims.KDim) + + def x_z_w_con_c_full_15_18(self): + return self._get_field("x_z_w_con_c_full_15_18", dims.CellDim, dims.KDim) + + def x_ddt_w_adv_15_18(self): + return self._get_field("x_ddt_w_adv_15_18", dims.CellDim, dims.KDim) + # TODO (magdalena) rename? class IconNHFinalExitSavepoint(IconSavepoint):