add numba_math.py; optimize coordinates.py and _small_angle_of_2_vectors

test/grid/geometry/test_intersections.py

@@ -3,7 +3,7 @@
 import uxarray as ux
 from uxarray.constants import ERROR_TOLERANCE
 from uxarray.grid.arcs import extreme_gca_z, on_minor_arc
-from uxarray.grid.coordinates import _lonlat_rad_to_xyz, _xyz_to_lonlat_rad,_xyz_to_lonlat_rad_scalar
+from uxarray.grid.coordinates import _lonlat_rad_to_xyz, _xyz_to_lonlat_rad
 from uxarray.grid.intersections import gca_gca_intersection, gca_const_lat_intersection, get_number_of_intersections
 
 def test_get_GCA_GCA_intersections_antimeridian():

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/arcs.py

@@ -62,7 +62,7 @@ def point_within_gca(pt_xyz, gca_a_xyz, gca_b_xyz):
     -----
     - The function ensures that the point lies on the same plane as the GCA before performing interval checks.
     - It assumes the input represents the smaller arc of the Great Circle.
-    - The `_angle_of_2_vectors` and `_xyz_to_lonlat_rad_scalar` functions are used for calculations.
+    - The `_angle_of_2_vectors` function is used for calculations.
     """
     # 1. Check if the input GCA spans exactly 180 degrees
     angle_ab = _angle_of_2_vectors(gca_a_xyz, gca_b_xyz)

uxarray/grid/coordinates.py

@@ -1,4 +1,3 @@
-import math
 import warnings
 
 import numpy as np
@@ -9,13 +8,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, np.ndarray, np.ndarray] | tuple[float, float, float]:
     """Converts Spherical latitude and longitude coordinates into Cartesian x,
     y, z coordinates."""
     x = np.cos(lon) * np.cos(lat)
@@ -66,35 +69,12 @@ def _xyz_to_lonlat_rad_no_norm(
     return lon, lat
 
 
-@njit(cache=True)
-def _xyz_to_lonlat_rad_scalar(x, y, z, normalize=True):
-    if normalize:
-        x, y, z = _normalize_xyz_scalar(x, y, z)
-
-    lon = np.arctan2(y, x)
-    lat = np.asin(z)
-
-    # Set longitude range to [0, 2*pi]
-    lon = lon % (2 * math.pi)
-
-    z_abs = abs(z)
-    if z_abs > (1.0 - ERROR_TOLERANCE):
-        lat = math.copysign(math.pi / 2, z)
-        lon = 0.0
-
-    lonlat = np.empty(2)
-    lonlat[0] = lon
-    lonlat[1] = lat
-
-    return lonlat
-
-
 def _xyz_to_lonlat_rad(
     x: np.ndarray | float,
     y: np.ndarray | float,
     z: np.ndarray | float,
     normalize: bool = True,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+) -> tuple[np.ndarray, np.ndarray] | tuple[float, float]:
     """Converts Cartesian x, y, z coordinates in Spherical longitude and
     latitude coordinates in radians.
 
@@ -188,11 +168,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 +320,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 +369,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 +410,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 +443,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 +476,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 +541,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 +577,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, prange
 
 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,63 +18,24 @@ def _small_angle_of_2_vectors(u, v):
 
     Parameters
     ----------
-    u : numpy.ndarray or iterable of length 3
+    u : iterable of length 3
         The first 3D vector.
-    v : numpy.ndarray or iterable of length 3
+    v : iterable of length 3
         The second 3D vector.
 
     Returns
     -------
     float
         The smallest angle between `u` and `v` in radians.
     """
-    # don't convert to numpy array if not already numpy array.
-    # The formula is: angle = 2 * arctan2(| |v|*u - |u|*v |, | |v|*u + |u|*v |)
-    v_norm = _numba_norm3(v)
-    u_norm = _numba_norm3(u)
-    v_norm_times_u = (v_norm * u[0], v_norm * u[1], v_norm * u[2])
-    u_norm_times_v = (u_norm * v[0], u_norm * v[1], u_norm * v[2])
-    vec_minus = (
-        v_norm_times_u[0] - u_norm_times_v[0],
-        v_norm_times_u[1] - u_norm_times_v[1],
-        v_norm_times_u[2] - u_norm_times_v[2],
-    )
-    vec_sum = (
-        v_norm_times_u[0] + u_norm_times_v[0],
-        v_norm_times_u[1] + u_norm_times_v[1],
-        v_norm_times_u[2] + u_norm_times_v[2],
-    )
-    norm_vec_minus = _numba_norm3(vec_minus)
-    norm_vec_sum = _numba_norm3(vec_sum)
-    angle_u_v_rad = 2 * np.arctan2(norm_vec_minus, 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
 
 
-# TODO: move _numba_norm3 to a higher-level utils file. For more details, see issue #1648.
-@njit(cache=True)
-def _numba_norm3(u):
-    """
-    Compute the Euclidean norm of a 3D vector.
-    Implementation is currently equivalent to np.linalg.norm:
-        sqrt(u[0]**2 + u[1]**2 + u[2]**2)
-
-    Does NOT internally convert u to a list or numpy array;
-    utilizing tuples in numba instead of many tiny lists/arrays
-    can improve performance significantly.
-
-    Parameters
-    ----------
-    u : iterable of length 3, possibly a numpy array
-        The 3D vector.
-
-    Returns
-    -------
-    float
-        The Euclidean norm of the vector `u`.
-    """
-    return (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5
-
-
 @njit(cache=True)
 def _angle_of_2_vectors(u, v):
     """