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 <[email protected]> Co-authored-by: Magdalena Luz <[email protected]>
C2SM · Sep 12, 2024 · df795b0 · df795b0
1 parent ee8490e
commit df795b0
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 42 deletions.
diff --git a/...osphere/diffusion/stencils/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/...osphere/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/...el/atmosphere/dycore/compute_horizontal_gradient_of_exner_pressure_for_multiple_levels.py b/...el/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/...osphere/dycore/src/icon4py/model/atmosphere/dycore/compute_hydrostatic_correction_term.py b/...osphere/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