Update Field.eval and VectorField eval to match new grid search API

Author: VeckoTheGecko

parcels/application_kernels/interpolation.py

@@ -1,9 +1,16 @@
 """Collection of pre-built interpolation kernels."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from parcels.field import Field
 
+if TYPE_CHECKING:
+    from parcels.uxgrid import _UXGRID_AXES
+
 __all__ = [
     "UXPiecewiseConstantFace",
     "UXPiecewiseLinearNode",
@@ -13,8 +20,7 @@
 def UXPiecewiseConstantFace(
     field: Field,
     ti: int,
-    ei: int,
-    bcoords: np.ndarray,
+    position: dict[_UXGRID_AXES, tuple[int, float | np.ndarray]],
     tau: np.float32 | np.float64,
     t: np.float32 | np.float64,
     z: np.float32 | np.float64,
@@ -26,15 +32,13 @@ def UXPiecewiseConstantFace(
     This interpolation method is appropriate for fields that are
     face registered, such as u,v in FESOM.
     """
-    zi, fi = field.grid.unravel_index(ei)
-    return field.data.values[ti, zi, fi]
+    return field.data.values[ti, position["Z"][0], position["CELL"][0]]
 
 
 def UXPiecewiseLinearNode(
     field: Field,
     ti: int,
-    ei: int,
-    bcoords: np.ndarray,
+    position: dict[_UXGRID_AXES, tuple[int, float | np.ndarray]],
     tau: np.float32 | np.float64,
     t: np.float32 | np.float64,
     z: np.float32 | np.float64,
@@ -47,7 +51,8 @@ def UXPiecewiseLinearNode(
     velocity W in FESOM2. Effectively, it applies barycentric interpolation in the lateral direction
     and piecewise linear interpolation in the vertical direction.
     """
-    k, fi = field.grid.unravel_index(ei)
+    k, fi = position["Z"][0], position["CELL"][0]
+    bcoords = position["CELL"][1]
     node_ids = field.grid.uxgrid.face_node_connectivity[fi, :]
     # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels.
     # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1.

parcels/field.py

@@ -105,8 +105,7 @@ class Field:
     def _interp_template(
         self,
         ti: int,
-        ei: int,
-        bcoords: np.ndarray,
+        position: dict[str, tuple[int, float | np.ndarray]],
         tau: np.float32 | np.float64,
         t: np.float32 | np.float64,
         z: np.float32 | np.float64,
@@ -316,8 +315,8 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
 
         try:
             tau, ti = _search_time_index(self, time)
-            bcoords, _ei = self.grid.search(z, y, x, ei=_ei)
-            value = self._interp_method(self, ti, _ei, bcoords, tau, time, z, y, x)
+            position = self.grid.search(z, y, x, ei=_ei)
+            value = self._interp_method(self, ti, position, tau, time, z, y, x)
 
             if np.isnan(value):
                 # Detect Out-of-bounds sampling and raise exception
@@ -445,14 +444,14 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
 
         try:
             tau, ti = _search_time_index(self.U, time)
-            bcoords, _ei = self.grid.search(z, y, x, ei=_ei)
+            position = self.grid.search(z, y, x, ei=_ei)
             if self._vector_interp_method is None:
-                u = self.U._interp_method(self.U, ti, _ei, bcoords, tau, time, z, y, x)
-                v = self.V._interp_method(self.V, ti, _ei, bcoords, tau, time, z, y, x)
+                u = self.U._interp_method(self.U, ti, position, tau, time, z, y, x)
+                v = self.V._interp_method(self.V, ti, position, tau, time, z, y, x)
                 if "3D" in self.vector_type:
-                    w = self.W._interp_method(self.W, ti, _ei, bcoords, tau, time, z, y, x)
+                    w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x)
             else:
-                (u, v, w) = self._vector_interp_method(self, ti, _ei, bcoords, time, z, y, x)
+                (u, v, w) = self._vector_interp_method(self, ti, position, time, z, y, x)
 
             # print(u,v)
             if applyConversion: