Combine divergence damping coefficient calculation with other stencils

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py

@@ -506,6 +506,10 @@ def __init__(
                 "iau_wgt_dyn": self._config.iau_wgt_dyn,
                 "is_iau_active": self._config.is_iau_active,
                 "limited_area": self._grid.limited_area,
+                "divdamp_order": gtx.int32(self._config.divdamp_order),
+                "mean_cell_area": self._grid.global_properties.mean_cell_area,
+                "max_nudging_coefficient": self._config.max_nudging_coefficient,
+                "dbl_eps": constants.DBL_EPS,
             },
             variants={
                 "apply_2nd_order_divergence_damping": [False, True],
@@ -684,16 +688,6 @@ def __init__(
                 "vertical_end": gtx.int32(self._grid.num_levels),
             },
         )
-        self._calculate_divdamp_fields = setup_program(
-            backend=backend,
-            program=dycore_utils.calculate_divdamp_fields,
-            constant_args={
-                "divdamp_order": gtx.int32(self._config.divdamp_order),
-                "mean_cell_area": self._grid.global_properties.mean_cell_area,
-                "max_nudging_coefficient": self._config.max_nudging_coefficient,
-                "dbl_eps": constants.DBL_EPS,
-            },
-        )
         self._compute_rayleigh_damping_factor = setup_program(
             backend=backend,
             program=dycore_utils.compute_rayleigh_damping_factor,
@@ -947,18 +941,6 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation
             self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator
         )
         """
-        Declared as enh_divdamp_fac in ICON.
-        """
-        self.reduced_fourth_order_divdamp_coeff_at_nest_boundary = data_alloc.zero_field(
-            self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator
-        )
-        """
-        Declared as bdy_divdamp in ICON.
-        """
-        self.fourth_order_divdamp_scaling_coeff = data_alloc.zero_field(
-            self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator
-        )
-        """
         Declared as scal_divdamp in ICON.
         """
         self.intermediate_fields = IntermediateFields.allocate(grid=self._grid, allocator=allocator)
@@ -1296,13 +1278,6 @@ def run_corrector_step(
             second_order_divdamp_factor * self._grid.global_properties.mean_cell_area
         )
 
