bitbudget 0.1.0__tar.gz

bitbudget-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Sean Moran
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bitbudget-0.1.0/PKG-INFO

@@ -0,0 +1,148 @@
+Metadata-Version: 2.4
+Name: bitbudget
+Version: 0.1.0
+Summary: How much retrieval quality do you keep per byte? A reproducible benchmark for embedding compression.
+Author: Sean Moran
+License: MIT
+Project-URL: Paper, https://arxiv.org/abs/2510.04127
+Project-URL: Leaderboard, https://github.com/sjmoran/bitbudget/blob/main/LEADERBOARD.md
+Keywords: retrieval,embeddings,quantisation,hashing,compression,ANN,RAG
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.21
+Provides-Extra: embed
+Requires-Dist: sentence-transformers>=2.2; extra == "embed"
+Provides-Extra: faiss
+Requires-Dist: faiss-cpu>=1.7.4; extra == "faiss"
+Provides-Extra: all
+Requires-Dist: sentence-transformers>=2.2; extra == "all"
+Requires-Dist: faiss-cpu>=1.7.4; extra == "all"
+Dynamic: license-file
+
+# BitBudget
+
+**How much retrieval quality do you keep per byte?**
+
+BitBudget is a small, reproducible benchmark for **embedding compression**. Give it an
+embedder and a corpus and it reports the retrieval quality (nDCG@10, recall@10) that each
+compression method retains against the **bytes it stores per vector** — the recall‑per‑byte
+frontier that every RAG and vector‑database deployment actually lives on.
+
+It is the companion benchmark to the survey *“Projection and Quantisation: A Unifying View of
+Learning to Hash, from Random Projections to the RAG Era”* and exists to answer one question
+that today is mostly answered by vendor blog posts: **when you binarise / int8 / RaBitQ /
+product‑quantise / Matryoshka‑truncate your embeddings, what do you actually lose?**
+
+## The headline finding
+
+> **Bits beat dimensions.** Spending a fixed byte budget on *more coarsely quantised*
+> coordinates beats spending it on *fewer full‑precision* coordinates, at every budget and
+> for every embedder we have tried. One‑bit codes with a cheap re‑ranking pass are **32×
+> smaller than float at no measurable loss**.
+
+```
+mxbai‑embed‑large (1024‑d), mean over 4 BEIR corpora
+  binary+rerank      128 B   nDCG 0.509   100% of float   ← 32× smaller, lossless
+  pq                 128 B   nDCG 0.488    96%
+  rabitq             128 B   nDCG 0.487    96%
+  matryoshka        1024 B   nDCG 0.439    86%             ← 4× smaller, projection axis
+  float32           4096 B   nDCG 0.508   100%
+```
+
+See **[LEADERBOARD.md](LEADERBOARD.md)** for the full table.
+
+## Install
+
+```bash
+pip install bitbudget            # evaluation only (numpy)
+pip install "bitbudget[all]"     # + sentence-transformers (embedding) + faiss
+```
+
+## Quickstart
+
+```bash
+bitbudget methods                                   # list compression methods
+bitbudget run --embedder mxbai --corpus scifact     # embed + evaluate, print a results card
+bitbudget leaderboard results/card_*.json           # render a markdown leaderboard
+
+bitbudget indexes                                   # list indexes (organisation axis)
+bitbudget bench-index --synthetic 100000 128        # recall vs QPS vs bytes: flat/hnsw/ivfpq/bittrie
+```
+
+`run` embeds (torch) and evaluates (numpy) in one process. The corpora auto‑download.
+
+### The organisation axis (`bench-index`)
+
+The compression leaderboard answers *quality per byte*; `bench-index` answers the orthogonal
+*recall per query-second*. It builds an index over the document vectors and reports recall@k,
+throughput (QPS) and bytes per vector, so HNSW and IVF‑PQ (which buy throughput and *add* bytes)
+can be compared against compact‑code indexes on one frontier. Run it on synthetic data, on a
+cached embedding (`--embedder mxbai --corpus scifact`), or on your own vectors (`--npz`). The
+faiss‑backed indexes need `pip install bitbudget[faiss]`; the numpy `bittrie` runs without it.
+
+The `bittrie` index ships a small C kernel (`_bittrie.c`) for the query hot‑path, compiled on
+first use and cached (no compiler needed to *install* — the wheel stays pure‑Python, and it falls
+back to numpy if no compiler is present). It builds **multithreaded** when OpenMP is available
+(GCC/clang on Linux, Homebrew `libomp` on macOS) and single‑threaded otherwise; results are
+bit‑identical to the numpy path, and recall/footprint are algorithmic and unchanged either way.
+
+Because faiss carries its own OpenMP runtime, it cannot share a process with the bit‑trie's
+`libomp` on macOS. `bench-index` therefore runs the faiss indexes and the bit‑trie in **separate
+subprocesses** and merges the results, so a single `bitbudget bench-index ...` works everywhere
+(pass `--no-split` to force one process, e.g. on Linux where both share one OpenMP runtime).
+
+> **macOS note.** torch and faiss each bundle their own OpenMP runtime and crash if imported
+> in the same process. The core methods are numpy‑only, so `run` is safe; if you add a
+> faiss‑backed method, run `bitbudget embed` (torch) and `bitbudget eval` (numpy/faiss)
+> as separate processes.
+
+## The protocol (frozen, so results are comparable)
+
+- **Corpora:** the BEIR subsets `scifact`, `nfcorpus`, `arguana`, `fiqa` (small enough to run
+  on a laptop, diverse enough to be honest). Numbers are the mean over corpora; `±` is the
+  standard deviation across them.
+- **Metrics:** `nDCG@10` against the graded BEIR judgements, and `recall@10` against the exact
+  floating‑point neighbours. `% of float` is nDCG relative to the uncompressed embedding.
+- **Memory:** bytes stored per document vector (`4D` float, `D` int8, `D/8` binary, `M` for an
+  `M`‑byte product code, `4·dim` for a truncated/PCA‑reduced vector).
+- **Embedders:** `minilm` (384‑d) and `mxbai` (1024‑d, Matryoshka) ship built in.
+
+## Add your method in five lines
+
+This is the point of the benchmark: drop in your compressor and it is scored against every
+built‑in on the same protocol.
+
+```python
+from bitbudget import method
+import numpy as np
+
+@method("my-2bit", bits=2)
+def my_2bit(demb, qemb):
+    codes = my_quantise(demb)                       # your compression
+    scores = qemb @ my_reconstruct(codes).T         # (queries x docs) similarity
+    return scores, demb.shape[1] * 2 / 8            # scores, bytes per stored vector
+```
+
+```bash
+bitbudget run --embedder mxbai --corpus scifact --methods my-2bit binary+rerank float32
+```
+
+Then open a pull request adding your row to [LEADERBOARD.md](LEADERBOARD.md). See
+[CONTRIBUTING.md](CONTRIBUTING.md).
+
+## Cite
+
+If BitBudget helps your work, please cite the survey:
+
+```bibtex
+@article{moran2025projection,
+  title   = {Projection and Quantisation: A Unifying View of Learning to Hash,
+             from Random Projections to the RAG Era},
+  author  = {Moran, Sean},
+  journal = {arXiv preprint arXiv:2510.04127},
+  year    = {2025}
+}
+```
+
+MIT licensed.

