Bugfixes due to wrong dimensions in greenline (#319)

(Fix) Part 2 of adapting to new `gt4py` field representation. - remove remaining `np.asarray` - fix bugs due to run dimension in serialbox input in `Diffusion` and `SolveNonHydro` Co-authored-by: Nina Burgdorfer <[email protected]> Co-authored-by: ninaburg <[email protected]>
C2SM · Nov 28, 2023 · 867c1c2 · 867c1c2
1 parent 2d2460a
commit 867c1c2
Show file tree

Hide file tree

Showing 13 changed files with 227 additions and 195 deletions.
diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py
@@ -264,7 +264,6 @@ def test_verify_diffusion_init_against_other_regular_savepoint(
     _verify_init_values_against_savepoint(diffusion_savepoint_init, diffusion)
 
 
-@pytest.mark.skip("#TODO (magdalena) assertion failure after update of gt4py")
 @pytest.mark.datatest
 @pytest.mark.parametrize(
     "step_date_init, step_date_exit",
@@ -320,7 +319,6 @@ def test_run_diffusion_single_step(
     verify_diffusion_fields(diagnostic_state, prognostic_state, diffusion_savepoint_exit)
 
 
-@pytest.mark.skip("#TODO (magdalena) assertion failure after update of gt4py")
 @pytest.mark.datatest
 @pytest.mark.parametrize("linit", [True])
 def test_run_diffusion_initial_step(

diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py
@@ -49,7 +49,7 @@ def test_diff_multfac_vn_and_smag_limit_for_initial_step(backend):
     smag_limit_init = zero_field(grid, KDim)
     k4 = 1.0
     efdt_ratio = 24.0
-    shape = np.asarray(diff_multfac_vn_init).shape
+    shape = diff_multfac_vn_init.asnumpy().shape
 
     expected_diff_multfac_vn_init = initial_diff_multfac_vn_numpy(shape, k4, efdt_ratio)
     expected_smag_limit_init = smag_limit_numpy(
@@ -71,7 +71,7 @@ def test_diff_multfac_vn_smag_limit_for_time_step_with_const_value(backend):
     k4 = 1.0
     substeps = 5.0
     efdt_ratio = 24.0
-    shape = np.asarray(diff_multfac_vn).shape
+    shape = diff_multfac_vn.asnumpy().shape
 
     expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps)
     expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps)
@@ -92,7 +92,7 @@ def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(backend):
     k4 = 0.003
     substeps = 1.0
 
-    shape = np.asarray(diff_multfac_vn).shape
+    shape = diff_multfac_vn.asnumpy().shape
     expected_diff_multfac_vn = diff_multfac_vn_numpy(shape, k4, substeps)
     expected_smag_limit = smag_limit_numpy(diff_multfac_vn_numpy, shape, k4, substeps)
     _setup_runtime_diff_multfac_vn.with_backend(backend)(

diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py
diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py
@@ -13,6 +13,7 @@
 from typing import Tuple
 
 import numpy as np
+from gt4py.next import as_field
 from gt4py.next.common import Dimension, Field
 from gt4py.next.ffront.decorator import field_operator, program
 from gt4py.next.ffront.fbuiltins import (  # noqa: A004 # import gt4py builtin
@@ -22,22 +23,21 @@
     maximum,
     minimum,
 )
-from gt4py.next.iterator.embedded import np_as_located_field
 
 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim
 
 
+def indices_field(dim: Dimension, grid, is_halfdim, dtype=int):
+    shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim]
+    return as_field((dim,), np.arange(shapex, dtype=dtype))
+
+
 def zero_field(grid, *dims: Dimension, is_halfdim=False, dtype=float):
     shapex = tuple(map(lambda x: grid.size[x], dims))
     if is_halfdim:
         assert len(shapex) == 2
         shapex = (shapex[0], shapex[1] + 1)
-    return np_as_located_field(*dims)(np.zeros(shapex, dtype=dtype))
-
-
-def indices_field(dim: Dimension, grid, is_halfdim, dtype=int):
-    shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim]
-    return np_as_located_field(dim)(np.arange(shapex, dtype=dtype))
+    return as_field(dims, np.zeros(shapex, dtype=dtype))
 
 
 def _allocate(*dims: Dimension, grid, is_halfdim=False, dtype=float):

diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py
@@ -69,7 +69,7 @@
 cached_backend = run_gtfn_cached
 compiled_backend = run_gtfn
 imperative_backend = run_gtfn_imperative
-backend = run_gtfn_cached
+backend = run_gtfn
 #
 
 
@@ -164,7 +164,7 @@ def run_predictor_step(
         )
 
         if not vn_only:
-            mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.with_backend(run_gtfn)(
+            mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.with_backend(backend)(
                 p_cell_in=prognostic_state.w,
                 c_intp=self.interpolation_state.c_intp,
                 p_vert_out=self.z_w_v,
@@ -177,7 +177,7 @@ def run_predictor_step(
                 },
             )
 
-        mo_math_divrot_rot_vertex_ri_dsl.with_backend(run_gtfn)(
+        mo_math_divrot_rot_vertex_ri_dsl.with_backend(backend)(
             vec_e=prognostic_state.vn,
             geofac_rot=self.interpolation_state.geofac_rot,
             rot_vec=self.zeta,
@@ -190,7 +190,7 @@ def run_predictor_step(
             },
         )
 
-        mo_velocity_advection_stencil_01.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_01.with_backend(backend)(
             vn=prognostic_state.vn,
             rbf_vec_coeff_e=self.interpolation_state.rbf_vec_coeff_e,
             vt=diagnostic_state.vt,
@@ -203,7 +203,7 @@ def run_predictor_step(
             },
         )
 
-        mo_velocity_advection_stencil_02.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_02.with_backend(backend)(
             wgtfac_e=self.metric_state.wgtfac_e,
             vn=prognostic_state.vn,
             vt=diagnostic_state.vt,
@@ -219,7 +219,7 @@ def run_predictor_step(
         )
 
         if not vn_only:
-            mo_velocity_advection_stencil_03.with_backend(run_gtfn)(
+            mo_velocity_advection_stencil_03.with_backend(backend)(
                 wgtfac_e=self.metric_state.wgtfac_e,
                 vt=diagnostic_state.vt,
                 z_vt_ie=z_vt_ie,
@@ -230,7 +230,7 @@ def run_predictor_step(
                 offset_provider={"Koff": KDim},
             )
 
-        velocity_prog.fused_stencils_4_5_6.with_backend(run_gtfn)(
+        velocity_prog.fused_stencils_4_5.with_backend(backend)(
             vn=prognostic_state.vn,
             vt=diagnostic_state.vt,
             vn_ie=diagnostic_state.vn_ie,
@@ -239,21 +239,28 @@ def run_predictor_step(
             ddxn_z_full=self.metric_state.ddxn_z_full,
             ddxt_z_full=self.metric_state.ddxt_z_full,
             z_w_concorr_me=z_w_concorr_me,
-            wgtfacq_e_dsl=self.metric_state.wgtfacq_e_dsl,
             k_field=self.k_field,
             nflatlev_startindex=self.vertical_params.nflatlev,
             nlev=self.grid.num_levels,
             horizontal_start=start_edge_lb_plus4,
             horizontal_end=end_edge_local_minus2,
             vertical_start=0,
+            vertical_end=self.grid.num_levels,
+            offset_provider={},
+        )
+        velocity_prog.mo_velocity_advection_stencil_06.with_backend(backend)(
+            wgtfacq_e=self.metric_state.wgtfacq_e_dsl,
+            vn=prognostic_state.vn,
+            vn_ie=diagnostic_state.vn_ie,
+            horizontal_start=start_edge_lb_plus4,
+            horizontal_end=end_edge_local_minus2,
+            vertical_start=self.grid.num_levels,
             vertical_end=self.grid.num_levels + 1,
-            offset_provider={
-                "Koff": KDim,
-            },
+            offset_provider={"Koff": KDim},
         )
 
         if not vn_only:
-            mo_velocity_advection_stencil_07.with_backend(run_gtfn)(
+            mo_velocity_advection_stencil_07.with_backend(backend)(
                 vn_ie=diagnostic_state.vn_ie,
                 inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths,
                 w=prognostic_state.w,
@@ -272,7 +279,7 @@ def run_predictor_step(
                 },
             )
 
-        mo_velocity_advection_stencil_08.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_08.with_backend(backend)(
             z_kin_hor_e=z_kin_hor_e,
             e_bln_c_s=self.interpolation_state.e_bln_c_s,
             z_ekinh=self.z_ekinh,
@@ -286,7 +293,7 @@ def run_predictor_step(
             },
         )
 
-        velocity_prog.fused_stencils_9_10.with_backend(run_gtfn)(
+        velocity_prog.fused_stencils_9_10.with_backend(backend)(
             z_w_concorr_me=z_w_concorr_me,
             e_bln_c_s=self.interpolation_state.e_bln_c_s,
             local_z_w_concorr_mc=self.z_w_concorr_mc,
@@ -306,7 +313,7 @@ def run_predictor_step(
             },
         )
 
-        velocity_prog.fused_stencils_11_to_13.with_backend(run_gtfn)(
+        velocity_prog.fused_stencils_11_to_13.with_backend(backend)(
             w=prognostic_state.w,
             w_concorr_c=diagnostic_state.w_concorr_c,
             local_z_w_con_c=self.z_w_con_c,
@@ -320,7 +327,7 @@ def run_predictor_step(
             offset_provider={},
         )
 
-        velocity_prog.fused_stencil_14.with_backend(run_gtfn)(
+        velocity_prog.fused_stencil_14.with_backend(backend)(
             ddqz_z_half=self.metric_state.ddqz_z_half,
             local_z_w_con_c=self.z_w_con_c,
             local_cfl_clipping=self.cfl_clipping,
@@ -336,7 +343,7 @@ def run_predictor_step(
 
         self._update_levmask_from_cfl_clipping()
 
-        mo_velocity_advection_stencil_15.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_15.with_backend(backend)(
             z_w_con_c=self.z_w_con_c,
             z_w_con_c_full=self.z_w_con_c_full,
             horizontal_start=start_cell_lb_plus3,
@@ -346,7 +353,7 @@ def run_predictor_step(
             offset_provider={"Koff": KDim},
         )
 
-        velocity_prog.fused_stencils_16_to_17.with_backend(run_gtfn)(
+        velocity_prog.fused_stencils_16_to_17.with_backend(backend)(
             w=prognostic_state.w,
             local_z_v_grad_w=self.z_v_grad_w,
             e_bln_c_s=self.interpolation_state.e_bln_c_s,
@@ -365,7 +372,7 @@ def run_predictor_step(
             },
         )
 
-        mo_velocity_advection_stencil_18.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_18.with_backend(backend)(
             levmask=self.levmask,
             cfl_clipping=self.cfl_clipping,
             owner_mask=self.c_owner_mask,
@@ -390,7 +397,7 @@ def run_predictor_step(
         # This behaviour needs to change for multiple blocks
         self.levelmask = self.levmask
 
-        mo_velocity_advection_stencil_19.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_19.with_backend(backend)(
             z_kin_hor_e=z_kin_hor_e,
             coeff_gradekin=self.metric_state.coeff_gradekin,
             z_ekinh=self.z_ekinh,
@@ -414,7 +421,7 @@ def run_predictor_step(
             },
         )
 
-        mo_velocity_advection_stencil_20.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_20.with_backend(backend)(
             levelmask=self.levelmask,
             c_lin_e=self.interpolation_state.c_lin_e,
             z_w_con_c_full=self.z_w_con_c_full,
@@ -494,7 +501,7 @@ def run_corrector_step(
         )
 
         if not vn_only:
-            mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.with_backend(run_gtfn)(
+            mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.with_backend(backend)(
                 p_cell_in=prognostic_state.w,
                 c_intp=self.interpolation_state.c_intp,
                 p_vert_out=self.z_w_v,
@@ -507,7 +514,7 @@ def run_corrector_step(
                 },
             )
 
-        mo_math_divrot_rot_vertex_ri_dsl.with_backend(run_gtfn)(
+        mo_math_divrot_rot_vertex_ri_dsl.with_backend(backend)(
             vec_e=prognostic_state.vn,
             geofac_rot=self.interpolation_state.geofac_rot,
             rot_vec=self.zeta,
@@ -521,7 +528,7 @@ def run_corrector_step(
         )
 
         if not vn_only:
-            mo_velocity_advection_stencil_07.with_backend(run_gtfn)(
+            mo_velocity_advection_stencil_07.with_backend(backend)(
                 vn_ie=diagnostic_state.vn_ie,
                 inv_dual_edge_length=self.edge_params.inverse_dual_edge_lengths,
                 w=prognostic_state.w,
@@ -540,7 +547,7 @@ def run_corrector_step(
                 },
             )
 
-        mo_velocity_advection_stencil_08.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_08.with_backend(backend)(
             z_kin_hor_e=z_kin_hor_e,
             e_bln_c_s=self.interpolation_state.e_bln_c_s,
             z_ekinh=self.z_ekinh,
@@ -554,7 +561,7 @@ def run_corrector_step(
             },
         )
 
-        velocity_prog.fused_stencils_11_to_13.with_backend(run_gtfn)(
+        velocity_prog.fused_stencils_11_to_13.with_backend(backend)(
             w=prognostic_state.w,
             w_concorr_c=diagnostic_state.w_concorr_c,
             local_z_w_con_c=self.z_w_con_c,
@@ -568,7 +575,7 @@ def run_corrector_step(
             offset_provider={},
         )
 
-        velocity_prog.fused_stencil_14.with_backend(run_gtfn)(
+        velocity_prog.fused_stencil_14.with_backend(backend)(
             ddqz_z_half=self.metric_state.ddqz_z_half,
             local_z_w_con_c=self.z_w_con_c,
             local_cfl_clipping=self.cfl_clipping,
@@ -584,7 +591,7 @@ def run_corrector_step(
 
         self._update_levmask_from_cfl_clipping()
 
-        mo_velocity_advection_stencil_15.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_15.with_backend(backend)(
             z_w_con_c=self.z_w_con_c,
             z_w_con_c_full=self.z_w_con_c_full,
             horizontal_start=start_cell_lb_plus3,
@@ -594,7 +601,7 @@ def run_corrector_step(
             offset_provider={"Koff": KDim},
         )
 
-        velocity_prog.fused_stencils_16_to_17.with_backend(run_gtfn)(
+        velocity_prog.fused_stencils_16_to_17.with_backend(backend)(
             w=prognostic_state.w,
             local_z_v_grad_w=self.z_v_grad_w,
             e_bln_c_s=self.interpolation_state.e_bln_c_s,
@@ -613,7 +620,7 @@ def run_corrector_step(
             },
         )
 
-        mo_velocity_advection_stencil_18.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_18.with_backend(backend)(
             levmask=self.levmask,
             cfl_clipping=self.cfl_clipping,
             owner_mask=self.c_owner_mask,
@@ -638,7 +645,7 @@ def run_corrector_step(
         # This behaviour needs to change for multiple blocks
         self.levelmask = self.levmask
 
-        mo_velocity_advection_stencil_19.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_19.with_backend(backend)(
             z_kin_hor_e=z_kin_hor_e,
             coeff_gradekin=self.metric_state.coeff_gradekin,
             z_ekinh=self.z_ekinh,
@@ -662,7 +669,7 @@ def run_corrector_step(
             },
         )
 
-        mo_velocity_advection_stencil_20.with_backend(run_gtfn)(
+        mo_velocity_advection_stencil_20.with_backend(backend)(
             levelmask=self.levelmask,
             c_lin_e=self.interpolation_state.c_lin_e,
             z_w_con_c_full=self.z_w_con_c_full,