Swap out spatialhash query for get_faces_containing_point

parcels/_datasets/unstructured/generic.py

@@ -208,7 +208,109 @@ def _fesom2_square_delaunay_uniform_z_coordinate():
     return ux.UxDataset({"U": u, "V": v, "W": w, "p": p}, uxgrid=uxgrid)
 
 
+def _fesom2_square_delaunay_antimeridian():
+    """
+    Delaunay grid that crosses the antimeridian with uniform z-coordinate, mimicking a FESOM2 dataset.
+    This dataset consists of a square domain with closed boundaries, where the grid is generated using Delaunay triangulation.
+    The bottom topography is flat and uniform, and the vertical grid spacing is constant with 10 layers spanning [0,1000.0]
+    The lateral velocity field components are non-zero constant, and the vertical velocity component is zero.
+    The pressure field is constant.
+    All fields are placed on location consistent with FESOM2 variable placement conventions
+    """
+    lon, lat = np.meshgrid(
+        np.linspace(-210.0, -150.0, Nx, dtype=np.float32), np.linspace(0, 60.0, Nx, dtype=np.float32)
+    )
+    # wrap longitude from [-180,180]
+    lon = np.where(lon < -180, lon + 360, lon)
+    lon_flat = lon.ravel()
+    lat_flat = lat.ravel()
+    zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float32)  # Vertical element faces
+    zc = 0.5 * (zf[:-1] + zf[1:])  # Vertical element centers
+    nz = zf.size
+    nz1 = zc.size
+
+    # mask any point on one of the boundaries
+    mask = (
+        np.isclose(lon_flat, 0.0) | np.isclose(lon_flat, 60.0) | np.isclose(lat_flat, 0.0) | np.isclose(lat_flat, 60.0)
+    )
+
+    boundary_points = np.flatnonzero(mask)
+
+    uxgrid = ux.Grid.from_points(
+        (lon_flat, lat_flat),
+        method="regional_delaunay",
+        boundary_points=boundary_points,
+    )
+    uxgrid.attrs["Conventions"] = "UGRID-1.0"
+
+    # Define arrays U (zonal), V (meridional) and P (sea surface height)
+    U = np.ones(
+        (T, nz1, uxgrid.n_face), dtype=np.float64
+    )  # Lateral velocity is on the element centers and face centers
+    V = np.ones(
+        (T, nz1, uxgrid.n_face), dtype=np.float64
+    )  # Lateral velocity is on the element centers and face centers
+    W = np.zeros(
+        (T, nz, uxgrid.n_node), dtype=np.float64
+    )  # Vertical velocity is on the element faces and face vertices
+    P = np.ones((T, nz1, uxgrid.n_node), dtype=np.float64)  # Pressure is on the element centers and face vertices
+
+    u = ux.UxDataArray(
+        data=U,
+        name="U",
+        uxgrid=uxgrid,
+        dims=["time", "nz1", "n_face"],
+        coords=dict(
+            time=(["time"], TIME),
+            nz1=(["nz1"], zc),
+        ),
+        attrs=dict(
+            description="zonal velocity", units="m/s", location="face", mesh="delaunay", Conventions="UGRID-1.0"
+        ),
+    )
+    v = ux.UxDataArray(
+        data=V,
+        name="V",
+        uxgrid=uxgrid,
+        dims=["time", "nz1", "n_face"],
+        coords=dict(
+            time=(["time"], TIME),
+            nz1=(["nz1"], zc),
+        ),
+        attrs=dict(
+            description="meridional velocity", units="m/s", location="face", mesh="delaunay", Conventions="UGRID-1.0"
+        ),
+    )
+    w = ux.UxDataArray(
+        data=W,
+        name="w",
+        uxgrid=uxgrid,
+        dims=["time", "nz", "n_node"],
+        coords=dict(
+            time=(["time"], TIME),
+            nz=(["nz"], zf),
+        ),
+        attrs=dict(
+            description="vertical velocity", units="m/s", location="node", mesh="delaunay", Conventions="UGRID-1.0"
+        ),
+    )
+    p = ux.UxDataArray(
+        data=P,
+        name="p",
+        uxgrid=uxgrid,
+        dims=["time", "nz1", "n_node"],
+        coords=dict(
+            time=(["time"], TIME),
+            nz1=(["nz1"], zc),
+        ),
+        attrs=dict(description="pressure", units="N/m^2", location="node", mesh="delaunay", Conventions="UGRID-1.0"),
+    )
+
+    return ux.UxDataset({"U": u, "V": v, "W": w, "p": p}, uxgrid=uxgrid)
+
+
 datasets = {
     "stommel_gyre_delaunay": _stommel_gyre_delaunay(),
     "fesom2_square_delaunay_uniform_z_coordinate": _fesom2_square_delaunay_uniform_z_coordinate(),
+    "fesom2_square_delaunay_antimeridian": _fesom2_square_delaunay_antimeridian(),
 }

