From 66452b9d98e17fb555b924b57826d6ae19eb0abf Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 29 Nov 2024 14:05:02 +0100 Subject: [PATCH] - move the poormans GridGeometry cache to common/test_utils/grid_utils.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 --- .../tests/diffusion_tests/test_diffusion.py | 2 +- .../icon4py/model/common/grid/grid_manager.py | 6 +- .../model/common/test_utils/grid_utils.py | 134 ++++++++++++++---- .../model/common/test_utils/pytest_config.py | 8 +- .../common/utils/gt4py_field_allocation.py | 19 ++- .../common/tests/grid_tests/test_geometry.py | 102 +++++-------- .../tests/grid_tests/test_grid_manager.py | 119 ++++++++-------- model/common/tests/grid_tests/test_icon.py | 3 +- model/common/tests/grid_tests/utils.py | 56 +------- model/common/tests/io_tests/test_io.py | 12 +- 10 files changed, 249 insertions(+), 212 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index b463a62e6..68e4b6bba 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/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( 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 95b3ef5b9..4f40a0131 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/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) . diff --git a/model/common/src/icon4py/model/common/test_utils/grid_utils.py b/model/common/src/icon4py/model/common/test_utils/grid_utils.py index dbb1628cd..35f4aa1e3 100644 --- a/model/common/src/icon4py/model/common/test_utils/grid_utils.py +++ b/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 diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py index b89512aed..7d17ecc4d 100644 --- a/model/common/src/icon4py/model/common/test_utils/pytest_config.py +++ b/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: diff --git a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index 137ead06b..638b91a8f 100644 --- a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py +++ b/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 diff --git a/model/common/tests/grid_tests/test_geometry.py b/model/common/tests/grid_tests/test_geometry.py index 7d5260eb2..31789099f 100644 --- a/model/common/tests/grid_tests/test_geometry.py +++ b/model/common/tests/grid_tests/test_geometry.py @@ -11,48 +11,20 @@ import pytest from icon4py.model.common import dimension as dims -from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import ( geometry as geometry, geometry_attributes as attrs, horizontal as h_grid, - icon, simple as simple, ) from icon4py.model.common.grid.geometry import as_sparse_field -from icon4py.model.common.test_utils import datatest_utils as dt_utils, helpers -from icon4py.model.common.utils import gt4py_field_allocation as alloc - -from . import utils - - -grid_geometries = {} - - -def get_grid_geometry(backend, grid_file) -> geometry.GridGeometry: - def construct_decomposition_info(grid: icon.IconGrid) -> definitions.DecompositionInfo: - edge_indices = alloc.allocate_indices(dims.EdgeDim, grid) - owner_mask = np.ones((grid.num_edges,), dtype=bool) - decomposition_info = definitions.DecompositionInfo(klevels=grid.num_levels) - decomposition_info.with_dimension(dims.EdgeDim, edge_indices.ndarray, owner_mask) - return decomposition_info - - def construct_grid_geometry(grid_file: str): - gm = utils.run_grid_manager(grid_file, backend=backend) - grid = gm.grid - decomposition_info = construct_decomposition_info(grid) - geometry_source = geometry.GridGeometry( - grid, decomposition_info, backend, gm.coordinates, gm.geometry, attrs.attrs - ) - return geometry_source - - if not grid_geometries.get(grid_file): - grid_geometries[grid_file] = construct_grid_geometry(grid_file) - return grid_geometries[grid_file] +from icon4py.model.common.test_utils import datatest_utils as dt_utils, grid_utils, helpers def test_geometry_raises_for_unknown_field(backend): - geometry = get_grid_geometry(backend, dt_utils.R02B04_GLOBAL) + geometry = grid_utils.get_grid_geometry( + backend, dt_utils.GLOBAL_EXPERIMENT, dt_utils.R02B04_GLOBAL + ) with pytest.raises(ValueError) as e: geometry.get("foo") assert "'foo'" in e.value @@ -67,9 +39,9 @@ def test_geometry_raises_for_unknown_field(backend): ], ) @pytest.mark.datatest -def test_edge_control_area(backend, grid_savepoint, grid_file, rtol): +def test_edge_control_area(backend, grid_savepoint, grid_file, experiment, rtol): expected = grid_savepoint.edge_areas() - geometry_source = get_grid_geometry(backend, grid_file) + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) result = geometry_source.get(attrs.EDGE_AREA) assert helpers.dallclose(expected.ndarray, result.ndarray, rtol) @@ -82,8 +54,8 @@ def test_edge_control_area(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_coriolis_parameter(backend, grid_savepoint, grid_file): - geometry_source = get_grid_geometry(backend, grid_file) +def test_coriolis_parameter(backend, grid_savepoint, grid_file, experiment): + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.f_e() result = geometry_source.get(attrs.CORIOLIS_PARAMETER) @@ -98,8 +70,8 @@ def test_coriolis_parameter(backend, grid_savepoint, grid_file): ], ) @pytest.mark.datatest -def test_compute_edge_length(backend, grid_savepoint, grid_file, rtol): - geometry_source = get_grid_geometry(backend, grid_file) +def test_compute_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.primal_edge_length() result = geometry_source.get(attrs.EDGE_LENGTH) assert helpers.dallclose(result.ndarray, expected.ndarray, rtol=rtol) @@ -113,9 +85,9 @@ def test_compute_edge_length(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_compute_inverse_edge_length(backend, grid_savepoint, grid_file, rtol): +def test_compute_inverse_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): expected = grid_savepoint.inverse_primal_edge_lengths() - geometry_source = get_grid_geometry(backend, grid_file) + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) computed = geometry_source.get(f"inverse_of_{attrs.EDGE_LENGTH}") assert helpers.dallclose(computed.ndarray, expected.ndarray, rtol=rtol) @@ -129,8 +101,8 @@ def test_compute_inverse_edge_length(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_compute_dual_edge_length(backend, grid_savepoint, grid_file, rtol): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_dual_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.dual_edge_length() result = grid_geometry.get(attrs.DUAL_EDGE_LENGTH) @@ -145,8 +117,8 @@ def test_compute_dual_edge_length(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_compute_inverse_dual_edge_length(backend, grid_savepoint, grid_file, rtol): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_inverse_dual_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.inv_dual_edge_length() result = grid_geometry.get(f"inverse_of_{attrs.DUAL_EDGE_LENGTH}") @@ -166,8 +138,8 @@ def test_compute_inverse_dual_edge_length(backend, grid_savepoint, grid_file, rt ], ) @pytest.mark.datatest -def test_compute_inverse_vertex_vertex_length(backend, grid_savepoint, grid_file, rtol): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_inverse_vertex_vertex_length(backend, grid_savepoint, grid_file, experiment, rtol): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.inv_vert_vert_length() result = grid_geometry.get(attrs.INVERSE_VERTEX_VERTEX_LENGTH) @@ -182,8 +154,10 @@ def test_compute_inverse_vertex_vertex_length(backend, grid_savepoint, grid_file (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_compute_coordinates_of_edge_tangent_and_normal(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_coordinates_of_edge_tangent_and_normal( + backend, grid_savepoint, grid_file, experiment +): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) x_normal = grid_geometry.get(attrs.EDGE_NORMAL_X) y_normal = grid_geometry.get(attrs.EDGE_NORMAL_Y) z_normal = grid_geometry.get(attrs.EDGE_NORMAL_Z) @@ -213,8 +187,8 @@ def test_compute_coordinates_of_edge_tangent_and_normal(backend, grid_savepoint, (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_compute_primal_normals(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_primal_normals(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) primal_normal_u = grid_geometry.get(attrs.EDGE_NORMAL_U) primal_normal_v = grid_geometry.get(attrs.EDGE_NORMAL_V) @@ -233,8 +207,8 @@ def test_compute_primal_normals(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_tangent_orientation(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_tangent_orientation(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) result = grid_geometry.get(attrs.TANGENT_ORIENTATION) expected = grid_savepoint.tangent_orientation() @@ -249,8 +223,8 @@ def test_tangent_orientation(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_cell_area(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_cell_area(backend, grid_savepoint, experiment, grid_file): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) result = grid_geometry.get(attrs.CELL_AREA) expected = grid_savepoint.cell_areas() @@ -265,8 +239,8 @@ def test_cell_area(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_primal_normal_cell(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_primal_normal_cell(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) primal_normal_cell_u_ref = grid_savepoint.primal_normal_cell_x().ndarray primal_normal_cell_v_ref = grid_savepoint.primal_normal_cell_y().ndarray primal_normal_cell_u = grid_geometry.get(attrs.EDGE_NORMAL_CELL_U) @@ -284,8 +258,8 @@ def test_primal_normal_cell(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_dual_normal_cell(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_dual_normal_cell(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) dual_normal_cell_u_ref = grid_savepoint.dual_normal_cell_x().ndarray dual_normal_cell_v_ref = grid_savepoint.dual_normal_cell_y().ndarray dual_normal_cell_u = grid_geometry.get(attrs.EDGE_TANGENT_CELL_U) @@ -303,8 +277,8 @@ def test_dual_normal_cell(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_primal_normal_vert(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_primal_normal_vert(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) primal_normal_vert_u_ref = grid_savepoint.primal_normal_vert_x().ndarray primal_normal_vert_v_ref = grid_savepoint.primal_normal_vert_y().ndarray primal_normal_vert_u = grid_geometry.get(attrs.EDGE_NORMAL_VERTEX_U) @@ -322,8 +296,8 @@ def test_primal_normal_vert(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_dual_normal_vert(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_dual_normal_vert(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) dual_normal_vert_u_ref = grid_savepoint.dual_normal_vert_x().ndarray dual_normal_vert_v_ref = grid_savepoint.dual_normal_vert_y().ndarray dual_normal_vert_u = grid_geometry.get(attrs.EDGE_TANGENT_VERTEX_U) @@ -357,7 +331,7 @@ def test_sparse_fields_creator(): ], ) def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file, backend=backend) + gm = grid_utils.get_grid_manager(grid_file, backend=backend, num_levels=1) grid = gm.grid coordinates = gm.coordinates @@ -366,7 +340,7 @@ def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_ edge_lat = coordinates[dims.EdgeDim]["lat"] edge_lon = coordinates[dims.EdgeDim]["lon"] lat_0, lon_0, lat_1, lon_1 = geometry.create_auxiliary_coordinate_arrays_for_orientation( - gm.grid, cell_lat, cell_lon, edge_lat, edge_lon + grid, cell_lat, cell_lon, edge_lat, edge_lon ) connectivity = grid.connectivities[dims.E2CDim] has_boundary_edges = np.count_nonzero(connectivity == -1) diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index 48dc1b508..19bef3ac0 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -13,8 +13,8 @@ import numpy as np import pytest +from gt4py.next import backend as gtx_backend -import icon4py.model.common.test_utils.datatest_utils as dt_utils from icon4py.model.common import dimension as dims from icon4py.model.common.grid import ( grid_manager as gm, @@ -22,9 +22,11 @@ vertical as v_grid, ) from icon4py.model.common.grid.grid_manager import GeometryName -from icon4py.model.common.test_utils import helpers - -from .utils import run_grid_manager +from icon4py.model.common.test_utils import ( + datatest_utils as dt_utils, + grid_utils as gridtest_utils, + helpers, +) if typing.TYPE_CHECKING: @@ -47,12 +49,21 @@ MCH_CH_RO4B09_GLOBAL_NUM_CELLS = 83886080 -zero_base = gm.ToZeroBasedIndexTransformation() +ZERO_BASE = gm.ToZeroBasedIndexTransformation() + +managers = {} + + +def _run_grid_manager(file: str, backend: gtx_backend.Backend) -> gm.GridManager: + if not managers.get(file): + manager = gridtest_utils.get_grid_manager(file, num_levels=1, backend=backend) + managers[file] = manager + return managers.get(file) @pytest.fixture def global_grid_file(): - return utils.resolve_file_from_gridfile_name(dt_utils.R02B04_GLOBAL) + return gridtest_utils.resolve_full_grid_file_name(dt_utils.R02B04_GLOBAL) @pytest.mark.with_netcdf @@ -79,7 +90,7 @@ def test_grid_file_dimension(global_grid_file): ], ) def test_grid_file_vertex_cell_edge_dimensions(grid_savepoint, grid_file): - file = utils.resolve_file_from_gridfile_name(grid_file) + file = gridtest_utils.resolve_full_grid_file_name(grid_file) parser = gm.GridFile(str(file)) try: parser.open() @@ -134,10 +145,9 @@ def test_grid_file_index_fields(global_grid_file, caplog, icon_grid): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_v2e(caplog, grid_savepoint, experiment, grid_file, backend): caplog.set_level(logging.DEBUG) - manager = run_grid_manager(grid_file, transformation=zero_base) - grid = manager.grid + grid = _run_grid_manager(grid_file, backend).grid seralized_v2e = grid_savepoint.v2e() # there are vertices at the boundary of a local domain or at a pentagon point that have less than # 6 neighbors hence there are "Missing values" in the grid file @@ -159,9 +169,8 @@ def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", [dims.CellDim, dims.EdgeDim, dims.VertexDim]) -def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): - manager = run_grid_manager(grid_file, transformation=zero_base) - refin_ctrl = manager.refinement +def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim, backend): + refin_ctrl = _run_grid_manager(grid_file, backend).refinement refin_ctrl_serialized = grid_savepoint.refin_ctrl(dim) assert np.all( refin_ctrl_serialized.ndarray @@ -179,9 +188,9 @@ def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_v2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_v2c(caplog, grid_savepoint, experiment, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_v2c = grid_savepoint.v2c() # there are vertices that have less than 6 neighboring cells: either pentagon points or # vertices at the boundary of the domain for a limited area mode @@ -231,9 +240,9 @@ def reset_invalid_index(index_array: np.ndarray): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_e2v = grid_savepoint.e2v() # all vertices in the system have to neighboring edges, there no edges that point nowhere @@ -266,7 +275,7 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): grid_file: name of grid file used """ - if utils.is_regional(grid_file): + if gridtest_utils.is_regional(grid_file): assert has_invalid_index(e2c_table) else: assert not has_invalid_index(e2c_table) @@ -282,9 +291,10 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + + grid = _run_grid_manager(grid_file, backend).grid serialized_e2c = grid_savepoint.e2c() e2c_table = grid.get_offset_provider("E2C").table assert_invalid_indices(serialized_e2c, grid_file) @@ -302,9 +312,9 @@ def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_c2e = grid_savepoint.c2e() # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the @@ -325,9 +335,9 @@ def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid assert np.allclose( grid.get_offset_provider("C2E2C").table, grid_savepoint.c2e2c(), @@ -343,9 +353,9 @@ def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) assert np.allclose( grid.get_offset_provider("C2E2CO").table, @@ -363,9 +373,9 @@ def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) serialized_e2c2e = serialized_grid.get_offset_provider("E2C2E").table serialized_e2c2eO = serialized_grid.get_offset_provider("E2C2EO").table @@ -395,10 +405,9 @@ def assert_unless_invalid(table, serialized_ref): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - gm = run_grid_manager(grid_file) - grid = gm.grid + grid = _run_grid_manager(grid_file, backend).grid # the "far" (adjacent to edge normal ) is not always there, because ICON only calculates those starting from # (lateral_boundary(dims.EdgeDim) + 1) to end(dims.EdgeDim) (see mo_intp_coeffs.f90) and only for owned cells serialized_ref = grid_savepoint.e2c2v() @@ -416,9 +425,9 @@ def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid c2v = grid.get_offset_provider("C2V").table assert np.allclose(c2v, grid_savepoint.c2v()) @@ -432,8 +441,8 @@ def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.with_netcdf -def test_grid_manager_grid_size(dim, size): - grid = run_grid_manager(utils.R02B04_GLOBAL).grid +def test_grid_manager_grid_size(dim, size, backend): + grid = _run_grid_manager(utils.R02B04_GLOBAL, backend=backend).grid assert size == grid.size[dim] @@ -472,8 +481,8 @@ def test_gt4py_transform_offset_by_1_where_valid(size): (dt_utils.REGIONAL_EXPERIMENT, MCH_CH_RO4B09_GLOBAL_NUM_CELLS), ], ) -def test_grid_manager_grid_level_and_root(grid_file, global_num_cells): - assert global_num_cells == run_grid_manager(grid_file, num_levels=1).grid.global_num_cells +def test_grid_manager_grid_level_and_root(grid_file, global_num_cells, backend): + assert global_num_cells == _run_grid_manager(grid_file, backend=backend).grid.global_num_cells @pytest.mark.datatest @@ -482,9 +491,9 @@ def test_grid_manager_grid_level_and_root(grid_file, global_num_cells): "grid_file, experiment", [(utils.R02B04_GLOBAL, dt_utils.JABW_EXPERIMENT)], ) -def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) assert np.allclose( grid.get_offset_provider("C2E2C2E").table, @@ -503,12 +512,10 @@ def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", utils.horizontal_dim()) -def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid): +def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid, backend): caplog.set_level(logging.INFO) serialized_grid = icon_grid - manager = run_grid_manager(grid_file, transformation=zero_base) - grid = manager.grid - + grid = _run_grid_manager(grid_file, backend).grid for domain in utils.global_grid_domains(dim): assert grid.start_index(domain) == serialized_grid.start_index( domain @@ -518,7 +525,7 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g ), f"end index wrong for domain {domain}" for domain in utils.valid_boundary_zones_for_dim(dim): - if not utils.is_regional(grid_file): + if not gridtest_utils.is_regional(grid_file): assert grid.start_index(domain) == 0 assert grid.end_index(domain) == 0 assert grid.start_index(domain) == serialized_grid.start_index( @@ -537,10 +544,10 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_read_geometry_fields(grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file) - cell_area = gm.geometry[GeometryName.CELL_AREA.value] - tangent_orientation = gm.geometry[GeometryName.TANGENT_ORIENTATION.value] +def test_read_geometry_fields(grid_savepoint, grid_file, backend): + manager = _run_grid_manager(grid_file, backend=backend) + cell_area = manager.geometry[GeometryName.CELL_AREA.value] + tangent_orientation = manager.geometry[GeometryName.TANGENT_ORIENTATION.value] assert helpers.dallclose(cell_area.asnumpy(), grid_savepoint.cell_areas().asnumpy()) assert helpers.dallclose( @@ -557,10 +564,10 @@ def test_read_geometry_fields(grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", (dims.CellDim, dims.EdgeDim, dims.VertexDim)) -def test_coordinates(grid_savepoint, grid_file, experiment, dim): - gm = utils.run_grid_manager(grid_file) - lat = gm.coordinates[dim]["lat"] - lon = gm.coordinates[dim]["lon"] +def test_coordinates(grid_savepoint, grid_file, experiment, dim, backend): + manager = _run_grid_manager(grid_file, backend=backend) + lat = manager.coordinates[dim]["lat"] + lon = manager.coordinates[dim]["lon"] assert helpers.dallclose(lat.asnumpy(), grid_savepoint.lat(dim).asnumpy()) assert helpers.dallclose(lon.asnumpy(), grid_savepoint.lon(dim).asnumpy()) @@ -573,10 +580,10 @@ def test_coordinates(grid_savepoint, grid_file, experiment, dim): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_tangent_orientation(experiment, grid_file, grid_savepoint): +def test_tangent_orientation(experiment, grid_file, grid_savepoint, backend): expected = grid_savepoint.tangent_orientation() - gm = utils.run_grid_manager(grid_file) - geometry_fields = gm.geometry + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry assert helpers.dallclose( - geometry_fields[GeometryName.TANGENT_ORIENTATION].ndarray, expected.ndarray + geometry_fields[GeometryName.TANGENT_ORIENTATION].asnumpy(), expected.asnumpy() ) diff --git a/model/common/tests/grid_tests/test_icon.py b/model/common/tests/grid_tests/test_icon.py index 153bff7f5..9dd2dc24b 100644 --- a/model/common/tests/grid_tests/test_icon.py +++ b/model/common/tests/grid_tests/test_icon.py @@ -17,13 +17,14 @@ icon, vertical as v_grid, ) +from icon4py.model.common.test_utils import grid_utils as gridtest_utils from . import utils @functools.cache def grid_from_file() -> icon.IconGrid: - file_name = utils.resolve_file_from_gridfile_name("mch_ch_r04b09_dsl") + file_name = gridtest_utils.resolve_full_grid_file_name("mch_ch_r04b09_dsl") manager = gm.GridManager( gm.ToZeroBasedIndexTransformation(), str(file_name), v_grid.VerticalGridConfig(1) ) diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index 9b14fba83..3b1463dc3 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -7,20 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -from pathlib import Path - -import gt4py.next as gtx -import gt4py.next.backend as gtx_backend - from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import grid_manager as gm, horizontal as h_grid, vertical as v_grid -from icon4py.model.common.test_utils import datatest_utils as dt_utils -from icon4py.model.common.test_utils.data_handling import download_and_extract +from icon4py.model.common.grid import horizontal as h_grid from icon4py.model.common.test_utils.datatest_utils import ( GRIDS_PATH, - MC_CH_R04B09_DSL_GRID_URI, R02B04_GLOBAL, - R02B04_GLOBAL_GRID_URI, REGIONAL_EXPERIMENT, ) @@ -32,31 +23,6 @@ r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.tar.gz").name -def resolve_file_from_gridfile_name(name: str) -> Path: - if name == REGIONAL_EXPERIMENT: - gridfile = r04b09_dsl_grid_path.joinpath("grid.nc") - if not gridfile.exists(): - download_and_extract( - MC_CH_R04B09_DSL_GRID_URI, - r04b09_dsl_grid_path, - r04b09_dsl_grid_path, - r04b09_dsl_data_file, - ) - return gridfile - elif name == R02B04_GLOBAL: - gridfile = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc") - if not gridfile.exists(): - download_and_extract( - R02B04_GLOBAL_GRID_URI, - r02b04_global_grid_path, - r02b04_global_grid_path, - r02b04_global_data_file, - ) - return gridfile - else: - raise ValueError(f"invalid name: use one of {R02B04_GLOBAL, REGIONAL_EXPERIMENT}") - - def horizontal_dim(): for dim in (dims.VertexDim, dims.EdgeDim, dims.CellDim): yield dim @@ -97,23 +63,3 @@ def valid_boundary_zones_for_dim(dim: dims.Dimension): ] yield from _domain(dim, zones) - - -def run_grid_manager( - experiment_name: str, - backend: gtx_backend.Backend = gtx.gtfn_cpu, - num_levels=65, - transformation=None, -) -> gm.GridManager: - if transformation is None: - transformation = gm.ToZeroBasedIndexTransformation() - file_name = resolve_file_from_gridfile_name(experiment_name) - with gm.GridManager( - transformation, file_name, v_grid.VerticalGridConfig(num_levels) - ) as grid_manager: - grid_manager(backend, limited_area=is_regional(experiment_name)) - return grid_manager - - -def is_regional(grid_file: str): - return grid_file == dt_utils.REGIONAL_EXPERIMENT diff --git a/model/common/tests/io_tests/test_io.py b/model/common/tests/io_tests/test_io.py index c5ebbfe00..4016871a1 100644 --- a/model/common/tests/io_tests/test_io.py +++ b/model/common/tests/io_tests/test_io.py @@ -40,7 +40,9 @@ grid_file = datatest_utils.GRIDS_PATH.joinpath( datatest_utils.R02B04_GLOBAL, grid_utils.GLOBAL_GRIDFILE ) -global_grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid +global_grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend +).grid def model_state(grid: base.BaseGrid) -> dict[str, xr.DataArray]: @@ -177,7 +179,9 @@ def test_io_monitor_write_ugrid_file(test_path): ) def test_io_monitor_write_and_read_ugrid_dataset(test_path, variables): path_name = test_path.absolute().as_posix() + "/output" - grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid + grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend + ).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config, @@ -227,7 +231,9 @@ def test_io_monitor_write_and_read_ugrid_dataset(test_path, variables): def test_fieldgroup_monitor_write_dataset_file_roll(test_path): - grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid + grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend + ).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config,