Combine divergence damping coefficient calculation with other stencils

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

@@ -6,7 +6,7 @@
 # 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, broadcast, maximum  # noqa: A004
+from gt4py.next import GridType, abs, broadcast, maximum  # noqa: A004
 
 from icon4py.model.common import field_type_aliases as fa
 from icon4py.model.common.dimension import EdgeDim, KDim
@@ -62,7 +62,7 @@ def _calculate_fourth_order_divdamp_scaling_coeff(
     return -interpolated_fourth_order_divdamp_factor * mean_cell_area**2
 
 
-@gtx.field_operator
+@gtx.field_operator(grid_type=GridType.UNSTRUCTURED)
 def _calculate_divdamp_fields(
     interpolated_fourth_order_divdamp_factor: fa.KField[float],
     divdamp_order: gtx.int32,
@@ -82,7 +82,7 @@ def _calculate_divdamp_fields(
             fourth_order_divdamp_scaling_coeff, max_nudging_coefficient, dbl_eps
         )
     )
-    return (fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary)
+    return fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary
 
 
 @gtx.program

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],
@@ -1296,13 +1300,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 +1349,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/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py

@@ -9,9 +9,10 @@
 
 import gt4py.next as gtx
 from gt4py.eve import utils as eve_utils
-from gt4py.next import broadcast
+from gt4py.next import GridType, broadcast
 from gt4py.next.experimental import concat_where
 
+from icon4py.model.atmosphere.dycore.dycore_utils import _calculate_divdamp_fields
 from icon4py.model.atmosphere.dycore.stencils.add_analysis_increments_to_vn import (
     _add_analysis_increments_to_vn,
 )
@@ -270,7 +271,7 @@ def _compute_theta_rho_face_values_and_pressure_gradient_and_update_vn(
     )
 
 
-@gtx.field_operator
+@gtx.field_operator(grid_type=GridType.UNSTRUCTURED)
 def _apply_divergence_damping_and_update_vn(
     horizontal_gradient_of_normal_wind_divergence: fa.EdgeKField[ta.vpfloat],
     next_vn: fa.EdgeKField[ta.wpfloat],
@@ -280,8 +281,6 @@ def _apply_divergence_damping_and_update_vn(
     corrector_normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat],
     normal_wind_tendency_due_to_slow_physics_process: fa.EdgeKField[ta.vpfloat],
     normal_wind_iau_increment: fa.EdgeKField[ta.vpfloat],
-    reduced_fourth_order_divdamp_coeff_at_nest_boundary: fa.KField[ta.wpfloat],
-    fourth_order_divdamp_scaling_coeff: fa.KField[ta.wpfloat],
     second_order_divdamp_scaling_coeff: ta.wpfloat,
     theta_v_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat],
     horizontal_pressure_gradient: fa.EdgeKField[ta.vpfloat],
@@ -298,7 +297,23 @@ def _apply_divergence_damping_and_update_vn(
     limited_area: bool,
     apply_2nd_order_divergence_damping: bool,
     apply_4th_order_divergence_damping: bool,
+    interpolated_fourth_order_divdamp_factor: fa.KField[ta.wpfloat],
+    divdamp_order: gtx.int32,
+    mean_cell_area: float,
+    second_order_divdamp_factor: float,
+    max_nudging_coefficient: float,
+    dbl_eps: float,
 ) -> fa.EdgeKField[ta.wpfloat]:
+    (fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary) = (
+        _calculate_divdamp_fields(
+            interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor,
+            divdamp_order=divdamp_order,
+            mean_cell_area=mean_cell_area,
+            second_order_divdamp_factor=second_order_divdamp_factor,
+            max_nudging_coefficient=max_nudging_coefficient,
+            dbl_eps=dbl_eps,
+        )
+    )
     # add dw/dz for divergence damping term. In ICON, this stencil starts from k = kstart_dd3d until k = nlev - 1.
     # Since scaling_factor_for_3d_divdamp is zero when k < kstart_dd3d, it is meaningless to execute computation
     # above level kstart_dd3d. But we have decided to remove this manual optimization in icon4py.
