diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py index c843d53bb..feef84c43 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py @@ -869,7 +869,6 @@ def time_step( f"running timestep: dtime = {dtime}, initial_timestep = {at_initial_timestep}, first_substep = {at_first_substep}, last_substep = {at_last_substep}, prep_adv = {lprep_adv}" ) - # # TODO: abishekg7 move this to tests if self.p_test_run: self._init_test_fields( self.intermediate_fields.z_rho_e, diff --git a/model/common/src/icon4py/model/common/field_type_aliases.py b/model/common/src/icon4py/model/common/field_type_aliases.py index 3ac4fdd6e..4b498aa3f 100644 --- a/model/common/src/icon4py/model/common/field_type_aliases.py +++ b/model/common/src/icon4py/model/common/field_type_aliases.py @@ -20,6 +20,7 @@ EdgeField: TypeAlias = Field[Dims[dims.EdgeDim], T] VertexField: TypeAlias = Field[Dims[dims.VertexDim], T] KField: TypeAlias = Field[Dims[dims.KDim], T] +KHalfField: TypeAlias = Field[Dims[dims.KHalfDim], T] CellKField: TypeAlias = Field[Dims[dims.CellDim, dims.KDim], T] EdgeKField: TypeAlias = Field[Dims[dims.EdgeDim, dims.KDim], T] diff --git a/model/common/src/icon4py/model/common/grid/geometry.py b/model/common/src/icon4py/model/common/grid/geometry.py index 9d9ba6cc3..df8dc2907 100644 --- a/model/common/src/icon4py/model/common/grid/geometry.py +++ b/model/common/src/icon4py/model/common/grid/geometry.py @@ -132,6 +132,7 @@ def __init__( { # TODO (@magdalena) rescaled by grid_length_rescale_factor (mo_grid_tools.f90) attrs.EDGE_CELL_DISTANCE: extra_fields[gm.GeometryName.EDGE_CELL_DISTANCE], + attrs.EDGE_VERTEX_DISTANCE: extra_fields[gm.GeometryName.EDGE_VERTEX_DISTANCE], attrs.CELL_AREA: extra_fields[gm.GeometryName.CELL_AREA], attrs.DUAL_AREA: extra_fields[gm.GeometryName.DUAL_AREA], attrs.TANGENT_ORIENTATION: extra_fields[gm.GeometryName.TANGENT_ORIENTATION], diff --git a/model/common/src/icon4py/model/common/grid/geometry_attributes.py b/model/common/src/icon4py/model/common/grid/geometry_attributes.py index ce9c59833..c83865ab7 100644 --- a/model/common/src/icon4py/model/common/grid/geometry_attributes.py +++ b/model/common/src/icon4py/model/common/grid/geometry_attributes.py @@ -31,6 +31,7 @@ EDGE_AREA: Final[str] = "edge_area" DUAL_AREA: Final[str] = "dual_area" EDGE_CELL_DISTANCE: Final[str] = "edge_midpoint_to_cell_center_distance" +EDGE_VERTEX_DISTANCE: Final[str] = "edge_midpoint_to_vertex_distance" TANGENT_ORIENTATION: Final[str] = "edge_orientation" CELL_NORMAL_ORIENTATION: Final[str] = "orientation_of_normal_to_cell_edges" VERTEX_EDGE_ORIENTATION: Final[str] = "orientation_of_edges_around_vertex" @@ -126,6 +127,14 @@ icon_var_name="t_grid_edges%edge_cell_length", dtype=ta.wpfloat, ), + EDGE_VERTEX_DISTANCE: dict( + standard_name=EDGE_VERTEX_DISTANCE, + long_name="distances between edge midpoint and adjacent vertices", + units="m", + dims=(dims.EdgeDim, dims.E2VDim), + icon_var_name="t_grid_edges%edge_vert_length", + dtype=ta.wpfloat, + ), DUAL_EDGE_LENGTH: dict( standard_name=DUAL_EDGE_LENGTH, long_name="length of the dual edge", 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 8e561e993..673cb6c71 100644 --- a/model/common/src/icon4py/model/common/grid/grid_manager.py +++ b/model/common/src/icon4py/model/common/grid/grid_manager.py @@ -193,6 +193,7 @@ class GeometryName(FieldName): EDGE_ORIENTATION_ON_VERTEX = "edge_orientation" # TODO (@halungge) compute from coordinates EDGE_CELL_DISTANCE = "edge_cell_distance" + EDGE_VERTEX_DISTANCE = "edge_vert_distance" class CoordinateName(FieldName): @@ -466,6 +467,10 @@ def _read_geometry_fields(self, backend: Optional[gtx_backend.Backend]): (dims.EdgeDim, dims.E2CDim), self._reader.variable(GeometryName.EDGE_CELL_DISTANCE, transpose=True), ), + GeometryName.EDGE_VERTEX_DISTANCE.value: gtx.as_field( + (dims.EdgeDim, dims.E2VDim), + self._reader.variable(GeometryName.EDGE_VERTEX_DISTANCE, transpose=True), + ), # TODO (@halungge) recompute from coordinates? field in gridfile contains NaN on boundary edges GeometryName.TANGENT_ORIENTATION.value: gtx.as_field( (dims.EdgeDim,), diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py index cd8de9139..3c8fa4fa8 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py @@ -20,6 +20,7 @@ GEOFAC_GRDIV: Final[str] = "geometrical_factor_for_gradient_of_divergence" GEOFAC_GRG_X: Final[str] = "geometrical_factor_for_green_gauss_gradient_x" GEOFAC_GRG_Y: Final[str] = "geometrical_factor_for_green_gauss_gradient_y" +CELL_AW_VERTS: Final[str] = "cell_to_vertex_interpolation_factor_by_area_weighting" attrs: dict[str, model.FieldMetaData] = { C_LIN_E: dict( @@ -86,4 +87,12 @@ icon_var_name="geofac_grg", dtype=ta.wpfloat, ), + CELL_AW_VERTS: dict( + standard_name=CELL_AW_VERTS, + long_name="coefficient for interpolation from cells to verts by area weighting", + units="", + dims=(dims.VertexDim, dims.V2CDim), + icon_var_name="cells_aw_verts", + dtype=ta.wpfloat, + ), } diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index 99ef65266..80b95976f 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -28,6 +28,7 @@ cell_domain = h_grid.domain(dims.CellDim) edge_domain = h_grid.domain(dims.EdgeDim) +vertex_domain = h_grid.domain(dims.VertexDim) class InterpolationFieldsFactory(factory.FieldSource, factory.GridProvider): @@ -181,9 +182,31 @@ def _register_computed_fields(self): ) }, ) - self.register_provider(geofac_grg) + cells_aw_verts = factory.NumpyFieldsProvider( + func=functools.partial(interpolation_fields.compute_cells_aw_verts), + fields=(attrs.CELL_AW_VERTS,), + domain=(dims.VertexDim, dims.V2CDim), + deps={ + "dual_area": geometry_attrs.DUAL_AREA, + "edge_vert_length": geometry_attrs.EDGE_VERTEX_DISTANCE, + "edge_cell_length": geometry_attrs.EDGE_CELL_DISTANCE, + }, + connectivities={ + "v2e": dims.V2EDim, + "e2v": dims.E2VDim, + "v2c": dims.V2CDim, + "e2c": dims.E2CDim, + }, + params={ + "horizontal_start": self.grid.start_index( + vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + }, + ) + self.register_provider(cells_aw_verts) + @property def metadata(self) -> dict[str, model.FieldMetaData]: return self._attrs 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 21828001b..61d481ead 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -828,7 +828,7 @@ def compute_cells_aw_verts( Args: dual_area: numpy array, representing a gtx.Field[gtx.Dims[VertexDim], ta.wpfloat] - edge_vert_length: \\ numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat] + edge_vert_length: \\ numpy array, representing a gtx.Field[gtx.Dims[EdgeDim, E2VDim], ta.wpfloat] edge_cell_length: // owner_mask: numpy array, representing a gtx.Field[gtx.Dims[VertexDim], bool] v2e: numpy array, representing a gtx.Field[gtx.Dims[VertexDim, V2EDim], gtx.int32] @@ -845,14 +845,18 @@ def compute_cells_aw_verts( cells_aw_verts[jv, :] = 0.0 for je in range(v2e.shape[1]): # INVALID_INDEX - if je > gm.GridFile.INVALID_INDEX and (je > 0 and v2e[jv, je] == v2e[jv, je - 1]): + if v2e[jv, je] == gm.GridFile.INVALID_INDEX or ( + 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]): + if v2c[jv, jc] == gm.GridFile.INVALID_INDEX or ( + jc > 0 and v2c[jv, jc] == v2c[jv, jc - 1] + ): continue if cell_offset_idx_0 == v2c[jv, jc]: cells_aw_verts[jv, jc] = ( diff --git a/model/common/src/icon4py/model/common/math/helpers.py b/model/common/src/icon4py/model/common/math/helpers.py index 1574071d5..83931b70d 100644 --- a/model/common/src/icon4py/model/common/math/helpers.py +++ b/model/common/src/icon4py/model/common/math/helpers.py @@ -7,11 +7,11 @@ # SPDX-License-Identifier: BSD-3-Clause from gt4py import next as gtx -from gt4py.next import field_operator +from gt4py.next import broadcast, field_operator from gt4py.next.ffront.fbuiltins import arccos, cos, sin, sqrt, where from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.dimension import E2C, E2V, Koff +from icon4py.model.common.dimension import E2C, E2V, CellDim, KDim, Koff from icon4py.model.common.type_alias import wpfloat @@ -529,3 +529,10 @@ def arc_length( """ return radius * arccos(dot_product_on_edges(x0, x1, y0, y1, z0, z1)) + + +@field_operator +def broadcast_cells_to_cell_k( + cell_value: wpfloat, +) -> fa.CellKField[wpfloat]: + return broadcast(cell_value, (CellDim, KDim)) diff --git a/model/common/src/icon4py/model/common/metrics/compute_coeff_gradekin.py b/model/common/src/icon4py/model/common/metrics/compute_coeff_gradekin.py index 176f8426a..c931c3e40 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_coeff_gradekin.py +++ b/model/common/src/icon4py/model/common/metrics/compute_coeff_gradekin.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from types import ModuleType + import numpy as np from icon4py.model.common import dimension as dims @@ -13,11 +15,12 @@ def compute_coeff_gradekin( - edge_cell_length: np.array, - inv_dual_edge_length: np.array, - horizontal_start: float, - horizontal_end: float, -) -> np.array: + edge_cell_length: data_alloc.NDArray, + inv_dual_edge_length: data_alloc.NDArray, + horizontal_start: int, + horizontal_end: int, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: """ Compute coefficients for improved calculation of kinetic energy gradient @@ -27,8 +30,8 @@ def compute_coeff_gradekin( horizontal_start: horizontal start index horizontal_end: horizontal end index """ - coeff_gradekin_0 = np.zeros_like(inv_dual_edge_length) - coeff_gradekin_1 = np.zeros_like(inv_dual_edge_length) + coeff_gradekin_0 = array_ns.zeros_like(inv_dual_edge_length) + coeff_gradekin_1 = array_ns.zeros_like(inv_dual_edge_length) for e in range(horizontal_start, horizontal_end): coeff_gradekin_0[e] = ( edge_cell_length[e, 1] / edge_cell_length[e, 0] * inv_dual_edge_length[e] @@ -36,5 +39,5 @@ def compute_coeff_gradekin( coeff_gradekin_1[e] = ( edge_cell_length[e, 0] / edge_cell_length[e, 1] * inv_dual_edge_length[e] ) - coeff_gradekin_full = np.column_stack((coeff_gradekin_0, coeff_gradekin_1)) - return data_alloc.numpy_to_1D_sparse_field(coeff_gradekin_full, dims.ECDim) + coeff_gradekin_full = array_ns.column_stack((coeff_gradekin_0, coeff_gradekin_1)) + return data_alloc.numpy_to_1D_sparse_field(coeff_gradekin_full, dims.ECDim).asnumpy() diff --git a/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py index 2a8db4fde..ed21439ff 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/compute_diffusion_metrics.py @@ -6,11 +6,25 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from types import ModuleType + import numpy as np from icon4py.model.common.utils import data_allocation as data_alloc +def compute_max_nbhgt_array_ns( + c2e2c: data_alloc.NDArray, + z_mc: data_alloc.NDArray, + nlev: int, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + z_mc_nlev = z_mc[:, nlev - 1] + max_nbhgt_0_1 = array_ns.maximum(z_mc_nlev[c2e2c[:, 0]], z_mc_nlev[c2e2c[:, 1]]) + max_nbhgt = array_ns.maximum(max_nbhgt_0_1, z_mc_nlev[c2e2c[:, 2]]) + return max_nbhgt + + def _compute_nbidx( k_range: range, z_mc: data_alloc.NDArray, @@ -62,8 +76,8 @@ def _compute_z_vintcoeff( def _compute_ls_params( k_start: list, k_end: list, - z_maxslp_avg: data_alloc.NDArray, - z_maxhgtd_avg: data_alloc.NDArray, + maxslp_avg: data_alloc.NDArray, + maxhgtd_avg: data_alloc.NDArray, c_owner_mask: data_alloc.NDArray, thslp_zdiffu: float, thhgtd_zdiffu: float, @@ -78,8 +92,7 @@ def _compute_ls_params( for jc in range(cell_nudging, n_cells): if ( - z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu - or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu + maxslp_avg[jc, nlev - 1] >= thslp_zdiffu or maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu ) and c_owner_mask[jc]: ji += 1 indlist[ji] = jc @@ -96,8 +109,8 @@ def _compute_ls_params( def _compute_k_start_end( z_mc: data_alloc.NDArray, max_nbhgt: data_alloc.NDArray, - z_maxslp_avg: data_alloc.NDArray, - z_maxhgtd_avg: data_alloc.NDArray, + maxslp_avg: data_alloc.NDArray, + maxhgtd_avg: data_alloc.NDArray, c_owner_mask: data_alloc.NDArray, thslp_zdiffu: float, thhgtd_zdiffu: float, @@ -109,8 +122,7 @@ def _compute_k_start_end( k_end = [None] * n_cells for jc in range(cell_nudging, n_cells): if ( - z_maxslp_avg[jc, nlev - 1] >= thslp_zdiffu - or z_maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu + maxslp_avg[jc, nlev - 1] >= thslp_zdiffu or maxhgtd_avg[jc, nlev - 1] >= thhgtd_zdiffu ) and c_owner_mask[jc]: for jk in reversed(range(nlev)): if z_mc[jc, jk] >= max_nbhgt[jc]: @@ -118,7 +130,7 @@ def _compute_k_start_end( break for jk in range(nlev): - if z_maxslp_avg[jc, jk] >= thslp_zdiffu or z_maxhgtd_avg[jc, jk] >= thhgtd_zdiffu: + if maxslp_avg[jc, jk] >= thslp_zdiffu or maxhgtd_avg[jc, jk] >= thhgtd_zdiffu: k_start[jc] = jk break @@ -129,29 +141,32 @@ def _compute_k_start_end( def compute_diffusion_metrics( + c2e2c: data_alloc.NDArray, z_mc: data_alloc.NDArray, - z_mc_off: data_alloc.NDArray, max_nbhgt: data_alloc.NDArray, c_owner_mask: data_alloc.NDArray, - nbidx: data_alloc.NDArray, - z_vintcoeff: data_alloc.NDArray, - z_maxslp_avg: data_alloc.NDArray, - z_maxhgtd_avg: data_alloc.NDArray, - mask_hdiff: data_alloc.NDArray, - zd_diffcoef_dsl: data_alloc.NDArray, - zd_intcoef_dsl: data_alloc.NDArray, - zd_vertoffset_dsl: data_alloc.NDArray, + maxslp_avg: data_alloc.NDArray, + maxhgtd_avg: data_alloc.NDArray, thslp_zdiffu: float, thhgtd_zdiffu: float, cell_nudging: int, - n_cells: int, nlev: int, + array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray, data_alloc.NDArray]: + n_cells = c2e2c.shape[0] + n_c2e2c = c2e2c.shape[1] + z_mc_off = z_mc[c2e2c] + nbidx = array_ns.ones(shape=(n_cells, n_c2e2c, nlev), dtype=int) + z_vintcoeff = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev)) + mask_hdiff = array_ns.zeros(shape=(n_cells, nlev), dtype=bool) + zd_vertoffset_dsl = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev)) + zd_intcoef_dsl = array_ns.zeros(shape=(n_cells, n_c2e2c, nlev)) + zd_diffcoef_dsl = array_ns.zeros(shape=(n_cells, nlev)) k_start, k_end = _compute_k_start_end( z_mc=z_mc, max_nbhgt=max_nbhgt, - z_maxslp_avg=z_maxslp_avg, - z_maxhgtd_avg=z_maxhgtd_avg, + maxslp_avg=maxslp_avg, + maxhgtd_avg=maxhgtd_avg, c_owner_mask=c_owner_mask, thslp_zdiffu=thslp_zdiffu, thhgtd_zdiffu=thhgtd_zdiffu, @@ -163,8 +178,8 @@ def compute_diffusion_metrics( indlist, listreduce, ji = _compute_ls_params( k_start=k_start, k_end=k_end, - z_maxslp_avg=z_maxslp_avg, - z_maxhgtd_avg=z_maxhgtd_avg, + maxslp_avg=maxslp_avg, + maxhgtd_avg=maxhgtd_avg, c_owner_mask=c_owner_mask, thslp_zdiffu=thslp_zdiffu, thhgtd_zdiffu=thhgtd_zdiffu, @@ -185,16 +200,30 @@ def compute_diffusion_metrics( ) zd_intcoef_dsl[jc, :, k_range] = z_vintcoeff[jc, :, k_range] - zd_vertoffset_dsl[jc, :, k_range] = nbidx[jc, :, k_range] - np.transpose([k_range] * 3) + zd_vertoffset_dsl[jc, :, k_range] = nbidx[jc, :, k_range] - array_ns.transpose( + [k_range] * 3 + ) mask_hdiff[jc, k_range] = True - zd_diffcoef_dsl_var = np.maximum( + zd_diffcoef_dsl_var = array_ns.maximum( 0.0, - np.maximum( - np.sqrt(np.maximum(0.0, z_maxslp_avg[jc, k_range] - thslp_zdiffu)) / 250.0, - 2.0e-4 * np.sqrt(np.maximum(0.0, z_maxhgtd_avg[jc, k_range] - thhgtd_zdiffu)), + array_ns.maximum( + array_ns.sqrt(array_ns.maximum(0.0, maxslp_avg[jc, k_range] - thslp_zdiffu)) + / 250.0, + 2.0e-4 + * array_ns.sqrt( + array_ns.maximum(0.0, maxhgtd_avg[jc, k_range] - thhgtd_zdiffu) + ), ), ) - zd_diffcoef_dsl[jc, k_range] = np.minimum(0.002, zd_diffcoef_dsl_var) + zd_diffcoef_dsl[jc, k_range] = array_ns.minimum(0.002, zd_diffcoef_dsl_var) + + # flatten first two dims: + zd_intcoef_dsl = zd_intcoef_dsl.reshape( + (zd_intcoef_dsl.shape[0] * zd_intcoef_dsl.shape[1],) + zd_intcoef_dsl.shape[2:] + ) + zd_vertoffset_dsl = zd_vertoffset_dsl.reshape( + (zd_vertoffset_dsl.shape[0] * zd_vertoffset_dsl.shape[1],) + zd_vertoffset_dsl.shape[2:] + ) return mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl diff --git a/model/common/src/icon4py/model/common/metrics/compute_flat_idx_max.py b/model/common/src/icon4py/model/common/metrics/compute_flat_idx_max.py new file mode 100644 index 000000000..b53cce6e5 --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/compute_flat_idx_max.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 + +from types import ModuleType + +import gt4py.next as gtx +import numpy as np + +from icon4py.model.common.utils import data_allocation as data_alloc + + +def compute_flat_idx_max( + e2c: data_alloc.NDArray, + z_mc: data_alloc.NDArray, + c_lin_e: data_alloc.NDArray, + z_ifc: data_alloc.NDArray, + k_lev: data_alloc.NDArray, + horizontal_lower: int, + horizontal_upper: int, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + z_me = array_ns.sum(z_mc[e2c] * array_ns.expand_dims(c_lin_e, axis=-1), axis=1) + z_ifc_e_0 = z_ifc[e2c[:, 0]] + z_ifc_e_k_0 = array_ns.roll(z_ifc_e_0, -1, axis=1) + z_ifc_e_1 = z_ifc[e2c[:, 1]] + z_ifc_e_k_1 = array_ns.roll(z_ifc_e_1, -1, axis=1) + flat_idx = array_ns.zeros_like(z_me) + for je in range(horizontal_lower, horizontal_upper): + for jk in range(k_lev.shape[0] - 1): + if ( + (z_me[je, jk] <= z_ifc_e_0[je, jk]) + and (z_me[je, jk] >= z_ifc_e_k_0[je, jk]) + and (z_me[je, jk] <= z_ifc_e_1[je, jk]) + and (z_me[je, jk] >= z_ifc_e_k_1[je, jk]) + ): + flat_idx[je, jk] = k_lev[jk] + flat_idx_max = array_ns.amax(flat_idx, axis=1) + return flat_idx_max.astype(gtx.int32) diff --git a/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py b/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py index d9fe3caa7..855c75453 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py +++ b/model/common/src/icon4py/model/common/metrics/compute_vwind_impl_wgt.py @@ -5,53 +5,53 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from types import ModuleType import numpy as np -import icon4py.model.common.field_type_aliases as fa -from icon4py.model.common.grid import base as grid -from icon4py.model.common.metrics.metric_fields import compute_vwind_impl_wgt_partial -from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc def compute_vwind_impl_wgt( - backend, - icon_grid: grid.BaseGrid, - vct_a: fa.KField[wpfloat], - z_ifc: fa.CellKField[wpfloat], - z_ddxn_z_half_e: fa.EdgeField[wpfloat], - z_ddxt_z_half_e: fa.EdgeField[wpfloat], - dual_edge_length: fa.EdgeField[wpfloat], - vwind_impl_wgt_full: fa.CellField[wpfloat], - vwind_impl_wgt_k: fa.CellField[wpfloat], - global_exp: str, - experiment: str, + c2e: data_alloc.NDArray, + vct_a: data_alloc.NDArray, + z_ifc: data_alloc.NDArray, + z_ddxn_z_half_e: data_alloc.NDArray, + z_ddxt_z_half_e: data_alloc.NDArray, + dual_edge_length: data_alloc.NDArray, vwind_offctr: float, + nlev: int, horizontal_start_cell: int, + n_cells: int, + array_ns: ModuleType = np, ) -> data_alloc.NDArray: - compute_vwind_impl_wgt_partial.with_backend(backend)( - z_ddxn_z_half_e=z_ddxn_z_half_e, - z_ddxt_z_half_e=z_ddxt_z_half_e, - dual_edge_length=dual_edge_length, - vct_a=vct_a, - z_ifc=z_ifc, - vwind_impl_wgt=vwind_impl_wgt_full, - vwind_impl_wgt_k=vwind_impl_wgt_k, - vwind_offctr=vwind_offctr, - horizontal_start=horizontal_start_cell, - horizontal_end=icon_grid.num_cells, - vertical_start=max(10, icon_grid.num_levels - 8), - vertical_end=icon_grid.num_levels, - offset_provider={ - "C2E": icon_grid.get_offset_provider("C2E"), - "Koff": icon_grid.get_offset_provider("Koff"), - }, - ) + init_val = 0.5 + vwind_offctr + vwind_impl_wgt = array_ns.full(z_ifc.shape[0], init_val) + for je in range(horizontal_start_cell, n_cells): + zn_off_0 = z_ddxn_z_half_e[c2e[je, 0], nlev] + zn_off_1 = z_ddxn_z_half_e[c2e[je, 1], nlev] + zn_off_2 = z_ddxn_z_half_e[c2e[je, 2], nlev] + zt_off_0 = z_ddxt_z_half_e[c2e[je, 0], nlev] + zt_off_1 = z_ddxt_z_half_e[c2e[je, 1], nlev] + zt_off_2 = z_ddxt_z_half_e[c2e[je, 2], nlev] + z_maxslope = max( + abs(zn_off_0), abs(zt_off_0), abs(zn_off_1), abs(zt_off_1), abs(zn_off_2), abs(zt_off_2) + ) + z_diff = max( + abs(zn_off_0 * dual_edge_length[c2e[je, 0]]), + abs(zn_off_1 * dual_edge_length[c2e[je, 1]]), + abs(zn_off_2 * dual_edge_length[c2e[je, 2]]), + ) - vwind_impl_wgt = ( - np.amin(vwind_impl_wgt_k.ndarray, axis=1) - if experiment == global_exp - else np.amax(vwind_impl_wgt_k.ndarray, axis=1) - ) + z_offctr = max( + vwind_offctr, 0.425 * z_maxslope ** (0.75), min(0.25, 0.00025 * (z_diff - 250.0)) + ) + z_offctr = min(max(vwind_offctr, 0.75), z_offctr) + vwind_impl_wgt[je] = 0.5 + z_offctr + + for jk in range(max(9, nlev - 9), nlev): + for je in range(horizontal_start_cell, n_cells): + z_diff_2 = (z_ifc[je, jk] - z_ifc[je, jk + 1]) / (vct_a[jk] - vct_a[jk + 1]) + if z_diff_2 < 0.6: + vwind_impl_wgt[je] = max(vwind_impl_wgt[je], 1.2 - z_diff_2) return vwind_impl_wgt diff --git a/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py b/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py index de0ae94cc..6f0853993 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py +++ b/model/common/src/icon4py/model/common/metrics/compute_wgtfacq.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from types import ModuleType + import numpy as np from icon4py.model.common.utils import data_allocation as data_alloc @@ -23,6 +25,7 @@ def _compute_z1_z2_z3( def compute_wgtfacq_c_dsl( z_ifc: data_alloc.NDArray, nlev: int, + array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ Compute weighting factor for quadratic interpolation to surface. @@ -33,8 +36,8 @@ def compute_wgtfacq_c_dsl( Returns: Field[CellDim, KDim] (full levels) """ - wgtfacq_c = np.zeros((z_ifc.shape[0], nlev + 1)) - wgtfacq_c_dsl = np.zeros((z_ifc.shape[0], nlev)) + wgtfacq_c = array_ns.zeros((z_ifc.shape[0], nlev + 1)) + wgtfacq_c_dsl = array_ns.zeros((z_ifc.shape[0], nlev)) z1, z2, z3 = _compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) wgtfacq_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) @@ -55,6 +58,7 @@ def compute_wgtfacq_e_dsl( wgtfacq_c_dsl: data_alloc.NDArray, n_edges: int, nlev: int, + array_ns: ModuleType = np, ): """ Compute weighting factor for quadratic interpolation to surface. @@ -69,8 +73,8 @@ def compute_wgtfacq_e_dsl( Returns: Field[EdgeDim, KDim] (full levels) """ - wgtfacq_e_dsl = np.zeros(shape=(n_edges, nlev + 1)) - z_aux_c = np.zeros((z_ifc.shape[0], 6)) + wgtfacq_e_dsl = array_ns.zeros(shape=(n_edges, nlev + 1)) + z_aux_c = array_ns.zeros((z_ifc.shape[0], 6)) z1, z2, z3 = _compute_z1_z2_z3(z_ifc, nlev, nlev - 1, nlev - 2, nlev - 3) z_aux_c[:, 2] = z1 * z2 / (z2 - z3) / (z1 - z3) z_aux_c[:, 1] = (z1 - wgtfacq_c_dsl[:, nlev - 3] * (z1 - z3)) / (z1 - z2) @@ -81,8 +85,8 @@ def compute_wgtfacq_e_dsl( z_aux_c[:, 4] = (z1 - z_aux_c[:, 5] * (z1 - z3)) / (z1 - z2) z_aux_c[:, 3] = 1.0 - (z_aux_c[:, 4] + z_aux_c[:, 5]) - c_lin_e = c_lin_e[:, :, np.newaxis] - z_aux_e = np.sum(c_lin_e * z_aux_c[e2c], axis=1) + c_lin_e = c_lin_e[:, :, array_ns.newaxis] + z_aux_e = array_ns.sum(c_lin_e * z_aux_c[e2c], axis=1) wgtfacq_e_dsl[:, nlev] = z_aux_e[:, 0] wgtfacq_e_dsl[:, nlev - 1] = z_aux_e[:, 1] diff --git a/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py b/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py index 13dced1f5..93be6640c 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py +++ b/model/common/src/icon4py/model/common/metrics/compute_zdiff_gradp_dsl.py @@ -6,6 +6,8 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from types import ModuleType + import numpy as np from icon4py.model.common.utils import data_allocation as data_alloc @@ -13,19 +15,24 @@ def compute_zdiff_gradp_dsl( e2c, - z_me: data_alloc.NDArray, z_mc: data_alloc.NDArray, + c_lin_e: data_alloc.NDArray, z_ifc: data_alloc.NDArray, flat_idx: data_alloc.NDArray, - z_aux2: data_alloc.NDArray, + z_ifc_sliced: data_alloc.NDArray, nlev: int, horizontal_start: int, horizontal_start_1: int, - nedges: int, + array_ns: ModuleType = np, ) -> data_alloc.NDArray: - zdiff_gradp = np.zeros_like(z_mc[e2c]) + nedges = e2c.shape[0] + z_me = array_ns.sum(z_mc[e2c] * array_ns.expand_dims(c_lin_e, axis=-1), axis=1) + z_aux1 = array_ns.maximum(z_ifc_sliced[e2c[:, 0]], z_ifc_sliced[e2c[:, 1]]) + z_aux2 = z_aux1 - 5.0 # extrapol_dist + zdiff_gradp = array_ns.zeros_like(z_mc[e2c]) zdiff_gradp[horizontal_start:, :, :] = ( - np.expand_dims(z_me, axis=1)[horizontal_start:, :, :] - z_mc[e2c][horizontal_start:, :, :] + array_ns.expand_dims(z_me, axis=1)[horizontal_start:, :, :] + - z_mc[e2c][horizontal_start:, :, :] ) """ First part for loop implementation with gt4py code @@ -44,8 +51,8 @@ def compute_zdiff_gradp_dsl( for jk in range(int(flat_idx[je]) + 1, nlev): """ Second part for loop implementation with gt4py code - >>> param_2 = as_field((KDim,), np.asarray([False] * nlev)) - >>> param_3 = as_field((KDim,), np.arange(nlev)) + >>> param_2 = as_field((KDim,), array_ns.asarray([False] * nlev)) + >>> param_3 = as_field((KDim,), array_ns.arange(nlev)) >>> z_ifc_off_e = as_field((KDim,), z_ifc[e2c[je, 0], :]) >>> _compute_param.with_backend(backend)( >>> z_me_jk=z_me[je, jk], @@ -56,7 +63,7 @@ def compute_zdiff_gradp_dsl( >>> out=(param_3, param_2), >>> offset_provider={} >>> ) - >>> zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], np.where(param_2.ndarray)[0][0]] + >>> zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], array_ns.where(param_2.ndarray)[0][0]] """ param = [False] * nlev @@ -68,7 +75,7 @@ def compute_zdiff_gradp_dsl( ): param[jk1] = True - zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], np.where(param)[0][0]] + zdiff_gradp[je, 0, jk] = z_me[je, jk] - z_mc[e2c[je, 0], array_ns.where(param)[0][0]] jk_start = int(flat_idx[je]) for jk in range(int(flat_idx[je]) + 1, nlev): @@ -109,4 +116,7 @@ def compute_zdiff_gradp_dsl( jk_start = jk1 break - return zdiff_gradp + zdiff_gradp_full_field = zdiff_gradp.reshape( + (zdiff_gradp.shape[0] * zdiff_gradp.shape[1],) + zdiff_gradp.shape[2:] + ) + return zdiff_gradp_full_field diff --git a/model/common/src/icon4py/model/common/metrics/metric_fields.py b/model/common/src/icon4py/model/common/metrics/metric_fields.py index 276eaebec..7281cd933 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -5,8 +5,6 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from dataclasses import dataclass -from typing import Final import gt4py.next as gtx from gt4py.next import ( @@ -16,6 +14,8 @@ broadcast, exp, field_operator, + int32, + log, maximum, minimum, neighbor_sum, @@ -38,10 +38,14 @@ from icon4py.model.common.interpolation.stencils.cell_2_edge_interpolation import ( _cell_2_edge_interpolation, ) +from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( + _compute_cell_2_vertex_interpolation, +) from icon4py.model.common.math.helpers import ( _grad_fd_tang, average_cell_kdim_level_up, average_edge_kdim_level_up, + broadcast_cells_to_cell_k, difference_k_level_up, grad_fd_norm, ) @@ -53,12 +57,6 @@ """ -@dataclass(frozen=True) -class MetricsConfig: - #: Temporal extrapolation of Exner for computation of horizontal pressure gradient, defined in `mo_nonhydrostatic_nml.f90` used only in metrics fields calculation. - exner_expol: Final[wpfloat] = 0.3333333333333 - - @program(grid_type=GridType.UNSTRUCTURED) def compute_z_mc( z_ifc: fa.CellKField[wpfloat], @@ -93,13 +91,15 @@ def compute_z_mc( ) +# TODO(@nfarabullini): ddqz_z_half vertical dimension is khalf, use K2KHalf once merged for z_ifc and z_mc +# TODO(@nfarabullini): change dimension type hint for ddqz_z_half to cell, khalf @field_operator def _compute_ddqz_z_half( z_ifc: fa.CellKField[wpfloat], z_mc: fa.CellKField[wpfloat], k: fa.KField[gtx.int32], nlev: gtx.int32, -): +): # -> Field[Dims[dims.CellDim, dims.KHalfDim], wpfloat]: # TODO: change this to concat_where once it's merged ddqz_z_half = where(k == 0, 2.0 * (z_ifc - z_mc), 0.0) ddqz_z_half = where((k > 0) & (k < nlev), z_mc(Koff[-1]) - z_mc, ddqz_z_half) @@ -504,9 +504,26 @@ def compute_ddxn_z_half_e( ) +@field_operator +def _compute_ddxt_z_half_e( + cell_in: fa.CellKField[wpfloat], + c_int: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], wpfloat], + inv_primal_edge_length: fa.EdgeField[wpfloat], + tangent_orientation: fa.EdgeField[wpfloat], +): + z_ifv = _compute_cell_2_vertex_interpolation(cell_in, c_int) + ddxt_z_half_e = _grad_fd_tang( + z_ifv, + inv_primal_edge_length, + tangent_orientation, + ) + return ddxt_z_half_e + + @program def compute_ddxt_z_half_e( - z_ifv: gtx.Field[gtx.Dims[dims.VertexDim, dims.KDim], float], + cell_in: fa.CellKField[wpfloat], + c_int: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], wpfloat], inv_primal_edge_length: fa.EdgeField[wpfloat], tangent_orientation: fa.EdgeField[wpfloat], ddxt_z_half_e: fa.EdgeKField[wpfloat], @@ -515,8 +532,9 @@ def compute_ddxt_z_half_e( vertical_start: gtx.int32, vertical_end: gtx.int32, ): - _grad_fd_tang( - z_ifv, + _compute_ddxt_z_half_e( + cell_in, + c_int, inv_primal_edge_length, tangent_orientation, out=ddxt_z_half_e, @@ -527,11 +545,23 @@ def compute_ddxt_z_half_e( ) -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_ddxn_z_full( - z_ddxnt_z_half_e: fa.EdgeKField[wpfloat], ddxn_z_full: fa.EdgeKField[wpfloat] + ddxnt_z_half_e: fa.EdgeKField[wpfloat], + ddxn_z_full: fa.EdgeKField[wpfloat], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): - average_edge_kdim_level_up(z_ddxnt_z_half_e, out=ddxn_z_full) + average_edge_kdim_level_up( + ddxnt_z_half_e, + out=ddxn_z_full, + domain={ + dims.EdgeDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) @field_operator @@ -572,23 +602,23 @@ def _compute_maxslp_maxhgtd( dual_edge_length: fa.EdgeField[wpfloat], ) -> tuple[fa.CellKField[wpfloat], fa.CellKField[wpfloat]]: z_maxslp_0_1 = maximum(abs(ddxn_z_full(C2E[0])), abs(ddxn_z_full(C2E[1]))) - z_maxslp = maximum(z_maxslp_0_1, abs(ddxn_z_full(C2E[2]))) + maxslp = maximum(z_maxslp_0_1, abs(ddxn_z_full(C2E[2]))) z_maxhgtd_0_1 = maximum( abs(ddxn_z_full(C2E[0]) * dual_edge_length(C2E[0])), abs(ddxn_z_full(C2E[1]) * dual_edge_length(C2E[1])), ) - z_maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) - return z_maxslp, z_maxhgtd + maxhgtd = maximum(z_maxhgtd_0_1, abs(ddxn_z_full(C2E[2]) * dual_edge_length(C2E[2]))) + return maxslp, maxhgtd @program def compute_maxslp_maxhgtd( ddxn_z_full: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], wpfloat], dual_edge_length: gtx.Field[gtx.Dims[dims.EdgeDim], wpfloat], - z_maxslp: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], - z_maxhgtd: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], + maxslp: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], + maxhgtd: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, @@ -602,8 +632,8 @@ def compute_maxslp_maxhgtd( Args: ddxn_z_full: dual_edge_length dual_edge_length: dual_edge_length - z_maxslp: output - z_maxhgtd: output + maxslp: output + maxhgtd: output horizontal_start: horizontal start index horizontal_end: horizontal end index vertical_start: vertical start index @@ -612,7 +642,7 @@ def compute_maxslp_maxhgtd( _compute_maxslp_maxhgtd( ddxn_z_full=ddxn_z_full, dual_edge_length=dual_edge_length, - out=(z_maxslp, z_maxhgtd), + out=(maxslp, maxhgtd), domain={ dims.CellDim: (horizontal_start, horizontal_end), dims.KDim: (vertical_start, vertical_end), @@ -664,6 +694,7 @@ def compute_exner_exfac( vertical_end: vertical end index """ + broadcast_cells_to_cell_k(exner_expol, out=exner_exfac) _compute_exner_exfac( ddxn_z_full=ddxn_z_full, dual_edge_length=dual_edge_length, @@ -806,10 +837,12 @@ def compute_wgtfac_e( @field_operator def _compute_flat_idx( - z_me: fa.EdgeKField[wpfloat], + z_mc: fa.CellKField[wpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], z_ifc: fa.CellKField[wpfloat], k_lev: fa.KField[gtx.int32], ) -> fa.EdgeKField[gtx.int32]: + z_me = _cell_2_edge_interpolation(in_field=z_mc, coeff=c_lin_e) z_ifc_e_0 = z_ifc(E2C[0]) z_ifc_e_k_0 = z_ifc_e_0(Koff[1]) z_ifc_e_1 = z_ifc(E2C[1]) @@ -822,6 +855,31 @@ def _compute_flat_idx( return flat_idx +@program(grid_type=GridType.UNSTRUCTURED) +def compute_flat_idx( + z_mc: fa.CellKField[wpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], + z_ifc: fa.CellKField[wpfloat], + k_lev: fa.KField[int32], + flat_idx: fa.EdgeKField[int32], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_flat_idx( + z_mc=z_mc, + c_lin_e=c_lin_e, + z_ifc=z_ifc, + k_lev=k_lev, + out=flat_idx, + domain={ + dims.EdgeDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) + + @field_operator def _compute_z_aux2( z_ifc: fa.CellField[wpfloat], @@ -833,18 +891,31 @@ def _compute_z_aux2( return z_aux2 +@program(grid_type=GridType.UNSTRUCTURED) +def compute_z_aux2( + z_ifc_sliced: fa.CellField[wpfloat], + z_aux2: fa.EdgeField[wpfloat], + horizontal_start: int32, + horizontal_end: int32, +): + _compute_z_aux2( + z_ifc=z_ifc_sliced, out=z_aux2, domain={dims.EdgeDim: (horizontal_start, horizontal_end)} + ) + + @field_operator def _compute_pg_edgeidx_vertidx( c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], float], z_ifc: fa.CellKField[wpfloat], - z_aux2: fa.EdgeField[wpfloat], + z_ifc_sliced: fa.CellField[wpfloat], e_owner_mask: fa.EdgeField[bool], - flat_idx_max: fa.EdgeField[gtx.int32], - e_lev: fa.EdgeField[gtx.int32], - k_lev: fa.KField[gtx.int32], - pg_edgeidx: fa.EdgeKField[gtx.int32], - pg_vertidx: fa.EdgeKField[gtx.int32], -) -> tuple[fa.EdgeKField[gtx.int32], fa.EdgeKField[gtx.int32]]: + flat_idx_max: fa.EdgeField[int32], + e_lev: fa.EdgeField[int32], + k_lev: fa.KField[int32], + pg_edgeidx: fa.EdgeKField[int32], + pg_vertidx: fa.EdgeKField[int32], +) -> tuple[fa.EdgeKField[int32], fa.EdgeKField[int32]]: + z_aux2 = _compute_z_aux2(z_ifc_sliced) e_lev = broadcast(e_lev, (dims.EdgeDim, dims.KDim)) k_lev = broadcast(k_lev, (dims.EdgeDim, dims.KDim)) z_mc = average_cell_kdim_level_up(z_ifc) @@ -858,15 +929,57 @@ def _compute_pg_edgeidx_vertidx( return pg_edgeidx, pg_vertidx +@program(grid_type=GridType.UNSTRUCTURED) +def compute_pg_edgeidx_vertidx( + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], float], + z_ifc: fa.CellKField[wpfloat], + z_ifc_sliced: fa.CellField[wpfloat], + e_owner_mask: fa.EdgeField[bool], + flat_idx_max: fa.EdgeField[int32], + e_lev: fa.EdgeField[int32], + k_lev: fa.KField[int32], + pg_edgeidx: fa.EdgeKField[int32], + pg_vertidx: fa.EdgeKField[int32], + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_pg_edgeidx_vertidx( + c_lin_e=c_lin_e, + z_ifc=z_ifc, + z_ifc_sliced=z_ifc_sliced, + e_owner_mask=e_owner_mask, + flat_idx_max=flat_idx_max, + e_lev=e_lev, + k_lev=k_lev, + pg_edgeidx=pg_edgeidx, + pg_vertidx=pg_vertidx, + out=(pg_edgeidx, pg_vertidx), + domain={ + dims.EdgeDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) + + @field_operator def _compute_pg_exdist_dsl( - z_me: fa.EdgeKField[wpfloat], - z_aux2: fa.EdgeField[wpfloat], + z_ifc_sliced: fa.CellField[wpfloat], + z_mc: fa.CellKField[wpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], e_owner_mask: fa.EdgeField[bool], - flat_idx_max: fa.EdgeField[gtx.int32], - k_lev: fa.KField[gtx.int32], + flat_idx_max: fa.EdgeField[int32], + k_lev: fa.KField[int32], + e_lev: fa.EdgeField[int32], pg_exdist_dsl: fa.EdgeKField[wpfloat], + h_start_zaux2: int32, + h_end_zaux2: int32, ) -> fa.EdgeKField[wpfloat]: + z_me = _cell_2_edge_interpolation(z_mc, c_lin_e) + z_aux2 = where( + (e_lev >= h_start_zaux2) & (e_lev < h_end_zaux2), _compute_z_aux2(z_ifc_sliced), 0.0 + ) k_lev = broadcast(k_lev, (dims.EdgeDim, dims.KDim)) pg_exdist_dsl = where( (k_lev >= (flat_idx_max + 1)) & (z_me < z_aux2) & e_owner_mask, @@ -878,16 +991,20 @@ def _compute_pg_exdist_dsl( @program def compute_pg_exdist_dsl( - z_aux2: fa.EdgeField[wpfloat], - z_me: fa.EdgeKField[wpfloat], + z_ifc_sliced: fa.CellField[wpfloat], + z_mc: fa.CellKField[wpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], e_owner_mask: fa.EdgeField[bool], - flat_idx_max: fa.EdgeField[gtx.int32], - k_lev: fa.KField[gtx.int32], + flat_idx_max: fa.EdgeField[int32], + k_lev: fa.KField[int32], + e_lev: fa.EdgeField[int32], pg_exdist_dsl: fa.EdgeKField[wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, + h_start_zaux2: int32, + h_end_zaux2: int32, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, ): """ Compute pg_edgeidx_dsl. @@ -895,8 +1012,9 @@ def compute_pg_exdist_dsl( See mo_vertical_grid.f90 Args: - z_aux2: Local field - z_me: Local field + z_ifc_sliced: z_ifc sliced field + z_mc: Local field + c_lin_e: interpolation field e_owner_mask: Field of booleans over edges flat_idx_max: Highest vertical index (counted from top to bottom) for which the edge point lies inside the cell box of the adjacent grid points k_lev: Field of K levels @@ -907,12 +1025,16 @@ def compute_pg_exdist_dsl( vertical_end: vertical end index """ _compute_pg_exdist_dsl( - z_me=z_me, - z_aux2=z_aux2, + z_ifc_sliced=z_ifc_sliced, + z_mc=z_mc, + c_lin_e=c_lin_e, e_owner_mask=e_owner_mask, flat_idx_max=flat_idx_max, k_lev=k_lev, + e_lev=e_lev, pg_exdist_dsl=pg_exdist_dsl, + h_start_zaux2=h_start_zaux2, + h_end_zaux2=h_end_zaux2, out=pg_exdist_dsl, domain={ dims.EdgeDim: (horizontal_start, horizontal_end), @@ -930,7 +1052,7 @@ def _compute_pg_edgeidx_dsl( return pg_edgeidx_dsl -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_pg_edgeidx_dsl( pg_edgeidx: fa.EdgeKField[gtx.int32], pg_vertidx: fa.EdgeKField[gtx.int32], @@ -973,7 +1095,7 @@ def _compute_mask_prog_halo_c( return mask_prog_halo_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_mask_prog_halo_c( c_refin_ctrl: fa.CellField[gtx.int32], mask_prog_halo_c: fa.CellField[bool], @@ -1001,14 +1123,13 @@ def compute_mask_prog_halo_c( @field_operator def _compute_bdy_halo_c( - c_refin_ctrl: fa.CellField[gtx.int32], - bdy_halo_c: fa.CellField[bool], + c_refin_ctrl: fa.CellField[int32], ) -> fa.CellField[bool]: - bdy_halo_c = where((c_refin_ctrl >= 1) & (c_refin_ctrl <= 4), True, bdy_halo_c) + bdy_halo_c = where((c_refin_ctrl >= 1) & (c_refin_ctrl <= 4), True, False) return bdy_halo_c -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_bdy_halo_c( c_refin_ctrl: fa.CellField[gtx.int32], bdy_halo_c: fa.CellField[bool], @@ -1018,7 +1139,7 @@ def compute_bdy_halo_c( """ Compute bdy_halo_c. - See mo_vertical_grid.f90. mask_prog_halo_c_dsl_low_refin in ICON + See mo_vertical_grid.f90. bdy_halo_c_dsl_low_refin in ICON Args: c_refin_ctrl: Cell field of refin_ctrl @@ -1028,7 +1149,41 @@ def compute_bdy_halo_c( """ _compute_bdy_halo_c( c_refin_ctrl, - bdy_halo_c, + out=bdy_halo_c, + domain={dims.CellDim: (horizontal_start, horizontal_end)}, + ) + + +@program(grid_type=GridType.UNSTRUCTURED) +def compute_mask_bdy_halo_c( + c_refin_ctrl: fa.CellField[int32], + mask_prog_halo_c: fa.CellField[bool], + bdy_halo_c: fa.CellField[bool], + horizontal_start: int32, + horizontal_end: int32, +): + """ + Compute bdy_halo_c. + Compute mask_prog_halo_c. + + + See mo_vertical_grid.f90. bdy_halo_c_dsl_low_refin in ICON + + Args: + c_refin_ctrl: Cell field of refin_ctrl + bdy_halo_c: output + horizontal_start: horizontal start index + horizontal_end: horizontal end index + """ + _compute_mask_prog_halo_c( + c_refin_ctrl, + mask_prog_halo_c, + out=mask_prog_halo_c, + domain={dims.CellDim: (horizontal_start, horizontal_end)}, + ) + + _compute_bdy_halo_c( + c_refin_ctrl, out=bdy_halo_c, domain={dims.CellDim: (horizontal_start, horizontal_end)}, ) @@ -1040,21 +1195,26 @@ def _compute_hmask_dd3d( grf_nudge_start_e: gtx.int32, grf_nudgezone_width: gtx.int32, ) -> fa.EdgeField[wpfloat]: - hmask_dd3d = ( - 1 - / (grf_nudgezone_width - 1) - * (e_refin_ctrl - (grf_nudge_start_e + grf_nudgezone_width - 1)) + e_refin_ctrl_wp = astype(e_refin_ctrl, wpfloat) + grf_nudge_start_e_wp = astype(grf_nudge_start_e, wpfloat) + grf_nudgezone_width_wp = astype(grf_nudgezone_width, wpfloat) + hmask_dd3d = where( + (e_refin_ctrl > (grf_nudge_start_e + grf_nudgezone_width - 1)), + 1.0 + / (grf_nudgezone_width_wp - 1.0) + * (e_refin_ctrl_wp - (grf_nudge_start_e_wp + grf_nudgezone_width_wp - 1.0)), + 0.0, ) hmask_dd3d = where( - (e_refin_ctrl <= 0) | (e_refin_ctrl >= (grf_nudge_start_e + 2 * (grf_nudgezone_width - 1))), - 1, + (e_refin_ctrl <= 0) + | (e_refin_ctrl_wp >= (grf_nudge_start_e_wp + 2.0 * (grf_nudgezone_width_wp - 1.0))), + 1.0, hmask_dd3d, ) - hmask_dd3d = where(e_refin_ctrl <= (grf_nudge_start_e + grf_nudgezone_width - 1), 0, hmask_dd3d) - return astype(hmask_dd3d, wpfloat) + return hmask_dd3d -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_hmask_dd3d( e_refin_ctrl: fa.EdgeField[gtx.int32], hmask_dd3d: fa.EdgeField[wpfloat], @@ -1087,27 +1247,27 @@ def compute_hmask_dd3d( @field_operator def _compute_weighted_cell_neighbor_sum( - field: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], + field: fa.CellKField[wpfloat], c_bln_avg: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat], -) -> gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat]: +) -> fa.CellKField[wpfloat]: field_avg = neighbor_sum(field(C2E2CO) * c_bln_avg, axis=C2E2CODim) return field_avg -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_weighted_cell_neighbor_sum( maxslp: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], maxhgtd: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], c_bln_avg: gtx.Field[gtx.Dims[dims.CellDim, C2E2CODim], wpfloat], - z_maxslp_avg: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], - z_maxhgtd_avg: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], + maxslp_avg: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], + maxhgtd_avg: gtx.Field[gtx.Dims[dims.CellDim, dims.KDim], wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, vertical_start: gtx.int32, vertical_end: gtx.int32, ): """ - Compute z_maxslp_avg and z_maxhgtd_avg. + Compute maxslp_avg and maxhgtd_avg. See mo_vertical_grid.f90. @@ -1115,8 +1275,8 @@ def compute_weighted_cell_neighbor_sum( maxslp: Max field over ddxn_z_full offset maxhgtd: Max field over ddxn_z_full offset*dual_edge_length offset c_bln_avg: Interpolation field - z_maxslp_avg: output - z_maxhgtd_avg: output + maxslp_avg: output + maxhgtd_avg: output horizontal_start: horizontal start index horizontal_end: horizontal end index vertical_start: vertical start index @@ -1126,7 +1286,7 @@ def compute_weighted_cell_neighbor_sum( _compute_weighted_cell_neighbor_sum( field=maxslp, c_bln_avg=c_bln_avg, - out=z_maxslp_avg, + out=maxslp_avg, domain={ dims.CellDim: (horizontal_start, horizontal_end), dims.KDim: (vertical_start, vertical_end), @@ -1136,7 +1296,7 @@ def compute_weighted_cell_neighbor_sum( _compute_weighted_cell_neighbor_sum( field=maxhgtd, c_bln_avg=c_bln_avg, - out=z_maxhgtd_avg, + out=maxhgtd_avg, domain={ dims.CellDim: (horizontal_start, horizontal_end), dims.KDim: (vertical_start, vertical_end), @@ -1146,17 +1306,17 @@ def compute_weighted_cell_neighbor_sum( @field_operator def _compute_max_nbhgt( - z_mc_nlev: gtx.Field[gtx.Dims[dims.CellDim], wpfloat], -) -> gtx.Field[gtx.Dims[dims.CellDim], wpfloat]: + z_mc_nlev: fa.CellField[wpfloat], +) -> fa.CellField[wpfloat]: max_nbhgt_0_1 = maximum(z_mc_nlev(C2E2C[0]), z_mc_nlev(C2E2C[1])) max_nbhgt = maximum(max_nbhgt_0_1, z_mc_nlev(C2E2C[2])) return max_nbhgt -@program +@program(grid_type=GridType.UNSTRUCTURED) def compute_max_nbhgt( - z_mc_nlev: gtx.Field[gtx.Dims[dims.CellDim], wpfloat], - max_nbhgt: gtx.Field[gtx.Dims[dims.CellDim], wpfloat], + z_mc_nlev: fa.CellField[wpfloat], + max_nbhgt: fa.CellField[wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, ): @@ -1196,7 +1356,68 @@ def _compute_param( @field_operator(grid_type=GridType.UNSTRUCTURED) def _compute_z_ifc_off_koff( - z_ifc_off: gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], wpfloat], -) -> gtx.Field[gtx.Dims[dims.EdgeDim, dims.KDim], wpfloat]: + z_ifc_off: fa.EdgeKField[wpfloat], +) -> fa.EdgeKField[wpfloat]: n = z_ifc_off(Koff[1]) return n + + +@field_operator +def _compute_theta_exner_ref_mc( + z_mc: fa.CellKField[wpfloat], + t0sl_bg: wpfloat, + del_t_bg: wpfloat, + h_scal_bg: wpfloat, + grav: wpfloat, + rd: wpfloat, + p0sl_bg: wpfloat, + rd_o_cpd: wpfloat, + p0ref: wpfloat, +): + z_aux1 = p0sl_bg * exp( + -grav + / rd + * h_scal_bg + / (t0sl_bg - del_t_bg) + * log((exp(z_mc / h_scal_bg) * (t0sl_bg - del_t_bg) + del_t_bg) / t0sl_bg) + ) + exner_ref_mc = (z_aux1 / p0ref) ** rd_o_cpd + z_temp = (t0sl_bg - del_t_bg) + del_t_bg * exp(-z_mc / h_scal_bg) + theta_ref_mc = z_temp / exner_ref_mc + return exner_ref_mc, theta_ref_mc + + +@program(grid_type=GridType.UNSTRUCTURED) +def compute_theta_exner_ref_mc( + z_mc: fa.CellKField[wpfloat], + exner_ref_mc: fa.CellKField[wpfloat], + theta_ref_mc: fa.CellKField[wpfloat], + t0sl_bg: wpfloat, + del_t_bg: wpfloat, + h_scal_bg: wpfloat, + grav: wpfloat, + rd: wpfloat, + p0sl_bg: wpfloat, + rd_o_cpd: wpfloat, + p0ref: wpfloat, + horizontal_start: int32, + horizontal_end: int32, + vertical_start: int32, + vertical_end: int32, +): + _compute_theta_exner_ref_mc( + z_mc=z_mc, + t0sl_bg=t0sl_bg, + del_t_bg=del_t_bg, + h_scal_bg=h_scal_bg, + grav=grav, + rd=rd, + p0sl_bg=p0sl_bg, + rd_o_cpd=rd_o_cpd, + p0ref=p0ref, + out=(exner_ref_mc, theta_ref_mc), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py new file mode 100644 index 000000000..78695931c --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py @@ -0,0 +1,400 @@ +# 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 + +from typing import Final + +import gt4py.next as gtx + +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.states import model + + +# TODO: revise names with domain scientists + +Z_MC: Final[str] = "height" +DDQZ_Z_HALF: Final[str] = "functional_determinant_of_metrics_on_interface_levels" +DDQZ_Z_FULL: Final[str] = "ddqz_z_full" +INV_DDQZ_Z_FULL: Final[str] = "inv_ddqz_z_full" +SCALFAC_DD3D: Final[str] = "scaling_factor_for_3d_divergence_damping" +RAYLEIGH_W: Final[str] = "rayleigh_w" +COEFF1_DWDZ: Final[str] = "coeff1_dwdz" +COEFF2_DWDZ: Final[str] = "coeff2_dwdz" +EXNER_REF_MC: Final[str] = "exner_ref_mc" +THETA_REF_MC: Final[str] = "theta_ref_mc" +D2DEXDZ2_FAC1_MC: Final[str] = "d2dexdz2_fac1_mc" +D2DEXDZ2_FAC2_MC: Final[str] = "d2dexdz2_fac2_mc" +VERT_OUT: Final[str] = "vert_out" +DDXT_Z_HALF_E: Final[str] = "ddxt_z_half_e" +DDXN_Z_HALF_E: Final[str] = "ddxn_z_half_e" +DDXN_Z_FULL: Final[str] = "ddxn_z_full" +VWIND_IMPL_WGT: Final[str] = "vwind_impl_wgt" +VWIND_EXPL_WGT: Final[str] = "vwind_expl_wgt" +EXNER_EXFAC: Final[str] = "exner_exfac" +WGTFAC_C: Final[str] = "wgtfac_c" +WGTFAC_E: Final[str] = "wgtfac_e" +FLAT_IDX_MAX: Final[str] = "flat_idx_max" +PG_EDGEIDX: Final[str] = "pg_edgeidx" +PG_VERTIDX: Final[str] = "pg_vertidx" +PG_EDGEIDX_DSL: Final[str] = "pg_edgeidx_dsl" +PG_EDGEDIST_DSL: Final[str] = "pg_exdist_dsl" +MASK_PROG_HALO_C: Final[str] = "mask_prog_halo_c" +BDY_HALO_C: Final[str] = "bdy_halo_c" +HMASK_DD3D: Final[str] = "hmask_dd3d" +ZDIFF_GRADP: Final[str] = "zdiff_gradp" +COEFF_GRADEKIN: Final[str] = "coeff_gradekin" +WGTFACQ_C: Final[str] = "weighting_factor_for_quadratic_interpolation_to_cell_surface" +WGTFACQ_E: Final[str] = "weighting_factor_for_quadratic_interpolation_to_edge_center" +MAXSLP: Final[str] = "maxslp" +MAXHGTD: Final[str] = "maxhgtd" +MAXSLP_AVG: Final[str] = "maxslp_avg" +MAXHGTD_AVG: Final[str] = "maxhgtd_avg" +MAX_NBHGT: Final[str] = "max_nbhgt" +MASK_HDIFF: Final[str] = "mask_hdiff" +ZD_DIFFCOEF_DSL: Final[str] = "zd_diffcoef_dsl" +ZD_INTCOEF_DSL: Final[str] = "zd_intcoef_dsl" +ZD_VERTOFFSET_DSL: Final[str] = "zd_vertoffset_dsl" + + +attrs: dict[str, model.FieldMetaData] = { + Z_MC: dict( + standard_name=Z_MC, + long_name="height", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="z_mc", + dtype=ta.wpfloat, + ), + DDQZ_Z_HALF: dict( + standard_name=DDQZ_Z_HALF, + long_name="functional_determinant_of_metrics_on_interface_levels", + units="", + dims=(dims.CellDim, dims.KHalfDim), + icon_var_name="ddqz_z_half", + dtype=ta.wpfloat, + ), + DDQZ_Z_FULL: dict( + standard_name=DDQZ_Z_FULL, + long_name="ddqz_z_full", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="ddqz_z_full", + dtype=ta.wpfloat, + ), + INV_DDQZ_Z_FULL: dict( + standard_name=INV_DDQZ_Z_FULL, + long_name="inv_ddqz_z_full", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="inv_ddqz_z_full", + dtype=ta.wpfloat, + ), + SCALFAC_DD3D: dict( + standard_name=SCALFAC_DD3D, + long_name="Scaling factor for 3D divergence damping terms", + units="", + dims=(dims.KDim,), + icon_var_name="scalfac_dd3d", + dtype=ta.wpfloat, + ), + RAYLEIGH_W: dict( + standard_name=RAYLEIGH_W, + long_name="rayleigh_w", + units="", + dims=(dims.KHalfDim,), + icon_var_name="rayleigh_w", + dtype=ta.wpfloat, + ), + COEFF1_DWDZ: dict( + standard_name=COEFF1_DWDZ, + long_name="coeff1_dwdz", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="coeff1_dwdz", + dtype=ta.wpfloat, + ), + COEFF2_DWDZ: dict( + standard_name=COEFF2_DWDZ, + long_name="coeff2_dwdz", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="coeff2_dwdz", + dtype=ta.wpfloat, + ), + EXNER_REF_MC: dict( + standard_name=EXNER_REF_MC, + long_name="exner_ref_mc", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="exner_ref_mc", + dtype=ta.wpfloat, + ), + THETA_REF_MC: dict( + standard_name=THETA_REF_MC, + long_name="theta_ref_mc", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="theta_ref_mc", + dtype=ta.wpfloat, + ), + D2DEXDZ2_FAC1_MC: dict( + standard_name=D2DEXDZ2_FAC1_MC, + long_name="d2dexdz2_fac1_mc", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="d2dexdz2_fac1_mc", + dtype=ta.wpfloat, + ), + D2DEXDZ2_FAC2_MC: dict( + standard_name=D2DEXDZ2_FAC2_MC, + long_name="d2dexdz2_fac2_mc", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="d2dexdz2_fac2_mc", + dtype=ta.wpfloat, + ), + VERT_OUT: dict( + standard_name=VERT_OUT, + long_name="vert_out", + units="", + dims=(dims.VertexDim, dims.KHalfDim), + icon_var_name="vert_out", + dtype=ta.wpfloat, + ), + DDXT_Z_HALF_E: dict( + standard_name=DDXT_Z_HALF_E, + long_name="ddxt_z_half_e", + units="", + dims=(dims.EdgeDim, dims.KHalfDim), + icon_var_name="ddxt_z_half_e", + dtype=ta.wpfloat, + ), + DDXN_Z_HALF_E: dict( + standard_name=DDXN_Z_HALF_E, + long_name="ddxn_z_half_e", + units="", + dims=(dims.EdgeDim, dims.KHalfDim), + icon_var_name="ddxn_z_half_e", + dtype=ta.wpfloat, + ), + DDXN_Z_FULL: dict( + standard_name=DDXN_Z_FULL, + long_name="ddxn_z_full", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="ddxn_z_full", + dtype=ta.wpfloat, + ), + VWIND_IMPL_WGT: dict( + standard_name=VWIND_IMPL_WGT, + long_name="vwind_impl_wgt", + units="", + dims=(dims.CellDim,), + icon_var_name="vwind_impl_wgt", + dtype=ta.wpfloat, + ), + VWIND_EXPL_WGT: dict( + standard_name=VWIND_EXPL_WGT, + long_name="vwind_expl_wgt", + units="", + dims=(dims.CellDim,), + icon_var_name="vwind_expl_wgt", + dtype=ta.wpfloat, + ), + EXNER_EXFAC: dict( + standard_name=EXNER_EXFAC, + long_name="exner_exfac", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="exner_exfac", + dtype=ta.wpfloat, + ), + WGTFAC_C: dict( + standard_name=WGTFAC_C, + long_name="wgtfac_c", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="wgtfac_c", + dtype=ta.wpfloat, + ), + WGTFAC_E: dict( + standard_name=WGTFAC_E, + long_name="wgtfac_e", + units="", + dims=(dims.EdgeDim, dims.KHalfDim), + icon_var_name="wgtfac_e", + dtype=ta.wpfloat, + ), + FLAT_IDX_MAX: dict( + standard_name=FLAT_IDX_MAX, + long_name="flat_idx_max", + units="", + dims=(dims.EdgeDim,), + icon_var_name="flat_idx_max", + dtype=ta.wpfloat, + ), + PG_EDGEIDX: dict( + standard_name=PG_EDGEIDX, + long_name="pg_edgeidx", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="pg_edgeidx", + dtype=gtx.int32, + ), + PG_VERTIDX: dict( + standard_name=PG_VERTIDX, + long_name="pg_vertidx", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="pg_vertidx", + dtype=gtx.int32, + ), + PG_EDGEIDX_DSL: dict( + standard_name=PG_EDGEIDX_DSL, + long_name="pg_edgeidx_dsl", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="pg_edgeidx_dsl", + dtype=bool, + ), + PG_EDGEDIST_DSL: dict( + standard_name=PG_EDGEDIST_DSL, + long_name="pg_exdist_dsl", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="pg_exdist_dsl", + dtype=ta.wpfloat, + ), + MASK_PROG_HALO_C: dict( + standard_name=MASK_PROG_HALO_C, + long_name="mask_prog_halo_c", + units="", + dims=(dims.CellDim,), + icon_var_name="mask_prog_halo_c", + dtype=bool, + ), + BDY_HALO_C: dict( + standard_name=BDY_HALO_C, + long_name="bdy_halo_c", + units="", + dims=(dims.CellDim,), + icon_var_name="bdy_halo_c", + dtype=bool, + ), + HMASK_DD3D: dict( + standard_name=HMASK_DD3D, + long_name="hmask_dd3d", + units="", + dims=(dims.EdgeDim,), + icon_var_name="hmask_dd3d", + dtype=ta.wpfloat, + ), + ZDIFF_GRADP: dict( + standard_name=ZDIFF_GRADP, + long_name="zdiff_gradp", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="zdiff_gradp", + dtype=ta.wpfloat, + ), + COEFF_GRADEKIN: dict( + standard_name=COEFF_GRADEKIN, + long_name="coeff_gradekin", + units="", + dims=(dims.EdgeDim,), + icon_var_name="coeff_gradekin", + dtype=ta.wpfloat, + ), + WGTFACQ_C: dict( + standard_name=WGTFACQ_C, + long_name="weighting_factor_for_quadratic_interpolation_to_cell_surface", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="weighting_factor_for_quadratic_interpolation_to_cell_surface", + dtype=ta.wpfloat, + ), + WGTFACQ_E: dict( + standard_name=WGTFACQ_E, + long_name="weighting_factor_for_quadratic_interpolation_to_edge_center", + units="", + dims=(dims.EdgeDim, dims.KDim), + icon_var_name="weighting_factor_for_quadratic_interpolation_to_edge_center", + dtype=ta.wpfloat, + ), + MAXSLP: dict( + standard_name=MAXSLP, + long_name="maxslp", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="maxslp", + dtype=ta.wpfloat, + ), + MAXHGTD: dict( + standard_name=MAXHGTD, + long_name="maxhgtd", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="maxhgtd", + dtype=ta.wpfloat, + ), + MAXSLP_AVG: dict( + standard_name=MAXSLP_AVG, + long_name="maxslp_avg", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="maxslp_avg", + dtype=ta.wpfloat, + ), + MAXHGTD_AVG: dict( + standard_name=MAXHGTD_AVG, + long_name="maxhgtd_avg", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="maxhgtd_avg", + dtype=ta.wpfloat, + ), + MAX_NBHGT: dict( + standard_name=MAX_NBHGT, + long_name="max_nbhgt", + units="", + dims=(dims.CellDim,), + icon_var_name="max_nbhgt", + dtype=ta.wpfloat, + ), + MASK_HDIFF: dict( + standard_name=MASK_HDIFF, + long_name="mask_hdiff", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="mask_hdiff", + dtype=ta.wpfloat, + ), + ZD_DIFFCOEF_DSL: dict( + standard_name=ZD_DIFFCOEF_DSL, + long_name="zd_diffcoef_dsl", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="zd_diffcoef_dsl", + dtype=ta.wpfloat, + ), + ZD_INTCOEF_DSL: dict( + standard_name=ZD_INTCOEF_DSL, + long_name="zd_intcoef_dsl", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="zd_intcoef_dsl", + dtype=ta.wpfloat, + ), + ZD_VERTOFFSET_DSL: dict( + standard_name=ZD_VERTOFFSET_DSL, + long_name="zd_vertoffset_dsl", + units="", + dims=(dims.CellDim, dims.KDim), + icon_var_name="zd_vertoffset_dsl", + dtype=ta.wpfloat, + ), +} diff --git a/model/common/src/icon4py/model/common/metrics/metrics_factory.py b/model/common/src/icon4py/model/common/metrics/metrics_factory.py new file mode 100644 index 000000000..a34071f15 --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/metrics_factory.py @@ -0,0 +1,790 @@ +# 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 functools +import math + +import gt4py.next as gtx +import numpy as np +from gt4py.next import backend as gtx_backend + +from icon4py.model.common import constants, dimension as dims +from icon4py.model.common.decomposition import definitions +from icon4py.model.common.grid import ( + geometry, + geometry_attributes as geometry_attrs, + horizontal as h_grid, + icon, + vertical as v_grid, +) +from icon4py.model.common.grid.vertical import VerticalGrid +from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory +from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( + compute_cell_2_vertex_interpolation, +) +from icon4py.model.common.metrics import ( + compute_coeff_gradekin, + compute_diffusion_metrics, + compute_flat_idx_max, + compute_vwind_impl_wgt, + compute_wgtfac_c, + compute_wgtfacq, + compute_zdiff_gradp_dsl, + metric_fields as mf, + metrics_attributes as attrs, +) +from icon4py.model.common.states import factory, model +from icon4py.model.common.utils import data_allocation as data_alloc + + +cell_domain = h_grid.domain(dims.CellDim) +edge_domain = h_grid.domain(dims.EdgeDim) +vertex_domain = h_grid.domain(dims.VertexDim) +vertical_domain = v_grid.domain(dims.KDim) +vertical_half_domain = v_grid.domain(dims.KHalfDim) + + +class MetricsFieldsFactory(factory.FieldSource, factory.GridProvider): + def __init__( + self, + grid: icon.IconGrid, + vertical_grid: VerticalGrid, + decomposition_info: definitions.DecompositionInfo, + geometry_source: geometry.GridGeometry, + interpolation_source: interpolation_factory.InterpolationFieldsFactory, + backend: gtx_backend.Backend, + metadata: dict[str, model.FieldMetaData], + interface_model_height: gtx.Field, + e_refin_ctrl: gtx.Field, + c_refin_ctrl: gtx.Field, + damping_height: float, + rayleigh_type: int, + rayleigh_coeff: float, + exner_expol: float, + vwind_offctr: float, + ): + self._backend = backend + self._xp = data_alloc.import_array_ns(backend) + self._allocator = gtx.constructors.zeros.partial(allocator=backend) + self._grid = grid + self._vertical_grid = vertical_grid + self._decomposition_info = decomposition_info + self._attrs = metadata + self._providers: dict[str, factory.FieldProvider] = {} + self._geometry = geometry_source + self._interpolation_source = interpolation_source + + vct_a = self._vertical_grid.vct_a + vct_a_1 = vct_a.asnumpy()[0] + self._config = { + "divdamp_trans_start": 12500.0, + "divdamp_trans_end": 17500.0, + "divdamp_type": 3, + "damping_height": damping_height, + "rayleigh_type": rayleigh_type, + "rayleigh_coeff": rayleigh_coeff, + "exner_expol": exner_expol, + "vwind_offctr": vwind_offctr, + "igradp_method": 3, + "igradp_constant": 3, + "thslp_zdiffu": 0.02, + "thhgtd_zdiffu": 125.0, + "vct_a_1": vct_a_1, + } + z_ifc_sliced = gtx.as_field( + (dims.CellDim,), interface_model_height.asnumpy()[:, self._grid.num_levels] + ) + k_index = gtx.as_field((dims.KDim,), np.arange(self._grid.num_levels + 1, dtype=gtx.int32)) + e_lev = gtx.as_field((dims.EdgeDim,), np.arange(self._grid.num_edges, dtype=gtx.int32)) + e_owner_mask = gtx.as_field( + (dims.EdgeDim,), self._decomposition_info.owner_mask(dims.EdgeDim) + ) + c_owner_mask = gtx.as_field( + (dims.CellDim,), self._decomposition_info.owner_mask(dims.CellDim) + ) + + self.register_provider( + factory.PrecomputedFieldProvider( + { + "height_on_interface_levels": interface_model_height, + "z_ifc_sliced": z_ifc_sliced, + "vct_a": vct_a, + "c_refin_ctrl": c_refin_ctrl, + "e_refin_ctrl": e_refin_ctrl, + "e_owner_mask": e_owner_mask, + "c_owner_mask": c_owner_mask, + "k_lev": k_index, + "e_lev": e_lev, + } + ) + ) + self._register_computed_fields() + + def __repr__(self): + return f"{self.__class__.__name__} on (grid={self._grid!r}) providing fields f{self.metadata.keys()}" + + @property + def _sources(self) -> factory.FieldSource: + return factory.CompositeSource(self, (self._geometry, self._interpolation_source)) + + def _register_computed_fields(self): + height = factory.ProgramFieldProvider( + func=mf.compute_z_mc.with_backend(self._backend), + domain={ + dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.END)), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={"z_mc": attrs.Z_MC}, + deps={"z_ifc": "height_on_interface_levels"}, + ) + self.register_provider(height) + + compute_ddqz_z_half = factory.ProgramFieldProvider( + func=mf.compute_ddqz_z_half.with_backend(self._backend), + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LOCAL), + cell_domain(h_grid.Zone.END), + ), + dims.KHalfDim: ( + vertical_half_domain(v_grid.Zone.TOP), + vertical_half_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={"ddqz_z_half": attrs.DDQZ_Z_HALF}, + deps={ + "z_ifc": "height_on_interface_levels", + "z_mc": attrs.Z_MC, + "k": "k_lev", + }, + params={"nlev": self._grid.num_levels}, + ) + self.register_provider(compute_ddqz_z_half) + + ddqz_z_full_and_inverse = factory.ProgramFieldProvider( + func=mf.compute_ddqz_z_full_and_inverse.with_backend(self._backend), + deps={"z_ifc": "height_on_interface_levels"}, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LOCAL), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={"ddqz_z_full": attrs.DDQZ_Z_FULL, "inv_ddqz_z_full": attrs.INV_DDQZ_Z_FULL}, + ) + self.register_provider(ddqz_z_full_and_inverse) + + compute_scalfac_dd3d = factory.ProgramFieldProvider( + func=mf.compute_scalfac_dd3d.with_backend(self._backend), + domain={ + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ) + }, + fields={"scalfac_dd3d": attrs.SCALFAC_DD3D}, + deps={"vct_a": "vct_a"}, + params={ + "divdamp_trans_start": self._config["divdamp_trans_start"], + "divdamp_trans_end": self._config["divdamp_trans_end"], + "divdamp_type": self._config["divdamp_type"], + }, + ) + self.register_provider(compute_scalfac_dd3d) + + compute_rayleigh_w = factory.ProgramFieldProvider( + func=mf.compute_rayleigh_w.with_backend(self._backend), + deps={"vct_a": "vct_a"}, + domain={ + dims.KHalfDim: ( + vertical_domain(v_grid.Zone.TOP), + v_grid.Domain(dims.KHalfDim, v_grid.Zone.DAMPING, 1), + ) + }, + fields={"rayleigh_w": attrs.RAYLEIGH_W}, + params={ + "damping_height": self._config["damping_height"], + "rayleigh_type": self._config["rayleigh_type"], + "rayleigh_classic": constants.RayleighType.CLASSIC, + "rayleigh_klemp": constants.RayleighType.KLEMP, + "rayleigh_coeff": self._config["rayleigh_coeff"], + "vct_a_1": self._config["vct_a_1"], + "pi_const": math.pi, + }, + ) + self.register_provider(compute_rayleigh_w) + + compute_coeff_dwdz = factory.ProgramFieldProvider( + func=mf.compute_coeff_dwdz.with_backend(self._backend), + deps={ + "ddqz_z_full": attrs.DDQZ_Z_FULL, + "z_ifc": "height_on_interface_levels", + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LOCAL), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + v_grid.Domain(dims.KHalfDim, v_grid.Zone.TOP, 1), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={"coeff1_dwdz": attrs.COEFF1_DWDZ, "coeff2_dwdz": attrs.COEFF2_DWDZ}, + ) + self.register_provider(compute_coeff_dwdz) + + compute_theta_exner_ref_mc = factory.ProgramFieldProvider( + func=mf.compute_theta_exner_ref_mc.with_backend(self._backend), + deps={ + "z_mc": attrs.Z_MC, + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LOCAL), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={"exner_ref_mc": attrs.EXNER_REF_MC, "theta_ref_mc": attrs.THETA_REF_MC}, + params={ + "t0sl_bg": constants.SEA_LEVEL_TEMPERATURE, + "del_t_bg": constants.DELTA_TEMPERATURE, + "h_scal_bg": constants._H_SCAL_BG, + "grav": constants.GRAV, + "rd": constants.RD, + "p0sl_bg": constants.SEAL_LEVEL_PRESSURE, + "rd_o_cpd": constants.RD_O_CPD, + "p0ref": constants.REFERENCE_PRESSURE, + }, + ) + self.register_provider(compute_theta_exner_ref_mc) + + compute_d2dexdz2_fac_mc = factory.ProgramFieldProvider( + func=mf.compute_d2dexdz2_fac_mc.with_backend(self._backend), + deps={ + "theta_ref_mc": attrs.THETA_REF_MC, + "inv_ddqz_z_full": attrs.INV_DDQZ_Z_FULL, + "exner_ref_mc": attrs.EXNER_REF_MC, + "z_mc": attrs.Z_MC, + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LOCAL), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={ + attrs.D2DEXDZ2_FAC1_MC: attrs.D2DEXDZ2_FAC1_MC, + attrs.D2DEXDZ2_FAC2_MC: attrs.D2DEXDZ2_FAC2_MC, + }, + params={ + "cpd": constants.CPD, + "grav": constants.GRAV, + "del_t_bg": constants.DEL_T_BG, + "h_scal_bg": constants._H_SCAL_BG, + "igradp_method": self._config["igradp_method"], + "igradp_constant": self._config["igradp_constant"], + }, + ) + self.register_provider(compute_d2dexdz2_fac_mc) + + compute_cell_to_vertex_interpolation = factory.ProgramFieldProvider( + func=compute_cell_2_vertex_interpolation.with_backend(self._backend), + deps={ + "cell_in": "height_on_interface_levels", + "c_int": interpolation_attributes.CELL_AW_VERTS, + }, + domain={ + dims.VertexDim: ( + vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), + vertex_domain(h_grid.Zone.INTERIOR), + ), + dims.KHalfDim: ( + vertical_half_domain(v_grid.Zone.TOP), + vertical_half_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.VERT_OUT: attrs.VERT_OUT}, + ) + self.register_provider(compute_cell_to_vertex_interpolation) + + compute_ddxt_z_half_e = factory.ProgramFieldProvider( + func=mf.compute_ddxt_z_half_e.with_backend(self._backend), + deps={ + "cell_in": "height_on_interface_levels", + "c_int": interpolation_attributes.CELL_AW_VERTS, + "inv_primal_edge_length": f"inverse_of_{geometry_attrs.EDGE_LENGTH}", + "tangent_orientation": geometry_attrs.TANGENT_ORIENTATION, + }, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3), + edge_domain(h_grid.Zone.INTERIOR), + ), + dims.KHalfDim: ( + vertical_half_domain(v_grid.Zone.TOP), + vertical_half_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.DDXT_Z_HALF_E: attrs.DDXT_Z_HALF_E}, + ) + self.register_provider(compute_ddxt_z_half_e) + + compute_ddxn_z_half_e = factory.ProgramFieldProvider( + func=mf.compute_ddxn_z_half_e.with_backend(self._backend), + deps={ + "z_ifc": "height_on_interface_levels", + "inv_dual_edge_length": f"inverse_of_{geometry_attrs.DUAL_EDGE_LENGTH}", + }, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), + edge_domain(h_grid.Zone.INTERIOR), + ), + dims.KHalfDim: ( + vertical_half_domain(v_grid.Zone.TOP), + vertical_half_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.DDXN_Z_HALF_E: attrs.DDXN_Z_HALF_E}, + ) + self.register_provider(compute_ddxn_z_half_e) + + compute_ddxn_z_full = factory.ProgramFieldProvider( + func=mf.compute_ddxn_z_full.with_backend(self._backend), + deps={ + "ddxnt_z_half_e": attrs.DDXN_Z_HALF_E, + }, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.LOCAL), + edge_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.DDXN_Z_FULL: attrs.DDXN_Z_FULL}, + ) + self.register_provider(compute_ddxn_z_full) + + compute_vwind_impl_wgt_np = factory.NumpyFieldsProvider( + func=functools.partial( + compute_vwind_impl_wgt.compute_vwind_impl_wgt, array_ns=self._xp + ), + domain=(dims.CellDim,), + connectivities={"c2e": dims.C2EDim}, + fields=(attrs.VWIND_IMPL_WGT,), + deps={ + "vct_a": "vct_a", + "z_ifc": "height_on_interface_levels", + "z_ddxn_z_half_e": attrs.DDXN_Z_HALF_E, + "z_ddxt_z_half_e": attrs.DDXT_Z_HALF_E, + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + }, + params={ + "vwind_offctr": self._config["vwind_offctr"], + "nlev": self._grid.num_levels, + "horizontal_start_cell": self._grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "n_cells": self._grid.num_cells, + }, + ) + self.register_provider(compute_vwind_impl_wgt_np) + + compute_vwind_expl_wgt = factory.ProgramFieldProvider( + func=mf.compute_vwind_expl_wgt.with_backend(self._backend), + deps={ + attrs.VWIND_IMPL_WGT: attrs.VWIND_IMPL_WGT, + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LOCAL), + cell_domain(h_grid.Zone.END), + ), + }, + fields={"vwind_expl_wgt": attrs.VWIND_EXPL_WGT}, + ) + self.register_provider(compute_vwind_expl_wgt) + + compute_exner_exfac = factory.ProgramFieldProvider( + func=mf.compute_exner_exfac.with_backend(self._backend), + deps={ + "ddxn_z_full": attrs.DDXN_Z_FULL, + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.EXNER_EXFAC: attrs.EXNER_EXFAC}, + params={"exner_expol": self._config["exner_expol"]}, + ) + self.register_provider(compute_exner_exfac) + + compute_wgtfac_c_np = factory.ProgramFieldProvider( + func=compute_wgtfac_c.compute_wgtfac_c.with_backend(self._backend), + deps={ + "z_ifc": "height_on_interface_levels", + "k": "k_lev", + }, + domain={ + dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.END)), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.WGTFAC_C: attrs.WGTFAC_C}, + params={"nlev": self._grid.num_levels}, + ) + self.register_provider(compute_wgtfac_c_np) + + compute_wgtfac_e = factory.ProgramFieldProvider( + func=mf.compute_wgtfac_e.with_backend(self._backend), + deps={ + attrs.WGTFAC_C: attrs.WGTFAC_C, + "c_lin_e": interpolation_attributes.C_LIN_E, + }, + domain={ + dims.CellDim: ( + edge_domain(h_grid.Zone.LOCAL), + edge_domain(h_grid.Zone.LOCAL), + ), + dims.KHalfDim: ( + vertical_half_domain(v_grid.Zone.TOP), + vertical_half_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.WGTFAC_E: attrs.WGTFAC_E}, + ) + self.register_provider(compute_wgtfac_e) + + compute_flat_idx_max_np = factory.NumpyFieldsProvider( + func=functools.partial(compute_flat_idx_max.compute_flat_idx_max, array_ns=self._xp), + domain=(dims.EdgeDim,), + fields=(attrs.FLAT_IDX_MAX,), + deps={ + "z_mc": attrs.Z_MC, + "c_lin_e": interpolation_attributes.C_LIN_E, + "z_ifc": "height_on_interface_levels", + "k_lev": "k_lev", + }, + connectivities={"e2c": dims.E2CDim}, + params={ + "horizontal_lower": self._grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) + ), + "horizontal_upper": self._grid.end_index(edge_domain(h_grid.Zone.LOCAL)), + }, + ) + self.register_provider(compute_flat_idx_max_np) + + compute_pg_edgeidx_vertidx = factory.ProgramFieldProvider( + func=mf.compute_pg_edgeidx_vertidx.with_backend(self._backend), + deps={ + "c_lin_e": interpolation_attributes.C_LIN_E, + "z_ifc": "height_on_interface_levels", + "z_ifc_sliced": "z_ifc_sliced", + "e_owner_mask": "e_owner_mask", + "flat_idx_max": attrs.FLAT_IDX_MAX, + "e_lev": "e_lev", + "k_lev": "k_lev", + }, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.NUDGING), + edge_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.PG_EDGEIDX: attrs.PG_EDGEIDX, attrs.PG_VERTIDX: attrs.PG_VERTIDX}, + ) + self.register_provider(compute_pg_edgeidx_vertidx) + + compute_pg_edgeidx_dsl = factory.ProgramFieldProvider( + func=mf.compute_pg_edgeidx_dsl.with_backend(self._backend), + deps={"pg_edgeidx": attrs.PG_EDGEIDX, "pg_vertidx": attrs.PG_VERTIDX}, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.NUDGING_LEVEL_2), + edge_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={"pg_edgeidx_dsl": attrs.PG_EDGEIDX_DSL}, + ) + self.register_provider(compute_pg_edgeidx_dsl) + + compute_pg_exdist_dsl = factory.ProgramFieldProvider( + func=mf.compute_pg_exdist_dsl.with_backend(self._backend), + deps={ + "z_ifc_sliced": "z_ifc_sliced", + "z_mc": attrs.Z_MC, + "c_lin_e": interpolation_attributes.C_LIN_E, + "e_owner_mask": "e_owner_mask", + "flat_idx_max": attrs.FLAT_IDX_MAX, + "k_lev": "k_lev", + "e_lev": "e_lev", + }, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.NUDGING), + edge_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + params={ + "h_start_zaux2": self._grid.end_index(edge_domain(h_grid.Zone.NUDGING)), + "h_end_zaux2": self._grid.end_index(edge_domain(h_grid.Zone.LOCAL)), + }, + fields={"pg_exdist_dsl": attrs.PG_EDGEDIST_DSL}, + ) + self.register_provider(compute_pg_exdist_dsl) + + compute_mask_bdy_halo_c = factory.ProgramFieldProvider( + func=mf.compute_mask_bdy_halo_c.with_backend(self._backend), + deps={ + "c_refin_ctrl": "c_refin_ctrl", + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.HALO), + cell_domain(h_grid.Zone.HALO), + ), + }, + fields={ + attrs.MASK_PROG_HALO_C: attrs.MASK_PROG_HALO_C, + attrs.BDY_HALO_C: attrs.BDY_HALO_C, + }, + ) + self.register_provider(compute_mask_bdy_halo_c) + + compute_hmask_dd3d = factory.ProgramFieldProvider( + func=mf.compute_hmask_dd3d.with_backend(self._backend), + deps={ + "e_refin_ctrl": "e_refin_ctrl", + }, + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), + edge_domain(h_grid.Zone.LOCAL), + ) + }, + fields={attrs.HMASK_DD3D: attrs.HMASK_DD3D}, + params={ + "grf_nudge_start_e": gtx.int32(h_grid._GRF_NUDGEZONE_START_EDGES), + "grf_nudgezone_width": gtx.int32(h_grid._GRF_NUDGEZONE_WIDTH), + }, + ) + self.register_provider(compute_hmask_dd3d) + + compute_zdiff_gradp_dsl_np = factory.NumpyFieldsProvider( + func=functools.partial( + compute_zdiff_gradp_dsl.compute_zdiff_gradp_dsl, array_ns=self._xp + ), + deps={ + "z_mc": attrs.Z_MC, + "c_lin_e": interpolation_attributes.C_LIN_E, + "z_ifc": "height_on_interface_levels", + "flat_idx": attrs.FLAT_IDX_MAX, + "z_ifc_sliced": "z_ifc_sliced", + }, + connectivities={"e2c": dims.E2CDim}, + domain=(dims.EdgeDim, dims.KDim), + fields=(attrs.ZDIFF_GRADP,), + params={ + "nlev": self._grid.num_levels, + "horizontal_start": self._grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_start_1": self._grid.start_index( + edge_domain(h_grid.Zone.NUDGING_LEVEL_2) + ), + }, + ) + self.register_provider(compute_zdiff_gradp_dsl_np) + + compute_coeff_gradekin_np = factory.NumpyFieldsProvider( + func=functools.partial( + compute_coeff_gradekin.compute_coeff_gradekin, array_ns=self._xp + ), + domain=(dims.EdgeDim,), + fields=(attrs.COEFF_GRADEKIN,), + deps={ + "edge_cell_length": geometry_attrs.EDGE_CELL_DISTANCE, + "inv_dual_edge_length": f"inverse_of_{geometry_attrs.DUAL_EDGE_LENGTH}", + }, + params={ + "horizontal_start": self._grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "horizontal_end": self._grid.num_edges, + }, + ) + self.register_provider(compute_coeff_gradekin_np) + + compute_wgtfacq_c = factory.NumpyFieldsProvider( + func=functools.partial(compute_wgtfacq.compute_wgtfacq_c_dsl, array_ns=self._xp), + domain=(dims.CellDim, dims.KDim), + fields=(attrs.WGTFACQ_C,), + deps={"z_ifc": "height_on_interface_levels"}, + params={"nlev": self._grid.num_levels}, + ) + + self.register_provider(compute_wgtfacq_c) + + compute_wgtfacq_e = factory.NumpyFieldsProvider( + func=functools.partial(compute_wgtfacq.compute_wgtfacq_e_dsl, array_ns=self._xp), + deps={ + "z_ifc": "height_on_interface_levels", + "c_lin_e": interpolation_attributes.C_LIN_E, + "wgtfacq_c_dsl": attrs.WGTFACQ_C, + }, + connectivities={"e2c": dims.E2CDim}, + domain=(dims.EdgeDim, dims.KDim), + fields=(attrs.WGTFACQ_E,), + params={"n_edges": self._grid.num_edges, "nlev": self._grid.num_levels}, + ) + + self.register_provider(compute_wgtfacq_e) + + compute_maxslp_maxhgtd = factory.ProgramFieldProvider( + func=mf.compute_maxslp_maxhgtd.with_backend(self._backend), + deps={ + "ddxn_z_full": attrs.DDXN_Z_FULL, + "dual_edge_length": geometry_attrs.DUAL_EDGE_LENGTH, + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.MAXSLP: attrs.MAXSLP, attrs.MAXHGTD: attrs.MAXHGTD}, + ) + self.register_provider(compute_maxslp_maxhgtd) + + compute_weighted_cell_neighbor_sum = factory.ProgramFieldProvider( + func=mf.compute_weighted_cell_neighbor_sum, + deps={ + "maxslp": attrs.MAXSLP, + "maxhgtd": attrs.MAXHGTD, + "c_bln_avg": interpolation_attributes.C_BLN_AVG, + }, + domain={ + dims.CellDim: ( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2), + cell_domain(h_grid.Zone.END), + ), + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={attrs.MAXSLP_AVG: attrs.MAXSLP_AVG, attrs.MAXHGTD_AVG: attrs.MAXHGTD_AVG}, + ) + self.register_provider(compute_weighted_cell_neighbor_sum) + + compute_max_nbhgt = factory.NumpyFieldsProvider( + func=functools.partial( + compute_diffusion_metrics.compute_max_nbhgt_array_ns, array_ns=self._xp + ), + deps={ + "z_mc": attrs.Z_MC, + }, + connectivities={"c2e2c": dims.C2E2CDim}, + domain=(dims.CellDim,), + fields=(attrs.MAX_NBHGT,), + params={ + "nlev": self._grid.num_levels, + }, + ) + self.register_provider(compute_max_nbhgt) + + compute_diffusion_metrics_np = factory.NumpyFieldsProvider( + func=functools.partial( + compute_diffusion_metrics.compute_diffusion_metrics, array_ns=self._xp + ), + deps={ + "z_mc": attrs.Z_MC, + "max_nbhgt": attrs.MAX_NBHGT, + "c_owner_mask": "c_owner_mask", + "maxslp_avg": attrs.MAXSLP_AVG, + "maxhgtd_avg": attrs.MAXHGTD_AVG, + }, + connectivities={"c2e2c": dims.C2E2CDim}, + domain=(dims.CellDim, dims.KDim), + fields=( + attrs.MASK_HDIFF, + attrs.ZD_DIFFCOEF_DSL, + attrs.ZD_INTCOEF_DSL, + attrs.ZD_VERTOFFSET_DSL, + ), + params={ + "thslp_zdiffu": self._config["thslp_zdiffu"], + "thhgtd_zdiffu": self._config["thhgtd_zdiffu"], + "cell_nudging": self._grid.start_index( + h_grid.domain(dims.CellDim)(h_grid.Zone.NUDGING) + ), + "nlev": self._grid.num_levels, + }, + ) + + self.register_provider(compute_diffusion_metrics_np) + + @property + def metadata(self) -> dict[str, model.FieldMetaData]: + return self._attrs + + @property + def backend(self) -> gtx_backend.Backend: + return self._backend + + @property + def grid(self): + return self._grid + + @property + def vertical_grid(self): + return self._vertical_grid diff --git a/model/common/src/icon4py/model/common/states/factory.py b/model/common/src/icon4py/model/common/states/factory.py index 1f517ac74..b3c6a3f28 100644 --- a/model/common/src/icon4py/model/common/states/factory.py +++ b/model/common/src/icon4py/model/common/states/factory.py @@ -419,7 +419,7 @@ def _allocate( self, backend: gtx_backend.Backend, grid: base_grid.BaseGrid, # TODO @halungge: change to vertical grid - dtype: state_utils.ScalarType = ta.wpfloat, + dtype: dict[str, state_utils.ScalarType], ) -> dict[str, state_utils.FieldType]: def _map_size(dim: gtx.Dimension, grid: base_grid.BaseGrid) -> int: if dim == dims.KHalfDim: @@ -433,7 +433,7 @@ def _map_dim(dim: gtx.Dimension) -> gtx.Dimension: allocate = gtx.constructors.zeros.partial(allocator=backend) field_domain = {_map_dim(dim): (0, _map_size(dim, grid)) for dim in self._dims} - return {k: allocate(field_domain, dtype=dtype) for k in self._fields.keys()} + return {k: allocate(field_domain, dtype=dtype[k]) for k in self._fields.keys()} # TODO (@halungge) this can be simplified when completely disentangling vertical and horizontal grid. # the IconGrid should then only contain horizontal connectivities and no longer any Koff which should be moved to the VerticalGrid @@ -500,9 +500,9 @@ def _compute( ) -> None: try: metadata = {v: factory.get(v, RetrievalType.METADATA) for k, v in self._output.items()} - dtype = metadata["dtype"] + dtype = {v: metadata[v]["dtype"] for v in self._output.values()} except (ValueError, KeyError): - dtype = ta.wpfloat + dtype = {v: ta.wpfloat for v in self._output.values()} self._fields = self._allocate(backend, grid_provider.grid, dtype=dtype) deps = {k: factory.get(v) for k, v in self._dependencies.items()} diff --git a/model/common/src/icon4py/model/common/states/metadata.py b/model/common/src/icon4py/model/common/states/metadata.py index 7101b821a..da43033a1 100644 --- a/model/common/src/icon4py/model/common/states/metadata.py +++ b/model/common/src/icon4py/model/common/states/metadata.py @@ -18,20 +18,12 @@ INTERFACE_LEVEL_STANDARD_NAME: Final[str] = "interface_model_level_number" attrs: Final[dict[str, model.FieldMetaData]] = { - "functional_determinant_of_metrics_on_interface_levels": dict( - standard_name="functional_determinant_of_metrics_on_interface_levels", - long_name="functional determinant of the metrics [sqrt(gamma)] on half levels", + "z_ifv": dict( + standard_name="z_ifv", + long_name="z_ifv", units="", - dims=(dims.CellDim, dims.KHalfDim), - dtype=ta.wpfloat, - icon_var_name="ddqz_z_half", - ), - "height": dict( - standard_name="height", - long_name="height", - units="m", - dims=(dims.CellDim, dims.KDim), - icon_var_name="z_mc", + dims=(dims.VertexDim, dims.KDim), + icon_var_name="z_ifv", dtype=ta.wpfloat, ), "height_on_interface_levels": dict( @@ -42,6 +34,14 @@ icon_var_name="z_ifc", dtype=ta.wpfloat, ), + "z_ifc_sliced": dict( + standard_name="z_ifc_sliced", + long_name="z_ifc_sliced", + units="m", + dims=(dims.CellDim), + icon_var_name="z_ifc_sliced", + dtype=ta.wpfloat, + ), "model_level_number": dict( standard_name="model_level_number", long_name="model level number", @@ -82,14 +82,6 @@ icon_var_name="c_lin_e", long_name="coefficients for cell to edge interpolation", ), - "scaling_factor_for_3d_divergence_damping": dict( - standard_name="scaling_factor_for_3d_divergence_damping", - units="", - dims=(dims.KDim), - dtype=ta.wpfloat, - icon_var_name="scalfac_dd3d", - long_name="Scaling factor for 3D divergence damping terms", - ), "model_interface_height": dict( standard_name="model_interface_height", long_name="height value of half levels without topography", @@ -115,4 +107,52 @@ dims=(dims.EdgeDim,), icon_var_name="refin_e_ctrl", ), + "c_refin_ctrl": dict( + standard_name="c_refin_ctrl", + units="", + dims=(dims.CellDim,), + dtype=ta.wpfloat, + icon_var_name="c_refin_ctrl", + long_name="refinement control field on cells", + ), + "e_refin_ctrl": dict( + standard_name="e_refin_ctrl", + units="", + dims=(dims.EdgeDim,), + dtype=ta.wpfloat, + icon_var_name="e_refin_ctrl", + long_name="refinement contorl fields on edges", + ), + "cells_aw_verts_field": dict( + standard_name="cells_aw_verts_field", + units="", + dims=(dims.VertexDim, dims.V2CDim), + dtype=ta.wpfloat, + icon_var_name="cells_aw_verts_field", + long_name="grid savepoint field", + ), + "e_lev": dict( + standard_name="e_lev", + long_name="e_lev", + units="", + dims=(dims.EdgeDim,), + icon_var_name="e_lev", + dtype=gtx.int32, + ), + "e_owner_mask": dict( + standard_name="e_owner_mask", + units="", + dims=(dims.EdgeDim), + dtype=bool, + icon_var_name="e_owner_mask", + long_name="grid savepoint field", + ), + "c_owner_mask": dict( + standard_name="c_owner_mask", + units="", + dims=(dims.CellDim), + dtype=bool, + icon_var_name="c_owner_mask", + long_name="grid savepoint field", + ), } diff --git a/model/common/src/icon4py/model/common/states/model.py b/model/common/src/icon4py/model/common/states/model.py index 65e99ae75..dd5647d8d 100644 --- a/model/common/src/icon4py/model/common/states/model.py +++ b/model/common/src/icon4py/model/common/states/model.py @@ -21,6 +21,7 @@ DimensionNames = Literal["cell", "edge", "vertex"] DimensionT = Union[gtx.Dimension, DimensionNames] BufferT = Union[np_t.ArrayLike, gtx.Field] +DTypeT = Union[ta.wpfloat, ta.vpfloat, gtx.int32, gtx.int64, gtx.float32, gtx.float64] class OptionalMetaData(TypedDict, total=False): diff --git a/model/common/tests/grid_tests/test_grid_manager.py b/model/common/tests/grid_tests/test_grid_manager.py index efc62b61f..43c42c7d6 100644 --- a/model/common/tests/grid_tests/test_grid_manager.py +++ b/model/common/tests/grid_tests/test_grid_manager.py @@ -605,12 +605,12 @@ def test_tangent_orientation(grid_file, grid_savepoint, backend): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_cell_normal_orientation(grid_file, grid_savepoint, backend): - expected = grid_savepoint.edge_orientation() +def test_edge_orientation_on_vertex(grid_file, grid_savepoint, backend): + expected = grid_savepoint.vertex_edge_orientation() manager = _run_grid_manager(grid_file, backend=backend) geometry_fields = manager.geometry assert helpers.dallclose( - geometry_fields[GeometryName.CELL_NORMAL_ORIENTATION].ndarray, expected.ndarray + geometry_fields[GeometryName.EDGE_ORIENTATION_ON_VERTEX].asnumpy(), expected.asnumpy() ) @@ -622,12 +622,30 @@ def test_cell_normal_orientation(grid_file, grid_savepoint, backend): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_edge_orientation_on_vertex(grid_file, grid_savepoint, backend): - expected = grid_savepoint.vertex_edge_orientation() +def test_dual_area(grid_file, grid_savepoint, backend): + expected = grid_savepoint.vertex_dual_area() manager = _run_grid_manager(grid_file, backend=backend) geometry_fields = manager.geometry + assert helpers.dallclose(geometry_fields[GeometryName.DUAL_AREA].asnumpy(), expected.asnumpy()) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +def test_edge_cell_distance(grid_file, grid_savepoint, backend): + expected = grid_savepoint.edge_cell_length() + manager = _run_grid_manager(grid_file, backend=backend) + geometry_fields = manager.geometry + assert helpers.dallclose( - geometry_fields[GeometryName.EDGE_ORIENTATION_ON_VERTEX].ndarray, expected.ndarray + geometry_fields[GeometryName.EDGE_CELL_DISTANCE].asnumpy(), + expected.asnumpy(), + equal_nan=True, ) @@ -639,11 +657,13 @@ def test_edge_orientation_on_vertex(grid_file, grid_savepoint, backend): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_dual_area(grid_file, grid_savepoint, backend): - expected = grid_savepoint.vertex_dual_area() +def test_cell_normal_orientation(grid_file, grid_savepoint, backend): + expected = grid_savepoint.edge_orientation() manager = _run_grid_manager(grid_file, backend=backend) geometry_fields = manager.geometry - assert helpers.dallclose(geometry_fields[GeometryName.DUAL_AREA].ndarray, expected.ndarray) + assert helpers.dallclose( + geometry_fields[GeometryName.CELL_NORMAL_ORIENTATION].asnumpy(), expected.asnumpy() + ) @pytest.mark.datatest @@ -654,13 +674,13 @@ def test_dual_area(grid_file, grid_savepoint, backend): (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), ], ) -def test_edge_cell_distance(grid_file, grid_savepoint, backend): - expected = grid_savepoint.edge_cell_length() +def test_edge_vertex_distance(grid_file, grid_savepoint, backend): + expected = grid_savepoint.edge_vert_length() manager = _run_grid_manager(grid_file, backend=backend) geometry_fields = manager.geometry assert helpers.dallclose( - geometry_fields[GeometryName.EDGE_CELL_DISTANCE].asnumpy(), + geometry_fields[GeometryName.EDGE_VERTEX_DISTANCE].asnumpy(), expected.asnumpy(), equal_nan=True, ) diff --git a/model/common/tests/interpolation_tests/test_interpolation_factory.py b/model/common/tests/interpolation_tests/test_interpolation_factory.py index 5b2ceed73..5c565d369 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_factory.py +++ b/model/common/tests/interpolation_tests/test_interpolation_factory.py @@ -224,3 +224,21 @@ def test_get_mass_conserving_cell_average_weight( assert field.shape == (grid.num_cells, 4) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) + + +@pytest.mark.parametrize( + "grid_file, experiment, rtol", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT, 5e-9), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT, 1e-11), + ], +) +@pytest.mark.datatest +def test_cells_aw_verts(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.c_intp() + factory = get_interpolation_factory(backend, experiment, grid_file) + grid = factory.grid + field = factory.get(attrs.CELL_AW_VERTS) + + assert field.shape == (grid.num_vertices, 6) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) diff --git a/model/common/tests/interpolation_tests/test_interpolation_fields.py b/model/common/tests/interpolation_tests/test_interpolation_fields.py index 3c77f5d81..82ebe676f 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_fields.py +++ b/model/common/tests/interpolation_tests/test_interpolation_fields.py @@ -298,7 +298,7 @@ def test_compute_cells_aw_verts( e2c=e2c, horizontal_start=horizontal_start_vertex, ) - assert test_helpers.dallclose(cells_aw_verts, cells_aw_verts_ref, atol=1e-3) + assert test_helpers.dallclose(cells_aw_verts, cells_aw_verts_ref) @pytest.mark.datatest diff --git a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py index 2917ed0f1..30bc4e1ee 100644 --- a/model/common/tests/metric_tests/test_compute_coeff_gradekin.py +++ b/model/common/tests/metric_tests/test_compute_coeff_gradekin.py @@ -28,4 +28,4 @@ def test_compute_coeff_gradekin(icon_grid, grid_savepoint, metrics_savepoint): coeff_gradekin_full = compute_coeff_gradekin( edge_cell_length, inv_dual_edge_length, horizontal_start, horizontal_end ) - assert helpers.dallclose(coeff_gradekin_ref.asnumpy(), coeff_gradekin_full.asnumpy()) + assert helpers.dallclose(coeff_gradekin_ref.asnumpy(), coeff_gradekin_full) diff --git a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py index d172bdbd1..9f2f0bd10 100644 --- a/model/common/tests/metric_tests/test_compute_diffusion_metrics.py +++ b/model/common/tests/metric_tests/test_compute_diffusion_metrics.py @@ -36,28 +36,16 @@ def test_compute_diffusion_metrics( if experiment == dt_utils.GLOBAL_EXPERIMENT: pytest.skip(f"Fields not computed for {experiment}") - mask_hdiff = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=bool).asnumpy() - zd_vertoffset_dsl = data_alloc.zero_field( - icon_grid, dims.CellDim, dims.C2E2CDim, dims.KDim - ).asnumpy() - z_vintcoeff = data_alloc.zero_field(icon_grid, dims.CellDim, dims.C2E2CDim, dims.KDim).asnumpy() - zd_intcoef_dsl = data_alloc.zero_field( - icon_grid, dims.CellDim, dims.C2E2CDim, dims.KDim - ).asnumpy() - z_maxslp_avg = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) - z_maxhgtd_avg = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) - zd_diffcoef_dsl = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim).asnumpy() + maxslp_avg = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) + maxhgtd_avg = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) maxslp = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) maxhgtd = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) max_nbhgt = data_alloc.zero_field(icon_grid, dims.CellDim) c2e2c = icon_grid.connectivities[dims.C2E2CDim] - nbidx = data_alloc.constant_field( - icon_grid, 1, dims.CellDim, dims.C2E2CDim, dims.KDim, dtype=int - ).asnumpy() c_bln_avg = interpolation_savepoint.c_bln_avg() thslp_zdiffu = 0.02 - thhgtd_zdiffu = 125 + thhgtd_zdiffu = 125.0 cell_nudging = icon_grid.start_index(h_grid.domain(dims.CellDim)(h_grid.Zone.NUDGING)) cell_lateral = icon_grid.start_index( @@ -69,8 +57,8 @@ def test_compute_diffusion_metrics( compute_maxslp_maxhgtd.with_backend(backend)( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), - z_maxslp=maxslp, - z_maxhgtd=maxhgtd, + maxslp=maxslp, + maxhgtd=maxhgtd, horizontal_start=cell_lateral, horizontal_end=icon_grid.num_cells, vertical_start=0, @@ -93,8 +81,8 @@ def test_compute_diffusion_metrics( maxslp=maxslp, maxhgtd=maxhgtd, c_bln_avg=c_bln_avg, - z_maxslp_avg=z_maxslp_avg, - z_maxhgtd_avg=z_maxhgtd_avg, + maxslp_avg=maxslp_avg, + maxhgtd_avg=maxhgtd_avg, horizontal_start=cell_lateral, horizontal_end=icon_grid.num_cells, vertical_start=0, @@ -113,40 +101,20 @@ def test_compute_diffusion_metrics( ) mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_metrics( + c2e2c=c2e2c, z_mc=z_mc.asnumpy(), - z_mc_off=z_mc.asnumpy()[c2e2c], max_nbhgt=max_nbhgt.asnumpy(), c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), - nbidx=nbidx, - z_vintcoeff=z_vintcoeff, - z_maxslp_avg=z_maxslp_avg.asnumpy(), - z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), - mask_hdiff=mask_hdiff, - zd_diffcoef_dsl=zd_diffcoef_dsl, - zd_intcoef_dsl=zd_intcoef_dsl, - zd_vertoffset_dsl=zd_vertoffset_dsl, + maxslp_avg=maxslp_avg.asnumpy(), + maxhgtd_avg=maxhgtd_avg.asnumpy(), thslp_zdiffu=thslp_zdiffu, thhgtd_zdiffu=thhgtd_zdiffu, cell_nudging=cell_nudging, - n_cells=icon_grid.num_cells, nlev=nlev, ) - zd_intcoef_dsl = data_alloc.flatten_first_two_dims( - dims.CECDim, - dims.KDim, - field=gtx.as_field((dims.CellDim, dims.C2E2CDim, dims.KDim), zd_intcoef_dsl), - ) - zd_vertoffset_dsl = data_alloc.flatten_first_two_dims( - dims.CECDim, - dims.KDim, - field=gtx.as_field((dims.CellDim, dims.C2E2CDim, dims.KDim), zd_vertoffset_dsl), - ) - assert helpers.dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy()) assert helpers.dallclose( zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11 ) - assert helpers.dallclose( - zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy() - ) - assert helpers.dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef().asnumpy()) + assert helpers.dallclose(zd_vertoffset_dsl, metrics_savepoint.zd_vertoffset().asnumpy()) + assert helpers.dallclose(zd_intcoef_dsl, metrics_savepoint.zd_intcoef().asnumpy()) diff --git a/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py b/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py index 3441e9f08..6a6a7ef42 100644 --- a/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py +++ b/model/common/tests/metric_tests/test_compute_zdiff_gradp_dsl.py @@ -17,10 +17,9 @@ from icon4py.model.common.metrics.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.metrics.metric_fields import ( _compute_flat_idx, - _compute_z_aux2, compute_z_mc, ) -from icon4py.model.common.utils.data_allocation import flatten_first_two_dims, zero_field +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.helpers import ( dallclose, is_roundtrip, @@ -32,10 +31,10 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav if is_roundtrip(backend): pytest.skip("skipping: slow backend") zdiff_gradp_ref = metrics_savepoint.zdiff_gradp() - z_mc = zero_field(icon_grid, dims.CellDim, dims.KDim) + z_mc = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) z_ifc = metrics_savepoint.z_ifc() k_lev = gtx.as_field((dims.KDim,), np.arange(icon_grid.num_levels, dtype=int)) - z_me = zero_field(icon_grid, dims.EdgeDim, dims.KDim) + z_me = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) edge_domain = h_grid.domain(dims.EdgeDim) horizontal_start_edge = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) start_nudging = icon_grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) @@ -54,9 +53,10 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav out=z_me, offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) - flat_idx = zero_field(icon_grid, dims.EdgeDim, dims.KDim) + flat_idx = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) _compute_flat_idx( - z_me=z_me, + z_mc=z_mc, + c_lin_e=interpolation_savepoint.c_lin_e(), z_ifc=z_ifc, k_lev=k_lev, out=flat_idx, @@ -71,29 +71,17 @@ def test_compute_zdiff_gradp_dsl(icon_grid, metrics_savepoint, interpolation_sav ) flat_idx_np = np.amax(flat_idx.asnumpy(), axis=1) z_ifc_sliced = gtx.as_field((dims.CellDim,), z_ifc.asnumpy()[:, icon_grid.num_levels]) - z_aux2 = zero_field(icon_grid, dims.EdgeDim) - _compute_z_aux2( - z_ifc=z_ifc_sliced, - out=z_aux2, - domain={dims.EdgeDim: (start_nudging, icon_grid.num_edges)}, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - zdiff_gradp_full_np = compute_zdiff_gradp_dsl( + zdiff_gradp_full_field = compute_zdiff_gradp_dsl( e2c=icon_grid.connectivities[dims.E2CDim], - z_me=z_me.asnumpy(), z_mc=z_mc.asnumpy(), + c_lin_e=interpolation_savepoint.c_lin_e().asnumpy(), z_ifc=metrics_savepoint.z_ifc().asnumpy(), flat_idx=flat_idx_np, - z_aux2=z_aux2.asnumpy(), + z_ifc_sliced=z_ifc_sliced.asnumpy(), nlev=icon_grid.num_levels, horizontal_start=horizontal_start_edge, horizontal_start_1=start_nudging, - nedges=icon_grid.num_edges, ) - zdiff_gradp_full_field = flatten_first_two_dims( - dims.ECDim, - dims.KDim, - field=gtx.as_field((dims.EdgeDim, dims.E2CDim, dims.KDim), zdiff_gradp_full_np), - ) - assert dallclose(zdiff_gradp_full_field.asnumpy(), zdiff_gradp_ref.asnumpy(), rtol=1.0e-5) + + assert dallclose(zdiff_gradp_full_field, zdiff_gradp_ref.asnumpy(), rtol=1.0e-5) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 097308ae6..2845d0950 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -24,10 +24,8 @@ compute_vwind_impl_wgt, ) from icon4py.model.common.metrics.metric_fields import ( - MetricsConfig, _compute_flat_idx, _compute_pg_edgeidx_vertidx, - _compute_z_aux2, compute_bdy_halo_c, compute_coeff_dwdz, compute_d2dexdz2_fac_mc, @@ -43,6 +41,7 @@ compute_pg_exdist_dsl, compute_rayleigh_w, compute_scalfac_dd3d, + compute_theta_exner_ref_mc, compute_vwind_expl_wgt, compute_wgtfac_e, compute_z_mc, @@ -156,7 +155,6 @@ def test_compute_ddqz_z_full_and_inverse(icon_grid, metrics_savepoint, backend): assert dallclose(inv_ddqz_z_full.asnumpy(), inv_ddqz_full_ref.asnumpy()) -# TODO: convert this to a stenciltest once it is possible to have only dims.KDim in domain @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_scalfac_dd3d(icon_grid, metrics_savepoint, grid_savepoint, backend): @@ -180,7 +178,6 @@ def test_compute_scalfac_dd3d(icon_grid, metrics_savepoint, grid_savepoint, back assert dallclose(scalfac_dd3d_ref.asnumpy(), scalfac_dd3d_full.asnumpy()) -# TODO: convert this to a stenciltest once it is possible to have only dims.KDim in domain @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT]) def test_compute_rayleigh_w(icon_grid, experiment, metrics_savepoint, grid_savepoint, backend): @@ -241,7 +238,6 @@ def test_compute_coeff_dwdz(icon_grid, metrics_savepoint, grid_savepoint, backen @pytest.mark.datatest @pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, backend): - backend = None if is_roundtrip(backend): pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() @@ -295,21 +291,12 @@ def test_compute_d2dexdz2_fac_mc(icon_grid, metrics_savepoint, grid_savepoint, b def test_compute_ddxt_z_full_e( grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint, backend ): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() - tangent_orientation = grid_savepoint.tangent_orientation() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() + ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy() horizontal_start_vertex = icon_grid.start_index( vertex_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2) ) horizontal_end_vertex = icon_grid.end_index(vertex_domain(horizontal.Zone.INTERIOR)) - horizontal_start_edge = icon_grid.start_index( - edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_3) - ) - - horizontal_end_edge = icon_grid.end_index(edge_domain(horizontal.Zone.INTERIOR)) vertical_start = 0 vertical_end = icon_grid.num_levels + 1 cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() @@ -324,26 +311,31 @@ def test_compute_ddxt_z_full_e( vertical_end=vertical_end, offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, ) - ddxt_z_half_e = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) - compute_ddxt_z_half_e.with_backend(backend)( - z_ifv=z_ifv, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - ddxt_z_half_e=ddxt_z_half_e, - horizontal_start=horizontal_start_edge, - horizontal_end=horizontal_end_edge, + ddxn_z_half_e = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) + compute_ddxn_z_half_e( + z_ifc=z_ifc, + inv_dual_edge_length=grid_savepoint.inv_dual_edge_length(), + ddxn_z_half_e=ddxn_z_half_e, + horizontal_start=icon_grid.start_index( + edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + horizontal_end=icon_grid.end_index(edge_domain(horizontal.Zone.INTERIOR)), vertical_start=vertical_start, vertical_end=vertical_end, - offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) - ddxt_z_full = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) + ddxn_z_full = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) compute_ddxn_z_full.with_backend(backend)( - z_ddxnt_z_half_e=ddxt_z_half_e, - ddxn_z_full=ddxt_z_full, + ddxnt_z_half_e=ddxn_z_half_e, + ddxn_z_full=ddxn_z_full, + horizontal_start=0, + horizontal_end=icon_grid.num_edges, + vertical_start=0, + vertical_end=icon_grid.num_levels, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - assert np.allclose(ddxt_z_full.asnumpy(), ddxt_z_full_ref) + assert dallclose(ddxn_z_full.asnumpy(), ddxn_z_full_ref) @pytest.mark.datatest @@ -419,8 +411,12 @@ def test_compute_ddxn_z_full( ) ddxn_z_full = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) compute_ddxn_z_full.with_backend(backend)( - z_ddxnt_z_half_e=ddxn_z_half_e, + ddxnt_z_half_e=ddxn_z_half_e, ddxn_z_full=ddxn_z_full, + horizontal_start=0, + horizontal_end=icon_grid.num_edges, + vertical_start=0, + vertical_end=icon_grid.num_levels, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -438,10 +434,6 @@ def test_compute_ddxt_z_full( tangent_orientation = grid_savepoint.tangent_orientation() inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy() - horizontal_start_vertex = icon_grid.start_index( - vertex_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2) - ) - horizontal_end_vertex = icon_grid.end_index(vertex_domain(horizontal.Zone.INTERIOR)) horizontal_start_edge = icon_grid.start_index( edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_3) ) @@ -449,20 +441,10 @@ def test_compute_ddxt_z_full( vertical_start = 0 vertical_end = icon_grid.num_levels + 1 cells_aw_verts = interpolation_savepoint.c_intp().asnumpy() - z_ifv = data_alloc.zero_field(icon_grid, dims.VertexDim, dims.KDim, extend={dims.KDim: 1}) - compute_cell_2_vertex_interpolation.with_backend(backend)( - z_ifc, - gtx.as_field((dims.VertexDim, dims.V2CDim), cells_aw_verts), - z_ifv, - offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, - horizontal_start=horizontal_start_vertex, - horizontal_end=horizontal_end_vertex, - vertical_start=vertical_start, - vertical_end=vertical_end, - ) ddxt_z_half_e = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1}) compute_ddxt_z_half_e.with_backend(backend)( - z_ifv=z_ifv, + cell_in=z_ifc, + c_int=gtx.as_field((dims.VertexDim, dims.V2CDim), cells_aw_verts), inv_primal_edge_length=inv_primal_edge_length, tangent_orientation=tangent_orientation, ddxt_z_half_e=ddxt_z_half_e, @@ -470,12 +452,19 @@ def test_compute_ddxt_z_full( horizontal_end=horizontal_end_edge, vertical_start=vertical_start, vertical_end=vertical_end, - offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, + offset_provider={ + "E2V": icon_grid.get_offset_provider("E2V"), + "V2C": icon_grid.get_offset_provider("V2C"), + }, ) ddxt_z_full = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) compute_ddxn_z_full.with_backend(backend)( - z_ddxnt_z_half_e=ddxt_z_half_e, + ddxnt_z_half_e=ddxt_z_half_e, ddxn_z_full=ddxt_z_full, + horizontal_start=0, + horizontal_end=icon_grid.num_edges, + vertical_start=0, + vertical_end=icon_grid.num_levels, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) @@ -487,23 +476,16 @@ def test_compute_ddxt_z_full( def test_compute_exner_exfac( grid_savepoint, experiment, interpolation_savepoint, icon_grid, metrics_savepoint, backend ): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") - horizontal_start = icon_grid.start_index(cell_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)) - config = ( - MetricsConfig(exner_expol=0.333) - if experiment == dt_utils.REGIONAL_EXPERIMENT - else MetricsConfig() - ) + exner_expol = 0.333 if experiment == dt_utils.REGIONAL_EXPERIMENT else 0.3333333333333 - exner_exfac = data_alloc.constant_field(icon_grid, config.exner_expol, dims.CellDim, dims.KDim) + exner_exfac = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) exner_exfac_ref = metrics_savepoint.exner_exfac() compute_exner_exfac.with_backend(backend)( ddxn_z_full=metrics_savepoint.ddxn_z_full(), dual_edge_length=grid_savepoint.dual_edge_length(), exner_exfac=exner_exfac, - exner_expol=config.exner_expol, + exner_expol=exner_expol, horizontal_start=horizontal_start, horizontal_end=icon_grid.num_cells, vertical_start=gtx.int32(0), @@ -515,12 +497,10 @@ def test_compute_exner_exfac( @pytest.mark.datatest -@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) +@pytest.mark.parametrize("experiment", [dt_utils.GLOBAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT]) def test_compute_vwind_impl_wgt( icon_grid, experiment, grid_savepoint, metrics_savepoint, interpolation_savepoint, backend ): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") z_ifc = metrics_savepoint.z_ifc() inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() z_ddxn_z_half_e = data_alloc.zero_field( @@ -531,7 +511,6 @@ def test_compute_vwind_impl_wgt( z_ddxt_z_half_e = data_alloc.zero_field( icon_grid, dims.EdgeDim, dims.KDim, extend={dims.KDim: 1} ) - z_ifv = data_alloc.zero_field(icon_grid, dims.VertexDim, dims.KDim, extend={dims.KDim: 1}) horizontal_start = icon_grid.start_index(edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)) horizontal_end = icon_grid.end_index(edge_domain(horizontal.Zone.INTERIOR)) @@ -555,23 +534,9 @@ def test_compute_vwind_impl_wgt( ) horizontal_end_edge = icon_grid.end_index(edge_domain(horizontal.Zone.INTERIOR)) - horizontal_start_vertex = icon_grid.start_index( - vertex_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2) - ) - horizontal_end_vertex = icon_grid.end_index(vertex_domain(horizontal.Zone.INTERIOR)) - compute_cell_2_vertex_interpolation( - z_ifc, - interpolation_savepoint.c_intp(), - z_ifv, - offset_provider={"V2C": icon_grid.get_offset_provider("V2C")}, - horizontal_start=horizontal_start_vertex, - horizontal_end=horizontal_end_vertex, - vertical_start=vertical_start, - vertical_end=vertical_end, - ) - compute_ddxt_z_half_e( - z_ifv=z_ifv, + cell_in=z_ifc, + c_int=interpolation_savepoint.c_intp(), inv_primal_edge_length=inv_primal_edge_length, tangent_orientation=tangent_orientation, ddxt_z_half_e=z_ddxt_z_half_e, @@ -579,7 +544,10 @@ def test_compute_vwind_impl_wgt( horizontal_end=horizontal_end_edge, vertical_start=vertical_start, vertical_end=vertical_end, - offset_provider={"E2V": icon_grid.get_offset_provider("E2V")}, + offset_provider={ + "E2V": icon_grid.get_offset_provider("E2V"), + "V2C": icon_grid.get_offset_provider("V2C"), + }, ) horizontal_start_cell = icon_grid.start_index( @@ -587,29 +555,19 @@ def test_compute_vwind_impl_wgt( ) vwind_impl_wgt_ref = metrics_savepoint.vwind_impl_wgt() dual_edge_length = grid_savepoint.dual_edge_length() - vwind_offctr = 0.2 - vwind_impl_wgt_full = data_alloc.constant_field(icon_grid, 0.5 + vwind_offctr, dims.CellDim) - init_val = 0.65 if experiment == dt_utils.GLOBAL_EXPERIMENT else 0.7 - vwind_impl_wgt_k = data_alloc.constant_field(icon_grid, init_val, dims.CellDim, dims.KDim) + vwind_offctr = 0.2 if experiment == dt_utils.REGIONAL_EXPERIMENT else 0.15 vwind_impl_wgt = compute_vwind_impl_wgt( - backend=backend, - icon_grid=icon_grid, - vct_a=grid_savepoint.vct_a(), - z_ifc=metrics_savepoint.z_ifc(), - z_ddxn_z_half_e=gtx.as_field( - (dims.EdgeDim,), z_ddxn_z_half_e.asnumpy()[:, icon_grid.num_levels] - ), - z_ddxt_z_half_e=gtx.as_field( - (dims.EdgeDim,), z_ddxt_z_half_e.asnumpy()[:, icon_grid.num_levels] - ), - dual_edge_length=dual_edge_length, - vwind_impl_wgt_full=vwind_impl_wgt_full, - vwind_impl_wgt_k=vwind_impl_wgt_k, - global_exp=dt_utils.GLOBAL_EXPERIMENT, - experiment=experiment, + c2e=icon_grid.connectivities[dims.C2EDim], + vct_a=grid_savepoint.vct_a().asnumpy(), + z_ifc=metrics_savepoint.z_ifc().asnumpy(), + z_ddxn_z_half_e=z_ddxn_z_half_e.asnumpy(), + z_ddxt_z_half_e=z_ddxt_z_half_e.asnumpy(), + dual_edge_length=dual_edge_length.asnumpy(), vwind_offctr=vwind_offctr, + nlev=icon_grid.num_levels, horizontal_start_cell=horizontal_start_cell, + n_cells=icon_grid.num_cells, ) assert dallclose(vwind_impl_wgt_ref.asnumpy(), vwind_impl_wgt) @@ -639,21 +597,18 @@ def test_compute_wgtfac_e(metrics_savepoint, interpolation_savepoint, icon_grid, def test_compute_pg_exdist_dsl( metrics_savepoint, interpolation_savepoint, icon_grid, grid_savepoint, backend ): - if is_roundtrip(backend): - pytest.skip("skipping: slow backend") pg_exdist_ref = metrics_savepoint.pg_exdist() nlev = icon_grid.num_levels k_lev = gtx.as_field((dims.KDim,), np.arange(nlev, dtype=gtx.int32)) pg_edgeidx = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) pg_vertidx = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) pg_exdist_dsl = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) - z_me = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) - z_aux2 = data_alloc.zero_field(icon_grid, dims.EdgeDim) z_mc = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) flat_idx = data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim) z_ifc = metrics_savepoint.z_ifc() z_ifc_sliced = gtx.as_field((dims.CellDim,), z_ifc.asnumpy()[:, nlev]) start_edge_nudging = icon_grid.end_index(edge_domain(horizontal.Zone.NUDGING)) + start_edge_nudging_2 = icon_grid.start_index(edge_domain(horizontal.Zone.NUDGING_LEVEL_2)) horizontal_start_edge = icon_grid.start_index( edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_3) ) @@ -663,25 +618,10 @@ def test_compute_pg_exdist_dsl( average_cell_kdim_level_up.with_backend(backend)( z_ifc, out=z_mc, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")} ) - cell_2_edge_interpolation.with_backend(backend)( - in_field=z_mc, - coeff=interpolation_savepoint.c_lin_e(), - out_field=z_me, - horizontal_start=0, - horizontal_end=icon_grid.num_edges, - vertical_start=0, - vertical_end=nlev, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) - _compute_z_aux2( - z_ifc=z_ifc_sliced, - out=z_aux2, - domain={dims.EdgeDim: (start_edge_nudging, icon_grid.num_edges)}, - offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, - ) _compute_flat_idx( - z_me=z_me, + z_mc=z_mc, + c_lin_e=interpolation_savepoint.c_lin_e(), z_ifc=z_ifc, k_lev=k_lev, out=flat_idx, @@ -695,25 +635,30 @@ def test_compute_pg_exdist_dsl( }, ) flat_idx_np = np.amax(flat_idx.asnumpy(), axis=1) + flat_idx_max = gtx.as_field((dims.EdgeDim,), flat_idx_np, dtype=gtx.int32) compute_pg_exdist_dsl.with_backend(backend)( - z_aux2=z_aux2, - z_me=z_me, + z_ifc_sliced=z_ifc_sliced, + z_mc=z_mc, + c_lin_e=interpolation_savepoint.c_lin_e(), e_owner_mask=grid_savepoint.e_owner_mask(), - flat_idx_max=gtx.as_field((dims.EdgeDim,), flat_idx_np, dtype=gtx.int32), + flat_idx_max=flat_idx_max, k_lev=k_lev, + e_lev=e_lev, pg_exdist_dsl=pg_exdist_dsl, + h_start_zaux2=start_edge_nudging, + h_end_zaux2=icon_grid.num_edges, horizontal_start=start_edge_nudging, horizontal_end=icon_grid.num_edges, vertical_start=0, vertical_end=nlev, - offset_provider={}, + offset_provider={"E2C": icon_grid.get_offset_provider("E2C")}, ) _compute_pg_edgeidx_vertidx( c_lin_e=interpolation_savepoint.c_lin_e(), z_ifc=z_ifc, - z_aux2=z_aux2, + z_ifc_sliced=z_ifc_sliced, e_owner_mask=grid_savepoint.e_owner_mask(), flat_idx_max=gtx.as_field((dims.EdgeDim,), flat_idx_np, dtype=gtx.int32), e_lev=e_lev, @@ -735,7 +680,7 @@ def test_compute_pg_exdist_dsl( pg_edgeidx=pg_edgeidx, pg_vertidx=pg_vertidx, pg_edgeidx_dsl=pg_edgeidx_dsl, - horizontal_start=int(0), + horizontal_start=start_edge_nudging_2, horizontal_end=icon_grid.num_edges, vertical_start=int(0), vertical_end=icon_grid.num_levels, @@ -789,9 +734,9 @@ def test_compute_bdy_halo_c(metrics_savepoint, icon_grid, grid_savepoint, backen def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backend): hmask_dd3d_full = data_alloc.zero_field(icon_grid, dims.EdgeDim) e_refin_ctrl = grid_savepoint.refin_ctrl(dims.EdgeDim) - horizontal_start = icon_grid.start_index(cell_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)) + horizontal_start = icon_grid.start_index(edge_domain(horizontal.Zone.LATERAL_BOUNDARY_LEVEL_2)) hmask_dd3d_ref = metrics_savepoint.hmask_dd3d() - compute_hmask_dd3d( + compute_hmask_dd3d.with_backend(backend)( e_refin_ctrl=e_refin_ctrl, hmask_dd3d=hmask_dd3d_full, grf_nudge_start_e=gtx.int32(horizontal._GRF_NUDGEZONE_START_EDGES), @@ -801,4 +746,47 @@ def test_compute_hmask_dd3d(metrics_savepoint, icon_grid, grid_savepoint, backen offset_provider={}, ) - dallclose(hmask_dd3d_full.asnumpy(), hmask_dd3d_ref.asnumpy()) + assert dallclose(hmask_dd3d_full.asnumpy(), hmask_dd3d_ref.asnumpy()) + + +@pytest.mark.datatest +@pytest.mark.parametrize("experiment", [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) +def test_compute_theta_exner_ref_mc(metrics_savepoint, icon_grid, backend): + exner_ref_mc_full = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) + theta_ref_mc_full = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) + t0sl_bg = constants.SEA_LEVEL_TEMPERATURE + del_t_bg = constants.DELTA_TEMPERATURE + h_scal_bg = constants._H_SCAL_BG + grav = constants.GRAV + rd = constants.RD + p0sl_bg = constants.SEAL_LEVEL_PRESSURE + rd_o_cpd = constants.RD_O_CPD + p0ref = constants.REFERENCE_PRESSURE + exner_ref_mc_ref = metrics_savepoint.exner_ref_mc() + theta_ref_mc_ref = metrics_savepoint.theta_ref_mc() + z_ifc = metrics_savepoint.z_ifc() + z_mc = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim) + average_cell_kdim_level_up.with_backend(backend)( + z_ifc, out=z_mc, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")} + ) + compute_theta_exner_ref_mc.with_backend(backend)( + z_mc=z_mc, + exner_ref_mc=exner_ref_mc_full, + theta_ref_mc=theta_ref_mc_full, + t0sl_bg=t0sl_bg, + del_t_bg=del_t_bg, + h_scal_bg=h_scal_bg, + grav=grav, + rd=rd, + p0sl_bg=p0sl_bg, + rd_o_cpd=rd_o_cpd, + p0ref=p0ref, + horizontal_start=int(0), + horizontal_end=icon_grid.num_cells, + vertical_start=int(0), + vertical_end=icon_grid.num_levels, + offset_provider={}, + ) + + assert dallclose(exner_ref_mc_ref.asnumpy(), exner_ref_mc_full.asnumpy()) + assert dallclose(theta_ref_mc_ref.asnumpy(), theta_ref_mc_full.asnumpy()) diff --git a/model/common/tests/metric_tests/test_metrics_factory.py b/model/common/tests/metric_tests/test_metrics_factory.py new file mode 100644 index 000000000..912ecbfc8 --- /dev/null +++ b/model/common/tests/metric_tests/test_metrics_factory.py @@ -0,0 +1,545 @@ +# 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 pytest + +from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import vertical as v_grid +from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory +from icon4py.model.common.metrics import ( + metrics_attributes as attrs, + metrics_factory, +) +from icon4py.model.testing import ( + datatest_utils as dt_utils, + grid_utils as gridtest_utils, + helpers as test_helpers, +) + + +metrics_factories = {} + + +def metrics_config(experiment: str) -> tuple: + rayleigh_coeff = 5.0 + lowest_layer_thickness = 50.0 + exner_expol = 0.333 + vwind_offctr = 0.2 + rayleigh_type = 2 + model_top_height = 23500.0 + stretch_factor = 1.0 + damping_height = 45000.0 + match experiment: + case dt_utils.REGIONAL_EXPERIMENT: + lowest_layer_thickness = 20.0 + model_top_height = 23000.0 + stretch_factor = 0.65 + damping_height = 12500.0 + case dt_utils.GLOBAL_EXPERIMENT: + model_top_height = 75000.0 + stretch_factor = 0.9 + damping_height = 50000.0 + rayleigh_coeff = 0.1 + exner_expol = 0.3333333333333 + vwind_offctr = 0.15 + + return ( + lowest_layer_thickness, + model_top_height, + stretch_factor, + damping_height, + rayleigh_coeff, + exner_expol, + vwind_offctr, + rayleigh_type, + ) + + +def get_metrics_factory( + backend, experiment, grid_file, grid_savepoint, metrics_savepoint +) -> metrics_factory.MetricsFieldsFactory: + name = experiment.join(backend.name) + factory = metrics_factories.get(name) + + if not factory: + geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) + ( + lowest_layer_thickness, + model_top_height, + stretch_factor, + damping_height, + rayleigh_coeff, + exner_expol, + vwind_offctr, + rayleigh_type, + ) = metrics_config(experiment) + + vertical_config = v_grid.VerticalGridConfig( + geometry.grid.num_levels, + lowest_layer_thickness=lowest_layer_thickness, + model_top_height=model_top_height, + stretch_factor=stretch_factor, + rayleigh_damping_height=damping_height, + ) + vertical_grid = v_grid.VerticalGrid( + vertical_config, grid_savepoint.vct_a(), grid_savepoint.vct_b() + ) + interpolation_fact = interpolation_factory.InterpolationFieldsFactory( + grid=geometry.grid, + decomposition_info=geometry._decomposition_info, + geometry_source=geometry, + backend=backend, + metadata=interpolation_attributes.attrs, + ) + factory = metrics_factory.MetricsFieldsFactory( + grid=geometry.grid, + vertical_grid=vertical_grid, + decomposition_info=geometry._decomposition_info, + geometry_source=geometry, + interpolation_source=interpolation_fact, + backend=backend, + metadata=attrs.attrs, + interface_model_height=metrics_savepoint.z_ifc(), + e_refin_ctrl=grid_savepoint.refin_ctrl(dims.EdgeDim), + c_refin_ctrl=grid_savepoint.refin_ctrl(dims.CellDim), + damping_height=damping_height, + rayleigh_type=rayleigh_type, + rayleigh_coeff=rayleigh_coeff, + exner_expol=exner_expol, + vwind_offctr=vwind_offctr, + ) + metrics_factories[name] = factory + return factory + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_inv_ddqz_z(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.inv_ddqz_z_full() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.INV_DDQZ_Z_FULL) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_ddqz_z_half(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.ddqz_z_half() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.DDQZ_Z_HALF) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_scalfac_dd3d(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.scalfac_dd3d() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.SCALFAC_DD3D) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_rayleigh_w(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.rayleigh_w() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.RAYLEIGH_W) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_coeffs_dwdz(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref_1 = metrics_savepoint.coeff1_dwdz() + field_ref_2 = metrics_savepoint.coeff2_dwdz() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field_1 = factory.get(attrs.COEFF1_DWDZ) + field_2 = factory.get(attrs.COEFF2_DWDZ) + assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy()) + assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_ref_mc(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref_1 = metrics_savepoint.theta_ref_mc() + field_ref_2 = metrics_savepoint.exner_ref_mc() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field_1 = factory.get(attrs.THETA_REF_MC) + field_2 = factory.get(attrs.EXNER_REF_MC) + assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy()) + assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_facs_mc(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref_1 = metrics_savepoint.d2dexdz2_fac1_mc() + field_ref_2 = metrics_savepoint.d2dexdz2_fac2_mc() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field_1 = factory.get(attrs.D2DEXDZ2_FAC1_MC) + field_2 = factory.get(attrs.D2DEXDZ2_FAC2_MC) + assert test_helpers.dallclose(field_1.asnumpy(), field_ref_1.asnumpy()) + assert test_helpers.dallclose(field_2.asnumpy(), field_ref_2.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_ddxn_z_full(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.ddxn_z_full() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.DDXN_Z_FULL) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-8) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + ( + dt_utils.REGIONAL_EXPERIMENT, + dt_utils.REGIONAL_EXPERIMENT, + ), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_vwind_impl_wgt(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.vwind_impl_wgt() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.VWIND_IMPL_WGT) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-9) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + ( + dt_utils.REGIONAL_EXPERIMENT, + dt_utils.REGIONAL_EXPERIMENT, + ), # TODO: check vwind_offctr value for regional + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_vwind_expl_wgt(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.vwind_expl_wgt() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.VWIND_EXPL_WGT) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-8) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_exner_exfac(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.exner_exfac() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.EXNER_EXFAC) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1.0e-5) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_pg_edgeidx_dsl(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.pg_edgeidx_dsl() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.PG_EDGEIDX_DSL) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_pg_exdist_dsl(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.pg_exdist() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.PG_EDGEDIST_DSL) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), atol=1.0e-5) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_mask_bdy_prog_halo_c( + grid_savepoint, metrics_savepoint, grid_file, experiment, backend +): + field_ref_1 = metrics_savepoint.mask_prog_halo_c() + field_ref_2 = metrics_savepoint.bdy_halo_c() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field_1 = factory.get(attrs.MASK_PROG_HALO_C) + field_2 = factory.get(attrs.BDY_HALO_C) + assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy()) + assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_hmask_dd3d(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.hmask_dd3d() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.HMASK_DD3D) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_zdiff_gradp(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.zdiff_gradp() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.ZDIFF_GRADP) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), atol=1.0e-5) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_coeff_gradekin(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref = metrics_savepoint.coeff_gradekin() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.COEFF_GRADEKIN) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-8) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), + ], +) +@pytest.mark.datatest +def test_factory_wgtfacq_e(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field = factory.get(attrs.WGTFACQ_E) + field_ref = metrics_savepoint.wgtfacq_e_dsl(field.shape[1]) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-9) + + +@pytest.mark.parametrize( + "grid_file, experiment", + [ + (dt_utils.REGIONAL_EXPERIMENT, dt_utils.REGIONAL_EXPERIMENT), + # (dt_utils.R02B04_GLOBAL, dt_utils.GLOBAL_EXPERIMENT), # zd_intcoef not present in dataset # noqa: ERA001 + ], +) +@pytest.mark.datatest +def test_factory_diffusion(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): + field_ref_1 = metrics_savepoint.mask_hdiff() + field_ref_2 = metrics_savepoint.zd_diffcoef() + field_ref_3 = metrics_savepoint.zd_intcoef() + field_ref_4 = metrics_savepoint.zd_vertoffset() + factory = get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + ) + field_1 = factory.get(attrs.MASK_HDIFF) + field_2 = factory.get(attrs.ZD_DIFFCOEF_DSL) + field_3 = factory.get(attrs.ZD_INTCOEF_DSL) + field_4 = factory.get(attrs.ZD_VERTOFFSET_DSL) + assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy()) + assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy(), rtol=1.0e-4) + assert test_helpers.dallclose(field_ref_3.asnumpy(), field_3.asnumpy()) + assert test_helpers.dallclose(field_ref_4.asnumpy(), field_4.asnumpy())