bitbudget-0.1.0/README.md

@@ -0,0 +1,126 @@
+# BitBudget
+
+**How much retrieval quality do you keep per byte?**
+
+BitBudget is a small, reproducible benchmark for **embedding compression**. Give it an
+embedder and a corpus and it reports the retrieval quality (nDCG@10, recall@10) that each
+compression method retains against the **bytes it stores per vector** — the recall‑per‑byte
+frontier that every RAG and vector‑database deployment actually lives on.
+
+It is the companion benchmark to the survey *“Projection and Quantisation: A Unifying View of
+Learning to Hash, from Random Projections to the RAG Era”* and exists to answer one question
+that today is mostly answered by vendor blog posts: **when you binarise / int8 / RaBitQ /
+product‑quantise / Matryoshka‑truncate your embeddings, what do you actually lose?**
+
+## The headline finding
+
+> **Bits beat dimensions.** Spending a fixed byte budget on *more coarsely quantised*
+> coordinates beats spending it on *fewer full‑precision* coordinates, at every budget and
+> for every embedder we have tried. One‑bit codes with a cheap re‑ranking pass are **32×
+> smaller than float at no measurable loss**.
+
+```
+mxbai‑embed‑large (1024‑d), mean over 4 BEIR corpora
+  binary+rerank      128 B   nDCG 0.509   100% of float   ← 32× smaller, lossless
+  pq                 128 B   nDCG 0.488    96%
+  rabitq             128 B   nDCG 0.487    96%
+  matryoshka        1024 B   nDCG 0.439    86%             ← 4× smaller, projection axis
+  float32           4096 B   nDCG 0.508   100%
+```
+
+See **[LEADERBOARD.md](LEADERBOARD.md)** for the full table.
+
+## Install
+
+```bash
+pip install bitbudget            # evaluation only (numpy)
+pip install "bitbudget[all]"     # + sentence-transformers (embedding) + faiss
+```
+
+## Quickstart
+
+```bash
+bitbudget methods                                   # list compression methods
+bitbudget run --embedder mxbai --corpus scifact     # embed + evaluate, print a results card
+bitbudget leaderboard results/card_*.json           # render a markdown leaderboard
+
+bitbudget indexes                                   # list indexes (organisation axis)
+bitbudget bench-index --synthetic 100000 128        # recall vs QPS vs bytes: flat/hnsw/ivfpq/bittrie
+```
+
+`run` embeds (torch) and evaluates (numpy) in one process. The corpora auto‑download.
+
+### The organisation axis (`bench-index`)
+
+The compression leaderboard answers *quality per byte*; `bench-index` answers the orthogonal
+*recall per query-second*. It builds an index over the document vectors and reports recall@k,
+throughput (QPS) and bytes per vector, so HNSW and IVF‑PQ (which buy throughput and *add* bytes)
+can be compared against compact‑code indexes on one frontier. Run it on synthetic data, on a
+cached embedding (`--embedder mxbai --corpus scifact`), or on your own vectors (`--npz`). The
+faiss‑backed indexes need `pip install bitbudget[faiss]`; the numpy `bittrie` runs without it.
+
+The `bittrie` index ships a small C kernel (`_bittrie.c`) for the query hot‑path, compiled on
+first use and cached (no compiler needed to *install* — the wheel stays pure‑Python, and it falls
+back to numpy if no compiler is present). It builds **multithreaded** when OpenMP is available
+(GCC/clang on Linux, Homebrew `libomp` on macOS) and single‑threaded otherwise; results are
+bit‑identical to the numpy path, and recall/footprint are algorithmic and unchanged either way.
+
+Because faiss carries its own OpenMP runtime, it cannot share a process with the bit‑trie's
+`libomp` on macOS. `bench-index` therefore runs the faiss indexes and the bit‑trie in **separate
+subprocesses** and merges the results, so a single `bitbudget bench-index ...` works everywhere
+(pass `--no-split` to force one process, e.g. on Linux where both share one OpenMP runtime).
+
+> **macOS note.** torch and faiss each bundle their own OpenMP runtime and crash if imported
+> in the same process. The core methods are numpy‑only, so `run` is safe; if you add a
+> faiss‑backed method, run `bitbudget embed` (torch) and `bitbudget eval` (numpy/faiss)
+> as separate processes.
+
+## The protocol (frozen, so results are comparable)
+
+- **Corpora:** the BEIR subsets `scifact`, `nfcorpus`, `arguana`, `fiqa` (small enough to run
+  on a laptop, diverse enough to be honest). Numbers are the mean over corpora; `±` is the
+  standard deviation across them.
+- **Metrics:** `nDCG@10` against the graded BEIR judgements, and `recall@10` against the exact
+  floating‑point neighbours. `% of float` is nDCG relative to the uncompressed embedding.
+- **Memory:** bytes stored per document vector (`4D` float, `D` int8, `D/8` binary, `M` for an
+  `M`‑byte product code, `4·dim` for a truncated/PCA‑reduced vector).
+- **Embedders:** `minilm` (384‑d) and `mxbai` (1024‑d, Matryoshka) ship built in.
+
+## Add your method in five lines
+
+This is the point of the benchmark: drop in your compressor and it is scored against every
+built‑in on the same protocol.
+
+```python
+from bitbudget import method
+import numpy as np
+
+@method("my-2bit", bits=2)
+def my_2bit(demb, qemb):
+    codes = my_quantise(demb)                       # your compression
+    scores = qemb @ my_reconstruct(codes).T         # (queries x docs) similarity
+    return scores, demb.shape[1] * 2 / 8            # scores, bytes per stored vector
+```
+
+```bash
+bitbudget run --embedder mxbai --corpus scifact --methods my-2bit binary+rerank float32
+```
+
+Then open a pull request adding your row to [LEADERBOARD.md](LEADERBOARD.md). See
+[CONTRIBUTING.md](CONTRIBUTING.md).
+
+## Cite
+
+If BitBudget helps your work, please cite the survey:
+
+```bibtex
+@article{moran2025projection,
+  title   = {Projection and Quantisation: A Unifying View of Learning to Hash,
+             from Random Projections to the RAG Era},
+  author  = {Moran, Sean},
+  journal = {arXiv preprint arXiv:2510.04127},
+  year    = {2025}
+}
+```
+
+MIT licensed.

