brainy/tests/benchmarks/lib/metrics.js

/**
 * @module tests/benchmarks/lib/metrics
 * @description Measurement helpers for scaling benchmarks: latency percentiles,
 * brute-force recall@k ground truth, and a memory snapshot. Generic MIT
 * benchmark tooling shared by the open-core leg and the A/B comparison so every
 * column is computed identically. Recall uses cosine distance to match Brainy's
 * default index metric (`this.distance = cosineDistance`).
 *
 * Not shipped in the npm package (tests/ is outside `files`).
 */

/**
 * @description The p-th percentile of an already-ascending-sorted array.
 * @param sorted - Ascending-sorted samples.
 * @param p - Percentile in [0, 100].
 * @returns The sample at the percentile (nearest-rank).
 */
export function percentile(sorted, p) {
  if (sorted.length === 0) return NaN
  const idx = Math.min(sorted.length - 1, Math.floor((p / 100) * sorted.length))
  return sorted[idx]
}

/**
 * @description Summarize a list of latency samples (milliseconds).
 * @param latMs - Latency samples in ms (any order).
 * @returns `{ p50, p95, p99, mean, n }`.
 */
export function summarize(latMs) {
  const s = [...latMs].sort((a, b) => a - b)
  const mean = s.reduce((acc, x) => acc + x, 0) / (s.length || 1)
  return { p50: percentile(s, 50), p95: percentile(s, 95), p99: percentile(s, 99), mean, n: s.length }
}

/** @description High-resolution elapsed-ms timer around an async function. */
export async function timed(fn) {
  const t = process.hrtime.bigint()
  const value = await fn()
  return { ms: Number(process.hrtime.bigint() - t) / 1e6, value }
}

/**
 * @description Cosine distance (`1 - cosineSimilarity`) — matches Brainy's
 * default index metric, so brute-force ground truth ranks identically to the index.
 * @param a - First vector.
 * @param b - Second vector.
 * @returns Cosine distance in [0, 2].
 */
export function cosineDistance(a, b) {
  let dot = 0, na = 0, nb = 0
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]
  }
  const denom = Math.sqrt(na) * Math.sqrt(nb)
  return denom === 0 ? 1 : 1 - dot / denom
}

/**
 * @description Exact top-k indices for `query` by ascending cosine distance over
 * the whole corpus — the recall ground truth. O(n·dim) per query, so call it on a
 * sample of queries, not all of them.
 * @param query - Query vector.
 * @param getVector - `(i) => number[]` corpus accessor.
 * @param n - Corpus size.
 * @param k - Neighbours to return.
 * @returns Array of the k nearest corpus indices, nearest first.
 */
export function bruteForceTopK(query, getVector, n, k) {
  const heap = [] // small: keep k best as {i, d}
  for (let i = 0; i < n; i++) {
    const d = cosineDistance(query, getVector(i))
    if (heap.length < k) {
      heap.push({ i, d })
      if (heap.length === k) heap.sort((a, b) => a.d - b.d)
    } else if (d < heap[k - 1].d) {
      // insert in order, drop the worst
      let pos = k - 1
      while (pos > 0 && heap[pos - 1].d > d) { heap[pos] = heap[pos - 1]; pos-- }
      heap[pos] = { i, d }
    }
  }
  return heap.slice(0, k).map((e) => e.i)
}

/**
 * @description recall@k = |approx ∩ truth| / |truth|. Keys must be comparable
 * (map entity ids to corpus indices, or vice-versa, before calling).
 * @param approxKeys - Keys the index returned.
 * @param truthKeys - Ground-truth keys.
 * @returns Recall in [0, 1].
 */
export function recallAtK(approxKeys, truthKeys) {
  if (truthKeys.length === 0) return 1
  const truth = new Set(truthKeys)
  let hit = 0
  for (const k of approxKeys) if (truth.has(k)) hit++
  return hit / truthKeys.length
}

/**
 * @description Process memory snapshot, optionally per-entity.
 * @param entityCount - Optional entity count for the per-entity figure.
 * @returns `{ rssGB, heapGB, rssBytesPerEntity }`.
 */
export function memSnapshot(entityCount) {
  const m = process.memoryUsage()
  return {
    rssGB: +(m.rss / 1e9).toFixed(3),
    heapGB: +(m.heapUsed / 1e9).toFixed(3),
    rssBytesPerEntity: entityCount ? Math.round(m.rss / entityCount) : null
  }
}