From 1452892b48272248ccc2005585e8bc15d795fbc5 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Fri, 29 Nov 2024 13:35:31 +0100 Subject: [PATCH 1/3] Interpolation fields edits (#592) bugfixes in interpolation coefficients --- .../interpolation/interpolation_fields.py | 323 +++++++++++------- .../test_interpolation_fields.py | 74 ++-- 2 files changed, 236 insertions(+), 161 deletions(-) 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/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() From 66452b9d98e17fb555b924b57826d6ae19eb0abf Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 29 Nov 2024 14:05:02 +0100 Subject: [PATCH 2/3] - move the poormans GridGeometry cache to common/test_utils/grid_utils.py (#613) - renameing of functions in grid_utils.py - add poormans grid_manager cache in test_grid_manager.py to speed up tests (replacement for @functools.cache that does no longer work) - add convenience as_numpy function --- .../tests/diffusion_tests/test_diffusion.py | 2 +- .../icon4py/model/common/grid/grid_manager.py | 6 +- .../model/common/test_utils/grid_utils.py | 134 ++++++++++++++---- .../model/common/test_utils/pytest_config.py | 8 +- .../common/utils/gt4py_field_allocation.py | 19 ++- .../common/tests/grid_tests/test_geometry.py | 102 +++++-------- .../tests/grid_tests/test_grid_manager.py | 119 ++++++++-------- model/common/tests/grid_tests/test_icon.py | 3 +- model/common/tests/grid_tests/utils.py | 56 +------- model/common/tests/io_tests/test_io.py | 12 +- 10 files changed, 249 insertions(+), 212 deletions(-) diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py index 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 95b3ef5b9d..4f40a01316 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -394,7 +394,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is FileNotFoundError: raise FileNotFoundError(f"gridfile {self._file_name} not found, aborting") - def __call__(self, backend: gtx_backend.Backend, limited_area=True): + def __call__(self, backend: Optional[gtx_backend.Backend], limited_area=True): if not self._reader: self.open() on_gpu = common_utils.gt4py_field_allocation.is_cupy_device(backend) @@ -403,7 +403,7 @@ def __call__(self, backend: gtx_backend.Backend, limited_area=True): self._coordinates = self._read_coordinates(backend) self._geometry = self._read_geometry_fields(backend) - def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict: + def _read_coordinates(self, backend: Optional[gtx_backend.Backend]) -> CoordinateDict: return { dims.CellDim: { "lat": gtx.as_field( @@ -449,7 +449,7 @@ def _read_coordinates(self, backend: gtx_backend.Backend) -> CoordinateDict: }, } - def _read_geometry_fields(self, backend: gtx_backend.Backend): + def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]): return { # TODO (@halungge) still needs to ported, values from "our" grid files contains (wrong) values: # based on bug in generator fixed with this [PR40](https://gitlab.dkrz.de/dwd-sw/dwd_icon_tools/-/merge_requests/40) . diff --git a/model/common/src/icon4py/model/common/test_utils/grid_utils.py b/model/common/src/icon4py/model/common/test_utils/grid_utils.py index 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 7d5260eb20..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, backend=backend) - grid = gm.grid - decomposition_info = construct_decomposition_info(grid) - geometry_source = geometry.GridGeometry( - grid, decomposition_info, backend, gm.coordinates, gm.geometry, attrs.attrs - ) - return geometry_source - - if not grid_geometries.get(grid_file): - grid_geometries[grid_file] = construct_grid_geometry(grid_file) - return grid_geometries[grid_file] +from icon4py.model.common.test_utils import datatest_utils as dt_utils, grid_utils, helpers def test_geometry_raises_for_unknown_field(backend): - geometry = get_grid_geometry(backend, dt_utils.R02B04_GLOBAL) + geometry = grid_utils.get_grid_geometry( + backend, dt_utils.GLOBAL_EXPERIMENT, dt_utils.R02B04_GLOBAL + ) with pytest.raises(ValueError) as e: geometry.get("foo") assert "'foo'" in e.value @@ -67,9 +39,9 @@ def test_geometry_raises_for_unknown_field(backend): ], ) @pytest.mark.datatest -def test_edge_control_area(backend, grid_savepoint, grid_file, rtol): +def test_edge_control_area(backend, grid_savepoint, grid_file, experiment, rtol): expected = grid_savepoint.edge_areas() - geometry_source = get_grid_geometry(backend, grid_file) + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) result = geometry_source.get(attrs.EDGE_AREA) assert helpers.dallclose(expected.ndarray, result.ndarray, rtol) @@ -82,8 +54,8 @@ def test_edge_control_area(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_coriolis_parameter(backend, grid_savepoint, grid_file): - geometry_source = get_grid_geometry(backend, grid_file) +def test_coriolis_parameter(backend, grid_savepoint, grid_file, experiment): + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.f_e() result = geometry_source.get(attrs.CORIOLIS_PARAMETER) @@ -98,8 +70,8 @@ def test_coriolis_parameter(backend, grid_savepoint, grid_file): ], ) @pytest.mark.datatest -def test_compute_edge_length(backend, grid_savepoint, grid_file, rtol): - geometry_source = get_grid_geometry(backend, grid_file) +def test_compute_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.primal_edge_length() result = geometry_source.get(attrs.EDGE_LENGTH) assert helpers.dallclose(result.ndarray, expected.ndarray, rtol=rtol) @@ -113,9 +85,9 @@ def test_compute_edge_length(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_compute_inverse_edge_length(backend, grid_savepoint, grid_file, rtol): +def test_compute_inverse_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): expected = grid_savepoint.inverse_primal_edge_lengths() - geometry_source = get_grid_geometry(backend, grid_file) + geometry_source = grid_utils.get_grid_geometry(backend, experiment, grid_file) computed = geometry_source.get(f"inverse_of_{attrs.EDGE_LENGTH}") assert helpers.dallclose(computed.ndarray, expected.ndarray, rtol=rtol) @@ -129,8 +101,8 @@ def test_compute_inverse_edge_length(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_compute_dual_edge_length(backend, grid_savepoint, grid_file, rtol): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_dual_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.dual_edge_length() result = grid_geometry.get(attrs.DUAL_EDGE_LENGTH) @@ -145,8 +117,8 @@ def test_compute_dual_edge_length(backend, grid_savepoint, grid_file, rtol): ], ) @pytest.mark.datatest -def test_compute_inverse_dual_edge_length(backend, grid_savepoint, grid_file, rtol): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_inverse_dual_edge_length(backend, grid_savepoint, grid_file, experiment, rtol): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.inv_dual_edge_length() result = grid_geometry.get(f"inverse_of_{attrs.DUAL_EDGE_LENGTH}") @@ -166,8 +138,8 @@ def test_compute_inverse_dual_edge_length(backend, grid_savepoint, grid_file, rt ], ) @pytest.mark.datatest -def test_compute_inverse_vertex_vertex_length(backend, grid_savepoint, grid_file, rtol): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_inverse_vertex_vertex_length(backend, grid_savepoint, grid_file, experiment, rtol): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) expected = grid_savepoint.inv_vert_vert_length() result = grid_geometry.get(attrs.INVERSE_VERTEX_VERTEX_LENGTH) @@ -182,8 +154,10 @@ def test_compute_inverse_vertex_vertex_length(backend, grid_savepoint, grid_file (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_compute_coordinates_of_edge_tangent_and_normal(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_coordinates_of_edge_tangent_and_normal( + backend, grid_savepoint, grid_file, experiment +): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) x_normal = grid_geometry.get(attrs.EDGE_NORMAL_X) y_normal = grid_geometry.get(attrs.EDGE_NORMAL_Y) z_normal = grid_geometry.get(attrs.EDGE_NORMAL_Z) @@ -213,8 +187,8 @@ def test_compute_coordinates_of_edge_tangent_and_normal(backend, grid_savepoint, (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_compute_primal_normals(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_compute_primal_normals(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) primal_normal_u = grid_geometry.get(attrs.EDGE_NORMAL_U) primal_normal_v = grid_geometry.get(attrs.EDGE_NORMAL_V) @@ -233,8 +207,8 @@ def test_compute_primal_normals(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_tangent_orientation(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_tangent_orientation(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) result = grid_geometry.get(attrs.TANGENT_ORIENTATION) expected = grid_savepoint.tangent_orientation() @@ -249,8 +223,8 @@ def test_tangent_orientation(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_cell_area(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_cell_area(backend, grid_savepoint, experiment, grid_file): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) result = grid_geometry.get(attrs.CELL_AREA) expected = grid_savepoint.cell_areas() @@ -265,8 +239,8 @@ def test_cell_area(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_primal_normal_cell(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_primal_normal_cell(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) primal_normal_cell_u_ref = grid_savepoint.primal_normal_cell_x().ndarray primal_normal_cell_v_ref = grid_savepoint.primal_normal_cell_y().ndarray primal_normal_cell_u = grid_geometry.get(attrs.EDGE_NORMAL_CELL_U) @@ -284,8 +258,8 @@ def test_primal_normal_cell(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_dual_normal_cell(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_dual_normal_cell(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) dual_normal_cell_u_ref = grid_savepoint.dual_normal_cell_x().ndarray dual_normal_cell_v_ref = grid_savepoint.dual_normal_cell_y().ndarray dual_normal_cell_u = grid_geometry.get(attrs.EDGE_TANGENT_CELL_U) @@ -303,8 +277,8 @@ def test_dual_normal_cell(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_primal_normal_vert(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_primal_normal_vert(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) primal_normal_vert_u_ref = grid_savepoint.primal_normal_vert_x().ndarray primal_normal_vert_v_ref = grid_savepoint.primal_normal_vert_y().ndarray primal_normal_vert_u = grid_geometry.get(attrs.EDGE_NORMAL_VERTEX_U) @@ -322,8 +296,8 @@ def test_primal_normal_vert(backend, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_dual_normal_vert(backend, grid_savepoint, grid_file): - grid_geometry = get_grid_geometry(backend, grid_file) +def test_dual_normal_vert(backend, grid_savepoint, grid_file, experiment): + grid_geometry = grid_utils.get_grid_geometry(backend, experiment, grid_file) dual_normal_vert_u_ref = grid_savepoint.dual_normal_vert_x().ndarray dual_normal_vert_v_ref = grid_savepoint.dual_normal_vert_y().ndarray dual_normal_vert_u = grid_geometry.get(attrs.EDGE_TANGENT_VERTEX_U) @@ -357,7 +331,7 @@ def test_sparse_fields_creator(): ], ) def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file, backend=backend) + gm = grid_utils.get_grid_manager(grid_file, backend=backend, num_levels=1) grid = gm.grid coordinates = gm.coordinates @@ -366,7 +340,7 @@ def test_create_auxiliary_orientation_coordinates(backend, grid_savepoint, grid_ edge_lat = coordinates[dims.EdgeDim]["lat"] edge_lon = coordinates[dims.EdgeDim]["lon"] lat_0, lon_0, lat_1, lon_1 = geometry.create_auxiliary_coordinate_arrays_for_orientation( - gm.grid, cell_lat, cell_lon, edge_lat, edge_lon + grid, cell_lat, cell_lon, edge_lat, edge_lon ) connectivity = grid.connectivities[dims.E2CDim] has_boundary_edges = np.count_nonzero(connectivity == -1) diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index 48dc1b5086..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, transformation=zero_base) - grid = manager.grid + grid = _run_grid_manager(grid_file, backend).grid seralized_v2e = grid_savepoint.v2e() # there are vertices at the boundary of a local domain or at a pentagon point that have less than # 6 neighbors hence there are "Missing values" in the grid file @@ -159,9 +169,8 @@ def test_grid_manager_eval_v2e(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", [dims.CellDim, dims.EdgeDim, dims.VertexDim]) -def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): - manager = run_grid_manager(grid_file, transformation=zero_base) - refin_ctrl = manager.refinement +def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim, backend): + refin_ctrl = _run_grid_manager(grid_file, backend).refinement refin_ctrl_serialized = grid_savepoint.refin_ctrl(dim) assert np.all( refin_ctrl_serialized.ndarray @@ -179,9 +188,9 @@ def test_grid_manager_refin_ctrl(grid_savepoint, grid_file, experiment, dim): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_v2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_v2c(caplog, grid_savepoint, experiment, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_v2c = grid_savepoint.v2c() # there are vertices that have less than 6 neighboring cells: either pentagon points or # vertices at the boundary of the domain for a limited area mode @@ -231,9 +240,9 @@ def reset_invalid_index(index_array: np.ndarray): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2v(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_e2v = grid_savepoint.e2v() # all vertices in the system have to neighboring edges, there no edges that point nowhere @@ -266,7 +275,7 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): grid_file: name of grid file used """ - if utils.is_regional(grid_file): + if gridtest_utils.is_regional(grid_file): assert has_invalid_index(e2c_table) else: assert not has_invalid_index(e2c_table) @@ -282,9 +291,10 @@ def assert_invalid_indices(e2c_table: np.ndarray, grid_file: str): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + + grid = _run_grid_manager(grid_file, backend).grid serialized_e2c = grid_savepoint.e2c() e2c_table = grid.get_offset_provider("E2C").table assert_invalid_indices(serialized_e2c, grid_file) @@ -302,9 +312,9 @@ def test_grid_manager_eval_e2c(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_c2e = grid_savepoint.c2e() # no cells with less than 3 neighboring edges exist, otherwise the cell is not there in the @@ -325,9 +335,9 @@ def test_grid_manager_eval_c2e(caplog, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid assert np.allclose( grid.get_offset_provider("C2E2C").table, grid_savepoint.c2e2c(), @@ -343,9 +353,9 @@ def test_grid_manager_eval_c2e2c(caplog, grid_savepoint, grid_file): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) assert np.allclose( grid.get_offset_provider("C2E2CO").table, @@ -363,9 +373,9 @@ def test_grid_manager_eval_c2e2cO(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c2e(caplog, grid_savepoint, grid_file, experiment, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) serialized_e2c2e = serialized_grid.get_offset_provider("E2C2E").table serialized_e2c2eO = serialized_grid.get_offset_provider("E2C2EO").table @@ -395,10 +405,9 @@ def assert_unless_invalid(table, serialized_ref): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - gm = run_grid_manager(grid_file) - grid = gm.grid + grid = _run_grid_manager(grid_file, backend).grid # the "far" (adjacent to edge normal ) is not always there, because ICON only calculates those starting from # (lateral_boundary(dims.EdgeDim) + 1) to end(dims.EdgeDim) (see mo_intp_coeffs.f90) and only for owned cells serialized_ref = grid_savepoint.e2c2v() @@ -416,9 +425,9 @@ def test_grid_manager_eval_e2c2v(caplog, grid_savepoint, grid_file): (utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid c2v = grid.get_offset_provider("C2V").table assert np.allclose(c2v, grid_savepoint.c2v()) @@ -432,8 +441,8 @@ def test_grid_manager_eval_c2v(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.with_netcdf -def test_grid_manager_grid_size(dim, size): - grid = run_grid_manager(utils.R02B04_GLOBAL).grid +def test_grid_manager_grid_size(dim, size, backend): + grid = _run_grid_manager(utils.R02B04_GLOBAL, backend=backend).grid assert size == grid.size[dim] @@ -472,8 +481,8 @@ def test_gt4py_transform_offset_by_1_where_valid(size): (dt_utils.REGIONAL_EXPERIMENT, MCH_CH_RO4B09_GLOBAL_NUM_CELLS), ], ) -def test_grid_manager_grid_level_and_root(grid_file, global_num_cells): - assert global_num_cells == run_grid_manager(grid_file, num_levels=1).grid.global_num_cells +def test_grid_manager_grid_level_and_root(grid_file, global_num_cells, backend): + assert global_num_cells == _run_grid_manager(grid_file, backend=backend).grid.global_num_cells @pytest.mark.datatest @@ -482,9 +491,9 @@ def test_grid_manager_grid_level_and_root(grid_file, global_num_cells): "grid_file, experiment", [(utils.R02B04_GLOBAL, dt_utils.JABW_EXPERIMENT)], ) -def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): +def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file, backend): caplog.set_level(logging.DEBUG) - grid = run_grid_manager(grid_file).grid + grid = _run_grid_manager(grid_file, backend).grid serialized_grid = grid_savepoint.construct_icon_grid(on_gpu=False) assert np.allclose( grid.get_offset_provider("C2E2C2E").table, @@ -503,12 +512,10 @@ def test_grid_manager_eval_c2e2c2e(caplog, grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", utils.horizontal_dim()) -def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid): +def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_grid, backend): caplog.set_level(logging.INFO) serialized_grid = icon_grid - manager = run_grid_manager(grid_file, transformation=zero_base) - grid = manager.grid - + grid = _run_grid_manager(grid_file, backend).grid for domain in utils.global_grid_domains(dim): assert grid.start_index(domain) == serialized_grid.start_index( domain @@ -518,7 +525,7 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g ), f"end index wrong for domain {domain}" for domain in utils.valid_boundary_zones_for_dim(dim): - if not utils.is_regional(grid_file): + if not gridtest_utils.is_regional(grid_file): assert grid.start_index(domain) == 0 assert grid.end_index(domain) == 0 assert grid.start_index(domain) == serialized_grid.start_index( @@ -537,10 +544,10 @@ def test_grid_manager_start_end_index(caplog, grid_file, experiment, dim, icon_g (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_read_geometry_fields(grid_savepoint, grid_file): - gm = utils.run_grid_manager(grid_file) - cell_area = gm.geometry[GeometryName.CELL_AREA.value] - tangent_orientation = gm.geometry[GeometryName.TANGENT_ORIENTATION.value] +def test_read_geometry_fields(grid_savepoint, grid_file, backend): + manager = _run_grid_manager(grid_file, backend=backend) + cell_area = manager.geometry[GeometryName.CELL_AREA.value] + tangent_orientation = manager.geometry[GeometryName.TANGENT_ORIENTATION.value] assert helpers.dallclose(cell_area.asnumpy(), grid_savepoint.cell_areas().asnumpy()) assert helpers.dallclose( @@ -557,10 +564,10 @@ def test_read_geometry_fields(grid_savepoint, grid_file): ], ) @pytest.mark.parametrize("dim", (dims.CellDim, dims.EdgeDim, dims.VertexDim)) -def test_coordinates(grid_savepoint, grid_file, experiment, dim): - gm = utils.run_grid_manager(grid_file) - lat = gm.coordinates[dim]["lat"] - lon = gm.coordinates[dim]["lon"] +def test_coordinates(grid_savepoint, grid_file, experiment, dim, backend): + manager = _run_grid_manager(grid_file, backend=backend) + lat = manager.coordinates[dim]["lat"] + lon = manager.coordinates[dim]["lon"] assert helpers.dallclose(lat.asnumpy(), grid_savepoint.lat(dim).asnumpy()) assert helpers.dallclose(lon.asnumpy(), grid_savepoint.lon(dim).asnumpy()) @@ -573,10 +580,10 @@ def test_coordinates(grid_savepoint, grid_file, experiment, dim): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_tangent_orientation(experiment, grid_file, grid_savepoint): +def test_tangent_orientation(experiment, grid_file, grid_savepoint, backend): expected = grid_savepoint.tangent_orientation() - gm = utils.run_grid_manager(grid_file) - geometry_fields = gm.geometry + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry assert helpers.dallclose( - geometry_fields[GeometryName.TANGENT_ORIENTATION].ndarray, expected.ndarray + geometry_fields[GeometryName.TANGENT_ORIENTATION].asnumpy(), expected.asnumpy() ) diff --git a/model/common/tests/grid_tests/test_icon.py b/model/common/tests/grid_tests/test_icon.py index 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/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, From 4db692345efe04f3baa416f4a6002a28311e31a4 Mon Sep 17 00:00:00 2001 From: David Strassmann <107603957+dastrm@users.noreply.github.com> Date: Mon, 2 Dec 2024 06:45:34 +0100 Subject: [PATCH 3/3] Add vertical advection granule with PPM (#605) * Add vertical advection granule with PPM --------- Co-authored-by: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> --- .../model/atmosphere/advection/advection.py | 29 +- .../atmosphere/advection/advection_states.py | 3 + .../advection/advection_vertical.py | 828 +++++++++++++++++- .../compute_ppm4gpu_courant_number.py | 174 ++++ .../compute_ppm4gpu_fractional_flux.py | 121 +++ .../stencils/compute_ppm4gpu_integer_flux.py | 78 ++ .../compute_vertical_tracer_flux_upwind.py | 43 + .../copy_cell_kdim_field_koff_plus1.py | 42 + ...limit_vertical_slope_semi_monotonically.py | 51 ++ ...test_apply_horizontal_density_increment.py | 12 +- ...apply_interpolated_tracer_time_tendency.py | 10 +- .../test_apply_vertical_density_increment.py | 10 +- ...st_average_horizontal_flux_subcycling_2.py | 6 +- ...st_average_horizontal_flux_subcycling_3.py | 8 +- ...e_antidiffusive_cell_fluxes_and_min_max.py | 36 +- ...test_compute_barycentric_backtrajectory.py | 54 +- ..._compute_barycentric_backtrajectory_alt.py | 38 +- .../test_compute_ffsl_backtrajectory.py | 100 +-- ...cktrajectory_counterclockwise_indicator.py | 14 +- ...te_ffsl_backtrajectory_length_indicator.py | 18 +- ...tal_tracer_flux_from_cubic_coefficients.py | 8 +- ...al_tracer_flux_from_linear_coefficients.py | 24 +- ...racer_flux_from_linear_coefficients_alt.py | 24 +- ...e_horizontal_multiplicative_flux_factor.py | 24 +- ...t_compute_ppm4gpu_parabola_coefficients.py | 4 +- .../test_compute_ppm_all_face_values.py | 18 +- .../test_compute_ppm_quadratic_face_values.py | 4 +- .../test_compute_ppm_quartic_face_values.py | 4 +- .../test_compute_ppm_slope.py | 8 +- .../test_compute_tendency.py | 8 +- ...ute_vertical_parabola_limiter_condition.py | 6 +- ...est_compute_vertical_tracer_flux_upwind.py | 58 ++ ...t_integrate_tracer_density_horizontally.py | 20 +- .../test_integrate_tracer_horizontally.py | 20 +- .../test_integrate_tracer_vertically.py | 22 +- ...it_vertical_parabola_semi_monotonically.py | 12 +- ...limit_vertical_slope_semi_monotonically.py | 52 ++ ...s_antidiffusive_cell_fluxes_and_min_max.py | 28 +- ...est_prepare_ffsl_flux_area_patches_list.py | 422 ++++----- ...cal_quadrature_for_cubic_reconstruction.py | 26 +- ...uadrature_list_for_cubic_reconstruction.py | 50 +- ...test_reconstruct_cubic_coefficients_svd.py | 112 +-- .../tests/advection_tests/test_advection.py | 30 +- .../advection/tests/advection_tests/utils.py | 67 +- .../common/test_utils/datatest_fixtures.py | 2 +- .../model/common/test_utils/datatest_utils.py | 2 +- 46 files changed, 2083 insertions(+), 647 deletions(-) create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py create mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py create mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index ae2aae72b2..410a40fd03 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -75,6 +75,10 @@ class VerticalAdvectionType(Enum): #: no vertical advection NO_ADVECTION = auto() + #: 1st order upwind + UPWIND_1ST_ORDER = auto() + #: 3rd order PPM + PPM_3RD_ORDER = auto() class VerticalAdvectionLimiter(Enum): @@ -84,6 +88,8 @@ class VerticalAdvectionLimiter(Enum): #: no vertical limiter NO_LIMITER = auto() + #: semi-monotonic vertical limiter + SEMI_MONOTONIC = auto() @dataclasses.dataclass(frozen=True) @@ -393,7 +399,7 @@ def convert_config_to_horizontal_vertical_advection( match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: - horizontal_advection = advection_horizontal.NoAdvection(grid=grid) + horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend) case HorizontalAdvectionType.LINEAR_2ND_ORDER: tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, @@ -417,13 +423,32 @@ def convert_config_to_horizontal_vertical_advection( match config.vertical_advection_limiter: case VerticalAdvectionLimiter.NO_LIMITER: - ... + vertical_limiter = advection_vertical.NoLimiter(grid=grid, backend=backend) + case VerticalAdvectionLimiter.SEMI_MONOTONIC: + vertical_limiter = advection_vertical.SemiMonotonicLimiter(grid=grid, backend=backend) case _: raise NotImplementedError(f"Unknown vertical advection limiter.") match config.vertical_advection_type: case VerticalAdvectionType.NO_ADVECTION: vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend) + case VerticalAdvectionType.UPWIND_1ST_ORDER: + boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) + vertical_advection = advection_vertical.FirstOrderUpwind( + boundary_conditions=boundary_conditions, + grid=grid, + metric_state=metric_state, + backend=backend, + ) + case VerticalAdvectionType.PPM_3RD_ORDER: + boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) + vertical_advection = advection_vertical.PiecewiseParabolicMethod( + boundary_conditions=boundary_conditions, + vertical_limiter=vertical_limiter, + grid=grid, + metric_state=metric_state, + backend=backend, + ) case _: raise NotImplementedError(f"Unknown vertical advection type.") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py index 9a038ecdb1..a09ed7dafe 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py @@ -85,3 +85,6 @@ class AdvectionMetricState: #: metrical modification factor for vertical part of divergence at full levels (KDim) deepatmo_divzu: fa.KField[ta.wpfloat] + + #: vertical grid spacing at full levels + ddqz_z_full: fa.CellKField[ta.wpfloat] diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index 0ddd9daf41..374db87949 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -10,13 +10,59 @@ import logging import icon4py.model.common.grid.states as grid_states +import gt4py.next as gtx from gt4py.next import backend from icon4py.model.atmosphere.advection import advection_states +from icon4py.model.atmosphere.advection.stencils.compute_ppm_quadratic_face_values import ( + compute_ppm_quadratic_face_values, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm_quartic_face_values import ( + compute_ppm_quartic_face_values, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm_slope import compute_ppm_slope +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_courant_number import ( + compute_ppm4gpu_courant_number, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( + compute_ppm4gpu_fractional_flux, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_integer_flux import ( + compute_ppm4gpu_integer_flux, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_parabola_coefficients import ( + compute_ppm4gpu_parabola_coefficients, +) +from icon4py.model.atmosphere.advection.stencils.compute_vertical_parabola_limiter_condition import ( + compute_vertical_parabola_limiter_condition, +) +from icon4py.model.atmosphere.advection.stencils.compute_vertical_tracer_flux_upwind import ( + compute_vertical_tracer_flux_upwind, +) from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field import copy_cell_kdim_field +from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field_koff_minus1 import ( + copy_cell_kdim_field_koff_minus1, +) +from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field_koff_plus1 import ( + copy_cell_kdim_field_koff_plus1, +) +from icon4py.model.atmosphere.advection.stencils.init_constant_cell_kdim_field import ( + init_constant_cell_kdim_field, +) +from icon4py.model.atmosphere.advection.stencils.integrate_tracer_vertically import ( + integrate_tracer_vertically, +) +from icon4py.model.atmosphere.advection.stencils.limit_vertical_parabola_semi_monotonically import ( + limit_vertical_parabola_semi_monotonically, +) +from icon4py.model.atmosphere.advection.stencils.limit_vertical_slope_semi_monotonically import ( + limit_vertical_slope_semi_monotonically, +) + from icon4py.model.common import ( + constants, dimension as dims, field_type_aliases as fa, type_alias as ta, @@ -34,6 +80,274 @@ log = logging.getLogger(__name__) +class BoundaryConditions(ABC): + """Class that sets the upper and lower boundary conditions.""" + + @abstractmethod + def run( + self, + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + """ + Set the vertical boundary conditions. + + Args: + p_mflx_tracer_v: output argument, field that contains new vertical tracer mass flux + horizontal_start: input argument, horizontal start index + horizontal_end: input argument, horizontal end index + + """ + ... + + +class NoFluxCondition(BoundaryConditions): + """Class that sets the upper and lower boundary fluxes to zero.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # stencils + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + + def run( + self, + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + log.debug("vertical boundary conditions computation - start") + + # set upper boundary conditions + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=p_mflx_tracer_v, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + # set lower boundary conditions + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=p_mflx_tracer_v, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + log.debug("vertical boundary conditions computation - end") + + +class VerticalLimiter(ABC): + """Class that limits the vertical reconstructed fields and the fluxes.""" + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + +class NoLimiter(VerticalLimiter): + """Class that implements no vertical parabola limiter.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # fields + self._l_limit = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + + # stencils + self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field_koff_plus1 = copy_cell_kdim_field_koff_plus1.with_backend( + self._backend + ) + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + # simply copy to up/low face values + log.debug("running stencil copy_cell_kdim_field - start") + self._copy_cell_kdim_field( + field_in=p_face, + field_out=p_face_up, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field - end") + + log.debug("running stencil copy_cell_kdim_field_koff_plus1 - start") + self._copy_cell_kdim_field_koff_plus1( + field_in=p_face, + field_out=p_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field_koff_plus1 - end") + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + +class SemiMonotonicLimiter(VerticalLimiter): + """Class that implements a semi-monotonic vertical parabola limiter.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._l_limit = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, dtype=gtx.int32, backend=self._backend + ) + + # stencils + self._limit_vertical_slope_semi_monotonically = ( + limit_vertical_slope_semi_monotonically.with_backend(self._backend) + ) + self._compute_vertical_parabola_limiter_condition = ( + compute_vertical_parabola_limiter_condition.with_backend(self._backend) + ) + self._limit_vertical_parabola_semi_monotonically = ( + limit_vertical_parabola_semi_monotonically.with_backend(self._backend) + ) + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + log.debug("running stencil limit_vertical_slope_semi_monotonically - start") + self._limit_vertical_slope_semi_monotonically( + p_cc=p_tracer_now, + z_slope=z_slope, + k=self._k_field, + elev=self._grid.num_levels - 1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil limit_vertical_slope_semi_monotonically - end") + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + # compute semi-monotonic limiter condition + log.debug("running stencil compute_vertical_parabola_limiter_condition - start") + self._compute_vertical_parabola_limiter_condition( + p_face=p_face, + p_cc=p_tracer_now, + l_limit=self._l_limit, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_vertical_parabola_limiter_condition - end") + + # apply semi-monotonic limiter condition and store to up/low face values + log.debug("running stencil limit_vertical_parabola_semi_monotonically - start") + self._limit_vertical_parabola_semi_monotonically( + l_limit=self._l_limit, + p_face=p_face, + p_cc=p_tracer_now, + p_face_up=p_face_up, + p_face_low=p_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil limit_vertical_parabola_semi_monotonically - end") + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + class VerticalAdvection(ABC): """Class that does one vertical advection step.""" @@ -62,6 +376,12 @@ def run( dtime: input argument, the time step even_timestep: input argument, determines whether halo points are included + Note: + Originally, if even step, vertical transport includes all halo points in order to avoid an additional synchronization step, i.e. + if lstep_even: i_rlstart = 2, i_rlend = min_rlcell + else: i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int + Horizontal advection is always called with the same indices though, i.e. + i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int """ ... @@ -90,6 +410,16 @@ def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): log.debug("vertical advection class init - end") + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + def run( self, prep_adv: advection_states.AdvectionPrepAdvState, @@ -103,15 +433,16 @@ def run( ): log.debug("vertical advection run - start") - horizontal_start = ( - self._start_cell_lateral_boundary_level_2 if even_timestep else self._start_cell_nudging + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep ) + log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( field_in=p_tracer_now, field_out=p_tracer_new, horizontal_start=horizontal_start, - horizontal_end=(self._end_cell_end if even_timestep else self._end_cell_local), + horizontal_end=horizontal_end, vertical_start=0, vertical_end=self._grid.num_levels, offset_provider=self._grid.offset_providers, @@ -135,13 +466,7 @@ def run( dtime: ta.wpfloat, even_timestep: bool = False, ): - log.debug("horizontal advection run - start") - - # TODO (dastrm): maybe change how the indices are handled here? originally: - # if even step, vertical transport includes all halo points in order to avoid an additional synchronization step, i.e. - # if lstep_even: i_rlstart = 2, i_rlend = min_rlcell - # else: i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int - # note: horizontal advection is always called with the same indices, i.e. i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int + log.debug("vertical advection run - start") self._compute_numerical_flux( prep_adv=prep_adv, @@ -149,6 +474,7 @@ def run( rhodz_now=rhodz_now, p_mflx_tracer_v=p_mflx_tracer_v, dtime=dtime, + even_timestep=even_timestep, ) self._update_unknowns( @@ -158,9 +484,10 @@ def run( rhodz_new=rhodz_new, p_mflx_tracer_v=p_mflx_tracer_v, dtime=dtime, + even_timestep=even_timestep, ) - log.debug("horizontal advection run - end") + log.debug("vertical advection run - end") @abstractmethod def _compute_numerical_flux( @@ -170,6 +497,7 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): ... @@ -182,35 +510,28 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): ... -class SemiLagrangian(FiniteVolume): - """Class that does one vertical semi-Lagrangian finite volume advection step.""" +class FirstOrderUpwind(FiniteVolume): + """Class that does one vertical first-order accurate upwind finite volume advection step.""" def __init__( self, + boundary_conditions: BoundaryConditions, grid: icon_grid.IconGrid, - interpolation_state: advection_states.AdvectionInterpolationState, - least_squares_state: advection_states.AdvectionLeastSquaresState, metric_state: advection_states.AdvectionMetricState, - edge_params: grid_states.EdgeParams, - cell_params: grid_states.CellParams, - backend: backend.Backend, - exchange: decomposition.ExchangeRuntime = decomposition.SingleNodeExchange(), + backend=backend, ): log.debug("vertical advection class init - start") # input arguments + self._boundary_conditions = boundary_conditions self._grid = grid - self._interpolation_state = interpolation_state - self._least_squares_state = least_squares_state self._metric_state = metric_state - self._edge_params = edge_params - self._cell_params = cell_params self._backend = backend - self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -221,8 +542,36 @@ def __init__( self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + + # stencils + self._compute_vertical_tracer_flux_upwind = ( + compute_vertical_tracer_flux_upwind.with_backend(self._backend) + ) + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + self._integrate_tracer_vertically = integrate_tracer_vertically.with_backend(self._backend) + + # misc + self._ivadv_tracer = 1 + self._iadv_slev_jt = 0 + log.debug("vertical advection class init - end") + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + def _compute_numerical_flux( self, prep_adv: advection_states.AdvectionPrepAdvState, @@ -230,9 +579,35 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): - # TODO (dastrm): implement this - ... + log.debug("vertical numerical flux computation - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + log.debug("running stencil compute_vertical_tracer_flux_upwind - start") + self._compute_vertical_tracer_flux_upwind( + p_cc=p_tracer_now, + p_mflx_contra_v=prep_adv.mass_flx_ic, + p_upflux=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_vertical_tracer_flux_upwind - end") + + # set boundary conditions + self._boundary_conditions.run( + p_mflx_tracer_v=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + log.debug("vertical numerical flux computation - end") def _update_unknowns( self, @@ -242,6 +617,403 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): - # TODO (dastrm): implement this - ... + log.debug("vertical unknowns update - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + # update tracer mass fraction + log.debug("running stencil integrate_tracer_vertically - start") + self._integrate_tracer_vertically( + tracer_now=p_tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=self._metric_state.deepatmo_divzl, + deepatmo_divzu=self._metric_state.deepatmo_divzu, + rhodz_new=rhodz_new, + tracer_new=p_tracer_new, + k=self._k_field, + p_dtime=dtime, + ivadv_tracer=self._ivadv_tracer, + iadv_slev_jt=self._iadv_slev_jt, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil integrate_tracer_vertically - end") + + log.debug("vertical unknowns update - end") + + +class PiecewiseParabolicMethod(FiniteVolume): + """Class that does one vertical PPM finite volume advection step.""" + + def __init__( + self, + boundary_conditions: BoundaryConditions, + vertical_limiter: VerticalLimiter, + grid: icon_grid.IconGrid, + metric_state: advection_states.AdvectionMetricState, + backend=backend, + ): + log.debug("vertical advection class init - start") + + # input arguments + self._boundary_conditions = boundary_conditions + self._vertical_limiter = vertical_limiter + self._grid = grid + self._metric_state = metric_state + self._backend = backend + + # cell indices + cell_domain = h_grid.domain(dims.CellDim) + self._start_cell_lateral_boundary_level_2 = self._grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + self._start_cell_nudging = self._grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) + + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_cfl = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, is_halfdim=True, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_slope = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_face = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, is_halfdim=True, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_face_up = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_face_low = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_delta_q = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_a1 = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + + # stencils + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + self._compute_ppm4gpu_courant_number = compute_ppm4gpu_courant_number.with_backend( + self._backend + ) + self._compute_ppm_slope = compute_ppm_slope.with_backend(self._backend) + self._compute_ppm_quadratic_face_values = compute_ppm_quadratic_face_values.with_backend( + self._backend + ) + self._compute_ppm_quartic_face_values = compute_ppm_quartic_face_values.with_backend( + self._backend + ) + self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field_koff_minus1 = copy_cell_kdim_field_koff_minus1.with_backend( + self._backend + ) + self._compute_ppm4gpu_parabola_coefficients = ( + compute_ppm4gpu_parabola_coefficients.with_backend(self._backend) + ) + self._compute_ppm4gpu_fractional_flux = compute_ppm4gpu_fractional_flux.with_backend( + self._backend + ) + self._compute_ppm4gpu_integer_flux = compute_ppm4gpu_integer_flux.with_backend( + self._backend + ) + self._integrate_tracer_vertically = integrate_tracer_vertically.with_backend(self._backend) + + # misc + self._slev = 0 + self._slevp1_ti = 1 + self._elev = self._grid.num_levels - 1 + self._nlev = self._grid.num_levels - 1 + self._ivadv_tracer = 1 + self._iadv_slev_jt = 0 + + log.debug("vertical advection class init - end") + + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + + def _compute_numerical_flux( + self, + prep_adv: advection_states.AdvectionPrepAdvState, + p_tracer_now: fa.CellKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + dtime: ta.wpfloat, + even_timestep: bool, + ): + log.debug("vertical numerical flux computation - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + ## compute density-weighted Courant number + + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=self._z_cfl, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + log.debug("running stencil compute_ppm4gpu_courant_number - start") + self._compute_ppm4gpu_courant_number( + p_mflx_contra_v=prep_adv.mass_flx_ic, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + k=self._k_field, + slevp1_ti=self._slevp1_ti, + nlev=self._nlev, + dbl_eps=constants.DBL_EPS, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_courant_number - end") + + ## reconstruct face values + + # compute slope + log.debug("running stencil compute_ppm_slope - start") + self._compute_ppm_slope( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + k=self._k_field, + z_slope=self._z_slope, + elev=self._elev, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_slope - end") + + # limit slope + self._vertical_limiter.limit_slope( + p_tracer_now=p_tracer_now, + z_slope=self._z_slope, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + # compute second highest face value + log.debug("running stencil compute_ppm_quadratic_face_values - start") + self._compute_ppm_quadratic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=2, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quadratic_face_values - end") + + # compute second lowest face value + log.debug("running stencil compute_ppm_quadratic_face_values - start") + self._compute_ppm_quadratic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels - 1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quadratic_face_values - end") + + # compute highest face value + log.debug("running stencil copy_cell_kdim_field - start") + self._copy_cell_kdim_field( + field_in=p_tracer_now, + field_out=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field - end") + + # compute lowest face value + log.debug("running stencil copy_cell_kdim_field_koff_minus1 - start") + self._copy_cell_kdim_field_koff_minus1( + field_in=p_tracer_now, + field_out=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field_koff_minus1 - end") + + # compute all other face values + log.debug("running stencil compute_ppm_quartic_face_values - start") + self._compute_ppm_quartic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + z_slope=self._z_slope, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=2, + vertical_end=self._grid.num_levels - 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quartic_face_values - end") + + ## limit reconstruction + + self._vertical_limiter.limit_parabola( + p_tracer_now=p_tracer_now, + p_face=self._z_face, + p_face_up=self._z_face_up, + p_face_low=self._z_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + ## compute fractional numerical flux + + log.debug("running stencil compute_ppm4gpu_parabola_coefficients - start") + self._compute_ppm4gpu_parabola_coefficients( + z_face_up=self._z_face_up, + z_face_low=self._z_face_low, + p_cc=p_tracer_now, + z_delta_q=self._z_delta_q, + z_a1=self._z_a1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_parabola_coefficients - end") + + log.debug("running stencil compute_ppm4gpu_fractional_flux - start") + self._compute_ppm4gpu_fractional_flux( + p_cc=p_tracer_now, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + z_delta_q=self._z_delta_q, + z_a1=self._z_a1, + p_upflux=p_mflx_tracer_v, + k=self._k_field, + slev=self._slev, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_fractional_flux - end") + + ## compute integer numerical flux + + log.debug("running stencil compute_ppm4gpu_integer_flux - start") + self._compute_ppm4gpu_integer_flux( + p_cc=p_tracer_now, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + p_upflux=p_mflx_tracer_v, + k=self._k_field, + slev=self._slev, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_integer_flux - end") + + ## set boundary conditions + + self._boundary_conditions.run( + p_mflx_tracer_v=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + ## apply flux limiter + + self._vertical_limiter.limit_fluxes( + horizontal_start=horizontal_start, horizontal_end=horizontal_end + ) + + log.debug("vertical numerical flux computation - end") + + def _update_unknowns( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_tracer_new: fa.CellKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + rhodz_new: fa.CellKField[ta.wpfloat], + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + dtime: ta.wpfloat, + even_timestep: bool, + ): + log.debug("vertical unknowns update - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + # update tracer mass fraction + log.debug("running stencil integrate_tracer_vertically - start") + self._integrate_tracer_vertically( + tracer_now=p_tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=self._metric_state.deepatmo_divzl, + deepatmo_divzu=self._metric_state.deepatmo_divzu, + rhodz_new=rhodz_new, + tracer_new=p_tracer_new, + k=self._k_field, + p_dtime=dtime, + ivadv_tracer=self._ivadv_tracer, + iadv_slev_jt=self._iadv_slev_jt, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil integrate_tracer_vertically - end") + + log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py new file mode 100644 index 0000000000..2c45621b29 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py @@ -0,0 +1,174 @@ +# 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.ffront.fbuiltins import abs, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _compute_courant_number_below( + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_mass: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + nlev: gtx.int32, + dbl_eps: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass_pos = z_mass > 0.0 + + in_bounds_p0 = k <= nlev - 1 + in_bounds_p1 = k <= nlev - 2 + in_bounds_p2 = k <= nlev - 3 + in_bounds_p3 = k <= nlev - 4 + + mass_gt_cellmass_p0 = where(z_mass_pos & in_bounds_p0, z_mass >= p_cellmass_now, False) + z_mass = z_mass - where(mass_gt_cellmass_p0, p_cellmass_now, 0.0) + + mass_gt_cellmass_p1 = mass_gt_cellmass_p0 & where( + z_mass_pos & in_bounds_p1, z_mass >= p_cellmass_now(Koff[1]), False + ) + z_mass = z_mass - where(mass_gt_cellmass_p1, p_cellmass_now(Koff[1]), 0.0) + + mass_gt_cellmass_p2 = mass_gt_cellmass_p1 & where( + z_mass_pos & in_bounds_p2, z_mass >= p_cellmass_now(Koff[2]), False + ) + z_mass = z_mass - where(mass_gt_cellmass_p2, p_cellmass_now(Koff[2]), 0.0) + + mass_gt_cellmass_p3 = mass_gt_cellmass_p2 & where( + z_mass_pos & in_bounds_p3, z_mass >= p_cellmass_now(Koff[3]), False + ) + z_mass = z_mass - where(mass_gt_cellmass_p3, p_cellmass_now(Koff[3]), 0.0) + + z_cfl = z_cfl + where(mass_gt_cellmass_p0, 1.0, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass_p1, 1.0, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass_p2, 1.0, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass_p3, 1.0, 0.0) + + p_cellmass_now_jks = p_cellmass_now + p_cellmass_now_jks = where(mass_gt_cellmass_p0, p_cellmass_now(Koff[1]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_p1, p_cellmass_now(Koff[2]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_p2, p_cellmass_now(Koff[3]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_p3, p_cellmass_now(Koff[4]), p_cellmass_now_jks) + + z_cflfrac = where(z_mass_pos, z_mass / p_cellmass_now_jks, 0.0) + z_cfl = z_cfl + where(z_cflfrac < 1.0, z_cflfrac, 1.0 - dbl_eps) + + return z_cfl + + +@gtx.field_operator +def _compute_courant_number_above( + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_mass: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + dbl_eps: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass_neg = z_mass <= 0.0 + + in_bounds_m0 = k >= slevp1_ti + 1 + in_bounds_m1 = k >= slevp1_ti + 2 + in_bounds_m2 = k >= slevp1_ti + 3 + in_bounds_m3 = k >= slevp1_ti + 4 + + mass_gt_cellmass_m0 = where( + z_mass_neg & in_bounds_m0, abs(z_mass) >= p_cellmass_now(Koff[-1]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m0, p_cellmass_now(Koff[-1]), 0.0) + + mass_gt_cellmass_m1 = mass_gt_cellmass_m0 & where( + z_mass_neg & in_bounds_m1, abs(z_mass) >= p_cellmass_now(Koff[-2]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m1, p_cellmass_now(Koff[-2]), 0.0) + + mass_gt_cellmass_m2 = mass_gt_cellmass_m1 & where( + z_mass_neg & in_bounds_m2, abs(z_mass) >= p_cellmass_now(Koff[-3]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m2, p_cellmass_now(Koff[-3]), 0.0) + + mass_gt_cellmass_m3 = mass_gt_cellmass_m2 & where( + z_mass_neg & in_bounds_m3, abs(z_mass) >= p_cellmass_now(Koff[-4]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m3, p_cellmass_now(Koff[-4]), 0.0) + + p_cellmass_now_jks = p_cellmass_now(Koff[-1]) + p_cellmass_now_jks = where(mass_gt_cellmass_m0, p_cellmass_now(Koff[-2]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_m1, p_cellmass_now(Koff[-3]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_m2, p_cellmass_now(Koff[-4]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_m3, p_cellmass_now(Koff[-5]), p_cellmass_now_jks) + + z_cfl = z_cfl - where(mass_gt_cellmass_m0, 1.0, 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass_m1, 1.0, 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass_m2, 1.0, 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass_m3, 1.0, 0.0) + + z_cflfrac = where(z_mass_neg, z_mass / p_cellmass_now_jks, 0.0) + z_cfl = z_cfl + where(abs(z_cflfrac) < 1.0, z_cflfrac, dbl_eps - 1.0) + + return z_cfl + + +@gtx.field_operator +def _compute_ppm4gpu_courant_number( + p_mflx_contra_v: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + nlev: gtx.int32, + dbl_eps: ta.wpfloat, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass = p_dtime * p_mflx_contra_v + + cfl_below = _compute_courant_number_below(p_cellmass_now, z_mass, z_cfl, k, nlev, dbl_eps) + cfl_above = _compute_courant_number_above(p_cellmass_now, z_mass, z_cfl, k, slevp1_ti, dbl_eps) + + z_cfl = cfl_below + cfl_above + + return z_cfl + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_courant_number( + p_mflx_contra_v: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + nlev: gtx.int32, + dbl_eps: ta.wpfloat, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_courant_number( + p_mflx_contra_v, + p_cellmass_now, + z_cfl, + k, + slevp1_ti, + nlev, + dbl_eps, + p_dtime, + out=z_cfl, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py new file mode 100644 index 0000000000..2ac11aa036 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py @@ -0,0 +1,121 @@ +# 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.ffront.fbuiltins import abs, astype, floor, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff +from icon4py.model.common.type_alias import wpfloat + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _sum_neighbor_contributions( + mask1: fa.CellKField[bool], + mask2: fa.CellKField[bool], + js: fa.CellKField[ta.wpfloat], + p_cc: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + p_cc_p0 = where(mask1 & (js == 0.0), p_cc, 0.0) + p_cc_p1 = where(mask1 & (js == 1.0), p_cc(Koff[1]), 0.0) + p_cc_p2 = where(mask1 & (js == 2.0), p_cc(Koff[2]), 0.0) + p_cc_p3 = where(mask1 & (js == 3.0), p_cc(Koff[3]), 0.0) + p_cc_p4 = where(mask1 & (js == 4.0), p_cc(Koff[4]), 0.0) + p_cc_m0 = where(mask2 & (js == 0.0), p_cc(Koff[-1]), 0.0) + p_cc_m1 = where(mask2 & (js == 1.0), p_cc(Koff[-2]), 0.0) + p_cc_m2 = where(mask2 & (js == 2.0), p_cc(Koff[-3]), 0.0) + p_cc_m3 = where(mask2 & (js == 3.0), p_cc(Koff[-4]), 0.0) + p_cc_m4 = where(mask2 & (js == 4.0), p_cc(Koff[-5]), 0.0) + p_cc_jks = ( + p_cc_p0 + + p_cc_p1 + + p_cc_p2 + + p_cc_p3 + + p_cc_p4 + + p_cc_m0 + + p_cc_m1 + + p_cc_m2 + + p_cc_m3 + + p_cc_m4 + ) + return p_cc_jks + + +@gtx.field_operator +def _compute_ppm4gpu_fractional_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + z_delta_q: fa.CellKField[ta.wpfloat], + z_a1: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + js = floor(abs(z_cfl)) + z_cflfrac = abs(z_cfl) - js + + z_cfl_pos = z_cfl > 0.0 + z_cfl_neg = not z_cfl_pos + wsign = where(z_cfl_pos, 1.0, -1.0) + + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) + + p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + z_delta_q_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_delta_q) + z_a1_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_a1) + + z_q_int = ( + p_cc_jks + + wsign * (z_delta_q_jks * (1.0 - z_cflfrac)) + - z_a1_jks * (1.0 - 3.0 * z_cflfrac + 2.0 * z_cflfrac * z_cflfrac) + ) + + p_upflux = where( + in_slev_bounds, wsign * p_cellmass_now_jks * z_cflfrac * z_q_int / p_dtime, 0.0 + ) + + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_fractional_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + z_delta_q: fa.CellKField[ta.wpfloat], + z_a1: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_fractional_flux( + p_cc, + p_cellmass_now, + z_cfl, + z_delta_q, + z_a1, + k, + slev, + p_dtime, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py new file mode 100644 index 0000000000..b6c5b8c7d0 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py @@ -0,0 +1,78 @@ +# 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.ffront.fbuiltins import abs, astype, floor, where + +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( + _sum_neighbor_contributions, +) +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.type_alias import wpfloat + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _compute_ppm4gpu_integer_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + js = floor(abs(z_cfl)) - 1.0 + + z_cfl_pos = z_cfl > 0.0 + z_cfl_neg = not z_cfl_pos + wsign = where(z_cfl_pos, 1.0, -1.0) + + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) + + p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + + z_iflx = wsign * p_cc_jks * p_cellmass_now_jks + + p_upflux = p_upflux + where(in_slev_bounds, z_iflx / p_dtime, 0.0) + + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_integer_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_integer_flux( + p_cc, + p_cellmass_now, + z_cfl, + p_upflux, + k, + slev, + p_dtime, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py new file mode 100644 index 0000000000..7e4a55d7ab --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py @@ -0,0 +1,43 @@ +# 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.ffront.fbuiltins import where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +@gtx.field_operator +def _compute_vertical_tracer_flux_upwind( + p_cc: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim +) -> fa.CellKField[ta.wpfloat]: + p_upflux = where(p_mflx_contra_v >= 0.0, p_cc, p_cc(Koff[-1])) * p_mflx_contra_v + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_vertical_tracer_flux_upwind( + p_cc: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_upflux: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_vertical_tracer_flux_upwind( + p_cc, + p_mflx_contra_v, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py new file mode 100644 index 0000000000..b3b296deed --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py @@ -0,0 +1,42 @@ +# 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 icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +# TODO (dastrm): move this highly generic stencil to common +# TODO (dastrm): this stencil has no test + + +@gtx.field_operator +def _copy_cell_kdim_field_koff_plus1( + field_in: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + return field_in(Koff[1]) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def copy_cell_kdim_field_koff_plus1( + field_in: fa.CellKField[ta.wpfloat], + field_out: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _copy_cell_kdim_field_koff_plus1( + field_in, + out=field_out, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py new file mode 100644 index 0000000000..4737343fa0 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py @@ -0,0 +1,51 @@ +# 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.ffront.fbuiltins import abs, minimum, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +@gtx.field_operator +def _limit_vertical_slope_semi_monotonically( + p_cc: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + elev: gtx.int32, +) -> fa.CellKField[ta.wpfloat]: + p_cc_min_last = minimum(p_cc(Koff[-1]), p_cc) + p_cc_min = where(k == elev, p_cc_min_last, minimum(p_cc_min_last, p_cc(Koff[1]))) + slope_l = minimum(abs(z_slope), 2.0 * (p_cc - p_cc_min)) + slope = where(z_slope >= 0.0, slope_l, -slope_l) + return slope + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def limit_vertical_slope_semi_monotonically( + p_cc: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + elev: gtx.int32, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _limit_vertical_slope_semi_monotonically( + p_cc, + z_slope, + k, + elev, + out=z_slope, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py index 33f793cb94..2492b7f8f6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ apply_horizontal_density_increment, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestApplyHorizontalDensityIncrement(helpers.StencilTest): @@ -30,15 +30,15 @@ class TestApplyHorizontalDensityIncrement(helpers.StencilTest): @staticmethod def reference( grid, - p_rhodz_new: xp.array, - p_mflx_contra_v: xp.array, - deepatmo_divzl: xp.array, - deepatmo_divzu: xp.array, + p_rhodz_new: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, p_dtime: float, **kwargs, ) -> dict: tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu - rhodz_ast2 = xp.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp + rhodz_ast2 = np.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp return dict(rhodz_ast2=rhodz_ast2) @pytest.fixture diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py index 778a5db7f0..0526018051 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ apply_interpolated_tracer_time_tendency, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestApplyInterpolatedTracerTimeTendency(helpers.StencilTest): @@ -24,13 +24,13 @@ class TestApplyInterpolatedTracerTimeTendency(helpers.StencilTest): @staticmethod def reference( grid, - p_tracer_now: xp.array, - p_grf_tend_tracer: xp.array, + p_tracer_now: np.array, + p_grf_tend_tracer: np.array, p_dtime, **kwargs, ) -> dict: p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer - p_tracer_new = xp.where(p_tracer_new < 0.0, 0.0, p_tracer_new) + p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) return dict(p_tracer_new=p_tracer_new) @@ -39,7 +39,7 @@ def input_data(self, grid) -> dict: p_tracer_now = helpers.random_field(grid, dims.CellDim, dims.KDim) p_grf_tend_tracer = helpers.random_field(grid, dims.CellDim, dims.KDim) p_tracer_new = helpers.random_field(grid, dims.CellDim, dims.KDim) - p_dtime = xp.float64(5.0) + p_dtime = np.float64(5.0) return dict( p_tracer_now=p_tracer_now, p_grf_tend_tracer=p_grf_tend_tracer, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py index 434ff83c37..dc512a74bd 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ apply_vertical_density_increment, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestApplyVerticalDensityIncrement(helpers.StencilTest): @@ -30,10 +30,10 @@ class TestApplyVerticalDensityIncrement(helpers.StencilTest): @staticmethod def reference( grid, - rhodz_ast: xp.array, - p_mflx_contra_v: xp.array, - deepatmo_divzl: xp.array, - deepatmo_divzu: xp.array, + rhodz_ast: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, p_dtime, **kwargs, ) -> dict: diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py index bd1b5ba5ef..c4b0da2ba6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ average_horizontal_flux_subcycling_2, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestAverageHorizontalFluxSubcycling2(helpers.StencilTest): @@ -24,8 +24,8 @@ class TestAverageHorizontalFluxSubcycling2(helpers.StencilTest): @staticmethod def reference( grid, - z_tracer_mflx_1_dsl: xp.array, - z_tracer_mflx_2_dsl: xp.array, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, **kwargs, ) -> dict: p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py index 2efcd635d7..44873bd564 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ average_horizontal_flux_subcycling_3, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestAverageHorizontalFluxSubcycling3(helpers.StencilTest): @@ -24,9 +24,9 @@ class TestAverageHorizontalFluxSubcycling3(helpers.StencilTest): @staticmethod def reference( grid, - z_tracer_mflx_1_dsl: xp.array, - z_tracer_mflx_2_dsl: xp.array, - z_tracer_mflx_3_dsl: xp.array, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + z_tracer_mflx_3_dsl: np.array, **kwargs, ) -> dict: p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py index 7b0c393770..490b4bb0e9 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_antidiffusive_cell_fluxes_and_min_max, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @@ -30,12 +30,12 @@ class TestComputeAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @staticmethod def reference( grid, - geofac_div: xp.ndarray, - p_rhodz_now: xp.ndarray, - p_rhodz_new: xp.ndarray, - z_mflx_low: xp.ndarray, - z_anti: xp.ndarray, - p_cc: xp.ndarray, + geofac_div: np.ndarray, + p_rhodz_now: np.ndarray, + p_rhodz_new: np.ndarray, + z_mflx_low: np.ndarray, + z_anti: np.ndarray, + p_cc: np.ndarray, p_dtime: float, **kwargs, ) -> dict: @@ -43,31 +43,31 @@ def reference( z_anti_c2e = z_anti[c2e] geofac_div = helpers.reshape(geofac_div, c2e.shape) - geofac_div = xp.expand_dims(geofac_div, axis=-1) + geofac_div = np.expand_dims(geofac_div, axis=-1) - zero_array = xp.zeros(p_rhodz_now.shape) + zero_array = np.zeros(p_rhodz_now.shape) z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] z_mflx_anti_in = -1.0 * ( - xp.minimum(zero_array, z_mflx_anti_1) - + xp.minimum(zero_array, z_mflx_anti_2) - + xp.minimum(zero_array, z_mflx_anti_3) + np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) ) z_mflx_anti_out = ( - xp.maximum(zero_array, z_mflx_anti_1) - + xp.maximum(zero_array, z_mflx_anti_2) - + xp.maximum(zero_array, z_mflx_anti_3) + np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) ) - z_fluxdiv_c = xp.sum(z_mflx_low[c2e] * geofac_div, axis=1) + z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new - z_tracer_max = xp.maximum(p_cc, z_tracer_new_low) - z_tracer_min = xp.minimum(p_cc, z_tracer_new_low) + z_tracer_max = np.maximum(p_cc, z_tracer_new_low) + z_tracer_min = np.minimum(p_cc, z_tracer_new_low) return dict( z_mflx_anti_in=z_mflx_anti_in, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py index bda355f300..ad8d7757a6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_barycentric_backtrajectory, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeBarycentricBacktrajectory(helpers.StencilTest): @@ -24,16 +24,16 @@ class TestComputeBarycentricBacktrajectory(helpers.StencilTest): @staticmethod def reference( grid, - p_vn: xp.array, - p_vt: xp.array, - cell_idx: xp.array, - cell_blk: xp.array, - pos_on_tplane_e_1: xp.array, - pos_on_tplane_e_2: xp.array, - primal_normal_cell_1: xp.array, - dual_normal_cell_1: xp.array, - primal_normal_cell_2: xp.array, - dual_normal_cell_2: xp.array, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + pos_on_tplane_e_1: np.array, + pos_on_tplane_e_2: np.array, + primal_normal_cell_1: np.array, + dual_normal_cell_1: np.array, + primal_normal_cell_2: np.array, + dual_normal_cell_2: np.array, p_dthalf: float, **kwargs, ) -> dict: @@ -48,27 +48,27 @@ def reference( dual_normal_cell_2 = dual_normal_cell_2.reshape(e2c.shape) lvn_pos = p_vn >= 0.0 - cell_idx = xp.expand_dims(cell_idx, axis=-1) - cell_blk = xp.expand_dims(cell_blk, axis=-1) - pos_on_tplane_e_1 = xp.expand_dims(pos_on_tplane_e_1, axis=-1) - pos_on_tplane_e_2 = xp.expand_dims(pos_on_tplane_e_2, axis=-1) - primal_normal_cell_1 = xp.expand_dims(primal_normal_cell_1, axis=-1) - dual_normal_cell_1 = xp.expand_dims(dual_normal_cell_1, axis=-1) - primal_normal_cell_2 = xp.expand_dims(primal_normal_cell_2, axis=-1) - dual_normal_cell_2 = xp.expand_dims(dual_normal_cell_2, axis=-1) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1) + pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1) + primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1) + dual_normal_cell_1 = np.expand_dims(dual_normal_cell_1, axis=-1) + primal_normal_cell_2 = np.expand_dims(primal_normal_cell_2, axis=-1) + dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) - p_cell_idx = xp.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_rel_idx_dsl = xp.where(lvn_pos, 0, 1) - p_cell_blk = xp.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, 0, 1) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) - p_distv_bary_1 = xp.where( + p_distv_bary_1 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], @@ -76,7 +76,7 @@ def reference( + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) - p_distv_bary_2 = xp.where( + p_distv_bary_2 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], @@ -96,7 +96,7 @@ def reference( def input_data(self, grid) -> dict: p_vn = helpers.random_field(grid, dims.EdgeDim, dims.KDim) p_vt = helpers.random_field(grid, dims.EdgeDim, dims.KDim) - cell_idx = xp.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) + cell_idx = np.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) cell_idx_new = helpers.numpy_to_1D_sparse_field(cell_idx, dims.ECDim) cell_blk = helpers.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32) cell_blk_new = helpers.as_1D_sparse_field(cell_blk, dims.ECDim) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py index 6a923c8c6b..ac65277bb2 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_barycentric_backtrajectory_alt, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeBarycentricBacktrajectoryAlt(helpers.StencilTest): @@ -24,14 +24,14 @@ class TestComputeBarycentricBacktrajectoryAlt(helpers.StencilTest): @staticmethod def reference( grid, - p_vn: xp.array, - p_vt: xp.array, - pos_on_tplane_e_1: xp.array, - pos_on_tplane_e_2: xp.array, - primal_normal_cell_1: xp.array, - dual_normal_cell_1: xp.array, - primal_normal_cell_2: xp.array, - dual_normal_cell_2: xp.array, + p_vn: np.array, + p_vt: np.array, + pos_on_tplane_e_1: np.array, + pos_on_tplane_e_2: np.array, + primal_normal_cell_1: np.array, + dual_normal_cell_1: np.array, + primal_normal_cell_2: np.array, + dual_normal_cell_2: np.array, p_dthalf: float, **kwargs, ) -> dict: @@ -44,21 +44,21 @@ def reference( dual_normal_cell_2 = dual_normal_cell_2.reshape(e2c.shape) lvn_pos = p_vn >= 0.0 - pos_on_tplane_e_1 = xp.expand_dims(pos_on_tplane_e_1, axis=-1) - pos_on_tplane_e_2 = xp.expand_dims(pos_on_tplane_e_2, axis=-1) - primal_normal_cell_1 = xp.expand_dims(primal_normal_cell_1, axis=-1) - dual_normal_cell_1 = xp.expand_dims(dual_normal_cell_1, axis=-1) - primal_normal_cell_2 = xp.expand_dims(primal_normal_cell_2, axis=-1) - dual_normal_cell_2 = xp.expand_dims(dual_normal_cell_2, axis=-1) + pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1) + pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1) + primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1) + dual_normal_cell_1 = np.expand_dims(dual_normal_cell_1, axis=-1) + primal_normal_cell_2 = np.expand_dims(primal_normal_cell_2, axis=-1) + dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) - p_distv_bary_1 = xp.where( + p_distv_bary_1 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], @@ -66,7 +66,7 @@ def reference( + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) - p_distv_bary_2 = xp.where( + p_distv_bary_2 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py index b3c81bf201..cb1e3aecc8 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ffsl_backtrajectory, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeFfslBacktrajectory(helpers.StencilTest): @@ -36,23 +36,23 @@ class TestComputeFfslBacktrajectory(helpers.StencilTest): @staticmethod def reference( grid, - p_vn: xp.array, - p_vt: xp.array, - cell_idx: xp.array, - cell_blk: xp.array, - edge_verts_1_x: xp.array, - edge_verts_2_x: xp.array, - edge_verts_1_y: xp.array, - edge_verts_2_y: xp.array, - pos_on_tplane_e_1_x: xp.array, - pos_on_tplane_e_2_x: xp.array, - pos_on_tplane_e_1_y: xp.array, - pos_on_tplane_e_2_y: xp.array, - primal_normal_cell_x: xp.array, - primal_normal_cell_y: xp.array, - dual_normal_cell_x: xp.array, - dual_normal_cell_y: xp.array, - lvn_sys_pos: xp.array, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + edge_verts_1_x: np.array, + edge_verts_2_x: np.array, + edge_verts_1_y: np.array, + edge_verts_2_y: np.array, + pos_on_tplane_e_1_x: np.array, + pos_on_tplane_e_2_x: np.array, + pos_on_tplane_e_1_y: np.array, + pos_on_tplane_e_2_y: np.array, + primal_normal_cell_x: np.array, + primal_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, + lvn_sys_pos: np.array, p_dt: float, **kwargs, ) -> dict: @@ -65,54 +65,54 @@ def reference( dual_normal_cell_y = helpers.reshape(dual_normal_cell_y, e2c_shape) lvn_pos = p_vn >= 0.0 - cell_idx = xp.expand_dims(cell_idx, axis=-1) - cell_blk = xp.expand_dims(cell_blk, axis=-1) - primal_normal_cell_x = xp.expand_dims(primal_normal_cell_x, axis=-1) - dual_normal_cell_x = xp.expand_dims(dual_normal_cell_x, axis=-1) - primal_normal_cell_y = xp.expand_dims(primal_normal_cell_y, axis=-1) - dual_normal_cell_y = xp.expand_dims(dual_normal_cell_y, axis=-1) - edge_verts_1_x = xp.expand_dims(edge_verts_1_x, axis=-1) - edge_verts_1_y = xp.expand_dims(edge_verts_1_y, axis=-1) - edge_verts_2_x = xp.expand_dims(edge_verts_2_x, axis=-1) - edge_verts_2_y = xp.expand_dims(edge_verts_2_y, axis=-1) - pos_on_tplane_e_1_x = xp.expand_dims(pos_on_tplane_e_1_x, axis=-1) - pos_on_tplane_e_1_y = xp.expand_dims(pos_on_tplane_e_1_y, axis=-1) - pos_on_tplane_e_2_x = xp.expand_dims(pos_on_tplane_e_2_x, axis=-1) - pos_on_tplane_e_2_y = xp.expand_dims(pos_on_tplane_e_2_y, axis=-1) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) + dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) + primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) + dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) + edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) + edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) + edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) + edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) + pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) + pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) + pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) + pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) - p_cell_idx = xp.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_blk = xp.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) - p_cell_rel_idx_dsl = xp.where(lvn_pos, 0, 1) + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, 0, 1) - depart_pts_1_x = xp.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt - depart_pts_1_y = xp.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt - depart_pts_2_x = xp.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt - depart_pts_2_y = xp.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt + depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt + depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt + depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt + depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt - pos_on_tplane_e_x = xp.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) - pos_on_tplane_e_y = xp.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y pos_dreg_vert_c_2_x = ( - xp.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x ) pos_dreg_vert_c_2_y = ( - xp.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y ) pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y pos_dreg_vert_c_4_x = ( - xp.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x ) pos_dreg_vert_c_4_y = ( - xp.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y ) - pn_cell_1 = xp.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) - pn_cell_2 = xp.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) - dn_cell_1 = xp.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) - dn_cell_2 = xp.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) + pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) + pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) + dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) + dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) p_coords_dreg_v_1_lon_dsl = ( pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 @@ -156,7 +156,7 @@ def reference( def input_data(self, grid) -> dict: p_vn = helpers.random_field(grid, dims.EdgeDim, dims.KDim) p_vt = helpers.random_field(grid, dims.EdgeDim, dims.KDim) - cell_idx = xp.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) + cell_idx = np.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) cell_idx_new = helpers.numpy_to_1D_sparse_field(cell_idx, dims.ECDim) cell_blk = helpers.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32) cell_blk_new = helpers.as_1D_sparse_field(cell_blk, dims.ECDim) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py index 13e112b2b0..42829df90a 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ffsl_backtrajectory_counterclockwise_indicator, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeFfslBacktrajectoryCounterclockwiseIndicator(helpers.StencilTest): @@ -23,17 +23,17 @@ class TestComputeFfslBacktrajectoryCounterclockwiseIndicator(helpers.StencilTest @staticmethod def reference( - grid, p_vn: xp.array, tangent_orientation: xp.array, lcounterclock: bool, **kwargs + grid, p_vn: np.array, tangent_orientation: np.array, lcounterclock: bool, **kwargs ) -> dict: - tangent_orientation = xp.expand_dims(tangent_orientation, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - tangent_orientation = xp.broadcast_to(tangent_orientation, p_vn.shape) + tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) - lvn_sys_pos_true = xp.where(tangent_orientation * p_vn >= 0.0, True, False) + lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) - mask_lcounterclock = xp.broadcast_to(lcounterclock, p_vn.shape) + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) - lvn_sys_pos = xp.where(mask_lcounterclock, lvn_sys_pos_true, False) + lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) return dict(lvn_sys_pos=lvn_sys_pos) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py index 302cbb0f1c..94efbdf63d 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ffsl_backtrajectory_length_indicator, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeFfslBacktrajectoryLengthIndicator(helpers.StencilTest): @@ -22,18 +22,18 @@ class TestComputeFfslBacktrajectoryLengthIndicator(helpers.StencilTest): OUTPUTS = ("opt_famask_dsl",) @staticmethod - def reference(grid, p_vn: xp.array, p_vt: xp.array, p_dt: float, **kwargs) -> dict: + def reference(grid, p_vn: np.array, p_vt: np.array, p_dt: float, **kwargs) -> dict: lvn_pos = p_vn >= 0.0 - traj_length = xp.sqrt(p_vn**2 + p_vt**2) * p_dt + traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt - edge_cell_length = xp.expand_dims( - xp.asarray(grid.connectivities[dims.E2CDim], dtype=float), axis=-1 + edge_cell_length = np.expand_dims( + np.asarray(grid.connectivities[dims.E2CDim], dtype=float), axis=-1 ) - e2c_length = xp.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) + e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) - opt_famask_dsl = xp.where( - traj_length > (1.25 * xp.broadcast_to(e2c_length, p_vn.shape)), + opt_famask_dsl = np.where( + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), 1, 0, ) @@ -44,7 +44,7 @@ def reference(grid, p_vn: xp.array, p_vt: xp.array, p_dt: float, **kwargs) -> di def input_data(self, grid) -> dict: p_vn = helpers.random_field(grid, dims.EdgeDim, dims.KDim) p_vt = helpers.random_field(grid, dims.EdgeDim, dims.KDim) - edge_cell_length = xp.asarray(grid.connectivities[dims.E2CDim], dtype=float) + edge_cell_length = np.asarray(grid.connectivities[dims.E2CDim], dtype=float) edge_cell_length_new = helpers.numpy_to_1D_sparse_field(edge_cell_length, dims.ECDim) opt_famask_dsl = helpers.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) p_dt = 1.0 diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py index 5b30b04957..9984ebe88c 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_horizontal_tracer_flux_from_cubic_coefficients, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeHorizontalTracerFluxFromCubicCoefficients(helpers.StencilTest): @@ -24,9 +24,9 @@ class TestComputeHorizontalTracerFluxFromCubicCoefficients(helpers.StencilTest): @staticmethod def reference( grid, - p_out_e_hybrid_2: xp.array, - p_mass_flx_e: xp.array, - z_dreg_area: xp.array, + p_out_e_hybrid_2: np.array, + p_mass_flx_e: np.array, + z_dreg_area: np.array, **kwargs, ) -> dict: p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py index f5aa12f44b..c26ba6212e 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -15,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.settings import xp class TestComputeHorizontalTracerFluxFromLinearCoefficients(helpers.StencilTest): @@ -25,14 +25,14 @@ class TestComputeHorizontalTracerFluxFromLinearCoefficients(helpers.StencilTest) @staticmethod def reference( grid, - z_lsq_coeff_1: xp.array, - z_lsq_coeff_2: xp.array, - z_lsq_coeff_3: xp.array, - distv_bary_1: xp.array, - distv_bary_2: xp.array, - p_mass_flx_e: xp.array, - cell_rel_idx_dsl: xp.array, - p_out_e: xp.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + p_out_e: np.array, **kwargs, ) -> dict: p_out_e_cp = p_out_e.copy() @@ -42,11 +42,11 @@ def reference( z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] p_out_e = ( - xp.where(cell_rel_idx_dsl == 1, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + np.where(cell_rel_idx_dsl == 1, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + distv_bary_1 - * xp.where(cell_rel_idx_dsl == 1, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + * np.where(cell_rel_idx_dsl == 1, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + distv_bary_2 - * xp.where(cell_rel_idx_dsl == 1, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) + * np.where(cell_rel_idx_dsl == 1, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) ) * p_mass_flx_e # restriction of execution domain diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py index 6bca883f67..21725ce7e6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -15,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.settings import xp class TestComputeHorizontalTracerFluxFromLinearCoefficientsAlt(helpers.StencilTest): @@ -25,14 +25,14 @@ class TestComputeHorizontalTracerFluxFromLinearCoefficientsAlt(helpers.StencilTe @staticmethod def reference( grid, - z_lsq_coeff_1: xp.array, - z_lsq_coeff_2: xp.array, - z_lsq_coeff_3: xp.array, - distv_bary_1: xp.array, - distv_bary_2: xp.array, - p_mass_flx_e: xp.array, - p_vn: xp.array, - p_out_e: xp.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + p_vn: np.array, + p_out_e: np.array, **kwargs, ) -> dict: p_out_e_cp = p_out_e.copy() @@ -44,9 +44,9 @@ def reference( lvn_pos_inv = p_vn < 0.0 p_out_e = ( - xp.where(lvn_pos_inv, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) - + distv_bary_1 * xp.where(lvn_pos_inv, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) - + distv_bary_2 * xp.where(lvn_pos_inv, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) + np.where(lvn_pos_inv, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + + distv_bary_1 * np.where(lvn_pos_inv, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + + distv_bary_2 * np.where(lvn_pos_inv, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) ) * p_mass_flx_e # restriction of execution domain diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py index 93e9cb70cb..b3d15bfc20 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_positive_definite_horizontal_multiplicative_flux_factor, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePositiveDefiniteHorizontalMultiplicativeFluxFactor(helpers.StencilTest): @@ -24,31 +24,31 @@ class TestComputePositiveDefiniteHorizontalMultiplicativeFluxFactor(helpers.Sten @staticmethod def reference( grid, - geofac_div: xp.ndarray, - p_cc: xp.ndarray, - p_rhodz_now: xp.ndarray, - p_mflx_tracer_h: xp.ndarray, + geofac_div: np.ndarray, + p_cc: np.ndarray, + p_rhodz_now: np.ndarray, + p_mflx_tracer_h: np.ndarray, p_dtime, dbl_eps, **kwargs, ) -> dict: geofac_div = helpers.reshape(geofac_div, grid.connectivities[dims.C2EDim].shape) - geofac_div = xp.expand_dims(geofac_div, axis=-1) - p_m_0 = xp.maximum( + geofac_div = np.expand_dims(geofac_div, axis=-1) + p_m_0 = np.maximum( 0.0, p_mflx_tracer_h[grid.connectivities[dims.C2EDim][:, 0]] * geofac_div[:, 0] * p_dtime, ) - p_m_1 = xp.maximum( + p_m_1 = np.maximum( 0.0, p_mflx_tracer_h[grid.connectivities[dims.C2EDim][:, 1]] * geofac_div[:, 1] * p_dtime, ) - p_m_2 = xp.maximum( + p_m_2 = np.maximum( 0.0, p_mflx_tracer_h[grid.connectivities[dims.C2EDim][:, 2]] * geofac_div[:, 2] * p_dtime, ) p_m = p_m_0 + p_m_1 + p_m_2 - r_m = xp.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) return dict(r_m=r_m) @@ -60,8 +60,8 @@ def input_data(self, grid) -> dict: p_rhodz_now = helpers.random_field(grid, dims.CellDim, dims.KDim) p_mflx_tracer_h = helpers.random_field(grid, dims.EdgeDim, dims.KDim) r_m = helpers.zero_field(grid, dims.CellDim, dims.KDim) - p_dtime = xp.float64(5) - dbl_eps = xp.float64(1e-9) + p_dtime = np.float64(5) + dbl_eps = np.float64(1e-9) return dict( geofac_div=geofac_div_new, p_cc=p_cc, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py index bcf8c0569f..b1455b6994 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ppm4gpu_parabola_coefficients, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpm4gpuParabolaCoefficients(helpers.StencilTest): @@ -23,7 +23,7 @@ class TestComputePpm4gpuParabolaCoefficients(helpers.StencilTest): @staticmethod def reference( - grid, z_face_up: xp.ndarray, z_face_low: xp.ndarray, p_cc: xp.ndarray, **kwargs + grid, z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray, **kwargs ) -> dict: z_delta_q = 0.5 * (z_face_up - z_face_low) z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py index 6483e06466..e6f25381dd 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest from gt4py.next import as_field @@ -15,7 +16,6 @@ compute_ppm_all_face_values, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpmAllFaceValues(helpers.StencilTest): @@ -25,10 +25,10 @@ class TestComputePpmAllFaceValues(helpers.StencilTest): @staticmethod def reference( grid, - p_cc: xp.array, - p_cellhgt_mc_now: xp.array, - p_face_in: xp.array, - k: xp.array, + p_cc: np.array, + p_cellhgt_mc_now: np.array, + p_face_in: np.array, + k: np.array, slev: gtx.int32, elev: gtx.int32, slevp1: gtx.int32, @@ -42,9 +42,9 @@ def reference( (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] ) - p_face = xp.where((k == slevp1) | (k == elev), p_face_a, p_face_in) - p_face = xp.where((k == slev), p_cc, p_face) - p_face[:, 1:] = xp.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) + p_face = np.where((k == slevp1) | (k == elev), p_face_a, p_face_in) + p_face = np.where((k == slev), p_cc, p_face) + p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) return dict(p_face=p_face) @pytest.fixture @@ -55,7 +55,7 @@ def input_data(self, grid) -> dict: p_face = helpers.zero_field(grid, dims.CellDim, dims.KDim) k = as_field( - (dims.KDim,), xp.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) + (dims.KDim,), np.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) ) slev = gtx.int32(1) slevp1 = gtx.int32(2) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py index f7aef94aaf..7ed3c06254 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ppm_quadratic_face_values, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp outslice = (slice(None), slice(1, None)) @@ -25,7 +25,7 @@ class TestComputePpmQuadraticFaceValues(helpers.StencilTest): OUTPUTS = (helpers.Output("p_face", refslice=outslice, gtslice=outslice),) @staticmethod - def reference(grid, p_cc: xp.array, p_cellhgt_mc_now: xp.array, **kwargs) -> dict: + def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs) -> dict: p_face = p_cc.copy() p_face[:, 1:] = p_cc[:, 1:] * ( 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py index 8be8b6ec73..dc65a506ce 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ppm_quartic_face_values, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpmQuarticFaceValues(helpers.StencilTest): @@ -23,7 +23,7 @@ class TestComputePpmQuarticFaceValues(helpers.StencilTest): @staticmethod def reference( - grid, p_cc: xp.array, p_cellhgt_mc_now: xp.array, z_slope: xp.array, **kwargs + grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.array, **kwargs ) -> dict: p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py index b9139805ef..f1b920944f 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest from gt4py.next import as_field @@ -15,7 +16,6 @@ compute_ppm_slope, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpmSlope(helpers.StencilTest): @@ -28,7 +28,7 @@ class TestComputePpmSlope(helpers.StencilTest): @staticmethod def reference( - grid, p_cc: xp.array, p_cellhgt_mc_now: xp.array, k: xp.array, elev: gtx.int32, **kwargs + grid, p_cc: np.array, p_cellhgt_mc_now: np.array, k: np.array, elev: gtx.int32, **kwargs ) -> dict: zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] @@ -56,7 +56,7 @@ def reference( + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 ) - z_slope = xp.where(k[1:-1] < elev, z_slope_a, z_slope_b) + z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) return dict(z_slope=z_slope) @pytest.fixture @@ -68,7 +68,7 @@ def input_data(self, grid) -> dict: ) k = as_field( (dims.KDim,), - xp.arange( + np.arange( 0, helpers._shape(grid, dims.KDim, extend={dims.KDim: 1})[0], dtype=gtx.int32 ), ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py index 30fd8f84f5..96179b99e6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_tendency, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeTendency(helpers.StencilTest): @@ -24,8 +24,8 @@ class TestComputeTendency(helpers.StencilTest): @staticmethod def reference( grid, - p_tracer_now: xp.array, - p_tracer_new: xp.array, + p_tracer_now: np.array, + p_tracer_new: np.array, p_dtime, **kwargs, ) -> dict: @@ -38,7 +38,7 @@ def input_data(self, grid) -> dict: p_tracer_now = helpers.random_field(grid, dims.CellDim, dims.KDim) p_tracer_new = helpers.random_field(grid, dims.CellDim, dims.KDim) opt_ddt_tracer_adv = helpers.zero_field(grid, dims.CellDim, dims.KDim) - p_dtime = xp.float64(5.0) + p_dtime = np.float64(5.0) return dict( p_tracer_now=p_tracer_now, p_tracer_new=p_tracer_new, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py index 56d27c46fb..54e87d6f33 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_vertical_parabola_limiter_condition, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeVerticalParabolaLimiterCondition(helpers.StencilTest): @@ -22,10 +22,10 @@ class TestComputeVerticalParabolaLimiterCondition(helpers.StencilTest): OUTPUTS = ("l_limit",) @staticmethod - def reference(grid, p_face: xp.array, p_cc: xp.array, **kwargs) -> dict: + def reference(grid, p_face: np.array, p_cc: np.array, **kwargs) -> dict: z_delta = p_face[:, :-1] - p_face[:, 1:] z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - l_limit = xp.where(xp.abs(z_delta) < -1 * z_a6i, 1, 0) + l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, 1, 0) return dict(l_limit=l_limit) @pytest.fixture diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py new file mode 100644 index 0000000000..f0e11e9a63 --- /dev/null +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py @@ -0,0 +1,58 @@ +# 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 +import numpy as np +import pytest + +import icon4py.model.common.test_utils.helpers as helpers +from icon4py.model.atmosphere.advection.stencils.compute_vertical_tracer_flux_upwind import ( + compute_vertical_tracer_flux_upwind, +) +from icon4py.model.common import dimension as dims + + +outslice = (slice(None), slice(1, None)) + + +class TestComputeVerticalTracerFluxUpwind(helpers.StencilTest): + PROGRAM = compute_vertical_tracer_flux_upwind + OUTPUTS = (helpers.Output("p_upflux", refslice=outslice, gtslice=outslice),) + + @staticmethod + def reference( + grid, + p_cc: np.array, + p_mflx_contra_v: np.array, + **kwargs, + ) -> dict: + p_upflux = p_cc.copy() + p_upflux[:, 1:] = ( + np.where(p_mflx_contra_v[:, 1:] >= 0.0, p_cc[:, 1:], p_cc[:, :-1]) + * p_mflx_contra_v[:, 1:] + ) + return dict(p_upflux=p_upflux) + + @pytest.fixture + def input_data(self, grid) -> dict: + p_cc = helpers.random_field(grid, dims.CellDim, dims.KDim) + p_mflx_contra_v = helpers.random_field( + grid, dims.CellDim, dims.KDim + ) # TODO (dastrm): should be KHalfDim + p_upflux = helpers.zero_field( + grid, dims.CellDim, dims.KDim + ) # TODO (dastrm): should be KHalfDim + return dict( + p_cc=p_cc, + p_mflx_contra_v=p_mflx_contra_v, + p_upflux=p_upflux, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=1, + vertical_end=gtx.int32(grid.num_levels), + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py index feb60a15e7..9ad0a6526b 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ integrate_tracer_density_horizontally, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestIntegrateTracerDensityHorizontally(helpers.StencilTest): @@ -29,25 +29,25 @@ class TestIntegrateTracerDensityHorizontally(helpers.StencilTest): @staticmethod def reference( grid, - p_mass_flx_e: xp.array, - geofac_div: xp.array, - z_rhofluxdiv_c: xp.array, - z_tracer_mflx: xp.array, - z_rho_now: xp.array, - z_tracer_now: xp.array, + p_mass_flx_e: np.array, + geofac_div: np.array, + z_rhofluxdiv_c: np.array, + z_tracer_mflx: np.array, + z_rho_now: np.array, + z_tracer_now: np.array, z_dtsub: float, nsub: gtx.int32, **kwargs, ) -> dict: c2e = grid.connectivities[dims.C2EDim] p_mass_flx_e_c2e = p_mass_flx_e[c2e] - geofac_div = xp.expand_dims(geofac_div, axis=-1) + geofac_div = np.expand_dims(geofac_div, axis=-1) z_tracer_mflx_c2e = z_tracer_mflx[c2e] z_rhofluxdiv_c_out = ( - xp.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == 1 else z_rhofluxdiv_c + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == 1 else z_rhofluxdiv_c ) - z_fluxdiv_c_dsl = xp.sum(z_tracer_mflx_c2e * geofac_div, axis=1) + z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py index a71ede7b21..f270f55c83 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ integrate_tracer_horizontally, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestIntegrateTracerHorizontally(helpers.StencilTest): @@ -24,22 +24,22 @@ class TestIntegrateTracerHorizontally(helpers.StencilTest): @staticmethod def reference( grid, - p_mflx_tracer_h: xp.array, - deepatmo_divh: xp.array, - tracer_now: xp.array, - rhodz_now: xp.array, - rhodz_new: xp.array, - geofac_div: xp.array, + p_mflx_tracer_h: np.array, + deepatmo_divh: np.array, + tracer_now: np.array, + rhodz_now: np.array, + rhodz_new: np.array, + geofac_div: np.array, p_dtime, **kwargs, ) -> dict: geofac_div = helpers.reshape(geofac_div, grid.connectivities[dims.C2EDim].shape) - geofac_div = xp.expand_dims(geofac_div, axis=-1) + geofac_div = np.expand_dims(geofac_div, axis=-1) tracer_new_hor = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh - * xp.sum(p_mflx_tracer_h[grid.connectivities[dims.C2EDim]] * geofac_div, axis=1) + * np.sum(p_mflx_tracer_h[grid.connectivities[dims.C2EDim]] * geofac_div, axis=1) ) / rhodz_new return dict(tracer_new_hor=tracer_new_hor) @@ -52,7 +52,7 @@ def input_data(self, grid) -> dict: rhodz_new = helpers.random_field(grid, dims.CellDim, dims.KDim) geofac_div = helpers.random_field(grid, dims.CellDim, dims.C2EDim) geofac_div_new = helpers.as_1D_sparse_field(geofac_div, dims.CEDim) - p_dtime = xp.float64(5.0) + p_dtime = np.float64(5.0) tracer_new_hor = helpers.zero_field(grid, dims.CellDim, dims.KDim) return dict( p_mflx_tracer_h=p_mflx_tracer_h, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py index 82e4f2ff5c..07ca62df76 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest from gt4py.next import as_field @@ -15,7 +16,6 @@ integrate_tracer_vertically, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestIntegrateTracerVertically(helpers.StencilTest): @@ -25,20 +25,20 @@ class TestIntegrateTracerVertically(helpers.StencilTest): @staticmethod def reference( grid, - tracer_now: xp.array, - rhodz_now: xp.array, - p_mflx_tracer_v: xp.array, - deepatmo_divzl: xp.array, - deepatmo_divzu: xp.array, - rhodz_new: xp.array, - k: xp.array, + tracer_now: np.array, + rhodz_now: np.array, + p_mflx_tracer_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + rhodz_new: np.array, + k: np.array, ivadv_tracer: gtx.int32, iadv_slev_jt: gtx.int32, p_dtime: float, **kwargs, ) -> dict: if ivadv_tracer != 0: - tracer_new = xp.where( + tracer_new = np.where( (iadv_slev_jt <= k), ( tracer_now * rhodz_now @@ -65,8 +65,8 @@ def input_data(self, grid) -> dict: deepatmo_divzu = helpers.random_field(grid, dims.KDim) rhodz_new = helpers.random_field(grid, dims.CellDim, dims.KDim) tracer_new = helpers.zero_field(grid, dims.CellDim, dims.KDim) - k = as_field((dims.KDim,), xp.arange(grid.num_levels, dtype=gtx.int32)) - p_dtime = xp.float64(5.0) + k = as_field((dims.KDim,), np.arange(grid.num_levels, dtype=gtx.int32)) + p_dtime = np.float64(5.0) ivadv_tracer = 1 iadv_slev_jt = 4 return dict( diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py index c587c51c9b..e1fc0066f9 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ limit_vertical_parabola_semi_monotonically, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestLimitVerticalParabolaSemiMonotonically(helpers.StencilTest): @@ -22,13 +22,13 @@ class TestLimitVerticalParabolaSemiMonotonically(helpers.StencilTest): OUTPUTS = ("p_face_up", "p_face_low") @staticmethod - def reference(grid, l_limit: xp.array, p_face: xp.array, p_cc: xp.array, **kwargs) -> dict: - q_face_up, q_face_low = xp.where( + def reference(grid, l_limit: np.array, p_face: np.array, p_cc: np.array, **kwargs) -> dict: + q_face_up, q_face_low = np.where( l_limit != 0, - xp.where( - (p_cc < xp.minimum(p_face[:, :-1], p_face[:, 1:])), + np.where( + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), (p_cc, p_cc), - xp.where( + np.where( p_face[:, :-1] > p_face[:, 1:], (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py new file mode 100644 index 0000000000..5db041d352 --- /dev/null +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py @@ -0,0 +1,52 @@ +# 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 +import numpy as np +import pytest +from gt4py.next import as_field + +import icon4py.model.common.test_utils.helpers as helpers +from icon4py.model.atmosphere.advection.stencils.limit_vertical_slope_semi_monotonically import ( + limit_vertical_slope_semi_monotonically, +) +from icon4py.model.common import dimension as dims + + +class TestLimitVerticalSlopeSemiMonotonically(helpers.StencilTest): + PROGRAM = limit_vertical_slope_semi_monotonically + OUTPUTS = (helpers.Output("z_slope", gtslice=(slice(None), slice(1, -1))),) + + @staticmethod + def reference( + grid, p_cc: np.array, z_slope: np.array, k: np.array, elev: gtx.int32, **kwargs + ) -> dict: + p_cc_min_last = np.minimum(p_cc[:, :-2], p_cc[:, 1:-1]) + p_cc_min = np.where(k[1:-1] == elev, p_cc_min_last, np.minimum(p_cc_min_last, p_cc[:, 2:])) + slope_l = np.minimum(np.abs(z_slope[:, 1:-1]), 2.0 * (p_cc[:, 1:-1] - p_cc_min)) + slope = np.where(z_slope[:, 1:-1] >= 0.0, slope_l, -slope_l) + return dict(z_slope=slope) + + @pytest.fixture + def input_data(self, grid) -> dict: + p_cc = helpers.random_field(grid, dims.CellDim, dims.KDim) + z_slope = helpers.random_field(grid, dims.CellDim, dims.KDim) + k = as_field( + (dims.KDim,), np.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) + ) + elev = k[-2].as_scalar() + return dict( + p_cc=p_cc, + z_slope=z_slope, + k=k, + elev=elev, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=1, + vertical_end=gtx.int32(grid.num_levels - 1), + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py index 3637ffa74e..fd7991126e 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ postprocess_antidiffusive_cell_fluxes_and_min_max, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestPostprocessAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @@ -24,27 +24,27 @@ class TestPostprocessAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @staticmethod def reference( grid, - refin_ctrl: xp.ndarray, - p_cc: xp.ndarray, - z_tracer_new_low: xp.ndarray, - z_tracer_max: xp.ndarray, - z_tracer_min: xp.ndarray, + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, lo_bound: float, hi_bound: float, **kwargs, ) -> dict: - refin_ctrl = xp.expand_dims(refin_ctrl, axis=1) - condition = xp.logical_or( - xp.equal(refin_ctrl, lo_bound * xp.ones(refin_ctrl.shape, dtype=gtx.int32)), - xp.equal(refin_ctrl, hi_bound * xp.ones(refin_ctrl.shape, dtype=gtx.int32)), + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=gtx.int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=gtx.int32)), ) - z_tracer_new_out = xp.where( + z_tracer_new_out = np.where( condition, - xp.minimum(1.1 * p_cc, xp.maximum(0.9 * p_cc, z_tracer_new_low)), + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), z_tracer_new_low, ) - z_tracer_max_out = xp.where(condition, xp.maximum(p_cc, z_tracer_new_out), z_tracer_max) - z_tracer_min_out = xp.where(condition, xp.minimum(p_cc, z_tracer_new_out), z_tracer_min) + z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) return dict( z_tracer_new_low=z_tracer_new_out, z_tracer_max=z_tracer_max_out, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py index 5ce482bd96..143aaaba60 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ prepare_ffsl_flux_area_patches_list, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp # Check whether lines inters. @@ -35,8 +35,8 @@ def _ccw_numpy( dx1dy2 = dx1 * dy2 dy1dx2 = dy1 * dx2 - lccw = xp.where(dx1dy2 > dy1dx2, True, False) - ccw_out = xp.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise return ccw_out @@ -81,7 +81,7 @@ def _lintersect_numpy( line1_p2_lon, line1_p2_lat, ) - lintersect_out = xp.where((intersect1 + intersect2) == -2, True, False) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) return lintersect_out @@ -98,10 +98,10 @@ def _line_intersect_numpy( line2_p2_lat, ): d1 = line1_p2_lon - line1_p1_lon - d1 = xp.where(d1 != 0.0, d1, line1_p2_lon) + d1 = np.where(d1 != 0.0, d1, line1_p2_lon) d2 = line2_p2_lon - line2_p1_lon - d2 = xp.where(d2 != 0.0, d2, line2_p2_lon) + d2 = np.where(d2 != 0.0, d2, line2_p2_lon) m1 = (line1_p2_lat - line1_p1_lat) / d1 m2 = (line2_p2_lat - line2_p1_lat) / d2 @@ -174,28 +174,28 @@ def _generate_flux_area_geometry( tri_line1_p1_lon = arrival_pts_1_lon_dsl tri_line1_p1_lat = arrival_pts_1_lat_dsl - tri_line1_p2_lon = xp.where( + tri_line1_p2_lon = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), ) - tri_line1_p2_lat = xp.where( + tri_line1_p2_lat = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), ) tri_line2_p1_lon = arrival_pts_2_lon_dsl tri_line2_p1_lat = arrival_pts_2_lat_dsl - tri_line2_p2_lon = xp.where( + tri_line2_p2_lon = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), ) - tri_line2_p2_lat = xp.where( + tri_line2_p2_lat = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), ) return ( @@ -240,26 +240,26 @@ def _apply_case1_patch0( ): dreg_patch0_1_lon_dsl = arrival_pts_1_lon_dsl dreg_patch0_1_lat_dsl = arrival_pts_1_lat_dsl - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_2_lon_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), arrival_pts_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), arrival_pts_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), depart_pts_1_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), depart_pts_1_lat_dsl, ) @@ -285,21 +285,21 @@ def _apply_case1_patch1( ps1_x, ps1_y, ): - dreg_patch1_1_lon_vmask = xp.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_1_lat_vmask = xp.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_4_lon_vmask = xp.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_4_lat_vmask = xp.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 ) - dreg_patch1_2_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 ) - dreg_patch1_3_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 ) - dreg_patch1_3_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 ) return ( @@ -325,21 +325,21 @@ def _apply_case1_patch2( ps2_y, ): # Case 1 - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_1_lat_vmask = xp.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_4_lon_vmask = xp.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_4_lat_vmask = xp.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_2_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 ) - dreg_patch2_2_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + dreg_patch2_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 ) - dreg_patch2_3_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + dreg_patch2_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 ) - dreg_patch2_3_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + dreg_patch2_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 ) return ( @@ -374,28 +374,28 @@ def _apply_case2a_patch0( dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) @@ -429,36 +429,36 @@ def _apply_case2a_patch1( dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, ): - dreg_patch1_1_lon_vmask = xp.where( + dreg_patch1_1_lon_vmask = np.where( mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask ) - dreg_patch1_1_lat_vmask = xp.where( + dreg_patch1_1_lat_vmask = np.where( mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask ) - dreg_patch1_4_lon_vmask = xp.where( + dreg_patch1_4_lon_vmask = np.where( mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask ) - dreg_patch1_4_lat_vmask = xp.where( + dreg_patch1_4_lat_vmask = np.where( mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask ) - dreg_patch1_2_lon_vmask = xp.where( + dreg_patch1_2_lon_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), dreg_patch1_2_lon_vmask, ) - dreg_patch1_2_lat_vmask = xp.where( + dreg_patch1_2_lat_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask, ) - dreg_patch1_3_lon_vmask = xp.where( + dreg_patch1_3_lon_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask, ) - dreg_patch1_3_lat_vmask = xp.where( + dreg_patch1_3_lat_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask, ) @@ -494,28 +494,28 @@ def _apply_case2b_patch0( dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) @@ -542,16 +542,16 @@ def _apply_case2b_patch1( dreg_patch1_4_lon_vmask, dreg_patch1_4_lat_vmask, ): - zeros_array = xp.zeros_like(mask_case2b) + zeros_array = np.zeros_like(mask_case2b) - dreg_patch1_1_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) return ( dreg_patch1_1_lon_vmask, @@ -583,36 +583,36 @@ def _apply_case2b_patch2( dreg_patch2_3_lon_vmask, dreg_patch2_3_lat_vmask, ): - dreg_patch2_1_lon_vmask = xp.where( + dreg_patch2_1_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) - dreg_patch2_1_lat_vmask = xp.where( + dreg_patch2_1_lat_vmask = np.where( mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask ) - dreg_patch2_4_lon_vmask = xp.where( + dreg_patch2_4_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask ) - dreg_patch2_4_lat_vmask = xp.where( + dreg_patch2_4_lat_vmask = np.where( mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask ) - dreg_patch2_2_lon_vmask = xp.where( + dreg_patch2_2_lon_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), dreg_patch2_2_lon_vmask, ) - dreg_patch2_2_lat_vmask = xp.where( + dreg_patch2_2_lat_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), dreg_patch2_2_lat_vmask, ) - dreg_patch2_3_lon_vmask = xp.where( + dreg_patch2_3_lon_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), dreg_patch2_3_lon_vmask, ) - dreg_patch2_3_lat_vmask = xp.where( + dreg_patch2_3_lat_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), dreg_patch2_3_lat_vmask, ) @@ -648,28 +648,28 @@ def _apply_case3a_patch0( dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) @@ -705,36 +705,36 @@ def _apply_case3a_patch1( dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, ): - dreg_patch1_1_lon_vmask = xp.where( + dreg_patch1_1_lon_vmask = np.where( mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask ) - dreg_patch1_1_lat_vmask = xp.where( + dreg_patch1_1_lat_vmask = np.where( mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask ) - dreg_patch1_2_lon_vmask = xp.where( + dreg_patch1_2_lon_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), dreg_patch1_2_lon_vmask, ) - dreg_patch1_2_lat_vmask = xp.where( + dreg_patch1_2_lat_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask, ) - dreg_patch1_3_lon_vmask = xp.where( + dreg_patch1_3_lon_vmask = np.where( mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask ) - dreg_patch1_3_lat_vmask = xp.where( + dreg_patch1_3_lat_vmask = np.where( mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask ) - dreg_patch1_4_lon_vmask = xp.where( + dreg_patch1_4_lon_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), dreg_patch1_4_lon_vmask, ) - dreg_patch1_4_lat_vmask = xp.where( + dreg_patch1_4_lat_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), dreg_patch1_4_lat_vmask, ) @@ -768,28 +768,28 @@ def _apply_case3b_patch0( dreg_patch0_3_lon_dsl, dreg_patch0_3_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = xp.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), dreg_patch0_3_lon_dsl, ) - dreg_patch0_3_lat_dsl = xp.where( + dreg_patch0_3_lat_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), dreg_patch0_3_lat_dsl, ) @@ -825,36 +825,36 @@ def _apply_case3b_patch2( dreg_patch2_4_lon_vmask, dreg_patch2_4_lat_vmask, ): - dreg_patch2_1_lon_vmask = xp.where( + dreg_patch2_1_lon_vmask = np.where( mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) - dreg_patch2_1_lat_vmask = xp.where( + dreg_patch2_1_lat_vmask = np.where( mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask ) - dreg_patch2_2_lon_vmask = xp.where( + dreg_patch2_2_lon_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), dreg_patch2_2_lon_vmask, ) - dreg_patch2_2_lat_vmask = xp.where( + dreg_patch2_2_lat_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask, ) - dreg_patch2_3_lon_vmask = xp.where( + dreg_patch2_3_lon_vmask = np.where( mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask ) - dreg_patch2_3_lat_vmask = xp.where( + dreg_patch2_3_lat_vmask = np.where( mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask ) - dreg_patch2_4_lon_vmask = xp.where( + dreg_patch2_4_lon_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), dreg_patch2_4_lon_vmask, ) - dreg_patch2_4_lat_vmask = xp.where( + dreg_patch2_4_lat_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), dreg_patch2_4_lat_vmask, ) @@ -890,10 +890,10 @@ def reference( ) -> dict: e2c = grid.connectivities[dims.E2CDim] ptr_v3_lon = helpers.reshape(ptr_v3_lon, e2c.shape) - ptr_v3_lon_e = xp.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) ptr_v3_lat = helpers.reshape(ptr_v3_lat, e2c.shape) - ptr_v3_lat_e = xp.expand_dims(ptr_v3_lat, axis=-1) - tangent_orientation_dsl = xp.expand_dims(tangent_orientation_dsl, axis=-1) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) result_tuple = cls._generate_flux_area_geometry( dreg_patch0_1_lon_dsl, @@ -955,13 +955,13 @@ def reference( tri_line2_p2_lat, ) - lvn_sys_pos = xp.where( - (p_vn * xp.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, + lvn_sys_pos = np.where( + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, True, False, ) - famask_bool = xp.where(famask_int == 1, True, False) - mask_case1 = xp.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + famask_bool = np.where(famask_int == 1, True, False) + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) ps1_x, ps1_y = _line_intersect_numpy( fl_line_p1_lon, fl_line_p1_lat, @@ -1052,8 +1052,8 @@ def reference( ) # ------------------------------------------------- Case 2a - mask_case2a = xp.logical_and.reduce( - [lintersect_line1, xp.logical_not(lintersect_line2), famask_bool] + mask_case2a = np.logical_and.reduce( + [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] ) # Case 2a - patch 0 result_tuple_patch0 = cls._apply_case2a_patch0( @@ -1118,18 +1118,18 @@ def reference( dreg_patch1_3_lat_vmask, ) = result_tuple_patch1 # Case 2a - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) # -------------------------------------------------- Case 2b - mask_case2b = xp.logical_and.reduce( - [lintersect_line2, xp.logical_not(lintersect_line1), famask_bool] + mask_case2b = np.logical_and.reduce( + [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] ) # Case 2b - patch 0 result_tuple_patch0_case2b = cls._apply_case2b_patch0( @@ -1241,7 +1241,7 @@ def reference( tri_line1_p2_lon, tri_line1_p2_lat, ) - mask_case3a = xp.logical_and(lintersect_e2_line1, famask_bool) + mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) pi1_x, pi1_y = _line_intersect_numpy( fl_e2_p1_lon, fl_e2_p1_lat, @@ -1316,14 +1316,14 @@ def reference( dreg_patch1_3_lat_vmask, ) # Case 3a - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) # ------------------------------------------------ Case 3b # Check whether flux area edge 1 intersects with triangle edge 2 @@ -1377,14 +1377,14 @@ def reference( dreg_patch0_3_lat_dsl, ) # Case 3b - patch 1 - dreg_patch1_1_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 ( @@ -1419,31 +1419,31 @@ def reference( # --------------------------------------------- Case 4 # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions # can't be overwritten by this new condition. - indices_previously_matched = xp.logical_or.reduce( + indices_previously_matched = np.logical_or.reduce( [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] ) # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices - mask_case4 = xp.logical_and.reduce( - [famask_bool, xp.logical_not(indices_previously_matched)] + mask_case4 = np.logical_and.reduce( + [famask_bool, np.logical_not(indices_previously_matched)] ) # Case 4 - patch 0 - no change # Case 4 - patch 1 - dreg_patch1_1_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) # Case 4 - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) return dict( dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py index 2e19beac72..08b651a24b 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ prepare_numerical_quadrature_for_cubic_reconstruction, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp @pytest.mark.slow_tests @@ -337,14 +337,14 @@ def _compute_vector_sums( def reference( cls, grid, - p_coords_dreg_v_1_x: xp.array, - p_coords_dreg_v_2_x: xp.array, - p_coords_dreg_v_3_x: xp.array, - p_coords_dreg_v_4_x: xp.array, - p_coords_dreg_v_1_y: xp.array, - p_coords_dreg_v_2_y: xp.array, - p_coords_dreg_v_3_y: xp.array, - p_coords_dreg_v_4_y: xp.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, shape_func_1_1: float, shape_func_2_1: float, shape_func_3_1: float, @@ -464,10 +464,10 @@ def reference( ) z_area = p_quad_vector_sum_1 - p_dreg_area_out = xp.where( + p_dreg_area_out = np.where( z_area >= 0.0, - xp.maximum(eps, xp.absolute(z_area)), - -xp.maximum(eps, xp.absolute(z_area)), + np.maximum(eps, np.absolute(z_area)), + -np.maximum(eps, np.absolute(z_area)), ) return dict( p_quad_vector_sum_1=p_quad_vector_sum_1, @@ -532,7 +532,7 @@ def input_data(self, grid) -> dict: wgt_zeta_2 = 0.003 wgt_eta_1 = 0.002 wgt_eta_2 = 0.007 - dbl_eps = xp.float64(0.1) + dbl_eps = np.float64(0.1) eps = 0.1 return dict( p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py index 9191ed3b29..4b78309142 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ prepare_numerical_quadrature_list_for_cubic_reconstruction, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp @pytest.mark.slow_tests @@ -89,18 +89,18 @@ def _compute_wgt_t_detjac( 1.0 + zeta_4, ) - famask_bool = xp.where(famask_int == 1, True, False) + famask_bool = np.where(famask_int == 1, True, False) - p_coords_dreg_v_1_x = xp.where(famask_bool, p_coords_dreg_v_1_x, 0.0) - p_coords_dreg_v_2_x = xp.where(famask_bool, p_coords_dreg_v_2_x, 0.0) - p_coords_dreg_v_3_x = xp.where(famask_bool, p_coords_dreg_v_3_x, 0.0) - p_coords_dreg_v_4_x = xp.where(famask_bool, p_coords_dreg_v_4_x, 0.0) - p_coords_dreg_v_1_y = xp.where(famask_bool, p_coords_dreg_v_1_y, 0.0) - p_coords_dreg_v_2_y = xp.where(famask_bool, p_coords_dreg_v_2_y, 0.0) - p_coords_dreg_v_3_y = xp.where(famask_bool, p_coords_dreg_v_3_y, 0.0) - p_coords_dreg_v_4_y = xp.where(famask_bool, p_coords_dreg_v_4_y, 0.0) + p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) - wgt_t_detjac_1 = xp.where( + wgt_t_detjac_1 = np.where( famask_bool, dbl_eps + z_wgt_1 @@ -125,7 +125,7 @@ def _compute_wgt_t_detjac( 0.0, ) - wgt_t_detjac_2 = xp.where( + wgt_t_detjac_2 = np.where( famask_bool, dbl_eps + z_wgt_2 @@ -150,7 +150,7 @@ def _compute_wgt_t_detjac( 0.0, ) - wgt_t_detjac_3 = xp.where( + wgt_t_detjac_3 = np.where( famask_bool, dbl_eps + z_wgt_3 @@ -174,7 +174,7 @@ def _compute_wgt_t_detjac( ), 0.0, ) - wgt_t_detjac_4 = xp.where( + wgt_t_detjac_4 = np.where( famask_bool, dbl_eps + z_wgt_4 @@ -374,16 +374,16 @@ def _compute_vector_sums( def reference( cls, grid, - famask_int: xp.array, - p_coords_dreg_v_1_x: xp.array, - p_coords_dreg_v_2_x: xp.array, - p_coords_dreg_v_3_x: xp.array, - p_coords_dreg_v_4_x: xp.array, - p_coords_dreg_v_1_y: xp.array, - p_coords_dreg_v_2_y: xp.array, - p_coords_dreg_v_3_y: xp.array, - p_coords_dreg_v_4_y: xp.array, - p_dreg_area_in: xp.array, + famask_int: np.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + p_dreg_area_in: np.array, shape_func_1_1: float, shape_func_2_1: float, shape_func_3_1: float, @@ -569,7 +569,7 @@ def input_data(self, grid) -> dict: wgt_zeta_2 = 0.003 wgt_eta_1 = 0.002 wgt_eta_2 = 0.007 - dbl_eps = xp.float64(0.1) + dbl_eps = np.float64(0.1) eps = 0.1 return dict( famask_int=famask_int, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py index 4bf6ad8f73..b4bc666c47 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -15,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.settings import xp @pytest.mark.slow_tests @@ -37,35 +37,35 @@ class TestReconstructCubicCoefficientsSvd(helpers.StencilTest): @staticmethod def reference( grid, - p_cc: xp.array, - lsq_pseudoinv_1: xp.array, - lsq_pseudoinv_2: xp.array, - lsq_pseudoinv_3: xp.array, - lsq_pseudoinv_4: xp.array, - lsq_pseudoinv_5: xp.array, - lsq_pseudoinv_6: xp.array, - lsq_pseudoinv_7: xp.array, - lsq_pseudoinv_8: xp.array, - lsq_pseudoinv_9: xp.array, - lsq_moments_1: xp.array, - lsq_moments_2: xp.array, - lsq_moments_3: xp.array, - lsq_moments_4: xp.array, - lsq_moments_5: xp.array, - lsq_moments_6: xp.array, - lsq_moments_7: xp.array, - lsq_moments_8: xp.array, - lsq_moments_9: xp.array, - p_coeff_1_dsl: xp.array, - p_coeff_2_dsl: xp.array, - p_coeff_3_dsl: xp.array, - p_coeff_4_dsl: xp.array, - p_coeff_5_dsl: xp.array, - p_coeff_6_dsl: xp.array, - p_coeff_7_dsl: xp.array, - p_coeff_8_dsl: xp.array, - p_coeff_9_dsl: xp.array, - p_coeff_10_dsl: xp.array, + p_cc: np.array, + lsq_pseudoinv_1: np.array, + lsq_pseudoinv_2: np.array, + lsq_pseudoinv_3: np.array, + lsq_pseudoinv_4: np.array, + lsq_pseudoinv_5: np.array, + lsq_pseudoinv_6: np.array, + lsq_pseudoinv_7: np.array, + lsq_pseudoinv_8: np.array, + lsq_pseudoinv_9: np.array, + lsq_moments_1: np.array, + lsq_moments_2: np.array, + lsq_moments_3: np.array, + lsq_moments_4: np.array, + lsq_moments_5: np.array, + lsq_moments_6: np.array, + lsq_moments_7: np.array, + lsq_moments_8: np.array, + lsq_moments_9: np.array, + p_coeff_1_dsl: np.array, + p_coeff_2_dsl: np.array, + p_coeff_3_dsl: np.array, + p_coeff_4_dsl: np.array, + p_coeff_5_dsl: np.array, + p_coeff_6_dsl: np.array, + p_coeff_7_dsl: np.array, + p_coeff_8_dsl: np.array, + p_coeff_9_dsl: np.array, + p_coeff_10_dsl: np.array, **kwargs, ) -> dict: p_coeff_1_dsl_cp = p_coeff_1_dsl.copy() @@ -80,41 +80,41 @@ def reference( p_coeff_10_dsl_cp = p_coeff_10_dsl.copy() c2e2c2e2c = grid.connectivities[dims.C2E2C2E2CDim] - lsq_moments_1 = xp.expand_dims(lsq_moments_1, axis=-1) - lsq_moments_2 = xp.expand_dims(lsq_moments_2, axis=-1) - lsq_moments_3 = xp.expand_dims(lsq_moments_3, axis=-1) - lsq_moments_4 = xp.expand_dims(lsq_moments_4, axis=-1) - lsq_moments_5 = xp.expand_dims(lsq_moments_5, axis=-1) - lsq_moments_6 = xp.expand_dims(lsq_moments_6, axis=-1) - lsq_moments_7 = xp.expand_dims(lsq_moments_7, axis=-1) - lsq_moments_8 = xp.expand_dims(lsq_moments_8, axis=-1) - lsq_moments_9 = xp.expand_dims(lsq_moments_9, axis=-1) - lsq_moments_1 = xp.broadcast_to(lsq_moments_1, p_cc.shape) - lsq_moments_2 = xp.broadcast_to(lsq_moments_2, p_cc.shape) - lsq_moments_3 = xp.broadcast_to(lsq_moments_3, p_cc.shape) - lsq_moments_4 = xp.broadcast_to(lsq_moments_4, p_cc.shape) - lsq_moments_5 = xp.broadcast_to(lsq_moments_5, p_cc.shape) - lsq_moments_6 = xp.broadcast_to(lsq_moments_6, p_cc.shape) - lsq_moments_7 = xp.broadcast_to(lsq_moments_7, p_cc.shape) - lsq_moments_8 = xp.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) lsq_pseudoinv_9 = helpers.reshape(lsq_pseudoinv_9, c2e2c2e2c.shape) - lsq_pseudoinv_9 = xp.expand_dims(lsq_pseudoinv_9, axis=-1) + lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) lsq_pseudoinv_8 = helpers.reshape(lsq_pseudoinv_8, c2e2c2e2c.shape) - lsq_pseudoinv_8 = xp.expand_dims(lsq_pseudoinv_8, axis=-1) + lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) lsq_pseudoinv_7 = helpers.reshape(lsq_pseudoinv_7, c2e2c2e2c.shape) - lsq_pseudoinv_7 = xp.expand_dims(lsq_pseudoinv_7, axis=-1) + lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) lsq_pseudoinv_6 = helpers.reshape(lsq_pseudoinv_6, c2e2c2e2c.shape) - lsq_pseudoinv_6 = xp.expand_dims(lsq_pseudoinv_6, axis=-1) + lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) lsq_pseudoinv_5 = helpers.reshape(lsq_pseudoinv_5, c2e2c2e2c.shape) - lsq_pseudoinv_5 = xp.expand_dims(lsq_pseudoinv_5, axis=-1) + lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) lsq_pseudoinv_4 = helpers.reshape(lsq_pseudoinv_4, c2e2c2e2c.shape) - lsq_pseudoinv_4 = xp.expand_dims(lsq_pseudoinv_4, axis=-1) + lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) lsq_pseudoinv_3 = helpers.reshape(lsq_pseudoinv_3, c2e2c2e2c.shape) - lsq_pseudoinv_3 = xp.expand_dims(lsq_pseudoinv_3, axis=-1) + lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) lsq_pseudoinv_2 = helpers.reshape(lsq_pseudoinv_2, c2e2c2e2c.shape) - lsq_pseudoinv_2 = xp.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) lsq_pseudoinv_1 = helpers.reshape(lsq_pseudoinv_1, c2e2c2e2c.shape) - lsq_pseudoinv_1 = xp.expand_dims(lsq_pseudoinv_1, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) p_coeff_10_dsl = ( lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) diff --git a/model/atmosphere/advection/tests/advection_tests/test_advection.py b/model/atmosphere/advection/tests/advection_tests/test_advection.py index 75ccbbbc5b..0afdba32f1 100644 --- a/model/atmosphere/advection/tests/advection_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection_tests/test_advection.py @@ -25,17 +25,13 @@ ) -# note about ntracer: The first tracer is always dry air which is not advected. Thus, originally -# ntracer=2 is the first tracer in transport_nml, ntracer=3 the second, and so on. -# Here though, ntracer=1 corresponds to the first tracer in transport_nml. - # ntracer legend for the serialization data used here in test_advection: # ------------------------------------ # ntracer | 1, 2, 3, 4, 5 | # ------------------------------------ # ivadv_tracer | 3, 0, 0, 2, 3 | -# itype_hlimit | 3, 3, 4, 0, 0 | -# itype_vlimit | 1, 0, 0, 1, 2 | +# itype_hlimit | 3, 4, 3, 0, 0 | +# itype_vlimit | 1, 0, 0, 2, 1 | # ihadv_tracer | 52, 2, 2, 0, 0 | # ------------------------------------ @@ -62,6 +58,24 @@ advection.VerticalAdvectionType.NO_ADVECTION, advection.VerticalAdvectionLimiter.NO_LIMITER, ), + ( + "2021-06-20T12:00:10.000", + False, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ( + "2021-06-20T12:00:20.000", + True, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), ], ) def test_advection_run_single_step( @@ -69,6 +83,7 @@ def test_advection_run_single_step( icon_grid, interpolation_savepoint, least_squares_savepoint, + metrics_savepoint, advection_init_savepoint, advection_exit_savepoint, data_provider, @@ -89,7 +104,7 @@ def test_advection_run_single_step( ) interpolation_state = construct_interpolation_state(interpolation_savepoint) least_squares_state = construct_least_squares_state(least_squares_savepoint) - metric_state = construct_metric_state(icon_grid) + metric_state = construct_metric_state(icon_grid, metrics_savepoint) edge_geometry = grid_savepoint.construct_edge_geometry() cell_geometry = grid_savepoint.construct_cell_geometry() @@ -132,4 +147,5 @@ def test_advection_run_single_step( diagnostic_state_ref=diagnostic_state_ref, p_tracer_new=p_tracer_new, p_tracer_new_ref=p_tracer_new_ref, + even_timestep=even_timestep, ) diff --git a/model/atmosphere/advection/tests/advection_tests/utils.py b/model/atmosphere/advection/tests/advection_tests/utils.py index 3e44294c74..dbd884c25f 100644 --- a/model/atmosphere/advection/tests/advection_tests/utils.py +++ b/model/atmosphere/advection/tests/advection_tests/utils.py @@ -8,6 +8,8 @@ import logging +import gt4py.next as gtx +import numpy as np from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta @@ -54,12 +56,16 @@ def construct_least_squares_state( ) -def construct_metric_state(icon_grid) -> advection_states.AdvectionMetricState: +def construct_metric_state( + icon_grid, savepoint: sb.MetricSavepoint +) -> advection_states.AdvectionMetricState: constant_f = helpers.constant_field(icon_grid, 1.0, dims.KDim) + ddqz_z_full_np = np.reciprocal(savepoint.inv_ddqz_z_full().asnumpy()) return advection_states.AdvectionMetricState( deepatmo_divh=constant_f, deepatmo_divzl=constant_f, deepatmo_divzu=constant_f, + ddqz_z_full=gtx.as_field((dims.CellDim, dims.KDim), ddqz_z_full_np), ) @@ -128,11 +134,17 @@ def verify_advection_fields( diagnostic_state_ref: advection_states.AdvectionDiagnosticState, p_tracer_new: fa.CellKField[ta.wpfloat], p_tracer_new_ref: fa.CellKField[ta.wpfloat], + even_timestep: bool, ): # cell indices cell_domain = h_grid.domain(dims.CellDim) start_cell_lateral_boundary = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY)) + start_cell_lateral_boundary_level_2 = grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + end_cell_end = grid.end_index(cell_domain(h_grid.Zone.END)) # edge indices edge_domain = h_grid.domain(dims.EdgeDim) @@ -141,46 +153,35 @@ def verify_advection_fields( ) end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) - # log advection output fields - log_dbg( - diagnostic_state.hfl_tracer.asnumpy()[start_edge_lateral_boundary_level_5:end_edge_halo, :], - "hfl_tracer", - ) - log_dbg( - diagnostic_state_ref.hfl_tracer.asnumpy()[ - start_edge_lateral_boundary_level_5:end_edge_halo, : - ], - "hfl_tracer_ref", - ) - log_dbg( - diagnostic_state.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "vfl_tracer", - ) - log_dbg( - diagnostic_state_ref.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "vfl_tracer_ref", - ) - log_dbg(p_tracer_new.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], "p_tracer_new") - log_dbg( - p_tracer_new_ref.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "p_tracer_new_ref", + hfl_tracer_range = np.arange(start_edge_lateral_boundary_level_5, end_edge_halo) + vfl_tracer_range = ( + np.arange(start_cell_lateral_boundary_level_2, end_cell_end) + if even_timestep + else np.arange(start_cell_nudging, end_cell_local) ) + p_tracer_new_range = np.arange(start_cell_lateral_boundary, end_cell_local) + + # log advection output fields + log_dbg(diagnostic_state.hfl_tracer.asnumpy()[hfl_tracer_range, :], "hfl_tracer") + log_dbg(diagnostic_state_ref.hfl_tracer.asnumpy()[hfl_tracer_range, :], "hfl_tracer_ref") + log_dbg(diagnostic_state.vfl_tracer.asnumpy()[vfl_tracer_range, :], "vfl_tracer") + log_dbg(diagnostic_state_ref.vfl_tracer.asnumpy()[vfl_tracer_range, :], "vfl_tracer_ref") + log_dbg(p_tracer_new.asnumpy()[p_tracer_new_range, :], "p_tracer_new") + log_dbg(p_tracer_new_ref.asnumpy()[p_tracer_new_range, :], "p_tracer_new_ref") # verify advection output fields assert helpers.dallclose( - diagnostic_state.hfl_tracer.asnumpy()[start_edge_lateral_boundary_level_5:end_edge_halo, :], - diagnostic_state_ref.hfl_tracer.asnumpy()[ - start_edge_lateral_boundary_level_5:end_edge_halo, : - ], + diagnostic_state.hfl_tracer.asnumpy()[hfl_tracer_range, :], + diagnostic_state_ref.hfl_tracer.asnumpy()[hfl_tracer_range, :], rtol=1e-10, ) - assert helpers.dallclose( # TODO (dastrm): adjust indices once there is vertical transport - diagnostic_state.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - diagnostic_state_ref.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], + assert helpers.dallclose( + diagnostic_state.vfl_tracer.asnumpy()[vfl_tracer_range, :], + diagnostic_state_ref.vfl_tracer.asnumpy()[vfl_tracer_range, :], rtol=1e-10, ) assert helpers.dallclose( - p_tracer_new.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - p_tracer_new_ref.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], + p_tracer_new.asnumpy()[p_tracer_new_range, :], + p_tracer_new_ref.asnumpy()[p_tracer_new_range, :], atol=1e-16, ) diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py index b84805ad52..c80ccd100f 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py @@ -167,7 +167,7 @@ def interpolation_savepoint(data_provider): # F811 @pytest.fixture def metrics_savepoint(data_provider): # F811 - """Load data from ICON mestric state savepoint.""" + """Load data from ICON metric state savepoint.""" return data_provider.from_metrics_savepoint() diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py index c1f348eaca..967bb91229 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py @@ -66,7 +66,7 @@ def get_test_data_root_path() -> pathlib.Path: DATA_URIS_JABW = {1: "https://polybox.ethz.ch/index.php/s/kp9Rab00guECrEd/download"} DATA_URIS_GAUSS3D = {1: "https://polybox.ethz.ch/index.php/s/IiRimdJH2ZBZ1od/download"} DATA_URIS_WK = {1: "https://polybox.ethz.ch/index.php/s/91DEUGmAkBgrXO6/download"} -DATA_URIS_ADVECTION = {1: "https://polybox.ethz.ch/index.php/s/KV6FYstcGysNDOj/download"} +DATA_URIS_ADVECTION = {1: "https://polybox.ethz.ch/index.php/s/3rTia1A41YW7KX9/download"} def get_global_grid_params(experiment: str) -> tuple[int, int]: