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 06b62413e..a858fb2cf 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 @@ -404,7 +404,10 @@ def _validate(self): raise NotImplementedError("divdamp_order can only be 24") if self.divdamp_type == DivergenceDampingType.TWO_DIMENSIONAL: - raise NotImplementedError("`DivergenceDampingType.TWO_DIMENSIONAL` (2) is not yet implemented") + raise NotImplementedError( + "`DivergenceDampingType.TWO_DIMENSIONAL` (2) is not yet implemented" + ) + class NonHydrostaticParams: """Calculates derived quantities depending on the NonHydrostaticConfig.""" 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 index 75c663628..d729e6a5f 100644 --- 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 @@ -43,7 +43,7 @@ def _compute_advection_in_horizontal_momentum_equation( d_time: ta.wpfloat, levelmask: fa.KField[bool], nlev: gtx.int32, - nrdmax: 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) @@ -61,7 +61,7 @@ def _compute_advection_in_horizontal_momentum_equation( ) normal_wind_advective_tendency = concat_where( - ((maximum(3, nrdmax - 2) - 1) <= dims.KDim) & (dims.KDim < (nlev - 4)), + ((maximum(3, end_index_of_damping_layer - 2) - 1) <= dims.KDim) & (dims.KDim < (nlev - 4)), _add_extra_diffusion_for_normal_wind_tendency_approaching_cfl( levelmask, c_lin_e, @@ -106,13 +106,49 @@ def compute_advection_in_horizontal_momentum_equation( d_time: ta.wpfloat, levelmask: fa.KField[bool], nlev: gtx.int32, - nrdmax: 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, ): - """Formerly known as fused_velocity_advection_stencil_19_to_20.""" + """ + 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 + - 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 + - d_time: time step + - levelmask: mask for valid vertical levels + - nlev: number of (full/model) vertical levels + - 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 + - vertical_end: end index in the vertical domain + + Returns: + - normal_wind_advective_tendency: horizontal advection tendency of the normal wind + + """ _compute_advection_in_horizontal_momentum_equation( vn, @@ -135,7 +171,7 @@ def compute_advection_in_horizontal_momentum_equation( d_time, levelmask, nlev, - nrdmax, + end_index_of_damping_layer, out=normal_wind_advective_tendency, domain={ dims.EdgeDim: (horizontal_start, horizontal_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 index 44a4650dc..794212fc9 100644 --- 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 @@ -42,7 +42,7 @@ def _compute_advective_vertical_wind_tendency_and_apply_diffusion( cfl_w_limit: ta.vpfloat, dtime: ta.wpfloat, nlev: gtx.int32, - nrdmax: gtx.int32, + end_index_of_damping_layer: gtx.int32, ) -> fa.CellKField[ta.vpfloat]: # TODO(havogt): we should get rid of the cell_lower_bound and cell_upper_bound, # they are only to protect write to halo (if I understand correctly) @@ -63,7 +63,7 @@ def _compute_advective_vertical_wind_tendency_and_apply_diffusion( vertical_wind_advective_tendency, ) vertical_wind_advective_tendency = concat_where( - ((maximum(3, nrdmax - 2) - 1) <= dims.KDim) & (dims.KDim < (nlev - 3)), + ((maximum(3, end_index_of_damping_layer - 2) - 1) <= dims.KDim) & (dims.KDim < (nlev - 3)), _add_extra_diffusion_for_w_con_approaching_cfl( cfl_clipping, owner_mask, @@ -102,7 +102,7 @@ def _compute_advection_in_vertical_momentum_equation( cfl_clipping: fa.CellKField[bool], owner_mask: fa.CellField[bool], nlev: gtx.int32, - nrdmax: gtx.int32, + end_index_of_damping_layer: gtx.int32, ) -> tuple[fa.CellKField[ta.vpfloat], fa.CellKField[ta.vpfloat]]: contravariant_corrected_w_at_cells_on_model_levels = ( _interpolate_contravariant_vertical_velocity_to_full_levels( @@ -127,7 +127,7 @@ def _compute_advection_in_vertical_momentum_equation( cfl_w_limit, dtime, nlev, - nrdmax, + end_index_of_damping_layer, ) if not skip_compute_predictor_vertical_advection else vertical_wind_advective_tendency @@ -156,13 +156,47 @@ def compute_advection_in_vertical_momentum_equation( cfl_clipping: fa.CellKField[bool], owner_mask: fa.CellField[bool], nlev: gtx.int32, - nrdmax: 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, ): - """Formerly known as fused_velocity_advection_stencil_15_to_18.""" + """ + Formerly known as fused_velocity_advection_stencil_15_to_18. + + This computes the vertical momentum advection in the vertical momentum equation + + Args: + - contravariant_corrected_w_at_cells_on_model_levels: contravariant-corrected vertical velocity at model levels + - vertical_wind_advective_tendency: vertical advection tendency of the vertical wind + - w: vertical wind at cell centers + - contravariant_corrected_w_at_cells_on_half_levels: contravariant-corrected vertical velocity at cells on half levels + - horizontal_advection_of_w_at_edges_on_half_levels: horizontal advection for vertical velocity at edges 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) + - e_bln_c_s: interpolation field (edge-to-cell interpolation weights) + - ddqz_z_half: metrics field + - area: cell area + - geofac_n2s: interpolation field + - 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 + - cfl_clipping: boolean field indicating CFL clipping applied per cell and level + - owner_mask: ownership mask for each cell + - nlev: number of (full/model) vertical 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 + - vertical_end: end index in the vertical dimension + + Returns: + - contravariant_corrected_w_at_cells_on_model_levels + - vertical_wind_advective_tendency + + """ _compute_advection_in_vertical_momentum_equation( vertical_wind_advective_tendency, @@ -182,7 +216,7 @@ def compute_advection_in_vertical_momentum_equation( cfl_clipping, owner_mask, nlev, - nrdmax, + end_index_of_damping_layer, out=(contravariant_corrected_w_at_cells_on_model_levels, vertical_wind_advective_tendency), domain={ dims.CellDim: (horizontal_start, horizontal_end), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_velocity_advection.py index 6581ada42..4dbc19d5e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_cell_diagnostics_for_velocity_advection.py @@ -98,7 +98,34 @@ def interpolate_horizontal_kinetic_energy_to_cells_and_compute_contravariant_ter vertical_start: gtx.int32, vertical_end: gtx.int32, ): - """Formerly known as fused_velocity_advection_stencil_8_to_13_predictor.""" + """ + Formerly known as fused_velocity_advection_stencil_8_to_13_predictor. + + This interpolates horizontal kinetic energy from edges to cells on models levels + and compute the contravariant correction term on half levels. + It also computes the vertical velocity with the contravariant correction term. + + Args: + - horizontal_kinetic_energy_at_cells_on_model_levels: horizontal kinetic energy computed at cell centers of model levels + - contravariant_correction_at_cells_on_half_levels: contravariant metric correction at cells on half levels + - contravariant_corrected_w_at_cells_on_half_levels: contravariant-corrected vertical velocity at cells on half levels + - w: vertical wind at cell centers + - 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 + - e_bln_c_s: interpolation field (edge-to-cell interpolation weights factors) + - wgtfac_c: metrics field + - nflatlev: number of flat levels + - nlev: total number of vertical 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: + - horizontal_kinetic_energy_at_cells_on_model_levels + - contravariant_correction_at_cells_on_half_levels + - contravariant_corrected_w_at_cells_on_half_levels + """ _interpolate_horizontal_kinetic_energy_to_cells_and_compute_contravariant_correction( horizontal_kinetic_energy_at_edges_on_model_levels, @@ -145,7 +172,31 @@ def interpolate_horizontal_kinetic_energy_to_cells_and_compute_contravariant_cor vertical_start: gtx.int32, vertical_end: gtx.int32, ): - """Formerly known as fused_velocity_advection_stencil_8_to_13_corrector.""" + """ + Formerly known as fused_velocity_advection_stencil_8_to_13_corrector. + + This interpolates horizontal kinetic energy from edges to cells on model levels + The contravariant correction term has been computed at the end of predictor in solve nonhydro + It also computes the contravariant-corrected vertical velocity at cell centers on half levels + + Args: + - horizontal_kinetic_energy_at_cells_on_model_levels: horizontal kinetic energy computed at cell centers of model levels + - contravariant_corrected_w_at_cells_on_half_levels: contravariant-corrected vertical velocity at cells on half levels + - w: vertical wind at cell centers + - horizontal_kinetic_energy_at_edges_on_model_levels: horizontal kinetic energy computed at edge of model levels + - contravariant_correction_at_cells_on_half_levels: contravariant metric correction at cells on half levels + - e_bln_c_s: interpolation field (edge-to-cell interpolation weights factors) + - nflatlev: number of flat levels + - nlev: total number of vertical 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: + - horizontal_kinetic_energy_at_cells_on_model_levels + - contravariant_corrected_w_at_cells_on_half_levels + """ _interpolate_to_cell_center( horizontal_kinetic_energy_at_edges_on_model_levels, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_velocity_advection.py index d953c086f..04a85c032 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_velocity_advection.py @@ -167,7 +167,48 @@ def compute_derived_horizontal_winds_and_ke_and_horizontal_advection_of_w_and_co vertical_start: gtx.int32, vertical_end: gtx.int32, ): - """Formerly known as fused_velocity_advection_stencil_1_to_7_predictor.""" + """ + 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 in the x-direction) + - ddxt_z_full: metrics field (derivative 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_derived_horizontal_winds_and_ke_and_horizontal_advection_of_w_and_contravariant_correction( tangential_wind_on_half_levels, contravariant_correction_at_edges_on_model_levels, @@ -223,7 +264,29 @@ def compute_horizontal_advection_of_w( vertical_start: gtx.int32, vertical_end: gtx.int32, ): - """Formerly known as fused_velocity_advection_stencil_1_to_7_corrector.""" + """ + Formerly known as fused_velocity_advection_stencil_1_to_7_corrector. + + This computes the horizontal advection of the vertical wind + + Agrs: + - horizontal_advection_of_w_at_edges_on_half_levels: horizontal advection for vertical velocity + - w: vertical wind at cell centers + - tangential_wind_on_half_levels: tangential wind interpolated to half levels + - vn_on_half_levels: normal wind on half levels + - 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 + - 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: + - horizontal_advection_of_w_at_edges_on_half_levels + """ _compute_horizontal_advection_of_w( w, 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 index dafcea4a5..12ed71a6a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity_advection.py @@ -262,7 +262,7 @@ def run_predictor_step( cfl_clipping=self.cfl_clipping, owner_mask=self.c_owner_mask, nlev=gtx.int32(self.grid.num_levels), - nrdmax=self.vertical_params.nrdmax, + end_index_of_damping_layer=self.vertical_params.end_index_of_damping_layer, horizontal_start=self._start_cell_lateral_boundary_level_4, horizontal_end=self._end_cell_halo, vertical_start=0, @@ -297,7 +297,7 @@ def run_predictor_step( d_time=dtime, levelmask=levmask, # TODO(havogt): can we get rid of the levelmask here? nlev=self.grid.num_levels, - nrdmax=self.vertical_params.nrdmax, + end_index_of_damping_layer=self.vertical_params.end_index_of_damping_layer, horizontal_start=self._start_edge_nudging_level_2, horizontal_end=self._end_edge_local, vertical_start=gtx.int32(0), @@ -401,7 +401,7 @@ def run_corrector_step( cfl_clipping=self.cfl_clipping, owner_mask=self.c_owner_mask, nlev=gtx.int32(self.grid.num_levels), - nrdmax=self.vertical_params.nrdmax, + end_index_of_damping_layer=self.vertical_params.end_index_of_damping_layer, horizontal_start=self._start_cell_lateral_boundary_level_4, horizontal_end=self._end_cell_halo, vertical_start=0, @@ -436,7 +436,7 @@ def run_corrector_step( d_time=dtime, levelmask=levmask, # TODO(havogt): can we get rid of the levelmask here? nlev=self.grid.num_levels, - nrdmax=self.vertical_params.nrdmax, + end_index_of_damping_layer=self.vertical_params.end_index_of_damping_layer, horizontal_start=self._start_edge_nudging_level_2, horizontal_end=self._end_edge_local, vertical_start=gtx.int32(0), 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 index 6bc5d3482..f9ecd14b5 100644 --- 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 @@ -64,7 +64,9 @@ def reference( normal_wind_advective_tendency_cp = normal_wind_advective_tendency.copy() k = np.arange(nlev) - upward_vorticity_at_vertices = mo_math_divrot_rot_vertex_ri_dsl_numpy(connectivities, vn, geofac_rot) + 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, @@ -79,7 +81,7 @@ def reference( vn_on_half_levels, ddqz_z_full_e, ) - + condition = (np.maximum(3, nrdmax - 2) - 1 <= k) & (k < nlev - 4) normal_wind_advective_tendency_extra_diffu = ( add_extra_diffusion_for_normal_wind_tendency_approaching_cfl_numpy( diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index b03dfd971..9b3ae5a3b 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -213,14 +213,9 @@ def nflatlev(self) -> gtx.int32: """Vertical index for bottom most level at which coordinate surfaces are flat.""" return gtx.int32(self.index(Domain(dims.KDim, Zone.FLAT))) - @functools.cached_property - def nrdmax(self) -> gtx.int32: - """Vertical index where damping starts.""" - return self.end_index_of_damping_layer - @functools.cached_property def end_index_of_damping_layer(self) -> gtx.int32: - """Vertical index where damping starts.""" + """Vertical index where damping ends.""" return self.index(Domain(dims.KDim, Zone.DAMPING)) @property