Skip to content

Combine divergence damping coefficient calculation with other stencils #951

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
edopao merged 19 commits into from
Jan 29, 2026
Merged

Combine divergence damping coefficient calculation with other stencils #951

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
2a0ba10
Combine divergence damping coefficient calculation with other stencils
muellch Nov 19, 2025
1364f3c
Merge branch 'main' into combine_divergence_damping
nfarabullini Jan 6, 2026
d0c0d5e
fix merge errors
edopao Jan 7, 2026
327d675
update gt4py to main
edopao Jan 9, 2026
636ed7f
Merge branch 'main' into combine_divergence_damping
edopao Jan 9, 2026
86e3be4
cleanup
edopao Jan 9, 2026
955bc10
remove unused code
edopao Jan 9, 2026
6f7d6a7
move _calculate_divdamp_fields into apply_4th_order_divergence_damping
edopao Jan 9, 2026
9366be9
inline _calculate_divdamp_fields
edopao Jan 9, 2026
77e5956
undo update uv lock
edopao Jan 9, 2026
d0555cd
undo extra change
edopao Jan 9, 2026
bc46cf9
fix
edopao Jan 9, 2026
eeebbcc
fix imports
edopao Jan 9, 2026
fff888d
edit test
edopao Jan 9, 2026
6797ce3
edit test
edopao Jan 10, 2026
a22e093
Merge branch 'main' into combine_divergence_damping
edopao Jan 29, 2026
6908b31
apply review comments
edopao Jan 29, 2026
e773485
fix rebase error
edopao Jan 29, 2026
d63b792
Apply suggestion from @edopao
edopao Jan 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 6 additions & 31 deletions model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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._cell_params.mean_cell_area,
"max_nudging_coefficient": self._config.max_nudging_coefficient,
"dbl_eps": constants.DBL_EPS,
},
variants={
"apply_2nd_order_divergence_damping": [False, True],
Expand Down Expand Up @@ -684,16 +688,6 @@ def __init__(
"vertical_end": gtx.int32(self._grid.num_levels),
},
)
self._calculate_divdamp_fields = setup_program(
backend=backend,
program=dycore_utils.calculate_divdamp_fields,
constant_args={
"divdamp_order": gtx.int32(self._config.divdamp_order),
"mean_cell_area": self._cell_params.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,
Expand Down Expand Up @@ -949,18 +943,6 @@ def _allocate_local_fields(self, allocator: gtx_allocators.FieldBufferAllocation
"""
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)

def _determine_local_domains(self):
Expand Down Expand Up @@ -1294,13 +1276,6 @@ def run_corrector_step(
second_order_divdamp_factor * self._cell_params.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,
Expand Down Expand Up @@ -1350,8 +1325,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,
Expand Down
179 changes: 0 additions & 179 deletions model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/solve_nonhydro_stencils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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],
Expand Down
29 changes: 23 additions & 6 deletions ...dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_4th_order_divergence_damping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,40 +6,57 @@
# Please, refer to the LICENSE file in the root directory.
# SPDX-License-Identifier: BSD-3-Clause
import gt4py.next as gtx
from gt4py.next import astype, broadcast
from gt4py.next import astype

from icon4py.model.atmosphere.dycore.dycore_utils import (
_calculate_fourth_order_divdamp_scaling_coeff,
)
from icon4py.model.common import field_type_aliases as fa
from icon4py.model.common.dimension import EdgeDim, KDim
from icon4py.model.common.type_alias import vpfloat, wpfloat


@gtx.field_operator
def _apply_4th_order_divergence_damping(
scal_divdamp: fa.KField[wpfloat],
interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat],
z_graddiv2_vn: fa.EdgeKField[vpfloat],
vn: fa.EdgeKField[wpfloat],
divdamp_order: gtx.int32,
mean_cell_area: float,
second_order_divdamp_factor: float,
) -> fa.EdgeKField[wpfloat]:
"""Formelry known as _mo_solve_nonhydro_4th_order_divdamp."""
"""Formerly known as _mo_solve_nonhydro_4th_order_divdamp."""
scal_divdamp = _calculate_fourth_order_divdamp_scaling_coeff(
interpolated_fourth_order_divdamp_factor,
divdamp_order,
mean_cell_area,
second_order_divdamp_factor,
)
z_graddiv2_vn_wp = astype(z_graddiv2_vn, wpfloat)
scal_divdamp = broadcast(scal_divdamp, (EdgeDim, KDim))
vn_wp = vn + (scal_divdamp * z_graddiv2_vn_wp)
return vn_wp


@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
def apply_4th_order_divergence_damping(
scal_divdamp: fa.KField[wpfloat],
interpolated_fourth_order_divdamp_factor: fa.KField[wpfloat],
z_graddiv2_vn: fa.EdgeKField[vpfloat],
vn: fa.EdgeKField[wpfloat],
divdamp_order: gtx.int32,
mean_cell_area: float,
second_order_divdamp_factor: float,
horizontal_start: gtx.int32,
horizontal_end: gtx.int32,
vertical_start: gtx.int32,
vertical_end: gtx.int32,
):
_apply_4th_order_divergence_damping(
scal_divdamp,
interpolated_fourth_order_divdamp_factor,
z_graddiv2_vn,
vn,
divdamp_order,
mean_cell_area,
second_order_divdamp_factor,
out=vn,
domain={
EdgeDim: (horizontal_start, horizontal_end),
Expand Down
Loading
Loading