Adapt to new sparse field handling in GT4Py (#168)

For field view the iterator IR changes are not visible, but the ones implemented in itir directly change.
C2SM · Mar 23, 2023 · 2a6a3c2 · 2a6a3c2
1 parent 981d50b
commit 2a6a3c2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py
@@ -16,6 +16,7 @@
 from gt4py.next.iterator.builtins import (
     deref,
     if_,
+    list_get,
     named_range,
     shift,
     unstructured_domain,
@@ -29,9 +30,9 @@
 
 @fundef
 def step(i, geofac_n2s_nbh, vcoef, theta_v, zd_vertoffset):
-    d_vcoef = deref(shift(i)(vcoef))
-    s_theta_v = shift(C2E2C, i, Koff, deref(shift(i)(zd_vertoffset)))(theta_v)
-    return deref(shift(i)(geofac_n2s_nbh)) * (
+    d_vcoef = list_get(i, deref(vcoef))
+    s_theta_v = shift(C2E2C, i, Koff, list_get(i, deref(zd_vertoffset)))(theta_v)
+    return list_get(i, deref(geofac_n2s_nbh)) * (
         d_vcoef * deref(s_theta_v) + (1.0 - d_vcoef) * deref(shift(Koff, 1)(s_theta_v))
     )
 

diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_20.py
@@ -15,6 +15,7 @@
 from gt4py.next.ffront import program_ast as past
 from gt4py.next.iterator.builtins import (
     deref,
+    list_get,
     named_range,
     shift,
     unstructured_domain,
@@ -35,11 +36,11 @@ def step(
     z_dexner_dz_c_1,
     z_dexner_dz_c_2,
 ):
-    d_ikoffset = deref(shift(i)(ikoffset))
+    d_ikoffset = list_get(i, deref(ikoffset))
     d_z_exner_exp_pr = deref(shift(Koff, d_ikoffset, E2C, i)(z_exner_ex_pr))
     d_z_dexner_dz_c_1 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_1))
     d_z_dexner_dz_c_2 = deref(shift(Koff, d_ikoffset, E2C, i)(z_dexner_dz_c_2))
-    d_zdiff_gradp = deref(shift(i)(zdiff_gradp))
+    d_zdiff_gradp = list_get(i, deref(zdiff_gradp))
 
     return d_z_exner_exp_pr + d_zdiff_gradp * (
         d_z_dexner_dz_c_1 + d_zdiff_gradp * d_z_dexner_dz_c_2

diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_solve_nonhydro_stencil_21.py
@@ -15,6 +15,7 @@
 from gt4py.next.ffront import program_ast as past
 from gt4py.next.iterator.builtins import (
     deref,
+    list_get,
     named_range,
     power,
     shift,
@@ -29,14 +30,14 @@
 
 @fundef
 def step(i, theta_v, ikoffset, zdiff_gradp, theta_v_ic, inv_ddqz_z_full):
-    d_ikoffset = deref(shift(i)(ikoffset))
+    d_ikoffset = list_get(i, deref(ikoffset))
 
     d_theta_v = deref(shift(Koff, d_ikoffset, E2C, i)(theta_v))
     s_theta_v_ic = shift(Koff, d_ikoffset, E2C, i)(theta_v_ic)
     d_theta_v_ic = deref(s_theta_v_ic)
     d_theta_v_ic_p1 = deref(shift(Koff, 1)(s_theta_v_ic))
     d_inv_ddqz_z_full = deref(shift(Koff, d_ikoffset, E2C, i)(inv_ddqz_z_full))
-    d_zdiff_gradp = deref(shift(i)(zdiff_gradp))
+    d_zdiff_gradp = list_get(i, deref(zdiff_gradp))
 
     return (
         d_theta_v + d_zdiff_gradp * (d_theta_v_ic - d_theta_v_ic_p1) * d_inv_ddqz_z_full