Fix parallel tests (#428)

- adapt mpi test to latest ghex version
- compute mean_cell_area from grid file properties
- remove datapath fixture

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,

model/atmosphere/diffusion/tests/diffusion_tests/mpi_tests/test_parallel_diffusion.py

@@ -18,27 +18,27 @@
 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,
     verify_diffusion_fields,
 )
 
 
-@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,

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

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,

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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later

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,

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]

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

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.")

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