diff --git a/test/test_geometry.py b/test/test_geometry.py index f87dcc887..83ba66d6c 100644 --- a/test/test_geometry.py +++ b/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(): diff --git a/test/test_grid.py b/test/test_grid.py index 89e11fcf2..45be30eee 100644 --- a/test/test_grid.py +++ b/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]) diff --git a/test/test_point_in_face.py b/test/test_point_in_face.py new file mode 100644 index 000000000..47c515af7 --- /dev/null +++ b/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]) diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index 25f04e7bb..bef5f1b4d 100644 --- a/uxarray/grid/coordinates.py +++ b/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 diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index e415b3926..7d50fb553 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -7,7 +7,7 @@ import shapely import spatialpandas from matplotlib.collections import LineCollection, PolyCollection -from numba import njit +from numba import njit, prange from shapely import Polygon from shapely import polygons as Polygons from spatialpandas.geometry import MultiPolygonArray, PolygonArray @@ -17,9 +17,6 @@ INT_DTYPE, INT_FILL_VALUE, ) -from uxarray.grid.arcs import ( - point_within_gca, -) from uxarray.grid.coordinates import _xyz_to_lonlat_rad from uxarray.grid.intersections import ( gca_gca_intersection, @@ -1011,200 +1008,78 @@ def inverse_stereographic_projection(x, y, central_lon, central_lat): return lon, lat -@njit(cache=True) -def point_in_face( - edges_xyz, - point_xyz, - inclusive=True, -): - """Determines if a point lies inside a face. - - Parameters - ---------- - edges_xyz : numpy.ndarray - Cartesian coordinates of each point in the face - point_xyz : numpy.ndarray - Cartesian coordinate of the point - inclusive : bool - Flag to determine whether to include points on the nodes and edges of the face - - Returns - ------- - bool - True if point is inside face, False otherwise - """ - - # Validate the inputs - if len(edges_xyz[0][0]) != 3: - raise ValueError("`edges_xyz` vertices must be in Cartesian coordinates.") - - if len(point_xyz) != 3: - raise ValueError("`point_xyz` must be a single [3] Cartesian coordinate.") - - # Initialize the intersection count - intersection_count = 0 - - # Set to hold unique intersections - unique_intersections = set() - - location = _classify_polygon_location(edges_xyz) - - if location == 1: - ref_point_xyz = REF_POINT_SOUTH_XYZ - elif location == -1: - ref_point_xyz = REF_POINT_NORTH_XYZ - else: - ref_point_xyz = REF_POINT_SOUTH_XYZ - - # Initialize the points arc between the point and the reference point - gca_cart = np.empty((2, 3), dtype=np.float64) - gca_cart[0] = point_xyz - gca_cart[1] = ref_point_xyz - - # Loop through the face's edges, checking each one for intersection - for ind in range(len(edges_xyz)): - # If the point lies on an edge, return True if inclusive - if point_within_gca( - point_xyz, - edges_xyz[ind][0], - edges_xyz[ind][1], - ): - if inclusive: - return True - else: - return False - - # Get the number of intersections between the edge and the point arc - intersections = gca_gca_intersection(edges_xyz[ind], gca_cart) - - # Add any unique intersections to the intersection_count - for intersection in intersections: - intersection_tuple = ( - intersection[0], - intersection[1], - intersection[2], - ) - if intersection_tuple not in unique_intersections: - unique_intersections.add(intersection_tuple) - intersection_count += 1 - - # Return True if the number of intersections is odd, False otherwise - return intersection_count % 2 == 1 - - -@njit(cache=True) -def _find_faces(face_edge_cartesian, point_xyz, inverse_indices): - """Finds the faces that contain a given point, inside a subset `face_edge_cartesian` - Parameters - ---------- - face_edge_cartesian : numpy.ndarray - Cartesian coordinates of all the faces according to their edges - point_xyz : numpy.ndarray - Cartesian coordinate of the point - inverse_indices : numpy.ndarray - The original indices of the subsetted grid - - Returns - ------- - index : array - The index of the face that contains the point - """ - - index = [] - - # Run for each face in the subset - for i, face in enumerate(inverse_indices): - # Check to see if the face contains the point - contains_point = point_in_face( - face_edge_cartesian[i], - point_xyz, - inclusive=True, - ) - - # If the point is found, add it to the index array - if contains_point: - index.append(face) +def _populate_max_face_radius(grid): + """Populates `max_face_radius`""" - # Return the index array - return index + face_node_connectivity = grid.face_node_connectivity.values + node_x = grid.node_x.values + node_y = grid.node_y.values + node_z = grid.node_z.values -def _populate_max_face_radius(self): - """Populates `max_face_radius` - - Returns - ------- - max_distance : np.float64 - The max distance from a node to a face center - """ - - # Parse all variables needed for `njit` functions - face_node_connectivity = self.face_node_connectivity.values - node_lats_rad = np.deg2rad(self.node_lat.values) - node_lons_rad = np.deg2rad(self.node_lon.values) - face_lats_rad = np.deg2rad(self.face_lat.values) - face_lons_rad = np.deg2rad(self.face_lon.values) + face_x = grid.face_x.values + face_y = grid.face_y.values + face_z = grid.face_z.values # Get the max distance max_distance = calculate_max_face_radius( face_node_connectivity, - node_lats_rad, - node_lons_rad, - face_lats_rad, - face_lons_rad, + node_x, + node_y, + node_z, + face_x, + face_y, + face_z, ) - # Return the max distance, which is the `max_face_radius` - return np.rad2deg(max_distance) + return max_distance -@njit(cache=True) +@njit(cache=True, parallel=True) def calculate_max_face_radius( - face_node_connectivity, node_lats_rad, node_lons_rad, face_lats_rad, face_lons_rad -): - """Finds the max face radius in the mesh. - Parameters - ---------- - face_node_connectivity : numpy.ndarray - Cartesian coordinates of all the faces according to their edges - node_lats_rad : numpy.ndarray - The `Grid.node_lat` array in radians - node_lons_rad : numpy.ndarray - The `Grid.node_lon` array in radians - face_lats_rad : numpy.ndarray - The `Grid.face_lat` array in radians - face_lons_rad : numpy.ndarray - The `Grid.face_lon` array in radians - - Returns - ------- - The max distance from a node to a face center - """ - - # Array to store all distances of each face to it's furthest node. - end_distances = np.zeros(len(face_node_connectivity)) - - # Loop over each face and its nodes - for ind, face in enumerate(face_node_connectivity): - # Filter out INT_FILL_VALUE - valid_nodes = face[face != INT_FILL_VALUE] - - # Get the face lat/lon of this face - face_lat = face_lats_rad[ind] - face_lon = face_lons_rad[ind] - - # Get the node lat/lon of this face - node_lat_rads = node_lats_rad[valid_nodes] - node_lon_rads = node_lons_rad[valid_nodes] - - # Calculate Haversine distances for all nodes in this face - distances = haversine_distance(node_lon_rads, node_lat_rads, face_lon, face_lat) - - # Store the max distance for this face - end_distances[ind] = np.max(distances) - - # Return the maximum distance found across all faces - return np.max(end_distances) + face_node_connectivity: np.ndarray, + node_x: np.ndarray, + node_y: np.ndarray, + node_z: np.ndarray, + face_x: np.ndarray, + face_y: np.ndarray, + face_z: np.ndarray, +) -> float: + n_faces, n_max_nodes = face_node_connectivity.shape + # workspace to hold the max squared-distance for each face + max2_per_face = np.empty(n_faces, dtype=np.float64) + + # parallel outer loop + for i in prange(n_faces): + fx = face_x[i] + fy = face_y[i] + fz = face_z[i] + + # track the max squared distance for this face + face_max2 = 0.0 + + # loop over all possible node slots + for j in range(n_max_nodes): + idx = face_node_connectivity[i, j] + if idx == INT_FILL_VALUE: + continue + dx = node_x[idx] - fx + dy = node_y[idx] - fy + dz = node_z[idx] - fz + d2 = dx * dx + dy * dy + dz * dz + if d2 > face_max2: + face_max2 = d2 + + max2_per_face[i] = face_max2 + + # now do a simple serial reduction + global_max2 = 0.0 + for i in range(n_faces): + if max2_per_face[i] > global_max2: + global_max2 = max2_per_face[i] + + # one sqrt at the end + return math.sqrt(global_max2) @njit(cache=True) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index e86dcf1a0..992fe9281 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -4,6 +4,7 @@ from typing import ( List, Optional, + Sequence, Set, Tuple, Union, @@ -18,7 +19,7 @@ from xarray.core.options import OPTIONS from xarray.core.utils import UncachedAccessor -from uxarray.constants import ERROR_TOLERANCE, INT_FILL_VALUE +from uxarray.constants import INT_FILL_VALUE from uxarray.conventions import ugrid from uxarray.cross_sections import GridCrossSectionAccessor from uxarray.formatting_html import grid_repr @@ -33,20 +34,18 @@ _populate_node_face_connectivity, ) from uxarray.grid.coordinates import ( - _lonlat_rad_to_xyz, _populate_edge_centroids, _populate_face_centerpoints, _populate_face_centroids, _populate_node_latlon, _populate_node_xyz, + _prepare_points_for_kdtree, _set_desired_longitude_range, - _xyz_to_lonlat_deg, prepare_points, ) from uxarray.grid.dual import construct_dual from uxarray.grid.geometry import ( _construct_boundary_edge_indices, - _find_faces, _grid_to_matplotlib_linecollection, _grid_to_matplotlib_polycollection, _grid_to_polygon_geodataframe, @@ -68,7 +67,7 @@ _populate_edge_face_distances, _populate_edge_node_distances, ) -from uxarray.grid.utils import _get_cartesian_face_edge_nodes_array +from uxarray.grid.point_in_face import _point_in_face_query from uxarray.grid.validation import ( _check_area, _check_connectivity, @@ -135,7 +134,7 @@ class Grid: A dataset of indices that correspond to the original grid, if the grid being constructed is a subset Examples - ---------- + -------- >>> import uxarray as ux >>> grid_path = "/path/to/grid.nc" @@ -1574,7 +1573,7 @@ def is_subset(self): @property def max_face_radius(self): - """Maximum face radius of the grid in degrees""" + """Maximum Euclidean distance from each face center to its nodes.""" if "max_face_radius" not in self._ds: self._ds["max_face_radius"] = _populate_max_face_radius(self) return self._ds["max_face_radius"] @@ -2566,116 +2565,85 @@ def get_faces_between_latitudes(self, lats: Tuple[float, float]): return faces_within_lat_bounds(lats, self.face_bounds_lat.values) def get_faces_containing_point( - self, point_xyz=None, point_lonlat=None, tolerance=ERROR_TOLERANCE - ): - """Identifies the indices of faces that contain a given point. + self, + *, + point_lonlat: Sequence[float] | np.ndarray = None, + point_xyz: Sequence[float] | np.ndarray = None, + 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. Parameters ---------- - point_xyz : numpy.ndarray - A point in cartesian coordinates. Best performance if - point_lonlat : numpy.ndarray - A point in spherical coordinates. - tolerance : numpy.ndarray - An optional error tolerance for points that lie on the nodes of a face. + 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). + 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. Returns ------- - index : numpy.ndarray - Array of the face indices containing point. Empty if no face is found. This function will typically return - a single face, unless the point falls directly on a corner or edge, where there will be multiple values. + If `return_counts=True`: + face_indices : np.ndarray, shape (N, M) or (N,) + - 2D array of face indices when `counts` contains values >1 or + when mixed counts exist; unused slots are filled with `INT_FILL_VALUE`. + - If **all** `counts == 1`, this will be squeezed to a 1-D array of shape `(N,)`. + counts : np.ndarray, shape (N,) + 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. + + 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`. Examples -------- - Open a grid from a file path: - >>> import uxarray as ux - >>> uxgrid = ux.open_grid("grid_filename.nc") - - Define a spherical point: - - >>> import numpy as np - >>> point_lonlat = np.array([45.2, 32.6], dtype=np.float64) + >>> grid_path = "/path/to/grid.nc" + >>> uxgrid = ux.open_grid(grid_path) - Define a cartesian point: + # 1. Query a Spherical (lon/lat) point + >>> indices, counts = uxgrid.get_faces_containing_point(point_lonlat=(0.0, 0.0)) - >>> point_xyz = np.array([0.0, 0.0, 1.0], dtype=np.float64) - - Find the indices of the faces that contain the given point: - - >>> lonlat_point_face_indices = uxgrid.get_faces_containing_point( - ... point_lonlat=point_lonlat - ... ) - >>> xyz_point_face_indices = uxgrid.get_faces_containing_point( - ... point_xyz=point_xyz + # 2. Query a Cartesian (xyz) point + >>> indices, counts = uxgrid.get_faces_containing_point( + ... point_xyz=(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 + ... ) """ - if point_xyz is None and point_lonlat is None: - raise ValueError("Either `point_xyz` or `point_lonlat` must be passed in.") - # Depending on the provided point coordinates, convert to get all needed coordinate systems - if point_xyz is None: - point_lonlat = np.asarray(point_lonlat, dtype=np.float64) - point_xyz = np.array( - _lonlat_rad_to_xyz(*np.deg2rad(point_lonlat)), dtype=np.float64 - ) - elif point_lonlat is None: - point_xyz = np.asarray(point_xyz, dtype=np.float64) - point_lonlat = np.array(_xyz_to_lonlat_deg(*point_xyz), dtype=np.float64) - - # Get the maximum face radius of the grid, plus a small adjustment for if the point is this exact radius away - max_face_radius = self.max_face_radius.values + 0.0001 - - # Try to find a subset in which the point resides - try: - subset = self.subset.bounding_circle( - r=max_face_radius, - center_coord=point_lonlat, - element="face centers", - inverse_indices=True, - ) - # If no subset is found, warn the user - except ValueError: - # If the grid is partial, let the user know the point likely lies outside the grid region - if self.partial_sphere_coverage: - warn( - "No faces found. The grid has partial spherical coverage, and the point may be outside the defined region of the grid." - ) - else: - warn("No faces found. Try adjusting the tolerance.") - return np.empty(0, dtype=np.int64) - - # Get the faces in terms of their edges - face_edge_nodes_xyz = _get_cartesian_face_edge_nodes_array( - subset.face_node_connectivity.values, - subset.n_face, - subset.n_max_face_nodes, - subset.node_x.values, - subset.node_y.values, - subset.node_z.values, - ) + pts = _prepare_points_for_kdtree(point_lonlat, point_xyz) - # Get the original face indices from the subset - inverse_indices = subset.inverse_indices.face.values + # Determine faces containing points + face_indices, counts = _point_in_face_query(source_grid=self, points=pts) - # Check to see if the point is on the nodes of any face - lies_on_node = np.isclose( - face_edge_nodes_xyz, - point_xyz[None, None, :], # Expands dimensions for broadcasting - rtol=tolerance, - atol=tolerance, - ) - - edge_matches = np.all(lies_on_node, axis=-1) - face_matches = np.any(edge_matches, axis=1) - face_indices = inverse_indices[np.any(face_matches, axis=1)] + # Return a list of lists if counts are not desired + if not return_counts: + output: List[List[int]] = [] + for i, c in enumerate(counts): + output.append(face_indices[i, :c].tolist()) + return output - # If a face is in face_indices, return that as the point was found to lie on a node - if len(face_indices) != 0: - return face_indices - else: - # Check if any of the faces in the subset contain the point - face_indices = _find_faces(face_edge_nodes_xyz, point_xyz, inverse_indices) + if (counts == 1).all(): + # collapse to a 1D array of length n_points + face_indices = face_indices[:, 0] - return face_indices + return face_indices, counts diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py new file mode 100644 index 000000000..9d5e93f1a --- /dev/null +++ b/uxarray/grid/point_in_face.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import numpy as np +from numba import njit, prange + +from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE +from uxarray.grid.arcs import point_within_gca +from uxarray.grid.utils import _get_cartesian_face_edge_nodes, _small_angle_of_2_vectors + +if TYPE_CHECKING: + from numpy.typing import ArrayLike + + from uxarray.grid.grid import Grid + + +@njit(cache=True) +def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: + """ + Determine whether a point lies within a face using the spherical winding-number method. + + This function sums the signed central angles between successive vertices of the face + as seen from `point`. If the total absolute winding exceeds π, the point is inside. + Points exactly on a node or edge also count as inside. + + Parameters + ---------- + face_edges : np.ndarray, shape (n_edges, 2, 3) + Cartesian coordinates (unit-vectors) of each great-circle edge of the face. + Each row is [start_xyz, end_xyz]. + point : np.ndarray, shape (3,) + 3D unit-vector of the query point on the unit sphere. + + Returns + ------- + inside : bool + True if the point is inside the face or lies exactly on a node/edge; False otherwise. + """ + # Check for an exact hit with any of the corner nodes + for e in range(face_edges.shape[0]): + if np.allclose( + face_edges[e, 0], point, rtol=ERROR_TOLERANCE, atol=ERROR_TOLERANCE + ): + return True + if np.allclose( + face_edges[e, 1], point, rtol=ERROR_TOLERANCE, atol=ERROR_TOLERANCE + ): + return True + if point_within_gca(point, face_edges[e, 0], face_edges[e, 1]): + return True + + n = face_edges.shape[0] + + total = 0.0 + p = point + for i in range(n): + a = face_edges[i, 0] + b = face_edges[i + 1, 0] if i + 1 < n else face_edges[0, 0] + + vi = a - p + vj = b - p + + # check if you’re right on a vertex + if np.linalg.norm(vi) < ERROR_TOLERANCE or np.linalg.norm(vj) < ERROR_TOLERANCE: + return True + + ang = _small_angle_of_2_vectors(vi, vj) + + # determine sign from cross + c = np.cross(vi, vj) + sign = 1.0 if (c[0] * p[0] + c[1] * p[1] + c[2] * p[2]) >= 0.0 else -1.0 + + total += sign * ang + + return np.abs(total) > np.pi + + +@njit(cache=True) +def _get_faces_containing_point( + point: np.ndarray, + candidate_indices: np.ndarray, + face_node_connectivity: np.ndarray, + n_nodes_per_face: np.ndarray, + node_x: np.ndarray, + node_y: np.ndarray, + node_z: np.ndarray, +) -> np.ndarray: + """ + Test each candidate face to see if it contains the query point. + + Parameters + ---------- + point : np.ndarray, shape (3,) + Cartesian unit-vector of the query point. + candidate_indices : np.ndarray, shape (k,) + Array of face indices to test (e.g., from a k-d tree cull). + face_node_connectivity : np.ndarray, shape (n_faces, max_nodes) + Node connectivity (node indices) for each face. + n_nodes_per_face : np.ndarray, shape (n_faces,) + Number of valid nodes per face. + node_x, node_y, node_z : np.ndarray, shape (n_nodes,) + Cartesian coordinates of each grid node. + + Returns + ------- + hits : np.ndarray, shape (h,) + Subset of `candidate_indices` for which the point is inside the face. + """ + hit_buf = np.empty(candidate_indices.shape[0], dtype=INT_DTYPE) + count = 0 + for k in range(candidate_indices.shape[0]): + fidx = candidate_indices[k] + face_edges = _get_cartesian_face_edge_nodes( + fidx, face_node_connectivity, n_nodes_per_face, node_x, node_y, node_z + ) + if _face_contains_point(face_edges, point): + hit_buf[count] = fidx + count += 1 + return hit_buf[:count] + + +@njit(cache=True, parallel=True) +def _batch_point_in_face( + points: np.ndarray, + flat_candidate_indices: np.ndarray, + offsets: np.ndarray, + face_node_connectivity: np.ndarray, + n_nodes_per_face: np.ndarray, + node_x: np.ndarray, + node_y: np.ndarray, + node_z: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: + """ + Parallel entry-point: for each point, test all candidate faces in batch. + + Parameters + ---------- + points : np.ndarray, shape (n_points, 3) + Cartesian coordinates of query points. + flat_candidate_indices : np.ndarray, shape (total_candidates,) + Flattened array of all candidate face indices across points. + offsets : np.ndarray, shape (n_points + 1,) + Offsets into `flat_candidate_indices` marking each point’s slice. + face_node_connectivity, n_nodes_per_face, node_x, node_y, node_z + As in `_get_faces_containing_point`. + + Returns + ------- + results : np.ndarray, shape (n_points, max_candidates) + Each row lists face indices containing the corresponding point; + unused entries are filled with `INT_FILL_VALUE`. + counts : np.ndarray, shape (n_points,) + Number of valid face-hits per point. + """ + n_points = offsets.shape[0] - 1 + results = np.full( + (n_points, face_node_connectivity.shape[1]), INT_FILL_VALUE, dtype=INT_DTYPE + ) + counts = np.zeros(n_points, dtype=INT_DTYPE) + + for i in prange(n_points): + start = offsets[i] + end = offsets[i + 1] + p = points[i] + cands = flat_candidate_indices[start:end] + + hits = _get_faces_containing_point( + p, cands, face_node_connectivity, n_nodes_per_face, node_x, node_y, node_z + ) + for j, fi in enumerate(hits): + results[i, j] = fi + counts[i] = hits.shape[0] + + return results, counts + + +def _point_in_face_query( + source_grid: Grid, points: ArrayLike +) -> Tuple[np.ndarray, np.ndarray]: + """ + Find grid faces that contain given Cartesian point(s) on the unit sphere. + + This function first uses a SciPy k-d tree (in Cartesian space) to cull + candidate faces within the maximum face‐radius, then calls the Numba‐ + accelerated winding‐number tester in parallel. + + Parameters + ---------- + source_grid : Grid + UXarray Grid object, which must provide: + - ._get_scipy_kd_tree(): a SciPy cKDTree over node‐centroids, + - .max_face_radius: float search radius, + - .face_node_connectivity, .n_nodes_per_face, .node_x, .node_y, .node_z + arrays for reconstructing face edges. + points : array_like, shape (3,) or (n_points, 3) + Cartesian coordinates of query point(s). + + Returns + ------- + face_indices : np.ndarray, shape (n_points, max_candidates) + 2D array of face indices containing each point; unused entries are + padded with `INT_FILL_VALUE`. + counts : np.ndarray, shape (n_points,) + Number of valid face‐hits per point. + """ + pts = np.asarray(points, dtype=np.float64) + if pts.ndim == 1: + pts = pts[np.newaxis, :] + # Cull with k-d tree + kdt = source_grid._get_scipy_kd_tree() + radius = source_grid.max_face_radius * 1.05 + cand_lists = kdt.query_ball_point(x=pts, r=radius, workers=-1) + + # Prepare flattened candidates and offsets + flat_cands = np.concatenate([np.array(lst, dtype=np.int64) for lst in cand_lists]) + lens = np.array([len(lst) for lst in cand_lists], dtype=np.int64) + offs = np.empty(len(lens) + 1, dtype=np.int64) + offs[0] = 0 + np.cumsum(lens, out=offs[1:]) + + # Perform the batch winding‐number test + return _batch_point_in_face( + pts, + flat_cands, + offs, + source_grid.face_node_connectivity.values, + source_grid.n_nodes_per_face.values, + source_grid.node_x.values, + source_grid.node_y.values, + source_grid.node_z.values, + )