diff --git a/model/testing/src/icon4py/model/testing/serialbox.py b/model/testing/src/icon4py/model/testing/serialbox.py index d015b62e58..eb0a6c3e21 100644 --- a/model/testing/src/icon4py/model/testing/serialbox.py +++ b/model/testing/src/icon4py/model/testing/serialbox.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import functools import logging +from collections.abc import Sequence from typing import Final, Literal, TypeAlias import gt4py.next as gtx @@ -84,8 +85,10 @@ def wrapper(self, *args, **kwargs): def log_meta_info(self): self.log.info(self.savepoint.metainfo) - def _get_field(self, name, *dimensions, dtype=float): + def _get_field(self, name, *dimensions, dtype=float, transpose: None | Sequence[int] = None): buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(dtype)) + if transpose is not None: + buffer = np.transpose(buffer, axes=transpose) buffer = self._reduce_to_dim_size(buffer, dimensions) self.log.debug(f"{name} {buffer.shape}") @@ -673,7 +676,7 @@ def pos_on_tplane_e_y(self): return self._get_field("pos_on_tplane_e_y", dims.EdgeDim, dims.E2CDim)[:, 0:2] def rbf_vec_coeff_e(self): - return self._get_field("rbf_vec_coeff_e", dims.EdgeDim, dims.E2C2EDim) + return self._get_field("rbf_vec_coeff_e", dims.EdgeDim, dims.E2C2EDim, transpose=(1, 0)) @IconSavepoint.optionally_registered() def rbf_vec_coeff_c1(self): @@ -694,10 +697,20 @@ def rbf_vec_coeff_c2(self): return gtx.as_field(dimensions, buffer, allocator=self.backend) def rbf_vec_coeff_v1(self): - return self._get_field("rbf_vec_coeff_v1", dims.VertexDim, dims.V2EDim) + dimensions = (dims.VertexDim, dims.V2EDim) + buffer = np.squeeze( + self.serializer.read("rbf_vec_coeff_v", self.savepoint).astype(float)[:, 0, :] + ).transpose() + buffer = self._reduce_to_dim_size(buffer, dimensions) + return gtx.as_field(dimensions, buffer, allocator=self.backend) def rbf_vec_coeff_v2(self): - return self._get_field("rbf_vec_coeff_v2", dims.VertexDim, dims.V2EDim) + dimensions = (dims.VertexDim, dims.V2EDim) + buffer = np.squeeze( + self.serializer.read("rbf_vec_coeff_v", self.savepoint).astype(float)[:, 1, :] + ).transpose() + buffer = self._reduce_to_dim_size(buffer, dimensions) + return gtx.as_field(dimensions, buffer, allocator=self.backend) def rbf_vec_idx_v(self): return self._get_field("rbf_vec_idx_v", dims.VertexDim, dims.V2EDim) diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/common.py b/tools/src/icon4py/tools/py2fgen/wrappers/common.py index 79785b991e..f6d908926c 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/common.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/common.py @@ -97,6 +97,16 @@ ), ] +Float64Array3D: TypeAlias = Annotated[ + data_alloc.NDArray, + py2fgen.ArrayParamDescriptor( + rank=3, + dtype=ts.ScalarKind.FLOAT64, + memory_space=py2fgen.MemorySpace.MAYBE_DEVICE, + is_optional=False, + ), +] + OptionalFloat64Array1D: TypeAlias = Annotated[ data_alloc.NDArray | None, py2fgen.ArrayParamDescriptor( diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py b/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py index 0c8044a1c3..ecfe925b2d 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py @@ -63,8 +63,7 @@ def diffusion_init( geofac_grg_y: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], gtx.float64], geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], gtx.float64], nudgecoeff_e: fa.EdgeField[wpfloat], - rbf_coeff_1: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], - rbf_coeff_2: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], + rbf_vec_coeff_v: wrapper_common.Float64Array3D, zd_cellidx: wrapper_common.OptionalInt32Array2D, zd_vertidx: wrapper_common.OptionalInt32Array2D, zd_intcoef: wrapper_common.OptionalFloat64Array2D, @@ -92,7 +91,8 @@ def diffusion_init( "Need to initialise grid using 'grid_init' before running 'diffusion_init'." ) - on_gpu = theta_ref_mc.array_ns != np # TODO(havogt): expose `on_gpu` from py2fgen + xp = theta_ref_mc.array_ns + on_gpu = xp != np # TODO(havogt): expose `on_gpu` from py2fgen actual_backend = wrapper_common.select_backend( wrapper_common.BackendIntEnum(backend), on_gpu=on_gpu ) @@ -134,7 +134,6 @@ def diffusion_init( dims.KDim: nlev, } ) - xp = wgtfac_c.array_ns if zd_cellidx is None: # then zdiffu_t is False or the list on that rank is empty, then all of the following are not initialized @@ -194,6 +193,15 @@ def diffusion_init( zd_diffcoef=zd_diffcoef, ) + # Create separate fields for the two components of the RBF vector coefficients and swap. + # TODO(havogt): we could use GT4Py's named collections. + rbf_coeff_1 = gtx.as_field( + [dims.VertexDim, dims.V2EDim], xp.transpose(rbf_vec_coeff_v[:, 0, :]), allocator=allocator + ) + rbf_coeff_2 = gtx.as_field( + [dims.VertexDim, dims.V2EDim], xp.transpose(rbf_vec_coeff_v[:, 1, :]), allocator=allocator + ) + # Interpolation state interpolation_state = DiffusionInterpolationState( e_bln_c_s=e_bln_c_s, diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py b/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py index 35cfdb4843..fc2da53cec 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py @@ -56,10 +56,9 @@ def solve_nh_init( geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], pos_on_tplane_e_1: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], gtx.float64], pos_on_tplane_e_2: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], gtx.float64], - rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], gtx.float64], + rbf_vec_coeff_e: wrapper_common.Float64Array2D, e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], gtx.float64], - rbf_coeff_1: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], - rbf_coeff_2: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], gtx.float64], + rbf_vec_coeff_v: wrapper_common.Float64Array3D, geofac_div: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], gtx.float64], geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], gtx.float64], geofac_grg_x: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], gtx.float64], @@ -129,7 +128,8 @@ def solve_nh_init( if grid_wrapper.grid_state is None: raise Exception("Need to initialise grid using 'grid_init' before running 'solve_nh_init'.") - on_gpu = c_lin_e.array_ns != np # TODO(havogt): expose `on_gpu` from py2fgen + xp = c_lin_e.array_ns + on_gpu = xp != np # TODO(havogt): expose `on_gpu` from py2fgen actual_backend = wrapper_common.select_backend( wrapper_common.BackendIntEnum(backend), on_gpu=on_gpu ) @@ -181,6 +181,19 @@ def solve_nh_init( ) nonhydro_params = solve_nonhydro.NonHydrostaticParams(config) + # Create separate fields for the two components of the RBF vector coefficients and swap. + # TODO(havogt): we could use GT4Py's named collections. + rbf_coeff_1 = gtx.as_field( + [dims.VertexDim, dims.V2EDim], xp.transpose(rbf_vec_coeff_v[:, 0, :]), allocator=allocator + ) + rbf_coeff_2 = gtx.as_field( + [dims.VertexDim, dims.V2EDim], xp.transpose(rbf_vec_coeff_v[:, 1, :]), allocator=allocator + ) + + # Swap indices in rbf_vec_coeff_e. TODO(havogt): Should eventually be done on the Fortran side. + rbf_vec_coeff_e_transposed = gtx.as_field( + [dims.EdgeDim, dims.E2C2EDim], xp.transpose(rbf_vec_coeff_e), allocator=allocator + ) interpolation_state = dycore_states.InterpolationState( c_lin_e=c_lin_e, c_intp=c_intp, @@ -189,7 +202,7 @@ def solve_nh_init( geofac_rot=geofac_rot, pos_on_tplane_e_1=pos_on_tplane_e_1[:, 0:2], pos_on_tplane_e_2=pos_on_tplane_e_2[:, 0:2], - rbf_vec_coeff_e=rbf_vec_coeff_e, + rbf_vec_coeff_e=rbf_vec_coeff_e_transposed, e_bln_c_s=e_bln_c_s, rbf_coeff_1=rbf_coeff_1, rbf_coeff_2=rbf_coeff_2, diff --git a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 index 9f306faa0c..e7e4dde5ce 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 +++ b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90 @@ -126,12 +126,10 @@ function diffusion_init_wrapper(theta_ref_mc, & geofac_n2s_size_1, & nudgecoeff_e, & nudgecoeff_e_size_0, & - rbf_coeff_1, & - rbf_coeff_1_size_0, & - rbf_coeff_1_size_1, & - rbf_coeff_2, & - rbf_coeff_2_size_0, & - rbf_coeff_2_size_1, & + rbf_vec_coeff_v, & + rbf_vec_coeff_v_size_0, & + rbf_vec_coeff_v_size_1, & + rbf_vec_coeff_v_size_2, & zd_cellidx, & zd_cellidx_size_0, & zd_cellidx_size_1, & @@ -210,17 +208,13 @@ function diffusion_init_wrapper(theta_ref_mc, & integer(c_int), value :: nudgecoeff_e_size_0 - type(c_ptr), value, target :: rbf_coeff_1 + type(c_ptr), value, target :: rbf_vec_coeff_v - integer(c_int), value :: rbf_coeff_1_size_0 + integer(c_int), value :: rbf_vec_coeff_v_size_0 - integer(c_int), value :: rbf_coeff_1_size_1 + integer(c_int), value :: rbf_vec_coeff_v_size_1 - type(c_ptr), value, target :: rbf_coeff_2 - - integer(c_int), value :: rbf_coeff_2_size_0 - - integer(c_int), value :: rbf_coeff_2_size_1 + integer(c_int), value :: rbf_vec_coeff_v_size_2 type(c_ptr), value, target :: zd_cellidx @@ -483,8 +477,7 @@ subroutine diffusion_init(theta_ref_mc, & geofac_grg_y, & geofac_n2s, & nudgecoeff_e, & - rbf_coeff_1, & - rbf_coeff_2, & + rbf_vec_coeff_v, & zd_cellidx, & zd_vertidx, & zd_intcoef, & @@ -525,9 +518,7 @@ subroutine diffusion_init(theta_ref_mc, & real(c_double), dimension(:), target :: nudgecoeff_e - real(c_double), dimension(:, :), target :: rbf_coeff_1 - - real(c_double), dimension(:, :), target :: rbf_coeff_2 + real(c_double), dimension(:, :, :), target :: rbf_vec_coeff_v integer(c_int), dimension(:, :), pointer :: zd_cellidx @@ -603,13 +594,11 @@ subroutine diffusion_init(theta_ref_mc, & integer(c_int) :: nudgecoeff_e_size_0 - integer(c_int) :: rbf_coeff_1_size_0 + integer(c_int) :: rbf_vec_coeff_v_size_0 - integer(c_int) :: rbf_coeff_1_size_1 + integer(c_int) :: rbf_vec_coeff_v_size_1 - integer(c_int) :: rbf_coeff_2_size_0 - - integer(c_int) :: rbf_coeff_2_size_1 + integer(c_int) :: rbf_vec_coeff_v_size_2 integer(c_int) :: zd_cellidx_size_0 @@ -652,8 +641,7 @@ subroutine diffusion_init(theta_ref_mc, & !$acc host_data use_device(geofac_grg_y) !$acc host_data use_device(geofac_n2s) !$acc host_data use_device(nudgecoeff_e) - !$acc host_data use_device(rbf_coeff_1) - !$acc host_data use_device(rbf_coeff_2) + !$acc host_data use_device(rbf_vec_coeff_v) !$acc host_data use_device(zd_cellidx) if(associated(zd_cellidx)) !$acc host_data use_device(zd_vertidx) if(associated(zd_vertidx)) !$acc host_data use_device(zd_intcoef) if(associated(zd_intcoef)) @@ -688,11 +676,9 @@ subroutine diffusion_init(theta_ref_mc, & nudgecoeff_e_size_0 = SIZE(nudgecoeff_e, 1) - rbf_coeff_1_size_0 = SIZE(rbf_coeff_1, 1) - rbf_coeff_1_size_1 = SIZE(rbf_coeff_1, 2) - - rbf_coeff_2_size_0 = SIZE(rbf_coeff_2, 1) - rbf_coeff_2_size_1 = SIZE(rbf_coeff_2, 2) + rbf_vec_coeff_v_size_0 = SIZE(rbf_vec_coeff_v, 1) + rbf_vec_coeff_v_size_1 = SIZE(rbf_vec_coeff_v, 2) + rbf_vec_coeff_v_size_2 = SIZE(rbf_vec_coeff_v, 3) if (associated(zd_cellidx)) then zd_cellidx_ptr = c_loc(zd_cellidx) @@ -740,12 +726,10 @@ subroutine diffusion_init(theta_ref_mc, & geofac_n2s_size_1=geofac_n2s_size_1, & nudgecoeff_e=c_loc(nudgecoeff_e), & nudgecoeff_e_size_0=nudgecoeff_e_size_0, & - rbf_coeff_1=c_loc(rbf_coeff_1), & - rbf_coeff_1_size_0=rbf_coeff_1_size_0, & - rbf_coeff_1_size_1=rbf_coeff_1_size_1, & - rbf_coeff_2=c_loc(rbf_coeff_2), & - rbf_coeff_2_size_0=rbf_coeff_2_size_0, & - rbf_coeff_2_size_1=rbf_coeff_2_size_1, & + rbf_vec_coeff_v=c_loc(rbf_vec_coeff_v), & + rbf_vec_coeff_v_size_0=rbf_vec_coeff_v_size_0, & + rbf_vec_coeff_v_size_1=rbf_vec_coeff_v_size_1, & + rbf_vec_coeff_v_size_2=rbf_vec_coeff_v_size_2, & zd_cellidx=zd_cellidx_ptr, & zd_cellidx_size_0=zd_cellidx_size_0, & zd_cellidx_size_1=zd_cellidx_size_1, & @@ -788,7 +772,6 @@ subroutine diffusion_init(theta_ref_mc, & !$acc end host_data !$acc end host_data !$acc end host_data - !$acc end host_data end subroutine diffusion_init end module \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h index 46010ce5d3..a1344d8239 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h +++ b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h @@ -14,15 +14,15 @@ extern int diffusion_init_wrapper( double *geofac_grg_x, int geofac_grg_x_size_0, int geofac_grg_x_size_1, double *geofac_grg_y, int geofac_grg_y_size_0, int geofac_grg_y_size_1, double *geofac_n2s, int geofac_n2s_size_0, int geofac_n2s_size_1, - double *nudgecoeff_e, int nudgecoeff_e_size_0, double *rbf_coeff_1, - int rbf_coeff_1_size_0, int rbf_coeff_1_size_1, double *rbf_coeff_2, - int rbf_coeff_2_size_0, int rbf_coeff_2_size_1, int *zd_cellidx, - int zd_cellidx_size_0, int zd_cellidx_size_1, int *zd_vertidx, - int zd_vertidx_size_0, int zd_vertidx_size_1, double *zd_intcoef, - int zd_intcoef_size_0, int zd_intcoef_size_1, double *zd_diffcoef, - int zd_diffcoef_size_0, int ndyn_substeps, int diffusion_type, int hdiff_w, - int hdiff_vn, int hdiff_smag_w, int zdiffu_t, int type_t_diffu, - int type_vn_diffu, double hdiff_efdt_ratio, double hdiff_w_efdt_ratio, + double *nudgecoeff_e, int nudgecoeff_e_size_0, double *rbf_vec_coeff_v, + int rbf_vec_coeff_v_size_0, int rbf_vec_coeff_v_size_1, + int rbf_vec_coeff_v_size_2, int *zd_cellidx, int zd_cellidx_size_0, + int zd_cellidx_size_1, int *zd_vertidx, int zd_vertidx_size_0, + int zd_vertidx_size_1, double *zd_intcoef, int zd_intcoef_size_0, + int zd_intcoef_size_1, double *zd_diffcoef, int zd_diffcoef_size_0, + int ndyn_substeps, int diffusion_type, int hdiff_w, int hdiff_vn, + int hdiff_smag_w, int zdiffu_t, int type_t_diffu, int type_vn_diffu, + double hdiff_efdt_ratio, double hdiff_w_efdt_ratio, double smagorinski_scaling_factor, int hdiff_temp, double denom_diffu_v, double nudge_max_coeff, int itype_sher, int ltkeshs, int backend, int on_gpu); \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py index dd26c83bfb..6c64ea4b15 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py +++ b/tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py @@ -365,12 +365,10 @@ def diffusion_init_wrapper( geofac_n2s_size_1, nudgecoeff_e, nudgecoeff_e_size_0, - rbf_coeff_1, - rbf_coeff_1_size_0, - rbf_coeff_1_size_1, - rbf_coeff_2, - rbf_coeff_2_size_0, - rbf_coeff_2_size_1, + rbf_vec_coeff_v, + rbf_vec_coeff_v_size_0, + rbf_vec_coeff_v_size_1, + rbf_vec_coeff_v_size_2, zd_cellidx, zd_cellidx_size_0, zd_cellidx_size_1, @@ -484,21 +482,12 @@ def diffusion_init_wrapper( nudgecoeff_e = (nudgecoeff_e, (nudgecoeff_e_size_0,), on_gpu, False) - rbf_coeff_1 = ( - rbf_coeff_1, + rbf_vec_coeff_v = ( + rbf_vec_coeff_v, ( - rbf_coeff_1_size_0, - rbf_coeff_1_size_1, - ), - on_gpu, - False, - ) - - rbf_coeff_2 = ( - rbf_coeff_2, - ( - rbf_coeff_2_size_0, - rbf_coeff_2_size_1, + rbf_vec_coeff_v_size_0, + rbf_vec_coeff_v_size_1, + rbf_vec_coeff_v_size_2, ), on_gpu, False, @@ -561,8 +550,7 @@ def diffusion_init_wrapper( geofac_grg_y=geofac_grg_y, geofac_n2s=geofac_n2s, nudgecoeff_e=nudgecoeff_e, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, + rbf_vec_coeff_v=rbf_vec_coeff_v, zd_cellidx=zd_cellidx, zd_vertidx=zd_vertidx, zd_intcoef=zd_intcoef, @@ -730,34 +718,18 @@ def diffusion_init_wrapper( ) logger.debug(msg) - rbf_coeff_1_arr = ( - _conversion.as_array(ffi, rbf_coeff_1, _definitions.FLOAT64) - if rbf_coeff_1 is not None - else None - ) - msg = "shape of rbf_coeff_1 after computation = %s" % str( - rbf_coeff_1_arr.shape if rbf_coeff_1 is not None else "None" - ) - logger.debug(msg) - msg = ( - "rbf_coeff_1 after computation: %s" % str(rbf_coeff_1_arr) - if rbf_coeff_1 is not None - else "None" - ) - logger.debug(msg) - - rbf_coeff_2_arr = ( - _conversion.as_array(ffi, rbf_coeff_2, _definitions.FLOAT64) - if rbf_coeff_2 is not None + rbf_vec_coeff_v_arr = ( + _conversion.as_array(ffi, rbf_vec_coeff_v, _definitions.FLOAT64) + if rbf_vec_coeff_v is not None else None ) - msg = "shape of rbf_coeff_2 after computation = %s" % str( - rbf_coeff_2_arr.shape if rbf_coeff_2 is not None else "None" + msg = "shape of rbf_vec_coeff_v after computation = %s" % str( + rbf_vec_coeff_v_arr.shape if rbf_vec_coeff_v is not None else "None" ) logger.debug(msg) msg = ( - "rbf_coeff_2 after computation: %s" % str(rbf_coeff_2_arr) - if rbf_coeff_2 is not None + "rbf_vec_coeff_v after computation: %s" % str(rbf_vec_coeff_v_arr) + if rbf_vec_coeff_v is not None else "None" ) logger.debug(msg) diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 index 6b29e00d95..025ce06a22 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90 @@ -382,12 +382,10 @@ function solve_nh_init_wrapper(c_lin_e, & e_bln_c_s, & e_bln_c_s_size_0, & e_bln_c_s_size_1, & - rbf_coeff_1, & - rbf_coeff_1_size_0, & - rbf_coeff_1_size_1, & - rbf_coeff_2, & - rbf_coeff_2_size_0, & - rbf_coeff_2_size_1, & + rbf_vec_coeff_v, & + rbf_vec_coeff_v_size_0, & + rbf_vec_coeff_v_size_1, & + rbf_vec_coeff_v_size_2, & geofac_div, & geofac_div_size_0, & geofac_div_size_1, & @@ -580,17 +578,13 @@ function solve_nh_init_wrapper(c_lin_e, & integer(c_int), value :: e_bln_c_s_size_1 - type(c_ptr), value, target :: rbf_coeff_1 + type(c_ptr), value, target :: rbf_vec_coeff_v - integer(c_int), value :: rbf_coeff_1_size_0 + integer(c_int), value :: rbf_vec_coeff_v_size_0 - integer(c_int), value :: rbf_coeff_1_size_1 + integer(c_int), value :: rbf_vec_coeff_v_size_1 - type(c_ptr), value, target :: rbf_coeff_2 - - integer(c_int), value :: rbf_coeff_2_size_0 - - integer(c_int), value :: rbf_coeff_2_size_1 + integer(c_int), value :: rbf_vec_coeff_v_size_2 type(c_ptr), value, target :: geofac_div @@ -1474,8 +1468,7 @@ subroutine solve_nh_init(c_lin_e, & pos_on_tplane_e_2, & rbf_vec_coeff_e, & e_bln_c_s, & - rbf_coeff_1, & - rbf_coeff_2, & + rbf_vec_coeff_v, & geofac_div, & geofac_n2s, & geofac_grg_x, & @@ -1562,9 +1555,7 @@ subroutine solve_nh_init(c_lin_e, & real(c_double), dimension(:, :), target :: e_bln_c_s - real(c_double), dimension(:, :), target :: rbf_coeff_1 - - real(c_double), dimension(:, :), target :: rbf_coeff_2 + real(c_double), dimension(:, :, :), target :: rbf_vec_coeff_v real(c_double), dimension(:, :), target :: geofac_div @@ -1734,13 +1725,11 @@ subroutine solve_nh_init(c_lin_e, & integer(c_int) :: e_bln_c_s_size_1 - integer(c_int) :: rbf_coeff_1_size_0 + integer(c_int) :: rbf_vec_coeff_v_size_0 - integer(c_int) :: rbf_coeff_1_size_1 + integer(c_int) :: rbf_vec_coeff_v_size_1 - integer(c_int) :: rbf_coeff_2_size_0 - - integer(c_int) :: rbf_coeff_2_size_1 + integer(c_int) :: rbf_vec_coeff_v_size_2 integer(c_int) :: geofac_div_size_0 @@ -1904,8 +1893,7 @@ subroutine solve_nh_init(c_lin_e, & !$acc host_data use_device(pos_on_tplane_e_2) !$acc host_data use_device(rbf_vec_coeff_e) !$acc host_data use_device(e_bln_c_s) - !$acc host_data use_device(rbf_coeff_1) - !$acc host_data use_device(rbf_coeff_2) + !$acc host_data use_device(rbf_vec_coeff_v) !$acc host_data use_device(geofac_div) !$acc host_data use_device(geofac_n2s) !$acc host_data use_device(geofac_grg_x) @@ -1979,11 +1967,9 @@ subroutine solve_nh_init(c_lin_e, & e_bln_c_s_size_0 = SIZE(e_bln_c_s, 1) e_bln_c_s_size_1 = SIZE(e_bln_c_s, 2) - rbf_coeff_1_size_0 = SIZE(rbf_coeff_1, 1) - rbf_coeff_1_size_1 = SIZE(rbf_coeff_1, 2) - - rbf_coeff_2_size_0 = SIZE(rbf_coeff_2, 1) - rbf_coeff_2_size_1 = SIZE(rbf_coeff_2, 2) + rbf_vec_coeff_v_size_0 = SIZE(rbf_vec_coeff_v, 1) + rbf_vec_coeff_v_size_1 = SIZE(rbf_vec_coeff_v, 2) + rbf_vec_coeff_v_size_2 = SIZE(rbf_vec_coeff_v, 3) geofac_div_size_0 = SIZE(geofac_div, 1) geofac_div_size_1 = SIZE(geofac_div, 2) @@ -2129,12 +2115,10 @@ subroutine solve_nh_init(c_lin_e, & e_bln_c_s=c_loc(e_bln_c_s), & e_bln_c_s_size_0=e_bln_c_s_size_0, & e_bln_c_s_size_1=e_bln_c_s_size_1, & - rbf_coeff_1=c_loc(rbf_coeff_1), & - rbf_coeff_1_size_0=rbf_coeff_1_size_0, & - rbf_coeff_1_size_1=rbf_coeff_1_size_1, & - rbf_coeff_2=c_loc(rbf_coeff_2), & - rbf_coeff_2_size_0=rbf_coeff_2_size_0, & - rbf_coeff_2_size_1=rbf_coeff_2_size_1, & + rbf_vec_coeff_v=c_loc(rbf_vec_coeff_v), & + rbf_vec_coeff_v_size_0=rbf_vec_coeff_v_size_0, & + rbf_vec_coeff_v_size_1=rbf_vec_coeff_v_size_1, & + rbf_vec_coeff_v_size_2=rbf_vec_coeff_v_size_2, & geofac_div=c_loc(geofac_div), & geofac_div_size_0=geofac_div_size_0, & geofac_div_size_1=geofac_div_size_1, & @@ -2319,7 +2303,6 @@ subroutine solve_nh_init(c_lin_e, & !$acc end host_data !$acc end host_data !$acc end host_data - !$acc end host_data end subroutine solve_nh_init end module \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h index 660a2fc279..f8cf1b9a09 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h @@ -47,52 +47,51 @@ extern int solve_nh_init_wrapper( double *pos_on_tplane_e_2, int pos_on_tplane_e_2_size_0, int pos_on_tplane_e_2_size_1, double *rbf_vec_coeff_e, int rbf_vec_coeff_e_size_0, int rbf_vec_coeff_e_size_1, double *e_bln_c_s, - int e_bln_c_s_size_0, int e_bln_c_s_size_1, double *rbf_coeff_1, - int rbf_coeff_1_size_0, int rbf_coeff_1_size_1, double *rbf_coeff_2, - int rbf_coeff_2_size_0, int rbf_coeff_2_size_1, double *geofac_div, - int geofac_div_size_0, int geofac_div_size_1, double *geofac_n2s, - int geofac_n2s_size_0, int geofac_n2s_size_1, double *geofac_grg_x, - int geofac_grg_x_size_0, int geofac_grg_x_size_1, double *geofac_grg_y, - int geofac_grg_y_size_0, int geofac_grg_y_size_1, double *nudgecoeff_e, - int nudgecoeff_e_size_0, int *mask_prog_halo_c, int mask_prog_halo_c_size_0, - double *rayleigh_w, int rayleigh_w_size_0, double *exner_exfac, - int exner_exfac_size_0, int exner_exfac_size_1, double *exner_ref_mc, - int exner_ref_mc_size_0, int exner_ref_mc_size_1, double *wgtfac_c, - int wgtfac_c_size_0, int wgtfac_c_size_1, double *wgtfacq_c, - int wgtfacq_c_size_0, int wgtfacq_c_size_1, double *inv_ddqz_z_full, - int inv_ddqz_z_full_size_0, int inv_ddqz_z_full_size_1, double *rho_ref_mc, - int rho_ref_mc_size_0, int rho_ref_mc_size_1, double *theta_ref_mc, - int theta_ref_mc_size_0, int theta_ref_mc_size_1, double *vwind_expl_wgt, - int vwind_expl_wgt_size_0, double *d_exner_dz_ref_ic, - int d_exner_dz_ref_ic_size_0, int d_exner_dz_ref_ic_size_1, - double *ddqz_z_half, int ddqz_z_half_size_0, int ddqz_z_half_size_1, - double *theta_ref_ic, int theta_ref_ic_size_0, int theta_ref_ic_size_1, - double *d2dexdz2_fac1_mc, int d2dexdz2_fac1_mc_size_0, - int d2dexdz2_fac1_mc_size_1, double *d2dexdz2_fac2_mc, - int d2dexdz2_fac2_mc_size_0, int d2dexdz2_fac2_mc_size_1, - double *rho_ref_me, int rho_ref_me_size_0, int rho_ref_me_size_1, - double *theta_ref_me, int theta_ref_me_size_0, int theta_ref_me_size_1, - double *ddxn_z_full, int ddxn_z_full_size_0, int ddxn_z_full_size_1, - double *zdiff_gradp, int zdiff_gradp_size_0, int zdiff_gradp_size_1, - int zdiff_gradp_size_2, int *vertoffset_gradp, int vertoffset_gradp_size_0, - int vertoffset_gradp_size_1, int vertoffset_gradp_size_2, int *pg_edgeidx, - int pg_edgeidx_size_0, int *pg_vertidx, int pg_vertidx_size_0, - double *pg_exdist, int pg_exdist_size_0, double *ddqz_z_full_e, - int ddqz_z_full_e_size_0, int ddqz_z_full_e_size_1, double *ddxt_z_full, - int ddxt_z_full_size_0, int ddxt_z_full_size_1, double *wgtfac_e, - int wgtfac_e_size_0, int wgtfac_e_size_1, double *wgtfacq_e, - int wgtfacq_e_size_0, int wgtfacq_e_size_1, double *vwind_impl_wgt, - int vwind_impl_wgt_size_0, double *hmask_dd3d, int hmask_dd3d_size_0, - double *scalfac_dd3d, int scalfac_dd3d_size_0, double *coeff1_dwdz, - int coeff1_dwdz_size_0, int coeff1_dwdz_size_1, double *coeff2_dwdz, - int coeff2_dwdz_size_0, int coeff2_dwdz_size_1, double *coeff_gradekin, - int coeff_gradekin_size_0, int coeff_gradekin_size_1, int *c_owner_mask, - int c_owner_mask_size_0, int itime_scheme, int iadv_rhotheta, - int igradp_method, int rayleigh_type, double rayleigh_coeff, - int divdamp_order, int is_iau_active, double iau_wgt_dyn, int divdamp_type, - double divdamp_trans_start, double divdamp_trans_end, int l_vert_nested, - int ldeepatmo, double rhotheta_offctr, double veladv_offctr, - double nudge_max_coeff, double divdamp_fac, double divdamp_fac2, - double divdamp_fac3, double divdamp_fac4, double divdamp_z, - double divdamp_z2, double divdamp_z3, double divdamp_z4, int nflat_gradp, - int backend, int on_gpu); \ No newline at end of file + int e_bln_c_s_size_0, int e_bln_c_s_size_1, double *rbf_vec_coeff_v, + int rbf_vec_coeff_v_size_0, int rbf_vec_coeff_v_size_1, + int rbf_vec_coeff_v_size_2, double *geofac_div, int geofac_div_size_0, + int geofac_div_size_1, double *geofac_n2s, int geofac_n2s_size_0, + int geofac_n2s_size_1, double *geofac_grg_x, int geofac_grg_x_size_0, + int geofac_grg_x_size_1, double *geofac_grg_y, int geofac_grg_y_size_0, + int geofac_grg_y_size_1, double *nudgecoeff_e, int nudgecoeff_e_size_0, + int *mask_prog_halo_c, int mask_prog_halo_c_size_0, double *rayleigh_w, + int rayleigh_w_size_0, double *exner_exfac, int exner_exfac_size_0, + int exner_exfac_size_1, double *exner_ref_mc, int exner_ref_mc_size_0, + int exner_ref_mc_size_1, double *wgtfac_c, int wgtfac_c_size_0, + int wgtfac_c_size_1, double *wgtfacq_c, int wgtfacq_c_size_0, + int wgtfacq_c_size_1, double *inv_ddqz_z_full, int inv_ddqz_z_full_size_0, + int inv_ddqz_z_full_size_1, double *rho_ref_mc, int rho_ref_mc_size_0, + int rho_ref_mc_size_1, double *theta_ref_mc, int theta_ref_mc_size_0, + int theta_ref_mc_size_1, double *vwind_expl_wgt, int vwind_expl_wgt_size_0, + double *d_exner_dz_ref_ic, int d_exner_dz_ref_ic_size_0, + int d_exner_dz_ref_ic_size_1, double *ddqz_z_half, int ddqz_z_half_size_0, + int ddqz_z_half_size_1, double *theta_ref_ic, int theta_ref_ic_size_0, + int theta_ref_ic_size_1, double *d2dexdz2_fac1_mc, + int d2dexdz2_fac1_mc_size_0, int d2dexdz2_fac1_mc_size_1, + double *d2dexdz2_fac2_mc, int d2dexdz2_fac2_mc_size_0, + int d2dexdz2_fac2_mc_size_1, double *rho_ref_me, int rho_ref_me_size_0, + int rho_ref_me_size_1, double *theta_ref_me, int theta_ref_me_size_0, + int theta_ref_me_size_1, double *ddxn_z_full, int ddxn_z_full_size_0, + int ddxn_z_full_size_1, double *zdiff_gradp, int zdiff_gradp_size_0, + int zdiff_gradp_size_1, int zdiff_gradp_size_2, int *vertoffset_gradp, + int vertoffset_gradp_size_0, int vertoffset_gradp_size_1, + int vertoffset_gradp_size_2, int *pg_edgeidx, int pg_edgeidx_size_0, + int *pg_vertidx, int pg_vertidx_size_0, double *pg_exdist, + int pg_exdist_size_0, double *ddqz_z_full_e, int ddqz_z_full_e_size_0, + int ddqz_z_full_e_size_1, double *ddxt_z_full, int ddxt_z_full_size_0, + int ddxt_z_full_size_1, double *wgtfac_e, int wgtfac_e_size_0, + int wgtfac_e_size_1, double *wgtfacq_e, int wgtfacq_e_size_0, + int wgtfacq_e_size_1, double *vwind_impl_wgt, int vwind_impl_wgt_size_0, + double *hmask_dd3d, int hmask_dd3d_size_0, double *scalfac_dd3d, + int scalfac_dd3d_size_0, double *coeff1_dwdz, int coeff1_dwdz_size_0, + int coeff1_dwdz_size_1, double *coeff2_dwdz, int coeff2_dwdz_size_0, + int coeff2_dwdz_size_1, double *coeff_gradekin, int coeff_gradekin_size_0, + int coeff_gradekin_size_1, int *c_owner_mask, int c_owner_mask_size_0, + int itime_scheme, int iadv_rhotheta, int igradp_method, int rayleigh_type, + double rayleigh_coeff, int divdamp_order, int is_iau_active, + double iau_wgt_dyn, int divdamp_type, double divdamp_trans_start, + double divdamp_trans_end, int l_vert_nested, int ldeepatmo, + double rhotheta_offctr, double veladv_offctr, double nudge_max_coeff, + double divdamp_fac, double divdamp_fac2, double divdamp_fac3, + double divdamp_fac4, double divdamp_z, double divdamp_z2, double divdamp_z3, + double divdamp_z4, int nflat_gradp, int backend, int on_gpu); \ No newline at end of file diff --git a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py index 15ead38e9a..74430e8b2a 100644 --- a/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py +++ b/tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py @@ -1197,12 +1197,10 @@ def solve_nh_init_wrapper( e_bln_c_s, e_bln_c_s_size_0, e_bln_c_s_size_1, - rbf_coeff_1, - rbf_coeff_1_size_0, - rbf_coeff_1_size_1, - rbf_coeff_2, - rbf_coeff_2_size_0, - rbf_coeff_2_size_1, + rbf_vec_coeff_v, + rbf_vec_coeff_v_size_0, + rbf_vec_coeff_v_size_1, + rbf_vec_coeff_v_size_2, geofac_div, geofac_div_size_0, geofac_div_size_1, @@ -1440,21 +1438,12 @@ def solve_nh_init_wrapper( False, ) - rbf_coeff_1 = ( - rbf_coeff_1, + rbf_vec_coeff_v = ( + rbf_vec_coeff_v, ( - rbf_coeff_1_size_0, - rbf_coeff_1_size_1, - ), - on_gpu, - False, - ) - - rbf_coeff_2 = ( - rbf_coeff_2, - ( - rbf_coeff_2_size_0, - rbf_coeff_2_size_1, + rbf_vec_coeff_v_size_0, + rbf_vec_coeff_v_size_1, + rbf_vec_coeff_v_size_2, ), on_gpu, False, @@ -1790,8 +1779,7 @@ def solve_nh_init_wrapper( pos_on_tplane_e_2=pos_on_tplane_e_2, rbf_vec_coeff_e=rbf_vec_coeff_e, e_bln_c_s=e_bln_c_s, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, + rbf_vec_coeff_v=rbf_vec_coeff_v, geofac_div=geofac_div, geofac_n2s=geofac_n2s, geofac_grg_x=geofac_grg_x, @@ -2019,34 +2007,18 @@ def solve_nh_init_wrapper( ) logger.debug(msg) - rbf_coeff_1_arr = ( - _conversion.as_array(ffi, rbf_coeff_1, _definitions.FLOAT64) - if rbf_coeff_1 is not None - else None - ) - msg = "shape of rbf_coeff_1 after computation = %s" % str( - rbf_coeff_1_arr.shape if rbf_coeff_1 is not None else "None" - ) - logger.debug(msg) - msg = ( - "rbf_coeff_1 after computation: %s" % str(rbf_coeff_1_arr) - if rbf_coeff_1 is not None - else "None" - ) - logger.debug(msg) - - rbf_coeff_2_arr = ( - _conversion.as_array(ffi, rbf_coeff_2, _definitions.FLOAT64) - if rbf_coeff_2 is not None + rbf_vec_coeff_v_arr = ( + _conversion.as_array(ffi, rbf_vec_coeff_v, _definitions.FLOAT64) + if rbf_vec_coeff_v is not None else None ) - msg = "shape of rbf_coeff_2 after computation = %s" % str( - rbf_coeff_2_arr.shape if rbf_coeff_2 is not None else "None" + msg = "shape of rbf_vec_coeff_v after computation = %s" % str( + rbf_vec_coeff_v_arr.shape if rbf_vec_coeff_v is not None else "None" ) logger.debug(msg) msg = ( - "rbf_coeff_2 after computation: %s" % str(rbf_coeff_2_arr) - if rbf_coeff_2 is not None + "rbf_vec_coeff_v after computation: %s" % str(rbf_vec_coeff_v_arr) + if rbf_vec_coeff_v is not None else "None" ) logger.debug(msg) diff --git a/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py index f85ec964b7..cd735794ff 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py @@ -96,8 +96,16 @@ def test_diffusion_wrapper_granule_inputs( geofac_grg_y = test_utils.array_to_array_info(geofac_grg_y_field.ndarray) geofac_n2s = test_utils.array_to_array_info(interpolation_savepoint.geofac_n2s().ndarray) nudgecoeff_e = test_utils.array_to_array_info(interpolation_savepoint.nudgecoeff_e().ndarray) - rbf_coeff_1 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v1().ndarray) - rbf_coeff_2 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v2().ndarray) + + rbf_vec_coeff_v_array = np.squeeze( + interpolation_savepoint.serializer.read( + "rbf_vec_coeff_v", interpolation_savepoint.savepoint + ).astype(float) + ) + rbf_vec_coeff_v_array = interpolation_savepoint._reduce_to_dim_size( + rbf_vec_coeff_v_array, [dims.V2EDim, dims.V2EDim, dims.VertexDim] + ) + rbf_vec_coeff_v = test_utils.array_to_array_info(rbf_vec_coeff_v_array) # --- Extract Diagnostic and Prognostic State Parameters --- hdef_ic = test_utils.array_to_array_info(savepoint_diffusion_init.hdef_ic().ndarray) @@ -158,8 +166,7 @@ def test_diffusion_wrapper_granule_inputs( geofac_grg_y=geofac_grg_y, geofac_n2s=geofac_n2s, nudgecoeff_e=nudgecoeff_e, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, + rbf_vec_coeff_v=rbf_vec_coeff_v, zd_cellidx=zd_cellidx, zd_vertidx=zd_vertidx, zd_intcoef=zd_intcoef, @@ -325,8 +332,15 @@ def test_diffusion_wrapper_single_step( geofac_grg_y = test_utils.array_to_array_info(geofac_grg_y_field.ndarray) geofac_n2s = test_utils.array_to_array_info(interpolation_savepoint.geofac_n2s().ndarray) nudgecoeff_e = test_utils.array_to_array_info(interpolation_savepoint.nudgecoeff_e().ndarray) - rbf_coeff_1 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v1().ndarray) - rbf_coeff_2 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v2().ndarray) + rbf_vec_coeff_v_array = np.squeeze( + interpolation_savepoint.serializer.read( + "rbf_vec_coeff_v", interpolation_savepoint.savepoint + ).astype(float) + ) + rbf_vec_coeff_v_array = interpolation_savepoint._reduce_to_dim_size( + rbf_vec_coeff_v_array, [dims.V2EDim, dims.V2EDim, dims.VertexDim] + ) + rbf_vec_coeff_v = test_utils.array_to_array_info(rbf_vec_coeff_v_array) # Diagnostic state parameters hdef_ic = test_utils.array_to_array_info(savepoint_diffusion_init.hdef_ic().ndarray) @@ -355,8 +369,7 @@ def test_diffusion_wrapper_single_step( geofac_grg_y=geofac_grg_y, geofac_n2s=geofac_n2s, nudgecoeff_e=nudgecoeff_e, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, + rbf_vec_coeff_v=rbf_vec_coeff_v, zd_cellidx=zd_cellidx, zd_vertidx=zd_vertidx, zd_intcoef=zd_intcoef, diff --git a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py index 726a86e7bd..b411f929fd 100644 --- a/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py +++ b/tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py @@ -134,11 +134,19 @@ def solve_nh_init( pos_on_tplane_e_2 = test_utils.array_to_array_info(pos_on_tplane_e_2_field.ndarray) rbf_vec_coeff_e = test_utils.array_to_array_info( - interpolation_savepoint.rbf_vec_coeff_e().ndarray + interpolation_savepoint._get_field("rbf_vec_coeff_e", dims.E2C2EDim, dims.EdgeDim).ndarray ) e_bln_c_s = test_utils.array_to_array_info(interpolation_savepoint.e_bln_c_s().ndarray) - rbf_coeff_1 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v1().ndarray) - rbf_coeff_2 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v2().ndarray) + + rbf_vec_coeff_v_array = np.squeeze( + interpolation_savepoint.serializer.read( + "rbf_vec_coeff_v", interpolation_savepoint.savepoint + ).astype(float) + ) + rbf_vec_coeff_v_array = interpolation_savepoint._reduce_to_dim_size( + rbf_vec_coeff_v_array, [dims.V2EDim, dims.V2EDim, dims.VertexDim] + ) + rbf_vec_coeff_v = test_utils.array_to_array_info(rbf_vec_coeff_v_array) geofac_div = test_utils.array_to_array_info(interpolation_savepoint.geofac_div().ndarray) geofac_n2s = test_utils.array_to_array_info(interpolation_savepoint.geofac_n2s().ndarray) geofac_grg_x = test_utils.array_to_array_info(interpolation_savepoint.geofac_grg()[0].ndarray) @@ -161,8 +169,7 @@ def solve_nh_init( pos_on_tplane_e_2=pos_on_tplane_e_2, rbf_vec_coeff_e=rbf_vec_coeff_e, e_bln_c_s=e_bln_c_s, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, + rbf_vec_coeff_v=rbf_vec_coeff_v, geofac_div=geofac_div, geofac_n2s=geofac_n2s, geofac_grg_x=geofac_grg_x, @@ -227,7 +234,7 @@ def solve_nh_init( divdamp_z3=divdamp_z3, divdamp_z4=divdamp_z4, nflat_gradp=nflat_gradp, - backend=wrapper_common.BackendIntEnum.DEFAULT, + backend=wrapper_common.BackendIntEnum.GTFN, ) @@ -374,11 +381,18 @@ def test_dycore_wrapper_granule_inputs( pos_on_tplane_e_2 = test_utils.array_to_array_info(pos_on_tplane_e_2_field.ndarray) rbf_vec_coeff_e = test_utils.array_to_array_info( - interpolation_savepoint.rbf_vec_coeff_e().ndarray + interpolation_savepoint._get_field("rbf_vec_coeff_e", dims.E2C2EDim, dims.EdgeDim).ndarray ) e_bln_c_s = test_utils.array_to_array_info(interpolation_savepoint.e_bln_c_s().ndarray) - rbf_coeff_1 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v1().ndarray) - rbf_coeff_2 = test_utils.array_to_array_info(interpolation_savepoint.rbf_vec_coeff_v2().ndarray) + rbf_vec_coeff_v_array = np.squeeze( + interpolation_savepoint.serializer.read( + "rbf_vec_coeff_v", interpolation_savepoint.savepoint + ).astype(float) + ) + rbf_vec_coeff_v_array = interpolation_savepoint._reduce_to_dim_size( + rbf_vec_coeff_v_array, [dims.V2EDim, dims.V2EDim, dims.VertexDim] + ) + rbf_vec_coeff_v = test_utils.array_to_array_info(rbf_vec_coeff_v_array) geofac_div = test_utils.array_to_array_info(interpolation_savepoint.geofac_div().ndarray) geofac_n2s = test_utils.array_to_array_info(interpolation_savepoint.geofac_n2s().ndarray) geofac_grg_x = test_utils.array_to_array_info(interpolation_savepoint.geofac_grg()[0].ndarray) @@ -582,8 +596,7 @@ def test_dycore_wrapper_granule_inputs( pos_on_tplane_e_2=pos_on_tplane_e_2, rbf_vec_coeff_e=rbf_vec_coeff_e, e_bln_c_s=e_bln_c_s, - rbf_coeff_1=rbf_coeff_1, - rbf_coeff_2=rbf_coeff_2, + rbf_vec_coeff_v=rbf_vec_coeff_v, geofac_div=geofac_div, geofac_n2s=geofac_n2s, geofac_grg_x=geofac_grg_x,