tidaldb/tidal/examples/ann_grid_search.rs

//! m12p3 ANN parameter grid search + quantization recall/memory frontier.
//!
//! Builds a `UsearchIndex` (HNSW) and an exact `BruteForceIndex` oracle over the
//! SAME deterministic, id-keyed corpus, then sweeps the HNSW parameters and the
//! quantization level — reporting, for each point, the **measured** recall@k vs
//! the exact oracle, the single-thread mean/p99 search latency, the build time,
//! and the **true** in-memory footprint (`UsearchIndex::memory_usage`, which
//! includes the proximity-graph links, not just the vectors).
//!
//! This is the tool that produces the documented `M`/`ef` and the F32/F16/Int8
//! recall+memory numbers the m12p3 exit gate requires — at the production shape
//! (1536-D). Recall here is a *property of the index given its parameters* and is
//! independent of load, so a deterministic single-thread harness is the right
//! instrument; the authoritative tail-latency-under-load number is the open-loop
//! `tidal-stress --verify-recall` ramp (m12p1).
//!
//! Run (local, 100k/1536-D — the exit-gate shape that fits one laptop):
//! ```bash
//! cargo run --release --example ann_grid_search -- \
//!     --corpus 100000 --dim 1536 --queries 200 --k 10
//! ```
//! The 1M shape needs ≈ 6 GB for the F32 oracle (1M × 1536 × 4 B) plus the HNSW;
//! run it on the k3s node with `--corpus 1000000`.

#![allow(
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss
)]

use std::{
    collections::HashSet,
    time::{Duration, Instant},
};

use tidaldb::storage::vector::{
    BruteForceIndex, DistanceMetric, QuantizationLevel, UsearchIndex, VectorId, VectorIndex,
    VectorIndexConfig,
};

// ---------------------------------------------------------------------------
// Deterministic corpus (SplitMix64) — reproducible recall ground truth
// ---------------------------------------------------------------------------

/// A deterministic, L2-normalized unit vector for `id` at dimensionality `dim`.
///
/// Same family the m12p1 recall harness uses (`SplitMix64`), so the corpus is
/// reproducible across runs and machines — the recall numbers are comparable.
fn unit_vector(id: u64, dim: usize) -> Vec<f32> {
    let mut state = id
        .wrapping_mul(0x9E37_79B9_7F4A_7C15)
        .wrapping_add(0x1234_5678_9ABC_DEF0);
    let mut next = || {
        state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = state;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^= z >> 31;
        // Map the high mantissa bits into [-0.5, 0.5).
        ((z >> 40) as f32 / (1u64 << 24) as f32) - 0.5
    };
    let mut v: Vec<f32> = (0..dim).map(|_| next()).collect();
    normalize(&mut v);
    v
}

/// L2-normalize in place (unit vector), falling back to the first axis if zero.
fn normalize(v: &mut [f32]) {
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > f32::EPSILON {
        for x in v.iter_mut() {
            *x /= norm;
        }
    } else if let Some(first) = v.first_mut() {
        *first = 1.0;
    }
}

/// A clustered (Gaussian-mixture) corpus vector for `id` — the faithful
/// representation of real embedding geometry, and the right instrument for an
/// honest recall@k.
///
/// Uniform-random high-dimensional vectors are PATHOLOGICAL for recall@k: by the
/// concentration of measure, every pair sits at cosine ≈ 0, so beyond a tiny
/// perturbation a query's "top-10" is an arbitrary draw from a thick equidistant
/// shell — recall@10 then measures impossible tie-breaking, not index quality,
/// and (worse) it gets *lower* as the corpus grows because the shell thickens.
/// (Measured: uniform-random recall@10 fell from ~0.97 at 10k to ~0.54 at 100k —
/// a corpus-size artifact, not an index regression.)
///
/// Real text/image embeddings instead live on a low-dimensional manifold with
/// clusters: a point's nearest neighbours are its cluster-mates, distinctly
/// closer than the bulk. We model that as a Gaussian mixture: `id` is assigned to
/// cluster `id % n_clusters`, and the vector is `center + spread · noise`,
/// normalized. With `spread = 0.5`, intra-cluster cosine ≈ 0.8 and inter-cluster
/// ≈ 0 — a clear neighbour structure, exactly the shape `related`/`for_you` reads
/// query against, where recall@k is a meaningful index-quality metric.
fn clustered_vector(centers: &[Vec<f32>], id: u64, dim: usize, spread: f32) -> Vec<f32> {
    let center = &centers[(id as usize) % centers.len()];
    let noise = unit_vector(id, dim);
    let mut v: Vec<f32> = center
        .iter()
        .zip(&noise)
        .map(|(c, nz)| c + spread * nz)
        .collect();
    normalize(&mut v);
    v
}