-        self._calculate_divdamp_fields(
-            interpolated_fourth_order_divdamp_factor=self.interpolated_fourth_order_divdamp_factor,
-            fourth_order_divdamp_scaling_coeff=self.fourth_order_divdamp_scaling_coeff,
-            reduced_fourth_order_divdamp_coeff_at_nest_boundary=self.reduced_fourth_order_divdamp_coeff_at_nest_boundary,
-            second_order_divdamp_factor=second_order_divdamp_factor,
-        )
-
         log.debug("corrector run velocity advection")
         self.velocity_advection.run_corrector_step(
             diagnostic_state=diagnostic_state_nh,
@@ -1352,8 +1327,8 @@ def run_corrector_step(
             normal_wind_iau_increment=diagnostic_state_nh.normal_wind_iau_increment,
             theta_v_at_edges_on_model_levels=z_fields.theta_v_at_edges_on_model_levels,
             horizontal_pressure_gradient=z_fields.horizontal_pressure_gradient,
-            reduced_fourth_order_divdamp_coeff_at_nest_boundary=self.reduced_fourth_order_divdamp_coeff_at_nest_boundary,
-            fourth_order_divdamp_scaling_coeff=self.fourth_order_divdamp_scaling_coeff,
+            interpolated_fourth_order_divdamp_factor=self.interpolated_fourth_order_divdamp_factor,
+            second_order_divdamp_factor=second_order_divdamp_factor,
             second_order_divdamp_scaling_coeff=second_order_divdamp_scaling_coeff,
             dtime=dtime,
             apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping,

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py

@@ -6,33 +6,13 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 import gt4py.next as gtx
-from gt4py.next.experimental import concat_where
 
 from icon4py.model.atmosphere.dycore.dycore_utils import (
     _broadcast_zero_to_three_edge_kdim_fields_wp,
 )
-from icon4py.model.atmosphere.dycore.stencils.compute_contravariant_correction import (
-    _compute_contravariant_correction,
-)
-from icon4py.model.atmosphere.dycore.stencils.compute_horizontal_kinetic_energy import (
-    _compute_horizontal_kinetic_energy,
-)
-from icon4py.model.atmosphere.dycore.stencils.compute_perturbation_of_rho_and_theta import (
-    _compute_perturbation_of_rho_and_theta,
-)
-from icon4py.model.atmosphere.dycore.stencils.compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers import (
-    _compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers,
-)
-from icon4py.model.atmosphere.dycore.stencils.compute_virtual_potential_temperatures_and_pressure_gradient import (
-    _compute_virtual_potential_temperatures_and_pressure_gradient,
-)
-from icon4py.model.atmosphere.dycore.stencils.extrapolate_at_top import _extrapolate_at_top
 from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import (
     _init_cell_kdim_field_with_zero_wp,
 )
-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.atmosphere.dycore.stencils.update_density_exner_wind import (
     _update_density_exner_wind,
 )
@@ -63,165 +43,6 @@ def init_test_fields(
     )
 
 
-@gtx.field_operator
-def _compute_pressure_gradient_and_perturbed_rho_and_potential_temperatures(
-    rho: fa.CellKField[float],
-    z_rth_pr_1: fa.CellKField[float],
-    z_rth_pr_2: fa.CellKField[float],
-    rho_ref_mc: fa.CellKField[float],
-    theta_v: fa.CellKField[float],
-    theta_ref_mc: fa.CellKField[float],
-    rho_ic: fa.CellKField[float],
-    wgtfac_c: fa.CellKField[float],
-    vwind_expl_wgt: fa.CellField[float],
-    exner_pr: fa.CellKField[float],
-    d_exner_dz_ref_ic: fa.CellKField[float],
-    ddqz_z_half: fa.CellKField[float],
-    z_theta_v_pr_ic: fa.CellKField[float],
-    theta_v_ic: fa.CellKField[float],
-    z_th_ddz_exner_c: fa.CellKField[float],
-) -> tuple[
-    fa.CellKField[float],
-    fa.CellKField[float],
-    fa.CellKField[float],
-    fa.CellKField[float],
-    fa.CellKField[float],
-    fa.CellKField[float],
-]:
-    (z_rth_pr_1, z_rth_pr_2) = concat_where(
-        dims.KDim == 0,
-        _compute_perturbation_of_rho_and_theta(rho, rho_ref_mc, theta_v, theta_ref_mc),
-        (z_rth_pr_1, z_rth_pr_2),
-    )
-
-    (rho_ic, z_rth_pr_1, z_rth_pr_2) = concat_where(
-        dims.KDim >= 1,
-        _compute_perturbation_of_rho_and_theta_and_rho_interface_cell_centers(
-            wgtfac_c, rho, rho_ref_mc, theta_v, theta_ref_mc
-        ),
-        (rho_ic, z_rth_pr_1, z_rth_pr_2),
-    )
-
-    (z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c) = concat_where(
-        dims.KDim >= 1,
-        _compute_virtual_potential_temperatures_and_pressure_gradient(
-            wgtfac_c,
-            z_rth_pr_2,
-            theta_v,
-            vwind_expl_wgt,
-            exner_pr,
-            d_exner_dz_ref_ic,
-            ddqz_z_half,
-        ),
-        (z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c),
-    )
-
-    return z_rth_pr_1, z_rth_pr_2, rho_ic, z_theta_v_pr_ic, theta_v_ic, z_th_ddz_exner_c
-
-
-@gtx.field_operator
-def _predictor_stencils_35_36(
-    vn: fa.EdgeKField[float],
-    ddxn_z_full: fa.EdgeKField[float],
-    ddxt_z_full: fa.EdgeKField[float],
-    vt: fa.EdgeKField[float],
-    z_w_concorr_me: fa.EdgeKField[float],
-    wgtfac_e: fa.EdgeKField[float],
-    vn_ie: fa.EdgeKField[float],
-    z_vt_ie: fa.EdgeKField[float],
-    z_kin_hor_e: fa.EdgeKField[float],
-    k_field: fa.KField[gtx.int32],
-    nflatlev_startindex: gtx.int32,
-) -> tuple[
-    fa.EdgeKField[float],
-    fa.EdgeKField[float],
-    fa.EdgeKField[float],
-    fa.EdgeKField[float],
-]:
-    z_w_concorr_me = concat_where(
-        dims.KDim >= nflatlev_startindex,
-        _compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, vt),
-        z_w_concorr_me,
-    )
-    (vn_ie, z_vt_ie, z_kin_hor_e) = concat_where(
-        dims.KDim >= 1,
-        _interpolate_vn_and_vt_to_ie_and_compute_ekin_on_edges(wgtfac_e, vn, vt),
-        (vn_ie, z_vt_ie, z_kin_hor_e),
-    )
-    return z_w_concorr_me, vn_ie, z_vt_ie, z_kin_hor_e
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def predictor_stencils_35_36(
-    vn: fa.EdgeKField[float],
-    ddxn_z_full: fa.EdgeKField[float],
-    ddxt_z_full: fa.EdgeKField[float],
-    vt: fa.EdgeKField[float],
-    z_w_concorr_me: fa.EdgeKField[float],
-    wgtfac_e: fa.EdgeKField[float],
-    vn_ie: fa.EdgeKField[float],
-    z_vt_ie: fa.EdgeKField[float],
-    z_kin_hor_e: fa.EdgeKField[float],
-    k_field: fa.KField[gtx.int32],
-    nflatlev_startindex: gtx.int32,
-    horizontal_start: gtx.int32,
-    horizontal_end: gtx.int32,
-    vertical_start: gtx.int32,
-    vertical_end: gtx.int32,
-):
-    _predictor_stencils_35_36(
-        vn,
-        ddxn_z_full,
-        ddxt_z_full,
-        vt,
-        z_w_concorr_me,
-        wgtfac_e,
-        vn_ie,
-        z_vt_ie,
-        z_kin_hor_e,
-        k_field,
-        nflatlev_startindex,
-        out=(z_w_concorr_me, vn_ie, z_vt_ie, z_kin_hor_e),
-        domain={
-            dims.EdgeDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_end),
-        },
-    )
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def predictor_stencils_37_38(
-    vn: fa.EdgeKField[float],
-    vt: fa.EdgeKField[float],
-    vn_ie: fa.EdgeKField[float],
-    z_vt_ie: fa.EdgeKField[float],
-    z_kin_hor_e: fa.EdgeKField[float],
-    wgtfacq_e_dsl: fa.EdgeKField[float],
-    horizontal_start: gtx.int32,
-    horizontal_end: gtx.int32,
-    vertical_start: gtx.int32,
-    vertical_end: gtx.int32,
-):
-    _compute_horizontal_kinetic_energy(
-        vn,
-        vt,
-        out=(vn_ie, z_vt_ie, z_kin_hor_e),
-        domain={
-            dims.EdgeDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_start + 1),
-        },
-    )
-    _extrapolate_at_top(
-        vn,
-        wgtfacq_e_dsl,
-        out=vn_ie,
-        domain={
-            dims.EdgeDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_end - 1, vertical_end),
-        },
-    )
-
-
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def stencils_61_62(
     rho_now: fa.CellKField[float],

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py

@@ -6,40 +6,57 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 import gt4py.next as gtx
-from gt4py.next import astype, broadcast
+from gt4py.next import astype
 
+from icon4py.model.atmosphere.dycore.dycore_utils import (
+    _calculate_fourth_order_divdamp_scaling_coeff,
+)
 from icon4py.model.common import field_type_aliases as fa
 from icon4py.model.common.dimension import EdgeDim, KDim
 from icon4py.model.common.type_alias import vpfloat, wpfloat
 
 
 @gtx.field_operator
 def _apply_4th_order_divergence_damping(
-    scal_divdamp: fa.KField[wpfloat],
+    interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat],
     z_graddiv2_vn: fa.EdgeKField[vpfloat],
     vn: fa.EdgeKField[wpfloat],
+    divdamp_order: gtx.int32,
+    mean_cell_area: float,
+    second_order_divdamp_factor: float,
 ) -> fa.EdgeKField[wpfloat]:
