Improve Point in Face Function Signature and Helper (#1272)

* update helpers and improve point in face function signature

* update example

Author: philipc2

test/test_point_in_face.py

@@ -35,7 +35,7 @@ def test_face_centers(grid):
 
     for fid, center in enumerate(centers_xyz):
         hits = grid.get_faces_containing_point(
-            point_xyz=center,
+            points=center,
             return_counts=False
         )
         assert isinstance(hits, list)
@@ -50,7 +50,7 @@ def test_face_centers(grid):
 
     for fid, (lon, lat) in enumerate(centers_lonlat):
         hits = grid.get_faces_containing_point(
-            point_lonlat=(lon, lat),
+            points=(lon, lat),
             return_counts=False
         )
         assert hits[0] == [fid]
@@ -61,7 +61,6 @@ def test_node_corners(grid):
     Cartesian and spherical (lon/lat) returns exactly the faces sharing it.
     """
 
-    print(grid.max_face_radius)
 
     node_coords = np.vstack([
         grid.node_x.values,
@@ -77,7 +76,7 @@ def test_node_corners(grid):
         expected = conn[nid, :counts[nid]].tolist()
 
         hits_xyz = grid.get_faces_containing_point(
-            point_xyz=(x, y, z),
+            points=(x, y, z),
             return_counts=False
         )[0]
         assert set(hits_xyz) == set(expected)
@@ -94,7 +93,7 @@ def test_node_corners(grid):
         expected = conn[nid, :counts[nid]].tolist()
 
         hits_ll = grid.get_faces_containing_point(
-            point_lonlat=(lon, lat),
+            points=(lon, lat),
             return_counts=False
         )[0]
         assert set(hits_ll) == set(expected)
@@ -110,24 +109,24 @@ def test_number_of_faces_found():
     # For a face center only one face should be found
     point_xyz = np.array([grid.face_x[100].values, grid.face_y[100].values, grid.face_z[100].values], dtype=np.float64)
 
-    assert len(grid.get_faces_containing_point(point_xyz=point_xyz, return_counts=False)[0]) == 1
+    assert len(grid.get_faces_containing_point(point_xyz, return_counts=False)[0]) == 1
 
     # For an edge two faces should be found
     point_xyz = np.array([grid.edge_x[100].values, grid.edge_y[100].values, grid.edge_z[100].values], dtype=np.float64)
 
-    assert len(grid.get_faces_containing_point(point_xyz=point_xyz, return_counts=False)[0]) == 2
+    assert len(grid.get_faces_containing_point(point_xyz, return_counts=False)[0]) == 2
 
     # For a node three faces should be found
     point_xyz = np.array([grid.node_x[100].values, grid.node_y[100].values, grid.node_z[100].values], dtype=np.float64)
 
-    assert len(grid.get_faces_containing_point(point_xyz=point_xyz, return_counts=False)[0]) == 3
+    assert len(grid.get_faces_containing_point(point_xyz, return_counts=False)[0]) == 3
 
     partial_grid.normalize_cartesian_coordinates()
 
     # Test for a node on the edge where only 2 faces should be found
     point_xyz = np.array([partial_grid.node_x[1].values, partial_grid.node_y[1].values, partial_grid.node_z[1].values], dtype=np.float64)
 
-    assert len(partial_grid.get_faces_containing_point(point_xyz=point_xyz, return_counts=False)[0]) == 2
+    assert len(partial_grid.get_faces_containing_point(point_xyz, return_counts=False)[0]) == 2
 
 def test_point_along_arc():
     node_lon = np.array([-40, -40, 40, 40])
@@ -137,9 +136,24 @@ def test_point_along_arc():
     uxgrid = ux.Grid.from_topology(node_lon, node_lat, face_node_connectivity)
 
     # point at exactly 20 degrees latitude
-    out1 = uxgrid.get_faces_containing_point(point_lonlat=np.array([0, 20], dtype=np.float64), return_counts=False)
+    out1 = uxgrid.get_faces_containing_point(np.array([0, 20], dtype=np.float64), return_counts=False)
 
     # point at 25.41 degrees latitude (max along the great circle arc)
-    out2 = uxgrid.get_faces_containing_point(point_lonlat=np.array([0, 25.41], dtype=np.float64), return_counts=False)
+    out2 = uxgrid.get_faces_containing_point(np.array([0, 25.41], dtype=np.float64), return_counts=False)
 
     nt.assert_array_equal(out1[0], out2[0])
+
+
+def test_coordinates(grid):
+
+    lonlat = np.vstack([grid.node_lon.values, grid.node_lat.values]).T
+    xyz = np.vstack([grid.node_x.values, grid.node_y.values, grid.node_z.values]).T
+
+    faces_from_lonlat, _ = grid.get_faces_containing_point(points=lonlat)
+    faces_from_xyz, _ = grid.get_faces_containing_point(points=xyz)
+
+    nt.assert_array_equal(faces_from_lonlat, faces_from_xyz)
+
+    with pytest.raises(ValueError):
+        dummy_points = np.ones((10, 4))
+        faces_query_both, _ = grid.get_faces_containing_point(points=dummy_points)

uxarray/grid/coordinates.py

@@ -868,19 +868,39 @@ def prepare_points(points, normalize):
     return np.vstack([x, y, z]).T
 
 
-def _prepare_points_for_kdtree(lonlat, xyz):
-    if (lonlat is not None) and (xyz is not None):
+def points_atleast_2d_xyz(points):
+    """
+    Ensure the input is at least 2D and return Cartesian (x, y, z) coordinates.
+
+    Parameters
+    ----------
+    points : array_like, shape (N, 2) or (N, 3)
+        - If shape is (N, 2), interpreted as [longitude, latitude] in degrees.
+        - If shape is (N, 3), interpreted as Cartesian [x, y, z] coordinates.
+
+    Returns
+    -------
+    points_xyz : ndarray, shape (N, 3)
+        Cartesian coordinates [x, y, z] for each input point.
+
+    Raises
+    ------
+    ValueError
+        If `points` (after `np.atleast_2d`) does not have 2 or 3 columns.
+
+    """
+
+    points = np.atleast_2d(points)
+
+    if points.shape[1] == 2:
+        points_lonlat_rad = np.deg2rad(points)
+        x, y, z = _lonlat_rad_to_xyz(points_lonlat_rad[:, 0], points_lonlat_rad[:, 1])
+        points_xyz = np.vstack([x, y, z]).T
+    elif points.shape[1] == 3:
+        points_xyz = points
+    else:
         raise ValueError(
-            "Both Cartesian (xyz) and Spherical (lonlat) coordinates were provided. One one can be "
-            "provided at a time."
+            "Points are neither Cartesian (shape N x 3) nor Spherical (shape N x 2)."
         )
 
-    # Convert to cartesian if points are spherical
-    if xyz is None:
-        lon, lat = map(np.deg2rad, lonlat)
-        xyz = _lonlat_rad_to_xyz(lon, lat)
-    pts = np.asarray(xyz, dtype=np.float64)
-    if pts.ndim == 1:
-        pts = pts[np.newaxis, :]
-
-    return pts
+    return points_xyz

uxarray/grid/grid.py

@@ -36,8 +36,8 @@
     _populate_face_centroids,
     _populate_node_latlon,
     _populate_node_xyz,
-    _prepare_points_for_kdtree,
     _set_desired_longitude_range,
+    points_atleast_2d_xyz,
     prepare_points,
 )
 from uxarray.grid.dual import construct_dual
@@ -2560,72 +2560,77 @@ def get_faces_between_latitudes(self, lats: Tuple[float, float]):
 
     def get_faces_containing_point(
         self,
-        *,
-        point_lonlat: Sequence[float] | np.ndarray = None,
-        point_xyz: Sequence[float] | np.ndarray = None,
+        points: Sequence[float] | np.ndarray,
         return_counts: bool = True,
     ) -> Tuple[np.ndarray, np.ndarray] | List[List[int]]:
         """
-        Identify which faces on the grid contain the given point(s).
-
-        Exactly one of `point_lonlat` or `point_xyz` must be provided.
+        Identify which grid faces contain the given point(s).
 
         Parameters
         ----------
-        point_lonlat : array_like of shape (2,), optional
-            Longitude and latitude in **degrees**: (lon, lat).
-        point_xyz : array_like of shape (3,), optional
-            Cartesian coordinates on the unit sphere: (x, y, z).
+        points : array_like, shape (N, 2) or (2,) or shape (N, 3) or (3,)
+            Query point(s) to locate on the grid.
+            - If last dimension is 2, interpreted as (longitude, latitude) in **degrees**.
+            - If last dimension is 3, interpreted as Cartesian coordinates on the unit sphere: (x, y, z).
+            You may pass a single point (shape `(2,)` or `(3,)`) or multiple points (shape `(N, 2)` or `(N, 3)`).
         return_counts : bool, default=True
-            - If True (default), returns a tuple `(face_indices, counts)`.
-            - If False, returns a `list` of per-point lists of face indices.
+            - If True, returns a tuple `(face_indices, counts)`.
+            - If False, returns a `list` of per-point lists of face indices (no padding).
 
         Returns
         -------
         If `return_counts=True`:
-          face_indices : np.ndarray, shape (N, M) or (N,)
-            - 2D array of face indices with unused slots are filled with `INT_FILL_VALUE`.
+          face_indices : np.ndarray, shape (N, M) or (N, 1)
+              2D array of face indices.  Rows are padded with `INT_FILL_VALUE` when a point
+              lies on corners of multiple faces.  If every queried point falls in exactly
+              one face, the result has shape `(N, 1)`.
           counts : np.ndarray, shape (N,)
-            Number of valid face indices in each row of `face_indices`.
+              Number of valid face indices in each row of `face_indices`.
 
         If `return_counts=False`:
           List[List[int]]
-            A Python list of length `N`, where each element is the
-            list of face indices (no padding) for that query point.
+              Python list of length `N`, where each element is the list of face
+              indices for that point (no padding, in natural order).
 
         Notes
         -----
-        - **Most** points will lie strictly inside exactly one face:
-          in that common case, `counts == 1` and the returned
-          `face_indices` can be collapsed to shape `(N,)`, with no padding.
-        - If a point falls exactly on a corner shared by multiple faces,
-          multiple face indices will appear in the first columns of each row;
-          the remainder of each row is filled with `INT_FILL_VALUE`.
+        - Most points will lie strictly inside exactly one face; in that case,
+          `counts == 1` and `face_indices` has one column.
+        - Points that lie exactly on a vertex or edge shared by multiple faces
+          return multiple indices in the first `counts[i]` columns of row `i`,
+          with any remaining columns filled by `INT_FILL_VALUE`.
+
 
         Examples
         --------
-        >>> import uxarray as ux
-        >>> grid_path = "/path/to/grid.nc"
-        >>> uxgrid = ux.open_grid(grid_path)
 
-        # 1. Query a Spherical (lon/lat) point
-        >>> indices, counts = uxgrid.get_faces_containing_point(point_lonlat=(0.0, 0.0))
+        Query a single spherical point
 
-        # 2. Query a Cartesian (xyz) point
-        >>> indices, counts = uxgrid.get_faces_containing_point(
-        ...     point_xyz=(0.0, 0.0, 1.0)
-        ... )
+        >>> face_indices, counts = uxgrid.get_faces_containing_point(points=(0.0, 0.0))
+
+        Query a single Cartesian point
+
+         >>> face_indices, counts = uxgrid.get_faces_containing_point(
+         ...     points=[0.0, 0.0, 1.0]
+         ... )
 
-        # 3. Return indices as a list of lists (no counts nor padding)
-        >>> indices = uxgrid.get_faces_containing_point(
-        ...     point_xyz=(0.0, 0.0, 1.0), return_counts=False
+        Query multiple points at once
+
+        >>> pts = [(0.0, 0.0), (10.0, 20.0)]
+        >>> face_indices, counts = uxgrid.get_faces_containing_point(points=pts)
+
+        Return a list of lists
+
+        >>> face_indices_list = uxgrid.get_faces_containing_point(
+        ...     points=[0.0, 0.0, 1.0], return_counts=False
         ... )
-        """
 
-        pts = _prepare_points_for_kdtree(point_lonlat, point_xyz)
+        """
 
         # Determine faces containing points
-        face_indices, counts = _point_in_face_query(source_grid=self, points=pts)
+        face_indices, counts = _point_in_face_query(
+            source_grid=self, points=points_atleast_2d_xyz(points)
+        )
 
         # Return a list of lists if counts are not desired
         if not return_counts: