Update UxGrid to use BaseGrid

Author: fluidnumericsJoe

parcels/basegrid.py

@@ -3,7 +3,7 @@
 
 class BaseGrid(ABC):
     @abstractmethod
-    def ravel_index(self, zi, yi, xi):
+    def ravel_index(self, zi: int, yi: int, xi: int):
         """Return the flat index of the given grid points.
 
         Parameters
@@ -23,7 +23,7 @@ def ravel_index(self, zi, yi, xi):
         ...
 
     @abstractmethod
-    def unravel_index(self, ei):
+    def unravel_index(self, ei: int):
         """Return the zi, yi, xi indices for a given flat index.
         Only used when working with fields on a structured grid.
 

parcels/field.py

@@ -29,7 +29,7 @@
     FieldSamplingError,
     _raise_field_out_of_bound_error,
 )
-from parcels.uxgrid import UxGrid, ensure_uxgrid
+from parcels.uxgrid import UxGrid
 from parcels.v4.grid import Grid
 from parcels.v4.gridadapter import GridAdapter
 
@@ -145,7 +145,7 @@ def __init__(
         self,
         name: str,
         data: xr.DataArray | ux.UxDataArray,
-        grid: ux.Grid | UxGrid | Grid,
+        grid: UxGrid | Grid,
         mesh_type: Mesh = "flat",
         interp_method: Callable | None = None,
     ):
