diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py index 0825affdcc..1adbe47c2d 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_08.py @@ -15,6 +15,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, astype, int32 +from icon4py.model.atmosphere.dycore.compute_pertubation_of_rho_and_theta import ( + _compute_pertubation_of_rho_and_theta, +) from icon4py.model.common.dimension import CellDim, KDim, Koff from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -31,13 +34,12 @@ def _mo_solve_nonhydro_stencil_08( Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat], ]: - wgtfac_c_wp, rho_ref_mc_wp, theta_ref_mc_wp = astype( - (wgtfac_c, rho_ref_mc, theta_ref_mc), wpfloat - ) + wgtfac_c_wp = astype(wgtfac_c, wpfloat) rho_ic = wgtfac_c_wp * rho + (wpfloat("1.0") - wgtfac_c_wp) * rho(Koff[-1]) - z_rth_pr_1 = astype(rho - rho_ref_mc_wp, vpfloat) - z_rth_pr_2 = astype(theta_v - theta_ref_mc_wp, vpfloat) + z_rth_pr_1, z_rth_pr_2 = _compute_pertubation_of_rho_and_theta( + rho=rho, rho_ref_mc=rho_ref_mc, theta_v=theta_v, theta_ref_mc=theta_ref_mc + ) return rho_ic, z_rth_pr_1, z_rth_pr_2 diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py index 8f65ab9ef8..22138e7ab6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_09.py @@ -15,6 +15,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, astype, int32 +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_05 import ( + _mo_solve_nonhydro_stencil_05, +) from icon4py.model.common.dimension import CellDim, KDim, Koff from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -37,7 +40,7 @@ def _mo_solve_nonhydro_stencil_09( (wgtfac_c, z_rth_pr_2, d_exner_dz_ref_ic, ddqz_z_half), wpfloat ) - z_theta_v_pr_ic_vp = wgtfac_c * z_rth_pr_2 + (vpfloat("1.0") - wgtfac_c) * z_rth_pr_2(Koff[-1]) + z_theta_v_pr_ic_vp = _mo_solve_nonhydro_stencil_05(wgtfac_c=wgtfac_c, z_exner_ex_pr=z_rth_pr_2) theta_v_ic_wp = wgtfac_c_wp * theta_v + (wpfloat("1.0") - wgtfac_c_wp) * theta_v(Koff[-1]) z_th_ddz_exner_c_wp = vwind_expl_wgt * theta_v_ic_wp * ( exner_pr(Koff[-1]) - exner_pr diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py index 3489658b53..984df48742 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_30.py @@ -15,7 +15,13 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, astype, int32, neighbor_sum -from icon4py.model.common.dimension import E2C2E, E2C2EO, E2C2EDim, E2C2EODim, EdgeDim, KDim +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_31 import ( + _mo_solve_nonhydro_stencil_31, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_01 import ( + _mo_velocity_advection_stencil_01, +) +from icon4py.model.common.dimension import E2C2EO, E2C2EDim, E2C2EODim, EdgeDim, KDim from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -30,9 +36,9 @@ def _mo_solve_nonhydro_stencil_30( Field[[EdgeDim, KDim], vpfloat], Field[[EdgeDim, KDim], vpfloat], ]: - z_vn_avg_wp = neighbor_sum(e_flx_avg * vn(E2C2EO), axis=E2C2EODim) + z_vn_avg_wp = _mo_solve_nonhydro_stencil_31(e_flx_avg=e_flx_avg, vn=vn) z_graddiv_vn_vp = astype(neighbor_sum(geofac_grdiv * vn(E2C2EO), axis=E2C2EODim), vpfloat) - vt_vp = astype(neighbor_sum(rbf_vec_coeff_e * vn(E2C2E), axis=E2C2EDim), vpfloat) + vt_vp = _mo_velocity_advection_stencil_01(vn=vn, rbf_vec_coeff_e=rbf_vec_coeff_e) return z_vn_avg_wp, z_graddiv_vn_vp, vt_vp diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py index 6a54552db6..a20b752761 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_36.py @@ -13,9 +13,15 @@ from gt4py.next.common import GridType from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, astype, int32 +from gt4py.next.ffront.fbuiltins import Field, int32 -from icon4py.model.common.dimension import EdgeDim, KDim, Koff +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_02 import ( + _mo_velocity_advection_stencil_02, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_03 import ( + _mo_velocity_advection_stencil_03, +) +from icon4py.model.common.dimension import EdgeDim, KDim from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -29,12 +35,9 @@ def _mo_solve_nonhydro_stencil_36( Field[[EdgeDim, KDim], vpfloat], Field[[EdgeDim, KDim], vpfloat], ]: - wgtfac_e_wp, vt_wp = astype((wgtfac_e, vt), wpfloat) - - vn_ie_wp = wgtfac_e_wp * vn + (wpfloat("1.0") - wgtfac_e_wp) * vn(Koff[-1]) - z_vt_ie_wp = astype(wgtfac_e * vt, wpfloat) + (wpfloat("1.0") - wgtfac_e_wp) * vt_wp(Koff[-1]) - z_kin_hor_e_wp = wpfloat("0.5") * (vn * vn + astype(vt * vt, wpfloat)) - return astype((vn_ie_wp, z_vt_ie_wp, z_kin_hor_e_wp), vpfloat) + z_vt_ie = _mo_velocity_advection_stencil_03(wgtfac_e=wgtfac_e, vt=vt) + vn_ie, z_kin_hor_e = _mo_velocity_advection_stencil_02(wgtfac_e=wgtfac_e, vn=vn, vt=vt) + return vn_ie, z_vt_ie, z_kin_hor_e @program(grid_type=GridType.UNSTRUCTURED) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py index eca9cc8954..ef87f67c43 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_solve_nonhydro_stencil_61.py @@ -15,6 +15,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, int32 +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_62 import ( + _mo_solve_nonhydro_stencil_62, +) from icon4py.model.common.dimension import CellDim, KDim from icon4py.model.common.type_alias import wpfloat @@ -35,7 +38,7 @@ def _mo_solve_nonhydro_stencil_61( ]: rho_new_wp = rho_now + dtime * grf_tend_rho exner_new_wp = theta_v_now + dtime * grf_tend_thv - w_new_wp = w_now + dtime * grf_tend_w + w_new_wp = _mo_solve_nonhydro_stencil_62(w_now=w_now, grf_tend_w=grf_tend_w, dtime=dtime) return rho_new_wp, exner_new_wp, w_new_wp diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py index eedc4283b6..a154e64db2 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_velocity_advection_stencil_02.py @@ -15,6 +15,9 @@ from gt4py.next.ffront.decorator import field_operator, program from gt4py.next.ffront.fbuiltins import Field, astype, int32 +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_05 import ( + _mo_velocity_advection_stencil_05, +) from icon4py.model.common.dimension import EdgeDim, KDim, Koff from icon4py.model.common.type_alias import vpfloat, wpfloat @@ -28,9 +31,9 @@ def _mo_velocity_advection_stencil_02( wgtfac_e_wp = astype(wgtfac_e, wpfloat) vn_ie_wp = wgtfac_e_wp * vn + (wpfloat("1.0") - wgtfac_e_wp) * vn(Koff[-1]) - z_kin_hor_e_wp = wpfloat("0.5") * (vn * vn + astype(vt * vt, wpfloat)) + _, _, z_kin_hor_e = _mo_velocity_advection_stencil_05(vn=vn, vt=vt) - return astype((vn_ie_wp, z_kin_hor_e_wp), vpfloat) + return astype(vn_ie_wp, vpfloat), z_kin_hor_e @program(grid_type=GridType.UNSTRUCTURED)