Merge branch 'v4-dev' into adding_circulation_models_datasets

Author: erikvansebille

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/_datasets/unstructured/generic.py

@@ -44,12 +44,12 @@ def _stommel_gyre_delaunay():
     uxgrid.attrs["Conventions"] = "UGRID-1.0"
 
     # Define arrays U (zonal), V (meridional) and P (sea surface height)
-    U = np.zeros((1, nz1, lat.size), dtype=np.float64)
-    V = np.zeros((1, nz1, lat.size), dtype=np.float64)
+    U = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64)
+    V = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64)
     W = np.zeros((1, nz, lat.size), dtype=np.float64)
-    P = np.zeros((1, nz1, lat.size), dtype=np.float64)
+    P = np.zeros((1, nz1, uxgrid.n_face), dtype=np.float64)
 
-    for i, (x, y) in enumerate(zip(lon_flat, lat_flat, strict=False)):
+    for i, (x, y) in enumerate(zip(uxgrid.face_lon, uxgrid.face_lat, strict=False)):
         xi = x / 60.0
         yi = y / 60.0
 
@@ -61,10 +61,10 @@ def _stommel_gyre_delaunay():
         data=U,
         name="U",
         uxgrid=uxgrid,
-        dims=["time", "nz1", "n_node"],
+        dims=["time", "nz1", "n_face"],
         coords=dict(
             time=(["time"], [TIME[0]]),
-            nz1=(["nz1"], [0]),
+            nz1=(["nz1"], zc),
         ),
         attrs=dict(
             description="zonal velocity", units="m/s", location="node", mesh="delaunay", Conventions="UGRID-1.0"
@@ -74,7 +74,7 @@ def _stommel_gyre_delaunay():
         data=V,
         name="V",
         uxgrid=uxgrid,
-        dims=["time", "nz1", "n_node"],
+        dims=["time", "nz1", "n_face"],
         coords=dict(
             time=(["time"], [TIME[0]]),
             nz1=(["nz1"], zc),
@@ -100,7 +100,7 @@ def _stommel_gyre_delaunay():
         data=P,
         name="p",
         uxgrid=uxgrid,
-        dims=["time", "nz1", "n_node"],
+        dims=["time", "nz1", "n_face"],
         coords=dict(
             time=(["time"], [TIME[0]]),
             nz1=(["nz1"], zc),

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
 
@@ -38,7 +40,7 @@ def _search_time_index(field: Field, time: datetime):
     if the sampled value is outside the time value range.
     """
     if field.time_interval is None:
-        return 0
+        return 0, 0
 
     if time not in field.time_interval:
         _raise_time_extrapolation_error(time, field=None)
@@ -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.
+        """
+        ...