diff --git a/benchmarks/geometry_kernels.py b/benchmarks/geometry_kernels.py index 08aea236e..3892c0911 100644 --- a/benchmarks/geometry_kernels.py +++ b/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).""" diff --git a/uxarray/grid/intersections.py b/uxarray/grid/intersections.py index 1e6100854..d858eca06 100644 --- a/uxarray/grid/intersections.py +++ b/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,22 +324,18 @@ 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 @@ -350,15 +343,14 @@ def _try_gca_gca_intersection(w0, w1, v0, v1): 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]