Update UxGrid ravelling to match BaseGrid API

Author: VeckoTheGecko

parcels/uxgrid.py

@@ -7,6 +7,7 @@
 from uxarray.grid.neighbors import _barycentric_coordinates
 
 from parcels.field import FieldOutOfBoundError  # Adjust import as necessary
+from parcels.xgrid import _search_1d_array
 
 from .basegrid import BaseGrid
 
@@ -62,18 +63,7 @@ def try_face(fid):
                 return bcoords, fid
             return None, None
 
-        def find_vertical_index() -> int:
-            nz = self.z.shape[0]
-            if nz == 1:
-                return 0
-            zf = self.z.values
-            # Return zi such that zf[zi] <= z < zf[zi+1]
-            zi = np.searchsorted(zf, z, side="right") - 1  # Search assumes that z is positive and increasing with i
-            if zi < 0 or zi >= nz - 1:
-                raise FieldOutOfBoundError(z, y, x)
-            return zi
-
-        zi = find_vertical_index()  # Find the vertical cell center nearest to z
+        zi, zeta = _search_1d_array(self.z.values, z)
 
         if ei is not None:
             _, fi = self.unravel_index(ei)
@@ -93,7 +83,7 @@ def find_vertical_index() -> int:
         if fi == -1:
             raise FieldOutOfBoundError(z, y, x)
 
-        return bcoords[0], self.ravel_index(zi, fi[0])
+        return {"Z": (zi, zeta), "CELL": (fi, bcoords[0])}
 
     def _get_barycentric_coordinates(self, y, x, fi):
         """Checks if a point is inside a given face id on a UxGrid."""
@@ -112,40 +102,10 @@ def _get_barycentric_coordinates(self, y, x, fi):
         err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs(np.dot(bcoord, nodes[:, 1]) - coord[1])
         return bcoord, err
 
-    def ravel_index(self, zi, fi):
-        """
-        Converts a face index and a vertical index into a single encoded index.
-
-        Parameters
-        ----------
-        zi : int
-            Vertical index (not used in unstructured grids, but kept for compatibility).
-        fi : int
-            Face index.
-
-        Returns
-        -------
-        int
-            Encoded index combining the face index and vertical index.
-        """
-        return fi + self.uxgrid.n_face * zi
+    def ravel_index(self, axis_indices: dict[_UXGRID_AXES, int]):
+        return axis_indices["CELL"] + self.uxgrid.n_face * axis_indices["Z"]
 
-    def unravel_index(self, ei):
-        """
-        Converts a single encoded index back into a vertical index and face index.
-
-        Parameters
-        ----------
-        ei : int
-            Encoded index to be unraveled.
-
-        Returns
-        -------
-        zi : int
-            Vertical index.
-        fi : int
-            Face index.
-        """
+    def unravel_index(self, ei) -> dict[_UXGRID_AXES, int]:
         zi = ei // self.uxgrid.n_face
         fi = ei % self.uxgrid.n_face
-        return zi, fi
+        return {"Z": zi, "CELL": fi}

parcels/xgrid.py

@@ -15,7 +15,6 @@
 
 _XGCM_AXIS_DIRECTION = Literal["X", "Y", "Z", "T"]
 _XGCM_AXIS_POSITION = Literal["center", "left", "right", "inner", "outer"]
-_AXIS_DIRECTION = Literal["X", "Y", "Z"]
 _XGCM_AXES = Mapping[_XGCM_AXIS_DIRECTION, xgcm.Axis]
 
 
@@ -196,13 +195,13 @@ def search(self, z, y, x, ei=None):
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
-    def ravel_index(self, axis_indices: dict[_AXIS_DIRECTION, int]) -> int:
+    def ravel_index(self, axis_indices: dict[_XGRID_AXES, int]) -> int:
         xi = axis_indices.get("X", 0)
         yi = axis_indices.get("Y", 0)
         zi = axis_indices.get("Z", 0)
         return xi + self.xdim * yi + self.xdim * self.ydim * zi
 
-    def unravel_index(self, ei) -> dict[_AXIS_DIRECTION, int]:
+    def unravel_index(self, ei) -> dict[_XGRID_AXES, int]:
         zi = ei // (self.xdim * self.ydim)
         ei = ei % (self.xdim * self.ydim)