# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The memory-bound arithmetic of brute-force search.

400,000 random 768-dimensional vectors (1.2 GB as float32). One brute-force
query (scores = vectors @ q, then argpartition for the top 10), timed as
float32, float16 and int8; effective bandwidth = bytes read ÷ seconds; and the
float32 time scaled to 10 million vectors.
"""
import argparse, json, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=400_000, help="Vectors.")
p.add_argument("--dim", type=int, default=768, help="Dimension.")
p.add_argument("--repeats", type=int, default=5, help="Queries per type (best kept).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0)
v32 = rng.standard_normal((args.n, args.dim), dtype=np.float32); q32 = rng.standard_normal(args.dim, dtype=np.float32)
rows = []
for name, v, q in (("float32", v32, q32), ("float16", v32.astype(np.float16), q32.astype(np.float16)),
                   ("int8", np.clip(v32 * 32, -127, 127).astype(np.int8), np.clip(q32 * 32, -127, 127).astype(np.int8))):
    ts = []
    for _ in range(args.repeats):
        t = time.perf_counter(); s = v @ q; top = np.argpartition(s, -10)[-10:]; ts.append(time.perf_counter() - t)
    el = min(ts); rows.append({"type": name, "gb": v.nbytes / 1e9, "ms": el * 1000, "gb_s": v.nbytes / el / 1e9})
    if name != "float32": del v
for r in rows: r["vs_float32"] = r["ms"] / rows[0]["ms"]
summary = {"ms_10m_float32": rows[0]["ms"] * 10_000_000 / args.n}
for r in rows: print(f"{r['type']:<8} {r['gb']:5.2f} GB  one query {r['ms']:8.1f} ms  ({r['gb_s']:5.1f} GB/s, {r['vs_float32']:.2f}x float32's time)")
print(f"float32 scaled to 10 million vectors: {summary['ms_10m_float32'] / 1000:.1f} s per query")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"types": rows}}, f)