bitbudget-0.1.0/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "bitbudget"
+version = "0.1.0"
+description = "How much retrieval quality do you keep per byte? A reproducible benchmark for embedding compression."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [{ name = "Sean Moran" }]
+keywords = ["retrieval", "embeddings", "quantisation", "hashing", "compression", "ANN", "RAG"]
+dependencies = [
+    "numpy>=1.21",
+]
+
+[project.optional-dependencies]
+# embedding a corpus (torch); kept optional so evaluation-only installs stay light and
+# avoid importing torch alongside faiss (they clash on macOS OpenMP -- see README).
+embed = ["sentence-transformers>=2.2"]
+# faiss-backed methods (graph/IVF) run in their own evaluation process.
+faiss = ["faiss-cpu>=1.7.4"]
+all = ["sentence-transformers>=2.2", "faiss-cpu>=1.7.4"]
+
+[project.scripts]
+bitbudget = "bitbudget.cli:main"
+
+[project.urls]
+Paper = "https://arxiv.org/abs/2510.04127"
+Leaderboard = "https://github.com/sjmoran/bitbudget/blob/main/LEADERBOARD.md"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+# Ship the bit-trie C kernel as data (compiled on demand at runtime, see _bittrie_build.py).
+# The wheel stays pure-Python (py3-none-any); no compiler is needed to install.
+[tool.setuptools.package-data]
+bitbudget = ["_bittrie.c"]

