From df795b0b857be6aae40968d970740a87e2a16f59 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 12 Sep 2024 13:32:25 +0200 Subject: [PATCH] Fix usage of as_offset (#542) Fix usage of as_offset in following stencils: * _compute_hydrostatic_correction_term * _compute_horizontal_gradient_of_exner_pressure_for_multiple_levels * _truly_horizontal_diffusion_nabla_of_theta_over_steep_points --------- Co-authored-by: Edoardo Paone Co-authored-by: Magdalena Luz --- ...fusion_nabla_of_theta_over_steep_points.py | 27 +++++---------- ...t_of_exner_pressure_for_multiple_levels.py | 20 +++++------ .../compute_hydrostatic_correction_term.py | 34 +++++++++++-------- 3 files changed, 39 insertions(+), 42 deletions(-) diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 039a45bb3a..6a276979e3 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -35,31 +35,22 @@ def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( ) -> fa.CellKField[vpfloat]: z_temp_wp = astype(z_temp, wpfloat) - theta_v_0 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[0]))) - theta_v_1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]))) - theta_v_2 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]))) + theta_v_0 = theta_v(C2E2C[0])(as_offset(Koff, zd_vertoffset(C2CEC[0]))) + theta_v_1 = theta_v(C2E2C[1])(as_offset(Koff, zd_vertoffset(C2CEC[1]))) + theta_v_2 = theta_v(C2E2C[2])(as_offset(Koff, zd_vertoffset(C2CEC[2]))) - theta_v_0_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[0]) + 1)) - theta_v_1_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[1]) + 1)) - theta_v_2_m1 = theta_v(as_offset(Koff, zd_vertoffset(C2CEC[2]) + 1)) + theta_v_0_m1 = theta_v(C2E2C[0])(as_offset(Koff, zd_vertoffset(C2CEC[0]) + 1)) + theta_v_1_m1 = theta_v(C2E2C[1])(as_offset(Koff, zd_vertoffset(C2CEC[1]) + 1)) + theta_v_2_m1 = theta_v(C2E2C[2])(as_offset(Koff, zd_vertoffset(C2CEC[2]) + 1)) sum_tmp = ( theta_v * geofac_n2s_c + geofac_n2s_nbh(C2CEC[0]) - * ( - vcoef(C2CEC[0]) * theta_v_0(C2E2C[0]) - + (wpfloat("1.0") - vcoef(C2CEC[0])) * theta_v_0_m1(C2E2C[0]) - ) + * (vcoef(C2CEC[0]) * theta_v_0 + (wpfloat("1.0") - vcoef(C2CEC[0])) * theta_v_0_m1) + geofac_n2s_nbh(C2CEC[1]) - * ( - vcoef(C2CEC[1]) * theta_v_1(C2E2C[1]) - + (wpfloat("1.0") - vcoef(C2CEC[1])) * theta_v_1_m1(C2E2C[1]) - ) + * (vcoef(C2CEC[1]) * theta_v_1 + (wpfloat("1.0") - vcoef(C2CEC[1])) * theta_v_1_m1) + geofac_n2s_nbh(C2CEC[2]) - * ( - vcoef(C2CEC[2]) * theta_v_2(C2E2C[2]) - + (wpfloat("1.0") - vcoef(C2CEC[2])) * theta_v_2_m1(C2E2C[2]) - ) + * (vcoef(C2CEC[2]) * theta_v_2 + (wpfloat("1.0") - vcoef(C2CEC[2])) * theta_v_2_m1) ) z_temp_wp = where( diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py index a13c509ae7..03faca6cc7 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py @@ -32,26 +32,26 @@ def _compute_horizontal_gradient_of_exner_pressure_for_multiple_levels( z_dexner_dz_c_2: fa.CellKField[vpfloat], ) -> fa.EdgeKField[vpfloat]: """Formerly known as _mo_solve_nonhydro_stencil_20.""" - z_exner_ex_pr_0 = z_exner_ex_pr(as_offset(Koff, ikoffset(E2EC[0]))) - z_exner_ex_pr_1 = z_exner_ex_pr(as_offset(Koff, ikoffset(E2EC[1]))) + z_exner_ex_pr_0 = z_exner_ex_pr(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + z_exner_ex_pr_1 = z_exner_ex_pr(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) - z_dexner_dz_c1_0 = z_dexner_dz_c_1(as_offset(Koff, ikoffset(E2EC[0]))) - z_dexner_dz_c1_1 = z_dexner_dz_c_1(as_offset(Koff, ikoffset(E2EC[1]))) + z_dexner_dz_c1_0 = z_dexner_dz_c_1(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + z_dexner_dz_c1_1 = z_dexner_dz_c_1(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) - z_dexner_dz_c2_0 = z_dexner_dz_c_2(as_offset(Koff, ikoffset(E2EC[0]))) - z_dexner_dz_c2_1 = z_dexner_dz_c_2(as_offset(Koff, ikoffset(E2EC[1]))) + z_dexner_dz_c2_0 = z_dexner_dz_c_2(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + z_dexner_dz_c2_1 = z_dexner_dz_c_2(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) z_gradh_exner_wp = inv_dual_edge_length * ( astype( ( - z_exner_ex_pr_1(E2C[1]) + z_exner_ex_pr_1 + zdiff_gradp(E2EC[1]) - * (z_dexner_dz_c1_1(E2C[1]) + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1(E2C[1])) + * (z_dexner_dz_c1_1 + zdiff_gradp(E2EC[1]) * z_dexner_dz_c2_1) ) - ( - z_exner_ex_pr_0(E2C[0]) + z_exner_ex_pr_0 + zdiff_gradp(E2EC[0]) - * (z_dexner_dz_c1_0(E2C[0]) + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0(E2C[0])) + * (z_dexner_dz_c1_0 + zdiff_gradp(E2EC[0]) * z_dexner_dz_c2_0) ), wpfloat, ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py index 478b224f48..a2b1bd4a50 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py @@ -35,24 +35,30 @@ def _compute_hydrostatic_correction_term( """Formerly known as _mo_solve_nonhydro_stencil_21.""" zdiff_gradp_wp = astype(zdiff_gradp, wpfloat) - theta_v_0 = theta_v(as_offset(Koff, ikoffset(E2EC[0]))) - theta_v_1 = theta_v(as_offset(Koff, ikoffset(E2EC[1]))) + theta_v_0 = theta_v(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + theta_v_1 = theta_v(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) - theta_v_ic_0 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[0]))) - theta_v_ic_1 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[1]))) + theta_v_ic_0 = theta_v_ic(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))) + theta_v_ic_1 = theta_v_ic(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))) - theta_v_ic_p1_0 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[0]) + 1)) - theta_v_ic_p1_1 = theta_v_ic(as_offset(Koff, ikoffset(E2EC[1]) + 1)) + theta_v_ic_p1_0 = theta_v_ic(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]) + 1)) + theta_v_ic_p1_1 = theta_v_ic(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]) + 1)) - inv_ddqz_z_full_0_wp = astype(inv_ddqz_z_full(as_offset(Koff, ikoffset(E2EC[0]))), wpfloat) - inv_ddqz_z_full_1_wp = astype(inv_ddqz_z_full(as_offset(Koff, ikoffset(E2EC[1]))), wpfloat) + inv_ddqz_z_full_0_wp = astype( + inv_ddqz_z_full(E2C[0])(as_offset(Koff, ikoffset(E2EC[0]))), wpfloat + ) + inv_ddqz_z_full_1_wp = astype( + inv_ddqz_z_full(E2C[1])(as_offset(Koff, ikoffset(E2EC[1]))), wpfloat + ) - z_theta_0 = theta_v_0(E2C[0]) + zdiff_gradp_wp(E2EC[0]) * ( - theta_v_ic_0(E2C[0]) - theta_v_ic_p1_0(E2C[0]) - ) * inv_ddqz_z_full_0_wp(E2C[0]) - z_theta_1 = theta_v_1(E2C[1]) + zdiff_gradp_wp(E2EC[1]) * ( - theta_v_ic_1(E2C[1]) - theta_v_ic_p1_1(E2C[1]) - ) * inv_ddqz_z_full_1_wp(E2C[1]) + z_theta_0 = ( + theta_v_0 + + zdiff_gradp_wp(E2EC[0]) * (theta_v_ic_0 - theta_v_ic_p1_0) * inv_ddqz_z_full_0_wp + ) + z_theta_1 = ( + theta_v_1 + + zdiff_gradp_wp(E2EC[1]) * (theta_v_ic_1 - theta_v_ic_p1_1) * inv_ddqz_z_full_1_wp + ) z_hydro_corr_wp = ( grav_o_cpd * inv_dual_edge_length