1import json
  2import math
  3import time
  4from pathlib import Path
  5
  6import numpy as np
  7from scipy.integrate import quad
  8from scipy.stats import beta
  9
 10SEED = 3044
 11rng = np.random.default_rng(SEED)
 12
 13
 14def exact_tau(m, d):
 15    """Theorem 7.1: expected Kendall correlation for one distance comparison."""
 16    f = lambda b: math.asin(math.sqrt(b)) * beta.pdf(b, m / 2.0, (d - m) / 2.0)
 17    return 2.0 / math.pi * quad(f, 0.0, 1.0, epsabs=2e-8)[0]
 18
 19
 20def asymptotic_tau(m, d):
 21    return 2.0 / math.pi * math.sqrt(m / d)
 22
 23
 24def make_projection(m, d, seed):
 25    # Haar random rank-m projector represented by an orthonormal-row matrix.
 26    local = np.random.default_rng(seed)
 27    q, _ = np.linalg.qr(local.normal(size=(d, m)), mode="reduced")
 28    return q.T.astype(np.float32)
 29
 30
 31def pairwise_order_check(d, m, n_pairs=30000, seed=0):
 32    local = np.random.default_rng(seed)
 33    r = make_projection(m, d, seed + 10000)
 34    # Independent triples are the exact theorem setup.
 35    x0 = local.normal(size=(n_pairs, d)).astype(np.float32)
 36    x1 = local.normal(size=(n_pairs, d)).astype(np.float32)
 37    x2 = local.normal(size=(n_pairs, d)).astype(np.float32)
 38    full = ((x1 - x0) ** 2).sum(1) - ((x2 - x0) ** 2).sum(1)
 39    y0, y1, y2 = x0 @ r.T, x1 @ r.T, x2 @ r.T
 40    proj = ((y1 - y0) ** 2).sum(1) - ((y2 - y0) ** 2).sum(1)
 41    agreement = np.mean(np.sign(full) == np.sign(proj))
 42    tau = 2 * agreement - 1
 43    return float(tau), float(agreement)
 44
 45
 46def recall_at(retrieved, truth, k):
 47    return float(np.mean([len(set(a[:k]) & set(b[:k])) / k for a, b in zip(retrieved, truth)]))
 48
 49
 50def adaptive_dimension(d, target_tau=0.5, safety_margin=0.0, max_rounds=8):
 51    """Select m by the proposed formula, then expand using held-out pairwise tau."""
 52    m = max(8, int(math.ceil(d * (math.pi * target_tau / 2) ** 2)))
 53    m = min(d, int(math.ceil(m / 8) * 8))
 54    history = []
 55    for round_id in range(max_rounds):
 56        tau, _ = pairwise_order_check(d, m, n_pairs=12000,
 57                                      seed=SEED + 5000 + round_id)
 58        history.append({"round": round_id, "m": m, "measured_tau": tau})
 59        if tau >= target_tau - safety_margin or m >= d:
 60            break
 61        m = min(d, max(m + 1, int(math.ceil(1.25 * m / 8) * 8)))
 62    return m, history
 63
 64def retrieval_trial(d=128, n=4000, nq=300, target_tau=0.5):
 65    local = np.random.default_rng(SEED + 77)
 66    # A mildly clustered retrieval-like data set: exact neighbors share a latent center.
 67    n_centers = 80
 68    centers = local.normal(size=(n_centers, d)).astype(np.float32)
 69    labels = local.integers(0, n_centers, size=n)
 70    data = centers[labels] + 0.8 * local.normal(size=(n, d)).astype(np.float32)
 71    qlabels = local.integers(0, n_centers, size=nq)
 72    queries = centers[qlabels] + 0.8 * local.normal(size=(nq, d)).astype(np.float32)
 73    # Full-space truth and candidate pair ranking metric.
 74    full_dist = ((queries[:, None, :] - data[None, :, :]) ** 2).sum(2)
 75    truth = np.argsort(full_dist, axis=1)
 76    # Standard conservative JL sizing (one common constant); cap at d because no
 77    # projection can improve on the original representation.
 78    eps = 0.5
 79    jl_m = min(d, int(math.ceil(8 * math.log(n) / (eps * eps))))
 80    rank_m = max(8, int(math.ceil(d * (math.pi * target_tau / 2) ** 2)))
 81    # Hardware-friendly rounding, while retaining the requested target budget.
 82    rank_m = min(d, int(math.ceil(rank_m / 8) * 8))
 83    rows = {"jl": jl_m, "ranking_rule": rank_m}
 84    out = {}
 85    for name, m in rows.items():
 86        t0 = time.perf_counter()
 87        r = make_projection(m, d, SEED + m)
 88        zd, zq = data @ r.T, queries @ r.T
 89        pd = ((zq[:, None, :] - zd[None, :, :]) ** 2).sum(2)
 90        got = np.argsort(pd, axis=1)
 91        # Kendall tau over the top-20 candidate ordering, measured against full space.
 92        taus = []
 93        for i in range(nq):
 94            ids = truth[i, :20]
 95            a = full_dist[i, ids]
 96            b = pd[i, ids]
 97            signs = np.sign((a[:, None] - a[None, :]) * (b[:, None] - b[None, :]))
 98            iu = np.triu_indices(len(ids), 1)
 99            taus.append(float(np.mean(signs[iu])))
100        out[name] = {
101            "m": m,
102            "memory_ratio_m_over_d": m / d,
103            "recall@1": recall_at(got, truth, 1),
104            "recall@10": recall_at(got, truth, 10),
105            "top20_kendall": float(np.mean(taus)),
106            "query_seconds": time.perf_counter() - t0,
107        }
108    return out
109
110
111def main():
112    verification = []
113    for d, m in [(64, 4), (64, 16), (64, 32), (128, 8), (128, 32), (128, 64), (256, 32), (256, 64)]:
114        empirical, agreement = pairwise_order_check(d, m, n_pairs=20000, seed=d + m)
115        verification.append({
116            "d": d, "m": m, "ratio": m / d,
117            "empirical_tau": empirical,
118            "empirical_agreement": agreement,
119            "exact_tau": exact_tau(m, d),
120            "asymptotic_tau": asymptotic_tau(m, d),
121        })
122    adaptive_m, adaptive_history = adaptive_dimension(128, target_tau=0.5)
123    result = {"seed": SEED, "verification": verification,
124              "adaptive_selection": {"final_m": adaptive_m, "history": adaptive_history},
125              "retrieval": retrieval_trial()}
126    Path("results.json").write_text(json.dumps(result, indent=2))
127    print(json.dumps(result, indent=2))
128
129
130if __name__ == "__main__":
131    main()