Implementing separate particle_positions and grid_positions dictionary as suggested in #2287

Author: reint-fischer

src/parcels/_core/field.py

@@ -216,13 +216,14 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
         else:
             _ei = particles.ei[:, self.igrid]
 
-        position = {"time": time, "z": z, "lat": y, "lon": x}
-        position.update(_search_time_index(self, time))
-        position.update(self.grid.search(z, y, x, ei=_ei))
-        _update_particles_ei(particles, position, self)
-        _update_particle_states_position(particles, position)
+        particle_positions = {"time": time, "z": z, "lat": y, "lon": x}
+        grid_positions = {}
+        grid_positions.update(_search_time_index(self, time))
+        grid_positions.update(self.grid.search(z, y, x, ei=_ei))
+        _update_particles_ei(particles, grid_positions, self)
+        _update_particle_states_position(particles, grid_positions)
 
-        value = self._interp_method(position, self)
+        value = self._interp_method(particle_positions, grid_positions, self)
 
         _update_particle_states_interp_value(particles, value)
 
@@ -301,21 +302,22 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
         else:
             _ei = particles.ei[:, self.igrid]
 
-        position = {"time": time, "z": z, "lat": y, "lon": x}
-        position.update(_search_time_index(self.U, time))
-        position.update(self.grid.search(z, y, x, ei=_ei))
-        _update_particles_ei(particles, position, self)
-        _update_particle_states_position(particles, position)
+        particle_positions = {"time": time, "z": z, "lat": y, "lon": x}
+        grid_positions = {}
+        grid_positions.update(_search_time_index(self.U, time))
+        grid_positions.update(self.grid.search(z, y, x, ei=_ei))
+        _update_particles_ei(particles, grid_positions, self)
+        _update_particle_states_position(particles, grid_positions)
 
         if self._vector_interp_method is None:
-            u = self.U._interp_method(position, self.U)
-            v = self.V._interp_method(position, self.V)
+            u = self.U._interp_method(particle_positions, grid_positions, self.U)
+            v = self.V._interp_method(particle_positions, grid_positions, self.V)
             if "3D" in self.vector_type:
-                w = self.W._interp_method(position, self.W)
+                w = self.W._interp_method(particle_positions, grid_positions, self.W)
             else:
                 w = 0.0
         else:
-            (u, v, w) = self._vector_interp_method(position, self)
+            (u, v, w) = self._vector_interp_method(particle_positions, grid_positions, self)
 
         if applyConversion:
             u = self.U.units.to_target(u, z, y, x)
@@ -341,45 +343,54 @@ def __getitem__(self, key):
             return _deal_with_errors(error, key, vector_type=self.vector_type)
 
 
-def _update_particles_ei(particles, position, field):
+def _update_particles_ei(particles, grid_positions, field):
     """Update the element index (ei) of the particles"""
     if particles is not None:
         if isinstance(field.grid, XGrid):
             particles.ei[:, field.igrid] = field.grid.ravel_index(
                 {
-                    "X": position["xi"],
-                    "Y": position["yi"],
-                    "Z": position["zi"],
+                    "X": grid_positions["X"]["index"],
+                    "Y": grid_positions["Y"]["index"],
+                    "Z": grid_positions["Z"]["index"],
                 }
             )
         elif isinstance(field.grid, UxGrid):
             particles.ei[:, field.igrid] = field.grid.ravel_index(
                 {
-                    "Z": position["Z"][0],
-                    "FACE": position["FACE"][0],
+                    "Z": grid_positions["Z"]["index"],
+                    "FACE": grid_positions["FACE"]["index"],
                 }
             )
 
 
-def _update_particle_states_position(particles, position):
+def _update_particle_states_position(particles, grid_positions):
     """Update the particle states based on the position dictionary."""
     if particles:  # TODO also support uxgrid search
-        for dim in ["xi", "yi"]:
-            if dim in position:
+        for dim in ["X", "Y"]:
+            if dim in grid_positions:
                 particles.state = np.maximum(
-                    np.where(position[dim] == -1, StatusCode.ErrorOutOfBounds, particles.state), particles.state
+                    np.where(grid_positions[dim]["index"] == -1, StatusCode.ErrorOutOfBounds, particles.state),
+                    particles.state,
                 )
                 particles.state = np.maximum(
-                    np.where(position[dim] == GRID_SEARCH_ERROR, StatusCode.ErrorGridSearching, particles.state),
+                    np.where(
+                        grid_positions[dim]["index"] == GRID_SEARCH_ERROR,
+                        StatusCode.ErrorGridSearching,
+                        particles.state,
+                    ),
                     particles.state,
                 )
