Skip to content

Commit

Permalink
put all input arguments back
Browse files Browse the repository at this point in the history
  • Loading branch information
nfarabullini committed Nov 20, 2023
1 parent 53880a6 commit 68b67c5
Showing 1 changed file with 161 additions and 156 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -409,76 +409,76 @@ def _fused_solve_nonhydro_stencil_15_to_28(
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],
# c_lin_e: Field[[EdgeDim, E2CDim], float],
# z_exner_ex_pr: Field[[CellDim, KDim], float],
# z_dexner_dz_c_1: Field[[CellDim, KDim], float],
# z_dexner_dz_c_2: Field[[CellDim, KDim], float],
# theta_v: Field[[CellDim, KDim], float],
# ikoffset: Field[[ECDim, KDim], int32],
# zdiff_gradp: Field[[ECDim, KDim], float],
# theta_v_ic: Field[[CellDim, KDim], float],
# inv_ddqz_z_full: Field[[CellDim, KDim], float],
# ipeidx_dsl: Field[[EdgeDim, KDim], bool],
# pg_exdist: Field[[EdgeDim, KDim], float],
# hmask_dd3d: Field[[EdgeDim], float],
# scalfac_dd3d: Field[[KDim], float],
# z_dwdz_dd: Field[[CellDim, KDim], float],
# inv_dual_edge_length: Field[[EdgeDim], float],
# ddt_vn_apc_ntl2: Field[[EdgeDim, KDim], float],
# vn_nnow: Field[[EdgeDim, KDim], float],
# ddt_vn_apc_ntl1: Field[[EdgeDim, KDim], float],
# ddt_vn_phy: Field[[EdgeDim, KDim], float],
ddxn_z_full: Field[[EdgeDim, KDim], float],
c_lin_e: Field[[EdgeDim, E2CDim], float],
z_exner_ex_pr: Field[[CellDim, KDim], float],
z_dexner_dz_c_1: Field[[CellDim, KDim], float],
z_dexner_dz_c_2: Field[[CellDim, KDim], float],
theta_v: Field[[CellDim, KDim], float],
ikoffset: Field[[ECDim, KDim], int32],
zdiff_gradp: Field[[ECDim, KDim], float],
theta_v_ic: Field[[CellDim, KDim], float],
inv_ddqz_z_full: Field[[CellDim, KDim], float],
ipeidx_dsl: Field[[EdgeDim, KDim], bool],
pg_exdist: Field[[EdgeDim, KDim], float],
hmask_dd3d: Field[[EdgeDim], float],
scalfac_dd3d: Field[[KDim], float],
z_dwdz_dd: Field[[CellDim, KDim], float],
inv_dual_edge_length: Field[[EdgeDim], float],
ddt_vn_apc_ntl2: Field[[EdgeDim, KDim], float],
vn_nnow: Field[[EdgeDim, KDim], float],
ddt_vn_apc_ntl1: Field[[EdgeDim, KDim], float],
ddt_vn_phy: Field[[EdgeDim, KDim], float],
z_graddiv_vn: Field[[EdgeDim, KDim], float],
# vn_incr: Field[[EdgeDim, KDim], float],
vn_incr: Field[[EdgeDim, KDim], float],
vn: Field[[EdgeDim, KDim], float],
z_rho_e: Field[[EdgeDim, KDim], float],
z_theta_v_e: Field[[EdgeDim, KDim], float],
z_gradh_exner: Field[[EdgeDim, KDim], float],
# z_graddiv2_vn: Field[[EdgeDim, KDim], float],
# z_hydro_corr: Field[[EdgeDim, KDim], float],
# geofac_grdiv: Field[[EdgeDim, E2C2EODim], float],
# scal_divdamp: Field[[KDim], float],
# bdy_divdamp: Field[[KDim], float],
# nudgecoeff_e: Field[[EdgeDim], float],
z_graddiv2_vn: Field[[EdgeDim, KDim], float],
z_hydro_corr: Field[[EdgeDim, KDim], float],
geofac_grdiv: Field[[EdgeDim, E2C2EODim], float],
scal_divdamp: Field[[KDim], float],
bdy_divdamp: Field[[KDim], float],
nudgecoeff_e: Field[[EdgeDim], float],
horz_idx: Field[[EdgeDim], int32],
# vert_idx: Field[[KDim], int32],
# grav_o_cpd: float,
vert_idx: Field[[KDim], int32],
grav_o_cpd: float,
p_dthalf: float,
idiv_method: int32,
# igradp_method: int32,
# wgt_nnow_vel: float,
# wgt_nnew_vel: float,
# dtime: float,
# cpd: float,
# iau_wgt_dyn: float,
# is_iau_active: bool,
# lhdiff_rcf: bool,
# divdamp_fac: float,
# divdamp_fac_o2: float,
# divdamp_order: int32,
# scal_divdamp_o2: float,
igradp_method: int32,
wgt_nnow_vel: float,
wgt_nnew_vel: float,
dtime: float,
cpd: float,
iau_wgt_dyn: float,
is_iau_active: bool,
lhdiff_rcf: bool,
divdamp_fac: float,
divdamp_fac_o2: float,
divdamp_order: int32,
scal_divdamp_o2: float,
limited_area: bool,
# itime_scheme: int32,
itime_scheme: int32,
istep: int32,
# horizontal_lower_0: int32,
# horizontal_upper_0: int32,
horizontal_lower_0: int32,
horizontal_upper_0: int32,
horizontal_lower_00: int32,
horizontal_upper_00: int32,
horizontal_lower_01: int32,
horizontal_upper_01: int32,
horizontal_lower_1: int32,
horizontal_upper_1: int32,
# horizontal_lower_2: int32,
# horizontal_upper_2: int32,
# horizontal_lower_3: int32,
# horizontal_upper_3: int32,
horizontal_lower_2: int32,
horizontal_upper_2: int32,
horizontal_lower_3: int32,
horizontal_upper_3: int32,
horizontal_lower_4: int32,
horizontal_upper_4: int32,
# kstart_dd3d: int32,
# nlev: int32,
# nflatlev: int32,
# nflat_gradp: int32,
kstart_dd3d: int32,
nlev: int32,
nflatlev: int32,
nflat_gradp: int32,
) -> tuple[
Field[[EdgeDim, KDim], float],
Field[[EdgeDim, KDim], float],
Expand Down Expand Up @@ -620,76 +620,80 @@ def fused_solve_nonhydro_stencil_15_to_28(
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],
# c_lin_e: Field[[EdgeDim, E2CDim], float],
# z_exner_ex_pr: Field[[CellDim, KDim], float],
# z_dexner_dz_c_1: Field[[CellDim, KDim], float],
# z_dexner_dz_c_2: Field[[CellDim, KDim], float],
# theta_v: Field[[CellDim, KDim], float],
# ikoffset: Field[[ECDim, KDim], int32],
# zdiff_gradp: Field[[ECDim, KDim], float],
# theta_v_ic: Field[[CellDim, KDim], float],
# inv_ddqz_z_full: Field[[CellDim, KDim], float],
# ipeidx_dsl: Field[[EdgeDim, KDim], bool],
# pg_exdist: Field[[EdgeDim, KDim], float],
# hmask_dd3d: Field[[EdgeDim], float],
# scalfac_dd3d: Field[[KDim], float],
# z_dwdz_dd: Field[[CellDim, KDim], float],
# inv_dual_edge_length: Field[[EdgeDim], float],
# ddt_vn_apc_ntl2: Field[[EdgeDim, KDim], float],
# vn_nnow: Field[[EdgeDim, KDim], float],
# ddt_vn_apc_ntl1: Field[[EdgeDim, KDim], float],
# ddt_vn_phy: Field[[EdgeDim, KDim], float],
ddxn_z_full: Field[[EdgeDim, KDim], float],
c_lin_e: Field[[EdgeDim, E2CDim], float],
z_exner_ex_pr: Field[[CellDim, KDim], float],
z_dexner_dz_c_1: Field[[CellDim, KDim], float],
z_dexner_dz_c_2: Field[[CellDim, KDim], float],
theta_v: Field[[CellDim, KDim], float],
ikoffset: Field[[ECDim, KDim], int32],
zdiff_gradp: Field[[ECDim, KDim], float],
theta_v_ic: Field[[CellDim, KDim], float],
inv_ddqz_z_full: Field[[CellDim, KDim], float],
ipeidx_dsl: Field[[EdgeDim, KDim], bool],
pg_exdist: Field[[EdgeDim, KDim], float],
hmask_dd3d: Field[[EdgeDim], float],
scalfac_dd3d: Field[[KDim], float],
z_dwdz_dd: Field[[CellDim, KDim], float],
inv_dual_edge_length: Field[[EdgeDim], float],
ddt_vn_apc_ntl2: Field[[EdgeDim, KDim], float],
vn_nnow: Field[[EdgeDim, KDim], float],
ddt_vn_apc_ntl1: Field[[EdgeDim, KDim], float],
ddt_vn_phy: Field[[EdgeDim, KDim], float],
z_graddiv_vn: Field[[EdgeDim, KDim], float],
# vn_incr: Field[[EdgeDim, KDim], float],
vn_incr: Field[[EdgeDim, KDim], float],
vn: Field[[EdgeDim, KDim], float],
z_theta_v_e: Field[[EdgeDim, KDim], float],
z_gradh_exner: Field[[EdgeDim, KDim], float],
# z_graddiv2_vn: Field[[EdgeDim, KDim], float],
# z_hydro_corr: Field[[EdgeDim, KDim], float],
z_graddiv2_vn: Field[[EdgeDim, KDim], float],
z_hydro_corr: Field[[EdgeDim, KDim], float],
z_rho_e: Field[[EdgeDim, KDim], float],
# geofac_grdiv: Field[[EdgeDim, E2C2EODim], float],
# scal_divdamp: Field[[KDim], float],
# bdy_divdamp: Field[[KDim], float],
# nudgecoeff_e: Field[[EdgeDim], float],
geofac_grdiv: Field[[EdgeDim, E2C2EODim], float],
scal_divdamp: Field[[KDim], float],
bdy_divdamp: Field[[KDim], float],
nudgecoeff_e: Field[[EdgeDim], float],
horz_idx: Field[[EdgeDim], int32],
# vert_idx: Field[[KDim], int32],
# grav_o_cpd: float,
vert_idx: Field[[KDim], int32],
grav_o_cpd: float,
p_dthalf: float,
idiv_method: int32,
# igradp_method: int32,
# wgt_nnow_vel: float,
# wgt_nnew_vel: float,
# dtime: float,
# cpd: float,
# iau_wgt_dyn: float,
# is_iau_active: bool,
# lhdiff_rcf: bool,
# divdamp_fac: float,
# divdamp_fac_o2: float,
# divdamp_order: int32,
# scal_divdamp_o2: float,
igradp_method: int32,
wgt_nnow_vel: float,
wgt_nnew_vel: float,
dtime: float,
cpd: float,
iau_wgt_dyn: float,
is_iau_active: bool,
lhdiff_rcf: bool,
divdamp_fac: float,
divdamp_fac_o2: float,
divdamp_order: int32,
scal_divdamp_o2: float,
limited_area: bool,
# itime_scheme: int32,
itime_scheme: int32,
istep: int32,
# horizontal_lower_0: int32,
# horizontal_upper_0: int32,
horizontal_lower_0: int32,
horizontal_upper_0: int32,
horizontal_lower_00: int32,
horizontal_upper_00: int32,
horizontal_lower_01: int32,
horizontal_upper_01: int32,
horizontal_lower_1: int32,
horizontal_upper_1: int32,
# horizontal_lower_2: int32,
# horizontal_upper_2: int32,
# horizontal_lower_3: int32,
# horizontal_upper_3: int32,
horizontal_lower_2: int32,
horizontal_upper_2: int32,
horizontal_lower_3: int32,
horizontal_upper_3: int32,
horizontal_lower_4: int32,
horizontal_upper_4: int32,
# kstart_dd3d: int32,
# nlev: int32,
# nflatlev: int32,
# nflat_gradp: int32,
kstart_dd3d: int32,
nlev: int32,
nflatlev: int32,
nflat_gradp: int32,
horizontal_lower: int32,
horizontal_upper: int32,
vertical_lower: int32,
vertical_upper: int32
):
_fused_solve_nonhydro_stencil_15_to_28(
geofac_grg_x,
Expand All @@ -706,75 +710,76 @@ def fused_solve_nonhydro_stencil_15_to_28(
theta_ref_me,
z_rth_pr_1,
z_rth_pr_2,
# ddxn_z_full,
# c_lin_e,
# z_exner_ex_pr,
# z_dexner_dz_c_1,
# z_dexner_dz_c_2,
# theta_v,
# ikoffset,
# zdiff_gradp,
# theta_v_ic,
# inv_ddqz_z_full,
# ipeidx_dsl,
# pg_exdist,
# hmask_dd3d,
# scalfac_dd3d,
# z_dwdz_dd,
# inv_dual_edge_length,
# ddt_vn_apc_ntl2,
# vn_nnow,
# ddt_vn_apc_ntl1,
# ddt_vn_phy,
ddxn_z_full,
c_lin_e,
z_exner_ex_pr,
z_dexner_dz_c_1,
z_dexner_dz_c_2,
theta_v,
ikoffset,
zdiff_gradp,
theta_v_ic,
inv_ddqz_z_full,
ipeidx_dsl,
pg_exdist,
hmask_dd3d,
scalfac_dd3d,
z_dwdz_dd,
inv_dual_edge_length,
ddt_vn_apc_ntl2,
vn_nnow,
ddt_vn_apc_ntl1,
ddt_vn_phy,
z_graddiv_vn,
# vn_incr,
vn_incr,
vn,
z_rho_e,
z_theta_v_e,
z_gradh_exner,
# z_graddiv2_vn,
# z_hydro_corr,
# geofac_grdiv,
# scal_divdamp,
# bdy_divdamp,
# nudgecoeff_e,
z_graddiv2_vn,
z_hydro_corr,
geofac_grdiv,
scal_divdamp,
bdy_divdamp,
nudgecoeff_e,
horz_idx,
# vert_idx,
# grav_o_cpd,
vert_idx,
grav_o_cpd,
p_dthalf,
idiv_method,
# igradp_method,
# wgt_nnow_vel,
# wgt_nnew_vel,
# dtime,
# cpd,
# iau_wgt_dyn,
# is_iau_active,
# lhdiff_rcf,
# divdamp_fac,
# divdamp_fac_o2,
# divdamp_order,
# scal_divdamp_o2,
igradp_method,
wgt_nnow_vel,
wgt_nnew_vel,
dtime,
cpd,
iau_wgt_dyn,
is_iau_active,
lhdiff_rcf,
divdamp_fac,
divdamp_fac_o2,
divdamp_order,
scal_divdamp_o2,
limited_area,
# itime_scheme,
itime_scheme,
istep,
# horizontal_lower_0,
# horizontal_upper_0,
horizontal_lower_0,
horizontal_upper_0,
horizontal_lower_00,
horizontal_upper_00,
horizontal_lower_01,
horizontal_upper_01,
horizontal_lower_1,
horizontal_upper_1,
# horizontal_lower_2,
# horizontal_upper_2,
# horizontal_lower_3,
# horizontal_upper_3,
horizontal_lower_2,
horizontal_upper_2,
horizontal_lower_3,
horizontal_upper_3,
horizontal_lower_4,
horizontal_upper_4,
# kstart_dd3d,
# nlev,
# nflatlev,
# nflat_gradp,
kstart_dd3d,
nlev,
nflatlev,
nflat_gradp,
out=(z_rho_e, z_theta_v_e, z_gradh_exner, vn, z_graddiv_vn),
)

0 comments on commit 68b67c5

Please sign in to comment.