From 58d9ecc7f3a74a7127121f573a6ed16f0a7bda7e Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:35:27 -0400 Subject: [PATCH 1/3] add numba_math.py; optimize coordinates.py E.g. locally seeing grid.construct_face_centers('welzl') time decrease from 14 s to ~150 ms, for a grid with 157k faces. --- uxarray/grid/coordinates.py | 48 +++++++++++-------- uxarray/grid/utils.py | 17 ++++--- uxarray/utils/numba_math.py | 91 +++++++++++++++++++++++++++++++++++++ 3 files changed, 129 insertions(+), 27 deletions(-) create mode 100644 uxarray/utils/numba_math.py diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index 9443e3769..8c61b65f7 100644 --- a/uxarray/grid/coordinates.py +++ b/uxarray/grid/coordinates.py @@ -9,13 +9,16 @@ from uxarray.conventions import ugrid from uxarray.errors import DimensionError from uxarray.grid.utils import _small_angle_of_2_vectors +from uxarray.utils.numba_math import ( + _numba_div3_scalar, _numba_norm3, +) @njit(cache=True) def _lonlat_rad_to_xyz( lon: np.ndarray | float, lat: np.ndarray | float, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[np.ndarray | float, np.ndarray | float, np.ndarray | float]: """Converts Spherical latitude and longitude coordinates into Cartesian x, y, z coordinates.""" x = np.cos(lon) * np.cos(lat) @@ -82,6 +85,10 @@ def _xyz_to_lonlat_rad_scalar(x, y, z, normalize=True): lat = math.copysign(math.pi / 2, z) lon = 0.0 + # TODO: constructing tiny numpy array inside numba function is sub-optimal, + # if function gets called many times. (see issue $1648). But, as of 2026-08-13, + # it looks like this function isn't being used anywhere, so maybe it should + # just be removed entirely, instead of being optimized? lonlat = np.empty(2) lonlat[0] = lon lonlat[1] = lat @@ -94,7 +101,7 @@ def _xyz_to_lonlat_rad( y: np.ndarray | float, z: np.ndarray | float, normalize: bool = True, -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[np.ndarray | float, np.ndarray | float]: """Converts Cartesian x, y, z coordinates in Spherical longitude and latitude coordinates in radians. @@ -188,11 +195,10 @@ def _normalize_xyz( @njit(cache=True) def _normalize_xyz_scalar(x: float, y: float, z: float): - denom = np.linalg.norm(np.asarray(np.array([x, y, z]), dtype=np.float64), ord=2) - x_norm = x / denom - y_norm = y / denom - z_norm = z / denom - return x_norm, y_norm, z_norm + """returns (x/|u|, y/|u|, z/|u|), where |u| = sqrt(x^2 + y^2 + z^2)""" + u = (x, y, z) + u_norm = _numba_norm3(u) + return _numba_div3_scalar(u, u_norm) def _populate_node_latlon(grid) -> None: @@ -341,6 +347,7 @@ def _construct_face_centroids(node_x, node_y, node_z, face_nodes, n_nodes_per_fa return centroid_x, centroid_y, centroid_z +@njit(cache=True) def _welzl_recursive(points, boundary, R): """Recursive helper function for Welzl's algorithm to find the smallest enclosing circle. @@ -389,6 +396,7 @@ def _welzl_recursive(points, boundary, R): return _welzl_recursive(temp_points, new_boundary, R) +@njit(cache=True) def _smallest_enclosing_circle(points): """Find the smallest circle that encloses all given points on a unit sphere using Welzl's algorithm. @@ -429,8 +437,8 @@ def _circle_from_two_points(p1, p2): center_lat = (p1[1] + p2[1]) / 2 center = (center_lon, center_lat) - v1 = np.array(_lonlat_rad_to_xyz(np.radians(p1[0]), np.radians(p1[1]))) - v2 = np.array(_lonlat_rad_to_xyz(np.radians(p2[0]), np.radians(p2[1]))) + v1 = _lonlat_rad_to_xyz(np.radians(p1[0]), np.radians(p1[1])) + v2 = _lonlat_rad_to_xyz(np.radians(p2[0]), np.radians(p2[1])) distance = _small_angle_of_2_vectors(v1, v2) radius = distance / 2 @@ -462,9 +470,9 @@ def _circle_from_three_points(p1, p2, p3): center_lat = (p1[1] + p2[1] + p3[1]) / 3 center = (center_lon, center_lat) - v1 = np.array(_lonlat_rad_to_xyz(np.radians(p1[0]), np.radians(p1[1]))) - v2 = np.array(_lonlat_rad_to_xyz(np.radians(p2[0]), np.radians(p2[1]))) - v3 = np.array(_lonlat_rad_to_xyz(np.radians(p3[0]), np.radians(p3[1]))) + v1 = _lonlat_rad_to_xyz(np.radians(p1[0]), np.radians(p1[1])) + v2 = _lonlat_rad_to_xyz(np.radians(p2[0]), np.radians(p2[1])) + v3 = _lonlat_rad_to_xyz(np.radians(p3[0]), np.radians(p3[1])) radius = ( max( @@ -495,8 +503,8 @@ def _is_inside_circle(circle, point): True if the point is inside the circle, False otherwise. """ center, radius = circle - v1 = np.array(_lonlat_rad_to_xyz(np.radians(center[0]), np.radians(center[1]))) - v2 = np.array(_lonlat_rad_to_xyz(np.radians(point[0]), np.radians(point[1]))) + v1 = _lonlat_rad_to_xyz(np.radians(center[0]), np.radians(center[1])) + v2 = _lonlat_rad_to_xyz(np.radians(point[0]), np.radians(point[1])) distance = _small_angle_of_2_vectors(v1, v2) return distance <= radius @@ -560,6 +568,7 @@ def _populate_face_centerpoints(grid, repopulate=False): ) +@njit(cache=True, parallel=True) def _construct_face_centerpoints(node_lon, node_lat, face_nodes, n_nodes_per_face): """Constructs the face centerpoint using Welzl's algorithm. @@ -595,12 +604,11 @@ def _construct_face_centerpoints(node_lon, node_lat, face_nodes, n_nodes_per_fac ] # Compute circles for all faces - circles = [_smallest_enclosing_circle(points) for points in points_arrays] - - # Extract centerpoints - ctrpt_lon, ctrpt_lat = zip(*[circle[0] for circle in circles]) - - return np.array(ctrpt_lon), np.array(ctrpt_lat) + for ipoint in prange(len(points_arrays)): + circle_lonlat = _smallest_enclosing_circle(points_arrays[ipoint])[0] + ctrpt_lon[ipoint] = circle_lonlat[0] + ctrpt_lat[ipoint] = circle_lonlat[1] + return ctrpt_lon, ctrpt_lat def _populate_edge_centroids(grid, repopulate=False): diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index d70962273..837fdd510 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -3,6 +3,9 @@ from numba import njit from uxarray.constants import INT_FILL_VALUE +from uxarray.utils.numba_math import ( + _numba_add3, _numba_sub3, _numba_mul3_scalar, _numba_norm3, +) @njit(cache=True) @@ -12,9 +15,9 @@ def _small_angle_of_2_vectors(u, v): Parameters ---------- - u : numpy.ndarray + u : iterable of length 3 The first 3D vector. - v : numpy.ndarray + v : iterable of length 3 The second 3D vector. Returns @@ -22,11 +25,11 @@ def _small_angle_of_2_vectors(u, v): float The smallest angle between `u` and `v` in radians. """ - v_norm_times_u = np.linalg.norm(v) * u - u_norm_times_v = np.linalg.norm(u) * v - vec_minus = v_norm_times_u - u_norm_times_v - vec_sum = v_norm_times_u + u_norm_times_v - angle_u_v_rad = 2 * np.arctan2(np.linalg.norm(vec_minus), np.linalg.norm(vec_sum)) + u_times_v_norm = _numba_mul3_scalar(v, _numba_norm3(u)) + v_times_u_norm = _numba_mul3_scalar(u, _numba_norm3(v)) + vec_minus = _numba_sub3(u_times_v_norm, v_times_u_norm) + vec_sum = _numba_add3(u_times_v_norm, v_times_u_norm) + angle_u_v_rad = 2 * np.arctan2(_numba_norm3(vec_minus), _numba_norm3(vec_sum)) return angle_u_v_rad diff --git a/uxarray/utils/numba_math.py b/uxarray/utils/numba_math.py new file mode 100644 index 000000000..fe9ce5a4e --- /dev/null +++ b/uxarray/utils/numba_math.py @@ -0,0 +1,91 @@ +""" +Purpose: numba math helpers/primitives + +Creating lots of tiny numpy arrays (or lists) costs a lot in numba; + it's much more efficient to use individual components + or write tuples, when looping across points. + E.g., instead of v1=np.array((x1,y1,z1)); v2=np.array((x2,y2,z2)); np.dot(v1,v2), + using v1=(x1,y1,z1); v2=(x2,y2,z2); _numba_dot3(v1,v2) is much faster, + because it doesn't need to allocate lists/arrays. + See issue #1648 for more details. +""" +import numpy as np +from numba import njit + +# ------- basic arithmetic with vectors ------- # + +@njit(cache=True) +def _numba_add3(u, v): + """component-wise addition of 3-vectors; returns (u[0] + v[0], u[1] + v[1], u[2] + v[2])""" + return u[0] + v[0], u[1] + v[1], u[2] + v[2] + +@njit(cache=True) +def _numba_add3_scalar(u, scalar): + """component-wise addition of 3-vector and scalar; returns (u[0] + scalar, u[1] + scalar, u[2] + scalar)""" + return u[0] + scalar, u[1] + scalar, u[2] + scalar + +@njit(cache=True) +def _numba_sub3(u, v): + """component-wise subtraction of 3-vectors; returns (u[0] - v[0], u[1] - v[1], u[2] - v[2])""" + return u[0] - v[0], u[1] - v[1], u[2] - v[2] + +# _numba_sub3_scalar not provided; just use _numba_add3_scalar with negative scalar + +@njit(cache=True) +def _numba_mul3(u, v): + """component-wise multiplication of 3-vectors; returns (u[0] * v[0], u[1] * v[1], u[2] * v[2])""" + return u[0] * v[0], u[1] * v[1], u[2] * v[2] + +@njit(cache=True) +def _numba_mul3_scalar(u, scalar): + """component-wise multiplication of 3-vector and scalar; returns (u[0] * scalar, u[1] * scalar, u[2] * scalar)""" + return u[0] * scalar, u[1] * scalar, u[2] * scalar + +@njit(cache=True) +def _numba_div3(u, v): + """component-wise division of 3-vectors; returns (u[0] / v[0], u[1] / v[1], u[2] / v[2])""" + return u[0] / v[0], u[1] / v[1], u[2] / v[2] + +@njit(cache=True) +def _numba_div3_scalar(u, scalar): + """component-wise division of 3-vector by scalar; returns (u[0] / scalar, u[1] / scalar, u[2] / scalar)""" + return u[0] / scalar, u[1] / scalar, u[2] / scalar + +@njit(cache=True) +def _numba_sqrt3(u): + """component-wise square root of 3-vector; returns (sqrt(u[0]), sqrt(u[1]), sqrt(u[2]))""" + return np.sqrt(u[0]), np.sqrt(u[1]), np.sqrt(u[2]) + + +# ------- vector arithmetic ------- # + +@njit(cache=True) +def _numba_norm3(u): + """Euclidean norm of a 3-vector; returns sqrt(u[0]**2 + u[1]**2 + u[2]**2)""" + return np.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + +@njit(cache=True) +def _numba_dot3(u, v): + """dot product of two 3-vectors; returns u[0]*v[0] + u[1]*v[1] + u[2]*v[2]""" + return u[0] * v[0] + u[1] * v[1] + u[2] * v[2] + +@njit(cache=True) +def _numba_cross3(u, v): + """cross product of two 3-vectors. + + Parameters + ---------- + u : iterable of length 3 + The first input vector. + v : iterable of length 3 + The second input vector. + + Examples + -------- + >>> _numba_cross3((1, 2, 30), (4, 0, 6)) + (12, 114, -8) # (2*6 - 30*0, 30*4 - 1*6, 1*0 - 2*4) + """ + cx = u[1] * v[2] - u[2] * v[1] + cy = u[2] * v[0] - u[0] * v[2] + cz = u[0] * v[1] - u[1] * v[0] + return (cx, cy, cz) From 65d2a4d37c4fc0b9c757f529cd33cdf3f6899bc4 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:39:35 -0400 Subject: [PATCH 2/3] forgot pre-commit ruff formatting --- uxarray/grid/coordinates.py | 3 ++- uxarray/grid/utils.py | 5 ++++- uxarray/utils/numba_math.py | 13 +++++++++++++ 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index 8c61b65f7..e9b944045 100644 --- a/uxarray/grid/coordinates.py +++ b/uxarray/grid/coordinates.py @@ -10,7 +10,8 @@ from uxarray.errors import DimensionError from uxarray.grid.utils import _small_angle_of_2_vectors from uxarray.utils.numba_math import ( - _numba_div3_scalar, _numba_norm3, + _numba_div3_scalar, + _numba_norm3, ) diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index 837fdd510..5b01e5570 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -4,7 +4,10 @@ from uxarray.constants import INT_FILL_VALUE from uxarray.utils.numba_math import ( - _numba_add3, _numba_sub3, _numba_mul3_scalar, _numba_norm3, + _numba_add3, + _numba_mul3_scalar, + _numba_norm3, + _numba_sub3, ) diff --git a/uxarray/utils/numba_math.py b/uxarray/utils/numba_math.py index fe9ce5a4e..3afafa3af 100644 --- a/uxarray/utils/numba_math.py +++ b/uxarray/utils/numba_math.py @@ -9,48 +9,58 @@ because it doesn't need to allocate lists/arrays. See issue #1648 for more details. """ + import numpy as np from numba import njit # ------- basic arithmetic with vectors ------- # + @njit(cache=True) def _numba_add3(u, v): """component-wise addition of 3-vectors; returns (u[0] + v[0], u[1] + v[1], u[2] + v[2])""" return u[0] + v[0], u[1] + v[1], u[2] + v[2] + @njit(cache=True) def _numba_add3_scalar(u, scalar): """component-wise addition of 3-vector and scalar; returns (u[0] + scalar, u[1] + scalar, u[2] + scalar)""" return u[0] + scalar, u[1] + scalar, u[2] + scalar + @njit(cache=True) def _numba_sub3(u, v): """component-wise subtraction of 3-vectors; returns (u[0] - v[0], u[1] - v[1], u[2] - v[2])""" return u[0] - v[0], u[1] - v[1], u[2] - v[2] + # _numba_sub3_scalar not provided; just use _numba_add3_scalar with negative scalar + @njit(cache=True) def _numba_mul3(u, v): """component-wise multiplication of 3-vectors; returns (u[0] * v[0], u[1] * v[1], u[2] * v[2])""" return u[0] * v[0], u[1] * v[1], u[2] * v[2] + @njit(cache=True) def _numba_mul3_scalar(u, scalar): """component-wise multiplication of 3-vector and scalar; returns (u[0] * scalar, u[1] * scalar, u[2] * scalar)""" return u[0] * scalar, u[1] * scalar, u[2] * scalar + @njit(cache=True) def _numba_div3(u, v): """component-wise division of 3-vectors; returns (u[0] / v[0], u[1] / v[1], u[2] / v[2])""" return u[0] / v[0], u[1] / v[1], u[2] / v[2] + @njit(cache=True) def _numba_div3_scalar(u, scalar): """component-wise division of 3-vector by scalar; returns (u[0] / scalar, u[1] / scalar, u[2] / scalar)""" return u[0] / scalar, u[1] / scalar, u[2] / scalar + @njit(cache=True) def _numba_sqrt3(u): """component-wise square root of 3-vector; returns (sqrt(u[0]), sqrt(u[1]), sqrt(u[2]))""" @@ -59,16 +69,19 @@ def _numba_sqrt3(u): # ------- vector arithmetic ------- # + @njit(cache=True) def _numba_norm3(u): """Euclidean norm of a 3-vector; returns sqrt(u[0]**2 + u[1]**2 + u[2]**2)""" return np.sqrt(u[0] ** 2 + u[1] ** 2 + u[2] ** 2) + @njit(cache=True) def _numba_dot3(u, v): """dot product of two 3-vectors; returns u[0]*v[0] + u[1]*v[1] + u[2]*v[2]""" return u[0] * v[0] + u[1] * v[1] + u[2] * v[2] + @njit(cache=True) def _numba_cross3(u, v): """cross product of two 3-vectors. From c323777a97dea98dd628b71ade02bfc36ad981e5 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Fri, 14 Aug 2026 11:14:39 -0400 Subject: [PATCH 3/3] add test_numba_math.py --- test/utils/test_numba_math.py | 81 +++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 test/utils/test_numba_math.py diff --git a/test/utils/test_numba_math.py b/test/utils/test_numba_math.py new file mode 100644 index 000000000..e2732043b --- /dev/null +++ b/test/utils/test_numba_math.py @@ -0,0 +1,81 @@ +""" +Purpose: testing routines from utils/numba_math.py +""" +import numpy as np +import pytest + +from uxarray.utils.numba_math import ( + _numba_add3, + _numba_add3_scalar, + _numba_sub3, + _numba_mul3, + _numba_mul3_scalar, + _numba_div3, + _numba_div3_scalar, + _numba_sqrt3, + _numba_norm3, + _numba_dot3, + _numba_cross3, +) + +def test_numba_add3(): + """ensure _numba_add3 and _numba_add3_scalar work as expected.""" + assert _numba_add3((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)) == (5.0, 7.0, 9.0) + assert _numba_add3((1,2,3), (4.0, 5.0, 6)) == (5.0, 7.0, 9) + assert _numba_add3([1,2,3], [4,5,-6]) == (5, 7, -3) + assert _numba_add3(np.array([1,2,3]), np.array([4,5,6])) == (5, 7, 9) + assert _numba_add3_scalar((1.0, 2.0, 3.0), 10.0) == (11.0, 12.0, 13.0) + assert _numba_add3_scalar([1,2,3], 10.0) == (11.0, 12.0, 13.0) + assert _numba_add3_scalar(np.array([1,2,3]), -10) == (-9, -8, -7) + with pytest.raises(TypeError, match="can't unbox heterogeneous list"): + _numba_add3_scalar([1,2,3.0], 10) # numba doesn't like [1,2,3.0] + assert _numba_add3_scalar(np.array([1,2,3]), 10) == (11, 12, 13) + assert _numba_add3_scalar(np.array([1,2,3.0]), 10.0) == (11.0, 12.0, 13.0) + +def test_numba_sub3(): + """ensure _numba_sub3 works as expected.""" + assert _numba_sub3((1.0, 2.0, 3.0), (4.0, 6.0, -5.0)) == (-3.0, -4.0, 8.0) + assert _numba_sub3([1,2,3], [4,5,-6]) == (-3, -3, 9) + assert _numba_sub3(np.array([1,2,3]), np.array([4, 6, -5])) == (-3, -4, 8) + +def test_numba_mul3(): + """ensure _numba_mul3 and _numba_mul3_scalar work as expected.""" + assert _numba_mul3((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)) == (4.0, 10.0, 18.0) + assert _numba_mul3([1,2,3], [4,5,-6]) == (4, 10, -18) + assert _numba_mul3(np.array([1,2,3]), np.array([4,5,-6])) == (4, 10, -18) + assert _numba_mul3_scalar((1.0, 2.0, 3.0), 10.0) == (10.0, 20.0, 30.0) + assert _numba_mul3_scalar([1,2,3], -10) == (-10, -20, -30) + assert _numba_mul3_scalar(np.array([1,2,3]), 10) == (10, 20, 30) + +def test_numba_div3(): + """ensure _numba_div3 and _numba_div3_scalar work as expected.""" + assert _numba_div3((4.0, 10.0, 18.0), (2.0, 5.0, 6.0)) == (2.0, 2.0, 3.0) + assert _numba_div3([4,10,18], [2,5,-6]) == (2, 2, -3) + assert _numba_div3(np.array([4,10,18]), np.array([2,5,-6])) == (2, 2, -3) + assert _numba_div3_scalar((10.0, 20.0, 30.0), 10.0) == (1.0, 2.0, 3.0) + assert _numba_div3_scalar([10,20,30], -10) == (-1, -2, -3) + assert _numba_div3_scalar(np.array([10,20,30]), 10) == (1, 2, 3) + +def test_numba_sqrt3(): + """ensure _numba_sqrt3 works as expected.""" + assert _numba_sqrt3((4.0, 9.0, 16.0)) == (2.0, 3.0, 4.0) + assert _numba_sqrt3([4,9,16]) == (2.0, 3.0, 4.0) + assert _numba_sqrt3(np.array([4,9,16])) == (2.0, 3.0, 4.0) + +def test_numba_norm3(): + """ensure _numba_norm3 works as expected.""" + assert _numba_norm3((3.0, 4.0, 12.0)) == 13.0 + assert _numba_norm3([0, 3, 4]) == 5.0 + assert _numba_norm3(np.array([0, -2, 0])) == 2.0 + +def test_numba_dot3(): + """ensure _numba_dot3 works as expected.""" + assert _numba_dot3((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)) == 32.0 + assert _numba_dot3([1,2,3], [4,5,-6]) == -4 + assert _numba_dot3(np.array([1,10,0]), np.array([0,-1,5])) == -10 + +def test_numba_cross3(): + """ensure _numba_cross3 works as expected.""" + assert _numba_cross3((1, 2, 30), (4, 0, 6)) == (12, 114, -8) + assert _numba_cross3([0,1,0],[0,0,10]) == (10,0,0) + assert _numba_cross3(np.array([1,10,0]), np.array([0,-1,5])) == (50, -5, -1)