Optimize Point in Face Computation (#1249)

* initial refactor

* optimize point in face

* use angle function

* optimize logic

* update comments

* remove debugging

* update max face radius

* add examples and clean up logic

* add examples and clean up logic

* clean up point util

* add example

* update typing

* better docstring

Author: philipc2

test/test_geometry.py

@@ -11,10 +11,12 @@
 from uxarray.grid.coordinates import _populate_node_latlon, _lonlat_rad_to_xyz, _normalize_xyz, _xyz_to_lonlat_rad, \
     _xyz_to_lonlat_deg, _xyz_to_lonlat_rad_scalar
 from uxarray.grid.arcs import extreme_gca_latitude, extreme_gca_z
-from uxarray.grid.utils import _get_cartesian_face_edge_nodes_array, _get_lonlat_rad_face_edge_nodes_array
+from uxarray.grid.utils import _get_cartesian_face_edge_nodes_array, _get_lonlat_rad_face_edge_nodes_array, _get_cartesian_face_edge_nodes
 
 from uxarray.grid.geometry import _pole_point_inside_polygon_cartesian, \
-    stereographic_projection, inverse_stereographic_projection, point_in_face, haversine_distance
+    stereographic_projection, inverse_stereographic_projection,  haversine_distance
+
+from uxarray.grid.point_in_face import _face_contains_point
 
 
 from uxarray.grid.bounds import _populate_face_bounds, _construct_face_bounds_array, _construct_face_bounds
@@ -1183,24 +1185,19 @@ def test_point_inside():
 
     # Open grid
     grid = ux.open_grid(grid_mpas_2)
+    grid.normalize_cartesian_coordinates()
 
-    # Get the face edges of all faces in the grid
-    faces_edges_cartesian = _get_cartesian_face_edge_nodes_array(
-        grid.face_node_connectivity.values,
-        grid.n_face,
-        grid.n_max_face_edges,
-        grid.node_x.values,
-        grid.node_y.values,
-        grid.node_z.values,
-    )
 
     # Loop through each face
     for i in range(grid.n_face):
+        face_edges = _get_cartesian_face_edge_nodes(
+            i, grid.face_node_connectivity.values, grid.n_nodes_per_face.values, grid.node_x.values, grid.node_y.values, grid.node_z.values
+        )
+
         # Set the point as the face center of the polygon
         point_xyz = np.array([grid.face_x[i].values, grid.face_y[i].values, grid.face_z[i].values])
-
         # Assert that the point is in the polygon
-        assert point_in_face(faces_edges_cartesian[i], point_xyz, inclusive=True)
+        assert _face_contains_point(face_edges, point_xyz)
 
 
 def test_point_outside():
@@ -1223,7 +1220,7 @@ def test_point_outside():
     point_xyz = np.array([grid.face_x[1].values, grid.face_y[1].values, grid.face_z[1].values])
 
     # Assert that the point is not in the face tested
-    assert not point_in_face(faces_edges_cartesian[0], point_xyz, inclusive=True)
+    assert not _face_contains_point(faces_edges_cartesian[0], point_xyz)
 
 
 def test_point_on_node():
@@ -1246,10 +1243,7 @@ def test_point_on_node():
     point_xyz = np.array([*faces_edges_cartesian[0][0][0]])
 
     # Assert that the point is in the face when inclusive is true
-    assert point_in_face(faces_edges_cartesian[0], point_xyz, inclusive=True)
-
-    # Assert that the point is not in the face when inclusive is false
-    assert not point_in_face(faces_edges_cartesian[0], point_xyz, inclusive=False)
+    assert _face_contains_point(faces_edges_cartesian[0], point_xyz)
 
 
 def test_point_inside_close():
@@ -1274,7 +1268,7 @@ def test_point_inside_close():
     )
 
     # Use point in face to determine if the point is inside or out of the face
