cortexm 0.3.0__py3-none-any.whl

cortexm/vsa/codecs.py

@@ -0,0 +1,397 @@
+"""Vector codecs — the cortexm-compress tier stack.
+
+Per-vector storage at 768 dims (plan appendix "From 768 MB down to 8–96 MB"):
+
+  codec    bytes  tier        binding/similarity
+  -------  -----  ----------  ----------------------------------------
+  int8      768   baseline    scalar symmetric quantization (Aeon-style)
+  binary     96   edge        bipolar ±1, packed bits; XOR/permutation
+                             ops map 1:1 to HDC hardware
+  rabitq     96   ultra-edge  JL-rotation + binarization (RaBitQ-style),
+                             provable angle preservation
+  pq          8   cloud       product quantization M=8 x 8-bit codes,
+                             ADC scoring via L1-resident lookup tables
+
+All codecs expose: encode_packed / decoded / query_vec / scores /
+to_bytes / from_bytes / corrupt, so the Palace and the tree index are
+codec-agnostic. INT8 additionally stores a per-vector float16 scale
+(aux). Binary supports triple-modular redundancy (TMR) for the
+self-healing memory feature.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from cortexm.errors import CodecError
+from cortexm.util import h64 as _h64
+
+
+def _norm(v: np.ndarray) -> np.ndarray:
+    n = float(np.linalg.norm(v))
+    return (v / n).astype(np.float32) if n > 0 else v.astype(np.float32)
+
+
+class BaseCodec:
+    name = "base"
+
+    @property
+    def bytes_per_vector(self) -> int:
+        raise NotImplementedError
+
+    def encode_packed(self, vec: np.ndarray) -> np.ndarray: ...
+    def decoded(self, packed: np.ndarray) -> np.ndarray: ...
+    def query_vec(self, q: np.ndarray) -> np.ndarray: ...
+    def scores(self, packed: np.ndarray, q: np.ndarray) -> np.ndarray: ...
+    def to_bytes(self, packed_row: np.ndarray) -> bytes: ...
+    def from_bytes(self, b: bytes) -> np.ndarray: ...
+    def corrupt(self, packed_row: np.ndarray, rate: float,
+                rng: np.random.Generator) -> np.ndarray: ...
+
+
+# ---------------------------------------------------------------------------
+class Int8Codec(BaseCodec):
+    """Symmetric scalar quantization: q = round(127 * v / max|v|)."""
+
+    name = "int8"
+    uses_aux = True
+
+    def __init__(self, dims: int) -> None:
+        self.dims = dims
+
+    @property
+    def bytes_per_vector(self) -> int:
+        return self.dims + 2  # + float16 scale
+
+    def encode_packed(self, vec: np.ndarray) -> np.ndarray:
+        v = np.asarray(vec, dtype=np.float32)
+        m = float(np.max(np.abs(v))) or 1.0
+        scale = 127.0 / m
+        q = np.clip(np.rint(v * scale), -127, 127).astype(np.int8)
+        return q
+
+    @staticmethod
+    def encode_scale(vec: np.ndarray) -> np.float16:
+        v = np.asarray(vec, dtype=np.float32)
+        m = float(np.max(np.abs(v))) or 1.0
+        return np.float16(m / 127.0)
+
+    def decoded(self, packed: np.ndarray, aux: np.ndarray | None = None) -> np.ndarray:
+        arr = np.atleast_2d(packed).astype(np.float32)
+        if aux is not None:
+            arr = arr * np.atleast_1d(aux).astype(np.float32)[:, None]
+        return arr
+
+    def query_vec(self, q: np.ndarray) -> np.ndarray:
+        return np.asarray(q, dtype=np.float32)
+
+    def scores(self, packed: np.ndarray, q: np.ndarray,
+               aux: np.ndarray | None = None) -> np.ndarray:
+        if aux is None:
+            raise CodecError("int8 scoring requires per-vector aux scales")
+        arr = packed.astype(np.float32)
+        if len(aux):
+            arr = arr * np.atleast_1d(aux).astype(np.float32)[:, None]
+        return arr @ np.asarray(q, dtype=np.float32)
+
+    def to_bytes(self, packed_row: np.ndarray, scale: np.float16) -> bytes:
+        return packed_row.tobytes() + np.float16(scale).tobytes()
+
+    def from_bytes(self, b: bytes) -> tuple[np.ndarray, np.float16]:
+        q = np.frombuffer(b[: self.dims], dtype=np.int8).copy()
+        scale = np.frombuffer(b[self.dims:], dtype=np.float16)[0]
+        return q, scale
+
+    def corrupt(self, packed_row: np.ndarray, rate: float,
+                rng: np.random.Generator) -> np.ndarray:
+        n = len(packed_row)
+        idx = rng.integers(0, n, size=max(1, int(n * rate)))
+        out = packed_row.copy()
+        noise = rng.integers(-30, 31, size=len(idx)).astype(np.int8)
+        out[idx] = np.clip(out[idx].astype(np.int16) + noise, -127, 127).astype(np.int8)
+        return out
+
+
+# ---------------------------------------------------------------------------
+class BinaryCodec(BaseCodec):
+    """Bipolar {±1} hypervectors, 1 bit/dim — the HDC hardware target.
+
+    Sparse embeddings binarize poorly (near-zero components produce
+    correlated sign noise), so a fixed JL rotation densifies energy
+    before binarization — the RaBitQ insight applied to the MAP model.
+    Set ``rotate=False`` for raw bipolar-MAP semantics.
+    """
+
+    name = "binary"
+    uses_aux = False
+
+    def __init__(self, dims: int, tmr: bool = False, rotate: bool = True,
+                 seed: int = 0x0C0FFEE) -> None:
+        self.dims = dims
+        self.words = dims // 8
+        self.tmr = tmr
+        self.rotate = rotate
+        self.R = None
+        if rotate:
+            rng = np.random.default_rng(_h64("binary-rot", seed) & 0xFFFFFFFF)
+            g = rng.standard_normal((dims, dims)).astype(np.float32)
+            q, _ = np.linalg.qr(g)
+            self.R = q.astype(np.float32)
+
+    @property
+    def bytes_per_vector(self) -> int:
+        return self.words * (3 if self.tmr else 1)
+
+    def encode_packed(self, vec: np.ndarray) -> np.ndarray:
+        v = np.asarray(vec, dtype=np.float32)
+        if self.R is not None:
+            v = self.R @ v
+        bits = (v > 0).astype(np.uint8)
+        packed = np.packbits(bits)
+        if not self.tmr:
+            return packed
+        return np.concatenate([packed, packed, packed])
+
+    def decoded(self, packed: np.ndarray) -> np.ndarray:
+        arr = np.atleast_2d(packed)
+        if self.tmr:
+            k = arr.shape[1] // 3
+            arr = _tmr_majority(arr[:, :k], arr[:, k:2 * k], arr[:, 2 * k:])
+        signs = (np.unpackbits(arr, axis=1, count=self.dims).astype(np.float32) * 2 - 1)
+        return signs / np.sqrt(self.dims)
+
+    def query_vec(self, q: np.ndarray) -> np.ndarray:
+        v = np.asarray(q, dtype=np.float32)
+        if self.R is not None:
+            v = self.R @ v
+        s = ((v > 0).astype(np.float32) * 2 - 1) / np.sqrt(self.dims)
+        return s
+
+    def scores(self, packed: np.ndarray, q: np.ndarray) -> np.ndarray:
+        v = np.asarray(q, dtype=np.float32)
+        if self.R is not None:
+            v = self.R @ v
+        qbits = np.packbits((v > 0).astype(np.uint8))
+        arr = np.atleast_2d(packed)
+        if self.tmr:
+            k = arr.shape[1] // 3
+            arr = _tmr_majority(arr[:, :k], arr[:, k:2 * k], arr[:, 2 * k:])
+        x = np.bitwise_xor(arr, qbits[None, :])
+        ham = np.bitwise_count(x).sum(axis=1).astype(np.float32)
+        return 1.0 - 2.0 * ham / self.dims
+
+    def to_bytes(self, packed_row: np.ndarray) -> bytes:
+        return packed_row.tobytes()
+
+    def from_bytes(self, b: bytes) -> np.ndarray:
+        return np.frombuffer(b, dtype=np.uint8).copy()
+
+    def corrupt(self, packed_row: np.ndarray, rate: float,
+                rng: np.random.Generator) -> np.ndarray:
+        out = packed_row.copy()
+        n = len(out)
+        idx = rng.integers(0, n, size=max(1, int(n * rate)))
+        out[idx] ^= rng.integers(1, 256, size=len(idx)).astype(np.uint8)
+        return out
+
+    # TMR helpers -------------------------------------------------------
+    def tmr_health(self, packed_row: np.ndarray) -> int:
+        """Number of bit positions where the 3 copies disagree (corruption)."""
+        k = len(packed_row) // 3
+        a, b, c = packed_row[:k], packed_row[k:2 * k], packed_row[2 * k:]
+        ab = np.bitwise_count(np.bitwise_xor(a, b)).sum()
+        ac = np.bitwise_count(np.bitwise_xor(a, c)).sum()
+        bc = np.bitwise_count(np.bitwise_xor(b, c)).sum()
+        return int((ab + ac + bc) // 6)
+
+
+def _tmr_majority(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> np.ndarray:
+    """Bit-level majority vote across 3 copies (self-healing memory)."""
+    ab = a == b
+    bc = b == c
+    ac = a == c
+    maj = np.where(ab, a, np.where(ac, a, c))
+    maj = np.where(ab | bc | ac, maj, a)
+    return maj.astype(np.uint8)
+
+
+# ---------------------------------------------------------------------------
+class RaBitQCodec(BaseCodec):
+    """RaBitQ-style: fixed JL rotation, then binarization.
+
+    The random orthogonal rotation concentrates energy so that sign
+    binarization preserves angles far better than raw binarization
+    (Wang et al., SIGMOD 2024/2025). Simplified here: unit-norm vectors,
+    single global rotation, Hamming-based cosine estimate.
+    """
+
+    name = "rabitq"
+    uses_aux = False
+
+    def __init__(self, dims: int, seed: int = 0x0C0FFEE) -> None:
+        self.dims = dims
+        self.words = dims // 8
+        rng = np.random.default_rng(_h64("rabitq-rot", seed) & 0xFFFFFFFF)
+        g = rng.standard_normal((dims, dims)).astype(np.float32)
+        q, _ = np.linalg.qr(g)
+        self.R = q.astype(np.float32)
+
+    @property
+    def bytes_per_vector(self) -> int:
+        return self.words
+
+    def encode_packed(self, vec: np.ndarray) -> np.ndarray:
+        rot = self.R @ np.asarray(vec, dtype=np.float32)
+        return np.packbits((rot > 0).astype(np.uint8))
+
+    def decoded(self, packed: np.ndarray) -> np.ndarray:
+        signs = np.unpackbits(np.atleast_2d(packed), axis=1,
+                              count=self.dims).astype(np.float32) * 2 - 1
+        return signs / np.sqrt(self.dims)
+
+    def query_vec(self, q: np.ndarray) -> np.ndarray:
+        rot = self.R @ np.asarray(q, dtype=np.float32)
+        s = ((rot > 0).astype(np.float32) * 2 - 1) / np.sqrt(self.dims)
+        return s
+
+    def scores(self, packed: np.ndarray, q: np.ndarray) -> np.ndarray:
+        return self.scores_packed(packed, self.query_packed(q))
+
+    def query_packed(self, q: np.ndarray) -> np.ndarray:
+        rot = self.R @ np.asarray(q, dtype=np.float32)
+        return np.packbits((rot > 0).astype(np.uint8))
+
+    def scores_packed(self, packed: np.ndarray, qbits: np.ndarray) -> np.ndarray:
+        x = np.bitwise_xor(np.atleast_2d(packed), qbits[None, :])
+        ham = np.bitwise_count(x).sum(axis=1).astype(np.float32)
+        return 1.0 - 2.0 * ham / self.dims
+
+    def to_bytes(self, packed_row: np.ndarray) -> bytes:
+        return packed_row.tobytes()
+
+    def from_bytes(self, b: bytes) -> np.ndarray:
+        return np.frombuffer(b, dtype=np.uint8).copy()
+
+    def corrupt(self, packed_row: np.ndarray, rate: float,
+                rng: np.random.Generator) -> np.ndarray:
+        out = packed_row.copy()
+        idx = rng.integers(0, len(out), size=max(1, int(len(out) * rate)))
+        out[idx] ^= rng.integers(1, 256, size=len(idx)).astype(np.uint8)
+        return out
+
+    def serialize_rotation(self) -> bytes:
+        return self.R.tobytes()
+
+
+# ---------------------------------------------------------------------------
+class PQCodec(BaseCodec):
+    """Product quantization: M subspaces x 8-bit codes (8 bytes/vector).
+
+    Cloud tier. Codebooks trained on the first N vectors (k-means per
+    subspace, frozen afterwards); asymmetric distance computation via
+    per-query lookup tables that stay L1-resident.
+    """
+
+    name = "pq"
+    uses_aux = False
+
+    def __init__(self, dims: int, m: int = 8, ks: int = 256,
+                 train_threshold: int = 4096, seed: int = 0x0C0FFEE) -> None:
+        if dims % m:
+            raise CodecError("dims must be divisible by M")
+        self.dims = dims
+        self.m = m
+        self.ks = ks
+        self.sub = dims // m
+        self.seed = seed
+        self.train_threshold = train_threshold
+        self.codebooks: np.ndarray | None = None   # (m, ks, sub)
+
+    @property
+    def trained(self) -> bool:
+        return self.codebooks is not None
+
+    @property
+    def bytes_per_vector(self) -> int:
+        return self.m
+
+    def train(self, vecs: np.ndarray) -> None:
+        from scipy.cluster.vq import kmeans2
+        rng = np.random.default_rng(self.seed)
+        data = np.asarray(vecs, dtype=np.float32)
+        n = min(self.ks * 4, len(data))
+        sample = data[rng.choice(len(data), size=n, replace=False)] if len(data) > n else data
+        books = np.zeros((self.m, self.ks, self.sub), dtype=np.float32)
+        for mi in range(self.m):
+            chunk = sample[:, mi * self.sub:(mi + 1) * self.sub]
+            k = min(self.ks, max(2, len(np.unique(chunk, axis=0))))
+            try:
+                cent, _ = kmeans2(chunk, k, iter=12, minit="++", seed=self.seed + mi)
+            except Exception:
+                cent = chunk[:k]
+            books[mi, : len(cent)] = cent
+        self.codebooks = books
+
+    def set_codebooks(self, books: np.ndarray) -> None:
+        self.codebooks = np.asarray(books, dtype=np.float32)
+
+    def encode_packed(self, vec: np.ndarray) -> np.ndarray:
+        if self.codebooks is None:
+            raise CodecError("PQ codec not trained")
+        v = np.asarray(vec, dtype=np.float32)
+        codes = np.zeros(self.m, dtype=np.uint8)
+        for mi in range(self.m):
+            chunk = v[mi * self.sub:(mi + 1) * self.sub]
+            sims = self.codebooks[mi] @ chunk
+            codes[mi] = int(np.argmax(sims))
+        return codes
+
+    def decoded(self, packed: np.ndarray) -> np.ndarray:
+        if self.codebooks is None:
+            raise CodecError("PQ codec not trained")
+        arr = np.atleast_2d(packed)
+        out = np.zeros((arr.shape[0], self.dims), dtype=np.float32)
+        for mi in range(self.m):
+            out[:, mi * self.sub:(mi + 1) * self.sub] = self.codebooks[mi][arr[:, mi]]
+        return out
+
+    def query_vec(self, q: np.ndarray) -> np.ndarray:
+        return np.asarray(q, dtype=np.float32)
+
+    def scores(self, packed: np.ndarray, q: np.ndarray) -> np.ndarray:
+        if self.codebooks is None:
+            raise CodecError("PQ codec not trained")
+        arr = np.atleast_2d(packed)
+        q = np.asarray(q, dtype=np.float32)
+        total = np.zeros(arr.shape[0], dtype=np.float32)
+        for mi in range(self.m):
+            qsub = q[mi * self.sub:(mi + 1) * self.sub]
+            lut = self.codebooks[mi] @ qsub          # (ks,)
+            total += lut[arr[:, mi]]
+        return total
+
+    def to_bytes(self, packed_row: np.ndarray) -> bytes:
+        return packed_row.tobytes()
+
+    def from_bytes(self, b: bytes) -> np.ndarray:
+        return np.frombuffer(b, dtype=np.uint8).copy()
+
+    def corrupt(self, packed_row: np.ndarray, rate: float,
+                rng: np.random.Generator) -> np.ndarray:
+        out = packed_row.copy()
+        idx = rng.integers(0, len(out), size=max(1, int(len(out) * rate)))
+        out[idx] = rng.integers(0, self.ks, size=len(idx)).astype(np.uint8)
+        return out
+
+
+def make_codec(name: str, dims: int, seed: int = 0x0C0FFEE,
+               tmr: bool = False, **kw) -> BaseCodec:
+    if name == "int8":
+        return Int8Codec(dims)
+    if name == "binary":
+        return BinaryCodec(dims, tmr=tmr, seed=seed)
+    if name == "rabitq":
+        return RaBitQCodec(dims, seed=seed)
+    if name == "pq":
+        return PQCodec(dims, seed=seed, **kw)
+    raise CodecError(f"unknown codec {name!r}")

cortexm/vsa/hologram_overlay.py

@@ -0,0 +1,139 @@
+"""HRR KG overlay — superposed holographic fact store for O(1) lookup.
+
+The Holographic Memory for Knowledge Graphs paper (arXiv:2606.24948,
+2026) shows that facts bound as (s, r, v) can be superposed into a
+single dense hologram M, and a query (s, r, ?) answered by unbinding
+the V slot from M:
+
+    M = Σ_f bind_V(v_f) * (bind_S(s_f) + bind_R(r_f))   # superposition
+    v_hat = unbind_V(M) → noisy residual ≈ Σ_f (bind_S(s_f) + bind_R(r_f)) * v_f
+    cleanup(v_hat) → snapped to nearest stored value vector
+
+Capacity bound: ~dims^2 / ln(dims) facts before saturation (Clarkson
+2023). At d=768 → ~93k facts; at d=16,384 → ~31M facts.
+
+Best used as a per-scope overlay on top of MemoryPalace — single-hop
+queries hit the overlay (O(1) + cleanup), multi-hop queries fall back
+to TreeIndex search. The overlay trades noise accumulation for lookup
+speed; saturation is detected by signal-to-noise ratio on the cleanup
+step.
+
+Pure numpy. Reuses VSA.bind/unbind and HopfieldCleanup.recall.
+
+arxiv research: arXiv:2606.24948 (2026 holographic KG); HolE AAAI 2016
+(Nickel, Rosasco, Poggio); Plate 1995 (HRR).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from cortexm.vsa.cleanup import HopfieldCleanup
+from cortexm.vsa.ops import VSA
+
+
+def _norm(v: np.ndarray) -> np.ndarray:
+    n = float(np.linalg.norm(v))
+    return v / n if n > 0 else v
+
+
+class HolographicFactOverlay:
+    """Per-scope superposed hologram for O(1) single-hop fact lookup.
+
+    Each (user_id, scope) tuple gets one dense hologram M. Adding a fact
+    binds (s, r, v) and adds to M. Querying unbinds the requested role
+    and snaps to the nearest stored item via Hopfield cleanup.
+
+    Falls back gracefully when M saturates — caller can switch to
+    MemoryPalace.search for the affected scope.
+    """
+
+    def __init__(self, vsa: VSA, cleanup: HopfieldCleanup,
+                 saturate_threshold: float = 0.55) -> None:
+        self.vsa = vsa
+        self.cleanup = cleanup
+        self.saturate_threshold = saturate_threshold
+        self._M: dict[tuple, np.ndarray] = {}     # scope → hologram
+        self._counts: dict[tuple, int] = {}        # scope → fact count
+        self._saturated: set[tuple] = set()
+
+    def add_fact(self, scope: tuple, s_vec: np.ndarray,
+                r_vec: np.ndarray, v_vec: np.ndarray) -> None:
+        """Add a fact (s, r, v) into the superposed hologram."""
+        # bind each role to its filler
+        bound = (self.vsa.bind("S", s_vec)
+                 + self.vsa.bind("R", r_vec)
+                 + self.vsa.bind("V", v_vec))
+        bound = _norm(bound)
+        h = self._M.get(scope)
+        if h is None:
+            self._M[scope] = bound.copy()
+            self._counts[scope] = 1
+        else:
+            self._M[scope] = _norm(h + bound)
+            self._counts[scope] = self._counts.get(scope, 0) + 1
+        # populate cleanup codebook with each filler
+        self.cleanup.add(f"s:{_hash_vec(s_vec)}", s_vec)
+        self.cleanup.add(f"r:{_hash_vec(r_vec)}", r_vec)
+        self.cleanup.add(f"v:{_hash_vec(v_vec)}", v_vec)
+
+    def query(self, scope: tuple, query_vec: np.ndarray,
+              target_role: str = "V",
+              fallback_embs: list[tuple[str, np.ndarray]] | None = None
+             ) -> tuple[str | None, float]:
+        """Query (s, r, ?) — unbind target_role and cleanup the residual.
+
+        Returns (item_key, confidence). Confidence below saturate_threshold
+        indicates either saturation or miss; caller should fall back to
+        MemoryPalace.search.
+
+        If fallback_embs is provided, those (key, vec) pairs are temporarily
+        added to the cleanup codebook before recall — useful for cross-
+        scope queries where the answer may be a value vector not yet in
+        the codebook.
+        """
+        M = self._M.get(scope)
+        if M is None:
+            return None, 0.0
+        if scope in self._saturated:
+            # signal saturation — caller should fallback to TreeIndex
+            return None, 0.0
+        # unbind the target role from M to get a noisy residual
+        residual = self.vsa.unbind(target_role, M)
+        # add fallback embeddings to cleanup if provided
+        if fallback_embs:
+            for k, v in fallback_embs:
+                self.cleanup.add(k, v)
+        self.cleanup.build()
+        item_key, conf = self.cleanup.recall(residual)
+        if conf < self.saturate_threshold:
+            # either miss or saturation — mark saturated if we have many
+            # facts and conf is consistently low
+            if self._counts.get(scope, 0) > 1000:
+                self._saturated.add(scope)
+            return None, conf
+        return item_key, conf
+
+    def stats(self) -> dict:
+        return {
+            "scopes": len(self._M),
+            "total_facts": sum(self._counts.values()),
+            "saturated_scopes": len(self._saturated),
+            "cleanup": self.cleanup.stats(),
+        }
+
+    def reset_scope(self, scope: tuple) -> None:
+        """Drop a saturated scope and let it rebuild from new adds."""
+        self._M.pop(scope, None)
+        self._counts.pop(scope, None)
+        self._saturated.discard(scope)
+
+
+def _hash_vec(v: np.ndarray) -> str:
+    """Stable hash for codebook keys."""
+    import hashlib
+    h = hashlib.blake2b(v.tobytes(), digest_size=8).hexdigest()
+    return h[:16]
+
+
+__all__ = ["HolographicFactOverlay"]

