pyturboquant-cpu 0.1.0__py3-none-any.whl

pyturboquant_cpu/__init__.py

@@ -0,0 +1,36 @@
+"""
+pyturboquant-cpu: CPU implementation of TurboQuant vector quantization.
+
+TurboQuant is a data-oblivious vector quantization algorithm from
+"TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate"
+(Zandieh et al., arXiv:2504.19874).
+
+Two quantization modes are provided:
+
+- **MSE-optimal** (``quantize_mse`` / ``dequantize_mse``):
+  Minimises mean-squared reconstruction error.
+- **Inner-product-optimal** (``quantize_prod`` / ``dequantize_prod``):
+  Gives an *unbiased* inner-product estimator by combining an MSE
+  quantizer with a 1-bit QJL residual correction.
+"""
+
+from pyturboquant_cpu.quantizer import (
+    QuantizedMSE,
+    QuantizedProd,
+    quantize_mse,
+    dequantize_mse,
+    quantize_prod,
+    dequantize_prod,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "quantize_mse",
+    "dequantize_mse",
+    "quantize_prod",
+    "dequantize_prod",
+    "QuantizedMSE",
+    "QuantizedProd",
+    "__version__",
+]

pyturboquant_cpu/codebook.py

@@ -0,0 +1,182 @@
+"""
+Lloyd-Max codebook computation for the Beta coordinate distribution.
+
+After a random orthogonal rotation, each coordinate of a unit-norm vector
+in R^d follows a Beta((d-1)/2, (d-1)/2) distribution on [-1, 1].
+
+This module solves the 1-D Lloyd-Max (k-means) quantisation problem for
+that distribution and caches the results for reuse.
+"""
+
+from __future__ import annotations
+
+import functools
+import math
+from typing import Tuple
+
+import numpy as np
+from scipy import integrate, special
+
+
+def beta_pdf(x: np.ndarray | float, dim: int) -> np.ndarray | float:
+    """Probability density of a single coordinate after random rotation.
+
+    For a unit vector uniformly distributed on S^{d-1}, each coordinate
+    follows the distribution:
+
+        f(x) = Γ(d/2) / (√π · Γ((d-1)/2)) · (1 - x²)^{(d-3)/2}
+
+    for x ∈ [-1, 1].
+
+    Parameters
+    ----------
+    x : array_like
+        Points at which to evaluate the PDF.
+    dim : int
+        Ambient dimension *d*.
+
+    Returns
+    -------
+    array_like
+        PDF values.
+    """
+    x = np.asarray(x, dtype=np.float64)
+    d = dim
+
+    # Use log-gamma for numerical stability
+    log_norm = (
+        special.gammaln(d / 2.0)
+        - 0.5 * math.log(math.pi)
+        - special.gammaln((d - 1) / 2.0)
+    )
+
+    exponent = (d - 3) / 2.0
+
+    # Mask values outside support
+    mask = np.abs(x) < 1.0
+    result = np.zeros_like(x, dtype=np.float64)
+    if exponent == 0:
+        # d = 3: uniform on [-1, 1]
+        result[mask] = math.exp(log_norm)
+    else:
+        safe_x = np.where(mask, x, 0.0)
+        log_body = exponent * np.log1p(-safe_x**2)
+        result[mask] = np.exp(log_norm + log_body[mask])
+
+    return result
+
+
+def _conditional_expectation(a: float, b: float, dim: int) -> float:
+    """Compute E[X | a ≤ X ≤ b] for the Beta coordinate distribution.
+
+    Parameters
+    ----------
+    a, b : float
+        Interval endpoints within [-1, 1].
+    dim : int
+        Ambient dimension *d*.
+
+    Returns
+    -------
+    float
+        Conditional mean.
+    """
+    if b - a < 1e-15:
+        return (a + b) / 2.0
+
+    # Numerator: ∫_a^b x · f(x) dx
+    numerator, _ = integrate.quad(lambda x: x * beta_pdf(x, dim), a, b)
+    # Denominator: ∫_a^b f(x) dx
+    denominator, _ = integrate.quad(lambda x: beta_pdf(x, dim), a, b)
+
+    if abs(denominator) < 1e-30:
+        return (a + b) / 2.0
+
+    return numerator / denominator
+
+
+def lloyd_max_codebook(
+    dim: int,
+    bits: int,
+    max_iter: int = 300,
+    tol: float = 1e-14,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute Lloyd-Max optimal codebook for the Beta coordinate distribution.
+
+    Solves the 1-D k-means problem for the distribution induced by random
+    rotation on the unit sphere in dimension *dim*, using *2^bits* centroids.
+
+    Parameters
+    ----------
+    dim : int
+        Vector dimension *d*.
+    bits : int
+        Bit-width per coordinate.  Produces ``2**bits`` centroids.
+    max_iter : int, optional
+        Maximum Lloyd iterations.
+    tol : float, optional
+        Convergence tolerance on centroid movement.
+
+    Returns
+    -------
+    centroids : ndarray, shape (2**bits,)
+        Sorted codebook centroids.
+    boundaries : ndarray, shape (2**bits - 1,)
+        Decision boundaries (midpoints between consecutive centroids).
+    """
+    n_levels = 1 << bits  # 2^bits
+
+    # Initialise centroids uniformly in [-1, 1]
+    centroids = np.linspace(-1.0, 1.0, n_levels + 2)[1:-1].copy()
+
+    for _ in range(max_iter):
+        # Boundaries = midpoints between consecutive centroids
+        boundaries = 0.5 * (centroids[:-1] + centroids[1:])
+
+        # Build full interval edges: [-1, b_1, b_2, ..., b_{n-1}, 1]
+        edges = np.empty(n_levels + 1, dtype=np.float64)
+        edges[0] = -1.0
+        edges[-1] = 1.0
+        edges[1:-1] = boundaries
+
+        # Update each centroid to the conditional expectation over its cell
+        new_centroids = np.empty_like(centroids)
+        for i in range(n_levels):
+            new_centroids[i] = _conditional_expectation(
+                edges[i], edges[i + 1], dim
+            )
+
+        delta = np.max(np.abs(new_centroids - centroids))
+        centroids = new_centroids
+
+        if delta < tol:
+            break
+
+    # Final boundaries
+    boundaries = 0.5 * (centroids[:-1] + centroids[1:])
+    return centroids.astype(np.float64), boundaries.astype(np.float64)
+
+
+@functools.lru_cache(maxsize=256)
+def get_codebook(
+    dim: int, bits: int
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Return a cached Lloyd-Max codebook for the given dimension and bit-width.
+
+    On the first call for a given ``(dim, bits)`` pair the codebook is
+    computed via :func:`lloyd_max_codebook` and then cached for all
+    subsequent calls.
+
+    Parameters
+    ----------
+    dim : int
+        Vector dimension.
+    bits : int
+        Bit-width per coordinate.
+
+    Returns
+    -------
+    centroids : ndarray, shape (2**bits,)
+    boundaries : ndarray, shape (2**bits - 1,)
+    """
+    return lloyd_max_codebook(dim, bits)

pyturboquant_cpu/packing.py

@@ -0,0 +1,258 @@
+"""
+Bit-packing and unpacking utilities.
+
+Quantised indices (b bits each) and QJL sign bits are packed into
+contiguous ``uint8`` arrays for compact storage.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def pack_indices(indices: np.ndarray, bits: int) -> np.ndarray:
+    """Pack an array of integer indices into a compact ``uint8`` byte array.
+
+    Each index value must be in ``[0, 2**bits)``.  The indices are stored
+    in a flat bit-stream, packed from the LSB of each output byte.
+
+    Parameters
+    ----------
+    indices : ndarray of int
+        Flat array of index values.
+    bits : int
+        Number of bits per index (1–8).
+
+    Returns
+    -------
+    packed : ndarray of uint8
+        Byte array containing the packed bits.
+    """
+    indices = np.asarray(indices, dtype=np.uint32).ravel()
+    n = indices.size
+    total_bits = n * bits
+    total_bytes = (total_bits + 7) // 8
+
+    packed = np.zeros(total_bytes, dtype=np.uint8)
+
+    # Build a flat bit stream
+    for b in range(bits):
+        # Extract bit `b` from each index
+        bit_values = ((indices >> b) & 1).astype(np.uint8)
+        for i, bv in enumerate(bit_values):
+            flat_bit_pos = i * bits + b
+            byte_idx = flat_bit_pos // 8
+            bit_idx = flat_bit_pos % 8
+            packed[byte_idx] |= bv << bit_idx
+
+    return packed
+
+
+def unpack_indices(packed: np.ndarray, bits: int, count: int) -> np.ndarray:
+    """Unpack integer indices from a packed ``uint8`` byte array.
+
+    Parameters
+    ----------
+    packed : ndarray of uint8
+        Byte array produced by :func:`pack_indices`.
+    bits : int
+        Number of bits per index (1–8).
+    count : int
+        Number of indices to unpack.
+
+    Returns
+    -------
+    indices : ndarray of uint32
+        Recovered index values.
+    """
+    packed = np.asarray(packed, dtype=np.uint8).ravel()
+    indices = np.zeros(count, dtype=np.uint32)
+
+    for i in range(count):
+        val = np.uint32(0)
+        for b in range(bits):
+            flat_bit_pos = i * bits + b
+            byte_idx = flat_bit_pos // 8
+            bit_idx = flat_bit_pos % 8
+            bit_val = (packed[byte_idx] >> bit_idx) & 1
+            val |= np.uint32(bit_val) << np.uint32(b)
+        indices[i] = val
+
+    return indices
+
+
+def pack_indices_fast(indices: np.ndarray, bits: int) -> np.ndarray:
+    """Vectorised bit-packing for common bit-widths.
+
+    Falls back to :func:`pack_indices` for uncommon widths.
+
+    Parameters
+    ----------
+    indices : ndarray of int
+        Flat array of index values.
+    bits : int
+        Number of bits per index.
+
+    Returns
+    -------
+    packed : ndarray of uint8
+    """
+    indices = np.asarray(indices, dtype=np.uint8).ravel()
+
+    if bits == 8:
+        return indices.copy()
+
+    if bits == 4:
+        # Pack two 4-bit values per byte
+        n = indices.size
+        padded = indices
+        if n % 2 != 0:
+            padded = np.append(indices, np.uint8(0))
+        low = padded[0::2] & 0x0F
+        high = (padded[1::2] & 0x0F) << 4
+        return (low | high).astype(np.uint8)
+
+    if bits == 2:
+        # Pack four 2-bit values per byte
+        n = indices.size
+        pad_len = (4 - n % 4) % 4
+        if pad_len:
+            padded = np.append(indices, np.zeros(pad_len, dtype=np.uint8))
+        else:
+            padded = indices
+        a = padded[0::4] & 0x03
+        b = (padded[1::4] & 0x03) << 2
+        c = (padded[2::4] & 0x03) << 4
+        d = (padded[3::4] & 0x03) << 6
+        return (a | b | c | d).astype(np.uint8)
+
+    if bits == 1:
+        return pack_signs_raw(indices)
+
+    # General fallback
+    return pack_indices(indices.astype(np.uint32), bits)
+
+
+def unpack_indices_fast(
+    packed: np.ndarray, bits: int, count: int
+) -> np.ndarray:
+    """Vectorised unpacking for common bit-widths.
+
+    Parameters
+    ----------
+    packed : ndarray of uint8
+    bits : int
+    count : int
+
+    Returns
+    -------
+    indices : ndarray of uint32
+    """
+    packed = np.asarray(packed, dtype=np.uint8).ravel()
+
+    if bits == 8:
+        return packed[:count].astype(np.uint32)
+
+    if bits == 4:
+        low = packed & 0x0F
+        high = (packed >> 4) & 0x0F
+        interleaved = np.empty(packed.size * 2, dtype=np.uint32)
+        interleaved[0::2] = low
+        interleaved[1::2] = high
+        return interleaved[:count]
+
+    if bits == 2:
+        a = packed & 0x03
+        b = (packed >> 2) & 0x03
+        c = (packed >> 4) & 0x03
+        d = (packed >> 6) & 0x03
+        interleaved = np.empty(packed.size * 4, dtype=np.uint32)
+        interleaved[0::4] = a
+        interleaved[1::4] = b
+        interleaved[2::4] = c
+        interleaved[3::4] = d
+        return interleaved[:count]
+
+    if bits == 1:
+        signs_raw = unpack_signs_raw(packed, count)
+        return signs_raw.astype(np.uint32)
+
+    return unpack_indices(packed, bits, count)
+
+
+def pack_signs_raw(bits_array: np.ndarray) -> np.ndarray:
+    """Pack a {0,1} bit array into bytes.
+
+    Parameters
+    ----------
+    bits_array : ndarray
+        Array of 0/1 values.
+
+    Returns
+    -------
+    packed : ndarray of uint8
+    """
+    bits_array = np.asarray(bits_array, dtype=np.uint8).ravel()
+    n = bits_array.size
+    pad_len = (8 - n % 8) % 8
+    if pad_len:
+        bits_array = np.append(
+            bits_array, np.zeros(pad_len, dtype=np.uint8)
+        )
+    reshaped = bits_array.reshape(-1, 8)
+    multipliers = (1 << np.arange(8, dtype=np.uint8))[np.newaxis, :]
+    return (reshaped * multipliers).sum(axis=1).astype(np.uint8)
+
+
+def unpack_signs_raw(packed: np.ndarray, count: int) -> np.ndarray:
+    """Unpack bytes into a {0,1} bit array.
+
+    Parameters
+    ----------
+    packed : ndarray of uint8
+    count : int
+
+    Returns
+    -------
+    bits_array : ndarray of uint8
+        Array of 0/1 values.
+    """
+    packed = np.asarray(packed, dtype=np.uint8).ravel()
+    bits = np.unpackbits(packed, bitorder="little")
+    return bits[:count]
+
+
+def pack_signs(signs: np.ndarray) -> np.ndarray:
+    """Pack a {-1, +1} sign array into bytes.
+
+    +1 is stored as bit 1, -1 as bit 0.
+
+    Parameters
+    ----------
+    signs : ndarray
+        Array of -1 / +1 values.
+
+    Returns
+    -------
+    packed : ndarray of uint8
+    """
+    bits = ((np.asarray(signs).ravel() + 1) // 2).astype(np.uint8)
+    return pack_signs_raw(bits)
+
+
+def unpack_signs(packed: np.ndarray, count: int) -> np.ndarray:
+    """Unpack bytes into a {-1, +1} sign array.
+
+    Parameters
+    ----------
+    packed : ndarray of uint8
+    count : int
+        Number of sign values to recover.
+
+    Returns
+    -------
+    signs : ndarray of float64
+        Array of -1.0 / +1.0 values.
+    """
+    bits = unpack_signs_raw(packed, count)
+    return bits.astype(np.float64) * 2.0 - 1.0

pyturboquant_cpu/qjl.py

@@ -0,0 +1,86 @@
+"""
+Quantized Johnson–Lindenstrauss (QJL) 1-bit inner-product quantiser.
+
+QJL provides an *unbiased* estimator for ⟨y, x⟩ using only 1 bit per
+coordinate.  It is used as the second stage of TurboQuant_Prod to
+correct the inner-product bias introduced by the MSE-optimal quantiser.
+
+Reference
+---------
+Zandieh et al., Definition 1 and Lemma 4 of arXiv:2504.19874.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+
+def generate_qjl_matrix(
+    dim: int, seed: int | None = None
+) -> np.ndarray:
+    """Generate a d×d random Gaussian matrix for QJL.
+
+    Each entry is drawn i.i.d. from N(0, 1).
+
+    Parameters
+    ----------
+    dim : int
+        Matrix dimension *d*.
+    seed : int or None, optional
+        Random seed for reproducibility.
+
+    Returns
+    -------
+    S : ndarray, shape (d, d)
+    """
+    rng = np.random.default_rng(seed)
+    return rng.standard_normal((dim, dim))
+
+
+def qjl_quantize(
+    vectors: np.ndarray, S: np.ndarray
+) -> np.ndarray:
+    """Apply QJL quantisation: sign(S · x).
+
+    Parameters
+    ----------
+    vectors : ndarray, shape (..., d)
+        Input vectors (should already be normalised / residuals).
+    S : ndarray, shape (d, d)
+        Random Gaussian matrix.
+
+    Returns
+    -------
+    signs : ndarray, shape (..., d)
+        Array of -1.0 / +1.0 values.
+    """
+    # vectors @ S.T gives (..., d)
+    projected = vectors @ S.T
+    signs = np.sign(projected)
+    # Replace exact zeros (extremely unlikely) with +1
+    signs[signs == 0] = 1.0
+    return signs
+
+
+def qjl_dequantize(
+    signs: np.ndarray, S: np.ndarray
+) -> np.ndarray:
+    """QJL inverse map: Q_qjl^{-1}(z) = √(π/2) · S^T · z / d.
+
+    Parameters
+    ----------
+    signs : ndarray, shape (..., d)
+        Array of -1 / +1 values.
+    S : ndarray, shape (d, d)
+        The same random matrix used during quantisation.
+
+    Returns
+    -------
+    reconstructed : ndarray, shape (..., d)
+    """
+    d = S.shape[0]
+    scale = math.sqrt(math.pi / 2.0) / d
+    # signs @ S gives (..., d)
+    return scale * (signs @ S)