-    """Formelry known as _mo_solve_nonhydro_4th_order_divdamp."""
+    """Formerly known as _mo_solve_nonhydro_4th_order_divdamp."""
+    scal_divdamp = _calculate_fourth_order_divdamp_scaling_coeff(
+        interpolated_fourth_order_divdamp_factor,
+        divdamp_order,
+        mean_cell_area,
+        second_order_divdamp_factor,
+    )
     z_graddiv2_vn_wp = astype(z_graddiv2_vn, wpfloat)
-    scal_divdamp = broadcast(scal_divdamp, (EdgeDim, KDim))
     vn_wp = vn + (scal_divdamp * z_graddiv2_vn_wp)
     return vn_wp
 
 
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def apply_4th_order_divergence_damping(
-    scal_divdamp: fa.KField[wpfloat],
+    interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat],
     z_graddiv2_vn: fa.EdgeKField[vpfloat],
     vn: fa.EdgeKField[wpfloat],
+    divdamp_order: gtx.int32,
+    mean_cell_area: float,
+    second_order_divdamp_factor: float,
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
     vertical_start: gtx.int32,
     vertical_end: gtx.int32,
 ):
     _apply_4th_order_divergence_damping(
-        scal_divdamp,
+        interpolated_fourth_order_divdamp_factor,
         z_graddiv2_vn,
         vn,
+        divdamp_order,
+        mean_cell_area,
+        second_order_divdamp_factor,
         out=vn,
         domain={
             EdgeDim: (horizontal_start, horizontal_end),