bitbudget-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

bitbudget-0.1.0/src/bitbudget/__init__.py

@@ -0,0 +1,28 @@
+"""BitBudget: how much retrieval quality do you keep per byte?
+
+A reproducible benchmark for embedding compression. Given an embedder and a corpus, it
+measures the retrieval quality (nDCG@10, recall@10) retained by each compression method
+against the bytes it stores per vector -- the recall-per-byte frontier.
+
+Add your own method in five lines::
+
+    from bitbudget import method
+    import numpy as np
+
+    @method("my-2bit", bits=2)
+    def my_2bit(demb, qemb):
+        # return (query x doc similarity scores, bytes per stored vector)
+        codes = my_quantise(demb)
+        return qemb @ my_dequantise(codes).T, demb.shape[1] * 2 / 8
+
+Then ``bitbudget run --embedder mxbai --corpus scifact`` scores it alongside the built-ins.
+"""
+from .methods import method, METHODS, list_methods
+from .embedders import embedder, EMBEDDERS, list_embedders
+from .eval import evaluate
+from .indexes import index, INDEXES, list_indexes, bench_indexes
+
+__version__ = "0.1.0"
+__all__ = ["method", "METHODS", "list_methods", "embedder", "EMBEDDERS",
+           "list_embedders", "evaluate", "index", "INDEXES", "list_indexes",
+           "bench_indexes", "__version__"]

