diff --git a/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py index bdc594a64b..c50885d158 100644 --- a/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion/mpi_tests/test_parallel_diffusion.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import logging from typing import Any import pytest @@ -21,6 +22,12 @@ from ..fixtures import * # noqa: F403 +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + +_log = logging.getLogger(__file__) + + @pytest.mark.mpi @pytest.mark.uses_concat_where @pytest.mark.parametrize( @@ -63,24 +70,24 @@ def test_parallel_diffusion( raise pytest.skip("This test is only executed for `dace` backends.") caplog.set_level("INFO") parallel_helpers.check_comm_size(processor_props) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: initializing diffusion for experiment '{definitions.Experiments.MCH_CH_R04B09}'" ) - print( + _log.info( f"local cells = {decomposition_info.global_index(dims.CellDim, decomposition.DecompositionInfo.EntryType.ALL).shape} " f"local edges = {decomposition_info.global_index(dims.EdgeDim, decomposition.DecompositionInfo.EntryType.ALL).shape} " f"local vertices = {decomposition_info.global_index(dims.VertexDim, decomposition.DecompositionInfo.EntryType.ALL).shape}" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells} Cells, {icon_grid.num_edges} Edges, {icon_grid.num_vertices} Vertices" ) config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps) dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" ) vertical_config = v_grid.VerticalGridConfig( @@ -113,7 +120,7 @@ def test_parallel_diffusion( orchestration=orchestration, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") diagnostic_state = diffusion_states.DiffusionDiagnosticState( hdef_ic=savepoint_diffusion_init.hdef_ic(), @@ -135,7 +142,7 @@ def test_parallel_diffusion( prognostic_state=prognostic_state, dtime=dtime, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") utils.verify_diffusion_fields( config=config, @@ -143,7 +150,7 @@ def test_parallel_diffusion( prognostic_state=prognostic_state, diffusion_savepoint=savepoint_diffusion_exit, ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE" ) @@ -190,19 +197,19 @@ def test_parallel_diffusion_multiple_steps( ###################################################################### caplog.set_level("INFO") parallel_helpers.check_comm_size(processor_props) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: initializing diffusion for experiment '{definitions.Experiments.MCH_CH_R04B09}'" ) - print( + _log.info( f"local cells = {decomposition_info.global_index(dims.CellDim, decomposition.DecompositionInfo.EntryType.ALL).shape} " f"local edges = {decomposition_info.global_index(dims.EdgeDim, decomposition.DecompositionInfo.EntryType.ALL).shape} " f"local vertices = {decomposition_info.global_index(dims.VertexDim, decomposition.DecompositionInfo.EntryType.ALL).shape}" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: using local grid with {icon_grid.num_cells} Cells, {icon_grid.num_edges} Edges, {icon_grid.num_vertices} Vertices" ) cell_geometry = grid_savepoint.construct_cell_geometry() @@ -218,7 +225,7 @@ def test_parallel_diffusion_multiple_steps( config = definitions.construct_diffusion_config(experiment, ndyn_substeps=ndyn_substeps) diffusion_params = diffusion_.DiffusionParams(config) dtime = savepoint_diffusion_init.get_metadata("dtime").get("dtime") - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" ) exchange = decomposition.create_exchange(processor_props, decomposition_info) @@ -245,7 +252,7 @@ def test_parallel_diffusion_multiple_steps( orchestration=False, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") diagnostic_state_dace_non_orch = diffusion_states.DiffusionDiagnosticState( hdef_ic=savepoint_diffusion_init.hdef_ic(), @@ -269,8 +276,8 @@ def test_parallel_diffusion_multiple_steps( prognostic_state=prognostic_state_dace_non_orch, dtime=dtime, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") - print( + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE" ) @@ -296,7 +303,7 @@ def test_parallel_diffusion_multiple_steps( backend=backend, orchestration=True, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") diagnostic_state_dace_orch = diffusion_states.DiffusionDiagnosticState( hdef_ic=savepoint_diffusion_init.hdef_ic(), @@ -320,8 +327,8 @@ def test_parallel_diffusion_multiple_steps( prognostic_state=prognostic_state_dace_orch, dtime=dtime, ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") - print( + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: running diffusion step - using {processor_props.comm_name} with {processor_props.comm_size} nodes - DONE" ) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py index e41ee37062..fbca95c910 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py @@ -32,7 +32,6 @@ from icon4py.model.common.metrics import metrics_attributes, metrics_factory from icon4py.model.common.states import prognostic_state as prognostics from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import grid_utils from icon4py.model.testing.fixtures.benchmark import ( geometry_field_source, interpolation_field_source, @@ -62,8 +61,6 @@ def solve_nonhydro( nonhydro_params = solve_nh.NonHydrostaticParams(config) - decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator) - vertical_config = v_grid.VerticalGridConfig( mesh.num_levels, lowest_layer_thickness=50, @@ -205,11 +202,7 @@ def solve_nonhydro( vertical_params=vertical_grid, edge_geometry=edge_geometry, cell_geometry=cell_geometry, - owner_mask=gtx.as_field( - (dims.CellDim,), - decomposition_info.owner_mask(dims.CellDim), # type: ignore[arg-type] # mypy not take the type of owner_mask - allocator=allocator, - ), + owner_mask=geometry_field_source.get("cell_owner_mask"), backend=backend_like, ) diff --git a/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py index ea57fde4c1..36354e3849 100644 --- a/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/mpi_tests/test_parallel_solve_nonhydro.py @@ -7,6 +7,8 @@ # SPDX-License-Identifier: BSD-3-Clause +import logging + import numpy as np import pytest from gt4py.next import typing as gtx_typing @@ -22,6 +24,9 @@ from ..fixtures import * # noqa: F403 +_log = logging.getLogger(__name__) + + @pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.datatest @pytest.mark.parametrize( @@ -67,25 +72,25 @@ def test_run_solve_nonhydro_single_step( pytest.xfail("ValueError: axes don't match array") parallel_helpers.check_comm_size(processor_props) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: inializing dycore for experiment 'mch_ch_r04_b09_dsl" ) - print( + _log.info( f"local cells = {decomposition_info.global_index(dims.CellDim, definitions.DecompositionInfo.EntryType.ALL).shape} " f"local edges = {decomposition_info.global_index(dims.EdgeDim, definitions.DecompositionInfo.EntryType.ALL).shape} " f"local vertices = {decomposition_info.global_index(dims.VertexDim, definitions.DecompositionInfo.EntryType.ALL).shape}" ) owned_cells = decomposition_info.owner_mask(dims.CellDim) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: GHEX context setup: from {processor_props.comm_name} with {processor_props.comm_size} nodes" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo edges {np.count_nonzero(np.invert( decomposition_info.owner_mask(dims.EdgeDim)))}" ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}" ) @@ -137,7 +142,7 @@ def test_run_solve_nonhydro_single_step( exchange=exchange, ) - print( + _log.info( f"rank={processor_props.rank}/{processor_props.comm_size}: entering : solve_nonhydro.time_step" ) @@ -153,7 +158,7 @@ def test_run_solve_nonhydro_single_step( at_first_substep=(substep_init == 1), at_last_substep=(substep_init == ndyn_substeps), ) - print(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ") + _log.info(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ") expected_theta_v = savepoint_nonhydro_step_final.theta_v_new().asnumpy() calculated_theta_v = prognostic_states.next.theta_v.asnumpy() diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 7696ae33ef..a763d2d841 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -52,6 +52,7 @@ io = [ "netcdf4>=1.6.1", "numpy>=1.23.3", "scikit-learn>=1.4.0", + "pymetis>2022.1", # TODO(halungge): there are failing tests starting from uxarray==2024.4.0: when a data file does not have # fields of a given dimension (eg 'edge') then something in uxarray goes wrong with the dimension # mapping. It is not yet clear whether this is a uxarray bug or on our side. diff --git a/model/common/src/icon4py/model/common/decomposition/decomposer.py b/model/common/src/icon4py/model/common/decomposition/decomposer.py new file mode 100644 index 0000000000..beaf30be14 --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/decomposer.py @@ -0,0 +1,82 @@ +# 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 typing import Protocol, runtime_checkable + +import gt4py.next as gtx + +from icon4py.model.common.utils import data_allocation as data_alloc + + +@runtime_checkable +class Decomposer(Protocol): + def __call__( + self, adjacency_matrix: data_alloc.NDArray, num_partitions: int + ) -> data_alloc.NDArray: + """ + Call the decomposition. + + Args: + adjacency_matrix: cell-to-cell connectivity matrix on the global (undecomposed) grid. In the Icon4py context this C2E2C + n_part: number of nodes + """ + ... + + +class MetisDecomposer(Decomposer): + """ + A simple decomposer using METIS for partitioning a grid topology. + + We use the simple pythonic interface to pymetis: just passing the adjacency matrix, which for ICON is + the full grid C2E2C neigbhor table. + if more control is needed (for example by using weights we need to switch to the C like interface) + https://documen.tician.de/pymetis/functionality.html + """ + + def __call__( + self, adjacency_matrix: data_alloc.NDArray, num_partitions: int + ) -> data_alloc.NDArray: + """ + Generate partition labels for this grid topology using METIS: + https://github.com/KarypisLab/METIS + + This method utilizes the pymetis Python bindings: + https://github.com/inducer/pymetis + + Args: + n_part: int, number of partitions to create + adjacency_matrix: nd array: neighbor table describing of the main dimension object to be distributed: for example cell -> cell neighbors + Returns: data_alloc.NDArray: array with partition label (int, rank number) for each cell + """ + + import pymetis # type: ignore [import-untyped] + + # Invalid indices are not allowed here. Metis will segfault or fail if + # there are any invalid indices in the adjacency matrix. + assert (adjacency_matrix >= 0).all() + + # The partitioning is done on all ranks, and this assumes that the + # partitioning is deterministic. + _, partition_index = pymetis.part_graph(nparts=num_partitions, adjacency=adjacency_matrix) + return data_alloc.array_ns_from_array(adjacency_matrix).array(partition_index) + + +class SingleNodeDecomposer(Decomposer): + def __call__( + self, adjacency_matrix: data_alloc.NDArray, num_partitions: int + ) -> data_alloc.NDArray: + """Dummy decomposer for single node: assigns all cells to rank = 0""" + if num_partitions != 1: + raise ValueError( + f"SingleNodeDecomposer can only be used for num_partitions=1, but got {num_partitions}" + ) + + return data_alloc.array_ns_from_array(adjacency_matrix).zeros( + adjacency_matrix.shape[0], + dtype=gtx.int32, # type: ignore [attr-defined] + ) diff --git a/model/common/src/icon4py/model/common/decomposition/definitions.py b/model/common/src/icon4py/model/common/decomposition/definitions.py index e02cb5a1a7..df0b98fb52 100644 --- a/model/common/src/icon4py/model/common/decomposition/definitions.py +++ b/model/common/src/icon4py/model/common/decomposition/definitions.py @@ -20,7 +20,8 @@ import gt4py.next as gtx import numpy as np -from icon4py.model.common import utils +from icon4py.model.common import dimension as dims, utils +from icon4py.model.common.grid import base from icon4py.model.common.orchestration.halo_exchange import DummyNestedSDFG from icon4py.model.common.states import utils as state_utils from icon4py.model.common.utils import data_allocation as data_alloc @@ -35,6 +36,9 @@ class ProcessProperties(Protocol): comm_name: str comm_size: int + def is_single_rank(self) -> bool: + return self.comm_size == 1 + @dataclasses.dataclass(frozen=True, init=False) class SingleNodeProcessProperties(ProcessProperties): @@ -71,41 +75,32 @@ def __call__(self) -> int: class DecompositionInfo: + def __init__( + self, + ) -> None: + self._global_index: dict[gtx.Dimension, data_alloc.NDArray] = {} + self._halo_levels: dict[gtx.Dimension, data_alloc.NDArray] = {} + self._owner_mask: dict[gtx.Dimension, data_alloc.NDArray] = {} + class EntryType(int, Enum): ALL = 0 OWNED = 1 HALO = 2 @utils.chainable - def with_dimension( - self, dim: gtx.Dimension, global_index: data_alloc.NDArray, owner_mask: data_alloc.NDArray + def set_dimension( + self, + dim: gtx.Dimension, + global_index: data_alloc.NDArray, + owner_mask: data_alloc.NDArray, + halo_levels: data_alloc.NDArray, ) -> None: self._global_index[dim] = global_index self._owner_mask[dim] = owner_mask + self._halo_levels[dim] = halo_levels - def __init__( - self, - num_cells: int | None = None, - num_edges: int | None = None, - num_vertices: int | None = None, - ): - self._global_index: dict[gtx.Dimension, data_alloc.NDArray] = {} - self._owner_mask: dict[gtx.Dimension, data_alloc.NDArray] = {} - self._num_vertices = num_vertices - self._num_cells = num_cells - self._num_edges = num_edges - - @property - def num_cells(self) -> int | None: - return self._num_cells - - @property - def num_edges(self) -> int | None: - return self._num_edges - - @property - def num_vertices(self) -> int | None: - return self._num_vertices + def is_distributed(self) -> bool: + return max(self._halo_levels[dims.CellDim]).item() > DecompositionFlag.OWNED def local_index( self, dim: gtx.Dimension, entry_type: EntryType = EntryType.ALL @@ -125,19 +120,15 @@ def local_index( def _to_local_index(self, dim: gtx.Dimension) -> data_alloc.NDArray: data = self._global_index[dim] assert data.ndim == 1 - if isinstance(data, np.ndarray): - import numpy as xp - else: - import cupy as xp # type: ignore[import-not-found, no-redef] - - xp.arange(data.shape[0]) - return xp.arange(data.shape[0]) + return data_alloc.array_ns_from_array(data).arange(data.shape[0]) def owner_mask(self, dim: gtx.Dimension) -> data_alloc.NDArray: return self._owner_mask[dim] def global_index( - self, dim: gtx.Dimension, entry_type: EntryType = EntryType.ALL + self, + dim: gtx.Dimension, + entry_type: DecompositionInfo.EntryType = EntryType.ALL, ) -> data_alloc.NDArray: match entry_type: case DecompositionInfo.EntryType.ALL: @@ -149,6 +140,24 @@ def global_index( case _: raise NotImplementedError() + def get_horizontal_size(self) -> base.HorizontalGridSize: + return base.HorizontalGridSize( + num_cells=self.global_index(dims.CellDim, self.EntryType.ALL).shape[0], + num_edges=self.global_index(dims.EdgeDim, self.EntryType.ALL).shape[0], + num_vertices=self.global_index(dims.VertexDim, self.EntryType.ALL).shape[0], + ) + + def get_halo_size(self, dim: gtx.Dimension, flag: DecompositionFlag) -> int: + level_mask = self.halo_level_mask(dim, flag) + return data_alloc.array_ns_from_array(level_mask).count_nonzero(level_mask) + + def halo_levels(self, dim: gtx.Dimension) -> data_alloc.NDArray: + return self._halo_levels[dim] + + def halo_level_mask(self, dim: gtx.Dimension, level: DecompositionFlag) -> data_alloc.NDArray: + levels = self._halo_levels[dim] + return data_alloc.array_ns_from_array(levels).where(levels == level.value, True, False) + class ExchangeResult(Protocol): def wait(self) -> None: ... @@ -414,5 +423,36 @@ def create_single_reduction_exchange(props: SingleNodeProcessProperties) -> Redu return SingleNodeReductions() +class DecompositionFlag(int, Enum): + UNDEFINED = -1 + OWNED = 0 + """used for locally owned cells, vertices, edges""" + + FIRST_HALO_LEVEL = 1 + """ + used for: + - cells that share 1 edge with an OWNED cell + - vertices that are on OWNED cell, but not owned + - edges that are on OWNED cell, but not owned + """ + + SECOND_HALO_LEVEL = 2 + """ + used for: + - cells that share one vertex with an OWNED cell + - vertices that are on a cell(FIRST_HALO_LINE) but not on an owned cell + - edges that have _exactly_ one vertex shared with and OWNED Cell + """ + + THIRD_HALO_LEVEL = 3 + """ + This type does not exist in ICON. It denotes the "closing/far" edges of the SECOND_HALO_LINE cells + used for: + - cells (NOT USED) + - vertices (NOT USED) + - edges that are only on the cell(SECOND_HALO_LINE) + """ + + single_node_default = SingleNodeExchange() single_node_reductions = SingleNodeReductions() diff --git a/model/common/src/icon4py/model/common/decomposition/halo.py b/model/common/src/icon4py/model/common/decomposition/halo.py new file mode 100644 index 0000000000..ffad08bd5e --- /dev/null +++ b/model/common/src/icon4py/model/common/decomposition/halo.py @@ -0,0 +1,468 @@ +# 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 functools +from types import ModuleType +from typing import Protocol, runtime_checkable + +import gt4py.next as gtx +import gt4py.next.typing as gtx_typing +import numpy as np + +from icon4py.model.common import dimension as dims, exceptions +from icon4py.model.common.decomposition import definitions as defs +from icon4py.model.common.grid import base, gridfile +from icon4py.model.common.utils import data_allocation as data_alloc + + +@runtime_checkable +class HaloConstructor(Protocol): + """Callable that takes a mapping from cells to ranks""" + + def __call__(self, cell_to_rank: data_alloc.NDArray) -> defs.DecompositionInfo: ... + + +class NoHalos(HaloConstructor): + def __init__( + self, + horizontal_size: base.HorizontalGridSize, + allocator: gtx_typing.FieldBufferAllocationUtil | None = None, + ): + self._size = horizontal_size + self._allocator = allocator + + def __call__(self, cell_to_rank: data_alloc.NDArray) -> defs.DecompositionInfo: + xp = data_alloc.import_array_ns(self._allocator) + create_arrays = functools.partial(self._create_dummy_decomposition_arrays, array_ns=xp) + decomposition_info = defs.DecompositionInfo() + + decomposition_info.set_dimension(dims.EdgeDim, *create_arrays(self._size.num_edges)) + decomposition_info.set_dimension(dims.CellDim, *create_arrays(self._size.num_cells)) + decomposition_info.set_dimension(dims.VertexDim, *create_arrays(self._size.num_vertices)) + return decomposition_info + + @staticmethod + def _create_dummy_decomposition_arrays( + size: int, array_ns: ModuleType = np + ) -> tuple[data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray]: + indices = array_ns.arange(size, dtype=gtx.int32) # type: ignore [attr-defined] + owner_mask = array_ns.full((size,), True, dtype=bool) + halo_levels = array_ns.full((size,), defs.DecompositionFlag.OWNED.value, dtype=gtx.int32) # type: ignore [attr-defined] + return indices, owner_mask, halo_levels + + +class IconLikeHaloConstructor(HaloConstructor): + """Creates necessary halo information for a given rank.""" + + def __init__( + self, + run_properties: defs.ProcessProperties, + connectivities: dict[gtx.FieldOffset | str, data_alloc.NDArray], + allocator: gtx_typing.FieldBufferAllocationUtil | None = None, + ): + """ + + Args: + run_properties: contains information on the communicator and local compute node. + connectivities: connectivity arrays needed to construct the halos + allocator: GT4Py buffer allocator + """ + self._xp = data_alloc.import_array_ns(allocator) + self._props = run_properties + self._connectivities = {self._value(k): v for k, v in connectivities.items()} + self._assert_all_neighbor_tables() + + @staticmethod + def _value(k: gtx.FieldOffset | str) -> str: + return str(k.value) if isinstance(k, gtx.FieldOffset) else k + + def _validate_mapping(self, cell_to_rank_mapping: data_alloc.NDArray) -> None: + # validate the distribution mapping: + num_cells = self._connectivity(dims.C2E2C).shape[0] + expected_shape = (num_cells,) + if cell_to_rank_mapping.shape != expected_shape: + raise exceptions.ValidationError( + "rank_mapping", + f"should have shape {expected_shape} but is {cell_to_rank_mapping.shape}", + ) + + # the decomposition should match the communicator size + if self._xp.max(cell_to_rank_mapping) > self._props.comm_size - 1: + raise exceptions.ValidationError( + "rank_mapping", + f"The distribution assumes more nodes than the current run is scheduled on {self._props} ", + ) + + def _assert_all_neighbor_tables(self) -> None: + # make sure we have all connectivity arrays used in the halo construction + relevant_dimension = [ + dims.C2E2C, + dims.E2C, + dims.C2E, + dims.C2V, + dims.V2C, + dims.V2E, + ] + for d in relevant_dimension: + assert ( + d.value in self._connectivities + ), f"Table for {d} is missing from the neighbor table array." + + def _connectivity(self, offset: gtx.FieldOffset | str) -> data_alloc.NDArray: + try: + return self._connectivities[self._value(offset)] + except KeyError as err: + raise exceptions.MissingConnectivityError( + f"Connectivity for offset {offset} is not available" + ) from err + + def _next_halo_line(self, cells: data_alloc.NDArray) -> data_alloc.NDArray: + """Returns the full-grid indices of the next halo line. + + Args: + cells: 1d array, full-grid indices of cells we want to find the neighbors of + Returns: + next_halo_cells: full-grid indices of the next halo line + """ + assert cells.ndim == 1, "input should be 1d array" + cell_neighbors = self._find_cell_neighbors(cells) + return self._xp.setdiff1d(cell_neighbors, cells, assume_unique=True) + + def _find_neighbors( + self, source_indices: data_alloc.NDArray, offset: gtx.FieldOffset | str + ) -> data_alloc.NDArray: + """Get a flattened list of all (unique) neighbors to a given global index list""" + assert source_indices.ndim == 1 + return self._xp.unique(self._connectivity(offset)[source_indices, :].flatten()) + + def _find_cell_neighbors(self, cells: data_alloc.NDArray) -> data_alloc.NDArray: + """Find all neighboring cells of a list of cells.""" + return self._find_neighbors(cells, dims.C2E2C) + + def _find_edge_neighbors_for_cells(self, cell_line: data_alloc.NDArray) -> data_alloc.NDArray: + return self._find_neighbors(cell_line, dims.C2E) + + def _find_edge_neighbors_for_vertices( + self, vertex_line: data_alloc.NDArray + ) -> data_alloc.NDArray: + return self._find_neighbors(vertex_line, dims.V2E) + + def _find_vertex_neighbors_for_cells(self, cell_line: data_alloc.NDArray) -> data_alloc.NDArray: + return self._find_neighbors(cell_line, dims.C2V) + + def _find_cell_neighbors_for_vertices( + self, vertex_line: data_alloc.NDArray + ) -> data_alloc.NDArray: + return self._find_neighbors(vertex_line, dims.V2C) + + def owned_cells(self, cell_to_rank: data_alloc.NDArray) -> data_alloc.NDArray: + """Returns the full-grid indices of the cells owned by this rank""" + assert cell_to_rank.ndim == 1 + return self._xp.where(cell_to_rank == self._props.rank)[0] + + def _update_owner_mask_by_max_rank_convention( + self, + cell_to_rank: data_alloc.NDArray, + owner_mask: data_alloc.NDArray, + all_indices: data_alloc.NDArray, + indices_on_cutting_line: data_alloc.NDArray, + target_connectivity: data_alloc.NDArray, + ) -> data_alloc.NDArray: + """ + In order to have unique ownership of edges (and vertices) among nodes there needs to be + a convention as to where those elements on the cutting line go: + according to a remark in `mo_decomposition_tools.f90` ICON puts them to the node + with the higher rank. + + Args: + owner_mask: owner mask for the dimension + all_indices: (global) indices of the dimension + indices_on_cutting_line: global indices of the elements on the cutting line + target_connectivity: connectivity matrix mapping the dimension d to cell + Returns: + updated owner mask + """ + updated_owner_mask = owner_mask.copy() + for index in indices_on_cutting_line: + local_index = self._xp.nonzero(all_indices == index)[0][0] + owning_ranks = cell_to_rank[target_connectivity[index]] + assert ( + self._xp.unique(owning_ranks).size > 1 + ), f"rank {self._props.rank}: all neighboring cells are owned by the same rank" + assert ( + self._props.rank in owning_ranks + ), f"rank {self._props.rank}: neither of the neighboring cells: {owning_ranks} is owned by me" + # assign the index to the rank with the higher rank + updated_owner_mask[local_index] = max(owning_ranks) <= self._props.rank + return updated_owner_mask + + def __call__(self, cell_to_rank: data_alloc.NDArray) -> defs.DecompositionInfo: + """ + Constructs the DecompositionInfo for the current rank. + + Args: + cell_to_rank: a mapping of cells to a rank + + The DecompositionInfo object is constructed for all horizontal + dimension starting from the cell distribution. Edges and vertices + are then handled through their connectivity to the distributed + cells. + + This constructs a halo similar to ICON which consists of + **exactly** 2 cell-halo lines. + + | /| /| /| + | / | / | / | + | / | / | cy / | + | / | / cx | / | + | / | / | / cz | rank 0 + |/ |/ |/ | + -----v0---e0----v1---e1----v2---e2----v3---e3-- cutting line + | /| /| /| + | c0 / | c1 / | c2 / | + | / | / | / | rank 1 + e4 e5 e6 e7 e8 e9 e10 + | / | / | / | / + | / c3 | / c4 | / c5 | / + |/ |/ |/ |/ + -v4---e11---v5---e12---v6----e13--v7--------- + | /| /| /| + | / | / | / | + | / | / | / | + | / | / | / | + | / | / | / | / + |/ |/ |/ |/ + + + Cells: + + - The "numbered" cells and edges are relevant for the halo + construction from the point of view of rank 0. + - Cells (c0, c1, c2) are the 1. HALO LEVEL: these are cells that + are neighbors of an owned cell. + - Cells (c3, c4. c5) are the 2. HALO LEVEL: cells that "close" the + hexagon of a vertex on the cutting line and do not share an edge + with an owned cell (that is are not LEVEL 1 cells). This is + _not_ the same as the neighboring cells of LEVEL 1 cells, and the + definition might be different from ICON. In the above picture if + the cut was along a corner (e0 -> e1 -> e7) then cz is part of + the 2. LEVEL because it is part of the hexagon on v2, but it is + not in c2e2c2e2c(c1). + + Note that this definition of 1. and 2. line differs from the + definition of boundary line counting used in [grid + refinement](grid_refinement.py), in terms of "distance" to the + cutting line all halo cells have a distance of 1. + + + Vertices: + + - 1. HALO LEVEL: are vertices on the cutting line that are not + owned, or in a different wording: all vertices of owned cells + that are not owned. + - In ICON every element in an array needs **exactly one owner** + (otherwise there would be duplicates and double-counting). For + elements on the cutting line (vertices and edges) there is no + clear indication which rank should own it, ICON uses the rank + with the higher rank value (see + (_update_owner_mask_by_max_rank_convention)) In the example above + (v0, v1, v2, v3) are in the 1. HALO LEVEL of rank 0 and owned by + rank 1. As a consequence, there are ranks that have no 1. HALO + LEVEL vertices. + - 2. HALO LEVEL: are vertices that are on HALO LEVEL cells, but not + on owned. For rank 0 these are (v4, v5, v6, v7). + + + Edges: + + - For edges a similar pattern is used as for the vertices. + - 1. HALO LEVEL: edges that are on owned cells but not owned + themselves (these are edges that share 2 vertices with a owned + cell). In terms of ownership the same convention is applied as + for the vertices: (e0, e1, e2, e3) are in the HALO LEVEL 1 of + rank 0, and are owned by rank 1. + - 2. HALO LEVEL: edges that share exactly one vertex with an owned + cell. The definition via vertices is important: see example above + with a cut along e0 -> e1 -> e7. For rank 0 these are the edges + (e4, e5, e6, e7, e8, e9, e10) in the example above. + - 3. HALO LEVEL: In **ICON4Py ONLY**, edges that "close" the halo + cells and share exactly 2 vertices with a HALO LEVEL 2 cell, but + none with an owned cell. These edges are **not** included in the + halo in ICON. These are (e11, e12, e13) for rank 0 in the example + above. This is the HALO LINE which makes the C2E connectivity + complete (= without skip value) for a distributed setup. + """ + + self._validate_mapping(cell_to_rank) + + #: cells + owned_cells = self.owned_cells(cell_to_rank) + first_halo_cells = self._next_halo_line(owned_cells) + #: vertices + vertex_on_owned_cells = self._find_vertex_neighbors_for_cells(owned_cells) + vertex_on_halo_cells = self._find_vertex_neighbors_for_cells( + self._xp.hstack( + ( + first_halo_cells, + (self._next_halo_line(self._xp.union1d(first_halo_cells, owned_cells))), + ) + ) + ) + vertex_on_cutting_line = self._xp.intersect1d(vertex_on_owned_cells, vertex_on_halo_cells) + vertex_second_halo = self._xp.setdiff1d(vertex_on_halo_cells, vertex_on_cutting_line) + all_vertices = self._xp.hstack((vertex_on_owned_cells, vertex_second_halo)) + + #: update cells to include all cells of the "dual cell" (hexagon) for nodes on the cutting line + dual_cells = self._find_cell_neighbors_for_vertices(vertex_on_cutting_line) + total_halo_cells = self._xp.setdiff1d(dual_cells, owned_cells) + second_halo_cells = self._xp.setdiff1d(total_halo_cells, first_halo_cells) + all_cells = self._xp.hstack((owned_cells, first_halo_cells, second_halo_cells)) + + #: edges + edge_on_owned_cells = self._find_edge_neighbors_for_cells(owned_cells) + edge_on_any_halo_line = self._find_edge_neighbors_for_cells(total_halo_cells) + + edge_on_cutting_line = self._xp.intersect1d(edge_on_owned_cells, edge_on_any_halo_line) + + # needs to be defined as vertex neighbor due to "corners" in the cut. + edge_second_level = self._xp.setdiff1d( + self._find_edge_neighbors_for_vertices(vertex_on_cutting_line), edge_on_owned_cells + ) + edge_third_level = self._xp.setdiff1d(edge_on_any_halo_line, edge_second_level) + edge_third_level = self._xp.setdiff1d(edge_third_level, edge_on_cutting_line) + + all_edges = self._xp.hstack((edge_on_owned_cells, edge_second_level, edge_third_level)) + #: construct decomposition info + decomp_info = defs.DecompositionInfo() + cell_owner_mask = self._xp.isin(all_cells, owned_cells) + cell_halo_levels = self._xp.full( + all_cells.size, + defs.DecompositionFlag.UNDEFINED.value, + dtype=gtx.int32, # type: ignore [attr-defined] + ) + cell_halo_levels[cell_owner_mask] = defs.DecompositionFlag.OWNED + cell_halo_levels[self._xp.isin(all_cells, first_halo_cells)] = ( + defs.DecompositionFlag.FIRST_HALO_LEVEL + ) + cell_halo_levels[self._xp.isin(all_cells, second_halo_cells)] = ( + defs.DecompositionFlag.SECOND_HALO_LEVEL + ) + decomp_info.set_dimension(dims.CellDim, all_cells, cell_owner_mask, cell_halo_levels) + vertex_owner_mask = self._xp.isin(all_vertices, vertex_on_owned_cells) + vertex_owner_mask = self._update_owner_mask_by_max_rank_convention( + cell_to_rank, + vertex_owner_mask, + all_vertices, + vertex_on_cutting_line, + self._connectivity(dims.V2C), + ) + vertex_second_level = self._xp.setdiff1d(vertex_on_halo_cells, vertex_on_owned_cells) + vertex_halo_levels = self._xp.full( + all_vertices.size, + defs.DecompositionFlag.UNDEFINED.value, + dtype=gtx.int32, # type: ignore [attr-defined] + ) + vertex_halo_levels[vertex_owner_mask] = defs.DecompositionFlag.OWNED + vertex_halo_levels[ + self._xp.logical_not(vertex_owner_mask) + & self._xp.isin(all_vertices, vertex_on_cutting_line) + ] = defs.DecompositionFlag.FIRST_HALO_LEVEL + + vertex_halo_levels[self._xp.isin(all_vertices, vertex_second_level)] = ( + defs.DecompositionFlag.SECOND_HALO_LEVEL + ) + decomp_info.set_dimension( + dims.VertexDim, all_vertices, vertex_owner_mask, vertex_halo_levels + ) + + edge_owner_mask = self._xp.isin(all_edges, edge_on_owned_cells) + edge_owner_mask = self._update_owner_mask_by_max_rank_convention( + cell_to_rank, + edge_owner_mask, + all_edges, + edge_on_cutting_line, + self._connectivity(dims.E2C), + ) + + edge_halo_levels = self._xp.full( + all_edges.shape, + defs.DecompositionFlag.UNDEFINED.value, + dtype=gtx.int32, # type: ignore [attr-defined] + ) + edge_halo_levels[edge_owner_mask] = defs.DecompositionFlag.OWNED + + # LEVEL_ONE edges are on an owned cell but are not owned: these are all edges on the cutting line that are not owned (by the convention) + edge_halo_levels[ + self._xp.logical_not(edge_owner_mask) & self._xp.isin(all_edges, edge_on_cutting_line) + ] = defs.DecompositionFlag.FIRST_HALO_LEVEL + + # LEVEL_TWO edges share exactly one vertex with an owned cell, they are on the first halo-line cells, but not on the cutting line + edge_halo_levels[self._xp.isin(all_edges, edge_second_level)] = ( + defs.DecompositionFlag.SECOND_HALO_LEVEL + ) + + # LEVEL_THREE edges are the "closing" edges of the second halo line, i.e. share no vertex with owned cells + edge_halo_levels[self._xp.isin(all_edges, edge_third_level)] = ( + defs.DecompositionFlag.THIRD_HALO_LEVEL + ) + decomp_info.set_dimension(dims.EdgeDim, all_edges, edge_owner_mask, edge_halo_levels) + return decomp_info + + +def get_halo_constructor( + run_properties: defs.ProcessProperties, + full_grid_size: base.HorizontalGridSize, + connectivities: dict[gtx.FieldOffset | str, data_alloc.NDArray], + allocator: gtx_typing.FieldBufferAllocationUtil | None, +) -> HaloConstructor: + """ + Factory method to create the halo constructor. + + Currently there is only one halo type (except for single node dummy). + If in the future we want to experiment with different halo types we should add an extra selection + parameter + Args: + processor_props: + full_grid_size + allocator: + connectivities: + + Returns: a HaloConstructor suitable for the run_properties + + """ + if run_properties.is_single_rank(): + return NoHalos( + horizontal_size=full_grid_size, + allocator=allocator, + ) + + return IconLikeHaloConstructor( + run_properties=run_properties, + connectivities=connectivities, + allocator=allocator, + ) + + +def global_to_local( + global_indices: data_alloc.NDArray, + indices_to_translate: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + """Translate an array of global indices into rank-local ones. + + Args: + global_indices: global indices owned on the rank: this is the implicit mapping encoding the local to global + indices_to_translate: the array to map to local indices + + """ + sorter = array_ns.argsort(global_indices) + + mask = array_ns.isin(indices_to_translate, global_indices) + positions = array_ns.searchsorted(global_indices, indices_to_translate, sorter=sorter) + local_neighbors = array_ns.full_like(indices_to_translate, gridfile.GridFile.INVALID_INDEX) + local_neighbors[mask] = sorter[positions[mask]] + return local_neighbors diff --git a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py index 7d77432a8f..a8715d229e 100644 --- a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py +++ b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass from types import ModuleType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Union +from typing import Any, ClassVar, Final, Union import dace # type: ignore[import-untyped] import numpy as np @@ -30,7 +30,7 @@ try: import ghex # type: ignore [import-not-found] - import mpi4py # type: ignore [import-not-found] + import mpi4py from ghex.context import make_context # type: ignore [import-not-found] from ghex.unstructured import ( # type: ignore [import-not-found] DomainDescriptor, @@ -45,13 +45,10 @@ mpi4py.rc.finalize = True except ImportError: - mpi4py = None + mpi4py = None # type: ignore [assignment] ghex = None unstructured = None -if TYPE_CHECKING: - import mpi4py.MPI # type: ignore [import-not-found] - CommId = Union[int, "mpi4py.MPI.Comm", None] log = logging.getLogger(__name__) @@ -109,7 +106,7 @@ def get_multinode_properties( @dataclass(frozen=True) class MPICommProcessProperties(definitions.ProcessProperties): - comm: mpi4py.MPI.Comm = None + comm: mpi4py.MPI.Comm @functools.cached_property def rank(self) -> int: # type: ignore [override] @@ -141,6 +138,10 @@ def __init__( dim: self._create_domain_descriptor(dim) for dim in dims.MAIN_HORIZONTAL_DIMENSIONS.values() } + self._field_size: dict[gtx.Dimension, int] = { + dim: self._decomposition_info.global_index(dim).shape[0] + for dim in dims.MAIN_HORIZONTAL_DIMENSIONS.values() + } log.info(f"domain descriptors for dimensions {self._domain_descriptors.keys()} initialized") self._patterns = { dim: self._create_pattern(dim) for dim in dims.MAIN_HORIZONTAL_DIMENSIONS.values() @@ -209,12 +210,8 @@ def _slice_field_based_on_dim(self, field: gtx.Field, dim: gtx.Dimension) -> dat This operation is *necessary* for the use inside FORTRAN as there fields are larger than the grid (nproma size). where it does not do anything in a purely Python setup. the granule context where fields otherwise have length nproma. """ - if dim == dims.VertexDim: - return field.ndarray[: self._decomposition_info.num_vertices] - elif dim == dims.EdgeDim: - return field.ndarray[: self._decomposition_info.num_edges] - elif dim == dims.CellDim: - return field.ndarray[: self._decomposition_info.num_cells] + if dim in dims.MAIN_HORIZONTAL_DIMENSIONS.values(): + return field.ndarray[: self._field_size[dim]] else: raise ValueError(f"Unknown dimension {dim}") diff --git a/model/common/src/icon4py/model/common/exceptions.py b/model/common/src/icon4py/model/common/exceptions.py index 3ade09da08..75f79dc7f1 100644 --- a/model/common/src/icon4py/model/common/exceptions.py +++ b/model/common/src/icon4py/model/common/exceptions.py @@ -16,9 +16,18 @@ class InvalidConfigError(Exception): class IncompleteStateError(Exception): - def __init__(self, field_name): + def __init__(self, field_name: str): super().__init__(f"Field '{field_name}' is missing.") +class ValidationError(Exception): + def __init__(self, name: str, msg: str): + super().__init__(f"'{name}': {msg}.") + + class IconGridError(RuntimeError): pass + + +class MissingConnectivityError(ValueError): + pass diff --git a/model/common/src/icon4py/model/common/grid/base.py b/model/common/src/icon4py/model/common/grid/base.py index b78c3efed5..e011295b18 100644 --- a/model/common/src/icon4py/model/common/grid/base.py +++ b/model/common/src/icon4py/model/common/grid/base.py @@ -15,7 +15,7 @@ import gt4py.next as gtx from gt4py.next import allocators as gtx_allocators, common as gtx_common -from icon4py.model.common import dimension as dims +from icon4py.model.common import dimension as dims, exceptions from icon4py.model.common.grid import horizontal as h_grid from icon4py.model.common.grid.gridfile import GridFile from icon4py.model.common.utils import data_allocation as data_alloc @@ -24,10 +24,6 @@ _log = logging.getLogger(__name__) -class MissingConnectivity(ValueError): - pass - - class GeometryType(enum.Enum): """Define geometries of the horizontal domain supported by the ICON grid. @@ -44,6 +40,9 @@ class HorizontalGridSize: num_edges: int num_cells: int + def __repr__(self): + return f"{self.__class__} (, , gtx_common.Neighbor if isinstance(offset, gtx.FieldOffset): offset = offset.value if offset not in self.connectivities: - raise MissingConnectivity( + raise exceptions.MissingConnectivityError( f"Missing connectivity for offset {offset} in grid {self.id}." ) connectivity = self.connectivities[offset] diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index 8f45803b31..d4fde35f3c 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -9,14 +9,19 @@ import logging import pathlib from types import ModuleType -from typing import Literal, Protocol, TypeAlias +from typing import Literal, TypeAlias import gt4py.next as gtx import gt4py.next.typing as gtx_typing import numpy as np from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.decomposition import definitions as decomposition +from icon4py.model.common.decomposition import ( + decomposer as decomp, + definitions as decomposition, + halo, +) +from icon4py.model.common.exceptions import InvalidConfigError from icon4py.model.common.grid import ( base, grid_refinement as refinement, @@ -28,41 +33,15 @@ _log = logging.getLogger(__name__) +_single_node_decomposer = decomp.SingleNodeDecomposer() +_single_process_props = decomposition.SingleNodeProcessProperties() +_fortran_to_python_transformer = gridfile.ToZeroBasedIndexTransformation() class IconGridError(RuntimeError): pass -class IndexTransformation(Protocol): - """Return a transformation field to be applied to index fields""" - - def __call__( - self, - array: data_alloc.NDArray, - ) -> data_alloc.NDArray: ... - - -class NoTransformation(IndexTransformation): - """Empty implementation of the Protocol. Just return zeros.""" - - def __call__(self, array: data_alloc.NDArray): - return np.zeros_like(array) - - -class ToZeroBasedIndexTransformation(IndexTransformation): - def __call__(self, array: data_alloc.NDArray): - """ - Calculate the index offset needed for usage with python. - - Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for - INVALID values which are marked with -1 in the grid file and are kept such. - """ - return np.asarray( - np.where(array == gridfile.GridFile.INVALID_INDEX, 0, -1), dtype=gtx.int32 - ) - - CoordinateDict: TypeAlias = dict[ gtx.Dimension, dict[Literal["lat", "lon", "x", "y", "z"], gtx.Field] ] @@ -84,24 +63,25 @@ class GridManager: def __init__( self, - transformation: IndexTransformation, grid_file: pathlib.Path | str, - config: v_grid.VerticalGridConfig, # TODO(halungge): remove to separate vertical and horizontal grid + config: v_grid.VerticalGridConfig, # TODO(msimberg): remove to separate vertical and horizontal grid + offset_transformation: gridfile.IndexTransformation = _fortran_to_python_transformer, global_reductions: decomposition.Reductions = decomposition.single_node_reductions, ): - self._transformation = transformation + self._offset_transformation = offset_transformation self._file_name = str(grid_file) self._vertical_config = config + # Output self._grid: icon.IconGrid | None = None self._decomposition_info: decomposition.DecompositionInfo | None = None self._geometry: GeometryDict = {} - self._reader = None self._coordinates: CoordinateDict = {} + self._reader = None self._global_reductions = global_reductions def open(self): """Open the gridfile resource for reading.""" - self._reader = gridfile.GridFile(self._file_name) + self._reader = gridfile.GridFile(self._file_name, self._offset_transformation) self._reader.open() def close(self): @@ -121,7 +101,18 @@ def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is FileNotFoundError: raise FileNotFoundError(f"gridfile {self._file_name} not found, aborting") - def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_values: bool): + def __call__( + self, + allocator: gtx_typing.FieldBufferAllocationUtil | None, + keep_skip_values: bool, + decomposer: decomp.Decomposer = _single_node_decomposer, + run_properties=_single_process_props, + ): + if not run_properties.is_single_rank() and isinstance( + decomposer, decomp.SingleNodeDecomposer + ): + raise InvalidConfigError("Need a Decomposer for multi node run") + if not self._reader: self.open() @@ -130,13 +121,16 @@ def __call__(self, allocator: gtx_typing.FieldBufferAllocationUtil, keep_skip_va else: geometry_type = base.GeometryType.ICOSAHEDRON - self._geometry = self._read_geometry_fields(allocator) - self._grid = self._construct_grid( + self._construct_decomposed_grid( allocator=allocator, - with_skip_values=keep_skip_values, + keep_skip_values=keep_skip_values, geometry_type=geometry_type, + decomposer=decomposer, + run_properties=run_properties, ) self._coordinates = self._read_coordinates(allocator, geometry_type) + self._geometry = self._read_geometry_fields(allocator) + self.close() def _read_coordinates( @@ -144,17 +138,24 @@ def _read_coordinates( allocator: gtx_typing.FieldBufferAllocationUtil, geometry_type: base.GeometryType, ) -> CoordinateDict: + my_cell_indices = self._decomposition_info.global_index(dims.CellDim) + my_edge_indices = self._decomposition_info.global_index(dims.EdgeDim) + my_vertex_indices = self._decomposition_info.global_index(dims.VertexDim) coordinates = { dims.CellDim: { "lat": gtx.as_field( (dims.CellDim,), - self._reader.variable(gridfile.CoordinateName.CELL_LATITUDE), + self._reader.variable( + gridfile.CoordinateName.CELL_LATITUDE, indices=my_cell_indices + ), dtype=ta.wpfloat, allocator=allocator, ), "lon": gtx.as_field( (dims.CellDim,), - self._reader.variable(gridfile.CoordinateName.CELL_LONGITUDE), + self._reader.variable( + gridfile.CoordinateName.CELL_LONGITUDE, indices=my_cell_indices + ), dtype=ta.wpfloat, allocator=allocator, ), @@ -162,13 +163,17 @@ def _read_coordinates( dims.EdgeDim: { "lat": gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.CoordinateName.EDGE_LATITUDE), + self._reader.variable( + gridfile.CoordinateName.EDGE_LATITUDE, indices=my_edge_indices + ), dtype=ta.wpfloat, allocator=allocator, ), "lon": gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.CoordinateName.EDGE_LONGITUDE), + self._reader.variable( + gridfile.CoordinateName.EDGE_LONGITUDE, indices=my_edge_indices + ), dtype=ta.wpfloat, allocator=allocator, ), @@ -176,13 +181,17 @@ def _read_coordinates( dims.VertexDim: { "lat": gtx.as_field( (dims.VertexDim,), - self._reader.variable(gridfile.CoordinateName.VERTEX_LATITUDE), + self._reader.variable( + gridfile.CoordinateName.VERTEX_LATITUDE, indices=my_vertex_indices + ), allocator=allocator, dtype=ta.wpfloat, ), "lon": gtx.as_field( (dims.VertexDim,), - self._reader.variable(gridfile.CoordinateName.VERTEX_LONGITUDE), + self._reader.variable( + gridfile.CoordinateName.VERTEX_LONGITUDE, indices=my_vertex_indices + ), allocator=allocator, dtype=ta.wpfloat, ), @@ -192,55 +201,55 @@ def _read_coordinates( if geometry_type == base.GeometryType.TORUS: coordinates[dims.CellDim]["x"] = gtx.as_field( (dims.CellDim,), - self._reader.variable(gridfile.CoordinateName.CELL_X), + self._reader.variable(gridfile.CoordinateName.CELL_X, indices=my_cell_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.CellDim]["y"] = gtx.as_field( (dims.CellDim,), - self._reader.variable(gridfile.CoordinateName.CELL_Y), + self._reader.variable(gridfile.CoordinateName.CELL_Y, indices=my_cell_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.CellDim]["z"] = gtx.as_field( (dims.CellDim,), - self._reader.variable(gridfile.CoordinateName.CELL_Z), + self._reader.variable(gridfile.CoordinateName.CELL_Z, indices=my_cell_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.EdgeDim]["x"] = gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.CoordinateName.EDGE_X), + self._reader.variable(gridfile.CoordinateName.EDGE_X, indices=my_edge_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.EdgeDim]["y"] = gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.CoordinateName.EDGE_Y), + self._reader.variable(gridfile.CoordinateName.EDGE_Y, indices=my_edge_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.EdgeDim]["z"] = gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.CoordinateName.EDGE_Z), + self._reader.variable(gridfile.CoordinateName.EDGE_Z, indices=my_edge_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.VertexDim]["x"] = gtx.as_field( (dims.VertexDim,), - self._reader.variable(gridfile.CoordinateName.VERTEX_X), + self._reader.variable(gridfile.CoordinateName.VERTEX_X, indices=my_vertex_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.VertexDim]["y"] = gtx.as_field( (dims.VertexDim,), - self._reader.variable(gridfile.CoordinateName.VERTEX_Y), + self._reader.variable(gridfile.CoordinateName.VERTEX_Y, indices=my_vertex_indices), dtype=ta.wpfloat, allocator=allocator, ) coordinates[dims.VertexDim]["z"] = gtx.as_field( (dims.VertexDim,), - self._reader.variable(gridfile.CoordinateName.VERTEX_Z), + self._reader.variable(gridfile.CoordinateName.VERTEX_Z, indices=my_vertex_indices), dtype=ta.wpfloat, allocator=allocator, ) @@ -251,57 +260,77 @@ def _read_geometry_fields( self, allocator: gtx_typing.FieldBufferAllocationUtil, ) -> GeometryDict: + my_cell_indices = self._decomposition_info.global_index(dims.CellDim) + my_edge_indices = self._decomposition_info.global_index(dims.EdgeDim) + my_vertex_indices = self._decomposition_info.global_index(dims.VertexDim) return { # TODO(halungge): still needs to ported, values from "our" grid files contains (wrong) values: # based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) . gridfile.GeometryName.CELL_AREA.value: gtx.as_field( (dims.CellDim,), - self._reader.variable(gridfile.GeometryName.CELL_AREA), + self._reader.variable(gridfile.GeometryName.CELL_AREA, indices=my_cell_indices), allocator=allocator, ), # TODO(halungge): easily computed from a neighbor_sum V2C over the cell areas? gridfile.GeometryName.DUAL_AREA.value: gtx.as_field( (dims.VertexDim,), - self._reader.variable(gridfile.GeometryName.DUAL_AREA), + self._reader.variable(gridfile.GeometryName.DUAL_AREA, indices=my_vertex_indices), allocator=allocator, ), gridfile.GeometryName.EDGE_LENGTH.value: gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.GeometryName.EDGE_LENGTH), + self._reader.variable(gridfile.GeometryName.EDGE_LENGTH, indices=my_edge_indices), allocator=allocator, ), gridfile.GeometryName.DUAL_EDGE_LENGTH.value: gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.GeometryName.DUAL_EDGE_LENGTH), + self._reader.variable( + gridfile.GeometryName.DUAL_EDGE_LENGTH, indices=my_edge_indices + ), allocator=allocator, ), gridfile.GeometryName.EDGE_CELL_DISTANCE.value: gtx.as_field( (dims.EdgeDim, dims.E2CDim), - self._reader.variable(gridfile.GeometryName.EDGE_CELL_DISTANCE, transpose=True), + self._reader.variable( + gridfile.GeometryName.EDGE_CELL_DISTANCE, + transpose=True, + indices=my_edge_indices, + ), allocator=allocator, ), gridfile.GeometryName.EDGE_VERTEX_DISTANCE.value: gtx.as_field( (dims.EdgeDim, dims.E2VDim), - self._reader.variable(gridfile.GeometryName.EDGE_VERTEX_DISTANCE, transpose=True), + self._reader.variable( + gridfile.GeometryName.EDGE_VERTEX_DISTANCE, + transpose=True, + indices=my_edge_indices, + ), allocator=allocator, ), # TODO(halungge): recompute from coordinates? field in gridfile contains NaN on boundary edges gridfile.GeometryName.TANGENT_ORIENTATION.value: gtx.as_field( (dims.EdgeDim,), - self._reader.variable(gridfile.GeometryName.TANGENT_ORIENTATION), + self._reader.variable( + gridfile.GeometryName.TANGENT_ORIENTATION, indices=my_edge_indices + ), allocator=allocator, ), gridfile.GeometryName.CELL_NORMAL_ORIENTATION.value: gtx.as_field( (dims.CellDim, dims.C2EDim), - self._reader.int_variable( - gridfile.GeometryName.CELL_NORMAL_ORIENTATION, transpose=True + self._reader.variable( + gridfile.GeometryName.CELL_NORMAL_ORIENTATION, + transpose=True, + indices=my_cell_indices, ), allocator=allocator, ), gridfile.GeometryName.EDGE_ORIENTATION_ON_VERTEX.value: gtx.as_field( (dims.VertexDim, dims.V2EDim), self._reader.int_variable( - gridfile.GeometryName.EDGE_ORIENTATION_ON_VERTEX, transpose=True + gridfile.GeometryName.EDGE_ORIENTATION_ON_VERTEX, + transpose=True, + apply_offset=False, + indices=my_vertex_indices, ), allocator=allocator, ), @@ -309,8 +338,6 @@ def _read_geometry_fields( def _read_grid_refinement_fields( self, - *, - decomposition_info: decomposition.DecompositionInfo | None = None, allocator: gtx_typing.FieldBufferAllocationUtil, ) -> dict[gtx.Dimension, gtx.Field]: """ @@ -333,7 +360,12 @@ def _read_grid_refinement_fields( refinement_control_fields = { dim: gtx.as_field( (dim,), - self._reader.int_variable(name, decomposition_info, transpose=False), + self._reader.int_variable( + name, + indices=self._decomposition_info.global_index(dim), + transpose=False, + apply_offset=False, + ), allocator=allocator, ) for dim, name in refinement_control_names.items() @@ -352,35 +384,127 @@ def geometry_fields(self) -> GeometryDict: def coordinates(self) -> CoordinateDict: return self._coordinates - def _construct_grid( + @property + def decomposition_info(self) -> decomposition.DecompositionInfo: + return self._decomposition_info + + def _construct_decomposed_grid( self, allocator: gtx_typing.FieldBufferAllocationUtil | None, - with_skip_values: bool, + keep_skip_values: bool, geometry_type: base.GeometryType, - ) -> icon.IconGrid: + decomposer: decomp.Decomposer, + run_properties: decomposition.ProcessProperties, + ) -> None: """Construct the grid topology from the icon grid file. - Reads connectivity fields from the grid file and constructs derived connectivities needed in - Icon4py from them. Adds constructed start/end index information to the grid. + Reads connectivity fields from the grid file and constructs derived + connectivities needed in Icon4py from them. Adds constructed start/end + index information to the grid. The grid will be distributed or not based + on run_properties. """ xp = data_alloc.import_array_ns(allocator) - refinement_fields = self._read_grid_refinement_fields(allocator=allocator) - limited_area = refinement.is_limited_area_grid( - refinement_fields[dims.CellDim].ndarray, array_ns=xp + ## FULL GRID PROPERTIES + cell_refinement = xp.asarray( + self._reader.variable(gridfile.GridRefinementName.CONTROL_CELLS) ) + global_size = self._read_full_grid_size() + global_params = self._construct_global_params(allocator, global_size, geometry_type) + limited_area = refinement.is_limited_area_grid(cell_refinement, array_ns=xp) + + if limited_area and not run_properties.is_single_rank(): + raise NotImplementedError("Limited-area grids are not supported in distributed runs") + + cell_to_cell_neighbors = self._get_index_field(gridfile.ConnectivityName.C2E2C, array_ns=xp) + global_neighbor_tables = { + dims.C2E2C: cell_to_cell_neighbors, + dims.C2E: self._get_index_field(gridfile.ConnectivityName.C2E, array_ns=xp), + dims.E2C: self._get_index_field(gridfile.ConnectivityName.E2C, array_ns=xp), + dims.V2E: self._get_index_field(gridfile.ConnectivityName.V2E, array_ns=xp), + dims.V2C: self._get_index_field(gridfile.ConnectivityName.V2C, array_ns=xp), + dims.C2V: self._get_index_field(gridfile.ConnectivityName.C2V, array_ns=xp), + dims.V2E2V: self._get_index_field(gridfile.ConnectivityName.V2E2V, array_ns=xp), + dims.E2V: self._get_index_field(gridfile.ConnectivityName.E2V, array_ns=xp), + } - num_cells = self._reader.dimension(gridfile.DimensionName.CELL_NAME) - num_edges = self._reader.dimension(gridfile.DimensionName.EDGE_NAME) - num_vertices = self._reader.dimension(gridfile.DimensionName.VERTEX_NAME) - uuid_ = self._reader.attribute(gridfile.MandatoryPropertyName.GRID_UUID) + cells_to_rank_mapping = decomposer(cell_to_cell_neighbors, run_properties.comm_size) + # HALO CONSTRUCTION + # TODO(halungge): reduce the set of neighbor tables used in the halo construction + # TODO(halungge): figure out where to do the host to device copies (xp.asarray...) + halo_constructor = halo.get_halo_constructor( + run_properties=run_properties, + full_grid_size=global_size, + connectivities=global_neighbor_tables, + allocator=allocator, + ) + + self._decomposition_info = halo_constructor(cells_to_rank_mapping) + distributed_size = self._decomposition_info.get_horizontal_size() + + neighbor_tables = self._get_local_connectivities(global_neighbor_tables, array_ns=xp) + + # COMPUTE remaining derived connectivities + neighbor_tables.update(_get_derived_connectivities(neighbor_tables, array_ns=xp)) + + refinement_fields = self._read_grid_refinement_fields(allocator) + + domain_bounds_constructor = functools.partial( + refinement.compute_domain_bounds, + refinement_fields=refinement_fields, + decomposition_info=self._decomposition_info, + array_ns=xp, + ) + start_index, end_index = icon.get_start_and_end_index(domain_bounds_constructor) + grid_config = base.GridConfig( + horizontal_config=distributed_size, + vertical_size=self._vertical_config.num_levels, + limited_area=limited_area, + keep_skip_values=keep_skip_values, + ) + + self._grid = icon.icon_grid( + self._reader.attribute(gridfile.MandatoryPropertyName.GRID_UUID), + allocator=allocator, + config=grid_config, + neighbor_tables=neighbor_tables, + start_index=start_index, + end_index=end_index, + global_properties=global_params, + refinement_control=refinement_fields, + ) + + def _get_local_connectivities( + self, + neighbor_tables_global: dict[gtx.FieldOffset, data_alloc.NDArray], + array_ns, + ) -> dict[gtx.FieldOffset, data_alloc.NDArray]: + global_to_local = functools.partial(halo.global_to_local, array_ns=array_ns) + if self.decomposition_info.is_distributed(): + return { + k: global_to_local( + self._decomposition_info.global_index(k.source), + v[self._decomposition_info.global_index(k.target[0])], + ) + for k, v in neighbor_tables_global.items() + } + else: + return neighbor_tables_global + + def _construct_global_params( + self, + allocator: gtx_typing.FieldBufferAllocationUtil, + global_size: base.HorizontalGridSize, + geometry_type: base.GeometryType, + ): + grid_root = self._reader.attribute(gridfile.MandatoryPropertyName.ROOT) + grid_level = self._reader.attribute(gridfile.MandatoryPropertyName.LEVEL) sphere_radius = self._reader.try_attribute(gridfile.MPIMPropertyName.SPHERE_RADIUS) domain_length = self._reader.try_attribute(gridfile.MPIMPropertyName.DOMAIN_LENGTH) domain_height = self._reader.try_attribute(gridfile.MPIMPropertyName.DOMAIN_HEIGHT) - grid_root = self._reader.attribute(gridfile.MandatoryPropertyName.ROOT) - grid_level = self._reader.attribute(gridfile.MandatoryPropertyName.LEVEL) - global_grid_params = icon.GlobalGridParams( + + return icon.GlobalGridParams( grid_shape=icon.GridShape( geometry_type=geometry_type, subdivision=icon.GridSubdivision(root=grid_root, level=grid_level), @@ -388,51 +512,37 @@ def _construct_grid( radius=sphere_radius, domain_length=domain_length, domain_height=domain_height, - num_cells=num_cells, + num_cells=global_size.num_cells, ) - grid_size = base.HorizontalGridSize( - num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells - ) - config = base.GridConfig( - horizontal_config=grid_size, - vertical_size=self._vertical_config.num_levels, - limited_area=limited_area, - keep_skip_values=with_skip_values, - ) + def _read_full_grid_size(self) -> base.HorizontalGridSize: + """ + Read the grid size propertes (cells, edges, vertices) from the grid file. - neighbor_tables = { - dims.C2E2C: xp.asarray(self._get_index_field(gridfile.ConnectivityName.C2E2C)), - dims.C2E: xp.asarray(self._get_index_field(gridfile.ConnectivityName.C2E)), - dims.E2C: xp.asarray(self._get_index_field(gridfile.ConnectivityName.E2C)), - dims.V2E: xp.asarray(self._get_index_field(gridfile.ConnectivityName.V2E)), - dims.E2V: xp.asarray(self._get_index_field(gridfile.ConnectivityName.E2V)), - dims.V2C: xp.asarray(self._get_index_field(gridfile.ConnectivityName.V2C)), - dims.C2V: xp.asarray(self._get_index_field(gridfile.ConnectivityName.C2V)), - dims.V2E2V: xp.asarray(self._get_index_field(gridfile.ConnectivityName.V2E2V)), - } - neighbor_tables.update(_get_derived_connectivities(neighbor_tables, array_ns=xp)) - domain_bounds_constructor = functools.partial( - refinement.compute_domain_bounds, refinement_fields=refinement_fields, array_ns=xp - ) - start_index, end_index = icon.get_start_and_end_index(domain_bounds_constructor) + As the grid file contains the _full_ (non-distributed) grid, these are the sizes of prior to distribution. - return icon.icon_grid( - id_=uuid_, - allocator=allocator, - config=config, - neighbor_tables=neighbor_tables, - start_index=start_index, - end_index=end_index, - global_properties=global_grid_params, - refinement_control=refinement_fields, + """ + num_cells = self._reader.dimension(gridfile.DynamicDimension.CELL_NAME) + num_edges = self._reader.dimension(gridfile.DynamicDimension.EDGE_NAME) + num_vertices = self._reader.dimension(gridfile.DynamicDimension.VERTEX_NAME) + full_grid_size = base.HorizontalGridSize( + num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells ) + return full_grid_size - def _get_index_field(self, field: gridfile.GridFileName, transpose=True, apply_offset=True): - field = self._reader.int_variable(field, transpose=transpose) - if apply_offset: - field = field + self._transformation(field) - return field + def _get_index_field( + self, + field: gridfile.GridFileName, + indices: data_alloc.NDArray | None = None, + transpose=True, + apply_offset=True, + array_ns: ModuleType = np, + ): + return array_ns.asarray( + self._reader.int_variable( + field, indices=indices, transpose=transpose, apply_offset=apply_offset + ) + ) def _get_derived_connectivities( diff --git a/model/common/src/icon4py/model/common/grid/grid_refinement.py b/model/common/src/icon4py/model/common/grid/grid_refinement.py index f279f8a25b..b720f4b8f3 100644 --- a/model/common/src/icon4py/model/common/grid/grid_refinement.py +++ b/model/common/src/icon4py/model/common/grid/grid_refinement.py @@ -13,7 +13,8 @@ from gt4py import next as gtx from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import horizontal as h_grid +from icon4py.model.common.decomposition import definitions as decomposition +from icon4py.model.common.grid import gridfile, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc @@ -34,11 +35,10 @@ """ _log = logging.getLogger(__name__) - -_MAX_BOUNDARY_DISTANCE: Final[dict[gtx.Dimension, int]] = { - dims.CellDim: 14, - dims.EdgeDim: 28, - dims.VertexDim: 14, +_MAX_ORDERED: Final[dict[gtx.Dimension, int]] = { + dims.CellDim: gridfile.FixedSizeDimension.CELL_GRF.size, + dims.EdgeDim: gridfile.FixedSizeDimension.EDGE_GRF.size, + dims.VertexDim: gridfile.FixedSizeDimension.VERTEX_GRF.size, } """ Grid points in the grid refinement fields are labeled with their distance to the lateral boundary. @@ -154,48 +154,135 @@ } +def _refinement_level_placed_with_halo(domain: h_grid.Domain) -> int: + """There is a speciality in the setup of the ICON halos: generally halo points are located at the end of the arrays after all + points owned by a node. This is true for global grids and for a local area model for all points with a + refinement control value larger than (6, 2) for HALO level 1 and (4,1) for HALO_LEVEL_2. + + That is for the local area grid some (but not all!) halo points that are lateral boundary points are placed with the lateral boundary + domains rather than the halo. The reason for this is mysterious (to me) as well as what advantage it might have. + + Disadvantage clearly is that in the LAM case the halos are **not contigous**. + """ + assert domain.zone.is_halo(), "Domain must be a halo Zone." + match domain.dim: + case dims.EdgeDim: + return 6 if domain.zone == h_grid.Zone.HALO else 4 + case dims.CellDim | dims.VertexDim: + return 2 if domain.zone == h_grid.Zone.HALO else 1 + case _: + raise ValueError(f"Invalid domain: {domain}, must be a HALO domain") + + def compute_domain_bounds( - dim: gtx.Dimension, refinement_fields: dict[gtx.Dimension, gtx.Field], array_ns: ModuleType = np + dim: gtx.Dimension, + refinement_fields: dict[gtx.Dimension, gtx.Field], + decomposition_info: decomposition.DecompositionInfo, + array_ns: ModuleType = np, ) -> tuple[dict[h_grid.Domain, gtx.int32], dict[h_grid.Domain, gtx.int32]]: # type: ignore [name-defined] - refinement_ctrl = refinement_fields[dim].ndarray - refinement_ctrl = convert_to_non_nested_refinement_values(refinement_ctrl, dim, array_ns) + """ + Compute the domain bounds (start_index, end_index) based on a grid Domain. + + In a local area model, ICON orders the field arrays according to their distance from the boundary. For each + dimension (cell, vertex, edge) points are "ordered" (moved to the beginning of the array) up to + the values defined in _GRID_REFINEMENT_BOUNDARY_WIDTH. We call these distance a Grid Zone. The `Dimension` + and the `Zone` determine a Grid `Domain`. For a given field values for a domain are located + in a contiguous section of the field array. + + This function can be used to deterine the start_index and end_index of a `Domain`in the arrays. + + For a global model grid all points are unordered. `Zone` that do not exist for a global model grid + return empty domains. + + For distributed grids halo points build their own `Domain`and are located at the end of the field arrays, with the exception of + some points in the lateral boundary as described in (_refinement_level_placed_with_halo) + + Args: + dim: Dimension one of `CellDim`. `VertexDim`, `EdgeDim` + refinement_fields: dict[Dimension, ndarray] containing the refinement_control values for each dimension + decomposition_info: DecompositionInfo needed to determine the HALO `Zone`s + array_ns: numpy or cupy + + """ + assert ( + dim in dims.MAIN_HORIZONTAL_DIMENSIONS.values() + ), f"Dimension must be one of {dims.MAIN_HORIZONTAL_DIMENSIONS.values()}" + refinement_ctrl = convert_to_non_nested_refinement_values( + refinement_fields[dim].ndarray, dim, array_ns + ) + owned = decomposition_info.owner_mask(dim) + halo_level_1 = decomposition_info.halo_level_mask( + dim, decomposition.DecompositionFlag.FIRST_HALO_LEVEL + ) + halo_level_2 = decomposition_info.halo_level_mask( + dim, decomposition.DecompositionFlag.SECOND_HALO_LEVEL + ) - domains = h_grid.get_domains_for_dim(dim) start_indices = {} end_indices = {} - for domain in domains: - start_index = 0 - end_index = refinement_ctrl.shape[0] - my_zone = domain.zone - if ( - my_zone is h_grid.Zone.END or my_zone.is_halo() - ): # TODO(halungge): implement for distributed - start_index = refinement_ctrl.shape[0] - end_index = refinement_ctrl.shape[0] - elif my_zone.is_lateral_boundary(): - found = array_ns.where(refinement_ctrl == my_zone.level)[0] - start_index, end_index = ( - (array_ns.min(found).item(), array_ns.max(found).item() + 1) - if found.size > 0 - else (0, 0) - ) - elif my_zone.is_nudging(): - value = _LAST_BOUNDARY[dim].level + my_zone.level - found = array_ns.where(refinement_ctrl == value)[0] - start_index, end_index = ( - (array_ns.min(found).item(), array_ns.max(found).item() + 1) - if found.size > 0 - else (0, 0) + + end_domain = h_grid.domain(dim)(h_grid.Zone.END) + start_indices[end_domain] = gtx.int32(refinement_ctrl.shape[0]) + end_indices[end_domain] = gtx.int32(refinement_ctrl.shape[0]) + + halo_domains = h_grid.get_halo_domains(dim) + for domain in halo_domains: + my_flag = decomposition.DecompositionFlag(domain.zone.level) + upper_boundary_level_1 = _refinement_level_placed_with_halo( + h_grid.domain(dim)(h_grid.Zone.HALO) + ) + not_lateral_boundary_1 = (refinement_ctrl < 1) | (refinement_ctrl > upper_boundary_level_1) + halo_region_1 = array_ns.where(halo_level_1 & not_lateral_boundary_1)[0] + not_lateral_boundary_2 = (refinement_ctrl < 1) | ( + refinement_ctrl + > _refinement_level_placed_with_halo(h_grid.domain(dim)(h_grid.Zone.HALO_LEVEL_2)) + ) + + halo_region_2 = array_ns.where(halo_level_2 & not_lateral_boundary_2)[0] + start_halo_2, end_halo_2 = ( + (array_ns.min(halo_region_2).item(), array_ns.max(halo_region_2).item() + 1) + if halo_region_2.size > 0 + else (refinement_ctrl.size, refinement_ctrl.size) + ) + if my_flag == h_grid.Zone.HALO.level: + start_index = ( + array_ns.min(halo_region_1).item() if halo_region_1.size > 0 else start_halo_2 ) - elif my_zone is h_grid.Zone.INTERIOR: - # for the Vertex and Edges the level after the nudging zones are not ordered anymore, so - # we rely on using the end index of the nudging zone for INTERIOR - value = get_nudging_refinement_value(dim) - found = array_ns.where(refinement_ctrl == value)[0] - start_index = array_ns.max(found).item() + 1 if found.size > 0 else 0 - end_index = refinement_ctrl.shape[0] + end_index = start_halo_2 + else: + start_index = start_halo_2 + end_index = end_halo_2 start_indices[domain] = gtx.int32(start_index) # type: ignore [attr-defined] end_indices[domain] = gtx.int32(end_index) # type: ignore [attr-defined] + + ordered_domains = h_grid.get_ordered_domains(dim) + for domain in ordered_domains: + value = ( + domain.zone.level + if domain.zone.is_lateral_boundary() + else _LAST_BOUNDARY[dim].level + domain.zone.level + ) + found = array_ns.where((refinement_ctrl == value) & owned)[0] + start_index, end_index = ( + (array_ns.min(found).item(), array_ns.max(found).item() + 1) + if found.size > 0 + else (0, 0) + ) + start_indices[domain] = gtx.int32(start_index) + end_indices[domain] = gtx.int32(end_index) + + interior_domain = h_grid.domain(dim)(h_grid.Zone.INTERIOR) + # for the Vertex and Edges the level after the nudging zones are not ordered anymore, so + # we rely on using the end index of the nudging zone for INTERIOR + nudging = h_grid.get_last_nudging(dim) + start_indices[interior_domain] = end_indices[nudging] + halo_1 = h_grid.domain(dim)(h_grid.Zone.HALO) + end_indices[interior_domain] = start_indices[halo_1] + + local_domain = h_grid.domain(dim)(h_grid.Zone.LOCAL) + start_indices[local_domain] = gtx.int32(0) + end_indices[local_domain] = start_indices[halo_1] + return start_indices, end_indices diff --git a/model/common/src/icon4py/model/common/grid/gridfile.py b/model/common/src/icon4py/model/common/grid/gridfile.py index 5873089d67..c51cdcc651 100644 --- a/model/common/src/icon4py/model/common/grid/gridfile.py +++ b/model/common/src/icon4py/model/common/grid/gridfile.py @@ -8,11 +8,13 @@ import enum import logging +from typing import Any, Protocol import numpy as np from gt4py import next as gtx from icon4py.model.common import exceptions +from icon4py.model.common.utils import data_allocation as data_alloc _log = logging.getLogger(__name__) @@ -29,6 +31,34 @@ def __init__(self, *args, **kwargs): raise ModuleNotFoundError("NetCDF4 is not installed.") +class IndexTransformation(Protocol): + """Return an offset field to be applied to index fields""" + + def __call__( + self, + array: data_alloc.NDArray, + ) -> data_alloc.NDArray: ... + + +class NoTransformation(IndexTransformation): + """Empty implementation of the Protocol. Just return zeros.""" + + def __call__(self, array: data_alloc.NDArray) -> data_alloc.NDArray: + return data_alloc.array_ns_from_array(array).zeros_like(array) + + +class ToZeroBasedIndexTransformation(IndexTransformation): + def __call__(self, array: data_alloc.NDArray) -> data_alloc.NDArray: + """ + Calculate the index offset needed for usage with python. + + Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for + INVALID values which are marked with -1 in the grid file and are kept such. + """ + xp = data_alloc.array_ns_from_array(array) + return xp.asarray(xp.where(array == GridFile.INVALID_INDEX, 0, -1), dtype=gtx.int32) + + class GridFileName(str, enum.Enum): pass @@ -106,36 +136,60 @@ class MandatoryPropertyName(PropertyName): ROOT = "grid_root" -class DimensionName(GridFileName): +class DimensionName(GridFileName): ... + + +class DynamicDimension(DimensionName): """Dimension values (sizes) used in grid file.""" #: number of vertices VERTEX_NAME = "vertex" - #: number of edges EDGE_NAME = "edge" #: number of cells CELL_NAME = "cell" + #: number of child domains (for nesting) + MAX_CHILD_DOMAINS = "max_chdom" + + +class FixedSizeDimension(DimensionName): + size: int + + def __new__(cls, value: str, size_: int): + obj = str.__new__(cls) + obj._value_ = value + obj.size = size_ + return obj + #: number of edges in a diamond: 4 - DIAMOND_EDGE_SIZE = "no" + DIAMOND_EDGE_SIZE = ("no", 4) #: number of edges/cells neighboring one vertex: 6 (for regular, non pentagons) - NEIGHBORS_TO_VERTEX_SIZE = "ne" + NEIGHBORS_TO_VERTEX_SIZE = ("ne", 6) #: number of cells edges, vertices and cells neighboring a cell: 3 - NEIGHBORS_TO_CELL_SIZE = "nv" + NEIGHBORS_TO_CELL_SIZE = ("nv", 3) #: number of vertices/cells neighboring an edge: 2 - NEIGHBORS_TO_EDGE_SIZE = "nc" - - #: number of child domains (for nesting) - MAX_CHILD_DOMAINS = "max_chdom" + NEIGHBORS_TO_EDGE_SIZE = ("nc", 2) #: Grid refinement: maximal number in grid-refinement (refin_ctl) array for each dimension - CELL_GRF = "cell_grf" - EDGE_GRF = "edge_grf" - VERTEX_GRF = "vert_grf" + CELL_GRF = ("cell_grf", 14) + EDGE_GRF = ("edge_grf", 28) + VERTEX_GRF = ("vert_grf", 14) + + def __str__(self): + return f"{self.name}({self.name}: {self.size})" + + def __hash__(self): + return hash((self.name, self.size)) + + def __eq__(self, other: Any) -> bool: + """Check equality based on zone name and level.""" + if not isinstance(other, FixedSizeDimension): + return False + return (self.name, self.size) == (other.name, other.size) class FieldName(GridFileName): ... @@ -175,15 +229,15 @@ class ConnectivityName(FieldName): class GeometryName(FieldName): + CELL_NORMAL_ORIENTATION = "orientation_of_normal" + TANGENT_ORIENTATION = "edge_system_orientation" + EDGE_ORIENTATION_ON_VERTEX = "edge_orientation" + # TODO(halungge): compute from coordinates CELL_AREA = "cell_area" - # TODO(halungge): compute from coordinates DUAL_AREA = "dual_area" EDGE_LENGTH = "edge_length" DUAL_EDGE_LENGTH = "dual_edge_length" - CELL_NORMAL_ORIENTATION = "orientation_of_normal" - TANGENT_ORIENTATION = "edge_system_orientation" - EDGE_ORIENTATION_ON_VERTEX = "edge_orientation" # TODO(halungge): compute from coordinates EDGE_CELL_DISTANCE = "edge_cell_distance" EDGE_VERTEX_DISTANCE = "edge_vert_distance" @@ -198,7 +252,7 @@ class GeometryName(FieldName): class CoordinateName(FieldName): """ Coordinates of cell centers, edge midpoints and vertices. - Units: radianfor both MPI-M and DWD + Units: radian for both MPI-M and DWD """ CELL_LONGITUDE = "clon" @@ -255,8 +309,9 @@ class GridFile: INVALID_INDEX = -1 - def __init__(self, file_name: str): + def __init__(self, file_name: str, offset_transformation: IndexTransformation): self._filename = file_name + self._offset_transformation = offset_transformation self._dataset = None def dimension(self, name: DimensionName) -> int: @@ -276,7 +331,11 @@ def try_attribute(self, name: PropertyName) -> str | int | float | None: return None def int_variable( - self, name: FieldName, indices: np.ndarray = None, transpose: bool = True + self, + name: FieldName, + indices: data_alloc.NDArray | None = None, + transpose: bool = True, + apply_offset: bool = True, ) -> np.ndarray: """Read a integer field from the grid file. @@ -284,36 +343,53 @@ def int_variable( Args: name: name of the field to read + indices: list of indices to read transpose: flag to indicate whether the file should be transposed (for 2d fields) + apply_offset: flag to indicate whether the offset should be applied + to the indices, defaults to True Returns: NDArray: field data """ - _log.debug(f"reading {name}: transposing = {transpose}") - return self.variable(name, indices, transpose=transpose, dtype=gtx.int32) + _log.debug(f"reading {name}: transposing = {transpose} apply_offset={apply_offset}") + variable = self.variable(name, indices, transpose=transpose, dtype=gtx.int32) + if apply_offset: + return variable + self._offset_transformation(variable) + return variable def variable( self, name: FieldName, - indices: np.ndarray = None, - transpose=False, + indices: data_alloc.NDArray | None = None, + transpose: bool = False, dtype: np.dtype = gtx.float64, ) -> np.ndarray: """Read a field from the grid file. - If a index array is given it only reads the values at those positions. + If an index array is given it only reads the values at those positions. Args: name: name of the field to read - indices: indices to read + indices: indices to read if requesting a restricted set of indices. We assume this be a 1d array it will be applied to the 1. dimension (after transposition) transpose: flag indicateing whether the array needs to be transposed to match icon4py dimension ordering, defaults to False dtype: datatype of the field """ + + assert indices is None or indices.ndim == 1, "indices must be 1 dimensional" + try: variable = self._dataset.variables[name] + variable_size = variable.ndim + n = (variable.shape[0],) if variable_size > 1 else () + target_shape = (*n, -1) + + slicer = [slice(None) for _ in range(variable_size)] + if indices is not None and indices.size > 0: + # apply the slicing to the correct dimension + slicer[(1 if transpose else 0)] = data_alloc.as_numpy(indices) _log.debug(f"reading {name}: transposing = {transpose}") - data = variable[:] if indices is None else variable[indices] - data = np.array(data, dtype=dtype) + data = np.asarray(variable[tuple(slicer)]) + data = np.array(data, dtype=dtype).ravel(order="K").reshape(target_shape) return np.transpose(data) if transpose else data except KeyError as err: msg = f"{name} does not exist in dataset" @@ -321,6 +397,13 @@ def variable( _log.debug(f"Error: {err}") raise exceptions.IconGridError(msg) from err + def __enter__(self): + self.open() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + def close(self): self._dataset.close() diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index 50d919166d..d9ca9724b6 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -115,23 +115,23 @@ class Zone(enum.Enum): ## CellDim - | ICON constant or value |python | ICON4py Name | - | from mo_impl_constants.f90 |index | | - |:------------------------------------- |:----:|:-------------------------- | - | `min_rlcell_int-3`, `min_rlcell` (-8) | 0 | `END` | - | `min_rlcell_int-3` (-7) | 1 | | - | `min_rlcell_int-2`, (-6) | 2 |`HALO_LEVEL_2` | - | `min_rlcell_int-1` (-5) | 3 |`HALO` | - | `min_rlcell_int`(-4) | 4 |`LOCAL` | - | (-3) | 5 | | unused in icon4py (relevant for nesting) - | (-2) | 6 | | unused in icon4py (relevant for nesting) - | (-1) | 7 | | unused in icon4py (relevant for nesting) - | `0` | 8 |`INTERIOR` | - | `1` | 9 |`LATERAL_BOUNDARY` | - | `2` |10 | `LATERAL_BOUNDARY_LEVEL_2` | - | `3` |11 | `LATERAL_BOUNDARY_LEVEL_3` | - | `grf_bdywidth_c` (4) |12 | `LATERAL_BOUNDARY_LEVEL_4` | - | `grf_bdywith_c +1`,max_rlcell (5) |13 | `NUDGING` | + | ICON constant or value | python| ICON4py Name | + | from mo_impl_constants.f90 | index | | + |:--------------------------------------|:------|:---------------------------| + | `min_rlcell_int-3`, `min_rlcell` (-8) | 0 | `END` | + | `min_rlcell_int-3` (-7) | 1 | | + | `min_rlcell_int-2`, (-6) | 2 | `HALO_LEVEL_2` | + | `min_rlcell_int-1` (-5) | 3 | `HALO` | + | `min_rlcell_int`(-4) | 4 | `LOCAL` | + | (-3) | 5 | | unused in icon4py (relevant for nesting) + | (-2) | 6 | | unused in icon4py (relevant for nesting) + | (-1) | 7 | | unused in icon4py (relevant for nesting) + | `0` | 8 | `INTERIOR` | + | `1` | 9 | `LATERAL_BOUNDARY` | + | `2` | 10 | `LATERAL_BOUNDARY_LEVEL_2` | + | `3` | 11 | `LATERAL_BOUNDARY_LEVEL_3` | + | `grf_bdywidth_c` (4) | 12 | `LATERAL_BOUNDARY_LEVEL_4` | + | `grf_bdywith_c +1`,max_rlcell (5) | 13 | `NUDGING` | Lateral boundary and nudging are only relevant for LAM runs, halo lines only for distributed domains. @@ -139,53 +139,55 @@ class Zone(enum.Enum): ## VertexDim - | ICON constant or value | python | ICON4Py Name | - | from mo_impl_constants.f90 | index | | - |:--------------------------------------- |:------|:-------------------------- | + | ICON constant or value | python| ICON4Py Name | + | from mo_impl_constants.f90 | index | | + |:----------------------------------------|:------|:---------------------------| | `min_rlvert` (-7) | 0 | `END` | - | `min_rlvert+1`, `min_rlvert_int-2` (-6) | 1 |`HALO_LEVEL_2` | - | `min_rlvert_int-1` (-5) | 2 |`HALO` | - | `min_rlvert_int` (-4) | 3 |`LOCAL` | + | `min_rlvert+1`, `min_rlvert_int-2` (-6) | 1 | `HALO_LEVEL_2` | + | `min_rlvert_int-1` (-5) | 2 | `HALO` | + | `min_rlvert_int` (-4) | 3 | `LOCAL` | | (-3) | 4 | | unused in icon4py (relevant for nesting) | (-2) | 5 | | unused in icon4py (relevant for nesting) | (-1) | 6 | | unused in icon4py (relevant for nesting) - | `0` | ` 7 |INTERIOR` | - | `1` | 8 |`LATERAL_BOUNDARY` | - | `2` | 9 |`LATERAL_BOUNDARY_LEVEL_2` | - | `3` | 10 |`LATERAL_BOUNDARY_LEVEL_3` | - | `4` | 11 |`LATERAL_BOUNDARY_LEVEL_4` | - | `max_rlvert` (5) | 12 |`NUDGING` | + | `0` | 7 | `INTERIOR` | + | `1` | 8 | `LATERAL_BOUNDARY` | + | `2` | 9 | `LATERAL_BOUNDARY_LEVEL_2` | + | `3` | 10 | `LATERAL_BOUNDARY_LEVEL_3` | + | `4` | 11 | `LATERAL_BOUNDARY_LEVEL_4` | + | `max_rlvert` (5) | 12 | `NUDGING` | For the meaning see above. ## EdgeDim - | ICON constant or value | python | ICON4Py Name | - | from mo_impl_constants.f90 | index | | - |:-------------------------------------- |:-------|:-------------------------- | - | `min_rledge` (-13) | 0 |`END` | - | `min_rledge_int-2` (-10) | 1 |`HALO_LEVEL_2` | - | `min_rledge_int-1` (-9) | 2 |`HALO` | - | `min_rledge_int` (-8) | 3 |`LOCAL` | - | (-7) | 4 | | unused in icon4py (relevant for nesting) - | (-6) | 5 | | unused in icon4py (relevant for nesting) - | (-5) | 6 | | unused in icon4py (relevant for nesting) - | (-4) | 7 | | unused in icon4py (relevant for nesting) - | (-3) | 8 | | unused in icon4py (relevant for nesting) - | (-2) | 9 | | unused in icon4py (relevant for nesting) - |(-1) | 10 | | unused in icon4py (relevant for nesting) - | `0` | 11 | `INTERIOR` | - | `1` | 12 | `LATERAL_BOUNDARY` | - | `2` | 13 | `LATERAL_BOUNDARY_LEVEL_2` | - | `3` | 14 |`LATERAL_BOUNDARY_LEVEL_3` | - | `4` | 15 |`LATERAL_BOUNDARY_LEVEL_4` | - | `5` | 16 |`LATERAL_BOUNDARY_LEVEL_5` | - | `6` | 17 |`LATERAL_BOUNDARY_LEVEL_6` | - | `7` | 18 | `LATERAL_BOUNDARY_LEVEL_7` | - | `8` | 19 | `LATERAL_BOUNDARY_LEVEL_8`| - | `grf_bdywidth_e` (9) | 20 | `NUDGING` | - | `grf_bdywidth_e+1`, `max_rledge` (10) | 21 | `NUDGING_LEVEL_2` | + | ICON constant or value | python | ICON4Py Name | + | from mo_impl_constants.f90 | index | | + |:---------------------------------------|:-------|:---------------------------| + | `min_rledge` (-13) | 0 | `END` | + | (-12) | 1 | | + | (-11) | 2 | | + | `min_rledge_int-2` (-10) | 3 | `HALO_LEVEL_2` | + | `min_rledge_int-1` (-9) | 4 | `HALO` | + | `min_rledge_int` (-8) | 5 | `LOCAL` | + | (-7) | 6 | | unused in icon4py (relevant for nesting) + | (-6) | 7 | | unused in icon4py (relevant for nesting) + | (-5) | 8 | | unused in icon4py (relevant for nesting) + | (-4) | 9 | | unused in icon4py (relevant for nesting) + | (-3) | 10 | | unused in icon4py (relevant for nesting) + | (-2) | 11 | | unused in icon4py (relevant for nesting) + | (-1) | 12 | | unused in icon4py (relevant for nesting) + | `0` | 13 | `INTERIOR` | + | `1` | 14 | `LATERAL_BOUNDARY` | + | `2` | 15 | `LATERAL_BOUNDARY_LEVEL_2` | + | `3` | 16 | `LATERAL_BOUNDARY_LEVEL_3` | + | `4` | 17 | `LATERAL_BOUNDARY_LEVEL_4` | + | `5` | 18 | `LATERAL_BOUNDARY_LEVEL_5` | + | `6` | 19 | `LATERAL_BOUNDARY_LEVEL_6` | + | `7` | 20 | `LATERAL_BOUNDARY_LEVEL_7` | + | `8` | 21 | `LATERAL_BOUNDARY_LEVEL_8` | + | `grf_bdywidth_e` (9) | 22 | `NUDGING` | + | `grf_bdywidth_e+1`, `max_rledge` (10) | 23 | `NUDGING_LEVEL_2` | """ @@ -295,6 +297,11 @@ def is_local(self) -> bool: EDGE_ZONES = tuple(Zone) + +def max_boundary_level(dim: gtx.Dimension) -> int: + return max((d.level for d in _get_zones_for_dim(dim) if d.is_lateral_boundary()), default=1) + + _ZONE_TO_INDEX_MAPPING = { Zone.END: lambda dim: _icon_domain_index(_ICON_END, dim), Zone.INTERIOR: lambda dim: _icon_domain_index(_ICON_INTERIOR, dim), @@ -411,6 +418,27 @@ def get_domains_for_dim(dim: gtx.Dimension) -> Iterator[Domain]: return domains +def get_halo_domains(dim: gtx.Dimension) -> Iterator[Domain]: + get_domain = domain(dim) + domains = (get_domain(zone) for zone in _get_zones_for_dim(dim) if zone.is_halo()) + return domains + + +def get_ordered_domains(dim: gtx.Dimension) -> Iterator[Domain]: + get_domain = domain(dim) + domains = ( + get_domain(zone) + for zone in _get_zones_for_dim(dim) + if zone.is_lateral_boundary() or zone.is_nudging() + ) + return domains + + +def get_last_nudging(dim): + zone = Zone.NUDGING if dim in (dims.VertexDim, dims.CellDim) else Zone.NUDGING_LEVEL_2 + return domain(dim)(zone) + + def get_start_end_idx_from_icon_arrays( dim: gtx.Dimension, start_indices: dict[gtx.Dimension, np.ndarray], diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 2e55da36d5..1822738edb 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -137,19 +137,18 @@ class IconGrid(base.Grid): ) -def _has_skip_values(offset: gtx.FieldOffset, limited_area: bool) -> bool: +def _has_skip_values(offset: gtx.FieldOffset, limited_area_or_distributed: bool) -> bool: """ For the icosahedral global grid skip values are only present for the pentagon points. - In the local area model there are also skip values at the boundaries when + In the local area model or a distributed grid there are also skip values at the boundaries or halos when accessing neighbouring cells or edges from vertices. """ dimension = offset.target[1] assert dimension.kind == gtx.DimensionKind.LOCAL, "only local dimensions can have skip values" - value = dimension in CONNECTIVITIES_ON_PENTAGONS or ( - limited_area and dimension in CONNECTIVITIES_ON_BOUNDARIES + return dimension in CONNECTIVITIES_ON_PENTAGONS or ( + limited_area_or_distributed and dimension in CONNECTIVITIES_ON_BOUNDARIES ) - return value def _should_replace_skip_values( @@ -190,14 +189,25 @@ def icon_grid( global_properties: GlobalGridParams, refinement_control: dict[gtx.Dimension, gtx.Field] | None = None, ) -> IconGrid: + # TODO(msimberg): What should we do about this. (The global) num_cells is + # not guaranteed to be set here when used through fortran. Should we: + # 1. Ignore distributed? + # 2. Compute num_cells with a reduction? + # 3. Use a ProcessProperties to detect it? + distributed = ( + config.num_cells < global_properties.num_cells + if global_properties.num_cells is not None + else False + ) + limited_area_or_distributed = config.limited_area or distributed connectivities = { offset.value: base.construct_connectivity( offset, data_alloc.import_array_ns(allocator).asarray(table), - skip_value=-1 if _has_skip_values(offset, config.limited_area) else None, + skip_value=-1 if _has_skip_values(offset, limited_area_or_distributed) else None, allocator=allocator, replace_skip_values=_should_replace_skip_values( - offset, config.keep_skip_values, config.limited_area + offset, config.keep_skip_values, limited_area_or_distributed ), ) for offset, table in neighbor_tables.items() 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..5e8246ea8f 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -92,7 +92,7 @@ def _sources(self) -> factory.FieldSource: def _register_computed_fields(self) -> None: nudging_coefficients_for_edges = factory.ProgramFieldProvider( - func=nudgecoeffs.compute_nudgecoeffs.with_backend(None), + func=nudgecoeffs.compute_nudgecoeffs, domain={ dims.EdgeDim: ( edge_domain(h_grid.Zone.NUDGING_LEVEL_2), diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py b/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py index c729c08d7f..1f4b519dcb 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py @@ -9,24 +9,24 @@ from gt4py.next import neighbor_sum import icon4py.model.common.type_alias as types -from icon4py.model.common import dimension as dims +from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import V2C, V2CDim @gtx.field_operator def _compute_cell_2_vertex_interpolation( - cell_in: gtx.Field[[dims.CellDim, dims.KDim], types.wpfloat], - c_int: gtx.Field[[dims.VertexDim, V2CDim], types.wpfloat], -) -> gtx.Field[[dims.VertexDim, dims.KDim], types.wpfloat]: + cell_in: fa.CellKField[types.wpfloat], + c_int: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], types.wpfloat], +) -> fa.VertexKField[types.wpfloat]: vert_out = neighbor_sum(c_int * cell_in(V2C), axis=V2CDim) return vert_out @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_cell_2_vertex_interpolation( - cell_in: gtx.Field[[dims.CellDim, dims.KDim], types.wpfloat], - c_int: gtx.Field[[dims.VertexDim, V2CDim], types.wpfloat], - vert_out: gtx.Field[[dims.VertexDim, dims.KDim], types.wpfloat], + cell_in: fa.CellKField[types.wpfloat], + c_int: gtx.Field[[dims.VertexDim, dims.V2CDim], types.wpfloat], + vert_out: fa.VertexKField[types.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, diff --git a/model/common/src/icon4py/model/common/utils/data_allocation.py b/model/common/src/icon4py/model/common/utils/data_allocation.py index c58a8cd827..d059153ee1 100644 --- a/model/common/src/icon4py/model/common/utils/data_allocation.py +++ b/model/common/src/icon4py/model/common/utils/data_allocation.py @@ -78,6 +78,15 @@ def array_ns(try_cupy: bool) -> ModuleType: return np +def array_ns_from_array(array: NDArray) -> ModuleType: + if isinstance(array, np.ndarray): + import numpy as xp + else: + import cupy as xp # type: ignore[no-redef] + + return xp + + def import_array_ns(allocator: gtx_allocators.FieldBufferAllocationUtil | None) -> ModuleType: """Import cupy or numpy depending on a chosen GT4Py backend DevicType.""" return array_ns(device_utils.is_cupy_device(allocator)) diff --git a/model/common/tests/common/decomposition/fixtures.py b/model/common/tests/common/decomposition/fixtures.py new file mode 100644 index 0000000000..ab78d81e9e --- /dev/null +++ b/model/common/tests/common/decomposition/fixtures.py @@ -0,0 +1,21 @@ +# 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 pytest +from gt4py.next import common as gtx_common + +from icon4py.model.common.grid import simple + + +@pytest.fixture(scope="session") +def simple_neighbor_tables(): + grid = simple.simple_grid() + neighbor_tables = { + k: v.ndarray for k, v in grid.connectivities.items() if gtx_common.is_neighbor_table(v) + } + return neighbor_tables diff --git a/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py b/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py index 0f0d4c59f1..9fa4adda38 100644 --- a/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py +++ b/model/common/tests/common/decomposition/mpi_tests/test_mpi_decomposition.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import functools +import logging from typing import Any import gt4py.next.typing as gtx_typing @@ -19,6 +20,7 @@ try: import mpi4py # import mpi4py to check for optional mpi dependency + from mpi4py import MPI except ImportError: pytest.skip("Skipping parallel on single node installation", allow_module_level=True) @@ -27,8 +29,8 @@ import icon4py.model.testing.test_utils as test_helpers from icon4py.model.common import dimension as dims, model_backends from icon4py.model.common.decomposition import definitions, mpi_decomposition -from icon4py.model.testing import definitions as test_defs, serialbox -from icon4py.model.testing.parallel_helpers import check_comm_size +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions as test_defs, parallel_helpers, serialbox from ...fixtures import ( backend, @@ -45,10 +47,15 @@ ) +_log = logging.getLogger(__name__) + + +@pytest.mark.mpi(min_size=2) @pytest.mark.parametrize("processor_props", [True], indirect=True) def test_props(processor_props: definitions.ProcessProperties) -> None: assert processor_props.comm assert processor_props.comm_size > 1 + assert 0 <= processor_props.rank < processor_props.comm_size @pytest.mark.mpi(min_size=2) @@ -78,7 +85,7 @@ def test_decomposition_info_masked( processor_props: definitions.ProcessProperties, experiment: test_defs.Experiment, ) -> None: - check_comm_size(processor_props, sizes=(2,)) + parallel_helpers.check_comm_size(processor_props, sizes=(2,)) my_rank = processor_props.rank all_indices = decomposition_info.global_index(dim, definitions.DecompositionInfo.EntryType.ALL) my_total = total[my_rank] @@ -133,7 +140,8 @@ def test_decomposition_info_local_index( processor_props: definitions.ProcessProperties, experiment: test_defs.Experiment, ): - check_comm_size(processor_props, sizes=(2,)) + caplog.set_level(logging.INFO) + parallel_helpers.check_comm_size(processor_props, sizes=(2,)) my_rank = processor_props.rank all_indices = decomposition_info.local_index(dim, definitions.DecompositionInfo.EntryType.ALL) my_total = total[my_rank] @@ -155,6 +163,45 @@ def test_decomposition_info_local_index( _assert_index_partitioning(all_indices, halo_indices, owned_indices) +@pytest.mark.datatest +@pytest.mark.mpi +@pytest.mark.parametrize("dim", (dims.CellDim, dims.EdgeDim, dims.VertexDim)) +def test_decomposition_info_halo_level_mask( + dim: gtx.Dimension, + experiment: test_defs.Experiment, + decomposition_info: definitions.DecompositionInfo, +) -> None: + first_halo_level = decomposition_info.halo_level_mask( + dim, definitions.DecompositionFlag.FIRST_HALO_LEVEL + ) + assert first_halo_level.ndim == 1 + assert np.count_nonzero(first_halo_level) == decomposition_info.get_halo_size( + dim, definitions.DecompositionFlag.FIRST_HALO_LEVEL + ) + second_halo_level = decomposition_info.halo_level_mask( + dim, definitions.DecompositionFlag.SECOND_HALO_LEVEL + ) + assert second_halo_level.ndim == 1 + assert np.count_nonzero(second_halo_level) == decomposition_info.get_halo_size( + dim, definitions.DecompositionFlag.SECOND_HALO_LEVEL + ) + assert np.count_nonzero(first_halo_level) + np.count_nonzero( + second_halo_level + ) == np.count_nonzero(~decomposition_info.owner_mask(dim)) + + +@pytest.mark.datatest +@pytest.mark.mpi +@pytest.mark.parametrize("dim", (dims.CellDim, dims.EdgeDim, dims.VertexDim)) +def test_decomposition_info_third_level_is_empty( + dim: gtx.Dimension, + experiment: test_defs.Experiment, + decomposition_info: definitions.DecompositionInfo, +) -> None: + level = decomposition_info.halo_level_mask(dim, definitions.DecompositionFlag.THIRD_HALO_LEVEL) + assert np.count_nonzero(level) == 0 + + @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) @pytest.mark.parametrize("num", [1, 2, 3, 4, 5, 6, 7, 8]) @@ -190,7 +237,7 @@ def test_decomposition_info_matches_gridsize( icon_grid: icon.IconGrid, processor_props: definitions.ProcessProperties, ) -> None: - check_comm_size(processor_props) + parallel_helpers.check_comm_size(processor_props) assert ( decomposition_info.global_index( dim=dims.CellDim, entry_type=definitions.DecompositionInfo.EntryType.ALL @@ -213,7 +260,7 @@ def test_decomposition_info_matches_gridsize( @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_create_multi_node_runtime_with_mpi( +def test_create_multi_rank_runtime_with_mpi( decomposition_info: definitions.DecompositionInfo, processor_props: definitions.ProcessProperties, ) -> None: @@ -265,13 +312,13 @@ def test_exchange_on_dummy_data( assert np.all(input_field.asnumpy() == number) exchange.exchange_and_wait(dimension, input_field) result = input_field.asnumpy() - print(f"rank={processor_props.rank} - num of halo points ={halo_points.shape}") - print( + _log.info(f"rank={processor_props.rank} - num of halo points ={halo_points.shape}") + _log.info( f" rank={processor_props.rank} - exchanged points: {np.sum(result != number)/grid.num_levels}" ) - print(f"rank={processor_props.rank} - halo points: {halo_points}") + _log.info(f"rank={processor_props.rank} - halo points: {halo_points}") changed_points = np.argwhere(result[:, 2] != number) - print(f"rank={processor_props.rank} - num changed points {changed_points.shape} ") + _log.info(f"rank={processor_props.rank} - num changed points {changed_points.shape} ") assert np.all(result[local_points, :] == number) assert np.all(result[halo_points, :] != number) @@ -292,7 +339,7 @@ def test_halo_exchange_for_sparse_field( edge_orientation = grid_savepoint.edge_orientation() area = grid_savepoint.cell_areas() field_ref = interpolation_savepoint.geofac_div() - print( + _log.info( f"{processor_props.rank}/{processor_props.comm_size}: size of reference field {field_ref.asnumpy().shape}" ) result = data_alloc.zero_field( @@ -308,7 +355,7 @@ def test_halo_exchange_for_sparse_field( out=result, offset_provider={"C2E": icon_grid.get_connectivity("C2E")}, ) - print( + _log.info( f"{processor_props.rank}/{processor_props.comm_size}: size of computed field {result.asnumpy().shape}" ) exchange.exchange_and_wait(dims.CellDim, result) diff --git a/model/common/tests/common/decomposition/mpi_tests/test_parallel_halo.py b/model/common/tests/common/decomposition/mpi_tests/test_parallel_halo.py new file mode 100644 index 0000000000..837aa99f8f --- /dev/null +++ b/model/common/tests/common/decomposition/mpi_tests/test_parallel_halo.py @@ -0,0 +1,87 @@ +# 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 logging + +import gt4py.next as gtx +import numpy as np +import pytest + +import icon4py.model.common.dimension as dims +from icon4py.model.common.decomposition import ( + definitions as decomposition_defs, + halo, + mpi_decomposition, +) +from icon4py.model.common.grid import simple +from icon4py.model.testing import parallel_helpers +from icon4py.model.testing.fixtures import processor_props + +from .. import utils +from ..fixtures import simple_neighbor_tables + + +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + + +_log = logging.getLogger(__name__) + +backend = None + + +def global_indices(dim: gtx.Dimension) -> np.ndarray: + mesh = simple.simple_grid() + return np.arange(mesh.size[dim], dtype=gtx.int32) + + +@pytest.mark.parametrize("dim", [dims.CellDim, dims.EdgeDim, dims.VertexDim]) +@pytest.mark.mpi(min_size=4) +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_element_ownership_is_unique( + dim, + processor_props, + simple_neighbor_tables, +): + parallel_helpers.check_comm_size(processor_props, sizes=[4]) + + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=processor_props, + allocator=backend, + ) + + decomposition_info = halo_generator(utils.SIMPLE_DISTRIBUTION) + owned = decomposition_info.global_index( + dim, decomposition_defs.DecompositionInfo.EntryType.OWNED + ) + _log.info(f"\nrank {processor_props.rank} owns {dim} : {owned} ") + # assert that each cell is only owned by one rank + comm = processor_props.comm + + my_size = owned.shape[0] + local_sizes = np.array(comm.gather(my_size, root=0)) + buffer_size = 27 + send_buf = np.full(buffer_size, -1, dtype=int) + send_buf[:my_size] = owned + _log.info(f"rank {processor_props.rank} send_buf: {send_buf}") + if processor_props.rank == 0: + _log.info(f"local_sizes: {local_sizes}") + recv_buffer = np.full((4, buffer_size), -1, dtype=int) + _log.info(f"{recv_buffer.shape}") + else: + recv_buffer = None + # Gatherv does not work if one of the buffers has size-0 (VertexDim) + comm.Gather(sendbuf=send_buf, recvbuf=recv_buffer, root=0) + if processor_props.rank == 0: + _log.info(f"global indices: {recv_buffer}") + # check there are no duplicates + values = recv_buffer[recv_buffer != -1] + assert values.size == len(np.unique(values)) + # check the buffer has all global indices + assert np.all(np.sort(values) == global_indices(dim)) diff --git a/model/common/tests/common/decomposition/unit_tests/test_definitions.py b/model/common/tests/common/decomposition/unit_tests/test_definitions.py index d966d5ba7b..3f8990909e 100644 --- a/model/common/tests/common/decomposition/unit_tests/test_definitions.py +++ b/model/common/tests/common/decomposition/unit_tests/test_definitions.py @@ -5,13 +5,22 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +import numpy as np import pytest +from gt4py.next import common as gtx_common +import icon4py.model.common.dimension as dims +import icon4py.model.common.utils.data_allocation as data_alloc +from icon4py.model.common.decomposition import definitions from icon4py.model.common.decomposition.definitions import ( DecompositionInfo, SingleNodeExchange, create_exchange, ) +from icon4py.model.common.grid import simple +from icon4py.model.testing import definitions as test_defs from icon4py.model.testing.fixtures.datatest import ( # import fixtures form test_utils backend, data_provider, @@ -22,14 +31,67 @@ processor_props, ) +from ...grid import utils as grid_utils + + +@pytest.mark.parametrize("processor_props", [False], indirect=True) +def test_create_single_node_runtime_without_mpi(processor_props): # fixture + decomposition_info = definitions.DecompositionInfo() + exchange = definitions.create_exchange(processor_props, decomposition_info) + + assert isinstance(exchange, definitions.SingleNodeExchange) + + +def get_neighbor_tables_for_simple_grid() -> dict[str, data_alloc.NDArray]: + grid = simple.simple_grid() + neighbor_tables = { + k: v.ndarray + for k, v in grid.connectivities.items() + if gtx_common.is_neighbor_connectivity(v) + } + return neighbor_tables + -@pytest.mark.datatest -def test_create_single_node_runtime_without_mpi(icon_grid, processor_props): - decomposition_info = DecompositionInfo( - num_cells=icon_grid.num_cells, - num_edges=icon_grid.num_edges, - num_vertices=icon_grid.num_vertices, +offsets = [dims.E2C, dims.E2V, dims.C2E, dims.C2E2C, dims.V2C, dims.V2E, dims.C2V, dims.E2C2V] + + +@pytest.mark.parametrize("dim", grid_utils.main_horizontal_dims()) +def test_decomposition_info_single_node_empty_halo(dim: gtx.Dimension) -> None: + manager = grid_utils.run_grid_manager( + test_defs.Grids.MCH_CH_R04B09_DSL, keep_skip_values=True, backend=None + ) + + decomposition_info = manager.decomposition_info + for level in ( + definitions.DecompositionFlag.FIRST_HALO_LEVEL, + definitions.DecompositionFlag.SECOND_HALO_LEVEL, + definitions.DecompositionFlag.THIRD_HALO_LEVEL, + ): + assert decomposition_info.get_halo_size(dim, level) == 0 + assert np.count_nonzero(decomposition_info.halo_level_mask(dim, level)) == 0 + assert ( + decomposition_info.get_halo_size(dim, definitions.DecompositionFlag.OWNED) + == manager.grid.size[dim] ) - exchange = create_exchange(processor_props, decomposition_info) - assert isinstance(exchange, SingleNodeExchange) + +@pytest.mark.parametrize( + "flag, expected", + [ + (definitions.DecompositionFlag.OWNED, False), + (definitions.DecompositionFlag.SECOND_HALO_LEVEL, True), + (definitions.DecompositionFlag.THIRD_HALO_LEVEL, True), + (definitions.DecompositionFlag.FIRST_HALO_LEVEL, True), + (definitions.DecompositionFlag.UNDEFINED, False), + ], +) +def test_decomposition_info_is_distributed(flag, expected) -> None: + mesh = simple.simple_grid(allocator=None, num_levels=10) + decomp = definitions.DecompositionInfo() + decomp.set_dimension( + dims.CellDim, + np.arange(mesh.num_cells), + np.arange(mesh.num_cells), + np.ones((mesh.num_cells,)) * flag, + ) + assert decomp.is_distributed() == expected diff --git a/model/common/tests/common/decomposition/unit_tests/test_halo.py b/model/common/tests/common/decomposition/unit_tests/test_halo.py new file mode 100644 index 0000000000..96e9917f73 --- /dev/null +++ b/model/common/tests/common/decomposition/unit_tests/test_halo.py @@ -0,0 +1,240 @@ +# 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 numpy as np +import pytest +from gt4py.next import common as gtx_common + +from icon4py.model.common import dimension as dims, exceptions, model_backends +from icon4py.model.common.decomposition import decomposer as decomp, definitions, halo +from icon4py.model.common.grid import base as base_grid, simple + +from ...fixtures import backend_like, processor_props +from .. import utils +from ..fixtures import simple_neighbor_tables +from ..utils import dummy_four_ranks +from .test_definitions import get_neighbor_tables_for_simple_grid, offsets + + +@pytest.mark.parametrize("rank", [0, 1, 2, 4]) +def test_halo_constructor_owned_cells(rank, simple_neighbor_tables, backend_like): + processor_props = utils.DummyProps(rank=rank) + allocator = model_backends.get_allocator(backend_like) + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=processor_props, + allocator=allocator, + ) + my_owned_cells = halo_generator.owned_cells(utils.SIMPLE_DISTRIBUTION) + + print(f"rank {processor_props.rank} owns {my_owned_cells} ") + assert my_owned_cells.size == len(utils._CELL_OWN[processor_props.rank]) + assert np.setdiff1d(my_owned_cells, utils._CELL_OWN[processor_props.rank]).size == 0 + + +@pytest.mark.parametrize("dim", [dims.CellDim, dims.VertexDim, dims.EdgeDim]) +@pytest.mark.parametrize("rank", [0, 1, 2, 4]) +def test_halo_constructor_decomposition_info_global_indices(rank, simple_neighbor_tables, dim): + processor_props = utils.dummy_four_ranks(rank=rank) + if processor_props.comm_size != 4: + pytest.skip( + f"This test requires exactly 4 MPI ranks, current run has {processor_props.comm_size}" + ) + + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=processor_props, + ) + + decomp_info = halo_generator(utils.SIMPLE_DISTRIBUTION) + my_halo = decomp_info.global_index(dim, definitions.DecompositionInfo.EntryType.HALO) + print(f"rank = {processor_props.rank}: has halo {dim} : {my_halo}") + expected = len(utils.HALO[dim][processor_props.rank]) + assert ( + my_halo.size == expected + ), f" rank = {processor_props.rank}: total halo size does not match for dim {dim}- expected {expected} bot was {my_halo.size}" + assert ( + missing := np.setdiff1d( + my_halo, utils.HALO[dim][processor_props.rank], assume_unique=True + ).size + == 0 + ), f"missing halo elements are {missing}" + my_owned = decomp_info.global_index(dim, definitions.DecompositionInfo.EntryType.OWNED) + print(f"rank {processor_props.rank} owns {dim} : {my_owned} ") + utils.assert_same_entries(dim, my_owned, utils.OWNED, processor_props.rank) + + +@pytest.mark.parametrize("dim", [dims.EdgeDim]) +@pytest.mark.parametrize("rank", (0, 1, 2, 3)) +def test_halo_constructor_decomposition_info_halo_levels(rank, dim, simple_neighbor_tables): + processor_props = utils.DummyProps(rank=rank) + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=processor_props, + ) + decomp_info = halo_generator(utils.SIMPLE_DISTRIBUTION) + my_halo_levels = decomp_info.halo_levels(dim) + print(f"{dim.value}: rank {processor_props.rank} has halo levels {my_halo_levels} ") + assert np.all( + my_halo_levels != definitions.DecompositionFlag.UNDEFINED + ), "All indices should have a defined DecompositionFlag" + + assert np.where(my_halo_levels == definitions.DecompositionFlag.OWNED)[0].size == len( + utils.OWNED[dim][processor_props.rank] + ) + owned_local_indices = decomp_info.local_index( + dim, definitions.DecompositionInfo.EntryType.OWNED + ) + assert np.all( + my_halo_levels[owned_local_indices] == definitions.DecompositionFlag.OWNED + ), "owned local indices should have DecompositionFlag.OWNED" + first_halo_level_local_index = np.where( + my_halo_levels == definitions.DecompositionFlag.FIRST_HALO_LEVEL + )[0] + first_halo_level_global_index = decomp_info.global_index( + dim, definitions.DecompositionInfo.EntryType.ALL + )[first_halo_level_local_index] + utils.assert_same_entries( + dim, first_halo_level_global_index, utils.FIRST_HALO_LINE, processor_props.rank + ) + second_halo_level_local_index = np.where( + my_halo_levels == definitions.DecompositionFlag.SECOND_HALO_LEVEL + )[0] + second_halo_level_global_index = decomp_info.global_index( + dim, definitions.DecompositionInfo.EntryType.ALL + )[second_halo_level_local_index] + utils.assert_same_entries( + dim, second_halo_level_global_index, utils.SECOND_HALO_LINE, processor_props.rank + ) + third_halo_level_index = np.where( + my_halo_levels == definitions.DecompositionFlag.THIRD_HALO_LEVEL + )[0] + third_halo_level_global_index = decomp_info.global_index( + dim, definitions.DecompositionInfo.EntryType.ALL + )[third_halo_level_index] + utils.assert_same_entries( + dim, third_halo_level_global_index, utils.THIRD_HALO_LINE, processor_props.rank + ) + + +def test_no_halo(): + grid_size = base_grid.HorizontalGridSize(num_cells=9, num_edges=14, num_vertices=6) + halo_generator = halo.NoHalos(horizontal_size=grid_size, allocator=None) + decomposition = decomp.SingleNodeDecomposer() + decomposition_info = halo_generator(decomposition(np.arange(grid_size.num_cells), 1)) + # cells + np.testing.assert_allclose( + np.arange(grid_size.num_cells), decomposition_info.global_index(dims.CellDim) + ) + assert np.all(decomposition_info.owner_mask(dims.CellDim)) + assert np.all( + decomposition_info.halo_levels(dims.CellDim) == definitions.DecompositionFlag.OWNED + ) + # edges + np.testing.assert_allclose( + np.arange(grid_size.num_edges), decomposition_info.global_index(dims.EdgeDim) + ) + assert np.all( + decomposition_info.halo_levels(dims.EdgeDim) == definitions.DecompositionFlag.OWNED + ) + assert np.all(decomposition_info.owner_mask(dims.EdgeDim)) + # vertices + np.testing.assert_allclose( + np.arange(grid_size.num_vertices), decomposition_info.global_index(dims.VertexDim) + ) + assert np.all( + decomposition_info.halo_levels(dims.VertexDim) == definitions.DecompositionFlag.OWNED + ) + assert np.all(decomposition_info.owner_mask(dims.VertexDim)) + + +def test_halo_constructor_validate_rank_mapping_wrong_shape(simple_neighbor_tables): + processor_props = utils.DummyProps(rank=2) + num_cells = simple_neighbor_tables["C2E2C"].shape[0] + with pytest.raises(exceptions.ValidationError) as e: + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=processor_props, + ) + halo_generator(np.zeros((num_cells, 3), dtype=int)) + assert f"should have shape ({num_cells},)" in e.value.args[0] + + +@pytest.mark.parametrize("rank", (0, 1, 2, 3)) +def test_halo_constructor_validate_number_of_node_mismatch(rank, simple_neighbor_tables): + processor_props = utils.DummyProps(rank=rank) + num_cells = simple_neighbor_tables["C2E2C"].shape[0] + distribution = (processor_props.comm_size + 1) * np.ones((num_cells,), dtype=int) + with pytest.raises(expected_exception=exceptions.ValidationError) as e: + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=processor_props, + ) + halo_generator(distribution) + assert "The distribution assumes more nodes than the current run" in e.value.args[0] + + +@pytest.mark.parametrize("offset", offsets) +@pytest.mark.parametrize("rank", [0, 1, 2, 3]) +def test_global_to_local_index(offset, rank): + grid = simple.simple_grid() + neighbor_tables = { + k: v.ndarray + for k, v in grid.connectivities.items() + if gtx_common.is_neighbor_connectivity(v) + } + props = dummy_four_ranks(rank) + halo_constructor = halo.IconLikeHaloConstructor(props, neighbor_tables) + decomposition_info = halo_constructor(utils.SIMPLE_DISTRIBUTION) + source_indices_on_local_grid = decomposition_info.global_index(offset.target[0]) + + offset_full_grid = grid.connectivities[offset.value].ndarray[source_indices_on_local_grid] + neighbor_dim = offset.source + neighbor_index_full_grid = decomposition_info.global_index(neighbor_dim) + + local_offset = halo.global_to_local( + decomposition_info.global_index(neighbor_dim), offset_full_grid + ) + + ## assert by backmapping + + for i in range(local_offset.shape[0]): + for k in range(local_offset.shape[1]): + k_ = local_offset[i][k] + if k_ == -1: + # global index is not on this local patch: + assert not np.isin(offset_full_grid[i][k], neighbor_index_full_grid) + else: + ( + neighbor_index_full_grid[k_] == offset_full_grid[i][k], + f"failed to map [{offset_full_grid[i]}] to local: [{local_offset[i]}]", + ) + + +@pytest.mark.parametrize("rank", (0, 1, 2, 3)) +def test_horizontal_size(rank): + simple_neighbor_tables = get_neighbor_tables_for_simple_grid() + props = dummy_four_ranks(rank) + halo_generator = halo.IconLikeHaloConstructor( + connectivities=simple_neighbor_tables, + run_properties=props, + ) + decomp_info = halo_generator(utils.SIMPLE_DISTRIBUTION) + horizontal_size = decomp_info.get_horizontal_size() + expected_verts = len(utils.OWNED[dims.VertexDim][rank]) + len(utils.HALO[dims.VertexDim][rank]) + assert ( + horizontal_size.num_vertices == expected_verts + ), f"local size mismatch on rank={rank} for {dims.VertexDim}: expected {expected_verts}, but was {horizontal_size.num_vertices}" + expected_edges = len(utils.OWNED[dims.EdgeDim][rank]) + len(utils.HALO[dims.EdgeDim][rank]) + assert ( + horizontal_size.num_edges == expected_edges + ), f"local size mismatch on rank={rank} for {dims.EdgeDim}: expected {expected_edges}, but was {horizontal_size.num_edges}" + expected_cells = len(utils.OWNED[dims.CellDim][rank]) + len(utils.HALO[dims.CellDim][rank]) + assert ( + horizontal_size.num_cells == expected_cells + ), f"local size mismatch on rank={rank} for {dims.CellDim}: expected {expected_cells}, but was {horizontal_size.num_cells}" diff --git a/model/common/tests/common/decomposition/utils.py b/model/common/tests/common/decomposition/utils.py new file mode 100644 index 0000000000..93b71da871 --- /dev/null +++ b/model/common/tests/common/decomposition/utils.py @@ -0,0 +1,190 @@ +# 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 dataclasses +from typing import Any + +import numpy as np +from gt4py import next as gtx + +from icon4py.model.common import dimension as dims +from icon4py.model.common.decomposition import definitions + + +""" +TESTDATA using the +[SimpleGrid](../../../src/icon4py/model/common/grid/simple.py) The distribution +maps all of the 18 cells of the simple grid to ranks 0..3. + +The dictionaries contain the mapping from rank to global (in the simple grid) +index of the dimension: +- _CELL_OWN: rank -> owned cells, essentially the inversion of the + SIMPLE_DISTRIBUTION +- _EDGE_OWN: rank -> owned edges + - _VERTEX_OWN: rank -> owned vertices + +The decision as to whether a "secondary" dimension (edge, vertices) is owned by +a rank are made according to the rules and conventions described in +(../../../src/icon4py/model/common/decomposition/halo.py). + + +_CELL_FIRST_HALO_LINE: +_CELL_SECOND_HALO_LINE: +_EDGE_FIRST_HALO_LINE: +_EDGE_SECOND_HALO_LINE: +_VERTEX_FIRST_HALO_LINE: +_VERTEX_SECOND_HALO_LINE: :mapping of rank to global indices that belongs to a ranks halo lines. +""" + + +SIMPLE_DISTRIBUTION = np.asarray( + [ + 0, # 0c + 1, # 1c + 1, # 2c + 0, # 3c + 0, # 4c + 1, # 5c + 0, # 6c + 0, # 7c + 2, # 8c + 2, # 9c + 0, # 10c + 2, # 11c + 3, # 12c + 3, # 13c + 1, # 14c + 3, # 15c + 3, # 16c + 1, # 17c + ] +) +_CELL_OWN = {0: [0, 3, 4, 6, 7, 10], 1: [1, 2, 5, 14, 17], 2: [8, 9, 11], 3: [12, 13, 15, 16]} +_CELL_FIRST_HALO_LINE = { + 0: [1, 11, 13, 9, 2, 15], + 1: [3, 8, 4, 11, 16, 13, 15], + 2: [5, 7, 6, 12, 14], + 3: [9, 10, 17, 14, 0, 1], +} +_CELL_SECOND_HALO_LINE = { + 0: [17, 5, 12, 14, 8, 16], + 1: [0, 7, 6, 9, 10, 12], + 2: [2, 1, 4, 3, 10, 15, 16, 17, 13], + 3: [6, 7, 8, 2, 3, 4, 5, 11], +} +_CELL_HALO = { + 0: _CELL_FIRST_HALO_LINE[0] + _CELL_SECOND_HALO_LINE[0], + 1: _CELL_FIRST_HALO_LINE[1] + _CELL_SECOND_HALO_LINE[1], + 2: _CELL_FIRST_HALO_LINE[2] + _CELL_SECOND_HALO_LINE[2], + 3: _CELL_FIRST_HALO_LINE[3] + _CELL_SECOND_HALO_LINE[3], +} +_EDGE_OWN = { + 0: [1, 5, 12, 13, 14, 9], + 1: [8, 7, 6, 25, 4, 2], + 2: [16, 11, 15, 17, 10, 24], + 3: [19, 23, 22, 26, 0, 3, 20, 18, 21], +} +_EDGE_FIRST_HALO_LINE = {0: [0, 4, 17, 21, 10, 2], 1: [3, 15, 20, 26, 24], 2: [18], 3: []} +_EDGE_SECOND_HALO_LINE = { + 0: [3, 6, 7, 8, 15, 24, 25, 26, 16, 22, 23, 18, 19, 20, 11], + 1: [0, 1, 5, 9, 12, 11, 10, 13, 16, 17, 18, 19, 21, 22, 23], + 2: [2, 9, 12, 4, 8, 7, 14, 21, 13, 19, 20, 22, 23, 25, 26], + 3: [11, 10, 14, 13, 16, 17, 24, 25, 6, 2, 1, 5, 4, 8, 7], +} +_EDGE_THIRD_HALO_LINE = { + 0: [], + 1: [14], + 2: [0, 1, 3, 5, 6], + 3: [9, 12, 15], +} +_EDGE_HALO = { + 0: _EDGE_FIRST_HALO_LINE[0] + _EDGE_SECOND_HALO_LINE[0] + _EDGE_THIRD_HALO_LINE[0], + 1: _EDGE_FIRST_HALO_LINE[1] + _EDGE_SECOND_HALO_LINE[1] + _EDGE_THIRD_HALO_LINE[1], + 2: _EDGE_FIRST_HALO_LINE[2] + _EDGE_SECOND_HALO_LINE[2] + _EDGE_THIRD_HALO_LINE[2], + 3: _EDGE_FIRST_HALO_LINE[3] + _EDGE_SECOND_HALO_LINE[3] + _EDGE_THIRD_HALO_LINE[3], +} +_VERTEX_OWN = { + 0: [4], + 1: [], + 2: [3, 5], + 3: [ + 0, + 1, + 2, + 6, + 7, + 8, + ], +} +_VERTEX_FIRST_HALO_LINE = { + 0: [0, 1, 5, 8, 7, 3], + 1: [1, 2, 0, 5, 3, 8, 6], + 2: [ + 6, + 8, + 7, + ], + 3: [], +} +_VERTEX_SECOND_HALO_LINE = { + 0: [2, 6], + 1: [7, 4], + 2: [4, 0, 2, 1], + 3: [3, 4, 5], +} +_VERTEX_HALO = { + 0: _VERTEX_FIRST_HALO_LINE[0] + _VERTEX_SECOND_HALO_LINE[0], + 1: _VERTEX_FIRST_HALO_LINE[1] + _VERTEX_SECOND_HALO_LINE[1], + 2: _VERTEX_FIRST_HALO_LINE[2] + _VERTEX_SECOND_HALO_LINE[2], + 3: _VERTEX_FIRST_HALO_LINE[3] + _VERTEX_SECOND_HALO_LINE[3], +} +OWNED = {dims.CellDim: _CELL_OWN, dims.EdgeDim: _EDGE_OWN, dims.VertexDim: _VERTEX_OWN} +HALO = {dims.CellDim: _CELL_HALO, dims.EdgeDim: _EDGE_HALO, dims.VertexDim: _VERTEX_HALO} +FIRST_HALO_LINE = { + dims.CellDim: _CELL_FIRST_HALO_LINE, + dims.VertexDim: _VERTEX_FIRST_HALO_LINE, + dims.EdgeDim: _EDGE_FIRST_HALO_LINE, +} +SECOND_HALO_LINE = { + dims.CellDim: _CELL_SECOND_HALO_LINE, + dims.VertexDim: _VERTEX_SECOND_HALO_LINE, + dims.EdgeDim: _EDGE_SECOND_HALO_LINE, +} +THIRD_HALO_LINE = { + dims.CellDim: {0: [], 1: [], 2: [], 3: []}, + dims.VertexDim: {0: [], 1: [], 2: [], 3: []}, + dims.EdgeDim: _EDGE_THIRD_HALO_LINE, +} + + +def assert_same_entries( + dim: gtx.Dimension, my_owned: np.ndarray, reference: dict[gtx.Dimension, dict], rank: int +) -> None: + print(f"myowned {my_owned}") + print(reference[dim][rank]) + assert ( + my_owned.size == len(reference[dim][rank]) + ), f"{dim}(rank = {rank}) : wrong size expected {len(reference[dim][rank])} but was {my_owned.size}" + assert np.setdiff1d(my_owned, reference[dim][rank], assume_unique=True).size == 0 + + +@dataclasses.dataclass(frozen=True, init=False) +class DummyProps(definitions.ProcessProperties): + comm: Any + comm_name: str + comm_size: int + rank: int + + def __init__(self, rank: int): + object.__setattr__(self, "rank", rank % 4) + object.__setattr__(self, "comm", None) + object.__setattr__(self, "comm_name", "dummy on 4") + object.__setattr__(self, "comm_size", 4) + + +def dummy_four_ranks(rank: int) -> definitions.ProcessProperties: + return DummyProps(rank) diff --git a/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py b/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py index 617b7163e6..321fbee6ab 100644 --- a/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py +++ b/model/common/tests/common/grid/mpi_tests/test_parallel_geometry.py @@ -16,7 +16,7 @@ import pytest from icon4py.model.common import dimension as dims -from icon4py.model.common.decomposition import definitions as decomposition +from icon4py.model.common.decomposition import definitions as decomposition, mpi_decomposition from icon4py.model.common.grid import geometry, geometry_attributes as attrs, horizontal as h_grid from icon4py.model.common.math import helpers as math_helpers from icon4py.model.common.utils import data_allocation as data_alloc @@ -39,6 +39,10 @@ if TYPE_CHECKING: from icon4py.model.testing import serialbox as sb +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + + edge_domain = h_grid.domain(dims.EdgeDim) lb_local = edge_domain(h_grid.Zone.LOCAL) lb_lateral = edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) @@ -73,13 +77,11 @@ def test_distributed_geometry_attrs( parallel_helpers.check_comm_size(processor_props) parallel_helpers.log_process_properties(processor_props) parallel_helpers.log_local_field_size(decomposition_info) - grid_geometry = geometry_from_savepoint field_ref = grid_savepoint.__getattribute__(grid_name)().asnumpy() - field = grid_geometry.get(attrs_name).asnumpy() + field = geometry_from_savepoint.get(attrs_name).asnumpy() assert test_utils.dallclose(field, field_ref, atol=1e-12) -@pytest.mark.xfail(reason="Wrong results") @pytest.mark.datatest @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) diff --git a/model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py b/model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py new file mode 100644 index 0000000000..49b08790ed --- /dev/null +++ b/model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py @@ -0,0 +1,582 @@ +# 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 functools +import logging +import operator + +import numpy as np +import pytest +from gt4py import next as gtx +from gt4py.next import common as gtx_common, typing as gtx_typing + +from icon4py.model.common import dimension as dims, exceptions, model_backends +from icon4py.model.common.decomposition import ( + decomposer as decomp, + definitions as decomp_defs, + mpi_decomposition, +) +from icon4py.model.common.grid import ( + base, + geometry, + geometry_attributes, + grid_manager as gm, + gridfile, + icon, + vertical as v_grid, +) +from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory +from icon4py.model.common.metrics import metrics_attributes, metrics_factory +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions as test_defs, grid_utils, test_utils +from icon4py.model.testing.fixtures.datatest import ( + backend, + experiment, + processor_props, + topography_savepoint, +) + +from . import utils + + +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + +_log = logging.getLogger(__file__) + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.mpi(min_size=2) +def test_grid_manager_validate_decomposer( + processor_props: decomp_defs.ProcessProperties, + experiment: test_defs.Experiment, +) -> None: + if experiment == test_defs.Experiments.MCH_CH_R04B09: + pytest.xfail("Limited-area grids not yet supported") + + file = grid_utils.resolve_full_grid_file_name(experiment.grid) + manager = gm.GridManager( + grid_file=file, + config=v_grid.VerticalGridConfig(num_levels=utils.NUM_LEVELS), + offset_transformation=gridfile.ToZeroBasedIndexTransformation(), + ) + with pytest.raises(exceptions.InvalidConfigError) as e: + manager( + keep_skip_values=True, + allocator=None, + run_properties=processor_props, + decomposer=decomp.SingleNodeDecomposer(), + ) + + assert "Need a Decomposer for multi" in e.value.args[0] + + +def _get_neighbor_tables(grid: base.Grid) -> dict: + return { + k: v.ndarray + for k, v in grid.connectivities.items() + if gtx_common.is_neighbor_connectivity(v) + } + + +def gather_field(field: np.ndarray, props: decomp_defs.ProcessProperties) -> tuple: + constant_dims = tuple(field.shape[1:]) + _log.info(f"gather_field on rank={props.rank} - gathering field of local shape {field.shape}") + constant_length = functools.reduce(operator.mul, constant_dims, 1) + local_sizes = np.array(props.comm.gather(field.size, root=0)) + if props.rank == 0: + recv_buffer = np.empty(np.sum(local_sizes), dtype=field.dtype) + _log.info( + f"gather_field on rank = {props.rank} - setup receive buffer with size {sum(local_sizes)} on rank 0" + ) + else: + recv_buffer = None + + props.comm.Gatherv(sendbuf=field, recvbuf=(recv_buffer, local_sizes), root=0) + if props.rank == 0: + local_first_dim = tuple(sz // constant_length for sz in local_sizes) + _log.info( + f" gather_field on rank = 0: computed local dims {local_first_dim} - constant dims {constant_dims}" + ) + gathered_field = recv_buffer.reshape((-1, *constant_dims)) # type: ignore [union-attr] + else: + gathered_field = None + local_first_dim = field.shape + return local_first_dim, gathered_field + + +def check_local_global_field( + decomposition_info: decomp_defs.DecompositionInfo, + processor_props: decomp_defs.ProcessProperties, # F811 # fixture + dim: gtx.Dimension, + global_reference_field: np.ndarray, + local_field: np.ndarray, +) -> None: + _log.info( + f" rank= {processor_props.rank}/{processor_props.comm_size}----exchanging field of main dim {dim}" + ) + assert ( + local_field.shape[0] + == decomposition_info.global_index(dim, decomp_defs.DecompositionInfo.EntryType.ALL).shape[ + 0 + ] + ) + + # Compare full local field, including halos, against global reference field. + # TODO(msimberg): Is halo always expected to be populated? + global_indices_local_field = decomposition_info.global_index( + dim, + decomp_defs.DecompositionInfo.EntryType.OWNED, # ALL if checking halos + ) + local_indices_local_field = decomposition_info.local_index( + dim, + decomp_defs.DecompositionInfo.EntryType.OWNED, # ALL if checking halos + ) + local_field_from_global_field = global_reference_field[global_indices_local_field] + local_field_from_local_field = local_field[local_indices_local_field] + np.testing.assert_allclose( + local_field_from_global_field, local_field_from_local_field, atol=1e-9, verbose=True + ) + + # Compare owned local field, excluding halos, against global reference + # field, by gathering owned entries to the first rank. This ensures that in + # total we have the full global field distributed on all ranks. + owned_entries = local_field[ + decomposition_info.local_index(dim, decomp_defs.DecompositionInfo.EntryType.OWNED) + ] + gathered_sizes, gathered_field = gather_field(owned_entries, processor_props) + + global_index_sizes, gathered_global_indices = gather_field( + decomposition_info.global_index(dim, decomp_defs.DecompositionInfo.EntryType.OWNED), + processor_props, + ) + + if processor_props.rank == 0: + _log.info(f"rank = {processor_props.rank}: asserting gathered fields: ") + + assert np.all( + gathered_sizes == global_index_sizes + ), f"gathered field sizes do not match: {dim} {gathered_sizes} - {global_index_sizes}" + _log.info( + f"rank = {processor_props.rank}: Checking field size on dim ={dim}: --- gathered sizes {gathered_sizes} = {sum(gathered_sizes)}" + ) + _log.info( + f"rank = {processor_props.rank}: --- gathered field has size {gathered_sizes}" + ) + sorted_ = np.zeros(global_reference_field.shape, dtype=gtx.float64) # type: ignore [attr-defined] + sorted_[gathered_global_indices] = gathered_field + _log.info( + f" rank = {processor_props.rank}: SHAPES: global reference field {global_reference_field.shape}, gathered = {gathered_field.shape}" + ) + + # TODO(msimberg): Is this true? Not true for RBF itnerpolation... why? + # We expect an exact match, since the starting point is the same (grid + # file) and we are doing the exact same computations in single rank and + # multi rank mode. + np.testing.assert_allclose(sorted_, global_reference_field, atol=1e-9, verbose=True) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize( + "attrs_name, dim", + [ + # TODO(msimberg): Get dim out of field? + (geometry_attributes.CELL_AREA, dims.CellDim), + (geometry_attributes.EDGE_LENGTH, dims.EdgeDim), + (geometry_attributes.VERTEX_LAT, dims.VertexDim), + (geometry_attributes.EDGE_NORMAL_VERTEX_U, dims.EdgeDim), + (geometry_attributes.EDGE_NORMAL_VERTEX_V, dims.EdgeDim), + (geometry_attributes.EDGE_NORMAL_CELL_U, dims.EdgeDim), + (geometry_attributes.EDGE_NORMAL_CELL_V, dims.EdgeDim), + (geometry_attributes.EDGE_TANGENT_X, dims.EdgeDim), + (geometry_attributes.EDGE_TANGENT_Y, dims.EdgeDim), + ], +) +def test_geometry_fields_compare_single_multi_rank( + processor_props: decomp_defs.ProcessProperties, + backend: gtx_typing.Backend | None, + experiment: test_defs.Experiment, + attrs_name: str, + dim: gtx.Dimension, +) -> None: + if experiment == test_defs.Experiments.MCH_CH_R04B09: + pytest.xfail("Limited-area grids not yet supported") + + if ( + test_utils.is_embedded(backend) + and attrs_name + in ( + geometry_attributes.EDGE_NORMAL_VERTEX_U, + geometry_attributes.EDGE_NORMAL_VERTEX_V, + geometry_attributes.EDGE_NORMAL_CELL_U, + geometry_attributes.EDGE_NORMAL_CELL_V, + ) + and experiment == test_defs.Experiments.EXCLAIM_APE + ): + pytest.xfail("IndexOutOfBounds with embedded backend") + + # TODO(msimberg): Add fixtures for single/multi-rank + # grid/geometry/interpolation/metrics factories. + file = grid_utils.resolve_full_grid_file_name(experiment.grid) + _log.info(f"running on {processor_props.comm} with {processor_props.comm_size} ranks") + single_rank_grid_manager = utils.run_grid_manager_for_single_rank(file) + single_rank_geometry = geometry.GridGeometry( + backend=backend, + grid=single_rank_grid_manager.grid, + coordinates=single_rank_grid_manager.coordinates, + decomposition_info=single_rank_grid_manager.decomposition_info, + extra_fields=single_rank_grid_manager.geometry_fields, + metadata=geometry_attributes.attrs, + ) + _log.info( + f"rank = {processor_props.rank} : single node grid has size {single_rank_grid_manager.decomposition_info.get_horizontal_size()!r}" + ) + + multi_rank_grid_manager = utils.run_grid_manager_for_multi_rank( + file=file, + run_properties=processor_props, + decomposer=decomp.MetisDecomposer(), + ) + _log.info( + f"rank = {processor_props.rank} : {multi_rank_grid_manager.decomposition_info.get_horizontal_size()!r}" + ) + _log.info( + f"rank = {processor_props.rank}: halo size for 'CellDim' " + f"(1: {multi_rank_grid_manager.decomposition_info.get_halo_size(dims.CellDim, decomp_defs.DecompositionFlag.FIRST_HALO_LEVEL)}), " + f"(2: {multi_rank_grid_manager.decomposition_info.get_halo_size(dims.CellDim, decomp_defs.DecompositionFlag.SECOND_HALO_LEVEL)})" + ) + multi_rank_geometry = geometry.GridGeometry( + backend=backend, + grid=multi_rank_grid_manager.grid, + coordinates=multi_rank_grid_manager.coordinates, + decomposition_info=multi_rank_grid_manager.decomposition_info, + extra_fields=multi_rank_grid_manager.geometry_fields, + metadata=geometry_attributes.attrs, + ) + + check_local_global_field( + decomposition_info=multi_rank_grid_manager.decomposition_info, + processor_props=processor_props, + dim=dim, + global_reference_field=single_rank_geometry.get(attrs_name).asnumpy(), + local_field=multi_rank_geometry.get(attrs_name).asnumpy(), + ) + + _log.info(f"rank = {processor_props.rank} - DONE") + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize( + "attrs_name, dim", + [ + (interpolation_attributes.GEOFAC_DIV, dims.CellDim), + (interpolation_attributes.GEOFAC_ROT, dims.VertexDim), + (interpolation_attributes.C_BLN_AVG, dims.CellDim), + (interpolation_attributes.RBF_VEC_COEFF_C1, dims.CellDim), + (interpolation_attributes.RBF_VEC_COEFF_C2, dims.CellDim), + (interpolation_attributes.RBF_VEC_COEFF_E, dims.EdgeDim), + (interpolation_attributes.RBF_VEC_COEFF_V1, dims.VertexDim), + (interpolation_attributes.RBF_VEC_COEFF_V2, dims.VertexDim), + ], +) +def test_interpolation_fields_compare_single_multi_rank( + processor_props: decomp_defs.ProcessProperties, + backend: gtx_typing.Backend | None, + experiment: test_defs.Experiment, + attrs_name: str, + dim: gtx.Dimension, +) -> None: + if experiment == test_defs.Experiments.MCH_CH_R04B09: + pytest.xfail("Limited-area grids not yet supported") + + file = grid_utils.resolve_full_grid_file_name(experiment.grid) + _log.info(f"running on {processor_props.comm} with {processor_props.comm_size} ranks") + single_rank_grid_manager = utils.run_grid_manager_for_single_rank(file) + single_rank_geometry = geometry.GridGeometry( + backend=backend, + grid=single_rank_grid_manager.grid, + coordinates=single_rank_grid_manager.coordinates, + decomposition_info=single_rank_grid_manager.decomposition_info, + extra_fields=single_rank_grid_manager.geometry_fields, + metadata=geometry_attributes.attrs, + ) + single_rank_interpolation = interpolation_factory.InterpolationFieldsFactory( + grid=single_rank_grid_manager.grid, + decomposition_info=single_rank_grid_manager.decomposition_info, + geometry_source=single_rank_geometry, + backend=backend, + metadata=interpolation_attributes.attrs, + exchange=decomp_defs.SingleNodeExchange(), + ) + _log.info( + f"rank = {processor_props.rank} : single node grid has size {single_rank_grid_manager.decomposition_info.get_horizontal_size()!r}" + ) + + multi_rank_grid_manager = utils.run_grid_manager_for_multi_rank( + file=file, + run_properties=processor_props, + decomposer=decomp.MetisDecomposer(), + ) + _log.info( + f"rank = {processor_props.rank} : {multi_rank_grid_manager.decomposition_info.get_horizontal_size()!r}" + ) + _log.info( + f"rank = {processor_props.rank}: halo size for 'CellDim' " + f"(1: {multi_rank_grid_manager.decomposition_info.get_halo_size(dims.CellDim, decomp_defs.DecompositionFlag.FIRST_HALO_LEVEL)}), " + f"(2: {multi_rank_grid_manager.decomposition_info.get_halo_size(dims.CellDim, decomp_defs.DecompositionFlag.SECOND_HALO_LEVEL)})" + ) + multi_rank_geometry = geometry.GridGeometry( + backend=backend, + grid=multi_rank_grid_manager.grid, + coordinates=multi_rank_grid_manager.coordinates, + decomposition_info=multi_rank_grid_manager.decomposition_info, + extra_fields=multi_rank_grid_manager.geometry_fields, + metadata=geometry_attributes.attrs, + ) + multi_rank_interpolation = interpolation_factory.InterpolationFieldsFactory( + grid=multi_rank_grid_manager.grid, + decomposition_info=multi_rank_grid_manager.decomposition_info, + geometry_source=multi_rank_geometry, + backend=backend, + metadata=interpolation_attributes.attrs, + exchange=mpi_decomposition.GHexMultiNodeExchange( + processor_props, multi_rank_grid_manager.decomposition_info + ), + ) + + field_ref = single_rank_interpolation.get(attrs_name).asnumpy() + field = multi_rank_interpolation.get(attrs_name).asnumpy() + + check_local_global_field( + decomposition_info=multi_rank_grid_manager.decomposition_info, + processor_props=processor_props, + dim=dim, + global_reference_field=field_ref, + local_field=field, + ) + + _log.info(f"rank = {processor_props.rank} - DONE") + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize("attrs_name, dim", [(metrics_attributes.DDXT_Z_HALF_E, dims.EdgeDim)]) +def test_metrics_fields_compare_single_multi_rank( + processor_props: decomp_defs.ProcessProperties, + backend: gtx_typing.Backend | None, + experiment: test_defs.Experiment, + attrs_name: str, + dim: gtx.Dimension, +) -> None: + if experiment == test_defs.Experiments.MCH_CH_R04B09: + pytest.xfail("Limited-area grids not yet supported") + + file = grid_utils.resolve_full_grid_file_name(experiment.grid) + + ( + lowest_layer_thickness, + model_top_height, + stretch_factor, + damping_height, + rayleigh_coeff, + exner_expol, + vwind_offctr, + rayleigh_type, + thslp_zdiffu, + thhgtd_zdiffu, + ) = test_defs.construct_metrics_config(experiment) + vertical_config = v_grid.VerticalGridConfig( + experiment.num_levels, + lowest_layer_thickness=lowest_layer_thickness, + model_top_height=model_top_height, + stretch_factor=stretch_factor, + rayleigh_damping_height=damping_height, + ) + # TODO(msimberg): Dummy vct_a? Taken from test_io.py. + xp = data_alloc.import_array_ns(backend) + allocator = model_backends.get_allocator(backend) + vertical_grid = v_grid.VerticalGrid( + config=vertical_config, + vct_a=gtx.as_field( + (dims.KDim,), + xp.linspace(12000.0, 0.0, experiment.num_levels + 1), + allocator=allocator, + ), + vct_b=gtx.as_field( + (dims.KDim,), + xp.linspace(12000.0, 0.0, experiment.num_levels + 1), + allocator=allocator, + ), + ) + + _log.info(f"running on {processor_props.comm} with {processor_props.comm_size} ranks") + single_rank_grid_manager = utils.run_grid_manager_for_single_rank(file, experiment.num_levels) + single_rank_geometry = geometry.GridGeometry( + backend=backend, + grid=single_rank_grid_manager.grid, + coordinates=single_rank_grid_manager.coordinates, + decomposition_info=single_rank_grid_manager.decomposition_info, + extra_fields=single_rank_grid_manager.geometry_fields, + metadata=geometry_attributes.attrs, + ) + single_rank_interpolation = interpolation_factory.InterpolationFieldsFactory( + grid=single_rank_grid_manager.grid, + decomposition_info=single_rank_grid_manager.decomposition_info, + geometry_source=single_rank_geometry, + backend=backend, + metadata=interpolation_attributes.attrs, + exchange=decomp_defs.SingleNodeExchange(), + ) + single_rank_metrics = metrics_factory.MetricsFieldsFactory( + grid=single_rank_geometry.grid, + vertical_grid=vertical_grid, + decomposition_info=single_rank_grid_manager.decomposition_info, + geometry_source=single_rank_geometry, + # TODO(msimberg): Valid dummy topography? + topography=( + gtx.as_field( + (dims.CellDim,), + xp.zeros(single_rank_geometry.grid.num_cells), + allocator=allocator, + ) + ), + interpolation_source=single_rank_interpolation, + backend=backend, + metadata=metrics_attributes.attrs, + rayleigh_type=rayleigh_type, + rayleigh_coeff=rayleigh_coeff, + exner_expol=exner_expol, + vwind_offctr=vwind_offctr, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + exchange=decomp_defs.SingleNodeExchange(), + ) + _log.info( + f"rank = {processor_props.rank} : single node grid has size {single_rank_grid_manager.decomposition_info.get_horizontal_size()!r}" + ) + + multi_rank_grid_manager = utils.run_grid_manager_for_multi_rank( + file=file, + run_properties=processor_props, + decomposer=decomp.MetisDecomposer(), + num_levels=experiment.num_levels, + ) + _log.info( + f"rank = {processor_props.rank} : {multi_rank_grid_manager.decomposition_info.get_horizontal_size()!r}" + ) + _log.info( + f"rank = {processor_props.rank}: halo size for 'CellDim' " + f"(1: {multi_rank_grid_manager.decomposition_info.get_halo_size(dims.CellDim, decomp_defs.DecompositionFlag.FIRST_HALO_LEVEL)}), " + f"(2: {multi_rank_grid_manager.decomposition_info.get_halo_size(dims.CellDim, decomp_defs.DecompositionFlag.SECOND_HALO_LEVEL)})" + ) + multi_rank_geometry = geometry.GridGeometry( + backend=backend, + grid=multi_rank_grid_manager.grid, + coordinates=multi_rank_grid_manager.coordinates, + decomposition_info=multi_rank_grid_manager.decomposition_info, + extra_fields=multi_rank_grid_manager.geometry_fields, + metadata=geometry_attributes.attrs, + ) + multi_rank_interpolation = interpolation_factory.InterpolationFieldsFactory( + grid=multi_rank_grid_manager.grid, + decomposition_info=multi_rank_grid_manager.decomposition_info, + geometry_source=multi_rank_geometry, + backend=backend, + metadata=interpolation_attributes.attrs, + exchange=mpi_decomposition.GHexMultiNodeExchange( + processor_props, multi_rank_grid_manager.decomposition_info + ), + ) + multi_rank_metrics = metrics_factory.MetricsFieldsFactory( + grid=multi_rank_geometry.grid, + vertical_grid=vertical_grid, + decomposition_info=multi_rank_grid_manager.decomposition_info, + geometry_source=multi_rank_geometry, + # TODO(msimberg): Valid dummy topography? + topography=( + gtx.as_field( + (dims.CellDim,), + xp.zeros(multi_rank_geometry.grid.num_cells), + allocator=allocator, + ) + ), + interpolation_source=multi_rank_interpolation, + backend=backend, + metadata=metrics_attributes.attrs, + rayleigh_type=rayleigh_type, + rayleigh_coeff=rayleigh_coeff, + exner_expol=exner_expol, + vwind_offctr=vwind_offctr, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + exchange=mpi_decomposition.GHexMultiNodeExchange( + processor_props, multi_rank_grid_manager.decomposition_info + ), + ) + + field_ref = single_rank_metrics.get(attrs_name).asnumpy() + field = multi_rank_metrics.get(attrs_name).asnumpy() + + check_local_global_field( + decomposition_info=multi_rank_grid_manager.decomposition_info, + processor_props=processor_props, + dim=dim, + global_reference_field=field_ref, + local_field=field, + ) + + _log.info(f"rank = {processor_props.rank} - DONE") + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_validate_skip_values_in_distributed_connectivities( + processor_props: decomp_defs.ProcessProperties, + experiment: test_defs.Experiment, +) -> None: + if experiment == test_defs.Experiments.MCH_CH_R04B09: + pytest.xfail("Limited-area grids not yet supported") + + file = grid_utils.resolve_full_grid_file_name(experiment.grid) + multi_rank_grid_manager = utils.run_grid_manager_for_multi_rank( + file=file, + run_properties=processor_props, + decomposer=decomp.MetisDecomposer(), + ) + distributed_grid = multi_rank_grid_manager.grid + for k, c in distributed_grid.connectivities.items(): + if gtx_common.is_neighbor_connectivity(c): + skip_values_in_table = np.count_nonzero(c.asnumpy() == c.skip_value) + found_skips = skip_values_in_table > 0 + assert ( + found_skips == (c.skip_value is not None) + ), f"rank={processor_props.rank} / {processor_props.comm_size}: {k} - # of skip values found in table = {skip_values_in_table}, skip value is {c.skip_value}" + if skip_values_in_table > 0: + dim = gtx.Dimension(k, gtx.DimensionKind.LOCAL) + assert ( + dim in icon.CONNECTIVITIES_ON_BOUNDARIES + or dim in icon.CONNECTIVITIES_ON_PENTAGONS + ), f"rank={processor_props.rank} / {processor_props.comm_size}: {k} has skip found in table, expected none" + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize("grid", [test_defs.Grids.MCH_CH_R04B09_DSL]) +def test_limited_area_raises( + processor_props: decomp_defs.ProcessProperties, + grid: test_defs.GridDescription, +) -> None: + with pytest.raises( + NotImplementedError, match="Limited-area grids are not supported in distributed runs" + ): + _ = utils.run_grid_manager_for_multi_rank( + file=grid_utils.resolve_full_grid_file_name(grid), + run_properties=processor_props, + decomposer=decomp.MetisDecomposer(), + ) diff --git a/model/common/tests/common/grid/mpi_tests/test_parallel_grid_refinement.py b/model/common/tests/common/grid/mpi_tests/test_parallel_grid_refinement.py new file mode 100644 index 0000000000..2b68be3e24 --- /dev/null +++ b/model/common/tests/common/grid/mpi_tests/test_parallel_grid_refinement.py @@ -0,0 +1,67 @@ +# 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 logging + +import gt4py.next as gtx +import pytest + +from icon4py.model.common.decomposition import definitions as decomposition, mpi_decomposition +from icon4py.model.common.grid import grid_refinement, horizontal as h_grid +from icon4py.model.testing import definitions, serialbox +from icon4py.model.testing.fixtures.datatest import ( + backend, + data_provider, + download_ser_data, + experiment, + grid_savepoint, + processor_props, +) + +from .. import utils + + +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + +_log = logging.getLogger(__name__) + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize("dim", utils.main_horizontal_dims()) +def test_compute_domain_bounds( + dim: gtx.Dimension, + experiment: definitions.Experiment, + grid_savepoint: serialbox.IconGridSavepoint, + processor_props: decomposition.ProcessProperties, +) -> None: + if processor_props.is_single_rank() and experiment == definitions.Experiments.EXCLAIM_APE: + pytest.xfail( + "end index data for single node APE are all 0 - re- serialization should fix that (patch%cells%end_index vs patch%cells%end_idx)" + ) + + ref_grid = grid_savepoint.construct_icon_grid(backend=None, keep_skip_values=True) + decomposition_info = grid_savepoint.construct_decomposition_info() + refin_ctrl = {dim: grid_savepoint.refin_ctrl(dim) for dim in utils.main_horizontal_dims()} + start_indices, end_indices = grid_refinement.compute_domain_bounds( + dim, refin_ctrl, decomposition_info + ) + for domain in h_grid.get_domains_for_dim(dim): + ref_start_index = ref_grid.start_index(domain) + ref_end_index = ref_grid.end_index(domain) + computed_start = start_indices[domain] + computed_end = end_indices[domain] + _log.info( + f"rank = {processor_props.rank}/{processor_props.comm_size}: domain={domain} : start = {computed_start} end = {computed_end} " + ) + assert ( + computed_start == ref_start_index + ), f"rank={processor_props.rank}/{processor_props.comm_size} - experiment = {experiment.name}: start_index for {domain} does not match: is {computed_start}, expected {ref_start_index}" + assert ( + computed_end == ref_end_index + ), f"rank={processor_props.rank}/{processor_props.comm_size} - experiment = {experiment.name}: end_index for {domain} does not match: is {computed_end}, expected {ref_end_index}" diff --git a/model/common/tests/common/grid/mpi_tests/test_parallel_icon.py b/model/common/tests/common/grid/mpi_tests/test_parallel_icon.py index 01b31899ea..d8bf542c55 100644 --- a/model/common/tests/common/grid/mpi_tests/test_parallel_icon.py +++ b/model/common/tests/common/grid/mpi_tests/test_parallel_icon.py @@ -8,24 +8,31 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING +import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.dimension as dims import icon4py.model.common.grid.horizontal as h_grid -from icon4py.model.common.decomposition import definitions as decomp_defs -from icon4py.model.testing import definitions as test_defs, parallel_helpers +from icon4py.model.common.decomposition import definitions as decomposition, mpi_decomposition +from icon4py.model.common.decomposition.decomposer import MetisDecomposer +from icon4py.model.common.grid import base as base_grid, gridfile, horizontal as h_grid +from icon4py.model.testing import definitions as test_defs, grid_utils, parallel_helpers from ...fixtures import ( backend, data_provider, download_ser_data, + experiment, grid_savepoint, icon_grid, processor_props, ) from .. import utils +from . import utils as parallel_utils if TYPE_CHECKING: @@ -33,26 +40,24 @@ from icon4py.model.common.grid import base as base_grid +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) -try: - import mpi4py # type: ignore[import-not-found] # F401: import mpi4py to check for optional mpi dependency +from mpi4py import MPI - from icon4py.model.common.decomposition import mpi_decomposition -except ImportError: - pytest.skip("Skipping parallel on single node installation", allow_module_level=True) + +_log = logging.getLogger(__name__) @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -def test_props(processor_props: decomp_defs.ProcessProperties) -> None: +def test_props(processor_props: decomposition.ProcessProperties) -> None: """dummy test to check whether the MPI initialization and GHEX setup works.""" import ghex.context as ghex # type: ignore[import-not-found] assert processor_props.comm - assert isinstance( - processor_props.comm, mpi4py.MPI.Comm - ), "comm needs to be an instance of MPI.Comm" + assert isinstance(processor_props.comm, MPI.Comm), "comm needs to be an instance of MPI.Comm" ghex.make_context(processor_props.comm) @@ -81,20 +86,20 @@ def test_props(processor_props: decomp_defs.ProcessProperties) -> None: ], ) @pytest.mark.parametrize("dim", utils.main_horizontal_dims()) -def test_distributed_local( - processor_props: decomp_defs.ProcessProperties, +def test_start_index_end_index_local_zone_on_distributed_lam_grid( + processor_props: decomposition.ProcessProperties, dim: gtx.Dimension, icon_grid: base_grid.Grid, experiment: test_defs.Experiment, ) -> None: parallel_helpers.check_comm_size(processor_props) domain = h_grid.domain(dim)(h_grid.Zone.LOCAL) - print( + _log.info( f"rank {processor_props.rank}/{processor_props.comm_size} dim = {dim} : {icon_grid.size[dim]}" ) # local still runs entire field: assert icon_grid.start_index(domain) == 0 - print( + _log.info( f"rank {processor_props.rank}/{processor_props.comm_size} dim = {dim} LOCAL : ({icon_grid.start_index(domain)}, {icon_grid.end_index(domain)})" ) assert ( @@ -152,8 +157,8 @@ def test_distributed_local( ], ) @pytest.mark.parametrize("zone, level", [(h_grid.Zone.HALO, 1), (h_grid.Zone.HALO_LEVEL_2, 2)]) -def test_distributed_halo( - processor_props: decomp_defs.ProcessProperties, +def test_start_index_end_index_halo_zones_on_distributed_lam_grid( + processor_props: decomposition.ProcessProperties, dim: gtx.Dimension, zone: h_grid.Zone, icon_grid: base_grid.Grid, @@ -165,7 +170,7 @@ def test_distributed_halo( start_index = icon_grid.start_index(domain) end_index = icon_grid.end_index(domain) rank = processor_props.rank - print( + _log.info( f"rank {rank}/{processor_props.comm_size} dim = {dim} {zone} : ({start_index}, {end_index})" ) expected = HALO_IDX[processor_props.comm_size][dim][rank][level - 1] diff --git a/model/common/tests/common/grid/mpi_tests/utils.py b/model/common/tests/common/grid/mpi_tests/utils.py new file mode 100644 index 0000000000..511ec82f77 --- /dev/null +++ b/model/common/tests/common/grid/mpi_tests/utils.py @@ -0,0 +1,47 @@ +# 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 pathlib + +from icon4py.model.common.decomposition import decomposer as decomp, definitions as decomp_defs +from icon4py.model.common.grid import grid_manager as gm, vertical as v_grid + + +def _grid_manager(file: pathlib.Path, num_levels: int) -> gm.GridManager: + return gm.GridManager( + grid_file=str(file), config=v_grid.VerticalGridConfig(num_levels=num_levels) + ) + + +NUM_LEVELS = 10 + + +def run_grid_manager_for_single_rank( + file: pathlib.Path, num_levels: int = NUM_LEVELS +) -> gm.GridManager: + manager = _grid_manager(file, num_levels) + manager( + keep_skip_values=True, + run_properties=decomp_defs.SingleNodeProcessProperties(), + decomposer=decomp.SingleNodeDecomposer(), + allocator=None, + ) + return manager + + +def run_grid_manager_for_multi_rank( + file: pathlib.Path, + run_properties: decomp_defs.ProcessProperties, + decomposer: decomp.Decomposer, + num_levels: int = NUM_LEVELS, +) -> gm.GridManager: + manager = _grid_manager(file, num_levels) + manager( + keep_skip_values=True, allocator=None, run_properties=run_properties, decomposer=decomposer + ) + return manager diff --git a/model/common/tests/common/grid/unit_tests/test_grid_manager.py b/model/common/tests/common/grid/unit_tests/test_grid_manager.py index 55dbfbb48d..24e915c41a 100644 --- a/model/common/tests/common/grid/unit_tests/test_grid_manager.py +++ b/model/common/tests/common/grid/unit_tests/test_grid_manager.py @@ -9,22 +9,27 @@ import logging import typing +from collections.abc import Iterator import gt4py.next as gtx import gt4py.next.typing as gtx_typing import numpy as np import pytest -from icon4py.model.common import dimension as dims +import icon4py.model.common.grid.gridfile +from icon4py.model.common import dimension as dims, model_backends +from icon4py.model.common.decomposition import decomposer as decomp, definitions as decomp_defs from icon4py.model.common.grid import ( base as base_grid, grid_manager as gm, grid_refinement as refin, gridfile, horizontal as h_grid, + icon, vertical as v_grid, ) -from icon4py.model.testing import definitions, test_utils +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import definitions, definitions as test_defs, grid_utils, test_utils if typing.TYPE_CHECKING: @@ -40,6 +45,7 @@ from icon4py.model.testing.fixtures import ( backend, + backend_like, cpu_allocator, data_provider, download_ser_data, @@ -48,15 +54,13 @@ processor_props, ) +from ...decomposition import utils as decomp_utils from .. import utils MCH_CH_RO4B09_GLOBAL_NUM_CELLS = 83886080 -ZERO_BASE = gm.ToZeroBasedIndexTransformation() - - # TODO @magdalena add test cases for hexagon vertices v2e2v # v2e2v: grid,??? @@ -342,7 +346,9 @@ def test_gridmanager_given_file_not_found_then_abort( fname = "./unknown_grid.nc" with pytest.raises(FileNotFoundError) as error: manager = gm.GridManager( - gm.NoTransformation(), fname, v_grid.VerticalGridConfig(num_levels=80) + grid_file=fname, + config=v_grid.VerticalGridConfig(num_levels=80), + offset_transformation=icon4py.model.common.grid.gridfile.NoTransformation(), ) manager(allocator=cpu_allocator, keep_skip_values=True) assert error.value == 1 @@ -351,7 +357,7 @@ def test_gridmanager_given_file_not_found_then_abort( @pytest.mark.parametrize("size", [100, 1500, 20000]) @pytest.mark.with_netcdf def test_gt4py_transform_offset_by_1_where_valid(size: int) -> None: - trafo = gm.ToZeroBasedIndexTransformation() + trafo = gridfile.ToZeroBasedIndexTransformation() rng = np.random.default_rng() input_field = rng.integers(-1, size, size) offset = trafo(input_field) @@ -552,3 +558,124 @@ def test_edge_vertex_distance( def test_limited_area_on_grid(grid_descriptor: definitions.GridDescription, expected: bool) -> None: grid = utils.run_grid_manager(grid_descriptor, keep_skip_values=True, backend=None).grid assert expected == grid.limited_area + + +@pytest.mark.datatest +@pytest.mark.parametrize("dim", utils.horizontal_dims()) +def test_decomposition_info_single_rank( + dim: gtx.Dimension, + experiment: definitions.Experiment, + grid_savepoint: serialbox.IconGridSavepoint, + backend: gtx_typing.Backend, +) -> None: + expected = grid_savepoint.construct_decomposition_info() + grid_file = experiment.grid + gm = utils.run_grid_manager(grid_file, keep_skip_values=True, backend=backend) + result = gm.decomposition_info + assert np.all(data_alloc.as_numpy(result.local_index(dim)) == expected.local_index(dim)) + assert np.all(data_alloc.as_numpy(result.global_index(dim)) == expected.global_index(dim)) + assert np.all(data_alloc.as_numpy(result.owner_mask(dim)) == expected.owner_mask(dim)) + assert np.all(data_alloc.as_numpy(result.halo_levels(dim)) == expected.halo_levels(dim)) + + +@pytest.mark.parametrize("rank", (0, 1, 2, 3), ids=lambda rank: f"rank{rank}") +@pytest.mark.parametrize( + "field_offset", + [ + dims.C2V, + dims.E2V, + dims.V2C, + dims.E2C, + dims.C2E, + dims.V2E, + dims.C2E2C, + dims.V2E2V, + dims.C2E2CO, + dims.C2E2C2E, + dims.C2E2C2E2C, + dims.E2C2V, + dims.E2C2E, + dims.E2C2EO, + ], + ids=lambda offset: offset.value, +) +def test_local_connectivity( + rank: int, + caplog: Iterator, + field_offset: gtx.FieldOffset, + backend_like: model_backends.BackendLike, +) -> None: + processor_props = decomp_utils.DummyProps(rank=rank) + caplog.set_level(logging.INFO) # type: ignore [attr-defined] + partitioner = decomp.MetisDecomposer() + allocator = model_backends.get_allocator(backend_like) + file = grid_utils.resolve_full_grid_file_name(test_defs.Grids.R02B04_GLOBAL) + manager = gm.GridManager(config=v_grid.VerticalGridConfig(num_levels=10), grid_file=file) + manager( + decomposer=partitioner, + allocator=allocator, + keep_skip_values=True, + run_properties=processor_props, + ) + grid = manager.grid + + decomposition_info = manager.decomposition_info + connectivity = grid.get_connectivity(field_offset).asnumpy() + + assert ( + connectivity.shape[0] + == decomposition_info.global_index( + field_offset.target[0], decomp_defs.DecompositionInfo.EntryType.ALL + ).size + ), "connectivity shapes do not match" + + # all neighbor indices are valid local indices + max_local_index = np.max( + decomposition_info.local_index( + field_offset.source, decomp_defs.DecompositionInfo.EntryType.ALL + ) + ) + assert ( + np.max(connectivity) == max_local_index + ), f"max value in the connectivity is {np.max(connectivity)} is larger than the local patch size {max_local_index}" + # - outer halo entries have SKIP_VALUE neighbors (depends on offsets) + neighbor_dim = field_offset.target[1] # type: ignore [misc] + dim = field_offset.target[0] + last_halo_level = ( + decomp_defs.DecompositionFlag.THIRD_HALO_LEVEL + if dim == dims.EdgeDim + else decomp_defs.DecompositionFlag.SECOND_HALO_LEVEL + ) + level_index = np.where( + data_alloc.as_numpy(decomposition_info.halo_levels(dim)) == last_halo_level.value + ) + if ( + neighbor_dim in icon.CONNECTIVITIES_ON_BOUNDARIES + or neighbor_dim in icon.CONNECTIVITIES_ON_PENTAGONS + ): + assert np.count_nonzero( + (connectivity[level_index] == gridfile.GridFile.INVALID_INDEX) > 0 + ), f"missing invalid index in {dim} - offset {field_offset}" + else: + assert np.count_nonzero( + (connectivity[level_index] == gridfile.GridFile.INVALID_INDEX) == 0 + ), f"have invalid index in {dim} - offset {field_offset} when none expected" + + +@pytest.mark.parametrize("ranks", (2, 3, 4)) +def test_decomposition_size( + ranks: int, + experiment: test_defs.Experiment, +) -> None: + if experiment == test_defs.Experiments.MCH_CH_R04B09: + pytest.xfail("Limited-area grids not yet supported") + + decomposer = decomp.MetisDecomposer() + file = grid_utils.resolve_full_grid_file_name(experiment.grid) + with gridfile.GridFile(str(file), gridfile.ToZeroBasedIndexTransformation()) as parser: + partitions = decomposer(parser.int_variable(gridfile.ConnectivityName.C2E2C), ranks) + sizes = [np.count_nonzero(partitions == r) for r in range(ranks)] + # Verify that sizes are close to each other. This is not a hard + # requirement, but simply a sanity check to make sure that partitions + # are relatively balanced. + assert max(sizes) - min(sizes) <= 2 diff --git a/model/common/tests/common/grid/unit_tests/test_grid_refinement.py b/model/common/tests/common/grid/unit_tests/test_grid_refinement.py index d1f7cae5d7..7ec901ca84 100644 --- a/model/common/tests/common/grid/unit_tests/test_grid_refinement.py +++ b/model/common/tests/common/grid/unit_tests/test_grid_refinement.py @@ -21,9 +21,6 @@ from .. import utils -_FALLBACK_FAIL = (-10, -10) - - @pytest.mark.parametrize("dim", utils.main_horizontal_dims()) @pytest.mark.parametrize( "grid_file, expected", @@ -48,7 +45,7 @@ def test_is_local_area_grid_for_grid_files( assert expected == limited_area -cell_bounds: dict[h_grid.Zone, tuple[int, int]] = { +MCH_OPR_R04B07_CELL_BOUNDS: dict[h_grid.Zone, tuple[int, int]] = { h_grid.Zone.LATERAL_BOUNDARY: (0, 629), h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2: (629, 1244), h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3: (1244, 1843), @@ -60,7 +57,7 @@ def test_is_local_area_grid_for_grid_files( h_grid.Zone.HALO: (10700, 10700), h_grid.Zone.HALO_LEVEL_2: (10700, 10700), } -edge_bounds: dict[h_grid.Zone, tuple[int, int]] = { +MCH_OPR_R04B07_EDGE_BOUNDS: dict[h_grid.Zone, tuple[int, int]] = { h_grid.Zone.LATERAL_BOUNDARY: (0, 318), h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2: (318, 947), h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3: (947, 1258), @@ -77,7 +74,7 @@ def test_is_local_area_grid_for_grid_files( h_grid.Zone.HALO: (16209, 16209), h_grid.Zone.HALO_LEVEL_2: (16209, 16209), } -vertex_bounds: dict[h_grid.Zone, tuple[int, int]] = { +MCH_OPR_R04B07_VERTEX_BOUNDS: dict[h_grid.Zone, tuple[int, int]] = { h_grid.Zone.LATERAL_BOUNDARY: (0, 318), h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2: (318, 629), h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3: (629, 933), @@ -91,20 +88,80 @@ def test_is_local_area_grid_for_grid_files( } +MCH_CH_R04B09_CELL_BOUNDS: dict[h_grid.Zone, tuple[int, int]] = { + h_grid.Zone.LATERAL_BOUNDARY: (0, 850), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2: (850, 1688), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3: (1688, 2511), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4: (2511, 3316), + h_grid.Zone.NUDGING: (3316, 4104), + h_grid.Zone.INTERIOR: (4104, 20896), + h_grid.Zone.LOCAL: (0, 20896), + h_grid.Zone.END: (20896, 20896), + h_grid.Zone.HALO: (20896, 20896), + h_grid.Zone.HALO_LEVEL_2: (20896, 20896), +} +MCH_CH_R04B09_EDGE_BOUNDS: dict[h_grid.Zone, tuple[int, int]] = { + h_grid.Zone.LATERAL_BOUNDARY: (0, 428), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2: (428, 1278), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3: (1278, 1700), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4: (1700, 2538), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5: (2538, 2954), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_6: (2954, 3777), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7: (3777, 4184), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_8: (4184, 4989), + h_grid.Zone.NUDGING: (4989, 5387), + h_grid.Zone.NUDGING_LEVEL_2: (5387, 6176), + h_grid.Zone.INTERIOR: (6176, 31558), + h_grid.Zone.LOCAL: (0, 31558), + h_grid.Zone.END: (31558, 31558), + h_grid.Zone.HALO: (31558, 31558), + h_grid.Zone.HALO_LEVEL_2: (31558, 31558), +} +MCH_CH_R04B09_VERTEX_BOUNDS: dict[h_grid.Zone, tuple[int, int]] = { + h_grid.Zone.LATERAL_BOUNDARY: (0, 428), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2: (428, 850), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3: (850, 1266), + h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4: (1266, 1673), + h_grid.Zone.NUDGING: (1673, 2071), + h_grid.Zone.INTERIOR: (2071, 10663), + h_grid.Zone.LOCAL: (0, 10663), + h_grid.Zone.END: (10663, 10663), + h_grid.Zone.HALO: (10663, 10663), + h_grid.Zone.HALO_LEVEL_2: (10663, 10663), +} +#: random invalid values to mark failure for start_/end_index tuple +_FALLBACK_FAIL = (-10, -10) + + @pytest.mark.parametrize( - "dim, expected", - [(dims.CellDim, cell_bounds), (dims.EdgeDim, edge_bounds), (dims.VertexDim, vertex_bounds)], + "grid_description, dim, expected", + [ + (test_defs.Grids.MCH_OPR_R04B07_DOMAIN01, dims.CellDim, MCH_OPR_R04B07_CELL_BOUNDS), + (test_defs.Grids.MCH_OPR_R04B07_DOMAIN01, dims.EdgeDim, MCH_OPR_R04B07_EDGE_BOUNDS), + (test_defs.Grids.MCH_OPR_R04B07_DOMAIN01, dims.VertexDim, MCH_OPR_R04B07_VERTEX_BOUNDS), + (test_defs.Grids.MCH_CH_R04B09_DSL, dims.CellDim, MCH_CH_R04B09_CELL_BOUNDS), + (test_defs.Grids.MCH_CH_R04B09_DSL, dims.EdgeDim, MCH_CH_R04B09_EDGE_BOUNDS), + (test_defs.Grids.MCH_CH_R04B09_DSL, dims.VertexDim, MCH_CH_R04B09_VERTEX_BOUNDS), + ], ) -def test_compute_start_index_for_limited_area_grid( +def test_compute_domain_bounds_for_limited_area_grid( + grid_description: test_defs.GridDescription, dim: gtx.Dimension, expected: dict[h_grid.Zone, tuple[int, int]], cpu_allocator: gtx_typing.FieldBufferAllocationUtil, ) -> None: - grid = grid_utils.get_grid_manager_from_identifier( - test_defs.Grids.MCH_OPR_R04B07_DOMAIN01, 1, True, cpu_allocator - ).grid + grid_manager = grid_utils.get_grid_manager_from_identifier( + grid_description, 1, True, cpu_allocator + ) + + grid = grid_manager.grid + assert grid.limited_area, "Test expects limited area grid" refinement_field = grid.refinement_control - start_index, end_index = refinement.compute_domain_bounds(dim, refinement_field, array_ns=np) + decomposition_info = grid_manager.decomposition_info + + start_index, end_index = refinement.compute_domain_bounds( + dim, refinement_field, decomposition_info=decomposition_info, array_ns=np + ) for d, v in start_index.items(): expected_value = expected.get(d.zone, _FALLBACK_FAIL)[0] @@ -126,9 +183,13 @@ def test_compute_domain_bounds_for_global_grid( dim: gtx.Dimension, cpu_allocator: gtx_typing.FieldBufferAllocationUtil, ) -> None: - grid = grid_utils.get_grid_manager_from_identifier(file, 1, True, cpu_allocator).grid + grid_manager = grid_utils.get_grid_manager_from_identifier(file, 1, True, cpu_allocator) + grid = grid_manager.grid refinement_fields = grid.refinement_control - start_index, end_index = refinement.compute_domain_bounds(dim, refinement_fields, array_ns=np) + decomposition_info = grid_manager.decomposition_info + start_index, end_index = refinement.compute_domain_bounds( + dim, refinement_fields, decomposition_info, array_ns=np + ) for k, v in start_index.items(): assert isinstance(v, gtx.int32) if k.zone.is_halo() or k.zone is h_grid.Zone.END: diff --git a/model/common/tests/common/grid/unit_tests/test_gridfile.py b/model/common/tests/common/grid/unit_tests/test_gridfile.py index 9797b941ef..c8a4ca034c 100644 --- a/model/common/tests/common/grid/unit_tests/test_gridfile.py +++ b/model/common/tests/common/grid/unit_tests/test_gridfile.py @@ -7,8 +7,10 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations +from collections.abc import Iterable from typing import TYPE_CHECKING +import numpy as np import pytest from icon4py.model.common import dimension as dims @@ -32,14 +34,19 @@ def test_grid_file_dimension() -> None: grid_descriptor = definitions.Grids.R02B04_GLOBAL global_grid_file = str(gridtest_utils.resolve_full_grid_file_name(grid_descriptor)) - parser = gridfile.GridFile(global_grid_file) + parser = gridfile.GridFile(global_grid_file, offset_transformation=gridfile.NoTransformation()) try: parser.open() - assert parser.dimension(gridfile.DimensionName.CELL_NAME) == grid_descriptor.sizes["cell"] assert ( - parser.dimension(gridfile.DimensionName.VERTEX_NAME) == grid_descriptor.sizes["vertex"] + parser.dimension(gridfile.DynamicDimension.CELL_NAME) == grid_descriptor.sizes["cell"] + ) + assert ( + parser.dimension(gridfile.DynamicDimension.VERTEX_NAME) + == grid_descriptor.sizes["vertex"] + ) + assert ( + parser.dimension(gridfile.DynamicDimension.EDGE_NAME) == grid_descriptor.sizes["edge"] ) - assert parser.dimension(gridfile.DimensionName.EDGE_NAME) == grid_descriptor.sizes["edge"] except Exception: pytest.fail() finally: @@ -52,19 +59,87 @@ def test_grid_file_vertex_cell_edge_dimensions( experiment: definitions.Experiment, grid_savepoint: serialbox.IconGridSavepoint ) -> None: file = gridtest_utils.resolve_full_grid_file_name(experiment.grid) - parser = gridfile.GridFile(str(file)) + parser = gridfile.GridFile(str(file), gridfile.ToZeroBasedIndexTransformation()) try: parser.open() - assert parser.dimension(gridfile.DimensionName.CELL_NAME) == grid_savepoint.num( + assert parser.dimension(gridfile.DynamicDimension.CELL_NAME) == grid_savepoint.num( dims.CellDim ) - assert parser.dimension(gridfile.DimensionName.VERTEX_NAME) == grid_savepoint.num( + assert parser.dimension(gridfile.DynamicDimension.VERTEX_NAME) == grid_savepoint.num( dims.VertexDim ) - assert parser.dimension(gridfile.DimensionName.EDGE_NAME) == grid_savepoint.num( + assert parser.dimension(gridfile.DynamicDimension.EDGE_NAME) == grid_savepoint.num( dims.EdgeDim ) except Exception as error: pytest.fail(f"reading of dimension from netcdf failed: {error}") finally: parser.close() + + +@pytest.mark.parametrize("apply_offset", (True, False)) +def test_int_variable(experiment: definitions.Experiment, apply_offset: bool) -> None: + file = gridtest_utils.resolve_full_grid_file_name(experiment.grid) + with gridfile.GridFile(str(file), gridfile.ToZeroBasedIndexTransformation()) as parser: + edge_dim = parser.dimension(gridfile.DynamicDimension.EDGE_NAME) + # use a test field that does not contain Pentagons + test_field = parser.int_variable(gridfile.ConnectivityName.C2E, apply_offset=apply_offset) + min_value = 0 if apply_offset else 1 + max_value = edge_dim - 1 if apply_offset else edge_dim + assert min_value == np.min(test_field) + assert max_value == np.max(test_field) + + +_index_selection: Iterable[list[int]] = [ + [0, 1, 2, 3, 4, 5], + [], + [0, 2, 4, 6, 7, 8, 24, 57], + [1, 2, 12, 13, 23, 24, 513], +] + + +@pytest.mark.parametrize( + "selection", + _index_selection, +) +def test_index_read_for_1d_fields(experiment: definitions.Experiment, selection: list[int]) -> None: + file = gridtest_utils.resolve_full_grid_file_name(experiment.grid) + with gridfile.GridFile(str(file), gridfile.ToZeroBasedIndexTransformation()) as parser: + indices_to_read = np.asarray(selection) if len(selection) > 0 else None + full_field = parser.variable(gridfile.CoordinateName.CELL_LATITUDE) + selective_field = parser.variable( + gridfile.CoordinateName.CELL_LATITUDE, indices=indices_to_read + ) + assert np.allclose(full_field[indices_to_read], selective_field) + + +@pytest.mark.parametrize( + "selection", + _index_selection, +) +@pytest.mark.parametrize( + "field", + ( + gridfile.ConnectivityName.V2E, + gridfile.ConnectivityName.V2C, + gridfile.ConnectivityName.E2V, + ), +) +@pytest.mark.parametrize("apply_offset", (True, False)) +def test_index_read_for_2d_connectivity( + experiment: definitions.Experiment, + selection: list[int], + field: gridfile.FieldName, + apply_offset: bool, +) -> None: + file = gridtest_utils.resolve_full_grid_file_name(experiment.grid) + with gridfile.GridFile(str(file), gridfile.ToZeroBasedIndexTransformation()) as parser: + indices_to_read = np.asarray(selection) if len(selection) > 0 else None + full_field = parser.int_variable(field, transpose=True, apply_offset=apply_offset) + selective_field = parser.int_variable( + field, + indices=indices_to_read, + transpose=True, + apply_offset=apply_offset, + ) + assert np.allclose(full_field[indices_to_read], selective_field) diff --git a/model/common/tests/common/grid/unit_tests/test_horizontal.py b/model/common/tests/common/grid/unit_tests/test_horizontal.py index bb2267d2bd..a859cef62e 100644 --- a/model/common/tests/common/grid/unit_tests/test_horizontal.py +++ b/model/common/tests/common/grid/unit_tests/test_horizontal.py @@ -58,3 +58,10 @@ def test_halo_zones(zone: h_grid.Zone) -> None: assert zone.is_halo() else: assert not zone.is_halo() + + +@pytest.mark.parametrize( + "dim, expected", [(dims.CellDim, 4), (dims.VertexDim, 4), (dims.EdgeDim, 8)] +) +def test_max_boundary_level(dim: gtx.Dimension, expected: int) -> None: + assert expected == h_grid.max_boundary_level(dim) diff --git a/model/common/tests/common/grid/unit_tests/test_icon.py b/model/common/tests/common/grid/unit_tests/test_icon.py index 0099f0725a..d930aaf2f5 100644 --- a/model/common/tests/common/grid/unit_tests/test_icon.py +++ b/model/common/tests/common/grid/unit_tests/test_icon.py @@ -195,9 +195,11 @@ def test_when_keep_skip_value_then_neighbor_table_matches_config( np.any(connectivity.asnumpy() == gridfile.GridFile.INVALID_INDEX).item() ) == icon._has_skip_values(offset, grid.config.limited_area) if not icon._has_skip_values(offset, grid.config.limited_area): - assert connectivity.skip_value is None + assert connectivity.skip_value is None, f"skip value for offset {offset} should be None" else: - assert connectivity.skip_value == gridfile.GridFile.INVALID_INDEX + assert ( + connectivity.skip_value == gridfile.GridFile.INVALID_INDEX + ), f"skip for offset {offset} value should be {gridfile.GridFile.INVALID_INDEX}" @pytest.mark.parametrize( 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 b3a6e699bc..d04ec95eeb 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 @@ -13,7 +13,7 @@ import pytest from icon4py.model.common import dimension as dims -from icon4py.model.common.decomposition import definitions as decomposition +from icon4py.model.common.decomposition import definitions as decomposition, mpi_decomposition from icon4py.model.common.grid import horizontal as h_grid from icon4py.model.common.interpolation import ( interpolation_attributes as attrs, @@ -41,6 +41,9 @@ from icon4py.model.testing import serialbox as sb +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + @pytest.mark.level("integration") @pytest.mark.datatest diff --git a/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py b/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py index 67900b8957..c9fa9aeaf3 100644 --- a/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py +++ b/model/common/tests/common/metrics/mpi_tests/test_parallel_metrics.py @@ -12,7 +12,7 @@ import pytest -from icon4py.model.common.decomposition import definitions as decomposition +from icon4py.model.common.decomposition import definitions as decomposition, mpi_decomposition from icon4py.model.common.metrics import metrics_attributes as attrs, metrics_factory from icon4py.model.testing import definitions as test_defs, parallel_helpers, test_utils @@ -39,6 +39,10 @@ from icon4py.model.testing import serialbox as sb +if mpi_decomposition.mpi4py is None: + pytest.skip("Skipping parallel tests on single node installation", allow_module_level=True) + + @pytest.mark.datatest @pytest.mark.mpi @pytest.mark.uses_concat_where diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py index ae02baa36e..cc6be3b747 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py @@ -35,6 +35,7 @@ geometry as grid_geometry, geometry_attributes as geometry_meta, grid_manager as gm, + gridfile, icon as icon_grid, states as grid_states, vertical as v_grid, @@ -64,7 +65,10 @@ def create_grid_manager( global_reductions: decomposition_defs.Reductions = decomposition_defs.single_node_reductions, ) -> gm.GridManager: grid_manager = gm.GridManager( - gm.ToZeroBasedIndexTransformation(), grid_file_path, vertical_grid_config, global_reductions + grid_file=grid_file_path, + config=vertical_grid_config, + offset_transformation=gridfile.ToZeroBasedIndexTransformation(), + global_reductions=global_reductions, ) grid_manager(allocator=allocator, keep_skip_values=True) diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index 0668ffca5a..7fb1139539 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -140,6 +140,14 @@ class Grids: file_name="Torus_Triangles_50000m_x_5000m_res500m.nc", uri="https://polybox.ethz.ch/index.php/s/eclzK00TM9nnLtE/download", ) + TORUS_1000X1000_250M: Final = GridDescription( + name="torus_1000x1000_res250", + description="Torus grid with a domain (1000x1000) vertices and a resolution (edge length) of 250m, generated by MPI-M GridGenerator", + sizes={"cell": 24, "vertex": 12, "edge": 36}, + shape=icon_grid.GridShape(geometry_type=base_grid.GeometryType.TORUS), + file_name="Torus_Triangles_1000m_x_1000m_res250m.nc", + uri="https://polybox.ethz.ch/index.php/s/eMDbDbdmKLkDiwp/download", + ) # Root URLs for downloading serialized data by communicator size diff --git a/model/testing/src/icon4py/model/testing/fixtures/benchmark.py b/model/testing/src/icon4py/model/testing/fixtures/benchmark.py index a2d6bb382e..86b20463d2 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/benchmark.py +++ b/model/testing/src/icon4py/model/testing/fixtures/benchmark.py @@ -25,7 +25,6 @@ from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory from icon4py.model.common.metrics import metrics_attributes, metrics_factory from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing import grid_utils @pytest.fixture( @@ -39,9 +38,8 @@ def geometry_field_source( pytest.skip("Incomplete grid Information for test, are you running with `simple_grid`?") mesh = grid_manager.grid - allocator = model_backends.get_allocator(backend_like) generic_concrete_backend = model_options.customize_backend(None, backend_like) - decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator) + decomposition_info = grid_manager.decomposition_info geometry_field_source = grid_geometry.GridGeometry( grid=mesh, @@ -64,9 +62,8 @@ def interpolation_field_source( ) -> Generator[interpolation_factory.InterpolationFieldsFactory, None, None]: mesh = grid_manager.grid - allocator = model_backends.get_allocator(backend_like) generic_concrete_backend = model_options.customize_backend(None, backend_like) - decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator) + decomposition_info = grid_manager.decomposition_info interpolation_field_source = interpolation_factory.InterpolationFieldsFactory( grid=mesh, @@ -91,7 +88,7 @@ def metrics_field_source( allocator = model_backends.get_allocator(backend_like) generic_concrete_backend = model_options.customize_backend(None, backend_like) - decomposition_info = grid_utils.construct_decomposition_info(mesh, allocator) + decomposition_info = grid_manager.decomposition_info vertical_config = v_grid.VerticalGridConfig( mesh.num_levels, diff --git a/model/testing/src/icon4py/model/testing/grid_utils.py b/model/testing/src/icon4py/model/testing/grid_utils.py index 7809625f6b..8f9f6a358e 100644 --- a/model/testing/src/icon4py/model/testing/grid_utils.py +++ b/model/testing/src/icon4py/model/testing/grid_utils.py @@ -7,19 +7,17 @@ # SPDX-License-Identifier: BSD-3-Clause import pathlib -import gt4py.next as gtx import gt4py.next.typing as gtx_typing -from icon4py.model.common import dimension as dims, model_backends -from icon4py.model.common.decomposition import definitions as decomposition_defs +from icon4py.model.common import model_backends from icon4py.model.common.grid import ( geometry, geometry_attributes as geometry_attrs, grid_manager as gm, - icon, + gridfile, vertical as v_grid, ) -from icon4py.model.common.utils import data_allocation as data_alloc, device_utils +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import config, data_handling, definitions @@ -52,7 +50,7 @@ def get_grid_manager_from_identifier( def get_grid_manager( - grid_file: pathlib.Path, + filename: pathlib.Path, num_levels: int, keep_skip_values: bool, allocator: gtx_typing.FieldBufferAllocationUtil, @@ -61,15 +59,15 @@ def get_grid_manager( Construct a GridManager instance for an ICON grid file. Args: - grid_file: full path to the file + filename: full path to the file num_levels: number of vertical levels, needed for IconGrid construction but independent from grid file keep_skip_values: whether to keep skip values backend: the gt4py Backend we are running on """ manager = gm.GridManager( - gm.ToZeroBasedIndexTransformation(), - grid_file, - v_grid.VerticalGridConfig(num_levels=num_levels), + grid_file=filename, + config=v_grid.VerticalGridConfig(num_levels=num_levels), + offset_transformation=gridfile.ToZeroBasedIndexTransformation(), ) manager(allocator=allocator, keep_skip_values=keep_skip_values) return manager @@ -99,26 +97,6 @@ def _download_grid_file(grid: definitions.GridDescription) -> pathlib.Path: return full_name -def construct_decomposition_info( - grid: icon.IconGrid, - allocator: gtx_typing.FieldBufferAllocationUtil | None = None, -) -> decomposition_defs.DecompositionInfo: - on_gpu = device_utils.is_cupy_device(allocator) - xp = data_alloc.array_ns(on_gpu) - - def _add_dimension(dim: gtx.Dimension) -> None: - indices = data_alloc.index_field(grid, dim, allocator=allocator) - owner_mask = xp.ones((grid.size[dim],), dtype=bool) - decomposition_info.with_dimension(dim, indices.ndarray, owner_mask) - - decomposition_info = decomposition_defs.DecompositionInfo() - _add_dimension(dims.EdgeDim) - _add_dimension(dims.VertexDim) - _add_dimension(dims.CellDim) - - return decomposition_info - - def get_grid_geometry( backend: gtx_typing.Backend | None, experiment: definitions.Experiment ) -> geometry.GridGeometry: @@ -127,15 +105,13 @@ def get_grid_geometry( def _construct_grid_geometry() -> geometry.GridGeometry: gm = get_grid_manager_from_identifier( experiment.grid, - keep_skip_values=True, num_levels=experiment.num_levels, + keep_skip_values=True, allocator=model_backends.get_allocator(backend), ) - grid = gm.grid - decomposition_info = construct_decomposition_info(grid, backend) return geometry.GridGeometry( - grid, - decomposition_info, + gm.grid, + gm.decomposition_info, backend, gm.coordinates, gm.geometry_fields, diff --git a/model/testing/src/icon4py/model/testing/parallel_helpers.py b/model/testing/src/icon4py/model/testing/parallel_helpers.py index 4837d1c711..b0ad1b0465 100644 --- a/model/testing/src/icon4py/model/testing/parallel_helpers.py +++ b/model/testing/src/icon4py/model/testing/parallel_helpers.py @@ -6,12 +6,11 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import logging -from collections.abc import Iterable import pytest +from icon4py.model.common import dimension as dims from icon4py.model.common.decomposition import definitions -from icon4py.model.common.decomposition.mpi_decomposition import get_multinode_properties log = logging.getLogger(__file__) @@ -24,21 +23,13 @@ def check_comm_size( pytest.xfail(f"wrong comm size: {props.comm_size}: test only works for comm-sizes: {sizes}") -@pytest.fixture(scope="session") -def processor_props(request: pytest.FixtureRequest) -> Iterable[definitions.ProcessProperties]: - runtype = definitions.get_runtype(with_mpi=True) - yield get_multinode_properties(runtype) - - -def log_process_properties( - props: definitions.ProcessProperties, level: int = logging.DEBUG -) -> None: +def log_process_properties(props: definitions.ProcessProperties) -> None: log.info(f"rank={props.rank}/{props.comm_size}") -def log_local_field_size( - decomposition_info: definitions.DecompositionInfo, level: int = logging.DEBUG -) -> None: +def log_local_field_size(decomposition_info: definitions.DecompositionInfo) -> None: log.info( - f"local grid size: cells={decomposition_info.num_cells}, edges={decomposition_info.num_edges}, vertices={decomposition_info.num_vertices}" + f"local grid size: cells={decomposition_info.global_index(dims.CellDim).size}, " + f"edges={decomposition_info.global_index(dims.EdgeDim).size}, " + f"vertices={decomposition_info.global_index(dims.VertexDim).size}" ) diff --git a/model/testing/src/icon4py/model/testing/serialbox.py b/model/testing/src/icon4py/model/testing/serialbox.py index 9cdb912e2d..93a4bc21ce 100644 --- a/model/testing/src/icon4py/model/testing/serialbox.py +++ b/model/testing/src/icon4py/model/testing/serialbox.py @@ -151,7 +151,10 @@ def __init__( ): super().__init__(sp, ser, size, backend) self._grid_id = grid_id - self.global_grid_params = icon.GlobalGridParams(grid_shape=grid_shape) + self.global_grid_params = icon.GlobalGridParams( + grid_shape=grid_shape, + num_cells=size[dims.CellDim], + ) def verts_vertex_lat(self): """vertex latituted""" @@ -452,34 +455,27 @@ def _read_field_for_dim(field_name, read_func, dim: gtx.Dimension): ) def owner_mask(self, dim: gtx.Dimension): - field_name = "owner_mask" - mask = self._read_field_for_dim(field_name, self._read_bool, dim) - return mask + return np.squeeze(self._read_field_for_dim("owner_mask", self._read_bool, dim)) def global_index(self, dim: gtx.Dimension): - field_name = "glb_index" - return self._read_field_for_dim(field_name, self._read_int32_shift1, dim) + return self._read_field_for_dim("glb_index", self._read_int32_shift1, dim) def decomp_domain(self, dim): - field_name = "decomp_domain" - return self._read_field_for_dim(field_name, self._read_int32, dim) + return self._read_field_for_dim("decomp_domain", self._read_int32, dim) - def construct_decomposition_info(self): + def construct_decomposition_info(self) -> decomposition.DecompositionInfo: return ( - decomposition.DecompositionInfo( - num_cells=self.num(dims.CellDim), - num_edges=self.num(dims.EdgeDim), - num_vertices=self.num(dims.VertexDim), - ) - .with_dimension(*self._get_decomp_fields(dims.CellDim)) - .with_dimension(*self._get_decomp_fields(dims.EdgeDim)) - .with_dimension(*self._get_decomp_fields(dims.VertexDim)) + decomposition.DecompositionInfo() + .set_dimension(*self._get_decomposition_fields(dims.CellDim)) + .set_dimension(*self._get_decomposition_fields(dims.EdgeDim)) + .set_dimension(*self._get_decomposition_fields(dims.VertexDim)) ) - def _get_decomp_fields(self, dim: gtx.Dimension): + def _get_decomposition_fields(self, dim: gtx.Dimension): global_index = self.global_index(dim) mask = self.owner_mask(dim)[0 : self.num(dim)] - return dim, global_index, mask + halo_levels = self.decomp_domain(dim)[0 : self.num(dim)] + return dim, global_index, mask, halo_levels def construct_icon_grid( self, diff --git a/pyproject.toml b/pyproject.toml index e349356eb7..5e0f3631a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,8 @@ test = [ typing = [ "mypy[faster-cache]>=1.13.0", "typing-extensions>=4.11.0", - "types-cffi>=1.16.0" + "types-cffi>=1.16.0", + "mpi4py>=4.0.0" ] # -- Standard project description options (PEP 621) -- diff --git a/pytest b/pytest new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tach.toml b/tach.toml index a9c93be531..902741f92e 100644 --- a/tach.toml +++ b/tach.toml @@ -23,10 +23,12 @@ exclude = [ "dace", "mpi4py", "netcdf4", + "pymetis", "xarray", "uxarray", "cftime", "viztracer", + ] rename = ["serialbox:serialbox4py"] diff --git a/tools/src/icon4py/tools/py2fgen/wrappers/common.py b/tools/src/icon4py/tools/py2fgen/wrappers/common.py index 2cc1b246a3..0e22449fac 100644 --- a/tools/src/icon4py/tools/py2fgen/wrappers/common.py +++ b/tools/src/icon4py/tools/py2fgen/wrappers/common.py @@ -236,12 +236,11 @@ def construct_decomposition( v_owner_mask = v_owner_mask[:num_vertices] decomposition_info = ( - definitions.DecompositionInfo( - num_cells=num_cells, num_edges=num_edges, num_vertices=num_vertices - ) - .with_dimension(dims.CellDim, c_glb_index, c_owner_mask) - .with_dimension(dims.EdgeDim, e_glb_index, e_owner_mask) - .with_dimension(dims.VertexDim, v_glb_index, v_owner_mask) + definitions.DecompositionInfo() + # TODO (halungge): last argument is called `decomp_domain` in icon, it is not needed in the granules should we pass it nevertheless? + .set_dimension(dims.CellDim, c_glb_index, c_owner_mask, None) + .set_dimension(dims.EdgeDim, e_glb_index, e_owner_mask, None) + .set_dimension(dims.VertexDim, v_glb_index, v_owner_mask, None) ) processor_props = definitions.get_processor_properties(definitions.MultiNodeRun(), comm_id) exchange = definitions.create_exchange(processor_props, decomposition_info) diff --git a/uv.lock b/uv.lock index 5c790b4963..77a922d41f 100644 --- a/uv.lock +++ b/uv.lock @@ -1579,6 +1579,7 @@ dev = [ { name = "coverage", extra = ["toml"] }, { name = "esbonio" }, { name = "icon4py-testing" }, + { name = "mpi4py" }, { name = "mypy", extra = ["faster-cache"] }, { name = "myst-parser" }, { name = "nox" }, @@ -1639,6 +1640,7 @@ test = [ { name = "pytest-xdist", extra = ["psutil"] }, ] typing = [ + { name = "mpi4py" }, { name = "mypy", extra = ["faster-cache"] }, { name = "types-cffi" }, { name = "typing-extensions" }, @@ -1678,6 +1680,7 @@ dev = [ { name = "coverage", extras = ["toml"], specifier = ">=7.5.0" }, { name = "esbonio", specifier = ">=0.16.0" }, { name = "icon4py-testing", editable = "model/testing" }, + { name = "mpi4py", specifier = ">=4.0.0" }, { name = "mypy", extras = ["faster-cache"], specifier = ">=1.13.0" }, { name = "myst-parser", specifier = ">=4.0.0" }, { name = "nox", git = "https://github.com/wntrblm/nox.git?rev=aa09595437608dfe21eb776d8a4bcc0bd5f9916b" }, @@ -1738,6 +1741,7 @@ test = [ { name = "pytest-xdist", extras = ["psutil"], specifier = ">=3.5.0" }, ] typing = [ + { name = "mpi4py", specifier = ">=4.0.0" }, { name = "mypy", extras = ["faster-cache"], specifier = ">=1.13.0" }, { name = "types-cffi", specifier = ">=1.16.0" }, { name = "typing-extensions", specifier = ">=4.11.0" }, @@ -1857,6 +1861,7 @@ all = [ { name = "mpi4py" }, { name = "netcdf4" }, { name = "numpy" }, + { name = "pymetis" }, { name = "scikit-learn" }, { name = "uxarray" }, { name = "xarray", extra = ["complete"] }, @@ -1880,6 +1885,7 @@ io = [ { name = "holoviews" }, { name = "netcdf4" }, { name = "numpy" }, + { name = "pymetis" }, { name = "scikit-learn" }, { name = "uxarray" }, { name = "xarray", extra = ["complete"] }, @@ -1904,6 +1910,7 @@ requires-dist = [ { name = "numpy", specifier = ">=1.23.3" }, { name = "numpy", marker = "extra == 'io'", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, + { name = "pymetis", marker = "extra == 'io'", specifier = ">2022.1" }, { name = "scikit-learn", marker = "extra == 'io'", specifier = ">=1.4.0" }, { name = "scipy", specifier = ">=1.14.1" }, { name = "typing-extensions", specifier = ">=4.11.0" }, @@ -3437,6 +3444,49 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f7/3f/01c8b82017c199075f8f788d0d906b9ffbbc5a47dc9918a945e13d5a2bda/pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a", size = 1205513, upload-time = "2024-05-04T13:41:57.345Z" }, ] +[[package]] +name = "pymetis" +version = "2025.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/cd/e8f6b7004abf2c7587fc760d0b8830d7c23f76b88f72dd3b412fd9f69e3d/pymetis-2025.1.1.tar.gz", hash = "sha256:c02bf0fdd9483ff9dac031c73219ab9aa7a90e110f09d3c00f1799f334d813d8", size = 5033014, upload-time = "2025-05-05T19:23:08.011Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1c/71/3f916e0590fc0efa56c9bd6e4e0f2f9cf0e6da3c04b690cadebc997b1f3e/pymetis-2025.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:6d857b37bffffe345c3b6314332dcd55ca9217255bba9e8941aa8a08b434662f", size = 244755, upload-time = "2025-05-05T19:22:17.317Z" }, + { url = "https://files.pythonhosted.org/packages/11/84/369a861c621fa91aa5a5f232dabf864fe349b0ea8daef7af5e63995ea626/pymetis-2025.1.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:00016ee75edac02da6c86df51dde1fb2e5abb815f1a334ecd273d33d9c4b2a0b", size = 217258, upload-time = "2025-05-05T19:22:19.177Z" }, + { url = "https://files.pythonhosted.org/packages/87/77/2eb3f059a784a478dcf142a8b941b46ed06baf921bbefdce72600a21f276/pymetis-2025.1.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:76852cdb0fbea473df31fb995c13366fa41a4969853e96acb566e075556c0559", size = 336565, upload-time = "2025-05-05T19:22:20.8Z" }, + { url = "https://files.pythonhosted.org/packages/04/0d/c7a71e9fa74b6402b961ae523a806495f480d3c6ee348c583830bd84d7ad/pymetis-2025.1.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83372d22b015c017dae859a95974c1ead9800772411e0ed3ac106d127c21db23", size = 274742, upload-time = "2025-05-05T19:22:22.09Z" }, + { url = "https://files.pythonhosted.org/packages/3a/96/ed0862d0789dd54c429e5635cdb6d1bb286261666d8f5c77d523a5c1a900/pymetis-2025.1.1-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:64137431e0f41515d2c784a57a68269607fcc930691a6661ad818f5804c5a9a9", size = 1461516, upload-time = "2025-05-05T19:22:24.809Z" }, + { url = "https://files.pythonhosted.org/packages/97/fd/7eae7d6a2462ce63669ed8813c43e8d8b1ea93ff440ceab38b0130d6a8cb/pymetis-2025.1.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:e0df438aca9fb832173644b2190b1c181bc10d8926ce8071752fb6c5cddf1117", size = 1290326, upload-time = "2025-05-05T19:22:26.29Z" }, + { url = "https://files.pythonhosted.org/packages/54/c5/567e26cb9be60b4f11eac3511d4889569a7fb7044dfc2392a3f905fb366c/pymetis-2025.1.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5bfbd3df3b3114c106054b15cc0455094526240fdd29c284084fdae3f4fdd289", size = 245886, upload-time = "2025-05-05T19:22:28.308Z" }, + { url = "https://files.pythonhosted.org/packages/ef/0b/35fc2fc5a16913af40a1eb2d85096117f2db2e2e82c477f7ab8fa91318ea/pymetis-2025.1.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:45c3478ff2a1d208e75cc5fe1189f0aac8d1efe6a61e12665be762347c91d6c0", size = 218446, upload-time = "2025-05-05T19:22:29.791Z" }, + { url = "https://files.pythonhosted.org/packages/d7/26/f86f9be6b62c23d7ac9141081b111d7215723063a0edc204951669f1d18d/pymetis-2025.1.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0b74810074c602fedb0bcece8fdcbf58773e3bcfa1e28884271d069598405dc5", size = 337537, upload-time = "2025-05-05T19:22:31.512Z" }, + { url = "https://files.pythonhosted.org/packages/5f/c1/daf8797d7af40fe5748366118c9383e60bda339cac5e74c48dc8e7136931/pymetis-2025.1.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3aeda5d9ee8df41db134c9cf0a10c51a805dae29c9dae7f9cb7ccad37e15d6ec", size = 276008, upload-time = "2025-05-05T19:22:33.332Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7f/ec0894623175ccc2d6de261c7f8fe7c4a99f7914ac6195bdff75883ad075/pymetis-2025.1.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:9b64ad48f45cb55d2077f965114f9818a10d24e2744d40c0ebc4d5a2065db10c", size = 1462382, upload-time = "2025-05-05T19:22:35.234Z" }, + { url = "https://files.pythonhosted.org/packages/66/8c/16cc223821ed25250270b6f50633b0fb0da43912007e068be6f52ff98185/pymetis-2025.1.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:408437b4aad82b67379f624b9faf9761b09caccdae083823a84b1bc7a470d86a", size = 1291900, upload-time = "2025-05-05T19:22:37.188Z" }, + { url = "https://files.pythonhosted.org/packages/84/8e/ac5d7110fea586fe660ef466e03abaca358dc2ed8282f01b4ae5cc93f5c1/pymetis-2025.1.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7dd5019162d8f19d03f1be9096547660d02a5925949e00529f8a5d7d4dfadb95", size = 246019, upload-time = "2025-05-05T19:22:39.039Z" }, + { url = "https://files.pythonhosted.org/packages/01/1e/74f3f9c9538a3435028b2bf00f659d84578861b6ca589acbd42d3623274c/pymetis-2025.1.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c65cb13972feda4f7f186c08c8cf619ac5f3f1264c9609466615b909219b9e58", size = 218388, upload-time = "2025-05-05T19:22:40.392Z" }, + { url = "https://files.pythonhosted.org/packages/a2/de/3462853dbb8ce25610db986cfc9f0e4fa3014f21412d7659bfa747596604/pymetis-2025.1.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d5b1d8c053dbd456bfc2f1fac5a2307c1c5cb614973a6485de966c049d48b18a", size = 337651, upload-time = "2025-05-05T19:22:42.145Z" }, + { url = "https://files.pythonhosted.org/packages/73/43/24c63cee0acd9d2ee9e1d657c933c33c16b83e638d7d36dca95437f4b979/pymetis-2025.1.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce64511e0fcb58b3012d6bd007303f55a4b3156025ff945ca5d4069349608334", size = 275477, upload-time = "2025-05-05T19:22:44.128Z" }, + { url = "https://files.pythonhosted.org/packages/07/b4/a6db323a67ac917ea902d1b52f64ca8defe1ae875890f55d7aefe55ceed5/pymetis-2025.1.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5fba99f74ea32e60abde9d8ad77c4ee9865fabe1fc2627b982d3bb359689360e", size = 1462678, upload-time = "2025-05-05T19:22:45.605Z" }, + { url = "https://files.pythonhosted.org/packages/c5/5f/0e94aa79362017f7464d249ab056245cc42b8cf2200e92fd593d89c77040/pymetis-2025.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ccc9f54edd33466a917ef5d2ea1c12e76787b4790e14754460daf427f1666247", size = 1292407, upload-time = "2025-05-05T19:22:47.504Z" }, + { url = "https://files.pythonhosted.org/packages/39/be/731625c087d720e533c2fd9350bfc1b581701e3ca507c6269d4c633b6653/pymetis-2025.1.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d4cdadbd49ba1f71a87c2a3c7e38241bff1cf16f9399ad05c1e856ce70c129d6", size = 246081, upload-time = "2025-05-05T19:22:48.826Z" }, + { url = "https://files.pythonhosted.org/packages/06/1e/d220c9090e0f2c81919a2c482bf912c59010a25558a5c3ef50ababe9cd7f/pymetis-2025.1.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d591bb9cd7d10b5f855b264ba713dc4a885492e61c5353c0a5ca34b9a28a4cf6", size = 218458, upload-time = "2025-05-05T19:22:50.209Z" }, + { url = "https://files.pythonhosted.org/packages/87/57/20fbab509ac524ec94bd608ddda41c0d027a5661549d7402e3204b52fbd1/pymetis-2025.1.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:18704c986d76725ccb86288be24b069618cfc5916672450b967946f41a52cf2b", size = 337713, upload-time = "2025-05-05T19:22:51.525Z" }, + { url = "https://files.pythonhosted.org/packages/b9/fe/bb12c811cededf3b9b483851622e4622e2a27b64b0bca720cd966f522d16/pymetis-2025.1.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:12ce2c20a0d30600b0d3fe049107ca76f3fdd9fdae1009b33b1ce168eb2aa089", size = 275622, upload-time = "2025-05-05T19:22:52.949Z" }, + { url = "https://files.pythonhosted.org/packages/1a/c8/87ac7974edf3940a7d3c6b8b1db20f2eb861c46634d8b5827e8b9c443d0f/pymetis-2025.1.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:33ac080d603a309cdaa235c83291a392f4c1e9e89d498b62a7c947ce6ab9f8f2", size = 1462651, upload-time = "2025-05-05T19:22:54.408Z" }, + { url = "https://files.pythonhosted.org/packages/57/18/763d5d4fa5da5c375b019d11251bc1955d005a4e4a6fccbf88e190b31aa6/pymetis-2025.1.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:c389c822fe38db2b5bff60674250a3f97448b372d060d5d4789fae961d425223", size = 1292442, upload-time = "2025-05-05T19:22:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/d9/71/c7a486a55ee0efa3cd611b585ed4ac0c92241ba7ee488f3a5bafdf2ad8b8/pymetis-2025.1.1-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:625da502859f7d9449ccc5702c5228396ded95446ee61ab02f3ba84d58bb1a3b", size = 244783, upload-time = "2025-05-05T19:22:57.205Z" }, + { url = "https://files.pythonhosted.org/packages/10/8d/f1170b961645e5a33926314b62c3fc9f5a5e472a0d735bcf1a2dd0bdf656/pymetis-2025.1.1-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:05d7737ffa5594b622e6be53c9ae11855f798e3cd67b345bc8cdc286a501fab5", size = 217473, upload-time = "2025-05-05T19:22:58.517Z" }, + { url = "https://files.pythonhosted.org/packages/e7/cc/bc3c5190500648f7fe7b1427a68bdf7fc9fa3e2f03a740fe4ee4a1730d50/pymetis-2025.1.1-pp310-pypy310_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f44b6cc411b395c04356305120b3fec19f4a0fc4e081d6f95c0acb3879fa1e24", size = 335596, upload-time = "2025-05-05T19:22:59.795Z" }, + { url = "https://files.pythonhosted.org/packages/65/97/fc54d7ac05cd4b3de8d9ff3ebacecee53249270c685f332d33f6aecfd79b/pymetis-2025.1.1-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c3b9c5eeaf0a4be0c6e158fe0b43de6187580d9dc9b0c40c6569891c420e2731", size = 274340, upload-time = "2025-05-05T19:23:01.111Z" }, + { url = "https://files.pythonhosted.org/packages/a0/9f/2a5c32a99a80b4a8d6ee880c57e760167a8caa91a9f7c12566081612bfb5/pymetis-2025.1.1-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0e4363afc9aee6517cb3bd67978fa2bd5e923015f2e5630aad8451fd17cb24ef", size = 246012, upload-time = "2025-05-05T19:23:02.437Z" }, + { url = "https://files.pythonhosted.org/packages/e0/84/0e26a51fc1e6327e5f3d8bd67cfc78f47873d0f7f1c2f852464b56052d17/pymetis-2025.1.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:3880b7ec42b5d5dc1cee5170304962caf6349a347a882feb8156f082922badce", size = 218677, upload-time = "2025-05-05T19:23:03.775Z" }, + { url = "https://files.pythonhosted.org/packages/64/7d/e77d5c7771a8c31e3cf00a5abb116ec8d70db5c8116692feba1fa0cc12ec/pymetis-2025.1.1-pp311-pypy311_pp73-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:03d71a0cf3d18c76c0a98131876d41d7810b89f8dcb2a21ef11eff9f290de688", size = 336761, upload-time = "2025-05-05T19:23:05.113Z" }, + { url = "https://files.pythonhosted.org/packages/5b/29/22137aeaef2b472aabddea09fe7f512423ad50e405bb8da6c82089ca6720/pymetis-2025.1.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:daf4d7ea3f147fbb5328566abd4f493b1911ea911000036693bb5c639c7db4a5", size = 275642, upload-time = "2025-05-05T19:23:06.543Z" }, +] + [[package]] name = "pyparsing" version = "3.2.0"