[py2f]: Add dycore wrapper and datatest

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py

@@ -301,6 +301,7 @@ def __init__(
         divdamp_z2: float = 40000.0,
         divdamp_z3: float = 60000.0,
         divdamp_z4: float = 80000.0,
+        htop_moist_proc: float = 22500.0,
     ):
         # parameters from namelist diffusion_nml
         self.itime_scheme: int = itime_scheme
@@ -361,6 +362,8 @@ def __init__(
         #: IAU weight for dynamics fields
         self.iau_wgt_dyn: float = iau_wgt_dyn
 
+        self.htop_moist_proc: float = htop_moist_proc
+
         self._validate()
 
     def _validate(self):
@@ -1055,9 +1058,7 @@ def run_predictor_step(
             )
         # TODO (Nikki) check when merging fused stencil
         lowest_level = self.grid.num_levels - 1
-        hydro_corr_horizontal = gtx.as_field(
-            (EdgeDim,), self.z_hydro_corr.asnumpy()[:, lowest_level]
-        )
+        hydro_corr_horizontal = gtx.as_field((EdgeDim,), self.z_hydro_corr.ndarray[:, lowest_level])
 
         if self.config.igradp_method == HorizontalPressureDiscretizationType.TAYLOR_HYDRO:
             apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure(

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py

@@ -12,7 +12,6 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import gt4py.next as gtx
-import numpy as np
 
 import icon4py.model.atmosphere.dycore.velocity.velocity_advection_program as velocity_prog
 from icon4py.model.atmosphere.dycore.add_extra_diffusion_for_normal_wind_tendency_approaching_cfl import (
@@ -50,6 +49,7 @@
 from icon4py.model.common import field_type_aliases as fa
 from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid
+from icon4py.model.common.settings import xp
 from icon4py.model.common.states import prognostic_state as prognostics
 from icon4py.model.common.utils import gt4py_field_allocation as field_alloc
 
@@ -414,7 +414,7 @@ def run_predictor_step(
 
     def _update_levmask_from_cfl_clipping(self):
         self.levmask = gtx.as_field(
-            domain=(KDim,), data=(np.any(self.cfl_clipping.asnumpy(), 0)), dtype=bool
+            domain=(KDim,), data=(xp.any(self.cfl_clipping.ndarray, 0)), dtype=bool
         )
 
     def _scale_factors_by_dtime(self, dtime):

model/common/src/icon4py/model/common/test_utils/helpers.py

@@ -21,6 +21,8 @@
 from gt4py.next import as_field, common as gt_common, constructors
 from gt4py.next.ffront.decorator import Program
 
+from icon4py.model.common.settings import xp
+
 from ..grid.base import BaseGrid
 from ..type_alias import wpfloat
 
@@ -104,7 +106,7 @@ def zero_field(
     dtype=wpfloat,
     extend: Optional[dict[gt_common.Dimension, int]] = None,
 ) -> gt_common.Field:
-    return as_field(dims, np.zeros(shape=_shape(grid, *dims, extend=extend), dtype=dtype))
+    return as_field(dims, xp.zeros(shape=_shape(grid, *dims, extend=extend), dtype=dtype))
 
 
 def constant_field(

model/common/src/icon4py/model/common/test_utils/serialbox_utils.py

@@ -15,14 +15,14 @@
 import uuid
 
 import gt4py.next as gtx
-import numpy as np
 import serialbox
 
 import icon4py.model.common.decomposition.definitions as decomposition
 import icon4py.model.common.field_type_aliases as fa
 import icon4py.model.common.test_utils.helpers as helpers
 from icon4py.model.common import dimension as dims
 from icon4py.model.common.grid import base, horizontal, icon
+from icon4py.model.common.settings import xp
 from icon4py.model.common.states import prognostic_state
 
 
@@ -49,7 +49,7 @@ def wrapper(self, *args, **kwargs):
                     )
                     if dims:
                         shp = tuple(self.sizes[d] for d in dims)
-                        return gtx.as_field(dims, np.zeros(shp))
+                        return gtx.as_field(dims, xp.zeros(shp))
                     else:
                         return None
 
@@ -61,16 +61,15 @@ def log_meta_info(self):
         self.log.info(self.savepoint.metainfo)
 
     def _get_field(self, name, *dimensions, dtype=float):
-        buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype))
+        buffer = xp.squeeze(self.serializer.read(name, self.savepoint).astype(dtype))
         buffer = self._reduce_to_dim_size(buffer, dimensions)
 
         self.log.debug(f"{name} {buffer.shape}")
         return gtx.as_field(dimensions, buffer)
 
-    def _get_field_component(self, name: str, ntnd: int, dims: tuple[gtx.Dimension, gtx.Dimension]):
-        buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(float))[
-            :, :, ntnd - 1
-        ]
+    def _get_field_component(self, name: str, ntnd: int, dims: tuple[gtx.Dimension, gtx]):
+        buffer = self.serializer.read(name, self.savepoint).astype(float)
+        buffer = xp.squeeze(buffer)[:, :, ntnd - 1]
         buffer = self._reduce_to_dim_size(buffer, dims)
         self.log.debug(f"{name} {buffer.shape}")
         return gtx.as_field(dims, buffer)
@@ -298,19 +297,20 @@ def e_owner_mask(self):
     def f_e(self):
         return self._get_field("f_e", dims.EdgeDim)
 
-    def print_connectivity_info(self, name: str, ar: np.ndarray):
+    def print_connectivity_info(self, name: str, ar: xp.ndarray):
         self.log.debug(f" connectivity {name} {ar.shape}")
 
     def c2e(self):
         return self._get_connectivity_array("c2e", dims.CellDim)
 
     def _get_connectivity_array(self, name: str, target_dim: gtx.Dimension, reverse: bool = False):
         if reverse:
-            connectivity = np.transpose(self._read_int32(name, offset=1))[
+            connectivity = xp.transpose(self._read_int32(name, offset=1))[
                 : self.sizes[target_dim], :
             ]
         else:
             connectivity = self._read_int32(name, offset=1)[: self.sizes[target_dim], :]
+        connectivity = xp.asarray(connectivity)
         self.log.debug(f" connectivity {name} : {connectivity.shape}")
         return connectivity
 
@@ -322,7 +322,7 @@ def e2c2e(self):
 
     def c2e2c2e(self):
         if self._c2e2c2e() is None:
-            return np.zeros((self.sizes[dims.CellDim], 9), dtype=int)
+            return xp.zeros((self.sizes[dims.CellDim], 9), dtype=int)
         else:
             return self._c2e2c2e()
 
@@ -359,7 +359,7 @@ def refin_ctrl(self, dim: gtx.Dimension):
         field_name = "refin_ctl"
         return gtx.as_field(
             (dim,),
-            np.squeeze(
+            xp.squeeze(
                 self._read_field_for_dim(field_name, self._read_int32, dim)[: self.num(dim)], 1
             ),
         )
@@ -384,7 +384,7 @@ def _read_field_for_dim(field_name, read_func, dim: gtx.Dimension):
     def owner_mask(self, dim: gtx.Dimension):
         field_name = "owner_mask"
         mask = self._read_field_for_dim(field_name, self._read_bool, dim)
-        return np.squeeze(mask)
+        return xp.squeeze(mask)
 
     def global_index(self, dim: gtx.Dimension):
         field_name = "glb_index"
@@ -427,8 +427,8 @@ def construct_icon_grid(self, on_gpu: bool) -> icon.IconGrid:
         )
         c2e2c = self.c2e2c()
         e2c2e = self.e2c2e()
-        c2e2c0 = np.column_stack(((np.asarray(range(c2e2c.shape[0]))), c2e2c))
-        e2c2e0 = np.column_stack(((np.asarray(range(e2c2e.shape[0]))), e2c2e))
+        c2e2c0 = xp.column_stack(((range(c2e2c.shape[0])), c2e2c))
+        e2c2e0 = xp.column_stack(((range(e2c2e.shape[0])), e2c2e))
         grid = (
             icon.IconGrid(self._grid_id)
             .with_config(config)
@@ -549,7 +549,7 @@ def geofac_grdiv(self):
         return self._get_field("geofac_grdiv", dims.EdgeDim, dims.E2C2EODim)
 
     def geofac_grg(self):
-        grg = np.squeeze(self.serializer.read("geofac_grg", self.savepoint))
+        grg = xp.squeeze(self.serializer.read("geofac_grg", self.savepoint))
         num_cells = self.sizes[dims.CellDim]
         return gtx.as_field((dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 0]), gtx.as_field(
             (dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 1]
@@ -577,21 +577,21 @@ def pos_on_tplane_e_y(self):
         return helpers.as_1D_sparse_field(field[:, 0:2], dims.ECDim)
 
     def rbf_vec_coeff_e(self):
-        buffer = np.squeeze(
+        buffer = xp.squeeze(
             self.serializer.read("rbf_vec_coeff_e", self.savepoint).astype(float)
         ).transpose()
         return gtx.as_field((dims.EdgeDim, dims.E2C2EDim), buffer)
 
     @IconSavepoint.optionally_registered()
     def rbf_vec_coeff_c1(self):
-        buffer = np.squeeze(
+        buffer = xp.squeeze(
             self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(float)
         ).transpose()
         return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer)
 
     @IconSavepoint.optionally_registered()
     def rbf_vec_coeff_c2(self):
-        buffer = np.squeeze(
+        buffer = xp.squeeze(
             self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(float)
         ).transpose()
         return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer)
@@ -725,9 +725,9 @@ def wgtfac_e(self):
     def wgtfacq_e_dsl(
         self, k_level
     ):  # TODO: @abishekg7 Simplify this after serialized data is fixed
-        ar = np.squeeze(self.serializer.read("wgtfacq_e", self.savepoint))
+        ar = xp.squeeze(self.serializer.read("wgtfacq_e", self.savepoint))
         k = k_level - 3
-        ar = np.pad(ar[:, ::-1], ((0, 0), (k, 0)), "constant", constant_values=(0.0,))
+        ar = xp.pad(ar[:, ::-1], ((0, 0), (k, 0)), "constant", constant_values=(0.0,))
         return self._get_field_from_ndarray(ar, dims.EdgeDim, dims.KDim)
 
     @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim)
@@ -742,17 +742,17 @@ def geopot(self):
         return self._get_field("geopot", dims.CellDim, dims.KDim)
 
     def _read_and_reorder_sparse_field(self, name: str, sparse_size=3):
-        ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[:, :, :]
+        ser_input = xp.squeeze(self.serializer.read(name, self.savepoint))[:, :, :]
         ser_input = self._reduce_to_dim_size(ser_input, (dims.CellDim, dims.C2E2CDim, dims.KDim))
         if ser_input.shape[1] != sparse_size:
-            ser_input = np.moveaxis(ser_input, 1, -1)
+            ser_input = xp.moveaxis(ser_input, 1, -1)
 
         return self._linearize_first_2dims(
             ser_input, sparse_size=sparse_size, target_dims=(dims.CECDim, dims.KDim)
         )
 
     def _linearize_first_2dims(
-        self, data: np.ndarray, sparse_size: int, target_dims: tuple[gtx.Dimension, ...]
+        self, data: xp.ndarray, sparse_size: int, target_dims: tuple[gtx.Dimension, ...]
     ):
         old_shape = data.shape
         assert old_shape[1] == sparse_size
@@ -763,10 +763,10 @@ def zd_vertoffset(self):
         return self._read_and_reorder_sparse_field("zd_vertoffset")
 
     def zd_vertidx(self):
-        return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint))
+        return xp.squeeze(self.serializer.read("zd_vertidx", self.savepoint))
 
     def zd_indlist(self):
-        return np.squeeze(self.serializer.read("zd_indlist", self.savepoint))
+        return xp.squeeze(self.serializer.read("zd_indlist", self.savepoint))
 
 
 class IconDiffusionInitSavepoint(IconSavepoint):
@@ -799,16 +799,16 @@ def exner(self):
         return self._get_field("exner", dims.CellDim, dims.KDim)
 
     def diff_multfac_smag(self):
-        return np.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint))
+        return xp.squeeze(self.serializer.read("diff_multfac_smag", self.savepoint))
 
     def enh_smag_fac(self):
-        return np.squeeze(self.serializer.read("enh_smag_fac", self.savepoint))
+        return xp.squeeze(self.serializer.read("enh_smag_fac", self.savepoint))
 
     def smag_limit(self):
-        return np.squeeze(self.serializer.read("smag_limit", self.savepoint))
+        return xp.squeeze(self.serializer.read("smag_limit", self.savepoint))
 
     def diff_multfac_n2w(self):
-        return np.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint))
+        return xp.squeeze(self.serializer.read("diff_multfac_n2w", self.savepoint))
 
     def nudgezone_diff(self) -> int:
         return self.serializer.read("nudgezone_diff", self.savepoint)[0]