add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors

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)

uxarray/grid/coordinates.py

@@ -9,13 +9,17 @@
 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 +86,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 +102,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 +196,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 +348,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 +397,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 +438,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 +471,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 +504,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 +569,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 +605,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):

uxarray/grid/utils.py

@@ -3,6 +3,12 @@
 from numba import njit
 
 from uxarray.constants import INT_FILL_VALUE
+from uxarray.utils.numba_math import (
+    _numba_add3,
+    _numba_mul3_scalar,
+    _numba_norm3,
+    _numba_sub3,
+)
 
 
 @njit(cache=True)
@@ -12,21 +18,21 @@ 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
     -------
     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
 
 

uxarray/utils/numba_math.py

@@ -0,0 +1,104 @@
+"""
+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)