@@ -166,10 +166,7 @@ def __init__(
 
         self.name = name
         self.data = data
-        if isinstance(grid, ux.Grid):
-            self.grid = ensure_uxgrid(grid)
-        else:
-            self.grid = grid
+        self.grid = grid
 
         try:
             self.time_interval = get_time_interval(data)

parcels/uxgrid.py

@@ -1,78 +1,30 @@
+from __future__ import annotations
+
 import numpy as np
 import uxarray as ux
 from uxarray.grid.neighbors import _barycentric_coordinates
 
 from parcels.field import FieldOutOfBoundError  # Adjust import as necessary
 
+from .basegrid import BaseGrid
+
 
-class UxGrid(ux.grid.Grid):
+class UxGrid(BaseGrid):
     """
     Extension of uxarray's Grid class that supports point-location search
     for interpolation on unstructured grids.
     """
 
-    @classmethod
-    def from_uxgrid(cls, grid: ux.grid.Grid) -> "UxGrid":
-        """
-        Create a UxGrid instance from an existing uxarray Grid instance.
-
-        Parameters
-        ----------
-        grid : uxarray.grid.Grid
-            A previously constructed uxarray Grid object.
-
-        Returns
-        -------
-        UxGrid
-            A new UxGrid object with the same internal state.
-        """
-        if isinstance(grid, cls):
-            return grid  # Already an extended grid
-
-        new = cls.__new__(cls)
-        new.__dict__.update(grid.__dict__)
-        return new
+    def __init__(self, grid: ux.grid.Grid) -> UxGrid:
+        self.uxgrid = grid
 
     def search(
         self, field, z: float, y: float, x: float, ei: int | None = None, search2D: bool = False
     ) -> tuple[np.ndarray, int]:
-        """
-        Locate the unstructured grid face containing the point (x, y),
-        returning interpolation weights and a face-based encoded index.
-
-        Parameters
-        ----------
-        field : parcels.Field
-            The field requesting the search. Used to access unravel_index(),
-            ravel_index(), and igrid metadata.
-        z : float
-            Vertical coordinate of the query point. Currently ignored.
-        y : float
-            Latitude of the query point.
-        x : float
-            Longitude of the query point.
-        ei : int, optional
-            Encoded index to test reuse of previous face. If valid, neighbors
-            of that face are also checked before falling back to global search.
-        search2D : bool, default=False
-            Ignored for now. Included for interface compatibility.
-
-        Returns
-        -------
-        bcoords : np.ndarray
-            Barycentric coordinates of the point in the containing face.
-        ei : int
-            Encoded index (e.g., raveled face index) corresponding to the face found.
-
-        Raises
-        ------
-        FieldOutOfBoundError
-            If no containing face is found within tolerance.
-        """
         tol = 1e-10
 
         def try_face(fid):
-            bcoords, err = self._get_barycentric_coordinates(y, x, fid)
+            bcoords, err = self.uxgrid._get_barycentric_coordinates(y, x, fid)
             if (bcoords >= 0).all() and (bcoords <= 1).all() and err < tol:
                 return bcoords, field.ravel_index(0, 0, fid)  # Z and time indices are 0 for now
             return None, None
@@ -84,15 +36,15 @@ def try_face(fid):
                 return bcoords, ei_new
 
             # Try neighbors of current face
-            for neighbor in self.face_face_connectivity[fi, :]:
+            for neighbor in self.uxgrid.face_face_connectivity[fi, :]:
                 if neighbor == -1:
                     continue
                 bcoords, ei_new = try_face(neighbor)
                 if bcoords is not None:
                     return bcoords, ei_new
 
         # Global fallback using spatial hash
-        fi, bcoords = self.get_spatial_hash().query([[x, y]])
+        fi, bcoords = self.uxgrid.get_spatial_hash().query([[x, y]])
         if fi == -1:
             raise FieldOutOfBoundError(z, y, x)
 
@@ -101,12 +53,12 @@ def try_face(fid):
     def _get_barycentric_coordinates(self, y, x, fi):
         """Checks if a point is inside a given face id on a UxGrid."""
         # Check if particle is in the same face, otherwise search again.
-        n_nodes = self.n_nodes_per_face[fi].to_numpy()
-        node_ids = self.face_node_connectivity[fi, 0:n_nodes]
+        n_nodes = self.uxgrid.n_nodes_per_face[fi].to_numpy()
+        node_ids = self.uxgrid.face_node_connectivity[fi, 0:n_nodes]
         nodes = np.column_stack(
             (
-                np.deg2rad(self.grid.node_lon[node_ids].to_numpy()),
-                np.deg2rad(self.grid.node_lat[node_ids].to_numpy()),
+                np.deg2rad(self.uxgrid.grid.node_lon[node_ids].to_numpy()),
+                np.deg2rad(self.uxgrid.grid.node_lat[node_ids].to_numpy()),
             )
         )
 
@@ -116,57 +68,9 @@ def _get_barycentric_coordinates(self, y, x, fi):
         return bcoord, err
 
     def ravel_index(self, zi, yi, xi):
-        """Return the flat index of the given grid points.
-
-        Parameters
-        ----------
-        zi : int
-            z index
-        yi : int
-            y index
-        xi : int
-            x index. When using an unstructured grid, this is the face index (fi)
-
-        Returns
-        -------
-        int
-            flat index
-        """
-        return xi + self.n_face * zi
+        return xi + self.uxgrid.n_face * zi
 
     def unravel_index(self, ei):
-        """Return the zi, yi, xi indices for a given flat index.
-        Only used when working with fields on a structured grid.
-
-        Parameters
-        ----------
-        ei : int
-            The flat index to be unraveled.
-
-        Returns
-        -------
-        zi : int
-            The z index.
-        yi : int
-            The y index.
-        xi : int
-            The x index.
-        """
-        zi = ei // self.n_face
-        fi = ei % self.n_face
+        zi = ei // self.uxgrid.n_face
+        fi = ei % self.uxgrid.n_face
         return zi, fi
-
-
-def ensure_uxgrid(grid: ux.grid.Grid) -> UxGrid:
-    """
-    Ensure a given uxarray grid is an instance of UxGrid.
-
-    Parameters
-    ----------
-    grid : uxarray.grid.Grid
-
-    Returns
-    -------
-    UxGrid
-    """
-    return UxGrid.from_uxgrid(grid)