/// A realistic query: a corpus point (spread across the catalog) perturbed by
/// deterministic noise of magnitude `noise` RELATIVE to the unit base — so the
/// query sits a small, controlled angle off a real corpus point. This is the
/// shape of a production `related`/`for_you` read (a seed/preference vector near
/// real content), and the same realistic model the authoritative m12p1 recall
/// harness (`tidal-stress` `QueryPool`) uses.
///
/// Uniform-random high-dim queries are the WRONG instrument: in 1536-D their
/// neighbours are near-equidistant (distance concentration), so recall looks
/// pathologically low for reasons unrelated to the index.
///
/// The noise is scaled by `1/√(dim/12)` so the additive perturbation has total
/// magnitude ≈ `noise` against the *unit* base (a raw `noise × U[-0.5,0.5]` per
/// component would be ≈ `noise·√(dim/12)` — ~11× too large at 1536-D, which would
/// push the query far off its base and understate recall). The query is NOT
/// re-normalized: the corpus is unit-norm, so L2 order over it equals cosine
/// order even for a non-unit query (`‖c‖=1` ⇒ argmin‖q−c‖² = argmax q·c) — the
/// exact metric the engine serves (it normalizes on write, not on query).
fn perturbed_query(corpus: &[Vec<f32>], i: usize, noise: f32) -> Vec<f32> {
    let n = corpus.len().max(1) as u64;
    let base = ((i as u64).wrapping_mul(2_654_435_761) % n) as usize;
    let mut q = corpus[base].clone();
    let dim = q.len().max(1);
    // Per-component coefficient so ‖noise_vec‖ ≈ noise for a unit base.
    let coef = noise / (dim as f32 / 12.0).sqrt();
    let mut state = (i as u64)
        .wrapping_mul(0x100_0000_01B3)
        .wrapping_add(0xABCD);
    let mut next = || {
        state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = state;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^= z >> 31;
        ((z >> 40) as f32 / (1u64 << 24) as f32) - 0.5
    };
    for c in &mut q {
        *c += coef * next();
    }
    q
}

// ---------------------------------------------------------------------------
// Measurement
// ---------------------------------------------------------------------------

/// recall@k of `got` against the exact `truth` id set.
fn recall(truth: &HashSet<VectorId>, got: &[VectorId], k: usize) -> f64 {
    let hits = got.iter().filter(|id| truth.contains(id)).count();
    hits as f64 / k as f64
}

/// One grid/quant point's measured outcome.
struct Point {
    label: String,
    recall: f64,
    mean_us: f64,
    p99_us: f64,
    build_s: f64,
    mem_mb: f64,
    mem_per_1m_gb: f64,
}

