add doc string on the combined stencils on velocity advection

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."""

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),

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),

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,

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,

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),

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(