bitbudget-0.1.0/src/bitbudget/_bittrie.c

@@ -0,0 +1,92 @@
+/* Fast batched bit-trie query: coarse-to-fine beam descent over sorted codes, then a
+ * full-precision re-ranking pass, looped over all queries in C. Plain C (no Python.h);
+ * loaded via ctypes from _bittrie_build.py and compiled on demand. Falls back to the
+ * numpy/Python path in bittrie.py if no compiler is available. */
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+
+typedef struct { double cost; long lo, hi; unsigned long long pre; } Node;
+
+static int cmp_node(const void *a, const void *b) {
+    double d = ((const Node *)a)->cost - ((const Node *)b)->cost;
+    return (d > 0) - (d < 0);
+}
+
+static long lower_bound_u64(const uint64_t *a, long lo, long hi, uint64_t key) {
+    while (lo < hi) { long mid = lo + (hi - lo) / 2; if (a[mid] < key) lo = mid + 1; else hi = mid; }
+    return lo;
+}
+
+/* out: nq*k int64 of original doc ids (top-k by inner product), -1 padding if fewer found. */
+/* Each query is independent and writes a disjoint slice of `out`, so the loop parallelises
+ * cleanly. Scratch buffers are allocated per thread inside the parallel region; when compiled
+ * without OpenMP the `omp` pragmas are ignored and this runs as one block, single-threaded. */
+void bt_query_batch(const uint64_t *codes, const int64_t *docids,
+                    const float *Xf, const float *qproj, const float *qf,
+                    int n, int d, int b, int nq, int depth, int beam, int k,
+                    int64_t *out) {
+    if (depth > b) depth = b;
+    if (beam < 1) beam = 1;
+    int cap = 2 * beam + 4;
+
+    #pragma omp parallel
+    {
+        Node *cur = (Node *)malloc(sizeof(Node) * cap);
+        Node *nxt = (Node *)malloc(sizeof(Node) * 2 * cap);
+        double *bscore = (double *)malloc(sizeof(double) * k);
+        long *bid = (long *)malloc(sizeof(long) * k);
+
+        #pragma omp for schedule(static)
+        for (int qi = 0; qi < nq; qi++) {
+            const float *qp = qproj + (size_t)qi * b;
+            const float *qv = qf + (size_t)qi * d;
+            int ncur = 1;
+            cur[0].cost = 0.0; cur[0].lo = 0; cur[0].hi = n; cur[0].pre = 0ULL;
+
+            for (int t = 0; t < depth; t++) {
+                int shift = b - t - 1;
+                int qb = qp[t] > 0.0f ? 1 : 0;
+                double c = fabs((double)qp[t]);
+                int nn = 0;
+                for (int i = 0; i < ncur; i++) {
+                    long lo = cur[i].lo, hi = cur[i].hi;
+                    unsigned long long pre = cur[i].pre;
+                    unsigned long long hi_part = (b - t) >= 64 ? 0ULL : (pre << (b - t));
+                    unsigned long long thresh = hi_part | (1ULL << shift);
+                    long m = lower_bound_u64(codes, lo, hi, (uint64_t)thresh);
+                    if (m > lo) { nxt[nn].cost = cur[i].cost + (qb == 0 ? 0.0 : c);
+                                  nxt[nn].lo = lo; nxt[nn].hi = m; nxt[nn].pre = (pre << 1); nn++; }
+                    if (hi > m) { nxt[nn].cost = cur[i].cost + (qb == 1 ? 0.0 : c);
+                                  nxt[nn].lo = m; nxt[nn].hi = hi; nxt[nn].pre = (pre << 1) | 1ULL; nn++; }
+                }
+                qsort(nxt, nn, sizeof(Node), cmp_node);
+                int keep = nn < beam ? nn : beam;
+                memcpy(cur, nxt, sizeof(Node) * keep);
+                ncur = keep;
+            }
+
+            for (int j = 0; j < k; j++) { bscore[j] = -1e300; bid[j] = -1; }
+            for (int i = 0; i < ncur; i++) {
+                for (long j = cur[i].lo; j < cur[i].hi; j++) {
+                    long id = (long)docids[j];
+                    const float *xr = Xf + (size_t)id * d;
+                    double s = 0.0;
+                    for (int dd = 0; dd < d; dd++) s += (double)qv[dd] * (double)xr[dd];
+                    int mn = 0;
+                    for (int t2 = 1; t2 < k; t2++) if (bscore[t2] < bscore[mn]) mn = t2;
+                    if (s > bscore[mn]) { bscore[mn] = s; bid[mn] = id; }
+                }
+            }
+            for (int a = 0; a < k; a++) {        /* selection-sort the k winners, score descending */
+                int mx = a;
+                for (int b2 = a + 1; b2 < k; b2++) if (bscore[b2] > bscore[mx]) mx = b2;
+                double ts = bscore[a]; bscore[a] = bscore[mx]; bscore[mx] = ts;
+                long ti = bid[a]; bid[a] = bid[mx]; bid[mx] = ti;
+            }
+            for (int j = 0; j < k; j++) out[(size_t)qi * k + j] = bid[j];
+        }
+        free(cur); free(nxt); free(bscore); free(bid);
+    }
+}

