Implementing use of particle.ei in curvilinear search

Author: erikvansebille

parcels/_index_search.py

@@ -75,20 +75,18 @@ def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray
 
 
 def _search_indices_curvilinear_2d(
-    grid: XGrid, y: np.ndarray, x: np.ndarray, yi_guess: np.ndarray | None = None, xi_guess: np.ndarray | None = None
+    grid: XGrid, y: np.ndarray, x: np.ndarray, yi: np.ndarray | None = None, xi: np.ndarray | None = None
 ):
-    yi_guess = np.array(yi_guess)
-    xi_guess = np.array(xi_guess)
-    xi = np.full(len(x), GRID_SEARCH_ERROR, dtype=np.int32)
-    yi = np.full(len(y), GRID_SEARCH_ERROR, dtype=np.int32)
-    if np.any(xi_guess):
+    if np.any(xi):
         # If an initial guess is provided, we first perform a point in cell check for all guessed indices
-        is_in_cell, coords = curvilinear_point_in_cell(grid, y, x, yi_guess, xi_guess)
+        is_in_cell, coords = curvilinear_point_in_cell(grid, y, x, yi, xi)
         y_check = y[is_in_cell == 0]
         x_check = x[is_in_cell == 0]
         zero_indices = np.where(is_in_cell == 0)[0]
     else:
         # Otherwise, we need to check all points
+        yi = np.full(len(y), GRID_SEARCH_ERROR, dtype=np.int32)
+        xi = np.full(len(x), GRID_SEARCH_ERROR, dtype=np.int32)
         y_check = y
         x_check = x
         coords = -1.0 * np.ones((len(y), 2), dtype=np.float32)

parcels/field.py

@@ -212,13 +212,14 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         conversion to the result. Note that we defer to
         scipy.interpolate to perform spatial interpolation.
         """
-        # if particle is None:
-        _ei = None
-        # else:
-        #    _ei = particle.ei[self.igrid]
+        if particle is None:
+            _ei = None
+        else:
+            _ei = particle.ei[:, self.igrid]
 
         tau, ti = _search_time_index(self, time)
         position = self.grid.search(z, y, x, ei=_ei)
+        _update_particles_ei(particle, position, self)
         _update_particle_states_position(particle, position)
 
         value = self._interp_method(self, ti, position, tau, time, z, y, x)
@@ -251,6 +252,7 @@ def __init__(
         self.V = V
         self.W = W
         self.grid = U.grid
+        self.igrid = U.igrid
 
         if W is None:
             _assert_same_time_interval((U, V))
@@ -294,13 +296,14 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
         conversion to the result. Note that we defer to
         scipy.interpolate to perform spatial interpolation.
         """
-        # if particle is None:
-        _ei = None
-        # else:
-        #    _ei = particle.ei[self.igrid]
+        if particle is None:
+            _ei = None
+        else:
+            _ei = particle.ei[:, self.igrid]
 
         tau, ti = _search_time_index(self.U, time)
         position = self.grid.search(z, y, x, ei=_ei)
+        _update_particles_ei(particle, position, self)
         _update_particle_states_position(particle, position)
 
         if self._vector_interp_method is None:
@@ -339,6 +342,18 @@ def __getitem__(self, key):
             return _deal_with_errors(error, key, vector_type=self.vector_type)
 
 
+def _update_particles_ei(particles, position, field):
+    """Update the element index (ei) of the particles"""
+    if particles is not None:
+        particles.ei[:, field.igrid] = field.grid.ravel_index(
+            {
+                "X": position["X"][0],
+                "Y": position["Y"][0],
+                "Z": position["Z"][0],
+            }
+        )
+
+
 def _update_particle_states_position(particle, position):
     """Update the particle states based on the position dictionary."""
     if particle:  # TODO also support uxgrid search