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)

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01814GKVzZVL3FuF3DeiZXhv

scripts/cloud/tasks/p1p2.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+# Phase P1/P2 validation on the 1000 Genomes Omni2.5 set + `simurg ingest` at scale.
+# Args: <data dir with bpm/egt/csv and data/>  <results dir>
+set -uo pipefail
+D=$1; OUT=$2
+.venv/bin/python scripts/validate_p1_large.py --root $D/data --bpm $D/HumanOmni2.5-4v1-Multi_B.bpm \
+  --egt $D/HumanOmni2.5-4v1-Multi_B.egt --csv $D/HumanOmni2.5-4v1_B.csv --out $OUT/summary.json \
+  --workers "$(nproc)" || echo "validation reported differences (see summary.json)"
+/usr/bin/time -v .venv/bin/simurg ingest --bpm $D/HumanOmni2.5-4v1-Multi_B.bpm --idats $D/data \
+  --out /work/ingest --threads "$(nproc)" --quiet 2> $OUT/ingest_time.txt || echo "ingest exit code $?"
+cp /work/ingest/ingest_summary.tsv $OUT/ingest_summary.tsv 2>/dev/null || true
+cp /work/ingest/run_manifest.json $OUT/ingest_run_manifest.json 2>/dev/null || true

scripts/cloud/tasks/p3.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# Phase P3: refit normalization transforms (1.2.0 and 1.1.2) for every vendor GTC and compare bit by bit.
+# Args: <data dir with bpm and data/>  <results dir>
+set -uo pipefail
+D=$1; OUT=$2
+/usr/bin/time -v .venv/bin/python scripts/validate_p3_large.py --root $D/data --bpm $D/HumanOmni2.5-4v1-Multi_B.bpm \
+  --out $OUT/summary.json --workers "$(nproc)" 2> $OUT/time.txt || echo "P3 validation reported differences"
+tail -25 $OUT/time.txt

scripts/cloud/trial.sh

@@ -18,6 +18,7 @@ MAX_PRICE=${MAX_PRICE:-0.30}        # USD/hour spot cap (on-demand is 0.6528)
 MAX_MINUTES=${MAX_MINUTES:-360}     # hard cap: 6 hours
 DISK_GB=${DISK_GB:-250}
 BUDGET_USD=${BUDGET_USD:-10}
+TASK=${TASK:-p3}                 # scripts/cloud/tasks/<TASK>.sh
 TAG=simurg-trial
 ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
 BUCKET=simurg-trial-$ACCOUNT
@@ -60,10 +61,12 @@ launch)
   echo "-> packing repository at $(git rev-parse --short HEAD) -> s3://$BUCKET/$RUN/"
   git archive --format=tar.gz -o /tmp/simurg-repo.tar.gz HEAD
   $AWS s3 cp /tmp/simurg-repo.tar.gz "s3://$BUCKET/$RUN/repo.tar.gz" --only-show-errors
-  sed -e "s/__BUCKET__/$BUCKET/" -e "s/__RUN__/$RUN/" -e "s/__MAX_MINUTES__/$MAX_MINUTES/" \
+  [[ -f scripts/cloud/tasks/$TASK.sh ]] || { echo "unknown TASK $TASK"; exit 2; }
+  sed -e "s/__BUCKET__/$BUCKET/" -e "s/__RUN__/$RUN/" -e "s/__MAX_MINUTES__/$MAX_MINUTES/" -e "s/__TASK__/$TASK/g" \
     scripts/cloud/userdata.sh > /tmp/simurg-userdata.sh
   AMI=$($AWS ssm get-parameter --name /aws/service/canonical/ubuntu/server/24.04/stable/current/arm64/hvm/ebs-gp3/ami-id \
     --query Parameter.Value --output text)