@@ -534,8 +549,6 @@ def apply_divergence_damping_and_update_vn(
     corrector_normal_wind_advective_tendency: fa.EdgeKField[ta.vpfloat],
     normal_wind_tendency_due_to_slow_physics_process: fa.EdgeKField[ta.vpfloat],
     normal_wind_iau_increment: fa.EdgeKField[ta.vpfloat],
-    reduced_fourth_order_divdamp_coeff_at_nest_boundary: fa.KField[ta.wpfloat],
-    fourth_order_divdamp_scaling_coeff: fa.KField[ta.wpfloat],
     second_order_divdamp_scaling_coeff: ta.wpfloat,
     theta_v_at_edges_on_model_levels: fa.EdgeKField[ta.wpfloat],
     horizontal_pressure_gradient: fa.EdgeKField[ta.vpfloat],
@@ -552,6 +565,12 @@ def apply_divergence_damping_and_update_vn(
     limited_area: bool,
     apply_2nd_order_divergence_damping: bool,
     apply_4th_order_divergence_damping: bool,
+    interpolated_fourth_order_divdamp_factor: fa.KField[ta.wpfloat],
+    divdamp_order: gtx.int32,
+    mean_cell_area: float,
+    second_order_divdamp_factor: float,
+    max_nudging_coefficient: float,
+    dbl_eps: float,
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
     vertical_start: gtx.int32,
@@ -574,8 +593,6 @@ def apply_divergence_damping_and_update_vn(
         - corrector_normal_wind_advective_tendency: horizontal advection tendency of the normal wind at corrector step [m s-2]
         - normal_wind_tendency_due_to_slow_physics_process: normal wind tendeny due to slow physics [m s-2]
         - normal_wind_iau_increment: iau increment to normal wind (data assimilation) [m s-1]
-        - reduced_fourth_order_divdamp_coeff_at_nest_boundary: fourth order divergence damping coefficient at nest boundary [m2 s2]
-        - fourth_order_divdamp_scaling_coeff: fourth order divergence damping coefficient [m2 s2]
         - second_order_divdamp_scaling_coeff: second order divergence damping coefficient [m s]
         - theta_v_at_edges_on_model_levels: virtual potential temperature at edges on model levels [K]
         - horizontal_pressure_gradient: horizontal pressure gradient at edges on model levels [Pa m-1]
@@ -611,8 +628,6 @@ def apply_divergence_damping_and_update_vn(
         normal_wind_iau_increment=normal_wind_iau_increment,
         theta_v_at_edges_on_model_levels=theta_v_at_edges_on_model_levels,
         horizontal_pressure_gradient=horizontal_pressure_gradient,
-        reduced_fourth_order_divdamp_coeff_at_nest_boundary=reduced_fourth_order_divdamp_coeff_at_nest_boundary,
-        fourth_order_divdamp_scaling_coeff=fourth_order_divdamp_scaling_coeff,
         second_order_divdamp_scaling_coeff=second_order_divdamp_scaling_coeff,
         horizontal_mask_for_3d_divdamp=horizontal_mask_for_3d_divdamp,
         scaling_factor_for_3d_divdamp=scaling_factor_for_3d_divdamp,
@@ -627,6 +642,12 @@ def apply_divergence_damping_and_update_vn(
         limited_area=limited_area,
         apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping,
         apply_4th_order_divergence_damping=apply_4th_order_divergence_damping,
+        interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor,
+        divdamp_order=divdamp_order,
+        mean_cell_area=mean_cell_area,
+        second_order_divdamp_factor=second_order_divdamp_factor,
+        max_nudging_coefficient=max_nudging_coefficient,
+        dbl_eps=dbl_eps,
         out=next_vn,
         domain={
             dims.EdgeDim: (horizontal_start, horizontal_end),

model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py

@@ -1625,20 +1625,19 @@ def test_apply_divergence_damping_and_update_vn(
     corrector_normal_wind_advective_tendency = sp_stencil_init.ddt_vn_apc_ntl(1)
     normal_wind_tendency_due_to_slow_physics_process = sp_stencil_init.ddt_vn_phy()
     normal_wind_iau_increment = sp_stencil_init.vn_incr()
-    reduced_fourth_order_divdamp_coeff_at_nest_boundary = sp_nh_init.bdy_divdamp()
-    fourth_order_divdamp_scaling_coeff = sp_nh_init.scal_divdamp()
+    interpolated_fourth_order_divdamp_factor = sp_nh_init.enh_divdamp_fac()
     theta_v_at_edges_on_model_levels = sp_stencil_init.z_theta_v_e()
     horizontal_pressure_gradient = sp_stencil_init.z_gradh_exner()
     current_vn = sp_stencil_init.vn()
     next_vn = savepoint_nonhydro_init.vn_new()
     horizontal_gradient_of_normal_wind_divergence = sp_nh_init.z_graddiv_vn()
     config = definitions.construct_nonhydrostatic_config(experiment)
+    mean_cell_area = grid_savepoint.mean_cell_area()
+    max_nudging_coefficient = grid_savepoint.nudge_max_coeff()
 
     iau_wgt_dyn = config.iau_wgt_dyn
     divdamp_order = config.divdamp_order
-    second_order_divdamp_scaling_coeff = (
-        sp_nh_init.divdamp_fac_o2() * grid_savepoint.mean_cell_area()
-    )
+    second_order_divdamp_scaling_coeff = sp_nh_init.divdamp_fac_o2() * mean_cell_area
     second_order_divdamp_factor = savepoint_nonhydro_init.divdamp_fac_o2()
     apply_2nd_order_divergence_damping = (
         divdamp_order == dycore_states.DivergenceDampingOrder.COMBINED
@@ -1668,8 +1667,6 @@ def test_apply_divergence_damping_and_update_vn(
         normal_wind_iau_increment=normal_wind_iau_increment,
         theta_v_at_edges_on_model_levels=theta_v_at_edges_on_model_levels,
         horizontal_pressure_gradient=horizontal_pressure_gradient,
-        reduced_fourth_order_divdamp_coeff_at_nest_boundary=reduced_fourth_order_divdamp_coeff_at_nest_boundary,
-        fourth_order_divdamp_scaling_coeff=fourth_order_divdamp_scaling_coeff,
         second_order_divdamp_scaling_coeff=second_order_divdamp_scaling_coeff,
         horizontal_mask_for_3d_divdamp=metrics_savepoint.hmask_dd3d(),
         scaling_factor_for_3d_divdamp=metrics_savepoint.scalfac_dd3d(),
@@ -1684,6 +1681,12 @@ def test_apply_divergence_damping_and_update_vn(
         limited_area=grid_savepoint.get_metadata("limited_area").get("limited_area"),
         apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping,
         apply_4th_order_divergence_damping=apply_4th_order_divergence_damping,
+        interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor,
+        divdamp_order=divdamp_order,
+        mean_cell_area=mean_cell_area,
+        second_order_divdamp_factor=second_order_divdamp_factor,
+        max_nudging_coefficient=max_nudging_coefficient,
+        dbl_eps=constants.DBL_EPS,
         horizontal_start=start_edge_nudging_level_2,
         horizontal_end=end_edge_local,
         vertical_start=gtx.int32(0),

model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py

@@ -20,6 +20,11 @@
 from icon4py.model.common.grid import base, horizontal as h_grid
 from icon4py.model.common.utils import data_allocation as data_alloc
 
+from .test_dycore_utils import (
+    calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy,
+    fourth_order_divdamp_scaling_coeff_numpy,
+)
+
 
 divergence_damp_order = DivergenceDampingOrder()
 
@@ -60,8 +65,6 @@ def reference(
         normal_wind_iau_increment: np.ndarray,
         theta_v_at_edges_on_model_levels: np.ndarray,
         horizontal_pressure_gradient: np.ndarray,
-        reduced_fourth_order_divdamp_coeff_at_nest_boundary: np.ndarray,
-        fourth_order_divdamp_scaling_coeff: np.ndarray,
         second_order_divdamp_scaling_coeff: ta.wpfloat,
         horizontal_mask_for_3d_divdamp: np.ndarray,
         scaling_factor_for_3d_divdamp: np.ndarray,
@@ -76,11 +79,30 @@ def reference(
         limited_area: gtx.int32,
         apply_2nd_order_divergence_damping: bool,
         apply_4th_order_divergence_damping: bool,
+        interpolated_fourth_order_divdamp_factor: np.ndarray,
+        divdamp_order: gtx.int32,
+        mean_cell_area: float,
+        second_order_divdamp_factor: float,
+        max_nudging_coefficient: float,
+        dbl_eps: float,
         horizontal_start: gtx.int32,
         horizontal_end: gtx.int32,
         vertical_start: gtx.int32,
         vertical_end: gtx.int32,
     ) -> dict:
+        fourth_order_divdamp_scaling_coeff = fourth_order_divdamp_scaling_coeff_numpy(
+            interpolated_fourth_order_divdamp_factor,
+            divdamp_order,
+            second_order_divdamp_factor,
+            mean_cell_area,
+        )
+
+        reduced_fourth_order_divdamp_coeff_at_nest_boundary = (
+            calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy(
+                fourth_order_divdamp_scaling_coeff, max_nudging_coefficient
+            )
+        )
+
         horz_idx = np.arange(horizontal_end)[:, np.newaxis]
 
         scaling_factor_for_3d_divdamp = np.expand_dims(scaling_factor_for_3d_divdamp, axis=0)
@@ -212,12 +234,13 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict:
         theta_v_at_edges_on_model_levels = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)
         horizontal_pressure_gradient = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)
         geofac_grdiv = data_alloc.random_field(grid, dims.EdgeDim, dims.E2C2EODim)
-        fourth_order_divdamp_scaling_coeff = data_alloc.random_field(grid, dims.KDim)
-        reduced_fourth_order_divdamp_coeff_at_nest_boundary = data_alloc.random_field(
-            grid, dims.KDim
-        )
+        interpolated_fourth_order_divdamp_factor = data_alloc.random_field(grid, dims.KDim)
         nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim)
 
+        mean_cell_area = 1000.0
+        max_nudging_coefficient = 0.3
+        dbl_eps = constants.DBL_EPS
+
         dtime = 0.9
         advection_implicit_weight_parameter = 0.75
         advection_explicit_weight_parameter = 0.25
@@ -254,8 +277,6 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict:
             normal_wind_iau_increment=normal_wind_iau_increment,
             theta_v_at_edges_on_model_levels=theta_v_at_edges_on_model_levels,
             horizontal_pressure_gradient=horizontal_pressure_gradient,
-            reduced_fourth_order_divdamp_coeff_at_nest_boundary=reduced_fourth_order_divdamp_coeff_at_nest_boundary,
-            fourth_order_divdamp_scaling_coeff=fourth_order_divdamp_scaling_coeff,
             second_order_divdamp_scaling_coeff=second_order_divdamp_scaling_coeff,
             horizontal_mask_for_3d_divdamp=horizontal_mask_for_3d_divdamp,
             scaling_factor_for_3d_divdamp=scaling_factor_for_3d_divdamp,
@@ -270,6 +291,12 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict:
             limited_area=limited_area,
             apply_2nd_order_divergence_damping=apply_2nd_order_divergence_damping,
             apply_4th_order_divergence_damping=apply_4th_order_divergence_damping,
+            interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor,
+            divdamp_order=divdamp_order,
+            mean_cell_area=mean_cell_area,
+            second_order_divdamp_factor=second_order_divdamp_factor,
+            max_nudging_coefficient=max_nudging_coefficient,
+            dbl_eps=dbl_eps,
             horizontal_start=start_edge_nudging_level_2,
             horizontal_end=end_edge_local,
             vertical_start=0,

model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py

@@ -27,15 +27,15 @@
 # TODO(): apply StencilTest structure to this test
 
 
-def fourth_order_divdamp_scaling_coeff_for_order_24_numpy(
-    a: np.ndarray, factor: float, mean_cell_area: float
+def fourth_order_divdamp_scaling_coeff_numpy(
+    a: np.ndarray, divdamp_order: int, factor: float, mean_cell_area: float
 ) -> np.ndarray:
-    a = np.maximum(0.0, a - 0.25 * factor)
-    return -a * mean_cell_area**2
+    b = np.maximum(0.0, a - 0.25 * factor) if divdamp_order == 24 else np.full_like(a, factor)
+    return -b * mean_cell_area**2
 
 
 def calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy(
-    coeff: float, field: np.ndarray
+    field: np.ndarray, coeff: float
 ) -> np.ndarray:
     return 0.75 / (coeff + constants.DBL_EPS) * np.abs(field)
 
@@ -61,8 +61,9 @@ def test_calculate_fourth_order_divdamp_scaling_coeff_order_24(
         offset_provider={},
     )
 
-    ref = fourth_order_divdamp_scaling_coeff_for_order_24_numpy(
+    ref = fourth_order_divdamp_scaling_coeff_numpy(
         interpolated_fourth_order_divdamp_factor.asnumpy(),
+        divdamp_order,
         second_order_divdamp_factor,
         mean_cell_area,
     )
@@ -106,7 +107,7 @@ def test_calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary(
     assert test_utils.dallclose(
         out.asnumpy(),
         calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy(
-            coeff, fourth_order_divdamp_scaling_coeff.asnumpy()
+            fourth_order_divdamp_scaling_coeff.asnumpy(), coeff
         ),
     )
 
@@ -123,13 +124,16 @@ def test_calculate_divdamp_fields(backend: gtx_typing.Backend) -> None:
     second_order_divdamp_factor = 0.7
     max_nudging_coefficient = 0.3
 
-    scaled_ref = fourth_order_divdamp_scaling_coeff_for_order_24_numpy(
-        divdamp_field.asnumpy(), second_order_divdamp_factor, mean_cell_area
+    scaled_ref = fourth_order_divdamp_scaling_coeff_numpy(
+        divdamp_field.asnumpy(),
+        divdamp_order,
+        second_order_divdamp_factor,
+        mean_cell_area,
     )
 
     reduced_fourth_order_divdamp_coeff_at_nest_boundary_ref = (
         calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy(
-            max_nudging_coefficient, scaled_ref
+            scaled_ref, max_nudging_coefficient
         )
     )
 

pyproject.toml

@@ -356,7 +356,7 @@ url = "https://test.pypi.org/simple/"
 [tool.uv.sources]
 dace = {git = "https://github.com/GridTools/dace", tag = "__gt4py-next-integration_2025_12_11"}
 ghex = {git = "https://github.com/msimberg/GHEX.git", branch = "async-mpi"}
-# gt4py = {git = "https://github.com/GridTools/gt4py", branch = "main"}
+gt4py = {git = "https://github.com/edopao/gt4py", branch = "gtfn_fix_canonicalize_domain"}
 # gt4py = {index = "test.pypi"}
 icon4py-atmosphere-advection = {workspace = true}
 icon4py-atmosphere-diffusion = {workspace = true}