Update xgrid and basegrid search, ravel, unravel API

Author: VeckoTheGecko

parcels/basegrid.py

@@ -1,15 +1,21 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
 
 
 class BaseGrid(ABC):
     @abstractmethod
-    def search(self, z: float, y: float, x: float, ei=None, search2D: bool = False):
+    def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int, float | np.ndarray]]:
         """
         Perform a spatial (and optionally vertical) search to locate the grid element
         that contains a given point (x, y, z).
 
         This method delegates to grid-type-specific logic (e.g., structured or unstructured)
-        to determine the appropriate indices and interpolation coordinates for evaluating a field.
+        to determine the appropriate indices and barycentric coordinates for evaluating a field.
 
         Parameters
         ----------
@@ -28,12 +34,21 @@ def search(self, z: float, y: float, x: float, ei=None, search2D: bool = False):
 
         Returns
         -------
-        bcoords : np.ndarray or tuple
-            Interpolation weights or barycentric coordinates within the containing cell/face.
-            The interpretation of `bcoords` depends on the grid type.
-        ei : int
-            Encoded index of the identified grid cell or face. This value can be cached for
-            future lookups to accelerate repeated searches.
+        dict
+            A dictionary mapping spatial axis names to tuples of (index, barycentric_coordinates).
+            The returned axes depend on the grid dimensionality and type:
+
+            - 3D structured grid: {"X": (xi, xsi), "Y": (yi, eta), "Z": (zi, zeta)}
+            - 2D structured grid: {"X": (xi, xsi), "Y": (yi, eta)}
+            - 1D structured grid (depth): {"Z": (zi, zeta)}
+            - Unstructured grid: {"Z": (zi, zeta), "FACE": (fi, bcoords)}
+
+            Where:
+            - index (int): The cell position of a particle along the given axis
+            - barycentric_coordinates (float or np.ndarray): The coordinates defining
+              a particle's position within the grid cell. For structured grids, this
+              is a single coordinate per axis; for unstructured grids, this can be
+              an array of coordinates for the face polygon.
 
         Raises
         ------
@@ -43,3 +58,72 @@ def search(self, z: float, y: float, x: float, ei=None, search2D: bool = False):
             Raised if the search method is not implemented for the current grid type.
         """
         ...
+
+    @abstractmethod
+    def ravel_index(self, axis_indices: dict[str, int]) -> int:
+        """
+        Convert a dictionary of axis indices to a single encoded index (ei).
+
+        This method takes the individual indices for each spatial axis and combines them
+        into a single integer that uniquely identifies a grid cell. This encoded
+        index can be used for efficient caching and lookup operations.
+
+        Parameters
+        ----------
+        axis_indices : dict[str, int]
+            A dictionary mapping axis names to their corresponding indices.
+            The expected keys depend on the grid dimensionality and type:
+
+            - 3D structured grid: {"X": xi, "Y": yi, "Z": zi}
+            - 2D structured grid: {"X": xi, "Y": yi}
+            - 1D structured grid: {"Z": zi}
+            - Unstructured grid: {"Z": zi, "FACE": fi}
+
+        Returns
+        -------
+        int
+            The encoded index (ei) representing the unique grid cell or face.
+
+        Raises
+        ------
+        KeyError
+            Raised when required axis keys are missing from axis_indices.
+        ValueError
+            Raised when index values are out of bounds for the grid.
+        NotImplementedError
+            Raised if the method is not implemented for the current grid type.
+        """
+        ...
+
+    @abstractmethod
+    def unravel_index(self, ei: int) -> dict[str, int]:
+        """
+        Convert a single encoded index (ei) back to a dictionary of axis indices.
+
+        This method is the inverse of ravel_index, taking an encoded index and
+        decomposing it back into the individual indices for each spatial axis.
+
+        Parameters
+        ----------
+        ei : int
+            The encoded index representing a unique grid cell or face.
+
+        Returns
+        -------
+        dict[str, int]
+            A dictionary mapping axis names to their corresponding indices.
+            The returned keys depend on the grid dimensionality and type:
+
+            - 3D structured grid: {"X": xi, "Y": yi, "Z": zi}
+            - 2D structured grid: {"X": xi, "Y": yi}
+            - 1D structured grid: {"Z": zi}
+            - Unstructured grid: {"Z": zi, "FACE": fi}
+
+        Raises
+        ------
+        ValueError
+            Raised when the encoded index is out of bounds or invalid for the grid.
+        NotImplementedError
+            Raised if the method is not implemented for the current grid type.
+        """
+        ...

parcels/xgrid.py

@@ -10,9 +10,10 @@
 from parcels.basegrid import BaseGrid
 from parcels.tools.converters import TimeConverter
 
-_AXIS_DIRECTION = Literal["X", "Y", "Z", "T"]
-_AXIS_POSITION = Literal["center", "left", "right", "inner", "outer"]
-_XGCM_AXES = Mapping[_AXIS_DIRECTION, xgcm.Axis]
+_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]
 
 
 def get_tracer_dimensionality(axis: xgcm.Axis | None) -> int:
@@ -180,85 +181,57 @@ def _gtype(self):
             else:
                 return GridType.CurvilinearSGrid
 
-    def search(self, z, y, x, ei=None, search2D=False):
+    def search(self, z, y, x, ei=None):
         ds = self.xgcm_grid._ds
 
