Use the field.igrid in the selection of which ei in unravel_index

Author: fluidnumerics-joe

parcels/_index_search.py

@@ -225,7 +225,7 @@ def _search_indices_rectilinear(
 
 def _search_indices_curvilinear(field: Field, time, z, y, x, ti=-1, particle=None, search2D=False):
     if particle:
-        zi, yi, xi = field.unravel_index(particle.ei[field.igrid])
+        zi, yi, xi = field.unravel_index(particle.ei)
     else:
         xi = int(field.grid.xdim / 2) - 1
         yi = int(field.grid.ydim / 2) - 1

parcels/field.py

@@ -1349,10 +1349,12 @@ def unravel_index(self, ei):
         xi : int
             The z index.
         """
-        zi = ei // (self.grid.xdim * self.grid.ydim)
-        ei = ei % (self.grid.xdim * self.grid.ydim)
-        yi = ei // self.grid.xdim
-        xi = ei % self.grid.xdim
+
+        _ei = ei[self.igrid]
+        zi = _ei // (self.grid.xdim * self.grid.ydim)
+        _ei = _ei % (self.grid.xdim * self.grid.ydim)
+        yi = _ei // self.grid.xdim
+        xi = _ei % self.grid.xdim
         return zi, yi, xi
 
 

tests/test_grids.py

@@ -92,11 +92,11 @@ def sampleTemp(particle, fieldset, time):  # pragma: no cover
     # check if particle xi and yi are different for the two grids
     # xi check from unraveled index
     assert np.all(
-        [fieldset.U.unravel_index(pset[i].ei)[2][0] != fieldset.U.unravel_index(pset[i].ei)[2][1] for i in range(3)]
+        [fieldset.U.unravel_index(pset[i].ei)[2] != fieldset.V.unravel_index(pset[i].ei)[2] for i in range(3)]
     )
     # yi check from unraveled index
     assert np.all(
-        [fieldset.U.unravel_index(pset[i].ei)[1][0] != fieldset.U.unravel_index(pset[i].ei)[1][1] for i in range(3)]
+        [fieldset.U.unravel_index(pset[i].ei)[1] != fieldset.V.unravel_index(pset[i].ei)[1] for i in range(3)]
     )
     # advect without updating temperature to test particle deletion
     pset.remove_indices(np.array([1]))

tests/test_particlefile.py

@@ -273,9 +273,9 @@ def Get_XiYi(particle, fieldset, time):  # pragma: no cover
         and that the first outputted value is zero.
         Be careful when using multiple grids, as the index may be different for the grids.
         """
-        particle.pxi0 = fieldset.U.unravel_index(particle.ei[0])[2]
-        particle.pxi1 = fieldset.U.unravel_index(particle.ei[1])[2]
-        particle.pyi = fieldset.U.unravel_index(particle.ei[0])[1]
+        particle.pxi0 = fieldset.U.unravel_index(particle.ei)[2]
+        particle.pxi1 = fieldset.P.unravel_index(particle.ei)[2]
+        particle.pyi = fieldset.U.unravel_index(particle.ei)[1]
 
     def SampleP(particle, fieldset, time):  # pragma: no cover
         if time > 5 * 3600: