Scalarized gca_gca_intersection

benchmarks/geometry_kernels.py

@@ -149,26 +149,26 @@ def setup(self):
         from uxarray.grid.intersections import (
             _accux_gca,
             _try_gca_gca_intersection,
-            gca_gca_intersection,
+            gca_gca_intersection
         )
 
-        self._accux_gca = _accux_gca
-        self._try_gca_gca_intersection = _try_gca_gca_intersection
         self.gca_gca_intersection = gca_gca_intersection
-
         self.gca_a = np.stack([_W0, _W1])
         self.gca_b = np.stack([_V0, _V1])
-        _accux_gca(_W0, _W1, _V0, _V1)
-        _try_gca_gca_intersection(_W0, _W1, _V0, _V1)
         gca_gca_intersection(self.gca_a, self.gca_b)
 
+        self._accux_gca = _accux_gca
+        self._try_gca_gca_intersection = _try_gca_gca_intersection
+        _accux_gca(*_W0, *_W1, *_V0, *_V1)
+        _try_gca_gca_intersection(*_W0, *_W1, *_V0, *_V1)
+
     def time_accux_gca_kernel(self):
         """Layer 1: pure numerical kernel."""
-        self._accux_gca(_W0, _W1, _V0, _V1)
+        self._accux_gca(*_W0, *_W1, *_V0, *_V1)
 
     def time_try_gca_gca_intersection(self):
         """Layer 2: batch/status layer."""
-        self._try_gca_gca_intersection(_W0, _W1, _V0, _V1)
+        self._try_gca_gca_intersection(*_W0, *_W1, *_V0, *_V1)
 
     def time_gca_gca_intersection(self):
         """Layer 3: dispatcher (full public API)."""

uxarray/grid/intersections.py

@@ -252,7 +252,7 @@ def faces_within_lat_bounds(lats, face_bounds_lat):
 
 
 @njit(cache=True, inline="always", error_model="numpy")
