diff --git a/bindings/src/icon4py/bindings/dycore_wrapper.py b/bindings/src/icon4py/bindings/dycore_wrapper.py index e14b966580..5b0f36c668 100644 --- a/bindings/src/icon4py/bindings/dycore_wrapper.py +++ b/bindings/src/icon4py/bindings/dycore_wrapper.py @@ -374,7 +374,7 @@ def solve_nh_run( # noqa: PLR0917 [too-many-positional-arguments] dynamical_vertical_volumetric_flux_at_cells_on_half_levels=vol_flx_ic, ) - # Make `max_vcfl` a 0-d array to avoid cupy synchronization, see `velocity_advection.py`. + # Make `max_vcfl` a 0-d array to avoid cupy synchronization, see `_update_max_vertical_cfl`. # Note, `max_vcfl` needs to be passed back to Fortran after the timestep. max_vcfl = data_alloc.scalar_like_array(max_vcfl_size1_array[0], xp) diff --git a/model/atmosphere/dycore/docs/dycore_numerics.rst b/model/atmosphere/dycore/docs/dycore_numerics.rst index dd994eafc3..a14592826b 100644 --- a/model/atmosphere/dycore/docs/dycore_numerics.rst +++ b/model/atmosphere/dycore/docs/dycore_numerics.rst @@ -9,5 +9,4 @@ described in detail. :maxdepth: 2 :caption: Dycore subcomponents: - dycore_numerics_advection dycore_numerics_nonhydro diff --git a/model/atmosphere/dycore/docs/dycore_numerics_advection.rst b/model/atmosphere/dycore/docs/dycore_numerics_advection.rst deleted file mode 100644 index 4b619ac07d..0000000000 --- a/model/atmosphere/dycore/docs/dycore_numerics_advection.rst +++ /dev/null @@ -1,7 +0,0 @@ -Dycore numerical documentation (velocity advection) -=================================================== - -This page documents the advection part of the dycore. - -.. autoscidoc:: icon4py.model.atmosphere.dycore.velocity_advection.VelocityAdvection.run_predictor_step - :no-index: 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 c37b5f9f1e..96bcb02f95 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 @@ -45,7 +45,12 @@ from icon4py.model.atmosphere.dycore.stencils.update_theta_and_exner_in_halo import ( update_theta_and_exner_in_halo, ) -from icon4py.model.atmosphere.dycore.velocity_advection import VelocityAdvection +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_corrector import ( + compute_velocity_advection_in_corrector_step, +) +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_predictor import ( + compute_velocity_advection_in_predictor_step, +) from icon4py.model.common import ( constants, dimension as dims, @@ -456,6 +461,22 @@ def __init__(self, config: NonHydrostaticConfig): """ +def _update_max_vertical_cfl( + diagnostic_state: nonhydro_states.DiagnosticStateNonHydro, + vertical_cfl: fa.CellKHalfField[ta.anyfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, +) -> None: + # Reductions should be performed on flat, contiguous arrays for best cupy performance + # as otherwise cupy won't use cub optimized kernels. + max_vertical_cfl = vertical_cfl.array_ns.max( # type: ignore[attr-defined] + vertical_cfl.ndarray[horizontal_start:horizontal_end, :].ravel(order="K") # type: ignore[attr-defined] + ) + diagnostic_state.max_vertical_cfl = vertical_cfl.array_ns.maximum( # type: ignore[attr-defined] + max_vertical_cfl, diagnostic_state.max_vertical_cfl + ) + + class SolveNonhydro: def __init__( self, @@ -886,15 +907,82 @@ def __init__( }, ) - self.velocity_advection = VelocityAdvection( - grid=grid, - metric_state=metric_state_nonhydro, - interpolation_state=interpolation_state, - vertical_params=vertical_params, - edge_params=edge_geometry, - owner_mask=owner_mask, + cell_horizontal_sizes = { + "start_cell_lateral_boundary_level_4": self._start_cell_lateral_boundary_level_4, + "end_cell_halo": self._end_cell_halo, + "start_edge_nudging_level_2": self._start_edge_nudging_level_2, + "end_edge_local": self._end_edge_local, + } + shared_constant_args: dict[str, gtx.Field | gtx_typing.Scalar] = { + "coeff1_dwdz": self._metric_state_nonhydro.coeff1_dwdz, + "coeff2_dwdz": self._metric_state_nonhydro.coeff2_dwdz, + "c_intp": self._interpolation_state.c_intp, + "inv_dual_edge_length": self._edge_geometry.inverse_dual_edge_lengths, + "inv_primal_edge_length": self._edge_geometry.inverse_primal_edge_lengths, + "tangent_orientation": self._edge_geometry.tangent_orientation, + "e_bln_c_s": self._interpolation_state.e_bln_c_s, + "ddqz_z_half": self._metric_state_nonhydro.ddqz_z_half, + "geofac_n2s": self._interpolation_state.geofac_n2s, + "owner_mask": owner_mask, + "coriolis_frequency": self._edge_geometry.coriolis_frequency, + "geofac_rot": self._interpolation_state.geofac_rot, + "coeff_gradekin": self._metric_state_nonhydro.coeff_gradekin, + "c_lin_e": self._interpolation_state.c_lin_e, + "ddqz_z_full_e": self._metric_state_nonhydro.ddqz_z_full_e, + "area_edge": self._edge_geometry.edge_areas, + "area": self._cell_params.area, + "geofac_grdiv": self._interpolation_state.geofac_grdiv, + } + + self._compute_velocity_advection_in_predictor_step = setup_program( backend=backend, + program=compute_velocity_advection_in_predictor_step, + constant_args={ + "rbf_vec_coeff_e": self._interpolation_state.rbf_vec_coeff_e, + "wgtfac_e": self._metric_state_nonhydro.wgtfac_e, + "wgtfacq_e": self._metric_state_nonhydro.wgtfacq_e, + "ddxn_z_full": self._metric_state_nonhydro.ddxn_z_full, + "ddxt_z_full": self._metric_state_nonhydro.ddxt_z_full, + "wgtfac_c": self._metric_state_nonhydro.wgtfac_c, + **shared_constant_args, + }, + variants={ + "skip_compute_predictor_vertical_advection": [True, False], + # Only True: deriving `apply_extra_diffusion_on_vn` from `max_vertical_cfl` would need a + # device synchronization, so the call site fixes it to True (see the TODO there). + "apply_extra_diffusion_on_vn": [True], + }, + horizontal_sizes={ + "start_edge_lateral_boundary_level_5": self._start_edge_lateral_boundary_level_5, + "end_edge_halo_level_2": self._end_edge_halo_level_2, + **cell_horizontal_sizes, + }, + vertical_sizes={ + "nflatlev": self._vertical_params.nflatlev, + "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer, + "vertical_start": gtx.int32(0), + "vertical_end": self._grid.num_levels, + }, + offset_provider=self._grid.connectivities, ) + + self._compute_velocity_advection_in_corrector_step = setup_program( + backend=backend, + program=compute_velocity_advection_in_corrector_step, + constant_args=shared_constant_args, + variants={ + # Only True, as for the predictor step above. + "apply_extra_diffusion_on_vn": [True], + }, + horizontal_sizes=cell_horizontal_sizes, + vertical_sizes={ + "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer, + "vertical_start": gtx.int32(0), + "vertical_end": self._grid.num_levels, + }, + offset_provider=self._grid.connectivities, + ) + self._allocate_local_fields(model_backends.get_allocator(backend)) self._en_smag_fac_for_zero_nshift( @@ -1008,6 +1096,9 @@ def _allocate_local_fields(self, allocator: gtx_typing.Allocator | None) -> None Declared as enh_divdamp_fac in ICON. """ self.intermediate_fields = IntermediateFields.allocate(grid=self._grid, allocator=allocator) + self._vertical_cfl = data_alloc.zero_field( + self._grid, dims.CellDim, dims.KHalfDim, allocator=allocator, dtype=ta.vpfloat + ) def _determine_local_domains(self) -> None: vertex_domain = h_grid.domain(dims.VertexDim) @@ -1021,6 +1112,9 @@ def _determine_local_domains(self) -> None: self._start_cell_lateral_boundary_level_3 = self._grid.start_index( cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) ) + self._start_cell_lateral_boundary_level_4 = self._grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4) + ) self._start_cell_nudging = self._grid.start_index(cell_domain(h_grid.Zone.NUDGING)) self._start_cell_local = self._grid.start_index(cell_domain(h_grid.Zone.LOCAL)) self._start_cell_halo = self._grid.start_index(cell_domain(h_grid.Zone.HALO)) @@ -1180,17 +1274,34 @@ def run_predictor_step( and not (at_initial_timestep and at_first_substep) ) - assert self._cell_params.area is not None + # Note, if we compute `apply_extra_diffusion_on_vn = max_vertical_cfl > VerticalCflConstants.W_LIMIT` + # from the reduction below, we would have to synchronize with the device before this call. + # TODO (Chia Rui): to decide whether make apply_extra_diffusion_on_vn a config parameter or remove it or always turn on extra diffusion + apply_extra_diffusion_on_vn = True - self.velocity_advection.run_predictor_step( - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - diagnostic_state=diagnostic_state_nh, - prognostic_state=prognostic_states.current, - contravariant_correction_at_edges_on_model_levels=self._contravariant_correction_at_edges_on_model_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels, + # TODO(havogt): however, our test data is probably not able to catch cfl_clipping conditions + self._compute_velocity_advection_in_predictor_step( + tangential_wind=diagnostic_state_nh.tangential_wind, tangential_wind_on_half_levels=z_fields.tangential_wind_on_half_levels, + vn_on_half_levels=diagnostic_state_nh.vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels=self._contravariant_correction_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels=diagnostic_state_nh.contravariant_correction_at_cells_on_half_levels, + vertical_wind_advective_tendency=diagnostic_state_nh.vertical_wind_advective_tendency.predictor, + vertical_cfl=self._vertical_cfl, + normal_wind_advective_tendency=diagnostic_state_nh.normal_wind_advective_tendency.predictor, + vn=prognostic_states.current.vn, + w=prognostic_states.current.w, dtime=dtime, - cell_areas=self._cell_params.area, + skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + ) + + _update_max_vertical_cfl( + diagnostic_state_nh, + self._vertical_cfl, + self._start_cell_lateral_boundary_level_4, + self._end_cell_halo, ) self._compute_perturbed_quantities_and_interpolation( @@ -1356,20 +1467,37 @@ def run_corrector_step( # scaling factor for second-order divergence damping: second_order_divdamp_factor_from_sfc_to_divdamp_z*delta_x**2 # delta_x**2 is approximated by the mean cell area # Coefficient for reduced fourth-order divergence d - assert self._cell_params.area is not None assert self._cell_params.mean_cell_area is not None second_order_divdamp_scaling_coeff = ( second_order_divdamp_factor * self._cell_params.mean_cell_area ) log.debug("corrector run velocity advection") - self.velocity_advection.run_corrector_step( - diagnostic_state=diagnostic_state_nh, - prognostic_state=prognostic_states.next, - horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels, + # Note, if we compute `apply_extra_diffusion_on_vn = max_vertical_cfl > VerticalCflConstants.W_LIMIT` + # from the reduction below, we would have to synchronize with the device before this call. + # TODO (Chia Rui): to decide whether make apply_extra_diffusion_on_vn a config parameter or remove it or always turn on extra diffusion + apply_extra_diffusion_on_vn = True + + self._compute_velocity_advection_in_corrector_step( + vertical_wind_advective_tendency=diagnostic_state_nh.vertical_wind_advective_tendency.corrector, + vertical_cfl=self._vertical_cfl, + normal_wind_advective_tendency=diagnostic_state_nh.normal_wind_advective_tendency.corrector, + vn=prognostic_states.next.vn, + w=prognostic_states.next.w, + tangential_wind=diagnostic_state_nh.tangential_wind, tangential_wind_on_half_levels=z_fields.tangential_wind_on_half_levels, + vn_on_half_levels=diagnostic_state_nh.vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=z_fields.horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels=diagnostic_state_nh.contravariant_correction_at_cells_on_half_levels, dtime=dtime, - cell_areas=self._cell_params.area, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + ) + + _update_max_vertical_cfl( + diagnostic_state_nh, + self._vertical_cfl, + self._start_cell_lateral_boundary_level_4, + self._end_cell_halo, ) self._compute_interpolation_and_nonhydro_buoy( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py index c6ca4999a5..3e67c5c14f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py @@ -5,9 +5,15 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +# Note: Currently unused, kept for implementing a version faithful to ICON. +# The ICON version requires the CFL reduction within velocity advection which +# drains the GPU kernel pipeline. + import gt4py.next as gtx from gt4py.next import abs, astype, minimum, neighbor_sum, where # noqa: A004 +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import VerticalCflConstants from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import E2C, E2C2EO, E2V from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -26,25 +32,24 @@ def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl( geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], vn: fa.EdgeKField[ta.wpfloat], ddt_vn_apc: fa.EdgeKField[ta.vpfloat], - cfl_w_limit: ta.vpfloat, - scalfac_exdiff: ta.wpfloat, dtime: ta.wpfloat, ) -> fa.EdgeKField[ta.vpfloat]: """Formerly known as _mo_velocity_advection_stencil_20.""" - z_w_con_c_full_wp, ddqz_z_full_e_wp, ddt_vn_apc_wp, cfl_w_limit_wp = astype( - (z_w_con_c_full, ddqz_z_full_e, ddt_vn_apc, cfl_w_limit), wpfloat + z_w_con_c_full_wp, ddqz_z_full_e_wp, ddt_vn_apc_wp = astype( + (z_w_con_c_full, ddqz_z_full_e, ddt_vn_apc), wpfloat ) w_con_e = neighbor_sum(c_lin_e * z_w_con_c_full_wp(E2C), axis=dims.E2CDim) - difcoef = scalfac_exdiff * minimum( - wpfloat("0.85") - cfl_w_limit_wp * dtime, - abs(w_con_e) * dtime / ddqz_z_full_e_wp - cfl_w_limit_wp * dtime, + vertical_cfl_number_at_edges = abs(w_con_e) * dtime / ddqz_z_full_e_wp + difcoef = (VerticalCflConstants.EXTRA_DIFFUSION_SCALING / dtime) * minimum( + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + vertical_cfl_number_at_edges - VerticalCflConstants.W_LIMIT, ) ddt_vn_apc_wp = where( # TODO(havogt): my guess is if the second condition is `True`, then # `(levelmask | levelmask(dims.KDim + 1))` is also `True` (levelmask | levelmask(dims.KDim + 1)) - & (abs(w_con_e) > astype(cfl_w_limit * ddqz_z_full_e, wpfloat)), + & (vertical_cfl_number_at_edges > VerticalCflConstants.W_LIMIT), ddt_vn_apc_wp + difcoef * area_edge @@ -72,8 +77,6 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl( geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], vn: fa.EdgeKField[ta.wpfloat], ddt_vn_apc: fa.EdgeKField[ta.vpfloat], - cfl_w_limit: ta.vpfloat, - scalfac_exdiff: ta.wpfloat, dtime: ta.wpfloat, horizontal_start: gtx.int32, horizontal_end: gtx.int32, @@ -92,8 +95,6 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl( geofac_grdiv=geofac_grdiv, vn=vn, ddt_vn_apc=ddt_vn_apc, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, dtime=dtime, out=ddt_vn_apc, domain={ diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_w_con_approaching_cfl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_w_con_approaching_cfl.py deleted file mode 100644 index 8e164584e8..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_extra_diffusion_for_w_con_approaching_cfl.py +++ /dev/null @@ -1,89 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import abs, astype, minimum, neighbor_sum, where # noqa: A004 - -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.dimension import C2E2CO -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _add_extra_diffusion_for_w_con_approaching_cfl( - cfl_clipping: fa.CellKHalfField[bool], - owner_mask: fa.CellField[bool], - z_w_con_c: fa.CellKHalfField[ta.vpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - w: fa.CellKHalfField[ta.wpfloat], - ddt_w_adv: fa.CellKHalfField[ta.vpfloat], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, -) -> fa.CellKHalfField[ta.vpfloat]: - """Formerly known as _mo_velocity_advection_stencil_18.""" - z_w_con_c_wp, ddqz_z_half_wp, ddt_w_adv_wp, cfl_w_limit_wp = astype( - (z_w_con_c, ddqz_z_half, ddt_w_adv, cfl_w_limit), wpfloat - ) - - difcoef = where( - cfl_clipping & owner_mask, - scalfac_exdiff - * minimum( - wpfloat("0.85") - cfl_w_limit_wp * dtime, - abs(z_w_con_c_wp) * dtime / ddqz_z_half_wp - cfl_w_limit_wp * dtime, - ), - wpfloat("0.0"), - ) - - ddt_w_adv_wp = where( - cfl_clipping & owner_mask, - ddt_w_adv_wp + difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=dims.C2E2CODim), - ddt_w_adv_wp, - ) - - return astype(ddt_w_adv_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def add_extra_diffusion_for_w_con_approaching_cfl( - cfl_clipping: fa.CellKHalfField[bool], - owner_mask: fa.CellField[bool], - z_w_con_c: fa.CellKHalfField[ta.vpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - w: fa.CellKHalfField[ta.wpfloat], - ddt_w_adv: fa.CellKHalfField[ta.vpfloat], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _add_extra_diffusion_for_w_con_approaching_cfl( - cfl_clipping=cfl_clipping, - owner_mask=owner_mask, - z_w_con_c=z_w_con_c, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - w=w, - ddt_w_adv=ddt_w_adv, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - out=ddt_w_adv, - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py deleted file mode 100644 index 50507f4334..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/add_interpolated_horizontal_advection_of_w.py +++ /dev/null @@ -1,47 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import astype, neighbor_sum - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.dimension import C2E -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _add_interpolated_horizontal_advection_of_w( - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], - z_v_grad_w: fa.EdgeKHalfField[vpfloat], - ddt_w_adv: fa.CellKHalfField[vpfloat], -) -> fa.CellKHalfField[vpfloat]: - """Formerly known as _mo_velocity_advection_stencil_17.""" - z_v_grad_w_wp, ddt_w_adv_wp = astype((z_v_grad_w, ddt_w_adv), wpfloat) - ddt_w_adv_wp = ddt_w_adv_wp + neighbor_sum(z_v_grad_w_wp(C2E) * e_bln_c_s, axis=dims.C2EDim) - return astype(ddt_w_adv_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def add_interpolated_horizontal_advection_of_w( - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], wpfloat], - z_v_grad_w: fa.EdgeKHalfField[vpfloat], - ddt_w_adv: fa.CellKHalfField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _add_interpolated_horizontal_advection_of_w( - e_bln_c_s=e_bln_c_s, - z_v_grad_w=z_v_grad_w, - ddt_w_adv=ddt_w_adv, - out=ddt_w_adv, - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_horizontal_momentum_equation.py deleted file mode 100644 index 621d49500f..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_horizontal_momentum_equation.py +++ /dev/null @@ -1,336 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import ( - abs, # noqa: A004 - astype, - maximum, - minimum, - neighbor_sum, - where, -) -from gt4py.next.experimental import concat_where - -from icon4py.model.atmosphere.dycore.stencils.mo_math_divrot_rot_vertex_ri_dsl import ( - _mo_math_divrot_rot_vertex_ri_dsl, -) -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.dimension import E2C, E2C2EO, E2V -from icon4py.model.common.interpolation.stencils.interpolate_to_cell_center_vp import ( - _interpolate_to_cell_center_vp, -) -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _compute_advective_normal_wind_tendency( - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - upward_vorticity_at_vertices_on_model_levels: fa.VertexKField[ta.vpfloat], - tangential_wind: fa.EdgeKField[ta.vpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - coriolis_frequency: fa.EdgeField[ta.wpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], - coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], - ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], -) -> fa.EdgeKField[ta.vpfloat]: - #: intermediate variable horizontal_kinetic_energy_at_cells_on_model_levels is originally declared as z_ekinh in ICON - horizontal_kinetic_energy_at_cells_on_model_levels = _interpolate_to_cell_center_vp( - horizontal_kinetic_energy_at_edges_on_model_levels, e_bln_c_s - ) - horizontal_kinetic_energy_at_cells_on_model_levels = astype( - horizontal_kinetic_energy_at_cells_on_model_levels, vpfloat - ) - - ( - contravariant_corrected_w_at_cells_on_model_levels_wp, - ddqz_z_full_e_wp, - tangential_wind_wp, - ) = astype( - (contravariant_corrected_w_at_cells_on_model_levels, ddqz_z_full_e, tangential_wind), - wpfloat, - ) - - horizontal_advection = ( - horizontal_kinetic_energy_at_edges_on_model_levels - * (coeff_gradekin[dims.E2CDim(0)] - coeff_gradekin[dims.E2CDim(1)]) - + coeff_gradekin[dims.E2CDim(1)] - * horizontal_kinetic_energy_at_cells_on_model_levels(E2C[1]) - - coeff_gradekin[dims.E2CDim(0)] - * horizontal_kinetic_energy_at_cells_on_model_levels(E2C[0]) - ) - - vertical_advection = ( - neighbor_sum( - c_lin_e * contravariant_corrected_w_at_cells_on_model_levels_wp(E2C), axis=dims.E2CDim - ) - * astype((vn_on_half_levels(dims.KDim - 0.5) - vn_on_half_levels(dims.KDim + 0.5)), wpfloat) - / ddqz_z_full_e_wp - ) - - coriolis_term = tangential_wind_wp * ( - coriolis_frequency - + astype( - vpfloat("0.5") - * neighbor_sum(upward_vorticity_at_vertices_on_model_levels(E2V), axis=dims.E2VDim), - wpfloat, - ) - ) - normal_wind_advective_tendency_wp = -(horizontal_advection + vertical_advection + coriolis_term) - - return astype(normal_wind_advective_tendency_wp, vpfloat) - - -@gtx.field_operator -def _compute_extra_diffusion( - vn: fa.EdgeKField[ta.wpfloat], - upward_vorticity_at_vertices_on_model_levels: fa.VertexKField[ta.vpfloat], - difcoef: fa.EdgeKField[ta.wpfloat], - area_edge: fa.EdgeField[ta.wpfloat], - geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], -) -> fa.EdgeKField[ta.wpfloat]: - gradient_of_divergence_of_vn = neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=dims.E2C2EODim) - - gradient_of_vorticity = ( - tangent_orientation - * inv_primal_edge_length - * astype( - upward_vorticity_at_vertices_on_model_levels(E2V[1]) - - upward_vorticity_at_vertices_on_model_levels(E2V[0]), - wpfloat, - ) - ) - - extra_diffusion_on_vn = ( - difcoef * area_edge * (gradient_of_divergence_of_vn + gradient_of_vorticity) - ) - - return extra_diffusion_on_vn - - -@gtx.field_operator -def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask( - c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], - contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], - area_edge: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - upward_vorticity_at_vertices_on_model_levels: fa.VertexKField[ta.vpfloat], - geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], - vn: fa.EdgeKField[ta.wpfloat], - normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat], - cfl_w_limit: ta.vpfloat, - scalfac_exdiff: ta.wpfloat, - dtime: ta.wpfloat, -) -> fa.EdgeKField[ta.vpfloat]: - ( - contravariant_corrected_w_at_cells_on_model_levels_wp, - ddqz_z_full_e_wp, - normal_wind_advective_tendency_wp, - cfl_w_limit_wp, - ) = astype( - ( - contravariant_corrected_w_at_cells_on_model_levels, - ddqz_z_full_e, - normal_wind_advective_tendency, - cfl_w_limit, - ), - wpfloat, - ) - - #: intermediate variable contravariant_corrected_w_at_edges_on_model_levels is originally declared as w_con_e in ICON - contravariant_corrected_w_at_edges_on_model_levels = neighbor_sum( - c_lin_e * contravariant_corrected_w_at_cells_on_model_levels_wp(E2C), axis=dims.E2CDim - ) - difcoef = scalfac_exdiff * minimum( - wpfloat("0.85") - cfl_w_limit_wp * dtime, - abs(contravariant_corrected_w_at_edges_on_model_levels) * dtime / ddqz_z_full_e_wp - - cfl_w_limit_wp * dtime, - ) - normal_wind_advective_tendency_wp = where( - abs(contravariant_corrected_w_at_edges_on_model_levels) - > astype(cfl_w_limit * ddqz_z_full_e, wpfloat), - normal_wind_advective_tendency_wp - + _compute_extra_diffusion( - vn=vn, - upward_vorticity_at_vertices_on_model_levels=upward_vorticity_at_vertices_on_model_levels, - difcoef=difcoef, - area_edge=area_edge, - geofac_grdiv=geofac_grdiv, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inv_primal_edge_length, - ), - normal_wind_advective_tendency_wp, - ) - return astype(normal_wind_advective_tendency_wp, vpfloat) - - -@gtx.field_operator -def _compute_advection_in_horizontal_momentum( - vn: fa.EdgeKField[ta.wpfloat], - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - tangential_wind: fa.EdgeKField[ta.vpfloat], - coriolis_frequency: fa.EdgeField[ta.wpfloat], - contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], - coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], - c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], - ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], - area_edge: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], - cfl_w_limit: ta.vpfloat, - scalfac_exdiff: ta.wpfloat, - dtime: ta.wpfloat, - apply_extra_diffusion_on_vn: bool, - nlev: gtx.int32, - end_index_of_damping_layer: gtx.int32, -) -> fa.EdgeKField[ta.vpfloat]: - upward_vorticity_at_vertices_on_model_levels = _mo_math_divrot_rot_vertex_ri_dsl(vn, geofac_rot) - upward_vorticity_at_vertices_on_model_levels = astype( - upward_vorticity_at_vertices_on_model_levels, vpfloat - ) - - normal_wind_advective_tendency = _compute_advective_normal_wind_tendency( - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - upward_vorticity_at_vertices_on_model_levels=upward_vorticity_at_vertices_on_model_levels, - tangential_wind=tangential_wind, - vn_on_half_levels=vn_on_half_levels, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - coriolis_frequency=coriolis_frequency, - e_bln_c_s=e_bln_c_s, - c_lin_e=c_lin_e, - coeff_gradekin=coeff_gradekin, - ddqz_z_full_e=ddqz_z_full_e, - ) - - if apply_extra_diffusion_on_vn: - normal_wind_advective_tendency = concat_where( - ((maximum(2, end_index_of_damping_layer - 2)) <= dims.KDim) & (dims.KDim < (nlev - 4)), - _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask( - c_lin_e=c_lin_e, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - ddqz_z_full_e=ddqz_z_full_e, - area_edge=area_edge, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inv_primal_edge_length, - upward_vorticity_at_vertices_on_model_levels=upward_vorticity_at_vertices_on_model_levels, - geofac_grdiv=geofac_grdiv, - vn=vn, - normal_wind_advective_tendency=normal_wind_advective_tendency, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - ), - normal_wind_advective_tendency, - ) - - return normal_wind_advective_tendency - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_advection_in_horizontal_momentum( - normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - tangential_wind: fa.EdgeKField[ta.vpfloat], - coriolis_frequency: fa.EdgeField[ta.wpfloat], - contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], - coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], - c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], - ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], - area_edge: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], - cfl_w_limit: ta.vpfloat, - scalfac_exdiff: ta.wpfloat, - dtime: ta.wpfloat, - apply_extra_diffusion_on_vn: bool, - end_index_of_damping_layer: gtx.int32, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - """ - Formerly known as fused_velocity_advection_stencil_19_to_20. - - This computes the horizontal advection in the horizontal momentum equation - - Args: - - normal_wind_advective_tendency: horizontal advection tendency of the normal wind - - vn: normal wind at edges - - horizontal_kinetic_energy_at_edges_on_model_levels: horizontal kinetic energy at edges on model levels - - horizontal_kinetic_energy_at_cells_on_model_levels: horizontal kinetic energy at cell centers on model levels - - tangential_wind: tangential wind at model levels - - coriolis_frequency: coriolis frequency parameter - - contravariant_corrected_w_at_cells_on_model_levels: contravariant-corrected vertical velocity at model levels - - vn_on_half_levels: normal wind on half levels - - e_bln_c_s: interpolation field (edge-to-cell interpolation weights) - - geofac_rot: metric field for rotor computation - - coeff_gradekin: metrics field/coefficient for the gradient of kinematic energy - - c_lin_e: metrics field for linear interpolation from cells to edges - - ddqz_z_full_e: metrics field equal to vertical spacing - - area_edge: area associated with each edge - - tangent_orientation: orientation of the edge with respect to the grid - - inv_primal_edge_length: inverse primal edge length - - geofac_grdiv: metrics field used to compute the gradient of a divergence (of vn) - - cfl_w_limit: CFL limit for vertical velocity - - scalfac_exdiff: scalar factor for external diffusion - - dtime: time step - - apply_extra_diffusion_on_vn: option to apply extra diffusion to vn - - end_index_of_damping_layer: vertical index where damping ends - - horizontal_start: start index in the horizontal domain - - horizontal_end: end index in the horizontal domain - - vertical_start: start index in the vertical domain at model top - - vertical_end: end index in the vertical domain at model bottom (or number of full/model vertical levels) - - Returns: - - normal_wind_advective_tendency: horizontal advection tendency of the normal wind - - """ - - _compute_advection_in_horizontal_momentum( - vn=vn, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - tangential_wind=tangential_wind, - coriolis_frequency=coriolis_frequency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vn_on_half_levels=vn_on_half_levels, - e_bln_c_s=e_bln_c_s, - geofac_rot=geofac_rot, - coeff_gradekin=coeff_gradekin, - c_lin_e=c_lin_e, - ddqz_z_full_e=ddqz_z_full_e, - area_edge=area_edge, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inv_primal_edge_length, - geofac_grdiv=geofac_grdiv, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, - nlev=vertical_end, - end_index_of_damping_layer=end_index_of_damping_layer, - out=normal_wind_advective_tendency, - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py deleted file mode 100644 index 4578777a83..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_advection_in_vertical_momentum_equation.py +++ /dev/null @@ -1,655 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import abs, astype, broadcast, maximum, where # noqa: A004 -from gt4py.next.experimental import concat_where - -from icon4py.model.atmosphere.dycore.stencils.add_extra_diffusion_for_w_con_approaching_cfl import ( - _add_extra_diffusion_for_w_con_approaching_cfl, -) -from icon4py.model.atmosphere.dycore.stencils.add_interpolated_horizontal_advection_of_w import ( - _add_interpolated_horizontal_advection_of_w, -) -from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_advection_term_for_vertical_velocity import ( - _compute_horizontal_advection_term_for_vertical_velocity, -) -from icon4py.model.atmosphere.dycore.stencils.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, -) -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_half_levels import ( - _interpolate_cell_field_to_half_levels_vp, -) -from icon4py.model.common.interpolation.stencils.interpolate_to_cell_center_vp import ( - _interpolate_to_cell_center_vp, -) -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _interpolate_contravariant_vertical_velocity_to_full_levels( - contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[vpfloat], - nlev: gtx.int32, -) -> fa.CellKField[vpfloat]: - # TODO(havogt): Note that `concat_where(dims.KDim == nlev-1, ...)` is currently broken - # because of insufficiency in the domain inference of GT4Py, - # see https://github.com/GridTools/gt4py/issues/2205. - return concat_where( - dims.KDim < nlev - 1, - vpfloat("0.5") - * ( - contravariant_corrected_w_at_cells_on_half_levels(dims.KDim - 0.5) - + contravariant_corrected_w_at_cells_on_half_levels(dims.KDim + 0.5) - ), - vpfloat("0.5") * contravariant_corrected_w_at_cells_on_half_levels(dims.KDim - 0.5), - ) - - -@gtx.field_operator -def _compute_horizontal_advection_of_w( - w: fa.CellKHalfField[ta.wpfloat], - tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], - inv_dual_edge_length: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], -) -> fa.EdgeKHalfField[ta.vpfloat]: - w_at_vertices = _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(w, c_intp) - - horizontal_advection_of_w_at_edges_on_half_levels = ( - _compute_horizontal_advection_term_for_vertical_velocity( - vn_on_half_levels, - inv_dual_edge_length, - w, - tangential_wind_on_half_levels, - inv_primal_edge_length, - tangent_orientation, - w_at_vertices, - ) - ) - - return astype(horizontal_advection_of_w_at_edges_on_half_levels, vpfloat) - - -@gtx.field_operator -def _add_vertical_advection_of_w_to_advective_vertical_wind_tendency( - contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[vpfloat], - w: fa.CellKHalfField[wpfloat], - coeff1_dwdz: fa.CellKField[vpfloat], - coeff2_dwdz: fa.CellKField[vpfloat], -) -> fa.CellKHalfField[vpfloat]: - contravariant_corrected_w_at_cells_on_half_levels_wp = astype( - contravariant_corrected_w_at_cells_on_half_levels, wpfloat - ) - coeff1_dwdz_at_half_levels = coeff1_dwdz(dims.KHalfDim + 0.5) - coeff2_dwdz_at_half_levels = coeff2_dwdz(dims.KHalfDim + 0.5) - coeff1_dwdz_wp, coeff2_dwdz_wp = astype( - (coeff1_dwdz_at_half_levels, coeff2_dwdz_at_half_levels), wpfloat - ) - - vertical_wind_advective_tendency_wp = -contravariant_corrected_w_at_cells_on_half_levels_wp * ( - w(dims.KHalfDim - 1) * coeff1_dwdz_wp - - w(dims.KHalfDim + 1) * coeff2_dwdz_wp - + w * astype(coeff2_dwdz_at_half_levels - coeff1_dwdz_at_half_levels, wpfloat) - ) - return astype(vertical_wind_advective_tendency_wp, vpfloat) - - -@gtx.field_operator -def _compute_maximum_cfl_and_clip_contravariant_vertical_velocity( - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, -) -> tuple[ - fa.CellKHalfField[ta.vpfloat], - fa.CellKHalfField[bool], - fa.CellKHalfField[ta.vpfloat], -]: - contravariant_corrected_w_at_cells_on_half_levels_wp, ddqz_z_half_wp = astype( - (contravariant_corrected_w_at_cells_on_half_levels, ddqz_z_half), wpfloat - ) - - cfl_clipping = where( - abs(contravariant_corrected_w_at_cells_on_half_levels) > cfl_w_limit * ddqz_z_half, - broadcast(True, (dims.CellDim, dims.KHalfDim)), - False, - ) - - vertical_cfl = where( - cfl_clipping, - contravariant_corrected_w_at_cells_on_half_levels_wp * dtime / ddqz_z_half_wp, - broadcast(wpfloat("0.0"), (dims.CellDim, dims.KHalfDim)), - ) - vertical_cfl_vp = astype(vertical_cfl, vpfloat) - - contravariant_corrected_w_at_cells_on_half_levels_wp = where( - (cfl_clipping) & (vertical_cfl_vp < -vpfloat("0.85")), - astype(-vpfloat("0.85") * ddqz_z_half, wpfloat) / dtime, - contravariant_corrected_w_at_cells_on_half_levels_wp, - ) - - contravariant_corrected_w_at_cells_on_half_levels_wp = where( - (cfl_clipping) & (vertical_cfl_vp > vpfloat("0.85")), - astype(vpfloat("0.85") * ddqz_z_half, wpfloat) / dtime, - contravariant_corrected_w_at_cells_on_half_levels_wp, - ) - - return ( - astype(contravariant_corrected_w_at_cells_on_half_levels_wp, vpfloat), - cfl_clipping, - vertical_cfl_vp, - ) - - -@gtx.field_operator -def _compute_contravariant_corrected_w( - w: fa.CellKHalfField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], -) -> fa.CellKHalfField[ta.vpfloat]: - contravariant_corrected_w_at_cells_on_half_levels = ( - astype(w, vpfloat) - contravariant_correction_at_cells_on_half_levels - ) - - return contravariant_corrected_w_at_cells_on_half_levels - - -@gtx.field_operator -def _compute_contravariant_corrected_w_and_cfl( - w: fa.CellKHalfField[ta.wpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, - nlev: gtx.int32, - end_index_of_damping_layer: gtx.int32, -) -> tuple[fa.CellKHalfField[ta.vpfloat], fa.CellKHalfField[bool], fa.CellKHalfField[ta.vpfloat]]: - #: intermediate variable contravariant_corrected_w_at_cells_on_half_levels is originally declared as z_w_con_c in ICON - contravariant_corrected_w_at_cells_on_half_levels = _compute_contravariant_corrected_w( - w, contravariant_correction_at_cells_on_half_levels - ) - - (contravariant_corrected_w_at_cells_on_half_levels, cfl_clipping, vertical_cfl) = concat_where( - (dims.KHalfDim >= maximum(2, end_index_of_damping_layer - 2)) & (dims.KHalfDim < nlev - 3), - _compute_maximum_cfl_and_clip_contravariant_vertical_velocity( - ddqz_z_half=ddqz_z_half, - contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - ), - ( - contravariant_corrected_w_at_cells_on_half_levels, - broadcast(False, (dims.CellDim, dims.KHalfDim)), - broadcast(vpfloat("0.0"), (dims.CellDim, dims.KHalfDim)), - ), - ) - - return contravariant_corrected_w_at_cells_on_half_levels, cfl_clipping, vertical_cfl - - -@gtx.field_operator -def _compute_advective_vertical_wind_tendency( - vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], - w: fa.CellKHalfField[ta.wpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[ta.wpfloat], - cfl_clipping: fa.CellKHalfField[bool], - coeff1_dwdz: fa.CellKField[ta.vpfloat], - coeff2_dwdz: fa.CellKField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - owner_mask: fa.CellField[bool], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, -) -> fa.CellKHalfField[ta.vpfloat]: - vertical_wind_advective_tendency = concat_where( - 1 <= dims.KHalfDim, - _add_vertical_advection_of_w_to_advective_vertical_wind_tendency( - contravariant_corrected_w_at_cells_on_half_levels, w, coeff1_dwdz, coeff2_dwdz - ), - vertical_wind_advective_tendency, - ) - - vertical_wind_advective_tendency = concat_where( - 1 <= dims.KHalfDim, - _add_interpolated_horizontal_advection_of_w( - e_bln_c_s, - horizontal_advection_of_w_at_edges_on_half_levels, - vertical_wind_advective_tendency, - ), - vertical_wind_advective_tendency, - ) - - vertical_wind_advective_tendency = _add_extra_diffusion_for_w_con_approaching_cfl( - cfl_clipping, - owner_mask, - contravariant_corrected_w_at_cells_on_half_levels, - ddqz_z_half, - area, - geofac_n2s, - w, - vertical_wind_advective_tendency, - scalfac_exdiff, - cfl_w_limit, - dtime, - ) - - return vertical_wind_advective_tendency - - -@gtx.field_operator -def _compute_advection_in_corrector_vertical_momentum( - vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], - w: fa.CellKHalfField[ta.wpfloat], - tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], - coeff1_dwdz: fa.CellKField[ta.vpfloat], - coeff2_dwdz: fa.CellKField[ta.vpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], - inv_dual_edge_length: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - owner_mask: fa.CellField[bool], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, - nlev: gtx.int32, - end_index_of_damping_layer: gtx.int32, -) -> tuple[fa.CellKHalfField[ta.vpfloat], fa.CellKField[ta.vpfloat], fa.CellKHalfField[ta.vpfloat]]: - #: intermediate variable horizontal_advection_of_w_at_edges_on_half_levels is originally declared as z_v_grad_w in ICON - horizontal_advection_of_w_at_edges_on_half_levels = _compute_horizontal_advection_of_w( - w=w, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - ) - - ( - contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping, - vertical_cfl, - ) = _compute_contravariant_corrected_w_and_cfl( - w=w, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - ddqz_z_half=ddqz_z_half, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=nlev, - end_index_of_damping_layer=end_index_of_damping_layer, - ) - - vertical_wind_advective_tendency = _compute_advective_vertical_wind_tendency( - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping=cfl_clipping, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - ) - - contravariant_corrected_w_at_cells_on_model_levels = ( - _interpolate_contravariant_vertical_velocity_to_full_levels( - contravariant_corrected_w_at_cells_on_half_levels, nlev - ) - ) - - return ( - vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl, - ) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_advection_in_corrector_vertical_momentum( - vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], - contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - vertical_cfl: fa.CellKHalfField[ta.vpfloat], - w: fa.CellKHalfField[ta.wpfloat], - tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], - coeff1_dwdz: fa.CellKField[ta.vpfloat], - coeff2_dwdz: fa.CellKField[ta.vpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], - inv_dual_edge_length: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - owner_mask: fa.CellField[bool], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, - end_index_of_damping_layer: gtx.int32, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - """ - This computes the vertical momentum advection in the vertical momentum equation - - Args: - - vertical_wind_advective_tendency: vertical advection tendency of the vertical wind - - contravariant_corrected_w_at_cells_on_model_levels: contravariant-corrected vertical velocity at cells on model levels - - vertical_cfl: vertical cfl number at cells on half levels - - w: vertical wind at cell centers - - tangential_wind: tangential wind at edges on model levels - - vn_on_half_levels: normal wind at edges on half levels - - contravariant_correction_at_edges_on_model_levels: contravariant correction at edges on model levels - - coeff1_dwdz: metrics field (first coefficient for vertical derivative of vertical wind) - - coeff2_dwdz: metrics field (second coefficient for vertical derivative of vertical wind) - - c_intp: interpolation field for cell-to-vertex interpolation - - inv_dual_edge_length: inverse dual edge length - - inv_primal_edge_length: inverse primal edge length - - tangent_orientation: orientation of the edge with respect to the grid - - e_bln_c_s: interpolation field (edge-to-cell interpolation weights) - - ddqz_z_half: metrics field - - area: cell area - - geofac_n2s: interpolation field - - owner_mask: ownership mask for each cell - - scalfac_exdiff: scalar factor for external diffusion - - cfl_w_limit: CFL limit for vertical velocity - - dtime: time step - - nflatlev: number of flat levels - - end_index_of_damping_layer: vertical index where damping ends - - horizontal_start: start index in the horizontal dimension - - horizontal_end: end index in the horizontal dimension - - vertical_start: start index in the vertical dimension at model top - - vertical_end: end index in the vertical dimension at model bottom (or number of full/model vertical levels) - - Returns: - - vertical_wind_advective_tendency - - contravariant_corrected_w_at_cells_on_model_levels - - vertical_cfl - """ - - _compute_advection_in_corrector_vertical_momentum( - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - w=w, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=vertical_end, - end_index_of_damping_layer=end_index_of_damping_layer, - out=( - vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl, - ), - domain=( - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ), - ) - - -@gtx.field_operator -def _interpolate_contravariant_correction_to_cells_on_half_levels( - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - wgtfac_c: fa.CellKHalfField[ta.vpfloat], - nflatlev: gtx.int32, -) -> fa.CellKHalfField[ta.vpfloat]: - contravariant_correction_at_cells_model_levels = _interpolate_to_cell_center_vp( - contravariant_correction_at_edges_on_model_levels, e_bln_c_s - ) - contravariant_correction_at_cells_model_levels = astype( - contravariant_correction_at_cells_model_levels, vpfloat - ) - - contravariant_correction_at_cells_on_half_levels = concat_where( - dims.KHalfDim >= nflatlev + 1, - _interpolate_cell_field_to_half_levels_vp( - wgtfac_c=wgtfac_c, interpolant=contravariant_correction_at_cells_model_levels - ), - broadcast(vpfloat("0.0"), (dims.CellDim, dims.KHalfDim)), - ) - - return contravariant_correction_at_cells_on_half_levels - - -@gtx.field_operator -def _compute_advection_in_predictor_vertical_momentum( - vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], - w: fa.CellKHalfField[ta.wpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - coeff1_dwdz: fa.CellKField[ta.vpfloat], - coeff2_dwdz: fa.CellKField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - wgtfac_c: fa.CellKHalfField[ta.vpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - owner_mask: fa.CellField[bool], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, - skip_compute_predictor_vertical_advection: bool, - nflatlev: gtx.int32, - nlev: gtx.int32, - end_index_of_damping_layer: gtx.int32, -) -> tuple[ - fa.CellKHalfField[ta.vpfloat], - fa.CellKHalfField[ta.vpfloat], - fa.CellKField[ta.vpfloat], - fa.CellKHalfField[ta.vpfloat], -]: - contravariant_correction_at_cells_on_half_levels = _interpolate_contravariant_correction_to_cells_on_half_levels( - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - nflatlev=nflatlev, - ) - - ( - contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping, - vertical_cfl, - ) = _compute_contravariant_corrected_w_and_cfl( - w=w, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - ddqz_z_half=ddqz_z_half, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=nlev, - end_index_of_damping_layer=end_index_of_damping_layer, - ) - - if not skip_compute_predictor_vertical_advection: - vertical_wind_advective_tendency = _compute_advective_vertical_wind_tendency( - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping=cfl_clipping, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - ) - - contravariant_corrected_w_at_cells_on_model_levels = ( - _interpolate_contravariant_vertical_velocity_to_full_levels( - contravariant_corrected_w_at_cells_on_half_levels, nlev - ) - ) - - return ( - contravariant_correction_at_cells_on_half_levels, - vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl, - ) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_advection_in_predictor_vertical_momentum( - contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], - vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], - contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], - vertical_cfl: fa.CellKHalfField[ta.vpfloat], - w: fa.CellKHalfField[ta.wpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - coeff1_dwdz: fa.CellKField[ta.vpfloat], - coeff2_dwdz: fa.CellKField[ta.vpfloat], - e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], - wgtfac_c: fa.CellKHalfField[ta.vpfloat], - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - area: fa.CellField[ta.wpfloat], - geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], - owner_mask: fa.CellField[bool], - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, - skip_compute_predictor_vertical_advection: bool, - nflatlev: gtx.int32, - end_index_of_damping_layer: gtx.int32, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - """ - This computes the vertical momentum advection in the vertical momentum equation - - Args: - - contravariant_correction_at_cells_on_half_levels: contravariant correction at cells on model levels - - vertical_wind_advective_tendency: vertical advection tendency of the vertical wind - - contravariant_corrected_w_at_cells_on_model_levels: contravariant-corrected vertical velocity at cells on model levels - - vertical_cfl: vertical cfl number at cells on half levels - - w: vertical wind at cells on half levels - - horizontal_advection_of_w_at_edges_on_half_levels: horizontal advection of w at edges on half levels - - contravariant_correction_at_edges_on_model_levels: contravariant correction at edges on model levels - - coeff1_dwdz: metrics field (first coefficient for vertical derivative of vertical wind) - - coeff2_dwdz: metrics field (second coefficient for vertical derivative of vertical wind) - - e_bln_c_s: interpolation field (edge-to-cell interpolation weights) - - wgtfac_c: metric coefficient for interpolating a cell variable from full to half levels - - ddqz_z_half: metrics field - - area: cell area - - geofac_n2s: interpolation field - - owner_mask: ownership mask for each cell - - scalfac_exdiff: scalar factor for external diffusion - - cfl_w_limit: CFL limit for vertical velocity - - dtime: time step - - skip_compute_predictor_vertical_advection: logical flag to skip the vertical advection - - nflatlev: number of flat levels - - end_index_of_damping_layer: vertical index where damping ends - - horizontal_start: start index in the horizontal dimension - - horizontal_end: end index in the horizontal dimension - - vertical_start: start index in the vertical dimension at model top - - vertical_end: end index in the vertical dimension at model bottom (or number of full/model vertical levels) - - Returns: - - contravariant_correction_at_cells_on_half_levels - - vertical_wind_advective_tendency - - contravariant_corrected_w_at_cells_on_model_levels - - vertical_cfl - """ - - _compute_advection_in_predictor_vertical_momentum( - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=nflatlev, - nlev=vertical_end, - end_index_of_damping_layer=end_index_of_damping_layer, - out=( - contravariant_correction_at_cells_on_half_levels, - vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl, - ), - domain=( - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - { - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ), - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_avg_vn_and_graddiv_vn_and_vt.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_avg_vn_and_graddiv_vn_and_vt.py index 551cb6aef5..bd56bf924f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_avg_vn_and_graddiv_vn_and_vt.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_avg_vn_and_graddiv_vn_and_vt.py @@ -14,7 +14,7 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C2EO from icon4py.model.common.interpolation.stencils.compute_tangential_wind import ( - _compute_tangential_wind_vp, + _compute_tangential_wind, ) from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -33,7 +33,7 @@ def _compute_avg_vn_and_graddiv_vn_and_vt( """Formerly known as _mo_solve_nonhydro_stencil_30.""" z_vn_avg_wp = _spatially_average_flux_or_velocity(e_flx_avg=e_flx_avg, flux_or_velocity=vn) z_graddiv_vn_vp = astype(neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=dims.E2C2EODim), vpfloat) - vt_vp = _compute_tangential_wind_vp(vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e) + vt_vp = astype(_compute_tangential_wind(vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e), vpfloat) return z_vn_avg_wp, z_graddiv_vn_vp, vt_vp diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_diagnostics_from_normal_wind.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_diagnostics_from_normal_wind.py index 87cf9ed44f..a0236a3f64 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_diagnostics_from_normal_wind.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_diagnostics_from_normal_wind.py @@ -9,24 +9,11 @@ from gt4py.next import astype from gt4py.next.experimental import concat_where -from icon4py.model.atmosphere.dycore.stencils.compute_contravariant_correction import ( - _compute_contravariant_correction, -) -from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_advection_term_for_vertical_velocity import ( - _compute_horizontal_advection_term_for_vertical_velocity, -) -from icon4py.model.atmosphere.dycore.stencils.extrapolate_at_top import _extrapolate_at_top -from icon4py.model.atmosphere.dycore.stencils.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, -) from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.interpolation.stencils.compute_tangential_wind import ( - _compute_tangential_wind_vp, -) from icon4py.model.common.type_alias import vpfloat, wpfloat -@gtx.field_operator +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) def _interpolate_to_half_levels( wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], x: fa.EdgeKField[ta.wpfloat], @@ -47,201 +34,3 @@ def _compute_horizontal_kinetic_energy( ) -> fa.EdgeKField[ta.vpfloat]: z_kin_hor_e_wp = wpfloat("0.5") * (vn * vn + astype(vt * vt, wpfloat)) return astype(z_kin_hor_e_wp, vpfloat) - - -@gtx.field_operator -def _compute_diagnostics_from_normal_wind( - tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - w: fa.CellKHalfField[ta.wpfloat], - rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], ta.wpfloat], - wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], - ddxn_z_full: fa.EdgeKField[ta.vpfloat], - ddxt_z_full: fa.EdgeKField[ta.vpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], - inv_dual_edge_length: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - skip_compute_predictor_vertical_advection: bool, - nflatlev: gtx.int32, -) -> tuple[ - fa.EdgeKField[ta.vpfloat], - fa.EdgeKHalfField[ta.vpfloat], - fa.EdgeKHalfField[ta.vpfloat], - fa.EdgeKField[ta.vpfloat], - fa.EdgeKField[ta.vpfloat], - fa.EdgeKHalfField[ta.vpfloat], -]: - tangential_wind = _compute_tangential_wind_vp(vn, rbf_vec_coeff_e) - horizontal_kinetic_energy_at_edges_on_model_levels = _compute_horizontal_kinetic_energy( - vn, tangential_wind - ) - vn_on_half_levels = _interpolate_to_half_levels(wgtfac_e, vn) - - tangential_wind_on_half_levels = ( - _interpolate_to_half_levels(wgtfac_e, tangential_wind) - if not skip_compute_predictor_vertical_advection - else tangential_wind_on_half_levels - ) - - contravariant_correction_at_edges_on_model_levels = concat_where( - nflatlev <= dims.KDim, - _compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, tangential_wind), - contravariant_correction_at_edges_on_model_levels, - ) - - w_at_vertices = _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(w, c_intp) - horizontal_advection_of_w_at_edges_on_half_levels = ( - _compute_horizontal_advection_term_for_vertical_velocity( - vn_on_half_levels, - inv_dual_edge_length, - w, - tangential_wind_on_half_levels, - inv_primal_edge_length, - tangent_orientation, - w_at_vertices, - ) - if not skip_compute_predictor_vertical_advection - else horizontal_advection_of_w_at_edges_on_half_levels - ) - - return ( - tangential_wind, - tangential_wind_on_half_levels, - vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels, - ) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_diagnostics_from_normal_wind( - tangential_wind: fa.EdgeKField[ta.vpfloat], - tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], - vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], - horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - w: fa.CellKHalfField[ta.wpfloat], - rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], ta.wpfloat], - wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], - ddxn_z_full: fa.EdgeKField[ta.vpfloat], - ddxt_z_full: fa.EdgeKField[ta.vpfloat], - wgtfacq_e: fa.EdgeKField[ta.vpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], - inv_dual_edge_length: fa.EdgeField[ta.wpfloat], - inv_primal_edge_length: fa.EdgeField[ta.wpfloat], - tangent_orientation: fa.EdgeField[ta.wpfloat], - skip_compute_predictor_vertical_advection: bool, - nflatlev: gtx.int32, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - """ - Formerly known as fused_velocity_advection_stencil_1_to_7_predictor. - - This computes the derived horizontal wind components, kinetic energy, horizontal advection of the vertical velocity, - and the contravariant correction. - It also extrapolates vertical velocity at the top level. - - Agrs: - - tangential_wind: tangential wind at model levels - - tangential_wind_on_half_levels: tangential wind interpolated to half levels - - vn_on_half_levels: normal wind interpolated to half levels - - horizontal_kinetic_energy_at_edges_on_model_levels: horizontal kinetic energy computed at edge of model levels - - contravariant_correction_at_edges_on_model_levels: contravariant metric correction at edge of model levels - - horizontal_advection_of_w_at_edges_on_half_levels: horizontal advection for vertical velocity - - vn: normal wind at edges - - w: vertical wind at cell centers - - rbf_vec_coeff_e: interpolation field (RBF vector coefficient on edges) - - wgtfac_e: metrics field - - ddxn_z_full: metrics field (derivative of topography in the normal direction) - - ddxt_z_full: metrics field (derivative of topography in the tangential direction) - - wgtfacq_e: metrics field (weights for interpolation) - - c_intp: interpolation field - - inv_dual_edge_length: inverse dual edge length - - inv_primal_edge_length: inverse primal edge length - - tangent_orientation: orientation of the edge with respect to the grid - - skip_compute_predictor_vertical_advection: logical flag to skip the vertical advection - - nflatlev: number of flat levels - - horizontal_start: start index in the horizontal direction - - horizontal_end: end index in the horizontal direction - - vertical_start: start index in the vertical direction - - vertical_end: end index in the vertical direction - - - Returns: - - tangential_wind - - tangential_wind_on_half_levels - - vn_on_half_levels - - horizontal_kinetic_energy_at_edges_on_model_levels - - contravariant_correction_at_edges_on_model_levels - - horizontal_advection_of_w_at_edges_on_half_levels - """ - - _compute_diagnostics_from_normal_wind( - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - vn=vn, - w=w, - rbf_vec_coeff_e=rbf_vec_coeff_e, - wgtfac_e=wgtfac_e, - ddxn_z_full=ddxn_z_full, - ddxt_z_full=ddxt_z_full, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=nflatlev, - out=( - tangential_wind, - tangential_wind_on_half_levels, - vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels, - ), - domain=( - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end - 1), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end - 1), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end - 1), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end - 1), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end - 1), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end - 1), - }, - ), - ) - _extrapolate_at_top( - wgtfacq_e, - vn, - out=vn_on_half_levels, - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_end - 1, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_term_for_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_term_for_vertical_velocity.py deleted file mode 100644 index 6c42df0468..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_advection_term_for_vertical_velocity.py +++ /dev/null @@ -1,65 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import astype - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.dimension import E2C, E2V -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _compute_horizontal_advection_term_for_vertical_velocity( - vn_ie: fa.EdgeKHalfField[vpfloat], - inv_dual_edge_length: fa.EdgeField[wpfloat], - w: fa.CellKHalfField[wpfloat], - z_vt_ie: fa.EdgeKHalfField[vpfloat], - inv_primal_edge_length: fa.EdgeField[wpfloat], - tangent_orientation: fa.EdgeField[wpfloat], - z_w_v: fa.VertexKHalfField[vpfloat], -) -> fa.EdgeKHalfField[vpfloat]: - """Formerly know as _mo_velocity_advection_stencil_07.""" - z_vt_ie_wp, vn_ie_wp = astype((z_vt_ie, vn_ie), wpfloat) - - z_v_grad_w_wp = vn_ie_wp * inv_dual_edge_length * ( - w(E2C[0]) - w(E2C[1]) - ) + z_vt_ie_wp * inv_primal_edge_length * tangent_orientation * astype( - z_w_v(E2V[0]) - z_w_v(E2V[1]), wpfloat - ) - return astype(z_v_grad_w_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_horizontal_advection_term_for_vertical_velocity( - vn_ie: fa.EdgeKHalfField[vpfloat], - inv_dual_edge_length: fa.EdgeField[wpfloat], - w: fa.CellKHalfField[wpfloat], - z_vt_ie: fa.EdgeKHalfField[vpfloat], - inv_primal_edge_length: fa.EdgeField[wpfloat], - tangent_orientation: fa.EdgeField[wpfloat], - z_w_v: fa.VertexKHalfField[vpfloat], - z_v_grad_w: fa.EdgeKHalfField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _compute_horizontal_advection_term_for_vertical_velocity( - vn_ie=vn_ie, - inv_dual_edge_length=inv_dual_edge_length, - w=w, - z_vt_ie=z_vt_ie, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - z_w_v=z_w_v, - out=z_v_grad_w, - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py deleted file mode 100644 index 890f5ea1a2..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_kinetic_energy.py +++ /dev/null @@ -1,63 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import astype - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _compute_horizontal_kinetic_energy( - vn: fa.EdgeKField[wpfloat], - vt: fa.EdgeKField[vpfloat], -) -> tuple[ - fa.EdgeKHalfField[vpfloat], - fa.EdgeKHalfField[vpfloat], - fa.EdgeKField[vpfloat], -]: - """Formerly known as _mo_solve_nonhydro_stencil_37 or _mo_velocity_advection_stencil_05.""" - # TODO(): This stencil doesn't only do what the name implies. It also - # assigns to vn_ie_wp and z_vt_ie_vp. These things should be separated. - vn_ie_wp = vn(dims.KHalfDim + 0.5) - z_vt_ie_vp = vt(dims.KHalfDim + 0.5) - z_kin_hor_e_wp = wpfloat("0.5") * (vn * vn + astype(vt * vt, wpfloat)) - return astype(vn_ie_wp, vpfloat), z_vt_ie_vp, astype(z_kin_hor_e_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_horizontal_kinetic_energy( - vn: fa.EdgeKField[wpfloat], - vt: fa.EdgeKField[vpfloat], - vn_ie: fa.EdgeKHalfField[vpfloat], - z_vt_ie: fa.EdgeKHalfField[vpfloat], - z_kin_hor_e: fa.EdgeKField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _compute_horizontal_kinetic_energy( - vn=vn, - vt=vt, - out=(vn_ie, z_vt_ie, z_kin_hor_e), - domain=( - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ), - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_velocity_quantities.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_velocity_quantities.py index 94d94bf646..7fea4140c8 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_velocity_quantities.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_horizontal_velocity_quantities.py @@ -29,7 +29,7 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import E2C2EO from icon4py.model.common.interpolation.stencils.compute_tangential_wind import ( - _compute_tangential_wind_vp, + _compute_tangential_wind, ) from icon4py.model.common.type_alias import vpfloat @@ -65,7 +65,9 @@ def _compute_horizontal_velocity_quantities_and_fluxes( horizontal_gradient_of_normal_wind_divergence = astype( neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=dims.E2C2EODim), vpfloat ) - tangential_wind = _compute_tangential_wind_vp(vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e) + tangential_wind = astype( + _compute_tangential_wind(vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e), vpfloat + ) ( mass_flux_at_edges_on_model_levels, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py deleted file mode 100644 index 97d9ea7ef3..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py +++ /dev/null @@ -1,76 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import abs, astype, broadcast, where # noqa: A004 - -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _compute_maximum_cfl_and_clip_contravariant_vertical_velocity( - ddqz_z_half: fa.CellKHalfField[ta.vpfloat], - z_w_con_c: fa.CellKHalfField[ta.vpfloat], - cfl_w_limit: ta.vpfloat, - dtime: ta.wpfloat, -) -> tuple[ - fa.CellKHalfField[bool], - fa.CellKHalfField[ta.vpfloat], - fa.CellKHalfField[ta.vpfloat], -]: - """Formerly know as _mo_velocity_advection_stencil_14.""" - z_w_con_c_wp, ddqz_z_half_wp = astype((z_w_con_c, ddqz_z_half), wpfloat) - - cfl_clipping = where( - abs(z_w_con_c) > cfl_w_limit * ddqz_z_half, - broadcast(True, (dims.CellDim, dims.KHalfDim)), - False, - ) - - vcfl = where(cfl_clipping, z_w_con_c_wp * dtime / ddqz_z_half_wp, wpfloat("0.0")) - vcfl_vp = astype(vcfl, vpfloat) - - z_w_con_c_wp = where( - (cfl_clipping) & (vcfl_vp < -vpfloat("0.85")), - astype(-vpfloat("0.85") * ddqz_z_half, wpfloat) / dtime, - z_w_con_c_wp, - ) - - z_w_con_c_wp = where( - (cfl_clipping) & (vcfl_vp > vpfloat("0.85")), - astype(vpfloat("0.85") * ddqz_z_half, wpfloat) / dtime, - z_w_con_c_wp, - ) - - return cfl_clipping, vcfl_vp, astype(z_w_con_c_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_maximum_cfl_and_clip_contravariant_vertical_velocity( - ddqz_z_half: fa.CellKHalfField[vpfloat], - z_w_con_c: fa.CellKHalfField[vpfloat], - cfl_clipping: fa.CellKHalfField[bool], - vcfl: fa.CellKHalfField[vpfloat], - cfl_w_limit: vpfloat, - dtime: wpfloat, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _compute_maximum_cfl_and_clip_contravariant_vertical_velocity( - ddqz_z_half=ddqz_z_half, - z_w_con_c=z_w_con_c, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - out=(cfl_clipping, vcfl, z_w_con_c), - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_contravariant_vertical_velocity_to_full_levels.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_contravariant_vertical_velocity_to_full_levels.py deleted file mode 100644 index 1d9d326226..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_contravariant_vertical_velocity_to_full_levels.py +++ /dev/null @@ -1,40 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.type_alias import vpfloat - - -@gtx.field_operator -def _interpolate_contravariant_vertical_velocity_to_full_levels( - z_w_con_c: fa.CellKHalfField[vpfloat], -) -> fa.CellKField[vpfloat]: - """Formerly know as _mo_velocity_advection_stencil_15.""" - z_w_con_c_full_vp = vpfloat("0.5") * (z_w_con_c(dims.KDim - 0.5) + z_w_con_c(dims.KDim + 0.5)) - return z_w_con_c_full_vp - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def interpolate_contravariant_vertical_velocity_to_full_levels( - z_w_con_c: fa.CellKHalfField[vpfloat], - z_w_con_c_full: fa.CellKField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - # TODO(): this should use a generic stencil - _interpolate_contravariant_vertical_velocity_to_full_levels( - z_w_con_c=z_w_con_c, - out=z_w_con_c_full, - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py deleted file mode 100644 index 9a38577495..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py +++ /dev/null @@ -1,69 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx - -from icon4py.model.atmosphere.dycore.stencils.interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges import ( - _interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges, -) -from icon4py.model.atmosphere.dycore.stencils.interpolate_vt_to_interface_edges import ( - _interpolate_vt_to_interface_edges, -) -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta - - -@gtx.field_operator -def _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges( - wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - vt: fa.EdgeKField[ta.vpfloat], -) -> tuple[ - fa.EdgeKHalfField[ta.vpfloat], - fa.EdgeKHalfField[ta.vpfloat], - fa.EdgeKField[ta.vpfloat], -]: - """Formerly known as _mo_solve_nonhydro_stencil_36.""" - z_vt_ie = _interpolate_vt_to_interface_edges(wgtfac_e=wgtfac_e, vt=vt) - vn_ie, z_kin_hor_e = _interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges( - wgtfac_e=wgtfac_e, vn=vn, vt=vt - ) - return vn_ie, z_vt_ie, z_kin_hor_e - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges( - wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - vt: fa.EdgeKField[ta.vpfloat], - vn_ie: fa.EdgeKHalfField[ta.vpfloat], - z_vt_ie: fa.EdgeKHalfField[ta.vpfloat], - z_kin_hor_e: fa.EdgeKField[ta.vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges( - wgtfac_e=wgtfac_e, - vn=vn, - vt=vt, - out=(vn_ie, z_vt_ie, z_kin_hor_e), - domain=( - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ), - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py deleted file mode 100644 index 2120d43263..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py +++ /dev/null @@ -1,65 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import astype - -from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_kinetic_energy import ( - _compute_horizontal_kinetic_energy, -) -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges( - wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - vt: fa.EdgeKField[ta.vpfloat], -) -> tuple[fa.EdgeKHalfField[ta.vpfloat], fa.EdgeKField[ta.vpfloat]]: - """Formerly known as _mo_velocity_advection_stencil_02.""" - # TODO(): This stencil fusion with the one below is not optimal: - # _compute_horizontal_kinetic_energy is wasting computation by assigning - # vn_ie and vt_ie which are thrown away. - wgtfac_e_wp = astype(wgtfac_e, wpfloat) - - vn_ie_wp = wgtfac_e_wp * vn(dims.KHalfDim + 0.5) + (wpfloat("1.0") - wgtfac_e_wp) * vn( - dims.KHalfDim - 0.5 - ) - _, _, z_kin_hor_e = _compute_horizontal_kinetic_energy(vn=vn, vt=vt) - - return astype(vn_ie_wp, vpfloat), z_kin_hor_e - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges( - wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], - vn: fa.EdgeKField[ta.wpfloat], - vt: fa.EdgeKField[ta.vpfloat], - vn_ie: fa.EdgeKHalfField[ta.vpfloat], - z_kin_hor_e: fa.EdgeKField[ta.vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges( - wgtfac_e=wgtfac_e, - vn=vn, - vt=vt, - out=(vn_ie, z_kin_hor_e), - domain=( - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - { - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ), - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py deleted file mode 100644 index 5c72b1ae41..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ /dev/null @@ -1,43 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import astype, neighbor_sum - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.dimension import V2C -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in: fa.CellKHalfField[wpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], wpfloat], -) -> fa.VertexKHalfField[vpfloat]: - p_vert_out_wp = neighbor_sum(p_cell_in(V2C) * c_intp, axis=dims.V2CDim) - return astype(p_vert_out_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in: fa.CellKHalfField[wpfloat], - c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], wpfloat], - p_vert_out: fa.VertexKHalfField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( - p_cell_in=p_cell_in, - c_intp=c_intp, - out=p_vert_out, - domain={ - dims.VertexDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_math_divrot_rot_vertex_ri_dsl.py deleted file mode 100644 index 5c61e2844b..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/mo_math_divrot_rot_vertex_ri_dsl.py +++ /dev/null @@ -1,43 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import astype, neighbor_sum - -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.dimension import V2E -from icon4py.model.common.type_alias import vpfloat, wpfloat - - -@gtx.field_operator -def _mo_math_divrot_rot_vertex_ri_dsl( - vec_e: fa.EdgeKField[wpfloat], - geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], wpfloat], -) -> fa.VertexKField[vpfloat]: - rot_vec_wp = neighbor_sum(vec_e(V2E) * geofac_rot, axis=dims.V2EDim) - return astype(rot_vec_wp, vpfloat) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def mo_math_divrot_rot_vertex_ri_dsl( - vec_e: fa.EdgeKField[wpfloat], - geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], wpfloat], - rot_vec: fa.VertexKField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _mo_math_divrot_rot_vertex_ri_dsl( - vec_e=vec_e, - geofac_rot=geofac_rot, - out=rot_vec, - domain={ - dims.VertexDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_corrector.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_corrector.py new file mode 100644 index 0000000000..556f1787f3 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_corrector.py @@ -0,0 +1,245 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx + +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import ( + _compute_advection_in_horizontal_momentum, + _compute_advection_in_vertical_momentum, +) +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta + + +@gtx.field_operator +def _compute_velocity_advection_in_corrector_step( + vn: fa.EdgeKField[ta.wpfloat], + w: fa.CellKHalfField[ta.wpfloat], + tangential_wind: fa.EdgeKField[ta.vpfloat], + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + coeff1_dwdz: fa.CellKField[ta.vpfloat], + coeff2_dwdz: fa.CellKField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + owner_mask: fa.CellField[bool], + coriolis_frequency: fa.EdgeField[ta.wpfloat], + geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], + coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + dtime: ta.wpfloat, + apply_extra_diffusion_on_vn: bool, + nlev: gtx.int32, + end_index_of_damping_layer: gtx.int32, +) -> tuple[ + fa.CellKHalfField[ta.vpfloat], + fa.CellKHalfField[ta.vpfloat], + fa.EdgeKField[ta.vpfloat], +]: + ( + vertical_wind_advective_tendency, + contravariant_corrected_w_at_cells_on_model_levels, + vertical_cfl, + ) = _compute_advection_in_vertical_momentum( + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + dtime=dtime, + skip_vertical_wind_advective_tendency=False, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + normal_wind_advective_tendency = _compute_advection_in_horizontal_momentum( + vn=vn, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + tangential_wind=tangential_wind, + coriolis_frequency=coriolis_frequency, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + vn_on_half_levels=vn_on_half_levels, + e_bln_c_s=e_bln_c_s, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + return ( + vertical_wind_advective_tendency, + vertical_cfl, + normal_wind_advective_tendency, + ) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_velocity_advection_in_corrector_step( + vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], + vertical_cfl: fa.CellKHalfField[ta.vpfloat], + normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat], + vn: fa.EdgeKField[ta.wpfloat], + w: fa.CellKHalfField[ta.wpfloat], + tangential_wind: fa.EdgeKField[ta.vpfloat], + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + coeff1_dwdz: fa.CellKField[ta.vpfloat], + coeff2_dwdz: fa.CellKField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + owner_mask: fa.CellField[bool], + coriolis_frequency: fa.EdgeField[ta.wpfloat], + geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], + coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + dtime: ta.wpfloat, + apply_extra_diffusion_on_vn: bool, + end_index_of_damping_layer: gtx.int32, + start_cell_lateral_boundary_level_4: gtx.int32, + end_cell_halo: gtx.int32, + start_edge_nudging_level_2: gtx.int32, + end_edge_local: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + """ + Compute the velocity advection tendencies of the corrector step. + + This is the port of ICON's `velocity_tendencies` (`mo_velocity_advection.f90`) for + `istep == 2`: the advection in the vertical momentum equation and the advection in the + horizontal momentum equation. The wind quantities derived from the normal wind are + computed in `compute_horizontal_velocity_quantities_and_fluxes` and passed in. + + Args: + - vertical_wind_advective_tendency: advective tendency of the vertical wind + - vertical_cfl: vertical cfl number at cells on half levels + - normal_wind_advective_tendency: advective tendency of the normal wind + - vn: normal wind at edges + - w: vertical wind at cell centers + - tangential_wind: tangential wind at edges on model levels + - tangential_wind_on_half_levels: tangential wind at edges on half levels + - vn_on_half_levels: normal wind at edges on half levels + - horizontal_kinetic_energy_at_edges_on_model_levels: horizontal kinetic energy at edges on model levels + - contravariant_correction_at_cells_on_half_levels: contravariant metric correction at cells on half levels + - coeff1_dwdz: metrics field (first coefficient for vertical derivative of vertical wind) + - coeff2_dwdz: metrics field (second coefficient for vertical derivative of vertical wind) + - c_intp: interpolation field for cell-to-vertex interpolation + - inv_dual_edge_length: inverse dual edge length + - inv_primal_edge_length: inverse primal edge length + - tangent_orientation: orientation of the edge with respect to the grid + - e_bln_c_s: interpolation field (edge-to-cell interpolation weights) + - ddqz_z_half: metrics field + - area: cell area + - geofac_n2s: interpolation field + - owner_mask: ownership mask for each cell + - coriolis_frequency: coriolis frequency parameter + - geofac_rot: metric field for rotor computation + - coeff_gradekin: metrics field/coefficient for the gradient of kinematic energy + - c_lin_e: metrics field for linear interpolation from cells to edges + - ddqz_z_full_e: metrics field equal to vertical spacing + - area_edge: area associated with each edge + - geofac_grdiv: metrics field used to compute the gradient of a divergence (of vn) + - dtime: time step + - apply_extra_diffusion_on_vn: option to apply extra diffusion to vn + - end_index_of_damping_layer: vertical index where damping ends + - start_cell_lateral_boundary_level_4: start index of the cell quantities + - end_cell_halo: end index of the cell quantities + - start_edge_nudging_level_2: start index of the normal wind tendency + - end_edge_local: end index of the normal wind tendency + - vertical_start: start index in the vertical dimension at model top + - vertical_end: end index in the vertical dimension at model bottom (number of model levels) + """ + + _compute_velocity_advection_in_corrector_step( + vn=vn, + w=w, + tangential_wind=tangential_wind, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + coriolis_frequency=coriolis_frequency, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + nlev=vertical_end, + end_index_of_damping_layer=end_index_of_damping_layer, + out=( + vertical_wind_advective_tendency, + vertical_cfl, + normal_wind_advective_tendency, + ), + domain=( + { + dims.CellDim: (start_cell_lateral_boundary_level_4, end_cell_halo), + dims.KHalfDim: (vertical_start + 1, vertical_end), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_4, end_cell_halo), + dims.KHalfDim: (vertical_start, vertical_end), + }, + { + dims.EdgeDim: (start_edge_nudging_level_2, end_edge_local), + dims.KDim: (vertical_start, vertical_end), + }, + ), + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_predictor.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_predictor.py new file mode 100644 index 0000000000..4b82712b6a --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_predictor.py @@ -0,0 +1,444 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +from gt4py.next import astype, broadcast +from gt4py.next.experimental import concat_where + +from icon4py.model.atmosphere.dycore.stencils.compute_contravariant_correction import ( + _compute_contravariant_correction, +) +from icon4py.model.atmosphere.dycore.stencils.compute_diagnostics_from_normal_wind import ( + _compute_horizontal_kinetic_energy, + _interpolate_to_half_levels, +) +from icon4py.model.atmosphere.dycore.stencils.extrapolate_at_top import _extrapolate_at_top +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import ( + _compute_advection_in_horizontal_momentum, + _compute_advection_in_vertical_momentum, +) +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.interpolation.stencils.compute_tangential_wind import ( + _compute_tangential_wind, +) +from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_half_levels import ( + _interpolate_cell_field_to_half_levels_vp, +) +from icon4py.model.common.interpolation.stencils.interpolate_to_cell_center_vp import ( + _interpolate_to_cell_center_vp, +) +from icon4py.model.common.type_alias import vpfloat + + +@gtx.field_operator +def _compute_diagnostics_from_normal_wind( + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + vn: fa.EdgeKField[ta.wpfloat], + rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], ta.wpfloat], + wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], + wgtfacq_e: fa.EdgeKField[ta.vpfloat], + ddxn_z_full: fa.EdgeKField[ta.vpfloat], + ddxt_z_full: fa.EdgeKField[ta.vpfloat], + skip_compute_predictor_vertical_advection: bool, + nlev: gtx.int32, +) -> tuple[ + fa.EdgeKField[ta.vpfloat], + fa.EdgeKHalfField[ta.vpfloat], + fa.EdgeKHalfField[ta.vpfloat], + fa.EdgeKField[ta.vpfloat], + fa.EdgeKField[ta.vpfloat], +]: + tangential_wind = astype(_compute_tangential_wind(vn, rbf_vec_coeff_e), vpfloat) + horizontal_kinetic_energy_at_edges_on_model_levels = _compute_horizontal_kinetic_energy( + vn, tangential_wind + ) + vn_on_half_levels = concat_where( + dims.KHalfDim < nlev, + _interpolate_to_half_levels(wgtfac_e, vn), + _extrapolate_at_top(wgtfacq_e, vn), + ) + + tangential_wind_on_half_levels = ( + _interpolate_to_half_levels(wgtfac_e, tangential_wind) + if not skip_compute_predictor_vertical_advection + else tangential_wind_on_half_levels + ) + + contravariant_correction_at_edges_on_model_levels = _compute_contravariant_correction( + vn, ddxn_z_full, ddxt_z_full, tangential_wind + ) + + return ( + tangential_wind, + tangential_wind_on_half_levels, + vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels, + ) + + +@gtx.field_operator +def _interpolate_contravariant_correction_to_cells_on_half_levels( + contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + wgtfac_c: fa.CellKHalfField[ta.vpfloat], + nflatlev: gtx.int32, +) -> fa.CellKHalfField[ta.vpfloat]: + contravariant_correction_at_cells_model_levels = _interpolate_to_cell_center_vp( + contravariant_correction_at_edges_on_model_levels, e_bln_c_s + ) + contravariant_correction_at_cells_model_levels = astype( + contravariant_correction_at_cells_model_levels, vpfloat + ) + + contravariant_correction_at_cells_on_half_levels = concat_where( + dims.KHalfDim >= nflatlev + 1, + _interpolate_cell_field_to_half_levels_vp( + wgtfac_c=wgtfac_c, interpolant=contravariant_correction_at_cells_model_levels + ), + broadcast(vpfloat("0.0"), (dims.CellDim, dims.KHalfDim)), + ) + + return contravariant_correction_at_cells_on_half_levels + + +@gtx.field_operator +def _compute_velocity_advection_in_predictor_step( + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], + vn: fa.EdgeKField[ta.wpfloat], + w: fa.CellKHalfField[ta.wpfloat], + rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], ta.wpfloat], + wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], + wgtfacq_e: fa.EdgeKField[ta.vpfloat], + ddxn_z_full: fa.EdgeKField[ta.vpfloat], + ddxt_z_full: fa.EdgeKField[ta.vpfloat], + coeff1_dwdz: fa.CellKField[ta.vpfloat], + coeff2_dwdz: fa.CellKField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + wgtfac_c: fa.CellKHalfField[ta.vpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + owner_mask: fa.CellField[bool], + coriolis_frequency: fa.EdgeField[ta.wpfloat], + geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], + coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + dtime: ta.wpfloat, + skip_compute_predictor_vertical_advection: bool, + apply_extra_diffusion_on_vn: bool, + nflatlev: gtx.int32, + nlev: gtx.int32, + end_index_of_damping_layer: gtx.int32, +) -> tuple[ + fa.EdgeKField[ta.vpfloat], + fa.EdgeKHalfField[ta.vpfloat], + fa.EdgeKHalfField[ta.vpfloat], + fa.EdgeKField[ta.vpfloat], + fa.EdgeKField[ta.vpfloat], + fa.CellKHalfField[ta.vpfloat], + fa.CellKHalfField[ta.vpfloat], + fa.CellKHalfField[ta.vpfloat], + fa.EdgeKField[ta.vpfloat], +]: + ( + tangential_wind, + tangential_wind_on_half_levels, + vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels, + ) = _compute_diagnostics_from_normal_wind( + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn=vn, + rbf_vec_coeff_e=rbf_vec_coeff_e, + wgtfac_e=wgtfac_e, + wgtfacq_e=wgtfacq_e, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, + nlev=nlev, + ) + + contravariant_correction_at_cells_on_half_levels = _interpolate_contravariant_correction_to_cells_on_half_levels( + contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, + e_bln_c_s=e_bln_c_s, + wgtfac_c=wgtfac_c, + nflatlev=nflatlev, + ) + + ( + maybe_vertical_wind_advective_tendency, # if `skip_compute_predictor_vertical_advection` this field will carry a dummy value + contravariant_corrected_w_at_cells_on_model_levels, + vertical_cfl, + ) = _compute_advection_in_vertical_momentum( + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + dtime=dtime, + skip_vertical_wind_advective_tendency=skip_compute_predictor_vertical_advection, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + # We forward the previous value if `skip_compute_predictor_vertical_advection`. + # This code looks weird because the `MOST_EFFICIENT` scheme skips computing the vertical_wind_advective_tendency, + # which the correct `EXPENSIVE` scheme would compute. + vertical_wind_advective_tendency = ( + maybe_vertical_wind_advective_tendency + if not skip_compute_predictor_vertical_advection + else vertical_wind_advective_tendency + ) + + normal_wind_advective_tendency = _compute_advection_in_horizontal_momentum( + vn=vn, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + tangential_wind=tangential_wind, + coriolis_frequency=coriolis_frequency, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + vn_on_half_levels=vn_on_half_levels, + e_bln_c_s=e_bln_c_s, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + return ( + tangential_wind, + tangential_wind_on_half_levels, + vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels, + vertical_wind_advective_tendency, + vertical_cfl, + normal_wind_advective_tendency, + ) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_velocity_advection_in_predictor_step( + tangential_wind: fa.EdgeKField[ta.vpfloat], + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + vertical_wind_advective_tendency: fa.CellKHalfField[ta.vpfloat], + vertical_cfl: fa.CellKHalfField[ta.vpfloat], + normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat], + vn: fa.EdgeKField[ta.wpfloat], + w: fa.CellKHalfField[ta.wpfloat], + rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], ta.wpfloat], + wgtfac_e: fa.EdgeKHalfField[ta.vpfloat], + wgtfacq_e: fa.EdgeKField[ta.vpfloat], + ddxn_z_full: fa.EdgeKField[ta.vpfloat], + ddxt_z_full: fa.EdgeKField[ta.vpfloat], + coeff1_dwdz: fa.CellKField[ta.vpfloat], + coeff2_dwdz: fa.CellKField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + wgtfac_c: fa.CellKHalfField[ta.vpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + owner_mask: fa.CellField[bool], + coriolis_frequency: fa.EdgeField[ta.wpfloat], + geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], + coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + dtime: ta.wpfloat, + skip_compute_predictor_vertical_advection: bool, + apply_extra_diffusion_on_vn: bool, + nflatlev: gtx.int32, + end_index_of_damping_layer: gtx.int32, + start_edge_lateral_boundary_level_5: gtx.int32, + end_edge_halo_level_2: gtx.int32, + start_cell_lateral_boundary_level_4: gtx.int32, + end_cell_halo: gtx.int32, + start_edge_nudging_level_2: gtx.int32, + end_edge_local: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + """ + Compute the velocity advection tendencies of the predictor step. + + This is the port of ICON's `velocity_tendencies` (`mo_velocity_advection.f90`) for + `istep == 1`: the wind quantities derived from the normal wind, the advection in the + vertical momentum equation and the advection in the horizontal momentum equation. + + Args: + - tangential_wind: tangential wind at edges on model levels + - tangential_wind_on_half_levels: tangential wind at edges on half levels + - vn_on_half_levels: normal wind at edges on half levels + - horizontal_kinetic_energy_at_edges_on_model_levels: horizontal kinetic energy at edges on model levels + - contravariant_correction_at_edges_on_model_levels: contravariant metric correction at edges on model levels + - contravariant_correction_at_cells_on_half_levels: contravariant metric correction at cells on half levels + - vertical_wind_advective_tendency: advective tendency of the vertical wind + - vertical_cfl: vertical cfl number at cells on half levels + - normal_wind_advective_tendency: advective tendency of the normal wind + - vn: normal wind at edges + - w: vertical wind at cell centers + - rbf_vec_coeff_e: interpolation field (RBF vector coefficient on edges) + - wgtfac_e: metrics field + - wgtfacq_e: metrics field (weights for interpolation) + - ddxn_z_full: metrics field (derivative of topography in the normal direction) + - ddxt_z_full: metrics field (derivative of topography in the tangential direction) + - coeff1_dwdz: metrics field (first coefficient for vertical derivative of vertical wind) + - coeff2_dwdz: metrics field (second coefficient for vertical derivative of vertical wind) + - c_intp: interpolation field for cell-to-vertex interpolation + - inv_dual_edge_length: inverse dual edge length + - inv_primal_edge_length: inverse primal edge length + - tangent_orientation: orientation of the edge with respect to the grid + - e_bln_c_s: interpolation field (edge-to-cell interpolation weights) + - wgtfac_c: metric coefficient for interpolating a cell variable from full to half levels + - ddqz_z_half: metrics field + - area: cell area + - geofac_n2s: interpolation field + - owner_mask: ownership mask for each cell + - coriolis_frequency: coriolis frequency parameter + - geofac_rot: metric field for rotor computation + - coeff_gradekin: metrics field/coefficient for the gradient of kinematic energy + - c_lin_e: metrics field for linear interpolation from cells to edges + - ddqz_z_full_e: metrics field equal to vertical spacing + - area_edge: area associated with each edge + - geofac_grdiv: metrics field used to compute the gradient of a divergence (of vn) + - dtime: time step + - skip_compute_predictor_vertical_advection: logical flag to skip the vertical advection + - apply_extra_diffusion_on_vn: option to apply extra diffusion to vn + - nflatlev: index of the first flat level + - end_index_of_damping_layer: vertical index where damping ends + - start_edge_lateral_boundary_level_5: start index of the edge quantities derived from vn + - end_edge_halo_level_2: end index of the edge quantities derived from vn + - start_cell_lateral_boundary_level_4: start index of the cell quantities + - end_cell_halo: end index of the cell quantities + - start_edge_nudging_level_2: start index of the normal wind tendency + - end_edge_local: end index of the normal wind tendency + - vertical_start: start index in the vertical dimension at model top + - vertical_end: end index in the vertical dimension at model bottom (number of model levels) + """ + + _compute_velocity_advection_in_predictor_step( + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vertical_wind_advective_tendency=vertical_wind_advective_tendency, + vn=vn, + w=w, + rbf_vec_coeff_e=rbf_vec_coeff_e, + wgtfac_e=wgtfac_e, + wgtfacq_e=wgtfacq_e, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + e_bln_c_s=e_bln_c_s, + wgtfac_c=wgtfac_c, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + coriolis_frequency=coriolis_frequency, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + nflatlev=nflatlev, + nlev=vertical_end, + end_index_of_damping_layer=end_index_of_damping_layer, + out=( + tangential_wind, + tangential_wind_on_half_levels, + vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels, + vertical_wind_advective_tendency, + vertical_cfl, + normal_wind_advective_tendency, + ), + domain=( + { + dims.EdgeDim: (start_edge_lateral_boundary_level_5, end_edge_halo_level_2), + dims.KDim: (vertical_start, vertical_end), + }, + { + dims.EdgeDim: (start_edge_lateral_boundary_level_5, end_edge_halo_level_2), + dims.KHalfDim: (vertical_start, vertical_end), + }, + { + dims.EdgeDim: (start_edge_lateral_boundary_level_5, end_edge_halo_level_2), + dims.KHalfDim: (vertical_start, vertical_end + 1), + }, + { + dims.EdgeDim: (start_edge_lateral_boundary_level_5, end_edge_halo_level_2), + dims.KDim: (vertical_start, vertical_end), + }, + { + dims.EdgeDim: (start_edge_lateral_boundary_level_5, end_edge_halo_level_2), + dims.KDim: (nflatlev, vertical_end), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_4, end_cell_halo), + dims.KHalfDim: (vertical_start, vertical_end), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_4, end_cell_halo), + dims.KHalfDim: (vertical_start + 1, vertical_end), + }, + { + dims.CellDim: (start_cell_lateral_boundary_level_4, end_cell_halo), + dims.KHalfDim: (vertical_start, vertical_end), + }, + { + dims.EdgeDim: (start_edge_nudging_level_2, end_edge_local), + dims.KDim: (vertical_start, vertical_end), + }, + ), + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_terms.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_terms.py new file mode 100644 index 0000000000..5e64c5cdc8 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/velocity_advection_terms.py @@ -0,0 +1,592 @@ +# 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 enum + +import gt4py.next as gtx +from gt4py.next import ( + abs, # noqa: A004 + astype, + broadcast, + maximum, + minimum, + neighbor_sum, + where, +) +from gt4py.next.experimental import concat_where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import C2E, C2E2CO, E2C, E2C2EO, E2V +from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_vertex import ( + _interpolate_cell_field_to_vertex, +) +from icon4py.model.common.interpolation.stencils.interpolate_to_cell_center_vp import ( + _interpolate_to_cell_center_vp, +) +from icon4py.model.common.math.stencils.compute_curl import _compute_curl +from icon4py.model.common.type_alias import vpfloat, wpfloat + + +class VerticalCflConstants(ta.wpfloat, enum.Enum): + #: w is clipped and extra diffusion is applied above this vertical CFL number + W_LIMIT = 0.65 + #: w is clipped to this vertical CFL number + W_MAX = 0.85 + EXTRA_DIFFUSION_SCALING = 0.05 / (W_MAX - W_LIMIT) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def _interpolate_contravariant_vertical_velocity_to_full_levels( + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[vpfloat], + nlev: gtx.int32, +) -> fa.CellKField[vpfloat]: + # TODO(havogt): Note that `concat_where(dims.KDim == nlev-1, ...)` is currently broken + # because of insufficiency in the domain inference of GT4Py, + # see https://github.com/GridTools/gt4py/issues/2205. + return concat_where( + dims.KDim < nlev - 1, + vpfloat("0.5") + * ( + contravariant_corrected_w_at_cells_on_half_levels(dims.KDim - 0.5) + + contravariant_corrected_w_at_cells_on_half_levels(dims.KDim + 0.5) + ), + vpfloat("0.5") * contravariant_corrected_w_at_cells_on_half_levels(dims.KDim - 0.5), + ) + + +@gtx.field_operator +def _compute_horizontal_advection_of_w( + w: fa.CellKHalfField[ta.wpfloat], + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], +) -> fa.EdgeKHalfField[ta.vpfloat]: + w_at_vertices = astype(_interpolate_cell_field_to_vertex(w, c_intp), vpfloat) + vn_on_half_levels_wp = astype(vn_on_half_levels, wpfloat) + + horizontal_advection_of_w_at_edges_on_half_levels = ( + vn_on_half_levels_wp * inv_dual_edge_length * (w(E2C[0]) - w(E2C[1])) + + tangential_wind_on_half_levels + * inv_primal_edge_length + * tangent_orientation + * astype(w_at_vertices(E2V[0]) - w_at_vertices(E2V[1]), wpfloat) + ) + + return astype(horizontal_advection_of_w_at_edges_on_half_levels, vpfloat) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def _compute_vertical_advection_of_w( + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[vpfloat], + w: fa.CellKHalfField[wpfloat], + coeff1_dwdz: fa.CellKField[vpfloat], + coeff2_dwdz: fa.CellKField[vpfloat], +) -> fa.CellKHalfField[vpfloat]: + contravariant_corrected_w_at_cells_on_half_levels_wp = astype( + contravariant_corrected_w_at_cells_on_half_levels, wpfloat + ) + coeff1_dwdz_at_half_levels = coeff1_dwdz(dims.KHalfDim + 0.5) + coeff2_dwdz_at_half_levels = coeff2_dwdz(dims.KHalfDim + 0.5) + coeff1_dwdz_wp, coeff2_dwdz_wp = astype( + (coeff1_dwdz_at_half_levels, coeff2_dwdz_at_half_levels), wpfloat + ) + + vertical_wind_advective_tendency_wp = -contravariant_corrected_w_at_cells_on_half_levels_wp * ( + w(dims.KHalfDim - 1) * coeff1_dwdz_wp + - w(dims.KHalfDim + 1) * coeff2_dwdz_wp + + w * astype(coeff2_dwdz_at_half_levels - coeff1_dwdz_at_half_levels, wpfloat) + ) + return astype(vertical_wind_advective_tendency_wp, vpfloat) + + +@gtx.field_operator +def _compute_interpolated_horizontal_advection_of_w( + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + horizontal_advection_of_w_at_edges_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], +) -> fa.CellKHalfField[ta.wpfloat]: + """Formerly known as _mo_velocity_advection_stencil_17.""" + horizontal_advection_of_w_at_edges_on_half_levels_wp = astype( + horizontal_advection_of_w_at_edges_on_half_levels, wpfloat + ) + return neighbor_sum( + horizontal_advection_of_w_at_edges_on_half_levels_wp(C2E) * e_bln_c_s, axis=dims.C2EDim + ) + + +@gtx.field_operator +def _compute_extra_diffusion_for_w( + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + w: fa.CellKHalfField[ta.wpfloat], + dtime: ta.wpfloat, +) -> fa.CellKHalfField[ta.wpfloat]: + """Formerly known as _mo_velocity_advection_stencil_18.""" + contravariant_corrected_w_at_cells_on_half_levels_wp, ddqz_z_half_wp = astype( + (contravariant_corrected_w_at_cells_on_half_levels, ddqz_z_half), wpfloat + ) + + difcoef = (VerticalCflConstants.EXTRA_DIFFUSION_SCALING / dtime) * minimum( + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + abs(contravariant_corrected_w_at_cells_on_half_levels_wp) * dtime / ddqz_z_half_wp + - VerticalCflConstants.W_LIMIT, + ) + + return difcoef * area * neighbor_sum(w(C2E2CO) * geofac_n2s, axis=dims.C2E2CODim) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def _compute_cfl( + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + dtime: ta.wpfloat, +) -> tuple[fa.CellKHalfField[bool], fa.CellKHalfField[ta.vpfloat]]: + contravariant_corrected_w_at_cells_on_half_levels_wp, ddqz_z_half_wp = astype( + (contravariant_corrected_w_at_cells_on_half_levels, ddqz_z_half), wpfloat + ) + + vertical_cfl_number = ( + contravariant_corrected_w_at_cells_on_half_levels_wp * dtime / ddqz_z_half_wp + ) + + cfl_clipping = where( + abs(vertical_cfl_number) > VerticalCflConstants.W_LIMIT, + broadcast(True, (dims.CellDim, dims.KHalfDim)), + False, + ) + + vertical_cfl = where( + cfl_clipping, + vertical_cfl_number, + broadcast(wpfloat("0.0"), (dims.CellDim, dims.KHalfDim)), + ) + + return cfl_clipping, astype(vertical_cfl, vpfloat) + + +@gtx.field_operator(grid_type=gtx.GridType.UNSTRUCTURED) +def _clip_contravariant_corrected_w( + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + cfl_clipping: fa.CellKHalfField[bool], + vertical_cfl: fa.CellKHalfField[ta.vpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + dtime: ta.wpfloat, +) -> fa.CellKHalfField[ta.vpfloat]: + contravariant_corrected_w_at_cells_on_half_levels_wp = astype( + contravariant_corrected_w_at_cells_on_half_levels, wpfloat + ) + + maximum_vertical_cfl = astype(VerticalCflConstants.W_MAX, vpfloat) + + contravariant_corrected_w_at_cells_on_half_levels_wp = where( + (cfl_clipping) & (vertical_cfl < -maximum_vertical_cfl), + astype(-maximum_vertical_cfl * ddqz_z_half, wpfloat) / dtime, + contravariant_corrected_w_at_cells_on_half_levels_wp, + ) + + contravariant_corrected_w_at_cells_on_half_levels_wp = where( + (cfl_clipping) & (vertical_cfl > maximum_vertical_cfl), + astype(maximum_vertical_cfl * ddqz_z_half, wpfloat) / dtime, + contravariant_corrected_w_at_cells_on_half_levels_wp, + ) + + return astype(contravariant_corrected_w_at_cells_on_half_levels_wp, vpfloat) + + +@gtx.field_operator +def _compute_contravariant_corrected_w( + w: fa.CellKHalfField[ta.wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], +) -> fa.CellKHalfField[ta.vpfloat]: + contravariant_corrected_w_at_cells_on_half_levels = ( + astype(w, vpfloat) - contravariant_correction_at_cells_on_half_levels + ) + + return contravariant_corrected_w_at_cells_on_half_levels + + +@gtx.field_operator +def _compute_contravariant_corrected_w_and_cfl( + w: fa.CellKHalfField[ta.wpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + dtime: ta.wpfloat, + nlev: gtx.int32, + end_index_of_damping_layer: gtx.int32, +) -> tuple[fa.CellKHalfField[ta.vpfloat], fa.CellKHalfField[bool], fa.CellKHalfField[ta.vpfloat]]: + #: intermediate variable contravariant_corrected_w_at_cells_on_half_levels is originally declared as z_w_con_c in ICON + contravariant_corrected_w_at_cells_on_half_levels = _compute_contravariant_corrected_w( + w, contravariant_correction_at_cells_on_half_levels + ) + + cfl_clipping, vertical_cfl = concat_where( + (dims.KHalfDim >= maximum(2, end_index_of_damping_layer - 2)) & (dims.KHalfDim < nlev - 3), + _compute_cfl( + ddqz_z_half=ddqz_z_half, + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + dtime=dtime, + ), + ( + broadcast(False, (dims.CellDim, dims.KHalfDim)), + broadcast(vpfloat("0.0"), (dims.CellDim, dims.KHalfDim)), + ), + ) + + contravariant_corrected_w_at_cells_on_half_levels = _clip_contravariant_corrected_w( + contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping, + vertical_cfl, + ddqz_z_half, + dtime, + ) + + return contravariant_corrected_w_at_cells_on_half_levels, cfl_clipping, vertical_cfl + + +@gtx.field_operator +def _compute_advective_vertical_wind_tendency( + w: fa.CellKHalfField[ta.wpfloat], + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + contravariant_corrected_w_at_cells_on_half_levels: fa.CellKHalfField[ta.wpfloat], + cfl_clipping: fa.CellKHalfField[bool], + coeff1_dwdz: fa.CellKField[ta.vpfloat], + coeff2_dwdz: fa.CellKField[ta.vpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + owner_mask: fa.CellField[bool], + dtime: ta.wpfloat, +) -> fa.CellKHalfField[ta.vpfloat]: + # TODO(havogt): the wp-vp roundtrips are here to be faithful to ICON's mixed precision. + # We assume it was not a deliberate decision in ICON and plan to drop the extra conversions + # in a separate PR. + vertical_advection_of_w = _compute_vertical_advection_of_w( + contravariant_corrected_w_at_cells_on_half_levels, w, coeff1_dwdz, coeff2_dwdz + ) + horizontal_advection_of_w_at_edges_on_half_levels = _compute_horizontal_advection_of_w( + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + ) + interpolated_horizontal_advection_of_w = _compute_interpolated_horizontal_advection_of_w( + e_bln_c_s, horizontal_advection_of_w_at_edges_on_half_levels + ) + extra_diffusion_for_w = _compute_extra_diffusion_for_w( + contravariant_corrected_w_at_cells_on_half_levels, + ddqz_z_half, + area, + geofac_n2s, + w, + dtime, + ) + + vertical_wind_advective_tendency = astype( + astype(vertical_advection_of_w, wpfloat) + interpolated_horizontal_advection_of_w, + vpfloat, + ) + + vertical_wind_advective_tendency_wp = astype(vertical_wind_advective_tendency, wpfloat) + vertical_wind_advective_tendency = astype( + where( + cfl_clipping & owner_mask, + vertical_wind_advective_tendency_wp + extra_diffusion_for_w, + vertical_wind_advective_tendency_wp, + ), + vpfloat, + ) + + return vertical_wind_advective_tendency + + +@gtx.field_operator +def _compute_advection_in_vertical_momentum( + w: fa.CellKHalfField[ta.wpfloat], + tangential_wind_on_half_levels: fa.EdgeKHalfField[ta.wpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + contravariant_correction_at_cells_on_half_levels: fa.CellKHalfField[ta.vpfloat], + coeff1_dwdz: fa.CellKField[ta.vpfloat], + coeff2_dwdz: fa.CellKField[ta.vpfloat], + c_intp: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], ta.wpfloat], + inv_dual_edge_length: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + ddqz_z_half: fa.CellKHalfField[ta.vpfloat], + area: fa.CellField[ta.wpfloat], + geofac_n2s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2CODim], ta.wpfloat], + owner_mask: fa.CellField[bool], + dtime: ta.wpfloat, + skip_vertical_wind_advective_tendency: bool, + nlev: gtx.int32, + end_index_of_damping_layer: gtx.int32, +) -> tuple[fa.CellKHalfField[ta.vpfloat], fa.CellKField[ta.vpfloat], fa.CellKHalfField[ta.vpfloat]]: + ( + contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping, + vertical_cfl, + ) = _compute_contravariant_corrected_w_and_cfl( + w=w, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + ddqz_z_half=ddqz_z_half, + dtime=dtime, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + vertical_wind_advective_tendency = ( + _compute_advective_vertical_wind_tendency( + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping=cfl_clipping, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + dtime=dtime, + ) + if not skip_vertical_wind_advective_tendency + # Skipped: a placeholder that must not be read; the caller keeps its previous tendency. + else broadcast(vpfloat("0.0"), (dims.CellDim, dims.KHalfDim)) + ) + + contravariant_corrected_w_at_cells_on_model_levels = ( + _interpolate_contravariant_vertical_velocity_to_full_levels( + contravariant_corrected_w_at_cells_on_half_levels, nlev + ) + ) + + return ( + vertical_wind_advective_tendency, + contravariant_corrected_w_at_cells_on_model_levels, + vertical_cfl, + ) + + +@gtx.field_operator +def _compute_advective_normal_wind_tendency( + horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + upward_vorticity_at_vertices_on_model_levels: fa.VertexKField[ta.vpfloat], + tangential_wind: fa.EdgeKField[ta.vpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + coriolis_frequency: fa.EdgeField[ta.wpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], +) -> fa.EdgeKField[ta.vpfloat]: + #: intermediate variable horizontal_kinetic_energy_at_cells_on_model_levels is originally declared as z_ekinh in ICON + horizontal_kinetic_energy_at_cells_on_model_levels = _interpolate_to_cell_center_vp( + horizontal_kinetic_energy_at_edges_on_model_levels, e_bln_c_s + ) + horizontal_kinetic_energy_at_cells_on_model_levels = astype( + horizontal_kinetic_energy_at_cells_on_model_levels, vpfloat + ) + + ( + contravariant_corrected_w_at_cells_on_model_levels_wp, + ddqz_z_full_e_wp, + tangential_wind_wp, + ) = astype( + (contravariant_corrected_w_at_cells_on_model_levels, ddqz_z_full_e, tangential_wind), + wpfloat, + ) + + horizontal_advection = ( + horizontal_kinetic_energy_at_edges_on_model_levels + * (coeff_gradekin[dims.E2CDim(0)] - coeff_gradekin[dims.E2CDim(1)]) + + coeff_gradekin[dims.E2CDim(1)] + * horizontal_kinetic_energy_at_cells_on_model_levels(E2C[1]) + - coeff_gradekin[dims.E2CDim(0)] + * horizontal_kinetic_energy_at_cells_on_model_levels(E2C[0]) + ) + + vertical_advection = ( + neighbor_sum( + c_lin_e * contravariant_corrected_w_at_cells_on_model_levels_wp(E2C), axis=dims.E2CDim + ) + * astype((vn_on_half_levels(dims.KDim - 0.5) - vn_on_half_levels(dims.KDim + 0.5)), wpfloat) + / ddqz_z_full_e_wp + ) + + coriolis_term = tangential_wind_wp * ( + coriolis_frequency + + astype( + vpfloat("0.5") + * neighbor_sum(upward_vorticity_at_vertices_on_model_levels(E2V), axis=dims.E2VDim), + wpfloat, + ) + ) + normal_wind_advective_tendency_wp = -(horizontal_advection + vertical_advection + coriolis_term) + + return astype(normal_wind_advective_tendency_wp, vpfloat) + + +@gtx.field_operator +def _compute_extra_diffusion( + vn: fa.EdgeKField[ta.wpfloat], + upward_vorticity_at_vertices_on_model_levels: fa.VertexKField[ta.vpfloat], + difcoef: fa.EdgeKField[ta.wpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], +) -> fa.EdgeKField[ta.wpfloat]: + gradient_of_divergence_of_vn = neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=dims.E2C2EODim) + + gradient_of_vorticity = ( + tangent_orientation + * inv_primal_edge_length + * astype( + upward_vorticity_at_vertices_on_model_levels(E2V[1]) + - upward_vorticity_at_vertices_on_model_levels(E2V[0]), + wpfloat, + ) + ) + + extra_diffusion_on_vn = ( + difcoef * area_edge * (gradient_of_divergence_of_vn + gradient_of_vorticity) + ) + + return extra_diffusion_on_vn + + +@gtx.field_operator +def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask( + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + upward_vorticity_at_vertices_on_model_levels: fa.VertexKField[ta.vpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + vn: fa.EdgeKField[ta.wpfloat], + normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat], + dtime: ta.wpfloat, +) -> fa.EdgeKField[ta.vpfloat]: + ( + contravariant_corrected_w_at_cells_on_model_levels_wp, + ddqz_z_full_e_wp, + normal_wind_advective_tendency_wp, + ) = astype( + ( + contravariant_corrected_w_at_cells_on_model_levels, + ddqz_z_full_e, + normal_wind_advective_tendency, + ), + wpfloat, + ) + + #: intermediate variable contravariant_corrected_w_at_edges_on_model_levels is originally declared as w_con_e in ICON + contravariant_corrected_w_at_edges_on_model_levels = neighbor_sum( + c_lin_e * contravariant_corrected_w_at_cells_on_model_levels_wp(E2C), axis=dims.E2CDim + ) + vertical_cfl_number_at_edges = ( + abs(contravariant_corrected_w_at_edges_on_model_levels) * dtime / ddqz_z_full_e_wp + ) + difcoef = (VerticalCflConstants.EXTRA_DIFFUSION_SCALING / dtime) * minimum( + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + vertical_cfl_number_at_edges - VerticalCflConstants.W_LIMIT, + ) + normal_wind_advective_tendency_wp = where( + vertical_cfl_number_at_edges > VerticalCflConstants.W_LIMIT, + normal_wind_advective_tendency_wp + + _compute_extra_diffusion( + vn=vn, + upward_vorticity_at_vertices_on_model_levels=upward_vorticity_at_vertices_on_model_levels, + difcoef=difcoef, + area_edge=area_edge, + geofac_grdiv=geofac_grdiv, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + ), + normal_wind_advective_tendency_wp, + ) + return astype(normal_wind_advective_tendency_wp, vpfloat) + + +@gtx.field_operator +def _compute_advection_in_horizontal_momentum( + vn: fa.EdgeKField[ta.wpfloat], + horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.vpfloat], + tangential_wind: fa.EdgeKField[ta.vpfloat], + coriolis_frequency: fa.EdgeField[ta.wpfloat], + contravariant_corrected_w_at_cells_on_model_levels: fa.CellKField[ta.vpfloat], + vn_on_half_levels: fa.EdgeKHalfField[ta.vpfloat], + e_bln_c_s: gtx.Field[gtx.Dims[dims.CellDim, dims.C2EDim], ta.wpfloat], + geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], ta.wpfloat], + coeff_gradekin: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.vpfloat], + c_lin_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim], ta.wpfloat], + ddqz_z_full_e: fa.EdgeKField[ta.vpfloat], + area_edge: fa.EdgeField[ta.wpfloat], + tangent_orientation: fa.EdgeField[ta.wpfloat], + inv_primal_edge_length: fa.EdgeField[ta.wpfloat], + geofac_grdiv: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + dtime: ta.wpfloat, + apply_extra_diffusion_on_vn: bool, + nlev: gtx.int32, + end_index_of_damping_layer: gtx.int32, +) -> fa.EdgeKField[ta.vpfloat]: + upward_vorticity_at_vertices_on_model_levels = astype(_compute_curl(vn, geofac_rot), vpfloat) + + normal_wind_advective_tendency = _compute_advective_normal_wind_tendency( + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + upward_vorticity_at_vertices_on_model_levels=upward_vorticity_at_vertices_on_model_levels, + tangential_wind=tangential_wind, + vn_on_half_levels=vn_on_half_levels, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + coriolis_frequency=coriolis_frequency, + e_bln_c_s=e_bln_c_s, + c_lin_e=c_lin_e, + coeff_gradekin=coeff_gradekin, + ddqz_z_full_e=ddqz_z_full_e, + ) + + if apply_extra_diffusion_on_vn: + normal_wind_advective_tendency = concat_where( + ((maximum(2, end_index_of_damping_layer - 2)) <= dims.KDim) & (dims.KDim < (nlev - 4)), + _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask( + c_lin_e=c_lin_e, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + upward_vorticity_at_vertices_on_model_levels=upward_vorticity_at_vertices_on_model_levels, + geofac_grdiv=geofac_grdiv, + vn=vn, + normal_wind_advective_tendency=normal_wind_advective_tendency, + dtime=dtime, + ), + normal_wind_advective_tendency, + ) + + return normal_wind_advective_tendency diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py deleted file mode 100644 index b2b37fb43a..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py +++ /dev/null @@ -1,392 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import gt4py.next as gtx -import gt4py.next.typing as gtx_typing - -from icon4py.model.atmosphere.dycore import dycore_states -from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_horizontal_momentum_equation import ( - compute_advection_in_horizontal_momentum, -) -from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_vertical_momentum_equation import ( - compute_advection_in_corrector_vertical_momentum, - compute_advection_in_predictor_vertical_momentum, -) -from icon4py.model.atmosphere.dycore.stencils.compute_diagnostics_from_normal_wind import ( - compute_diagnostics_from_normal_wind, -) -from icon4py.model.common import ( - dimension as dims, - field_type_aliases as fa, - model_backends, - type_alias as ta, -) -from icon4py.model.common.grid import ( - horizontal as h_grid, - icon as icon_grid, - states as grid_states, - vertical as v_grid, -) -from icon4py.model.common.model_options import setup_program -from icon4py.model.common.states import nonhydro_states, prognostic_state as prognostics -from icon4py.model.common.utils import data_allocation as data_alloc - - -class VelocityAdvection: - def __init__( - self, - *, - grid: icon_grid.IconGrid, - metric_state: dycore_states.MetricStateNonHydro, - interpolation_state: dycore_states.InterpolationState, - vertical_params: v_grid.VerticalGrid, - edge_params: grid_states.EdgeParams, - owner_mask: fa.CellField[bool], - backend: gtx_typing.Backend - | model_backends.DeviceType - | model_backends.BackendDescriptor - | None, - ): - self._grid: icon_grid.IconGrid = grid - self._metric_state: dycore_states.MetricStateNonHydro = metric_state - self._interpolation_state: dycore_states.InterpolationState = interpolation_state - self._vertical_params: v_grid.VerticalGrid = vertical_params - self._edge_params: grid_states.EdgeParams = edge_params - self._c_owner_mask: fa.CellField[bool] = owner_mask - - self._cfl_w_limit: float = 0.65 - self._scalfac_exdiff: float = 0.05 - self._allocate_local_fields(model_backends.get_allocator(backend)) - self._determine_local_domains() - - self._compute_diagnostics_from_normal_wind = setup_program( - backend=backend, - program=compute_diagnostics_from_normal_wind, - constant_args={ - "rbf_vec_coeff_e": self._interpolation_state.rbf_vec_coeff_e, - "wgtfac_e": self._metric_state.wgtfac_e, - "ddxn_z_full": self._metric_state.ddxn_z_full, - "ddxt_z_full": self._metric_state.ddxt_z_full, - "wgtfacq_e": self._metric_state.wgtfacq_e, - "c_intp": self._interpolation_state.c_intp, - "inv_dual_edge_length": self._edge_params.inverse_dual_edge_lengths, - "inv_primal_edge_length": self._edge_params.inverse_primal_edge_lengths, - "tangent_orientation": self._edge_params.tangent_orientation, - }, - variants={ - "skip_compute_predictor_vertical_advection": [True, False], - }, - horizontal_sizes={ - "horizontal_start": self._start_edge_lateral_boundary_level_5, - "horizontal_end": self._end_edge_halo_level_2, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(self._grid.num_levels + 1), - "nflatlev": self._vertical_params.nflatlev, - }, - offset_provider=self._grid.connectivities, - ) - - self._compute_advection_in_predictor_vertical_momentum = setup_program( - backend=backend, - program=compute_advection_in_predictor_vertical_momentum, - constant_args={ - "coeff1_dwdz": self._metric_state.coeff1_dwdz, - "coeff2_dwdz": self._metric_state.coeff2_dwdz, - "e_bln_c_s": self._interpolation_state.e_bln_c_s, - "wgtfac_c": self._metric_state.wgtfac_c, - "ddqz_z_half": self._metric_state.ddqz_z_half, - "geofac_n2s": self._interpolation_state.geofac_n2s, - "owner_mask": self._c_owner_mask, - }, - variants={ - "skip_compute_predictor_vertical_advection": [True, False], - }, - vertical_sizes={ - "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer, - "nflatlev": self._vertical_params.nflatlev, - "vertical_start": gtx.int32(0), - "vertical_end": self._grid.num_levels, - }, - horizontal_sizes={ - "horizontal_start": self._start_cell_lateral_boundary_level_4, - "horizontal_end": self._end_cell_halo, - }, - offset_provider=self._grid.connectivities, - ) - - self._compute_advection_in_corrector_vertical_momentum = setup_program( - backend=backend, - program=compute_advection_in_corrector_vertical_momentum, - constant_args={ - "coeff1_dwdz": self._metric_state.coeff1_dwdz, - "coeff2_dwdz": self._metric_state.coeff2_dwdz, - "c_intp": self._interpolation_state.c_intp, - "inv_dual_edge_length": self._edge_params.inverse_dual_edge_lengths, - "inv_primal_edge_length": self._edge_params.inverse_primal_edge_lengths, - "tangent_orientation": self._edge_params.tangent_orientation, - "e_bln_c_s": self._interpolation_state.e_bln_c_s, - "ddqz_z_half": self._metric_state.ddqz_z_half, - "geofac_n2s": self._interpolation_state.geofac_n2s, - "owner_mask": self._c_owner_mask, - }, - horizontal_sizes={ - "horizontal_start": self._start_cell_lateral_boundary_level_4, - "horizontal_end": self._end_cell_halo, - }, - vertical_sizes={ - "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer, - "vertical_start": gtx.int32(0), - "vertical_end": self._grid.num_levels, - }, - offset_provider=self._grid.connectivities, - ) - - self._compute_advection_in_horizontal_momentum = setup_program( - backend=backend, - program=compute_advection_in_horizontal_momentum, - constant_args={ - "e_bln_c_s": self._interpolation_state.e_bln_c_s, - "geofac_rot": self._interpolation_state.geofac_rot, - "coeff_gradekin": self._metric_state.coeff_gradekin, - "c_lin_e": self._interpolation_state.c_lin_e, - "ddqz_z_full_e": self._metric_state.ddqz_z_full_e, - "area_edge": self._edge_params.edge_areas, - "tangent_orientation": self._edge_params.tangent_orientation, - "inv_primal_edge_length": self._edge_params.inverse_primal_edge_lengths, - "geofac_grdiv": self._interpolation_state.geofac_grdiv, - "coriolis_frequency": self._edge_params.coriolis_frequency, - }, - variants={ - "apply_extra_diffusion_on_vn": [False, True], - }, - horizontal_sizes={ - "horizontal_start": self._start_edge_nudging_level_2, - "horizontal_end": self._end_edge_local, - }, - vertical_sizes={ - "end_index_of_damping_layer": self._vertical_params.end_index_of_damping_layer, - "vertical_start": gtx.int32(0), - "vertical_end": self._grid.num_levels, - }, - offset_provider=self._grid.connectivities, - ) - - def _allocate_local_fields(self, allocator: gtx_typing.Allocator | None) -> None: - self._horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KHalfDim, allocator=allocator, dtype=ta.vpfloat - ) - """ - Declared as z_v_grad_w in ICON. vn dw/dn + vt dw/dt. The bottom half level is never - computed because w[nlevp1-1] is diagnostic. - """ - - self._contravariant_corrected_w_at_cells_on_model_levels = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.vpfloat - ) - """ - Declared as z_w_con_c_full in ICON. w - (vn dz/dn + vt dz/dt), z is topography height - """ - - self._vertical_cfl = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KHalfDim, allocator=allocator, dtype=ta.vpfloat - ) - - def _determine_local_domains(self) -> None: - vertex_domain = h_grid.domain(dims.VertexDim) - edge_domain = h_grid.domain(dims.EdgeDim) - cell_domain = h_grid.domain(dims.CellDim) - self._start_vertex_lateral_boundary_level_2 = self._grid.start_index( - vertex_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) - ) - self._end_vertex_halo = self._grid.end_index(vertex_domain(h_grid.Zone.HALO)) - - self._start_edge_lateral_boundary_level_5 = self._grid.start_index( - edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) - ) - self._start_edge_lateral_boundary_level_7 = self._grid.start_index( - edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7) - ) - self._start_edge_nudging_level_2 = self._grid.start_index( - edge_domain(h_grid.Zone.NUDGING_LEVEL_2) - ) - - self._end_edge_local = self._grid.end_index(edge_domain(h_grid.Zone.LOCAL)) - self._end_edge_halo = self._grid.end_index(edge_domain(h_grid.Zone.HALO)) - self._end_edge_halo_level_2 = self._grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) - - self._start_cell_lateral_boundary_level_3 = self._grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_3) - ) - self._start_cell_lateral_boundary_level_4 = self._grid.start_index( - cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4) - ) - self._start_cell_nudging = self._grid.start_index(cell_domain(h_grid.Zone.NUDGING)) - self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) - self._end_cell_halo = self._grid.end_index(cell_domain(h_grid.Zone.HALO)) - - def run_predictor_step( - self, - *, - skip_compute_predictor_vertical_advection: bool, - diagnostic_state: nonhydro_states.DiagnosticStateNonHydro, - prognostic_state: prognostics.PrognosticState, - contravariant_correction_at_edges_on_model_levels: fa.EdgeKField[ta.anyfloat], - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.anyfloat], - tangential_wind_on_half_levels: fa.EdgeKField[ta.anyfloat], - dtime: ta.wpfloat, - cell_areas: fa.CellField[ta.wpfloat], - ) -> None: - """ - Compute some diagnostic variables that are used in the predictor step - of the dycore and advective tendency of normal and vertical winds. - - Args: - skip_compute_predictor_vertical_advection: Option to skip computation of advective tendency of vertical wind - diagnostic_state: DiagnosticStateNonHydro class - prognostic_state: PrognosticState class - contravariant_correction_at_edges_on_model_levels: Contravariant corrected vertical wind at edge [m s-1] - horizontal_kinetic_energy_at_edges_on_model_levels: Horizontal kinetic energy at edge [m^2 s-2] - tangential_wind_on_half_levels: tangential wind at edge on k-half levels [m s-1] - dtime: time step [m s-1] - cell_areas: cell area [m^2] - """ - - cfl_w_limit, scalfac_exdiff = self._scale_factors_by_dtime(dtime) - - self._compute_diagnostics_from_normal_wind( - tangential_wind=diagnostic_state.tangential_wind, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=diagnostic_state.vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels=self._horizontal_advection_of_w_at_edges_on_half_levels, - vn=prognostic_state.vn, - w=prognostic_state.w, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - ) - - # TODO(havogt): however, our test data is probably not able to catch cfl_clipping conditons - self._compute_advection_in_predictor_vertical_momentum( - contravariant_correction_at_cells_on_half_levels=diagnostic_state.contravariant_correction_at_cells_on_half_levels, - vertical_wind_advective_tendency=diagnostic_state.vertical_wind_advective_tendency.predictor, - contravariant_corrected_w_at_cells_on_model_levels=self._contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl=self._vertical_cfl, - w=prognostic_state.w, - horizontal_advection_of_w_at_edges_on_half_levels=self._horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - area=cell_areas, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - ) - - # Reductions should be performed on flat, contiguous arrays for best cupy performance - # as otherwise cupy won't use cub optimized kernels. - max_vertical_cfl = self._vertical_cfl.array_ns.max( # type: ignore[attr-defined] - self._vertical_cfl.ndarray[ - self._start_cell_lateral_boundary_level_4 : self._end_cell_halo, : - ].ravel(order="K") # type: ignore[attr-defined] - ) - diagnostic_state.max_vertical_cfl = self._vertical_cfl.array_ns.maximum( # type: ignore[attr-defined] - max_vertical_cfl, diagnostic_state.max_vertical_cfl - ) - - # Note, if we compute `apply_extra_diffusion_on_vn = max_vertical_cfl > cfl_w_limit * dtime` here, - # we would have to synchronize with the device already here to get the value of `max_vertical_cfl`. - # TODO (Chia Rui): to decide whether make apply_extra_diffusion_on_vn a config parameter or remove it or always turn on extra diffusion - apply_extra_diffusion_on_vn = True - self._compute_advection_in_horizontal_momentum( - normal_wind_advective_tendency=diagnostic_state.normal_wind_advective_tendency.predictor, - vn=prognostic_state.vn, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - tangential_wind=diagnostic_state.tangential_wind, - contravariant_corrected_w_at_cells_on_model_levels=self._contravariant_corrected_w_at_cells_on_model_levels, - vn_on_half_levels=diagnostic_state.vn_on_half_levels, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, - ) - - def _scale_factors_by_dtime(self, dtime: float) -> tuple[float, float]: - scaled_cfl_w_limit = self._cfl_w_limit / dtime - scalfac_exdiff = self._scalfac_exdiff / (dtime * (0.85 - scaled_cfl_w_limit * dtime)) - return scaled_cfl_w_limit, scalfac_exdiff - - def run_corrector_step( - self, - *, - diagnostic_state: nonhydro_states.DiagnosticStateNonHydro, - prognostic_state: prognostics.PrognosticState, - horizontal_kinetic_energy_at_edges_on_model_levels: fa.EdgeKField[ta.anyfloat], - tangential_wind_on_half_levels: fa.EdgeKField[ta.anyfloat], - dtime: ta.wpfloat, - cell_areas: fa.CellField[ta.wpfloat], - ) -> None: - """ - Compute some diagnostic variables that are used in the corrector step - of the dycore and advective tendency of normal and vertical winds. - - Args: - diagnostic_state: DiagnosticStateNonHydro class - prognostic_state: PrognosticState class - horizontal_kinetic_energy_at_edges_on_model_levels: Horizontal kinetic energy at edge [m^2 s-2] - tangential_wind_on_half_levels: tangential wind at edge on k-half levels [m s-1] - dtime: time step [m s-1] - cell_areas: cell area [m^2] - """ - - cfl_w_limit, scalfac_exdiff = self._scale_factors_by_dtime(dtime) - - self._compute_advection_in_corrector_vertical_momentum( - vertical_wind_advective_tendency=diagnostic_state.vertical_wind_advective_tendency.corrector, - contravariant_corrected_w_at_cells_on_model_levels=self._contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl=self._vertical_cfl, - w=prognostic_state.w, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=diagnostic_state.vn_on_half_levels, - contravariant_correction_at_cells_on_half_levels=diagnostic_state.contravariant_correction_at_cells_on_half_levels, - area=cell_areas, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - ) - - # Reductions should be performed on flat, contiguous arrays for best cupy performance - # as otherwise cupy won't use cub optimized kernels. - max_vertical_cfl = self._vertical_cfl.array_ns.max( # type: ignore[attr-defined] - self._vertical_cfl.ndarray[ - self._start_cell_lateral_boundary_level_4 : self._end_cell_halo, : - ].ravel(order="K") # type: ignore[attr-defined] - ) - - diagnostic_state.max_vertical_cfl = self._vertical_cfl.array_ns.maximum( # type: ignore[attr-defined] - max_vertical_cfl, diagnostic_state.max_vertical_cfl - ) - - # Note, if we compute `apply_extra_diffusion_on_vn = max_vertical_cfl > cfl_w_limit * dtime` here, - # we would have to synchronize with the device already here to get the value of `max_vertical_cfl`. - # TODO (Chia Rui): to decide whether make apply_extra_diffusion_on_vn a config parameter or remove it or always turn on extra diffusion - apply_extra_diffusion_on_vn = True - self._compute_advection_in_horizontal_momentum( - normal_wind_advective_tendency=diagnostic_state.normal_wind_advective_tendency.corrector, - vn=prognostic_state.vn, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - tangential_wind=diagnostic_state.tangential_wind, - contravariant_corrected_w_at_cells_on_model_levels=self._contravariant_corrected_w_at_cells_on_model_levels, - vn_on_half_levels=diagnostic_state.vn_on_half_levels, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, - ) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py index 7ac95b4a0a..d22c1a9d50 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_velocity_advection.py @@ -12,17 +12,12 @@ import pytest from gt4py.next import typing as gtx_typing -from icon4py.model.atmosphere.dycore import dycore_states, velocity_advection as advection -from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_horizontal_momentum_equation import ( - compute_advection_in_horizontal_momentum, -) -from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_vertical_momentum_equation import ( - compute_advection_in_corrector_vertical_momentum, - compute_advection_in_predictor_vertical_momentum, -) -from icon4py.model.atmosphere.dycore.stencils.compute_diagnostics_from_normal_wind import ( - compute_diagnostics_from_normal_wind, +from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro +from icon4py.model.atmosphere.dycore.stencils import ( + velocity_advection_corrector, + velocity_advection_predictor, ) +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import VerticalCflConstants from icon4py.model.common import dimension as dims, type_alias as ta, utils as common_utils from icon4py.model.common.grid import ( horizontal as h_grid, @@ -77,71 +72,13 @@ def create_vertical_params( (test_defs.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000"), ], ) -def test_verify_velocity_init_against_savepoint( # noqa: PLR0917 [too-many-positional-arguments] - interpolation_savepoint: serialbox.InterpolationSavepoint, - step_date_init: str, - grid_savepoint: serialbox.IconGridSavepoint, - icon_grid: icon.IconGrid, - metrics_savepoint: serialbox.MetricSavepoint, - experiment: test_defs.Experiment, - backend: gtx_typing.Backend | None, -) -> None: - interpolation_state = utils.construct_interpolation_state(interpolation_savepoint) - metric_state_nonhydro = utils.construct_metric_state(metrics_savepoint, grid_savepoint) - vertical_config = experiment.config.vertical_grid - vertical_params = create_vertical_params(vertical_config, grid_savepoint) - - velocity_advection = advection.VelocityAdvection( - grid=icon_grid, - metric_state=metric_state_nonhydro, - interpolation_state=interpolation_state, - vertical_params=vertical_params, - edge_params=grid_savepoint.construct_edge_geometry(), - owner_mask=grid_savepoint.c_owner_mask(), - backend=backend, - ) - assert velocity_advection._cfl_w_limit == 0.65 - assert velocity_advection._scalfac_exdiff == 0.05 - assert test_utils.dallclose(velocity_advection._vertical_cfl.asnumpy(), 0.0) - - -@pytest.mark.embedded_static_args -@pytest.mark.datatest -@pytest.mark.parametrize( - "experiment_description, step_date_init", - [ - (test_defs.Experiments.MCH_CH_R04B09, "2021-06-20T12:00:10.000"), - (test_defs.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000"), - ], -) -def test_scale_factors_by_dtime( # noqa: PLR0917 [too-many-positional-arguments] - interpolation_savepoint, - metrics_savepoint, - experiment, - step_date_init, - savepoint_velocity_init, - icon_grid, - grid_savepoint, - backend, -): +def test_extra_diffusion_constants_match_icon(experiment, step_date_init, savepoint_velocity_init): + # ICON serializes both constants per second, divided by dtime. dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") - interpolation_state = utils.construct_interpolation_state(interpolation_savepoint) - metric_state_nonhydro = utils.construct_metric_state(metrics_savepoint, grid_savepoint) - vertical_config = experiment.config.vertical_grid - vertical_params = create_vertical_params(vertical_config, grid_savepoint) - - velocity_advection = advection.VelocityAdvection( - grid=icon_grid, - metric_state=metric_state_nonhydro, - interpolation_state=interpolation_state, - vertical_params=vertical_params, - edge_params=grid_savepoint.construct_edge_geometry(), - owner_mask=grid_savepoint.c_owner_mask(), - backend=backend, + assert VerticalCflConstants.W_LIMIT / dtime == savepoint_velocity_init.cfl_w_limit() + assert VerticalCflConstants.EXTRA_DIFFUSION_SCALING / dtime == pytest.approx( + savepoint_velocity_init.scalfac_exdiff(), rel=1e-14 ) - (cfl_w_limit, scalfac_exdiff) = velocity_advection._scale_factors_by_dtime(dtime) - assert cfl_w_limit == savepoint_velocity_init.cfl_w_limit() - assert scalfac_exdiff == savepoint_velocity_init.scalfac_exdiff() @pytest.mark.embedded_remap_error @@ -222,25 +159,88 @@ def test_velocity_predictor_step( # noqa: PLR0917 [too-many-positional-argument vertical_config = experiment.config.vertical_grid vertical_params = create_vertical_params(vertical_config, grid_savepoint) - velocity_advection = advection.VelocityAdvection( - grid=icon_grid, - metric_state=metric_state_nonhydro, - interpolation_state=interpolation_state, - vertical_params=vertical_params, - edge_params=edge_geometry, - owner_mask=grid_savepoint.c_owner_mask(), - backend=backend, + edge_domain = h_grid.domain(dims.EdgeDim) + cell_domain = h_grid.domain(dims.CellDim) + end_cell_halo = icon_grid.end_index(cell_domain(h_grid.Zone.HALO)) + end_edge_halo_level_2 = icon_grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) + end_edge_local = icon_grid.end_index(edge_domain(h_grid.Zone.LOCAL)) + start_cell_lateral_boundary_level_4 = icon_grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4) + ) + start_edge_lateral_boundary_level_5 = icon_grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) + ) + start_edge_nudging_level_2 = icon_grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) + vertical_cfl = data_alloc.zero_field( + icon_grid, dims.CellDim, dims.KHalfDim, allocator=backend, dtype=ta.vpfloat ) - velocity_advection.run_predictor_step( - skip_compute_predictor_vertical_advection=vn_only, - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - contravariant_correction_at_edges_on_model_levels=init_savepoint.z_w_concorr_me(), - horizontal_kinetic_energy_at_edges_on_model_levels=init_savepoint.z_kin_hor_e(), - tangential_wind_on_half_levels=init_savepoint.z_vt_ie(), + contravariant_correction_at_edges_on_model_levels = init_savepoint.z_w_concorr_me() + horizontal_kinetic_energy_at_edges_on_model_levels = init_savepoint.z_kin_hor_e() + tangential_wind_on_half_levels = init_savepoint.z_vt_ie() + + velocity_advection_predictor.compute_velocity_advection_in_predictor_step.with_backend(backend)( + tangential_wind=diagnostic_state.tangential_wind, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=diagnostic_state.vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels=diagnostic_state.contravariant_correction_at_cells_on_half_levels, + vertical_wind_advective_tendency=diagnostic_state.vertical_wind_advective_tendency.predictor, + vertical_cfl=vertical_cfl, + normal_wind_advective_tendency=diagnostic_state.normal_wind_advective_tendency.predictor, + vn=prognostic_state.vn, + w=prognostic_state.w, + rbf_vec_coeff_e=interpolation_state.rbf_vec_coeff_e, + wgtfac_e=metric_state_nonhydro.wgtfac_e, + wgtfacq_e=metric_state_nonhydro.wgtfacq_e, + ddxn_z_full=metric_state_nonhydro.ddxn_z_full, + ddxt_z_full=metric_state_nonhydro.ddxt_z_full, + coeff1_dwdz=metric_state_nonhydro.coeff1_dwdz, + coeff2_dwdz=metric_state_nonhydro.coeff2_dwdz, + c_intp=interpolation_state.c_intp, + inv_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, + inv_primal_edge_length=edge_geometry.inverse_primal_edge_lengths, + tangent_orientation=edge_geometry.tangent_orientation, + e_bln_c_s=interpolation_state.e_bln_c_s, + wgtfac_c=metric_state_nonhydro.wgtfac_c, + ddqz_z_half=metric_state_nonhydro.ddqz_z_half, + area=cell_geometry.area, + geofac_n2s=interpolation_state.geofac_n2s, + owner_mask=grid_savepoint.c_owner_mask(), + coriolis_frequency=edge_geometry.coriolis_frequency, + geofac_rot=interpolation_state.geofac_rot, + coeff_gradekin=metric_state_nonhydro.coeff_gradekin, + c_lin_e=interpolation_state.c_lin_e, + ddqz_z_full_e=metric_state_nonhydro.ddqz_z_full_e, + area_edge=edge_geometry.edge_areas, + geofac_grdiv=interpolation_state.geofac_grdiv, dtime=dtime, - cell_areas=cell_geometry.area, + skip_compute_predictor_vertical_advection=vn_only, + apply_extra_diffusion_on_vn=True, + nflatlev=vertical_params.nflatlev, + end_index_of_damping_layer=vertical_params.end_index_of_damping_layer, + start_edge_lateral_boundary_level_5=start_edge_lateral_boundary_level_5, + end_edge_halo_level_2=end_edge_halo_level_2, + start_cell_lateral_boundary_level_4=start_cell_lateral_boundary_level_4, + end_cell_halo=end_cell_halo, + start_edge_nudging_level_2=start_edge_nudging_level_2, + end_edge_local=end_edge_local, + vertical_start=gtx.int32(0), + vertical_end=icon_grid.num_levels, + offset_provider={ + "C2E": icon_grid.get_connectivity("C2E"), + "C2E2CO": icon_grid.get_connectivity("C2E2CO"), + "E2C": icon_grid.get_connectivity("E2C"), + "E2C2E": icon_grid.get_connectivity("E2C2E"), + "E2C2EO": icon_grid.get_connectivity("E2C2EO"), + "E2V": icon_grid.get_connectivity("E2V"), + "V2C": icon_grid.get_connectivity("V2C"), + "V2E": icon_grid.get_connectivity("V2E"), + }, + ) + solve_nonhydro._update_max_vertical_cfl( + diagnostic_state, vertical_cfl, start_cell_lateral_boundary_level_4, end_cell_halo ) icon_result_ddt_vn_apc_pc = savepoint_velocity_exit.ddt_vn_apc_pc(0).asnumpy() @@ -284,8 +284,43 @@ def test_velocity_predictor_step( # noqa: PLR0917 [too-many-positional-argument atol=1.0e-15, ) + # ICON sets z_vt_ie on the top half level unconditionally + # (mo_velocity_advection.f90:300) whereas icon4py leaves the whole field untouched + # when the predictor's vertical advection is skipped. + first_comparable_half_level = 1 if vn_only else 0 + assert test_utils.dallclose( + tangential_wind_on_half_levels.asnumpy()[:, first_comparable_half_level:], + savepoint_velocity_exit.z_vt_ie().asnumpy()[:, first_comparable_half_level:], + rtol=1.0e-14, + atol=1.0e-14, + ) + + assert test_utils.dallclose( + horizontal_kinetic_energy_at_edges_on_model_levels.asnumpy(), + savepoint_velocity_exit.z_kin_hor_e().asnumpy(), + rtol=1.0e-14, + atol=1.0e-14, + ) + + assert test_utils.dallclose( + contravariant_correction_at_edges_on_model_levels.asnumpy(), + savepoint_velocity_exit.z_w_concorr_me().asnumpy(), + rtol=1.0e-15, + atol=1.0e-15, + ) + assert diagnostic_state.max_vertical_cfl == icon_result_max_vcfl_dyn + _compare_cfl( + vertical_cfl=vertical_cfl.asnumpy(), + icon_result_cfl_clipping=savepoint_velocity_exit.cfl_clipping().asnumpy(), + icon_result_max_vcfl_dyn=icon_result_max_vcfl_dyn, + horizontal_start=start_cell_lateral_boundary_level_4, + horizontal_end=end_cell_halo, + vertical_start=max(2, grid_savepoint.nrdmax() - 2), + vertical_end=icon_grid.num_levels - 3, + ) + @pytest.mark.embedded_remap_error @pytest.mark.datatest @@ -370,23 +405,68 @@ def test_velocity_corrector_step( # noqa: PLR0917 [too-many-positional-argument vertical_config = experiment.config.vertical_grid vertical_params = create_vertical_params(vertical_config, grid_savepoint) - velocity_advection = advection.VelocityAdvection( - grid=icon_grid, - metric_state=metric_state_nonhydro, - interpolation_state=interpolation_state, - vertical_params=vertical_params, - edge_params=edge_geometry, - owner_mask=grid_savepoint.c_owner_mask(), - backend=backend, + edge_domain = h_grid.domain(dims.EdgeDim) + cell_domain = h_grid.domain(dims.CellDim) + end_cell_halo = icon_grid.end_index(cell_domain(h_grid.Zone.HALO)) + end_edge_local = icon_grid.end_index(edge_domain(h_grid.Zone.LOCAL)) + start_cell_lateral_boundary_level_4 = icon_grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4) + ) + start_edge_nudging_level_2 = icon_grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) + vertical_cfl = data_alloc.zero_field( + icon_grid, dims.CellDim, dims.KHalfDim, allocator=backend, dtype=ta.vpfloat ) - velocity_advection.run_corrector_step( - diagnostic_state=diagnostic_state, - prognostic_state=prognostic_state, - horizontal_kinetic_energy_at_edges_on_model_levels=init_savepoint.z_kin_hor_e(), + velocity_advection_corrector.compute_velocity_advection_in_corrector_step.with_backend(backend)( + vertical_wind_advective_tendency=diagnostic_state.vertical_wind_advective_tendency.corrector, + vertical_cfl=vertical_cfl, + normal_wind_advective_tendency=diagnostic_state.normal_wind_advective_tendency.corrector, + vn=prognostic_state.vn, + w=prognostic_state.w, + tangential_wind=diagnostic_state.tangential_wind, tangential_wind_on_half_levels=init_savepoint.z_vt_ie(), + vn_on_half_levels=diagnostic_state.vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=init_savepoint.z_kin_hor_e(), + contravariant_correction_at_cells_on_half_levels=diagnostic_state.contravariant_correction_at_cells_on_half_levels, + coeff1_dwdz=metric_state_nonhydro.coeff1_dwdz, + coeff2_dwdz=metric_state_nonhydro.coeff2_dwdz, + c_intp=interpolation_state.c_intp, + inv_dual_edge_length=edge_geometry.inverse_dual_edge_lengths, + inv_primal_edge_length=edge_geometry.inverse_primal_edge_lengths, + tangent_orientation=edge_geometry.tangent_orientation, + e_bln_c_s=interpolation_state.e_bln_c_s, + ddqz_z_half=metric_state_nonhydro.ddqz_z_half, + area=cell_geometry.area, + geofac_n2s=interpolation_state.geofac_n2s, + owner_mask=grid_savepoint.c_owner_mask(), + coriolis_frequency=edge_geometry.coriolis_frequency, + geofac_rot=interpolation_state.geofac_rot, + coeff_gradekin=metric_state_nonhydro.coeff_gradekin, + c_lin_e=interpolation_state.c_lin_e, + ddqz_z_full_e=metric_state_nonhydro.ddqz_z_full_e, + area_edge=edge_geometry.edge_areas, + geofac_grdiv=interpolation_state.geofac_grdiv, dtime=dtime, - cell_areas=cell_geometry.area, + apply_extra_diffusion_on_vn=True, + end_index_of_damping_layer=vertical_params.end_index_of_damping_layer, + start_cell_lateral_boundary_level_4=start_cell_lateral_boundary_level_4, + end_cell_halo=end_cell_halo, + start_edge_nudging_level_2=start_edge_nudging_level_2, + end_edge_local=end_edge_local, + vertical_start=gtx.int32(0), + vertical_end=icon_grid.num_levels, + offset_provider={ + "C2E": icon_grid.get_connectivity("C2E"), + "C2E2CO": icon_grid.get_connectivity("C2E2CO"), + "E2C": icon_grid.get_connectivity("E2C"), + "E2C2EO": icon_grid.get_connectivity("E2C2EO"), + "E2V": icon_grid.get_connectivity("E2V"), + "V2C": icon_grid.get_connectivity("V2C"), + "V2E": icon_grid.get_connectivity("V2E"), + }, + ) + solve_nonhydro._update_max_vertical_cfl( + diagnostic_state, vertical_cfl, start_cell_lateral_boundary_level_4, end_cell_halo ) icon_result_ddt_vn_apc_pc = savepoint_velocity_exit.ddt_vn_apc_pc(1).asnumpy() @@ -409,527 +489,12 @@ def test_velocity_corrector_step( # noqa: PLR0917 [too-many-positional-argument assert diagnostic_state.max_vertical_cfl == icon_result_max_vcfl_dyn - -@pytest.mark.datatest -@pytest.mark.embedded_remap_error -@pytest.mark.parametrize( - "experiment_description, step_date_init, step_date_exit", - [ - ( - test_defs.Experiments.MCH_CH_R04B09, - "2021-06-20T12:00:10.000", - "2021-06-20T12:00:10.000", - ), - (test_defs.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), - ], -) -def test_compute_diagnostics_from_normal_wind( # noqa: PLR0917 [too-many-positional-arguments] - experiment, - step_date_init, - step_date_exit, - icon_grid, - grid_savepoint, - interpolation_savepoint, - metrics_savepoint, - savepoint_velocity_init, - savepoint_velocity_exit, - backend, -): - edge_domain = h_grid.domain(dims.EdgeDim) - - tangential_wind_on_half_levels = savepoint_velocity_init.z_vt_ie() - tangential_wind = savepoint_velocity_init.vt() - vn_on_half_levels = savepoint_velocity_init.vn_ie() - horizontal_kinetic_energy_at_edges_on_model_levels = savepoint_velocity_init.z_kin_hor_e() - horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( - icon_grid, dims.EdgeDim, dims.KHalfDim, allocator=backend - ) - vn = savepoint_velocity_init.vn() - w = savepoint_velocity_init.w() - - rbf_vec_coeff_e = interpolation_savepoint.rbf_vec_coeff_e() - wgtfac_e = metrics_savepoint.wgtfac_e() - ddxn_z_full = metrics_savepoint.ddxn_z_full() - ddxt_z_full = metrics_savepoint.ddxt_z_full() - contravariant_correction_at_edges_on_model_levels = savepoint_velocity_init.z_w_concorr_me() - wgtfacq_e = metrics_savepoint.wgtfacq_e() - nflatlev = grid_savepoint.nflatlev() - c_intp = interpolation_savepoint.c_intp() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - tangent_orientation = grid_savepoint.tangent_orientation() - - skip_compute_predictor_vertical_advection = savepoint_velocity_init.vn_only() - # TODO(havogt): we need a test where skip_compute_predictor_vertical_advection is True! - - horizontal_start = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) - horizontal_end = icon_grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) - - icon_result_vt = savepoint_velocity_exit.vt() - icon_result_z_vt_ie = savepoint_velocity_exit.z_vt_ie() - icon_result_vn_ie = savepoint_velocity_exit.vn_ie() - icon_result_z_kin_hor_e = savepoint_velocity_exit.z_kin_hor_e() - icon_result_z_w_concorr_me = savepoint_velocity_exit.z_w_concorr_me() - icon_result_z_v_grad_w = savepoint_velocity_exit.z_v_grad_w() - - compute_diagnostics_from_normal_wind.with_backend(backend)( - tangential_wind=tangential_wind, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - vn=vn, - w=w, - rbf_vec_coeff_e=rbf_vec_coeff_e, - wgtfac_e=wgtfac_e, - ddxn_z_full=ddxn_z_full, - ddxt_z_full=ddxt_z_full, - wgtfacq_e=wgtfacq_e, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=gtx.int32(nflatlev), - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=gtx.int32(0), - vertical_end=gtx.int32(icon_grid.num_levels + 1), - offset_provider={ - "E2C": icon_grid.get_connectivity("E2C"), - "E2V": icon_grid.get_connectivity("E2V"), - "V2C": icon_grid.get_connectivity("V2C"), - "E2C2E": icon_grid.get_connectivity("E2C2E"), - }, - ) - - assert test_utils.dallclose( - icon_result_vt.asnumpy(), tangential_wind.asnumpy(), rtol=1.0e-14, atol=1.0e-14 - ) - assert test_utils.dallclose( - icon_result_z_vt_ie.asnumpy(), - tangential_wind_on_half_levels.asnumpy(), - rtol=1.0e-14, - atol=1.0e-14, - ) - assert test_utils.dallclose( - icon_result_vn_ie.asnumpy(), vn_on_half_levels.asnumpy(), rtol=1.0e-15, atol=1.0e-15 - ) - assert test_utils.dallclose( - icon_result_z_kin_hor_e.asnumpy(), - horizontal_kinetic_energy_at_edges_on_model_levels.asnumpy(), - rtol=1.0e-14, - atol=1.0e-14, - ) - assert test_utils.dallclose( - icon_result_z_w_concorr_me.asnumpy(), - contravariant_correction_at_edges_on_model_levels.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, - ) - # the restriction is ok, as this is a velocity advection temporary - lateral_boundary_7 = icon_grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7)) - halo_1 = icon_grid.end_index(edge_domain(h_grid.Zone.HALO)) - # ICON's z_v_grad_w stores only nlev levels, so the bottom half level has no reference - assert test_utils.dallclose( - icon_result_z_v_grad_w.asnumpy()[lateral_boundary_7:halo_1, :], - horizontal_advection_of_w_at_edges_on_half_levels.asnumpy()[lateral_boundary_7:halo_1, :-1], - rtol=1.0e-15, - atol=1.0e-15, - ) - - -@pytest.mark.datatest -@pytest.mark.uses_concat_where -@pytest.mark.parametrize( - "experiment_description, step_date_init, step_date_exit", - [ - ( - test_defs.Experiments.MCH_CH_R04B09, - "2021-06-20T12:00:10.000", - "2021-06-20T12:00:10.000", - ), - ( - test_defs.Experiments.MCH_CH_R04B09, - "2021-06-20T12:00:20.000", - "2021-06-20T12:00:20.000", - ), - (test_defs.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), - ], -) -@pytest.mark.parametrize("istep_init, istep_exit", [(1, 1)]) -def test_compute_advection_in_predictor_vertical_momentum( # noqa: PLR0917 [too-many-positional-arguments] - experiment, - step_date_init, - step_date_exit, - istep_init, - istep_exit, - icon_grid, - grid_savepoint, - interpolation_savepoint, - metrics_savepoint, - savepoint_velocity_exit, - backend, - savepoint_velocity_init, -): - scalfac_exdiff = savepoint_velocity_init.scalfac_exdiff() - cfl_w_limit = savepoint_velocity_init.cfl_w_limit() - ddqz_z_half = metrics_savepoint.ddqz_z_half() - contravariant_correction_at_edges_on_model_levels = savepoint_velocity_exit.z_w_concorr_me() - contravariant_correction_at_cells_on_half_levels = savepoint_velocity_init.w_concorr_c() - w = savepoint_velocity_init.w() - horizontal_advection_of_w_at_edges_on_half_levels = savepoint_velocity_exit.z_v_grad_w() - vertical_wind_advective_tendency = savepoint_velocity_init.ddt_w_adv_pc(istep_init - 1) - contravariant_corrected_w_at_cells_on_model_levels = savepoint_velocity_init.z_w_con_c_full() - vertical_cfl = data_alloc.zero_field( - icon_grid, dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat, allocator=backend - ) - skip_compute_predictor_vertical_advection = savepoint_velocity_init.lvn_only() - - coeff1_dwdz = metrics_savepoint.coeff1_dwdz() - coeff2_dwdz = metrics_savepoint.coeff2_dwdz() - e_bln_c_s = interpolation_savepoint.e_bln_c_s() - wgtfac_c = metrics_savepoint.wgtfac_c() - owner_mask = grid_savepoint.c_owner_mask() - area = grid_savepoint.cell_areas() - geofac_n2s = interpolation_savepoint.geofac_n2s() - - icon_result_z_w_con_c_full = savepoint_velocity_exit.z_w_con_c_full() - icon_result_ddt_w_adv = savepoint_velocity_exit.ddt_w_adv_pc(istep_exit - 1) - icon_result_w_concorr_c = savepoint_velocity_exit.w_concorr_c() - icon_result_cfl_clipping = savepoint_velocity_exit.cfl_clipping() - icon_result_max_vcfl_dyn = savepoint_velocity_exit.max_vcfl_dyn() - - end_index_of_damping_layer = grid_savepoint.nrdmax() - - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") - cell_domain = h_grid.domain(dims.CellDim) - start_cell_nudging_for_vertical_wind_advective_tendency = icon_grid.start_index( - cell_domain(h_grid.Zone.NUDGING) - ) - end_cell_local_for_vertical_wind_advective_tendency = icon_grid.end_index( - cell_domain(h_grid.Zone.LOCAL) - ) - horizontal_start = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) - horizontal_end = icon_grid.end_index(cell_domain(h_grid.Zone.HALO)) - vertical_start = 0 - vertical_end = icon_grid.num_levels - compute_advection_in_predictor_vertical_momentum.with_backend(backend)( - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl=vertical_cfl, - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=grid_savepoint.nflatlev(), - end_index_of_damping_layer=end_index_of_damping_layer, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={ - "C2E": icon_grid.get_connectivity("C2E"), - "C2E2CO": icon_grid.get_connectivity("C2E2CO"), - "V2C": icon_grid.get_connectivity("V2C"), - "E2C": icon_grid.get_connectivity("E2C"), - "E2V": icon_grid.get_connectivity("E2V"), - }, - ) - - assert test_utils.dallclose( - icon_result_w_concorr_c.asnumpy(), - contravariant_correction_at_cells_on_half_levels.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, - ) - - assert test_utils.dallclose( - icon_result_z_w_con_c_full.asnumpy(), - contravariant_corrected_w_at_cells_on_model_levels.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, - ) - assert test_utils.dallclose( - icon_result_ddt_w_adv.asnumpy()[ - start_cell_nudging_for_vertical_wind_advective_tendency:end_cell_local_for_vertical_wind_advective_tendency, - :, - ], - vertical_wind_advective_tendency.asnumpy()[ - start_cell_nudging_for_vertical_wind_advective_tendency:end_cell_local_for_vertical_wind_advective_tendency, - :, - ], - rtol=1.0e-15, - atol=1.0e-15, - ) - - # TODO(OngChia): currently direct comparison of vcfl_dsl is not possible because it is not properly updated in icon run _compare_cfl( vertical_cfl=vertical_cfl.asnumpy(), - icon_result_cfl_clipping=icon_result_cfl_clipping.asnumpy(), + icon_result_cfl_clipping=savepoint_velocity_exit.cfl_clipping().asnumpy(), icon_result_max_vcfl_dyn=icon_result_max_vcfl_dyn, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=max(2, end_index_of_damping_layer - 2), + horizontal_start=start_cell_lateral_boundary_level_4, + horizontal_end=end_cell_halo, + vertical_start=max(2, grid_savepoint.nrdmax() - 2), vertical_end=icon_grid.num_levels - 3, ) - - -@pytest.mark.datatest -@pytest.mark.embedded_remap_error -@pytest.mark.parametrize( - "experiment_description, step_date_init, step_date_exit", - [ - ( - test_defs.Experiments.MCH_CH_R04B09, - "2021-06-20T12:00:10.000", - "2021-06-20T12:00:10.000", - ), - ( - test_defs.Experiments.MCH_CH_R04B09, - "2021-06-20T12:00:20.000", - "2021-06-20T12:00:20.000", - ), - (test_defs.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), - ], -) -@pytest.mark.parametrize("istep_init, istep_exit", [(2, 2)]) -def test_compute_advection_in_corrector_vertical_momentum( # noqa: PLR0917 [too-many-positional-arguments] - experiment, - step_date_init, - step_date_exit, - istep_init, - istep_exit, - icon_grid, - grid_savepoint, - interpolation_savepoint, - metrics_savepoint, - savepoint_velocity_exit, - savepoint_velocity_init, - backend, -): - scalfac_exdiff = savepoint_velocity_init.scalfac_exdiff() - cfl_w_limit = savepoint_velocity_init.cfl_w_limit() - ddqz_z_half = metrics_savepoint.ddqz_z_half() - contravariant_correction_at_cells_on_half_levels = savepoint_velocity_exit.w_concorr_c() - w = savepoint_velocity_init.w() - tangential_wind_on_half_levels = savepoint_velocity_exit.z_vt_ie() - vn_on_half_levels = savepoint_velocity_exit.vn_ie() - vertical_wind_advective_tendency = savepoint_velocity_init.ddt_w_adv_pc(istep_init - 1) - contravariant_corrected_w_at_cells_on_model_levels = savepoint_velocity_init.z_w_con_c_full() - vertical_cfl = data_alloc.zero_field( - icon_grid, dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat, allocator=backend - ) - - coeff1_dwdz = metrics_savepoint.coeff1_dwdz() - coeff2_dwdz = metrics_savepoint.coeff2_dwdz() - c_intp = interpolation_savepoint.c_intp() - inv_dual_edge_length = grid_savepoint.inv_dual_edge_length() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - tangent_orientation = grid_savepoint.tangent_orientation() - e_bln_c_s = interpolation_savepoint.e_bln_c_s() - owner_mask = grid_savepoint.c_owner_mask() - area = grid_savepoint.cell_areas() - geofac_n2s = interpolation_savepoint.geofac_n2s() - - icon_result_z_w_con_c_full = savepoint_velocity_exit.z_w_con_c_full() - icon_result_ddt_w_adv = savepoint_velocity_exit.ddt_w_adv_pc(istep_exit - 1) - icon_result_cfl_clipping = savepoint_velocity_exit.cfl_clipping() - icon_result_max_vcfl_dyn = savepoint_velocity_exit.max_vcfl_dyn() - - end_index_of_damping_layer = grid_savepoint.nrdmax() - - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") - cell_domain = h_grid.domain(dims.CellDim) - start_cell_nudging_for_vertical_wind_advective_tendency = icon_grid.start_index( - cell_domain(h_grid.Zone.NUDGING) - ) - end_cell_local_for_vertical_wind_advective_tendency = icon_grid.end_index( - cell_domain(h_grid.Zone.LOCAL) - ) - horizontal_start = icon_grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) - horizontal_end = icon_grid.end_index(cell_domain(h_grid.Zone.HALO)) - vertical_start = 0 - vertical_end = icon_grid.num_levels - compute_advection_in_corrector_vertical_momentum.with_backend(backend)( - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl=vertical_cfl, - w=w, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - end_index_of_damping_layer=end_index_of_damping_layer, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - offset_provider={ - "C2E": icon_grid.get_connectivity("C2E"), - "C2E2CO": icon_grid.get_connectivity("C2E2CO"), - "V2C": icon_grid.get_connectivity("V2C"), - "E2C": icon_grid.get_connectivity("E2C"), - "E2V": icon_grid.get_connectivity("E2V"), - }, - ) - - assert test_utils.dallclose( - icon_result_z_w_con_c_full.asnumpy(), - contravariant_corrected_w_at_cells_on_model_levels.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, - ) - assert test_utils.dallclose( - icon_result_ddt_w_adv.asnumpy()[ - start_cell_nudging_for_vertical_wind_advective_tendency:end_cell_local_for_vertical_wind_advective_tendency, - :, - ], - vertical_wind_advective_tendency.asnumpy()[ - start_cell_nudging_for_vertical_wind_advective_tendency:end_cell_local_for_vertical_wind_advective_tendency, - :, - ], - rtol=1.0e-15, - atol=1.0e-15, - ) - - # TODO(OngChia): currently direct comparison of vcfl_dsl is not possible because it is not properly updated in icon run - _compare_cfl( - vertical_cfl=vertical_cfl.asnumpy(), - icon_result_cfl_clipping=icon_result_cfl_clipping.asnumpy(), - icon_result_max_vcfl_dyn=icon_result_max_vcfl_dyn, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=max(2, end_index_of_damping_layer - 2), - vertical_end=icon_grid.num_levels - 3, - ) - - -@pytest.mark.datatest -@pytest.mark.embedded_remap_error -@pytest.mark.parametrize( - "experiment_description, step_date_init, step_date_exit", - [ - ( - test_defs.Experiments.MCH_CH_R04B09, - "2021-06-20T12:00:10.000", - "2021-06-20T12:00:10.000", - ), - (test_defs.Experiments.EXCLAIM_APE, "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000"), - ], -) -@pytest.mark.parametrize("istep_init, istep_exit", [(1, 1), (2, 2)]) -def test_compute_advection_in_horizontal_momentum( # noqa: PLR0917 [too-many-positional-arguments] - experiment, - step_date_init, - step_date_exit, - istep_init, - istep_exit, - icon_grid, - grid_savepoint, - interpolation_savepoint, - metrics_savepoint, - backend, - savepoint_velocity_init, - savepoint_velocity_exit, -): - vn = savepoint_velocity_init.vn() - horizontal_kinetic_energy_at_edges_on_model_levels = savepoint_velocity_exit.z_kin_hor_e() - tangential_wind = savepoint_velocity_exit.vt() - contravariant_corrected_w_at_cells_on_model_levels = savepoint_velocity_exit.z_w_con_c_full() - vn_on_half_levels = savepoint_velocity_exit.vn_ie() - normal_wind_advective_tendency = savepoint_velocity_init.ddt_vn_apc_pc(istep_init - 1) - - e_bln_c_s = interpolation_savepoint.e_bln_c_s() - geofac_rot = interpolation_savepoint.geofac_rot() - coeff_gradekin = metrics_savepoint.coeff_gradekin() - coriolis_frequency = grid_savepoint.f_e() - c_lin_e = interpolation_savepoint.c_lin_e() - ddqz_z_full_e = metrics_savepoint.ddqz_z_full_e() - area_edge = grid_savepoint.edge_areas() - tangent_orientation = grid_savepoint.tangent_orientation() - inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths() - geofac_grdiv = interpolation_savepoint.geofac_grdiv() - - edge_domain = h_grid.domain(dims.EdgeDim) - - start_edge_nudging_level_2 = icon_grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) - end_edge_local = icon_grid.end_index(edge_domain(h_grid.Zone.LOCAL)) - - dtime = savepoint_velocity_init.get_metadata("dtime").get("dtime") - end_index_of_damping_layer = grid_savepoint.nrdmax() - - icon_result_ddt_vn_apc = savepoint_velocity_exit.ddt_vn_apc_pc(istep_exit - 1) - - scalfac_exdiff = savepoint_velocity_init.scalfac_exdiff() - cfl_w_limit = savepoint_velocity_init.cfl_w_limit() - max_vertical_cfl = savepoint_velocity_exit.max_vcfl_dyn() - apply_extra_diffusion_on_vn = max_vertical_cfl > cfl_w_limit * dtime - - compute_advection_in_horizontal_momentum.with_backend(backend)( - normal_wind_advective_tendency=normal_wind_advective_tendency, - vn=vn, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - tangential_wind=tangential_wind, - coriolis_frequency=coriolis_frequency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vn_on_half_levels=vn_on_half_levels, - e_bln_c_s=e_bln_c_s, - geofac_rot=geofac_rot, - coeff_gradekin=coeff_gradekin, - c_lin_e=c_lin_e, - ddqz_z_full_e=ddqz_z_full_e, - area_edge=area_edge, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inv_primal_edge_length, - geofac_grdiv=geofac_grdiv, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, - end_index_of_damping_layer=end_index_of_damping_layer, - horizontal_start=start_edge_nudging_level_2, - horizontal_end=end_edge_local, - vertical_start=0, - vertical_end=icon_grid.num_levels, - offset_provider={ - "V2E": icon_grid.get_connectivity("V2E"), - "E2V": icon_grid.get_connectivity("E2V"), - "E2C": icon_grid.get_connectivity("E2C"), - "E2C2EO": icon_grid.get_connectivity("E2C2EO"), - "C2E": icon_grid.get_connectivity("C2E"), - }, - ) - - assert test_utils.dallclose( - icon_result_ddt_vn_apc.asnumpy(), - normal_wind_advective_tendency.asnumpy(), - rtol=1.0e-15, - atol=1.0e-15, - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py index 83d4e308fd..47a4ff99a6 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl.py @@ -15,6 +15,7 @@ from icon4py.model.atmosphere.dycore.stencils.add_extra_diffusion_for_normal_wind_tendency_approaching_cfl import ( add_extra_diffusion_for_normal_wind_tendency_approaching_cfl, ) +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import VerticalCflConstants from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base from icon4py.model.common.states import utils as state_utils @@ -35,8 +36,6 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( geofac_grdiv: np.ndarray, vn: np.ndarray, ddt_vn_apc: np.ndarray, - cfl_w_limit: ta.wpfloat, - scalfac_exdiff: ta.wpfloat, dtime: ta.wpfloat, ) -> np.ndarray: w_con_e = np.zeros_like(vn) @@ -66,11 +65,12 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( ) difcoef = np.where( ((levelmask_offset_0) | (levelmask_offset_1)) - & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), - scalfac_exdiff + & (np.abs(w_con_e) * dtime / ddqz_z_full_e > VerticalCflConstants.W_LIMIT), + VerticalCflConstants.EXTRA_DIFFUSION_SCALING + / dtime * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(w_con_e) * dtime / ddqz_z_full_e - cfl_w_limit * dtime, + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + np.abs(w_con_e) * dtime / ddqz_z_full_e - VerticalCflConstants.W_LIMIT, ), difcoef, ) @@ -78,7 +78,7 @@ def add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( e2c2eo = connectivities[dims.E2C2EO] ddt_vn_apc = np.where( ((levelmask_offset_0) | (levelmask_offset_1)) - & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), + & (np.abs(w_con_e) * dtime / ddqz_z_full_e > VerticalCflConstants.W_LIMIT), ddt_vn_apc + difcoef * area_edge @@ -118,8 +118,6 @@ def input_data( geofac_grdiv = data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) vn = data_alloc.random_field(dims.EdgeDim, dims.KDim) ddt_vn_apc = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - cfl_w_limit = ta.vpfloat("4.0") - scalfac_exdiff = 6.0 dtime = 2.0 return dict( levelmask=levelmask, @@ -131,8 +129,6 @@ def input_data( geofac_grdiv=geofac_grdiv, vn=vn, ddt_vn_apc=ddt_vn_apc, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, dtime=dtime, c_lin_e=c_lin_e, z_w_con_c_full=z_w_con_c_full, @@ -157,8 +153,6 @@ def reference( geofac_grdiv: np.ndarray, vn: np.ndarray, ddt_vn_apc: np.ndarray, - cfl_w_limit: ta.wpfloat, - scalfac_exdiff: ta.wpfloat, dtime: ta.wpfloat, **kwargs: Any, ) -> dict: @@ -176,8 +170,6 @@ def reference( geofac_grdiv=geofac_grdiv, vn=vn, ddt_vn_apc=ddt_vn_apc, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, dtime=dtime, ) return dict(ddt_vn_apc=ddt_vn_apc) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py deleted file mode 100644 index ef95e0dd2a..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_extra_diffusion_for_w_con_approaching_cfl.py +++ /dev/null @@ -1,145 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -import icon4py.model.common.type_alias as ta -from icon4py.model.atmosphere.dycore.stencils.add_extra_diffusion_for_w_con_approaching_cfl import ( - add_extra_diffusion_for_w_con_approaching_cfl, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -def add_extra_diffusion_for_w_con_approaching_cfl_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - cfl_clipping: np.ndarray, - owner_mask: np.ndarray, - z_w_con_c: np.ndarray, - ddqz_z_half: np.ndarray, - area: np.ndarray, - geofac_n2s: np.ndarray, - w: np.ndarray, - ddt_w_adv: np.ndarray, - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, -) -> np.ndarray: - owner_mask = np.expand_dims(owner_mask, axis=-1) - area = np.expand_dims(area, axis=-1) - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - - difcoef = np.where( - (cfl_clipping == 1) & (owner_mask == 1), - scalfac_exdiff - * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(z_w_con_c) * dtime / ddqz_z_half - cfl_w_limit * dtime, - ), - 0, - ) - - c2e2cO = connectivities[dims.C2E2CO] - ddt_w_adv = np.where( - (cfl_clipping == 1) & (owner_mask == 1), - ddt_w_adv - + difcoef - * area - * np.sum( - np.where( - (c2e2cO != -1)[:, :, np.newaxis], - w[c2e2cO] * geofac_n2s, - 0, - ), - axis=1, - ), - ddt_w_adv, - ) - return ddt_w_adv - - -@pytest.mark.embedded_remap_error -class TestAddExtraDiffusionForWConApproachingCfl(stencil_tests.StencilTest): - PROGRAM = add_extra_diffusion_for_w_con_approaching_cfl - OUTPUTS = ("ddt_w_adv",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - cfl_clipping: np.ndarray, - owner_mask: np.ndarray, - z_w_con_c: np.ndarray, - ddqz_z_half: np.ndarray, - area: np.ndarray, - geofac_n2s: np.ndarray, - w: np.ndarray, - ddt_w_adv: np.ndarray, - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - ddt_w_adv = add_extra_diffusion_for_w_con_approaching_cfl_numpy( - connectivities=connectivities, - cfl_clipping=cfl_clipping, - owner_mask=owner_mask, - z_w_con_c=z_w_con_c, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - w=w, - ddt_w_adv=ddt_w_adv, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - ) - return dict(ddt_w_adv=ddt_w_adv) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - cfl_clipping = data_alloc.random_mask(dims.CellDim, dims.KHalfDim) - owner_mask = data_alloc.random_mask(dims.CellDim) - z_w_con_c = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=vpfloat) - ddqz_z_half = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=vpfloat) - area = data_alloc.random_field(dims.CellDim, dtype=wpfloat) - geofac_n2s = data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=wpfloat) - w = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=wpfloat) - ddt_w_adv = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=vpfloat) - scalfac_exdiff = wpfloat("10.0") - cfl_w_limit = vpfloat("3.0") - dtime = wpfloat("2.0") - - return dict( - cfl_clipping=cfl_clipping, - owner_mask=owner_mask, - z_w_con_c=z_w_con_c, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - w=w, - ddt_w_adv=ddt_w_adv, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels + 1), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py deleted file mode 100644 index 19936f06df..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_add_interpolated_horizontal_advection_of_w.py +++ /dev/null @@ -1,76 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.add_interpolated_horizontal_advection_of_w import ( - add_interpolated_horizontal_advection_of_w, -) -from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.testing import stencil_tests - - -def add_interpolated_horizontal_advection_of_w_numpy( - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - e_bln_c_s: np.ndarray, - z_v_grad_w: np.ndarray, - ddt_w_adv: np.ndarray, - **kwargs: Any, -) -> np.ndarray: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - c2e = connectivities[dims.C2E] - - ddt_w_adv = ddt_w_adv + np.sum( - z_v_grad_w[c2e] * e_bln_c_s, - axis=1, - ) - return ddt_w_adv - - -class TestAddInterpolatedHorizontalAdvectionOfW(stencil_tests.StencilTest): - PROGRAM = add_interpolated_horizontal_advection_of_w - OUTPUTS = ("ddt_w_adv",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - e_bln_c_s: np.ndarray, - z_v_grad_w: np.ndarray, - ddt_w_adv: np.ndarray, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - ddt_w_adv = add_interpolated_horizontal_advection_of_w_numpy( - connectivities, e_bln_c_s, z_v_grad_w, ddt_w_adv - ) - return dict(ddt_w_adv=ddt_w_adv) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_v_grad_w = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=ta.vpfloat) - e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) - ddt_w_adv = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat) - - return dict( - e_bln_c_s=e_bln_c_s, - z_v_grad_w=z_v_grad_w, - ddt_w_adv=ddt_w_adv, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels + 1), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py deleted file mode 100644 index cc5ecc40f4..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_horizontal_momentum_equation.py +++ /dev/null @@ -1,322 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_horizontal_momentum_equation import ( - compute_advection_in_horizontal_momentum, -) -from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.states import utils as state_utils -from icon4py.model.testing import stencil_tests -from icon4py.model.testing.reference_funcs import interpolate_to_cell_center_numpy - -from .test_mo_math_divrot_rot_vertex_ri_dsl import mo_math_divrot_rot_vertex_ri_dsl_numpy - - -def _compute_advective_normal_wind_tendency_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, - coeff_gradekin: np.ndarray, - horizontal_kinetic_energy_at_cells_on_model_levels: np.ndarray, - upward_vorticity_at_vertices: np.ndarray, - tangential_wind: np.ndarray, - coriolis_frequency: np.ndarray, - c_lin_e: np.ndarray, - contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, - vn_on_half_levels: np.ndarray, - ddqz_z_full_e: np.ndarray, -) -> np.ndarray: - e2c = connectivities[dims.E2C] - horizontal_kinetic_energy_at_cells_on_model_levels_e2c = ( - horizontal_kinetic_energy_at_cells_on_model_levels[e2c] - ) - coeff_gradekin = np.expand_dims(coeff_gradekin, axis=-1) - coriolis_frequency = np.expand_dims(coriolis_frequency, axis=-1) - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - - normal_wind_advective_tendency = -( - (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) - * horizontal_kinetic_energy_at_edges_on_model_levels - + ( - -coeff_gradekin[:, 0] * horizontal_kinetic_energy_at_cells_on_model_levels_e2c[:, 0] - + coeff_gradekin[:, 1] * horizontal_kinetic_energy_at_cells_on_model_levels_e2c[:, 1] - ) - + tangential_wind - * ( - coriolis_frequency - + 0.5 * np.sum(upward_vorticity_at_vertices[connectivities[dims.E2V]], axis=1) - ) - + np.sum(contravariant_corrected_w_at_cells_on_model_levels[e2c] * c_lin_e, axis=1) - * (vn_on_half_levels[:, :-1] - vn_on_half_levels[:, 1:]) - / ddqz_z_full_e - ) - return normal_wind_advective_tendency - - -def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - c_lin_e: np.ndarray, - contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, - ddqz_z_full_e: np.ndarray, - area_edge: np.ndarray, - tangent_orientation: np.ndarray, - inv_primal_edge_length: np.ndarray, - upward_vorticity_at_vertices: np.ndarray, - geofac_grdiv: np.ndarray, - vn: np.ndarray, - normal_wind_advective_tendency: np.ndarray, - cfl_w_limit: ta.wpfloat, - scalfac_exdiff: ta.wpfloat, - dtime: ta.wpfloat, -) -> np.ndarray: - contravariant_corrected_w_at_edges_on_model_levels = np.zeros_like(vn) - difcoef = np.zeros_like(vn) - - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) - area_edge = np.expand_dims(area_edge, axis=-1) - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - - e2c = connectivities[dims.E2C] - contravariant_corrected_w_at_edges_on_model_levels = np.sum( - np.where( - (e2c != -1)[:, :, np.newaxis], - c_lin_e * contravariant_corrected_w_at_cells_on_model_levels[e2c], - 0, - ), - axis=1, - ) - - difcoef = np.where( - (np.abs(contravariant_corrected_w_at_edges_on_model_levels) > cfl_w_limit * ddqz_z_full_e), - scalfac_exdiff - * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(contravariant_corrected_w_at_edges_on_model_levels) * dtime / ddqz_z_full_e - - cfl_w_limit * dtime, - ), - difcoef, - ) - e2v = connectivities[dims.E2V] - e2c2eo = connectivities[dims.E2C2EO] - normal_wind_advective_tendency = np.where( - (np.abs(contravariant_corrected_w_at_edges_on_model_levels) > cfl_w_limit * ddqz_z_full_e), - normal_wind_advective_tendency - + difcoef - * area_edge - * ( - np.sum( - np.where( - (e2c2eo != -1)[:, :, np.newaxis], - geofac_grdiv * vn[e2c2eo], - 0, - ), - axis=1, - ) - + tangent_orientation - * inv_primal_edge_length - * (upward_vorticity_at_vertices[e2v][:, 1] - upward_vorticity_at_vertices[e2v][:, 0]) - ), - normal_wind_advective_tendency, - ) - return normal_wind_advective_tendency - - -@pytest.mark.embedded_remap_error -@pytest.mark.continuous_benchmarking -class TestFusedVelocityAdvectionStencilsHMomentum(stencil_tests.StencilTest): - PROGRAM = compute_advection_in_horizontal_momentum - OUTPUTS = ("normal_wind_advective_tendency",) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "end_index_of_damping_layer", - "vertical_start", - "vertical_end", - "apply_extra_diffusion_on_vn", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "end_index_of_damping_layer", - "vertical_start", - "vertical_end", - "apply_extra_diffusion_on_vn", - ), - } - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - normal_wind_advective_tendency: np.ndarray, - vn: np.ndarray, - horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, - tangential_wind: np.ndarray, - coriolis_frequency: np.ndarray, - contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, - vn_on_half_levels: np.ndarray, - e_bln_c_s: np.ndarray, - geofac_rot: np.ndarray, - coeff_gradekin: np.ndarray, - c_lin_e: np.ndarray, - ddqz_z_full_e: np.ndarray, - area_edge: np.ndarray, - tangent_orientation: np.ndarray, - inv_primal_edge_length: np.ndarray, - geofac_grdiv: np.ndarray, - cfl_w_limit: ta.wpfloat, - scalfac_exdiff: ta.wpfloat, - dtime: ta.wpfloat, - apply_extra_diffusion_on_vn: bool, - end_index_of_damping_layer: int, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - normal_wind_advective_tendency_cp = normal_wind_advective_tendency.copy() - nlev = kwargs["vertical_end"] - k = np.arange(nlev) - - horizontal_kinetic_energy_at_cells_on_model_levels = interpolate_to_cell_center_numpy( - connectivities, horizontal_kinetic_energy_at_edges_on_model_levels, e_bln_c_s - ) - - upward_vorticity_at_vertices = mo_math_divrot_rot_vertex_ri_dsl_numpy( - connectivities, vn, geofac_rot - ) - - normal_wind_advective_tendency = _compute_advective_normal_wind_tendency_numpy( - connectivities=connectivities, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - coeff_gradekin=coeff_gradekin, - horizontal_kinetic_energy_at_cells_on_model_levels=horizontal_kinetic_energy_at_cells_on_model_levels, - upward_vorticity_at_vertices=upward_vorticity_at_vertices, - tangential_wind=tangential_wind, - coriolis_frequency=coriolis_frequency, - c_lin_e=c_lin_e, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vn_on_half_levels=vn_on_half_levels, - ddqz_z_full_e=ddqz_z_full_e, - ) - - condition = (np.maximum(2, end_index_of_damping_layer - 2) <= k) & (k < nlev - 4) - - if apply_extra_diffusion_on_vn: - normal_wind_advective_tendency_extra_diffu = _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask_numpy( - connectivities=connectivities, - c_lin_e=c_lin_e, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - ddqz_z_full_e=ddqz_z_full_e, - area_edge=area_edge, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inv_primal_edge_length, - upward_vorticity_at_vertices=upward_vorticity_at_vertices, - geofac_grdiv=geofac_grdiv, - vn=vn, - normal_wind_advective_tendency=normal_wind_advective_tendency, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - ) - - normal_wind_advective_tendency = np.where( - condition, - normal_wind_advective_tendency_extra_diffu, - normal_wind_advective_tendency, - ) - - # restriction of execution domain - normal_wind_advective_tendency[0 : kwargs["horizontal_start"], :] = ( - normal_wind_advective_tendency_cp[0 : kwargs["horizontal_start"], :] - ) - normal_wind_advective_tendency[kwargs["horizontal_end"] :, :] = ( - normal_wind_advective_tendency_cp[kwargs["horizontal_end"] :, :] - ) - - return dict(normal_wind_advective_tendency=normal_wind_advective_tendency) - - @stencil_tests.input_data_fixture( - params=[ - {"apply_extra_diffusion_on_vn": value} for value in [True, False] - ], # True for testing, False for benchmarking - ids=lambda param: f"apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]", - ) - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, - grid: base.Grid, - request: pytest.FixtureRequest, - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - normal_wind_advective_tendency = data_alloc.zero_field(dims.EdgeDim, dims.KDim) - vn = data_alloc.random_field(dims.EdgeDim, dims.KDim) - horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( - dims.EdgeDim, dims.KDim - ) - tangential_wind = data_alloc.random_field(dims.EdgeDim, dims.KDim) - coriolis_frequency = data_alloc.random_field(dims.EdgeDim) - contravariant_corrected_w_at_cells_on_model_levels = data_alloc.random_field( - dims.CellDim, dims.KDim - ) - vn_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) - coeff_gradekin = data_alloc.random_field(dims.EdgeDim, dims.E2CDim) - e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim) - c_lin_e = data_alloc.random_field(dims.EdgeDim, dims.E2CDim) - ddqz_z_full_e = data_alloc.random_field( - dims.EdgeDim, dims.KDim, low=0.0 - ) # this makes sure that the simplified stencil produces the same result as the numpy version - area_edge = data_alloc.random_field(dims.EdgeDim) - tangent_orientation = data_alloc.random_field(dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(dims.EdgeDim) - geofac_grdiv = data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) - - geofac_rot = data_alloc.random_field(dims.VertexDim, dims.V2EDim) - scalfac_exdiff = 0.6 - dtime = 2.0 - cfl_w_limit = 0.65 / dtime - apply_extra_diffusion_on_vn = request.param["apply_extra_diffusion_on_vn"] - - end_index_of_damping_layer = 12 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime - edge_domain = h_grid.domain(dims.EdgeDim) - horizontal_start = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) - horizontal_end = grid.end_index(edge_domain(h_grid.Zone.LOCAL)) - - return dict( - normal_wind_advective_tendency=normal_wind_advective_tendency, - vn=vn, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - tangential_wind=tangential_wind, - coriolis_frequency=coriolis_frequency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vn_on_half_levels=vn_on_half_levels, - e_bln_c_s=e_bln_c_s, - geofac_rot=geofac_rot, - coeff_gradekin=coeff_gradekin, - c_lin_e=c_lin_e, - ddqz_z_full_e=ddqz_z_full_e, - area_edge=area_edge, - tangent_orientation=tangent_orientation, - inv_primal_edge_length=inv_primal_edge_length, - geofac_grdiv=geofac_grdiv, - cfl_w_limit=cfl_w_limit, - scalfac_exdiff=scalfac_exdiff, - dtime=dtime, - apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, - end_index_of_damping_layer=end_index_of_damping_layer, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py deleted file mode 100644 index 867f08d069..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_advection_in_vertical_momentum_equation.py +++ /dev/null @@ -1,742 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.compute_advection_in_vertical_momentum_equation import ( - compute_advection_in_corrector_vertical_momentum, - compute_advection_in_predictor_vertical_momentum, -) -from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.states import utils as state_utils -from icon4py.model.testing import stencil_tests -from icon4py.model.testing.reference_funcs import interpolate_to_cell_center_numpy - -from .test_add_interpolated_horizontal_advection_of_w import ( - add_interpolated_horizontal_advection_of_w_numpy, -) -from .test_compute_horizontal_advection_term_for_vertical_velocity import ( - compute_horizontal_advection_term_for_vertical_velocity_numpy, -) -from .test_interpolate_cell_field_to_half_levels import ( - interpolate_cell_field_to_half_levels_vp_numpy, -) -from .test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( - mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy, -) - - -def interpolate_contravariant_correction_to_cells_on_half_levels_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - contravariant_correction_at_cells_on_half_levels: np.ndarray, - contravariant_correction_at_edges_on_model_levels: np.ndarray, - e_bln_c_s: np.ndarray, - wgtfac_c: np.ndarray, - nflatlev: int, - nlev: int, -) -> np.ndarray: - k = np.arange(nlev) - - contravariant_correction_at_cells_model_levels = interpolate_to_cell_center_numpy( - connectivities, contravariant_correction_at_edges_on_model_levels, e_bln_c_s - ) - - condition = k >= nflatlev + 1 - contravariant_correction_at_cells_on_half_levels = np.where( - condition, - interpolate_cell_field_to_half_levels_vp_numpy( - wgtfac_c=wgtfac_c, interpolant=contravariant_correction_at_cells_model_levels - )[:, :-1], - np.zeros_like(contravariant_correction_at_cells_on_half_levels), - ) - - return contravariant_correction_at_cells_on_half_levels - - -def interpolate_contravariant_vertical_velocity_to_full_levels_numpy( - contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, -) -> np.ndarray: - num_rows, num_cols = contravariant_corrected_w_at_cells_on_half_levels.shape - contravariant_corrected_w_with_surface = np.zeros((num_rows, num_cols + 1)) - contravariant_corrected_w_with_surface[:, :-1] = ( - contravariant_corrected_w_at_cells_on_half_levels - ) - contravariant_corrected_w_at_cells_on_model_levels = 0.5 * ( - contravariant_corrected_w_with_surface[:, :-1] - + contravariant_corrected_w_with_surface[:, 1:] - ) - return contravariant_corrected_w_at_cells_on_model_levels - - -def compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( - *, - w: np.ndarray, - contravariant_correction_at_cells_on_half_levels: np.ndarray, - ddqz_z_half: np.ndarray, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, - nlev: int, - end_index_of_damping_layer: int, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - num_rows, num_cols = contravariant_correction_at_cells_on_half_levels.shape - - k = np.arange(num_cols) - condition = (np.maximum(2, end_index_of_damping_layer - 2) <= k) & (k < nlev - 3) - - contravariant_corrected_w_at_cells_on_half_levels = ( - w - contravariant_correction_at_cells_on_half_levels - ) - - cfl_clipping = np.where( - (np.abs(contravariant_corrected_w_at_cells_on_half_levels) > cfl_w_limit * ddqz_z_half) - & condition, - np.ones([num_rows, num_cols]), - np.zeros_like(contravariant_corrected_w_at_cells_on_half_levels), - ) - vertical_cfl = np.where( - cfl_clipping == 1.0, - contravariant_corrected_w_at_cells_on_half_levels * dtime / ddqz_z_half, - 0.0, - ) - contravariant_corrected_w_at_cells_on_half_levels = np.where( - (cfl_clipping == 1.0) & (vertical_cfl < -0.85), - -0.85 * ddqz_z_half / dtime, - contravariant_corrected_w_at_cells_on_half_levels, - ) - contravariant_corrected_w_at_cells_on_half_levels = np.where( - (cfl_clipping == 1.0) & (vertical_cfl > 0.85), - 0.85 * ddqz_z_half / dtime, - contravariant_corrected_w_at_cells_on_half_levels, - ) - - return contravariant_corrected_w_at_cells_on_half_levels, cfl_clipping, vertical_cfl - - -def compute_horizontal_advection_of_w( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - w: np.ndarray, - tangential_wind_on_half_levels: np.ndarray, - vn_on_half_levels: np.ndarray, - c_intp: np.ndarray, - inv_dual_edge_length: np.ndarray, - inv_primal_edge_length: np.ndarray, - tangent_orientation: np.ndarray, -) -> np.ndarray: - w_at_vertices = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - connectivities, w, c_intp - ) - - horizontal_advection_of_w_at_edges_on_half_levels = ( - compute_horizontal_advection_term_for_vertical_velocity_numpy( - connectivities=connectivities, - vn_ie=vn_on_half_levels, - inv_dual_edge_length=inv_dual_edge_length, - w=w, - z_vt_ie=tangential_wind_on_half_levels, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - z_w_v=w_at_vertices, - ) - ) - - return horizontal_advection_of_w_at_edges_on_half_levels - - -def add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - cfl_clipping: np.ndarray, - owner_mask: np.ndarray, - contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, - ddqz_z_half: np.ndarray, - area: np.ndarray, - geofac_n2s: np.ndarray, - w: np.ndarray, - vertical_wind_advective_tendency: np.ndarray, - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, -) -> np.ndarray: - owner_mask = np.expand_dims(owner_mask, axis=-1) - area = np.expand_dims(area, axis=-1) - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - - difcoef = np.where( - (cfl_clipping == 1) & (owner_mask == 1), - scalfac_exdiff - * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(contravariant_corrected_w_at_cells_on_half_levels) * dtime / ddqz_z_half - - cfl_w_limit * dtime, - ), - 0, - ) - - c2e2cO = connectivities[dims.C2E2CO] - vertical_wind_advective_tendency = np.where( - (cfl_clipping == 1) & (owner_mask == 1), - vertical_wind_advective_tendency - + difcoef - * area - * np.sum( - np.where( - (c2e2cO != -1)[:, :, np.newaxis], - w[c2e2cO] * geofac_n2s, - 0, - ), - axis=1, - ), - vertical_wind_advective_tendency, - ) - return vertical_wind_advective_tendency - - -def compute_advective_vertical_wind_tendency_numpy( - z_w_con_c: np.ndarray, - w: np.ndarray, - coeff1_dwdz: np.ndarray, - coeff2_dwdz: np.ndarray, - **kwargs: Any, -) -> np.ndarray: - # coeff*_dwdz live on model levels; model level k pairs with half level k - nlev = coeff1_dwdz.shape[1] - ddt_w_adv = np.zeros((z_w_con_c.shape[0], nlev + 1)) - c1, c2 = coeff1_dwdz[:, 1:nlev], coeff2_dwdz[:, 1:nlev] - ddt_w_adv[:, 1:nlev] = -z_w_con_c[:, 1:nlev] * ( - w[:, 0 : nlev - 1] * c1 - w[:, 2 : nlev + 1] * c2 + w[:, 1:nlev] * (c2 - c1) - ) - return ddt_w_adv - - -def compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - vertical_wind_advective_tendency: np.ndarray, - w: np.ndarray, - horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, - contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, - cfl_clipping: np.ndarray, - coeff1_dwdz: np.ndarray, - coeff2_dwdz: np.ndarray, - e_bln_c_s: np.ndarray, - ddqz_z_half: np.ndarray, - area: np.ndarray, - geofac_n2s: np.ndarray, - owner_mask: np.ndarray, - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, - nlev: int, - end_index_of_damping_layer: int, -) -> np.ndarray: - k = np.arange(nlev) - - condition1 = k >= 1 - vertical_wind_advective_tendency = np.where( - condition1, - compute_advective_vertical_wind_tendency_numpy( - contravariant_corrected_w_at_cells_on_half_levels, w, coeff1_dwdz, coeff2_dwdz - )[:, :-1], - vertical_wind_advective_tendency, - ) - - vertical_wind_advective_tendency = np.where( - condition1, - add_interpolated_horizontal_advection_of_w_numpy( - connectivities, - e_bln_c_s, - horizontal_advection_of_w_at_edges_on_half_levels[:, :-1], - vertical_wind_advective_tendency, - ), - vertical_wind_advective_tendency, - ) - - condition2 = (np.maximum(2, end_index_of_damping_layer - 2) <= k) & (k < nlev - 3) - - vertical_wind_advective_tendency = np.where( - condition2, - add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( - connectivities=connectivities, - cfl_clipping=cfl_clipping, - owner_mask=owner_mask, - contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - w=w[:, :-1], - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - ), - vertical_wind_advective_tendency, - ) - - return vertical_wind_advective_tendency - - -@pytest.mark.embedded_remap_error -@pytest.mark.continuous_benchmarking -class TestFusedVelocityAdvectionStencilVMomentum(stencil_tests.StencilTest): - PROGRAM = compute_advection_in_corrector_vertical_momentum - OUTPUTS = ( - "vertical_wind_advective_tendency", - "contravariant_corrected_w_at_cells_on_model_levels", - "vertical_cfl", - ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "end_index_of_damping_layer", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "end_index_of_damping_layer", - ), - } - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - vertical_wind_advective_tendency: np.ndarray, - contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, - vertical_cfl: np.ndarray, - w: np.ndarray, - tangential_wind_on_half_levels: np.ndarray, - vn_on_half_levels: np.ndarray, - contravariant_correction_at_cells_on_half_levels: np.ndarray, - coeff1_dwdz: np.ndarray, - coeff2_dwdz: np.ndarray, - c_intp: np.ndarray, - inv_dual_edge_length: np.ndarray, - inv_primal_edge_length: np.ndarray, - tangent_orientation: np.ndarray, - e_bln_c_s: np.ndarray, - ddqz_z_half: np.ndarray, - area: np.ndarray, - geofac_n2s: np.ndarray, - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, - owner_mask: np.ndarray, - end_index_of_damping_layer: int, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - nlev = kwargs["vertical_end"] - - horizontal_advection_of_w_at_edges_on_half_levels = compute_horizontal_advection_of_w( - connectivities=connectivities, - w=w, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - ) - - # We need to store the initial return field, because we only compute on a subdomain. - contravariant_corrected_w_at_cells_on_model_levels_ret = ( - contravariant_corrected_w_at_cells_on_model_levels.copy() - ) - vertical_wind_advective_tendency_ret = vertical_wind_advective_tendency.copy() - vertical_cfl_ret = vertical_cfl.copy() - - ( - contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping, - vertical_cfl, - ) = compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( - w=w[:, :-1], - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels[ - :, :-1 - ], - ddqz_z_half=ddqz_z_half[:, :-1], - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=nlev, - end_index_of_damping_layer=end_index_of_damping_layer, - ) - - vertical_wind_advective_tendency = compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( - connectivities=connectivities, - vertical_wind_advective_tendency=vertical_wind_advective_tendency[:, :-1], - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping=cfl_clipping, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half[:, :-1], - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=nlev, - end_index_of_damping_layer=end_index_of_damping_layer, - ) - - contravariant_corrected_w_at_cells_on_model_levels = ( - interpolate_contravariant_vertical_velocity_to_full_levels_numpy( - contravariant_corrected_w_at_cells_on_half_levels - ) - ) - - # Apply the slicing. - horizontal_start = kwargs["horizontal_start"] - horizontal_end = kwargs["horizontal_end"] - vertical_start = kwargs["vertical_start"] - vertical_end = kwargs["vertical_end"] - - contravariant_corrected_w_at_cells_on_model_levels_ret[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] = contravariant_corrected_w_at_cells_on_model_levels[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] - vertical_wind_advective_tendency_ret[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] = vertical_wind_advective_tendency[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] - vertical_cfl_ret[horizontal_start:horizontal_end, vertical_start:vertical_end] = ( - vertical_cfl[horizontal_start:horizontal_end, vertical_start:vertical_end] - ) - - return dict( - vertical_wind_advective_tendency=vertical_wind_advective_tendency_ret, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels_ret, - vertical_cfl=vertical_cfl_ret, - ) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - contravariant_corrected_w_at_cells_on_model_levels = data_alloc.zero_field( - dims.CellDim, dims.KDim - ) - vertical_wind_advective_tendency = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) - w = data_alloc.random_field(dims.CellDim, dims.KHalfDim) - tangential_wind_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) - vn_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) - contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( - dims.CellDim, dims.KHalfDim - ) - - coeff1_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) - coeff2_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) - - c_intp = data_alloc.random_field(dims.VertexDim, dims.V2CDim) - inv_dual_edge_length = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) - inv_primal_edge_length = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) - tangent_orientation = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) - e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim) - - vertical_cfl = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) - owner_mask = data_alloc.random_mask(dims.CellDim) - ddqz_z_half = data_alloc.random_field(dims.CellDim, dims.KHalfDim) - area = data_alloc.random_field(dims.CellDim) - geofac_n2s = data_alloc.random_field(dims.CellDim, dims.C2E2CODim) - - scalfac_exdiff = 10.0 - dtime = 2.0 - cfl_w_limit = 0.65 / dtime - - end_index_of_damping_layer = 12 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime - - cell_domain = h_grid.domain(dims.CellDim) - horizontal_start = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) - horizontal_end = grid.end_index(cell_domain(h_grid.Zone.HALO)) - vertical_start = 0 - vertical_end = grid.num_levels - - return dict( - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl=vertical_cfl, - w=w, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - end_index_of_damping_layer=end_index_of_damping_layer, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - ) - - -@pytest.mark.embedded_remap_error -@pytest.mark.continuous_benchmarking -class TestFusedVelocityAdvectionStencilVMomentumAndContravariant(stencil_tests.StencilTest): - PROGRAM = compute_advection_in_predictor_vertical_momentum - OUTPUTS = ( - "contravariant_correction_at_cells_on_half_levels", - "vertical_wind_advective_tendency", - "contravariant_corrected_w_at_cells_on_model_levels", - "vertical_cfl", - ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "nflatlev", - "skip_compute_predictor_vertical_advection", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "end_index_of_damping_layer", - "nflatlev", - "skip_compute_predictor_vertical_advection", - ), - } - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - contravariant_correction_at_cells_on_half_levels: np.ndarray, - vertical_wind_advective_tendency: np.ndarray, - contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, - vertical_cfl: np.ndarray, - w: np.ndarray, - horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, - contravariant_correction_at_edges_on_model_levels: np.ndarray, - coeff1_dwdz: np.ndarray, - coeff2_dwdz: np.ndarray, - e_bln_c_s: np.ndarray, - wgtfac_c: np.ndarray, - ddqz_z_half: np.ndarray, - area: np.ndarray, - geofac_n2s: np.ndarray, - scalfac_exdiff: ta.wpfloat, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, - owner_mask: np.ndarray, - nflatlev: int, - end_index_of_damping_layer: int, - skip_compute_predictor_vertical_advection: bool, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - nlev = kwargs["vertical_end"] - - # We need to store the initial return field, because we only compute on a subdomain. - contravariant_correction_at_cells_on_half_levels_ret = ( - contravariant_correction_at_cells_on_half_levels.copy() - ) - contravariant_corrected_w_at_cells_on_model_levels_ret = ( - contravariant_corrected_w_at_cells_on_model_levels.copy() - ) - vertical_wind_advective_tendency_ret = vertical_wind_advective_tendency.copy() - vertical_cfl_ret = vertical_cfl.copy() - - contravariant_correction_at_cells_on_half_levels_nlev = interpolate_contravariant_correction_to_cells_on_half_levels_numpy( - connectivities=connectivities, - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels[ - :, :-1 - ], - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - nflatlev=nflatlev, - nlev=nlev, - ) - - ( - contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping, - vertical_cfl, - ) = compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( - w=w[:, :-1], - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels_nlev, - ddqz_z_half=ddqz_z_half[:, :-1], - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=nlev, - end_index_of_damping_layer=end_index_of_damping_layer, - ) - if not skip_compute_predictor_vertical_advection: - vertical_wind_advective_tendency = compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( - connectivities=connectivities, - vertical_wind_advective_tendency=vertical_wind_advective_tendency[:, :-1], - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, - cfl_clipping=cfl_clipping, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - ddqz_z_half=ddqz_z_half[:, :-1], - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - nlev=nlev, - end_index_of_damping_layer=end_index_of_damping_layer, - ) - - contravariant_corrected_w_at_cells_on_model_levels = ( - interpolate_contravariant_vertical_velocity_to_full_levels_numpy( - contravariant_corrected_w_at_cells_on_half_levels - ) - ) - - # Apply the slicing. - horizontal_start = kwargs["horizontal_start"] - horizontal_end = kwargs["horizontal_end"] - vertical_start = kwargs["vertical_start"] - vertical_end = kwargs["vertical_end"] - - contravariant_correction_at_cells_on_half_levels_ret[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] = contravariant_correction_at_cells_on_half_levels_nlev[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] - contravariant_corrected_w_at_cells_on_model_levels_ret[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] = contravariant_corrected_w_at_cells_on_model_levels[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] - vertical_wind_advective_tendency_ret[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] = vertical_wind_advective_tendency[ - horizontal_start:horizontal_end, vertical_start:vertical_end - ] - vertical_cfl_ret[horizontal_start:horizontal_end, vertical_start:vertical_end] = ( - vertical_cfl[horizontal_start:horizontal_end, vertical_start:vertical_end] - ) - - return dict( - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels_ret, - vertical_wind_advective_tendency=vertical_wind_advective_tendency_ret, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels_ret, - vertical_cfl=vertical_cfl_ret, - ) - - @stencil_tests.input_data_fixture( - params=[ - {"skip_compute_predictor_vertical_advection": value} for value in [True, False] - ], # True for benchmarking, False for testing - ids=lambda param: ( - f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]" - ), - ) - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, - grid: base.Grid, - request: pytest.FixtureRequest, - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - contravariant_corrected_w_at_cells_on_model_levels = data_alloc.zero_field( - dims.CellDim, dims.KDim - ) - vertical_wind_advective_tendency = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) - w = data_alloc.random_field(dims.CellDim, dims.KHalfDim) - horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.random_field( - dims.EdgeDim, dims.KHalfDim - ) - contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( - dims.EdgeDim, dims.KDim - ) - contravariant_correction_at_cells_on_half_levels = data_alloc.zero_field( - dims.CellDim, dims.KHalfDim - ) - - coeff1_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) - coeff2_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) - - e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim) - wgtfac_c = data_alloc.random_field(dims.CellDim, dims.KHalfDim) - - vertical_cfl = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) - owner_mask = data_alloc.random_mask(dims.CellDim) - ddqz_z_half = data_alloc.random_field(dims.CellDim, dims.KHalfDim) - area = data_alloc.random_field(dims.CellDim) - geofac_n2s = data_alloc.random_field(dims.CellDim, dims.C2E2CODim) - - scalfac_exdiff = 10.0 - dtime = 2.0 - cfl_w_limit = 0.65 / dtime - - skip_compute_predictor_vertical_advection = request.param[ - "skip_compute_predictor_vertical_advection" - ] - - nflatlev = 5 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime - end_index_of_damping_layer = 12 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime - - cell_domain = h_grid.domain(dims.CellDim) - horizontal_start = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4)) - horizontal_end = grid.end_index(cell_domain(h_grid.Zone.HALO)) - vertical_start = 0 - vertical_end = grid.num_levels - - return dict( - contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, - vertical_wind_advective_tendency=vertical_wind_advective_tendency, - contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, - vertical_cfl=vertical_cfl, - w=w, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - coeff1_dwdz=coeff1_dwdz, - coeff2_dwdz=coeff2_dwdz, - e_bln_c_s=e_bln_c_s, - wgtfac_c=wgtfac_c, - ddqz_z_half=ddqz_z_half, - area=area, - geofac_n2s=geofac_n2s, - owner_mask=owner_mask, - scalfac_exdiff=scalfac_exdiff, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=nflatlev, - end_index_of_damping_layer=end_index_of_damping_layer, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py deleted file mode 100644 index 1856d14f58..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_diagnostics_from_normal_wind.py +++ /dev/null @@ -1,371 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from collections.abc import Mapping - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.compute_diagnostics_from_normal_wind import ( - compute_diagnostics_from_normal_wind, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.states import utils as state_utils -from icon4py.model.testing import reference_funcs, stencil_tests - -from .test_compute_contravariant_correction import compute_contravariant_correction_numpy -from .test_compute_horizontal_advection_term_for_vertical_velocity import ( - compute_horizontal_advection_term_for_vertical_velocity_numpy, -) -from .test_compute_horizontal_kinetic_energy import compute_horizontal_kinetic_energy_numpy -from .test_extrapolate_at_top import extrapolate_at_top_numpy -from .test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges import ( - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_numpy, - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_vn_ie_numpy, - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_z_kin_hor_e_numpy, -) -from .test_interpolate_vt_to_interface_edges import interpolate_vt_to_interface_edges_numpy -from .test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( - mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy, -) - - -def compute_diagnostics_from_normal_wind_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - tangential_wind_on_half_levels: np.ndarray, - tangential_wind: np.ndarray, - vn_on_half_levels: np.ndarray, - horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, - contravariant_correction_at_edges_on_model_levels: np.ndarray, - vn: np.ndarray, - rbf_vec_coeff_e: np.ndarray, - wgtfac_e: np.ndarray, - ddxn_z_full: np.ndarray, - ddxt_z_full: np.ndarray, - wgtfacq_e: np.ndarray, - horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, - c_intp: np.ndarray, - inv_dual_edge_length: np.ndarray, - w: np.ndarray, - inv_primal_edge_length: np.ndarray, - tangent_orientation: np.ndarray, - skip_compute_predictor_vertical_advection: bool, - nflatlev: int, - nlevp1: int, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, -) -> tuple[np.ndarray, ...]: - k: np.ndarray = np.arange(nlevp1) - k = k[np.newaxis, :] - k_nlev = k[:, :-1] - - initial_vn_on_half_levels = vn_on_half_levels.copy() - - tangential_wind = np.where( - k_nlev >= vertical_start, - reference_funcs.compute_tangential_wind_numpy(connectivities, vn, rbf_vec_coeff_e), - tangential_wind, - ) - - horizontal_kinetic_energy_at_edges_on_model_levels = np.where( - k_nlev >= vertical_start, - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_z_kin_hor_e_numpy( - vn, tangential_wind - ), - horizontal_kinetic_energy_at_edges_on_model_levels, - ) - - vn_on_half_levels = ( - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_vn_ie_numpy(wgtfac_e, vn) - ) - - if not skip_compute_predictor_vertical_advection: - tangential_wind_on_half_levels[:, : nlevp1 - 1] = np.where( - k_nlev >= vertical_start, - interpolate_vt_to_interface_edges_numpy(wgtfac_e, tangential_wind)[:, : nlevp1 - 1], - tangential_wind_on_half_levels[:, : nlevp1 - 1], - ) - - contravariant_correction_at_edges_on_model_levels = np.where( - k_nlev >= nflatlev, - compute_contravariant_correction_numpy(vn, ddxn_z_full, ddxt_z_full, tangential_wind), - contravariant_correction_at_edges_on_model_levels, - ) - - if not skip_compute_predictor_vertical_advection: - w_at_vertices = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - connectivities, w, c_intp - ) - horizontal_advection_of_w_at_edges_on_half_levels[:, : nlevp1 - 1] = np.where( - k_nlev >= vertical_start, - compute_horizontal_advection_term_for_vertical_velocity_numpy( - connectivities=connectivities, - vn_ie=vn_on_half_levels, - inv_dual_edge_length=inv_dual_edge_length, - w=w, - z_vt_ie=tangential_wind_on_half_levels, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - z_w_v=w_at_vertices, - )[:, : nlevp1 - 1], - horizontal_advection_of_w_at_edges_on_half_levels[:, : nlevp1 - 1], - ) - - vn_on_half_levels[:, -1] = extrapolate_to_surface_numpy(wgtfacq_e, vn) - - vn_on_half_levels[:horizontal_start, :] = initial_vn_on_half_levels[:horizontal_start, :] - vn_on_half_levels[horizontal_end:, :] = initial_vn_on_half_levels[horizontal_end:, :] - - return ( - tangential_wind, - tangential_wind_on_half_levels, - vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels, - ) - - -def extrapolate_to_surface_numpy(wgtfacq_e: np.ndarray, vn: np.ndarray) -> np.ndarray: - vn_k_minus_1 = vn[:, -1] - vn_k_minus_2 = vn[:, -2] - vn_k_minus_3 = vn[:, -3] - wgtfacq_e_k_minus_1 = wgtfacq_e[:, -1] - wgtfacq_e_k_minus_2 = wgtfacq_e[:, -2] - wgtfacq_e_k_minus_3 = wgtfacq_e[:, -3] - vn_at_surface = ( - wgtfacq_e_k_minus_1 * vn_k_minus_1 - + wgtfacq_e_k_minus_2 * vn_k_minus_2 - + wgtfacq_e_k_minus_3 * vn_k_minus_3 - ) - return vn_at_surface - - -@pytest.mark.embedded_remap_error -@pytest.mark.continuous_benchmarking -class TestComputeDerivedHorizontalWindsAndKEAndHorizontalAdvectionofWAndContravariantCorrection( - stencil_tests.StencilTest -): - PROGRAM = compute_diagnostics_from_normal_wind - OUTPUTS = ( - "tangential_wind", - "tangential_wind_on_half_levels", - "vn_on_half_levels", - "horizontal_kinetic_energy_at_edges_on_model_levels", - "contravariant_correction_at_edges_on_model_levels", - "horizontal_advection_of_w_at_edges_on_half_levels", - ) - STATIC_PARAMS = { - stencil_tests.StandardStaticVariants.NONE: (), - stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( - "horizontal_start", - "horizontal_end", - "vertical_start", - "vertical_end", - "nflatlev", - ), - stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( - "vertical_start", - "vertical_end", - "nflatlev", - ), - } - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - tangential_wind: np.ndarray, - tangential_wind_on_half_levels: np.ndarray, - vn_on_half_levels: np.ndarray, - horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, - contravariant_correction_at_edges_on_model_levels: np.ndarray, - horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, - vn: np.ndarray, - w: np.ndarray, - rbf_vec_coeff_e: np.ndarray, - wgtfac_e: np.ndarray, - ddxn_z_full: np.ndarray, - ddxt_z_full: np.ndarray, - wgtfacq_e: np.ndarray, - c_intp: np.ndarray, - inv_dual_edge_length: np.ndarray, - inv_primal_edge_length: np.ndarray, - tangent_orientation: np.ndarray, - skip_compute_predictor_vertical_advection: bool, - nflatlev: int, - horizontal_start: int, - horizontal_end: int, - vertical_start: int, - vertical_end: int, - **kwargs: object, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - initial_tangential_wind = tangential_wind.copy() - initial_tangential_wind_on_half_levels = tangential_wind_on_half_levels.copy() - initial_horizontal_kinetic_energy_at_edges_on_model_levels = ( - horizontal_kinetic_energy_at_edges_on_model_levels.copy() - ) - initial_contravariant_correction_at_edges_on_model_levels = ( - contravariant_correction_at_edges_on_model_levels.copy() - ) - initial_horizontal_advection_of_w_at_edges_on_half_levels = ( - horizontal_advection_of_w_at_edges_on_half_levels.copy() - ) - ( - tangential_wind, - tangential_wind_on_half_levels, - vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels, - ) = compute_diagnostics_from_normal_wind_numpy( - connectivities=connectivities, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - tangential_wind=tangential_wind, - vn_on_half_levels=vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - vn=vn, - rbf_vec_coeff_e=rbf_vec_coeff_e, - wgtfac_e=wgtfac_e, - ddxn_z_full=ddxn_z_full, - ddxt_z_full=ddxt_z_full, - wgtfacq_e=wgtfacq_e, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - w=w, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=nflatlev, - nlevp1=vertical_end, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - ) - - tangential_wind[:horizontal_start, :] = initial_tangential_wind[:horizontal_start, :] - tangential_wind[horizontal_end:, :] = initial_tangential_wind[horizontal_end:, :] - - tangential_wind_on_half_levels[:horizontal_start, :] = ( - initial_tangential_wind_on_half_levels[:horizontal_start, :] - ) - tangential_wind_on_half_levels[horizontal_end:, :] = initial_tangential_wind_on_half_levels[ - horizontal_end:, : - ] - horizontal_kinetic_energy_at_edges_on_model_levels[:horizontal_start, :] = ( - initial_horizontal_kinetic_energy_at_edges_on_model_levels[:horizontal_start, :] - ) - horizontal_kinetic_energy_at_edges_on_model_levels[horizontal_end:, :] = ( - initial_horizontal_kinetic_energy_at_edges_on_model_levels[horizontal_end:, :] - ) - - contravariant_correction_at_edges_on_model_levels[:horizontal_start, :] = ( - initial_contravariant_correction_at_edges_on_model_levels[:horizontal_start, :] - ) - contravariant_correction_at_edges_on_model_levels[horizontal_end:, :] = ( - initial_contravariant_correction_at_edges_on_model_levels[horizontal_end:, :] - ) - - horizontal_advection_of_w_at_edges_on_half_levels[:horizontal_start, :] = ( - initial_horizontal_advection_of_w_at_edges_on_half_levels[:horizontal_start, :] - ) - horizontal_advection_of_w_at_edges_on_half_levels[horizontal_end:, :] = ( - initial_horizontal_advection_of_w_at_edges_on_half_levels[horizontal_end:, :] - ) - - return dict( - tangential_wind=tangential_wind, - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - vn_on_half_levels=vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - ) - - @stencil_tests.input_data_fixture( - params=[ - {"skip_compute_predictor_vertical_advection": value} for value in [True, False] - ], # True for benchmarking, False for testing - ids=lambda param: ( - f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]" - ), - ) - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, - grid: base.Grid, - request: pytest.FixtureRequest, - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.zero_field( - dims.EdgeDim, dims.KHalfDim - ) - tangential_wind = data_alloc.random_field(dims.EdgeDim, dims.KDim) - tangential_wind_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) - vn_on_half_levels = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim) - horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( - dims.EdgeDim, dims.KDim - ) - contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( - dims.EdgeDim, dims.KDim - ) - vn = data_alloc.random_field(dims.EdgeDim, dims.KDim) - w = data_alloc.random_field(dims.CellDim, dims.KHalfDim) - rbf_vec_coeff_e = data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim) - wgtfac_e = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) - ddxn_z_full = data_alloc.random_field(dims.EdgeDim, dims.KDim) - ddxt_z_full = data_alloc.random_field(dims.EdgeDim, dims.KDim) - inv_dual_edge_length = data_alloc.random_field(dims.EdgeDim) - inv_primal_edge_length = data_alloc.random_field(dims.EdgeDim) - tangent_orientation = data_alloc.random_field(dims.EdgeDim) - wgtfacq_e = data_alloc.random_field(dims.EdgeDim, dims.KDim) - c_intp = data_alloc.random_field(dims.VertexDim, dims.V2CDim) - - nlev = grid.num_levels - nflatlev = 5 # value is set to reflect the MCH ch1 experiment. Changing this value will change the expected runtime - - skip_compute_predictor_vertical_advection = request.param[ - "skip_compute_predictor_vertical_advection" - ] - - edge_domain = h_grid.domain(dims.EdgeDim) - horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5)) - horizontal_end = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) - vertical_start = 0 - vertical_end = nlev + 1 - - return dict( - tangential_wind_on_half_levels=tangential_wind_on_half_levels, - tangential_wind=tangential_wind, - vn_on_half_levels=vn_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, - contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, - horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, - vn=vn, - w=w, - rbf_vec_coeff_e=rbf_vec_coeff_e, - wgtfac_e=wgtfac_e, - ddxn_z_full=ddxn_z_full, - ddxt_z_full=ddxt_z_full, - wgtfacq_e=wgtfacq_e, - c_intp=c_intp, - inv_dual_edge_length=inv_dual_edge_length, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, - nflatlev=nflatlev, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=vertical_start, - vertical_end=vertical_end, - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py deleted file mode 100644 index 71a1980918..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_advection_term_for_vertical_velocity.py +++ /dev/null @@ -1,118 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_advection_term_for_vertical_velocity import ( - compute_horizontal_advection_term_for_vertical_velocity, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base, horizontal as h_grid -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -def compute_horizontal_advection_term_for_vertical_velocity_numpy( - *, - connectivities: Mapping[gtx.FieldOffset, np.ndarray], - vn_ie: np.ndarray, - inv_dual_edge_length: np.ndarray, - w: np.ndarray, - z_vt_ie: np.ndarray, - inv_primal_edge_length: np.ndarray, - tangent_orientation: np.ndarray, - z_w_v: np.ndarray, -) -> np.ndarray: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - - w_e2c = w[connectivities[dims.E2C]] - z_w_v_e2v = z_w_v[connectivities[dims.E2V]] - - red_w = w_e2c[:, 0] - w_e2c[:, 1] - red_z_w_v = z_w_v_e2v[:, 0] - z_w_v_e2v[:, 1] - - z_v_grad_w = ( - vn_ie * inv_dual_edge_length * red_w - + z_vt_ie * inv_primal_edge_length * tangent_orientation * red_z_w_v - ) - return z_v_grad_w - - -class TestComputeHorizontalAdvectionTermForVerticalVelocity(stencil_tests.StencilTest): - PROGRAM = compute_horizontal_advection_term_for_vertical_velocity - OUTPUTS = ("z_v_grad_w",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - vn_ie: np.ndarray, - inv_dual_edge_length: np.ndarray, - w: np.ndarray, - z_vt_ie: np.ndarray, - inv_primal_edge_length: np.ndarray, - tangent_orientation: np.ndarray, - z_w_v: np.ndarray, - z_v_grad_w: np.ndarray, - horizontal_start: int, - horizontal_end: int, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - z_v_grad_w[horizontal_start:horizontal_end, :] = ( - compute_horizontal_advection_term_for_vertical_velocity_numpy( - connectivities=connectivities, - vn_ie=vn_ie, - inv_dual_edge_length=inv_dual_edge_length, - w=w, - z_vt_ie=z_vt_ie, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - z_w_v=z_w_v, - ) - )[horizontal_start:horizontal_end, :] - return dict(z_v_grad_w=z_v_grad_w) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn_ie = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - inv_dual_edge_length = data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) - w = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=wpfloat) - z_vt_ie = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - inv_primal_edge_length = data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) - tangent_orientation = data_alloc.random_field(dims.EdgeDim, dtype=wpfloat) - z_w_v = data_alloc.random_field(dims.VertexDim, dims.KHalfDim, dtype=vpfloat) - z_v_grad_w = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - - edge_domain = h_grid.domain(dims.EdgeDim) - horizontal_start = grid.start_index(edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7)) - horizontal_end = grid.end_index(edge_domain(h_grid.Zone.HALO)) - - return dict( - vn_ie=vn_ie, - inv_dual_edge_length=inv_dual_edge_length, - w=w, - z_vt_ie=z_vt_ie, - inv_primal_edge_length=inv_primal_edge_length, - tangent_orientation=tangent_orientation, - z_w_v=z_w_v, - z_v_grad_w=z_v_grad_w, - horizontal_start=horizontal_start, - horizontal_end=horizontal_end, - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels + 1), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py deleted file mode 100644 index 813808ef6c..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_kinetic_energy.py +++ /dev/null @@ -1,70 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_kinetic_energy import ( - compute_horizontal_kinetic_energy, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -def compute_horizontal_kinetic_energy_numpy(vn: np.ndarray, vt: np.ndarray) -> tuple: - nlev = vn.shape[1] - vn_ie = np.zeros((vn.shape[0], nlev + 1)) - vn_ie[:, :nlev] = vn - z_vt_ie = np.zeros((vt.shape[0], nlev + 1)) - z_vt_ie[:, :nlev] = vt - z_kin_hor_e = 0.5 * ((vn * vn) + (vt * vt)) - return vn_ie, z_vt_ie, z_kin_hor_e - - -class TestComputeHorizontalKineticEnergy(stencil_tests.StencilTest): - PROGRAM = compute_horizontal_kinetic_energy - OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - vn: np.ndarray, - vt: np.ndarray, - **kwargs: Any, - ) -> dict: - vn_ie, z_vt_ie, z_kin_hor_e = compute_horizontal_kinetic_energy_numpy(vn, vt) - return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - vn = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) - vt = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - - vn_ie = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - z_vt_ie = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - z_kin_hor_e = data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - - return dict( - vn=vn, - vt=vt, - vn_ie=vn_ie, - z_vt_ie=z_vt_ie, - z_kin_hor_e=z_kin_hor_e, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_edges), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py index 93ddd0fe29..fa1b193d97 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_horizontal_velocity_quantities_and_fluxes.py @@ -22,13 +22,13 @@ from .test_compute_avg_vn_and_graddiv_vn_and_vt import compute_avg_vn_and_graddiv_vn_and_vt_numpy from .test_compute_contravariant_correction import compute_contravariant_correction_numpy -from .test_compute_diagnostics_from_normal_wind import extrapolate_to_surface_numpy -from .test_compute_horizontal_kinetic_energy import compute_horizontal_kinetic_energy_numpy from .test_compute_mass_flux import compute_mass_flux_numpy -from .test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges import ( - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_numpy, -) from .test_interpolate_vt_to_interface_edges import interpolate_vt_to_interface_edges_numpy +from .test_velocity_advection_terms import ( + compute_horizontal_kinetic_energy_at_edges_numpy, + extrapolate_to_surface_numpy, + interpolate_vn_to_half_levels_numpy, +) def compute_vt_vn_on_half_levels_and_kinetic_energy_numpy( @@ -38,51 +38,17 @@ def compute_vt_vn_on_half_levels_and_kinetic_energy_numpy( tangential_wind: np.ndarray, vn_on_half_levels: np.ndarray, tangential_wind_on_half_levels: np.ndarray, - horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, wgtfac_e: np.ndarray, wgtfacq_e: np.ndarray, - nlevp1: int, ) -> tuple[np.ndarray, ...]: - k = np.arange(nlevp1)[np.newaxis, :] - k_nlev = k[:, :-1] - - vn_ie, z_kin_hor_e = interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_numpy( - wgtfac_e, vn, tangential_wind - ) - vn_on_half_levels[:, :-1], horizontal_kinetic_energy_at_edges_on_model_levels = np.where( - k_nlev >= 1, - (vn_ie[:, :-1], z_kin_hor_e), - (vn_on_half_levels[:, :-1], horizontal_kinetic_energy_at_edges_on_model_levels), - ) - - tangential_wind_on_half_levels[:, :-1] = np.where( - k_nlev >= 1, - interpolate_vt_to_interface_edges_numpy(wgtfac_e, tangential_wind)[:, :-1], - tangential_wind_on_half_levels[:, :-1], - ) - - vn_ie_at_top, tangential_wind_on_half_levels_at_top, z_kin_hor_e_at_top = ( - compute_horizontal_kinetic_energy_numpy(vn, tangential_wind) - ) - ( - vn_on_half_levels[:, :-1], - tangential_wind_on_half_levels[:, :-1], - horizontal_kinetic_energy_at_edges_on_model_levels, - ) = np.where( - k_nlev == 0, - ( - vn_ie_at_top[:, :-1], - tangential_wind_on_half_levels_at_top[:, :-1], - z_kin_hor_e_at_top, - ), - ( - vn_on_half_levels[:, :-1], - tangential_wind_on_half_levels[:, :-1], - horizontal_kinetic_energy_at_edges_on_model_levels, - ), - ) - + vn_on_half_levels[:, :-1] = interpolate_vn_to_half_levels_numpy(wgtfac_e, vn)[:, :-1] vn_on_half_levels[:, -1] = extrapolate_to_surface_numpy(wgtfacq_e, vn) + tangential_wind_on_half_levels[:, :-1] = interpolate_vt_to_interface_edges_numpy( + wgtfac_e, tangential_wind + )[:, :-1] + horizontal_kinetic_energy_at_edges_on_model_levels = ( + compute_horizontal_kinetic_energy_at_edges_numpy(vn, tangential_wind) + ) return ( vn_on_half_levels, @@ -211,10 +177,8 @@ def reference( tangential_wind=tangential_wind, vn_on_half_levels=vn_on_half_levels, tangential_wind_on_half_levels=tangential_wind_on_half_levels, - horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, wgtfac_e=wgtfac_e, wgtfacq_e=wgtfacq_e, - nlevp1=vertical_end, ) spatially_averaged_vn[:horizontal_start, :] = initial_spatially_averaged_vn[ diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py deleted file mode 100644 index 19fbcefe46..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_compute_maximum_cfl_and_clip_contravariant_vertical_velocity.py +++ /dev/null @@ -1,95 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.compute_maximum_cfl_and_clip_contravariant_vertical_velocity import ( - compute_maximum_cfl_and_clip_contravariant_vertical_velocity, -) -from icon4py.model.common import dimension as dims, type_alias as ta -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.testing import stencil_tests - - -def compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( - ddqz_z_half: np.ndarray, z_w_con_c: np.ndarray, cfl_w_limit: ta.wpfloat, dtime: ta.wpfloat -) -> tuple: - num_rows, num_cols = z_w_con_c.shape - cfl_clipping = np.where( - np.abs(z_w_con_c) > cfl_w_limit * ddqz_z_half, - np.ones([num_rows, num_cols]), - np.zeros_like(z_w_con_c), - ) - vcfl = np.where(cfl_clipping == 1.0, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = np.where( - (cfl_clipping == 1.0) & (vcfl < -0.85), - -0.85 * ddqz_z_half / dtime, - z_w_con_c, - ) - z_w_con_c = np.where( - (cfl_clipping == 1.0) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c - ) - - return cfl_clipping, vcfl, z_w_con_c - - -class TestComputeMaximumCflAndClipContravariantVerticalVelocity(stencil_tests.StencilTest): - PROGRAM = compute_maximum_cfl_and_clip_contravariant_vertical_velocity - OUTPUTS = ("cfl_clipping", "vcfl", "z_w_con_c") - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - ddqz_z_half: np.ndarray, - z_w_con_c: np.ndarray, - cfl_w_limit: ta.wpfloat, - dtime: ta.wpfloat, - **kwargs: Any, - ) -> dict: - ( - cfl_clipping, - vcfl, - z_w_con_c, - ) = compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( - ddqz_z_half, z_w_con_c, cfl_w_limit, dtime - ) - - return dict( - cfl_clipping=cfl_clipping, - vcfl=vcfl, - z_w_con_c=z_w_con_c, - ) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - ddqz_z_half = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat) - z_w_con_c = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat) - cfl_clipping = data_alloc.random_mask(dims.CellDim, dims.KHalfDim, dtype=bool) - vcfl = data_alloc.zero_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat) - cfl_w_limit = ta.vpfloat("5.0") - dtime = ta.wpfloat("9.0") - - return dict( - ddqz_z_half=ddqz_z_half, - z_w_con_c=z_w_con_c, - cfl_clipping=cfl_clipping, - vcfl=vcfl, - cfl_w_limit=cfl_w_limit, - dtime=dtime, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels + 1), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py deleted file mode 100644 index 121c654a6d..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_contravariant_vertical_velocity_to_full_levels.py +++ /dev/null @@ -1,60 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.interpolate_contravariant_vertical_velocity_to_full_levels import ( - interpolate_contravariant_vertical_velocity_to_full_levels, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat -from icon4py.model.testing import stencil_tests - - -def interpolate_contravariant_vertical_velocity_to_full_levels_numpy( - z_w_con_c: np.ndarray, -) -> np.ndarray: - z_w_con_c_full = 0.5 * (z_w_con_c[:, :-1] + z_w_con_c[:, 1:]) - return z_w_con_c_full - - -class TestInterpolateContravariantVerticalVelocityToFullLevels(stencil_tests.StencilTest): - PROGRAM = interpolate_contravariant_vertical_velocity_to_full_levels - OUTPUTS = ("z_w_con_c_full",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - z_w_con_c: np.ndarray, - **kwargs: Any, - ) -> dict: - z_w_con_c_full = interpolate_contravariant_vertical_velocity_to_full_levels_numpy(z_w_con_c) - return dict(z_w_con_c_full=z_w_con_c_full) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - z_w_con_c = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=vpfloat) - - z_w_con_c_full = data_alloc.zero_field(dims.CellDim, dims.KDim, dtype=vpfloat) - - return dict( - z_w_con_c=z_w_con_c, - z_w_con_c_full=z_w_con_c_full, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_cells), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_half_levels.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_half_levels.py new file mode 100644 index 0000000000..3e7cc305e3 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_to_half_levels.py @@ -0,0 +1,52 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Any + +import gt4py.next as gtx +import numpy as np + +from icon4py.model.atmosphere.dycore.stencils.compute_diagnostics_from_normal_wind import ( + _interpolate_to_half_levels, +) +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.grid import base +from icon4py.model.common.states import utils as state_utils +from icon4py.model.testing import stencil_tests + +from .test_velocity_advection_terms import interpolate_vn_to_half_levels_numpy + + +class TestInterpolateToHalfLevels(stencil_tests.StencilTest): + PROGRAM = _interpolate_to_half_levels + # the operator reads x at K+0.5, so the surface half level is out of reach + OUTPUTS = ( + stencil_tests.Output( + "out", refslice=(slice(None), slice(0, -1)), gtslice=(slice(None), slice(0, -1)) + ), + ) + + @stencil_tests.static_reference + def reference(grid: base.Grid, *, wgtfac_e: np.ndarray, x: np.ndarray, **kwargs: Any) -> dict: + return dict(out=interpolate_vn_to_half_levels_numpy(wgtfac_e, x)) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + return dict( + wgtfac_e=data_alloc.random_field( + dims.EdgeDim, dims.KHalfDim, low=0.0, high=1.0, dtype=ta.vpfloat + ), + x=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat), + out=data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=ta.vpfloat), + domain={ + dims.EdgeDim: (0, gtx.int32(grid.num_edges)), + dims.KHalfDim: (0, gtx.int32(grid.num_levels)), + }, + ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py deleted file mode 100644 index 03d5ade611..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges.py +++ /dev/null @@ -1,70 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges import ( - interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -class TestInterpolateVnAndVtToIeAndComputeEkinOnEdges(stencil_tests.StencilTest): - PROGRAM = interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges - OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - wgtfac_e: np.ndarray, - vn: np.ndarray, - vt: np.ndarray, - **kwargs: Any, - ) -> dict: - nlev = vn.shape[1] - w = wgtfac_e[:, 1:nlev] - - vn_ie = np.zeros((vn.shape[0], nlev + 1)) - vn_ie[:, 1:nlev] = w * vn[:, 1:nlev] + (1.0 - w) * vn[:, 0 : nlev - 1] - z_vt_ie = np.zeros((vt.shape[0], nlev + 1)) - z_vt_ie[:, 1:nlev] = w * vt[:, 1:nlev] + (1.0 - w) * vt[:, 0 : nlev - 1] - z_kin_hor_e = 0.5 * (vn**2 + vt**2) - z_kin_hor_e[:, 0] = 0 - return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_e = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - vn = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) - vt = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn_ie = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - z_vt_ie = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - z_kin_hor_e = data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - - return dict( - wgtfac_e=wgtfac_e, - vn=vn, - vt=vt, - vn_ie=vn_ie, - z_vt_ie=z_vt_ie, - z_kin_hor_e=z_kin_hor_e, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_edges), - vertical_start=1, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py deleted file mode 100644 index 4bd90c5a08..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges.py +++ /dev/null @@ -1,111 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges import ( - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -def interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_vn_ie_numpy( - wgtfac_e: np.ndarray, vn: np.ndarray -) -> np.ndarray: - _nlev = vn.shape[1] - vn_ie = np.zeros((vn.shape[0], _nlev + 1)) - _w = wgtfac_e[:, 1:_nlev] - vn_ie[:, 1:_nlev] = _w * vn[:, 1:_nlev] + (1.0 - _w) * vn[:, 0 : _nlev - 1] - vn_ie[:, 0] = vn[:, 0] - return vn_ie - - -def interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_z_kin_hor_e_numpy( - vn: np.ndarray, vt: np.ndarray -) -> np.ndarray: - z_kin_hor_e = 0.5 * (vn * vn + vt * vt) - return z_kin_hor_e - - -def interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_numpy( - wgtfac_e: np.ndarray, vn: np.ndarray, vt: np.ndarray, **kwargs: Any -) -> tuple[np.ndarray, np.ndarray]: - vn_ie = interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_vn_ie_numpy( - wgtfac_e, vn - ) - z_kin_hor_e = ( - interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_z_kin_hor_e_numpy(vn, vt) - ) - return ( - vn_ie, - z_kin_hor_e, - ) - - -class TestInterpolateVnToHalfLevelsAndComputeKineticEnergyOnEdges(stencil_tests.StencilTest): - PROGRAM = interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges - OUTPUTS = ("vn_ie", "z_kin_hor_e") - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - wgtfac_e: np.ndarray, - vn: np.ndarray, - vt: np.ndarray, - vn_ie: np.ndarray, - z_kin_hor_e: np.ndarray, - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, - **kwargs: Any, - ) -> dict: - subset = (slice(horizontal_start, horizontal_end), slice(vertical_start, vertical_end)) - vn_ie, z_kin_hor_e = vn_ie.copy(), z_kin_hor_e.copy() - vn_ie[subset], z_kin_hor_e[subset] = ( - x[subset] - for x in interpolate_vn_to_half_levels_and_compute_kinetic_energy_on_edges_numpy( - wgtfac_e, vn, vt - ) - ) - - return dict( - vn_ie=vn_ie, - z_kin_hor_e=z_kin_hor_e, - ) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - wgtfac_e = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - vn = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) - vt = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - - vn_ie = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=vpfloat) - z_kin_hor_e = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=vpfloat) - - return dict( - wgtfac_e=wgtfac_e, - vn=vn, - vt=vt, - vn_ie=vn_ie, - z_kin_hor_e=z_kin_hor_e, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_edges), - vertical_start=1, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py deleted file mode 100644 index b65fc8bd3b..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ /dev/null @@ -1,70 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( - mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - connectivities: Mapping[gtx.FieldOffset, np.ndarray], p_cell_in: np.ndarray, c_intp: np.ndarray -) -> np.ndarray: - v2c = connectivities[dims.V2C] - c_intp = np.expand_dims(c_intp, axis=-1) - p_vert_out = np.sum(np.where((v2c != -1)[:, :, np.newaxis], p_cell_in[v2c] * c_intp, 0), axis=1) - return p_vert_out - - -class TestMoIconInterpolationScalarCells2vertsScalarRiDsl(stencil_tests.StencilTest): - PROGRAM = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl - OUTPUTS = ("p_vert_out",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - p_cell_in: np.ndarray, - c_intp: np.ndarray, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - p_vert_out = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( - connectivities, p_cell_in, c_intp - ) - return dict( - p_vert_out=p_vert_out, - ) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - p_cell_in = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=wpfloat) - c_intp = data_alloc.random_field(dims.VertexDim, dims.V2CDim, dtype=wpfloat) - p_vert_out = data_alloc.zero_field(dims.VertexDim, dims.KHalfDim, dtype=vpfloat) - - return dict( - p_cell_in=p_cell_in, - c_intp=c_intp, - p_vert_out=p_vert_out, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_vertices), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels + 1), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py deleted file mode 100644 index 7d249683be..0000000000 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py +++ /dev/null @@ -1,66 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -from collections.abc import Mapping -from typing import Any - -import gt4py.next as gtx -import numpy as np -import pytest - -from icon4py.model.atmosphere.dycore.stencils.mo_math_divrot_rot_vertex_ri_dsl import ( - mo_math_divrot_rot_vertex_ri_dsl, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.grid import base -from icon4py.model.common.states import utils as state_utils -from icon4py.model.common.type_alias import vpfloat, wpfloat -from icon4py.model.testing import stencil_tests - - -def mo_math_divrot_rot_vertex_ri_dsl_numpy( - connectivities: Mapping[gtx.FieldOffset, np.ndarray], vec_e: np.ndarray, geofac_rot: np.ndarray -) -> np.ndarray: - v2e = connectivities[dims.V2E] - geofac_rot = np.expand_dims(geofac_rot, axis=-1) - rot_vec = np.sum(np.where((v2e != -1)[:, :, np.newaxis], vec_e[v2e] * geofac_rot, 0), axis=1) - return rot_vec - - -class TestMoMathDivrotRotVertexRiDsl(stencil_tests.StencilTest): - PROGRAM = mo_math_divrot_rot_vertex_ri_dsl - OUTPUTS = ("rot_vec",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - vec_e: np.ndarray, - geofac_rot: np.ndarray, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - rot_vec = mo_math_divrot_rot_vertex_ri_dsl_numpy(connectivities, vec_e, geofac_rot) - return dict(rot_vec=rot_vec) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, gtx.Field | state_utils.ScalarType]: - vec_e = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) - geofac_rot = data_alloc.random_field(dims.VertexDim, dims.V2EDim, dtype=wpfloat) - rot_vec = data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=vpfloat) - - return dict( - vec_e=vec_e, - geofac_rot=geofac_rot, - rot_vec=rot_vec, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_vertices), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_corrector.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_corrector.py new file mode 100644 index 0000000000..11b37de506 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_corrector.py @@ -0,0 +1,301 @@ +# 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 Any + +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_corrector import ( + compute_velocity_advection_in_corrector_step, +) +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.grid import base, horizontal as h_grid +from icon4py.model.common.states import utils as state_utils +from icon4py.model.testing import stencil_tests + +from .test_velocity_advection_terms import ( + _restore_outside, + compute_advection_in_horizontal_momentum_numpy, + compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy, + compute_horizontal_advection_of_w_numpy, + compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy, + interpolate_contravariant_vertical_velocity_to_full_levels_numpy, +) + + +@pytest.mark.embedded_remap_error +@pytest.mark.uses_concat_where +@pytest.mark.continuous_benchmarking +class TestComputeVelocityAdvectionInCorrectorStep(stencil_tests.StencilTest): + PROGRAM = compute_velocity_advection_in_corrector_step + OUTPUTS = ( + "vertical_wind_advective_tendency", + "vertical_cfl", + "normal_wind_advective_tendency", + ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "end_index_of_damping_layer", + "apply_extra_diffusion_on_vn", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "start_cell_lateral_boundary_level_4", + "end_cell_halo", + "start_edge_nudging_level_2", + "end_edge_local", + "vertical_start", + "vertical_end", + "end_index_of_damping_layer", + "apply_extra_diffusion_on_vn", + ), + } + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + vertical_wind_advective_tendency: np.ndarray, + vertical_cfl: np.ndarray, + normal_wind_advective_tendency: np.ndarray, + vn: np.ndarray, + w: np.ndarray, + tangential_wind: np.ndarray, + tangential_wind_on_half_levels: np.ndarray, + vn_on_half_levels: np.ndarray, + horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, + contravariant_correction_at_cells_on_half_levels: np.ndarray, + coeff1_dwdz: np.ndarray, + coeff2_dwdz: np.ndarray, + c_intp: np.ndarray, + inv_dual_edge_length: np.ndarray, + inv_primal_edge_length: np.ndarray, + tangent_orientation: np.ndarray, + e_bln_c_s: np.ndarray, + ddqz_z_half: np.ndarray, + area: np.ndarray, + geofac_n2s: np.ndarray, + owner_mask: np.ndarray, + coriolis_frequency: np.ndarray, + geofac_rot: np.ndarray, + coeff_gradekin: np.ndarray, + c_lin_e: np.ndarray, + ddqz_z_full_e: np.ndarray, + area_edge: np.ndarray, + geofac_grdiv: np.ndarray, + dtime: ta.wpfloat, + apply_extra_diffusion_on_vn: bool, + end_index_of_damping_layer: int, + start_cell_lateral_boundary_level_4: int, + end_cell_halo: int, + start_edge_nudging_level_2: int, + end_edge_local: int, + vertical_start: int, + vertical_end: int, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + nlev = vertical_end + + horizontal_advection_of_w_at_edges_on_half_levels = compute_horizontal_advection_of_w_numpy( + connectivities=connectivities, + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + ) + + ( + contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping, + vertical_cfl_new, + ) = compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( + w=w[:, :-1], + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels[ + :, :-1 + ], + ddqz_z_half=ddqz_z_half[:, :-1], + dtime=dtime, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + vertical_wind_advective_tendency_new = compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( + connectivities=connectivities, + vertical_wind_advective_tendency=vertical_wind_advective_tendency[:, :-1], + w=w, + horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping=cfl_clipping, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half[:, :-1], + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + dtime=dtime, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + contravariant_corrected_w_at_cells_on_model_levels = ( + interpolate_contravariant_vertical_velocity_to_full_levels_numpy( + contravariant_corrected_w_at_cells_on_half_levels + ) + ) + + normal_wind_advective_tendency_new = compute_advection_in_horizontal_momentum_numpy( + connectivities=connectivities, + vn=vn, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + tangential_wind=tangential_wind, + coriolis_frequency=coriolis_frequency, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + vn_on_half_levels=vn_on_half_levels, + e_bln_c_s=e_bln_c_s, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + cell_slice = (start_cell_lateral_boundary_level_4, end_cell_halo) + return dict( + vertical_wind_advective_tendency=_restore_outside( + vertical_wind_advective_tendency_new, + vertical_wind_advective_tendency, + cell_slice, + # ICON computes the tendency over jk = 2..nlev + # (mo_velocity_advection.f90:598), so the top half level is untouched. + (vertical_start + 1, vertical_end), + ), + vertical_cfl=_restore_outside( + vertical_cfl_new, vertical_cfl, cell_slice, (vertical_start, vertical_end) + ), + normal_wind_advective_tendency=_restore_outside( + normal_wind_advective_tendency_new, + normal_wind_advective_tendency, + (start_edge_nudging_level_2, end_edge_local), + (vertical_start, vertical_end), + ), + ) + + @stencil_tests.input_data_fixture( + params=[{"apply_extra_diffusion_on_vn": value} for value in [True, False]], + ids=lambda param: f"apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]", + ) + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, + grid: base.Grid, + request: pytest.FixtureRequest, + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + vertical_wind_advective_tendency = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) + vertical_cfl = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) + normal_wind_advective_tendency = data_alloc.zero_field(dims.EdgeDim, dims.KDim) + + vn = data_alloc.random_field(dims.EdgeDim, dims.KDim) + w = data_alloc.random_field(dims.CellDim, dims.KHalfDim) + tangential_wind = data_alloc.random_field(dims.EdgeDim, dims.KDim) + tangential_wind_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) + vn_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) + horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( + dims.EdgeDim, dims.KDim + ) + contravariant_correction_at_cells_on_half_levels = data_alloc.random_field( + dims.CellDim, dims.KHalfDim + ) + + coeff1_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) + coeff2_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) + c_intp = data_alloc.random_field(dims.VertexDim, dims.V2CDim) + inv_dual_edge_length = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + inv_primal_edge_length = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + tangent_orientation = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim) + # positive thicknesses: the CFL clipping compares |w| * dtime / ddqz_z_half with the limit + ddqz_z_half = data_alloc.random_field(dims.CellDim, dims.KHalfDim, low=0.0) + area = data_alloc.random_field(dims.CellDim) + geofac_n2s = data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + owner_mask = data_alloc.random_mask(dims.CellDim) + coriolis_frequency = data_alloc.random_field(dims.EdgeDim) + geofac_rot = data_alloc.random_field(dims.VertexDim, dims.V2EDim) + coeff_gradekin = data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + c_lin_e = data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + # low=0.0 makes sure the simplified stencil produces the same result as the numpy version + ddqz_z_full_e = data_alloc.random_field(dims.EdgeDim, dims.KDim, low=0.0) + area_edge = data_alloc.random_field(dims.EdgeDim) + geofac_grdiv = data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + + dtime = 2.0 + + # value is set to reflect the MCH ch1 experiment. Changing it changes the runtime + end_index_of_damping_layer = 12 + + edge_domain = h_grid.domain(dims.EdgeDim) + cell_domain = h_grid.domain(dims.CellDim) + start_cell_lateral_boundary_level_4 = grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4) + ) + end_cell_halo = grid.end_index(cell_domain(h_grid.Zone.HALO)) + start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) + end_edge_local = grid.end_index(edge_domain(h_grid.Zone.LOCAL)) + + return dict( + vertical_wind_advective_tendency=vertical_wind_advective_tendency, + vertical_cfl=vertical_cfl, + normal_wind_advective_tendency=normal_wind_advective_tendency, + vn=vn, + w=w, + tangential_wind=tangential_wind, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + coriolis_frequency=coriolis_frequency, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + apply_extra_diffusion_on_vn=request.param["apply_extra_diffusion_on_vn"], + end_index_of_damping_layer=end_index_of_damping_layer, + start_cell_lateral_boundary_level_4=start_cell_lateral_boundary_level_4, + end_cell_halo=end_cell_halo, + start_edge_nudging_level_2=start_edge_nudging_level_2, + end_edge_local=end_edge_local, + vertical_start=0, + vertical_end=gtx.int32(grid.num_levels), + ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_predictor.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_predictor.py new file mode 100644 index 0000000000..06514f19e8 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_predictor.py @@ -0,0 +1,426 @@ +# 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 Any + +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_predictor import ( + compute_velocity_advection_in_predictor_step, +) +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.grid import base, horizontal as h_grid +from icon4py.model.common.states import utils as state_utils +from icon4py.model.testing import stencil_tests + +from .test_velocity_advection_terms import ( + _restore_outside, + compute_advection_in_horizontal_momentum_numpy, + compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy, + compute_diagnostics_from_normal_wind_numpy, + compute_horizontal_advection_of_w_numpy, + compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy, + interpolate_contravariant_correction_to_cells_on_half_levels_numpy, + interpolate_contravariant_vertical_velocity_to_full_levels_numpy, +) + + +@pytest.mark.embedded_remap_error +@pytest.mark.uses_concat_where +@pytest.mark.continuous_benchmarking +class TestComputeVelocityAdvectionInPredictorStep(stencil_tests.StencilTest): + PROGRAM = compute_velocity_advection_in_predictor_step + OUTPUTS = ( + "tangential_wind", + "tangential_wind_on_half_levels", + "vn_on_half_levels", + "horizontal_kinetic_energy_at_edges_on_model_levels", + "contravariant_correction_at_edges_on_model_levels", + "contravariant_correction_at_cells_on_half_levels", + "vertical_wind_advective_tendency", + "vertical_cfl", + "normal_wind_advective_tendency", + ) + STATIC_PARAMS = { + stencil_tests.StandardStaticVariants.NONE: (), + stencil_tests.StandardStaticVariants.COMPILE_TIME_VERTICAL: ( + "vertical_start", + "vertical_end", + "nflatlev", + "end_index_of_damping_layer", + "skip_compute_predictor_vertical_advection", + "apply_extra_diffusion_on_vn", + ), + stencil_tests.StandardStaticVariants.COMPILE_TIME_DOMAIN: ( + "start_edge_lateral_boundary_level_5", + "end_edge_halo_level_2", + "start_cell_lateral_boundary_level_4", + "end_cell_halo", + "start_edge_nudging_level_2", + "end_edge_local", + "vertical_start", + "vertical_end", + "nflatlev", + "end_index_of_damping_layer", + "skip_compute_predictor_vertical_advection", + "apply_extra_diffusion_on_vn", + ), + } + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + tangential_wind: np.ndarray, + tangential_wind_on_half_levels: np.ndarray, + vn_on_half_levels: np.ndarray, + horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, + contravariant_correction_at_edges_on_model_levels: np.ndarray, + contravariant_correction_at_cells_on_half_levels: np.ndarray, + vertical_wind_advective_tendency: np.ndarray, + vertical_cfl: np.ndarray, + normal_wind_advective_tendency: np.ndarray, + vn: np.ndarray, + w: np.ndarray, + rbf_vec_coeff_e: np.ndarray, + wgtfac_e: np.ndarray, + wgtfacq_e: np.ndarray, + ddxn_z_full: np.ndarray, + ddxt_z_full: np.ndarray, + coeff1_dwdz: np.ndarray, + coeff2_dwdz: np.ndarray, + c_intp: np.ndarray, + inv_dual_edge_length: np.ndarray, + inv_primal_edge_length: np.ndarray, + tangent_orientation: np.ndarray, + e_bln_c_s: np.ndarray, + wgtfac_c: np.ndarray, + ddqz_z_half: np.ndarray, + area: np.ndarray, + geofac_n2s: np.ndarray, + owner_mask: np.ndarray, + coriolis_frequency: np.ndarray, + geofac_rot: np.ndarray, + coeff_gradekin: np.ndarray, + c_lin_e: np.ndarray, + ddqz_z_full_e: np.ndarray, + area_edge: np.ndarray, + geofac_grdiv: np.ndarray, + dtime: ta.wpfloat, + skip_compute_predictor_vertical_advection: bool, + apply_extra_diffusion_on_vn: bool, + nflatlev: int, + end_index_of_damping_layer: int, + start_edge_lateral_boundary_level_5: int, + end_edge_halo_level_2: int, + start_cell_lateral_boundary_level_4: int, + end_cell_halo: int, + start_edge_nudging_level_2: int, + end_edge_local: int, + vertical_start: int, + vertical_end: int, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + nlev = vertical_end + + ( + tangential_wind_new, + tangential_wind_on_half_levels_new, + vn_on_half_levels_new, + horizontal_kinetic_energy_new, + contravariant_correction_at_edges_new, + ) = compute_diagnostics_from_normal_wind_numpy( + connectivities=connectivities, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn=vn, + rbf_vec_coeff_e=rbf_vec_coeff_e, + wgtfac_e=wgtfac_e, + wgtfacq_e=wgtfacq_e, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + skip_compute_predictor_vertical_advection=skip_compute_predictor_vertical_advection, + nlev=nlev, + ) + + contravariant_correction_at_cells_on_half_levels_new = interpolate_contravariant_correction_to_cells_on_half_levels_numpy( + connectivities=connectivities, + contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_new, + e_bln_c_s=e_bln_c_s, + wgtfac_c=wgtfac_c, + nflatlev=nflatlev, + nlev=nlev, + ) + + ( + contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping, + vertical_cfl_new, + ) = compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( + w=w[:, :-1], + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels_new, + ddqz_z_half=ddqz_z_half[:, :-1], + dtime=dtime, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + vertical_wind_advective_tendency_new = vertical_wind_advective_tendency[:, :-1] + if not skip_compute_predictor_vertical_advection: + horizontal_advection_of_w_at_edges_on_half_levels = ( + compute_horizontal_advection_of_w_numpy( + connectivities=connectivities, + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels_new, + vn_on_half_levels=vn_on_half_levels_new, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + ) + ) + vertical_wind_advective_tendency_new = compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( + connectivities=connectivities, + vertical_wind_advective_tendency=vertical_wind_advective_tendency_new, + w=w, + horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + cfl_clipping=cfl_clipping, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + e_bln_c_s=e_bln_c_s, + ddqz_z_half=ddqz_z_half[:, :-1], + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + dtime=dtime, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + contravariant_corrected_w_at_cells_on_model_levels = ( + interpolate_contravariant_vertical_velocity_to_full_levels_numpy( + contravariant_corrected_w_at_cells_on_half_levels + ) + ) + + normal_wind_advective_tendency_new = compute_advection_in_horizontal_momentum_numpy( + connectivities=connectivities, + vn=vn, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_new, + tangential_wind=tangential_wind_new, + coriolis_frequency=coriolis_frequency, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + vn_on_half_levels=vn_on_half_levels_new, + e_bln_c_s=e_bln_c_s, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + apply_extra_diffusion_on_vn=apply_extra_diffusion_on_vn, + nlev=nlev, + end_index_of_damping_layer=end_index_of_damping_layer, + ) + + edge_slice = (start_edge_lateral_boundary_level_5, end_edge_halo_level_2) + cell_slice = (start_cell_lateral_boundary_level_4, end_cell_halo) + + return dict( + tangential_wind=_restore_outside( + tangential_wind_new, tangential_wind, edge_slice, (vertical_start, vertical_end) + ), + tangential_wind_on_half_levels=_restore_outside( + tangential_wind_on_half_levels_new, + tangential_wind_on_half_levels, + edge_slice, + (vertical_start, vertical_end), + ), + vn_on_half_levels=_restore_outside( + vn_on_half_levels_new, + vn_on_half_levels, + edge_slice, + (vertical_start, vertical_end + 1), + ), + horizontal_kinetic_energy_at_edges_on_model_levels=_restore_outside( + horizontal_kinetic_energy_new, + horizontal_kinetic_energy_at_edges_on_model_levels, + edge_slice, + (vertical_start, vertical_end), + ), + contravariant_correction_at_edges_on_model_levels=_restore_outside( + contravariant_correction_at_edges_new, + contravariant_correction_at_edges_on_model_levels, + edge_slice, + (nflatlev, vertical_end), + ), + contravariant_correction_at_cells_on_half_levels=_restore_outside( + contravariant_correction_at_cells_on_half_levels_new, + contravariant_correction_at_cells_on_half_levels, + cell_slice, + (vertical_start, vertical_end), + ), + vertical_wind_advective_tendency=_restore_outside( + vertical_wind_advective_tendency_new, + vertical_wind_advective_tendency, + cell_slice, + # ICON computes the tendency over jk = 2..nlev + # (mo_velocity_advection.f90:598), so the top half level is untouched. + (vertical_start + 1, vertical_end), + ), + vertical_cfl=_restore_outside( + vertical_cfl_new, vertical_cfl, cell_slice, (vertical_start, vertical_end) + ), + normal_wind_advective_tendency=_restore_outside( + normal_wind_advective_tendency_new, + normal_wind_advective_tendency, + (start_edge_nudging_level_2, end_edge_local), + (vertical_start, vertical_end), + ), + ) + + @stencil_tests.input_data_fixture( + params=[ + { + "skip_compute_predictor_vertical_advection": skip, + "apply_extra_diffusion_on_vn": diffu, + } + for skip, diffu in ((False, True), (True, False)) + ], + ids=lambda param: ( + f"skip_compute_predictor_vertical_advection[{param['skip_compute_predictor_vertical_advection']}]" + f"-apply_extra_diffusion_on_vn[{param['apply_extra_diffusion_on_vn']}]" + ), + ) + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, + grid: base.Grid, + request: pytest.FixtureRequest, + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + tangential_wind = data_alloc.random_field(dims.EdgeDim, dims.KDim) + tangential_wind_on_half_levels = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) + vn_on_half_levels = data_alloc.zero_field(dims.EdgeDim, dims.KHalfDim) + horizontal_kinetic_energy_at_edges_on_model_levels = data_alloc.random_field( + dims.EdgeDim, dims.KDim + ) + contravariant_correction_at_edges_on_model_levels = data_alloc.random_field( + dims.EdgeDim, dims.KDim + ) + contravariant_correction_at_cells_on_half_levels = data_alloc.zero_field( + dims.CellDim, dims.KHalfDim + ) + vertical_wind_advective_tendency = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) + vertical_cfl = data_alloc.zero_field(dims.CellDim, dims.KHalfDim) + normal_wind_advective_tendency = data_alloc.zero_field(dims.EdgeDim, dims.KDim) + + vn = data_alloc.random_field(dims.EdgeDim, dims.KDim) + w = data_alloc.random_field(dims.CellDim, dims.KHalfDim) + + rbf_vec_coeff_e = data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim) + wgtfac_e = data_alloc.random_field(dims.EdgeDim, dims.KHalfDim) + wgtfacq_e = data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddxn_z_full = data_alloc.random_field(dims.EdgeDim, dims.KDim) + ddxt_z_full = data_alloc.random_field(dims.EdgeDim, dims.KDim) + coeff1_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) + coeff2_dwdz = data_alloc.random_field(dims.CellDim, dims.KDim) + c_intp = data_alloc.random_field(dims.VertexDim, dims.V2CDim) + inv_dual_edge_length = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + inv_primal_edge_length = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + tangent_orientation = data_alloc.random_field(dims.EdgeDim, low=1.0e-5) + e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim) + wgtfac_c = data_alloc.random_field(dims.CellDim, dims.KHalfDim) + # positive thicknesses: the CFL clipping compares |w| * dtime / ddqz_z_half with the limit + ddqz_z_half = data_alloc.random_field(dims.CellDim, dims.KHalfDim, low=0.0) + area = data_alloc.random_field(dims.CellDim) + geofac_n2s = data_alloc.random_field(dims.CellDim, dims.C2E2CODim) + owner_mask = data_alloc.random_mask(dims.CellDim) + coriolis_frequency = data_alloc.random_field(dims.EdgeDim) + geofac_rot = data_alloc.random_field(dims.VertexDim, dims.V2EDim) + coeff_gradekin = data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + c_lin_e = data_alloc.random_field(dims.EdgeDim, dims.E2CDim) + # low=0.0 makes sure the simplified stencil produces the same result as the numpy version + ddqz_z_full_e = data_alloc.random_field(dims.EdgeDim, dims.KDim, low=0.0) + area_edge = data_alloc.random_field(dims.EdgeDim) + geofac_grdiv = data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim) + + dtime = 2.0 + + # values are set to reflect the MCH ch1 experiment. Changing them changes the runtime + nflatlev = 5 + end_index_of_damping_layer = 12 + + edge_domain = h_grid.domain(dims.EdgeDim) + cell_domain = h_grid.domain(dims.CellDim) + start_edge_lateral_boundary_level_5 = grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_5) + ) + end_edge_halo_level_2 = grid.end_index(edge_domain(h_grid.Zone.HALO_LEVEL_2)) + start_cell_lateral_boundary_level_4 = grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_4) + ) + end_cell_halo = grid.end_index(cell_domain(h_grid.Zone.HALO)) + start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) + end_edge_local = grid.end_index(edge_domain(h_grid.Zone.LOCAL)) + + return dict( + tangential_wind=tangential_wind, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels=contravariant_correction_at_edges_on_model_levels, + contravariant_correction_at_cells_on_half_levels=contravariant_correction_at_cells_on_half_levels, + vertical_wind_advective_tendency=vertical_wind_advective_tendency, + vertical_cfl=vertical_cfl, + normal_wind_advective_tendency=normal_wind_advective_tendency, + vn=vn, + w=w, + rbf_vec_coeff_e=rbf_vec_coeff_e, + wgtfac_e=wgtfac_e, + wgtfacq_e=wgtfacq_e, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + e_bln_c_s=e_bln_c_s, + wgtfac_c=wgtfac_c, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + owner_mask=owner_mask, + coriolis_frequency=coriolis_frequency, + geofac_rot=geofac_rot, + coeff_gradekin=coeff_gradekin, + c_lin_e=c_lin_e, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + geofac_grdiv=geofac_grdiv, + dtime=dtime, + skip_compute_predictor_vertical_advection=request.param[ + "skip_compute_predictor_vertical_advection" + ], + apply_extra_diffusion_on_vn=request.param["apply_extra_diffusion_on_vn"], + nflatlev=nflatlev, + end_index_of_damping_layer=end_index_of_damping_layer, + start_edge_lateral_boundary_level_5=start_edge_lateral_boundary_level_5, + end_edge_halo_level_2=end_edge_halo_level_2, + start_cell_lateral_boundary_level_4=start_cell_lateral_boundary_level_4, + end_cell_halo=end_cell_halo, + start_edge_nudging_level_2=start_edge_nudging_level_2, + end_edge_local=end_edge_local, + vertical_start=0, + vertical_end=gtx.int32(grid.num_levels), + ) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_terms.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_terms.py new file mode 100644 index 0000000000..41b29473cf --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_velocity_advection_terms.py @@ -0,0 +1,1150 @@ +# 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 collections.abc import Mapping +from typing import Any + +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.dycore.stencils.velocity_advection_terms import ( + VerticalCflConstants, + _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask, + _clip_contravariant_corrected_w, + _compute_advective_normal_wind_tendency, + _compute_cfl, + _compute_extra_diffusion, + _compute_extra_diffusion_for_w, + _compute_horizontal_advection_of_w, + _compute_interpolated_horizontal_advection_of_w, + _compute_vertical_advection_of_w, + _interpolate_contravariant_vertical_velocity_to_full_levels, +) +from icon4py.model.common import dimension as dims, type_alias as ta +from icon4py.model.common.grid import base, horizontal as h_grid +from icon4py.model.common.states import utils as state_utils +from icon4py.model.testing import reference_funcs, stencil_tests +from icon4py.model.testing.reference_funcs import ( + compute_curl_numpy, + interpolate_cell_field_to_vertex_numpy, + interpolate_to_cell_center_numpy, +) + +from .test_compute_contravariant_correction import compute_contravariant_correction_numpy +from .test_interpolate_cell_field_to_half_levels import ( + interpolate_cell_field_to_half_levels_vp_numpy, +) +from .test_interpolate_vt_to_interface_edges import interpolate_vt_to_interface_edges_numpy + + +def interpolate_vn_to_half_levels_numpy(wgtfac_e: np.ndarray, vn: np.ndarray) -> np.ndarray: + nlev = vn.shape[1] + vn_ie = np.zeros((vn.shape[0], nlev + 1)) + w = wgtfac_e[:, 1:nlev] + vn_ie[:, 1:nlev] = w * vn[:, 1:nlev] + (1.0 - w) * vn[:, 0 : nlev - 1] + vn_ie[:, 0] = vn[:, 0] + return vn_ie + + +def compute_horizontal_kinetic_energy_at_edges_numpy(vn: np.ndarray, vt: np.ndarray) -> np.ndarray: + return 0.5 * (vn * vn + vt * vt) + + +def extrapolate_to_surface_numpy(wgtfacq_e: np.ndarray, vn: np.ndarray) -> np.ndarray: + vn_k_minus_1 = vn[:, -1] + vn_k_minus_2 = vn[:, -2] + vn_k_minus_3 = vn[:, -3] + wgtfacq_e_k_minus_1 = wgtfacq_e[:, -1] + wgtfacq_e_k_minus_2 = wgtfacq_e[:, -2] + wgtfacq_e_k_minus_3 = wgtfacq_e[:, -3] + vn_at_surface = ( + wgtfacq_e_k_minus_1 * vn_k_minus_1 + + wgtfacq_e_k_minus_2 * vn_k_minus_2 + + wgtfacq_e_k_minus_3 * vn_k_minus_3 + ) + return vn_at_surface + + +def compute_diagnostics_from_normal_wind_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + tangential_wind_on_half_levels: np.ndarray, + vn: np.ndarray, + rbf_vec_coeff_e: np.ndarray, + wgtfac_e: np.ndarray, + wgtfacq_e: np.ndarray, + ddxn_z_full: np.ndarray, + ddxt_z_full: np.ndarray, + skip_compute_predictor_vertical_advection: bool, + nlev: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + tangential_wind = reference_funcs.compute_tangential_wind_numpy( + connectivities, vn, rbf_vec_coeff_e + ) + horizontal_kinetic_energy_at_edges_on_model_levels = ( + compute_horizontal_kinetic_energy_at_edges_numpy(vn, tangential_wind) + ) + vn_on_half_levels = interpolate_vn_to_half_levels_numpy(wgtfac_e, vn) + vn_on_half_levels[:, nlev] = extrapolate_to_surface_numpy(wgtfacq_e, vn) + + tangential_wind_on_half_levels = tangential_wind_on_half_levels.copy() + if not skip_compute_predictor_vertical_advection: + tangential_wind_on_half_levels[:, :nlev] = interpolate_vt_to_interface_edges_numpy( + wgtfac_e, tangential_wind + )[:, :nlev] + + contravariant_correction_at_edges_on_model_levels = compute_contravariant_correction_numpy( + vn, ddxn_z_full, ddxt_z_full, tangential_wind + ) + + return ( + tangential_wind, + tangential_wind_on_half_levels, + vn_on_half_levels, + horizontal_kinetic_energy_at_edges_on_model_levels, + contravariant_correction_at_edges_on_model_levels, + ) + + +def interpolate_contravariant_correction_to_cells_on_half_levels_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + contravariant_correction_at_edges_on_model_levels: np.ndarray, + e_bln_c_s: np.ndarray, + wgtfac_c: np.ndarray, + nflatlev: int, + nlev: int, +) -> np.ndarray: + k = np.arange(nlev) + + contravariant_correction_at_cells_model_levels = interpolate_to_cell_center_numpy( + connectivities, contravariant_correction_at_edges_on_model_levels, e_bln_c_s + ) + + return np.where( + k >= nflatlev + 1, + interpolate_cell_field_to_half_levels_vp_numpy( + wgtfac_c=wgtfac_c, interpolant=contravariant_correction_at_cells_model_levels + )[:, :-1], + 0.0, + ) + + +def interpolate_contravariant_vertical_velocity_to_full_levels_numpy( + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, +) -> np.ndarray: + num_rows, num_cols = contravariant_corrected_w_at_cells_on_half_levels.shape + contravariant_corrected_w_with_surface = np.zeros((num_rows, num_cols + 1)) + contravariant_corrected_w_with_surface[:, :-1] = ( + contravariant_corrected_w_at_cells_on_half_levels + ) + return 0.5 * ( + contravariant_corrected_w_with_surface[:, :-1] + + contravariant_corrected_w_with_surface[:, 1:] + ) + + +def compute_maximum_cfl_and_clip_contravariant_vertical_velocity_numpy( + *, + w: np.ndarray, + contravariant_correction_at_cells_on_half_levels: np.ndarray, + ddqz_z_half: np.ndarray, + dtime: ta.wpfloat, + nlev: int, + end_index_of_damping_layer: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + num_rows, num_cols = contravariant_correction_at_cells_on_half_levels.shape + + k = np.arange(num_cols) + condition = (np.maximum(2, end_index_of_damping_layer - 2) <= k) & (k < nlev - 3) + + contravariant_corrected_w_at_cells_on_half_levels = ( + w - contravariant_correction_at_cells_on_half_levels + ) + + cfl_clipping = np.where( + ( + np.abs(contravariant_corrected_w_at_cells_on_half_levels) * dtime / ddqz_z_half + > VerticalCflConstants.W_LIMIT + ) + & condition, + np.ones([num_rows, num_cols]), + np.zeros_like(contravariant_corrected_w_at_cells_on_half_levels), + ) + vertical_cfl = np.where( + cfl_clipping == 1.0, + contravariant_corrected_w_at_cells_on_half_levels * dtime / ddqz_z_half, + 0.0, + ) + contravariant_corrected_w_at_cells_on_half_levels = np.where( + (cfl_clipping == 1.0) & (vertical_cfl < -0.85), + -0.85 * ddqz_z_half / dtime, + contravariant_corrected_w_at_cells_on_half_levels, + ) + contravariant_corrected_w_at_cells_on_half_levels = np.where( + (cfl_clipping == 1.0) & (vertical_cfl > 0.85), + 0.85 * ddqz_z_half / dtime, + contravariant_corrected_w_at_cells_on_half_levels, + ) + + return contravariant_corrected_w_at_cells_on_half_levels, cfl_clipping, vertical_cfl + + +def compute_horizontal_advection_of_w_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + w: np.ndarray, + tangential_wind_on_half_levels: np.ndarray, + vn_on_half_levels: np.ndarray, + c_intp: np.ndarray, + inv_dual_edge_length: np.ndarray, + inv_primal_edge_length: np.ndarray, + tangent_orientation: np.ndarray, +) -> np.ndarray: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + e2c = connectivities[dims.E2C] + e2v = connectivities[dims.E2V] + w_at_vertices = interpolate_cell_field_to_vertex_numpy(connectivities, w, c_intp) + + return vn_on_half_levels * inv_dual_edge_length * (w[e2c[:, 0]] - w[e2c[:, 1]]) + ( + tangential_wind_on_half_levels + * inv_primal_edge_length + * tangent_orientation + * (w_at_vertices[e2v[:, 0]] - w_at_vertices[e2v[:, 1]]) + ) + + +def add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + cfl_clipping: np.ndarray, + owner_mask: np.ndarray, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + ddqz_z_half: np.ndarray, + area: np.ndarray, + geofac_n2s: np.ndarray, + w: np.ndarray, + vertical_wind_advective_tendency: np.ndarray, + dtime: ta.wpfloat, +) -> np.ndarray: + owner_mask = np.expand_dims(owner_mask, axis=-1) + area = np.expand_dims(area, axis=-1) + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + + difcoef = np.where( + (cfl_clipping == 1) & (owner_mask == 1), + VerticalCflConstants.EXTRA_DIFFUSION_SCALING + / dtime + * np.minimum( + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + np.abs(contravariant_corrected_w_at_cells_on_half_levels) * dtime / ddqz_z_half + - VerticalCflConstants.W_LIMIT, + ), + 0, + ) + + c2e2cO = connectivities[dims.C2E2CO] + return np.where( + (cfl_clipping == 1) & (owner_mask == 1), + vertical_wind_advective_tendency + + difcoef + * area + * np.sum( + np.where( + (c2e2cO != -1)[:, :, np.newaxis], + w[c2e2cO] * geofac_n2s, + 0, + ), + axis=1, + ), + vertical_wind_advective_tendency, + ) + + +def compute_advective_vertical_wind_tendency_numpy( + z_w_con_c: np.ndarray, + w: np.ndarray, + coeff1_dwdz: np.ndarray, + coeff2_dwdz: np.ndarray, + **kwargs: Any, +) -> np.ndarray: + # coeff*_dwdz live on model levels; model level k pairs with half level k + nlev = coeff1_dwdz.shape[1] + ddt_w_adv = np.zeros((z_w_con_c.shape[0], nlev + 1)) + c1, c2 = coeff1_dwdz[:, 1:nlev], coeff2_dwdz[:, 1:nlev] + ddt_w_adv[:, 1:nlev] = -z_w_con_c[:, 1:nlev] * ( + w[:, 0 : nlev - 1] * c1 - w[:, 2 : nlev + 1] * c2 + w[:, 1:nlev] * (c2 - c1) + ) + return ddt_w_adv + + +def compute_advective_vertical_wind_tendency_and_apply_diffusion_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + vertical_wind_advective_tendency: np.ndarray, + w: np.ndarray, + horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + cfl_clipping: np.ndarray, + coeff1_dwdz: np.ndarray, + coeff2_dwdz: np.ndarray, + e_bln_c_s: np.ndarray, + ddqz_z_half: np.ndarray, + area: np.ndarray, + geofac_n2s: np.ndarray, + owner_mask: np.ndarray, + dtime: ta.wpfloat, + nlev: int, + end_index_of_damping_layer: int, +) -> np.ndarray: + k = np.arange(nlev) + + condition1 = k >= 1 + vertical_wind_advective_tendency = np.where( + condition1, + compute_advective_vertical_wind_tendency_numpy( + contravariant_corrected_w_at_cells_on_half_levels, w, coeff1_dwdz, coeff2_dwdz + )[:, :-1], + vertical_wind_advective_tendency, + ) + + vertical_wind_advective_tendency = np.where( + condition1, + vertical_wind_advective_tendency + + compute_interpolated_horizontal_advection_of_w_numpy( + connectivities, + e_bln_c_s, + horizontal_advection_of_w_at_edges_on_half_levels[:, :-1], + ), + vertical_wind_advective_tendency, + ) + + condition2 = (np.maximum(2, end_index_of_damping_layer - 2) <= k) & (k < nlev - 3) + + return np.where( + condition2, + add_extra_diffusion_for_w_approaching_cfl_wihtout_levmask_numpy( + connectivities=connectivities, + cfl_clipping=cfl_clipping, + owner_mask=owner_mask, + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + w=w[:, :-1], + vertical_wind_advective_tendency=vertical_wind_advective_tendency, + dtime=dtime, + ), + vertical_wind_advective_tendency, + ) + + +def _compute_advective_normal_wind_tendency_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, + coeff_gradekin: np.ndarray, + horizontal_kinetic_energy_at_cells_on_model_levels: np.ndarray, + upward_vorticity_at_vertices: np.ndarray, + tangential_wind: np.ndarray, + coriolis_frequency: np.ndarray, + c_lin_e: np.ndarray, + contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, + vn_on_half_levels: np.ndarray, + ddqz_z_full_e: np.ndarray, +) -> np.ndarray: + e2c = connectivities[dims.E2C] + horizontal_kinetic_energy_at_cells_on_model_levels_e2c = ( + horizontal_kinetic_energy_at_cells_on_model_levels[e2c] + ) + coeff_gradekin = np.expand_dims(coeff_gradekin, axis=-1) + coriolis_frequency = np.expand_dims(coriolis_frequency, axis=-1) + c_lin_e = np.expand_dims(c_lin_e, axis=-1) + + return -( + (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) + * horizontal_kinetic_energy_at_edges_on_model_levels + + ( + -coeff_gradekin[:, 0] * horizontal_kinetic_energy_at_cells_on_model_levels_e2c[:, 0] + + coeff_gradekin[:, 1] * horizontal_kinetic_energy_at_cells_on_model_levels_e2c[:, 1] + ) + + tangential_wind + * ( + coriolis_frequency + + 0.5 * np.sum(upward_vorticity_at_vertices[connectivities[dims.E2V]], axis=1) + ) + + np.sum(contravariant_corrected_w_at_cells_on_model_levels[e2c] * c_lin_e, axis=1) + * (vn_on_half_levels[:, :-1] - vn_on_half_levels[:, 1:]) + / ddqz_z_full_e + ) + + +def _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + c_lin_e: np.ndarray, + contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, + ddqz_z_full_e: np.ndarray, + area_edge: np.ndarray, + tangent_orientation: np.ndarray, + inv_primal_edge_length: np.ndarray, + upward_vorticity_at_vertices: np.ndarray, + geofac_grdiv: np.ndarray, + vn: np.ndarray, + normal_wind_advective_tendency: np.ndarray, + dtime: ta.wpfloat, +) -> np.ndarray: + c_lin_e = np.expand_dims(c_lin_e, axis=-1) + geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) + area_edge = np.expand_dims(area_edge, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) + + e2c = connectivities[dims.E2C] + contravariant_corrected_w_at_edges_on_model_levels = np.sum( + np.where( + (e2c != -1)[:, :, np.newaxis], + c_lin_e * contravariant_corrected_w_at_cells_on_model_levels[e2c], + 0, + ), + axis=1, + ) + + difcoef = np.where( + ( + np.abs(contravariant_corrected_w_at_edges_on_model_levels) * dtime / ddqz_z_full_e + > VerticalCflConstants.W_LIMIT + ), + VerticalCflConstants.EXTRA_DIFFUSION_SCALING + / dtime + * np.minimum( + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + np.abs(contravariant_corrected_w_at_edges_on_model_levels) * dtime / ddqz_z_full_e + - VerticalCflConstants.W_LIMIT, + ), + np.zeros_like(vn), + ) + e2v = connectivities[dims.E2V] + e2c2eo = connectivities[dims.E2C2EO] + return np.where( + ( + np.abs(contravariant_corrected_w_at_edges_on_model_levels) * dtime / ddqz_z_full_e + > VerticalCflConstants.W_LIMIT + ), + normal_wind_advective_tendency + + difcoef + * area_edge + * ( + np.sum( + np.where( + (e2c2eo != -1)[:, :, np.newaxis], + geofac_grdiv * vn[e2c2eo], + 0, + ), + axis=1, + ) + + tangent_orientation + * inv_primal_edge_length + * (upward_vorticity_at_vertices[e2v][:, 1] - upward_vorticity_at_vertices[e2v][:, 0]) + ), + normal_wind_advective_tendency, + ) + + +def compute_advection_in_horizontal_momentum_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + vn: np.ndarray, + horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, + tangential_wind: np.ndarray, + coriolis_frequency: np.ndarray, + contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, + vn_on_half_levels: np.ndarray, + e_bln_c_s: np.ndarray, + geofac_rot: np.ndarray, + coeff_gradekin: np.ndarray, + c_lin_e: np.ndarray, + ddqz_z_full_e: np.ndarray, + area_edge: np.ndarray, + tangent_orientation: np.ndarray, + inv_primal_edge_length: np.ndarray, + geofac_grdiv: np.ndarray, + dtime: ta.wpfloat, + apply_extra_diffusion_on_vn: bool, + nlev: int, + end_index_of_damping_layer: int, +) -> np.ndarray: + k = np.arange(nlev) + + horizontal_kinetic_energy_at_cells_on_model_levels = interpolate_to_cell_center_numpy( + connectivities, horizontal_kinetic_energy_at_edges_on_model_levels, e_bln_c_s + ) + upward_vorticity_at_vertices = compute_curl_numpy(connectivities, vn, geofac_rot) + + normal_wind_advective_tendency = _compute_advective_normal_wind_tendency_numpy( + connectivities=connectivities, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + coeff_gradekin=coeff_gradekin, + horizontal_kinetic_energy_at_cells_on_model_levels=horizontal_kinetic_energy_at_cells_on_model_levels, + upward_vorticity_at_vertices=upward_vorticity_at_vertices, + tangential_wind=tangential_wind, + coriolis_frequency=coriolis_frequency, + c_lin_e=c_lin_e, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + vn_on_half_levels=vn_on_half_levels, + ddqz_z_full_e=ddqz_z_full_e, + ) + + if apply_extra_diffusion_on_vn: + condition = (np.maximum(2, end_index_of_damping_layer - 2) <= k) & (k < nlev - 4) + normal_wind_advective_tendency = np.where( + condition, + _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask_numpy( + connectivities=connectivities, + c_lin_e=c_lin_e, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + upward_vorticity_at_vertices=upward_vorticity_at_vertices, + geofac_grdiv=geofac_grdiv, + vn=vn, + normal_wind_advective_tendency=normal_wind_advective_tendency, + dtime=dtime, + ), + normal_wind_advective_tendency, + ) + + return normal_wind_advective_tendency + + +def _restore_outside( + computed: np.ndarray, + initial: np.ndarray, + horizontal: tuple[int, int], + vertical: tuple[int, int], +) -> np.ndarray: + """Return `computed` on the given domain and `initial` everywhere else.""" + domain = (slice(*horizontal), slice(*vertical)) + result = initial.copy() + result[domain] = computed[domain] + return result + + +def compute_interpolated_horizontal_advection_of_w_numpy( + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + e_bln_c_s: np.ndarray, + horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, + **kwargs: Any, +) -> np.ndarray: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + c2e = connectivities[dims.C2E] + return np.sum( + horizontal_advection_of_w_at_edges_on_half_levels[c2e] * e_bln_c_s, + axis=1, + ) + + +def compute_extra_diffusion_for_w_numpy( + *, + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + ddqz_z_half: np.ndarray, + area: np.ndarray, + geofac_n2s: np.ndarray, + w: np.ndarray, + dtime: ta.wpfloat, +) -> np.ndarray: + area = np.expand_dims(area, axis=-1) + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + + difcoef = (VerticalCflConstants.EXTRA_DIFFUSION_SCALING / dtime) * np.minimum( + VerticalCflConstants.W_MAX - VerticalCflConstants.W_LIMIT, + np.abs(contravariant_corrected_w_at_cells_on_half_levels) * dtime / ddqz_z_half + - VerticalCflConstants.W_LIMIT, + ) + + c2e2cO = connectivities[dims.C2E2CO] + return ( + difcoef + * area + * np.sum( + np.where( + (c2e2cO != -1)[:, :, np.newaxis], + w[c2e2cO] * geofac_n2s, + 0, + ), + axis=1, + ) + ) + + +class TestComputeInterpolatedHorizontalAdvectionOfW(stencil_tests.StencilTest): + PROGRAM = _compute_interpolated_horizontal_advection_of_w + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + e_bln_c_s: np.ndarray, + horizontal_advection_of_w_at_edges_on_half_levels: np.ndarray, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + return dict( + out=compute_interpolated_horizontal_advection_of_w_numpy( + connectivities, + e_bln_c_s, + horizontal_advection_of_w_at_edges_on_half_levels, + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + e_bln_c_s = data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat) + horizontal_advection_of_w_at_edges_on_half_levels = data_alloc.random_field( + dims.EdgeDim, dims.KHalfDim, dtype=ta.vpfloat + ) + interpolated_horizontal_advection_of_w = data_alloc.random_field( + dims.CellDim, dims.KHalfDim, dtype=ta.wpfloat + ) + + return dict( + e_bln_c_s=e_bln_c_s, + horizontal_advection_of_w_at_edges_on_half_levels=horizontal_advection_of_w_at_edges_on_half_levels, + out=interpolated_horizontal_advection_of_w, + domain={ + dims.CellDim: (0, gtx.int32(grid.num_cells)), + dims.KHalfDim: (0, gtx.int32(grid.num_levels + 1)), + }, + ) + + +@pytest.mark.embedded_remap_error +class TestComputeExtraDiffusionForW(stencil_tests.StencilTest): + PROGRAM = _compute_extra_diffusion_for_w + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + ddqz_z_half: np.ndarray, + area: np.ndarray, + geofac_n2s: np.ndarray, + w: np.ndarray, + dtime: ta.wpfloat, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + return dict( + out=compute_extra_diffusion_for_w_numpy( + connectivities=connectivities, + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + w=w, + dtime=dtime, + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + contravariant_corrected_w_at_cells_on_half_levels = data_alloc.random_field( + dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat + ) + ddqz_z_half = data_alloc.random_field( + dims.CellDim, dims.KHalfDim, low=0.5, high=1.5, dtype=ta.vpfloat + ) + area = data_alloc.random_field(dims.CellDim, dtype=ta.wpfloat) + geofac_n2s = data_alloc.random_field(dims.CellDim, dims.C2E2CODim, dtype=ta.wpfloat) + w = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.wpfloat) + extra_diffusion = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.wpfloat) + + return dict( + contravariant_corrected_w_at_cells_on_half_levels=contravariant_corrected_w_at_cells_on_half_levels, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + w=w, + dtime=ta.wpfloat("2.0"), + out=extra_diffusion, + domain={ + dims.CellDim: (0, gtx.int32(grid.num_cells)), + dims.KHalfDim: (0, gtx.int32(grid.num_levels + 1)), + }, + ) + + +class TestComputeCfl(stencil_tests.StencilTest): + PROGRAM = _compute_cfl + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + ddqz_z_half: np.ndarray, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + dtime: ta.wpfloat, + **kwargs: Any, + ) -> dict: + cfl_clipping = ( + np.abs(contravariant_corrected_w_at_cells_on_half_levels) * dtime / ddqz_z_half + > VerticalCflConstants.W_LIMIT + ) + vertical_cfl = np.where( + cfl_clipping, + contravariant_corrected_w_at_cells_on_half_levels * dtime / ddqz_z_half, + 0.0, + ) + return dict(out=(cfl_clipping, vertical_cfl)) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + return dict( + ddqz_z_half=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, low=0.5, high=1.5, dtype=ta.vpfloat + ), + contravariant_corrected_w_at_cells_on_half_levels=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat + ), + dtime=ta.wpfloat("2.0"), + out=( + data_alloc.random_mask(dims.CellDim, dims.KHalfDim), + data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat), + ), + domain={ + dims.CellDim: (0, gtx.int32(grid.num_cells)), + dims.KHalfDim: (0, gtx.int32(grid.num_levels + 1)), + }, + ) + + +class TestClipContravariantCorrectedW(stencil_tests.StencilTest): + PROGRAM = _clip_contravariant_corrected_w + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + cfl_clipping: np.ndarray, + vertical_cfl: np.ndarray, + ddqz_z_half: np.ndarray, + dtime: ta.wpfloat, + **kwargs: Any, + ) -> dict: + clipped = np.where( + cfl_clipping & (vertical_cfl < -0.85), + -0.85 * ddqz_z_half / dtime, + contravariant_corrected_w_at_cells_on_half_levels, + ) + clipped = np.where( + cfl_clipping & (vertical_cfl > 0.85), 0.85 * ddqz_z_half / dtime, clipped + ) + return dict(out=clipped) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + return dict( + contravariant_corrected_w_at_cells_on_half_levels=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat + ), + cfl_clipping=data_alloc.random_mask(dims.CellDim, dims.KHalfDim), + vertical_cfl=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, low=-2.0, high=2.0, dtype=ta.vpfloat + ), + ddqz_z_half=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, low=0.5, high=1.5, dtype=ta.vpfloat + ), + dtime=ta.wpfloat("2.0"), + out=data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat), + domain={ + dims.CellDim: (0, gtx.int32(grid.num_cells)), + dims.KHalfDim: (0, gtx.int32(grid.num_levels + 1)), + }, + ) + + +class TestComputeExtraDiffusion(stencil_tests.StencilTest): + PROGRAM = _compute_extra_diffusion + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + vn: np.ndarray, + upward_vorticity_at_vertices_on_model_levels: np.ndarray, + difcoef: np.ndarray, + area_edge: np.ndarray, + geofac_grdiv: np.ndarray, + tangent_orientation: np.ndarray, + inv_primal_edge_length: np.ndarray, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + e2c2eo = connectivities[dims.E2C2EO] + e2v = connectivities[dims.E2V] + gradient_of_divergence_of_vn = np.sum( + np.where( + (e2c2eo != -1)[:, :, np.newaxis], + vn[e2c2eo] * np.expand_dims(geofac_grdiv, axis=-1), + 0.0, + ), + axis=1, + ) + gradient_of_vorticity = np.expand_dims( + tangent_orientation * inv_primal_edge_length, axis=-1 + ) * ( + upward_vorticity_at_vertices_on_model_levels[e2v[:, 1]] + - upward_vorticity_at_vertices_on_model_levels[e2v[:, 0]] + ) + return dict( + out=difcoef + * np.expand_dims(area_edge, axis=-1) + * (gradient_of_divergence_of_vn + gradient_of_vorticity) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + return dict( + vn=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat), + upward_vorticity_at_vertices_on_model_levels=data_alloc.random_field( + dims.VertexDim, dims.KDim, dtype=ta.vpfloat + ), + difcoef=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat), + area_edge=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + geofac_grdiv=data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim, dtype=ta.wpfloat), + tangent_orientation=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + inv_primal_edge_length=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + out=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat), + domain={ + dims.EdgeDim: (0, gtx.int32(grid.num_edges)), + dims.KDim: (0, gtx.int32(grid.num_levels)), + }, + ) + + +class TestComputeVerticalAdvectionOfW(stencil_tests.StencilTest): + PROGRAM = _compute_vertical_advection_of_w + # the operator reads w at K-1 and K+1, so it is only defined on the interior half levels + OUTPUTS = ( + stencil_tests.Output( + "out", refslice=(slice(None), slice(1, -1)), gtslice=(slice(None), slice(1, -1)) + ), + ) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + w: np.ndarray, + coeff1_dwdz: np.ndarray, + coeff2_dwdz: np.ndarray, + **kwargs: Any, + ) -> dict: + return dict( + out=compute_advective_vertical_wind_tendency_numpy( + z_w_con_c=contravariant_corrected_w_at_cells_on_half_levels, + w=w, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + return dict( + contravariant_corrected_w_at_cells_on_half_levels=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat + ), + w=data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.wpfloat), + coeff1_dwdz=data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat), + coeff2_dwdz=data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat), + out=data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat), + domain={ + dims.CellDim: (0, gtx.int32(grid.num_cells)), + dims.KHalfDim: (1, gtx.int32(grid.num_levels)), + }, + ) + + +class TestInterpolateContravariantVerticalVelocityToFullLevels(stencil_tests.StencilTest): + PROGRAM = _interpolate_contravariant_vertical_velocity_to_full_levels + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + contravariant_corrected_w_at_cells_on_half_levels: np.ndarray, + **kwargs: Any, + ) -> dict: + return dict( + out=interpolate_contravariant_vertical_velocity_to_full_levels_numpy( + contravariant_corrected_w_at_cells_on_half_levels[:, : grid.num_levels] + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + return dict( + contravariant_corrected_w_at_cells_on_half_levels=data_alloc.random_field( + dims.CellDim, dims.KHalfDim, dtype=ta.vpfloat + ), + nlev=gtx.int32(grid.num_levels), + out=data_alloc.random_field(dims.CellDim, dims.KDim, dtype=ta.vpfloat), + domain={ + dims.CellDim: (0, gtx.int32(grid.num_cells)), + dims.KDim: (0, gtx.int32(grid.num_levels)), + }, + ) + + +class TestComputeAdvectiveNormalWindTendency(stencil_tests.StencilTest): + PROGRAM = _compute_advective_normal_wind_tendency + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + horizontal_kinetic_energy_at_edges_on_model_levels: np.ndarray, + upward_vorticity_at_vertices_on_model_levels: np.ndarray, + tangential_wind: np.ndarray, + vn_on_half_levels: np.ndarray, + contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, + coriolis_frequency: np.ndarray, + e_bln_c_s: np.ndarray, + c_lin_e: np.ndarray, + coeff_gradekin: np.ndarray, + ddqz_z_full_e: np.ndarray, + out: np.ndarray, + domain: dict, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + normal_wind_advective_tendency = _compute_advective_normal_wind_tendency_numpy( + connectivities=connectivities, + horizontal_kinetic_energy_at_edges_on_model_levels=horizontal_kinetic_energy_at_edges_on_model_levels, + coeff_gradekin=coeff_gradekin, + horizontal_kinetic_energy_at_cells_on_model_levels=interpolate_to_cell_center_numpy( + connectivities, horizontal_kinetic_energy_at_edges_on_model_levels, e_bln_c_s + ), + upward_vorticity_at_vertices=upward_vorticity_at_vertices_on_model_levels, + tangential_wind=tangential_wind, + coriolis_frequency=coriolis_frequency, + c_lin_e=c_lin_e, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + vn_on_half_levels=vn_on_half_levels, + ddqz_z_full_e=ddqz_z_full_e, + ) + return dict( + out=_restore_outside( + normal_wind_advective_tendency, out, domain[dims.EdgeDim], domain[dims.KDim] + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + # The operator reads both E2C neighbours unmasked, so it runs where they exist. + edge_domain = h_grid.domain(dims.EdgeDim) + start_edge_nudging_level_2 = grid.start_index(edge_domain(h_grid.Zone.NUDGING_LEVEL_2)) + end_edge_local = grid.end_index(edge_domain(h_grid.Zone.LOCAL)) + + return dict( + horizontal_kinetic_energy_at_edges_on_model_levels=data_alloc.random_field( + dims.EdgeDim, dims.KDim, dtype=ta.vpfloat + ), + upward_vorticity_at_vertices_on_model_levels=data_alloc.random_field( + dims.VertexDim, dims.KDim, dtype=ta.vpfloat + ), + tangential_wind=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat), + vn_on_half_levels=data_alloc.random_field( + dims.EdgeDim, dims.KHalfDim, dtype=ta.vpfloat + ), + contravariant_corrected_w_at_cells_on_model_levels=data_alloc.random_field( + dims.CellDim, dims.KDim, dtype=ta.vpfloat + ), + coriolis_frequency=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + e_bln_c_s=data_alloc.random_field(dims.CellDim, dims.C2EDim, dtype=ta.wpfloat), + c_lin_e=data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat), + coeff_gradekin=data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dtype=ta.vpfloat), + ddqz_z_full_e=data_alloc.random_field( + dims.EdgeDim, dims.KDim, low=0.5, high=1.5, dtype=ta.vpfloat + ), + out=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat), + domain={ + dims.EdgeDim: (start_edge_nudging_level_2, end_edge_local), + dims.KDim: (0, gtx.int32(grid.num_levels)), + }, + ) + + +class TestAddExtraDiffusionForNormalWindTendencyWithoutLevelmask(stencil_tests.StencilTest): + PROGRAM = _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + c_lin_e: np.ndarray, + contravariant_corrected_w_at_cells_on_model_levels: np.ndarray, + ddqz_z_full_e: np.ndarray, + area_edge: np.ndarray, + tangent_orientation: np.ndarray, + inv_primal_edge_length: np.ndarray, + upward_vorticity_at_vertices_on_model_levels: np.ndarray, + geofac_grdiv: np.ndarray, + vn: np.ndarray, + normal_wind_advective_tendency: np.ndarray, + dtime: ta.wpfloat, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + return dict( + out=_add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_without_levelmask_numpy( + connectivities=connectivities, + c_lin_e=c_lin_e, + contravariant_corrected_w_at_cells_on_model_levels=contravariant_corrected_w_at_cells_on_model_levels, + ddqz_z_full_e=ddqz_z_full_e, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + upward_vorticity_at_vertices=upward_vorticity_at_vertices_on_model_levels, + geofac_grdiv=geofac_grdiv, + vn=vn, + normal_wind_advective_tendency=normal_wind_advective_tendency, + dtime=dtime, + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + dtime = ta.wpfloat("2.0") + return dict( + c_lin_e=data_alloc.random_field(dims.EdgeDim, dims.E2CDim, dtype=ta.wpfloat), + contravariant_corrected_w_at_cells_on_model_levels=data_alloc.random_field( + dims.CellDim, dims.KDim, dtype=ta.vpfloat + ), + ddqz_z_full_e=data_alloc.random_field( + dims.EdgeDim, dims.KDim, low=0.5, high=1.5, dtype=ta.vpfloat + ), + area_edge=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + tangent_orientation=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + inv_primal_edge_length=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + upward_vorticity_at_vertices_on_model_levels=data_alloc.random_field( + dims.VertexDim, dims.KDim, dtype=ta.vpfloat + ), + geofac_grdiv=data_alloc.random_field(dims.EdgeDim, dims.E2C2EODim, dtype=ta.wpfloat), + vn=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat), + normal_wind_advective_tendency=data_alloc.random_field( + dims.EdgeDim, dims.KDim, dtype=ta.vpfloat + ), + dtime=dtime, + out=data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat), + domain={ + dims.EdgeDim: (0, gtx.int32(grid.num_edges)), + dims.KDim: (0, gtx.int32(grid.num_levels)), + }, + ) + + +class TestComputeHorizontalAdvectionOfW(stencil_tests.StencilTest): + PROGRAM = _compute_horizontal_advection_of_w + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + w: np.ndarray, + tangential_wind_on_half_levels: np.ndarray, + vn_on_half_levels: np.ndarray, + c_intp: np.ndarray, + inv_dual_edge_length: np.ndarray, + inv_primal_edge_length: np.ndarray, + tangent_orientation: np.ndarray, + out: np.ndarray, + domain: dict, + **kwargs: Any, + ) -> dict: + horizontal_advection_of_w = compute_horizontal_advection_of_w_numpy( + connectivities=stencil_tests.connectivities_asnumpy(grid), + w=w, + tangential_wind_on_half_levels=tangential_wind_on_half_levels, + vn_on_half_levels=vn_on_half_levels, + c_intp=c_intp, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + ) + return dict( + out=_restore_outside( + horizontal_advection_of_w, out, domain[dims.EdgeDim], domain[dims.KHalfDim] + ) + ) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + # The operator reads both E2C neighbours unmasked, so it runs where they exist. + edge_domain = h_grid.domain(dims.EdgeDim) + start_edge_lateral_boundary_level_7 = grid.start_index( + edge_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_7) + ) + end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) + + return dict( + w=data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=ta.wpfloat), + tangential_wind_on_half_levels=data_alloc.random_field( + dims.EdgeDim, dims.KHalfDim, dtype=ta.wpfloat + ), + vn_on_half_levels=data_alloc.random_field( + dims.EdgeDim, dims.KHalfDim, dtype=ta.vpfloat + ), + c_intp=data_alloc.random_field(dims.VertexDim, dims.V2CDim, dtype=ta.wpfloat), + inv_dual_edge_length=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + inv_primal_edge_length=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + tangent_orientation=data_alloc.random_field(dims.EdgeDim, dtype=ta.wpfloat), + out=data_alloc.random_field(dims.EdgeDim, dims.KHalfDim, dtype=ta.vpfloat), + domain={ + dims.EdgeDim: (start_edge_lateral_boundary_level_7, end_edge_halo), + dims.KHalfDim: (0, gtx.int32(grid.num_levels + 1)), + }, + ) diff --git a/model/atmosphere/tracer_advection/src/icon4py/model/atmosphere/tracer_advection/tracer_advection_horizontal.py b/model/atmosphere/tracer_advection/src/icon4py/model/atmosphere/tracer_advection/tracer_advection_horizontal.py index aff8f0a4de..9a775a5c3a 100644 --- a/model/atmosphere/tracer_advection/src/icon4py/model/atmosphere/tracer_advection/tracer_advection_horizontal.py +++ b/model/atmosphere/tracer_advection/src/icon4py/model/atmosphere/tracer_advection/tracer_advection_horizontal.py @@ -45,7 +45,7 @@ from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid from icon4py.model.common.interpolation.stencils.compute_tangential_wind import ( - compute_tangential_wind_wp, + compute_tangential_wind, ) from icon4py.model.common.math.stencils import generic_math_operations from icon4py.model.common.states import tracer_prep_adv_states as prep_adv_states @@ -652,9 +652,9 @@ def __init__( ) # stencils - self._compute_tangential_wind_wp = model_options.setup_program( + self._compute_tangential_wind = model_options.setup_program( backend=self._backend, - program=compute_tangential_wind_wp, + program=compute_tangential_wind, constant_args={ "rbf_vec_coeff_e": self._interpolation_state.rbf_vec_coeff_e, }, @@ -724,12 +724,12 @@ def _compute_numerical_flux( ## tracer-independent part # compute tangential velocity - log.debug("running stencil compute_tangential_wind_wp - start") - self._compute_tangential_wind_wp( + log.debug("running stencil compute_tangential_wind - start") + self._compute_tangential_wind( vn=prep_adv.vn_traj, vt=self._z_real_vt, ) - log.debug("running stencil compute_tangential_wind_wp - end") + log.debug("running stencil compute_tangential_wind - end") # backtrajectory calculation log.debug("running stencil compute_barycentric_backtrajectory_alt - start") diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py b/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py deleted file mode 100644 index 1e31e37b6b..0000000000 --- a/model/common/src/icon4py/model/common/interpolation/stencils/compute_cell_2_vertex_interpolation.py +++ /dev/null @@ -1,55 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import gt4py.next as gtx -from gt4py.next import neighbor_sum - -import icon4py.model.common.type_alias as types -from icon4py.model.common import dimension as dims, field_type_aliases as fa -from icon4py.model.common.dimension import V2C - - -@gtx.field_operator -def _compute_cell_2_vertex_interpolation( - cell_in: fa.CellKHalfField[types.wpfloat], - c_int: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], types.wpfloat], -) -> fa.VertexKHalfField[types.wpfloat]: - vert_out = neighbor_sum(c_int * cell_in(V2C), axis=dims.V2CDim) - return vert_out - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_cell_2_vertex_interpolation( - cell_in: fa.CellKHalfField[types.wpfloat], - c_int: gtx.Field[[dims.VertexDim, dims.V2CDim], types.wpfloat], - vert_out: fa.VertexKHalfField[types.wpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - """ - Compute the interpolation from cell to vertex field. - - Args: - cell_in: input cell field - c_int: interpolation coefficients - vert_out: (output) vertex field - horizontal_start: horizontal start index - horizontal_end: horizontal end index - vertical_start: vertical start index - vertical_end: vertical end index - """ - _compute_cell_2_vertex_interpolation( - cell_in=cell_in, - c_int=c_int, - out=vert_out, - domain={ - dims.VertexDim: (horizontal_start, horizontal_end), - dims.KHalfDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/compute_tangential_wind.py b/model/common/src/icon4py/model/common/interpolation/stencils/compute_tangential_wind.py index d972c0a0bb..35cdfb0f3b 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/compute_tangential_wind.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/compute_tangential_wind.py @@ -6,15 +6,15 @@ # 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 astype, neighbor_sum +from gt4py.next import neighbor_sum from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import E2C2E -from icon4py.model.common.type_alias import vpfloat, wpfloat +from icon4py.model.common.type_alias import wpfloat @gtx.field_operator -def _compute_tangential_wind_wp( +def _compute_tangential_wind( vn: fa.EdgeKField[wpfloat], rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], wpfloat], ) -> fa.EdgeKField[wpfloat]: @@ -37,19 +37,8 @@ def _compute_tangential_wind_wp( return neighbor_sum(rbf_vec_coeff_e * vn(E2C2E), axis=dims.E2C2EDim) -@gtx.field_operator -def _compute_tangential_wind_vp( - vn: fa.EdgeKField[wpfloat], - rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], wpfloat], -) -> fa.EdgeKField[vpfloat]: - """ - Variable-precision variant of ``_compute_tangential_wind_wp``. - """ - return astype(_compute_tangential_wind_wp(vn, rbf_vec_coeff_e), vpfloat) - - @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_tangential_wind_wp( +def compute_tangential_wind( vn: fa.EdgeKField[wpfloat], rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], wpfloat], vt: fa.EdgeKField[wpfloat], @@ -58,28 +47,7 @@ def compute_tangential_wind_wp( vertical_start: gtx.int32, vertical_end: gtx.int32, ) -> None: - _compute_tangential_wind_wp( - vn=vn, - rbf_vec_coeff_e=rbf_vec_coeff_e, - out=vt, - domain={ - dims.EdgeDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def compute_tangential_wind_vp( - vn: fa.EdgeKField[wpfloat], - rbf_vec_coeff_e: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2C2EDim], wpfloat], - vt: fa.EdgeKField[vpfloat], - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -) -> None: - _compute_tangential_wind_vp( + _compute_tangential_wind( vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e, out=vt, diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/interpolate_cell_field_to_vertex.py b/model/common/src/icon4py/model/common/interpolation/stencils/interpolate_cell_field_to_vertex.py new file mode 100644 index 0000000000..1dd2bdf744 --- /dev/null +++ b/model/common/src/icon4py/model/common/interpolation/stencils/interpolate_cell_field_to_vertex.py @@ -0,0 +1,22 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +from gt4py.next import neighbor_sum + +import icon4py.model.common.type_alias as types +from icon4py.model.common import dimension as dims, field_type_aliases as fa +from icon4py.model.common.dimension import V2C + + +@gtx.field_operator +def _interpolate_cell_field_to_vertex( + cell_in: fa.CellKHalfField[types.wpfloat], + c_int: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2CDim], types.wpfloat], +) -> fa.VertexKHalfField[types.wpfloat]: + vert_out = neighbor_sum(c_int * cell_in(V2C), axis=dims.V2CDim) + return vert_out diff --git a/model/common/src/icon4py/model/common/math/stencils/compute_curl.py b/model/common/src/icon4py/model/common/math/stencils/compute_curl.py new file mode 100644 index 0000000000..6a0cfd8321 --- /dev/null +++ b/model/common/src/icon4py/model/common/math/stencils/compute_curl.py @@ -0,0 +1,22 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next import neighbor_sum + +from icon4py.model.common import dimension as dims, field_type_aliases as fa +from icon4py.model.common.dimension import V2E +from icon4py.model.common.type_alias import wpfloat + + +@gtx.field_operator +def _compute_curl( + vec_e: fa.EdgeKField[wpfloat], + geofac_rot: gtx.Field[gtx.Dims[dims.VertexDim, dims.V2EDim], wpfloat], +) -> fa.VertexKField[wpfloat]: + return neighbor_sum(vec_e(V2E) * geofac_rot, axis=dims.V2EDim) 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 cb2b622e8b..464037a97c 100644 --- a/model/common/src/icon4py/model/common/metrics/metric_fields.py +++ b/model/common/src/icon4py/model/common/metrics/metric_fields.py @@ -35,8 +35,8 @@ _cell_2_edge_interpolation, _cell_2_edge_interpolation_on_half_levels, ) -from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( - _compute_cell_2_vertex_interpolation, +from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_vertex import ( + _interpolate_cell_field_to_vertex, ) from icon4py.model.common.math.gradient import _grad_fd_tang, grad_fd_norm from icon4py.model.common.math.vertical_operations import ( @@ -361,7 +361,7 @@ def _compute_ddxt_z_half_e( inv_primal_edge_length: fa.EdgeField[wpfloat], tangent_orientation: fa.EdgeField[wpfloat], ): - z_ifv = _compute_cell_2_vertex_interpolation(cell_in, c_int) + z_ifv = _interpolate_cell_field_to_vertex(cell_in, c_int) ddxt_z_half_e = _grad_fd_tang( z_ifv, inv_primal_edge_length, diff --git a/model/common/tests/common/interpolation/stencil_tests/test_compute_tangential_wind.py b/model/common/tests/common/interpolation/stencil_tests/test_compute_tangential_wind.py index 46232535e6..c197a0b4bd 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_compute_tangential_wind.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_compute_tangential_wind.py @@ -15,8 +15,7 @@ from icon4py.model.common import dimension as dims, type_alias as ta from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.interpolation.stencils.compute_tangential_wind import ( - compute_tangential_wind_vp, - compute_tangential_wind_wp, + compute_tangential_wind, ) from icon4py.model.testing import reference_funcs, stencil_tests @@ -76,7 +75,7 @@ def tangential_wind_input_data( class TestComputeTangentialWindWpHalfLevels(stencil_tests.StencilTest): """Half-level input (nlev + 1 rows).""" - PROGRAM = compute_tangential_wind_wp + PROGRAM = compute_tangential_wind OUTPUTS = ("vt",) @stencil_tests.static_reference @@ -93,7 +92,7 @@ def input_data( class TestComputeTangentialWindWpFullLevels(stencil_tests.StencilTest): """Full-level input (nlev rows).""" - PROGRAM = compute_tangential_wind_wp + PROGRAM = compute_tangential_wind OUTPUTS = ("vt",) @stencil_tests.static_reference @@ -105,41 +104,3 @@ def input_data( data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid ) -> dict[str, Any]: return tangential_wind_input_data(data_alloc, grid, on_half_levels=False) - - -@pytest.mark.embedded_remap_error -class TestComputeTangentialWind(stencil_tests.StencilTest): - """Variable-precision variant used by the dycore velocity advection.""" - - PROGRAM = compute_tangential_wind_vp - OUTPUTS = ("vt",) - - @stencil_tests.static_reference - def reference( - grid: base.Grid, - *, - vn: np.ndarray, - rbf_vec_coeff_e: np.ndarray, - **kwargs: Any, - ) -> dict: - connectivities = stencil_tests.connectivities_asnumpy(grid) - vt = reference_funcs.compute_tangential_wind_numpy(connectivities, vn, rbf_vec_coeff_e) - return dict(vt=vt) - - @stencil_tests.input_data_fixture - def input_data( - data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid - ) -> dict[str, Any]: - vn = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=ta.wpfloat) - rbf_vec_coeff_e = data_alloc.random_field(dims.EdgeDim, dims.E2C2EDim, dtype=ta.wpfloat) - vt = data_alloc.zero_field(dims.EdgeDim, dims.KDim, dtype=ta.vpfloat) - - return dict( - vn=vn, - rbf_vec_coeff_e=rbf_vec_coeff_e, - vt=vt, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_edges), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels), - ) diff --git a/model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py b/model/common/tests/common/interpolation/stencil_tests/test_interpolate_cell_field_to_vertex.py similarity index 59% rename from model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py rename to model/common/tests/common/interpolation/stencil_tests/test_interpolate_cell_field_to_vertex.py index 8dd80941f1..63e7ac1fda 100644 --- a/model/common/tests/common/interpolation/stencil_tests/test_compute_cell_2_vertex_interpolation.py +++ b/model/common/tests/common/interpolation/stencil_tests/test_interpolate_cell_field_to_vertex.py @@ -14,16 +14,16 @@ import icon4py.model.common.type_alias as types from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base -from icon4py.model.common.interpolation.stencils.compute_cell_2_vertex_interpolation import ( - compute_cell_2_vertex_interpolation, +from icon4py.model.common.interpolation.stencils.interpolate_cell_field_to_vertex import ( + _interpolate_cell_field_to_vertex, ) -from icon4py.model.testing import stencil_tests +from icon4py.model.testing import reference_funcs, stencil_tests @pytest.mark.skip_value_error -class TestComputeCells2VertsInterpolation(stencil_tests.StencilTest): - PROGRAM = compute_cell_2_vertex_interpolation - OUTPUTS = ("vert_out",) +class TestInterpolateCellFieldToVertex(stencil_tests.StencilTest): + PROGRAM = _interpolate_cell_field_to_vertex + OUTPUTS = ("out",) @stencil_tests.static_reference def reference( @@ -34,26 +34,22 @@ def reference( **kwargs: Any, ) -> dict: connectivities = stencil_tests.connectivities_asnumpy(grid) - v2c = connectivities[dims.V2C] - c_int = np.expand_dims(c_int, axis=-1) - out_field = np.sum(cell_in[v2c] * c_int, axis=1) - return dict( - vert_out=out_field, + out=reference_funcs.interpolate_cell_field_to_vertex_numpy( + connectivities, cell_in, c_int + ), ) @stencil_tests.input_data_fixture def input_data(data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid) -> dict: cell_in = data_alloc.random_field(dims.CellDim, dims.KHalfDim, dtype=types.wpfloat) c_int = data_alloc.random_field(dims.VertexDim, dims.V2CDim, dtype=types.wpfloat) - vert_out = data_alloc.zero_field(dims.VertexDim, dims.KHalfDim, dtype=types.wpfloat) - return dict( cell_in=cell_in, c_int=c_int, - vert_out=vert_out, - horizontal_start=0, - horizontal_end=gtx.int32(grid.num_vertices), - vertical_start=0, - vertical_end=gtx.int32(grid.num_levels + 1), + out=data_alloc.zero_field(dims.VertexDim, dims.KHalfDim, dtype=types.wpfloat), + domain={ + dims.VertexDim: (0, gtx.int32(grid.num_vertices)), + dims.KHalfDim: (0, gtx.int32(grid.num_levels + 1)), + }, ) diff --git a/model/common/tests/common/math/stencil_tests/test_compute_curl.py b/model/common/tests/common/math/stencil_tests/test_compute_curl.py new file mode 100644 index 0000000000..3e7203d69f --- /dev/null +++ b/model/common/tests/common/math/stencil_tests/test_compute_curl.py @@ -0,0 +1,50 @@ +# 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 Any + +import gt4py.next as gtx +import numpy as np + +from icon4py.model.common import dimension as dims +from icon4py.model.common.grid import base +from icon4py.model.common.math.stencils.compute_curl import _compute_curl +from icon4py.model.common.states import utils as state_utils +from icon4py.model.common.type_alias import wpfloat +from icon4py.model.testing import reference_funcs, stencil_tests + + +class TestComputeCurl(stencil_tests.StencilTest): + PROGRAM = _compute_curl + OUTPUTS = ("out",) + + @stencil_tests.static_reference + def reference( + grid: base.Grid, + *, + vec_e: np.ndarray, + geofac_rot: np.ndarray, + **kwargs: Any, + ) -> dict: + connectivities = stencil_tests.connectivities_asnumpy(grid) + return dict(out=reference_funcs.compute_curl_numpy(connectivities, vec_e, geofac_rot)) + + @stencil_tests.input_data_fixture + def input_data( + data_alloc: stencil_tests.DataAllocationWrapper, grid: base.Grid + ) -> dict[str, gtx.Field | state_utils.ScalarType]: + vec_e = data_alloc.random_field(dims.EdgeDim, dims.KDim, dtype=wpfloat) + geofac_rot = data_alloc.random_field(dims.VertexDim, dims.V2EDim, dtype=wpfloat) + return dict( + vec_e=vec_e, + geofac_rot=geofac_rot, + out=data_alloc.zero_field(dims.VertexDim, dims.KDim, dtype=wpfloat), + domain={ + dims.VertexDim: (0, gtx.int32(grid.num_vertices)), + dims.KDim: (0, gtx.int32(grid.num_levels)), + }, + ) diff --git a/model/testing/src/icon4py/model/testing/reference_funcs.py b/model/testing/src/icon4py/model/testing/reference_funcs.py index 0e015aa21d..7aebd9979c 100644 --- a/model/testing/src/icon4py/model/testing/reference_funcs.py +++ b/model/testing/src/icon4py/model/testing/reference_funcs.py @@ -73,3 +73,23 @@ def interpolate_to_cell_center_numpy( e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) c2e = connectivities[dims.C2E] return np.sum(interpolant[c2e] * e_bln_c_s, axis=1) + + +def interpolate_cell_field_to_vertex_numpy( + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + cell_field: np.ndarray, + c_intp: np.ndarray, +) -> np.ndarray: + v2c = connectivities[dims.V2C] + c_intp = np.expand_dims(c_intp, axis=-1) + return np.sum(np.where((v2c != -1)[:, :, np.newaxis], cell_field[v2c] * c_intp, 0), axis=1) + + +def compute_curl_numpy( + connectivities: Mapping[gtx.FieldOffset, np.ndarray], + edge_field: np.ndarray, + geofac_rot: np.ndarray, +) -> np.ndarray: + v2e = connectivities[dims.V2E] + geofac_rot = np.expand_dims(geofac_rot, axis=-1) + return np.sum(np.where((v2e != -1)[:, :, np.newaxis], edge_field[v2e] * geofac_rot, 0), axis=1) diff --git a/scripts/python/compare_icon4py_openacc.py b/scripts/python/compare_icon4py_openacc.py index 0faf601863..67a996d93f 100755 --- a/scripts/python/compare_icon4py_openacc.py +++ b/scripts/python/compare_icon4py_openacc.py @@ -46,6 +46,9 @@ # besides the gt4py stencil name, a dictionary of static arguments that should be # matched in the gt4py timer report. If the value is `None`, we do not check the # static arguments and assume the stencil name is the same. +# The velocity advection kernels have no entry: 'compute_velocity_advection_in_predictor_step' +# and 'compute_velocity_advection_in_corrector_step' each aggregate several openacc kernels, +# which this one-icon4py-variant-per-fortran-kernel mapping cannot express. fortran_to_icon4py: dict[str, VariantDescriptor | None] = { "apply_diffusion_to_theta_and_exner": None, "apply_diffusion_to_vn": None, @@ -56,14 +59,6 @@ "calculate_diagnostic_quantities_for_turbulence": None, "calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools": None, "calculate_nabla2_and_smag_coefficients_for_vn": None, - "compute_advection_in_horizontal_momentum_equation": ( - "compute_advection_in_horizontal_momentum", - {}, - ), - "compute_advection_in_vertical_momentum_equation": ( - "compute_advection_in_corrector_vertical_momentum", - {}, - ), "compute_averaged_vn_and_fluxes_and_prepare_tracer_advection": ( "compute_averaged_vn_and_fluxes", { @@ -78,26 +73,6 @@ "prepare_fluxes_for_advection": True, }, ), - "compute_advection_in_predictor_vertical_momentum": ( - "compute_advection_in_predictor_vertical_momentum", - { - "skip_compute_predictor_vertical_advection": False, - }, - ), - "compute_advection_in_predictor_vertical_momentum_skip": ( - "compute_advection_in_predictor_vertical_momentum", - { - "skip_compute_predictor_vertical_advection": True, - }, - ), - "compute_diagnostics_from_normal_wind": ( - "compute_diagnostics_from_normal_wind", - {"skip_compute_predictor_vertical_advection": False}, - ), - "compute_diagnostics_from_normal_wind_skip": ( - "compute_diagnostics_from_normal_wind", - {"skip_compute_predictor_vertical_advection": True}, - ), "compute_horizontal_velocity_quantities_and_fluxes": None, "compute_perturbed_quantities_and_interpolation": None, "compute_rho_theta_pgrad_and_update_vn": (