P3: vectorised nearest-neighbour search; task-driven cloud trials

- fit.py: exact and binned nearest-neighbour searches use one distance
  matrix per searched bin (first-minimum rule kept); fit 2.6 s -> 1.4 s
  per GSA sample, still bit-identical to the oracle
- scripts/validate_p3_large.py: refits 1.2.0 and 1.1.2 transforms for
  every vendor GTC and reports ULP distances per field
- cloud trial runs scripts/cloud/tasks/<TASK>.sh (TASK=p3 default;
  p1p2 = previous validation + ingest)
This commit is contained in:
kemal committed 2026-09-27 22:57:03 +01:00
1 parent c0d556d4af
commit 7deabd907c
6 files changed
+197 -44

No files matched your search

+27 -28
View File
@@ -85,17 +85,18 @@ def closest_points(xref: FloatArr, yref: FloatArr, x: FloatArr, y: FloatArr) ->
def _closest_slow(xref: FloatArr, yref: FloatArr, x: FloatArr, y: FloatArr) -> IdxArr:
out = np.empty(len(x), dtype=np.int64)
for i in range(len(x)):
dx = x[i] - xref
dy = y[i] - yref
d = dx * dx + dy * dy # float32, as in the reference
out[i] = int(np.argmin(d)) # first minimum wins (strict <)
return out
"""Exact search: squared distances in float32; the first minimum wins (C strict ``<``)."""
dx = x[:, None] - xref[None, :]
dy = y[:, None] - yref[None, :]
return np.argmin(dx * dx + dy * dy, axis=1).astype(np.int64)
def _closest_binned(raw_x: FloatArr, raw_y: FloatArr, axis_x: FloatArr, axis_y: FloatArr) -> IdxArr:
"""QUIRK: the vendor's binned approximate search (``findClosestSitesToPointsAlongAxis``)."""
"""QUIRK: the vendor's binned approximate search (``findClosestSitesToPointsAlongAxis``).
Vectorised: sweep points that search the same bin are handled together with one distance matrix; within a bin
candidates keep ascending locus order, so ``argmin`` reproduces the reference's first-minimum rule.
"""
if (axis_x > F32(0.0001)).any():
if (axis_y > F32(0.0001)).any():
return np.full(len(axis_x), -1, dtype=np.int64)
@@ -107,30 +108,28 @@ def _closest_binned(raw_x: FloatArr, raw_y: FloatArr, axis_x: FloatArr, axis_y:
qa = (raw_a[eligible] / bin_width).astype(np.float32)
b1 = np.clip(trunc_int(qa), 0, 11)
b2 = np.clip(trunc_int((qa - F32(0.5)).astype(np.float32)), 0, 10)
bins = [eligible[b1 == k] for k in range(12)]
shifted = [eligible[b2 == k] for k in range(11)]
# which bin each sweep point searches: 0..11 = regular bins, 12..22 = shifted bins
q = (axis_a / bin_width).astype(np.float32)
k_raw = trunc_int(q)
rem = (q - k_raw.astype(np.float32)).astype(np.float32)
k = np.clip(k_raw, 0, 11)
mid = (rem >= F32(0.25)) & (rem <= F32(0.75))
low = rem < F32(0.25)
choice = np.where(mid, k, np.where(low, np.where(k == 0, k, 12 + k - 1), np.where(k == 11, k, 12 + k)))
out = np.full(len(axis_a), -1, dtype=np.int64)
for i in range(len(axis_a)):
q = F32(axis_a[i] / bin_width)
k = int(np.trunc(q))
rem = F32(q - F32(k))
k = min(max(k, 0), 11)
if F32(0.25) <= rem <= F32(0.75):
cand = bins[k]
elif rem < F32(0.25):
cand = bins[k] if k == 0 else shifted[k - 1]
elif k == 11:
cand = bins[k]
else:
cand = shifted[k]
for c in np.unique(choice):
cand = eligible[b1 == c] if c < 12 else eligible[b2 == c - 12]
if len(cand) == 0:
continue
dx = raw_x[cand] - axis_x[i]
dy = raw_y[cand] - axis_y[i]
pts = np.flatnonzero(choice == c)
dx = raw_x[cand][None, :] - axis_x[pts][:, None]
dy = raw_y[cand][None, :] - axis_y[pts][:, None]
d = dx * dx + dy * dy
j = int(np.argmin(d))
if d[j] < 1e20:
out[i] = int(cand[j])
j = np.argmin(d, axis=1)
best = d[np.arange(len(pts)), j]
out[pts] = np.where(best < 1e20, cand[j], -1)
return out