-def _accux_gca(w0, w1, v0, v1):
+def _accux_gca(w00, w01, w02, w10, w11, w12, v00, v01, v02, v10, v11, v12):
     """Compute the candidate intersection points of two great-circle arcs.
 
     Pure numerical kernel (mirrors AccuSphGeom ``accux_gca``).
@@ -262,21 +262,21 @@ def _accux_gca(w0, w1, v0, v1):
 
     Parameters
     ----------
-    w0, w1 : np.ndarray, shape (3,)
+    w00, w01, w02, w10, w11, w12 : float
         Cartesian endpoints of the first arc.
-    v0, v1 : np.ndarray, shape (3,)
+    v00, v01, v02, v10, v11, v12 : float
         Cartesian endpoints of the second arc.
 
     Returns
     -------
-    pos, neg : np.ndarray, shape (3,)
+    pos_x, pos_y, pos_z, neg_x, neg_y, neg_z : float
         Two antipodal candidate unit vectors.
     """
     n1x_hi, n1y_hi, n1z_hi, n1x_lo, n1y_lo, n1z_lo = accucross(
-        w0[0], w0[1], w0[2], w1[0], w1[1], w1[2]
+        w00, w01, w02, w10, w11, w12
     )
     n2x_hi, n2y_hi, n2z_hi, n2x_lo, n2y_lo, n2z_lo = accucross(
-        v0[0], v0[1], v0[2], v1[0], v1[1], v1[2]
+        v00, v01, v02, v10, v11, v12
     )
     vx_hi, vy_hi, vz_hi, vx_lo, vy_lo, vz_lo = accucross_pair(
         n1x_hi,
@@ -300,21 +300,18 @@ def _accux_gca(w0, w1, v0, v1):
     sum_hi, sum_lo = _sum_of_squares_c((vx_hi, vy_hi, vz_hi), (vx_lo, vy_lo, vz_lo))
     vn, _ = acc_sqrt_re(sum_hi, sum_lo)
     # vn==0 (coplanar arcs) yields inf via IEEE division under error_model="numpy",
-    # so pos/neg become non-finite and the status layer masks them out. Branch-free.
+    # so the candidates become non-finite and the status layer masks them out.
     inv = 1.0 / vn
-    pos = np.empty(3)
-    pos[0] = vx * inv
-    pos[1] = vy * inv
-    pos[2] = vz * inv
-    neg = np.empty(3)
-    neg[0] = -pos[0]
-    neg[1] = -pos[1]
-    neg[2] = -pos[2]
-    return pos, neg
+    pos_x = vx * inv
+    pos_y = vy * inv
+    pos_z = vz * inv
+    return pos_x, pos_y, pos_z, -pos_x, -pos_y, -pos_z
 
 
-@njit(cache=True, error_model="numpy")
-def _try_gca_gca_intersection(w0, w1, v0, v1):
+@njit(cache=True, inline="always", error_model="numpy")
+def _try_gca_gca_intersection(
+    w00, w01, w02, w10, w11, w12, v00, v01, v02, v10, v11, v12
+):
     """Select the valid great-circle intersection and report a status code.
 
     Batch/status layer (mirrors AccuSphGeom ``try_gca_gca_intersection``).
@@ -327,38 +324,33 @@ def _try_gca_gca_intersection(w0, w1, v0, v1):
         1  both candidates are valid
         2  neither candidate is valid  (includes coplanar/parallel case)
     """
-    pos, neg = _accux_gca(w0, w1, v0, v1)
-
-    pos_fin = (
-        int(math.isfinite(pos[0]))
-        * int(math.isfinite(pos[1]))
-        * int(math.isfinite(pos[2]))
+    px, py, pz, ngx, ngy, ngz = _accux_gca(
+        w00, w01, w02, w10, w11, w12, v00, v01, v02, v10, v11, v12
     )
+
+    pos_fin = int(math.isfinite(px)) * int(math.isfinite(py)) * int(math.isfinite(pz))
     neg_fin = (
-        int(math.isfinite(neg[0]))
-        * int(math.isfinite(neg[1]))
-        * int(math.isfinite(neg[2]))
+        int(math.isfinite(ngx)) * int(math.isfinite(ngy)) * int(math.isfinite(ngz))
     )
-    pos_on_a = pos_fin * on_minor_arc(pos, w0, w1)
-    pos_on_b = pos_fin * on_minor_arc(pos, v0, v1)
-    neg_on_a = neg_fin * on_minor_arc(neg, w0, w1)
-    neg_on_b = neg_fin * on_minor_arc(neg, v0, v1)
+    pos_on_a = pos_fin * _on_minor_arc_xyz(px, py, pz, w00, w01, w02, w10, w11, w12)
+    pos_on_b = pos_fin * _on_minor_arc_xyz(px, py, pz, v00, v01, v02, v10, v11, v12)
+    neg_on_a = neg_fin * _on_minor_arc_xyz(ngx, ngy, ngz, w00, w01, w02, w10, w11, w12)
+    neg_on_b = neg_fin * _on_minor_arc_xyz(ngx, ngy, ngz, v00, v01, v02, v10, v11, v12)
 
     pos_valid = pos_fin * pos_on_a * pos_on_b
     neg_valid = neg_fin * neg_on_a * neg_on_b
 
     pos_mask = pos_valid * (1 - neg_valid)
     neg_mask = neg_valid * (1 - pos_valid)
 
-    point = np.empty(3)
-    point[0] = pos_mask * pos[0] + neg_mask * neg[0]
-    point[1] = pos_mask * pos[1] + neg_mask * neg[1]
-    point[2] = pos_mask * pos[2] + neg_mask * neg[2]
+    point_x = pos_mask * px + neg_mask * ngx
+    point_y = pos_mask * py + neg_mask * ngy
+    point_z = pos_mask * pz + neg_mask * ngz
 
     both = pos_valid * neg_valid
     none = (1 - pos_valid) * (1 - neg_valid)
     status = both + none * 2
-    return point, status, pos, neg
+    return point_x, point_y, point_z, status, px, py, pz, ngx, ngy, ngz
 
 
 @njit(cache=True, error_model="numpy")
@@ -396,40 +388,62 @@ def gca_gca_intersection(gca_a_xyz, gca_b_xyz):
     if gca_a_xyz.shape[1] != 3 or gca_b_xyz.shape[1] != 3:
         raise DimensionError("The two GCAs must be in the cartesian [x, y, z] format")
 
-    w0 = gca_a_xyz[0]
-    w1 = gca_a_xyz[1]
-    v0 = gca_b_xyz[0]
-    v1 = gca_b_xyz[1]
-
-    point, status, pos, neg = _try_gca_gca_intersection(w0, w1, v0, v1)
+    # Unpack to scalars and run the allocation-free scalar chain.
+    w00 = gca_a_xyz[0, 0]
+    w01 = gca_a_xyz[0, 1]
+    w02 = gca_a_xyz[0, 2]
+    w10 = gca_a_xyz[1, 0]
+    w11 = gca_a_xyz[1, 1]
+    w12 = gca_a_xyz[1, 2]
+    v00 = gca_b_xyz[0, 0]
+    v01 = gca_b_xyz[0, 1]
+    v02 = gca_b_xyz[0, 2]
+    v10 = gca_b_xyz[1, 0]
+    v11 = gca_b_xyz[1, 1]
+    v12 = gca_b_xyz[1, 2]
+
+    (
+        point_x,
+        point_y,
+        point_z,
+        status,
+        pos_x,
+        pos_y,
+        pos_z,
+        neg_x,
+        neg_y,
+        neg_z,
+    ) = _try_gca_gca_intersection(
+        w00, w01, w02, w10, w11, w12, v00, v01, v02, v10, v11, v12
+    )
 
     res = np.empty((2, 3))
     count = 0
     if status == 0:
-        res[0, 0] = point[0]
-        res[0, 1] = point[1]
-        res[0, 2] = point[2]
+        res[0, 0] = point_x
+        res[0, 1] = point_y
+        res[0, 2] = point_z
         count = 1
     elif status == 1:
-        res[0, 0] = pos[0]
-        res[0, 1] = pos[1]
-        res[0, 2] = pos[2]
-        res[1, 0] = neg[0]
-        res[1, 1] = neg[1]
-        res[1, 2] = neg[2]
+        res[0, 0] = pos_x
+        res[0, 1] = pos_y
+        res[0, 2] = pos_z
+        res[1, 0] = neg_x
+        res[1, 1] = neg_y
+        res[1, 2] = neg_z
         count = 2
     else:
         # status == 2: no candidate on both arcs.
         # Check for coplanar overlap (shared endpoints) outside the kernel.
-        if on_minor_arc(v0, w0, w1):
-            res[count, 0] = v0[0]
-            res[count, 1] = v0[1]
-            res[count, 2] = v0[2]
+        if _on_minor_arc_xyz(v00, v01, v02, w00, w01, w02, w10, w11, w12):
+            res[count, 0] = v00
+            res[count, 1] = v01
+            res[count, 2] = v02
             count += 1
-        if on_minor_arc(v1, w0, w1):
-            res[count, 0] = v1[0]
-            res[count, 1] = v1[1]
-            res[count, 2] = v1[2]
+        if _on_minor_arc_xyz(v10, v11, v12, w00, w01, w02, w10, w11, w12):
+            res[count, 0] = v10
+            res[count, 1] = v11
+            res[count, 2] = v12
             count += 1
     return res[:count]