+  echo "-> task: $TASK"
   echo "-> launching $TYPE spot (cap \$$MAX_PRICE/h, hard limit ${MAX_MINUTES} min, ${DISK_GB} GB disk deleted on exit)"
   ID=$($AWS ec2 run-instances --image-id "$AMI" --instance-type "$TYPE" --count 1 \
     --iam-instance-profile Name=$ROLE \

scripts/cloud/userdata.sh

@@ -54,20 +54,10 @@ echo "=== streaming 107 GB IDAT+GTC archive (no copy of the .tgz is kept) $(date
 $S3 cp --no-sign-request "$SRC/Omni25_idats_gtcs_2141_samples.tgz" - | pigz -dc | tar -x -C $D/data
 echo "=== extracted $(du -sh $D/data | cut -f1) at $(date -u +%T) ==="
 
-echo "=== validation $(date -u +%T) ==="
+echo "=== unit tests $(date -u +%T) ==="
 .venv/bin/python -m pytest tests/unit -q -p no:cacheprovider || echo "unit tests reported failures (continuing)"
-.venv/bin/python scripts/validate_p1_large.py --root $D/data --bpm $D/HumanOmni2.5-4v1-Multi_B.bpm \
-  --egt $D/HumanOmni2.5-4v1-Multi_B.egt --csv $D/HumanOmni2.5-4v1_B.csv --out /work/summary.json \
-  --workers "$(nproc)" || echo "validation reported differences (see summary.json)"
-$S3 cp /work/summary.json "s3://$BUCKET/$RUN/summary.json" --only-show-errors
-
-echo "=== simurg ingest at scale $(date -u +%T) ==="
-/usr/bin/time -v .venv/bin/simurg ingest --bpm $D/HumanOmni2.5-4v1-Multi_B.bpm --idats $D/data \
-  --out /work/ingest --threads "$(nproc)" --quiet 2> /work/ingest_time.txt || echo "ingest exit code $?"
-tail -25 /work/ingest_time.txt
-for f in ingest_summary.tsv run_manifest.json; do
-  $S3 cp "/work/ingest/$f" "s3://$BUCKET/$RUN/ingest_$f" --only-show-errors || true
-done
-$S3 cp /work/ingest_time.txt "s3://$BUCKET/$RUN/ingest_time.txt" --only-show-errors || true
-echo "ingest output size: $(du -sh /work/ingest | cut -f1)"
+echo "=== task __TASK__ $(date -u +%T) ==="
+mkdir -p /work/results
+bash scripts/cloud/tasks/__TASK__.sh $D /work/results || echo "task exit code $?"
+$S3 cp /work/results "s3://$BUCKET/$RUN/" --recursive --only-show-errors || true
 echo "=== done $(date -u +%T) ==="

scripts/validate_p3_large.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Phase P3 validation on a large set of vendor GTCs: refit every sample's normalization transforms from the GTC's
+raw X/Y (normalization 1.2.0 and 1.1.2) and compare them with the transforms stored in the GTC.
+
+Usage: validate_p3_large.py --root DIR --bpm F --out summary.json [--workers N] [--limit N]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import multiprocessing as mp
+import os
+import sys
+import time
+import traceback
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+FIELDS = ("offset_x", "offset_y", "scale_x", "scale_y", "shear", "theta")
+VERSIONS = ("1.2.0", "1.1.2")
+_S: dict[str, Any] = {}
+
+
+def _init(bins: np.ndarray) -> None:  # type: ignore[type-arg]
+    _S["bins"] = bins
+
+
+def _ulps(a: np.ndarray, b: np.ndarray) -> np.ndarray:  # type: ignore[type-arg]
+    ai = a.astype(np.float32).view(np.int32).astype(np.int64)
+    bi = b.astype(np.float32).view(np.int32).astype(np.int64)
+    # map sign-magnitude to a monotonic integer line so ULP distance works across zero
+    ai = np.where(ai < 0, -(ai & 0x7FFFFFFF), ai)
+    bi = np.where(bi < 0, -(bi & 0x7FFFFFFF), bi)
+    return np.abs(ai - bi)
+
+
+def check(path: str) -> dict[str, Any]:
+    from simurg.formats.gtc import read_gtc
+    from simurg.signal.normalization.fit import fit_transforms
+
+    res: dict[str, Any] = {"gtc": path}
+    try:
+        g = read_gtc(path, check_write_complete=False)
+        ref = g.normalization_transforms
+        refm = np.array([[r[f] for f in FIELDS] for r in ref], dtype=np.float32)
+        refr = np.array([r["reserved"] for r in ref], dtype=np.float32)
+        res["autocall_version"] = g.autocall_version
+        for v in VERSIONS:
+            xf = fit_transforms(g.raw_x, g.raw_y, _S["bins"], v)
+            mine = np.array([[getattr(x, f) for f in FIELDS] for x in xf], dtype=np.float32)
+            mres = np.array([[x.cvx, x.cvy, x.nn12, x.rr12, x.taa, x.tbb] for x in xf], dtype=np.float32)
+            if mine.shape != refm.shape:
+                res[v] = {"error": f"{len(xf)} bins vs {len(ref)} in GTC"}
+                continue
+            u = _ulps(mine, refm)
+            ur = _ulps(mres, refr)
+            res[v] = {
+                "bins_exact": int((u.max(axis=1) == 0).sum()),
+                "n_bins": len(ref),
+                "max_ulp": int(u.max()),
+                "max_ulp_per_field": dict(zip(FIELDS, u.max(axis=0).tolist(), strict=True)),
+                "reserved_exact": bool((ur == 0).all()),
+                "reserved_ref_all_zero": bool((refr == 0).all()),
+            }
+    except Exception as e:
+        res["error"] = f"{type(e).__name__}: {e}"
+        res["trace"] = traceback.format_exc(limit=3)
+    return res
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", required=True, type=Path)
+    ap.add_argument("--bpm", required=True, type=Path)
+    ap.add_argument("--out", required=True, type=Path)
+    ap.add_argument("--workers", type=int, default=os.cpu_count() or 1)
+    ap.add_argument("--limit", type=int, default=0)
+    a = ap.parse_args()
+
+    from simurg.formats.bpm import read_bpm
+    from simurg.formats.discovery import walk_files
+
+    t0 = time.time()
+    gtcs = [str(p) for p in walk_files(a.root, (".gtc",))]
+    if a.limit:
+        gtcs = gtcs[: a.limit]
+    bins = read_bpm(a.bpm).norm_bin_ids()
+    with ProcessPoolExecutor(a.workers, mp_context=mp.get_context("spawn"), initializer=_init, initargs=(bins,)) as ex:
+        results = list(ex.map(check, gtcs, chunksize=2))
+
+    summary: dict[str, Any] = {"n_gtc": len(gtcs), "errors": [r for r in results if "error" in r][:20]}
+    for v in VERSIONS:
+        rs = [r[v] for r in results if isinstance(r.get(v), dict) and "bins_exact" in r[v]]
+        summary[v] = {
+            "samples_all_bins_exact": sum(r["bins_exact"] == r["n_bins"] for r in rs),
+            "bins_exact": sum(r["bins_exact"] for r in rs),
+            "bins_total": sum(r["n_bins"] for r in rs),
+            "max_ulp": max((r["max_ulp"] for r in rs), default=None),
+            "ulp_histogram": dict(
+                sorted(
+                    {
+                        str(k): int(n)
+                        for k, n in zip(*np.unique([r["max_ulp"] for r in rs], return_counts=True), strict=True)
+                    }.items(),
+                    key=lambda kv: int(kv[0]),
+                )
+            ),
+            "max_ulp_per_field": {f: max((r["max_ulp_per_field"][f] for r in rs), default=None) for f in FIELDS},
+            "reserved_exact_samples": sum(r["reserved_exact"] for r in rs),
+            "reserved_ref_all_zero_samples": sum(r["reserved_ref_all_zero"] for r in rs),
+        }
+    summary["autocall_versions"] = sorted({r.get("autocall_version", "") for r in results})
+    worst = sorted(
+        (r for r in results if isinstance(r.get("1.1.2"), dict) and "max_ulp" in r["1.1.2"]),
+        key=lambda r: -min(r["1.1.2"]["max_ulp"], r["1.2.0"].get("max_ulp", 1 << 40)),
+    )[:10]
+    summary["worst_samples"] = worst
+    summary["seconds"] = round(time.time() - t0, 1)
+    best = max(VERSIONS, key=lambda v: summary[v]["bins_exact"])
+    summary["best_matching_version"] = best
+    summary["passed"] = summary[best]["bins_exact"] == summary[best]["bins_total"] > 0
+    a.out.parent.mkdir(parents=True, exist_ok=True)
+    a.out.write_text(json.dumps(summary, indent=2, default=str) + "\n")
+    print(
+        json.dumps(
+            {k: summary[k] for k in ("n_gtc", "best_matching_version", "passed", "seconds")}
+            | {
+                v: {k: summary[v][k] for k in ("samples_all_bins_exact", "bins_exact", "bins_total", "max_ulp")}
+                for v in VERSIONS
+            }
+        )
+    )
+    return 0 if summary["passed"] else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/simurg/signal/normalization/fit.py

@@ -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