remove code duplication

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
 
 

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

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
 
 

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)

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
 
 

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)