Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/main' into fix_symbolic_grid_sizes
Browse files Browse the repository at this point in the history
  • Loading branch information
@tehrengruber
tehrengruber committed Feb 1, 2024
2 parents 827c341 + 4cc9eef commit a2a6362
Show file tree
Hide file tree
Showing 45 changed files with 480 additions and 396 deletions.
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_18.py → ...ra_diffusion_for_w_con_approaching_cfl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@


@field_operator
def _mo_velocity_advection_stencil_18(
def _add_extra_diffusion_for_w_con_approaching_cfl(
levmask: Field[[KDim], bool],
cfl_clipping: Field[[CellDim, KDim], bool],
owner_mask: Field[[CellDim], bool],
Expand All @@ -42,6 +42,7 @@ def _mo_velocity_advection_stencil_18(
cfl_w_limit: vpfloat,
dtime: wpfloat,
) -> Field[[CellDim, KDim], vpfloat]:
"""Formerly known as _mo_velocity_advection_stencil_18."""
z_w_con_c_wp, ddqz_z_half_wp, ddt_w_adv_wp, cfl_w_limit_wp = astype(
(z_w_con_c, ddqz_z_half, ddt_w_adv, cfl_w_limit), wpfloat
)
Expand All @@ -66,7 +67,7 @@ def _mo_velocity_advection_stencil_18(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_18(
def add_extra_diffusion_for_w_con_approaching_cfl(
levmask: Field[[KDim], bool],
cfl_clipping: Field[[CellDim, KDim], bool],
owner_mask: Field[[CellDim], bool],
Expand All @@ -84,7 +85,7 @@ def mo_velocity_advection_stencil_18(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_18(
_add_extra_diffusion_for_w_con_approaching_cfl(
levmask,
cfl_clipping,
owner_mask,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_20.py → ...extra_diffusion_for_wn_approaching_cfl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@


@field_operator
def _mo_velocity_advection_stencil_20(
def _add_extra_diffusion_for_wn_approaching_cfl(
levelmask: Field[[KDim], bool],
c_lin_e: Field[[EdgeDim, E2CDim], wpfloat],
z_w_con_c_full: Field[[CellDim, KDim], vpfloat],
Expand All @@ -58,6 +58,7 @@ def _mo_velocity_advection_stencil_20(
scalfac_exdiff: wpfloat,
dtime: wpfloat,
) -> Field[[EdgeDim, KDim], vpfloat]:
"""Formerly known as _mo_velocity_advection_stencil_20."""
z_w_con_c_full_wp, ddqz_z_full_e_wp, ddt_vn_apc_wp, cfl_w_limit_wp = astype(
(z_w_con_c_full, ddqz_z_full_e, ddt_vn_apc, cfl_w_limit), wpfloat
)
Expand Down Expand Up @@ -97,7 +98,7 @@ def _mo_velocity_advection_stencil_20(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_20(
def add_extra_diffusion_for_wn_approaching_cfl(
levelmask: Field[[KDim], bool],
c_lin_e: Field[[EdgeDim, E2CDim], wpfloat],
z_w_con_c_full: Field[[CellDim, KDim], vpfloat],
Expand All @@ -117,7 +118,7 @@ def mo_velocity_advection_stencil_20(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_20(
_add_extra_diffusion_for_wn_approaching_cfl(
levelmask,
c_lin_e,
z_w_con_c_full,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_17.py → ...interpolated_horizontal_advection_of_w.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,18 +20,19 @@


@field_operator
def _mo_velocity_advection_stencil_17(
def _add_interpolated_horizontal_advection_of_w(
e_bln_c_s: Field[[CEDim], wpfloat],
z_v_grad_w: Field[[EdgeDim, KDim], vpfloat],
ddt_w_adv: Field[[CellDim, KDim], vpfloat],
) -> Field[[CellDim, KDim], vpfloat]:
"""Formerly known as _mo_velocity_advection_stencil_17."""
z_v_grad_w_wp, ddt_w_adv_wp = astype((z_v_grad_w, ddt_w_adv), wpfloat)
ddt_w_adv_wp = ddt_w_adv_wp + neighbor_sum(z_v_grad_w_wp(C2E) * e_bln_c_s(C2CE), axis=C2EDim)
return astype(ddt_w_adv_wp, vpfloat)


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_17(
def add_interpolated_horizontal_advection_of_w(
e_bln_c_s: Field[[CEDim], wpfloat],
z_v_grad_w: Field[[EdgeDim, KDim], vpfloat],
ddt_w_adv: Field[[CellDim, KDim], vpfloat],
Expand All @@ -40,7 +41,7 @@ def mo_velocity_advection_stencil_17(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_17(
_add_interpolated_horizontal_advection_of_w(
e_bln_c_s,
z_v_grad_w,
ddt_w_adv,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_19.py → ...compute_advective_normal_wind_tendency.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@


@field_operator
def _mo_velocity_advection_stencil_19(
def _compute_advective_normal_wind_tendency(
z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat],
coeff_gradekin: Field[[ECDim], vpfloat],
z_ekinh: Field[[CellDim, KDim], vpfloat],
Expand All @@ -44,6 +44,7 @@ def _mo_velocity_advection_stencil_19(
vn_ie: Field[[EdgeDim, KDim], vpfloat],
ddqz_z_full_e: Field[[EdgeDim, KDim], vpfloat],
) -> Field[[EdgeDim, KDim], vpfloat]:
"""Formerly known as _mo_velocity_advection_stencil_19."""
vt_wp, z_w_con_c_full_wp, ddqz_z_full_e_wp = astype(
(vt, z_w_con_c_full, ddqz_z_full_e), wpfloat
)
Expand All @@ -65,7 +66,7 @@ def _mo_velocity_advection_stencil_19(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_19(
def compute_advective_normal_wind_tendency(
z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat],
coeff_gradekin: Field[[ECDim], vpfloat],
z_ekinh: Field[[CellDim, KDim], vpfloat],
Expand All @@ -82,7 +83,7 @@ def mo_velocity_advection_stencil_19(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_19(
_compute_advective_normal_wind_tendency(
z_kin_hor_e,
coeff_gradekin,
z_ekinh,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_16.py → ...mpute_advective_vertical_wind_tendency.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,13 @@


@field_operator
def _mo_velocity_advection_stencil_16(
def _compute_advective_vertical_wind_tendency(
z_w_con_c: Field[[CellDim, KDim], vpfloat],
w: Field[[CellDim, KDim], wpfloat],
coeff1_dwdz: Field[[CellDim, KDim], vpfloat],
coeff2_dwdz: Field[[CellDim, KDim], vpfloat],
) -> Field[[CellDim, KDim], vpfloat]:
"""Formerly known as _mo_velocity_advection_stencil_16."""
z_w_con_c_wp = astype(z_w_con_c, wpfloat)
coeff1_dwdz_wp, coeff2_dwdz_wp = astype((coeff1_dwdz, coeff2_dwdz), wpfloat)

Expand All @@ -38,7 +39,7 @@ def _mo_velocity_advection_stencil_16(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_16(
def compute_advective_vertical_wind_tendency(
z_w_con_c: Field[[CellDim, KDim], vpfloat],
w: Field[[CellDim, KDim], wpfloat],
coeff1_dwdz: Field[[CellDim, KDim], vpfloat],
Expand All @@ -49,7 +50,7 @@ def mo_velocity_advection_stencil_16(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_16(
_compute_advective_vertical_wind_tendency(
z_w_con_c,
w,
coeff1_dwdz,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_07.py → ...l_advection_term_for_vertical_velocity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@


@field_operator
def _mo_velocity_advection_stencil_07(
def _compute_horizontal_advection_term_for_vertical_velocity(
vn_ie: Field[[EdgeDim, KDim], vpfloat],
inv_dual_edge_length: Field[[EdgeDim], wpfloat],
w: Field[[CellDim, KDim], wpfloat],
Expand All @@ -29,6 +29,7 @@ def _mo_velocity_advection_stencil_07(
tangent_orientation: Field[[EdgeDim], wpfloat],
z_w_v: Field[[VertexDim, KDim], vpfloat],
) -> Field[[EdgeDim, KDim], vpfloat]:
"""Formerly know as _mo_velocity_advection_stencil_07."""
z_vt_ie_wp, vn_ie_wp = astype((z_vt_ie, vn_ie), wpfloat)

z_v_grad_w_wp = vn_ie_wp * inv_dual_edge_length * (
Expand All @@ -40,7 +41,7 @@ def _mo_velocity_advection_stencil_07(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_07(
def compute_horizontal_advection_term_for_vertical_velocity(
vn_ie: Field[[EdgeDim, KDim], vpfloat],
inv_dual_edge_length: Field[[EdgeDim], wpfloat],
w: Field[[CellDim, KDim], wpfloat],
Expand All @@ -54,7 +55,7 @@ def mo_velocity_advection_stencil_07(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_07(
_compute_horizontal_advection_term_for_vertical_velocity(
vn_ie,
inv_dual_edge_length,
w,
Expand Down
8 changes: 4 additions & 4 deletions ...re/dycore/mo_solve_nonhydro_stencil_37.py → ...core/compute_horizontal_kinetic_energy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,23 +20,23 @@


@field_operator
def _mo_solve_nonhydro_stencil_37(
def _compute_horizontal_kinetic_energy(
vn: Field[[EdgeDim, KDim], wpfloat],
vt: Field[[EdgeDim, KDim], vpfloat],
) -> tuple[
Field[[EdgeDim, KDim], vpfloat],
Field[[EdgeDim, KDim], vpfloat],
Field[[EdgeDim, KDim], vpfloat],
]:

"""Formerly known as _mo_solve_nonhydro_stencil_37 or _mo_velocity_advection_stencil_05."""
vn_ie_wp = vn
z_vt_ie_vp = vt
z_kin_hor_e_wp = wpfloat("0.5") * (vn * vn + astype(vt * vt, wpfloat))
return astype(vn_ie_wp, vpfloat), z_vt_ie_vp, astype(z_kin_hor_e_wp, vpfloat)


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_37(
def compute_horizontal_kinetic_energy(
vn: Field[[EdgeDim, KDim], wpfloat],
vt: Field[[EdgeDim, KDim], vpfloat],
vn_ie: Field[[EdgeDim, KDim], vpfloat],
Expand All @@ -47,7 +47,7 @@ def mo_solve_nonhydro_stencil_37(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_37(
_compute_horizontal_kinetic_energy(
vn,
vt,
out=(vn_ie, z_vt_ie, z_kin_hor_e),
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_14.py → ...d_clip_contravariant_vertical_velocity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@


@field_operator
def _mo_velocity_advection_stencil_14(
def _compute_maximum_cfl_and_clip_contravariant_vertical_velocity(
ddqz_z_half: Field[[CellDim, KDim], vpfloat],
z_w_con_c: Field[[CellDim, KDim], vpfloat],
cfl_w_limit: vpfloat,
Expand All @@ -37,6 +37,7 @@ def _mo_velocity_advection_stencil_14(
Field[[CellDim, KDim], vpfloat],
Field[[CellDim, KDim], vpfloat],
]:
"""Formerly know as _mo_velocity_advection_stencil_14."""
z_w_con_c_wp, ddqz_z_half_wp = astype((z_w_con_c, ddqz_z_half), wpfloat)

cfl_clipping = where(
Expand Down Expand Up @@ -64,7 +65,7 @@ def _mo_velocity_advection_stencil_14(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_14(
def compute_maximum_cfl_and_clip_contravariant_vertical_velocity(
ddqz_z_half: Field[[CellDim, KDim], vpfloat],
z_w_con_c: Field[[CellDim, KDim], vpfloat],
cfl_clipping: Field[[CellDim, KDim], bool],
Expand All @@ -76,7 +77,7 @@ def mo_velocity_advection_stencil_14(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_14(
_compute_maximum_cfl_and_clip_contravariant_vertical_velocity(
ddqz_z_half,
z_w_con_c,
cfl_w_limit,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_01.py → ...osphere/dycore/compute_tangential_wind.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,16 +20,17 @@


@field_operator
def _mo_velocity_advection_stencil_01(
def _compute_tangential_wind(
vn: Field[[EdgeDim, KDim], wpfloat],
rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], wpfloat],
) -> Field[[EdgeDim, KDim], vpfloat]:
"""Formerly knowan as _mo_velocity_advection_stencil_01."""
vt_wp = neighbor_sum(rbf_vec_coeff_e * vn(E2C2E), axis=E2C2EDim)
return astype(vt_wp, vpfloat)


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_01(
def compute_tangential_wind(
vn: Field[[EdgeDim, KDim], wpfloat],
rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], wpfloat],
vt: Field[[EdgeDim, KDim], vpfloat],
Expand All @@ -38,7 +39,7 @@ def mo_velocity_advection_stencil_01(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_01(
_compute_tangential_wind(
vn,
rbf_vec_coeff_e,
out=vt,
Expand Down
7 changes: 4 additions & 3 deletions ...ycore/mo_velocity_advection_stencil_13.py → ...orrect_contravariant_vertical_velocity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,24 +20,25 @@


@field_operator
def _mo_velocity_advection_stencil_13(
def _correct_contravariant_vertical_velocity(
z_w_con_c: Field[[CellDim, KDim], vpfloat],
w_concorr_c: Field[[CellDim, KDim], vpfloat],
) -> Field[[CellDim, KDim], vpfloat]:
"""Formerly known as _mo_velocity_advection_stencil_13."""
z_w_con_c_vp = z_w_con_c - w_concorr_c
return z_w_con_c_vp


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_13(
def correct_contravariant_vertical_velocity(
w_concorr_c: Field[[CellDim, KDim], vpfloat],
z_w_con_c: Field[[CellDim, KDim], vpfloat],
horizontal_start: int32,
horizontal_end: int32,
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_13(
_correct_contravariant_vertical_velocity(
z_w_con_c,
w_concorr_c,
out=z_w_con_c,
Expand Down
24 changes: 12 additions & 12 deletions ...e/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,17 @@
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, where

from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_15 import (
_mo_velocity_advection_stencil_15,
from icon4py.model.atmosphere.dycore.add_extra_diffusion_for_w_con_approaching_cfl import (
_add_extra_diffusion_for_w_con_approaching_cfl,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_16 import (
_mo_velocity_advection_stencil_16,
from icon4py.model.atmosphere.dycore.add_interpolated_horizontal_advection_of_w import (
_add_interpolated_horizontal_advection_of_w,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_17 import (
_mo_velocity_advection_stencil_17,
from icon4py.model.atmosphere.dycore.compute_advective_vertical_wind_tendency import (
_compute_advective_vertical_wind_tendency,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_18 import (
_mo_velocity_advection_stencil_18,
from icon4py.model.atmosphere.dycore.interpolate_contravatiant_vertical_verlocity_to_full_levels import (
_interpolate_contravatiant_vertical_verlocity_to_full_levels,
)
from icon4py.model.common.dimension import C2E2CODim, CEDim, CellDim, EdgeDim, KDim
from icon4py.model.common.type_alias import vpfloat, wpfloat
Expand Down Expand Up @@ -60,19 +60,19 @@ def _fused_velocity_advection_stencil_16_to_18(

ddt_w_adv = where(
(cell_lower_bound <= cell < cell_upper_bound) & (int32(1) <= k),
_mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz),
_compute_advective_vertical_wind_tendency(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz),
ddt_w_adv,
)
ddt_w_adv = where(
(cell_lower_bound <= cell < cell_upper_bound) & (int32(1) <= k),
_mo_velocity_advection_stencil_17(e_bln_c_s, z_v_grad_w, ddt_w_adv),
_add_interpolated_horizontal_advection_of_w(e_bln_c_s, z_v_grad_w, ddt_w_adv),
ddt_w_adv,
)
ddt_w_adv = (
where(
(cell_lower_bound <= cell < cell_upper_bound)
& (maximum(2, nrdmax - 2) <= k < nlev - 3),
_mo_velocity_advection_stencil_18(
_add_extra_diffusion_for_w_con_approaching_cfl(
levelmask,
cfl_clipping,
owner_mask,
Expand Down Expand Up @@ -122,7 +122,7 @@ def _fused_velocity_advection_stencil_15_to_18(
lvn_only: bool,
extra_diffu: bool,
) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]:
z_w_con_c_full = _mo_velocity_advection_stencil_15(z_w_con_c)
z_w_con_c_full = _interpolate_contravatiant_vertical_verlocity_to_full_levels(z_w_con_c)
ddt_w_adv = (
_fused_velocity_advection_stencil_16_to_18(
z_w_con_c,
Expand Down
Loading

0 comments on commit a2a6362

Please sign in to comment.