Skip to content

Commit

Permalink
Remove and rename duplicate stencils (#332)
Browse files Browse the repository at this point in the history
Remove and rename duplicate stencils
  • Loading branch information
huppd authored Dec 21, 2023
1 parent 6e153b5 commit f7e7c24
Show file tree
Hide file tree
Showing 24 changed files with 95 additions and 578 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -20,20 +20,21 @@


@field_operator
def _mo_velocity_advection_stencil_04(
def _compute_contravariant_correction(
vn: Field[[EdgeDim, KDim], wpfloat],
ddxn_z_full: Field[[EdgeDim, KDim], vpfloat],
ddxt_z_full: Field[[EdgeDim, KDim], vpfloat],
vt: Field[[EdgeDim, KDim], vpfloat],
) -> Field[[EdgeDim, KDim], vpfloat]:
"""Formerly known as _mo_solve_nonhydro_stencil_35 or mo_velocity_advection_stencil_04."""
ddxn_z_full_wp = astype(ddxn_z_full, wpfloat)

z_w_concorr_me_wp = vn * ddxn_z_full_wp + astype(vt * ddxt_z_full, wpfloat)
return astype(z_w_concorr_me_wp, vpfloat)


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_04(
def compute_contravariant_correction(
vn: Field[[EdgeDim, KDim], wpfloat],
ddxn_z_full: Field[[EdgeDim, KDim], vpfloat],
ddxt_z_full: Field[[EdgeDim, KDim], vpfloat],
Expand All @@ -44,7 +45,7 @@ def mo_velocity_advection_stencil_04(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_04(
_compute_contravariant_correction(
vn,
ddxn_z_full,
ddxt_z_full,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,13 @@


@field_operator
def _mo_solve_nonhydro_stencil_07(
def _compute_pertubation_of_rho_and_theta(
rho: Field[[CellDim, KDim], wpfloat],
rho_ref_mc: Field[[CellDim, KDim], vpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
theta_ref_mc: Field[[CellDim, KDim], vpfloat],
) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]:
"""Formerly known as mo_solve_nonhydro_stencil_07 or mo_solve_nonhydro_stencil_13."""
rho_ref_mc_wp, theta_ref_mc_wp = astype((rho_ref_mc, theta_ref_mc), wpfloat)

z_rth_pr_1_wp = rho - rho_ref_mc_wp
Expand All @@ -34,7 +35,7 @@ def _mo_solve_nonhydro_stencil_07(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_07(
def compute_pertubation_of_rho_and_theta(
rho: Field[[CellDim, KDim], wpfloat],
rho_ref_mc: Field[[CellDim, KDim], vpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
Expand All @@ -46,7 +47,7 @@ def mo_solve_nonhydro_stencil_07(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_07(
_compute_pertubation_of_rho_and_theta(
rho,
rho_ref_mc,
theta_v,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,11 @@


@field_operator
def _mo_velocity_advection_stencil_06(
def _extrapolate_at_top(
wgtfacq_e: Field[[EdgeDim, KDim], vpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
) -> Field[[EdgeDim, KDim], vpfloat]:
"""Formerly known as mo_velocity_advection_stencil_06 or mo_solve_nonhydro_stencil_38."""
wgtfacq_e_wp = astype(wgtfacq_e, wpfloat)

vn_ie_wp = (
Expand All @@ -36,7 +37,7 @@ def _mo_velocity_advection_stencil_06(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_06(
def extrapolate_at_top(
wgtfacq_e: Field[[EdgeDim, KDim], vpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
vn_ie: Field[[EdgeDim, KDim], vpfloat],
Expand All @@ -45,7 +46,7 @@ def mo_velocity_advection_stencil_06(
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_06(
_extrapolate_at_top(
wgtfacq_e,
vn,
out=vn_ie,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import broadcast, int32, where

from icon4py.model.atmosphere.dycore.compute_contravariant_correction import (
_compute_contravariant_correction,
)
from icon4py.model.atmosphere.dycore.extrapolate_at_top import _extrapolate_at_top
from icon4py.model.atmosphere.dycore.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import (
_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl,
)
Expand All @@ -26,15 +30,9 @@
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_03 import (
_mo_velocity_advection_stencil_03,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_04 import (
_mo_velocity_advection_stencil_04,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_05 import (
_mo_velocity_advection_stencil_05,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_06 import (
_mo_velocity_advection_stencil_06,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_07 import (
_mo_velocity_advection_stencil_07,
)
Expand Down Expand Up @@ -93,11 +91,11 @@ def _fused_velocity_advection_stencil_1_to_6(
(vn_ie, z_vt_ie, z_kin_hor_e),
)

vn_ie = where(k == nlevp1, _mo_velocity_advection_stencil_06(wgtfacq_e_dsl, vn), vn_ie)
vn_ie = where(k == nlevp1, _extrapolate_at_top(wgtfacq_e_dsl, vn), vn_ie)

z_w_concorr_me = where(
nflatlev < k < nlevp1,
_mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt),
_compute_contravariant_correction(vn, ddxn_z_full, ddxt_z_full, vt),
z_w_concorr_me,
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,7 @@
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import int32, maximum, where

from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_08 import (
_mo_velocity_advection_stencil_08,
)
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_09 import (
_mo_velocity_advection_stencil_09,
)
from icon4py.model.atmosphere.dycore.interpolate_to_cell_center import _interpolate_to_cell_center
from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_10 import (
_mo_velocity_advection_stencil_10,
)
Expand Down Expand Up @@ -70,14 +65,14 @@ def _fused_velocity_advection_stencil_8_to_14(
]:
z_ekinh = where(
k < nlev,
_mo_velocity_advection_stencil_08(z_kin_hor_e, e_bln_c_s),
_interpolate_to_cell_center(z_kin_hor_e, e_bln_c_s),
z_ekinh,
)

z_w_concorr_mc = (
where(
nflatlev < k < nlev,
_mo_velocity_advection_stencil_09(z_w_concorr_me, e_bln_c_s),
_interpolate_to_cell_center(z_w_concorr_me, e_bln_c_s),
z_w_concorr_mc,
)
if istep == 1
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,30 +20,30 @@


@field_operator
def _mo_velocity_advection_stencil_08(
z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat],
def _interpolate_to_cell_center(
interpolant: Field[[EdgeDim, KDim], vpfloat],
e_bln_c_s: Field[[CEDim], wpfloat],
) -> Field[[CellDim, KDim], vpfloat]:
z_kin_hor_e_wp = astype(z_kin_hor_e, wpfloat)

z_ekinh_wp = neighbor_sum(z_kin_hor_e_wp(C2E) * e_bln_c_s(C2CE), axis=C2EDim)
return astype(z_ekinh_wp, vpfloat)
"""Formerly known as mo_velocity_advection_stencil_08 or mo_velocity_advection_stencil_09."""
interpolant_wp = astype(interpolant, wpfloat)
interpolation_wp = neighbor_sum(e_bln_c_s(C2CE) * interpolant_wp(C2E), axis=C2EDim)
return astype(interpolation_wp, vpfloat)


@program(grid_type=GridType.UNSTRUCTURED)
def mo_velocity_advection_stencil_08(
z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat],
def interpolate_to_cell_center(
interpolant: Field[[EdgeDim, KDim], vpfloat],
e_bln_c_s: Field[[CEDim], wpfloat],
z_ekinh: Field[[CellDim, KDim], vpfloat],
interpolation: Field[[CellDim, KDim], vpfloat],
horizontal_start: int32,
horizontal_end: int32,
vertical_start: int32,
vertical_end: int32,
):
_mo_velocity_advection_stencil_08(
z_kin_hor_e,
_interpolate_to_cell_center(
interpolant,
e_bln_c_s,
out=z_ekinh,
out=interpolation,
domain={
CellDim: (horizontal_start, horizontal_end),
KDim: (vertical_start, vertical_end),
Expand Down

This file was deleted.

This file was deleted.

This file was deleted.

Loading

0 comments on commit f7e7c24

Please sign in to comment.