Merge remote-tracking branch 'origin/v4-dev' into minimal-v4-particleset

Author: fluidnumericsJoe

docs/reference/grids.rst

@@ -1,13 +1,6 @@
 Gridsets and grids
 ==================
 
-parcels.gridset module
-----------------------
-
-.. automodule:: parcels.gridset
-    :members:
-    :show-inheritance:
-
 parcels.grid module
 -------------------
 

parcels/__init__.py

@@ -6,7 +6,6 @@
 from parcels.field import *
 from parcels.fieldset import *
 from parcels.grid import *
-from parcels.gridset import *
 from parcels.interaction import *
 from parcels.kernel import *
 from parcels.particle import *

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/_reprs.py

@@ -65,7 +65,7 @@ def particleset_repr(pset: ParticleSet) -> str:
 
 def fieldset_repr(fieldset: FieldSet) -> str:  # TODO v4: Rework or remove entirely
     """Return a pretty repr for FieldSet"""
-    fields_repr = "\n".join([repr(f) for f in fieldset.get_fields()])
+    fields_repr = "\n".join([repr(f) for f in fieldset.fields.values()])
 
     out = f"""<{type(fieldset).__name__}>
     fields:

parcels/application_kernels/advection.py

@@ -175,7 +175,7 @@ def AdvectionAnalytical(particle, fieldset, time):  # pragma: no cover
     tol = 1e-10
     I_s = 10  # number of intermediate time steps
     direction = 1.0 if particle.dt > 0 else -1.0
-    withW = True if "W" in [f.name for f in fieldset.get_fields()] else False
+    withW = True if "W" in [f.name for f in fieldset.fields.values()] else False
     withTime = True if len(fieldset.U.grid.time) > 1 else False
     tau, zeta, eta, xsi, ti, zi, yi, xi = fieldset.U._search_indices(
         time, particle.depth, particle.lat, particle.lon, particle=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.
+        """
+        ...

parcels/field.py

@@ -169,7 +169,7 @@ def __init__(
             self.time_interval = get_time_interval(data)
         except ValueError as e:
             e.add_note(
-                f"Error getting time interval for field {name!r}. Are you sure that the time dimension on the xarray dataset is stored as datetime or cftime datetime objects?"
+                f"Error getting time interval for field {name!r}. Are you sure that the time dimension on the xarray dataset is stored as timedelta, datetime or cftime datetime objects?"
             )
             raise e
 

parcels/fieldset.py

@@ -16,6 +16,7 @@
 
 if TYPE_CHECKING:
     from parcels._typing import DatetimeLike
+    from parcels.basegrid import BaseGrid
 __all__ = ["FieldSet"]
 
 
@@ -109,10 +110,6 @@ def dimrange(self, dim):
 
         return maxleft, minright
 
-    @property
-    def gridset_size(self):
-        return len(self.fields)
-
     def add_field(self, field: Field, name: str | None = None):
         """Add a :class:`parcels.field.Field` object to the FieldSet.
 
@@ -164,15 +161,6 @@ def add_constant_field(self, name: str, value, mesh: Mesh = "flat"):
         da = xr.DataArray(
             data=np.full((1, 1, 1, 1), value),
             dims=["time", "ZG", "YG", "XG"],
-            coords={
-                "ZG": (["ZG"], np.arange(1), {"axis": "Z"}),
-                "YG": (["YG"], np.arange(1), {"axis": "Y"}),
-                "XG": (["XG"], np.arange(1), {"axis": "X"}),
-                "lon": (["XG"], np.arange(1), {"axis": "X"}),
-                "lat": (["YG"], np.arange(1), {"axis": "Y"}),
-                "depth": (["ZG"], np.arange(1), {"axis": "Z"}),
-                "time": (["time"], np.arange(1), {"axis": "T"}),
-            },
         )
         grid = XGrid(xgcm.Grid(da))
         self.add_field(
@@ -184,17 +172,6 @@ def add_constant_field(self, name: str, value, mesh: Mesh = "flat"):
             )
         )
 
-    def get_fields(self) -> list[Field | VectorField]:
-        """Returns a list of all the :class:`parcels.field.Field` and :class:`parcels.field.VectorField`
-        objects associated with this FieldSet.
-        """
-        fields = []
-        for v in self.__dict__.values():
-            if type(v) in [Field, VectorField]:
-                if v not in fields:
-                    fields.append(v)
-        return fields
-
     def add_constant(self, name, value):
         """Add a constant to the FieldSet. Note that all constants are
         stored as 32-bit floats.
@@ -219,6 +196,14 @@ def add_constant(self, name, value):
 
         self.constants[name] = np.float32(value)
 
+    @property
+    def gridset(self) -> list[BaseGrid]:
+        grids = []
+        for field in self.fields.values():
+            if field.grid not in grids:
+                grids.append(field.grid)
+        return grids
+
     # def computeTimeChunk(self, time=0.0, dt=1):
     #     """Load a chunk of three data time steps into the FieldSet.
     #     This is used when FieldSet uses data imported from netcdf,