parcels/uxgrid.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import uxarray as ux
+from uxarray.grid.coordinates import _lonlat_rad_to_xyz
 from uxarray.grid.neighbors import _barycentric_coordinates
 
 from parcels.field import FieldOutOfBoundError  # Adjust import as necessary
@@ -67,49 +68,125 @@ def get_axis_dim(self, axis: _UXGRID_AXES) -> int:
         elif axis == "FACE":
             return self.uxgrid.n_face
 
-    def search(self, z, y, x, ei=None):
-        tol = 1e-10
-
+    def search(self, z, y, x, ei=None, tol=1e-6):
         def try_face(fid):
-            bcoords, err = self.uxgrid._get_barycentric_coordinates(y, x, fid)
+            bcoords, err = self._get_barycentric_coordinates_latlon(y, x, fid)
             if (bcoords >= 0).all() and (bcoords <= 1).all() and err < tol:
-                return bcoords, fid
-            return None, None
+                return bcoords
+            else:
+                bcoords = self._get_barycentric_coordinates_cartesian(y, x, fid)
+                if (bcoords >= 0).all() and (bcoords <= 1).all():
+                    return bcoords
+
+            return None
 
         zi, zeta = _search_1d_array(self.z.values, z)
 
         if ei is not None:
             _, fi = self.unravel_index(ei)
-            bcoords, fi_new = try_face(fi)
+            bcoords = try_face(fi)
             if bcoords is not None:
-                return bcoords, self.ravel_index(zi, fi_new)
+                return bcoords, self.ravel_index(zi, fi)
             # Try neighbors of current face
             for neighbor in self.uxgrid.face_face_connectivity[fi, :]:
                 if neighbor == -1:
                     continue
-                bcoords, fi_new = try_face(neighbor)
+                bcoords = try_face(neighbor)
                 if bcoords is not None:
-                    return bcoords, self.ravel_index(zi, fi_new)
+                    return bcoords, self.ravel_index(zi, neighbor)
 
-        # Global fallback using spatial hash
-        fi, bcoords = self.uxgrid.get_spatial_hash().query([[x, y]])
+        # Global fallback as last ditch effort
+        face_ids = self.uxgrid.get_faces_containing_point([x, y], return_counts=False)[0]
+        fi = face_ids[0] if len(face_ids) > 0 else -1
         if fi == -1:
             raise FieldOutOfBoundError(z, y, x)
-        return {"Z": (zi, zeta), "FACE": (fi[0], bcoords[0])}
+        bcoords = try_face(fi)
+        if bcoords is None:
+            raise FieldOutOfBoundError(z, y, x)
 
-    def _get_barycentric_coordinates(self, y, x, fi):
+        return {"Z": (zi, zeta), "FACE": (fi, bcoords)}
+
+    def _get_barycentric_coordinates_latlon(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.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.uxgrid.grid.node_lon[node_ids].to_numpy()),
-                np.deg2rad(self.uxgrid.grid.node_lat[node_ids].to_numpy()),
+                np.deg2rad(self.uxgrid.node_lon[node_ids].to_numpy()),
+                np.deg2rad(self.uxgrid.node_lat[node_ids].to_numpy()),
             )
         )
 
         coord = np.deg2rad([x, y])
         bcoord = np.asarray(_barycentric_coordinates(nodes, coord))
         err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs(np.dot(bcoord, nodes[:, 1]) - coord[1])
         return bcoord, err
