Move transformations (index2offset, transpose) from Fortran to Python by havogt · Pull Request #1119 · C2SM/icon4py

havogt · 2026-03-20T09:13:47Z

Combines #1118 and #1115, because both require v3 serialized data:

Move computation of vertoffset_gradp from vertidx_gradp to Python (index2offset)
Move transpose (and decomposition) of rbf_vec_coeff_v and rbf_vec_coeff_e to Python

Additional:

Add missing halo exchange to compute_zdiff_gradp (test failure triggered by new serialized data, most likely unrelated to the other changes)
Refactor kflip_wgtfacq to a more general flip on fields

Co-authored-by: Hannes Vogt <[email protected]>

…t_gradp_and_dim_swap3

havogt · 2026-03-20T14:14:17Z

cscs-ci run default

havogt · 2026-03-20T14:14:21Z

cscs-ci run distributed

havogt · 2026-03-20T14:24:56Z

cscs-ci run default

havogt · 2026-03-20T15:28:08Z

cscs-ci run default

havogt · 2026-03-20T15:28:17Z

cscs-ci run default

…icon4py into tmp_vertoffset_gradp_and_dim_swap3

havogt · 2026-03-23T06:37:54Z

cscs-ci run default

havogt · 2026-03-23T10:21:34Z

cscs-ci run default

havogt · 2026-03-23T10:23:25Z

cscs-ci run distributed

havogt · 2026-03-23T10:30:46Z

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

