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]