+
+    def _get_barycentric_coordinates_cartesian(self, y, x, fi):
+        n_nodes = self.uxgrid.n_nodes_per_face[fi].to_numpy()
+        node_ids = self.uxgrid.face_node_connectivity[fi, 0:n_nodes]
+
+        coord = np.deg2rad([x, y])
+        x, y, z = _lonlat_rad_to_xyz(coord[0], coord[1])
+        cart_coord = np.array([x, y, z]).T
+        # Second attempt to find barycentric coordinates using cartesian coordinates
+        nodes = np.stack(
+            (
+                self.uxgrid.node_x[node_ids].values,
+                self.uxgrid.node_y[node_ids].values,
+                self.uxgrid.node_z[node_ids].values,
+            ),
+            axis=-1,
+        )
+
+        bcoord = np.asarray(_barycentric_coordinates_cartesian(nodes, cart_coord))
+
+        return bcoord
+
+
+def _barycentric_coordinates_cartesian(nodes, point, min_area=1e-8):
+    """
+    Compute the barycentric coordinates of a point P inside a convex polygon using area-based weights.
+    So that this method generalizes to n-sided polygons, we use the Waschpress points as the generalized
+    barycentric coordinates, which is only valid for convex polygons.
+
+    Parameters
+    ----------
+        nodes : numpy.ndarray
+            Cartesian coordinates (x,y,z) of each corner node of a face
+        point : numpy.ndarray
+            Cartesian coordinates (x,y,z) of the point
+
+    Returns
+    -------
+    numpy.ndarray
+        Barycentric coordinates corresponding to each vertex.
+
+    """
+    n = len(nodes)
+    sum_wi = 0
+    w = []
+
+    for i in range(0, n):
+        vim1 = nodes[i - 1]
+        vi = nodes[i]
+        vi1 = nodes[(i + 1) % n]
+        a0 = _triangle_area_cartesian(vim1, vi, vi1)
+        a1 = max(_triangle_area_cartesian(point, vim1, vi), min_area)
+        a2 = max(_triangle_area_cartesian(point, vi, vi1), min_area)
+        sum_wi += a0 / (a1 * a2)
+        w.append(a0 / (a1 * a2))
+
+    barycentric_coords = [w_i / sum_wi for w_i in w]
+
+    return barycentric_coords
+
+
+def _triangle_area_cartesian(A, B, C):
+    """Compute the area of a triangle given by three points."""
+    d1 = B - A
+    d2 = C - A
+    d3 = np.cross(d1, d2)
+    return 0.5 * np.linalg.norm(d3)

tests/v4/test_uxarray_fieldset.py

@@ -115,10 +115,26 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval():
         V=Field(name="V", data=ds.V, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace),
         W=Field(name="W", data=ds.W, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode),
     )
-    P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseConstantFace)
+    P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode)
     fieldset = FieldSet([UVW, P, UVW.U, UVW.V, UVW.W])
 
     assert fieldset.U.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0
     assert fieldset.V.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0
     assert fieldset.W.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 0.0
     assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0
+
+
+def test_fesom2_square_delaunay_antimeridian_eval():
+    """
+    Test the evaluation of a fieldset with a FESOM2 square Delaunay grid that crosses the antimeridian.
+    Ensures that the fieldset can be created and evaluated correctly.
+    Since the underlying data is constant, we can check that the values are as expected.
+    """
+    ds = datasets_unstructured["fesom2_square_delaunay_antimeridian"]
+    P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode)
+    fieldset = FieldSet([P])
+
+    assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=-170.0, applyConversion=False) == 1.0
+    assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=-180.0, applyConversion=False) == 1.0
+    assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=180.0, applyConversion=False) == 1.0
+    assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=170.0, applyConversion=False) == 1.0