e8_utils.py

"""
e8_utils.py — fully self-contained E8 lattice utilities for the LatticeMemory HF Space.

No imports from liora_core. All math is inlined.
"""
from __future__ import annotations

import math
from itertools import combinations

import torch
import torch.nn.functional as F


# ---------------------------------------------------------------------------
# Core E8 lattice math
# ---------------------------------------------------------------------------

def _decode_d8(x: torch.Tensor) -> torch.Tensor:
    """Round x to the nearest D8 lattice point (even-sum integer coordinates)."""
    z = torch.round(x)
    parity = z.sum(dim=-1) % 2
    diff = x - z
    worst = diff.abs().argmax(dim=-1)
    adj = torch.sign(diff)
    adj = torch.where(adj == 0, torch.ones_like(adj), adj)
    mask = torch.zeros_like(x)
    mask.scatter_(-1, worst.unsqueeze(-1), 1.0)
    z_fixed = z + adj * mask
    return torch.where(parity.unsqueeze(-1) == 1, z_fixed, z)


def _e8_nearest(x: torch.Tensor) -> torch.Tensor:
    """Return the nearest E8 lattice point to x (batched, last dim = 8)."""
    z0 = _decode_d8(x)
    z1 = _decode_d8(x - 0.5) + 0.5
    d0 = (x - z0).pow(2).sum(dim=-1, keepdim=True)
    d1 = (x - z1).pow(2).sum(dim=-1, keepdim=True)
    return torch.where(d0 <= d1, z0, z1)


# ---------------------------------------------------------------------------
# Shell-1 codebook  (240 vectors, shape [240, 8])
# ---------------------------------------------------------------------------

def _build_shell1_codebook(device: torch.device = torch.device("cpu")) -> torch.Tensor:
    """Build the 240-vector E8 shell-1 codebook.

    E8 shell-1 consists of:
      - 112 vectors of the form (±1, ±1, 0, 0, 0, 0, 0, 0) in all permutations
      - 128 vectors of the form (±½, ±½, …, ±½) with an even number of minus signs
    Total: 112 + 128 = 240 vectors.
    """
    vecs: list[list[float]] = []
    # ±1 in two coordinates, rest 0
    for i, j in combinations(range(8), 2):
        for si in (1.0, -1.0):
            for sj in (1.0, -1.0):
                v = [0.0] * 8
                v[i] = si
                v[j] = sj
                vecs.append(v)
    # (±½)^8 with even number of minus signs
    for mask in range(256):
        signs = [1.0 if (mask >> bit) & 1 == 0 else -1.0 for bit in range(8)]
        if signs.count(-1.0) % 2 == 0:
            vecs.append([s * 0.5 for s in signs])
    codebook = torch.tensor(vecs, dtype=torch.float32, device=device)
    if codebook.shape != (240, 8):
        raise RuntimeError(
            f"expected E8 shell-1 codebook shape (240, 8), got {tuple(codebook.shape)}"
        )
    return codebook


# ---------------------------------------------------------------------------
# RF-Snap: batch snap embeddings to E8
# ---------------------------------------------------------------------------

@torch.no_grad()
def nestquant_snap(embeddings: torch.Tensor) -> torch.Tensor:
    """Snap a batch of embeddings [B, D] (D divisible by 8) to E8 lattice points.

    Each 8-dim block is independently scaled, snapped to the nearest E8 point,
    and rescaled back.  The operation is a no-op in the limit of small beta.

    Args:
        embeddings: float32 tensor of shape [B, D].

    Returns:
        Snapped embeddings of shape [B, D].
    """
    if embeddings.dim() != 2:
        raise ValueError(f"nestquant_snap expects [B, D], got {tuple(embeddings.shape)}")
    B, D = embeddings.shape
    if D % 8 != 0:
        raise ValueError(f"D={D} must be divisible by 8")

    blocks = embeddings.float().reshape(B, D // 8, 8)           # [B, n_blocks, 8]
    beta = blocks.norm(p=2, dim=-1).clamp_min(1e-8) / math.sqrt(2.0)  # [B, n_blocks]
    snapped = _e8_nearest(blocks / beta.unsqueeze(-1))          # [B, n_blocks, 8]
    return (snapped * beta.unsqueeze(-1)).reshape(B, D)


# ---------------------------------------------------------------------------
# E8 address: encode a single embedding as a hex string
# ---------------------------------------------------------------------------

@torch.no_grad()
def embedding_to_e8_address(embedding: torch.Tensor, device: torch.device = torch.device("cpu")) -> str:
    """Convert a single embedding vector to its E8 lattice address (hex string).

    Each 8-dim block is mapped to an index in [0, 239] (the closest shell-1
    codebook vector after unit-normalisation), then packed as a byte.  The
    resulting byte string is returned as a lowercase hex string.

    Expected address length: (D // 8) * 2 hex chars.  For D=1024 → 256 chars.

    Args:
        embedding: 1-D float tensor of shape [D].
        device: torch device for codebook (CPU by default).

    Returns:
        Lowercase hex string of length D // 4.
    """
    if embedding.dim() != 1:
        raise ValueError(f"embedding_to_e8_address expects 1-D tensor, got {tuple(embedding.shape)}")
    D = embedding.numel()
    if D % 8 != 0:
        raise ValueError(f"D={D} must be divisible by 8")

    codebook = _build_shell1_codebook(device)  # [240, 8]

    vector = embedding.float().to(device)
    blocks = vector.reshape(-1, 8)                                       # [n_blocks, 8]
    beta = blocks.norm(p=2, dim=-1).clamp_min(1e-8) / math.sqrt(2.0)   # [n_blocks]
    snapped = _e8_nearest(blocks / beta.unsqueeze(-1))                   # [n_blocks, 8]

    # Unit-normalise snapped blocks and project onto codebook
    snapped_unit = snapped / snapped.norm(dim=-1, keepdim=True).clamp_min(1e-8)
    dots = snapped_unit @ codebook.T / math.sqrt(2.0)                   # [n_blocks, 240]
    indices = dots.argmax(dim=-1)                                        # [n_blocks], values in [0,239]

    return bytes(indices.tolist()).hex()


# ---------------------------------------------------------------------------
# Index size comparison
# ---------------------------------------------------------------------------

def index_size_bytes(n_docs: int, d_model: int) -> dict[str, int]:
    """Return byte counts for different index formats.

    Keys:
      float32  — raw fp32 storage (4 bytes/float)
      int4     — 4-bit quantization (0.5 bytes/float)
      rfsnap   — RF-Snap E8 (3 bytes per 8-dim block: 1 byte index + 2 bytes fp16 scale)

    Args:
        n_docs:  number of document vectors.
        d_model: embedding dimension (must be divisible by 8).

    Returns:
        Dict with keys 'float32', 'int4', 'rfsnap' and integer byte values.
    """
    if d_model % 8 != 0:
        raise ValueError(f"d_model={d_model} must be divisible by 8")
    n_blocks = d_model // 8
    return {
        "float32": n_docs * d_model * 4,
        "int4":    n_docs * d_model // 2,
        "rfsnap":  n_docs * n_blocks * 3,
    }