From 3d87c46a5a9b8e15d48e1e5afe49f2755648d492 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 30 Apr 2024 12:46:28 +0200 Subject: [PATCH] Fix parallel tests (#428) - adapt mpi test to latest ghex version - compute mean_cell_area from grid file properties - remove datapath fixture --- .../tests/diffusion_tests/conftest.py | 1 - .../mpi_tests/test_parallel_diffusion.py | 20 +++--- .../dycore/compute_theta_and_exner.py | 2 +- .../dycore/tests/dycore_tests/conftest.py | 1 - .../tests/dycore_tests/mpi_tests/__init__.py | 12 ++++ .../mpi_tests/test_parallel_solve_nonhydro.py | 54 +++++++-------- model/common/pyproject.toml | 2 +- .../src/icon4py/model/common/constants.py | 3 + .../common/decomposition/mpi_decomposition.py | 21 ++++-- .../src/icon4py/model/common/grid/base.py | 13 ++-- .../icon4py/model/common/grid/grid_manager.py | 11 +++- .../icon4py/model/common/grid/horizontal.py | 42 +++++++++++- .../src/icon4py/model/common/grid/icon.py | 27 ++++++++ .../src/icon4py/model/common/grid/simple.py | 3 +- .../common/test_utils/datatest_fixtures.py | 25 +++---- .../model/common/test_utils/datatest_utils.py | 23 ++++++- .../model/common/test_utils/helpers.py | 7 +- .../common/test_utils/serialbox_utils.py | 24 +++++-- .../test_mpi_decomposition.py | 66 ++++++++++++++++--- model/common/tests/grid_tests/conftest.py | 1 - .../tests/grid_tests/test_grid_manager.py | 17 +++++ .../tests/grid_tests/test_horizontal.py | 30 +++++++++ .../tests/interpolation_tests/conftest.py | 1 - .../test_interpolation_fields.py | 1 - model/common/tests/metric_tests/conftest.py | 1 - .../metric_tests/test_compute_nudgecoeffs.py | 1 - model/driver/tests/conftest.py | 1 - requirements-dev-opt.txt | 2 +- tools/src/icon4pytools/icon4pygen/backend.py | 4 +- 29 files changed, 313 insertions(+), 103 deletions(-) create mode 100644 model/atmosphere/dycore/tests/dycore_tests/mpi_tests/__init__.py create mode 100644 model/common/tests/grid_tests/test_horizontal.py diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/conftest.py b/model/atmosphere/diffusion/tests/diffusion_tests/conftest.py index 6fb11ae2b2..b878349e14 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/conftest.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/conftest.py @@ -16,7 +16,6 @@ from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package damping_height, data_provider, - datapath, decomposition_info, download_ser_data, experiment, diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py index 5d651335d7..5afdf229b2 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py @@ -18,12 +18,14 @@ from icon4py.model.common.decomposition import definitions from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.test_utils.datatest_utils import REGIONAL_EXPERIMENT from icon4py.model.common.test_utils.parallel_helpers import ( # noqa: F401 # import fixtures from test_utils package check_comm_size, processor_props, ) from ..utils import ( + construct_config, construct_diagnostics, construct_interpolation_state, construct_metric_state, @@ -31,14 +33,12 @@ ) -@pytest.mark.xfail( - "TODO(@halungge) fails due to expectation of field allocation (vertical ~ contiguous) in ghex." -) @pytest.mark.mpi +@pytest.mark.parametrize("experiment", [REGIONAL_EXPERIMENT]) @pytest.mark.parametrize("ndyn_substeps", [2]) @pytest.mark.parametrize("linit", [True, False]) def test_parallel_diffusion( - r04b09_diffusion_config, + experiment, step_date_init, linit, ndyn_substeps, @@ -54,7 +54,7 @@ def test_parallel_diffusion( ): check_comm_size(processor_props) print( - f"rank={processor_props.rank}/{processor_props.comm_size}: inializing diffusion for experiment 'mch_ch_r04_b09_dsl" + f"rank={processor_props.rank}/{processor_props.comm_size}: inializing diffusion for experiment '{REGIONAL_EXPERIMENT}'" ) print( f"rank={processor_props.rank}/{processor_props.comm_size}: decomposition info : klevels = {decomposition_info.klevels}, " @@ -73,8 +73,8 @@ def test_parallel_diffusion( cell_geometry = grid_savepoint.construct_cell_geometry() edge_geometry = grid_savepoint.construct_edge_geometry() interpolation_state = construct_interpolation_state(interpolation_savepoint) - - diffusion_params = DiffusionParams(r04b09_diffusion_config) + config = construct_config(experiment, ndyn_substeps=ndyn_substeps) + diffusion_params = DiffusionParams(config) dtime = diffusion_savepoint_init.get_metadata("dtime").get("dtime") print( f"rank={processor_props.rank}/{processor_props.comm_size}: setup: using {processor_props.comm_name} with {processor_props.comm_size} nodes" @@ -85,7 +85,7 @@ def test_parallel_diffusion( diffusion.init( grid=icon_grid, - config=r04b09_diffusion_config, + config=config, params=diffusion_params, vertical_params=VerticalModelParams(grid_savepoint.vct_a(), damping_height), metric_state=metric_state, @@ -94,7 +94,7 @@ def test_parallel_diffusion( cell_params=cell_geometry, ) print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ") - diagnostic_state = construct_diagnostics(diffusion_savepoint_init, grid_savepoint) + diagnostic_state = construct_diagnostics(diffusion_savepoint_init) prognostic_state = diffusion_savepoint_init.construct_prognostics() if linit: diffusion.initial_run( @@ -111,7 +111,7 @@ def test_parallel_diffusion( print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion run ") verify_diffusion_fields( - config=r04b09_diffusion_config, + config=config, diagnostic_state=diagnostic_state, prognostic_state=prognostic_state, diffusion_savepoint=diffusion_savepoint_exit, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py index da5df4473c..1974e1ed88 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_theta_and_exner.py @@ -29,7 +29,7 @@ def _compute_theta_and_exner( rd_o_cvd: wpfloat, rd_o_p0ref: wpfloat, ) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]: - """Formelry known as _mo_solve_nonhydro_stencil_66.""" + """Formerly known as _mo_solve_nonhydro_stencil_66.""" theta_v_wp = where(bdy_halo_c, exner, theta_v) exner_wp = where(bdy_halo_c, exp(rd_o_cvd * log(rd_o_p0ref * rho * exner)), exner) return theta_v_wp, exner_wp diff --git a/model/atmosphere/dycore/tests/dycore_tests/conftest.py b/model/atmosphere/dycore/tests/dycore_tests/conftest.py index 6ae72dd249..444ebe60f3 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/conftest.py +++ b/model/atmosphere/dycore/tests/dycore_tests/conftest.py @@ -15,7 +15,6 @@ from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa F401 damping_height, data_provider, - datapath, download_ser_data, experiment, grid_savepoint, diff --git a/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/__init__.py b/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/__init__.py new file mode 100644 index 0000000000..15dfdb0098 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py index 9663071fbc..13289beba2 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py @@ -16,7 +16,6 @@ import pytest from icon4py.model.atmosphere.dycore.nh_solve.solve_nonhydro import ( - NonHydrostaticConfig, NonHydrostaticParams, SolveNonhydro, ) @@ -28,7 +27,6 @@ from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams -from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa : F401 fixture decomposition_info, ) @@ -38,10 +36,14 @@ processor_props, ) - -@pytest.mark.xfail( - "TODO(@halungge) fails due to expectation of field allocation (vertical ~ contiguous) in ghex." +from ..test_solve_nonhydro import create_prognostic_states +from ..utils import ( + construct_config, + construct_interpolation_state_for_nonhydro, + construct_nh_metric_state, ) + + @pytest.mark.datatest @pytest.mark.parametrize( "istep_init, jstep_init, step_date_init,istep_exit, jstep_exit, step_date_exit", @@ -55,6 +57,8 @@ def test_run_solve_nonhydro_single_step( jstep_exit, step_date_init, step_date_exit, + experiment, + ndyn_substeps, icon_grid, savepoint_nonhydro_init, damping_height, @@ -91,7 +95,7 @@ def test_run_solve_nonhydro_single_step( f"rank={processor_props.rank}/{processor_props.comm_size}: number of halo cells {np.count_nonzero(np.invert(owned_cells))}" ) - config = NonHydrostaticConfig() + config = construct_config(experiment, ndyn_substeps=ndyn_substeps) sp = savepoint_nonhydro_init sp_step_exit = savepoint_nonhydro_step_exit nonhydro_params = NonHydrostaticParams(config) @@ -116,7 +120,6 @@ def test_run_solve_nonhydro_single_step( nnew = 1 recompute = sp_v.get_metadata("recompute").get("recompute") linit = sp_v.get_metadata("linit").get("linit") - dyn_timestep = sp_v.get_metadata("dyn_timestep").get("dyn_timestep") diagnostic_state_nh = DiagnosticStateNonHydro( theta_v_ic=sp.theta_v_ic(), @@ -139,31 +142,20 @@ def test_run_solve_nonhydro_single_step( rho_incr=None, # sp.rho_incr(), vn_incr=None, # sp.vn_incr(), exner_incr=None, # sp.exner_incr(), + exner_dyn_incr=sp.exner_dyn_incr(), ) - - prognostic_state_nnow = PrognosticState( - w=sp.w_now(), - vn=sp.vn_now(), - theta_v=sp.theta_v_now(), - rho=sp.rho_now(), - exner=sp.exner_now(), - ) - - prognostic_state_nnew = PrognosticState( - w=sp.w_new(), - vn=sp.vn_new(), - theta_v=sp.theta_v_new(), - rho=sp.rho_new(), - exner=sp.exner_new(), - ) - - interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro() - metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels) + initial_divdamp_fac = sp.divdamp_fac_o2() + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() + prognostic_state_ls = create_prognostic_states(sp) + prognostic_state_nnew = prognostic_state_ls[1] + exchange = definitions.create_exchange(processor_props, decomposition_info) + solve_nonhydro = SolveNonhydro(exchange) solve_nonhydro.init( grid=icon_grid, @@ -177,7 +169,6 @@ def test_run_solve_nonhydro_single_step( owner_mask=grid_savepoint.c_owner_mask(), ) - prognostic_state_ls = [prognostic_state_nnow, prognostic_state_nnew] print( f"rank={processor_props.rank}/{processor_props.comm_size}: entering : solve_nonhydro.time_step" ) @@ -186,20 +177,21 @@ def test_run_solve_nonhydro_single_step( diagnostic_state_nh=diagnostic_state_nh, prognostic_state_ls=prognostic_state_ls, prep_adv=prep_adv, - divdamp_fac_o2=0.032, + divdamp_fac_o2=initial_divdamp_fac, dtime=dtime, - idyn_timestep=dyn_timestep, l_recompute=recompute, l_init=linit, nnew=nnew, nnow=nnow, lclean_mflx=clean_mflx, lprep_adv=lprep_adv, + at_first_substep=jstep_init == 0, + at_last_substep=jstep_init == (ndyn_substeps - 1), ) print(f"rank={processor_props.rank}/{processor_props.comm_size}: dycore step run ") - expected_theta_v = np.asarray(sp_step_exit.theta_v_new()) - calculated_theta_v = np.asarray(prognostic_state_nnew.theta_v) + expected_theta_v = sp_step_exit.theta_v_new().asnumpy() + calculated_theta_v = prognostic_state_nnew.theta_v.asnumpy() assert dallclose( expected_theta_v, calculated_theta_v, diff --git a/model/common/pyproject.toml b/model/common/pyproject.toml index 3bd4622234..69ce8ba965 100644 --- a/model/common/pyproject.toml +++ b/model/common/pyproject.toml @@ -30,7 +30,7 @@ requires-python = ">=3.10" [project.optional-dependencies] all = ["icon4py-common[ghex,netcdf]"] -ghex = ["pyghex>=0.3.0", "mpi4py<=3.1.4"] +ghex = ["ghex", "mpi4py"] netcdf = ["netcdf4>=1.6.0"] [project.urls] diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index ee1af2b906..a931a02e4e 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -47,6 +47,9 @@ SEAL_LEVEL_PRESSURE: Final[wpfloat] = 101325.0 P0SL_BG: Final[wpfloat] = SEAL_LEVEL_PRESSURE +# average earth radius in [m] +EARTH_RADIUS: Final[float] = 6.371229e6 + #: sea level temperature for reference atmosphere [K] SEA_LEVEL_TEMPERATURE: Final[wpfloat] = 288.15 T0SL_BG: Final[wpfloat] = SEA_LEVEL_TEMPERATURE 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 bcc3d98b19..d6fb764ee0 100644 --- a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py +++ b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py @@ -25,8 +25,15 @@ try: import ghex - import ghex.unstructured as unstructured import mpi4py + from ghex.context import make_context + from ghex.unstructured import ( + DomainDescriptor, + HaloGenerator, + make_communication_object, + make_field_descriptor, + make_pattern, + ) mpi4py.rc.initialize = False mpi4py.rc.finalize = True @@ -120,7 +127,7 @@ def __init__( props: definitions.ProcessProperties, domain_decomposition: definitions.DecompositionInfo, ): - self._context = ghex.context(ghex.mpi_comm(props.comm), False) + self._context = make_context(props.comm, False) self._domain_id_gen = definitions.DomainDescriptorIdGenerator(props) self._decomposition_info = domain_decomposition self._domain_descriptors = { @@ -140,7 +147,7 @@ def __init__( EdgeDim: self._create_pattern(EdgeDim), } log.info(f"patterns for dimensions {self._patterns.keys()} initialized ") - self._comm = unstructured.make_co(self._context) + self._comm = make_communication_object(self._context) log.info("communication object initialized") def _domain_descriptor_info(self, descr): @@ -162,7 +169,7 @@ def _create_domain_descriptor(self, dim: Dimension): # first arg is the domain ID which builds up an MPI Tag. # if those ids are not different for all domain descriptors the system might deadlock # if two parallel exchanges with the same domain id are done - domain_desc = unstructured.domain_descriptor( + domain_desc = DomainDescriptor( self._domain_id_gen(), all_global.tolist(), local_halo.tolist() ) log.debug( @@ -176,9 +183,9 @@ def _create_pattern(self, horizontal_dim: Dimension): global_halo_idx = self._decomposition_info.global_index( horizontal_dim, definitions.DecompositionInfo.EntryType.HALO ) - halo_generator = unstructured.halo_generator_with_gids(global_halo_idx) + halo_generator = HaloGenerator.from_gids(global_halo_idx) log.debug(f"halo generator for dim='{horizontal_dim.value}' created") - pattern = unstructured.make_pattern( + pattern = make_pattern( self._context, halo_generator, [self._domain_descriptors[horizontal_dim]], @@ -195,7 +202,7 @@ def exchange(self, dim: definitions.Dimension, *fields: Sequence[Field]): domain_descriptor = self._domain_descriptors[dim] assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found" applied_patterns = [ - pattern(unstructured.field_descriptor(domain_descriptor, f.asnumpy())) for f in fields + pattern(make_field_descriptor(domain_descriptor, f.asnumpy())) for f in fields ] handle = self._comm.exchange(applied_patterns) log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.") diff --git a/model/common/src/icon4py/model/common/grid/base.py b/model/common/src/icon4py/model/common/grid/base.py index 1373c9c25b..ede8df26d3 100644 --- a/model/common/src/icon4py/model/common/grid/base.py +++ b/model/common/src/icon4py/model/common/grid/base.py @@ -22,7 +22,6 @@ from gt4py.next.iterator.embedded import NeighborTableOffsetProvider from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim -from icon4py.model.common.grid.horizontal import HorizontalGridSize from icon4py.model.common.grid.utils import neighbortable_offset_provider_for_1d_sparse_fields from icon4py.model.common.grid.vertical import VerticalGridSize from icon4py.model.common.utils import builder @@ -32,14 +31,20 @@ class MissingConnectivity(ValueError): pass -@dataclass( - frozen=True, -) +@dataclass(frozen=True) +class HorizontalGridSize: + num_vertices: int + num_edges: int + num_cells: int + + +@dataclass(frozen=True, kw_only=True) class GridConfig: horizontal_config: HorizontalGridSize vertical_config: VerticalGridSize limited_area: bool = True n_shift_total: int = 0 + length_rescale_factor: float = 1.0 lvertnest: bool = False on_gpu: bool = False 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 6b4db3e51e..12827cdeba 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -51,9 +51,8 @@ def __init__(self, *args, **kwargs): V2EDim, VertexDim, ) -from icon4py.model.common.grid.base import GridConfig, VerticalGridSize -from icon4py.model.common.grid.horizontal import HorizontalGridSize -from icon4py.model.common.grid.icon import IconGrid +from icon4py.model.common.grid.base import GridConfig, HorizontalGridSize, VerticalGridSize +from icon4py.model.common.grid.icon import GlobalGridParams, IconGrid class GridFileName(str, Enum): @@ -81,6 +80,8 @@ class GridFile: class PropertyName(GridFileName): GRID_ID = "uuidOfHGrid" PARENT_GRID_ID = "uuidOfParHGrid" + LEVEL = "grid_level" + ROOT = "grid_root" class OffsetName(GridFileName): """Names for connectivities used in the grid file.""" @@ -361,6 +362,9 @@ def _from_grid_dataset(self, dataset: Dataset, on_gpu: bool, limited_area=True) num_cells = reader.dimension(GridFile.DimensionName.CELL_NAME) num_edges = reader.dimension(GridFile.DimensionName.EDGE_NAME) num_vertices = reader.dimension(GridFile.DimensionName.VERTEX_NAME) + grid_level = dataset.getncattr(GridFile.PropertyName.LEVEL) + grid_root = dataset.getncattr(GridFile.PropertyName.ROOT) + global_params = GlobalGridParams(level=grid_level, root=grid_root) grid_size = HorizontalGridSize( num_vertices=num_vertices, num_edges=num_edges, num_cells=num_cells @@ -396,6 +400,7 @@ def _from_grid_dataset(self, dataset: Dataset, on_gpu: bool, limited_area=True) icon_grid = ( IconGrid() .with_config(config) + .with_global_params(global_params) .with_connectivities( { C2EDim: c2e, diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py index 01d7178d93..7b85323ed5 100644 --- a/model/common/src/icon4py/model/common/grid/horizontal.py +++ b/model/common/src/icon4py/model/common/grid/horizontal.py @@ -10,13 +10,15 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +import math from dataclasses import dataclass +from functools import cached_property from typing import ClassVar, Final from gt4py.next import Dimension, Field, GridType, field_operator, neighbor_sum, program from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.common import dimension +from icon4py.model.common import constants, dimension from icon4py.model.common.dimension import ( E2C, V2C, @@ -302,8 +304,44 @@ def __init__( class CellParams: #: Area of a cell, defined in ICON in mo_model_domain.f90:t_grid_cells%area area: Field[[CellDim], float] - + #: Mean area of a cell [m^2] mean_cell_area: float + length_rescale_factor: float = 1.0 + + @classmethod + def from_global_num_cells( + cls, + area: Field[[CellDim], float], + global_num_cells: int, + length_rescale_factor: float = 1.0, + ): + mean_cell_area = cls._compute_mean_cell_area(constants.EARTH_RADIUS, global_num_cells) + return cls( + area=area, mean_cell_area=mean_cell_area, length_rescale_factor=length_rescale_factor + ) + + @cached_property + def characteristic_length(self): + return math.sqrt(self.mean_cell_area) + + @cached_property + def mean_cell_area(self): + return self.mean_cell_area + + @staticmethod + def _compute_mean_cell_area(radius, num_cells): + """ + Compute the mean cell area. + + Computes the mean cell area by dividing the sphere by the number of cells in the + global grid. + + Args: + radius: average earth radius, might be rescaled by a scaling parameter + num_cells: number of cells on the global grid + Returns: mean area of one cell [m^2] + """ + return 4.0 * math.pi * radius**2 / num_cells @field_operator diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index b383b2723c..2f3eb6879c 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -10,6 +10,8 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later +from dataclasses import dataclass +from functools import cached_property import numpy as np from gt4py.next.common import Dimension, DimensionKind @@ -42,12 +44,23 @@ from icon4py.model.common.utils import builder +@dataclass(frozen=True) +class GlobalGridParams: + root: int + level: int + + @cached_property + def num_cells(self): + return 20.0 * self.root**2 * 4.0**self.level + + class IconGrid(BaseGrid): def __init__(self): """Instantiate a grid according to the ICON model.""" super().__init__() self.start_indices = {} self.end_indices = {} + self.global_properties = None self.offset_provider_mapping = { "C2E": (self._get_offset_provider, C2EDim, CellDim, EdgeDim), "E2C": (self._get_offset_provider, E2CDim, EdgeDim, CellDim), @@ -80,6 +93,10 @@ def with_start_end_indices( self.start_indices[dim] = start_indices.astype(int32) self.end_indices[dim] = end_indices.astype(int32) + @builder + def with_global_params(self, global_params: GlobalGridParams): + self.global_properties = global_params + @property def num_levels(self): return self.config.num_levels if self.config else 0 @@ -88,6 +105,16 @@ def num_levels(self): def num_cells(self): return self.config.num_cells if self.config else 0 + @property + def global_num_cells(self): + """ + Return the number of cells in the global grid. + + If the global grid parameters are not set, it assumes that we are in a one node scenario + and returns the local number of cells. + """ + return self.global_properties.num_cells if self.global_properties else self.num_cells + @property def num_vertices(self): return self.config.num_vertices if self.config else 0 diff --git a/model/common/src/icon4py/model/common/grid/simple.py b/model/common/src/icon4py/model/common/grid/simple.py index 73c27df076..96c7b0f5d1 100644 --- a/model/common/src/icon4py/model/common/grid/simple.py +++ b/model/common/src/icon4py/model/common/grid/simple.py @@ -39,7 +39,7 @@ V2EDim, VertexDim, ) -from icon4py.model.common.grid.base import BaseGrid, GridConfig +from icon4py.model.common.grid.base import BaseGrid, GridConfig, HorizontalGridSize # periodic # @@ -59,7 +59,6 @@ # |20e \ |23e \ |26e \ # | 15c \ | 16c \ | 17c \ # 0v 1v 2v 0v -from icon4py.model.common.grid.horizontal import HorizontalGridSize from icon4py.model.common.grid.vertical import VerticalGridSize diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py index d13873b984..b2d449575c 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py @@ -23,6 +23,7 @@ SERIALIZED_DATA_PATH, create_icon_serial_data_provider, get_datapath_for_experiment, + get_global_grid_params, get_processor_properties_for_run, get_ranked_data_path, ) @@ -43,11 +44,6 @@ def ranked_data_path(processor_props): return get_ranked_data_path(SERIALIZED_DATA_PATH, processor_props) -@pytest.fixture -def datapath(ranked_data_path, experiment): - return get_datapath_for_experiment(ranked_data_path, experiment) - - @pytest.fixture def download_ser_data(request, processor_props, ranked_data_path, experiment, pytestconfig): """ @@ -83,13 +79,15 @@ def download_ser_data(request, processor_props, ranked_data_path, experiment, py @pytest.fixture -def data_provider(download_ser_data, datapath, processor_props): - return create_icon_serial_data_provider(datapath, processor_props) +def data_provider(download_ser_data, ranked_data_path, experiment, processor_props): + data_path = get_datapath_for_experiment(ranked_data_path, experiment) + return create_icon_serial_data_provider(data_path, processor_props) @pytest.fixture -def grid_savepoint(data_provider): - return data_provider.from_savepoint_grid() +def grid_savepoint(data_provider, experiment): + root, level = get_global_grid_params(experiment) + return data_provider.from_savepoint_grid(root, level) def is_regional(experiment_name): @@ -97,7 +95,7 @@ def is_regional(experiment_name): @pytest.fixture -def icon_grid(grid_savepoint, experiment): +def icon_grid(grid_savepoint): """ Load the icon grid from an ICON savepoint. @@ -107,8 +105,11 @@ def icon_grid(grid_savepoint, experiment): @pytest.fixture -def decomposition_info(data_provider): - return data_provider.from_savepoint_grid().construct_decomposition_info() +def decomposition_info(data_provider, experiment): + root, level = get_global_grid_params(experiment) + return data_provider.from_savepoint_grid( + grid_root=root, grid_level=level + ).construct_decomposition_info() @pytest.fixture diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py index d2eba8917a..d97b6a431b 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import os +import re from pathlib import Path from icon4py.model.common.decomposition.definitions import get_processor_properties @@ -47,12 +48,30 @@ def get_test_data_root_path() -> Path: DATA_URIS = { 1: "https://polybox.ethz.ch/index.php/s/xhooaubvGffG8Qy/download", - 2: "https://polybox.ethz.ch/index.php/s/YyC5qDJWyC39y7u/download", - 4: "https://polybox.ethz.ch/index.php/s/UIHOVJs6FVPpz9V/download", + 2: "https://polybox.ethz.ch/index.php/s/P6F6ZbzWHI881dZ/download", + 4: "https://polybox.ethz.ch/index.php/s/NfES3j9no15A0aX/download", } DATA_URIS_APE = {1: "https://polybox.ethz.ch/index.php/s/y9WRP1mpPlf2BtM/download"} +def get_global_grid_params(experiment: str) -> tuple[int, int]: + """Get the grid root and level from the experiment name. + + Reads the level and root parameters from a string in the canonical ICON gridfile format + RxyBab where 'xy' and 'ab' are numbers and denote the root and level of the icosahedron grid construction. + + Args: experiment: str: The experiment name. + Returns: tuple[int, int]: The grid root and level. + """ + try: + root, level = map(int, re.search("[Rr](\d+)[Bb](\d+)", experiment).groups()) + return root, level + except AttributeError as err: + raise ValueError( + f"Could not parse grid_root and grid_level from experiment: {experiment} no 'rXbY'pattern." + ) from err + + def get_processor_properties_for_run(run_instance): return get_processor_properties(run_instance) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index 7de2eea602..470d30503a 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -111,7 +111,8 @@ def constant_field( grid: BaseGrid, value: float, *dims: gt_common.Dimension, dtype=wpfloat ) -> gt_common.Field: return as_field( - dims, value * np.ones(shape=tuple(map(lambda x: grid.size[x], dims)), dtype=dtype) + dims, + value * np.ones(shape=tuple(map(lambda x: grid.size[x], dims)), dtype=dtype), ) @@ -191,7 +192,9 @@ def _test_validation(self, grid, backend, input_data): ) assert np.allclose( - input_data[name].asnumpy()[gtslice], reference_outputs[name][refslice], equal_nan=True + input_data[name].asnumpy()[gtslice], + reference_outputs[name][refslice], + equal_nan=True, ), f"Validation failed for '{name}'" diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 6856b6efb5..b45289c556 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -43,9 +43,9 @@ V2EDim, VertexDim, ) -from icon4py.model.common.grid.base import GridConfig, VerticalGridSize -from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize -from icon4py.model.common.grid.icon import IconGrid +from icon4py.model.common.grid.base import GridConfig, HorizontalGridSize, VerticalGridSize +from icon4py.model.common.grid.horizontal import CellParams, EdgeParams +from icon4py.model.common.grid.icon import GlobalGridParams, IconGrid from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims @@ -142,6 +142,10 @@ def _read(self, name: str, offset=0, dtype=int): class IconGridSavepoint(IconSavepoint): + def __init__(self, sp: ser.Savepoint, ser: ser.Serializer, size: dict, root: int, level: int): + super().__init__(sp, ser, size) + self.global_grid_params = GlobalGridParams(root, level) + def v_dual_area(self): return self._get_field("v_dual_area", VertexDim) @@ -338,6 +342,7 @@ def construct_icon_grid(self, on_gpu: bool) -> IconGrid: vertex_ends = self.vertex_end_index() edge_starts = self.edge_start_index() edge_ends = self.edge_end_index() + config = GridConfig( horizontal_config=HorizontalGridSize( num_vertices=self.num(VertexDim), @@ -355,6 +360,7 @@ def construct_icon_grid(self, on_gpu: bool) -> IconGrid: grid = ( IconGrid() .with_config(config) + .with_global_params(self.global_grid_params) .with_start_end_indices(VertexDim, vertex_starts, vertex_ends) .with_start_end_indices(EdgeDim, edge_starts, edge_ends) .with_start_end_indices(CellDim, cell_starts, cell_ends) @@ -426,7 +432,11 @@ def construct_edge_geometry(self) -> EdgeParams: ) def construct_cell_geometry(self) -> CellParams: - return CellParams(area=self.cell_areas(), mean_cell_area=self.mean_cell_area()) + return CellParams.from_global_num_cells( + area=self.cell_areas(), + global_num_cells=self.global_grid_params.num_cells, + length_rescale_factor=1.0, + ) class InterpolationSavepoint(IconSavepoint): @@ -1156,9 +1166,11 @@ def _grid_size(self): } return grid_sizes - def from_savepoint_grid(self) -> IconGridSavepoint: + def from_savepoint_grid(self, grid_root, grid_level) -> IconGridSavepoint: savepoint = self._get_icon_grid_savepoint() - return IconGridSavepoint(savepoint, self.serializer, size=self.grid_size) + return IconGridSavepoint( + savepoint, self.serializer, size=self.grid_size, root=grid_root, level=grid_level + ) def _get_icon_grid_savepoint(self): savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint() diff --git a/model/common/tests/decomposition_tests/test_mpi_decomposition.py b/model/common/tests/decomposition_tests/test_mpi_decomposition.py index 7c43001aff..d062dd1527 100644 --- a/model/common/tests/decomposition_tests/test_mpi_decomposition.py +++ b/model/common/tests/decomposition_tests/test_mpi_decomposition.py @@ -14,9 +14,11 @@ import numpy as np import pytest +from icon4py.model.common.test_utils.helpers import constant_field + try: - import mpi4py # noqa: F401 # test for optional dependency + import mpi4py # noqa: F401 # import mpi4py to check for optional mpi dependency except ImportError: pytest.skip("Skipping parallel on single node installation", allow_module_level=True) @@ -27,15 +29,15 @@ create_exchange, ) from icon4py.model.common.decomposition.mpi_decomposition import GHexMultiNodeExchange -from icon4py.model.common.dimension import CellDim, EdgeDim, VertexDim +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils data_provider, - datapath, decomposition_info, download_ser_data, experiment, grid_savepoint, icon_grid, + metrics_savepoint, ranked_data_path, ) from icon4py.model.common.test_utils.parallel_helpers import ( # noqa: F401 # import fixtures from test_utils package @@ -147,14 +149,17 @@ def test_decomposition_info_local_index( @pytest.mark.mpi @pytest.mark.parametrize("processor_props", [True], indirect=True) -@pytest.mark.parametrize("num", [1, 2, 3]) -def test_domain_descriptor_id_are_globally_unique(num, processor_props): # noqa: F811 # fixture +@pytest.mark.parametrize("num", [1, 2, 3, 4, 5, 6, 7, 8]) +def test_domain_descriptor_id_are_globally_unique( + num, + processor_props, # noqa F811 #fixture +): props = processor_props size = props.comm_size id_gen = DomainDescriptorIdGenerator(parallel_props=props) id1 = id_gen() assert id1 == props.comm_size * props.rank - assert id1 < props.comm_size * (props.rank + 1) + assert id1 < props.comm_size * (props.rank + 2) ids = [] ids.append(id1) for _ in range(1, num * size): @@ -170,6 +175,7 @@ def test_domain_descriptor_id_are_globally_unique(num, processor_props): # noqa @pytest.mark.mpi @pytest.mark.datatest +@pytest.mark.parametrize("processor_props", [True], indirect=True) def test_decomposition_info_matches_gridsize( caplog, download_ser_data, # noqa: F811 #fixture @@ -182,7 +188,7 @@ def test_decomposition_info_matches_gridsize( decomposition_info.global_index( dim=CellDim, entry_type=DecompositionInfo.EntryType.ALL ).shape[0] - == icon_grid.n_cells() + == icon_grid.num_cells ) assert ( decomposition_info.global_index(VertexDim, DecompositionInfo.EntryType.ALL).shape[0] @@ -209,11 +215,51 @@ def test_create_multi_node_runtime_with_mpi( @pytest.mark.parametrize("processor_props", [False], indirect=True) +@pytest.mark.mpi_skip() def test_create_single_node_runtime_without_mpi( processor_props, # noqa: F811 # fixture decomposition_info, # noqa: F811 # fixture ): - props = processor_props - exchange = create_exchange(props, decomposition_info) - + exchange = create_exchange(processor_props, decomposition_info) assert isinstance(exchange, SingleNodeExchange) + + +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.parametrize("dimension", (CellDim, VertexDim, EdgeDim)) +def test_exchange_on_dummy_data( + processor_props, # noqa: F811 # fixture + decomposition_info, # noqa: F811 # fixture + grid_savepoint, # noqa: F811 # fixture + metrics_savepoint, # noqa: F811 # fixture + dimension, +): + exchange = create_exchange(processor_props, decomposition_info) + grid = grid_savepoint.construct_icon_grid(on_gpu=False) + + number = processor_props.rank + 10.0 + input_field = constant_field( + grid, + number, + dimension, + KDim, + ) + + halo_points = decomposition_info.local_index(dimension, DecompositionInfo.EntryType.HALO) + local_points = decomposition_info.local_index(dimension, DecompositionInfo.EntryType.OWNED) + assert np.all(input_field == 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( + f" rank={processor_props.rank} - exchanged points: {np.sum(result != number)/grid.num_levels}" + ) + print(f"rank={processor_props.rank} - halo points: {halo_points}") + + assert np.all(result[local_points, :] == number) + assert np.all(result[halo_points, :] != number) + + changed_points = np.argwhere(result[:, 2] != number) + print(f"rank={processor_props.rank} - num changed points {changed_points.shape} ") + + print(f"rank={processor_props.rank} - changed points {changed_points} ") diff --git a/model/common/tests/grid_tests/conftest.py b/model/common/tests/grid_tests/conftest.py index 3b03703cd1..f3484bbaf6 100644 --- a/model/common/tests/grid_tests/conftest.py +++ b/model/common/tests/grid_tests/conftest.py @@ -15,7 +15,6 @@ from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 damping_height, data_provider, - datapath, decomposition_info, download_ser_data, experiment, diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index 6e80133759..680e862d57 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -70,6 +70,7 @@ MCH_CH_04B09_NUM_VERTICES = 10663 MCH_CH_R04B09_LOCAL_NUM_EDGES = 31558 MCH_CH_RO4B09_LOCAL_NUM_CELLS = 20896 +MCH_CH_RO4B09_GLOBAL_NUM_CELLS = 83886080 MCH_CH_R04B09_CELL_DOMAINS = { @@ -113,8 +114,11 @@ def simple_grid_gridfile(tmp_path): path = tmp_path.joinpath(SIMPLE_GRID_NC).absolute() grid = SimpleGrid() + dataset = netCDF4.Dataset(path, "w", format="NETCDF4") dataset.setncattr(GridFile.PropertyName.GRID_ID, str(uuid4())) + dataset.setncattr(GridFile.PropertyName.LEVEL, 0) + dataset.setncattr(GridFile.PropertyName.ROOT, 0) dataset.createDimension(GridFile.DimensionName.VERTEX_NAME, size=grid.num_vertices) dataset.createDimension(GridFile.DimensionName.EDGE_NAME, size=grid.num_edges) @@ -935,3 +939,16 @@ def test_get_start_end_index_for_global_grid( from_grid_file = init_grid_manager(file, num_levels=num_levels).get_grid() assert from_grid_file.get_start_index(dim, marker) == start_index assert from_grid_file.get_end_index(dim, marker) == end_index + + +@pytest.mark.parametrize( + "grid_file, global_num_cells", + [ + (R02B04_GLOBAL, R02B04_GLOBAL_NUM_CELLS), + (REGIONAL_EXPERIMENT, MCH_CH_RO4B09_GLOBAL_NUM_CELLS), + ], +) +def test_grid_level_and_root(grid_file, global_num_cells): + file = resolve_file_from_gridfile_name(grid_file) + grid = init_grid_manager(file, num_levels=10).get_grid() + assert global_num_cells == grid.global_num_cells diff --git a/model/common/tests/grid_tests/test_horizontal.py b/model/common/tests/grid_tests/test_horizontal.py new file mode 100644 index 0000000000..08f6108501 --- /dev/null +++ b/model/common/tests/grid_tests/test_horizontal.py @@ -0,0 +1,30 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import pytest + +from icon4py.model.common import constants +from icon4py.model.common.grid.horizontal import CellParams +from icon4py.model.common.grid.icon import GlobalGridParams + + +@pytest.mark.parametrize( + "grid_root,grid_level,expected", + [ + (2, 4, 24907282236.708576), + (4, 9, 6080879.45232143), + ], +) +def test_mean_cell_area_calculation(grid_root, grid_level, expected): + params = GlobalGridParams(grid_root, grid_level) + assert expected == CellParams._compute_mean_cell_area(constants.EARTH_RADIUS, params.num_cells) diff --git a/model/common/tests/interpolation_tests/conftest.py b/model/common/tests/interpolation_tests/conftest.py index 831b6b4731..1657376fa4 100644 --- a/model/common/tests/interpolation_tests/conftest.py +++ b/model/common/tests/interpolation_tests/conftest.py @@ -13,7 +13,6 @@ from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, - datapath, download_ser_data, experiment, grid_savepoint, diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields.py b/model/common/tests/interpolation_tests/test_interpolation_fields.py index cb42c70fb6..0666111900 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields.py @@ -43,7 +43,6 @@ ) from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, - datapath, download_ser_data, experiment, processor_props, diff --git a/model/common/tests/metric_tests/conftest.py b/model/common/tests/metric_tests/conftest.py index 65b35da272..4e97877a56 100644 --- a/model/common/tests/metric_tests/conftest.py +++ b/model/common/tests/metric_tests/conftest.py @@ -13,7 +13,6 @@ from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, - datapath, download_ser_data, experiment, grid_savepoint, diff --git a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py index 59dccd4756..aaa8981362 100644 --- a/model/common/tests/metric_tests/test_compute_nudgecoeffs.py +++ b/model/common/tests/metric_tests/test_compute_nudgecoeffs.py @@ -32,7 +32,6 @@ from icon4py.model.common.metrics.stencils.compute_nudgecoeffs import compute_nudgecoeffs from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, - datapath, download_ser_data, experiment, grid_savepoint, diff --git a/model/driver/tests/conftest.py b/model/driver/tests/conftest.py index 3e6786593c..9317e0718e 100644 --- a/model/driver/tests/conftest.py +++ b/model/driver/tests/conftest.py @@ -19,7 +19,6 @@ from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 damping_height, data_provider, - datapath, download_ser_data, experiment, grid_savepoint, diff --git a/requirements-dev-opt.txt b/requirements-dev-opt.txt index b5d0b8bd15..8c70c768b5 100644 --- a/requirements-dev-opt.txt +++ b/requirements-dev-opt.txt @@ -1,4 +1,4 @@ -git+https://github.com/ghex-org/GHEX.git#subdirectory=bindings/python +git+https://github.com/ghex-org/GHEX.git@master#subdirectory=bindings/python -r base-requirements-dev.txt # icon4py model -e ./model/common[all] diff --git a/tools/src/icon4pytools/icon4pygen/backend.py b/tools/src/icon4pytools/icon4pygen/backend.py index 25445d1a17..82259a021c 100644 --- a/tools/src/icon4pytools/icon4pygen/backend.py +++ b/tools/src/icon4pytools/icon4pygen/backend.py @@ -33,7 +33,9 @@ GRID_SIZE_ARGS = ["num_cells", "num_edges", "num_vertices"] -def transform_and_configure_fencil(fencil: itir.FencilDefinition) -> itir.FencilDefinition: +def transform_and_configure_fencil( + fencil: itir.FencilDefinition, +) -> itir.FencilDefinition: """Transform the domain representation and configure the FencilDefinition parameters.""" grid_size_symbols = [itir.Sym(id=arg) for arg in GRID_SIZE_ARGS]