-        if search2D:
-            zi = 0
-        else:
-            zi, _ = _search_1d_array(ds.depth.values, z)
+        zi, zeta = _search_1d_array(ds.depth.values, z)
 
         if ds.lon.ndim == 1:
             yi, eta = _search_1d_array(ds.lat.values, y)
             xi, xsi = _search_1d_array(ds.lon.values, x)
-            return np.array([eta, xsi, 1 - eta, 1 - xsi]), self.ravel_index(zi, yi, xi)
+            return {"X": (xi, xsi), "Y": (yi, eta), "Z": (zi, zeta)}
 
         yi, xi = None, None
         if ei is not None:
-            _, yi, xi = self.unravel_index(ei)
+            axis_indices = self.unravel_index(ei)
+            xi = axis_indices.get("X")
+            yi = axis_indices.get("Y")
 
         if ds.lon.ndim == 2:
             eta, xsi, yi, xi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
 
-            return np.array([eta, xsi, 1 - eta, 1 - xsi]), self.ravel_index(zi, yi, xi)
+            return {"X": (xi, xsi), "Y": (yi, eta), "Z": (zi, zeta)}
 
         raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.")
 
-    def ravel_index(self, zi, yi, xi):
-        """
-        Converts a z, y, and x index into a single encoded index.
-
-        Parameters
-        ----------
-        zi : int
-            Vertical index.
-        yi : int
-            Latitude index.
-        xi : int
-            Longitude index.
-
-        Returns
-        -------
-        int
-            Encoded index.
-        """
+    def ravel_index(self, axis_indices: dict[_AXIS_DIRECTION, 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):
-        """
-        Converts a single encoded index back into a Z, Y, and X indices.
-
-        Parameters
-        ----------
-        ei : int
-            Encoded index to be unraveled.
-
-        Returns
-        -------
-        zi : int
-            Vertical index.
-        yi : int
-            Latitude index.
-        xi : int
-            Longitude index.
-        """
+    def unravel_index(self, ei) -> dict[_AXIS_DIRECTION, int]:
         zi = ei // (self.xdim * self.ydim)
         ei = ei % (self.xdim * self.ydim)
 
         yi = ei // self.xdim
         xi = ei % self.xdim
-        return zi, yi, xi
+        return {
+            "X": xi,
+            "Y": yi,
+            "Z": zi,
+        }
 
 
-def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _AXIS_DIRECTION | None:
+def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None:
     """For a given dimension name in a grid, returns the direction axis it is on."""
     for axis_name, axis in axes.items():
         if dim in axis.coords.values():
             return axis_name
     return None
 
 
-def get_position_from_dim_name(axes: _XGCM_AXES, dim: str) -> _AXIS_POSITION | None:
+def get_position_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_POSITION | None:
     """For a given dimension, returns the position of the variable in the grid."""
     for axis in axes.values():
         var_to_position = {var: position for position, var in axis.coords.items()}
@@ -287,7 +260,7 @@ def assert_valid_field_array(da: xr.DataArray, axes: _XGCM_AXES):
     assert_all_dimensions_correspond_with_axis(da, axes)
 
     dim_to_axis = {dim: get_axis_from_dim_name(axes, dim) for dim in da.dims}
-    dim_to_axis = cast(dict[Hashable, _AXIS_DIRECTION], dim_to_axis)
+    dim_to_axis = cast(dict[Hashable, _XGCM_AXIS_DIRECTION], dim_to_axis)
 
     # Assert all dimensions are present
     if set(dim_to_axis.values()) != {"T", "Z", "Y", "X"}:

tests/v4/test_xgrid.py

@@ -129,11 +129,10 @@ def test_xgrid_ravel_unravel_index():
     for xi in range(xdim):
         for yi in range(ydim):
             for zi in range(zdim):
-                ei = grid.ravel_index(zi, yi, xi)
-                zi_test, yi_test, xi_test = grid.unravel_index(ei)
-                assert xi == xi_test, f"Expected xi {xi} but got {xi_test} for ei {ei}"
-                assert yi == yi_test, f"Expected yi {yi} but got {yi_test} for ei {ei}"
-                assert zi == zi_test, f"Expected zi {zi} but got {zi_test} for ei {ei}"
+                axis_indices = {"X": xi, "Y": yi, "Z": zi}
+                ei = grid.ravel_index(axis_indices)
+                axis_indices_test = grid.unravel_index(ei)
+                assert axis_indices_test == axis_indices
                 encountered_eis.append(ei)
 
     encountered_eis = sorted(encountered_eis)
@@ -156,12 +155,12 @@ def test_xgrid_search_cpoints(ds):
 
     for xi in range(grid.xdim - 1):
         for yi in range(grid.ydim - 1):
+            axis_indices = {"X": xi, "Y": yi, "Z": 0}
+
             lat, lon = lat_array[yi, xi], lon_array[yi, xi]
-            bcoords, ei = grid.search(0, lat, lon, ei=None, search2D=True)
-            zi_test, yi_test, xi_test = grid.unravel_index(ei)
-            assert xi == xi_test
-            assert yi == yi_test
-            assert zi_test == 0
+            axis_indices_bcoords = grid.search(0, lat, lon, ei=None)
+            axis_indices_test = {k: v[0] for k, v in axis_indices_bcoords.items()}
+            assert axis_indices == axis_indices_test
 
             # assert np.isclose(bcoords[0], 0.5) #? Should this not be the case with the cell center points?
             # assert np.isclose(bcoords[1], 0.5)