diff --git a/model/common/src/icon4py/model/common/grid/topography.py b/model/common/src/icon4py/model/common/grid/topography.py index 800c35b592..e6ba6d317d 100644 --- a/model/common/src/icon4py/model/common/grid/topography.py +++ b/model/common/src/icon4py/model/common/grid/topography.py @@ -6,84 +6,57 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import backend as gtx_backend +from types import ModuleType + +import numpy as np -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.grid import base -from icon4py.model.common.math.stencils.compute_nabla2_on_cell import compute_nabla2_on_cell from icon4py.model.common.utils import data_allocation as data_alloc -@gtx.field_operator -def _update_smoothed_topography( - smoothed_topography: fa.CellField[ta.wpfloat], - nabla2_topo: fa.CellField[ta.wpfloat], - cell_areas: fa.CellField[ta.wpfloat], -) -> fa.CellField[ta.wpfloat]: +def compute_nabla2_on_cell( + psi_c: data_alloc.NDArray, + geofac_n2s: data_alloc.NDArray, + c2e2co: data_alloc.NDArray, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: """ - Updates the smoothed topography field inside the loop. + Computes the Laplacian (nabla squared) of a scalar field defined on cell + centres. """ - return smoothed_topography + 0.125 * nabla2_topo * cell_areas + nabla2_psi_c = array_ns.sum(psi_c[c2e2co] * geofac_n2s, axis=1) + return nabla2_psi_c -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def update_smoothed_topography( - nabla2_topo: fa.CellField[ta.wpfloat], - cell_areas: fa.CellField[ta.wpfloat], - smoothed_topography: fa.CellField[ta.wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, -): - _update_smoothed_topography( - smoothed_topography=smoothed_topography, - nabla2_topo=nabla2_topo, - cell_areas=cell_areas, - out=smoothed_topography, - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - }, - ) + smoothed_topography: np.ndarray, + nabla2_topo: np.ndarray, + cell_areas: np.ndarray, +) -> data_alloc.NDArray: + """ + Updates the smoothed topography field inside the loop. (Numpy version) + """ + return smoothed_topography + 0.125 * nabla2_topo * cell_areas def smooth_topography( - topography: fa.CellField[ta.wpfloat], - cell_areas: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - grid: base.BaseGrid, - backend: gtx_backend.Backend, + topography: data_alloc.NDArray, + cell_areas: data_alloc.NDArray, + geofac_n2s: data_alloc.NDArray, + c2e2co: data_alloc.NDArray, num_iterations: int = 25, -) -> fa.CellField[ta.wpfloat]: + array_ns: ModuleType = np, +) -> data_alloc.NDArray: """ Computes the smoothed (laplacian-filtered) topography needed by the SLEVE coordinate. """ - smoothed_topography = gtx.as_field( - (dims.CellDim,), topography.ndarray, allocator=backend - ) # TODO (@halungge) this should copy? - - nabla2_topo = data_alloc.zero_field(grid, dims.CellDim, backend=backend) + smoothed_topography = topography.copy() for _ in range(num_iterations): - compute_nabla2_on_cell.with_backend(backend)( - psi_c=smoothed_topography, - geofac_n2s=geofac_n2s, - nabla2_psi_c=nabla2_topo, - horizontal_start=0, - horizontal_end=grid.num_cells, - offset_provider={ - "C2E2CO": grid.get_connectivity("C2E2CO"), - }, - ) - - update_smoothed_topography.with_backend(backend)( - nabla2_topo=nabla2_topo, - cell_areas=cell_areas, - smoothed_topography=smoothed_topography, - horizontal_start=0, - horizontal_end=grid.num_cells, - offset_provider={}, + nabla2_topo = compute_nabla2_on_cell(smoothed_topography, geofac_n2s, c2e2co, array_ns) + smoothed_topography = update_smoothed_topography( + smoothed_topography, nabla2_topo, cell_areas ) return smoothed_topography diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index f5866051ec..42d72e9e4a 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -19,8 +19,9 @@ from gt4py.next import backend as gtx_backend import icon4py.model.common.states.metadata as data +import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa -from icon4py.model.common.grid import base, icon as icon_grid, topography as topo +from icon4py.model.common.grid import topography as topo from icon4py.model.common.utils import data_allocation as data_alloc @@ -99,7 +100,7 @@ class VerticalGridConfig: #: Defined in ICON namelist nonhydrostatic_nml. Height [m] above which moist physics and advection of cloud and precipitation variables are turned off. htop_moist_proc: Final[float] = 22500.0 #: file name containing vct_a and vct_b table - file_path: pathlib.Path = None + file_path: Optional[pathlib.Path] = None # Parameters for setting up the decay function of the topographic signal for # SLEVE. Default values from mo_sleve_nml. @@ -553,9 +554,12 @@ def _compute_SLEVE_coordinate_from_vcta_and_topography( topography: data_alloc.NDArray, cell_areas: data_alloc.NDArray, geofac_n2s: data_alloc.NDArray, - grid: base.BaseGrid, - vertical_geometry: VerticalGrid, - backend: gtx_backend.Backend, + c2e2co: data_alloc.NDArray, + nflatlev: int, + model_top_height: ta.wpfloat, + SLEVE_decay_scale_1: ta.wpfloat, + SLEVE_decay_exponent: ta.wpfloat, + SLEVE_decay_scale_2: ta.wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -567,56 +571,49 @@ def _compute_SLEVE_coordinate_from_vcta_and_topography( blends smothly into the surface layer at num_lev + 1 which is the topography. """ + num_cells = cell_areas.shape[0] + num_levels = vct_a.shape[0] - 1 def _decay_func( vct_a: data_alloc.NDArray, - model_top_height: float, - decay_scale: float, - decay_exponent: float, + model_top_height: ta.wpfloat, + decay_scale: ta.wpfloat, + decay_exponent: ta.wpfloat, ) -> data_alloc.NDArray: return array_ns.sinh( (model_top_height / decay_scale) ** decay_exponent - (vct_a / decay_scale) ** decay_exponent ) / array_ns.sinh((model_top_height / decay_scale) ** decay_exponent) - assert topography.shape == (grid.num_cells,) - topo_field = gtx.as_field((dims.CellDim,), topography, allocator=backend) - assert cell_areas.shape == (grid.num_cells,) - cell_area_field = gtx.as_field((dims.CellDim,), cell_areas, allocator=backend) - assert geofac_n2s.shape == (grid.num_cells, 4) - geofac_n2s_field = gtx.as_field((dims.CellDim, dims.C2E2CODim), geofac_n2s, allocator=backend) - smoothed_topography = topo.smooth_topography( - topography=topo_field, - cell_areas=cell_area_field, - geofac_n2s=geofac_n2s_field, - grid=grid, - backend=backend, - ).ndarray + topography=topography, + cell_areas=cell_areas, + geofac_n2s=geofac_n2s, + c2e2co=c2e2co, + ) - vertical_coordinate = array_ns.zeros((grid.num_cells, grid.num_levels + 1), dtype=float) - vertical_coordinate[:, grid.num_levels] = topography + vertical_coordinate = array_ns.zeros((num_cells, num_levels + 1), dtype=float) + vertical_coordinate[:, num_levels] = topography # Small-scale topography (i.e. full topo - smooth topo) small_scale_topography = topography - smoothed_topography - k = range(vertical_geometry.nflatlev + 1) + k = range(nflatlev + 1) vertical_coordinate[:, k] = vct_a[k] - k = range(vertical_geometry.nflatlev + 1, grid.num_levels) - vertical_config = vertical_geometry.config + k = range(nflatlev + 1, num_levels) # Scaling factors for large-scale and small-scale topography z_fac1 = _decay_func( vct_a[k], - vertical_config.model_top_height, - vertical_config.SLEVE_decay_scale_1, - vertical_config.SLEVE_decay_exponent, + model_top_height, + SLEVE_decay_scale_1, + SLEVE_decay_exponent, ) z_fac2 = _decay_func( vct_a[k], - vertical_config.model_top_height, - vertical_config.SLEVE_decay_scale_2, - vertical_config.SLEVE_decay_exponent, + model_top_height, + SLEVE_decay_scale_2, + SLEVE_decay_exponent, ) vertical_coordinate[:, k] = ( vct_a[k][array_ns.newaxis, :] @@ -630,48 +627,43 @@ def _decay_func( def _check_and_correct_layer_thickness( vertical_coordinate: data_alloc.NDArray, vct_a: data_alloc.NDArray, - vertical_config: VerticalGridConfig, - grid: base.BaseGrid, + SLEVE_minimum_layer_thickness_1: ta.wpfloat, + SLEVE_minimum_relative_layer_thickness_1: ta.wpfloat, + SLEVE_minimum_layer_thickness_2: ta.wpfloat, + SLEVE_minimum_relative_layer_thickness_2: ta.wpfloat, + lowest_layer_thickness: ta.wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: - num_levels = vertical_config.num_levels - num_cells = grid.num_cells - ktop_thicklimit = array_ns.asarray(num_cells * [num_levels], dtype=float) + num_cells = vertical_coordinate.shape[0] + num_levels = vertical_coordinate.shape[1] - 1 + ktop_thicklimit = array_ns.asarray(num_cells * [num_levels], dtype=int) # Ensure that layer thicknesses are not too small; this would potentially # cause instabilities in vertical advection for k in reversed(range(num_levels)): delta_vct_a = vct_a[k] - vct_a[k + 1] - if delta_vct_a < vertical_config.SLEVE_minimum_layer_thickness_1: + if delta_vct_a < SLEVE_minimum_layer_thickness_1: # limit layer thickness to SLEVE_minimum_relative_layer_thickness_1 times its nominal value - minimum_layer_thickness = ( - vertical_config.SLEVE_minimum_relative_layer_thickness_1 * delta_vct_a - ) - elif delta_vct_a < vertical_config.SLEVE_minimum_layer_thickness_2: + minimum_layer_thickness = SLEVE_minimum_relative_layer_thickness_1 * delta_vct_a + elif delta_vct_a < SLEVE_minimum_layer_thickness_2: # limitation factor changes from SLEVE_minimum_relative_layer_thickness_1 to SLEVE_minimum_relative_layer_thickness_2 layer_thickness_adjustment_factor = ( - (vertical_config.SLEVE_minimum_layer_thickness_2 - delta_vct_a) - / ( - vertical_config.SLEVE_minimum_layer_thickness_2 - - vertical_config.SLEVE_minimum_layer_thickness_1 - ) + (SLEVE_minimum_layer_thickness_2 - delta_vct_a) + / (SLEVE_minimum_layer_thickness_2 - SLEVE_minimum_layer_thickness_1) ) ** 2 minimum_layer_thickness = ( - vertical_config.SLEVE_minimum_relative_layer_thickness_1 - * layer_thickness_adjustment_factor - + vertical_config.SLEVE_minimum_relative_layer_thickness_2 + SLEVE_minimum_relative_layer_thickness_1 * layer_thickness_adjustment_factor + + SLEVE_minimum_relative_layer_thickness_2 * (1.0 - layer_thickness_adjustment_factor) ) * delta_vct_a else: # limitation factor decreases again minimum_layer_thickness = ( - vertical_config.SLEVE_minimum_relative_layer_thickness_2 - * vertical_config.SLEVE_minimum_layer_thickness_2 - * (delta_vct_a / vertical_config.SLEVE_minimum_layer_thickness_2) ** (1.0 / 3.0) + SLEVE_minimum_relative_layer_thickness_2 + * SLEVE_minimum_layer_thickness_2 + * (delta_vct_a / SLEVE_minimum_layer_thickness_2) ** (1.0 / 3.0) ) - minimum_layer_thickness = max( - minimum_layer_thickness, min(50, vertical_config.lowest_layer_thickness) - ) + minimum_layer_thickness = max(minimum_layer_thickness, min(50, lowest_layer_thickness)) # Ensure that the layer thickness is not too small, if so fix it and # save the layer index @@ -713,14 +705,12 @@ def _check_and_correct_layer_thickness( # Check if ktop_thicklimit is sufficiently far away from the model top if not array_ns.all(ktop_thicklimit > 2): - if vertical_config.num_levels > 6: + if num_levels > 6: raise exceptions.InvalidConfigError( - f"Model top is too low and num_levels, {vertical_config.num_levels}, > 6." + f"Model top is too low and num_levels, {num_levels}, > 6." ) else: - log.warning( - f"Model top is too low. But num_levels, {vertical_config.num_levels}, <= 6. " - ) + log.warning(f"Model top is too low. But num_levels, {num_levels}, <= 6. ") return vertical_coordinate @@ -728,14 +718,11 @@ def _check_and_correct_layer_thickness( def _check_flatness_of_flat_level( vertical_coordinate: data_alloc.NDArray, vct_a: data_alloc.NDArray, - vertical_geometry: VerticalGrid, + nflatlev: int, array_ns: ModuleType = np, ) -> None: # Check if level nflatlev is still flat - if not array_ns.all( - vertical_coordinate[:, vertical_geometry.nflatlev - 1] - == vct_a[vertical_geometry.nflatlev - 1] - ): + if not array_ns.all(vertical_coordinate[:, nflatlev - 1] == vct_a[nflatlev - 1]): raise exceptions.InvalidComputationError("Level nflatlev is not flat") @@ -744,9 +731,17 @@ def compute_vertical_coordinate( topography: data_alloc.NDArray, cell_areas: data_alloc.NDArray, geofac_n2s: data_alloc.NDArray, - grid: icon_grid.IconGrid, - vertical_geometry: VerticalGrid, - backend: gtx_backend.Backend, + c2e2co: data_alloc.NDArray, + nflatlev: int, + model_top_height: ta.wpfloat, + SLEVE_decay_scale_1: ta.wpfloat, + SLEVE_decay_exponent: ta.wpfloat, + SLEVE_decay_scale_2: ta.wpfloat, + SLEVE_minimum_layer_thickness_1: ta.wpfloat, + SLEVE_minimum_relative_layer_thickness_1: ta.wpfloat, + SLEVE_minimum_layer_thickness_2: ta.wpfloat, + SLEVE_minimum_relative_layer_thickness_2: ta.wpfloat, + lowest_layer_thickness: ta.wpfloat, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -770,29 +765,35 @@ def compute_vertical_coordinate( exceptions.InvalidConfigError: If model top is too low and num_levels > 6. """ - vertical_config = vertical_geometry.config - vertical_coordinate = _compute_SLEVE_coordinate_from_vcta_and_topography( + vertical_coordinates_on_half_levels = _compute_SLEVE_coordinate_from_vcta_and_topography( vct_a=vct_a, topography=topography, cell_areas=cell_areas, geofac_n2s=geofac_n2s, - grid=grid, - vertical_geometry=vertical_geometry, - backend=backend, + c2e2co=c2e2co, + nflatlev=nflatlev, + model_top_height=model_top_height, + SLEVE_decay_scale_1=SLEVE_decay_scale_1, + SLEVE_decay_exponent=SLEVE_decay_exponent, + SLEVE_decay_scale_2=SLEVE_decay_scale_2, array_ns=array_ns, ) + vertical_coordinate = _check_and_correct_layer_thickness( - vertical_coordinate=vertical_coordinate, + vertical_coordinate=vertical_coordinates_on_half_levels, vct_a=vct_a, - vertical_config=vertical_config, - grid=grid, + SLEVE_minimum_layer_thickness_1=SLEVE_minimum_layer_thickness_1, + SLEVE_minimum_relative_layer_thickness_1=SLEVE_minimum_relative_layer_thickness_1, + SLEVE_minimum_layer_thickness_2=SLEVE_minimum_layer_thickness_2, + SLEVE_minimum_relative_layer_thickness_2=SLEVE_minimum_relative_layer_thickness_2, + lowest_layer_thickness=lowest_layer_thickness, array_ns=array_ns, ) _check_flatness_of_flat_level( - vertical_coordinate=vertical_coordinate, + vertical_coordinate=vertical_coordinates_on_half_levels, vct_a=vct_a, - vertical_geometry=vertical_geometry, + nflatlev=nflatlev, array_ns=array_ns, ) 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 f9f2095d1b..cf566cfdc0 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py @@ -25,6 +25,7 @@ POS_ON_TPLANE_E_X: Final[str] = "pos_on_tplane_e_x" POS_ON_TPLANE_E_Y: Final[str] = "pos_on_tplane_e_y" CELL_AW_VERTS: Final[str] = "cell_to_vertex_interpolation_factor_by_area_weighting" +NUDGECOEFFS_E: Final[str] = "nudging_coefficients_for_edges" RBF_VEC_COEFF_C1: Final[str] = "rbf_interpolation_coefficient_cell_1" RBF_VEC_COEFF_C2: Final[str] = "rbf_interpolation_coefficient_cell_2" RBF_VEC_COEFF_E: Final[str] = "rbf_interpolation_coefficient_edge" @@ -40,6 +41,14 @@ icon_var_name="c_lin_e", dtype=ta.wpfloat, ), + NUDGECOEFFS_E: dict( + standard_name=NUDGECOEFFS_E, + long_name="nudging_coefficients_for_edges", + units="", # TODO (Yilu) : need to check unit + dims=(dims.EdgeDim,), + icon_var_name="nudgecoeffs_e", + dtype=ta.wpfloat, + ), C_BLN_AVG: dict( standard_name=C_BLN_AVG, long_name="mass conserving bilinear cell average weight", 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 e9f8a66370..05259e44a8 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -5,7 +5,6 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause - import functools import logging from typing import Optional @@ -13,6 +12,7 @@ import gt4py.next as gtx from gt4py.next import backend as gtx_backend +import icon4py.model.common.metrics.compute_nudgecoeffs as common_metrics from icon4py.model.common import constants, dimension as dims from icon4py.model.common.decomposition import definitions from icon4py.model.common.grid import ( @@ -20,6 +20,7 @@ geometry_attributes as geometry_attrs, horizontal as h_grid, icon, + refinement, ) from icon4py.model.common.interpolation import ( interpolation_attributes as attrs, @@ -59,6 +60,9 @@ def __init__( self._config = { "divavg_cntrwgt": 0.5, "weighting_factor": 0.0, + "nudge_max_coeffs": 0.375, + "nudge_efold_width": 2.0, + "nudge_zone_width": 10, "rbf_kernel_cell": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.CELL], "rbf_kernel_edge": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.EDGE], "rbf_kernel_vertex": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.VERTEX], @@ -76,6 +80,15 @@ def __init__( f"initialized interpolation factory for backend = '{self._backend_name()}' and grid = '{self._grid}'" ) log.debug(f"using array_ns {self._xp} ") + + self.register_provider( + factory.PrecomputedFieldProvider( + { + "refinement_control_at_edges": self._grid.refinement_control[dims.EdgeDim], + } + ) + ) + self._register_computed_fields() def __repr__(self): @@ -86,6 +99,29 @@ def _sources(self) -> factory.FieldSource: return factory.CompositeSource(self, (self._geometry,)) def _register_computed_fields(self): + nudging_coefficients_for_edges = factory.ProgramFieldProvider( + func=common_metrics.compute_nudgecoeffs.with_backend(None), + domain={ + dims.EdgeDim: ( + edge_domain(h_grid.Zone.NUDGING_LEVEL_2), + edge_domain(h_grid.Zone.END), + ) + }, + fields={attrs.NUDGECOEFFS_E: attrs.NUDGECOEFFS_E}, + deps={ + "refin_ctrl": "refinement_control_at_edges", + }, + params={ + "grf_nudge_start_e": refinement.refine_control_value( + dims.EdgeDim, h_grid.Zone.NUDGING + ).value, + "nudge_max_coeffs": self._config["nudge_max_coeffs"], + "nudge_efold_width": self._config["nudge_efold_width"], + "nudge_zone_width": self._config["nudge_zone_width"], + }, + ) + self.register_provider(nudging_coefficients_for_edges) + geofac_div = factory.EmbeddedFieldOperatorProvider( # needs to be computed on fieldview-embedded backend func=interpolation_fields.compute_geofac_div.with_backend(None), diff --git a/model/common/src/icon4py/model/common/metrics/compute_nudgecoeffs.py b/model/common/src/icon4py/model/common/metrics/compute_nudgecoeffs.py index bd178722b4..da1dae0e99 100644 --- a/model/common/src/icon4py/model/common/metrics/compute_nudgecoeffs.py +++ b/model/common/src/icon4py/model/common/metrics/compute_nudgecoeffs.py @@ -33,7 +33,7 @@ def _compute_nudgecoeffs( # TODO (@halungge) not registered in factory @program(grid_type=GridType.UNSTRUCTURED) def compute_nudgecoeffs( - nudgecoeffs_e: fa.EdgeField[wpfloat], + nudging_coefficients_for_edges: fa.EdgeField[wpfloat], refin_ctrl: fa.EdgeField[gtx.int32], grf_nudge_start_e: gtx.int32, nudge_max_coeffs: wpfloat, @@ -49,6 +49,6 @@ def compute_nudgecoeffs( nudge_max_coeffs, nudge_efold_width, nudge_zone_width, - out=nudgecoeffs_e, + out=nudging_coefficients_for_edges, domain={dims.EdgeDim: (horizontal_start, horizontal_end)}, ) 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 03c4cf5ddd..39b16766f7 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 @@ -19,7 +19,7 @@ def compute_zdiff_gradp_dsl( c_lin_e: data_alloc.NDArray, z_ifc: data_alloc.NDArray, flat_idx: data_alloc.NDArray, - z_ifc_sliced: data_alloc.NDArray, + topography: data_alloc.NDArray, nlev: int, horizontal_start: int, horizontal_start_1: int, @@ -27,7 +27,7 @@ def compute_zdiff_gradp_dsl( ) -> data_alloc.NDArray: 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_aux1 = array_ns.maximum(topography[e2c[:, 0]], topography[e2c[:, 1]]) z_aux2 = z_aux1 - 5.0 # extrapol_dist zdiff_gradp = array_ns.zeros_like(z_mc[e2c]) zdiff_gradp[horizontal_start:, :, :] = ( 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 ef6cf60e39..7ca22635c9 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -652,7 +652,7 @@ def _compute_downward_extrapolation_distance( def _compute_pressure_gradient_downward_extrapolation_mask_distance( z_mc: fa.CellKField[wpfloat], c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], wpfloat], - z_ifc_sliced: fa.CellField[wpfloat], + topography: fa.CellField[wpfloat], e_owner_mask: fa.EdgeField[bool], flat_idx_max: fa.EdgeField[gtx.int32], e_lev: fa.EdgeField[gtx.int32], @@ -668,7 +668,7 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance( Args: z_mc: height of cells [m] c_lin_e: interpolation coefficient from cells to edges - z_ifc_sliced: ground level height of cells [m] + topography: ground level height of cells [m] e_owner_mask: mask edges owned by PE. flat_idx_max: level from where edge levels start to become flat e_lev: edge indices @@ -685,7 +685,7 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance( e_lev = broadcast(e_lev, (dims.EdgeDim, dims.KDim)) k_lev = broadcast(k_lev, (dims.EdgeDim, dims.KDim)) z_me = _cell_2_edge_interpolation(in_field=z_mc, coeff=c_lin_e) - downward_distance = _compute_downward_extrapolation_distance(z_ifc_sliced) + downward_distance = _compute_downward_extrapolation_distance(topography) extrapolation_distance = concat_where( (horizontal_start_distance <= dims.EdgeDim) & (dims.EdgeDim < horizontal_end_distance), downward_distance, @@ -708,7 +708,7 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance( def compute_pressure_gradient_downward_extrapolation_mask_distance( z_mc: fa.CellKField[wpfloat], c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], float], - z_ifc_sliced: fa.CellField[wpfloat], + topography: fa.CellField[wpfloat], e_owner_mask: fa.EdgeField[bool], flat_idx_max: fa.EdgeField[gtx.int32], e_lev: fa.EdgeField[gtx.int32], @@ -725,7 +725,7 @@ def compute_pressure_gradient_downward_extrapolation_mask_distance( _compute_pressure_gradient_downward_extrapolation_mask_distance( z_mc=z_mc, c_lin_e=c_lin_e, - z_ifc_sliced=z_ifc_sliced, + topography=topography, flat_idx_max=flat_idx_max, e_owner_mask=e_owner_mask, e_lev=e_lev, diff --git a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py index 489cc8ab67..e5ad2c0b8b 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py +++ b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py @@ -63,6 +63,7 @@ ZD_DIFFCOEF_DSL: Final[str] = "zd_diffcoef_dsl" ZD_INTCOEF_DSL: Final[str] = "zd_intcoef_dsl" ZD_VERTOFFSET_DSL: Final[str] = "zd_vertoffset_dsl" +CELL_HEIGHT_ON_HALF_LEVEL: Final[str] = "vertical_coordinates_on_half_levels" attrs: dict[str, model.FieldMetaData] = { @@ -410,5 +411,12 @@ icon_var_name="zd_vertoffset_dsl", dtype=ta.wpfloat, ), + CELL_HEIGHT_ON_HALF_LEVEL: dict( + standard_name=CELL_HEIGHT_ON_HALF_LEVEL, + long_name="vertical_coordinates_on_half_levels", + units="m", + dims=(dims.CellDim, dims.KHalfDim), + icon_var_name="z_ifc", + dtype=ta.wpfloat, + ), } -CELL_HEIGHT_ON_INTERFACE_LEVEL = "height_on_interface_levels" diff --git a/model/common/src/icon4py/model/common/metrics/metrics_factory.py b/model/common/src/icon4py/model/common/metrics/metrics_factory.py index 0cf91f25de..cb1e019bd5 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_factory.py +++ b/model/common/src/icon4py/model/common/metrics/metrics_factory.py @@ -23,7 +23,6 @@ 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 import cell_2_edge_interpolation from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( @@ -53,16 +52,15 @@ class MetricsFieldsFactory(factory.FieldSource, factory.GridProvider): def __init__( self, grid: icon.IconGrid, - vertical_grid: VerticalGrid, + vertical_grid: v_grid.VerticalGrid, decomposition_info: definitions.DecompositionInfo, geometry_source: geometry.GridGeometry, + topography: gtx.Field, 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, @@ -88,7 +86,7 @@ def __init__( "divdamp_trans_start": 12500.0, "divdamp_trans_end": 17500.0, "divdamp_type": 3, - "damping_height": damping_height, + "damping_height": vertical_grid.config.rayleigh_damping_height, "rayleigh_type": rayleigh_type, "rayleigh_coeff": rayleigh_coeff, "exner_expol": exner_expol, @@ -99,9 +97,7 @@ def __init__( "thhgtd_zdiffu": 125.0, "vct_a_1": vct_a_1, } - z_ifc_sliced = gtx.as_field( - (dims.CellDim,), interface_model_height.ndarray[:, self._grid.num_levels] - ) + k_index = data_alloc.index_field( self._grid, dims.KDim, extend={dims.KDim: 1}, backend=self._backend ) @@ -116,8 +112,7 @@ def __init__( self.register_provider( factory.PrecomputedFieldProvider( { - attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL: interface_model_height, - "z_ifc_sliced": z_ifc_sliced, + "topography": topography, "vct_a": vct_a, "c_refin_ctrl": c_refin_ctrl, "e_refin_ctrl": e_refin_ctrl, @@ -138,6 +133,38 @@ def _sources(self) -> factory.FieldSource: return factory.CompositeSource(self, (self._geometry, self._interpolation_source)) def _register_computed_fields(self): + vertical_coordinates_on_half_levels = factory.NumpyFieldsProvider( + func=functools.partial( + v_grid.compute_vertical_coordinate, + array_ns=self._xp, + ), + fields=(attrs.CELL_HEIGHT_ON_HALF_LEVEL,), + domain={ + dims.CellDim: (0, cell_domain(h_grid.Zone.END)), + dims.KDim: (vertical_domain(v_grid.Zone.TOP), vertical_domain(v_grid.Zone.BOTTOM)), + }, + deps={ + "vct_a": "vct_a", + "topography": "topography", + "cell_areas": geometry_attrs.CELL_AREA, + "geofac_n2s": interpolation_attributes.GEOFAC_N2S, + }, + connectivities={"c2e2co": dims.C2E2CODim}, + params={ + "nflatlev": self._vertical_grid.nflatlev, + "model_top_height": self._vertical_grid.config.model_top_height, + "SLEVE_decay_scale_1": self.vertical_grid.config.SLEVE_decay_scale_1, + "SLEVE_decay_exponent": self._vertical_grid.config.SLEVE_decay_exponent, + "SLEVE_decay_scale_2": self._vertical_grid.config.SLEVE_decay_scale_2, + "SLEVE_minimum_layer_thickness_1": self._vertical_grid.config.SLEVE_minimum_layer_thickness_1, + "SLEVE_minimum_relative_layer_thickness_1": self._vertical_grid.config.SLEVE_minimum_relative_layer_thickness_1, + "SLEVE_minimum_layer_thickness_2": self._vertical_grid.config.SLEVE_minimum_layer_thickness_2, + "SLEVE_minimum_relative_layer_thickness_2": self._vertical_grid.config.SLEVE_minimum_relative_layer_thickness_2, + "lowest_layer_thickness": self._vertical_grid.config.lowest_layer_thickness, + }, + ) + self.register_provider(vertical_coordinates_on_half_levels) + height = factory.ProgramFieldProvider( func=math_helpers.average_two_vertical_levels_downwards_on_cells.with_backend( self._backend @@ -150,7 +177,7 @@ def _register_computed_fields(self): ), }, fields={"average": attrs.Z_MC}, - deps={"input_field": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL}, + deps={"input_field": attrs.CELL_HEIGHT_ON_HALF_LEVEL}, ) self.register_provider(height) @@ -168,7 +195,7 @@ def _register_computed_fields(self): }, fields={"ddqz_z_half": attrs.DDQZ_Z_HALF}, deps={ - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "z_mc": attrs.Z_MC, }, params={"nlev": self._grid.num_levels}, @@ -177,7 +204,7 @@ def _register_computed_fields(self): ddqz_z_full_and_inverse = factory.ProgramFieldProvider( func=mf.compute_ddqz_z_full_and_inverse.with_backend(self._backend), - deps={"z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL}, + deps={"z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL}, domain={ dims.CellDim: ( cell_domain(h_grid.Zone.LOCAL), @@ -191,6 +218,7 @@ def _register_computed_fields(self): 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) + ddqz_full_on_edges = factory.ProgramFieldProvider( func=cell_2_edge_interpolation.cell_2_edge_interpolation.with_backend(self._backend), deps={"in_field": attrs.DDQZ_Z_FULL, "coeff": interpolation_attributes.C_LIN_E}, @@ -250,7 +278,7 @@ def _register_computed_fields(self): func=mf.compute_coeff_dwdz.with_backend(self._backend), deps={ "ddqz_z_full": attrs.DDQZ_Z_FULL, - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, }, domain={ dims.CellDim: ( @@ -329,7 +357,7 @@ def _register_computed_fields(self): compute_cell_to_vertex_interpolation = factory.ProgramFieldProvider( func=compute_cell_2_vertex_interpolation.with_backend(self._backend), deps={ - "cell_in": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "cell_in": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "c_int": interpolation_attributes.CELL_AW_VERTS, }, domain={ @@ -349,7 +377,7 @@ def _register_computed_fields(self): compute_ddxt_z_half_e = factory.ProgramFieldProvider( func=mf.compute_ddxt_z_half_e.with_backend(self._backend), deps={ - "cell_in": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "cell_in": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "c_int": interpolation_attributes.CELL_AW_VERTS, "inv_primal_edge_length": f"inverse_of_{geometry_attrs.EDGE_LENGTH}", "tangent_orientation": geometry_attrs.TANGENT_ORIENTATION, @@ -371,7 +399,7 @@ def _register_computed_fields(self): compute_ddxn_z_half_e = factory.ProgramFieldProvider( func=mf.compute_ddxn_z_half_e.with_backend(self._backend), deps={ - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "inv_dual_edge_length": f"inverse_of_{geometry_attrs.DUAL_EDGE_LENGTH}", }, domain={ @@ -436,7 +464,7 @@ def _register_computed_fields(self): fields=(attrs.EXNER_W_IMPLICIT_WEIGHT_PARAMETER,), deps={ "vct_a": "vct_a", - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "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, @@ -495,7 +523,7 @@ def _register_computed_fields(self): wgtfac_c_provider = factory.ProgramFieldProvider( func=weight_factors.compute_wgtfac_c.with_backend(self._backend), deps={ - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, }, domain={ dims.CellDim: (cell_domain(h_grid.Zone.LOCAL), cell_domain(h_grid.Zone.END)), @@ -533,7 +561,7 @@ def _register_computed_fields(self): deps={ "z_mc": attrs.Z_MC, "c_lin_e": interpolation_attributes.C_LIN_E, - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "k_lev": "k_lev", }, domain={ @@ -566,7 +594,7 @@ def _register_computed_fields(self): deps={ "z_mc": attrs.Z_MC, "c_lin_e": interpolation_attributes.C_LIN_E, - "z_ifc_sliced": "z_ifc_sliced", + "topography": "topography", "e_owner_mask": "e_owner_mask", "flat_idx_max": attrs.FLAT_IDX_MAX, "e_lev": "e_lev", @@ -634,9 +662,9 @@ def _register_computed_fields(self): deps={ "z_mc": attrs.Z_MC, "c_lin_e": interpolation_attributes.C_LIN_E, - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "flat_idx": attrs.FLAT_IDX_MAX, - "z_ifc_sliced": "z_ifc_sliced", + "topography": "topography", }, connectivities={"e2c": dims.E2CDim}, domain=(dims.EdgeDim, dims.KDim), @@ -676,7 +704,7 @@ def _register_computed_fields(self): func=functools.partial(weight_factors.compute_wgtfacq_c_dsl, array_ns=self._xp), domain=(dims.CellDim, dims.KDim), fields=(attrs.WGTFACQ_C,), - deps={"z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL}, + deps={"z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL}, params={"nlev": self._grid.num_levels}, ) @@ -685,7 +713,7 @@ def _register_computed_fields(self): compute_wgtfacq_e = factory.NumpyFieldsProvider( func=functools.partial(weight_factors.compute_wgtfacq_e_dsl, array_ns=self._xp), deps={ - "z_ifc": attrs.CELL_HEIGHT_ON_INTERFACE_LEVEL, + "z_ifc": attrs.CELL_HEIGHT_ON_HALF_LEVEL, "c_lin_e": interpolation_attributes.C_LIN_E, "wgtfacq_c_dsl": attrs.WGTFACQ_C, }, diff --git a/model/common/tests/grid_tests/test_topography.py b/model/common/tests/grid_tests/test_topography.py index c8addaffe4..8fd8a03086 100644 --- a/model/common/tests/grid_tests/test_topography.py +++ b/model/common/tests/grid_tests/test_topography.py @@ -9,6 +9,7 @@ import pytest from icon4py.model.common.grid import geometry, topography as topo +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import datatest_utils as dt_utils, helpers @@ -17,7 +18,8 @@ @pytest.mark.parametrize( "experiment", [ - (dt_utils.GAUSS3D_EXPERIMENT), + dt_utils.GAUSS3D_EXPERIMENT, + dt_utils.REGIONAL_EXPERIMENT, ], ) def test_topography_smoothing_with_serialized_data( @@ -33,17 +35,16 @@ def test_topography_smoothing_with_serialized_data( num_iterations = 25 topography = topography_savepoint.topo_c() - topography_smoothed_verif_np = topography_savepoint.topo_smt_c().asnumpy() + xp = data_alloc.import_array_ns(backend) + topography_smoothed_ref = topography_savepoint.topo_smt_c().asnumpy() topography_smoothed = topo.smooth_topography( - topography=topography, - grid=icon_grid, - cell_areas=cell_geometry.area, - geofac_n2s=geofac_n2s, - backend=backend, + topography=topography.ndarray, + cell_areas=cell_geometry.area.ndarray, + geofac_n2s=geofac_n2s.ndarray, + c2e2co=icon_grid.get_connectivity("C2E2CO").ndarray, num_iterations=num_iterations, + array_ns=xp, ) - assert helpers.dallclose( - topography_smoothed_verif_np, topography_smoothed.asnumpy(), atol=1.0e-14 - ) + assert helpers.dallclose(topography_smoothed_ref, topography_smoothed, atol=1.0e-14) diff --git a/model/common/tests/grid_tests/test_vertical.py b/model/common/tests/grid_tests/test_vertical.py index 084412fdae..f9e41bf738 100644 --- a/model/common/tests/grid_tests/test_vertical.py +++ b/model/common/tests/grid_tests/test_vertical.py @@ -297,9 +297,12 @@ def test_vct_a_vct_b_calculation_from_icon_input( assert helpers.dallclose(vct_b.asnumpy(), grid_savepoint.vct_b().asnumpy()) -@pytest.mark.embedded_remap_error +@pytest.mark.level("unit") @pytest.mark.datatest -@pytest.mark.parametrize("experiment", [dt_utils.GAUSS3D_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT]) +@pytest.mark.parametrize( + "experiment", + [dt_utils.REGIONAL_EXPERIMENT, dt_utils.GAUSS3D_EXPERIMENT, dt_utils.GLOBAL_EXPERIMENT], +) def test_compute_vertical_coordinate( grid_savepoint, metrics_savepoint, @@ -307,44 +310,70 @@ def test_compute_vertical_coordinate( interpolation_savepoint, icon_grid, experiment, + model_top_height, backend, ): xp = data_alloc.array_ns(data_alloc.is_cupy_device(backend)) vct_a = grid_savepoint.vct_a() vct_b = grid_savepoint.vct_b() cell_geometry = grid_savepoint.construct_cell_geometry() + + specific_values = ( + {"rayleigh_damping_height": 12500.0, "stretch_factor": 0.65, "lowest_layer_thickness": 20.0} + if experiment == dt_utils.REGIONAL_EXPERIMENT + else { + "rayleigh_damping_height": 45000.0, + "stretch_factor": 1.0, + "lowest_layer_thickness": 50.0, + } + ) + vertical_config = v_grid.VerticalGridConfig( num_levels=grid_savepoint.num(dims.KDim), + flat_height=16000.0, + htop_moist_proc=22500.0, + maximal_layer_thickness=25000.0, + **specific_values, ) + vertical_geometry = v_grid.VerticalGrid( config=vertical_config, vct_a=vct_a, vct_b=vct_b, ) - if experiment == dt_utils.GAUSS3D_EXPERIMENT: + assert vertical_geometry.nflatlev == grid_savepoint.nflatlev() + + if experiment in (dt_utils.REGIONAL_EXPERIMENT, dt_utils.GAUSS3D_EXPERIMENT): topography = topography_savepoint.topo_c() elif experiment == dt_utils.GLOBAL_EXPERIMENT: topography = data_alloc.zero_field( icon_grid, dims.CellDim, backend=backend, dtype=ta.wpfloat ) - else: - raise ValueError(f"Unsupported experiment: {experiment}") geofac_n2s = interpolation_savepoint.geofac_n2s() - vertical_coordinates_on_cell_khalf = v_grid.compute_vertical_coordinate( + vertical_coordinates_on_half_levels = v_grid.compute_vertical_coordinate( vct_a=vct_a.ndarray, topography=topography.ndarray, - geofac_n2s=geofac_n2s.ndarray, cell_areas=cell_geometry.area.ndarray, - grid=icon_grid, - vertical_geometry=vertical_geometry, - backend=backend, + geofac_n2s=geofac_n2s.ndarray, + c2e2co=icon_grid.get_connectivity("C2E2CO").ndarray, + nflatlev=vertical_geometry.nflatlev, + model_top_height=model_top_height, + SLEVE_decay_scale_1=4000.0, + SLEVE_decay_exponent=1.2, + SLEVE_decay_scale_2=2500.0, + SLEVE_minimum_layer_thickness_1=100.0, + SLEVE_minimum_relative_layer_thickness_1=1.0 / 3.0, + SLEVE_minimum_layer_thickness_2=500.0, + SLEVE_minimum_relative_layer_thickness_2=0.5, + lowest_layer_thickness=vertical_config.lowest_layer_thickness, array_ns=xp, ) assert helpers.dallclose( - data_alloc.as_numpy(vertical_coordinates_on_cell_khalf), + data_alloc.as_numpy(vertical_coordinates_on_half_levels), metrics_savepoint.z_ifc().asnumpy(), + rtol=1e-9, atol=1e-13, ) diff --git a/model/common/tests/interpolation_tests/test_interpolation_factory.py b/model/common/tests/interpolation_tests/test_interpolation_factory.py index fd3573594a..895f7f14ca 100644 --- a/model/common/tests/interpolation_tests/test_interpolation_factory.py +++ b/model/common/tests/interpolation_tests/test_interpolation_factory.py @@ -316,6 +316,23 @@ def test_cells_aw_verts(interpolation_savepoint, grid_file, experiment, backend, assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) +@pytest.mark.level("integration") +@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_nudgecoeffs(interpolation_savepoint, grid_file, experiment, backend, rtol): + field_ref = interpolation_savepoint.nudgecoeff_e() + factory = _get_interpolation_factory(backend, experiment, grid_file) + field = factory.get(attrs.NUDGECOEFFS_E) + + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=rtol) + + @pytest.mark.level("integration") @pytest.mark.parametrize( "grid_file, experiment, atol", diff --git a/model/common/tests/metric_tests/conftest.py b/model/common/tests/metric_tests/conftest.py index bb98485929..65166bee11 100644 --- a/model/common/tests/metric_tests/conftest.py +++ b/model/common/tests/metric_tests/conftest.py @@ -16,4 +16,5 @@ metrics_savepoint, processor_props, ranked_data_path, + topography_savepoint, ) diff --git a/model/common/tests/metric_tests/test_metrics_factory.py b/model/common/tests/metric_tests/test_metrics_factory.py index 5ae3d7a64f..6942a5c434 100644 --- a/model/common/tests/metric_tests/test_metrics_factory.py +++ b/model/common/tests/metric_tests/test_metrics_factory.py @@ -69,12 +69,16 @@ def _get_metrics_factory( backend: Optional[gtx_backend.Backend], experiment: str, grid_file: str, + icon_grid, grid_savepoint: serialbox.IconGridSavepoint, + topography_savepoint: serialbox.TopographySavepoint, metrics_savepoint: serialbox.MetricSavepoint, ) -> metrics_factory.MetricsFieldsFactory: registry_name = "_".join((experiment, data_alloc.backend_name(backend))) factory = metrics_factories.get(registry_name) + topography = topography_savepoint.topo_c() + if not factory: geometry = gridtest_utils.get_grid_geometry(backend, experiment, grid_file) ( @@ -110,13 +114,12 @@ def _get_metrics_factory( vertical_grid=vertical_grid, decomposition_info=geometry._decomposition_info, geometry_source=geometry, + topography=topography, interpolation_source=interpolation_field_source, 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, @@ -135,17 +138,27 @@ def _get_metrics_factory( ], ) @pytest.mark.datatest -def test_factory_z_mc(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_z_mc( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.z_mc() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.Z_MC) - assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-10) @pytest.mark.level("integration") @@ -158,7 +171,13 @@ def test_factory_z_mc(grid_savepoint, metrics_savepoint, grid_file, experiment, ) @pytest.mark.datatest def test_factory_ddqz_z_and_inverse( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): inverse_field_ref = metrics_savepoint.inv_ddqz_z_full() field_ref = metrics_savepoint.ddqz_z_full() @@ -166,13 +185,15 @@ def test_factory_ddqz_z_and_inverse( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) inverse_field = factory.get(attrs.INV_DDQZ_Z_FULL) field = factory.get(attrs.DDQZ_Z_FULL) - assert test_helpers.dallclose(inverse_field_ref.asnumpy(), inverse_field.asnumpy()) - assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + assert test_helpers.dallclose(inverse_field_ref.asnumpy(), inverse_field.asnumpy(), atol=1e-10) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-7) @pytest.mark.parametrize( @@ -183,14 +204,24 @@ def test_factory_ddqz_z_and_inverse( ], ) @pytest.mark.datatest -def test_factory_ddqz_full_e(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_ddqz_full_e( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.ddqz_z_full_e().asnumpy() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.DDQZ_Z_FULL_E) assert test_helpers.dallclose(field_ref, field.asnumpy(), rtol=1e-8) @@ -205,17 +236,27 @@ def test_factory_ddqz_full_e(grid_savepoint, metrics_savepoint, grid_file, exper ], ) @pytest.mark.datatest -def test_factory_ddqz_z_half(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_ddqz_z_half( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.ddqz_z_half() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.DDQZ_Z_HALF) - assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-9) @pytest.mark.level("integration") @@ -228,15 +269,23 @@ def test_factory_ddqz_z_half(grid_savepoint, metrics_savepoint, grid_file, exper ) @pytest.mark.datatest def test_factory_scaling_factor_for_3d_divdamp( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref = metrics_savepoint.scalfac_dd3d() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.SCALING_FACTOR_FOR_3D_DIVDAMP) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) @@ -251,14 +300,24 @@ def test_factory_scaling_factor_for_3d_divdamp( ], ) @pytest.mark.datatest -def test_factory_rayleigh_w(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_rayleigh_w( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.rayleigh_w() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.RAYLEIGH_W) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) @@ -273,20 +332,30 @@ def test_factory_rayleigh_w(grid_savepoint, metrics_savepoint, grid_file, experi ], ) @pytest.mark.datatest -def test_factory_coeffs_dwdz(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_coeffs_dwdz( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + 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, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_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()) + assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy(), atol=1e-11) + assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy(), atol=1e-11) @pytest.mark.level("integration") @@ -298,20 +367,30 @@ def test_factory_coeffs_dwdz(grid_savepoint, metrics_savepoint, grid_file, exper ], ) @pytest.mark.datatest -def test_factory_ref_mc(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_ref_mc( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + 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, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_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()) + assert test_helpers.dallclose(field_ref_1.asnumpy(), field_1.asnumpy(), atol=1e-9) + assert test_helpers.dallclose(field_ref_2.asnumpy(), field_2.asnumpy(), atol=1e-10) @pytest.mark.level("integration") @@ -324,7 +403,13 @@ def test_factory_ref_mc(grid_savepoint, metrics_savepoint, grid_file, experiment ) @pytest.mark.datatest def test_factory_d2dexdz2_facs_mc( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref_1 = metrics_savepoint.d2dexdz2_fac1_mc() field_ref_2 = metrics_savepoint.d2dexdz2_fac2_mc() @@ -332,13 +417,15 @@ def test_factory_d2dexdz2_facs_mc( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_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()) + assert test_helpers.dallclose(field_1.asnumpy(), field_ref_1.asnumpy(), atol=1e-12) + assert test_helpers.dallclose(field_2.asnumpy(), field_ref_2.asnumpy(), atol=1e-12) @pytest.mark.parametrize( @@ -349,17 +436,27 @@ def test_factory_d2dexdz2_facs_mc( ], ) @pytest.mark.datatest -def test_factory_ddxn_z_full(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_ddxn_z_full( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.ddxn_z_full() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.DDXN_Z_FULL) - assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-8) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), atol=1e-8) @pytest.mark.level("integration") @@ -372,15 +469,24 @@ def test_factory_ddxn_z_full(grid_savepoint, metrics_savepoint, grid_file, exper ) @pytest.mark.datatest def test_factory_ddxt_z_full( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend, caplog + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, + caplog, ): field_ref = metrics_savepoint.ddxt_z_full().asnumpy() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.DDXT_Z_FULL) caplog.set_level(logging.DEBUG) @@ -401,15 +507,23 @@ def test_factory_ddxt_z_full( ) @pytest.mark.datatest def test_factory_exner_w_implicit_weight_parameter( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref = metrics_savepoint.vwind_impl_wgt() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.EXNER_W_IMPLICIT_WEIGHT_PARAMETER) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-9) @@ -428,15 +542,23 @@ def test_factory_exner_w_implicit_weight_parameter( ) @pytest.mark.datatest def test_factory_exner_w_explicit_weight_parameter( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref = metrics_savepoint.vwind_expl_wgt() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.EXNER_W_EXPLICIT_WEIGHT_PARAMETER) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-8) @@ -452,17 +574,27 @@ def test_factory_exner_w_explicit_weight_parameter( ], ) @pytest.mark.datatest -def test_factory_exner_exfac(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_exner_exfac( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.exner_exfac() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.EXNER_EXFAC) - assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1.0e-5) + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), atol=1e-8) @pytest.mark.level("integration") @@ -476,7 +608,13 @@ def test_factory_exner_exfac(grid_savepoint, metrics_savepoint, grid_file, exper ) @pytest.mark.datatest def test_factory_pressure_gradient_fields( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_1_ref = metrics_savepoint.pg_exdist() field_2_ref = metrics_savepoint.pg_edgeidx_dsl() @@ -484,8 +622,10 @@ def test_factory_pressure_gradient_fields( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field_1 = factory.get(attrs.PG_EDGEDIST_DSL) assert test_helpers.dallclose(field_1_ref.asnumpy(), field_1.asnumpy(), atol=1.0e-5) @@ -502,7 +642,13 @@ def test_factory_pressure_gradient_fields( ) @pytest.mark.datatest def test_factory_mask_bdy_prog_halo_c( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref_1 = metrics_savepoint.mask_prog_halo_c() field_ref_2 = metrics_savepoint.bdy_halo_c() @@ -510,8 +656,10 @@ def test_factory_mask_bdy_prog_halo_c( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field_1 = factory.get(attrs.MASK_PROG_HALO_C) field_2 = factory.get(attrs.BDY_HALO_C) @@ -529,15 +677,23 @@ def test_factory_mask_bdy_prog_halo_c( ) @pytest.mark.datatest def test_factory_horizontal_mask_for_3d_divdamp( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref = metrics_savepoint.hmask_dd3d() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.HORIZONTAL_MASK_FOR_3D_DIVDAMP) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy()) @@ -554,14 +710,24 @@ def test_factory_horizontal_mask_for_3d_divdamp( ], ) @pytest.mark.datatest -def test_factory_zdiff_gradp(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_zdiff_gradp( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.zdiff_gradp() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.ZDIFF_GRADP) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), atol=1.0e-5) @@ -576,14 +742,24 @@ def test_factory_zdiff_gradp(grid_savepoint, metrics_savepoint, grid_file, exper ], ) @pytest.mark.datatest -def test_factory_coeff_gradekin(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_coeff_gradekin( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): field_ref = metrics_savepoint.coeff_gradekin() factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, ) field = factory.get(attrs.COEFF_GRADEKIN) assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-8) @@ -598,19 +774,61 @@ def test_factory_coeff_gradekin(grid_savepoint, metrics_savepoint, grid_file, ex ], ) @pytest.mark.datatest -def test_factory_wgtfacq_e(grid_savepoint, metrics_savepoint, grid_file, experiment, backend): +def test_factory_wgtfacq_e( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): factory = _get_metrics_factory( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_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.level("integration") +@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_vertical_coordinates_on_half_levels( + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, +): + factory = _get_metrics_factory( + backend=backend, + experiment=experiment, + grid_file=grid_file, + icon_grid=icon_grid, + grid_savepoint=grid_savepoint, + metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_savepoint, + ) + field = factory.get(attrs.CELL_HEIGHT_ON_HALF_LEVEL) + field_ref = metrics_savepoint.z_ifc() + assert test_helpers.dallclose(field_ref.asnumpy(), field.asnumpy(), rtol=1e-9) + + @pytest.mark.level("integration") @pytest.mark.embedded_remap_error @pytest.mark.parametrize( @@ -621,7 +839,13 @@ def test_factory_wgtfacq_e(grid_savepoint, metrics_savepoint, grid_file, experim ) @pytest.mark.datatest def test_factory_compute_diffusion_metrics( - grid_savepoint, metrics_savepoint, grid_file, experiment, backend + grid_savepoint, + metrics_savepoint, + topography_savepoint, + grid_file, + icon_grid, + experiment, + backend, ): field_ref_1 = metrics_savepoint.mask_hdiff() field_ref_2 = metrics_savepoint.zd_diffcoef() @@ -631,14 +855,16 @@ def test_factory_compute_diffusion_metrics( backend=backend, experiment=experiment, grid_file=grid_file, + icon_grid=icon_grid, grid_savepoint=grid_savepoint, metrics_savepoint=metrics_savepoint, + topography_savepoint=topography_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_2.asnumpy(), field_2.asnumpy(), atol=1.0e-10) + assert test_helpers.dallclose(field_ref_3.asnumpy(), field_3.asnumpy(), atol=1.0e-8) assert test_helpers.dallclose(field_ref_4.asnumpy(), field_4.asnumpy())