From 2a0ba108e81eafbb5485787f5e7dfe7e7f488e00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 19 Nov 2025 11:28:35 +0100 Subject: [PATCH 01/16] Combine divergence damping coefficient calculation with other stencils --- .../model/atmosphere/dycore/dycore_utils.py | 6 +-- .../model/atmosphere/dycore/solve_nonhydro.py | 20 +++++---- ...ge_diagnostics_for_dycore_and_update_vn.py | 41 +++++++++++++----- .../integration_tests/test_solve_nonhydro.py | 15 ++++--- ..._apply_divergence_damping_and_update_vn.py | 42 +++++++++++++++---- .../dycore/stencil_tests/test_dycore_utils.py | 16 ++++--- 6 files changed, 99 insertions(+), 41 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py index 9a8ac84a84..c7efa672b6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py +++ b/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 abs, broadcast, maximum, GridType # 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 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 8b939ed6d6..e7c69e0fda 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 @@ -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,12 +1300,12 @@ 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, - ) + #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( @@ -1352,8 +1356,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, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index 341de7eb15..0c6a18f55b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/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 broadcast, GridType 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), diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index e2711f3c82..8374c820ce 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1625,19 +1625,20 @@ 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() + sp_nh_init.divdamp_fac_o2() * mean_cell_area ) second_order_divdamp_factor = savepoint_nonhydro_init.divdamp_fac_o2() apply_2nd_order_divergence_damping = ( @@ -1668,8 +1669,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 +1683,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), diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index d49922254d..463425e725 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -19,7 +19,9 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc - +from model.atmosphere.dycore.tests.dycore.stencil_tests.test_dycore_utils import \ + fourth_order_divdamp_scaling_coeff_numpy, \ + calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy divergence_damp_order = DivergenceDampingOrder() @@ -42,8 +44,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, @@ -58,11 +58,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) @@ -177,12 +196,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 @@ -219,8 +239,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, @@ -235,6 +253,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, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py index bde56aba25..d39a589d1a 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py @@ -27,15 +27,18 @@ # 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: np.int32, factor: float, mean_cell_area: float ) -> np.ndarray: - a = np.maximum(0.0, a - 0.25 * factor) - return -a * mean_cell_area**2 + if divdamp_order == 24: + b = np.maximum(0.0, a - 0.25 * factor) + else: + b = 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 +64,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, ) From d0c0d5e423d0374eb60842cddb581563e1cc6970 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Wed, 7 Jan 2026 14:28:28 +0100 Subject: [PATCH 02/16] fix merge errors --- .../model/atmosphere/dycore/dycore_utils.py | 2 +- .../model/atmosphere/dycore/solve_nonhydro.py | 7 ------- ...dge_diagnostics_for_dycore_and_update_vn.py | 6 +++--- .../integration_tests/test_solve_nonhydro.py | 4 +--- ...t_apply_divergence_damping_and_update_vn.py | 11 +++++++---- .../dycore/stencil_tests/test_dycore_utils.py | 18 +++++++++--------- 6 files changed, 21 insertions(+), 27 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py index c7efa672b6..921f4be4af 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py +++ b/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, GridType # 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 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 e7c69e0fda..48066d2227 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 @@ -1300,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, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index 0c6a18f55b..f61f8c79c5 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -9,7 +9,7 @@ import gt4py.next as gtx from gt4py.eve import utils as eve_utils -from gt4py.next import broadcast, GridType +from gt4py.next import GridType, broadcast from gt4py.next.experimental import concat_where from icon4py.model.atmosphere.dycore.dycore_utils import _calculate_divdamp_fields @@ -304,8 +304,7 @@ def _apply_divergence_damping_and_update_vn( max_nudging_coefficient: float, dbl_eps: float, ) -> fa.EdgeKField[ta.wpfloat]: - - (fourth_order_divdamp_scaling_coeff, reduced_fourth_order_divdamp_coeff_at_nest_boundary) = \ + (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, @@ -314,6 +313,7 @@ def _apply_divergence_damping_and_update_vn( 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. diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 8374c820ce..c6bbc89dcf 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1637,9 +1637,7 @@ def test_apply_divergence_damping_and_update_vn( iau_wgt_dyn = config.iau_wgt_dyn divdamp_order = config.divdamp_order - second_order_divdamp_scaling_coeff = ( - sp_nh_init.divdamp_fac_o2() * 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 diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 255a3fd5f6..2c11a6857d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -19,9 +19,12 @@ from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import base, horizontal as h_grid from icon4py.model.common.utils import data_allocation as data_alloc -from model.atmosphere.dycore.tests.dycore.stencil_tests.test_dycore_utils import \ - fourth_order_divdamp_scaling_coeff_numpy, \ - calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy + +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() @@ -91,7 +94,7 @@ def reference( interpolated_fourth_order_divdamp_factor, divdamp_order, second_order_divdamp_factor, - mean_cell_area + mean_cell_area, ) reduced_fourth_order_divdamp_coeff_at_nest_boundary = ( diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py index d39a589d1a..cae87978ba 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py @@ -28,12 +28,9 @@ def fourth_order_divdamp_scaling_coeff_numpy( - a: np.ndarray, divdamp_order: np.int32, factor: float, mean_cell_area: float + a: np.ndarray, divdamp_order: int, factor: float, mean_cell_area: float ) -> np.ndarray: - if divdamp_order == 24: - b = np.maximum(0.0, a - 0.25 * factor) - else: - b = factor + 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 @@ -110,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 ), ) @@ -127,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 ) ) From 327d675a43a29fb5eaf98eac75d1afd4651d0b47 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 09:27:47 +0100 Subject: [PATCH 03/16] update gt4py to main --- pyproject.toml | 2 +- uv.lock | 36 ++++++++++++++++-------------------- 2 files changed, 17 insertions(+), 21 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9d8a5abeac..9d81dc118e 100644 --- a/pyproject.toml +++ b/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} diff --git a/uv.lock b/uv.lock index 0e6ebf090a..dd66c4b123 100644 --- a/uv.lock +++ b/uv.lock @@ -1399,8 +1399,8 @@ wheels = [ [[package]] name = "gt4py" -version = "1.1.2" -source = { registry = "https://pypi.org/simple" } +version = "1.1.2+unknown.version.details" +source = { git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain#f9d6d707667d5ee937bf854a04c7d4422ee9a6ad" } dependencies = [ { name = "attrs" }, { name = "black" }, @@ -1430,10 +1430,6 @@ dependencies = [ { name = "versioningit" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/77/74/8a388e4be44fec6713b034df7950edfa8f1eb9c5e34feb639aaf38c4897c/gt4py-1.1.2.tar.gz", hash = "sha256:4f4999236e26a32d9f6d74aebdb21094fda2511d07e4f73615e396c35b13c78a", size = 773452, upload-time = "2025-12-08T12:15:48.908Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/9c/19c6263d0897a44542a97371df0a9d243e37654053d8b2f6d067fb8ed1a8/gt4py-1.1.2-py3-none-any.whl", hash = "sha256:44665eb59c4d174887fb787a01b8848ddf38607109648dfa1922d234e31153e2", size = 978719, upload-time = "2025-12-08T12:15:46.682Z" }, -] [package.optional-dependencies] cuda11 = [ @@ -1752,7 +1748,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1769,7 +1765,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1786,7 +1782,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1803,7 +1799,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1820,7 +1816,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1885,9 +1881,9 @@ requires-dist = [ { name = "dace", git = "https://github.com/GridTools/dace?tag=__gt4py-next-integration_2025_12_11" }, { name = "datashader", marker = "extra == 'io'", specifier = ">=0.16.1" }, { name = "ghex", marker = "extra == 'distributed'", git = "https://github.com/msimberg/GHEX.git?branch=async-mpi" }, - { name = "gt4py", specifier = "==1.1.2" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "holoviews", marker = "extra == 'io'", specifier = ">=1.16.0" }, { name = "icon4py-common", extras = ["distributed", "io"], marker = "extra == 'all'", editable = "model/common" }, { name = "mpi4py", marker = "extra == 'distributed'", specifier = ">=3.1.5" }, @@ -1923,7 +1919,7 @@ dependencies = [ requires-dist = [ { name = "click", specifier = ">=8.0.1" }, { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -1951,7 +1947,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -1980,7 +1976,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "filelock", specifier = ">=3.18.0" }, - { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, @@ -2029,9 +2025,9 @@ requires-dist = [ { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" }, { name = "cupy-cuda12x", marker = "extra == 'cuda12'", specifier = ">=13.0" }, { name = "fprettify", specifier = ">=0.3.7" }, - { name = "gt4py", specifier = "==1.1.2" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, + { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, { name = "icon4py-atmosphere-advection", editable = "model/atmosphere/advection" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, From 86e3be4c51030f82e4785916aff2c947b2d94127 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 11:13:22 +0100 Subject: [PATCH 04/16] cleanup --- .../model/atmosphere/dycore/solve_nonhydro.py | 10 ---------- ...dge_diagnostics_for_dycore_and_update_vn.py | 18 ++++++++++-------- ...t_apply_divergence_damping_and_update_vn.py | 6 +++--- 3 files changed, 13 insertions(+), 21 deletions(-) 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 48066d2227..db9e1f09d3 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 @@ -688,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, diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index f61f8c79c5..d02822c74a 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -289,6 +289,7 @@ def _apply_divergence_damping_and_update_vn( inv_dual_edge_length: fa.EdgeField[ta.wpfloat], nudgecoeff_e: fa.EdgeField[ta.wpfloat], geofac_grdiv: gtx.Field[[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + interpolated_fourth_order_divdamp_factor: fa.KField[ta.wpfloat], advection_explicit_weight_parameter: ta.wpfloat, advection_implicit_weight_parameter: ta.wpfloat, dtime: ta.wpfloat, @@ -297,7 +298,6 @@ 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, @@ -557,6 +557,7 @@ def apply_divergence_damping_and_update_vn( inv_dual_edge_length: fa.EdgeField[ta.wpfloat], nudgecoeff_e: fa.EdgeField[ta.wpfloat], geofac_grdiv: gtx.Field[[dims.EdgeDim, dims.E2C2EODim], ta.wpfloat], + interpolated_fourth_order_divdamp_factor: fa.KField[ta.wpfloat], advection_explicit_weight_parameter: ta.wpfloat, advection_implicit_weight_parameter: ta.wpfloat, dtime: ta.wpfloat, @@ -565,7 +566,6 @@ 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, @@ -601,18 +601,20 @@ def apply_divergence_damping_and_update_vn( - inv_dual_edge_length: inverse dual edge length - nudgecoeff_e: nudging coefficient for fourth order divergence damping at nest boundary - geofac_grdiv: metric coefficient for computation of horizontal gradient of divergence - - fourth_order_divdamp_factor: scaling factor for fourth order divergence damping - - second_order_divdamp_factor: scaling factor for second order divergence damping + - interpolated_fourth_order_divdamp_factor - advection_explicit_weight_parameter: explicitness weight of normal_wind_advective_tendency - advection_implicit_weight_parameter: implicitness weight of normal_wind_advective_tendency - dtime: time step [s] - iau_wgt_dyn: a scaling factor for iau increment - is_iau_active: option for iau increment analysis - - itime_scheme: ICON itime scheme (see ICON tutorial) - limited_area: option indicating the grid is limited area or not + - apply_2nd_order_divergence_damping: scaling factor for second order divergence damping + - apply_4th_order_divergence_damping: scaling factor for fourth order divergence damping - divdamp_order: divergence damping order (see the class DivergenceDampingOrder) - - start_edge_nudging_level_2: start index of second nudging level zone for edges - - end_edge_local: end index of local zone for edges + - mean_cell_area + - second_order_divdamp_factor + - max_nudging_coefficient + - dbl_eps Returns: - next_vn: normal wind to be updated [m s-1] @@ -634,6 +636,7 @@ def apply_divergence_damping_and_update_vn( inv_dual_edge_length=inv_dual_edge_length, nudgecoeff_e=nudgecoeff_e, geofac_grdiv=geofac_grdiv, + interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, advection_explicit_weight_parameter=advection_explicit_weight_parameter, advection_implicit_weight_parameter=advection_implicit_weight_parameter, dtime=dtime, @@ -642,7 +645,6 @@ 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, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 2c11a6857d..86793816d9 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -75,8 +75,8 @@ def reference( advection_implicit_weight_parameter: ta.wpfloat, dtime: ta.wpfloat, iau_wgt_dyn: ta.wpfloat, - is_iau_active: gtx.int32, - limited_area: gtx.int32, + is_iau_active: bool, + limited_area: bool, apply_2nd_order_divergence_damping: bool, apply_4th_order_divergence_damping: bool, interpolated_fourth_order_divdamp_factor: np.ndarray, @@ -283,6 +283,7 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict: inv_dual_edge_length=inv_dual_edge_length, nudgecoeff_e=nudgecoeff_e, geofac_grdiv=geofac_grdiv, + interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, advection_explicit_weight_parameter=advection_explicit_weight_parameter, advection_implicit_weight_parameter=advection_implicit_weight_parameter, dtime=dtime, @@ -291,7 +292,6 @@ 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, From 955bc10359c7b661e785c7e82368c06c1e98cf50 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 11:18:01 +0100 Subject: [PATCH 05/16] remove unused code --- .../dycore/solve_nonhydro_stencils.py | 179 ------------------ 1 file changed, 179 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py index bdb18e294c..2f9bf6efcd 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py +++ b/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], From 6f7d6a7389e9b058005ff3022a7266a60e083af2 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 11:50:14 +0100 Subject: [PATCH 06/16] move _calculate_divdamp_fields into apply_4th_order_divergence_damping --- ...ge_diagnostics_for_dycore_and_update_vn.py | 21 +++++----- ..._apply_divergence_damping_and_update_vn.py | 40 +++++++++---------- 2 files changed, 29 insertions(+), 32 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index d02822c74a..bfcdc5316f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -304,16 +304,6 @@ def _apply_divergence_damping_and_update_vn( 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. @@ -350,6 +340,17 @@ def _apply_divergence_damping_and_update_vn( squared_horizontal_gradient_of_total_divergence = _compute_graddiv2_of_vn( geofac_grdiv=geofac_grdiv, z_graddiv_vn=horizontal_gradient_of_total_divergence ) + ( + 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, + ) if limited_area: next_vn = _apply_weighted_2nd_and_4th_order_divergence_damping( scal_divdamp=fourth_order_divdamp_scaling_coeff, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 86793816d9..32ef6bca28 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -90,19 +90,6 @@ def reference( 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) @@ -135,6 +122,14 @@ def reference( next_vn, ) + if apply_2nd_order_divergence_damping: + next_vn = np.where( + (horizontal_start <= horz_idx) & (horz_idx < horizontal_end), + next_vn + + (second_order_divdamp_scaling_coeff * horizontal_gradient_of_total_divergence), + next_vn, + ) + if apply_4th_order_divergence_damping: e2c2eO = connectivities[dims.E2C2EODim] # verified for e-10 @@ -151,16 +146,17 @@ def reference( ), np.zeros_like(horizontal_gradient_of_total_divergence), ) - - if apply_2nd_order_divergence_damping: - next_vn = np.where( - (horizontal_start <= horz_idx) & (horz_idx < horizontal_end), - next_vn - + (second_order_divdamp_scaling_coeff * horizontal_gradient_of_total_divergence), - next_vn, + 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 + ) ) - - if apply_4th_order_divergence_damping: if limited_area: next_vn = np.where( (horizontal_start <= horz_idx) & (horz_idx < horizontal_end), From 9366be98bc97ff9a5f198714d37e6ac28d7106c6 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 16:28:14 +0100 Subject: [PATCH 07/16] inline _calculate_divdamp_fields --- .../apply_4th_order_divergence_damping.py | 31 +++++++++--- ...ed_2nd_and_4th_order_divergence_damping.py | 44 ++++++++++++----- ...ge_diagnostics_for_dycore_and_update_vn.py | 27 +++++------ ...test_apply_4th_order_divergence_damping.py | 31 +++++++++--- ..._apply_divergence_damping_and_update_vn.py | 21 +++++---- ...ed_2nd_and_4th_order_divergence_damping.py | 47 ++++++++++++++----- 6 files changed, 140 insertions(+), 61 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py index 79c86fa3a0..2917771c26 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py @@ -6,8 +6,11 @@ # 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 @@ -15,31 +18,45 @@ @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.""" + fourth_order_divdamp_scaling_coeff = _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) + vn_wp = vn + (fourth_order_divdamp_scaling_coeff * 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), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_weighted_2nd_and_4th_order_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_weighted_2nd_and_4th_order_divergence_damping.py index b8d69dadeb..f99353e001 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_weighted_2nd_and_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_weighted_2nd_and_4th_order_divergence_damping.py @@ -6,47 +6,69 @@ # 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, + _calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary, +) 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 _apply_weighted_2nd_and_4th_order_divergence_damping( - scal_divdamp: fa.KField[wpfloat], - bdy_divdamp: fa.KField[wpfloat], + interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat], nudgecoeff_e: fa.EdgeField[wpfloat], z_graddiv2_vn: fa.EdgeKField[vpfloat], vn: fa.EdgeKField[wpfloat], + divdamp_order: gtx.int32, + mean_cell_area: float, + second_order_divdamp_factor: float, + max_nudging_coefficient: float, + dbl_eps: float, ) -> fa.EdgeKField[wpfloat]: - """Formelry known as _mo_solve_nonhydro_stencil_27.""" + """Formerly known as _mo_solve_nonhydro_stencil_27.""" + scal_divdamp = _calculate_fourth_order_divdamp_scaling_coeff( + interpolated_fourth_order_divdamp_factor, + divdamp_order, + mean_cell_area, + second_order_divdamp_factor, + ) + bdy_divdamp = _calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary( + scal_divdamp, max_nudging_coefficient, dbl_eps + ) z_graddiv2_vn_wp = astype(z_graddiv2_vn, wpfloat) - - scal_divdamp = broadcast(scal_divdamp, (dims.EdgeDim, dims.KDim)) - bdy_divdamp = broadcast(bdy_divdamp, (dims.EdgeDim, dims.KDim)) vn_wp = vn + (scal_divdamp + bdy_divdamp * nudgecoeff_e) * z_graddiv2_vn_wp return vn_wp @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def apply_weighted_2nd_and_4th_order_divergence_damping( - scal_divdamp: fa.KField[wpfloat], - bdy_divdamp: fa.KField[wpfloat], + interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat], nudgecoeff_e: fa.EdgeField[wpfloat], z_graddiv2_vn: fa.EdgeKField[vpfloat], vn: fa.EdgeKField[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, vertical_end: gtx.int32, ): _apply_weighted_2nd_and_4th_order_divergence_damping( - scal_divdamp, - bdy_divdamp, + interpolated_fourth_order_divdamp_factor, nudgecoeff_e, z_graddiv2_vn, vn, + divdamp_order, + mean_cell_area, + second_order_divdamp_factor, + max_nudging_coefficient, + dbl_eps, out=vn, domain={ dims.EdgeDim: (horizontal_start, horizontal_end), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index bfcdc5316f..5e5a3c85b7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -12,7 +12,6 @@ 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, ) @@ -340,30 +339,28 @@ def _apply_divergence_damping_and_update_vn( squared_horizontal_gradient_of_total_divergence = _compute_graddiv2_of_vn( geofac_grdiv=geofac_grdiv, z_graddiv_vn=horizontal_gradient_of_total_divergence ) - ( - 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, - ) if limited_area: next_vn = _apply_weighted_2nd_and_4th_order_divergence_damping( - scal_divdamp=fourth_order_divdamp_scaling_coeff, - bdy_divdamp=reduced_fourth_order_divdamp_coeff_at_nest_boundary, + interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, nudgecoeff_e=nudgecoeff_e, z_graddiv2_vn=squared_horizontal_gradient_of_total_divergence, vn=next_vn, + 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, ) else: next_vn = _apply_4th_order_divergence_damping( - scal_divdamp=fourth_order_divdamp_scaling_coeff, + interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, z_graddiv2_vn=squared_horizontal_gradient_of_total_divergence, vn=next_vn, + 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, ) if is_iau_active: diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py index 82148b6461..8b1cfbec2d 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py @@ -10,15 +10,16 @@ import gt4py.next as gtx import numpy as np import pytest +from model.atmosphere.dycore.tests.dycore.stencil_tests import test_dycore_utils from icon4py.model.atmosphere.dycore.stencils.apply_4th_order_divergence_damping import ( apply_4th_order_divergence_damping, ) -from icon4py.model.common import dimension as dims +from icon4py.model.common import constants, 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.common.utils.data_allocation import random_field +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest @@ -39,24 +40,40 @@ class TestApply4thOrderDivergenceDamping(StencilTest): @staticmethod def reference( connectivities: dict[gtx.Dimension, np.ndarray], - scal_divdamp: np.ndarray, + interpolated_fourth_order_divdamp_factor: np.ndarray, z_graddiv2_vn: np.ndarray, vn: np.ndarray, + divdamp_order: gtx.int32, + mean_cell_area: float, + second_order_divdamp_factor: float, **kwargs: Any, ) -> dict: + scal_divdamp = test_dycore_utils.fourth_order_divdamp_scaling_coeff_numpy( + interpolated_fourth_order_divdamp_factor, + divdamp_order, + second_order_divdamp_factor, + mean_cell_area, + ) vn = apply_4th_order_divergence_damping_numpy(scal_divdamp, z_graddiv2_vn, vn) return dict(vn=vn) @pytest.fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - scal_divdamp = random_field(grid, dims.KDim, dtype=wpfloat) - z_graddiv2_vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + interpolated_fourth_order_divdamp_factor = data_alloc.random_field(grid, dims.KDim) + z_graddiv2_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + + divdamp_order = 24 + mean_cell_area = 1000.0 + second_order_divdamp_factor = 3.0 return dict( - scal_divdamp=scal_divdamp, + interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, z_graddiv2_vn=z_graddiv2_vn, vn=vn, + divdamp_order=divdamp_order, + mean_cell_area=mean_cell_area, + second_order_divdamp_factor=second_order_divdamp_factor, horizontal_start=0, horizontal_end=gtx.int32(grid.num_edges), vertical_start=0, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 32ef6bca28..5778d3f558 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -9,6 +9,7 @@ import gt4py.next as gtx import numpy as np import pytest +from model.atmosphere.dycore.tests.dycore.stencil_tests import test_dycore_utils import icon4py.model.common.type_alias as ta import icon4py.model.testing.stencil_tests as test_helpers @@ -71,6 +72,7 @@ def reference( inv_dual_edge_length: np.ndarray, nudgecoeff_e: np.ndarray, geofac_grdiv: np.ndarray, + interpolated_fourth_order_divdamp_factor: np.ndarray, advection_explicit_weight_parameter: ta.wpfloat, advection_implicit_weight_parameter: ta.wpfloat, dtime: ta.wpfloat, @@ -79,7 +81,6 @@ def reference( limited_area: bool, 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, @@ -146,17 +147,17 @@ def reference( ), np.zeros_like(horizontal_gradient_of_total_divergence), ) - 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 + fourth_order_divdamp_scaling_coeff = ( + test_dycore_utils.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 = test_dycore_utils.calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy( + fourth_order_divdamp_scaling_coeff, max_nudging_coefficient + ) if limited_area: next_vn = np.where( (horizontal_start <= horz_idx) & (horz_idx < horizontal_end), diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py index 31de7f151f..2d4ad1c29e 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py @@ -10,15 +10,16 @@ import gt4py.next as gtx import numpy as np import pytest +from model.atmosphere.dycore.tests.dycore.stencil_tests import test_dycore_utils from icon4py.model.atmosphere.dycore.stencils.apply_weighted_2nd_and_4th_order_divergence_damping import ( apply_weighted_2nd_and_4th_order_divergence_damping, ) -from icon4py.model.common import dimension as dims +from icon4py.model.common import constants, 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.common.utils.data_allocation import random_field +from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest @@ -41,13 +42,28 @@ class TestApplyWeighted2ndAnd4thOrderDivergenceDamping(StencilTest): @staticmethod def reference( connectivities: dict[gtx.Dimension, np.ndarray], - scal_divdamp: np.ndarray, - bdy_divdamp: np.ndarray, + interpolated_fourth_order_divdamp_factor: np.ndarray, nudgecoeff_e: np.ndarray, z_graddiv2_vn: np.ndarray, vn: np.ndarray, + divdamp_order: gtx.int32, + mean_cell_area: float, + second_order_divdamp_factor: float, + max_nudging_coefficient: float, + dbl_eps: float, **kwargs: Any, ) -> dict: + scal_divdamp = test_dycore_utils.fourth_order_divdamp_scaling_coeff_numpy( + interpolated_fourth_order_divdamp_factor, + divdamp_order, + second_order_divdamp_factor, + mean_cell_area, + ) + bdy_divdamp = ( + test_dycore_utils.calculate_reduced_fourth_order_divdamp_coeff_at_nest_boundary_numpy( + scal_divdamp, max_nudging_coefficient + ) + ) vn = apply_weighted_2nd_and_4th_order_divergence_damping_numpy( scal_divdamp, bdy_divdamp, @@ -59,18 +75,27 @@ def reference( @pytest.fixture def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]: - scal_divdamp = random_field(grid, dims.KDim, dtype=wpfloat) - bdy_divdamp = random_field(grid, dims.KDim, dtype=wpfloat) - nudgecoeff_e = random_field(grid, dims.EdgeDim, dtype=wpfloat) - z_graddiv2_vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) - vn = random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + interpolated_fourth_order_divdamp_factor = data_alloc.random_field(grid, dims.KDim) + nudgecoeff_e = data_alloc.random_field(grid, dims.EdgeDim, dtype=wpfloat) + z_graddiv2_vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=vpfloat) + vn = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim, dtype=wpfloat) + + divdamp_order = 24 + mean_cell_area = 1000.0 + second_order_divdamp_factor = 3.0 + max_nudging_coefficient = 0.3 + dbl_eps = constants.DBL_EPS return dict( - scal_divdamp=scal_divdamp, - bdy_divdamp=bdy_divdamp, + interpolated_fourth_order_divdamp_factor=interpolated_fourth_order_divdamp_factor, nudgecoeff_e=nudgecoeff_e, z_graddiv2_vn=z_graddiv2_vn, vn=vn, + 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=0, horizontal_end=gtx.int32(grid.num_edges), vertical_start=0, From 77e5956c488bf4b1f867445061b31bfe7df9e08a Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 16:28:39 +0100 Subject: [PATCH 08/16] undo update uv lock --- pyproject.toml | 2 +- uv.lock | 36 ++++++++++++++++++++---------------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9d81dc118e..9d8a5abeac 100644 --- a/pyproject.toml +++ b/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/edopao/gt4py", branch = "gtfn_fix_canonicalize_domain"} +# gt4py = {git = "https://github.com/GridTools/gt4py", branch = "main"} # gt4py = {index = "test.pypi"} icon4py-atmosphere-advection = {workspace = true} icon4py-atmosphere-diffusion = {workspace = true} diff --git a/uv.lock b/uv.lock index dd66c4b123..0e6ebf090a 100644 --- a/uv.lock +++ b/uv.lock @@ -1399,8 +1399,8 @@ wheels = [ [[package]] name = "gt4py" -version = "1.1.2+unknown.version.details" -source = { git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain#f9d6d707667d5ee937bf854a04c7d4422ee9a6ad" } +version = "1.1.2" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "black" }, @@ -1430,6 +1430,10 @@ dependencies = [ { name = "versioningit" }, { name = "xxhash" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/77/74/8a388e4be44fec6713b034df7950edfa8f1eb9c5e34feb639aaf38c4897c/gt4py-1.1.2.tar.gz", hash = "sha256:4f4999236e26a32d9f6d74aebdb21094fda2511d07e4f73615e396c35b13c78a", size = 773452, upload-time = "2025-12-08T12:15:48.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bd/9c/19c6263d0897a44542a97371df0a9d243e37654053d8b2f6d067fb8ed1a8/gt4py-1.1.2-py3-none-any.whl", hash = "sha256:44665eb59c4d174887fb787a01b8848ddf38607109648dfa1922d234e31153e2", size = 978719, upload-time = "2025-12-08T12:15:46.682Z" }, +] [package.optional-dependencies] cuda11 = [ @@ -1748,7 +1752,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1765,7 +1769,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1782,7 +1786,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1799,7 +1803,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1816,7 +1820,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1881,9 +1885,9 @@ requires-dist = [ { name = "dace", git = "https://github.com/GridTools/dace?tag=__gt4py-next-integration_2025_12_11" }, { name = "datashader", marker = "extra == 'io'", specifier = ">=0.16.1" }, { name = "ghex", marker = "extra == 'distributed'", git = "https://github.com/msimberg/GHEX.git?branch=async-mpi" }, - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, { name = "holoviews", marker = "extra == 'io'", specifier = ">=1.16.0" }, { name = "icon4py-common", extras = ["distributed", "io"], marker = "extra == 'all'", editable = "model/common" }, { name = "mpi4py", marker = "extra == 'distributed'", specifier = ">=3.1.5" }, @@ -1919,7 +1923,7 @@ dependencies = [ requires-dist = [ { name = "click", specifier = ">=8.0.1" }, { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -1947,7 +1951,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -1976,7 +1980,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "filelock", specifier = ">=3.18.0" }, - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, @@ -2025,9 +2029,9 @@ requires-dist = [ { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" }, { name = "cupy-cuda12x", marker = "extra == 'cuda12'", specifier = ">=13.0" }, { name = "fprettify", specifier = ">=0.3.7" }, - { name = "gt4py", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/edopao/gt4py?branch=gtfn_fix_canonicalize_domain" }, + { name = "gt4py", specifier = "==1.1.2" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, { name = "icon4py-atmosphere-advection", editable = "model/atmosphere/advection" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, From d0555cdedbed27821b63c20aeb8b28284ff4ccd1 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 16:33:47 +0100 Subject: [PATCH 09/16] undo extra change --- .../src/icon4py/model/atmosphere/dycore/dycore_utils.py | 6 +++--- .../dycore/stencils/apply_4th_order_divergence_damping.py | 4 ++-- .../compute_edge_diagnostics_for_dycore_and_update_vn.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py index 921f4be4af..9a8ac84a84 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/dycore_utils.py +++ b/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 GridType, abs, broadcast, maximum # noqa: A004 +from gt4py.next import 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(grid_type=GridType.UNSTRUCTURED) +@gtx.field_operator 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 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py index 2917771c26..d468c9f4bc 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py @@ -26,14 +26,14 @@ def _apply_4th_order_divergence_damping( second_order_divdamp_factor: float, ) -> fa.EdgeKField[wpfloat]: """Formerly known as _mo_solve_nonhydro_4th_order_divdamp.""" - fourth_order_divdamp_scaling_coeff = _calculate_fourth_order_divdamp_scaling_coeff( + 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) - vn_wp = vn + (fourth_order_divdamp_scaling_coeff * z_graddiv2_vn_wp) + vn_wp = vn + (scal_divdamp * z_graddiv2_vn_wp) return vn_wp diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index 5e5a3c85b7..57a2c1787b 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -9,7 +9,7 @@ import gt4py.next as gtx from gt4py.eve import utils as eve_utils -from gt4py.next import GridType, broadcast +from gt4py.next import broadcast from gt4py.next.experimental import concat_where from icon4py.model.atmosphere.dycore.stencils.add_analysis_increments_to_vn import ( @@ -270,7 +270,7 @@ def _compute_theta_rho_face_values_and_pressure_gradient_and_update_vn( ) -@gtx.field_operator(grid_type=GridType.UNSTRUCTURED) +@gtx.field_operator def _apply_divergence_damping_and_update_vn( horizontal_gradient_of_normal_wind_divergence: fa.EdgeKField[ta.vpfloat], next_vn: fa.EdgeKField[ta.wpfloat], From bc46cf976dedff1cdc08e125b0d7d7038542e98d Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 16:38:10 +0100 Subject: [PATCH 10/16] fix --- .../compute_edge_diagnostics_for_dycore_and_update_vn.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py index 57a2c1787b..a8e2e79385 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_edge_diagnostics_for_dycore_and_update_vn.py @@ -359,8 +359,6 @@ def _apply_divergence_damping_and_update_vn( 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, ) if is_iau_active: From eeebbcc494a743e76cb5b0d19ae7ed22676212a9 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 16:46:02 +0100 Subject: [PATCH 11/16] fix imports --- .../test_apply_4th_order_divergence_damping.py | 3 ++- .../test_apply_divergence_damping_and_update_vn.py | 6 +----- ...t_apply_weighted_2nd_and_4th_order_divergence_damping.py | 3 ++- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py index 8b1cfbec2d..b4ca1f1d56 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_4th_order_divergence_damping.py @@ -10,7 +10,6 @@ import gt4py.next as gtx import numpy as np import pytest -from model.atmosphere.dycore.tests.dycore.stencil_tests import test_dycore_utils from icon4py.model.atmosphere.dycore.stencils.apply_4th_order_divergence_damping import ( apply_4th_order_divergence_damping, @@ -22,6 +21,8 @@ from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest +from . import test_dycore_utils + def apply_4th_order_divergence_damping_numpy( scal_divdamp: np.ndarray, diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py index 5778d3f558..94849c4b3b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_divergence_damping_and_update_vn.py @@ -9,7 +9,6 @@ import gt4py.next as gtx import numpy as np import pytest -from model.atmosphere.dycore.tests.dycore.stencil_tests import test_dycore_utils import icon4py.model.common.type_alias as ta import icon4py.model.testing.stencil_tests as test_helpers @@ -21,10 +20,7 @@ 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, -) +from . import test_dycore_utils divergence_damp_order = DivergenceDampingOrder() diff --git a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py index 2d4ad1c29e..33505b461b 100644 --- a/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py +++ b/model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_weighted_2nd_and_4th_order_divergence_damping.py @@ -10,7 +10,6 @@ import gt4py.next as gtx import numpy as np import pytest -from model.atmosphere.dycore.tests.dycore.stencil_tests import test_dycore_utils from icon4py.model.atmosphere.dycore.stencils.apply_weighted_2nd_and_4th_order_divergence_damping import ( apply_weighted_2nd_and_4th_order_divergence_damping, @@ -22,6 +21,8 @@ from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest +from . import test_dycore_utils + def apply_weighted_2nd_and_4th_order_divergence_damping_numpy( scal_divdamp: np.ndarray, From fff888d54fd2f648e95cc696ebf0f886e17248c8 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 9 Jan 2026 17:41:21 +0100 Subject: [PATCH 12/16] edit test --- .../tests/dycore/integration_tests/test_solve_nonhydro.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index c6bbc89dcf..d57028ce14 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1633,7 +1633,6 @@ def test_apply_divergence_damping_and_update_vn( 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 @@ -1685,7 +1684,7 @@ def test_apply_divergence_damping_and_update_vn( divdamp_order=divdamp_order, mean_cell_area=mean_cell_area, second_order_divdamp_factor=second_order_divdamp_factor, - max_nudging_coefficient=max_nudging_coefficient, + max_nudging_coefficient=config.max_nudging_coefficient, dbl_eps=constants.DBL_EPS, horizontal_start=start_edge_nudging_level_2, horizontal_end=end_edge_local, From 6797ce32c44bf170cddf2e3e2ba3f10a7db70bf4 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Sat, 10 Jan 2026 08:25:56 +0100 Subject: [PATCH 13/16] edit test --- .../model/atmosphere/dycore/solve_nonhydro.py | 12 ------------ .../dycore/integration_tests/test_solve_nonhydro.py | 10 ---------- 2 files changed, 22 deletions(-) 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 db9e1f09d3..12780f8554 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 @@ -941,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) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index d57028ce14..d0118054a1 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -585,16 +585,6 @@ def test_nonhydro_corrector_step( at_last_substep=at_last_substep, ) - if icon_grid.limited_area: - assert test_utils.dallclose( - solve_nonhydro.reduced_fourth_order_divdamp_coeff_at_nest_boundary.asnumpy(), - init_savepoint.bdy_divdamp().asnumpy(), - ) - - assert test_utils.dallclose( - solve_nonhydro.fourth_order_divdamp_scaling_coeff.asnumpy(), - init_savepoint.scal_divdamp().asnumpy(), - ) # stencil 10 assert test_utils.dallclose( diagnostic_state_nh.rho_at_cells_on_half_levels.asnumpy(), From 6908b31ac8889033ab5f03401fd7d1fb3d423345 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Thu, 29 Jan 2026 11:44:15 +0100 Subject: [PATCH 14/16] apply review comments --- .../model/atmosphere/dycore/solve_nonhydro.py | 2 +- .../integration_tests/test_solve_nonhydro.py | 22 ++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) 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 326ba26f1c..8eaaad99c7 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 @@ -941,7 +941,7 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation self._grid, dims.KDim, dtype=ta.wpfloat, allocator=allocator ) """ - Declared as scal_divdamp in ICON. + Declared as enh_divdamp_fac in ICON. """ self.intermediate_fields = IntermediateFields.allocate(grid=self._grid, allocator=allocator) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 9259a2a4f7..8f1f0f3425 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1615,7 +1615,6 @@ 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() - 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() @@ -1624,6 +1623,13 @@ def test_apply_divergence_damping_and_update_vn( config = definitions.construct_nonhydrostatic_config(experiment) mean_cell_area = grid_savepoint.mean_cell_area() + # TODO(): Use serialized data ('enh_divdamp_fac' in icon) instead of computing 'interpolated_fourth_order_divdamp_factor' + interpolated_fourth_order_divdamp_factor = data_alloc.zero_field( + icon_grid, + dims.KDim, + allocator=backend, + ) + iau_wgt_dyn = config.iau_wgt_dyn divdamp_order = config.divdamp_order second_order_divdamp_scaling_coeff = sp_nh_init.divdamp_fac_o2() * mean_cell_area @@ -1643,6 +1649,20 @@ def test_apply_divergence_damping_and_update_vn( vn_ref = sp_nh_exit.vn_new() + smagorinsky.en_smag_fac_for_zero_nshift.with_backend(backend)( + grid_savepoint.vct_a(), + config.fourth_order_divdamp_factor, + config.fourth_order_divdamp_factor2, + config.fourth_order_divdamp_factor3, + config.fourth_order_divdamp_factor4, + config.fourth_order_divdamp_z, + config.fourth_order_divdamp_z2, + config.fourth_order_divdamp_z3, + config.fourth_order_divdamp_z4, + interpolated_fourth_order_divdamp_factor, + offset_provider={"Koff": dims.KDim}, + ) + compute_edge_diagnostics_for_dycore_and_update_vn.apply_divergence_damping_and_update_vn.with_backend( backend )( From e773485a03c8822db4567ec66342b896819e2086 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Thu, 29 Jan 2026 13:20:17 +0100 Subject: [PATCH 15/16] fix rebase error --- .../src/icon4py/model/atmosphere/dycore/solve_nonhydro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8eaaad99c7..b1b2c2dade 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 @@ -507,7 +507,7 @@ def __init__( "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, + "mean_cell_area": self._cell_params.mean_cell_area, "max_nudging_coefficient": self._config.max_nudging_coefficient, "dbl_eps": constants.DBL_EPS, }, From d63b792bde0fb9b8d598a58484afbfe60136ed67 Mon Sep 17 00:00:00 2001 From: edopao Date: Thu, 29 Jan 2026 15:24:08 +0100 Subject: [PATCH 16/16] Apply suggestion from @edopao --- .../tests/dycore/integration_tests/test_solve_nonhydro.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py index 8f1f0f3425..f00a231235 100644 --- a/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py @@ -1623,7 +1623,7 @@ def test_apply_divergence_damping_and_update_vn( config = definitions.construct_nonhydrostatic_config(experiment) mean_cell_area = grid_savepoint.mean_cell_area() - # TODO(): Use serialized data ('enh_divdamp_fac' in icon) instead of computing 'interpolated_fourth_order_divdamp_factor' + # TODO: Use serialized data ('enh_divdamp_fac' in icon) instead of computing 'interpolated_fourth_order_divdamp_factor' interpolated_fourth_order_divdamp_factor = data_alloc.zero_field( icon_grid, dims.KDim,