@@ -129,6 +129,5 @@ def compute_zdiff_gradp_dsl(  # noqa: PLR0912 [too-many-branches]
                        break

    vertoffset_gradp = vertidx_gradp - vertoffset_gradp
-    vertoffset_gradp[:horizontal_start_1, :, :] = 0.0


I don't see why this line was added...

could this be related to the test that verifies this over the entire domain instead of dallclose(vertoffset_gradp[horizontal_start_1:, :, :] ...) similarly to the problems in test_solve_nonhydro.py with exner_at_cells_on_half_levels and temporal_extrapolation_of_perturbed_exner?

As discussed, I couldn't find this line in Fortran. My best guess is that is made the test work with the previous serialized data, but not sure...

…p_and_dim_swap3

havogt · 2026-03-24T08:25:46Z

cscs-ci run default

havogt · 2026-03-24T09:26:46Z

cscs-ci run default

havogt · 2026-03-24T09:26:52Z

cscs-ci run distributed

Copilot

Pull request overview

This PR updates ICON4Py’s Python/py2fgen wrapper and test infrastructure to consume v3 serialized data, moving a few previously-Fortran-side layout/index transformations into Python.

Changes:

Move vertoffset_gradp derivation to Python by passing vertidx_gradp through wrappers/savepoints and converting via index2offset.
Move rbf_vec_coeff_{e,v} transpose/decomposition logic to Python (new 3D rbf_vec_coeff_v plumbing; swap/transpose done in wrappers/savepoints).
Introduce field_utils (flip, index2offset), remove kflip_wgtfacq, and update/extend unit + integration + MPI tests accordingly (serialized Experiment.version bumped to 3).

Reviewed changes

Copilot reviewed 36 out of 37 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py	Update dycore wrapper tests to use v3 fields (`vertidx_gradp`, new RBF coeff layout, raw wgtfacq read).
tools/tests/tools/py2fgen/wrappers/test_diffusion_wrapper.py	Update diffusion wrapper tests to use 3D `rbf_vec_coeff_v` instead of `v1/v2`.
tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.py	Regenerate reference Python wrapper signature/packing to accept `rbf_vec_coeff_v` and `vertidx_gradp`.
tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.h	Update C header for new dycore wrapper argument list (3D RBF + vertidx).
tools/tests/tools/py2fgen/wrappers/references/dycore/dycore.f90	Update Fortran wrapper interface to match new argument list and shapes.
tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.py	Regenerate reference Python wrapper signature/packing to accept `rbf_vec_coeff_v`.
tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.h	Update C header for new diffusion wrapper argument list (3D RBF).
tools/tests/tools/py2fgen/wrappers/references/diffusion/diffusion.f90	Update Fortran wrapper interface to match new argument list and shapes.
tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py	Implement Python-side RBF decomposition/transposes, wgtfacq flipping, and `vertidx→vertoffset` conversion.
tools/src/icon4py/tools/py2fgen/wrappers/diffusion_wrapper.py	Accept 3D RBF input and decompose into `rbf_coeff_{1,2}` in Python.
tools/src/icon4py/tools/py2fgen/wrappers/common.py	Add `Float64Array3D` type alias for py2fgen wrapper signatures.
model/testing/src/icon4py/model/testing/test_utils.py	Add `assert_dallclose` helper (numpy.testing-based).
model/testing/src/icon4py/model/testing/serialbox.py	Add `_get_field(..., transpose=...)`, flip `wgtfacq_*`, transpose `rbf_vec_coeff_e`, derive `vertoffset_gradp` from `vertidx_gradp`.
model/testing/src/icon4py/model/testing/definitions.py	Bump serialized experiment version from 2 → 3.
model/driver/tests/driver/integration_tests/test_icon4py.py	Switch to `metrics_savepoint.wgtfacq_{c,e}()` (no `_dsl`).
model/driver/src/icon4py/model/driver/initialization_utils.py	Switch to `metrics_savepoint.wgtfacq_{c,e}()` (no `_dsl`).
model/common/tests/common/utils_test.py	Add shared `dummy_exchange` for common tests (renamed import target).
model/common/tests/common/utils/unit_tests/test_field_utils.py	New unit tests for `field_utils.flip` and `field_utils.index2offset`.
model/common/tests/common/utils/unit_tests/init.py	Package init for new common utils unit tests.
model/common/tests/common/utils/init.py	Package init for common utils tests.
model/common/tests/common/metrics/unit_tests/test_metrics_factory.py	Adjust comparisons for uninitialized lateral boundary regions; use new `assert_dallclose` and `wgtfacq_e()`.
model/common/tests/common/metrics/unit_tests/test_metric_fields.py	Update imports to `utils_test.dummy_exchange`.
model/common/tests/common/metrics/unit_tests/test_compute_zdiff_gradp.py	Switch test to new `compute_zdiff_gradp` implementation and updated utils import.
model/common/tests/common/metrics/unit_tests/test_compute_weight_factors.py	Use flipped `wgtfacq_{c,e}()` and update utils import.
model/common/tests/common/metrics/unit_tests/test_compute_coeff_gradekin.py	Update utils import to `utils_test`.
model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py	Update distributed metrics comparisons with optional slicing for known uninitialized regions; use `assert_dallclose`; switch `wgtfacq_e()`.
model/common/tests/common/interpolation/unit_tests/test_rbf_interpolation.py	Update utils import to `utils_test`.
model/common/tests/common/interpolation/unit_tests/test_interpolation_fields.py	Update utils import to `utils_test`.
model/common/tests/common/grid/unit_tests/test_vertical.py	Update utils import to `utils_test`.
model/common/tests/common/grid/unit_tests/test_topography.py	Update utils import to `utils_test`.
model/common/src/icon4py/model/common/utils/field_utils.py	New helpers: `flip` and `index2offset` (array-api compatible).
model/common/src/icon4py/model/common/utils/data_allocation.py	Remove `kflip_wgtfacq` (replaced by `field_utils.flip`).
model/common/src/icon4py/model/common/metrics/metrics_factory.py	Switch provider from `compute_zdiff_gradp_dsl` to `compute_zdiff_gradp`.
model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp.py	Rename to `compute_zdiff_gradp` and add explicit exchanges for outputs (partial).
model/atmosphere/dycore/tests/dycore/utils.py	Switch to flipped `wgtfacq_{c,e}()` accessors.
model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py	Switch to `wgtfacq_e()` accessor.
model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py	Switch to flipped `wgtfacq_{c,e}()`; adjust array comparisons to avoid uninitialized boundary regions.

Comments suppressed due to low confidence (2)

model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp.py:135

compute_zdiff_gradp manually exchanges zdiff_gradp, but it returns vertoffset_gradp as a second output without exchanging it. Since vertoffset_gradp is later used as the ikoffset field in dycore stencils, its halo/boundary values can be inconsistent across ranks. Add an exchange for vertoffset_gradp (or ensure NumpyDataProvider exchanges both outputs).
model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp.py:134
Typo in comment: “implict exchange” → “implicit exchange”.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

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

tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py

Co-authored-by: Copilot <[email protected]>

jcanton

some small suggestions

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

jcanton · 2026-03-24T10:36:47Z

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

@@ -129,6 +129,5 @@ def compute_zdiff_gradp_dsl(  # noqa: PLR0912 [too-many-branches]
                        break

    vertoffset_gradp = vertidx_gradp - vertoffset_gradp
-    vertoffset_gradp[:horizontal_start_1, :, :] = 0.0


could this be related to the test that verifies this over the entire domain instead of dallclose(vertoffset_gradp[horizontal_start_1:, :, :] ...) similarly to the problems in test_solve_nonhydro.py with exner_at_cells_on_half_levels and temporal_extrapolation_of_perturbed_exner?

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

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

model/common/tests/common/utils/unit_tests/test_field_utils.py

model/testing/src/icon4py/model/testing/serialbox.py

jcanton · 2026-03-24T14:07:00Z

model/testing/src/icon4py/model/testing/test_utils.py

+def assert_dallclose(
+    actual: npt.ArrayLike,
+    desired: npt.ArrayLike,
+    rtol: float = 1.0e-12,
+    atol: float = 0.0,
+    equal_nan: bool = False,
+) -> None:
+    np_testing.assert_allclose(actual, desired, rtol=rtol, atol=atol, equal_nan=equal_nan)  # type: ignore[arg-type]
+


no need for the helper, use np.testing.assert_allclose directly as in other tests?

Changed the other tests, the reason for helper is the defaults for double precision.

tools/tests/tools/py2fgen/wrappers/test_dycore_wrapper.py

jcanton · 2026-03-24T15:57:31Z

ah, remember to close the other two PRs that were merged into this one

havogt · 2026-03-24T17:22:23Z

cscs-ci run default

havogt · 2026-03-24T17:22:32Z

cscs-ci run distributed

…icon4py into tmp_vertoffset_gradp_and_dim_swap3

github-actions · 2026-03-24T18:43:05Z

Mandatory Tests

Please make sure you run these tests via comment before you merge!

cscs-ci run default
cscs-ci run distributed

Optional Tests

To run benchmarks you can use:

cscs-ci run benchmark-bencher

To run tests and benchmarks with the DaCe backend you can use:

cscs-ci run dace

To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:

cscs-ci run extra

For more detailed information please look at CI in the EXCLAIM universe.

havogt · 2026-03-24T18:44:01Z

cscs-ci run default

havogt · 2026-03-24T18:44:55Z

cscs-ci run distributed

jcanton

looks very nice now, and you are lucky it's not blocked by the distributed CI so go ahead and merge!

havogt · 2026-03-26T13:36:33Z

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

    vertoffset_gradp = vertidx_gradp - vertoffset_gradp
-    vertoffset_gradp[:horizontal_start_1, :, :] = 0.0
+
+    # TODO(havogt): NumpyDataProvider needs to be extended to support implict exchange.


Suggested change

# TODO(havogt): NumpyDataProvider needs to be extended to support implict exchange.

# TODO(havogt): NumpyDataProvider needs to be extended to support implicit exchange.

:(

havogt and others added 14 commits March 18, 2026 11:08

bindings: move swap from Fortran to Python

dc52a12

Apply suggestions from code review

073f425

Co-authored-by: Hannes Vogt <[email protected]>

update references

37c21b2

Move vertoffset_gradp to Python

66c1bf1

cleanup flip

6f91160

fix references

47e9c7a

add tests and fix bugs

b085cb3

fixes

6bfed5d

fix test

491344c

Merge branch 'main' into blueline_swap_rbf_vec_coeff_e_and_v

1a54450

fix test

bc8a73f

fix tests

cd4b75c

Merge branch 'blueline_swap_rbf_vec_coeff_e_and_v' into tmp_vertoffse…

fa8def5

…t_gradp_and_dim_swap3

switch to serialized data v3

2a23643

Update base.yml

742336f

resolve module name conflict

64d3e59

Merge branch 'main' into tmp_vertoffset_gradp_and_dim_swap3

96a90bc

havogt added 2 commits March 20, 2026 22:18

limit test

9d865bd

Merge branch 'tmp_vertoffset_gradp_and_dim_swap3' of github.com:C2SM/…

27106a3

…icon4py into tmp_vertoffset_gradp_and_dim_swap3

havogt commented Mar 23, 2026

View reviewed changes

havogt added 2 commits March 23, 2026 15:50

Merge remote-tracking branch 'upstream/main' into tmp_vertoffset_grad…

c6a8ac9

…p_and_dim_swap3

switch on/off download

f6ba342

Merge branch 'main' into tmp_vertoffset_gradp_and_dim_swap3

b34fe50

havogt added 2 commits March 24, 2026 10:23

cleanup exchange in zdiff_gradp

5382941

renames

3d81bbd

havogt marked this pull request as ready for review March 24, 2026 10:07

havogt requested review from Copilot and jcanton March 24, 2026 10:07

Copilot started reviewing on behalf of havogt March 24, 2026 10:08 View session

Copilot AI reviewed Mar 24, 2026

View reviewed changes

model/common/src/icon4py/model/common/utils/field_utils.py Outdated Show resolved Hide resolved

tools/src/icon4py/tools/py2fgen/wrappers/dycore_wrapper.py Outdated Show resolved Hide resolved

havogt changed the title ~~Testing new serialized data for vertoffset gradp and dim swap3~~ Move transformations (index2offset, transpose) from Fortran to Python Mar 24, 2026

Apply suggestions from code review

404bf7d

Co-authored-by: Copilot <[email protected]>

jcanton requested changes Mar 24, 2026

View reviewed changes

havogt added 3 commits March 24, 2026 17:57

address review comments

de1c4aa

slightly improve dummy_exchange

378234f

add doctest to index2offset

dff89b7

Merge branch 'tmp_vertoffset_gradp_and_dim_swap3' of github.com:C2SM/…

0057ef2

…icon4py into tmp_vertoffset_gradp_and_dim_swap3

havogt requested a review from jcanton March 24, 2026 19:44

jcanton approved these changes Mar 25, 2026

View reviewed changes

havogt merged commit 76caa74 into main Mar 25, 2026
54 checks passed

This was referenced Mar 25, 2026

Move vertoffset_gradp to Python #1118

Closed

bindings: move swap from Fortran to Python #1115

Closed

havogt commented Mar 26, 2026

View reviewed changes

	# TODO(havogt): NumpyDataProvider needs to be extended to support implict exchange.
	# TODO(havogt): NumpyDataProvider needs to be extended to support implicit exchange.

Conversation

havogt commented Mar 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

havogt commented Mar 20, 2026

Uh oh!

havogt commented Mar 20, 2026

Uh oh!

havogt commented Mar 20, 2026

Uh oh!

havogt commented Mar 20, 2026

Uh oh!

havogt commented Mar 20, 2026

Uh oh!

havogt commented Mar 23, 2026

Uh oh!

havogt commented Mar 23, 2026

Uh oh!

havogt commented Mar 23, 2026

Uh oh!

havogt Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

jcanton Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

jcanton left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jcanton Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jcanton Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Mar 24, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jcanton commented Mar 24, 2026

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

github-actions bot commented Mar 24, 2026

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

havogt commented Mar 24, 2026

Uh oh!

jcanton left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

havogt Mar 26, 2026

Choose a reason for hiding this comment

havogt commented Mar 20, 2026 •

edited

Loading