Structured grids barycentric coordinates

parcels/_index_search.py

@@ -21,6 +21,8 @@
 from .grid import GridType
 
 if TYPE_CHECKING:
+    from parcels.xgrid import XGrid
+
     from .field import Field
     # from .grid import Grid
 
@@ -270,6 +272,89 @@ def _search_indices_rectilinear(
     return (zeta, eta, xsi, _ei)
 
 
+def _search_indices_curvilinear_2d(
+    grid: XGrid, y: float, x: float, yi_guess: int | None = None, xi_guess: int | None = None
+):
+    yi, xi = yi_guess, xi_guess
+    if yi is None:
+        yi = int(grid.ydim / 2) - 1
+
+    if xi is None:
+        xi = int(grid.xdim / 2) - 1
+
+    xsi = eta = -1.0
+    invA = np.array(
+        [
+            [1, 0, 0, 0],
+            [-1, 1, 0, 0],
+            [-1, 0, 0, 1],
+            [1, -1, 1, -1],
+        ]
+    )
+    maxIterSearch = 1e6
+    it = 0
+    tol = 1.0e-10
+
+    # # ! Error handling for out of bounds
+    # TODO: Re-enable in some capacity
+    # if x < field.lonlat_minmax[0] or x > field.lonlat_minmax[1]:
+    #     if grid.lon[0, 0] < grid.lon[0, -1]:
+    #         _raise_field_out_of_bound_error(y, x)
+    #     elif x < grid.lon[0, 0] and x > grid.lon[0, -1]:  # This prevents from crashing in [160, -160]
+    #         _raise_field_out_of_bound_error(z, y, x)
+
+    # if y < field.lonlat_minmax[2] or y > field.lonlat_minmax[3]:
+    #     _raise_field_out_of_bound_error(z, y, x)
+
+    while xsi < -tol or xsi > 1 + tol or eta < -tol or eta > 1 + tol:
+        px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]])
+
+        py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]])
+        a = np.dot(invA, px)
+        b = np.dot(invA, py)
+
+        aa = a[3] * b[2] - a[2] * b[3]
+        bb = a[3] * b[0] - a[0] * b[3] + a[1] * b[2] - a[2] * b[1] + x * b[3] - y * a[3]
+        cc = a[1] * b[0] - a[0] * b[1] + x * b[1] - y * a[1]
+        if abs(aa) < 1e-12:  # Rectilinear cell, or quasi
+            eta = -cc / bb
+        else:
+            det2 = bb * bb - 4 * aa * cc
+            if det2 > 0:  # so, if det is nan we keep the xsi, eta from previous iter
+                det = np.sqrt(det2)
+                eta = (-bb + det) / (2 * aa)
+        if abs(a[1] + a[3] * eta) < 1e-12:  # this happens when recti cell rotated of 90deg
+            xsi = ((y - py[0]) / (py[1] - py[0]) + (y - py[3]) / (py[2] - py[3])) * 0.5
+        else:
+            xsi = (x - a[0] - a[2] * eta) / (a[1] + a[3] * eta)
+        if xsi < 0 and eta < 0 and xi == 0 and yi == 0:
+            _raise_field_out_of_bound_error(0, y, x)
+        if xsi > 1 and eta > 1 and xi == grid.xdim - 1 and yi == grid.ydim - 1:
+            _raise_field_out_of_bound_error(0, y, x)
+        if xsi < -tol:
+            xi -= 1
+        elif xsi > 1 + tol:
+            xi += 1
+        if eta < -tol:
+            yi -= 1
+        elif eta > 1 + tol:
+            yi += 1
+        (yi, xi) = _reconnect_bnd_indices(yi, xi, grid.ydim, grid.xdim, grid.mesh)
+        it += 1
+        if it > maxIterSearch:
+            print(f"Correct cell not found after {maxIterSearch} iterations")
+            _raise_field_out_of_bound_error(0, y, x)
+    xsi = max(0.0, xsi)
+    eta = max(0.0, eta)
+    xsi = min(1.0, xsi)
+    eta = min(1.0, eta)
+
+    if not ((0 <= xsi <= 1) and (0 <= eta <= 1)):
+        _raise_field_sampling_error(y, x)
+
+    return (yi, eta, xi, xsi)
+
+
 ## TODO :  Still need to implement the search_indices_curvilinear
 def _search_indices_curvilinear(field: Field, time, z, y, x, ti, particle=None, search2D=False):
     if particle:

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: {"Z": (zi, zeta), "Y": (yi, eta), "X": (xi, xsi)}
+            - 2D structured grid: {"Y": (yi, eta), "X": (xi, xsi)}
+            - 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: {"Z": zi, "Y": yi, "X": xi}
+            - 2D structured grid: {"Y": yi, "X": xi}
+            - 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: {"Z": zi, "Y": yi, "X": xi}
+            - 2D structured grid: {"Y": yi, "X": xi}
+            - 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.
+        """
+        ...