Merge pull request #2051 from OceanParcels/cleanup-vectorfield-and-uxgrid

Cleanup vectorfield and uxgrid

Author: fluidnumericsJoe

parcels/field.py

@@ -309,10 +309,10 @@ 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]
 
         try:
             tau, ti = _search_time_index(self, time)
@@ -386,6 +386,7 @@ def __init__(
         self.U = U
         self.V = V
         self.W = W
+        self.grid = U.grid
 
         if W is None:
             assert_same_time_interval((U, V))
@@ -430,41 +431,49 @@ def vector_interp_method(self, method: Callable):
     #         and np.allclose(grid1.depth, grid2.depth)
     #         and np.allclose(grid1.time, grid2.time)
     #     )
-    def _interpolate(self, time, z, y, x, ei):
-        bcoords, _ei, ti = self._search_indices(time, z, y, x, ei=ei)
+    def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True):
+        """Interpolate field values in space and time.
 
-        if self._vector_interp_method is None:
-            u = self.U.eval(time, z, y, x, _ei, applyConversion=False)
-            v = self.V.eval(time, z, y, x, _ei, applyConversion=False)
-            if "3D" in self.vector_type:
-                w = self.W.eval(time, z, y, x, _ei, applyConversion=False)
-                return (u, v, w)
-            else:
-                return (u, v, 0)
-        else:
-            (u, v, w) = self._vector_interp_method(ti, _ei, bcoords, time, z, y, x)
-            return (u, v, w)
+        We interpolate linearly in time and apply implicit unit
+        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]
 
-    def eval(self, time, z, y, x, ei=None, applyConversion=True):
-        if ei is None:
-            _ei = 0
-        else:
-            _ei = ei[self.igrid]
+        try:
+            tau, ti = _search_time_index(self.U, time)
+            bcoords, _ei = 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)
+                if "3D" in self.vector_type:
+                    w = self.W._interp_method(self.W, ti, _ei, bcoords, 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._interpolate(time, z, y, x, _ei)
+            # print(u,v)
+            if applyConversion:
+                u = self.U.units.to_target(u, z, y, x)
+                v = self.V.units.to_target(v, z, y, x)
+                if "3D" in self.vector_type:
+                    w = self.W.units.to_target(w, z, y, x) if self.W else 0.0
 
-        if applyConversion:
-            u = self.U.units.to_target(u, z, y, x)
-            v = self.V.units.to_target(v, z, y, x)
             if "3D" in self.vector_type:
-                w = self.W.units.to_target(w, z, y, x)
+                return (u, v, w)
+            else:
+                return (u, v)
 
-        return (u, v, w)
+        except (FieldSamplingError, FieldOutOfBoundError, FieldOutOfBoundSurfaceError) as e:
+            e.add_note(f"Error interpolating field '{self.name}'.")
+            raise e
 
     def __getitem__(self, key):
         try:
             if _isParticle(key):
-                return self.eval(key.time, key.depth, key.lat, key.lon, key.ei)
+                return self.eval(key.time, key.depth, key.lat, key.lon, key)
             else:
                 return self.eval(*key)
         except tuple(AllParcelsErrorCodes.keys()) as error:

parcels/uxgrid.py

@@ -35,6 +35,20 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray) -> UxGrid:
             raise ValueError("z must be a 1D array of vertical coordinates")
         self.z = z
 
+    @property
+    def depth(self):
+        """
+        Note
+        ----
+        Included for compatibility with v3 codebase. May be removed in future.
+        TODO v4: Evaluate
+        """
+        try:
+            _ = self.z.values
+        except KeyError:
+            return np.zeros(1)
+        return self.z.values
+
     def search(
         self, z: float, y: float, x: float, ei: int | None = None, search2D: bool = False
     ) -> tuple[np.ndarray, int]: