Move transformations (index2offset, transpose) from Fortran to Python

model/atmosphere/advection/tests/advection/integration_tests/test_advection.py

@@ -21,6 +21,7 @@
 from icon4py.model.common.utils import data_allocation as data_alloc
 from icon4py.model.testing import (
     definitions,
+    exchange_utils,
     grid_utils,
     grid_utils as gridtest_utils,
     serialbox as sb,
@@ -47,7 +48,6 @@
     construct_least_squares_state,
     construct_metric_state,
     construct_prep_adv,
-    dummy_exchange,
     log_serialized,
     verify_advection_fields,
 )
@@ -162,7 +162,7 @@ def test_advection_run_single_step(
         ),
         min_rlcell_int=icon_grid.end_index(h_grid.domain(dims.CellDim)(h_grid.Zone.LOCAL)),
         geometry_type=icon_grid.geometry_type,
-        exchange=dummy_exchange,
+        exchange=exchange_utils.dummy_exchange_with_bound_dim,
     )
 
     least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend)
@@ -270,7 +270,7 @@ def test_compute_lsq_coeffs(
         start_idx,
         min_rlcell_int,
         icon_grid.geometry_type,
-        exchange=dummy_exchange,
+        exchange=exchange_utils.dummy_exchange_with_bound_dim,
     )
 
     assert test_helpers.dallclose(

model/atmosphere/advection/tests/advection/utils.py

@@ -197,7 +197,3 @@ def verify_advection_fields(
         p_tracer_new_ref.asnumpy()[p_tracer_new_range, :],
         atol=1e-16,
     )
-
-
-def dummy_exchange(*field: data_alloc.NDArray) -> None:
-    return None

model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py

@@ -1094,7 +1094,7 @@ def test_compute_perturbed_quantities_and_interpolation(
     reference_theta_at_cells_on_half_levels = metrics_savepoint.theta_ref_ic()
     d2dexdz2_fac1_mc = metrics_savepoint.d2dexdz2_fac1_mc()
     d2dexdz2_fac2_mc = metrics_savepoint.d2dexdz2_fac2_mc()
-    wgtfacq_c = metrics_savepoint.wgtfacq_c_dsl()
+    wgtfacq_c = metrics_savepoint.wgtfacq_c()
     wgtfac_c = metrics_savepoint.wgtfac_c()
     exner_w_explicit_weight_parameter = metrics_savepoint.vwind_expl_wgt()
     ddz_of_reference_exner_at_cells_on_half_levels = metrics_savepoint.d_exner_dz_ref_ic()
@@ -1166,17 +1166,24 @@ def test_compute_perturbed_quantities_and_interpolation(
     assert test_utils.dallclose(
         perturbed_theta_v_at_cells_on_model_levels.asnumpy(), z_rth_pr_2_ref.asnumpy()
     )
+    # `z_exner_ex_pr` is only computed in a subset of the whole domain, reference may contain garbage outside this range
     assert test_utils.dallclose(
-        temporal_extrapolation_of_perturbed_exner.asnumpy(), z_exner_ex_pr_ref.asnumpy()
+        temporal_extrapolation_of_perturbed_exner.asnumpy()[
+            start_cell_lateral_boundary_level_3:end_cell_halo, :
+        ],
+        z_exner_ex_pr_ref.asnumpy()[start_cell_lateral_boundary_level_3:end_cell_halo, :],
     )
     assert test_utils.dallclose(
         perturbed_exner_at_cells_on_model_levels.asnumpy(), exner_pr_ref.asnumpy()
     )
     assert test_utils.dallclose(rho_at_cells_on_half_levels.asnumpy(), rho_ic_ref.asnumpy())
 
+    # `exner_at_cells_on_half_levels` is only computed in a subset of the whole domain, reference may contain garbage outside this range
     assert test_utils.dallclose(
-        exner_at_cells_on_half_levels.asnumpy()[:, nflatlev:],
-        z_exner_ic_ref.asnumpy()[:, nflatlev:],
+        exner_at_cells_on_half_levels.asnumpy()[
+            start_cell_lateral_boundary_level_3:end_cell_halo, nflatlev:
+        ],
+        z_exner_ic_ref.asnumpy()[start_cell_lateral_boundary_level_3:end_cell_halo, nflatlev:],
         rtol=1e-11,
     )
 
@@ -1768,7 +1775,7 @@ def test_compute_horizontal_velocity_quantities_and_fluxes(
     ddxn_z_full = metrics_savepoint.ddxn_z_full()
     ddxt_z_full = metrics_savepoint.ddxt_z_full()
     wgtfac_e = metrics_savepoint.wgtfac_e()
-    wgtfacq_e = metrics_savepoint.wgtfacq_e_dsl()
+    wgtfacq_e = metrics_savepoint.wgtfacq_e()
     rbf_vec_coeff_e = interpolation_savepoint.rbf_vec_coeff_e()
     geofac_grdiv = interpolation_savepoint.geofac_grdiv()
     nflatlev = vertical_params.nflatlev
@@ -2152,7 +2159,7 @@ def test_vertically_implicit_solver_at_predictor_step(
         reference_exner_at_cells_on_model_levels=metrics_savepoint.exner_ref_mc(),
         e_bln_c_s=interpolation_savepoint.e_bln_c_s(),
         wgtfac_c=metrics_savepoint.wgtfac_c(),
-        wgtfacq_c=metrics_savepoint.wgtfacq_c_dsl(),
+        wgtfacq_c=metrics_savepoint.wgtfacq_c(),
         iau_wgt_dyn=iau_wgt_dyn,
         dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"),
         is_iau_active=is_iau_active,

model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py

@@ -494,7 +494,7 @@ def test_compute_diagnostics_from_normal_wind(
     ddxn_z_full = metrics_savepoint.ddxn_z_full()
     ddxt_z_full = metrics_savepoint.ddxt_z_full()
     contravariant_correction_at_edges_on_model_levels = savepoint_velocity_init.z_w_concorr_me()
-    wgtfacq_e = metrics_savepoint.wgtfacq_e_dsl()
+    wgtfacq_e = metrics_savepoint.wgtfacq_e()
     nflatlev = grid_savepoint.nflatlev()
     c_intp = interpolation_savepoint.c_intp()
     inv_dual_edge_length = grid_savepoint.inv_dual_edge_length()

model/atmosphere/dycore/tests/dycore/utils.py

@@ -51,7 +51,7 @@ def construct_metric_state(
         time_extrapolation_parameter_for_exner=metrics_savepoint.exner_exfac(),
         reference_exner_at_cells_on_model_levels=metrics_savepoint.exner_ref_mc(),
         wgtfac_c=metrics_savepoint.wgtfac_c(),
-        wgtfacq_c=metrics_savepoint.wgtfacq_c_dsl(),
+        wgtfacq_c=metrics_savepoint.wgtfacq_c(),
         inv_ddqz_z_full=metrics_savepoint.inv_ddqz_z_full(),
         reference_rho_at_cells_on_model_levels=metrics_savepoint.rho_ref_mc(),
         reference_theta_at_cells_on_model_levels=metrics_savepoint.theta_ref_mc(),
@@ -71,7 +71,7 @@ def construct_metric_state(
         ddqz_z_full_e=metrics_savepoint.ddqz_z_full_e(),
         ddxt_z_full=metrics_savepoint.ddxt_z_full(),
         wgtfac_e=metrics_savepoint.wgtfac_e(),
-        wgtfacq_e=metrics_savepoint.wgtfacq_e_dsl(),
+        wgtfacq_e=metrics_savepoint.wgtfacq_e(),
         exner_w_implicit_weight_parameter=metrics_savepoint.vwind_impl_wgt(),
         horizontal_mask_for_3d_divdamp=metrics_savepoint.hmask_dd3d(),
         scaling_factor_for_3d_divdamp=metrics_savepoint.scalfac_dd3d(),

model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py

@@ -16,6 +16,7 @@
 
 from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, run_full_muphys
 from icon4py.model.common import dimension as dims, model_backends
+from icon4py.model.testing import test_utils
 from icon4py.model.testing.fixtures.datatest import backend_like
 
 from . import utils
@@ -114,10 +115,10 @@ def test_full_muphys(
     rtol = 1e-14
     atol = 1e-16
 
-    np.testing.assert_allclose(ref.qv.asnumpy(), out.qv.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qc.asnumpy(), out.qc.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qi.asnumpy(), out.qi.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qr.asnumpy(), out.qr.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qs.asnumpy(), out.qs.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qg.asnumpy(), out.qg.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.t.asnumpy(), out.t.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qv.asnumpy(), out.qv.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qc.asnumpy(), out.qc.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qi.asnumpy(), out.qi.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qr.asnumpy(), out.qr.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qs.asnumpy(), out.qs.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qg.asnumpy(), out.qg.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.t.asnumpy(), out.t.asnumpy(), atol=atol, rtol=rtol)

model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py

@@ -16,6 +16,7 @@
 
 from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, run_graupel_only
 from icon4py.model.common import dimension as dims, model_backends
+from icon4py.model.testing import test_utils
 from icon4py.model.testing.fixtures.datatest import backend_like
 
 from . import utils
@@ -106,10 +107,10 @@ def test_graupel_only(
     rtol = 1e-14
     atol = 1e-16
 
-    np.testing.assert_allclose(ref.qv.asnumpy(), out.qv.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qc.asnumpy(), out.qc.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qi.asnumpy(), out.qi.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qr.asnumpy(), out.qr.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qs.asnumpy(), out.qs.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.qg.asnumpy(), out.qg.asnumpy(), atol=atol, rtol=rtol)
-    np.testing.assert_allclose(ref.t.asnumpy(), out.t.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qv.asnumpy(), out.qv.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qc.asnumpy(), out.qc.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qi.asnumpy(), out.qi.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qr.asnumpy(), out.qr.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qs.asnumpy(), out.qs.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.qg.asnumpy(), out.qg.asnumpy(), atol=atol, rtol=rtol)
+    test_utils.assert_dallclose(ref.t.asnumpy(), out.t.asnumpy(), atol=atol, rtol=rtol)

model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py → model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp.py

@@ -15,7 +15,7 @@
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
-def compute_zdiff_gradp_dsl(  # noqa: PLR0912 [too-many-branches]
+def compute_zdiff_gradp(  # noqa: PLR0912 [too-many-branches]
     e2c,
     z_mc: data_alloc.NDArray,
     c_lin_e: data_alloc.NDArray,
@@ -129,6 +129,9 @@ def compute_zdiff_gradp_dsl(  # noqa: PLR0912 [too-many-branches]
                         break
 
     vertoffset_gradp = vertidx_gradp - vertoffset_gradp
-    vertoffset_gradp[:horizontal_start_1, :, :] = 0.0
+
+    # TODO(havogt): NumpyDataProvider needs to be extended to support implict exchange.
+    exchange(zdiff_gradp[:, 0, :])
+    exchange(zdiff_gradp[:, 1, :])
 
     return zdiff_gradp, vertoffset_gradp

model/common/src/icon4py/model/common/metrics/metrics_factory.py

@@ -36,7 +36,7 @@
     compute_advection_metrics,
     compute_coeff_gradekin,
     compute_diffusion_metrics,
-    compute_zdiff_gradp_dsl,
+    compute_zdiff_gradp,
     metric_fields as mf,
     metrics_attributes as attrs,
     reference_atmosphere,
@@ -762,9 +762,9 @@ def _register_computed_fields(self) -> None:  # noqa: PLR0915 [too-many-statemen
         )
         self.register_provider(compute_horizontal_mask_for_3d_divdamp)
 
-        compute_zdiff_gradp_dsl_np = factory.NumpyDataProvider(
+        compute_zdiff_gradp_np = factory.NumpyDataProvider(
             func=functools.partial(
-                compute_zdiff_gradp_dsl.compute_zdiff_gradp_dsl,
+                compute_zdiff_gradp.compute_zdiff_gradp,
                 array_ns=self._xp,
                 exchange=functools.partial(
                     self._exchange.exchange, dims.EdgeDim, stream=decomposition.BLOCK
@@ -793,7 +793,7 @@ def _register_computed_fields(self) -> None:  # noqa: PLR0915 [too-many-statemen
                 ),
             },
         )
-        self.register_provider(compute_zdiff_gradp_dsl_np)
+        self.register_provider(compute_zdiff_gradp_np)
 
         coeff_gradekin = factory.NumpyDataProvider(
             func=functools.partial(

model/common/src/icon4py/model/common/utils/data_allocation.py

@@ -223,14 +223,6 @@ def list2field(
     return gtx.as_field(domain, arr, allocator=allocator)
 
 
-def kflip_wgtfacq(
-    arr: NDArray,
-    domain: gtx.Domain,
-    allocator: gtx_typing.Allocator,
-) -> gtx.Field:
-    return gtx.as_field(domain, arr[:, ::-1], allocator=allocator)  # type: ignore [arg-type] # type "ndarray[Any, dtype[Any] | Any"; expected "NDArrayObject"
-
-
 def adjust_fortran_indices(inp: NDArray) -> NDArray:
     """For some Fortran arrays we need to subtract 1 to be compatible with Python indexing."""
     return inp - 1

model/common/src/icon4py/model/common/utils/field_utils.py

@@ -0,0 +1,74 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import array_api_compat
+from gt4py import next as gtx
+from gt4py.next import typing as gtx_typing
+
+
+def flip(field: gtx.Field, dim: gtx.Dimension, allocator: gtx_typing.Allocator) -> gtx.Field:
+    """Flip a field along a given dimension.
+
+    Args:
+        field: The field to flip.
+        dim: The dimension along which to flip the field.
+        allocator: Allocator to use for the output field.
+    """
+    # Note: `allocator` needs to be passed explicitly since GT4Py fields currently don't persist how they were allocated.
+    xp = array_api_compat.array_namespace(field.ndarray)
+    flipped_array = xp.flip(field.ndarray, axis=field.domain.dims.index(dim))
+    return gtx.as_field(field.domain, flipped_array, allocator=allocator)
+
+
+def index2offset(
+    index_field: gtx.Field, dim: gtx.Dimension, allocator: gtx_typing.Allocator
+) -> gtx.Field:
+    """Convert an index field to an offset field.
+
+    Note: Additionally clips negative indices to become zero offset as Fortran initializes some indices with `0` (which corresponds to `-1` in Python) to indicate that they are not used.
+    As GT4Py's unstructured domain inference is incomplete and runs over the whole domain we might use out-of-bounds offsets in intermediate computations.
+
+    Args:
+        index_field: Index field in Python indexing (0-based).
+        dim: The dimension along which to convert indices to offsets.
+        allocator: Allocator to use for the output field.
+
+    Example:
+        >>> import numpy as np
+        >>> from gt4py import next as gtx
+        >>> from icon4py.model.common import dimension as dims
+        >>> index_field = gtx.as_field(
+        ...     {dims.CellDim: range(2, 6)}, np.array([5, 3, 4, 2], dtype=gtx.int32)
+        ... )
+        >>> result = index2offset(index_field, dims.CellDim, allocator=np)
+        >>> result.domain
+        Domain(dims=(Dimension(value='Cell', kind=<DimensionKind.HORIZONTAL: 'horizontal'>),), ranges=(UnitRange(2, 6),))
+        >>> result.ndarray
+        array([ 3,  0,  0, -3], dtype=int32)
+    """
+    # Note: `allocator` needs to be passed explicitly since GT4Py fields currently don't persist how they were allocated.
+    xp = array_api_compat.array_namespace(index_field.ndarray)
+
+    current_index = gtx.as_field(
+        gtx.Domain(index_field.domain[dim]),
+        xp.arange(
+            index_field.domain[dim].unit_range.start,
+            index_field.domain[dim].unit_range.stop,
+            dtype=index_field.ndarray.dtype,
+        ),
+        allocator=allocator,
+    )
+    # use GT4Py's broadcasting and field arithmetic (includes clipping)
+    offset_field = gtx.where(  # type: ignore[attr-defined]
+        index_field >= 0,  # type: ignore[operator]
+        index_field - current_index,
+        0,
+    )
+
+    # if GT4Py embedded would propagate the allocator, we could avoid this extra conversion.
+    return gtx.as_field(offset_field.domain, offset_field.ndarray, allocator=allocator)

model/common/tests/common/decomposition/unit_tests/test_halo.py

@@ -130,23 +130,23 @@ def test_no_halo():
     decomposition = decomp.SingleNodeDecomposer()
     decomposition_info = halo_generator(decomposition(np.arange(grid_size.num_cells), 1))
     # cells
-    np.testing.assert_allclose(
+    test_utils.assert_dallclose(
         np.arange(grid_size.num_cells), decomposition_info.global_index(dims.CellDim)
     )
     assert np.all(decomposition_info.owner_mask(dims.CellDim))
     assert np.all(
         decomposition_info.halo_levels(dims.CellDim) == definitions.DecompositionFlag.OWNED
     )
     # edges
-    np.testing.assert_allclose(
+    test_utils.assert_dallclose(
         np.arange(grid_size.num_edges), decomposition_info.global_index(dims.EdgeDim)
     )
     assert np.all(
         decomposition_info.halo_levels(dims.EdgeDim) == definitions.DecompositionFlag.OWNED
     )
     assert np.all(decomposition_info.owner_mask(dims.EdgeDim))
     # vertices
-    np.testing.assert_allclose(
+    test_utils.assert_dallclose(
         np.arange(grid_size.num_vertices), decomposition_info.global_index(dims.VertexDim)
     )
     assert np.all(

model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py

@@ -123,7 +123,7 @@ def check_local_global_field(
     check_halos: bool,
 ) -> None:
     if dim == dims.KDim:
-        np.testing.assert_allclose(global_reference_field, local_field)
+        test_utils.assert_dallclose(global_reference_field, local_field)
         return
 
     _log.info(
@@ -138,7 +138,7 @@ def check_local_global_field(
 
     # Compare halo against global reference field
     if check_halos:
-        np.testing.assert_allclose(
+        test_utils.assert_dallclose(
             global_reference_field[
                 decomposition_info.global_index(dim, decomp_defs.DecompositionInfo.EntryType.HALO)
             ],
@@ -180,7 +180,7 @@ def check_local_global_field(
             f" rank = {processor_props.rank}: SHAPES: global reference field {global_reference_field.shape}, gathered = {gathered_field.shape}"
         )
 
-        np.testing.assert_allclose(sorted_, global_reference_field, atol=1e-9, verbose=True)
+        test_utils.assert_dallclose(sorted_, global_reference_field, atol=1e-9, verbose=True)
 
 
 # These fields can't be computed with the embedded backend for one reason or

model/common/tests/common/grid/unit_tests/test_topography.py

@@ -13,11 +13,9 @@
 
 from icon4py.model.common.grid import topography as topo
 from icon4py.model.common.utils import data_allocation as data_alloc
-from icon4py.model.testing import definitions, test_utils
+from icon4py.model.testing import exchange_utils, test_utils
 from icon4py.model.testing.fixtures import *  # noqa: F403
 
-from ... import utils
-
 
 if TYPE_CHECKING:
     import gt4py.next.typing as gtx_typing
@@ -53,7 +51,7 @@ def test_topography_smoothing_with_serialized_data(
         c2e2co=icon_grid.get_connectivity("C2E2CO").ndarray,
         num_iterations=num_iterations,
         array_ns=xp,
-        exchange=utils.dummy_exchange,
+        exchange=exchange_utils.dummy_exchange_with_bound_dim,
     )
 
     assert test_utils.dallclose(topography_smoothed_ref, topography_smoothed, atol=1.0e-14)