-    assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=False)
+    assert _face_contains_point(faces_edges_cartesian[0], point)
 
 
 def test_point_outside_close():
@@ -1299,7 +1293,7 @@ def test_point_outside_close():
     )
 
     # Use point in face to determine if the point is inside or out of the face
-    assert not point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=False)
+    assert not _face_contains_point(faces_edges_cartesian[0], point)
 
 
 def test_face_at_pole():
@@ -1322,7 +1316,7 @@ def test_face_at_pole():
         grid.node_z.values,
     )
 
-    assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=True)
+    assert _face_contains_point(faces_edges_cartesian[0], point)
 
 
 def test_face_at_antimeridian():
@@ -1344,7 +1338,7 @@ def test_face_at_antimeridian():
         grid.node_z.values,
     )
 
-    assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=True)
+    assert _face_contains_point(faces_edges_cartesian[0], point)
 
 
 def test_face_normal_face():
@@ -1367,7 +1361,7 @@ def test_face_normal_face():
         grid.node_z.values,
     )
 
-    assert point_in_face(faces_edges_cartesian[0], point_xyz=point, inclusive=True)
+    assert _face_contains_point(faces_edges_cartesian[0], point)
 
 
 def test_stereographic_projection_stereographic_projection():

test/test_grid.py

@@ -770,76 +770,10 @@ def test_normalize_existing_coordinates_norm_initial():
     assert _check_normalization(uxgrid)
 
 
-def test_number_of_faces_found():
-    """Test function for `self.get_face_containing_point`,
-    to ensure the correct number of faces is found, depending on where the point is."""
-    grid = ux.open_grid(gridfile_mpas)
-    partial_grid = ux.open_grid(quad_hex_grid_path)
 
-    # 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)) == 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)) == 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)) == 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)) == 2
-
-
-def test_whole_grid():
-    """Tests `self.get_faces_containing_point`on an entire grid,
-    checking that for each face center, one face is found to contain it"""
-
-    grid = ux.open_grid(gridfile_mpas_two)
-    grid.normalize_cartesian_coordinates()
-
-    # Ensure a face is found on the grid for every face center
-    for i in range(len(grid.face_x.values)):
-        point_xyz = np.array([grid.face_x[i].values, grid.face_y[i].values, grid.face_z[i].values], dtype=np.float64)
-
-        assert len(grid.get_faces_containing_point(point_xyz=point_xyz)) == 1
-
-def test_point_types():
-    """Tests that `self.get_faces_containing_point` works with cartesian and lonlat"""
-
-    # Open the grid
-    grid = ux.open_grid(gridfile_mpas)
-
-    # Assign a cartesian point and a lon/lat point
-    point_xyz = np.array([grid.node_x[100].values, grid.node_y[100].values, grid.node_z[100].values], dtype=np.float64)
-    point_lonlat = np.array([grid.node_lon[100].values, grid.node_lat[100].values])
-
-    # Test both points find faces
-    assert len(grid.get_faces_containing_point(point_xyz=point_xyz)) != 0
-    assert len(grid.get_faces_containing_point(point_lonlat=point_lonlat)) !=0
-
-def test_point_along_arc():
-    node_lon = np.array([-40, -40, 40, 40])
-    node_lat = np.array([-20, 20, 20, -20])
-    face_node_connectivity = np.array([[0, 1, 2, 3]], dtype=np.int64)
-
-    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))
-
-    # 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))
-
-    nt.assert_array_equal(out1, out2)
 
 def test_from_topology():
     node_lon = np.array([-20.0, 0.0, 20.0, -20, -40])

test/test_point_in_face.py

