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 b463a62e66..68e4b6bba5 100644 --- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py @@ -61,7 +61,7 @@ def _construct_minimal_decomposition_info(grid: icon.IconGrid): return decomposition_info if not grid_functionality[experiment].get(name): - gm = grid_utils.get_icon_grid_from_gridfile(experiment, backend) + gm = grid_utils.get_grid_manager_for_experiment(experiment, backend) grid = gm.grid decomposition_info = _construct_minimal_decomposition_info(grid) geometry_ = geometry.GridGeometry( diff --git a/model/common/src/icon4py/model/common/grid/grid_manager.py b/model/common/src/icon4py/model/common/grid/grid_manager.py index a39269ce0d..442b12810c 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -395,7 +395,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is FileNotFoundError: raise FileNotFoundError(f"gridfile {self._file_name} not found, aborting") - def __call__(self, backend: gtx_backend.Backend, limited_area=True): + def __call__(self, backend: Optional[gtx_backend.Backend], limited_area=True): if not self._reader: self.open() on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) @@ -404,7 +404,7 @@ def __call__(self, backend: gtx_backend.Backend, limited_area=True): self._coordinates = self._read_coordinates(backend) self._geometry = self._read_geometry_fields(backend) - def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict: + def _read_coordinates(self, backend: Optional[gtx_backend.Backend]) -> CoordinateDict: return { dims.CellDim: { "lat": gtx.as_field( @@ -450,7 +450,7 @@ def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict: }, } - def _read_geometry_fields(self, backend: gtx_backend.Backend): + def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]): return { # TODO (@halungge) still needs to ported, values from "our" grid files contains (wrong) values: # based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) . diff --git a/model/common/src/icon4py/model/common/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 dbb1628cd7..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 gt4py.next as gtx import gt4py.next.backend as gtx_backend import pytest import icon4py.model.common.grid.grid_manager as gm -from icon4py.model.common.grid import vertical as v_grid +from icon4py.model.common import dimension as dims +from icon4py.model.common.decomposition import definitions +from icon4py.model.common.grid import ( + geometry, + geometry_attributes as geometry_attrs, + icon, + vertical as v_grid, +) from icon4py.model.common.test_utils import data_handling, datatest_utils as dt_utils +from icon4py.model.common.utils import gt4py_field_allocation as alloc REGIONAL_GRIDFILE = "grid.nc" @@ -23,62 +32,139 @@ MCH_CH_R04B09_LEVELS = 65 +grid_geometries = {} + -def get_icon_grid_from_gridfile( +def get_grid_manager_for_experiment( experiment: str, backend: gtx_backend.Backend = None ) -> gm.GridManager: if experiment == dt_utils.GLOBAL_EXPERIMENT: - return _download_and_load_from_gridfile( + return _download_and_load_gridfile( dt_utils.R02B04_GLOBAL, - GLOBAL_GRIDFILE, num_levels=GLOBAL_NUM_LEVELS, - limited_area=False, backend=backend, ) elif experiment == dt_utils.REGIONAL_EXPERIMENT: - return _download_and_load_from_gridfile( + return _download_and_load_gridfile( dt_utils.REGIONAL_EXPERIMENT, - REGIONAL_GRIDFILE, num_levels=MCH_CH_R04B09_LEVELS, - limited_area=True, backend=backend, ) else: raise ValueError(f"Unknown experiment: {experiment}") -def download_grid_file(file_path: str, filename: str): - grid_file = dt_utils.GRIDS_PATH.joinpath(file_path, filename) - if not grid_file.exists(): +def get_grid_manager( + grid_file: str, num_levels: int, backend: gtx_backend.Backend +) -> gm.GridManager: + return _download_and_load_gridfile(grid_file, num_levels=num_levels, backend=backend) + + +def _file_name(grid_file: str): + match grid_file: + case dt_utils.REGIONAL_EXPERIMENT: + return REGIONAL_GRIDFILE + case dt_utils.R02B04_GLOBAL: + return GLOBAL_GRIDFILE + case _: + raise NotImplementedError(f"Add grid path for experiment '{grid_file}'") + + +def resolve_full_grid_file_name(grid_file_str: str) -> pathlib.Path: + return dt_utils.GRIDS_PATH.joinpath(grid_file_str, _file_name(grid_file_str)) + + +def _download_grid_file(file_path: str) -> pathlib.Path: + full_name = resolve_full_grid_file_name(file_path) + if not full_name.exists(): data_handling.download_and_extract( dt_utils.GRID_URIS[file_path], - grid_file.parent, - grid_file.parent, + full_name.parent, + full_name.parent, ) - return grid_file + return full_name -def load_grid_from_file( - grid_file: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool +def _run_grid_manager_for_file( + file: str, num_levels: int, backend: gtx_backend.Backend ) -> gm.GridManager: + """ + Load a grid file. + Args: + file: full path to the file (file + path) + num_levels: number of vertical levels, needed for IconGrid construction but independent from grid file + backend: the gt4py Backend we are running on + + Returns: + + """ + limited_area = is_regional(str(file)) + transformation = gm.ToZeroBasedIndexTransformation() manager = gm.GridManager( - gm.ToZeroBasedIndexTransformation(), - str(grid_file), + transformation, + file, v_grid.VerticalGridConfig(num_levels=num_levels), ) manager(backend=backend, limited_area=limited_area) + manager.close() return manager -def _download_and_load_from_gridfile( - file_path: str, filename: str, num_levels: int, backend: gtx_backend.Backend, limited_area: bool +def _download_and_load_gridfile( + file_path: str, num_levels: int, backend: gtx_backend.Backend ) -> gm.GridManager: - grid_file = download_grid_file(file_path, filename) - - gm = load_grid_from_file(grid_file, num_levels, backend, limited_area) + grid_file = _download_grid_file(file_path) + gm = _run_grid_manager_for_file(str(grid_file), num_levels, backend) return gm +def is_regional(experiment_or_file: str): + return ( + dt_utils.REGIONAL_EXPERIMENT in experiment_or_file + or REGIONAL_GRIDFILE in experiment_or_file + ) + + +def get_num_levels(experiment: str): + return MCH_CH_R04B09_LEVELS if experiment == dt_utils.REGIONAL_EXPERIMENT else GLOBAL_NUM_LEVELS + + +def get_grid_geometry( + backend: gtx_backend.Backend, experiment: str, grid_file: str +) -> geometry.GridGeometry: + on_gpu = alloc.is_cupy_device(backend) + xp = alloc.array_ns(on_gpu) + num_levels = get_num_levels(experiment) + + def construct_decomposition_info(grid: icon.IconGrid) -> definitions.DecompositionInfo: + def _add_dimension(dim: gtx.Dimension): + indices = alloc.allocate_indices(dim, grid) + owner_mask = xp.ones((grid.size[dim],), dtype=bool) + decomposition_info.with_dimension(dim, indices.ndarray, owner_mask) + + decomposition_info = definitions.DecompositionInfo(klevels=grid.num_levels) + _add_dimension(dims.EdgeDim) + _add_dimension(dims.VertexDim) + _add_dimension(dims.CellDim) + + return decomposition_info + + def construct_grid_geometry(grid_file: str): + gm = _run_grid_manager_for_file(grid_file, backend=backend, num_levels=num_levels) + grid = gm.grid + decomposition_info = construct_decomposition_info(grid) + geometry_source = geometry.GridGeometry( + grid, decomposition_info, backend, gm.coordinates, gm.geometry, geometry_attrs.attrs + ) + return geometry_source + + if not grid_geometries.get(grid_file): + grid_geometries[grid_file] = construct_grid_geometry( + str(resolve_full_grid_file_name(grid_file)) + ) + return grid_geometries[grid_file] + + @pytest.fixture def grid(request): return request.param diff --git a/model/common/src/icon4py/model/common/test_utils/pytest_config.py b/model/common/src/icon4py/model/common/test_utils/pytest_config.py index b89512aed8..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 @@ -163,18 +163,18 @@ def pytest_generate_tests(metafunc): grid_instance = SimpleGrid() elif selected_grid_type == "icon_grid": from icon4py.model.common.test_utils.grid_utils import ( - get_icon_grid_from_gridfile, + get_grid_manager_for_experiment, ) - grid_instance = get_icon_grid_from_gridfile( + grid_instance = get_grid_manager_for_experiment( REGIONAL_EXPERIMENT, backend=selected_backend ).grid elif selected_grid_type == "icon_grid_global": from icon4py.model.common.test_utils.grid_utils import ( - get_icon_grid_from_gridfile, + get_grid_manager_for_experiment, ) - grid_instance = get_icon_grid_from_gridfile( + grid_instance = get_grid_manager_for_experiment( GLOBAL_EXPERIMENT, backend=selected_backend ).grid else: diff --git a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index 137ead06b1..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 @@ -6,10 +6,11 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import logging as log -from typing import Optional +from typing import Optional, TypeAlias, Union import gt4py._core.definitions as gt_core_defs import gt4py.next as gtx +import numpy as np from gt4py.next import backend from icon4py.model.common import dimension, type_alias as ta @@ -26,6 +27,22 @@ ) +try: + import cupy as xp +except ImportError: + import numpy as xp + + +NDArrayInterface: TypeAlias = Union[np.ndarray, xp.ndarray, gtx.Field] + + +def as_numpy(array: NDArrayInterface): + if isinstance(array, np.ndarray): + return array + else: + return array.asnumpy() + + def is_cupy_device(backend: backend.Backend) -> bool: if backend is not None: return backend.allocator.__gt_device_type__ in CUDA_DEVICE_TYPES diff --git a/model/common/tests/grid_tests/test_geometry.py b/model/common/tests/grid_tests/test_geometry.py index 4ba023cb45..31789099fb 100644 --- a/model/common/tests/grid_tests/test_geometry.py +++ b/model/common/tests/grid_tests/test_geometry.py @@ -7,59 +7,24 @@ # SPDX-License-Identifier: BSD-3-Clause import functools -import gt4py.next as gtx import numpy as np 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: - def _add_dimension(dim: gtx.Dimension): - indices = alloc.allocate_indices(dim, grid) - owner_mask = np.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 = utils.run_grid_manager(grid_file, backend=backend) - grid = gm.grid - decomposition_info = construct_decomposition_info(grid) - geometry_source = geometry.GridGeometry( - grid, decomposition_info, backend, gm.coordinates, gm.geometry, attrs.attrs - ) - return geometry_source - - if not grid_geometries.get(grid_file): - grid_geometries[grid_file] = construct_grid_geometry(grid_file) - return grid_geometries[grid_file] +from icon4py.model.common.test_utils import datatest_utils as dt_utils, grid_utils, helpers def test_geometry_raises_for_unknown_field(backend): - geometry = get_grid_geometry(backend, dt_utils.R02B04_GLOBAL) + geometry = grid_utils.get_grid_geometry( + backend, dt_utils.GLOBAL_EXPERIMENT, dt_utils.R02B04_GLOBAL + ) with pytest.raises(ValueError) as e: geometry.get("foo") assert "'foo'" in e.value @@ -74,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) @@ -89,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) @@ -105,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) @@ -120,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) @@ -136,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) @@ -152,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}") @@ -173,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) @@ -189,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) @@ -220,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) @@ -240,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() @@ -256,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() @@ -272,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) @@ -291,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) @@ -310,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) @@ -329,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) @@ -364,7 +331,7 @@ def test_sparse_fields_creator(): ], ) def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file, backend=backend) + gm = grid_utils.get_grid_manager(grid_file, backend=backend, num_levels=1) grid = gm.grid coordinates = gm.coordinates @@ -373,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 8b147e849d..6fd93eb1c1 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -13,8 +13,8 @@ import numpy as np import pytest +from gt4py.next import backend as gtx_backend -import icon4py.model.common.test_utils.datatest_utils as dt_utils from icon4py.model.common import dimension as dims from icon4py.model.common.grid import ( grid_manager as gm, @@ -22,9 +22,11 @@ vertical as v_grid, ) from icon4py.model.common.grid.grid_manager import GeometryName -from icon4py.model.common.test_utils import helpers - -from .utils import run_grid_manager +from icon4py.model.common.test_utils import ( + datatest_utils as dt_utils, + grid_utils as gridtest_utils, + helpers, +) if typing.TYPE_CHECKING: @@ -47,12 +49,21 @@ MCH_CH_RO4B09_GLOBAL_NUM_CELLS = 83886080 -zero_base = gm.ToZeroBasedIndexTransformation() +ZERO_BASE = gm.ToZeroBasedIndexTransformation() + +managers = {} + + +def _run_grid_manager(file: str, backend: gtx_backend.Backend) -> gm.GridManager: + if not managers.get(file): + manager = gridtest_utils.get_grid_manager(file, num_levels=1, backend=backend) + managers[file] = manager + return managers.get(file) @pytest.fixture def global_grid_file(): - return utils.resolve_file_from_gridfile_name(dt_utils.R02B04_GLOBAL) + return gridtest_utils.resolve_full_grid_file_name(dt_utils.R02B04_GLOBAL) @pytest.mark.with_netcdf @@ -79,7 +90,7 @@ def test_grid_file_dimension(global_grid_file): ], ) def test_grid_file_vertex_cell_edge_dimensions(grid_savepoint, grid_file): - file = utils.resolve_file_from_gridfile_name(grid_file) + file = gridtest_utils.resolve_full_grid_file_name(grid_file) parser = gm.GridFile(str(file)) try: parser.open() @@ -134,10 +145,9 @@ def test_grid_file_index_fields(global_grid_file, caplog, icon_grid): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_v2e(caplog, grid_savepoint, experiment, grid_file, backend): caplog.set_level(logging.DEBUG) - manager = run_grid_manager(grid_file, transformation=zero_base) - grid = manager.grid + grid = _run_grid_manager(grid_file, backend).grid seralized_v2e = grid_savepoint.v2e() # there are vertices at the boundary of a local domain or at a pentagon point that have less than # 6 neighbors hence there are "Missing values" in the grid file @@ -159,9 +169,8 @@ def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", [dims.CellDim, dims.EdgeDim, dims.VertexDim]) -def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): - manager = run_grid_manager(grid_file, transformation=zero_base) - refin_ctrl = manager.refinement +def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim, backend): + refin_ctrl = _run_grid_manager(grid_file, backend).refinement refin_ctrl_serialized = grid_savepoint.refin_ctrl(dim) assert np.all( refin_ctrl_serialized.ndarray @@ -179,9 +188,9 @@ def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_v2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_v2c(caplog, grid_savepoint, experiment, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_v2c = grid_savepoint.v2c() # there are vertices that have less than 6 neighboring cells: either pentagon points or # vertices at the boundary of the domain for a limited area mode @@ -231,9 +240,9 @@ def reset_invalid_index(index_array: np.ndarray): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_e2v = grid_savepoint.e2v() # all vertices in the system have to neighboring edges, there no edges that point nowhere @@ -266,7 +275,7 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): grid_file: name of grid file used """ - if utils.is_regional(grid_file): + if gridtest_utils.is_regional(grid_file): assert has_invalid_index(e2c_table) else: assert not has_invalid_index(e2c_table) @@ -282,9 +291,10 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + + grid = _run_grid_manager(grid_file, backend).grid serialized_e2c = grid_savepoint.e2c() e2c_table = grid.get_offset_provider("E2C").table assert_invalid_indices(serialized_e2c, grid_file) @@ -302,9 +312,9 @@ def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_c2e = grid_savepoint.c2e() # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the @@ -325,9 +335,9 @@ def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid assert np.allclose( grid.get_offset_provider("C2E2C").table, grid_savepoint.c2e2c(), @@ -343,9 +353,9 @@ def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) assert np.allclose( grid.get_offset_provider("C2E2CO").table, @@ -363,9 +373,9 @@ def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) serialized_e2c2e = serialized_grid.get_offset_provider("E2C2E").table serialized_e2c2eO = serialized_grid.get_offset_provider("E2C2EO").table @@ -395,10 +405,9 @@ def assert_unless_invalid(table, serialized_ref): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - gm = run_grid_manager(grid_file) - grid = gm.grid + grid = _run_grid_manager(grid_file, backend).grid # the "far" (adjacent to edge normal ) is not always there, because ICON only calculates those starting from # (lateral_boundary(dims.EdgeDim) + 1) to end(dims.EdgeDim) (see mo_intp_coeffs.f90) and only for owned cells serialized_ref = grid_savepoint.e2c2v() @@ -416,9 +425,9 @@ def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid c2v = grid.get_offset_provider("C2V").table assert np.allclose(c2v, grid_savepoint.c2v()) @@ -432,8 +441,8 @@ def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.with_netcdf -def test_grid_manager_grid_size(dim, size): - grid = run_grid_manager(utils.R02B04_GLOBAL).grid +def test_grid_manager_grid_size(dim, size, backend): + grid = _run_grid_manager(utils.R02B04_GLOBAL, backend=backend).grid assert size == grid.size[dim] @@ -472,8 +481,8 @@ def test_gt4py_transform_offset_by_1_where_valid(size): (dt_utils.REGIONAL_EXPERIMENT, MCH_CH_RO4B09_GLOBAL_NUM_CELLS), ], ) -def test_grid_manager_grid_level_and_root(grid_file, global_num_cells): - assert global_num_cells == run_grid_manager(grid_file, num_levels=1).grid.global_num_cells +def test_grid_manager_grid_level_and_root(grid_file, global_num_cells, backend): + assert global_num_cells == _run_grid_manager(grid_file, backend=backend).grid.global_num_cells @pytest.mark.datatest @@ -482,9 +491,9 @@ def test_grid_manager_grid_level_and_root(grid_file, global_num_cells): "grid_file, experiment", [(utils.R02B04_GLOBAL, dt_utils.JABW_EXPERIMENT)], ) -def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) assert np.allclose( grid.get_offset_provider("C2E2C2E").table, @@ -503,12 +512,10 @@ def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", utils.horizontal_dim()) -def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid): +def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid, backend): caplog.set_level(logging.INFO) serialized_grid = icon_grid - manager = run_grid_manager(grid_file, transformation=zero_base) - grid = manager.grid - + grid = _run_grid_manager(grid_file, backend).grid for domain in utils.global_grid_domains(dim): assert grid.start_index(domain) == serialized_grid.start_index( domain @@ -518,7 +525,7 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g ), f"end index wrong for domain {domain}" for domain in utils.valid_boundary_zones_for_dim(dim): - if not utils.is_regional(grid_file): + if not gridtest_utils.is_regional(grid_file): assert grid.start_index(domain) == 0 assert grid.end_index(domain) == 0 assert grid.start_index(domain) == serialized_grid.start_index( @@ -537,10 +544,10 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_read_geometry_fields(grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file) - cell_area = gm.geometry[GeometryName.CELL_AREA.value] - tangent_orientation = gm.geometry[GeometryName.TANGENT_ORIENTATION.value] +def test_read_geometry_fields(grid_savepoint, grid_file, backend): + manager = _run_grid_manager(grid_file, backend=backend) + cell_area = manager.geometry[GeometryName.CELL_AREA.value] + tangent_orientation = manager.geometry[GeometryName.TANGENT_ORIENTATION.value] assert helpers.dallclose(cell_area.asnumpy(), grid_savepoint.cell_areas().asnumpy()) assert helpers.dallclose( @@ -557,10 +564,10 @@ def test_read_geometry_fields(grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", (dims.CellDim, dims.EdgeDim, dims.VertexDim)) -def test_coordinates(grid_savepoint, grid_file, experiment, dim): - gm = utils.run_grid_manager(grid_file) - lat = gm.coordinates[dim]["lat"] - lon = gm.coordinates[dim]["lon"] +def test_coordinates(grid_savepoint, grid_file, experiment, dim, backend): + manager = _run_grid_manager(grid_file, backend=backend) + lat = manager.coordinates[dim]["lat"] + lon = manager.coordinates[dim]["lon"] assert helpers.dallclose(lat.asnumpy(), grid_savepoint.lat(dim).asnumpy()) assert helpers.dallclose(lon.asnumpy(), grid_savepoint.lon(dim).asnumpy()) @@ -573,12 +580,12 @@ 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(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() ) @@ -590,10 +597,10 @@ def test_tangent_orientation(experiment, grid_file, grid_savepoint): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_cell_normal_orientation(experiment, grid_file, grid_savepoint): +def test_cell_normal_orientation(grid_file, grid_savepoint, backend): expected = grid_savepoint.edge_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.CELL_NORMAL_ORIENTATION].ndarray, expected.ndarray ) @@ -607,10 +614,10 @@ def test_cell_normal_orientation(experiment, grid_file, grid_savepoint): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_edge_orientation_on_vertex(experiment, grid_file, grid_savepoint): +def test_edge_orientation_on_vertex(grid_file, grid_savepoint, backend): expected = grid_savepoint.vertex_edge_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.EDGE_ORIENTATION_ON_VERTEX].ndarray, expected.ndarray ) @@ -624,8 +631,8 @@ def test_edge_orientation_on_vertex(experiment, grid_file, grid_savepoint): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_dual_area(experiment, grid_file, grid_savepoint): +def test_dual_area(grid_file, grid_savepoint, backend): expected = grid_savepoint.vertex_dual_area() - 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.DUAL_AREA].ndarray, expected.ndarray) diff --git a/model/common/tests/grid_tests/test_icon.py b/model/common/tests/grid_tests/test_icon.py index 153bff7f5f..9dd2dc24bc 100644 --- a/model/common/tests/grid_tests/test_icon.py +++ b/model/common/tests/grid_tests/test_icon.py @@ -17,13 +17,14 @@ icon, vertical as v_grid, ) +from icon4py.model.common.test_utils import grid_utils as gridtest_utils from . import utils @functools.cache def grid_from_file() -> icon.IconGrid: - file_name = utils.resolve_file_from_gridfile_name("mch_ch_r04b09_dsl") + file_name = gridtest_utils.resolve_full_grid_file_name("mch_ch_r04b09_dsl") manager = gm.GridManager( gm.ToZeroBasedIndexTransformation(), str(file_name), v_grid.VerticalGridConfig(1) ) diff --git a/model/common/tests/grid_tests/utils.py b/model/common/tests/grid_tests/utils.py index 9b14fba832..3b1463dc36 100644 --- a/model/common/tests/grid_tests/utils.py +++ b/model/common/tests/grid_tests/utils.py @@ -7,20 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations -from pathlib import Path - -import gt4py.next as gtx -import gt4py.next.backend as gtx_backend - from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import grid_manager as gm, horizontal as h_grid, vertical as v_grid -from icon4py.model.common.test_utils import datatest_utils as dt_utils -from icon4py.model.common.test_utils.data_handling import download_and_extract +from icon4py.model.common.grid import horizontal as h_grid from icon4py.model.common.test_utils.datatest_utils import ( GRIDS_PATH, - MC_CH_R04B09_DSL_GRID_URI, R02B04_GLOBAL, - R02B04_GLOBAL_GRID_URI, REGIONAL_EXPERIMENT, ) @@ -32,31 +23,6 @@ r02b04_global_data_file = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.tar.gz").name -def resolve_file_from_gridfile_name(name: str) -> Path: - if name == REGIONAL_EXPERIMENT: - gridfile = r04b09_dsl_grid_path.joinpath("grid.nc") - if not gridfile.exists(): - download_and_extract( - MC_CH_R04B09_DSL_GRID_URI, - r04b09_dsl_grid_path, - r04b09_dsl_grid_path, - r04b09_dsl_data_file, - ) - return gridfile - elif name == R02B04_GLOBAL: - gridfile = r02b04_global_grid_path.joinpath("icon_grid_0013_R02B04_R.nc") - if not gridfile.exists(): - download_and_extract( - R02B04_GLOBAL_GRID_URI, - r02b04_global_grid_path, - r02b04_global_grid_path, - r02b04_global_data_file, - ) - return gridfile - else: - raise ValueError(f"invalid name: use one of {R02B04_GLOBAL, REGIONAL_EXPERIMENT}") - - def horizontal_dim(): for dim in (dims.VertexDim, dims.EdgeDim, dims.CellDim): yield dim @@ -97,23 +63,3 @@ def valid_boundary_zones_for_dim(dim: dims.Dimension): ] yield from _domain(dim, zones) - - -def run_grid_manager( - experiment_name: str, - backend: gtx_backend.Backend = gtx.gtfn_cpu, - num_levels=65, - transformation=None, -) -> gm.GridManager: - if transformation is None: - transformation = gm.ToZeroBasedIndexTransformation() - file_name = resolve_file_from_gridfile_name(experiment_name) - with gm.GridManager( - transformation, file_name, v_grid.VerticalGridConfig(num_levels) - ) as grid_manager: - grid_manager(backend, limited_area=is_regional(experiment_name)) - return grid_manager - - -def is_regional(grid_file: str): - return grid_file == dt_utils.REGIONAL_EXPERIMENT diff --git a/model/common/tests/interpolation_tests/test_call_field_operator.py b/model/common/tests/interpolation_tests/test_call_field_operator.py deleted file mode 100644 index 42b6884c77..0000000000 --- a/model/common/tests/interpolation_tests/test_call_field_operator.py +++ /dev/null @@ -1,41 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -import gt4py.next as gtx -from gt4py.next import neighbor_sum - -import icon4py.model.common.test_utils.helpers as test_utils -from icon4py.model.common import dimension as dims -from icon4py.model.common.dimension import C2E, C2EDim -from icon4py.model.common.grid import simple - - -@gtx.field_operator -def field_op( - in_field: gtx.Field[gtx.Dims[dims.EdgeDim], float], - coeff: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], float], -) -> gtx.Field[gtx.Dims[dims.CellDim], float]: - return neighbor_sum(in_field(C2E) * coeff, axis=C2EDim) - - -def test_call_field_operator(backend): - grid = simple.SimpleGrid() - hstart = 0 - hend = grid.num_cells - coefficient = test_utils.constant_field(grid, 0.8, dims.CellDim, dims.C2EDim, dtype=float) - in_field = test_utils.constant_field(grid, 1.0, dims.EdgeDim, dtype=float) - out_field = test_utils.zero_field(grid, dims.CellDim, dtype=float) - expected = test_utils.constant_field(grid, 2.4, dims.CellDim, dtype=float) - field_op.with_backend(backend)( - in_field=in_field, - coeff=coefficient, - out=out_field, - offset_provider={"C2E": grid.get_offset_provider("C2E")}, - domain={dims.CellDim: (hstart, hend)}, - ) - test_utils.dallclose(out_field.asnumpy(), expected.asnumpy()) diff --git a/model/common/tests/interpolation_tests/test_interpolation_factory.py b/model/common/tests/interpolation_tests/test_interpolation_factory.py index 42caf7ef6d..461e4cb3f7 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_factory.py +++ b/model/common/tests/interpolation_tests/test_interpolation_factory.py @@ -9,13 +9,11 @@ import pytest import icon4py.model.common.states.factory as factory -import icon4py.model.common.test_utils.datatest_utils as dt_utils from icon4py.model.common.interpolation import ( interpolation_attributes as attrs, interpolation_factory, ) - -from .. import utils +from icon4py.model.common.test_utils import datatest_utils as dt_utils, grid_utils as gridtest_utils C2E_SIZE = 3 @@ -30,7 +28,7 @@ ) @pytest.mark.datatest def test_factory_raises_error_on_unknown_field(grid_file, experiment, backend, decomposition_info): - geometry = utils.get_grid_geometry(backend, grid_file) + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) interpolation_source = interpolation_factory.InterpolationFieldsFactory( grid=geometry.grid, decomposition_info=decomposition_info, @@ -52,7 +50,7 @@ def test_factory_raises_error_on_unknown_field(grid_file, experiment, backend, d ) @pytest.mark.datatest def test_get_c_lin_e(grid_file, experiment, backend, decomposition_info): - geometry = utils.get_grid_geometry(backend, grid_file) + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) grid = geometry.grid factory = interpolation_factory.InterpolationFieldsFactory( grid=grid, @@ -74,7 +72,7 @@ def test_get_c_lin_e(grid_file, experiment, backend, decomposition_info): ) @pytest.mark.datatest def test_get_geofac_div(grid_file, experiment, backend, decomposition_info): - geometry = utils.get_grid_geometry(backend, grid_file) + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) grid = geometry.grid factory = interpolation_factory.InterpolationFieldsFactory( grid=grid, @@ -96,7 +94,7 @@ def test_get_geofac_div(grid_file, experiment, backend, decomposition_info): ) @pytest.mark.datatest def test_get_geofac_rot(grid_file, experiment, backend, decomposition_info): - geometry = utils.get_grid_geometry(backend, grid_file) + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) grid = geometry.grid factory = interpolation_factory.InterpolationFieldsFactory( grid=grid, diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields.py b/model/common/tests/interpolation_tests/test_interpolation_fields.py index 90c72b8b38..ba4120b40c 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields.py @@ -6,7 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx import numpy as np import pytest @@ -15,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, @@ -44,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() @@ -61,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, backend): if backend is not None: pytest.xfail("writes a sparse fields: only runs in field view embedded") @@ -77,11 +78,11 @@ def test_compute_geofac_div(grid_savepoint, interpolation_savepoint, icon_grid, out=(geofac_div), offset_provider={"C2E": mesh.get_offset_provider("C2E")}, ) - gtx.as_field(geofac_div.domain, geofac_div.ndarray, allocator=backend) assert test_helpers.dallclose(geofac_div.asnumpy(), geofac_div_ref.asnumpy()) @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, backend): if backend is not None: pytest.xfail("writes a sparse fields: only runs in field view embedded") @@ -108,9 +109,10 @@ 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.compute_geofac_div() + geofac_div = interpolation_savepoint.geofac_div() geofac_n2s_ref = interpolation_savepoint.geofac_n2s() c2e = icon_grid.connectivities[dims.C2EDim] e2c = icon_grid.connectivities[dims.E2CDim] @@ -128,10 +130,11 @@ 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() - geofac_div = interpolation_savepoint.compute_geofac_div() + geofac_div = interpolation_savepoint.geofac_div() c_lin_e = interpolation_savepoint.c_lin_e() geofac_grg_ref = interpolation_savepoint.geofac_grg() owner_mask = grid_savepoint.c_owner_mask() @@ -165,8 +168,9 @@ 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.compute_geofac_div() + geofac_div = interpolation_savepoint.geofac_div() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() geofac_grdiv_ref = interpolation_savepoint.geofac_grdiv() owner_mask = grid_savepoint.c_owner_mask() @@ -187,42 +191,43 @@ 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() - geofac_div = interpolation_savepoint.compute_geofac_div().asnumpy() + geofac_div = interpolation_savepoint.geofac_div().asnumpy() owner_mask = grid_savepoint.e_owner_mask().asnumpy() primal_cart_normal_x = grid_savepoint.primal_cart_normal_x().asnumpy() primal_cart_normal_y = grid_savepoint.primal_cart_normal_y().asnumpy() @@ -252,6 +257,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 ): @@ -259,7 +265,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] @@ -269,40 +274,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 c5ebbfe006..4016871a13 100644 --- a/model/common/tests/io_tests/test_io.py +++ b/model/common/tests/io_tests/test_io.py @@ -40,7 +40,9 @@ grid_file = datatest_utils.GRIDS_PATH.joinpath( datatest_utils.R02B04_GLOBAL, grid_utils.GLOBAL_GRIDFILE ) -global_grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid +global_grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend +).grid def model_state(grid: base.BaseGrid) -> dict[str, xr.DataArray]: @@ -177,7 +179,9 @@ def test_io_monitor_write_ugrid_file(test_path): ) def test_io_monitor_write_and_read_ugrid_dataset(test_path, variables): path_name = test_path.absolute().as_posix() + "/output" - grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid + grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend + ).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config, @@ -227,7 +231,9 @@ def test_io_monitor_write_and_read_ugrid_dataset(test_path, variables): def test_fieldgroup_monitor_write_dataset_file_roll(test_path): - grid = grid_utils.get_icon_grid_from_gridfile(datatest_utils.GLOBAL_EXPERIMENT, backend).grid + grid = grid_utils.get_grid_manager_for_experiment( + datatest_utils.GLOBAL_EXPERIMENT, backend + ).grid vertical_config = v_grid.VerticalGridConfig(num_levels=grid.num_levels) vertical_params = v_grid.VerticalGrid( config=vertical_config, diff --git a/model/common/tests/states_test/test_factory.py b/model/common/tests/states_test/test_factory.py index eb835efbd7..621d9daef5 100644 --- a/model/common/tests/states_test/test_factory.py +++ b/model/common/tests/states_test/test_factory.py @@ -11,7 +11,7 @@ import gt4py.next as gtx import pytest -from icon4py.model.common import dimension as dims +from icon4py.model.common import dimension as dims, utils as common_utils from icon4py.model.common.grid import horizontal as h_grid, icon, vertical as v_grid from icon4py.model.common.math import helpers as math_helpers from icon4py.model.common.metrics import metric_fields as metrics @@ -62,7 +62,7 @@ def backend(self): @pytest.fixture def cell_coordinate_source(grid_savepoint, backend): - on_gpu = test_helpers.is_gpu(backend) + on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) grid = grid_savepoint.construct_icon_grid(on_gpu) lat = grid_savepoint.lat(dims.CellDim) lon = grid_savepoint.lon(dims.CellDim) @@ -77,7 +77,7 @@ def cell_coordinate_source(grid_savepoint, backend): @pytest.fixture def height_coordinate_source(metrics_savepoint, grid_savepoint, backend): - on_gpu = test_helpers.is_gpu(backend) + on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) grid = grid_savepoint.construct_icon_grid(on_gpu) z_ifc = metrics_savepoint.z_ifc() vct_a = grid_savepoint.vct_a() diff --git a/model/common/tests/utils.py b/model/common/tests/utils.py deleted file mode 100644 index 02b7b3774b..0000000000 --- a/model/common/tests/utils.py +++ /dev/null @@ -1,83 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -import logging as log - -import gt4py._core.definitions as gtcore_defs -import gt4py.next as gtx -import gt4py.next.backend as gtx_backend - -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 -from icon4py.model.common.utils import gt4py_field_allocation as alloc - -from .grid_tests import utils as gridtest_utils - - -def is_cupy_device(backend: gtx_backend.Backend) -> bool: - cuda_device_types = ( - gtcore_defs.DeviceType.CUDA, - gtcore_defs.DeviceType.CUDA_MANAGED, - gtcore_defs.DeviceType.ROCM, - ) - return backend.allocator.__gt_device_type__ in cuda_device_types - - -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: gtx_backend.Backend): - is_cupy_device(backend) - return array_ns(is_cupy_device(backend)) - - -grid_geometries = {} - - -# TODO @halungge: copied from test_geometry.py: should be remove from there. -# also check the imports. Should it rather go to the test_utils package? -def get_grid_geometry(backend: gtx_backend.Backend, grid_file: str) -> geometry.GridGeometry: - on_gpu = is_cupy_device(backend) - xp = array_ns(on_gpu) - - 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 = gridtest_utils.run_grid_manager(grid_file, on_gpu=on_gpu) - 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(grid_file) - return grid_geometries[grid_file] 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: ../../