bitbudget-0.1.0/src/bitbudget/_bittrie_build.py

@@ -0,0 +1,94 @@
+"""Compile the bit-trie C kernel on demand and load it via ctypes.
+
+The package ships pure-Python on PyPI (the .c is data, not a built extension), so install never
+needs a compiler. The first time the fast path is requested we compile _bittrie.c into a cached
+shared library with the system C compiler; if no compiler is available we return None and the
+caller falls back to the numpy/Python path. The compiled library is cached by source hash, so the
+build happens once per machine.
+"""
+import ctypes
+import hashlib
+import os
+import subprocess
+
+_LIB = None
+_TRIED = False
+
+
+def _cache_dir():
+    d = os.environ.get("BITBUDGET_CACHE") or os.path.join(os.path.expanduser("~"), ".cache", "bitbudget")
+    os.makedirs(d, exist_ok=True)
+    return d
+
+
+def _brew_libomp():
+    """Return the Homebrew libomp prefix on macOS, or None."""
+    try:
+        p = subprocess.run(["brew", "--prefix", "libomp"], capture_output=True, text=True)
+        prefix = p.stdout.strip()
+        if p.returncode == 0 and prefix and os.path.isdir(prefix):
+            return prefix
+    except Exception:
+        pass
+    return None
+
+
+def _variants():
+    """(cflags, lflags) compile attempts, OpenMP first, single-thread C last."""
+    v = []
+    omp = _brew_libomp()
+    if omp:                                                    # macOS clang + Homebrew libomp (try first)
+        v.append((["-Xpreprocessor", "-fopenmp", f"-I{omp}/include"],
+                  [f"-L{omp}/lib", "-lomp", f"-Wl,-rpath,{omp}/lib"]))
+    v.append((["-fopenmp"], ["-fopenmp"]))                     # gcc / OpenMP-capable clang (Linux)
+    v.append(([], []))                                         # no OpenMP: single-threaded C
+    return v
+
+
+def _set_sig(lib):
+    f = lib.bt_query_batch
+    f.restype = None
+    P = ctypes.POINTER
+    f.argtypes = [P(ctypes.c_uint64), P(ctypes.c_int64), P(ctypes.c_float),
+                  P(ctypes.c_float), P(ctypes.c_float),
+                  ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int,
+                  ctypes.c_int, ctypes.c_int, ctypes.c_int, P(ctypes.c_int64)]
+    return lib
+
+
+def get_lib():
+    """Return the loaded ctypes library exposing bt_query_batch, or None if unavailable.
+
+    Tries to build a multithreaded (OpenMP) kernel first, falling back to single-threaded C,
+    then to None (the caller uses the numpy/Python path). The compiled library is cached by
+    source hash; set BITBUDGET_NO_OMP=1 to force the single-threaded build.
+    """
+    global _LIB, _TRIED
+    if _TRIED:
+        return _LIB
+    _TRIED = True
+    src = os.path.join(os.path.dirname(__file__), "_bittrie.c")
+    if not os.path.exists(src):
+        return None
+    h = hashlib.sha1(open(src, "rb").read()).hexdigest()[:12]
+    so = os.path.join(_cache_dir(), f"_bittrie_{h}.so")
+    if os.path.exists(so):                                     # cached: load directly
+        try:
+            _LIB = _set_sig(ctypes.CDLL(so)); return _LIB
+        except Exception:
+            try: os.remove(so)
+            except Exception: pass
+    cc = os.environ.get("CC") or "cc"
+    variants = [([], [])] if os.environ.get("BITBUDGET_NO_OMP") else _variants()
+    for cflags, lflags in variants:
+        cmd = [cc, "-O3", "-shared", "-fPIC", *cflags, "-o", so, src, *lflags, "-lm"]
+        try:
+            subprocess.run(cmd, check=True, capture_output=True)
+        except Exception:
+            continue
+        try:                                                  # must actually load (libomp present at runtime)
+            _LIB = _set_sig(ctypes.CDLL(so)); return _LIB
+        except Exception:
+            try: os.remove(so)
+            except Exception: pass
+    return None