cortexm/vsa/index.py

@@ -0,0 +1,163 @@
+"""Page-Clustered Vector Index — hierarchical tree, O(log N) retrieval.
+
+Per the plan (Aeon-style): leaf pages hold up to ``leaf_size`` quantized
+vectors; internal nodes store an fp32 centroid + radius (1 - min cos).
+Search is best-first over the bound ``centroid_sim - radius`` with exact
+scoring inside visited pages only — sub-millisecond at 100K+ vectors
+while brute-force scans the whole palace.
+"""
+
+from __future__ import annotations
+
+import heapq
+import numpy as np
+
+from cortexm.vsa.codecs import BaseCodec
+
+
+class _Node:
+    __slots__ = ("centroid", "radius", "children", "rows", "is_leaf")
+
+    def __init__(self) -> None:
+        self.centroid: np.ndarray | None = None
+        self.radius: float = 1.0
+        self.children: list["_Node"] = []
+        self.rows: np.ndarray | None = None
+        self.is_leaf = True
+
+
+class TreeIndex:
+    def __init__(self, codec: BaseCodec, rows_getter, n: int, *,
+                 branch: int = 8, leaf: int = 512, seed: int = 0x0C0FFEE,
+                 sample_limit: int = 20000) -> None:
+        """``rows_getter(indices) -> (packed, aux)`` fetches codec rows."""
+        self.codec = codec
+        self.get_rows = rows_getter
+        self.n = n
+        self.branch = branch
+        self.leaf = leaf
+        self.seed = seed
+        self.sample_limit = sample_limit
+        self.root: _Node | None = None
+        self.leaf_rows_scanned = 0
+
+    # ------------------------------------------------------------- build
+    def build(self) -> None:
+        if self.n == 0:
+            return
+        rng = np.random.default_rng(self.seed)
+        self.root = self._build(np.arange(self.n, dtype=np.int64), rng, depth=0)
+
+    def _decode(self, idx: np.ndarray):
+        packed, aux = self.get_rows(idx)
+        return self.codec.decoded(packed, aux) if self.codec.uses_aux else self.codec.decoded(packed)
+
+    def _build(self, idx: np.ndarray, rng: np.random.Generator, depth: int) -> _Node:
+        node = _Node()
+        # sample for centroid + kmeans
+        take = min(len(idx), self.sample_limit if depth == 0 else 4096)
+        sample_idx = idx if len(idx) <= take else idx[rng.choice(len(idx), take, replace=False)]
+        sample = self._decode(sample_idx).astype(np.float32)
+        centroid = sample.mean(axis=0)
+        cn = float(np.linalg.norm(centroid))
+        centroid = centroid / cn if cn > 0 else centroid
+        node.centroid = centroid
+        # EXACT radius w.r.t. this node's own centroid (streamed over all rows)
+        max_dist = 0.0
+        for i0 in range(0, len(idx), 8192):
+            batch = self._decode(idx[i0:i0 + 8192]).astype(np.float32)
+            sims = batch @ centroid
+            if len(sims):
+                max_dist = max(max_dist, float(1.0 - sims.min()))
+        node.radius = max_dist + 1e-6
+        if len(idx) <= self.leaf or depth > 24:
+            node.rows = idx
+            node.is_leaf = True
+            return node
+        k = min(self.branch, len(idx))
+        cent = _kmeans(sample, k, rng)
+        # assign all rows (streamed)
+        parts: list[list[int]] = [[] for _ in range(k)]
+        for i0 in range(0, len(idx), 8192):
+            batch_idx = idx[i0:i0 + 8192]
+            batch = self._decode(batch_idx).astype(np.float32)
+            sims = batch @ cent.T                     # (b, k)
+            assign = np.argmax(sims, axis=1)
+            for j, a in enumerate(assign):
+                parts[int(a)].append(int(batch_idx[j]))
+        node.is_leaf = False
+        for p in parts:
+            if p:
+                node.children.append(
+                    self._build(np.array(p, dtype=np.int64), rng, depth + 1))
+        if len(node.children) == 1:
+            return node.children[0]
+        return node
+
+    # ------------------------------------------------------------ search
+    def search(self, q: np.ndarray, k: int, beam: int = 4) -> tuple[np.ndarray, np.ndarray]:
+        """Return (row_indices, scores) of approximate top-k."""
+        if self.root is None or self.n == 0:
+            return np.array([], dtype=np.int64), np.array([], dtype=np.float32)
+        qv = self.codec.query_vec(q)
+        heap: list[tuple[float, int, _Node]] = []
+        counter = 0
+        csim = float(qv @ self.root.centroid)
+        heapq.heappush(heap, (-(csim - self.root.radius), counter, self.root))
+        best: list[tuple[float, int]] = []            # (score, row)
+        visited_leaves = 0
+        while heap:
+            negbound, _, node = heapq.heappop(heap)
+            bound = -negbound
+            if len(best) >= k and bound <= best[0][0] and visited_leaves >= beam:
+                break
+            if node.is_leaf:
+                rows = node.rows
+                packed, aux = self.get_rows(rows)
+                sc = (self.codec.scores(packed, q, aux) if self.codec.uses_aux
+                      else self.codec.scores(packed, q))
+                for r, s in zip(rows.tolist(), sc.tolist()):
+                    if len(best) < k:
+                        heapq.heappush(best, (float(s), int(r)))
+                    elif float(s) > best[0][0]:
+                        heapq.heapreplace(best, (float(s), int(r)))
+                visited_leaves += 1
+                self.leaf_rows_scanned += len(rows)
+                if visited_leaves >= max(beam, 1) * 8:
+                    break
+            else:
+                for ch in node.children:
+                    c = float(qv @ ch.centroid)
+                    heapq.heappush(heap, (-(c - ch.radius), counter := counter + 1, ch))
+        order = sorted(best, key=lambda t: -t[0])
+        return (np.array([r for _, r in order], dtype=np.int64),
+                np.array([s for s, _ in order], dtype=np.float32))
+
+
+def _kmeans(data: np.ndarray, k: int, rng: np.random.Generator,
+            iters: int = 8) -> np.ndarray:
+    """Small deterministic k-means (kmeans++ style init)."""
+    n = len(data)
+    if n <= k:
+        return data.copy()
+    # kmeans++ init
+    cent = [data[rng.integers(n)]]
+    d = 1.0 - data @ cent[0]
+    d = np.maximum(d, 1e-8)
+    for _ in range(1, k):
+        probs = d / d.sum()
+        cent.append(data[rng.choice(n, p=probs)])
+        d = np.minimum(d, np.maximum(1.0 - data @ cent[-1], 1e-8))
+    C = np.stack(cent).astype(np.float32)
+    for _ in range(iters):
+        sims = data @ C.T
+        assign = np.argmax(sims, axis=1)
+        for j in range(k):
+            mask = assign == j
+            if mask.any():
+                C[j] = data[mask].mean(axis=0)
+            else:
+                C[j] = data[rng.integers(n)]
+        cn = np.linalg.norm(C, axis=1, keepdims=True)
+        C = C / np.maximum(cn, 1e-8)
+    return C