From f7ff27d8770734983e3a0b721590d947ffe264fe Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 10:27:49 +0100 Subject: [PATCH 01/25] place exchange in lsq_pseudoinv coefficients --- .../advection/advection_lsq_coeffs.py | 22 ++++++---- .../interpolation/interpolation_attributes.py | 9 ++++ .../interpolation/interpolation_factory.py | 42 +++++++++++++++++++ .../mpi_tests/test_parallel_interpolation.py | 24 +++++++++++ 4 files changed, 89 insertions(+), 8 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 01c0cc80bf..249e70d994 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -5,6 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from collections.abc import Callable +from types import ModuleType import numpy as np import scipy @@ -86,16 +88,18 @@ def lsq_compute_coeffs( lsq_dim_stencil: int, start_idx: int, min_rlcell_int: int, - geometry_type: base_grid.GeometryType, + geometry_type: int, + exchange: Callable[[data_alloc.NDArray], None], + array_ns: ModuleType = np, ) -> data_alloc.NDArray: - lsq_weights_c = np.zeros((min_rlcell_int, lsq_dim_stencil)) - lsq_pseudoinv = np.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) - z_lsq_mat_c = np.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) + lsq_weights_c = array_ns.zeros((min_rlcell_int, lsq_dim_stencil)) + lsq_pseudoinv = array_ns.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) + z_lsq_mat_c = array_ns.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) for jc in range(start_idx, min_rlcell_int): - match geometry_type: + match base_grid.GeometryType(geometry_type): case base_grid.GeometryType.ICOSAHEDRON: - z_dist_g = np.zeros((lsq_dim_c, 2)) + z_dist_g = array_ns.zeros((lsq_dim_c, 2)) for js in range(lsq_dim_stencil): z_dist_g[js, :] = gnomonic_proj_single_val( cell_lon[jc], cell_lat[jc], cell_lon[c2e2c[jc, js]], cell_lat[c2e2c[jc, js]] @@ -107,7 +111,7 @@ def lsq_compute_coeffs( z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 case base_grid.GeometryType.TORUS: ilc_s = c2e2c[jc, :lsq_dim_stencil] - cc_cell = np.zeros((lsq_dim_stencil, 2)) + cc_cell = array_ns.zeros((lsq_dim_stencil, 2)) cc_cv = (cell_center_x[jc], cell_center_y[jc]) for js in range(lsq_dim_stencil): @@ -127,7 +131,7 @@ def lsq_compute_coeffs( z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c ) - + exchange(lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, lsq_pseudoinv, @@ -138,5 +142,7 @@ def lsq_compute_coeffs( lsq_dim_unk, lsq_dim_c, ) + exchange(lsq_pseudoinv[:, 0, :]) + exchange(lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py index d50765ea3f..5ec0f04224 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py @@ -34,6 +34,7 @@ RBF_SCALE_CELL: Final[str] = "rbf_scale_cell" RBF_SCALE_EDGE: Final[str] = "rbf_scale_edge" RBF_SCALE_VERTEX: Final[str] = "rbf_scale_vertex" +LSQ_PSEUDOINV: Final[str] = "lsq_interpolation_coefficient" attrs: dict[str, model.FieldMetaData] = { C_LIN_E: dict( @@ -212,4 +213,12 @@ icon_var_name="rbf_vec_scale_v", dtype=ta.wpfloat, ), + LSQ_PSEUDOINV: dict( + standard_name=LSQ_PSEUDOINV, + long_name="lsq_pseudoinv", + units="", + dims=(dims.CellDim, dims.C2E2CDim), + icon_var_name="ptr_int_lsq%lsq_pseudoinv", + dtype=ta.wpfloat, + ), } diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index a22b909fd8..65972ac2aa 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -12,6 +12,7 @@ import gt4py.next.typing as gtx_typing import icon4py.model.common.interpolation.stencils.compute_nudgecoeffs as nudgecoeffs +from icon4py.model.atmosphere.advection import advection_lsq_coeffs from icon4py.model.common import constants, dimension as dims from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import ( @@ -67,6 +68,10 @@ def __init__( "rbf_kernel_cell": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.CELL], "rbf_kernel_edge": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.EDGE], "rbf_kernel_vertex": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.VERTEX], + "lsq_dim_unk": 2, + "lsq_dim_c": 3, + "lsq_wgt_exp": 2, + "lsq_dim_stencil": 3, } log.info( f"initialized interpolation factory for backend = '{self._backend_name()}' and grid = '{self._grid}'" @@ -228,6 +233,43 @@ def _register_computed_fields(self) -> None: ) self.register_provider(rbf_scale_vertex_np) + lsq_pseudoinv = factory.NumpyDataProvider( + func=functools.partial( + advection_lsq_coeffs.lsq_compute_coeffs, + exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + array_ns=self._xp, + ), + fields=(attrs.LSQ_PSEUDOINV,), + domain=(), + deps={ + "cell_center_x": geometry_attrs.CELL_CENTER_X, + "cell_center_y": geometry_attrs.CELL_CENTER_Y, + "cell_lat": geometry_attrs.CELL_LAT, + "cell_lon": geometry_attrs.CELL_LON, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c": dims.C2E2CDim}, + params={ + "domain_length": self.grid.global_properties.domain_length + if self.grid.global_properties.domain_length is not None + else 0.0, + "domain_height": self.grid.global_properties.domain_height + if self.grid.global_properties.domain_height is not None + else 0.0, + "grid_sphere_radius": constants.EARTH_RADIUS, + "lsq_dim_unk": self._config["lsq_dim_unk"], + "lsq_dim_c": self._config["lsq_dim_c"], + "lsq_wgt_exp": self._config["lsq_wgt_exp"], + "lsq_dim_stencil": self._config["lsq_dim_stencil"], + "start_idx": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "min_rlcell_int": self.grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)), + "geometry_type": self.grid.global_properties.geometry_type.value, + }, + ) + self.register_provider(lsq_pseudoinv) + match self.grid.global_properties.geometry_type: case base.GeometryType.ICOSAHEDRON: cell_average_weight = factory.NumpyDataProvider( diff --git a/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py b/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py index 712a69f490..f2c9afd3f7 100644 --- a/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py +++ b/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py @@ -214,3 +214,27 @@ def test_distributed_interpolation_rbf( field_ref = interpolation_savepoint.__getattribute__(intrp_name)().asnumpy() field = factory.get(attrs_name).asnumpy() test_utils.dallclose(field, field_ref, atol=atol) + + +@pytest.mark.datatest +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_distributed_interpolation_lsq_pseudoinv( + backend: gtx_typing.Backend, + interpolation_savepoint: sb.InterpolationSavepoint, + grid_savepoint: sb.IconGridSavepoint, + experiment: test_defs.Experiment, + processor_props: decomposition.ProcessProperties, + decomposition_info: decomposition.DecompositionInfo, + interpolation_factory_from_savepoint: interpolation_factory.InterpolationFieldsFactory, +) -> None: + parallel_helpers.check_comm_size(processor_props) + parallel_helpers.log_process_properties(processor_props) + parallel_helpers.log_local_field_size(decomposition_info) + factory = interpolation_factory_from_savepoint + field_ref_1 = interpolation_savepoint.__getattribute__("lsq_pseudoinv_1")().asnumpy() + field_ref_2 = interpolation_savepoint.__getattribute__("lsq_pseudoinv_2")().asnumpy() + field_1 = factory.get(attrs.LSQ_PSEUDOINV)[:, 0, :] + field_2 = factory.get(attrs.LSQ_PSEUDOINV)[:, 1, :] + test_utils.dallclose(field_1, field_ref_1, atol=1e-15) # type: ignore[arg-type] # mypy does not recognize sliced array as still an array + test_utils.dallclose(field_2, field_ref_2, atol=1e-15) # type: ignore[arg-type] # mypy does not recognize sliced array as still an array From fa62906162ee3d159305760ab2923f76514f9ab8 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 13:56:34 +0100 Subject: [PATCH 02/25] mpi and exchange edits --- .../model/atmosphere/advection/advection.py | 4 +- .../advection/advection_lsq_coeffs.py | 8 +- .../advection/tests/advection/fixtures.py | 2 + .../integration_tests/test_advection.py | 4 +- .../tests/advection/mpi_tests/__init__.py | 7 + .../mpi_tests/test_parallel_advection.py | 291 ++++++++++++++++++ .../interpolation/interpolation_factory.py | 12 +- 7 files changed, 317 insertions(+), 11 deletions(-) create mode 100644 model/atmosphere/advection/tests/advection/mpi_tests/__init__.py create mode 100644 model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index 23318563de..2894b5aa88 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -281,10 +281,12 @@ def run( if self._even_timestep else (diagnostic_state.airmass_new, self._start_cell_lateral_boundary_level_3) ) + log.debug("running stencil apply_density_increment - start") + self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) self._apply_density_increment( rhodz_in=rhodz_in, - p_mflx_contra_v=prep_adv.mass_flx_ic, + p_mflx_contra_v=prep_adv.mass_flx_ic, # exchange deepatmo_divzl=self._metric_state.deepatmo_divzl, deepatmo_divzu=self._metric_state.deepatmo_divzu, rhodz_out=self._rhodz_ast2, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 249e70d994..c09a5a5534 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -11,6 +11,7 @@ import numpy as np import scipy +from icon4py.model.common.dimension import CellDim from icon4py.model.common.grid import base as base_grid from icon4py.model.common.math.projection import ( gnomonic_proj_single_val, @@ -131,7 +132,8 @@ def lsq_compute_coeffs( z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c ) - exchange(lsq_weights_c) + + exchange(lsq_weights_c, dim=CellDim) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, lsq_pseudoinv, @@ -142,7 +144,7 @@ def lsq_compute_coeffs( lsq_dim_unk, lsq_dim_c, ) - exchange(lsq_pseudoinv[:, 0, :]) - exchange(lsq_pseudoinv[:, 1, :]) + exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) return lsq_pseudoinv diff --git a/model/atmosphere/advection/tests/advection/fixtures.py b/model/atmosphere/advection/tests/advection/fixtures.py index aabae893bb..8aef7d508e 100644 --- a/model/atmosphere/advection/tests/advection/fixtures.py +++ b/model/atmosphere/advection/tests/advection/fixtures.py @@ -44,3 +44,5 @@ def advection_exit_savepoint(data_provider, date): fixture, passing 'date='. """ return data_provider.from_advection_exit_savepoint(size=data_provider.grid_size, date=date) + +from icon4py.model.testing.fixtures.datatest import decomposition_info diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index ee43d3ae3d..9f051f3d0d 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -12,7 +12,7 @@ import icon4py.model.testing.test_utils as test_helpers from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.atmosphere.advection.advection_lsq_coeffs import lsq_compute_coeffs -from icon4py.model.common import constants, dimension as dims +from icon4py.model.common import constants, dimension as dims, utils from icon4py.model.common.grid import ( base as base_grid, geometry_attributes as geometry_attrs, @@ -37,6 +37,7 @@ metrics_savepoint, processor_props, ) +from model.common.tests.common.utils import dummy_exchange from ..fixtures import advection_exit_savepoint, advection_init_savepoint from ..utils import ( @@ -179,6 +180,7 @@ def test_advection_run_single_step( cell_params=cell_geometry, even_timestep=even_timestep, backend=backend, + exchange=dummy_exchange ) diagnostic_state = construct_diagnostic_init_state( diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/__init__.py b/model/atmosphere/advection/tests/advection/mpi_tests/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/atmosphere/advection/tests/advection/mpi_tests/__init__.py @@ -0,0 +1,7 @@ +# 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 diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py new file mode 100644 index 0000000000..4d88de67e9 --- /dev/null +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -0,0 +1,291 @@ +# 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 gt4py.next.typing as gtx_typing +import pytest + +from icon4py.model.testing import test_utils +from icon4py.model.atmosphere.advection import advection, advection_states +from icon4py.model.atmosphere.advection.advection_lsq_coeffs import lsq_compute_coeffs +from icon4py.model.common import constants, dimension as dims, utils +from icon4py.model.common.grid import ( + base as base_grid, + geometry_attributes as geometry_attrs, + horizontal as h_grid, +) +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import ( + definitions as test_defs, + grid_utils, + grid_utils as gridtest_utils, + serialbox as sb, +) +from icon4py.model.testing.fixtures.datatest import ( + backend, + backend_like, + data_provider, + download_ser_data, + experiment, + grid_savepoint, + icon_grid, + interpolation_savepoint, + metrics_savepoint, + processor_props, +) +from icon4py.model.common.decomposition import definitions +from ..fixtures import advection_exit_savepoint, advection_init_savepoint +from ..utils import ( + construct_config, + construct_diagnostic_exit_state, + construct_diagnostic_init_state, + construct_interpolation_state, + construct_least_squares_state, + construct_metric_state, + construct_prep_adv, + log_serialized, + verify_advection_fields, +) + +import numpy as np +import pytest +from gt4py.next import typing as gtx_typing + +from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as nh +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.decomposition import definitions, mpi_decomposition +from icon4py.model.common.grid import icon, states as grid_states, vertical as v_grid +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions as test_defs, parallel_helpers, serialbox, test_utils + +from .. import utils +from ..fixtures import * # noqa: F403 + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.datatest +@pytest.mark.parametrize("experiment", [test_defs.Experiments.MCH_CH_R04B09]) +@pytest.mark.parametrize( + "date, even_timestep, ntracer, horizontal_advection_type, horizontal_advection_limiter, vertical_advection_type, vertical_advection_limiter", + [ + ( + "2021-06-20T12:00:10.000", + False, + 1, + advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, + advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, + advection.VerticalAdvectionType.NO_ADVECTION, + advection.VerticalAdvectionLimiter.NO_LIMITER, + ), + ( + "2021-06-20T12:00:20.000", + True, + 1, + advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, + advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, + advection.VerticalAdvectionType.NO_ADVECTION, + advection.VerticalAdvectionLimiter.NO_LIMITER, + ), + ( + "2021-06-20T12:00:10.000", + False, + 4, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ( + "2021-06-20T12:00:20.000", + True, + 4, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ], +) +@pytest.mark.mpi +def test_advection_run_single_step( + date, + even_timestep, + ntracer, + horizontal_advection_type, + horizontal_advection_limiter, + vertical_advection_type, + vertical_advection_limiter, + *, + grid_savepoint, + icon_grid, + interpolation_savepoint, + metrics_savepoint, + # data_provider, + backend, + advection_init_savepoint, + advection_exit_savepoint, + experiment: test_defs.Experiment, + processor_props: definitions.ProcessProperties, + decomposition_info: definitions.DecompositionInfo, # : F811 fixture +): + if test_utils.is_embedded(backend): + # https://github.com/GridTools/gt4py/issues/1583 + pytest.xfail("ValueError: axes don't match array") + # TODO(OngChia): the last datatest fails on GPU (or even CPU) backend when there is no advection because the horizontal flux is not zero. Further check required. + if ( + even_timestep + and horizontal_advection_type == advection.HorizontalAdvectionType.NO_ADVECTION + ): + pytest.xfail( + "This test is skipped until the cause of nonzero horizontal advection if revealed." + ) + config = construct_config( + horizontal_advection_type=horizontal_advection_type, + horizontal_advection_limiter=horizontal_advection_limiter, + vertical_advection_type=vertical_advection_type, + vertical_advection_limiter=vertical_advection_limiter, + ) + exchange = definitions.create_exchange(processor_props, decomposition_info) + interpolation_state = construct_interpolation_state(interpolation_savepoint, backend=backend) + geometry = gridtest_utils.get_grid_geometry(backend, experiment) + least_squares_coeffs = lsq_compute_coeffs( + cell_center_x=geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy(), + cell_center_y=geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy(), + cell_lat=geometry.get(geometry_attrs.CELL_LAT).asnumpy(), + cell_lon=geometry.get(geometry_attrs.CELL_LON).asnumpy(), + c2e2c=icon_grid.connectivities["C2E2C"].asnumpy(), + cell_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), + domain_length=geometry.grid.global_properties.domain_length, + domain_height=geometry.grid.global_properties.domain_height, + grid_sphere_radius=constants.EARTH_RADIUS, + lsq_dim_unk=2, + lsq_dim_c=3, + lsq_wgt_exp=2, + lsq_dim_stencil=3, + start_idx=icon_grid.start_index( + h_grid.domain(dims.CellDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + min_rlcell_int=icon_grid.end_index(h_grid.domain(dims.CellDim)(h_grid.Zone.HALO_LEVEL_2)), + geometry_type=icon_grid.geometry_type, + exchange=exchange, + ) + + least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend) + + metric_state = construct_metric_state(icon_grid, metrics_savepoint, backend=backend) + edge_geometry = grid_savepoint.construct_edge_geometry() + cell_geometry = grid_savepoint.construct_cell_geometry() + exchange = definitions.create_exchange(processor_props, decomposition_info) + + advection_granule = advection.convert_config_to_advection( + config=config, + grid=icon_grid, + interpolation_state=interpolation_state, + least_squares_state=least_squares_state, + metric_state=metric_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + even_timestep=even_timestep, + backend=backend, + exchange=exchange + ) + + diagnostic_state = construct_diagnostic_init_state( + icon_grid, advection_init_savepoint, ntracer, backend=backend + ) + prep_adv = construct_prep_adv(advection_init_savepoint) + p_tracer_now = advection_init_savepoint.tracer(ntracer) + p_tracer_new = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) + dtime = advection_init_savepoint.get_metadata("dtime").get("dtime") + + log_serialized(diagnostic_state, prep_adv, p_tracer_now, dtime) + + advection_granule.run( + diagnostic_state=diagnostic_state, + prep_adv=prep_adv, + p_tracer_now=p_tracer_now, + p_tracer_new=p_tracer_new, + dtime=dtime, + ) + + diagnostic_state_ref = construct_diagnostic_exit_state( + icon_grid, advection_exit_savepoint, ntracer, backend=backend + ) + p_tracer_new_ref = advection_exit_savepoint.tracer(ntracer) + + verify_advection_fields( + grid=icon_grid, + diagnostic_state=diagnostic_state, + diagnostic_state_ref=diagnostic_state_ref, + p_tracer_new=p_tracer_new, + p_tracer_new_ref=p_tracer_new_ref, + even_timestep=even_timestep, + ) + + +# @pytest.mark.level("unit") +# @pytest.mark.datatest +# def test_lsq_compute_coeffs( +# icon_grid: base_grid.Grid, +# grid_savepoint: sb.IconGridSavepoint, +# backend: gtx_typing.Backend, +# interpolation_savepoint: sb.InterpolationSavepoint, +# experiment: test_defs.Experiment, +# ) -> None: +# gm = grid_utils.get_grid_manager_from_identifier( +# experiment.grid, +# num_levels=1, +# keep_skip_values=True, +# allocator=backend, +# ) +# +# c2e2c = gm.grid.connectivities["C2E2C"].asnumpy() +# cell_owner_mask = grid_savepoint.c_owner_mask().asnumpy() +# grid_sphere_radius = constants.EARTH_RADIUS +# lsq_dim_unk = 2 +# lsq_dim_c = 3 +# lsq_wgt_exp = 2 +# cell_domain = h_grid.domain(dims.CellDim) +# +# min_rlcell_int = gm.grid.end_index(cell_domain(h_grid.Zone.LOCAL)) +# start_idx = gm.grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) +# +# grid_geometry = grid_utils.get_grid_geometry(backend, experiment) +# cell_center_x = grid_geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy() +# cell_center_y = grid_geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy() +# domain_length = gm.grid.global_properties.domain_length +# domain_height = gm.grid.global_properties.domain_height +# lsq_dim_stencil = 3 +# +# coordinates = gm.coordinates +# cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() +# cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() +# lsq_pseudoinv = lsq_compute_coeffs( +# cell_center_x, +# cell_center_y, +# cell_lat, +# cell_lon, +# c2e2c, +# cell_owner_mask, +# domain_length, +# domain_height, +# grid_sphere_radius, +# lsq_dim_unk, +# lsq_dim_c, +# lsq_wgt_exp, +# lsq_dim_stencil, +# start_idx, +# min_rlcell_int, +# exchange, +# icon_grid.geometry_type, +# ) +# +# assert test_helpers.dallclose( +# interpolation_savepoint.lsq_pseudoinv_1().asnumpy(), lsq_pseudoinv[:, 0, :], atol=1e-15 +# ) +# assert test_helpers.dallclose( +# interpolation_savepoint.lsq_pseudoinv_2().asnumpy(), lsq_pseudoinv[:, 1, :], atol=1e-15 +# ) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index 65972ac2aa..2e984bc3d1 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -58,6 +58,8 @@ def __init__( self._providers: dict[str, factory.FieldProvider] = {} self._geometry = geometry_source self._exchange = exchange + domain_length = self.grid.global_properties.domain_length + domain_height = self.grid.global_properties.domain_height # TODO @halungge: Dummy config dict - to be replaced by real configuration self._config = { "divergence_averaging_central_cell_weight": 0.5, # divavg_cntrwgt in ICON @@ -72,6 +74,8 @@ def __init__( "lsq_dim_c": 3, "lsq_wgt_exp": 2, "lsq_dim_stencil": 3, + "domain_length": 0.0 if domain_length is None else domain_length, + "domain_height": 0.0 if domain_height is None else domain_height, } log.info( f"initialized interpolation factory for backend = '{self._backend_name()}' and grid = '{self._grid}'" @@ -250,12 +254,8 @@ def _register_computed_fields(self) -> None: }, connectivities={"c2e2c": dims.C2E2CDim}, params={ - "domain_length": self.grid.global_properties.domain_length - if self.grid.global_properties.domain_length is not None - else 0.0, - "domain_height": self.grid.global_properties.domain_height - if self.grid.global_properties.domain_height is not None - else 0.0, + "domain_length": self._config["domain_length"], + "domain_height": self._config["domain_height"], "grid_sphere_radius": constants.EARTH_RADIUS, "lsq_dim_unk": self._config["lsq_dim_unk"], "lsq_dim_c": self._config["lsq_dim_c"], From 836395780cea334c9a990f67604c165fd21e53f9 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 14:06:57 +0100 Subject: [PATCH 03/25] small test update --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index c09a5a5534..140ce8650d 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -133,7 +133,7 @@ def lsq_compute_coeffs( cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c ) - exchange(lsq_weights_c, dim=CellDim) + exchange(lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, lsq_pseudoinv, @@ -144,7 +144,7 @@ def lsq_compute_coeffs( lsq_dim_unk, lsq_dim_c, ) - exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 0, :]) + exchange(lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv From 95391116a253c3be9ec9f77f266606ff26a8c5d8 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 14:17:40 +0100 Subject: [PATCH 04/25] small edit to distributed tests on Ci --- ci/distributed.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ci/distributed.yml b/ci/distributed.yml index 916afe1734..c0a38d8558 100644 --- a/ci/distributed.yml +++ b/ci/distributed.yml @@ -79,7 +79,7 @@ build_distributed_cpu: - ci/scripts/ci-mpi-wrapper.sh pytest -sv -k mpi_tests --with-mpi --backend=$BACKEND model/$COMPONENT parallel: matrix: - - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, common] + - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, atmosphere/advection, common] BACKEND: [embedded, gtfn_cpu, dace_cpu] rules: - if: $COMPONENT == 'atmosphere/diffusion' @@ -91,6 +91,9 @@ build_distributed_cpu: - if: $COMPONENT == 'atmosphere/dycore' variables: SLURM_TIMELIMIT: '00:15:00' + - if: $COMPONENT == 'atmosphere/advection' + variables: + SLURM_TIMELIMIT: '00:15:00' - when: on_success variables: SLURM_TIMELIMIT: '00:30:00' From 0328df20ddb628211f30841c10e40c64a984933c Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 14:37:48 +0100 Subject: [PATCH 05/25] small edit to mpi test --- .../tests/advection/mpi_tests/test_parallel_advection.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index 4d88de67e9..58b9de5327 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -65,6 +65,13 @@ from .. import utils from ..fixtures import * # noqa: F403 +try: + import mpi4py + + mpi_decomposition.init_mpi() +except ImportError: + pytest.skip("Skipping parallel on single node installation", allow_module_level=True) + @pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.datatest @pytest.mark.parametrize("experiment", [test_defs.Experiments.MCH_CH_R04B09]) From 3306e78308b59e56cd2456747da6269eb10a49e4 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 14:44:57 +0100 Subject: [PATCH 06/25] small import edit --- .../tests/advection/integration_tests/test_advection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index 9f051f3d0d..7c00f098b7 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -37,7 +37,7 @@ metrics_savepoint, processor_props, ) -from model.common.tests.common.utils import dummy_exchange +from icon4py.model.common.tests.common.utils import dummy_exchange from ..fixtures import advection_exit_savepoint, advection_init_savepoint from ..utils import ( From 0655216dbf382fbdf3b891c79db02ac45664490c Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 14:53:56 +0100 Subject: [PATCH 07/25] removed unused import --- .../tests/advection/integration_tests/test_advection.py | 1 - 1 file changed, 1 deletion(-) diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index 7c00f098b7..be47b2e5bb 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -37,7 +37,6 @@ metrics_savepoint, processor_props, ) -from icon4py.model.common.tests.common.utils import dummy_exchange from ..fixtures import advection_exit_savepoint, advection_init_savepoint from ..utils import ( From 011ab594a75facbe735d8dfd2144bb749570ba08 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 6 Feb 2026 15:00:11 +0100 Subject: [PATCH 08/25] placed dim back --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 140ce8650d..c09a5a5534 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -133,7 +133,7 @@ def lsq_compute_coeffs( cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c ) - exchange(lsq_weights_c) + exchange(lsq_weights_c, dim=CellDim) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, lsq_pseudoinv, @@ -144,7 +144,7 @@ def lsq_compute_coeffs( lsq_dim_unk, lsq_dim_c, ) - exchange(lsq_pseudoinv[:, 0, :]) - exchange(lsq_pseudoinv[:, 1, :]) + exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) return lsq_pseudoinv From 9bdccd2519c9ca253650df6369d61b511baf8115 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 9 Feb 2026 13:03:15 +0100 Subject: [PATCH 09/25] added setup program for programs in general advection --- .../model/atmosphere/advection/advection.py | 69 +++++++++++++------ 1 file changed, 48 insertions(+), 21 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index 2894b5aa88..fd10304d14 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -11,6 +11,7 @@ from abc import ABC, abstractmethod from enum import Enum +import gt4py.next as gtx import gt4py.next.typing as gtx_typing import icon4py.model.common.grid.states as grid_states @@ -29,6 +30,7 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid +from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import data_allocation as data_alloc @@ -175,7 +177,19 @@ def __init__( self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) # stencils - self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field = setup_program( + backend=self._backend, + program=copy_cell_kdim_field, + horizontal_sizes={ + "horizontal_start": self._start_cell_nudging, + "horizontal_end": self._end_cell_local, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": self._grid.num_levels, + }, + offset_provider=self._grid.connectivities, + ) def run( self, @@ -188,18 +202,13 @@ def run( log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) + #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("communication of prep_adv cell field: mass_flx_ic - end") log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( field_in=p_tracer_now, field_out=p_tracer_new, - horizontal_start=self._start_cell_nudging, - horizontal_end=self._end_cell_local, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.connectivities, ) log.debug("running stencil copy_cell_kdim_field - end") @@ -237,9 +246,33 @@ def __init__( ) self._determine_local_domains() # stencils - self._apply_density_increment = apply_density_increment.with_backend(self._backend) - self._apply_interpolated_tracer_time_tendency = ( - apply_interpolated_tracer_time_tendency.with_backend(self._backend) + self._apply_density_increment = setup_program( + backend=self._backend, + program=apply_density_increment, + constant_args={ + "deepatmo_divzl": self._metric_state.deepatmo_divzl, + "deepatmo_divzu": self._metric_state.deepatmo_divzu, + }, + horizontal_sizes={ + "horizontal_end": self._end_cell_end, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, + offset_provider=self._grid.connectivities, + ) + self._apply_interpolated_tracer_time_tendency = setup_program( + backend=self._backend, + program=apply_interpolated_tracer_time_tendency, + horizontal_sizes={ + "horizontal_start": self._start_cell_lateral_boundary, + "horizontal_end": self._end_cell_lateral_boundary_level_4, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, ) log.debug("advection class init - end") @@ -287,16 +320,10 @@ def run( self._apply_density_increment( rhodz_in=rhodz_in, p_mflx_contra_v=prep_adv.mass_flx_ic, # exchange - deepatmo_divzl=self._metric_state.deepatmo_divzl, - deepatmo_divzu=self._metric_state.deepatmo_divzu, rhodz_out=self._rhodz_ast2, p_dtime=dtime, even_timestep=self._even_timestep, horizontal_start=horizontal_start, - horizontal_end=self._end_cell_end, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.connectivities, ) log.debug("running stencil apply_density_increment - end") @@ -324,6 +351,7 @@ def run( p_mflx_tracer_h=diagnostic_state.hfl_tracer, dtime=dtime, ) + #self._exchange.exchange_and_wait(dims.EdgeDim, diagnostic_state.hfl_tracer) else: # horizontal transport @@ -336,6 +364,7 @@ def run( p_mflx_tracer_h=diagnostic_state.hfl_tracer, dtime=dtime, ) + #self._exchange.exchange_and_wait(dims.EdgeDim, diagnostic_state.hfl_tracer) # vertical transport self._vertical_advection.run( @@ -357,11 +386,6 @@ def run( p_grf_tend_tracer=diagnostic_state.grf_tend_tracer, p_tracer_new=p_tracer_new, p_dtime=dtime, - horizontal_start=self._start_cell_lateral_boundary, - horizontal_end=self._end_cell_lateral_boundary_level_4, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.connectivities, ) log.debug("running stencil apply_interpolated_tracer_time_tendency - end") @@ -411,6 +435,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- least_squares_state=least_squares_state, horizontal_limiter=horizontal_limiter, backend=backend, + exchange=exchange, ) horizontal_advection = advection_horizontal.SemiLagrangian( tracer_flux=tracer_flux, @@ -443,6 +468,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- grid=grid, metric_state=metric_state, backend=backend, + exchange=exchange, ) case VerticalAdvectionType.PPM_3RD_ORDER: boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) @@ -452,6 +478,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- grid=grid, metric_state=metric_state, backend=backend, + exchange=exchange, ) case _: raise NotImplementedError("Unknown vertical advection type.") From 81bee974a88bfec68311276c0dc0f37a606d2dd2 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 9 Feb 2026 13:11:16 +0100 Subject: [PATCH 10/25] edits to tests --- .../integration_tests/test_advection.py | 2 +- .../mpi_tests/test_parallel_advection.py | 164 +++++++++--------- 2 files changed, 85 insertions(+), 81 deletions(-) diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index be47b2e5bb..722c520890 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -179,7 +179,7 @@ def test_advection_run_single_step( cell_params=cell_geometry, even_timestep=even_timestep, backend=backend, - exchange=dummy_exchange + exchange=None, ) diagnostic_state = construct_diagnostic_init_state( diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index 58b9de5327..858226c6da 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -7,23 +7,33 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next.typing as gtx_typing +import numpy as np import pytest +from gt4py.next import typing as gtx_typing -from icon4py.model.testing import test_utils +import icon4py.model.testing.test_utils as test_helpers from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.atmosphere.advection.advection_lsq_coeffs import lsq_compute_coeffs -from icon4py.model.common import constants, dimension as dims, utils +from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as nh +from icon4py.model.common import constants, dimension as dims, type_alias as ta, utils +from icon4py.model.common.decomposition import definitions, mpi_decomposition from icon4py.model.common.grid import ( base as base_grid, geometry_attributes as geometry_attrs, horizontal as h_grid, + icon, + states as grid_states, + vertical as v_grid, ) from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import ( definitions as test_defs, grid_utils, grid_utils as gridtest_utils, + parallel_helpers, + serialbox, serialbox as sb, + test_utils, ) from icon4py.model.testing.fixtures.datatest import ( backend, @@ -37,7 +47,9 @@ metrics_savepoint, processor_props, ) -from icon4py.model.common.decomposition import definitions + +from .. import utils +from ..fixtures import * # noqa: F403 from ..fixtures import advection_exit_savepoint, advection_init_savepoint from ..utils import ( construct_config, @@ -51,19 +63,6 @@ verify_advection_fields, ) -import numpy as np -import pytest -from gt4py.next import typing as gtx_typing - -from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as nh -from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.decomposition import definitions, mpi_decomposition -from icon4py.model.common.grid import icon, states as grid_states, vertical as v_grid -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import definitions as test_defs, parallel_helpers, serialbox, test_utils - -from .. import utils -from ..fixtures import * # noqa: F403 try: import mpi4py @@ -72,6 +71,7 @@ except ImportError: pytest.skip("Skipping parallel on single node installation", allow_module_level=True) + @pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.datatest @pytest.mark.parametrize("experiment", [test_defs.Experiments.MCH_CH_R04B09]) @@ -197,7 +197,7 @@ def test_advection_run_single_step( cell_params=cell_geometry, even_timestep=even_timestep, backend=backend, - exchange=exchange + exchange=exchange, ) diagnostic_state = construct_diagnostic_init_state( @@ -233,66 +233,70 @@ def test_advection_run_single_step( ) -# @pytest.mark.level("unit") -# @pytest.mark.datatest -# def test_lsq_compute_coeffs( -# icon_grid: base_grid.Grid, -# grid_savepoint: sb.IconGridSavepoint, -# backend: gtx_typing.Backend, -# interpolation_savepoint: sb.InterpolationSavepoint, -# experiment: test_defs.Experiment, -# ) -> None: -# gm = grid_utils.get_grid_manager_from_identifier( -# experiment.grid, -# num_levels=1, -# keep_skip_values=True, -# allocator=backend, -# ) -# -# c2e2c = gm.grid.connectivities["C2E2C"].asnumpy() -# cell_owner_mask = grid_savepoint.c_owner_mask().asnumpy() -# grid_sphere_radius = constants.EARTH_RADIUS -# lsq_dim_unk = 2 -# lsq_dim_c = 3 -# lsq_wgt_exp = 2 -# cell_domain = h_grid.domain(dims.CellDim) -# -# min_rlcell_int = gm.grid.end_index(cell_domain(h_grid.Zone.LOCAL)) -# start_idx = gm.grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) -# -# grid_geometry = grid_utils.get_grid_geometry(backend, experiment) -# cell_center_x = grid_geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy() -# cell_center_y = grid_geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy() -# domain_length = gm.grid.global_properties.domain_length -# domain_height = gm.grid.global_properties.domain_height -# lsq_dim_stencil = 3 -# -# coordinates = gm.coordinates -# cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() -# cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() -# lsq_pseudoinv = lsq_compute_coeffs( -# cell_center_x, -# cell_center_y, -# cell_lat, -# cell_lon, -# c2e2c, -# cell_owner_mask, -# domain_length, -# domain_height, -# grid_sphere_radius, -# lsq_dim_unk, -# lsq_dim_c, -# lsq_wgt_exp, -# lsq_dim_stencil, -# start_idx, -# min_rlcell_int, -# exchange, -# icon_grid.geometry_type, -# ) -# -# assert test_helpers.dallclose( -# interpolation_savepoint.lsq_pseudoinv_1().asnumpy(), lsq_pseudoinv[:, 0, :], atol=1e-15 -# ) -# assert test_helpers.dallclose( -# interpolation_savepoint.lsq_pseudoinv_2().asnumpy(), lsq_pseudoinv[:, 1, :], atol=1e-15 -# ) +@pytest.mark.level("unit") +@pytest.mark.datatest +@pytest.mark.mpi +def test_lsq_compute_coeffs( + icon_grid: base_grid.Grid, + grid_savepoint: sb.IconGridSavepoint, + backend: gtx_typing.Backend, + interpolation_savepoint: sb.InterpolationSavepoint, + experiment: test_defs.Experiment, + processor_props: definitions.ProcessProperties, + decomposition_info: definitions.DecompositionInfo, # : F811 fixture +) -> None: + gm = grid_utils.get_grid_manager_from_identifier( + experiment.grid, + num_levels=1, + keep_skip_values=True, + allocator=backend, + ) + + c2e2c = gm.grid.connectivities["C2E2C"].asnumpy() + cell_owner_mask = grid_savepoint.c_owner_mask().asnumpy() + grid_sphere_radius = constants.EARTH_RADIUS + lsq_dim_unk = 2 + lsq_dim_c = 3 + lsq_wgt_exp = 2 + cell_domain = h_grid.domain(dims.CellDim) + + min_rlcell_int = gm.grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + start_idx = gm.grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) + + grid_geometry = grid_utils.get_grid_geometry(backend, experiment) + cell_center_x = grid_geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy() + cell_center_y = grid_geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy() + domain_length = gm.grid.global_properties.domain_length + domain_height = gm.grid.global_properties.domain_height + lsq_dim_stencil = 3 + exchange = definitions.create_exchange(processor_props, decomposition_info) + + coordinates = gm.coordinates + cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() + cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() + lsq_pseudoinv = lsq_compute_coeffs( + cell_center_x, + cell_center_y, + cell_lat, + cell_lon, + c2e2c, + cell_owner_mask, + domain_length, + domain_height, + grid_sphere_radius, + lsq_dim_unk, + lsq_dim_c, + lsq_wgt_exp, + lsq_dim_stencil, + start_idx, + min_rlcell_int, + icon_grid.geometry_type.value, + exchange, + ) + + assert test_helpers.dallclose( + interpolation_savepoint.lsq_pseudoinv_1().asnumpy(), lsq_pseudoinv[:, 0, :], atol=1e-15 + ) + assert test_helpers.dallclose( + interpolation_savepoint.lsq_pseudoinv_2().asnumpy(), lsq_pseudoinv[:, 1, :], atol=1e-15 + ) From 1cf9d0136d4428b1fcbb25f3284ed60ed8756af4 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 9 Feb 2026 14:26:26 +0100 Subject: [PATCH 11/25] small edits --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 9 ++++++--- .../tests/advection/integration_tests/test_advection.py | 2 ++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index c09a5a5534..0624e16c49 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -133,7 +133,9 @@ def lsq_compute_coeffs( cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c ) - exchange(lsq_weights_c, dim=CellDim) + if exchange is not None: + exchange(lsq_weights_c, dim=CellDim) + lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, lsq_pseudoinv, @@ -144,7 +146,8 @@ def lsq_compute_coeffs( lsq_dim_unk, lsq_dim_c, ) - exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) + if exchange is not None: + exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) return lsq_pseudoinv diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index 722c520890..9e358df272 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -161,6 +161,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=None ) least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend) @@ -269,6 +270,7 @@ def test_lsq_compute_coeffs( start_idx, min_rlcell_int, icon_grid.geometry_type, + exchange=None ) assert test_helpers.dallclose( From 090120ee079f6cc40cb2780e7db027c97fd0da96 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 10 Feb 2026 09:44:24 +0100 Subject: [PATCH 12/25] more edits for exchange and tests --- .../icon4py/model/atmosphere/advection/advection.py | 8 +------- .../atmosphere/advection/advection_horizontal.py | 2 ++ .../model/atmosphere/advection/advection_vertical.py | 11 +++++++++-- .../atmosphere/advection/tests/advection/fixtures.py | 4 +--- .../advection/integration_tests/test_advection.py | 4 ++-- .../advection/mpi_tests/test_parallel_advection.py | 10 +++------- 6 files changed, 18 insertions(+), 21 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index fd10304d14..004356eb2e 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -202,7 +202,6 @@ def run( log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") - #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("communication of prep_adv cell field: mass_flx_ic - end") log.debug("running stencil copy_cell_kdim_field - start") @@ -305,7 +304,6 @@ def run( log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("communication of prep_adv cell field: mass_flx_ic - end") # reintegrate density for conservation of mass @@ -316,10 +314,9 @@ def run( ) log.debug("running stencil apply_density_increment - start") - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) self._apply_density_increment( rhodz_in=rhodz_in, - p_mflx_contra_v=prep_adv.mass_flx_ic, # exchange + p_mflx_contra_v=prep_adv.mass_flx_ic, rhodz_out=self._rhodz_ast2, p_dtime=dtime, even_timestep=self._even_timestep, @@ -351,7 +348,6 @@ def run( p_mflx_tracer_h=diagnostic_state.hfl_tracer, dtime=dtime, ) - #self._exchange.exchange_and_wait(dims.EdgeDim, diagnostic_state.hfl_tracer) else: # horizontal transport @@ -364,7 +360,6 @@ def run( p_mflx_tracer_h=diagnostic_state.hfl_tracer, dtime=dtime, ) - #self._exchange.exchange_and_wait(dims.EdgeDim, diagnostic_state.hfl_tracer) # vertical transport self._vertical_advection.run( @@ -391,7 +386,6 @@ def run( # exchange updated tracer values, originally happens only if iforcing /= inwp log.debug("communication of advection cell field: p_tracer_new - start") - self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("communication of advection cell field: p_tracer_new - end") # finalize step diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index 200bed37db..4bade94bff 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -196,11 +196,13 @@ def __init__( least_squares_state: advection_states.AdvectionLeastSquaresState, backend: model_backends.BackendLike, horizontal_limiter: HorizontalFluxLimiter | None = None, + exchange: decomposition.ExchangeRuntime | None = None, ): self._grid = grid self._least_squares_state = least_squares_state self._backend = backend self._horizontal_limiter = horizontal_limiter or HorizontalFluxLimiter() + self._exchange = exchange or decomposition.SingleNodeExchange() # cell indices cell_domain = h_grid.domain(dims.CellDim) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index af595436fa..dc278a59a6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -63,6 +63,7 @@ field_type_aliases as fa, type_alias as ta, ) +from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import data_allocation as data_alloc @@ -538,6 +539,7 @@ def __init__( grid: icon_grid.IconGrid, metric_state: advection_states.AdvectionMetricState, backend: gtx_typing.Backend | None, + exchange: decomposition.ExchangeRuntime | None = None, ): log.debug("vertical advection class init - start") @@ -546,6 +548,7 @@ def __init__( self._grid = grid self._metric_state = metric_state self._backend = backend + self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -625,7 +628,7 @@ def _compute_numerical_flux( horizontal_start, horizontal_end = self._get_horizontal_start_end( even_timestep=even_timestep ) - + self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("running stencil compute_vertical_tracer_flux_upwind - start") self._compute_vertical_tracer_flux_upwind( p_cc=p_tracer_now, @@ -674,7 +677,7 @@ def _update_unknowns( horizontal_end=horizontal_end, ) log.debug("running stencil integrate_tracer_vertically - end") - + self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("vertical unknowns update - end") @@ -688,6 +691,7 @@ def __init__( grid: icon_grid.IconGrid, metric_state: advection_states.AdvectionMetricState, backend: gtx_typing.Backend | None, + exchange: decomposition.ExchangeRuntime | None = None, ): log.debug("vertical advection class init - start") @@ -697,6 +701,7 @@ def __init__( self._grid = grid self._metric_state = metric_state self._backend = backend + self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -913,6 +918,7 @@ def _compute_numerical_flux( log.debug("running stencil init_constant_cell_kdim_field - end") log.debug("running stencil compute_ppm4gpu_courant_number - start") + self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) self._compute_ppm4gpu_courant_number( p_mflx_contra_v=prep_adv.mass_flx_ic, p_cellmass_now=rhodz_now, @@ -1095,6 +1101,7 @@ def _update_unknowns( horizontal_start=horizontal_start, horizontal_end=horizontal_end, ) + self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("running stencil integrate_tracer_vertically - end") log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/tests/advection/fixtures.py b/model/atmosphere/advection/tests/advection/fixtures.py index 8aef7d508e..bee4cd2229 100644 --- a/model/atmosphere/advection/tests/advection/fixtures.py +++ b/model/atmosphere/advection/tests/advection/fixtures.py @@ -10,7 +10,7 @@ import pytest -from icon4py.model.testing.fixtures.datatest import data_provider +from icon4py.model.testing.fixtures.datatest import data_provider, decomposition_info @pytest.fixture @@ -44,5 +44,3 @@ def advection_exit_savepoint(data_provider, date): fixture, passing 'date='. """ return data_provider.from_advection_exit_savepoint(size=data_provider.grid_size, date=date) - -from icon4py.model.testing.fixtures.datatest import decomposition_info diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index 9e358df272..25bb1009e0 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -161,7 +161,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=None + exchange=None, ) least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend) @@ -270,7 +270,7 @@ def test_lsq_compute_coeffs( start_idx, min_rlcell_int, icon_grid.geometry_type, - exchange=None + exchange=None, ) assert test_helpers.dallclose( diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index 858226c6da..daab43e86f 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -223,14 +223,10 @@ def test_advection_run_single_step( ) p_tracer_new_ref = advection_exit_savepoint.tracer(ntracer) - verify_advection_fields( - grid=icon_grid, - diagnostic_state=diagnostic_state, - diagnostic_state_ref=diagnostic_state_ref, - p_tracer_new=p_tracer_new, - p_tracer_new_ref=p_tracer_new_ref, - even_timestep=even_timestep, + assert test_helpers.dallclose( + diagnostic_state.hfl_tracer.asnumpy(), diagnostic_state_ref.hfl_tracer.asnumpy(), atol=1e-8 ) + assert test_helpers.dallclose(p_tracer_new_ref.asnumpy(), p_tracer_new.asnumpy(), atol=1e-10) @pytest.mark.level("unit") From aabab23b056b65e0f96268e8f48843def2788312 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:13:32 +0100 Subject: [PATCH 13/25] small edit to mpi test --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 0624e16c49..8c4c33facf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -134,7 +134,7 @@ def lsq_compute_coeffs( ) if exchange is not None: - exchange(lsq_weights_c, dim=CellDim) + exchange(CellDim, lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -147,7 +147,7 @@ def lsq_compute_coeffs( lsq_dim_c, ) if exchange is not None: - exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) + exchange(CellDim, lsq_pseudoinv[:, 0, :]) + exchange(CellDim, lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv From eb8b307d3ead27c24d2edab072ba0b44a34131b8 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 10 Feb 2026 10:41:32 +0100 Subject: [PATCH 14/25] small edit to mpi test --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 8c4c33facf..09f8190aa6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -134,7 +134,7 @@ def lsq_compute_coeffs( ) if exchange is not None: - exchange(CellDim, lsq_weights_c) + exchange(lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -147,7 +147,7 @@ def lsq_compute_coeffs( lsq_dim_c, ) if exchange is not None: - exchange(CellDim, lsq_pseudoinv[:, 0, :]) - exchange(CellDim, lsq_pseudoinv[:, 1, :]) + exchange(lsq_pseudoinv[:, 0, :]) + exchange(lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv From 04f7f72a44c73ac2f2fb4fac32fb0eab81a67e8f Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 10 Feb 2026 11:47:30 +0100 Subject: [PATCH 15/25] small edit to mpi test --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 09f8190aa6..8c4c33facf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -134,7 +134,7 @@ def lsq_compute_coeffs( ) if exchange is not None: - exchange(lsq_weights_c) + exchange(CellDim, lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -147,7 +147,7 @@ def lsq_compute_coeffs( lsq_dim_c, ) if exchange is not None: - exchange(lsq_pseudoinv[:, 0, :]) - exchange(lsq_pseudoinv[:, 1, :]) + exchange(CellDim, lsq_pseudoinv[:, 0, :]) + exchange(CellDim, lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv From 738afc29816fff48ead17c7c33a0adebbd32277b Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 10 Feb 2026 11:50:19 +0100 Subject: [PATCH 16/25] small edit to mpi test --- .../model/atmosphere/advection/advection_lsq_coeffs.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 8c4c33facf..0624e16c49 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -134,7 +134,7 @@ def lsq_compute_coeffs( ) if exchange is not None: - exchange(CellDim, lsq_weights_c) + exchange(lsq_weights_c, dim=CellDim) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -147,7 +147,7 @@ def lsq_compute_coeffs( lsq_dim_c, ) if exchange is not None: - exchange(CellDim, lsq_pseudoinv[:, 0, :]) - exchange(CellDim, lsq_pseudoinv[:, 1, :]) + exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) return lsq_pseudoinv From e897fcc2648e3e3d81189e222efd3cc962970342 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Tue, 10 Feb 2026 16:32:08 +0100 Subject: [PATCH 17/25] cleanups in lsq coeffs funcs --- .../advection/advection_lsq_coeffs.py | 45 +++++++++++-------- .../icon4py/model/common/math/projection.py | 16 +------ 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 0624e16c49..66f0214574 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -13,10 +13,7 @@ from icon4py.model.common.dimension import CellDim from icon4py.model.common.grid import base as base_grid -from icon4py.model.common.math.projection import ( - gnomonic_proj_single_val, - plane_torus_closest_coordinates, -) +from icon4py.model.common.math.projection import diff_on_edges_torus_numpy, gnomonic_proj from icon4py.model.common.utils import data_allocation as data_alloc @@ -96,27 +93,30 @@ def lsq_compute_coeffs( lsq_weights_c = array_ns.zeros((min_rlcell_int, lsq_dim_stencil)) lsq_pseudoinv = array_ns.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) z_lsq_mat_c = array_ns.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) - - for jc in range(start_idx, min_rlcell_int): - match base_grid.GeometryType(geometry_type): - case base_grid.GeometryType.ICOSAHEDRON: - z_dist_g = array_ns.zeros((lsq_dim_c, 2)) - for js in range(lsq_dim_stencil): - z_dist_g[js, :] = gnomonic_proj_single_val( - cell_lon[jc], cell_lat[jc], cell_lon[c2e2c[jc, js]], cell_lat[c2e2c[jc, js]] + z_dist_g = array_ns.zeros((min_rlcell_int, lsq_dim_c, 2)) + match base_grid.GeometryType(geometry_type): + case base_grid.GeometryType.ICOSAHEDRON: + for js in range(lsq_dim_stencil): + z_dist_g[:, js, :] = np.asarray( + gnomonic_proj( + cell_lon[:], cell_lat[:], cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]] ) - z_dist_g *= grid_sphere_radius + ).T + + z_dist_g *= grid_sphere_radius + min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) + for jc in range(start_idx, min_rlcell_int): if cell_owner_mask[jc]: z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 - case base_grid.GeometryType.TORUS: + case base_grid.GeometryType.TORUS: + for jc in range(start_idx, min_rlcell_int): ilc_s = c2e2c[jc, :lsq_dim_stencil] cc_cell = array_ns.zeros((lsq_dim_stencil, 2)) cc_cv = (cell_center_x[jc], cell_center_y[jc]) for js in range(lsq_dim_stencil): - cc_cell[js, :] = plane_torus_closest_coordinates( + cc_cell[js, :] = diff_on_edges_torus_numpy( cell_center_x[jc], cell_center_y[jc], cell_center_x[ilc_s][js], @@ -124,13 +124,20 @@ def lsq_compute_coeffs( domain_length, domain_height, ) - z_dist_g = cc_cell - cc_cv + z_dist_g[jc, :, :] = cc_cell - cc_cv + for jc in range(start_idx, min_rlcell_int): lsq_weights_c[jc, :] = compute_lsq_weights_c( - z_dist_g, lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp + z_dist_g[jc, :, :], lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp ) z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( - cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c + cell_owner_mask, + z_lsq_mat_c, + lsq_weights_c, + z_dist_g[jc, :, :], + jc, + lsq_dim_unk, + lsq_dim_c, ) if exchange is not None: diff --git a/model/common/src/icon4py/model/common/math/projection.py b/model/common/src/icon4py/model/common/math/projection.py index 9df9650ee1..fcec8fbc5f 100644 --- a/model/common/src/icon4py/model/common/math/projection.py +++ b/model/common/src/icon4py/model/common/math/projection.py @@ -6,7 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import math import numpy as np @@ -48,20 +47,7 @@ def gnomonic_proj( return x, y -def gnomonic_proj_single_val( - lon_c: float, lat_c: float, lon: float, lat: float -) -> tuple[float, float]: - cosc = math.sin(lat_c) * math.sin(lat) + math.cos(lat_c) * math.cos(lat) * math.cos(lon - lon_c) - zk = 1.0 / cosc - - x = zk * math.cos(lat) * math.sin(lon - lon_c) - y = zk * ( - math.cos(lat_c) * math.sin(lat) - math.sin(lat_c) * math.cos(lat) * math.cos(lon - lon_c) - ) - return x, y - - -def plane_torus_closest_coordinates( +def diff_on_edges_torus_numpy( cc_cv_x: float, cc_cv_y: float, cc_cell_x: float, From 8583d8de6cad957d808dd1c9b0b9e292cb8dffb9 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 11 Feb 2026 13:41:26 +0100 Subject: [PATCH 18/25] updates --- .../model/atmosphere/advection/advection.py | 8 +++- .../advection/advection_horizontal.py | 8 +++- .../advection/advection_lsq_coeffs.py | 18 ++++---- .../advection/advection_vertical.py | 9 ++-- .../advection/tests/advection/fixtures.py | 12 +++++ .../mpi_tests/test_parallel_advection.py | 44 +++++++------------ 6 files changed, 54 insertions(+), 45 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index 004356eb2e..a73dae483b 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -422,7 +422,9 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: - horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend) + horizontal_advection = advection_horizontal.NoAdvection( + grid=grid, backend=backend, exchange=exchange + ) case HorizontalAdvectionType.LINEAR_2ND_ORDER: tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, @@ -454,7 +456,9 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- match config.vertical_advection_type: case VerticalAdvectionType.NO_ADVECTION: - vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend) + vertical_advection = advection_vertical.NoAdvection( + grid=grid, backend=backend, exchange=exchange + ) case VerticalAdvectionType.UPWIND_1ST_ORDER: boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) vertical_advection = advection_vertical.FirstOrderUpwind( diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index 4bade94bff..a3b884df0d 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -166,6 +166,7 @@ def apply_flux_limiter( r_m=self._r_m, p_mflx_tracer_h=p_mflx_tracer_h, ) + # self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) log.debug( "running stencil apply_positive_definite_horizontal_multiplicative_flux_factor - end" ) @@ -298,6 +299,7 @@ def compute_tracer_flux( log.debug( "running stencil compute_horizontal_tracer_flux_from_linear_coefficients_alt - end" ) + # self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) self._horizontal_limiter.apply_flux_limiter( p_tracer_now=p_tracer_now, @@ -342,7 +344,7 @@ def run( class NoAdvection(HorizontalAdvection): """Class that implements disabled horizontal advection.""" - def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike): + def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike, exchange): log.debug("horizontal advection class init - start") # input arguments @@ -352,6 +354,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike cell_domain = h_grid.domain(dims.CellDim) self._start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) self._end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + self._exchange = exchange # stencils self._copy_cell_kdim_field = setup_program( @@ -388,7 +391,7 @@ def run( field_out=p_tracer_new, ) log.debug("running stencil copy_cell_kdim_field - end") - + self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("horizontal advection run - end") @@ -423,6 +426,7 @@ def run( p_mflx_tracer_h=p_mflx_tracer_h, dtime=dtime, ) + # self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) log.debug("horizontal advection run - end") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 66f0214574..f2d7d738a9 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -31,11 +31,11 @@ def compute_lsq_pseudoinv( for jjk in range(lsq_dim_unk): for jc in range(start_idx, min_rlcell_int): u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :]) - if cell_owner_mask[jc]: - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] - + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] - ) + # if cell_owner_mask[jc]: + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] + + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] + ) return lsq_pseudoinv @@ -61,8 +61,8 @@ def compute_z_lsq_mat_c( lsq_dim_c: int, ) -> data_alloc.NDArray: min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + # if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 for js in range(lsq_dim_c): z_lsq_mat_c[jc, js, :lsq_dim_unk] = lsq_weights_c[jc, js] * z_dist_g[js, :] @@ -107,8 +107,8 @@ def lsq_compute_coeffs( min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) for jc in range(start_idx, min_rlcell_int): - if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + # if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 case base_grid.GeometryType.TORUS: for jc in range(start_idx, min_rlcell_int): ilc_s = c2e2c[jc, :lsq_dim_stencil] diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index dc278a59a6..d3fdc9c9f2 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -401,12 +401,13 @@ def run( class NoAdvection(VerticalAdvection): """Class that implements disabled vertical advection.""" - def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None): + def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None, exchange): log.debug("vertical advection class init - start") # input arguments self._grid = grid self._backend = backend + self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -465,7 +466,7 @@ def run( horizontal_end=horizontal_end, ) log.debug("running stencil copy_cell_kdim_field - end") - + # self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("vertical advection run - end") @@ -628,7 +629,7 @@ def _compute_numerical_flux( horizontal_start, horizontal_end = self._get_horizontal_start_end( even_timestep=even_timestep ) - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) + # self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("running stencil compute_vertical_tracer_flux_upwind - start") self._compute_vertical_tracer_flux_upwind( p_cc=p_tracer_now, @@ -1103,5 +1104,5 @@ def _update_unknowns( ) self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("running stencil integrate_tracer_vertically - end") - + # self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/tests/advection/fixtures.py b/model/atmosphere/advection/tests/advection/fixtures.py index bee4cd2229..f039d29ed4 100644 --- a/model/atmosphere/advection/tests/advection/fixtures.py +++ b/model/atmosphere/advection/tests/advection/fixtures.py @@ -10,6 +10,8 @@ import pytest +from icon4py.model.atmosphere.advection import advection_states +from icon4py.model.testing import serialbox from icon4py.model.testing.fixtures.datatest import data_provider, decomposition_info @@ -44,3 +46,13 @@ def advection_exit_savepoint(data_provider, date): fixture, passing 'date='. """ return data_provider.from_advection_exit_savepoint(size=data_provider.grid_size, date=date) + + +@pytest.fixture +def advection_lsq_state( + interpolation_savepoint: serialbox.InterpolationSavepoint, +) -> advection_states.AdvectionLeastSquaresState: + return advection_states.AdvectionLeastSquaresState( + lsq_pseudoinv_1=interpolation_savepoint.lsq_pseudoinv_1(), + lsq_pseudoinv_2=interpolation_savepoint.lsq_pseudoinv_2(), + ) diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index daab43e86f..96c9c54596 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -25,6 +25,7 @@ states as grid_states, vertical as v_grid, ) +from icon4py.model.common.interpolation import interpolation_attributes as intp_attrs from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import ( definitions as test_defs, @@ -50,7 +51,7 @@ from .. import utils from ..fixtures import * # noqa: F403 -from ..fixtures import advection_exit_savepoint, advection_init_savepoint +from ..fixtures import advection_exit_savepoint, advection_init_savepoint, advection_lsq_state from ..utils import ( construct_config, construct_diagnostic_exit_state, @@ -137,6 +138,7 @@ def test_advection_run_single_step( experiment: test_defs.Experiment, processor_props: definitions.ProcessProperties, decomposition_info: definitions.DecompositionInfo, # : F811 fixture + advection_lsq_state, ): if test_utils.is_embedded(backend): # https://github.com/GridTools/gt4py/issues/1583 @@ -149,38 +151,18 @@ def test_advection_run_single_step( pytest.xfail( "This test is skipped until the cause of nonzero horizontal advection if revealed." ) + + parallel_helpers.check_comm_size(processor_props) + parallel_helpers.log_process_properties(processor_props) + parallel_helpers.log_local_field_size(decomposition_info) config = construct_config( horizontal_advection_type=horizontal_advection_type, horizontal_advection_limiter=horizontal_advection_limiter, vertical_advection_type=vertical_advection_type, vertical_advection_limiter=vertical_advection_limiter, ) - exchange = definitions.create_exchange(processor_props, decomposition_info) - interpolation_state = construct_interpolation_state(interpolation_savepoint, backend=backend) - geometry = gridtest_utils.get_grid_geometry(backend, experiment) - least_squares_coeffs = lsq_compute_coeffs( - cell_center_x=geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy(), - cell_center_y=geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy(), - cell_lat=geometry.get(geometry_attrs.CELL_LAT).asnumpy(), - cell_lon=geometry.get(geometry_attrs.CELL_LON).asnumpy(), - c2e2c=icon_grid.connectivities["C2E2C"].asnumpy(), - cell_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), - domain_length=geometry.grid.global_properties.domain_length, - domain_height=geometry.grid.global_properties.domain_height, - grid_sphere_radius=constants.EARTH_RADIUS, - lsq_dim_unk=2, - lsq_dim_c=3, - lsq_wgt_exp=2, - lsq_dim_stencil=3, - start_idx=icon_grid.start_index( - h_grid.domain(dims.CellDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) - ), - min_rlcell_int=icon_grid.end_index(h_grid.domain(dims.CellDim)(h_grid.Zone.HALO_LEVEL_2)), - geometry_type=icon_grid.geometry_type, - exchange=exchange, - ) - least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend) + interpolation_state = construct_interpolation_state(interpolation_savepoint, backend=backend) metric_state = construct_metric_state(icon_grid, metrics_savepoint, backend=backend) edge_geometry = grid_savepoint.construct_edge_geometry() @@ -189,9 +171,9 @@ def test_advection_run_single_step( advection_granule = advection.convert_config_to_advection( config=config, - grid=icon_grid, + grid=icon_grid, # gm.grid, interpolation_state=interpolation_state, - least_squares_state=least_squares_state, + least_squares_state=advection_lsq_state, # [field_1, field_2],#least_squares_state, metric_state=metric_state, edge_params=edge_geometry, cell_params=cell_geometry, @@ -205,6 +187,7 @@ def test_advection_run_single_step( ) prep_adv = construct_prep_adv(advection_init_savepoint) p_tracer_now = advection_init_savepoint.tracer(ntracer) + p_tracer_new = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) dtime = advection_init_savepoint.get_metadata("dtime").get("dtime") @@ -226,6 +209,11 @@ def test_advection_run_single_step( assert test_helpers.dallclose( diagnostic_state.hfl_tracer.asnumpy(), diagnostic_state_ref.hfl_tracer.asnumpy(), atol=1e-8 ) + assert test_utils.dallclose( + diagnostic_state.vfl_tracer.asnumpy(), + diagnostic_state_ref.vfl_tracer.asnumpy(), + rtol=1e-10, + ) assert test_helpers.dallclose(p_tracer_new_ref.asnumpy(), p_tracer_new.asnumpy(), atol=1e-10) From e0c8325e2698811732f7e09f517fdcd74525f5bb Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 11 Feb 2026 15:23:00 +0100 Subject: [PATCH 19/25] updates --- .../src/icon4py/model/atmosphere/advection/advection.py | 4 ++++ .../model/atmosphere/advection/advection_horizontal.py | 5 ++--- .../icon4py/model/atmosphere/advection/advection_vertical.py | 2 +- .../tests/advection/mpi_tests/test_parallel_advection.py | 5 +++-- 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index a73dae483b..7072874d64 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -202,6 +202,7 @@ def run( log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") + #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("communication of prep_adv cell field: mass_flx_ic - end") log.debug("running stencil copy_cell_kdim_field - start") @@ -386,6 +387,7 @@ def run( # exchange updated tracer values, originally happens only if iforcing /= inwp log.debug("communication of advection cell field: p_tracer_new - start") + #self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("communication of advection cell field: p_tracer_new - end") # finalize step @@ -422,10 +424,12 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: + #breakpoint() horizontal_advection = advection_horizontal.NoAdvection( grid=grid, backend=backend, exchange=exchange ) case HorizontalAdvectionType.LINEAR_2ND_ORDER: + #breakpoint() tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, least_squares_state=least_squares_state, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index a3b884df0d..ce91d4a033 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -166,7 +166,7 @@ def apply_flux_limiter( r_m=self._r_m, p_mflx_tracer_h=p_mflx_tracer_h, ) - # self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) + #self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) log.debug( "running stencil apply_positive_definite_horizontal_multiplicative_flux_factor - end" ) @@ -426,8 +426,7 @@ def run( p_mflx_tracer_h=p_mflx_tracer_h, dtime=dtime, ) - # self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) - + # self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("horizontal advection run - end") @abstractmethod diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index d3fdc9c9f2..a6ed0c2dca 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -919,7 +919,7 @@ def _compute_numerical_flux( log.debug("running stencil init_constant_cell_kdim_field - end") log.debug("running stencil compute_ppm4gpu_courant_number - start") - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) + #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) self._compute_ppm4gpu_courant_number( p_mflx_contra_v=prep_adv.mass_flx_ic, p_cellmass_now=rhodz_now, diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index 96c9c54596..e5ea34ca20 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -171,9 +171,9 @@ def test_advection_run_single_step( advection_granule = advection.convert_config_to_advection( config=config, - grid=icon_grid, # gm.grid, + grid=icon_grid, interpolation_state=interpolation_state, - least_squares_state=advection_lsq_state, # [field_1, field_2],#least_squares_state, + least_squares_state=advection_lsq_state, metric_state=metric_state, edge_params=edge_geometry, cell_params=cell_geometry, @@ -214,6 +214,7 @@ def test_advection_run_single_step( diagnostic_state_ref.vfl_tracer.asnumpy(), rtol=1e-10, ) + #breakpoint() assert test_helpers.dallclose(p_tracer_new_ref.asnumpy(), p_tracer_new.asnumpy(), atol=1e-10) From 101544c620ab89e77264dbed810805327b8e5786 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Wed, 11 Feb 2026 15:44:58 +0100 Subject: [PATCH 20/25] small edit to test --- .../atmosphere/advection/advection_lsq_coeffs.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index f2d7d738a9..75800cce42 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -141,7 +141,10 @@ def lsq_compute_coeffs( ) if exchange is not None: - exchange(lsq_weights_c, dim=CellDim) + try: + exchange(lsq_weights_c, dim=CellDim) + except TypeError: + exchange(lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -154,7 +157,11 @@ def lsq_compute_coeffs( lsq_dim_c, ) if exchange is not None: - exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) + try: + exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) + except TypeError: + exchange(lsq_pseudoinv[:, 0, :]) + exchange(lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv From 8884347920c79dc8519f5b621f509b2ca89d45af Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 12 Feb 2026 16:41:53 +0100 Subject: [PATCH 21/25] small edits to exchange --- .../icon4py/model/atmosphere/advection/advection_vertical.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index a6ed0c2dca..09ab1b246d 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -629,7 +629,7 @@ def _compute_numerical_flux( horizontal_start, horizontal_end = self._get_horizontal_start_end( even_timestep=even_timestep ) - # self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) + #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("running stencil compute_vertical_tracer_flux_upwind - start") self._compute_vertical_tracer_flux_upwind( p_cc=p_tracer_now, @@ -1059,7 +1059,6 @@ def _compute_numerical_flux( log.debug("running stencil compute_ppm4gpu_integer_flux - end") ## set boundary conditions - self._boundary_conditions.run( p_mflx_tracer_v=p_mflx_tracer_v, horizontal_start=horizontal_start, @@ -1067,7 +1066,6 @@ def _compute_numerical_flux( ) ## apply flux limiter - self._vertical_limiter.limit_fluxes( horizontal_start=horizontal_start, horizontal_end=horizontal_end ) From edb08341f3da2b1838ae9ab028b9cbc886090232 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 12 Feb 2026 18:09:45 +0100 Subject: [PATCH 22/25] small edit --- .../advection/advection_lsq_coeffs.py | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index 75800cce42..cb1432d099 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -140,11 +140,10 @@ def lsq_compute_coeffs( lsq_dim_c, ) - if exchange is not None: - try: - exchange(lsq_weights_c, dim=CellDim) - except TypeError: - exchange(lsq_weights_c) + try: + exchange(lsq_weights_c, dim=CellDim) + except TypeError: + exchange(lsq_weights_c) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -156,12 +155,12 @@ def lsq_compute_coeffs( lsq_dim_unk, lsq_dim_c, ) - if exchange is not None: - try: - exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) - except TypeError: - exchange(lsq_pseudoinv[:, 0, :]) - exchange(lsq_pseudoinv[:, 1, :]) + + try: + exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) + except TypeError: + exchange(lsq_pseudoinv[:, 0, :]) + exchange(lsq_pseudoinv[:, 1, :]) return lsq_pseudoinv From 33990e2b68276db550153dd251b150f311226b71 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 13 Feb 2026 10:45:14 +0100 Subject: [PATCH 23/25] cleanup --- .../model/atmosphere/advection/advection.py | 11 ----------- .../advection/advection_horizontal.py | 3 --- .../advection/advection_lsq_coeffs.py | 18 +++++++++--------- .../atmosphere/advection/advection_vertical.py | 4 ---- .../mpi_tests/test_parallel_advection.py | 1 - 5 files changed, 9 insertions(+), 28 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index 7072874d64..a35274acfd 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -201,10 +201,6 @@ def run( ): log.debug("advection run - start") - log.debug("communication of prep_adv cell field: mass_flx_ic - start") - #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) - log.debug("communication of prep_adv cell field: mass_flx_ic - end") - log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( field_in=p_tracer_now, @@ -385,11 +381,6 @@ def run( ) log.debug("running stencil apply_interpolated_tracer_time_tendency - end") - # exchange updated tracer values, originally happens only if iforcing /= inwp - log.debug("communication of advection cell field: p_tracer_new - start") - #self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) - log.debug("communication of advection cell field: p_tracer_new - end") - # finalize step self._even_timestep = not self._even_timestep @@ -424,12 +415,10 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: - #breakpoint() horizontal_advection = advection_horizontal.NoAdvection( grid=grid, backend=backend, exchange=exchange ) case HorizontalAdvectionType.LINEAR_2ND_ORDER: - #breakpoint() tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, least_squares_state=least_squares_state, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index ce91d4a033..8be6071ce3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -166,7 +166,6 @@ def apply_flux_limiter( r_m=self._r_m, p_mflx_tracer_h=p_mflx_tracer_h, ) - #self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) log.debug( "running stencil apply_positive_definite_horizontal_multiplicative_flux_factor - end" ) @@ -299,7 +298,6 @@ def compute_tracer_flux( log.debug( "running stencil compute_horizontal_tracer_flux_from_linear_coefficients_alt - end" ) - # self._exchange.exchange_and_wait(dims.EdgeDim, p_mflx_tracer_h) self._horizontal_limiter.apply_flux_limiter( p_tracer_now=p_tracer_now, @@ -426,7 +424,6 @@ def run( p_mflx_tracer_h=p_mflx_tracer_h, dtime=dtime, ) - # self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("horizontal advection run - end") @abstractmethod diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py index cb1432d099..d1bc4ad2d5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py @@ -31,11 +31,11 @@ def compute_lsq_pseudoinv( for jjk in range(lsq_dim_unk): for jc in range(start_idx, min_rlcell_int): u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :]) - # if cell_owner_mask[jc]: - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] - + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] - ) + if cell_owner_mask[jc]: + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] + + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] + ) return lsq_pseudoinv @@ -61,8 +61,8 @@ def compute_z_lsq_mat_c( lsq_dim_c: int, ) -> data_alloc.NDArray: min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - # if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 for js in range(lsq_dim_c): z_lsq_mat_c[jc, js, :lsq_dim_unk] = lsq_weights_c[jc, js] * z_dist_g[js, :] @@ -107,8 +107,8 @@ def lsq_compute_coeffs( min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) for jc in range(start_idx, min_rlcell_int): - # if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 case base_grid.GeometryType.TORUS: for jc in range(start_idx, min_rlcell_int): ilc_s = c2e2c[jc, :lsq_dim_stencil] diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index 09ab1b246d..c4b70028d6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -466,7 +466,6 @@ def run( horizontal_end=horizontal_end, ) log.debug("running stencil copy_cell_kdim_field - end") - # self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("vertical advection run - end") @@ -629,7 +628,6 @@ def _compute_numerical_flux( horizontal_start, horizontal_end = self._get_horizontal_start_end( even_timestep=even_timestep ) - #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("running stencil compute_vertical_tracer_flux_upwind - start") self._compute_vertical_tracer_flux_upwind( p_cc=p_tracer_now, @@ -919,7 +917,6 @@ def _compute_numerical_flux( log.debug("running stencil init_constant_cell_kdim_field - end") log.debug("running stencil compute_ppm4gpu_courant_number - start") - #self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) self._compute_ppm4gpu_courant_number( p_mflx_contra_v=prep_adv.mass_flx_ic, p_cellmass_now=rhodz_now, @@ -1102,5 +1099,4 @@ def _update_unknowns( ) self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("running stencil integrate_tracer_vertically - end") - # self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index e5ea34ca20..51ac9e6e60 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -214,7 +214,6 @@ def test_advection_run_single_step( diagnostic_state_ref.vfl_tracer.asnumpy(), rtol=1e-10, ) - #breakpoint() assert test_helpers.dallclose(p_tracer_new_ref.asnumpy(), p_tracer_new.asnumpy(), atol=1e-10) From df7f20cd4a911989734b4265a2836cd85d655013 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 13 Feb 2026 11:06:13 +0100 Subject: [PATCH 24/25] further cleanup and edits --- model/atmosphere/advection/pyproject.toml | 2 - .../advection/advection_lsq_coeffs.py | 166 ------------------ .../integration_tests/test_advection.py | 12 +- .../mpi_tests/test_parallel_advection.py | 21 +-- .../interpolation/interpolation_factory.py | 3 +- .../interpolation/interpolation_fields.py | 153 +++++++++++++++- uv.lock | 4 - 7 files changed, 164 insertions(+), 197 deletions(-) delete mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml index b55044c165..86aaa63dec 100644 --- a/model/atmosphere/advection/pyproject.toml +++ b/model/atmosphere/advection/pyproject.toml @@ -27,8 +27,6 @@ dependencies = [ "icon4py-common>=0.0.6", # external dependencies "gt4py==1.1.3", - 'numpy>=1.23.3', - 'scipy>=1.14.1', 'packaging>=20.0' ] description = "ICON advection." diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py deleted file mode 100644 index d1bc4ad2d5..0000000000 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ /dev/null @@ -1,166 +0,0 @@ -# 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 -from collections.abc import Callable -from types import ModuleType - -import numpy as np -import scipy - -from icon4py.model.common.dimension import CellDim -from icon4py.model.common.grid import base as base_grid -from icon4py.model.common.math.projection import diff_on_edges_torus_numpy, gnomonic_proj -from icon4py.model.common.utils import data_allocation as data_alloc - - -def compute_lsq_pseudoinv( - cell_owner_mask: data_alloc.NDArray, - lsq_pseudoinv: data_alloc.NDArray, - z_lsq_mat_c: data_alloc.NDArray, - lsq_weights_c: data_alloc.NDArray, - start_idx: int, - min_rlcell_int: int, - lsq_dim_unk: int, - lsq_dim_c: int, -) -> data_alloc.NDArray: - for jjb in range(lsq_dim_c): - for jjk in range(lsq_dim_unk): - for jc in range(start_idx, min_rlcell_int): - u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :]) - if cell_owner_mask[jc]: - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] - + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] - ) - return lsq_pseudoinv - - -def compute_lsq_weights_c( - z_dist_g: data_alloc.NDArray, - lsq_weights_c_jc: data_alloc.NDArray, - lsq_dim_stencil: int, - lsq_wgt_exp: int, -) -> data_alloc.NDArray: - for js in range(lsq_dim_stencil): - z_norm = np.sqrt(np.dot(z_dist_g[js, :], z_dist_g[js, :])) - lsq_weights_c_jc[js] = 1.0 / (z_norm**lsq_wgt_exp) - return lsq_weights_c_jc / np.max(lsq_weights_c_jc) - - -def compute_z_lsq_mat_c( - cell_owner_mask: data_alloc.NDArray, - z_lsq_mat_c: data_alloc.NDArray, - lsq_weights_c: data_alloc.NDArray, - z_dist_g: data_alloc.NDArray, - jc: int, - lsq_dim_unk: int, - lsq_dim_c: int, -) -> data_alloc.NDArray: - min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 - - for js in range(lsq_dim_c): - z_lsq_mat_c[jc, js, :lsq_dim_unk] = lsq_weights_c[jc, js] * z_dist_g[js, :] - - return z_lsq_mat_c[jc, js, :lsq_dim_unk] - - -def lsq_compute_coeffs( - cell_center_x: data_alloc.NDArray, - cell_center_y: data_alloc.NDArray, - cell_lat: data_alloc.NDArray, - cell_lon: data_alloc.NDArray, - c2e2c: data_alloc.NDArray, - cell_owner_mask: data_alloc.NDArray, - domain_length: float, - domain_height: float, - grid_sphere_radius: float, - lsq_dim_unk: int, - lsq_dim_c: int, - lsq_wgt_exp: int, - lsq_dim_stencil: int, - start_idx: int, - min_rlcell_int: int, - geometry_type: int, - exchange: Callable[[data_alloc.NDArray], None], - array_ns: ModuleType = np, -) -> data_alloc.NDArray: - lsq_weights_c = array_ns.zeros((min_rlcell_int, lsq_dim_stencil)) - lsq_pseudoinv = array_ns.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) - z_lsq_mat_c = array_ns.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) - z_dist_g = array_ns.zeros((min_rlcell_int, lsq_dim_c, 2)) - match base_grid.GeometryType(geometry_type): - case base_grid.GeometryType.ICOSAHEDRON: - for js in range(lsq_dim_stencil): - z_dist_g[:, js, :] = np.asarray( - gnomonic_proj( - cell_lon[:], cell_lat[:], cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]] - ) - ).T - - z_dist_g *= grid_sphere_radius - min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - - for jc in range(start_idx, min_rlcell_int): - if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 - case base_grid.GeometryType.TORUS: - for jc in range(start_idx, min_rlcell_int): - ilc_s = c2e2c[jc, :lsq_dim_stencil] - cc_cell = array_ns.zeros((lsq_dim_stencil, 2)) - - cc_cv = (cell_center_x[jc], cell_center_y[jc]) - for js in range(lsq_dim_stencil): - cc_cell[js, :] = diff_on_edges_torus_numpy( - cell_center_x[jc], - cell_center_y[jc], - cell_center_x[ilc_s][js], - cell_center_y[ilc_s][js], - domain_length, - domain_height, - ) - z_dist_g[jc, :, :] = cc_cell - cc_cv - - for jc in range(start_idx, min_rlcell_int): - lsq_weights_c[jc, :] = compute_lsq_weights_c( - z_dist_g[jc, :, :], lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp - ) - z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( - cell_owner_mask, - z_lsq_mat_c, - lsq_weights_c, - z_dist_g[jc, :, :], - jc, - lsq_dim_unk, - lsq_dim_c, - ) - - try: - exchange(lsq_weights_c, dim=CellDim) - except TypeError: - exchange(lsq_weights_c) - - lsq_pseudoinv = compute_lsq_pseudoinv( - cell_owner_mask, - lsq_pseudoinv, - z_lsq_mat_c, - lsq_weights_c, - start_idx, - min_rlcell_int, - lsq_dim_unk, - lsq_dim_c, - ) - - try: - exchange(lsq_pseudoinv[:, 0, :], dim=CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=CellDim) - except TypeError: - exchange(lsq_pseudoinv[:, 0, :]) - exchange(lsq_pseudoinv[:, 1, :]) - - return lsq_pseudoinv diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index 25bb1009e0..e8ee1804ac 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -10,14 +10,14 @@ import pytest import icon4py.model.testing.test_utils as test_helpers -from icon4py.model.atmosphere.advection import advection, advection_states -from icon4py.model.atmosphere.advection.advection_lsq_coeffs import lsq_compute_coeffs -from icon4py.model.common import constants, dimension as dims, utils +from icon4py.model.atmosphere.advection import advection +from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import ( base as base_grid, geometry_attributes as geometry_attrs, horizontal as h_grid, ) +from icon4py.model.common.interpolation.interpolation_fields import compute_lsq_coeffs from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import ( definitions, @@ -142,7 +142,7 @@ def test_advection_run_single_step( ) interpolation_state = construct_interpolation_state(interpolation_savepoint, backend=backend) geometry = gridtest_utils.get_grid_geometry(backend, experiment) - least_squares_coeffs = lsq_compute_coeffs( + least_squares_coeffs = compute_lsq_coeffs( cell_center_x=geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy(), cell_center_y=geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy(), cell_lat=geometry.get(geometry_attrs.CELL_LAT).asnumpy(), @@ -218,7 +218,7 @@ def test_advection_run_single_step( @pytest.mark.level("unit") @pytest.mark.datatest -def test_lsq_compute_coeffs( +def test_compute_lsq_coeffs( icon_grid: base_grid.Grid, grid_savepoint: sb.IconGridSavepoint, backend: gtx_typing.Backend, @@ -253,7 +253,7 @@ def test_lsq_compute_coeffs( coordinates = gm.coordinates cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() - lsq_pseudoinv = lsq_compute_coeffs( + lsq_pseudoinv = compute_lsq_coeffs( cell_center_x, cell_center_y, cell_lat, diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py index 51ac9e6e60..3b36a57891 100644 --- a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -6,33 +6,24 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import gt4py.next.typing as gtx_typing -import numpy as np import pytest from gt4py.next import typing as gtx_typing import icon4py.model.testing.test_utils as test_helpers -from icon4py.model.atmosphere.advection import advection, advection_states -from icon4py.model.atmosphere.advection.advection_lsq_coeffs import lsq_compute_coeffs -from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as nh -from icon4py.model.common import constants, dimension as dims, type_alias as ta, utils +from icon4py.model.atmosphere.advection import advection +from icon4py.model.common import constants, dimension as dims from icon4py.model.common.decomposition import definitions, mpi_decomposition from icon4py.model.common.grid import ( base as base_grid, geometry_attributes as geometry_attrs, horizontal as h_grid, - icon, - states as grid_states, - vertical as v_grid, ) -from icon4py.model.common.interpolation import interpolation_attributes as intp_attrs +from icon4py.model.common.interpolation.interpolation_fields import compute_lsq_coeffs from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import ( definitions as test_defs, grid_utils, - grid_utils as gridtest_utils, parallel_helpers, - serialbox, serialbox as sb, test_utils, ) @@ -49,9 +40,7 @@ processor_props, ) -from .. import utils from ..fixtures import * # noqa: F403 -from ..fixtures import advection_exit_savepoint, advection_init_savepoint, advection_lsq_state from ..utils import ( construct_config, construct_diagnostic_exit_state, @@ -220,7 +209,7 @@ def test_advection_run_single_step( @pytest.mark.level("unit") @pytest.mark.datatest @pytest.mark.mpi -def test_lsq_compute_coeffs( +def test_compute_lsq_coeffs( icon_grid: base_grid.Grid, grid_savepoint: sb.IconGridSavepoint, backend: gtx_typing.Backend, @@ -258,7 +247,7 @@ def test_lsq_compute_coeffs( coordinates = gm.coordinates cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() - lsq_pseudoinv = lsq_compute_coeffs( + lsq_pseudoinv = compute_lsq_coeffs( cell_center_x, cell_center_y, cell_lat, diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index 2e984bc3d1..343ce8a8a1 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -12,7 +12,6 @@ import gt4py.next.typing as gtx_typing import icon4py.model.common.interpolation.stencils.compute_nudgecoeffs as nudgecoeffs -from icon4py.model.atmosphere.advection import advection_lsq_coeffs from icon4py.model.common import constants, dimension as dims from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import ( @@ -239,7 +238,7 @@ def _register_computed_fields(self) -> None: lsq_pseudoinv = factory.NumpyDataProvider( func=functools.partial( - advection_lsq_coeffs.lsq_compute_coeffs, + interpolation_fields.compute_lsq_coeffs, exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), array_ns=self._xp, ), diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index a4e30d4659..704f1ff0f5 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -13,6 +13,7 @@ from typing import Final import numpy as np +import scipy from gt4py import next as gtx from gt4py.next import where @@ -21,8 +22,9 @@ import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims from icon4py.model.common.dimension import C2E, V2E -from icon4py.model.common.grid import gridfile +from icon4py.model.common.grid import base as base_grid, gridfile from icon4py.model.common.grid.geometry_stencils import compute_primal_cart_normal +from icon4py.model.common.math.projection import diff_on_edges_torus_numpy, gnomonic_proj from icon4py.model.common.utils import data_allocation as data_alloc @@ -1149,3 +1151,152 @@ def compute_pos_on_tplane_e_x_y_torus( exchange(pos_on_tplane_e_x, pos_on_tplane_e_y) return pos_on_tplane_e_x, pos_on_tplane_e_y + + +def compute_lsq_pseudoinv( + cell_owner_mask: data_alloc.NDArray, + lsq_pseudoinv: data_alloc.NDArray, + z_lsq_mat_c: data_alloc.NDArray, + lsq_weights_c: data_alloc.NDArray, + start_idx: int, + min_rlcell_int: int, + lsq_dim_unk: int, + lsq_dim_c: int, +) -> data_alloc.NDArray: + for jjb in range(lsq_dim_c): + for jjk in range(lsq_dim_unk): + for jc in range(start_idx, min_rlcell_int): + u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :]) + if cell_owner_mask[jc]: + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] + + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] + ) + return lsq_pseudoinv + + +def compute_lsq_weights_c( + z_dist_g: data_alloc.NDArray, + lsq_weights_c_jc: data_alloc.NDArray, + lsq_dim_stencil: int, + lsq_wgt_exp: int, +) -> data_alloc.NDArray: + for js in range(lsq_dim_stencil): + z_norm = np.sqrt(np.dot(z_dist_g[js, :], z_dist_g[js, :])) + lsq_weights_c_jc[js] = 1.0 / (z_norm**lsq_wgt_exp) + return lsq_weights_c_jc / np.max(lsq_weights_c_jc) + + +def compute_z_lsq_mat_c( + cell_owner_mask: data_alloc.NDArray, + z_lsq_mat_c: data_alloc.NDArray, + lsq_weights_c: data_alloc.NDArray, + z_dist_g: data_alloc.NDArray, + jc: int, + lsq_dim_unk: int, + lsq_dim_c: int, +) -> data_alloc.NDArray: + min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) + if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + + for js in range(lsq_dim_c): + z_lsq_mat_c[jc, js, :lsq_dim_unk] = lsq_weights_c[jc, js] * z_dist_g[js, :] + + return z_lsq_mat_c[jc, js, :lsq_dim_unk] + + +def compute_lsq_coeffs( + cell_center_x: data_alloc.NDArray, + cell_center_y: data_alloc.NDArray, + cell_lat: data_alloc.NDArray, + cell_lon: data_alloc.NDArray, + c2e2c: data_alloc.NDArray, + cell_owner_mask: data_alloc.NDArray, + domain_length: float, + domain_height: float, + grid_sphere_radius: float, + lsq_dim_unk: int, + lsq_dim_c: int, + lsq_wgt_exp: int, + lsq_dim_stencil: int, + start_idx: int, + min_rlcell_int: int, + geometry_type: int, + exchange: Callable[[data_alloc.NDArray], None], + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + lsq_weights_c = array_ns.zeros((min_rlcell_int, lsq_dim_stencil)) + lsq_pseudoinv = array_ns.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) + z_lsq_mat_c = array_ns.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) + z_dist_g = array_ns.zeros((min_rlcell_int, lsq_dim_c, 2)) + match base_grid.GeometryType(geometry_type): + case base_grid.GeometryType.ICOSAHEDRON: + for js in range(lsq_dim_stencil): + z_dist_g[:, js, :] = np.asarray( + gnomonic_proj( + cell_lon[:], cell_lat[:], cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]] + ) + ).T + + z_dist_g *= grid_sphere_radius + min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) + + for jc in range(start_idx, min_rlcell_int): + if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + case base_grid.GeometryType.TORUS: + for jc in range(start_idx, min_rlcell_int): + ilc_s = c2e2c[jc, :lsq_dim_stencil] + cc_cell = array_ns.zeros((lsq_dim_stencil, 2)) + + cc_cv = (cell_center_x[jc], cell_center_y[jc]) + for js in range(lsq_dim_stencil): + cc_cell[js, :] = diff_on_edges_torus_numpy( + cell_center_x[jc], + cell_center_y[jc], + cell_center_x[ilc_s][js], + cell_center_y[ilc_s][js], + domain_length, + domain_height, + ) + z_dist_g[jc, :, :] = cc_cell - cc_cv + + for jc in range(start_idx, min_rlcell_int): + lsq_weights_c[jc, :] = compute_lsq_weights_c( + z_dist_g[jc, :, :], lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp + ) + z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( + cell_owner_mask, + z_lsq_mat_c, + lsq_weights_c, + z_dist_g[jc, :, :], + jc, + lsq_dim_unk, + lsq_dim_c, + ) + + try: + exchange(lsq_weights_c, dim=dims.CellDim) + except TypeError: + exchange(lsq_weights_c) + + lsq_pseudoinv = compute_lsq_pseudoinv( + cell_owner_mask, + lsq_pseudoinv, + z_lsq_mat_c, + lsq_weights_c, + start_idx, + min_rlcell_int, + lsq_dim_unk, + lsq_dim_c, + ) + + try: + exchange(lsq_pseudoinv[:, 0, :], dim=dims.CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=dims.CellDim) + except TypeError: + exchange(lsq_pseudoinv[:, 0, :]) + exchange(lsq_pseudoinv[:, 1, :]) + + return lsq_pseudoinv diff --git a/uv.lock b/uv.lock index 5c790b4963..cf518295a9 100644 --- a/uv.lock +++ b/uv.lock @@ -1750,18 +1750,14 @@ source = { editable = "model/atmosphere/advection" } dependencies = [ { name = "gt4py" }, { name = "icon4py-common" }, - { name = "numpy" }, { name = "packaging" }, - { name = "scipy" }, ] [package.metadata] requires-dist = [ { name = "gt4py", specifier = "==1.1.3" }, { name = "icon4py-common", editable = "model/common" }, - { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, - { name = "scipy", specifier = ">=1.14.1" }, ] [[package]] From bb8cd41fe350ab0757c2e9c792411f431e0b9a9e Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Mon, 16 Feb 2026 11:20:59 +0100 Subject: [PATCH 25/25] edits to exchange default value --- .../model/atmosphere/advection/advection.py | 6 +-- .../integration_tests/test_advection.py | 3 -- .../interpolation/interpolation_fields.py | 38 ++++++++----------- 3 files changed, 19 insertions(+), 28 deletions(-) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index a35274acfd..2ee13c0da5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -162,7 +162,7 @@ def __init__( self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None, - exchange: decomposition.ExchangeRuntime | None = None, + exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, ): log.debug("advection class init - start") @@ -221,7 +221,7 @@ def __init__( grid: icon_grid.IconGrid, metric_state: advection_states.AdvectionMetricState, backend: gtx_typing.Backend | None, - exchange: decomposition.ExchangeRuntime | None = None, + exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, even_timestep: bool = False, ): log.debug("advection class init - start") @@ -486,7 +486,7 @@ def convert_config_to_advection( edge_params: grid_states.EdgeParams, cell_params: grid_states.CellParams, backend: gtx_typing.Backend | None, - exchange: decomposition.ExchangeRuntime | None = None, + exchange: decomposition.ExchangeRuntime = decomposition.single_node_default, even_timestep: bool = False, ) -> Advection: exchange = exchange or decomposition.SingleNodeExchange() diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index e8ee1804ac..16931b41ad 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -161,7 +161,6 @@ 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=None, ) least_squares_state = construct_least_squares_state(least_squares_coeffs, backend=backend) @@ -180,7 +179,6 @@ def test_advection_run_single_step( cell_params=cell_geometry, even_timestep=even_timestep, backend=backend, - exchange=None, ) diagnostic_state = construct_diagnostic_init_state( @@ -270,7 +268,6 @@ def test_compute_lsq_coeffs( start_idx, min_rlcell_int, icon_grid.geometry_type, - exchange=None, ) assert test_helpers.dallclose( diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index 704f1ff0f5..ee4f661657 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -21,6 +21,7 @@ import icon4py.model.common.math.projection as proj import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims +from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.dimension import C2E, V2E from icon4py.model.common.grid import base as base_grid, gridfile from icon4py.model.common.grid.geometry_stencils import compute_primal_cart_normal @@ -38,7 +39,7 @@ def compute_c_lin_e( inv_dual_edge_length: data_alloc.NDArray, edge_owner_mask: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -112,7 +113,7 @@ def compute_geofac_n2s( e2c: data_alloc.NDArray, c2e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -173,7 +174,7 @@ def compute_geofac_grg( e2c: data_alloc.NDArray, c2e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: owned = array_ns.stack((owner_mask, owner_mask, owner_mask)).T @@ -220,7 +221,7 @@ def compute_geofac_grdiv( e2c: data_alloc.NDArray, e2c2e: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -441,7 +442,7 @@ def _force_mass_conservation_to_c_bln_avg( cell_owner_mask: data_alloc.NDArray, divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, niter: int = 1000, ) -> data_alloc.NDArray: @@ -619,7 +620,7 @@ def compute_mass_conserving_bilinear_cell_average_weight( divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, horizontal_start_level_3: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: c_bln_avg = _compute_c_bln_avg( @@ -652,7 +653,7 @@ def compute_mass_conserving_bilinear_cell_average_weight_torus( divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, horizontal_start_level_3: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: c_bln_avg = _compute_uniform_c_bln_avg( @@ -723,7 +724,7 @@ def compute_e_flx_avg( e2c2e: data_alloc.NDArray, horizontal_start_p3: gtx.int32, horizontal_start_p4: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -888,7 +889,7 @@ def compute_cells_aw_verts( v2c: data_alloc.NDArray, e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -1021,7 +1022,7 @@ def compute_pos_on_tplane_e_x_y( owner_mask: data_alloc.NDArray, e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: """ @@ -1108,7 +1109,7 @@ def compute_pos_on_tplane_e_x_y( def compute_pos_on_tplane_e_x_y_torus( dual_edge_length: data_alloc.NDArray, e2c: data_alloc.NDArray, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -1223,7 +1224,7 @@ def compute_lsq_coeffs( start_idx: int, min_rlcell_int: int, geometry_type: int, - exchange: Callable[[data_alloc.NDArray], None], + exchange: decomposition.ExchangeRuntime = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: lsq_weights_c = array_ns.zeros((min_rlcell_int, lsq_dim_stencil)) @@ -1276,10 +1277,7 @@ def compute_lsq_coeffs( lsq_dim_c, ) - try: - exchange(lsq_weights_c, dim=dims.CellDim) - except TypeError: - exchange(lsq_weights_c) + exchange(lsq_weights_c, dim=dims.CellDim) lsq_pseudoinv = compute_lsq_pseudoinv( cell_owner_mask, @@ -1292,11 +1290,7 @@ def compute_lsq_coeffs( lsq_dim_c, ) - try: - exchange(lsq_pseudoinv[:, 0, :], dim=dims.CellDim) - exchange(lsq_pseudoinv[:, 1, :], dim=dims.CellDim) - except TypeError: - exchange(lsq_pseudoinv[:, 0, :]) - exchange(lsq_pseudoinv[:, 1, :]) + exchange(lsq_pseudoinv[:, 0, :], dim=dims.CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=dims.CellDim) return lsq_pseudoinv