brainy/tests/benchmarks/lib/corpus.js

/**
 * @module tests/benchmarks/lib/corpus
 * @description Deterministic synthetic-corpus generator for scaling benchmarks —
 * the single source of truth so the open-core leg and the (proprietary) A/B
 * comparison leg measure the IDENTICAL workload. Vectors follow a
 * mixture-of-clusters distribution (clusters of related points + small noise),
 * which is closer to real embedding geometry than uniform-random and is the
 * worst case to avoid for product quantization — matching the methodology used
 * by the native-provider benchmark suite so columns are comparable.
 *
 * Everything is recomputed on demand from a seed + index (no large arrays held),
 * so a 1M / 10M corpus costs O(clusters·dim) memory, not O(n·dim). Same seed ⇒
 * byte-identical corpus across runs, machines, and the two A/B legs.
 *
 * Cortex-free by construction: this is generic MIT benchmark tooling, usable by
 * any open-core consumer to benchmark their own deployment. It is NOT shipped in
 * the npm package (tests/ is outside `files`).
 */

/**
 * @description Deterministic 32-bit LCG (Numerical Recipes constants). Returns a
 * closure producing floats in [0, 1). Seeded so every run is reproducible.
 * @param seed - 32-bit unsigned seed.
 * @returns A function returning the next pseudo-random float in [0, 1).
 */
export function makeRng(seed = 0x12345678) {
  let state = seed >>> 0
  return () => {
    state = (Math.imul(state, 1664525) + 1013904223) >>> 0
    return state / 4294967296
  }
}

/**
 * @description Number of cluster centres for a corpus of `n` points:
 * `min(1024, max(64, floor(sqrt(n))))`. Sub-linear so clusters stay dense as the
 * corpus grows (matches the native-suite generator).
 * @param n - Corpus size.
 * @returns The cluster count.
 */
export function clusterCount(n) {
  return Math.min(1024, Math.max(64, Math.floor(Math.sqrt(n))))
}

/**
 * @description Build a deterministic corpus descriptor for `n` entities of
 * dimension `dim`. Nothing large is allocated up front: cluster centres
 * (clusters·dim) are materialized once; every entity vector, its metadata, the
 * query set, and the hub→neighbour edge set are recomputed on demand from the
 * seed and the index.
 *
 * @param opts - Corpus parameters.
 * @param opts.n - Number of entities.
 * @param opts.dim - Vector dimension (default 384, the all-MiniLM-L6-v2 size).
 * @param opts.seed - Master seed (default 0x12345678).
 * @param opts.noise - Per-component uniform noise added to a cluster centre (default 0.1).
 * @param opts.categories - Distinct values for the `category` metadata field (default 10).
 * @param opts.hubs - Entities that get an out-edge neighbourhood (default min(n/10, 1000)).
 * @param opts.fanout - Out-edges per hub (default 100).
 * @returns A descriptor exposing `vector(i)`, `metadata(i)`, `queryVector(j)`,
 *   `neighborIndex(hub, f)`, plus the resolved parameters.
 */
export function makeCorpus(opts) {
  const n = opts.n
  const dim = opts.dim ?? 384
  const seed = (opts.seed ?? 0x12345678) >>> 0
  const noise = opts.noise ?? 0.1
  const categories = opts.categories ?? 10
  const clusters = clusterCount(n)
  const hubs = opts.hubs ?? Math.min(Math.floor(n / 10) || 1, 1000)
  const fanout = opts.fanout ?? 100

  // Materialize cluster centres once (clusters · dim floats — small).
  // Components are CENTERED in [-1, 1) so clusters separate by both DIRECTION
  // (cosine — Brainy's default metric) and DISTANCE (L2 — the native/pgvector
  // metric). A positive-orthant corpus would be degenerate under cosine and
  // tank recall, so the shared corpus must be fair to both metrics.
  const centerRng = makeRng(seed ^ 0x9e3779b9)
  const centers = new Array(clusters)
  for (let c = 0; c < clusters; c++) {
    const v = new Array(dim)
    for (let d = 0; d < dim; d++) v[d] = centerRng() * 2 - 1
    centers[c] = v
  }

  /** Deterministic per-index noise stream (decorrelated from the centre stream). */
  const noiseAt = (i) => makeRng((seed + 0x85ebca6b * (i + 1)) >>> 0)

  return {
    n, dim, seed, clusters, hubs, fanout, categories,

    /**
     * @returns A deterministic, collision-free, UUID-shaped id for entity `i`
     * (the index is encoded in the node field). Used so the caller never depends
     * on `addMany` completion order to know which id holds which vector — the
     * id↔index map is exact, which recall and the edge graph both require.
     */
    id(i) {
      return `00000000-0000-4000-8000-${i.toString(16).padStart(12, '0')}`
    },

    /** @returns The vector for entity `i`: its cluster centre + seeded noise. */
    vector(i) {
      const center = centers[i % clusters]
      const r = noiseAt(i)
      const v = new Array(dim)
      for (let d = 0; d < dim; d++) v[d] = center[d] + (r() * 2 - 1) * noise
      return v
    },

    /** @returns Deterministic metadata for entity `i`. */
    metadata(i) {
      const r = noiseAt(i ^ 0x1234)
      return { idx: i, category: i % categories, score: Math.floor(r() * 1000), active: (i & 1) === 0 }
    },

    /**
     * @returns A query vector for query `j`: near a deterministically chosen
     * cluster (drawn from the SAME distribution as the corpus, distinct stream).
     */
    queryVector(j) {
      const c = j % clusters
      const center = centers[c]
      const r = makeRng((seed + 0xc2b2ae35 * (j + 1)) >>> 0)
      const v = new Array(dim)
      for (let d = 0; d < dim; d++) v[d] = center[d] + (r() * 2 - 1) * noise
      return v
    },

    /** @returns The global index of out-neighbour `f` of `hub` (deterministic, wraps within n). */
    neighborIndex(hub, f) {
      return (hub * fanout + f + hubs) % n
    }
  }
}