diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_15_to_28.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_15_to_28.py index 01a2fffe6e..58c5349f0e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_15_to_28.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_solve_nonhydro_stencil_15_to_28.py @@ -21,9 +21,9 @@ # from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_4th_order_divdamp import ( # _mo_solve_nonhydro_4th_order_divdamp, # ) -# from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1 import ( -# _mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1, -# ) +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1 import ( + _mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1, +) # from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_17 import ( # _mo_solve_nonhydro_stencil_17, # ) @@ -76,16 +76,16 @@ def _fused_solve_nonhydro_stencil_15_to_28_predictor( geofac_grg_x: Field[[CellDim, C2E2CODim], float], geofac_grg_y: Field[[CellDim, C2E2CODim], float], - # p_vn: Field[[EdgeDim, KDim], float], - # p_vt: Field[[EdgeDim, KDim], float], - # pos_on_tplane_e_1: Field[[ECDim], float], - # pos_on_tplane_e_2: Field[[ECDim], float], - # primal_normal_cell_1: Field[[ECDim], float], - # dual_normal_cell_1: Field[[ECDim], float], - # primal_normal_cell_2: Field[[ECDim], float], - # dual_normal_cell_2: Field[[ECDim], float], - # rho_ref_me: Field[[EdgeDim, KDim], float], - # theta_ref_me: Field[[EdgeDim, KDim], float], + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + pos_on_tplane_e_1: Field[[ECDim], float], + pos_on_tplane_e_2: Field[[ECDim], float], + primal_normal_cell_1: Field[[ECDim], float], + dual_normal_cell_1: Field[[ECDim], float], + primal_normal_cell_2: Field[[ECDim], float], + dual_normal_cell_2: Field[[ECDim], float], + rho_ref_me: Field[[EdgeDim, KDim], float], + theta_ref_me: Field[[EdgeDim, KDim], float], z_rth_pr_1: Field[[CellDim, KDim], float], z_rth_pr_2: Field[[CellDim, KDim], float], # ddxn_z_full: Field[[EdgeDim, KDim], float], @@ -115,7 +115,7 @@ def _fused_solve_nonhydro_stencil_15_to_28_predictor( # grav_o_cpd: float, # dtime: float, # cpd: float, - # p_dthalf: float, + p_dthalf: float, # iau_wgt_dyn: float, # is_iau_active: bool, limited_area: bool, @@ -127,8 +127,8 @@ def _fused_solve_nonhydro_stencil_15_to_28_predictor( horizontal_upper_00: int32, horizontal_lower_01: int32, horizontal_upper_01: int32, - # horizontal_lower_1: int32, - # horizontal_upper_1: int32, + horizontal_lower_1: int32, + horizontal_upper_1: int32, # horizontal_lower_3: int32, # horizontal_upper_3: int32, horizontal_lower_4: int32, @@ -169,29 +169,29 @@ def _fused_solve_nonhydro_stencil_15_to_28_predictor( (z_rho_e, z_theta_v_e), ) if limited_area else (z_rho_e, z_theta_v_e) - # (z_rho_e, z_theta_v_e) = where( - # (horizontal_lower_1 <= horz_idx < horizontal_upper_1), - # _mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1( - # p_vn=p_vn, - # p_vt=p_vt, - # pos_on_tplane_e_1=pos_on_tplane_e_1, - # pos_on_tplane_e_2=pos_on_tplane_e_2, - # primal_normal_cell_1=primal_normal_cell_1, - # dual_normal_cell_1=dual_normal_cell_1, - # primal_normal_cell_2=primal_normal_cell_2, - # dual_normal_cell_2=dual_normal_cell_2, - # p_dthalf=p_dthalf, - # rho_ref_me=rho_ref_me, - # theta_ref_me=theta_ref_me, - # z_grad_rth_1=z_grad_rth_1, - # z_grad_rth_2=z_grad_rth_2, - # z_grad_rth_3=z_grad_rth_3, - # z_grad_rth_4=z_grad_rth_4, - # z_rth_pr_1=z_rth_pr_1, - # z_rth_pr_2=z_rth_pr_2, - # ), - # (z_rho_e, z_theta_v_e), - # ) + (z_rho_e, z_theta_v_e) = where( + (horizontal_lower_1 <= horz_idx < horizontal_upper_1), + _mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1( + p_vn=p_vn, + p_vt=p_vt, + pos_on_tplane_e_1=pos_on_tplane_e_1, + pos_on_tplane_e_2=pos_on_tplane_e_2, + primal_normal_cell_1=primal_normal_cell_1, + dual_normal_cell_1=dual_normal_cell_1, + primal_normal_cell_2=primal_normal_cell_2, + dual_normal_cell_2=dual_normal_cell_2, + p_dthalf=p_dthalf, + rho_ref_me=rho_ref_me, + theta_ref_me=theta_ref_me, + z_grad_rth_1=z_grad_rth_1, + z_grad_rth_2=z_grad_rth_2, + z_grad_rth_3=z_grad_rth_3, + z_grad_rth_4=z_grad_rth_4, + z_rth_pr_1=z_rth_pr_1, + z_rth_pr_2=z_rth_pr_2, + ), + (z_rho_e, z_theta_v_e), + ) # z_gradh_exner = where( # (horizontal_lower_0 <= horz_idx < horizontal_upper_0) & (vert_idx < nflatlev), @@ -397,16 +397,16 @@ def _fused_solve_nonhydro_stencil_15_to_28_predictor( def _fused_solve_nonhydro_stencil_15_to_28( geofac_grg_x: Field[[CellDim, C2E2CODim], float], geofac_grg_y: Field[[CellDim, C2E2CODim], float], - # p_vn: Field[[EdgeDim, KDim], float], - # p_vt: Field[[EdgeDim, KDim], float], - # pos_on_tplane_e_1: Field[[ECDim], float], - # pos_on_tplane_e_2: Field[[ECDim], float], - # primal_normal_cell_1: Field[[ECDim], float], - # dual_normal_cell_1: Field[[ECDim], float], - # primal_normal_cell_2: Field[[ECDim], float], - # dual_normal_cell_2: Field[[ECDim], float], - # rho_ref_me: Field[[EdgeDim, KDim], float], - # theta_ref_me: Field[[EdgeDim, KDim], float], + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + pos_on_tplane_e_1: Field[[ECDim], float], + pos_on_tplane_e_2: Field[[ECDim], float], + primal_normal_cell_1: Field[[ECDim], float], + dual_normal_cell_1: Field[[ECDim], float], + primal_normal_cell_2: Field[[ECDim], float], + dual_normal_cell_2: Field[[ECDim], float], + rho_ref_me: Field[[EdgeDim, KDim], float], + theta_ref_me: Field[[EdgeDim, KDim], float], z_rth_pr_1: Field[[CellDim, KDim], float], z_rth_pr_2: Field[[CellDim, KDim], float], # ddxn_z_full: Field[[EdgeDim, KDim], float], @@ -444,7 +444,7 @@ def _fused_solve_nonhydro_stencil_15_to_28( horz_idx: Field[[EdgeDim], int32], # vert_idx: Field[[KDim], int32], # grav_o_cpd: float, - # p_dthalf: float, + p_dthalf: float, idiv_method: int32, # igradp_method: int32, # wgt_nnow_vel: float, @@ -467,8 +467,8 @@ def _fused_solve_nonhydro_stencil_15_to_28( horizontal_upper_00: int32, horizontal_lower_01: int32, horizontal_upper_01: int32, - # horizontal_lower_1: int32, - # horizontal_upper_1: int32, + horizontal_lower_1: int32, + horizontal_upper_1: int32, # horizontal_lower_2: int32, # horizontal_upper_2: int32, # horizontal_lower_3: int32, @@ -490,18 +490,18 @@ def _fused_solve_nonhydro_stencil_15_to_28( ( z_rho_e, z_theta_v_e, z_gradh_exner, vn, ) = (_fused_solve_nonhydro_stencil_15_to_28_predictor( geofac_grg_x=geofac_grg_x, geofac_grg_y=geofac_grg_y, - # p_vn=p_vn, - # p_vt=p_vt, - # pos_on_tplane_e_1=pos_on_tplane_e_1, - # pos_on_tplane_e_2=pos_on_tplane_e_2, - # primal_normal_cell_1=primal_normal_cell_1, - # dual_normal_cell_1=dual_normal_cell_1, - # primal_normal_cell_2=primal_normal_cell_2, - # dual_normal_cell_2=dual_normal_cell_2, - # rho_ref_me=rho_ref_me, - # theta_ref_me=theta_ref_me, - z_rth_pr_1=z_rth_pr_1, - z_rth_pr_2=z_rth_pr_2, + p_vn=p_vn, + p_vt=p_vt, + pos_on_tplane_e_1=pos_on_tplane_e_1, + pos_on_tplane_e_2=pos_on_tplane_e_2, + primal_normal_cell_1=primal_normal_cell_1, + dual_normal_cell_1=dual_normal_cell_1, + primal_normal_cell_2=primal_normal_cell_2, + dual_normal_cell_2=dual_normal_cell_2, + rho_ref_me=rho_ref_me, + theta_ref_me=theta_ref_me, + z_rth_pr_1=z_rth_pr_1, + z_rth_pr_2=z_rth_pr_2, # ddxn_z_full=ddxn_z_full, # c_lin_e=c_lin_e, # z_exner_ex_pr=z_exner_ex_pr, @@ -529,7 +529,7 @@ def _fused_solve_nonhydro_stencil_15_to_28( # grav_o_cpd=grav_o_cpd, # dtime=dtime, # cpd=cpd, - # p_dthalf=p_dthalf, + p_dthalf=p_dthalf, # iau_wgt_dyn=iau_wgt_dyn, # is_iau_active=is_iau_active, limited_area=limited_area, @@ -541,8 +541,8 @@ def _fused_solve_nonhydro_stencil_15_to_28( horizontal_upper_00=horizontal_upper_00, horizontal_lower_01=horizontal_lower_01, horizontal_upper_01=horizontal_upper_01, - # horizontal_lower_1=horizontal_lower_1, - # horizontal_upper_1=horizontal_upper_1, + horizontal_lower_1=horizontal_lower_1, + horizontal_upper_1=horizontal_upper_1, # horizontal_lower_3=horizontal_lower_3, # horizontal_upper_3=horizontal_upper_3, horizontal_lower_4=horizontal_lower_4, @@ -608,16 +608,16 @@ def _fused_solve_nonhydro_stencil_15_to_28( def fused_solve_nonhydro_stencil_15_to_28( geofac_grg_x: Field[[CellDim, C2E2CODim], float], geofac_grg_y: Field[[CellDim, C2E2CODim], float], - # p_vn: Field[[EdgeDim, KDim], float], - # p_vt: Field[[EdgeDim, KDim], float], - # pos_on_tplane_e_1: Field[[ECDim], float], - # pos_on_tplane_e_2: Field[[ECDim], float], - # primal_normal_cell_1: Field[[ECDim], float], - # dual_normal_cell_1: Field[[ECDim], float], - # primal_normal_cell_2: Field[[ECDim], float], - # dual_normal_cell_2: Field[[ECDim], float], - # rho_ref_me: Field[[EdgeDim, KDim], float], - # theta_ref_me: Field[[EdgeDim, KDim], float], + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + pos_on_tplane_e_1: Field[[ECDim], float], + pos_on_tplane_e_2: Field[[ECDim], float], + primal_normal_cell_1: Field[[ECDim], float], + dual_normal_cell_1: Field[[ECDim], float], + primal_normal_cell_2: Field[[ECDim], float], + dual_normal_cell_2: Field[[ECDim], float], + rho_ref_me: Field[[EdgeDim, KDim], float], + theta_ref_me: Field[[EdgeDim, KDim], float], z_rth_pr_1: Field[[CellDim, KDim], float], z_rth_pr_2: Field[[CellDim, KDim], float], # ddxn_z_full: Field[[EdgeDim, KDim], float], @@ -655,7 +655,7 @@ def fused_solve_nonhydro_stencil_15_to_28( horz_idx: Field[[EdgeDim], int32], # vert_idx: Field[[KDim], int32], # grav_o_cpd: float, - # p_dthalf: float, + p_dthalf: float, idiv_method: int32, # igradp_method: int32, # wgt_nnow_vel: float, @@ -678,8 +678,8 @@ def fused_solve_nonhydro_stencil_15_to_28( horizontal_upper_00: int32, horizontal_lower_01: int32, horizontal_upper_01: int32, - # horizontal_lower_1: int32, - # horizontal_upper_1: int32, + horizontal_lower_1: int32, + horizontal_upper_1: int32, # horizontal_lower_2: int32, # horizontal_upper_2: int32, # horizontal_lower_3: int32, @@ -694,16 +694,16 @@ def fused_solve_nonhydro_stencil_15_to_28( _fused_solve_nonhydro_stencil_15_to_28( geofac_grg_x, geofac_grg_y, - # p_vn, - # p_vt, - # pos_on_tplane_e_1, - # pos_on_tplane_e_2, - # primal_normal_cell_1, - # dual_normal_cell_1, - # primal_normal_cell_2, - # dual_normal_cell_2, - # rho_ref_me, - # theta_ref_me, + p_vn, + p_vt, + pos_on_tplane_e_1, + pos_on_tplane_e_2, + primal_normal_cell_1, + dual_normal_cell_1, + primal_normal_cell_2, + dual_normal_cell_2, + rho_ref_me, + theta_ref_me, z_rth_pr_1, z_rth_pr_2, # ddxn_z_full, @@ -741,7 +741,7 @@ def fused_solve_nonhydro_stencil_15_to_28( horz_idx, # vert_idx, # grav_o_cpd, - # p_dthalf, + p_dthalf, idiv_method, # igradp_method, # wgt_nnow_vel, @@ -764,8 +764,8 @@ def fused_solve_nonhydro_stencil_15_to_28( horizontal_upper_00, horizontal_lower_01, horizontal_upper_01, - # horizontal_lower_1, - # horizontal_upper_1, + horizontal_lower_1, + horizontal_upper_1, # horizontal_lower_2, # horizontal_upper_2, # horizontal_lower_3,