- move the poormans GridGeometry cache to common/test_utils/grid_util…

…s.py (#613)

- renameing of functions in grid_utils.py
- add poormans grid_manager cache in test_grid_manager.py to speed up tests (replacement for @functools.cache that does no longer work)
- add convenience as_numpy function

model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py

@@ -61,7 +61,7 @@ def _construct_minimal_decomposition_info(grid: icon.IconGrid):
         return decomposition_info
 
     if not grid_functionality[experiment].get(name):
-        gm = grid_utils.get_icon_grid_from_gridfile(experiment, backend)
+        gm = grid_utils.get_grid_manager_for_experiment(experiment, backend)
         grid = gm.grid
         decomposition_info = _construct_minimal_decomposition_info(grid)
         geometry_ = geometry.GridGeometry(

model/common/src/icon4py/model/common/grid/grid_manager.py

@@ -394,7 +394,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_type is FileNotFoundError:
             raise FileNotFoundError(f"gridfile {self._file_name} not found, aborting")
 
-    def __call__(self, backend: gtx_backend.Backend, limited_area=True):
+    def __call__(self, backend: Optional[gtx_backend.Backend], limited_area=True):
         if not self._reader:
             self.open()
         on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend)
@@ -403,7 +403,7 @@ def __call__(self, backend: gtx_backend.Backend, limited_area=True):
         self._coordinates = self._read_coordinates(backend)
         self._geometry = self._read_geometry_fields(backend)
 
-    def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict:
+    def _read_coordinates(self, backend: Optional[gtx_backend.Backend]) -> CoordinateDict:
         return {
             dims.CellDim: {
                 "lat": gtx.as_field(
@@ -449,7 +449,7 @@ def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict:
             },
         }
 
-    def _read_geometry_fields(self, backend: gtx_backend.Backend):
+    def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]):
         return {
             # TODO (@halungge) still needs to ported, values from "our" grid files contains (wrong) values:
             #   based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) .

model/common/src/icon4py/model/common/test_utils/grid_utils.py

@@ -5,14 +5,23 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+import pathlib
 
-
+import gt4py.next as gtx
 import gt4py.next.backend as gtx_backend
 import pytest
 
 import icon4py.model.common.grid.grid_manager as gm
-from icon4py.model.common.grid import vertical as v_grid
+from icon4py.model.common import dimension as dims
+from icon4py.model.common.decomposition import definitions
+from icon4py.model.common.grid import (
+    geometry,
+    geometry_attributes as geometry_attrs,
+    icon,
+    vertical as v_grid,
+)
 from icon4py.model.common.test_utils import data_handling, datatest_utils as dt_utils
+from icon4py.model.common.utils import gt4py_field_allocation as alloc
 
 
 REGIONAL_GRIDFILE = "grid.nc"
@@ -23,62 +32,139 @@
 
 MCH_CH_R04B09_LEVELS = 65
 
+grid_geometries = {}
+
 
-def get_icon_grid_from_gridfile(
+def get_grid_manager_for_experiment(
     experiment: str, backend: gtx_backend.Backend = None
 ) -> gm.GridManager:
     if experiment == dt_utils.GLOBAL_EXPERIMENT:
-        return _download_and_load_from_gridfile(
+        return _download_and_load_gridfile(
             dt_utils.R02B04_GLOBAL,
-            GLOBAL_GRIDFILE,
             num_levels=GLOBAL_NUM_LEVELS,
-            limited_area=False,
             backend=backend,
         )
     elif experiment == dt_utils.REGIONAL_EXPERIMENT:
-        return _download_and_load_from_gridfile(
+        return _download_and_load_gridfile(
             dt_utils.REGIONAL_EXPERIMENT,
-            REGIONAL_GRIDFILE,
             num_levels=MCH_CH_R04B09_LEVELS,
-            limited_area=True,
             backend=backend,
         )
     else:
         raise ValueError(f"Unknown experiment: {experiment}")
 
 
-def download_grid_file(file_path: str, filename: str):
-    grid_file = dt_utils.GRIDS_PATH.joinpath(file_path, filename)
-    if not grid_file.exists():
+def get_grid_manager(
+    grid_file: str, num_levels: int, backend: gtx_backend.Backend
+) -> gm.GridManager:
+    return _download_and_load_gridfile(grid_file, num_levels=num_levels, backend=backend)
+
+
+def _file_name(grid_file: str):
+    match grid_file:
+        case dt_utils.REGIONAL_EXPERIMENT:
+            return REGIONAL_GRIDFILE
+        case dt_utils.R02B04_GLOBAL:
+            return GLOBAL_GRIDFILE
+        case _:
+            raise NotImplementedError(f"Add grid path for experiment '{grid_file}'")
+
+
+def resolve_full_grid_file_name(grid_file_str: str) -> pathlib.Path:
+    return dt_utils.GRIDS_PATH.joinpath(grid_file_str, _file_name(grid_file_str))
+
+
+def _download_grid_file(file_path: str) -> pathlib.Path:
+    full_name = resolve_full_grid_file_name(file_path)
+    if not full_name.exists():
         data_handling.download_and_extract(
             dt_utils.GRID_URIS[file_path],
-            grid_file.parent,
-            grid_file.parent,
+            full_name.parent,
+            full_name.parent,
         )
-    return grid_file
+    return full_name
 
 
-def load_grid_from_file(
-    grid_file: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool
+def _run_grid_manager_for_file(
+    file: str, num_levels: int, backend: gtx_backend.Backend
 ) -> gm.GridManager:
+    """
+    Load a grid file.
+    Args:
+        file: full path to the file (file + path)
+        num_levels: number of vertical levels, needed for IconGrid construction but independent from grid file
+        backend: the gt4py Backend we are running on
+
+    Returns:
+
+    """
+    limited_area = is_regional(str(file))
+    transformation = gm.ToZeroBasedIndexTransformation()
     manager = gm.GridManager(
-        gm.ToZeroBasedIndexTransformation(),
-        str(grid_file),
+        transformation,
+        file,
         v_grid.VerticalGridConfig(num_levels=num_levels),
     )
     manager(backend=backend, limited_area=limited_area)
+    manager.close()
     return manager
 
 
-def _download_and_load_from_gridfile(
-    file_path: str, filename: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool
+def _download_and_load_gridfile(
+    file_path: str, num_levels: int, backend: gtx_backend.Backend
 ) -> gm.GridManager:
-    grid_file = download_grid_file(file_path, filename)
-
-    gm = load_grid_from_file(grid_file, num_levels, backend, limited_area)
+    grid_file = _download_grid_file(file_path)
+    gm = _run_grid_manager_for_file(str(grid_file), num_levels, backend)
     return gm
 
 
+def is_regional(experiment_or_file: str):
+    return (
+        dt_utils.REGIONAL_EXPERIMENT in experiment_or_file
+        or REGIONAL_GRIDFILE in experiment_or_file
+    )
+
+
+def get_num_levels(experiment: str):
+    return MCH_CH_R04B09_LEVELS if experiment == dt_utils.REGIONAL_EXPERIMENT else GLOBAL_NUM_LEVELS
+
+
+def get_grid_geometry(
+    backend: gtx_backend.Backend, experiment: str, grid_file: str
+) -> geometry.GridGeometry:
+    on_gpu = alloc.is_cupy_device(backend)
+    xp = alloc.array_ns(on_gpu)
+    num_levels = get_num_levels(experiment)
+
+    def construct_decomposition_info(grid: icon.IconGrid) -> definitions.DecompositionInfo:
+        def _add_dimension(dim: gtx.Dimension):
+            indices = alloc.allocate_indices(dim, grid)
+            owner_mask = xp.ones((grid.size[dim],), dtype=bool)
+            decomposition_info.with_dimension(dim, indices.ndarray, owner_mask)
+
+        decomposition_info = definitions.DecompositionInfo(klevels=grid.num_levels)
+        _add_dimension(dims.EdgeDim)
+        _add_dimension(dims.VertexDim)
+        _add_dimension(dims.CellDim)
+
+        return decomposition_info
+
+    def construct_grid_geometry(grid_file: str):
+        gm = _run_grid_manager_for_file(grid_file, backend=backend, num_levels=num_levels)
+        grid = gm.grid
+        decomposition_info = construct_decomposition_info(grid)
+        geometry_source = geometry.GridGeometry(
+            grid, decomposition_info, backend, gm.coordinates, gm.geometry, geometry_attrs.attrs
+        )
+        return geometry_source
+
+    if not grid_geometries.get(grid_file):
+        grid_geometries[grid_file] = construct_grid_geometry(
+            str(resolve_full_grid_file_name(grid_file))
+        )
+    return grid_geometries[grid_file]
+
+
 @pytest.fixture
 def grid(request):
     return request.param

model/common/src/icon4py/model/common/test_utils/pytest_config.py

@@ -163,18 +163,18 @@ def pytest_generate_tests(metafunc):
                 grid_instance = SimpleGrid()
             elif selected_grid_type == "icon_grid":
                 from icon4py.model.common.test_utils.grid_utils import (
-                    get_icon_grid_from_gridfile,
+                    get_grid_manager_for_experiment,
                 )
 
-                grid_instance = get_icon_grid_from_gridfile(
+                grid_instance = get_grid_manager_for_experiment(
                     REGIONAL_EXPERIMENT, backend=selected_backend
                 ).grid
             elif selected_grid_type == "icon_grid_global":
                 from icon4py.model.common.test_utils.grid_utils import (
-                    get_icon_grid_from_gridfile,
+                    get_grid_manager_for_experiment,
                 )
 
-                grid_instance = get_icon_grid_from_gridfile(
+                grid_instance = get_grid_manager_for_experiment(
                     GLOBAL_EXPERIMENT, backend=selected_backend
                 ).grid
             else:

model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py

@@ -6,10 +6,11 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 import logging as log
-from typing import Optional
+from typing import Optional, TypeAlias, Union
 
 import gt4py._core.definitions as gt_core_defs
 import gt4py.next as gtx
+import numpy as np
 from gt4py.next import backend
 
 from icon4py.model.common import dimension, type_alias as ta
@@ -26,6 +27,22 @@
 )
 
 
+try:
+    import cupy as xp
+except ImportError:
+    import numpy as xp
+
+
+NDArrayInterface: TypeAlias = Union[np.ndarray, xp.ndarray, gtx.Field]
+
+
+def as_numpy(array: NDArrayInterface):
+    if isinstance(array, np.ndarray):
+        return array
+    else:
+        return array.asnumpy()
+
+
 def is_cupy_device(backend: backend.Backend) -> bool:
     if backend is not None:
         return backend.allocator.__gt_device_type__ in CUDA_DEVICE_TYPES