bitbudget-0.1.0/src/bitbudget/bittrie.py

@@ -0,0 +1,108 @@
+"""Bit-trie index: a van Emde Boas / PATRICIA-style radix trie over compact codes.
+
+The sorted array of packed sign-codes *is* the trie: a node is a contiguous range of codes
+sharing a bit prefix. Search is coarse-to-fine beam descent (the most discriminative bit is the
+shallowest), then a re-ranking pass. Unsupervised (random rotation + sign), numpy-only, one knob
+(``beam``). It stores compact codes rather than full vectors, so it sits on the organisation axis
+at a fraction of a graph's footprint. Research entry to the BitBudget index benchmark; see the
+projection-quantisation-organisation survey for the motivation.
+
+This reference implementation packs into a single uint64, so ``n_bits <= 64``.
+"""
+import numpy as np
+
+
+class BitTrieIndex:
+    def __init__(self, n_bits=64, seed=0):
+        assert n_bits <= 64, "this reference packs into one uint64; use a word-array for >64 bits"
+        self.b = n_bits
+        self.seed = seed
+
+    def fit(self, X):
+        X = np.ascontiguousarray(X, dtype=np.float32)
+        n, d = X.shape
+        rng = np.random.RandomState(self.seed)
+        R = rng.randn(d, self.b).astype(np.float32)               # unsupervised random rotation (RaBitQ)
+        R, _ = np.linalg.qr(R) if self.b <= d else (R, None)
+        P = X @ R
+        order = np.argsort(-P.var(axis=0))                        # coarse-to-fine: MSB = highest variance
+        self.R = R[:, order]
+        P = P[:, order]
+        bits = (P > 0).astype(np.uint64)
+        weights = (np.uint64(1) << np.arange(self.b - 1, -1, -1, dtype=np.uint64))
+        codes = bits @ weights
+        order_idx = np.argsort(codes, kind="stable")
+        self.codes = codes[order_idx]                             # sorted: the implicit trie
+        self.docids = order_idx.astype(np.int64)
+        self.X = X                                                # kept only for optional float re-rank (cold)
+        self.n, self.d = n, d
+        return self
+
+    def _descend(self, qbits, qconf, beam, depth):
+        """Confidence-ordered beam over the sorted codes. Stop at ``depth`` bits so each surviving
+        prefix is still a fat bucket (the coarse-to-fine / anytime property)."""
+        b = self.b
+        state = [(0.0, 0, self.n, 0, 0)]                          # (cost, lo, hi, prefix_int, depth)
+        for t in range(depth):
+            shift = b - t - 1
+            nxt = []
+            for cost, lo, hi, pre, _ in state:
+                thresh = (pre << (b - t)) | (1 << shift)
+                m = lo + int(np.searchsorted(self.codes[lo:hi], np.uint64(thresh), "left"))
+                qb = int(qbits[t]); c = float(qconf[t])
+                left = (lo, m, (pre << 1))
+                right = (m, hi, (pre << 1) | 1)
+                for child_bit, (clo, chi, cpre) in ((0, left), (1, right)):
+                    if chi <= clo:
+                        continue
+                    add = 0.0 if child_bit == qb else c           # sibling costs its margin
+                    nxt.append((cost + add, clo, chi, cpre, t + 1))
+            nxt.sort(key=lambda e: e[0])
+            state = nxt[:beam]
+        return state
+
+    def query(self, q, topk=10, beam=64, depth=28, rerank="float"):
+        q = np.asarray(q, dtype=np.float32)
+        p = q @ self.R
+        qbits = (p > 0).astype(int)
+        qconf = np.abs(p)
+        ranges = self._descend(qbits, qconf, beam, depth)
+        cand = np.concatenate([self.docids[lo:hi] for _, lo, hi, _, _ in ranges]) if ranges else np.empty(0, np.int64)
+        if cand.size == 0:
+            return cand
+        if rerank == "float":                                     # DiskANN-style: cold float for re-rank only
+            s = self.X[cand] @ q
+        else:                                                     # asymmetric: float query vs sign code, no float
+            shifts = np.arange(self.b - 1, -1, -1, dtype=np.uint64)
+            sign = np.where(((self.codes_full(cand)[:, None] >> shifts) & np.uint64(1)) > 0, 1.0, -1.0)
+            s = sign @ (q @ self.R)
+        return cand[np.argsort(-s)[:topk]]
+
+    def query_batch(self, Q, topk=10, beam=64, depth=28):
+        """Top-k for every query in Q. Uses the compiled C kernel when available (one to two
+        orders of magnitude faster), else falls back to the pure-Python per-query path."""
+        Q = np.ascontiguousarray(Q, dtype=np.float32)
+        depth = min(depth, self.b)
+        from ._bittrie_build import get_lib
+        lib = get_lib()
+        if lib is None:
+            return np.array([self.query(q, topk, beam, depth, rerank="float") for q in Q])
+        import ctypes
+        qproj = np.ascontiguousarray(Q @ self.R, dtype=np.float32)
+        Xf = np.ascontiguousarray(self.X, dtype=np.float32)
+        out = np.full((len(Q), topk), -1, dtype=np.int64)
+        cast = lambda a, t: a.ctypes.data_as(ctypes.POINTER(t))
+        lib.bt_query_batch(
+            cast(self.codes, ctypes.c_uint64), cast(self.docids, ctypes.c_int64),
+            cast(Xf, ctypes.c_float), cast(qproj, ctypes.c_float), cast(Q, ctypes.c_float),
+            self.n, self.d, self.b, len(Q), int(depth), int(beam), int(topk),
+            cast(out, ctypes.c_int64))
+        return out
+
+    def codes_full(self, cand):
+        inv = np.empty(self.n, np.int64); inv[self.docids] = np.arange(self.n)
+        return self.codes[inv[cand]]
+
+    def index_bytes(self):
+        """Bytes that must be in RAM to route (codes + leaf ids); the trie is implicit in the sort."""
+        return self.codes.nbytes + self.docids.nbytes