From 92ce7804938c1864c3dc9f75603705c86f5ecc9f Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 27 Nov 2024 17:33:46 +0100 Subject: [PATCH 1/8] add backend to grid_manager - allocate geometry and coordinates on specified backend - determine on_gpu for icon_grid from backend --- .../icon4py/model/common/grid/grid_manager.py | 120 +++++++++++------- .../model/common/test_utils/grid_utils.py | 17 ++- .../model/common/test_utils/pytest_config.py | 22 ++-- .../common/tests/grid_tests/test_geometry.py | 4 +- .../tests/grid_tests/test_grid_manager.py | 8 +- model/common/tests/grid_tests/test_icon.py | 2 +- model/common/tests/grid_tests/utils.py | 14 +- 7 files changed, 112 insertions(+), 75 deletions(-) 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 ace6af29b0..ce8f3fbce0 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -11,13 +11,21 @@ from typing import Literal, Optional, Protocol, TypeAlias, Union import gt4py.next as gtx +import gt4py.next.backend as gtx_backend +import numpy as np -from icon4py.model.common import dimension as dims, exceptions +import icon4py.model.common.utils as common_utils +from icon4py.model.common import dimension as dims, exceptions, type_alias as ta from icon4py.model.common.decomposition import ( definitions as decomposition, ) from icon4py.model.common.grid import base, icon, vertical as v_grid -from icon4py.model.common.settings import xp + + +try: + import cupy as xp +except ImportError: + import numpy as xp try: @@ -31,6 +39,8 @@ def __init__(self, *args, **kwargs): raise ModuleNotFoundError("NetCDF4 is not installed.") +NDArray: TypeAlias = Union[np.ndarray, xp.ndarray] + _log = logging.getLogger(__name__) @@ -251,8 +261,8 @@ def attribute(self, name: PropertyName) -> Union[str, int, float]: return self._dataset.getncattr(name) def int_variable( - self, name: FieldName, indices: xp.ndarray = None, transpose: bool = True - ) -> xp.ndarray: + self, name: FieldName, indices: np.ndarray = None, transpose: bool = True + ) -> np.ndarray: """Read a integer field from the grid file. Reads as gtx.int32. @@ -261,16 +271,16 @@ def int_variable( name: name of the field to read transpose: flag to indicate whether the file should be transposed (for 2d fields) Returns: - xp.ndarray: field data + np.ndarray: field data """ _log.debug(f"reading {name}: transposing = {transpose}") data = self.variable(name, indices, dtype=gtx.int32) - return xp.transpose(data) if transpose else data + return np.transpose(data) if transpose else data def variable( - self, name: FieldName, indices: xp.ndarray = None, dtype: xp.dtype = gtx.float64 - ) -> xp.ndarray: + self, name: FieldName, indices: np.ndarray = None, dtype: np.dtype = gtx.float64 + ) -> np.ndarray: """Read a field from the grid file. If a index array is given it only reads the values at those positions. @@ -283,7 +293,7 @@ def variable( variable = self._dataset.variables[name] _log.debug(f"reading {name}: {variable}") data = variable[:] if indices is None else variable[indices] - data = xp.array(data, dtype=dtype) + data = np.array(data, dtype=dtype) return data except KeyError as err: msg = f"{name} does not exist in dataset" @@ -308,20 +318,20 @@ class IndexTransformation(Protocol): def __call__( self, - array: xp.ndarray, - ) -> xp.ndarray: + array: NDArray, + ) -> NDArray: ... class NoTransformation(IndexTransformation): """Empty implementation of the Protocol. Just return zeros.""" - def __call__(self, array: xp.ndarray): + def __call__(self, array: NDArray): return xp.zeros_like(array) class ToZeroBasedIndexTransformation(IndexTransformation): - def __call__(self, array: xp.ndarray): + def __call__(self, array: NDArray): """ Calculate the index offset needed for usage with python. @@ -384,67 +394,80 @@ 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, on_gpu: bool = False, limited_area=True): + def __call__(self, backend: gtx_backend.Backend, limited_area=True): if not self._reader: self.open() + on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) self._grid = self._construct_grid(on_gpu=on_gpu, limited_area=limited_area) - self._refinement = self._read_grid_refinement_fields() - self._coordinates = self._read_coordinates() - self._geometry = self._read_geometry_fields() + self._refinement = self._read_grid_refinement_fields(backend) + self._coordinates = self._read_coordinates(backend) + self._geometry = self._read_geometry_fields(backend) - def _read_coordinates(self): + def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict: return { dims.CellDim: { "lat": gtx.as_field( (dims.CellDim,), self._reader.variable(CoordinateName.CELL_LATITUDE), - dtype=float, + dtype=ta.wpfloat, + allocator=backend, ), "lon": gtx.as_field( (dims.CellDim,), self._reader.variable(CoordinateName.CELL_LONGITUDE), - dtype=float, + dtype=ta.wpfloat, + allocator=backend, ), }, dims.EdgeDim: { "lat": gtx.as_field( - (dims.EdgeDim,), self._reader.variable(CoordinateName.EDGE_LATITUDE) + (dims.EdgeDim,), + self._reader.variable(CoordinateName.EDGE_LATITUDE), + dtype=ta.wpfloat, + allocator=backend, ), "lon": gtx.as_field( - (dims.EdgeDim,), self._reader.variable(CoordinateName.EDGE_LONGITUDE) + (dims.EdgeDim,), + self._reader.variable(CoordinateName.EDGE_LONGITUDE), + dtype=ta.wpfloat, + allocator=backend, ), }, dims.VertexDim: { "lat": gtx.as_field( (dims.VertexDim,), self._reader.variable(CoordinateName.VERTEX_LATITUDE), - dtype=float, + allocator=backend, + dtype=ta.wpfloat, ), "lon": gtx.as_field( (dims.VertexDim,), self._reader.variable(CoordinateName.VERTEX_LONGITUDE), - dtype=float, + allocator=backend, + dtype=ta.wpfloat, ), }, } - def _read_geometry_fields(self): + def _read_geometry_fields(self, backend: 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) . GeometryName.CELL_AREA.value: gtx.as_field( - (dims.CellDim,), self._reader.variable(GeometryName.CELL_AREA) + (dims.CellDim,), self._reader.variable(GeometryName.CELL_AREA), allocator=backend ), GeometryName.TANGENT_ORIENTATION.value: gtx.as_field( - (dims.EdgeDim,), self._reader.variable(GeometryName.TANGENT_ORIENTATION) + (dims.EdgeDim,), + self._reader.variable(GeometryName.TANGENT_ORIENTATION), + allocator=backend, ), } def _read_start_end_indices( self, ) -> tuple[ - dict[dims.Dimension : xp.ndarray], - dict[dims.Dimension : xp.ndarray], + dict[dims.Dimension : np.ndarray], + dict[dims.Dimension : np.ndarray], dict[dims.Dimension : gtx.int32], ]: """ " @@ -496,27 +519,30 @@ def _read_start_end_indices( return start_indices, end_indices, grid_refinement_dimensions def _read_grid_refinement_fields( - self, decomposition_info: Optional[decomposition.DecompositionInfo] = None - ) -> tuple[dict[dims.Dimension : xp.ndarray]]: + self, + backend: gtx_backend.Backend, + decomposition_info: Optional[decomposition.DecompositionInfo] = None, + ) -> tuple[dict[dims.Dimension : NDArray]]: """ Reads the refinement control fields from the grid file. Refinement control contains the classification of each entry in a field to predefined horizontal grid zones as for example the distance to the boundaries, see [refinement.py](refinement.py) """ + xp = common_utils.gt4py_field_allocation.import_array_ns(backend) refinement_control_names = { dims.CellDim: GridRefinementName.CONTROL_CELLS, dims.EdgeDim: GridRefinementName.CONTROL_EDGES, dims.VertexDim: GridRefinementName.CONTROL_VERTICES, } refinement_control_fields = { - dim: self._reader.int_variable(name, decomposition_info, transpose=False) + dim: xp.asarray(self._reader.int_variable(name, decomposition_info, transpose=False)) for dim, name in refinement_control_names.items() } return refinement_control_fields @property - def grid(self): + def grid(self) -> icon.IconGrid: return self._grid @property @@ -653,7 +679,7 @@ def _update_size_for_1d_sparse_dims(grid): ) -def _construct_diamond_vertices(e2v: xp.ndarray, c2v: xp.ndarray, e2c: xp.ndarray) -> xp.ndarray: +def _construct_diamond_vertices(e2v: NDArray, c2v: NDArray, e2c: NDArray) -> NDArray: r""" Construct the connectivity table for the vertices of a diamond in the ICON triangular grid. @@ -675,11 +701,11 @@ def _construct_diamond_vertices(e2v: xp.ndarray, c2v: xp.ndarray, e2c: xp.ndarra Ordering is the same as ICON uses. Args: - e2v: xp.ndarray containing the connectivity table for edge-to-vertex - c2v: xp.ndarray containing the connectivity table for cell-to-vertex - e2c: xp.ndarray containing the connectivity table for edge-to-cell + e2v: ndarray containing the connectivity table for edge-to-vertex + c2v: ndarray containing the connectivity table for cell-to-vertex + e2c: ndarray containing the connectivity table for edge-to-cell - Returns: xp.ndarray containing the connectivity table for edge-to-vertex on the diamond + Returns: ndarray containing the connectivity table for edge-to-vertex on the diamond """ dummy_c2v = _patch_with_dummy_lastline(c2v) expanded = dummy_c2v[e2c, :] @@ -692,7 +718,7 @@ def _construct_diamond_vertices(e2v: xp.ndarray, c2v: xp.ndarray, e2c: xp.ndarra return xp.hstack((e2v, far_indices)) -def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: +def _construct_diamond_edges(e2c: NDArray, c2e: NDArray) -> NDArray: r""" Construct the connectivity table for the edges of a diamond in the ICON triangular grid. @@ -712,10 +738,10 @@ def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: Args: - e2c: xp.ndarray containing the connectivity table for edge-to-cell - c2e: xp.ndarray containing the connectivity table for cell-to-edge + e2c: ndarray containing the connectivity table for edge-to-cell + c2e: ndarray containing the connectivity table for cell-to-edge - Returns: xp.ndarray containing the connectivity table for central edge-to- boundary edges + Returns: ndarray containing the connectivity table for central edge-to- boundary edges on the diamond """ dummy_c2e = _patch_with_dummy_lastline(c2e) @@ -731,7 +757,7 @@ def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: return e2c2e -def _construct_triangle_edges(c2e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: +def _construct_triangle_edges(c2e2c: NDArray, c2e: NDArray) -> NDArray: r"""Compute the connectivity from a central cell to all neighboring edges of its cell neighbors. ----e3---- ----e7---- @@ -752,7 +778,7 @@ def _construct_triangle_edges(c2e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: c2e2c: shape (n_cells, 3) connectivity table from a central cell to its cell neighbors c2e: shape (n_cells, 3), connectivity table from a cell to its neighboring edges Returns: - xp.ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring + ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring edges of its cell neighbors """ dummy_c2e = _patch_with_dummy_lastline(c2e) @@ -760,7 +786,7 @@ def _construct_triangle_edges(c2e2c: xp.ndarray, c2e: xp.ndarray) -> xp.ndarray: return table -def _construct_butterfly_cells(c2e2c: xp.ndarray) -> xp.ndarray: +def _construct_butterfly_cells(c2e2c: NDArray) -> NDArray: r"""Compute the connectivity from a central cell to all neighboring cells of its cell neighbors. / \ / \ @@ -784,7 +810,7 @@ def _construct_butterfly_cells(c2e2c: xp.ndarray) -> xp.ndarray: Args: c2e2c: shape (n_cells, 3) connectivity table from a central cell to its cell neighbors Returns: - xp.ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring cells of its cell neighbors + ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring cells of its cell neighbors """ dummy_c2e2c = _patch_with_dummy_lastline(c2e2c) c2e2c2e2c = xp.reshape(dummy_c2e2c[c2e2c, :], (c2e2c.shape[0], 9)) @@ -799,7 +825,7 @@ def _patch_with_dummy_lastline(ar): encountering a -1 = GridFile.INVALID_INDEX value Args: - ar: xp.ndarray connectivity array to be patched + ar: ndarray connectivity array to be patched Returns: same array with an additional line containing only GridFile.INVALID_INDEX 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 6815a674d1..645e4842f4 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 @@ -8,6 +8,7 @@ import functools +import gt4py.next.backend as gtx_backend import pytest import icon4py.model.common.grid.grid_manager as gm @@ -25,22 +26,24 @@ @functools.cache -def get_icon_grid_from_gridfile(experiment: str, on_gpu: bool = False) -> gm.GridManager: +def get_icon_grid_from_gridfile( + experiment: str, backend: gtx_backend.Backend = False +) -> gm.GridManager: if experiment == dt_utils.GLOBAL_EXPERIMENT: return _download_and_load_from_gridfile( dt_utils.R02B04_GLOBAL, GLOBAL_GRIDFILE, num_levels=GLOBAL_NUM_LEVELS, - on_gpu=on_gpu, limited_area=False, + backend=backend, ) elif experiment == dt_utils.REGIONAL_EXPERIMENT: return _download_and_load_from_gridfile( dt_utils.REGIONAL_EXPERIMENT, REGIONAL_GRIDFILE, num_levels=MCH_CH_R04B09_LEVELS, - on_gpu=on_gpu, limited_area=True, + backend=backend, ) else: raise ValueError(f"Unknown experiment: {experiment}") @@ -58,23 +61,23 @@ def download_grid_file(file_path: str, filename: str): def load_grid_from_file( - grid_file: str, num_levels: int, on_gpu: bool, limited_area: bool + grid_file: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool ) -> gm.GridManager: manager = gm.GridManager( gm.ToZeroBasedIndexTransformation(), str(grid_file), v_grid.VerticalGridConfig(num_levels=num_levels), ) - manager(on_gpu=on_gpu, limited_area=limited_area) + manager(backend=backend, limited_area=limited_area) return manager def _download_and_load_from_gridfile( - file_path: str, filename: str, num_levels: int, on_gpu: bool, limited_area: bool + file_path: str, filename: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool ) -> gm.GridManager: grid_file = download_grid_file(file_path, filename) - gm = load_grid_from_file(grid_file, num_levels, on_gpu, limited_area) + gm = load_grid_from_file(grid_file, num_levels, backend, limited_area) return gm 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 29ebf41da0..b89512aed8 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 @@ -18,6 +18,8 @@ ) +DEFAULT_BACKEND = "roundtrip" + backends = { "embedded": None, "roundtrip": itir_python, @@ -96,7 +98,7 @@ def pytest_addoption(parser): parser.addoption( "--backend", action="store", - default="roundtrip", + default=DEFAULT_BACKEND, help="GT4Py backend to use when executing stencils. Defaults to roundtrip backend, other options include gtfn_cpu, gtfn_gpu, and embedded", ) except ValueError: @@ -140,19 +142,15 @@ def pytest_runtest_setup(item): def pytest_generate_tests(metafunc): - on_gpu = False + selected_backend = backends[DEFAULT_BACKEND] # parametrise backend if "backend" in metafunc.fixturenames: backend_option = metafunc.config.getoption("backend") check_backend_validity(backend_option) - if backend_option in gpu_backends: - on_gpu = True - - metafunc.parametrize( - "backend", [backends[backend_option]], ids=[f"backend={backend_option}"] - ) + selected_backend = backends[backend_option] + metafunc.parametrize("backend", [selected_backend], ids=[f"backend={backend_option}"]) # parametrise grid if "grid" in metafunc.fixturenames: @@ -168,13 +166,17 @@ def pytest_generate_tests(metafunc): get_icon_grid_from_gridfile, ) - grid_instance = get_icon_grid_from_gridfile(REGIONAL_EXPERIMENT, on_gpu).grid + grid_instance = get_icon_grid_from_gridfile( + 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, ) - grid_instance = get_icon_grid_from_gridfile(GLOBAL_EXPERIMENT, on_gpu).grid + grid_instance = get_icon_grid_from_gridfile( + GLOBAL_EXPERIMENT, backend=selected_backend + ).grid else: raise ValueError(f"Unknown grid type: {selected_grid_type}") metafunc.parametrize("grid", [grid_instance], ids=[f"grid={selected_grid_type}"]) diff --git a/model/common/tests/grid_tests/test_geometry.py b/model/common/tests/grid_tests/test_geometry.py index ff651113cc..7d5260eb20 100644 --- a/model/common/tests/grid_tests/test_geometry.py +++ b/model/common/tests/grid_tests/test_geometry.py @@ -38,7 +38,7 @@ def construct_decomposition_info(grid: icon.IconGrid) -> definitions.Decompositi return decomposition_info def construct_grid_geometry(grid_file: str): - gm = utils.run_grid_manager(grid_file) + gm = utils.run_grid_manager(grid_file, backend=backend) grid = gm.grid decomposition_info = construct_decomposition_info(grid) geometry_source = geometry.GridGeometry( @@ -357,7 +357,7 @@ def test_sparse_fields_creator(): ], ) def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file) + gm = utils.run_grid_manager(grid_file, backend=backend) grid = gm.grid coordinates = gm.coordinates diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index 2a2e11433f..48dc1b5086 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -136,7 +136,7 @@ def test_grid_file_index_fields(global_grid_file, caplog, icon_grid): ) def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): caplog.set_level(logging.DEBUG) - manager = run_grid_manager(grid_file, zero_base) + manager = run_grid_manager(grid_file, transformation=zero_base) grid = manager.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 @@ -160,7 +160,7 @@ 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, zero_base) + manager = run_grid_manager(grid_file, transformation=zero_base) refin_ctrl = manager.refinement refin_ctrl_serialized = grid_savepoint.refin_ctrl(dim) assert np.all( @@ -450,7 +450,7 @@ def test_gridmanager_given_file_not_found_then_abort(): manager = gm.GridManager( gm.NoTransformation(), fname, v_grid.VerticalGridConfig(num_levels=80) ) - manager() + manager(None) assert error.value == 1 @@ -506,7 +506,7 @@ def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid): caplog.set_level(logging.INFO) serialized_grid = icon_grid - manager = run_grid_manager(grid_file, zero_base) + manager = run_grid_manager(grid_file, transformation=zero_base) grid = manager.grid for domain in utils.global_grid_domains(dim): diff --git a/model/common/tests/grid_tests/test_icon.py b/model/common/tests/grid_tests/test_icon.py index c020f4aff4..153bff7f5f 100644 --- a/model/common/tests/grid_tests/test_icon.py +++ b/model/common/tests/grid_tests/test_icon.py @@ -27,7 +27,7 @@ def grid_from_file() -> icon.IconGrid: manager = gm.GridManager( gm.ToZeroBasedIndexTransformation(), str(file_name), v_grid.VerticalGridConfig(1) ) - manager() + manager(backend=None) return manager.grid diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index 29b1e8f137..ad4d5b5954 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -7,9 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -import functools from pathlib import Path +import gt4py.next.backend as gtx_backend +from gt4py.next.program_processors.runners import gtfn + 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 @@ -97,15 +99,19 @@ def valid_boundary_zones_for_dim(dim: dims.Dimension): yield from _domain(dim, zones) -@functools.cache -def run_grid_manager(experiment_name: str, num_levels=65, transformation=None) -> gm.GridManager: +def run_grid_manager( + experiment_name: str, + backend: gtx_backend.Backend = gtfn.run_gtfn_cached, + 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(limited_area=is_regional(experiment_name)) + grid_manager(backend, limited_area=is_regional(experiment_name)) return grid_manager From 512e3cfc57ba182d40a142fa0ae3de803d8718ac Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 27 Nov 2024 17:37:39 +0100 Subject: [PATCH 2/8] fix backend import from gt4py.next --- model/common/tests/grid_tests/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index ad4d5b5954..cd471b9433 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -9,8 +9,8 @@ from pathlib import Path +import gt4py.next as gtx import gt4py.next.backend as gtx_backend -from gt4py.next.program_processors.runners import gtfn 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 @@ -101,7 +101,7 @@ def valid_boundary_zones_for_dim(dim: dims.Dimension): def run_grid_manager( experiment_name: str, - backend: gtx_backend.Backend = gtfn.run_gtfn_cached, + backend: gtx_backend.Backend = gtx.run_gtfn_cached, num_levels=65, transformation=None, ) -> gm.GridManager: From 4bc8ad186356a83f78d70d36b51daedbc69db29e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 27 Nov 2024 21:53:32 +0100 Subject: [PATCH 3/8] remove cache - might make it slower but should run --- .../diffusion/tests/diffusion_tests/test_diffusion.py | 4 +--- .../common/src/icon4py/model/common/test_utils/grid_utils.py | 2 -- model/common/src/icon4py/model/common/test_utils/helpers.py | 4 ---- 3 files changed, 1 insertion(+), 9 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index 8b73824ce8..b463a62e66 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py @@ -61,9 +61,7 @@ def _construct_minimal_decomposition_info(grid: icon.IconGrid): return decomposition_info if not grid_functionality[experiment].get(name): - on_gpu = helpers.is_gpu(backend) - gm = grid_utils.get_icon_grid_from_gridfile(experiment, on_gpu) - gm() + gm = grid_utils.get_icon_grid_from_gridfile(experiment, backend) grid = gm.grid decomposition_info = _construct_minimal_decomposition_info(grid) geometry_ = geometry.GridGeometry( 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 645e4842f4..f6cffb203d 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 @@ -6,7 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import functools import gt4py.next.backend as gtx_backend import pytest @@ -25,7 +24,6 @@ MCH_CH_R04B09_LEVELS = 65 -@functools.cache def get_icon_grid_from_gridfile( experiment: str, backend: gtx_backend.Backend = False ) -> gm.GridManager: 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 b46ce85eae..7a70372c3a 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -44,10 +44,6 @@ def is_embedded(backend) -> bool: return backend is None -def is_gpu(backend) -> bool: - return "gpu" in backend.name if backend else False - - def is_roundtrip(backend) -> bool: return backend.name == "roundtrip" if backend else False From c58812f448d1d1ce94d943f693e07d9a48a4853a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 28 Nov 2024 10:05:50 +0100 Subject: [PATCH 4/8] use numpy in grid construction --- .../icon4py/model/common/grid/grid_manager.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) 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 ce8f3fbce0..f378723b08 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -327,7 +327,7 @@ class NoTransformation(IndexTransformation): """Empty implementation of the Protocol. Just return zeros.""" def __call__(self, array: NDArray): - return xp.zeros_like(array) + return np.zeros_like(array) class ToZeroBasedIndexTransformation(IndexTransformation): @@ -338,7 +338,7 @@ def __call__(self, array: NDArray): Fortran indices are 1-based, hence the offset is -1 for 0-based ness of python except for INVALID values which are marked with -1 in the grid file and are kept such. """ - return xp.asarray(xp.where(array == GridFile.INVALID_INDEX, 0, -1), dtype=gtx.int32) + return np.asarray(np.where(array == GridFile.INVALID_INDEX, 0, -1), dtype=gtx.int32) CoordinateDict: TypeAlias = dict[dims.Dimension, dict[Literal["lat", "lon"], gtx.Field]] @@ -529,14 +529,14 @@ def _read_grid_refinement_fields( Refinement control contains the classification of each entry in a field to predefined horizontal grid zones as for example the distance to the boundaries, see [refinement.py](refinement.py) """ - xp = common_utils.gt4py_field_allocation.import_array_ns(backend) + np = common_utils.gt4py_field_allocation.import_array_ns(backend) refinement_control_names = { dims.CellDim: GridRefinementName.CONTROL_CELLS, dims.EdgeDim: GridRefinementName.CONTROL_EDGES, dims.VertexDim: GridRefinementName.CONTROL_VERTICES, } refinement_control_fields = { - dim: xp.asarray(self._reader.int_variable(name, decomposition_info, transpose=False)) + dim: np.asarray(self._reader.int_variable(name, decomposition_info, transpose=False)) for dim, name in refinement_control_names.items() } return refinement_control_fields @@ -642,14 +642,14 @@ def _add_derived_connectivities(grid: icon.IconGrid) -> icon.IconGrid: e2c2e = _construct_diamond_edges( grid.connectivities[dims.E2CDim], grid.connectivities[dims.C2EDim] ) - e2c2e0 = xp.column_stack((xp.asarray(range(e2c2e.shape[0])), e2c2e)) + e2c2e0 = np.column_stack((np.asarray(range(e2c2e.shape[0])), e2c2e)) c2e2c2e = _construct_triangle_edges( grid.connectivities[dims.C2E2CDim], grid.connectivities[dims.C2EDim] ) - c2e2c0 = xp.column_stack( + c2e2c0 = np.column_stack( ( - xp.asarray(range(grid.connectivities[dims.C2E2CDim].shape[0])), + np.asarray(range(grid.connectivities[dims.C2E2CDim].shape[0])), (grid.connectivities[dims.C2E2CDim]), ) ) @@ -711,11 +711,11 @@ def _construct_diamond_vertices(e2v: NDArray, c2v: NDArray, e2c: NDArray) -> NDA expanded = dummy_c2v[e2c, :] sh = expanded.shape flat = expanded.reshape(sh[0], sh[1] * sh[2]) - far_indices = xp.zeros_like(e2v) + far_indices = np.zeros_like(e2v) # TODO (magdalena) vectorize speed this up? for i in range(sh[0]): - far_indices[i, :] = flat[i, ~xp.isin(flat[i, :], e2v[i, :])][:2] - return xp.hstack((e2v, far_indices)) + far_indices[i, :] = flat[i, ~np.isin(flat[i, :], e2v[i, :])][:2] + return np.hstack((e2v, far_indices)) def _construct_diamond_edges(e2c: NDArray, c2e: NDArray) -> NDArray: @@ -750,9 +750,9 @@ def _construct_diamond_edges(e2c: NDArray, c2e: NDArray) -> NDArray: flattened = expanded.reshape(sh[0], sh[1] * sh[2]) diamond_sides = 4 - e2c2e = GridFile.INVALID_INDEX * xp.ones((sh[0], diamond_sides), dtype=gtx.int32) + e2c2e = GridFile.INVALID_INDEX * np.ones((sh[0], diamond_sides), dtype=gtx.int32) for i in range(sh[0]): - var = flattened[i, (~xp.isin(flattened[i, :], xp.asarray([i, GridFile.INVALID_INDEX])))] + var = flattened[i, (~np.isin(flattened[i, :], np.asarray([i, GridFile.INVALID_INDEX])))] e2c2e[i, : var.shape[0]] = var return e2c2e @@ -782,7 +782,7 @@ def _construct_triangle_edges(c2e2c: NDArray, c2e: NDArray) -> NDArray: edges of its cell neighbors """ dummy_c2e = _patch_with_dummy_lastline(c2e) - table = xp.reshape(dummy_c2e[c2e2c, :], (c2e2c.shape[0], 9)) + table = np.reshape(dummy_c2e[c2e2c, :], (c2e2c.shape[0], 9)) return table @@ -813,7 +813,7 @@ def _construct_butterfly_cells(c2e2c: NDArray) -> NDArray: ndarray: shape(n_cells, 9) connectivity table from a central cell to all neighboring cells of its cell neighbors """ dummy_c2e2c = _patch_with_dummy_lastline(c2e2c) - c2e2c2e2c = xp.reshape(dummy_c2e2c[c2e2c, :], (c2e2c.shape[0], 9)) + c2e2c2e2c = np.reshape(dummy_c2e2c[c2e2c, :], (c2e2c.shape[0], 9)) return c2e2c2e2c @@ -830,9 +830,9 @@ def _patch_with_dummy_lastline(ar): Returns: same array with an additional line containing only GridFile.INVALID_INDEX """ - patched_ar = xp.append( + patched_ar = np.append( ar, - GridFile.INVALID_INDEX * xp.ones((1, ar.shape[1]), dtype=gtx.int32), + GridFile.INVALID_INDEX * np.ones((1, ar.shape[1]), dtype=gtx.int32), axis=0, ) return patched_ar From f89099d53c71a3ea96226e464ed61a549322c190 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 28 Nov 2024 10:32:11 +0100 Subject: [PATCH 5/8] use numpy in grid construction in grid_manager fix io tests --- .../icon4py/model/common/test_utils/grid_utils.py | 2 +- model/common/tests/io_tests/test_io.py | 14 +++++--------- 2 files changed, 6 insertions(+), 10 deletions(-) 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 f6cffb203d..dbb1628cd7 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 @@ -25,7 +25,7 @@ def get_icon_grid_from_gridfile( - experiment: str, backend: gtx_backend.Backend = False + experiment: str, backend: gtx_backend.Backend = None ) -> gm.GridManager: if experiment == dt_utils.GLOBAL_EXPERIMENT: return _download_and_load_from_gridfile( diff --git a/model/common/tests/io_tests/test_io.py b/model/common/tests/io_tests/test_io.py index 4ce82bba87..85306c57fb 100644 --- a/model/common/tests/io_tests/test_io.py +++ b/model/common/tests/io_tests/test_io.py @@ -32,15 +32,15 @@ from icon4py.model.common.test_utils import datatest_utils, grid_utils, helpers +# fixing backend to fieldview embedded here. +backend = None UNLIMITED = None simple_grid = simple.SimpleGrid() 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, on_gpu=False -).grid +global_grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid def model_state(grid: base.BaseGrid) -> dict[str, xr.DataArray]: @@ -177,9 +177,7 @@ 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, on_gpu=False - ).grid + grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config, @@ -229,9 +227,7 @@ 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, on_gpu=False - ).grid + grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config, From 0bd93a0a26664a43ef212c46733c36361d578250 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 28 Nov 2024 11:18:13 +0100 Subject: [PATCH 6/8] fix wrong import name for default gtfn_cpu backend --- model/common/tests/grid_tests/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index cd471b9433..9b14fba832 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -101,7 +101,7 @@ def valid_boundary_zones_for_dim(dim: dims.Dimension): def run_grid_manager( experiment_name: str, - backend: gtx_backend.Backend = gtx.run_gtfn_cached, + backend: gtx_backend.Backend = gtx.gtfn_cpu, num_levels=65, transformation=None, ) -> gm.GridManager: From df259d96a0c5d0e8b390dbc7c8041a83c7662048 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Thu, 28 Nov 2024 13:50:28 +0100 Subject: [PATCH 7/8] Update model/common/tests/io_tests/test_io.py Co-authored-by: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> --- model/common/tests/io_tests/test_io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/common/tests/io_tests/test_io.py b/model/common/tests/io_tests/test_io.py index 85306c57fb..c5ebbfe006 100644 --- a/model/common/tests/io_tests/test_io.py +++ b/model/common/tests/io_tests/test_io.py @@ -32,7 +32,7 @@ from icon4py.model.common.test_utils import datatest_utils, grid_utils, helpers -# fixing backend to fieldview embedded here. +# setting backend to fieldview embedded here. backend = None UNLIMITED = None simple_grid = simple.SimpleGrid() From 78c60275bb0b993782f2e5c95d9075fbf1e66d28 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 28 Nov 2024 14:09:39 +0100 Subject: [PATCH 8/8] rename np to xp --- model/common/src/icon4py/model/common/grid/grid_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 f378723b08..95b3ef5b9d 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -529,14 +529,14 @@ def _read_grid_refinement_fields( Refinement control contains the classification of each entry in a field to predefined horizontal grid zones as for example the distance to the boundaries, see [refinement.py](refinement.py) """ - np = common_utils.gt4py_field_allocation.import_array_ns(backend) + xp = common_utils.gt4py_field_allocation.import_array_ns(backend) refinement_control_names = { dims.CellDim: GridRefinementName.CONTROL_CELLS, dims.EdgeDim: GridRefinementName.CONTROL_EDGES, dims.VertexDim: GridRefinementName.CONTROL_VERTICES, } refinement_control_fields = { - dim: np.asarray(self._reader.int_variable(name, decomposition_info, transpose=False)) + dim: xp.asarray(self._reader.int_variable(name, decomposition_info, transpose=False)) for dim, name in refinement_control_names.items() } return refinement_control_fields