From 28e983fc8167f864f81c8138b94e7cde77178466 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Mon, 12 May 2025 13:44:43 -0600 Subject: [PATCH 01/13] initial refactor --- uxarray/grid/geometry.py | 82 ++++++++++++++--------------------- uxarray/grid/grid.py | 5 ++- uxarray/grid/point_in_face.py | 21 +++++++++ 3 files changed, 58 insertions(+), 50 deletions(-) create mode 100644 uxarray/grid/point_in_face.py diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index e415b3926..508d0dbb9 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -1129,79 +1129,63 @@ def _find_faces(face_edge_cartesian, point_xyz, inverse_indices): return index -def _populate_max_face_radius(self): - """Populates `max_face_radius` +def _populate_max_face_radius(grid): + """Populates `max_face_radius`""" - Returns - ------- - max_distance : np.float64 - The max distance from a node to a face center - """ + face_node_connectivity = grid.face_node_connectivity.values - # 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) + node_x = grid.node_x.values + node_y = grid.node_y.values + node_z = grid.node_z.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) def calculate_max_face_radius( - face_node_connectivity, node_lats_rad, node_lons_rad, face_lats_rad, face_lons_rad + face_node_connectivity, node_x, node_y, node_z, face_x, face_y, face_z ): - """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 - """ + """Finds the max face radius in the grid.""" # 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): + for face_idx, cur_face in enumerate(face_node_connectivity): # Filter out INT_FILL_VALUE - valid_nodes = face[face != INT_FILL_VALUE] + node_indices = cur_face[cur_face != INT_FILL_VALUE] - # Get the face lat/lon of this face - face_lat = face_lats_rad[ind] - face_lon = face_lons_rad[ind] + fx = face_x[face_idx] + fy = face_y[face_idx] + fz = face_z[face_idx] - # Get the node lat/lon of this face - node_lat_rads = node_lats_rad[valid_nodes] - node_lon_rads = node_lons_rad[valid_nodes] + nx = node_x[node_indices] + ny = node_y[node_indices] + nz = node_z[node_indices] - # Calculate Haversine distances for all nodes in this face - distances = haversine_distance(node_lon_rads, node_lat_rads, face_lon, face_lat) + # vectorized distance calculation + dx = nx - fx + dy = ny - fy + dz = nz - fz + distances = np.sqrt(dx * dx + dy * dy + dz * dz) # Store the max distance for this face - end_distances[ind] = np.max(distances) + end_distances[face_idx] = np.max(distances) # Return the maximum distance found across all faces return np.max(end_distances) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index a7669265e..e75c0dd7a 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1574,7 +1574,7 @@ def is_subset(self): @property def max_face_radius(self): - """Maximum face radius of the grid in degrees""" + """TODO: Maximum face radius of the grid in degrees""" if "max_face_radius" not in self._ds: self._ds["max_face_radius"] = _populate_max_face_radius(self) return self._ds["max_face_radius"] @@ -2624,8 +2624,11 @@ def get_faces_containing_point( 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 + # TODO: Use a smaller eps max_face_radius = self.max_face_radius.values + 0.0001 + # tree = self._get_scipy_kd_tree() + # Try to find a subset in which the point resides try: subset = self.subset.bounding_circle( diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py new file mode 100644 index 000000000..9bde094f4 --- /dev/null +++ b/uxarray/grid/point_in_face.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +# from typing import TYPE_CHECKING +import numpy as np + +# import xarray as xr +# +# if TYPE_CHECKING: +# from uxarray.core.dataarray import UxDataArray +# from uxarray.core.dataset import UxDataset +# from uxarray.grid.grid import Grid + + +def _candidate_faces(): + pass + + +def _point_in_face(point: np.ndarray): + # + + pass From bd37fa30c115ab4ea91c62e042956320fd084b21 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Mon, 12 May 2025 20:56:45 -0600 Subject: [PATCH 02/13] optimize point in face --- test/test_grid.py | 66 --------- test/test_point_in_face.py | 142 +++++++++++++++++++ uxarray/grid/geometry.py | 2 + uxarray/grid/grid.py | 156 +++++++------------- uxarray/grid/point_in_face.py | 258 ++++++++++++++++++++++++++++++++-- 5 files changed, 442 insertions(+), 182 deletions(-) create mode 100644 test/test_point_in_face.py 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..e32f558c8 --- /dev/null +++ b/test/test_point_in_face.py @@ -0,0 +1,142 @@ +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. + """ + 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 (as before)… + 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 + + 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/geometry.py b/uxarray/grid/geometry.py index 508d0dbb9..d75a2a9f6 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -25,6 +25,8 @@ gca_gca_intersection, ) +from uxarray.grid.utils import _get_cartesian_face_edge_nodes + POLE_POINTS_XYZ = { "North": np.array([0.0, 0.0, 1.0]), "South": np.array([0.0, 0.0, -1.0]), diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index e75c0dd7a..fdcc35345 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -7,6 +7,7 @@ Set, Tuple, Union, +Sequence, ) from warnings import warn @@ -104,6 +105,7 @@ from uxarray.plot.accessor import GridPlotAccessor from uxarray.subset import GridSubsetAccessor from uxarray.utils.numba import is_numba_function_cached +from uxarray.grid.point_in_face import _point_in_face_query class Grid: @@ -2565,119 +2567,65 @@ 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: Optional[Sequence[float]] = None, + point_xyz: Optional[Sequence[float]] = None, + tolerance: float = ERROR_TOLERANCE, + return_counts: bool = True, + ) -> Union[Tuple[np.ndarray, np.ndarray], List[List[int]]]: + """ + Identify which faces (cells) 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). + tolerance : float, default=ERROR_TOLERANCE + Distance tolerance for classifying points on edges or nodes. + 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. - - 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) - - Define a cartesian point: - - >>> 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 - ... ) - - """ - 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 `return_counts=True`: + face_indices : np.ndarray, shape (N, M) + 2D array of face indices; unused slots are filled with `INT_FILL_VALUE`. + 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. + """ + # 1) validate + if (point_lonlat is None) == (point_xyz is None): + raise ValueError("Exactly one of `point_lonlat` or `point_xyz` must be provided.") + + # 2) convert to cartesian 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 - # TODO: Use a smaller eps - max_face_radius = self.max_face_radius.values + 0.0001 - - # tree = self._get_scipy_kd_tree() - - # 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, - ) + lon, lat = map(np.deg2rad, point_lonlat) + point_xyz = _lonlat_rad_to_xyz(lon, lat) + pts = np.asarray(point_xyz, dtype=np.float64) + if pts.ndim == 1: + pts = pts[np.newaxis, :] - # Get the original face indices from the subset - inverse_indices = subset.inverse_indices.face.values + # 3) run the Numba query + 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, - ) + # 4) if caller only wants the indices, strip out fill‐values + if not return_counts: + output: List[List[int]] = [] + for i, c in enumerate(counts): + output.append(face_indices[i, :c].tolist()) + return output - 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 face_indices, counts - # 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) - return face_indices diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index 9bde094f4..19735eab2 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -1,21 +1,255 @@ from __future__ import annotations -# from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Tuple import numpy as np -# import xarray as xr -# -# if TYPE_CHECKING: -# from uxarray.core.dataarray import UxDataArray -# from uxarray.core.dataset import UxDataset -# from uxarray.grid.grid import Grid +from numba import njit, prange +from uxarray.grid.utils import _get_cartesian_face_edge_nodes +from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE +from uxarray.grid.arcs import point_within_gca -def _candidate_faces(): - pass +if TYPE_CHECKING: + from uxarray.grid.grid import Grid + from numpy.typing import ArrayLike -def _point_in_face(point: np.ndarray): - # +@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. - pass + 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. + """ + # 1a) quick tests for exact node‐ or edge‐hits + 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 + + # 1b) build closed loop of vertices + n = face_edges.shape[0] + verts = np.empty((n + 1, 3), dtype=np.float64) + for i in range(n): + verts[i] = face_edges[i, 0] + verts[n] = face_edges[0, 0] + + # 2) accumulate signed angle + total = 0.0 + p = point + for i in range(n): + vi = verts[i] - p + vj = verts[i+1] - p + + ni = np.linalg.norm(vi) + nj = np.linalg.norm(vj) + # degenerate on-vertex case + if ni < ERROR_TOLERANCE or nj < ERROR_TOLERANCE: + return True + + ui = vi / ni + uj = vj / nj + + cosang = ui.dot(uj) + # clamp for numerical safety + if cosang > 1.0: cosang = 1.0 + if cosang < -1.0: cosang = -1.0 + ang = np.arccos(cosang) + + # sign = sign of (ui × uj)·p + cx = ui[1]*uj[2] - ui[2]*uj[1] + cy = ui[2]*uj[0] - ui[0]*uj[2] + cz = ui[0]*uj[1] - ui[1]*uj[0] + sign = 1.0 if (cx*p[0] + cy*p[1] + cz*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. + """ + hits = [] + for idx in candidate_indices: + face_edges = _get_cartesian_face_edge_nodes( + idx, + face_node_connectivity, + n_nodes_per_face, + node_x, node_y, node_z + ) + if _face_contains_point(face_edges, point): + hits.append(idx) + return np.asarray(hits, dtype=INT_DTYPE) + + +@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, :] + + # 1) cull with k-d tree + kdt = source_grid._get_scipy_kd_tree() + radius = source_grid.max_face_radius + cand_lists = kdt.query_ball_point(x=pts, r=radius) + + # 2) prepare flattened candidates + 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:]) + + # 3) 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, + ) From 980c3abf7949e80d2aad9e2f1f1f14b6e0d8ca27 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Mon, 12 May 2025 21:17:14 -0600 Subject: [PATCH 03/13] use angle function --- test/test_point_in_face.py | 1 - uxarray/grid/geometry.py | 2 - uxarray/grid/grid.py | 25 ++++---- uxarray/grid/point_in_face.py | 109 ++++++++++++++-------------------- 4 files changed, 54 insertions(+), 83 deletions(-) diff --git a/test/test_point_in_face.py b/test/test_point_in_face.py index e32f558c8..318baaac7 100644 --- a/test/test_point_in_face.py +++ b/test/test_point_in_face.py @@ -139,4 +139,3 @@ def test_point_along_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/geometry.py b/uxarray/grid/geometry.py index d75a2a9f6..508d0dbb9 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -25,8 +25,6 @@ gca_gca_intersection, ) -from uxarray.grid.utils import _get_cartesian_face_edge_nodes - POLE_POINTS_XYZ = { "North": np.array([0.0, 0.0, 1.0]), "South": np.array([0.0, 0.0, -1.0]), diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index fdcc35345..baa09dfc2 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -4,10 +4,10 @@ from typing import ( List, Optional, + Sequence, Set, Tuple, Union, -Sequence, ) from warnings import warn @@ -41,13 +41,11 @@ _populate_node_latlon, _populate_node_xyz, _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, @@ -69,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, @@ -105,7 +103,6 @@ from uxarray.plot.accessor import GridPlotAccessor from uxarray.subset import GridSubsetAccessor from uxarray.utils.numba import is_numba_function_cached -from uxarray.grid.point_in_face import _point_in_face_query class Grid: @@ -2567,12 +2564,12 @@ 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_lonlat: Optional[Sequence[float]] = None, - point_xyz: Optional[Sequence[float]] = None, - tolerance: float = ERROR_TOLERANCE, - return_counts: bool = True, + self, + *, + point_lonlat: Optional[Sequence[float]] = None, + point_xyz: Optional[Sequence[float]] = None, + tolerance: float = ERROR_TOLERANCE, + return_counts: bool = True, ) -> Union[Tuple[np.ndarray, np.ndarray], List[List[int]]]: """ Identify which faces (cells) on the grid contain the given point(s). @@ -2606,7 +2603,9 @@ def get_faces_containing_point( """ # 1) validate if (point_lonlat is None) == (point_xyz is None): - raise ValueError("Exactly one of `point_lonlat` or `point_xyz` must be provided.") + raise ValueError( + "Exactly one of `point_lonlat` or `point_xyz` must be provided." + ) # 2) convert to cartesian if point_xyz is None: @@ -2627,5 +2626,3 @@ def get_faces_containing_point( return output return face_indices, counts - - diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index 19735eab2..576752b6d 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -1,18 +1,19 @@ from __future__ import annotations from typing import TYPE_CHECKING, Tuple -import numpy as np +import numpy as np from numba import njit, prange -from uxarray.grid.utils import _get_cartesian_face_edge_nodes 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 uxarray.grid.grid import Grid 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: @@ -36,70 +37,57 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: inside : bool True if the point is inside the face or lies exactly on a node/edge; False otherwise. """ - # 1a) quick tests for exact node‐ or edge‐hits + # 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): + 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): + 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]): + if point_within_gca(point, face_edges[e, 0], face_edges[e, 1]): return True - # 1b) build closed loop of vertices + # Build a closed loop of vertices n = face_edges.shape[0] verts = np.empty((n + 1, 3), dtype=np.float64) for i in range(n): verts[i] = face_edges[i, 0] verts[n] = face_edges[0, 0] - # 2) accumulate signed angle total = 0.0 p = point for i in range(n): - vi = verts[i] - p - vj = verts[i+1] - p + vi = verts[i] - p + vj = verts[i + 1] - p - ni = np.linalg.norm(vi) - nj = np.linalg.norm(vj) - # degenerate on-vertex case - if ni < ERROR_TOLERANCE or nj < ERROR_TOLERANCE: + # check if you’re right on a vertex + if np.linalg.norm(vi) < ERROR_TOLERANCE or np.linalg.norm(vj) < ERROR_TOLERANCE: return True - ui = vi / ni - uj = vj / nj - - cosang = ui.dot(uj) - # clamp for numerical safety - if cosang > 1.0: cosang = 1.0 - if cosang < -1.0: cosang = -1.0 - ang = np.arccos(cosang) + ang = _small_angle_of_2_vectors(vi, vj) - # sign = sign of (ui × uj)·p - cx = ui[1]*uj[2] - ui[2]*uj[1] - cy = ui[2]*uj[0] - ui[0]*uj[2] - cz = ui[0]*uj[1] - ui[1]*uj[0] - sign = 1.0 if (cx*p[0] + cy*p[1] + cz*p[2]) >= 0.0 else -1.0 + # 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: + 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. @@ -124,10 +112,7 @@ def _get_faces_containing_point( hits = [] for idx in candidate_indices: face_edges = _get_cartesian_face_edge_nodes( - idx, - face_node_connectivity, - n_nodes_per_face, - node_x, node_y, node_z + idx, face_node_connectivity, n_nodes_per_face, node_x, node_y, node_z ) if _face_contains_point(face_edges, point): hits.append(idx) @@ -136,15 +121,15 @@ def _get_faces_containing_point( @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]: + 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. @@ -169,9 +154,7 @@ def _batch_point_in_face( """ n_points = offsets.shape[0] - 1 results = np.full( - (n_points, face_node_connectivity.shape[1]), - INT_FILL_VALUE, - dtype=INT_DTYPE + (n_points, face_node_connectivity.shape[1]), INT_FILL_VALUE, dtype=INT_DTYPE ) counts = np.zeros(n_points, dtype=INT_DTYPE) @@ -182,11 +165,7 @@ def _batch_point_in_face( 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 + 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 @@ -195,11 +174,9 @@ def _batch_point_in_face( return results, counts - def _point_in_face_query( - source_grid: Grid, - points: ArrayLike - ) -> Tuple[np.ndarray, np.ndarray]: + source_grid: Grid, points: ArrayLike +) -> Tuple[np.ndarray, np.ndarray]: """ Find grid faces that contain given Cartesian point(s) on the unit sphere. From 7c61a614c93b061441726f1539313f6386e62504 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Mon, 12 May 2025 21:46:56 -0600 Subject: [PATCH 04/13] optimize logic --- uxarray/grid/grid.py | 2 +- uxarray/grid/point_in_face.py | 32 ++++++++++++++++++++------------ 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index baa09dfc2..977fcfea8 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1573,7 +1573,7 @@ def is_subset(self): @property def max_face_radius(self): - """TODO: 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"] diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index 576752b6d..ccf4d877c 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -50,18 +50,23 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: if point_within_gca(point, face_edges[e, 0], face_edges[e, 1]): return True - # Build a closed loop of vertices + # # Build a closed loop of vertices n = face_edges.shape[0] - verts = np.empty((n + 1, 3), dtype=np.float64) - for i in range(n): - verts[i] = face_edges[i, 0] - verts[n] = face_edges[0, 0] + # verts = np.empty((n + 1, 3), dtype=np.float64) + # for i in range(n): + # verts[i] = face_edges[i, 0] + # verts[n] = face_edges[0, 0] total = 0.0 p = point for i in range(n): - vi = verts[i] - p - vj = verts[i + 1] - p + 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 + # vi = verts[i] - p + # vj = verts[i + 1] - p # check if you’re right on a vertex if np.linalg.norm(vi) < ERROR_TOLERANCE or np.linalg.norm(vj) < ERROR_TOLERANCE: @@ -109,14 +114,17 @@ def _get_faces_containing_point( hits : np.ndarray, shape (h,) Subset of `candidate_indices` for which the point is inside the face. """ - hits = [] - for idx in candidate_indices: + 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( - idx, face_node_connectivity, n_nodes_per_face, node_x, node_y, node_z + fidx, face_node_connectivity, n_nodes_per_face, node_x, node_y, node_z ) if _face_contains_point(face_edges, point): - hits.append(idx) - return np.asarray(hits, dtype=INT_DTYPE) + hit_buf[count] = fidx + count += 1 + return hit_buf[:count] @njit(cache=True, parallel=True) From efe6d3f6ae67cfd58e20c5ab25bdf71774b7d227 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Tue, 13 May 2025 08:19:39 -0600 Subject: [PATCH 05/13] update comments --- test/test_point_in_face.py | 5 +++-- uxarray/grid/grid.py | 14 +++++--------- uxarray/grid/point_in_face.py | 20 ++++++++++---------- 3 files changed, 18 insertions(+), 21 deletions(-) diff --git a/test/test_point_in_face.py b/test/test_point_in_face.py index 318baaac7..398aec765 100644 --- a/test/test_point_in_face.py +++ b/test/test_point_in_face.py @@ -66,10 +66,10 @@ def test_node_corners(grid): grid.node_z.values, ]).T - conn = grid.node_face_connectivity.values + conn = grid.node_face_connectivity.values counts = np.sum(conn != INT_FILL_VALUE, axis=1) - # Cartesian tests (as before)… + # Cartesian tests for nid, (x, y, z) in enumerate(node_coords): expected = conn[nid, :counts[nid]].tolist() @@ -86,6 +86,7 @@ def test_node_corners(grid): grid.node_lat.values, ]).T + # Spherical tests for nid, (lon, lat) in enumerate(node_lonlat): expected = conn[nid, :counts[nid]].tolist() diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 977fcfea8..52ed162bc 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -19,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 @@ -2568,11 +2568,10 @@ def get_faces_containing_point( *, point_lonlat: Optional[Sequence[float]] = None, point_xyz: Optional[Sequence[float]] = None, - tolerance: float = ERROR_TOLERANCE, return_counts: bool = True, ) -> Union[Tuple[np.ndarray, np.ndarray], List[List[int]]]: """ - Identify which faces (cells) on the grid contain the given point(s). + Identify which faces on the grid contain the given point(s). Exactly one of `point_lonlat` or `point_xyz` must be provided. @@ -2582,8 +2581,6 @@ def get_faces_containing_point( Longitude and latitude in **degrees**: (lon, lat). point_xyz : array_like of shape (3,), optional Cartesian coordinates on the unit sphere: (x, y, z). - tolerance : float, default=ERROR_TOLERANCE - Distance tolerance for classifying points on edges or nodes. 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. @@ -2601,13 +2598,12 @@ def get_faces_containing_point( A Python list of length `N`, where each element is the list of face-indices (no padding) for that query point. """ - # 1) validate if (point_lonlat is None) == (point_xyz is None): raise ValueError( "Exactly one of `point_lonlat` or `point_xyz` must be provided." ) - # 2) convert to cartesian + # Convert to cartesian if points are spherical if point_xyz is None: lon, lat = map(np.deg2rad, point_lonlat) point_xyz = _lonlat_rad_to_xyz(lon, lat) @@ -2615,10 +2611,10 @@ def get_faces_containing_point( if pts.ndim == 1: pts = pts[np.newaxis, :] - # 3) run the Numba query + # Determine faces containing points face_indices, counts = _point_in_face_query(source_grid=self, points=pts) - # 4) if caller only wants the indices, strip out fill‐values + # Return a list of lists if counts are not desired if not return_counts: output: List[List[int]] = [] for i, c in enumerate(counts): diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index ccf4d877c..f4e37e934 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -1,5 +1,6 @@ from __future__ import annotations +import time from typing import TYPE_CHECKING, Tuple import numpy as np @@ -50,12 +51,7 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: if point_within_gca(point, face_edges[e, 0], face_edges[e, 1]): return True - # # Build a closed loop of vertices n = face_edges.shape[0] - # verts = np.empty((n + 1, 3), dtype=np.float64) - # for i in range(n): - # verts[i] = face_edges[i, 0] - # verts[n] = face_edges[0, 0] total = 0.0 p = point @@ -65,8 +61,6 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: vi = a - p vj = b - p - # vi = verts[i] - p - # vj = verts[i + 1] - p # check if you’re right on a vertex if np.linalg.norm(vi) < ERROR_TOLERANCE or np.linalg.norm(vj) < ERROR_TOLERANCE: @@ -214,18 +208,24 @@ def _point_in_face_query( pts = np.asarray(points, dtype=np.float64) if pts.ndim == 1: pts = pts[np.newaxis, :] - + start = time.time() # 1) cull with k-d tree kdt = source_grid._get_scipy_kd_tree() radius = source_grid.max_face_radius - cand_lists = kdt.query_ball_point(x=pts, r=radius) + cand_lists = kdt.query_ball_point(x=pts, r=radius, workers=-1) + end = time.time() + print(f"Cull Time: {end - start} s") - # 2) prepare flattened candidates + offsets + start = time.time() + # 2) 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:]) + end = time.time() + + print(f"Flatten Time: {end - start} s") # 3) perform the batch winding‐number test return _batch_point_in_face( From 5ea1a6aa7a1cf9a102819cc1dffa51631f420a27 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Tue, 13 May 2025 08:22:19 -0600 Subject: [PATCH 06/13] remove debugging --- uxarray/grid/point_in_face.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index f4e37e934..f7cfacf9e 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -1,6 +1,5 @@ from __future__ import annotations -import time from typing import TYPE_CHECKING, Tuple import numpy as np @@ -208,26 +207,19 @@ def _point_in_face_query( pts = np.asarray(points, dtype=np.float64) if pts.ndim == 1: pts = pts[np.newaxis, :] - start = time.time() - # 1) cull with k-d tree + # Cull with k-d tree kdt = source_grid._get_scipy_kd_tree() radius = source_grid.max_face_radius cand_lists = kdt.query_ball_point(x=pts, r=radius, workers=-1) - end = time.time() - print(f"Cull Time: {end - start} s") - start = time.time() - # 2) prepare flattened candidates and offsets + # 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:]) - end = time.time() - print(f"Flatten Time: {end - start} s") - - # 3) perform the batch winding‐number test + # Perform the batch winding‐number test return _batch_point_in_face( pts, flat_cands, From 0c0700aec98b45bef540fa818d20b9819c7f527c Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Tue, 13 May 2025 09:47:02 -0600 Subject: [PATCH 07/13] update max face radius --- test/test_point_in_face.py | 3 ++ uxarray/grid/geometry.py | 78 ++++++++++++++++++++--------------- uxarray/grid/point_in_face.py | 2 +- 3 files changed, 49 insertions(+), 34 deletions(-) diff --git a/test/test_point_in_face.py b/test/test_point_in_face.py index 398aec765..47c515af7 100644 --- a/test/test_point_in_face.py +++ b/test/test_point_in_face.py @@ -60,6 +60,9 @@ 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, diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index 508d0dbb9..384d28be5 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 @@ -1156,39 +1156,51 @@ def _populate_max_face_radius(grid): return max_distance -@njit(cache=True) +@njit(cache=True, parallel=True) def calculate_max_face_radius( - face_node_connectivity, node_x, node_y, node_z, face_x, face_y, face_z -): - """Finds the max face radius in the grid.""" - - # 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 face_idx, cur_face in enumerate(face_node_connectivity): - # Filter out INT_FILL_VALUE - node_indices = cur_face[cur_face != INT_FILL_VALUE] - - fx = face_x[face_idx] - fy = face_y[face_idx] - fz = face_z[face_idx] - - nx = node_x[node_indices] - ny = node_y[node_indices] - nz = node_z[node_indices] - - # vectorized distance calculation - dx = nx - fx - dy = ny - fy - dz = nz - fz - distances = np.sqrt(dx * dx + dy * dy + dz * dz) - - # Store the max distance for this face - end_distances[face_idx] = 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/point_in_face.py b/uxarray/grid/point_in_face.py index f7cfacf9e..9d5e93f1a 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -209,7 +209,7 @@ def _point_in_face_query( pts = pts[np.newaxis, :] # Cull with k-d tree kdt = source_grid._get_scipy_kd_tree() - radius = source_grid.max_face_radius + 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 From 609f03c4fef8770066e8e478ba51f32a05cba30f Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 15 May 2025 15:13:59 -0600 Subject: [PATCH 08/13] add examples and clean up logic --- test/test_geometry.py | 40 +++++------ uxarray/grid/geometry.py | 121 ---------------------------------- uxarray/grid/point_in_face.py | 1 + 3 files changed, 18 insertions(+), 144 deletions(-) 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/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index 384d28be5..7d50fb553 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -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,124 +1008,6 @@ 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) - - # Return the index array - return index - - def _populate_max_face_radius(grid): """Populates `max_face_radius`""" diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index 9d5e93f1a..0142f3f66 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -73,6 +73,7 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: total += sign * ang + print(np.abs(total)) return np.abs(total) > np.pi From 983dd1535cdaec9135bb01e8a211ec8a07c30f78 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 15 May 2025 15:14:42 -0600 Subject: [PATCH 09/13] add examples and clean up logic --- uxarray/grid/point_in_face.py | 1 - 1 file changed, 1 deletion(-) diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index 0142f3f66..9d5e93f1a 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -73,7 +73,6 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: total += sign * ang - print(np.abs(total)) return np.abs(total) > np.pi From c9cd0d2da05024a267f525320d86806816cb046c Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 15 May 2025 18:56:22 -0600 Subject: [PATCH 10/13] clean up point util --- uxarray/grid/coordinates.py | 18 ++++++++++++++++++ uxarray/grid/grid.py | 14 ++------------ 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index 25f04e7bb..7404ee7e0 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 None) == (xyz is 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/grid.py b/uxarray/grid/grid.py index 5ce690a32..ea3a3d512 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -34,12 +34,12 @@ _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, prepare_points, ) @@ -2599,18 +2599,8 @@ def get_faces_containing_point( A Python list of length `N`, where each element is the list of face-indices (no padding) for that query point. """ - if (point_lonlat is None) == (point_xyz is None): - raise ValueError( - "Exactly one of `point_lonlat` or `point_xyz` must be provided." - ) - # Convert to cartesian if points are spherical - if point_xyz is None: - lon, lat = map(np.deg2rad, point_lonlat) - point_xyz = _lonlat_rad_to_xyz(lon, lat) - pts = np.asarray(point_xyz, dtype=np.float64) - if pts.ndim == 1: - pts = pts[np.newaxis, :] + 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) From cafc73c66ce9938ebe583e521b670630711c96f8 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 15 May 2025 19:03:34 -0600 Subject: [PATCH 11/13] add example --- uxarray/grid/grid.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index ea3a3d512..b062d5f21 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -134,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" @@ -2598,6 +2598,31 @@ def get_faces_containing_point( List[List[int]] A Python list of length `N`, where each element is the list of face-indices (no padding) for that query point. + + Examples + -------- + + >>> import uxarray as ux + >>> grid_path = "/path/to/grid.nc" + >>> uxgrid = ux.open_grid(grid_path) + + 1. Query a Spherical (lonlat) point + + >>> indices, counts = uxgrid.get_faces_containing_point(point_lonlat=(0.0, 0.0)) + + 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 + + >>> indices = uxgrid.get_faces_containing_point( + ... point_xyz=(0.0, 0.0, 1.0), return_counts=False + ... ) + + """ pts = _prepare_points_for_kdtree(point_lonlat, point_xyz) From 5ffe430a01c4d06fd002f8ea83eb5d50cdd6ce3d Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 15 May 2025 19:18:54 -0600 Subject: [PATCH 12/13] update typing --- uxarray/grid/coordinates.py | 2 +- uxarray/grid/grid.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index 7404ee7e0..bef5f1b4d 100644 --- a/uxarray/grid/coordinates.py +++ b/uxarray/grid/coordinates.py @@ -869,7 +869,7 @@ def prepare_points(points, normalize): def _prepare_points_for_kdtree(lonlat, xyz): - if (lonlat is None) == (xyz is None): + 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." diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index b062d5f21..01e1b6e7a 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -2567,10 +2567,10 @@ def get_faces_between_latitudes(self, lats: Tuple[float, float]): def get_faces_containing_point( self, *, - point_lonlat: Optional[Sequence[float]] = None, - point_xyz: Optional[Sequence[float]] = None, + point_lonlat: Sequence[float] | np.ndarray = None, + point_xyz: Sequence[float] | np.ndarray = None, return_counts: bool = True, - ) -> Union[Tuple[np.ndarray, np.ndarray], List[List[int]]]: + ) -> Tuple[np.ndarray, np.ndarray] | List[List[int]]: """ Identify which faces on the grid contain the given point(s). From 3f3d60ad49b4cdffe126e7fee4474bf07a3ed45d Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 15 May 2025 20:07:38 -0600 Subject: [PATCH 13/13] better docstring --- uxarray/grid/grid.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 01e1b6e7a..992fe9281 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -2589,40 +2589,45 @@ def get_faces_containing_point( Returns ------- If `return_counts=True`: - face_indices : np.ndarray, shape (N, M) - 2D array of face indices; unused slots are filled with `INT_FILL_VALUE`. + 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. + 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 -------- - >>> import uxarray as ux >>> grid_path = "/path/to/grid.nc" >>> uxgrid = ux.open_grid(grid_path) - 1. Query a Spherical (lonlat) point - + # 1. Query a Spherical (lon/lat) point >>> indices, counts = uxgrid.get_faces_containing_point(point_lonlat=(0.0, 0.0)) - 2. Query a Cartesian (xyz) point - + # 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 - + # 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 ... ) - - """ pts = _prepare_points_for_kdtree(point_lonlat, point_xyz) @@ -2637,4 +2642,8 @@ def get_faces_containing_point( output.append(face_indices[i, :c].tolist()) return output + if (counts == 1).all(): + # collapse to a 1D array of length n_points + face_indices = face_indices[:, 0] + return face_indices, counts