/// Build a `UsearchIndex` with the given parameters over `vectors`, returning the
/// built index and the wall-clock build time.
fn build_hnsw(
    vectors: &[Vec<f32>],
    dim: usize,
    quant: QuantizationLevel,
    connectivity: usize,
    ef_construction: usize,
    ef_search: usize,
) -> (UsearchIndex, Duration) {
    let index = UsearchIndex::new(VectorIndexConfig {
        dimensions: dim,
        metric: DistanceMetric::L2,
        quantization: quant,
        connectivity,
        ef_construction,
        ef_search,
    })
    .expect("usearch index construction");
    // Parallel build, done CORRECTLY: reserve one writer slot per thread up front
    // (`reserve_with_threads`), THEN insert distinct ids concurrently. A plain
    // `reserve` allocates a single slot, so concurrent `add` would corrupt the
    // graph (recall collapses ~0.95→~0.1) — that was tried and rejected. With the
    // per-thread reservation, USearch's own test does exactly this, and recall
    // matches a sequential build while the 100k/1536-D build drops from ~10 min to
    // well under a minute.
    let threads = std::thread::available_parallelism().map_or(8, std::num::NonZeroUsize::get);
    index
        .reserve_with_threads(vectors.len(), threads)
        .expect("reserve");
    let start = Instant::now();
    let chunk = vectors.len().div_ceil(threads).max(1);
    std::thread::scope(|s| {
        for (c, slice) in vectors.chunks(chunk).enumerate() {
            let index = &index;
            let base = c * chunk;
            s.spawn(move || {
                for (i, v) in slice.iter().enumerate() {
                    index.insert((base + i) as VectorId, v).expect("insert");
                }
            });
        }
    });
    (index, start.elapsed())
}

/// Run `queries` against `index` at `ef_search`, scoring recall@k vs `truths`
/// and recording per-query latency. Returns `(avg_recall, mean_us, p99_us)`.
fn measure(
    index: &UsearchIndex,
    queries: &[Vec<f32>],
    truths: &[HashSet<VectorId>],
    k: usize,
    ef_search: usize,
) -> (f64, f64, f64) {
    let mut recalls = 0.0;
    let mut lat_us: Vec<f64> = Vec::with_capacity(queries.len());
    for (q, truth) in queries.iter().zip(truths) {
        let start = Instant::now();
        let res = index.search(q, k, ef_search).expect("search");
        lat_us.push(start.elapsed().as_secs_f64() * 1e6);
        let ids: Vec<VectorId> = res.iter().map(|r| r.id).collect();
        recalls += recall(truth, &ids, k);
    }
    let mean = lat_us.iter().sum::<f64>() / lat_us.len() as f64;
    lat_us.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    let p99 = lat_us[((lat_us.len() as f64 * 0.99) as usize).min(lat_us.len() - 1)];
    (recalls / queries.len() as f64, mean, p99)
}

/// Footprint in MB now, and extrapolated to 1M vectors (GB), from the measured
/// `memory_usage` of the built index over `n` vectors.
fn footprint(index: &UsearchIndex, n: usize) -> (f64, f64) {
    let bytes = index.memory_usage();
    let mb = bytes as f64 / (1024.0 * 1024.0);
    let per_1m_gb = (bytes as f64 / n as f64) * 1_000_000.0 / (1024.0 * 1024.0 * 1024.0);
    (mb, per_1m_gb)
}

const fn quant_name(q: QuantizationLevel) -> &'static str {
    match q {
        QuantizationLevel::F32 => "F32",
        QuantizationLevel::F16 => "F16",
        QuantizationLevel::Int8 => "Int8",
    }
}

// ---------------------------------------------------------------------------
// Arg parsing (tiny, dependency-free)
// ---------------------------------------------------------------------------

fn arg(args: &[String], flag: &str, default: usize) -> usize {
    args.iter()
        .position(|a| a == flag)
        .and_then(|i| args.get(i + 1))
        .and_then(|v| v.parse().ok())
        .unwrap_or(default)
}

