diff --git a/jenkins/spack-PR b/jenkins/spack-PR index e69c02b14b..483c2c5e91 100644 --- a/jenkins/spack-PR +++ b/jenkins/spack-PR @@ -53,7 +53,7 @@ pipeline { steps { script{ sh """ - . ./spack-c2sm/setup-env.sh + . ./spack-c2sm/setup-env.sh /mch-environment/v7 cd local_copy spack env activate spack/gt4py-main spack install -v --test=root diff --git a/jenkins/spack-PR-stable b/jenkins/spack-PR-stable index 4398d1b36e..16b0c0929d 100644 --- a/jenkins/spack-PR-stable +++ b/jenkins/spack-PR-stable @@ -53,7 +53,7 @@ pipeline { steps { script{ sh """ - . ./spack-c2sm/setup-env.sh + . ./spack-c2sm/setup-env.sh /mch-environment/v7 cd local_copy spack env activate spack/gt4py-stable spack install -v --test=root diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index 8b73824ce8..68e4b6bba5 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_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 ace6af29b0..4f40a01316 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,27 +318,27 @@ 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): - return xp.zeros_like(array) + def __call__(self, array: NDArray): + return np.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. 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]] @@ -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: Optional[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: Optional[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: 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) . 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 @@ -616,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]), ) ) @@ -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,24 +701,24 @@ 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, :] 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: 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) @@ -724,14 +750,14 @@ def _construct_diamond_edges(e2c: xp.ndarray, c2e: xp.ndarray) -> xp.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 -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,15 +778,15 @@ 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) - table = xp.reshape(dummy_c2e[c2e2c, :], (c2e2c.shape[0], 9)) + table = np.reshape(dummy_c2e[c2e2c, :], (c2e2c.shape[0], 9)) 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,10 +810,10 @@ 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)) + c2e2c2e2c = np.reshape(dummy_c2e2c[c2e2c, :], (c2e2c.shape[0], 9)) return c2e2c2e2c @@ -799,14 +825,14 @@ 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 """ - 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 diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index d358f65f90..497c275b33 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -14,6 +14,7 @@ import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims from icon4py.model.common.dimension import C2E, V2E +from icon4py.model.common.grid import grid_manager as gm def compute_c_lin_e( @@ -369,34 +370,34 @@ def weighting_factors( x[i] = np.where(x[i] < -3.5, x[i] + np.pi * 2, x[i]) mask = np.logical_and(abs(x[1] - x[0]) > 1.0e-11, abs(y[2] - y[0]) > 1.0e-11) + wgt_1_no_mask = ( + 1.0 + / ((y[1] - y[0]) - (x[1] - x[0]) * (y[2] - y[0]) / (x[2] - x[0])) + * (1.0 - wgt_loc) + * (-y[0] + x[0] * (y[2] - y[0]) / (x[2] - x[0])) + ) wgt[2] = np.where( mask, 1.0 / ((y[2] - y[0]) - (x[2] - x[0]) * (y[1] - y[0]) / (x[1] - x[0])) * (1.0 - wgt_loc) * (-y[0] + x[0] * (y[1] - y[0]) / (x[1] - x[0])), - 1.0 - / ((y[1] - y[0]) - (x[1] - x[0]) * (y[2] - y[0]) / (x[2] - x[0])) - * (1.0 - wgt_loc) - * (-y[0] + x[0] * (y[2] - y[0]) / (x[2] - x[0])), + (-(1.0 - wgt_loc) * x[0] - wgt_1_no_mask * (x[1] - x[0])) / (x[2] - x[0]), ) wgt[1] = np.where( mask, (-(1.0 - wgt_loc) * x[0] - wgt[2] * (x[2] - x[0])) / (x[1] - x[0]), - (-(1.0 - wgt_loc) * x[0] - wgt[1] * (x[1] - x[0])) / (x[2] - x[0]), + wgt_1_no_mask, ) - wgt[1], wgt[2] = np.where(mask, (wgt[1], wgt[2]), (wgt[2], wgt[1])) - wgt[0] = 1.0 - wgt_loc - wgt[1] - wgt[2] - + wgt[0] = 1.0 - wgt[1] - wgt[2] if wgt_loc == 0.0 else 1.0 - wgt_loc - wgt[1] - wgt[2] return wgt -def compute_c_bln_avg( - divavg_cntrwgt: ta.wpfloat, - owner_mask: np.ndarray, +def _compute_c_bln_avg( c2e2c: np.ndarray, lat: np.ndarray, lon: np.ndarray, + divavg_cntrwgt: ta.wpfloat, horizontal_start: np.int32, ) -> np.ndarray: """ @@ -413,120 +414,184 @@ def compute_c_bln_avg( Returns: c_bln_avg: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] """ - llb = horizontal_start num_cells = c2e2c.shape[0] - c_bln_avg = np.zeros([num_cells, 4]) - wgt_loc = divavg_cntrwgt - yloc = np.zeros(num_cells) - xloc = np.zeros(num_cells) - yloc[llb:] = lat[llb:] - xloc[llb:] = lon[llb:] - ytemp = np.zeros([3, num_cells]) - xtemp = np.zeros([3, num_cells]) - - for i in range(3): - ytemp[i, llb:] = lat[c2e2c[llb:, i]] - xtemp[i, llb:] = lon[c2e2c[llb:, i]] + ytemp = np.zeros([c2e2c.shape[1], num_cells - horizontal_start]) + xtemp = np.zeros([c2e2c.shape[1], num_cells - horizontal_start]) + + for i in range(ytemp.shape[0]): + ytemp[i] = lat[c2e2c[horizontal_start:, i]] + xtemp[i] = lon[c2e2c[horizontal_start:, i]] wgt = weighting_factors( - ytemp[:, llb:], - xtemp[:, llb:], - yloc[llb:], - xloc[llb:], - wgt_loc, + ytemp, + xtemp, + lat[horizontal_start:], + lon[horizontal_start:], + divavg_cntrwgt, ) - - c_bln_avg[llb:, 0] = np.where(owner_mask[llb:], wgt_loc, c_bln_avg[llb:, 0]) - for i in range(3): - c_bln_avg[llb:, i + 1] = np.where(owner_mask[llb:], wgt[i], c_bln_avg[llb:, i + 1]) - + c_bln_avg = np.zeros((c2e2c.shape[0], c2e2c.shape[1] + 1)) + c_bln_avg[horizontal_start:, 0] = divavg_cntrwgt + c_bln_avg[horizontal_start:, 1] = wgt[0] + c_bln_avg[horizontal_start:, 2] = wgt[1] + c_bln_avg[horizontal_start:, 3] = wgt[2] return c_bln_avg -def compute_force_mass_conservation_to_c_bln_avg( +def _force_mass_conservation_to_c_bln_avg( + c2e2c0: np.ndarray, c_bln_avg: np.ndarray, - divavg_cntrwgt: ta.wpfloat, - owner_mask: np.ndarray, - c2e2c: np.ndarray, cell_areas: np.ndarray, + cell_owner_mask: np.ndarray, + divavg_cntrwgt: ta.wpfloat, horizontal_start: np.int32, - horizontal_start_p3: np.int32, - niter: np.ndarray = 1000, + niter: int = 1000, ) -> np.ndarray: """ - Compute the weighting coefficients for cell averaging with variable interpolation factors. + Iteratively enforce mass conservation to the input field c_bln_avg. - The weighting factors are based on the requirement that sum(w(i)*x(i)) = 0 + Mass conservation is enforced by the following condition: + The three point divergence calculated on any given grid point is used with a total factor of 1. - and sum(w(i)*y(i)) = 0, which ensures that linear horizontal gradients are not aliased into a checkerboard pattern between upward- and downward directed cells. The third condition is sum(w(i)) = 1., and the weight of the local point is 0.5. + Practically, the sum of the bilinear cell weights applied to a cell from all neighbors times its area should be exactly one. + + The weights are adjusted iteratively by the condition up to a max of niter iterations - force_mass_conservation_to_bilinear_cellavg_wgt Args: - c_bln_avg: bilinear cellavg wgt, numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] - divavg_cntrwgt: - owner_mask: numpy array, representing a gtx.Field[gtx.Dims[CellDim], bool] - c2e2c: numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, C2E2CDim], gtx.int32] - cell_areas: numpy array, representing a gtx.Field[gtx.Dims[CellDim], ta.wpfloat] + c2e2c0: cell to cell connectivity + c_bln_avg: input field: bilinear cell weight average + cell_areas: area of cells + cell_owner_mask: + divavg_cntrwgt: configured central weight horizontal_start: - horizontal_start_p3: - niter: number of iterations until convergence is assumed + niter: max number of iterations Returns: - c_bln_avg: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat] + """ - llb = horizontal_start - llb2 = horizontal_start_p3 - num_cells = c2e2c.shape[0] - index = np.arange(llb, num_cells) - inv_neighbor_id = -np.ones([num_cells, 3], dtype=int) - for i in range(3): - for j in range(3): - inv_neighbor_id[llb:, j] = np.where( - np.logical_and(c2e2c[c2e2c[llb:, j], i] == index, c2e2c[llb:, j] >= 0), - i, - inv_neighbor_id[llb:, j], - ) + def _compute_local_weights(c_bln_avg, cell_areas, c2e2c0, inverse_neighbor_idx) -> np.ndarray: + """ + Compute the total weight which each local point contributes to the sum. - relax_coeff = 0.46 - maxwgt_loc = divavg_cntrwgt + 0.003 - minwgt_loc = divavg_cntrwgt - 0.003 - # TODO: in this function halo cell exchanges (sync) are missing, here for inv_neighbor_id, but also within the iteration for several variables - for iteration in range(niter): - wgt_loc_sum = c_bln_avg[llb:, 0] * cell_areas[llb:] + np.sum( - c_bln_avg[c2e2c[llb:], inv_neighbor_id[llb:] + 1] * cell_areas[c2e2c[llb:]], axis=1 + Args: + c_bln_avg: ndarray representing a weight field of (CellDim, C2E2C0Dim) + inverse_neighbor_index: Sequence of to access all weights of a local cell in a field of shape (CellDim, C2E2C0Dim) + + Returns: ndarray of CellDim, containing the sum of weigh contributions for each local cell index + + """ + weights = np.sum(c_bln_avg[c2e2c0, inverse_neighbor_idx] * cell_areas[c2e2c0], axis=1) + return weights + + def _compute_residual_to_mass_conservation( + owner_mask: np.ndarray, local_weight: np.ndarray, cell_area: np.ndarray + ) -> np.ndarray: + """The local_weight weighted by the area should be 1. We compute how far we are off that weight.""" + horizontal_size = local_weight.shape[0] + assert horizontal_size == owner_mask.shape[0], "Fields do not have the same shape" + assert horizontal_size == cell_area.shape[0], "Fields do not have the same shape" + residual = np.where(owner_mask, local_weight / cell_area - 1.0, 0.0) + return residual + + def _apply_correction( + c_bln_avg: np.ndarray, + residual: np.ndarray, + c2e2c0: np.ndarray, + divavg_cntrwgt: float, + horizontal_start: gtx.int32, + ) -> np.ndarray: + """Apply correction to local weigths based on the computed residuals.""" + maxwgt_loc = divavg_cntrwgt + 0.003 + minwgt_loc = divavg_cntrwgt - 0.003 + relax_coeff = 0.46 + c_bln_avg[horizontal_start:, :] = ( + c_bln_avg[horizontal_start:, :] - relax_coeff * residual[c2e2c0][horizontal_start:, :] ) - resid = wgt_loc_sum[llb2 - llb :] / cell_areas[llb2:] - 1.0 - if iteration < niter - 1: - c_bln_avg[llb2:, 0] = np.where( - owner_mask[llb2:], c_bln_avg[llb2:, 0] - relax_coeff * resid, c_bln_avg[llb2:, 0] - ) - for i in range(3): - c_bln_avg[llb2:, i + 1] = np.where( - owner_mask[llb2:], - c_bln_avg[llb2:, i + 1] - relax_coeff * resid[c2e2c[llb2:, i] - llb2], - c_bln_avg[llb2:, i + 1], - ) - wgt_loc_sum = np.sum(c_bln_avg[llb2:], axis=1) - 1.0 - for i in range(4): - c_bln_avg[llb2:, i] = c_bln_avg[llb2:, i] - 0.25 * wgt_loc_sum - c_bln_avg[llb2:, 0] = np.where( - owner_mask[llb2:], - np.where(c_bln_avg[llb2:, 0] > minwgt_loc, c_bln_avg[llb2:, 0], minwgt_loc), - c_bln_avg[llb2:, 0], - ) - c_bln_avg[llb2:, 0] = np.where( - owner_mask[llb2:], - np.where(c_bln_avg[llb2:, 0] < maxwgt_loc, c_bln_avg[llb2:, 0], maxwgt_loc), - c_bln_avg[llb2:, 0], - ) - else: - c_bln_avg[llb2:, 0] = np.where( - owner_mask[llb2:], c_bln_avg[llb2:, 0] - resid, c_bln_avg[llb2:, 0] + local_weight = np.sum(c_bln_avg, axis=1) - 1.0 + + c_bln_avg[horizontal_start:, :] = c_bln_avg[horizontal_start:, :] - ( + 0.25 * local_weight[horizontal_start:, np.newaxis] + ) + + # avoid runaway condition: + c_bln_avg[horizontal_start:, 0] = np.maximum(c_bln_avg[horizontal_start:, 0], minwgt_loc) + c_bln_avg[horizontal_start:, 0] = np.minimum(c_bln_avg[horizontal_start:, 0], maxwgt_loc) + return c_bln_avg + + def _enforce_mass_conservation( + c_bln_avg: np.ndarray, + residual: np.ndarray, + owner_mask: np.ndarray, + horizontal_start: gtx.int32, + ) -> np.ndarray: + """Enforce the mass conservation condition on the local cells by forcefully subtracting the + residual from the central field contribution.""" + c_bln_avg[horizontal_start:, 0] = np.where( + owner_mask[horizontal_start:], + c_bln_avg[horizontal_start:, 0] - residual[horizontal_start:], + c_bln_avg[horizontal_start:, 0], + ) + return c_bln_avg + + local_summed_weights = np.zeros(c_bln_avg.shape[0]) + residual = np.zeros(c_bln_avg.shape[0]) + inverse_neighbor_idx = create_inverse_neighbor_index(c2e2c0) + + for iteration in range(niter): + local_summed_weights[horizontal_start:] = _compute_local_weights( + c_bln_avg, cell_areas, c2e2c0, inverse_neighbor_idx + )[horizontal_start:] + + residual[horizontal_start:] = _compute_residual_to_mass_conservation( + cell_owner_mask, local_summed_weights, cell_areas + )[horizontal_start:] + + max_ = np.max(residual) + if iteration >= (niter - 1) or max_ < 1e-9: + print(f"number of iterations: {iteration} - max residual={max_}") + c_bln_avg = _enforce_mass_conservation( + c_bln_avg, residual, cell_owner_mask, horizontal_start ) + return c_bln_avg + + c_bln_avg = _apply_correction( + c_bln_avg=c_bln_avg, + residual=residual, + c2e2c0=c2e2c0, + divavg_cntrwgt=divavg_cntrwgt, + horizontal_start=horizontal_start, + ) + return c_bln_avg +def compute_mass_conserving_bilinear_cell_average_weight( + c2e2c0: np.ndarray, + lat: np.ndarray, + lon: np.ndarray, + cell_areas: np.ndarray, + cell_owner_mask: np.ndarray, + divavg_cntrwgt: ta.wpfloat, + horizontal_start: np.int32, + horizontal_start_level_3, +) -> np.ndarray: + c_bln_avg = _compute_c_bln_avg(c2e2c0[:, 1:], lat, lon, divavg_cntrwgt, horizontal_start) + return _force_mass_conservation_to_c_bln_avg( + c2e2c0, c_bln_avg, cell_areas, cell_owner_mask, divavg_cntrwgt, horizontal_start_level_3 + ) + + +def create_inverse_neighbor_index(c2e2c0): + inv_neighbor_idx = -1 * np.ones(c2e2c0.shape, dtype=int) + + for jc in range(c2e2c0.shape[0]): + for i in range(c2e2c0.shape[1]): + if c2e2c0[jc, i] >= 0: + inv_neighbor_idx[jc, i] = np.argwhere(c2e2c0[c2e2c0[jc, i], :] == jc)[0, 0] + + return inv_neighbor_idx + + def compute_e_flx_avg( c_bln_avg: np.ndarray, geofac_div: np.ndarray, @@ -673,12 +738,11 @@ def compute_cells_aw_verts( dual_area: np.ndarray, edge_vert_length: np.ndarray, edge_cell_length: np.ndarray, - owner_mask: np.ndarray, v2e: np.ndarray, e2v: np.ndarray, v2c: np.ndarray, e2c: np.ndarray, - horizontal_start: np.int32, + horizontal_start_vertex: ta.wpfloat, ) -> np.ndarray: """ Compute cells_aw_verts. @@ -697,37 +761,47 @@ def compute_cells_aw_verts( Returns: aw_verts: numpy array, representing a gtx.Field[gtx.Dims[VertexDim, 6], ta.wpfloat] """ - llb = horizontal_start - cells_aw_verts = np.zeros([v2e.shape[0], 6]) - for i in range(e2c.shape[1]): + cells_aw_verts = np.zeros(v2e.shape) + for jv in range(horizontal_start_vertex, cells_aw_verts.shape[0]): + cells_aw_verts[jv, :] = 0.0 for je in range(v2e.shape[1]): - for jc in range(v2c.shape[1]): - mask = np.where( - np.logical_and(v2e[llb:, je] >= 0, e2c[v2e[llb:, je], i] == v2c[llb:, jc]), - owner_mask[llb:], - False, - ) - index = np.arange(llb, v2e.shape[0]) - idx_ve = np.where(e2v[v2e[llb:, je], 0] == index, 0, 1) - cells_aw_verts[llb:, jc] = np.where( - mask, - cells_aw_verts[llb:, jc] - + 0.5 - / dual_area[llb:] - * edge_vert_length[v2e[llb:, je], idx_ve] - * edge_cell_length[v2e[llb:, je], i], - cells_aw_verts[llb:, jc], - ) + # INVALID_INDEX + if je > gm.GridFile.INVALID_INDEX and (je > 0 and v2e[jv, je] == v2e[jv, je - 1]): + continue + ile = v2e[jv, je] + idx_ve = 0 if e2v[ile, 0] == jv else 1 + cell_offset_idx_0 = e2c[ile, 0] + cell_offset_idx_1 = e2c[ile, 1] + for jc in range(v2e.shape[1]): + if jc > gm.GridFile.INVALID_INDEX and (jc > 0 and v2c[jv, jc] == v2c[jv, jc - 1]): + continue + if cell_offset_idx_0 == v2c[jv, jc]: + cells_aw_verts[jv, jc] = ( + cells_aw_verts[jv, jc] + + 0.5 + / dual_area[jv] + * edge_vert_length[ile, idx_ve] + * edge_cell_length[ile, 0] + ) + elif cell_offset_idx_1 == v2c[jv, jc]: + cells_aw_verts[jv, jc] = ( + cells_aw_verts[jv, jc] + + 0.5 + / dual_area[jv] + * edge_vert_length[ile, idx_ve] + * edge_cell_length[ile, 1] + ) + return cells_aw_verts def compute_e_bln_c_s( - owner_mask: np.ndarray, c2e: np.ndarray, cells_lat: np.ndarray, cells_lon: np.ndarray, edges_lat: np.ndarray, edges_lon: np.ndarray, + weighting_factor: float, ) -> np.ndarray: """ Compute e_bln_c_s. @@ -760,11 +834,12 @@ def compute_e_bln_c_s( xtemp, yloc, xloc, - 0.0, + weighting_factor, ) - for i in range(wgt.shape[0]): - e_bln_c_s[llb:, i] = np.where(owner_mask[llb:], wgt[i], e_bln_c_s[llb:, i]) + e_bln_c_s[:, 0] = wgt[0] + e_bln_c_s[:, 1] = wgt[1] + e_bln_c_s[:, 2] = wgt[2] return e_bln_c_s 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..35f4aa1e3a 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 functools - +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,61 +32,139 @@ MCH_CH_R04B09_LEVELS = 65 +grid_geometries = {} + -@functools.cache -def get_icon_grid_from_gridfile(experiment: str, on_gpu: bool = False) -> gm.GridManager: +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, - on_gpu=on_gpu, - 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, - on_gpu=on_gpu, - 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, on_gpu: bool, 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(on_gpu=on_gpu, limited_area=limited_area) + 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, on_gpu: bool, 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, on_gpu, 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/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 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..7d17ecc4d9 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: @@ -165,16 +163,20 @@ 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(REGIONAL_EXPERIMENT, on_gpu).grid + 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(GLOBAL_EXPERIMENT, on_gpu).grid + grid_instance = get_grid_manager_for_experiment( + 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/src/icon4py/model/common/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index 1799fe8cbb..638b91a8f0 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 @@ -5,13 +5,78 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Optional +import logging as log +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 type_alias as ta -from icon4py.model.common.settings import xp +from icon4py.model.common import dimension, type_alias as ta + + +""" Enum values from Enum values taken from DLPack reference implementation at: + https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h + via GT4Py +""" +CUDA_DEVICE_TYPES = ( + gt_core_defs.DeviceType.CUDA, + gt_core_defs.DeviceType.CUDA_MANAGED, + gt_core_defs.DeviceType.ROCM, +) + + +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 + else: + return False + + +def array_ns(try_cupy: bool): + if try_cupy: + try: + import cupy as cp + + return cp + except ImportError: + log.warn("No cupy installed, falling back to numpy for array_ns") + import numpy as np + + return np + + +def import_array_ns(backend: backend.Backend): + """Import cupy or numpy depending on a chosen GT4Py backend DevicType.""" + return array_ns(is_cupy_device(backend)) + + +def as_field(field: gtx.Field, backend: backend.Backend) -> gtx.Field: + """Convenience function to transfer an existing Field to a given backend.""" + return gtx.as_field(field.domain, field.ndarray, allocator=backend) + + +def _size(grid, dim: gtx.Dimension, is_half_dim: bool) -> int: + if dim == dimension.KDim and is_half_dim: + return grid.size[dim] + 1 + return grid.size[dim] def allocate_zero_field( @@ -20,12 +85,9 @@ def allocate_zero_field( is_halfdim=False, dtype=ta.wpfloat, backend: Optional[backend.Backend] = None, -): - shapex = tuple(map(lambda x: grid.size[x], dims)) - if is_halfdim: - assert len(shapex) == 2 - shapex = (shapex[0], shapex[1] + 1) - return gtx.as_field(dims, xp.zeros(shapex, dtype=dtype), allocator=backend) +) -> gtx.Field: + dimensions = {d: range(_size(grid, d, is_halfdim)) for d in dims} + return gtx.zeros(dimensions, dtype=dtype, allocator=backend) def allocate_indices( @@ -34,6 +96,7 @@ def allocate_indices( is_halfdim=False, dtype=gtx.int32, backend: Optional[backend.Backend] = None, -): - shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim] +) -> gtx.Field: + xp = import_array_ns(backend) + shapex = _size(grid, dim, is_halfdim) return gtx.as_field((dim,), xp.arange(shapex, dtype=dtype), allocator=backend) diff --git a/model/common/tests/grid_tests/test_geometry.py b/model/common/tests/grid_tests/test_geometry.py index ff651113cc..31789099fb 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) - 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) + 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 2a2e11433f..19bef3ac04 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, 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, 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] @@ -450,7 +459,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 @@ -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, 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 c020f4aff4..9dd2dc24bc 100644 --- a/model/common/tests/grid_tests/test_icon.py +++ b/model/common/tests/grid_tests/test_icon.py @@ -17,17 +17,18 @@ 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) ) - 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..3b1463dc36 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -7,18 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -import functools -from pathlib import Path - 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, ) @@ -30,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 @@ -95,19 +63,3 @@ 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: - 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)) - return grid_manager - - -def is_regional(grid_file: str): - return grid_file == dt_utils.REGIONAL_EXPERIMENT diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields.py b/model/common/tests/interpolation_tests/test_interpolation_fields.py index 498dae15ad..24ef28264a 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields.py @@ -14,20 +14,20 @@ import icon4py.model.common.test_utils.helpers as test_helpers from icon4py.model.common import constants from icon4py.model.common.interpolation.interpolation_fields import ( - compute_c_bln_avg, compute_c_lin_e, compute_cells_aw_verts, compute_e_bln_c_s, compute_e_flx_avg, - compute_force_mass_conservation_to_c_bln_avg, compute_geofac_div, compute_geofac_grdiv, compute_geofac_grg, compute_geofac_n2s, compute_geofac_rot, + compute_mass_conserving_bilinear_cell_average_weight, compute_pos_on_tplane_e_x_y, compute_primal_normal_ec, ) +from icon4py.model.common.test_utils import datatest_utils as dt_utils from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa: F401 # import fixtures from test_utils package data_provider, download_ser_data, @@ -43,6 +43,7 @@ @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): # fixture inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() edge_cell_length = grid_savepoint.edge_cell_length() @@ -60,6 +61,7 @@ def test_compute_c_lin_e(grid_savepoint, interpolation_savepoint, icon_grid): # @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_geofac_div(grid_savepoint, interpolation_savepoint, icon_grid): mesh = icon_grid primal_edge_length = grid_savepoint.primal_edge_length() @@ -79,6 +81,7 @@ def test_compute_geofac_div(grid_savepoint, interpolation_savepoint, icon_grid): @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_geofac_rot(grid_savepoint, interpolation_savepoint, icon_grid): mesh = icon_grid dual_edge_length = grid_savepoint.dual_edge_length() @@ -102,6 +105,7 @@ def test_compute_geofac_rot(grid_savepoint, interpolation_savepoint, icon_grid): @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_geofac_n2s(grid_savepoint, interpolation_savepoint, icon_grid): dual_edge_length = grid_savepoint.dual_edge_length() geofac_div = interpolation_savepoint.geofac_div() @@ -122,6 +126,7 @@ def test_compute_geofac_n2s(grid_savepoint, interpolation_savepoint, icon_grid): @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_geofac_grg(grid_savepoint, interpolation_savepoint, icon_grid): primal_normal_cell_x = grid_savepoint.primal_normal_cell_x().asnumpy() primal_normal_cell_y = grid_savepoint.primal_normal_cell_y().asnumpy() @@ -159,6 +164,7 @@ def test_compute_geofac_grg(grid_savepoint, interpolation_savepoint, icon_grid): @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_geofac_grdiv(grid_savepoint, interpolation_savepoint, icon_grid): geofac_div = interpolation_savepoint.geofac_div() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() @@ -181,38 +187,39 @@ def test_compute_geofac_grdiv(grid_savepoint, interpolation_savepoint, icon_grid @pytest.mark.datatest -def test_compute_c_bln_avg(grid_savepoint, interpolation_savepoint, icon_grid): +@pytest.mark.parametrize( + "experiment, atol", + [(dt_utils.REGIONAL_EXPERIMENT, 1e-10), (dt_utils.GLOBAL_EXPERIMENT, 1e-10)], +) +def test_compute_c_bln_avg(grid_savepoint, interpolation_savepoint, icon_grid, atol): cell_areas = grid_savepoint.cell_areas().asnumpy() + # both experiment use the default value divavg_cntrwgt = 0.5 c_bln_avg_ref = interpolation_savepoint.c_bln_avg().asnumpy() - owner_mask = grid_savepoint.c_owner_mask().asnumpy() - c2e2c = icon_grid.connectivities[dims.C2E2CDim] horizontal_start = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) horizontal_start_p2 = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3)) lat = grid_savepoint.cell_center_lat().asnumpy() lon = grid_savepoint.cell_center_lon().asnumpy() - c_bln_avg = compute_c_bln_avg( - divavg_cntrwgt, - owner_mask, - c2e2c, + cell_owner_mask = grid_savepoint.c_owner_mask() + + c2e2c0 = icon_grid.connectivities[dims.C2E2CODim] + + c_bln_avg = compute_mass_conserving_bilinear_cell_average_weight( + c2e2c0, lat, lon, - horizontal_start, - ) - c_bln_avg = compute_force_mass_conservation_to_c_bln_avg( - c_bln_avg, - divavg_cntrwgt, - owner_mask, - c2e2c, cell_areas, + cell_owner_mask, + divavg_cntrwgt, horizontal_start, horizontal_start_p2, ) - assert test_helpers.dallclose(c_bln_avg, c_bln_avg_ref, atol=1e-4, rtol=1e-5) + assert test_helpers.dallclose(c_bln_avg, c_bln_avg_ref, atol=atol) @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_e_flx_avg(grid_savepoint, interpolation_savepoint, icon_grid): e_flx_avg_ref = interpolation_savepoint.e_flx_avg().asnumpy() c_bln_avg = interpolation_savepoint.c_bln_avg().asnumpy() @@ -246,6 +253,7 @@ def test_compute_e_flx_avg(grid_savepoint, interpolation_savepoint, icon_grid): @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_cells_aw_verts( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint ): @@ -253,7 +261,6 @@ def test_compute_cells_aw_verts( dual_area = grid_savepoint.vertex_dual_area().asnumpy() edge_vert_length = grid_savepoint.edge_vert_length().asnumpy() edge_cell_length = grid_savepoint.edge_cell_length().asnumpy() - owner_mask = grid_savepoint.v_owner_mask() e2c = icon_grid.connectivities[dims.E2CDim] v2c = icon_grid.connectivities[dims.V2CDim] v2e = icon_grid.connectivities[dims.V2EDim] @@ -263,40 +270,33 @@ def test_compute_cells_aw_verts( ) cells_aw_verts = compute_cells_aw_verts( - dual_area, - edge_vert_length, - edge_cell_length, - owner_mask, - v2e, - e2v, - v2c, - e2c, - horizontal_start_vertex, + dual_area=dual_area, + edge_vert_length=edge_vert_length, + edge_cell_length=edge_cell_length, + v2e=v2e, + e2v=e2v, + v2c=v2c, + e2c=e2c, + horizontal_start_vertex=horizontal_start_vertex, ) - assert test_helpers.dallclose(cells_aw_verts, cells_aw_verts_ref) + assert test_helpers.dallclose(cells_aw_verts, cells_aw_verts_ref, atol=1e-3) @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_e_bln_c_s(grid_savepoint, interpolation_savepoint, icon_grid): e_bln_c_s_ref = interpolation_savepoint.e_bln_c_s().asnumpy() - owner_mask = grid_savepoint.c_owner_mask().asnumpy() c2e = icon_grid.connectivities[dims.C2EDim] cells_lat = grid_savepoint.cell_center_lat().asnumpy() cells_lon = grid_savepoint.cell_center_lon().asnumpy() edges_lat = grid_savepoint.edges_center_lat().asnumpy() edges_lon = grid_savepoint.edges_center_lon().asnumpy() - e_bln_c_s = compute_e_bln_c_s( - owner_mask, - c2e, - cells_lat, - cells_lon, - edges_lat, - edges_lon, - ) + e_bln_c_s = compute_e_bln_c_s(c2e, cells_lat, cells_lon, edges_lat, edges_lon, 0.0) assert test_helpers.dallclose(e_bln_c_s, e_bln_c_s_ref, atol=1e-6, rtol=1e-7) @pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_pos_on_tplane_e(grid_savepoint, interpolation_savepoint, icon_grid): pos_on_tplane_e_x_ref = interpolation_savepoint.pos_on_tplane_e_x().asnumpy() pos_on_tplane_e_y_ref = interpolation_savepoint.pos_on_tplane_e_y().asnumpy() diff --git a/model/common/tests/io_tests/test_io.py b/model/common/tests/io_tests/test_io.py index 4ce82bba87..4016871a13 100644 --- a/model/common/tests/io_tests/test_io.py +++ b/model/common/tests/io_tests/test_io.py @@ -32,14 +32,16 @@ from icon4py.model.common.test_utils import datatest_utils, grid_utils, helpers +# setting 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 +global_grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend ).grid @@ -177,8 +179,8 @@ 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_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( @@ -229,8 +231,8 @@ 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_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( diff --git a/spack/gt4py-main/spack.yaml b/spack/gt4py-main/spack.yaml index b9af8e396f..bdf31620b1 100644 --- a/spack/gt4py-main/spack.yaml +++ b/spack/gt4py-main/spack.yaml @@ -1,10 +1,10 @@ spack: specs: - - py-icon4py@main%gcc ^py-gt4py@main%gcc + - py-icon4py@main%gcc ^py-gt4py@main%gcc ^python@3.10.4 view: false concretizer: unify: true develop: py-icon4py: - spec: py-icon4py@main%gcc ^py-gt4py@main%gcc + spec: py-icon4py@main%gcc ^py-gt4py@main%gcc ^python@3.10.4 path: ../../ diff --git a/spack/gt4py-stable/spack.yaml b/spack/gt4py-stable/spack.yaml index 0e6ac50fdd..c5268a97c7 100644 --- a/spack/gt4py-stable/spack.yaml +++ b/spack/gt4py-stable/spack.yaml @@ -1,10 +1,10 @@ spack: specs: - - py-icon4py@main%gcc ^py-gt4py@1.0.3.9%gcc + - py-icon4py@main%gcc ^py-gt4py@1.0.3.10%gcc ^python@3.10.4 view: false concretizer: unify: true develop: py-icon4py: - spec: py-icon4py@main%gcc ^py-gt4py@1.0.3.9%gcc + spec: py-icon4py@main%gcc ^py-gt4py@1.0.3.10%gcc ^python@3.10.4 path: ../../