first edits for fused_stencils datatest implementation

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

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())
+

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):