diff --git a/benchmarks/to_raster.py b/benchmarks/to_raster.py new file mode 100644 index 000000000..fb0cb45b6 --- /dev/null +++ b/benchmarks/to_raster.py @@ -0,0 +1,165 @@ +import cartopy.crs as ccrs +import matplotlib +matplotlib.use("Agg") # default backend causes ASV to crash on MacOS. +import matplotlib.pyplot as plt +import numpy as np +from scipy.spatial import ConvexHull +import uxarray as ux + + + +from .helpers._memsize import dataset_nbytes, grid_nbytes +from .helpers._peakmem import numba_threads, peak_allocated + + +def make_sphere_grid(n_faces, area_ratio, *, seed=0): + """ + Returns a variable-resolution triangular grid on the unit sphere. + + n_faces : target number of triangular faces (result is approximate) + area_ratio : (face area at poles) / (face area at equator) + seed : RNG seed for the per-ring longitude offsets + + Returns + ------- + node_lon, node_lat : (n_nodes,) degrees (lon in [-180, 180)) + face_node_connectivity : (n_faces, 3) int array, counter-clockwise + when viewed from outside the sphere. + """ + # note: Claude made this function, aside from some small edits by hand. + # The goal was to generate ANY "reasonable-looking" grid with varying face areas, + # because to_raster() performance was bad for grids with significant area variations. + + # You can check it produces reasonable-looking results by doing something like: + # areas = make_sphere_grid(1e4, 100).compute_face_areas(as_uxarray=True) + # areas.plot() # to view the actual areas + # (areas*0 + np.arange(len(areas.n_face))).plot() # to view face indices. + + edge_ratio = np.sqrt(area_ratio) # edge length ~ sqrt(area) + + def h_of(lat, h_pole): # target edge length vs latitude + return (h_pole / edge_ratio) * edge_ratio ** (np.abs(lat) / (np.pi / 2)) + + def layout(h_pole): + """Ring latitudes (all, south to north) and nodes per ring.""" + lats = [0.0] + while True: + nxt = lats[-1] + 0.866 * h_of(lats[-1], h_pole) + if nxt > np.pi / 2 - 0.6 * h_pole: # leave room for the pole node + break + lats.append(nxt) + lats = np.array(lats) + lats = np.concatenate([-lats[:0:-1], lats]) + counts = np.maximum( + 3, np.round(2 * np.pi * np.cos(lats) / h_of(lats, h_pole)) + ).astype(int) + return lats, counts + + # --- choose h_pole so that the face count is close to n_faces + # analytic estimate: faces ~ C / h_pole^2 (equilateral triangles) + phi = np.linspace(-np.pi / 2, np.pi / 2, 20001) + C = np.sum(2 * np.pi * np.cos(phi) / ((np.sqrt(3) / 4) * h_of(phi, 1.0) ** 2)) \ + * (phi[1] - phi[0]) + h0 = np.sqrt(C / n_faces) + + def faces_for(h_pole): # triangles = 2*nodes - 4 on a sphere + return 2 * (layout(h_pole)[1].sum() + 2) - 4 + + lo, hi = np.log(h0 / 2), np.log(h0 * 2) # refine with the true ring counts + for _ in range(30): + mid = 0.5 * (lo + hi) + if faces_for(np.exp(mid)) > n_faces: + lo = mid # too many faces -> larger h + else: + hi = mid + h_pole = np.exp(0.5 * (lo + hi)) + + # --- nodes on each ring + rng = np.random.default_rng(seed) + lats, counts = layout(h_pole) + lon_list, lat_list = [], [] + for lat, n in zip(lats, counts): + lon_list.append(rng.uniform(0, 2 * np.pi) + 2 * np.pi * np.arange(n) / n) + lat_list.append(np.full(n, lat)) + lon_list += [np.array([0.0]), np.array([0.0])] # poles + lat_list += [np.array([np.pi / 2]), np.array([-np.pi / 2])] + + lon = np.concatenate(lon_list) + lat = np.concatenate(lat_list) + + # --- spherical Delaunay via 3D convex hull + xyz = np.column_stack([np.cos(lat) * np.cos(lon), + np.cos(lat) * np.sin(lon), + np.sin(lat)]) + faces = ConvexHull(xyz).simplices.copy() + + # orient counter-clockwise as seen from outside + a, b, c = xyz[faces[:, 0]], xyz[faces[:, 1]], xyz[faces[:, 2]] + flip = np.einsum("ij,ij->i", np.cross(b - a, c - a), a + b + c) < 0 + faces[flip] = faces[flip][:, [0, 2, 1]] + + node_lon = (np.degrees(lon) + 180.0) % 360.0 - 180.0 + node_lat = np.degrees(lat) + connectivity = faces.astype(np.int64) + return ux.Grid.from_topology(node_lon, node_lat, connectivity) + + + +class ToRaster: + """Benchmark the time it takes to call to_raster(), + on a variable-resolution grid. + """ + param_names = ["n_faces_and_area_ratio",] + params = [[ + ## commented-out below are some additional reasonable parameter combinations, + ## excluded from default benchmark suite for brevity, but kept because they + ## could be uncommented for local runs or more detailed profiling. + # (1e3, 10.0), + # (1e3, 100.0), + ## never do (1e3, 1000.0), it isn't properly resolved; + # it would just make a bunch of faces stretching from pole to equator. + (1e4, 10.0), + (1e4, 100.0), + # (1e4, 1000.0), + # (1e5, 10.0), + (1e5, 100.0), + # (1e5, 1000.0), + ]] + + # if any combination takes longer than 3 mins, give up. + timeout = 180 + + def _warmup(self): + # warm up numba, using a small example + grid = make_sphere_grid(1e3, 10.0) + data = grid.compute_face_areas(as_uxarray=True) + # data could be any values; let's use areas for fun :) + fig, ax = plt.subplots(subplot_kw={"projection": ccrs.Robinson()}) + ax.set_global() + data.to_raster(ax=ax) + plt.close() + + def setup(self, n_faces_and_area_ratio): + n_faces, area_ratio = n_faces_and_area_ratio + grid = make_sphere_grid(n_faces, area_ratio) + self.data = grid.compute_face_areas(as_uxarray=True) + self._warmup() + self.fig, self.ax = plt.subplots(subplot_kw={"projection": ccrs.Robinson()}) + self.ax.set_global() + + def teardown(self, n_faces_and_area_ratio): + del self.data + del self.fig + del self.ax + plt.close() + + def time_to_raster(self, n_faces_and_area_ratio): + """Time to call to_raster() on a variable-resolution grid.""" + self.data.to_raster(ax=self.ax) + + def track_peakmem(self, n_faces_and_area_ratio): + """High-water allocation of to_raster()""" + with numba_threads(1): + return peak_allocated(lambda: self.data.to_raster(ax=self.ax)) + + track_peakmem.unit = "bytes" diff --git a/test/grid/geometry/test_geometry.py b/test/grid/geometry/test_geometry.py index 64dc6c41a..106412bd8 100644 --- a/test/grid/geometry/test_geometry.py +++ b/test/grid/geometry/test_geometry.py @@ -6,7 +6,7 @@ from uxarray.constants import ERROR_TOLERANCE, INT_FILL_VALUE from uxarray.grid.coordinates import _lonlat_rad_to_xyz, _normalize_xyz from uxarray.grid.geometry import haversine_distance, _pole_point_inside_polygon_cartesian -from uxarray.grid.point_in_face import _face_contains_point +from uxarray.grid.point_in_face import _point_in_face_from_grid from uxarray.grid.utils import _get_cartesian_face_edge_nodes_array @@ -114,25 +114,14 @@ def test_cache_and_override_geodataframe(gridpath): def test_face_at_antimeridian(): - """Test the function `point_in_face`, where the face crosses the antimeridian""" + """Test the function `_point_in_face`, where the face crosses the antimeridian""" # Generate a face crossing the antimeridian vertices_lonlat = [[350, 60.0], [350, 10.0], [50.0, 10.0], [50.0, 60.0]] vertices_lonlat = np.array(vertices_lonlat) point = np.array(_lonlat_rad_to_xyz(np.deg2rad(25), np.deg2rad(30))) - - # Create the grid and face edges grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - 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, - ) - - assert _face_contains_point(faces_edges_cartesian[0], point) + assert _point_in_face_from_grid(point, grid, fidx=0) def test_face_at_pole(): @@ -141,21 +130,9 @@ def test_face_at_pole(): # Generate a face that is at a pole vertices_lonlat = [[10.0, 90.0], [10.0, 10.0], [50.0, 10.0], [50.0, 60.0]] vertices_lonlat = np.array(vertices_lonlat) - point = np.array(_lonlat_rad_to_xyz(np.deg2rad(25), np.deg2rad(30))) - - # Create the grid and face edges grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - 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, - ) - - assert _face_contains_point(faces_edges_cartesian[0], point) + assert _point_in_face_from_grid(point, grid, fidx=0) def test_face_normal_face(): @@ -166,19 +143,8 @@ def test_face_normal_face(): vertices_lonlat = [[10.0, 60.0], [10.0, 10.0], [50.0, 10.0], [50.0, 60.0]] vertices_lonlat = np.array(vertices_lonlat) point = np.array(_lonlat_rad_to_xyz(np.deg2rad(25), np.deg2rad(30))) - - # Create the grid and face edges grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - 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, - ) - - assert _face_contains_point(faces_edges_cartesian[0], point) + assert _point_in_face_from_grid(point, grid, fidx=0) def test_haversine_distance_creation(): diff --git a/test/grid/geometry/test_point_in_face.py b/test/grid/geometry/test_point_in_face.py index 2a0cafc87..1829c7eeb 100644 --- a/test/grid/geometry/test_point_in_face.py +++ b/test/grid/geometry/test_point_in_face.py @@ -6,72 +6,45 @@ from uxarray.constants import INT_FILL_VALUE from uxarray.grid.coordinates import _lonlat_rad_to_xyz from uxarray.grid.utils import _get_cartesian_face_edge_nodes_array, _get_cartesian_face_edge_nodes -from uxarray.grid.point_in_face import _face_contains_point +from uxarray.grid.point_in_face import _point_in_face_from_grid def test_point_inside(gridpath): """Test the function `point_in_face`, where the points are all inside the face""" - # Open grid grid = ux.open_grid(gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc")) grid.normalize_cartesian_coordinates() # 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 _face_contains_point(face_edges, point_xyz) + assert _point_in_face_from_grid(point_xyz, grid, i) def test_point_outside(gridpath): """Test the function `point_in_face`, where the point is outside the face""" - # Open grid grid = ux.open_grid(gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc")) - # 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, - ) - # Set the point as the face center of a different face than the face tested 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 _face_contains_point(faces_edges_cartesian[0], point_xyz) + assert not _point_in_face_from_grid(point_xyz, grid, fidx=0) def test_point_on_node(gridpath): """Test the function `point_in_face`, when the point is on the node of the polygon""" - # Open grid grid = ux.open_grid(gridpath("mpas", "QU", "mesh.QU.1920km.151026.nc")) - # 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, - ) - # Set the point as a node - point_xyz = np.array([*faces_edges_cartesian[0][0][0]]) + point_xyz = np.array([grid.node_x[0].values, grid.node_y[0].values, grid.node_z[0].values]) - # Assert that the point is in the face when inclusive is true - assert _face_contains_point(faces_edges_cartesian[0], point_xyz) + # Assert that the point is in the face + assert _point_in_face_from_grid(point_xyz, grid, fidx=0) def test_point_inside_close(): @@ -86,17 +59,8 @@ def test_point_inside_close(): # Create the grid and face edges grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - 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, - ) - # Use point in face to determine if the point is inside or out of the face - assert _face_contains_point(faces_edges_cartesian[0], point) + assert _point_in_face_from_grid(point, grid, fidx=0) def test_point_outside_close(): @@ -111,14 +75,5 @@ def test_point_outside_close(): # Create the grid and face edges grid = ux.Grid.from_face_vertices(vertices_lonlat, latlon=True) - 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, - ) - - # Use point in face to determine if the point is inside or out of the face - assert not _face_contains_point(faces_edges_cartesian[0], point) + + assert not _point_in_face_from_grid(point, grid, fidx=0) diff --git a/uxarray/grid/arcs.py b/uxarray/grid/arcs.py index 9ec3df40d..dfe410988 100644 --- a/uxarray/grid/arcs.py +++ b/uxarray/grid/arcs.py @@ -9,6 +9,11 @@ ) from uxarray.grid.utils import _angle_of_2_vectors from uxarray.utils.computing import _cdp8, accucross +from uxarray.utils.numba_math import ( + _numba_cross3, + _numba_dot3, + _numba_sub3, +) # Magnitude below which orient3d_on_sphere classifies a result as zero. For # double-precision unit-vector inputs this covers rounding error in the @@ -40,11 +45,11 @@ def point_within_gca(pt_xyz, gca_a_xyz, gca_b_xyz): Parameters ---------- - pt_xyz : numpy.ndarray + pt_xyz : iterable of length 3 Cartesian coordinates of the point. - gca_a_xyz : numpy.ndarray + gca_a_xyz : iterable of length 3 Cartesian coordinates of the first endpoint of the Great Circle Arc. - gca_b_xyz : numpy.ndarray + gca_b_xyz : iterable of length 3 Cartesian coordinates of the second endpoint of the Great Circle Arc. Returns @@ -66,25 +71,28 @@ def point_within_gca(pt_xyz, gca_a_xyz, gca_b_xyz): """ # 1. Check if the input GCA spans exactly 180 degrees angle_ab = _angle_of_2_vectors(gca_a_xyz, gca_b_xyz) - if np.allclose(angle_ab, np.pi, rtol=0.0, atol=MACHINE_EPSILON): + if np.isclose(angle_ab, np.pi, rtol=0.0, atol=MACHINE_EPSILON): raise ValueError( "The input Great Circle Arc spans exactly 180 degrees, which can correspond to multiple planes. " "Consider breaking the Great Circle Arc into two smaller arcs." ) # (numba complains about f-strings, so don't put actual values in message.) # 2. Verify if the point lies on the plane of the GCA - cross_product = np.cross(gca_a_xyz, gca_b_xyz) - if not np.allclose( - np.dot(cross_product, pt_xyz), 0, rtol=MACHINE_EPSILON, atol=MACHINE_EPSILON + cross_product = _numba_cross3(gca_a_xyz, gca_b_xyz) + if not np.isclose( + _numba_dot3(cross_product, pt_xyz), + 0, + rtol=MACHINE_EPSILON, + atol=MACHINE_EPSILON, ): return False # 3. Check if the point lies within the Great Circle Arc interval - pt_a = gca_a_xyz - pt_xyz - pt_b = gca_b_xyz - pt_xyz + pt_a = _numba_sub3(gca_a_xyz, pt_xyz) + pt_b = _numba_sub3(gca_b_xyz, pt_xyz) # Use the dot product to determine the sign of the angle between pt_a and pt_b - cos_theta = np.dot(pt_a, pt_b) + cos_theta = _numba_dot3(pt_a, pt_b) # Return True if the point lies within the interval (smaller arc) if cos_theta < 0: diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index 31481956d..26f8defc9 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -13,8 +13,7 @@ from uxarray.grid.intersections import ( gca_gca_intersection, ) -from uxarray.grid.point_in_face import _face_contains_point -from uxarray.grid.utils import _get_cartesian_face_edge_nodes +from uxarray.grid.point_in_face import _point_in_face from uxarray.utils.imports import _raise_hint_if_optional_deps_missing POLE_POINTS_XYZ = { @@ -1260,20 +1259,14 @@ def barycentric_coordinates_cartesian(polygon_xyz, point_xyz): polygon_xyz[i + 2], ) - # Get the triangle in terms of its edges for the `point_in_face` check - face_edge = _get_cartesian_face_edge_nodes( - face_idx=0, - face_node_connectivity=np.array([[0, 1, 2]]), - n_edges_per_face=np.array([3]), - node_x=np.array([node_0[0], node_1[0], node_2[0]], dtype=np.float64), - node_y=np.array([node_0[1], node_1[1], node_2[1]], dtype=np.float64), - node_z=np.array([node_0[2], node_1[2], node_2[2]], dtype=np.float64), - ) + node_x = np.array([node_0[0], node_1[0], node_2[0]], dtype=np.float64) + node_y = np.array([node_0[1], node_1[1], node_2[1]], dtype=np.float64) + node_z = np.array([node_0[2], node_1[2], node_2[2]], dtype=np.float64) + nodes_idx = np.array([0, 1, 2]) # Check to see if the point lies within the current triangle - contains_point = _face_contains_point( - face_edge, - point_xyz, + contains_point = _point_in_face( + point_xyz, nodes_idx, node_x, node_y, node_z ) # If the point is in the current triangle, get the weights for that triangle diff --git a/uxarray/grid/point_in_face.py b/uxarray/grid/point_in_face.py index ed6faf7fd..d2f361aec 100644 --- a/uxarray/grid/point_in_face.py +++ b/uxarray/grid/point_in_face.py @@ -7,7 +7,14 @@ 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 +from uxarray.grid.utils import _small_angle_of_2_vectors +from uxarray.utils.numba_math import ( + _numba_allclose3, + _numba_cross3, + _numba_dot3, + _numba_norm3, + _numba_sub3, +) if TYPE_CHECKING: from numpy.typing import ArrayLike @@ -15,61 +22,103 @@ from uxarray.grid.grid import Grid -@njit(cache=True) -def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: +def _point_in_face_from_grid(point: np.ndarray, grid: Grid, fidx: int): + """Returns whether this point lies within the indicated face of this grid. + Helper function providing convenient entry point into `_point_in_face`; + see `_point_in_face` for full docstring. """ - Determine whether a point lies within a face using the spherical winding-number method. + n_nodes_in_face = grid.n_nodes_per_face[fidx].item() + nodes_idx = grid.face_node_connectivity[fidx][:n_nodes_in_face].values + nodes_x = grid.node_x.values + nodes_y = grid.node_y.values + nodes_z = grid.node_z.values + return _point_in_face(point, nodes_idx, nodes_x, nodes_y, nodes_z) + - This function sums the signed central angles between successive vertices of the face +@njit(cache=True) +def _point_in_face( + point: np.ndarray | tuple[float, float, float], + nodes_idx: np.ndarray, + node_x: np.ndarray, + node_y: np.ndarray, + node_z: np.ndarray, +) -> bool: + """Returns whether this point lies within the face formed by these nodes. + + Uses the spherical winding-number method, which + 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,) + point : iterable of length 3 3D unit-vector of the query point on the unit sphere. + nodes_idx : np.ndarray, shape (n_nodes,) + Node indices (within node_x, node_y, node_z) for precisely all nodes in this face. + Likely from `face_node_connectivity[fidx][:n_nodes_per_face[fidx]]`. + node_x, node_y, node_z : np.ndarray, shape (n_nodes,) + Cartesian coordinates of all nodes. + (This method uses the values at indices indicated by nodes_idx.) Returns ------- inside : bool True if the point is inside the face or lies exactly on a node/edge; False otherwise. """ + # Rewritten from _face_contains_point_from_edges to avoids creating tiny numpy arrays. + # Creating tiny numpy arrays from scratch inside numba is very inefficient. + # (Creating tiny numpy arrays from indexing larger arrays is fine, though.) + # This is the main reason to provide inputs as nodes_idx and node_x, ..., instead of + # simply asking to provide a single array of x, y, z coordinates for all nodes; + # the former avoids any need to create a tiny numpy array to store the coordinates. + + n_nodes = len(nodes_idx) + max_i_node = n_nodes - 1 + # 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 + for i in range(n_nodes): + node_idx = nodes_idx[i] + node_xyz = (node_x[node_idx], node_y[node_idx], node_z[node_idx]) + if _numba_allclose3( + node_xyz, 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] + # Check whether point lies on any edge of the face + # (edges are great-circle arcs between successive nodes) + for i in range(n_nodes): + # edge is formed by nodes (a, b) + ai = nodes_idx[i] + bi = nodes_idx[i + 1] if i < max_i_node else nodes_idx[0] + a = (node_x[ai], node_y[ai], node_z[ai]) + b = (node_x[bi], node_y[bi], node_z[bi]) + if point_within_gca(point, a, b): + return True + + # Apply spherical winding-number method: 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] + for i in range(n_nodes): + # edge is formed by nodes (a, b). + ai = nodes_idx[i] + bi = nodes_idx[i + 1] if i < max_i_node else nodes_idx[0] - vi = a - p - vj = b - p + a = (node_x[ai], node_y[ai], node_z[ai]) + b = (node_x[bi], node_y[bi], node_z[bi]) + + vi = _numba_sub3(a, point) + vj = _numba_sub3(b, point) # check if you’re right on a vertex - if np.linalg.norm(vi) < ERROR_TOLERANCE or np.linalg.norm(vj) < ERROR_TOLERANCE: + if _numba_norm3(vi) < ERROR_TOLERANCE or _numba_norm3(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 + c = _numba_cross3(vi, vj) + sign = 1.0 if _numba_dot3(c, point) >= 0.0 else -1.0 total += sign * ang @@ -77,21 +126,32 @@ def _face_contains_point(face_edges: np.ndarray, point: np.ndarray) -> bool: @njit(cache=True) -def _get_faces_containing_point( - point: np.ndarray, +def _set_faces_containing_point( + result: np.ndarray, + i: int, + point: np.ndarray | tuple[float, float, float], 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: +) -> int: """ - Test each candidate face to see if it contains the query point. + Test each candidate face to see if it contains the query point, + setting result[i, j] = candidate_indices[k] for all k where the + point is inside the face. j starts at 0 and increments by 1 for + each hit. Returns the total number of hits. Parameters ---------- - point : np.ndarray, shape (3,) + result : np.ndarray, shape (n_points, max_possible_hits) + Preallocated array to store face indices for each point. + max_possible_hits can be much less than max_candidates; + the maximum number of hits is n_max_face_nodes, because + the "worst case" of point being a node would lead to hits + of all faces it is a part of, but nothing else. + point : iterable of length 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). @@ -104,20 +164,17 @@ def _get_faces_containing_point( Returns ------- - hits : np.ndarray, shape (h,) - Subset of `candidate_indices` for which the point is inside the face. + n_hits : int + Number of candidate faces that contain the point. """ - 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 + nodes_idx = face_node_connectivity[fidx][: n_nodes_per_face[fidx]] + if _point_in_face(point, nodes_idx, node_x, node_y, node_z): + result[i, count] = fidx count += 1 - return hit_buf[:count] + return count @njit(cache=True, parallel=True, nogil=True) @@ -162,15 +219,21 @@ def _batch_point_in_face( for i in prange(n_points): start = offsets[i] end = offsets[i + 1] - p = points[i] + p = (points[i][0], points[i][1], points[i][2]) 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 + n_hits = _set_faces_containing_point( + results, + i, + 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] + counts[i] = n_hits return results, counts diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index 4d01e5049..bfa389371 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -5,6 +5,8 @@ from uxarray.constants import EDGE_NODE_SORT_THRESHOLD, INT_DTYPE, INT_FILL_VALUE from uxarray.utils.numba_math import ( _numba_add3, + _numba_cross3, + _numba_dot3, _numba_mul3_scalar, _numba_norm3, _numba_sub3, @@ -46,9 +48,9 @@ def _angle_of_2_vectors(u, v): Parameters ---------- - u : numpy.ndarray + u : iterable of length 3 The first 3D vector (float), originating from the center of the unit sphere. - v : numpy.ndarray + v : iterable of length 3 The second 3D vector (float), originating from the center of the unit sphere. Returns @@ -62,13 +64,14 @@ def _angle_of_2_vectors(u, v): - Special cases such as vectors aligned along the same longitude are handled explicitly. """ # Compute the cross product to determine the direction of the normal - normal = np.cross(u, v) + normal = _numba_cross3(u, v) # Calculate the angle using arctangent of cross and dot products - angle_u_v_rad = np.arctan2(np.linalg.norm(normal), np.dot(u, v)) + angle_u_v_rad = np.arctan2(_numba_norm3(normal), _numba_dot3(u, v)) # Determine the direction of the angle - normal_z = np.dot(normal, np.array([0.0, 0.0, 1.0])) + normal_z = normal[2] + if normal_z > 0: # Counterclockwise direction return angle_u_v_rad diff --git a/uxarray/utils/numba_math.py b/uxarray/utils/numba_math.py index c2f0ad673..7f75f2da2 100644 --- a/uxarray/utils/numba_math.py +++ b/uxarray/utils/numba_math.py @@ -124,3 +124,13 @@ def _numba_allfinite3(u): int(math.isfinite(u[0])) * int(math.isfinite(u[1])) * int(math.isfinite(u[2])) ) # (use `*` instead of `and` to avoid branching logic) + + +@njit(cache=True) +def _numba_allclose3(u, v, rtol=1e-05, atol=1e-08): + """Return (as 1 or 0) whether all components of two 3-vectors are close to each other.""" + return ( + int(np.isclose(u[0], v[0], rtol=rtol, atol=atol)) + * int(np.isclose(u[1], v[1], rtol=rtol, atol=atol)) + * int(np.isclose(u[2], v[2], rtol=rtol, atol=atol)) + )