@@ -0,0 +1,145 @@
+import os
+from pathlib import Path
+
+import numpy as np
+import numpy.testing as nt
+import pytest
+
+import uxarray as ux
+from uxarray.constants import INT_FILL_VALUE
+from .test_gradient import quad_hex_grid_path
+
+current_path = Path(os.path.dirname(os.path.realpath(__file__)))
+gridfile_mpas = current_path / "meshfiles" / "mpas" / "QU" / "mesh.QU.1920km.151026.nc"
+
+@pytest.fixture(params=["healpix", "quadhex"])
+def grid(request):
+    if request.param == "healpix":
+        g = ux.Grid.from_healpix(zoom=0)
+    else:
+        g = ux.open_grid(quad_hex_grid_path)
+    g.normalize_cartesian_coordinates()
+    return g
+
+def test_face_centers(grid):
+    """
+    For each face of the grid, verify that
+      - querying its Cartesian center returns exactly that face
+      - querying its lon/lat center returns exactly that face
+    """
+    centers_xyz = np.vstack([
+        grid.face_x.values,
+        grid.face_y.values,
+        grid.face_z.values,
+    ]).T
+
+    for fid, center in enumerate(centers_xyz):
+        hits = grid.get_faces_containing_point(
+            point_xyz=center,
+            return_counts=False
+        )
+        assert isinstance(hits, list)
+        assert len(hits) == 1
+        assert hits[0] == [fid]
+
+
+    centers_lonlat = np.vstack([
+        grid.face_lon.values,
+        grid.face_lat.values,
+    ]).T
+
+    for fid, (lon, lat) in enumerate(centers_lonlat):
+        hits = grid.get_faces_containing_point(
+            point_lonlat=(lon, lat),
+            return_counts=False
+        )
+        assert hits[0] == [fid]
+
+def test_node_corners(grid):
+    """
+    For each corner node, verify that querying its coordinate in both
+    Cartesian and spherical (lon/lat) returns exactly the faces sharing it.
+    """
+
+    print(grid.max_face_radius)
+
+    node_coords = np.vstack([
+        grid.node_x.values,
+        grid.node_y.values,
+        grid.node_z.values,
+    ]).T
+
+    conn = grid.node_face_connectivity.values
+    counts = np.sum(conn != INT_FILL_VALUE, axis=1)
+
+    # Cartesian tests
+    for nid, (x, y, z) in enumerate(node_coords):
+        expected = conn[nid, :counts[nid]].tolist()
+
+        hits_xyz = grid.get_faces_containing_point(
+            point_xyz=(x, y, z),
+            return_counts=False
+        )[0]
+        assert set(hits_xyz) == set(expected)
+        assert len(hits_xyz) == len(expected)
+
+
+    node_lonlat = np.vstack([
+        grid.node_lon.values,
+        grid.node_lat.values,
+    ]).T
+
+    # Spherical tests
+    for nid, (lon, lat) in enumerate(node_lonlat):
+        expected = conn[nid, :counts[nid]].tolist()
+
+        hits_ll = grid.get_faces_containing_point(
+            point_lonlat=(lon, lat),
+            return_counts=False
+        )[0]
+        assert set(hits_ll) == set(expected)
+        assert len(hits_ll) == len(expected)
+
+
+def test_number_of_faces_found():
+    """Test function for `self.get_face_containing_point`,
+    to ensure the correct number of faces is found, depending on where the point is."""
+    grid = ux.open_grid(gridfile_mpas)
+    partial_grid = ux.open_grid(quad_hex_grid_path)
+
+    # 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
+
+    # 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
+
+    # 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
+
+    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
+
+def test_point_along_arc():
+    node_lon = np.array([-40, -40, 40, 40])
+    node_lat = np.array([-20, 20, 20, -20])
+    face_node_connectivity = np.array([[0, 1, 2, 3]], dtype=np.int64)
+
+    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)
+
+    # 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)
+
+    nt.assert_array_equal(out1[0], out2[0])

uxarray/grid/coordinates.py

@@ -866,3 +866,21 @@ 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):
+        raise ValueError(
+            "Both Cartesian (xyz) and Spherical (lonlat) coordinates were provided. One one can be "
+            "provided at a time."
+        )
+
+    # 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