-        if "zi" in position:
+        if "Z" in grid_positions:
             particles.state = np.maximum(
-                np.where(position["zi"][0] == RIGHT_OUT_OF_BOUNDS, StatusCode.ErrorOutOfBounds, particles.state),
+                np.where(
+                    grid_positions["Z"]["index"] == RIGHT_OUT_OF_BOUNDS, StatusCode.ErrorOutOfBounds, particles.state
+                ),
                 particles.state,
             )
             particles.state = np.maximum(
-                np.where(position["zi"][0] == LEFT_OUT_OF_BOUNDS, StatusCode.ErrorThroughSurface, particles.state),
+                np.where(
+                    grid_positions["Z"]["index"] == LEFT_OUT_OF_BOUNDS, StatusCode.ErrorThroughSurface, particles.state
+                ),
                 particles.state,
             )
 

src/parcels/_core/index_search.py

@@ -75,14 +75,19 @@ def _search_time_index(field: Field, time: datetime):
     if the sampled value is outside the time value range.
     """
     if field.time_interval is None:
-        return {"ti": np.zeros(shape=time.shape, dtype=np.int32), "tau": np.zeros(shape=time.shape, dtype=np.float32)}
+        return {
+            "T": {
+                "index": np.zeros(shape=time.shape, dtype=np.int32),
+                "bcoord": np.zeros(shape=time.shape, dtype=np.float32),
+            }
+        }
 
     if not field.time_interval.is_all_time_in_interval(time):
         _raise_time_extrapolation_error(time, field=None)
 
     ti = np.searchsorted(field.data.time.data, time, side="right") - 1
     tau = (time - field.data.time.data[ti]) / (field.data.time.data[ti + 1] - field.data.time.data[ti])
-    return {"ti": np.atleast_1d(ti), "tau": np.atleast_1d(tau)}
+    return {"T": {"index": np.atleast_1d(ti), "bcoord": np.atleast_1d(tau)}}
 
 
 def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: np.ndarray):

src/parcels/_core/particleset.py

@@ -295,12 +295,12 @@ def _neighbors_by_coor(self, coor):
     def populate_indices(self):
         """Pre-populate guesses of particle ei (element id) indices"""
         for i, grid in enumerate(self.fieldset.gridset):
-            position = grid.search(self.z, self.lat, self.lon)
+            grid_positions = grid.search(self.z, self.lat, self.lon)
             self._data["ei"][:, i] = grid.ravel_index(
                 {
-                    "X": position["xi"],
-                    "Y": position["yi"],
-                    "Z": position["zi"],
+                    "X": grid_positions["X"]["index"],
+                    "Y": grid_positions["Y"]["index"],
+                    "Z": grid_positions["Z"]["index"],
                 }
             )
 

src/parcels/_core/uxgrid.py

@@ -118,4 +118,4 @@ def search(self, z, y, x, ei=None, tol=1e-6):
             coords[zero_indices, :] = coords_q
             fi[zero_indices] = face_ids_q
 
-        return {"Z": (zi, zeta), "FACE": (fi, coords)}
+        return {"Z": {"index":zi, "bcoord":zeta}, "FACE": {"index":fi, "bcoord":coords}}

src/parcels/_core/xgrid.py

@@ -289,7 +289,11 @@ def search(self, z, y, x, ei=None):
         if ds.lon.ndim == 1:
             yi, eta = _search_1d_array(ds.lat.values, y)
             xi, xsi = _search_1d_array(ds.lon.values, x)
-            return {"zi": zi, "zeta": zeta, "yi": yi, "eta": eta, "xi": xi, "xsi": xsi}
+            return {
+                "Z": {"index": zi, "bcoord": zeta},
+                "Y": {"index": yi, "bcoord": eta},
+                "X": {"index": xi, "bcoord": xsi},
+            }
 
         yi, xi = None, None
         if ei is not None:
@@ -300,7 +304,11 @@ def search(self, z, y, x, ei=None):
         if ds.lon.ndim == 2:
             yi, eta, xi, xsi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
 
-            return {"zi": zi, "zeta": zeta, "yi": yi, "eta": eta, "xi": xi, "xsi": xsi}
+            return {
+                "Z": {"index": zi, "bcoord": zeta},
+                "Y": {"index": yi, "bcoord": eta},
+                "X": {"index": xi, "bcoord": xsi},
+            }
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")