Scalarized gca_gca_intersection

benchmarks/geometry_kernels.py

@@ -147,28 +147,28 @@ class GCAGCAIntersection:
 
     def setup(self):
         from uxarray.grid.intersections import (
-            _accux_gca,
-            _try_gca_gca_intersection,
-            gca_gca_intersection,
+            _accux_gca_scalar,
+            _try_gca_gca_intersection_scalar,
+            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_scalar = _accux_gca_scalar
+        self._try_gca_gca_intersection_scalar = _try_gca_gca_intersection_scalar
+        _accux_gca_scalar(*_W0, *_W1, *_V0, *_V1)
+        _try_gca_gca_intersection_scalar(*_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_scalar(*_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_scalar(*_W0, *_W1, *_V0, *_V1)
 
     def time_gca_gca_intersection(self):
         """Layer 3: dispatcher (full public API)."""

uxarray/grid/intersections.py

@@ -252,8 +252,10 @@ def faces_within_lat_bounds(lats, face_bounds_lat):
 
 
 @njit(cache=True, inline="always", error_model="numpy")
-def _accux_gca(w0, w1, v0, v1):
-    """Compute the candidate intersection points of two great-circle arcs.
+def _accux_gca_scalar(w00, w01, w02, w10, w11, w12, v00, v01, v02, v10, v11, v12):
+    """Scalar-argument form of :func:`_accux_gca`: returns the six components.
+
+    Compute the candidate intersection points of two great-circle arcs.
 
     Pure numerical kernel (mirrors AccuSphGeom ``accux_gca``).
 
@@ -262,21 +264,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,22 +302,67 @@ 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_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, inline="always", error_model="numpy")
+def _accux_gca(w0, w1, v0, v1):
+    """Compute the candidate intersection points of two great-circle arcs.
+
+    Pure numerical kernel (mirrors AccuSphGeom ``accux_gca``).
+
+    Computes the two antipodal candidate intersection points of the great-circle
+    arcs w0-w1 and v0-v1.  No branching, no validity filtering.
+
+    Parameters
+    ----------
+    w0, w1 : np.ndarray, shape (3,)
+        Cartesian endpoints of the first arc.
+    v0, v1 : np.ndarray, shape (3,)
+        Cartesian endpoints of the second arc.
+
+    Returns
+    -------
+    pos, neg : np.ndarray, shape (3,)
+        Two antipodal candidate unit vectors.
+    """
+    pos_x, pos_y, pos_z, neg_x, neg_y, neg_z = _accux_gca_scalar(
+        w0[0],
+        w0[1],
+        w0[2],
+        w1[0],
+        w1[1],
+        w1[2],
+        v0[0],
+        v0[1],
+        v0[2],
+        v1[0],
+        v1[1],
+        v1[2],
+    )
     pos = np.empty(3)
-    pos[0] = vx * inv
-    pos[1] = vy * inv
-    pos[2] = vz * inv
+    pos[0] = pos_x
+    pos[1] = pos_y
+    pos[2] = pos_z
     neg = np.empty(3)
-    neg[0] = -pos[0]
-    neg[1] = -pos[1]
-    neg[2] = -pos[2]
+    neg[0] = neg_x
+    neg[1] = neg_y
+    neg[2] = neg_z
     return pos, neg
 
 
-@njit(cache=True, error_model="numpy")
-def _try_gca_gca_intersection(w0, w1, v0, v1):
-    """Select the valid great-circle intersection and report a status code.
+@njit(cache=True, inline="always", error_model="numpy")
+def _try_gca_gca_intersection_scalar(
+    w00, w01, w02, w10, w11, w12, v00, v01, v02, v10, v11, v12
+):
+    """Scalar-argument form of :func:`_try_gca_gca_intersection`.
+
+    Select the valid great-circle intersection and report a status code.
 
     Batch/status layer (mirrors AccuSphGeom ``try_gca_gca_intersection``).
 
@@ -327,37 +374,87 @@ 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_scalar(
+        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_x, point_y, point_z, status, px, py, pz, ngx, ngy, ngz
+
+
+@njit(cache=True, error_model="numpy")
+def _try_gca_gca_intersection(w0, w1, v0, v1):
+    """Select the valid great-circle intersection and report a status code.
+
+    Batch/status layer (mirrors AccuSphGeom ``try_gca_gca_intersection``).
+
+    Calls the pure numerical kernel, applies integer mask arithmetic to determine
+    validity, selects the output point without if/else branching in the hot path.
+
+    Status codes mirror AccuSphGeom:
+        0  exactly one candidate is valid
+        1  both candidates are valid
+        2  neither candidate is valid  (includes coplanar/parallel case)
+    """
+    (
+        point_x,
+        point_y,
+        point_z,
+        status,
+        pos_x,
+        pos_y,
+        pos_z,
+        neg_x,
+        neg_y,
+        neg_z,
+    ) = _try_gca_gca_intersection_scalar(
+        w0[0],
+        w0[1],
+        w0[2],
+        w1[0],
+        w1[1],
+        w1[2],
+        v0[0],
+        v0[1],
+        v0[2],
+        v1[0],
+        v1[1],
+        v1[2],
+    )
+
+    point = np.empty(3)
+    point[0] = point_x
+    point[1] = point_y
+    point[2] = point_z
+    pos = np.empty(3)
+    pos[0] = pos_x
+    pos[1] = pos_y
+    pos[2] = pos_z
+    neg = np.empty(3)
+    neg[0] = neg_x
+    neg[1] = neg_y
+    neg[2] = neg_z
     return point, status, pos, neg
 
 
@@ -396,40 +493,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_scalar(
+        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]