#[allow(clippy::too_many_lines)] // a linear measurement script reads best top-to-bottom
fn main() {
    let args: Vec<String> = std::env::args().collect();
    let corpus = arg(&args, "--corpus", 20_000);
    let dim = arg(&args, "--dim", 1536);
    let n_queries = arg(&args, "--queries", 200);
    let k = arg(&args, "--k", 10);
    // Query perturbation noise (× 1000 on the CLI to keep the arg an integer):
    // `--noise-milli 50` ⇒ 0.05 relative — a query a few degrees off a real point.
    let noise = arg(&args, "--noise-milli", 50) as f32 / 1000.0;
    // Clustered corpus shape (Gaussian mixture — see `clustered_vector`):
    // `--clusters` defaults to ~100 points/cluster; `--spread-milli 500` ⇒ 0.5.
    let n_clusters = arg(&args, "--clusters", (corpus / 100).max(1));
    let spread = arg(&args, "--spread-milli", 500) as f32 / 1000.0;

    eprintln!(
        "[grid] building clustered corpus: {corpus} vectors × {dim}-D, {n_clusters} clusters \
         (spread {spread:.2}), {n_queries} queries (noise {noise:.3}), recall@{k}"
    );
    let corpus_start = Instant::now();
    let centers: Vec<Vec<f32>> = (0..n_clusters as u64)
        .map(|c| unit_vector(c.wrapping_add(0x00C0_FFEE), dim))
        .collect();
    let vectors: Vec<Vec<f32>> = (0..corpus as u64)
        .map(|id| clustered_vector(&centers, id, dim, spread))
        .collect();
    // Queries are corpus points perturbed by small noise — they stay in-cluster,
    // so each has a well-defined nearest cluster (the production read shape, the
    // same model the m12p1 harness uses; see `perturbed_query`).
    let queries: Vec<Vec<f32>> = (0..n_queries)
        .map(|q| perturbed_query(&vectors, q, noise))
        .collect();
    eprintln!(
        "[grid] corpus built in {:.1}s",
        corpus_start.elapsed().as_secs_f64()
    );

    // Exact ground truth, computed ONCE (independent of HNSW parameters) and
    // reused across every grid/quant point. F32 brute force == the true answer.
    eprintln!("[grid] computing exact brute-force ground truth (once)…");
    let gt_start = Instant::now();
    let oracle = BruteForceIndex::new(VectorIndexConfig {
        dimensions: dim,
        metric: DistanceMetric::L2,
        quantization: QuantizationLevel::F32,
        connectivity: 16,
        ef_construction: 400,
        ef_search: 400,
    });
    for (id, v) in vectors.iter().enumerate() {
        oracle.insert(id as VectorId, v).expect("oracle insert");
    }
    let truths: Vec<HashSet<VectorId>> = queries
        .iter()
        .map(|q| {
            oracle
                .search(q, k, 0)
                .expect("oracle search")
                .iter()
                .map(|r| r.id)
                .collect()
        })
        .collect();
    eprintln!(
        "[grid] ground truth ready in {:.1}s",
        gt_start.elapsed().as_secs_f64()
    );

    // -------------------------------------------------------------------
    // Sweep 1 — HNSW graph parameters at F16 (the production quantization).
    // For each (M, ef_construction) we build the graph ONCE and sweep ef_search
    // on it (search-time only — no rebuild), since ef_search is a per-query knob.
    // -------------------------------------------------------------------
    let grid_graphs: &[(usize, usize)] = &[(16, 400), (24, 400), (32, 400)];
    let ef_searches: &[usize] = &[128, 200, 400, 600];
    let mut grid_points: Vec<Point> = Vec::new();
    for &(m, ef_c) in grid_graphs {
        eprintln!("[grid] building HNSW M={m} ef_c={ef_c} (F16)…");
        let (index, build) = build_hnsw(&vectors, dim, QuantizationLevel::F16, m, ef_c, 200);
        let (mem_mb, mem_1m) = footprint(&index, corpus);
        for &ef_s in ef_searches {
            let (rec, mean, p99) = measure(&index, &queries, &truths, k, ef_s);
            grid_points.push(Point {
                label: format!("M={m}, ef_c={ef_c}, ef_s={ef_s}"),
                recall: rec,
                mean_us: mean,
                p99_us: p99,
                build_s: build.as_secs_f64(),
                mem_mb,
                mem_per_1m_gb: mem_1m,
            });
        }
    }

    // -------------------------------------------------------------------
    // Sweep 2 — quantization at a graph + beam that clears the 0.95 gate
    // (M=24, ef_c=400, ef_s=400), so the F32→F16→Int8 recall penalty and the
    // memory saving are compared on a config that actually meets the target.
    // -------------------------------------------------------------------
    let quants = [
        QuantizationLevel::F32,
        QuantizationLevel::F16,
        QuantizationLevel::Int8,
    ];
    let quant_m = 24usize;
    let quant_ef_construction = 400usize;
    let quant_beam = 400usize;
    let mut quant_points: Vec<Point> = Vec::new();
    for q in quants {
        eprintln!(
            "[grid] building HNSW {} M={quant_m} ef_c={quant_ef_construction}…",
            quant_name(q)
        );
        let (index, build) =
            build_hnsw(&vectors, dim, q, quant_m, quant_ef_construction, quant_beam);
        let (mem_mb, mem_1m) = footprint(&index, corpus);
        let (rec, mean, p99) = measure(&index, &queries, &truths, k, quant_beam);
        quant_points.push(Point {
            label: format!(
                "{} (M={quant_m}, ef_c={quant_ef_construction}, ef_s={quant_beam})",
                quant_name(q)
            ),
            recall: rec,
            mean_us: mean,
            p99_us: p99,
            build_s: build.as_secs_f64(),
            mem_mb,
            mem_per_1m_gb: mem_1m,
        });
    }

    // -------------------------------------------------------------------
    // Report — markdown tables ready to paste into docs/profiling/.
    // -------------------------------------------------------------------
    println!(
        "\n## ANN grid search — measured (corpus={corpus}, dim={dim}, recall@{k}, queries={n_queries})\n"
    );
    println!("### HNSW parameter sweep (F16)\n");
    println!(
        "| M / ef_c / ef_s | recall@{k} | mean (µs) | p99 (µs) | build (s) | mem (MB) | mem/1M (GB) |"
    );
    println!("|---|---|---|---|---|---|---|");
    for p in &grid_points {
        println!(
            "| {} | {:.4} | {:.0} | {:.0} | {:.1} | {:.0} | {:.2} |",
            p.label, p.recall, p.mean_us, p.p99_us, p.build_s, p.mem_mb, p.mem_per_1m_gb
        );
    }

    println!("\n### Quantization sweep (M=16, ef_c=400, ef_s=200)\n");
    println!("| quantization | recall@{k} | mean (µs) | p99 (µs) | mem (MB) | mem/1M (GB) |");
    println!("|---|---|---|---|---|---|");
    for p in &quant_points {
        println!(
            "| {} | {:.4} | {:.0} | {:.0} | {:.0} | {:.2} |",
            p.label, p.recall, p.mean_us, p.p99_us, p.mem_mb, p.mem_per_1m_gb
        );
    }

    // Recommend the cheapest grid point (min mean latency) that clears recall ≥ 0.95.
    let best = grid_points
        .iter()
        .filter(|p| p.recall >= 0.95)
        .min_by(|a, b| {
            a.mean_us
                .partial_cmp(&b.mean_us)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
    println!("\n### Recommendation\n");
    match best {
        Some(p) => println!(
            "- **Frontier point:** `{}` → recall@{k} {:.4}, mean {:.0} µs, p99 {:.0} µs, {:.2} GB/1M.",
            p.label, p.recall, p.mean_us, p.p99_us, p.mem_per_1m_gb
        ),
        None => println!(
            "- ⚠ NO grid point reached recall@{k} ≥ 0.95 at this shape — widen ef or raise M."
        ),
    }
    let smallest_ok = quant_points
        .iter()
        .filter(|p| p.recall >= 0.95)
        .min_by(|a, b| {
            a.mem_per_1m_gb
                .partial_cmp(&b.mem_per_1m_gb)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
    if let Some(p) = smallest_ok {
        println!(
            "- **Smallest quantization clearing recall ≥ 0.95:** `{}` → {:.4} recall, {:.2} GB/1M.",
            p.label, p.recall, p.mem_per_1m_gb
        );
    }
}