pyosq

pyosq.py

"""

python port of Apache Lucene's OptimizedScalarQuantizer.

source: lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java

overview
------------------
OSQ is a *per-vector* scalar quantizer based on three main ideas:

1. **centroid-centering**: Each vector is translated by the segment centroid
   before quantization.

2. **per-vector interval initialisation**: The quantization interval [a, b] is
   initialised from the per-vector mean and standard deviation, scaled by the
   analytically-optimal MSE-minimising interval for each bit depth (derived
   for a unit-normal distribution).

3. **anisotropic loss + coordinate descent**: The interval is then refined by
   minimising a weighted loss holding the reconstruction error.  it penalises
   error parallel to the document vector (which hurts nearest-neighbour
   recall most), controlled by lambda ∈ [0, 1].

supported similarity metrics: DOT_PRODUCT, COSINE, EUCLIDEAN, MAX_INNER_PRODUCT.
"""

from __future__ import annotations

import enum
import math
from dataclasses import dataclass
from typing import Optional

import numpy as np

class Similarity(enum.Enum):
    DOT_PRODUCT       = "dot_product"
    COSINE            = "cosine"
    EUCLIDEAN         = "euclidean"
    MAX_INNER_PRODUCT = "max_inner_product"


# MINIMUM_MSE_GRID[bits-1] = [a, b] optimal for unit-normal distribution.
MINIMUM_MSE_GRID: list[tuple[float, float]] = [
    (-0.798,  0.798),   # 1 bit
    (-1.493,  1.493),   # 2 bits
    (-2.051,  2.051),   # 3 bits
    (-2.514,  2.514),   # 4 bits
    (-2.916,  2.916),   # 5 bits
    (-3.278,  3.278),   # 6 bits
    (-3.611,  3.611),   # 7 bits
    (-3.922,  3.922),   # 8 bits
]

DEFAULT_LAMBDA:float=0.1 # weight of isotropic vs anisotropic loss
DEFAULT_ITERS:int=5 # coordinate-descent iterations


@dataclass(frozen=True)
class QuantizationResult:
    """
    lower_interval, upper_interval
        Per-vector quantization bounds [a, b] (after centering).
    additional_correction
        For EUCLIDEAN  : squared L2 norm of the *centered* vector.
        For all others : dot product of the *original* (un-centered) vector
                         with the segment centroid.
        Used at query time to reconstruct the true dot product from the
        integer dot product of the quantized vectors.
    quantized_component_sum
        Sum of the integer-quantized components (1ᵀ x̂).
        Stored alongside the quantized vector; used for the corrective offset
        in dot-product reconstruction without an extra accumulation pass.
    """
    lower_interval:          float
    upper_interval:          float
    additional_correction:   float
    quantized_component_sum: int


class OptimizedScalarQuantizer:
    """
    parameters
    ----------
    similarity : Similarity
        the distance metric the index will use.
    lambda_ : float
        anisotropy weight.  0 = pure parallel-error minimisation;
        1 = pure isotropic MSE minimisation.  Lucene default = 0.1.
    iters : int
        maximum coordinate-descent iterations.  Lucene default = 5.
    """

    def __init__(
        self,
        similarity: Similarity = Similarity.DOT_PRODUCT,
        lambda_: float = DEFAULT_LAMBDA,
        iters:   int   = DEFAULT_ITERS,
    ) -> None:
        self.similarity = similarity
        self.lambda_    = lambda_
        self.iters      = iters


    def scalar_quantize(
        self,
        vector:      np.ndarray,          # float32, shape (d,)
        bits:        int,                  # 1 .. 8
        centroid:    np.ndarray,           # float32, shape (d,)
        destination: Optional[np.ndarray] = None,  # uint8/int, shape (d,)
    ) -> QuantizationResult:
        """
        quantize *vector* to *bits* bits, centering on *centroid*.

        returns a QuantizationResult and writes quantized values into
        *destination* (allocated if not provided).
        """
        assert 1 <= bits <= 8, f"bits must be 1..8, got {bits}"
        d = len(vector)
        if destination is None:
            destination = np.empty(d, dtype=np.int32)

        if self.similarity != Similarity.EUCLIDEAN:
            centroid_dot = float(np.dot(vector, centroid))
        else:
            centroid_dot = 0.0

        # center the vector on the segment centroid
        centered = vector.astype(np.float64) - centroid.astype(np.float64)

        # per-vector statistics (online Welford for mean+var)
        vec_mean, vec_var = _welford_mean_var(centered)
        vec_std = math.sqrt(max(vec_var, 0.0))
        norm2   = float(np.dot(centered, centered))
        vmin    = float(centered.min())
        vmax    = float(centered.max())

        points = 1 << bits
        # initialise interval from MSE grid scaled to this vector's distribution
        grid_a, grid_b = MINIMUM_MSE_GRID[bits - 1]
        a = _clamp(grid_a * vec_std + vec_mean, vmin, vmax)
        b = _clamp(grid_b * vec_std + vec_mean, vmin, vmax)
        interval = [a, b]

        # anisotropic coordinate-descent refinement
        self._optimize_intervals(interval, centered, norm2, points)

        # quantize
        a, b     = interval
        n_steps  = float((1 << bits) - 1)
        step     = (b - a) / n_steps if (b - a) > 0 else 1.0
        quantized, qsum = _quantize_vector(centered, a, b, step, n_steps)
        destination[:d] = quantized

        additional = norm2 if self.similarity == Similarity.EUCLIDEAN else centroid_dot

        return QuantizationResult(
            lower_interval          = float(a),
            upper_interval          = float(b),
            additional_correction   = float(additional),
            quantized_component_sum = int(qsum),
        )

    def multi_scalar_quantize(
        self,
        vector:    np.ndarray,        # float32, shape (d,)
        bits_list: list[int],          # e.g. [1, 2, 4]
        centroid:  np.ndarray,         # float32, shape (d,)
    ) -> tuple[list[np.ndarray], list[QuantizationResult]]:
        """
        quantize *vector* at multiple bit levels in a single pass.

        returns (destinations, results) where each destination is an int32
        ndarray of shape (d,).
        """
        assert all(1 <= b <= 8 for b in bits_list)

        if self.similarity != Similarity.EUCLIDEAN:
            centroid_dot = float(np.dot(vector, centroid))
        else:
            centroid_dot = 0.0

        centered = vector.astype(np.float64) - centroid.astype(np.float64)

        vec_mean, vec_var = _welford_mean_var(centered)
        vec_std = math.sqrt(max(vec_var, 0.0))
        norm2   = float(np.dot(centered, centered))
        vmin    = float(centered.min())
        vmax    = float(centered.max())

        destinations: list[np.ndarray]      = []
        results:      list[QuantizationResult] = []

        for bits in bits_list:
            points  = 1 << bits
            grid_a, grid_b = MINIMUM_MSE_GRID[bits - 1]
            a = _clamp(grid_a * vec_std + vec_mean, vmin, vmax)
            b = _clamp(grid_b * vec_std + vec_mean, vmin, vmax)
            interval = [a, b]

            self._optimize_intervals(interval, centered, norm2, points)

            a, b    = interval
            n_steps = float((1 << bits) - 1)
            step    = (b - a) / n_steps if (b - a) > 0 else 1.0
            quantized, qsum = _quantize_vector(centered, a, b, step, n_steps)

            additional = norm2 if self.similarity == Similarity.EUCLIDEAN else centroid_dot

            destinations.append(quantized)
            results.append(QuantizationResult(
                lower_interval          = float(a),
                upper_interval          = float(b),
                additional_correction   = float(additional),
                quantized_component_sum = int(qsum),
            ))

        return destinations, results


    def _loss(
        self,
        centered: np.ndarray,
        interval: list[float],
        points:   int,
        norm2:    float,
    ) -> float:
        """
        anisotropic quantization loss.

        L = (1-lamnda) * (x_T * eps)^2 / norm2(x) + lamnda * norm2(eps)

        where eps is the per-component reconstruction error after quantize-then-dequantize.
        """
        a, b    = interval
        step    = (b - a) / (points - 1.0)
        step_inv = 1.0 / step if step != 0 else 0.0

        # quantize-dequantize: x_q = a + step * round((clamp(x, a, b) - a) / step)
        clamped = np.clip(centered, a, b)
        x_q     = a + step * np.round((clamped - a) * step_inv)
        err     = centered - x_q          # ε

        # dot product of raw centered vector with error
        xe = float(np.dot(centered, err))
        e  = float(np.dot(err, err))

        anisotropic = (1.0 - self.lambda_) * xe * xe / norm2 if norm2 > 0 else 0.0
        isotropic   = self.lambda_ * e
        return anisotropic + isotropic

    def _optimize_intervals(
        self,
        init_interval: list[float],  # mutated in place
        centered:      np.ndarray,
        norm2:         float,
        points:        int,
    ) -> None:
        """
        refine [a, b] via coordinate descent minimising the anisotropic loss.
        """
        initial_loss = self._loss(centered, init_interval, points, norm2)
        scale = (1.0 - self.lambda_) / norm2 if norm2 > 0 else 0.0
        if not math.isfinite(scale):
            return

        for _ in range(self.iters):
            a, b    = init_interval
            step_inv = (points - 1.0) / (b - a) if (b - a) > 0 else 0.0

            clamped = np.clip(centered, a, b)
            k       = np.round((clamped - a) * step_inv)
            s       = k / (points - 1.0)         # shape (d,)
            one_m_s = 1.0 - s

            daa = float(np.dot(one_m_s, one_m_s))
            dab = float(np.dot(one_m_s, s))
            dbb = float(np.dot(s, s))
            dax = float(np.dot(centered, one_m_s))
            dbx = float(np.dot(centered, s))

            m0 = scale * dax * dax + self.lambda_ * daa
            m1 = scale * dax * dbx + self.lambda_ * dab
            m2 = scale * dbx * dbx + self.lambda_ * dbb

            det = m0 * m2 - m1 * m1
            if det == 0.0:
                return

            a_opt = (m2 * dax - m1 * dbx) / det
            b_opt = (m0 * dbx - m1 * dax) / det

            if abs(init_interval[0] - a_opt) < 1e-8 and abs(init_interval[1] - b_opt) < 1e-8:
                return

            new_loss = self._loss(centered, [a_opt, b_opt], points, norm2)
            if new_loss > initial_loss:
                return

            init_interval[0] = a_opt
            init_interval[1] = b_opt
            initial_loss      = new_loss


def dequantize(
    quantized:      np.ndarray,   # int/uint8, shape (d,)
    bits:           int,
    lower_interval: float,
    upper_interval: float,
    centroid:       np.ndarray,   # float32/64, shape (d,)
) -> np.ndarray:
    """
    reconstruct the approximate float vector from its quantized representation.
    """
    n_steps = (1 << bits) - 1
    step    = (upper_interval - lower_interval) / n_steps
    return (quantized.astype(np.float64) * step + lower_interval + centroid).astype(np.float32)


def pack_as_binary(quantized: np.ndarray) -> np.ndarray:
    """
    pack a 1-bit quantized vector (values 0 or 1, shape (d,)) into ceil(d/8)
    bytes, MSB first.
    """
    assert np.all((quantized == 0) | (quantized == 1)), "Values must be 0 or 1"
    d    = len(quantized)
    n    = (d + 7) // 8
    out  = np.zeros(n, dtype=np.uint8)
    for i, val in enumerate(quantized):
        byte_idx = i // 8
        bit_pos  = 7 - (i % 8)   # MSB first
        out[byte_idx] |= (int(val) & 1) << bit_pos
    return out


def unpack_binary(packed: np.ndarray, length: int) -> np.ndarray:
    """
    inverse of pack_as_binary.
    """
    out = np.empty(length, dtype=np.int32)
    idx = 0
    for byte_val in packed:
        for bit_pos in range(7, -1, -1):
            if idx >= length:
                break
            out[idx] = (int(byte_val) >> bit_pos) & 1
            idx += 1
    return out


def transpose_half_byte(q: np.ndarray) -> np.ndarray:
    """
    bit-plane transpose of a 4-bit (nibble) quantized query vector.

    takes a vector of values in [0, 15] and reorganises its bits so that
    bit-plane 0 of all dimensions is packed first, then bit-plane 1, etc.

    based on RaBitQ bit decomposition (Gao & Long, 2024).
    """
    assert np.all((q >= 0) & (q <= 15)), "Values must be in [0, 15]"
    d        = len(q)
    stripe   = (d + 7) // 8
    out      = np.zeros(4 * stripe, dtype=np.uint8)

    i = 0
    while i < d:
        lower_byte        = 0
        lower_middle_byte = 0
        upper_middle_byte = 0
        upper_byte        = 0
        for j in range(7, -1, -1):
            if i >= d:
                break
            v                  = int(q[i])
            lower_byte        |= (v & 1)         << j
            lower_middle_byte |= ((v >> 1) & 1)  << j
            upper_middle_byte |= ((v >> 2) & 1)  << j
            upper_byte        |= ((v >> 3) & 1)  << j
            i += 1
        index = ((i + 7) // 8) - 1
        out[index]               = lower_byte        & 0xFF
        out[index +     stripe]  = lower_middle_byte & 0xFF
        out[index + 2 * stripe]  = upper_middle_byte & 0xFF
        out[index + 3 * stripe]  = upper_byte        & 0xFF

    return out


def transpose_dibit(q: np.ndarray) -> np.ndarray:
    """
    bit-plane transpose for 2-bit (dibit) quantized vectors (values 0-3).
    layout: [bit-0 stripe | bit-1 stripe], each stripe is ceil(d/8) bytes.
    """
    assert np.all((q >= 0) & (q <= 3)), "Values must be in [0, 3]"
    d      = len(q)
    stripe = (d + 7) // 8
    out    = np.zeros(2 * stripe, dtype=np.uint8)

    i = 0
    index = 0
    while i < d:
        lo = 0
        hi = 0
        for j in range(7, -1, -1):
            if i >= d:
                break
            v   = int(q[i])
            lo |= (v & 1)        << j
            hi |= ((v >> 1) & 1) << j
            i += 1
        out[index]          = lo & 0xFF
        out[index + stripe] = hi & 0xFF
        index += 1

    return out


def untranspose_dibit(packed: np.ndarray, length: int) -> np.ndarray:
    """
    inverse of transpose_dibit.
    """
    stripe = len(packed) // 2
    out    = np.empty(length, dtype=np.int32)
    i      = 0
    index  = 0
    while i < length:
        lo = int(packed[index])
        hi = int(packed[index + stripe])
        for j in range(7, -1, -1):
            if i >= length:
                break
            out[i] = ((lo >> j) & 1) | (((hi >> j) & 1) << 1)
            i += 1
        index += 1
    return out


def discretize(value: int, bucket: int) -> int:
    """
    round *value* up to the nearest multiple of *bucket*.
    """
    return ((value + bucket - 1) // bucket) * bucket


def reconstruct_dot_product(
    query_quantized:     np.ndarray,   # int, shape (d,) — quantized query
    doc_quantized:       np.ndarray,   # int, shape (d,) — quantized doc
    query_result:        QuantizationResult,
    doc_result:          QuantizationResult,
    centroid:            np.ndarray,   # float, shape (d,)
    similarity:          Similarity    = Similarity.DOT_PRODUCT,
) -> float:
    """
    reconstruct the approximate original dot product from two OSQ-quantized
    vectors and their QuantizationResult metadata.

    the dot product between the original float vectors y and x is:

        y_T * x = (y−m)_T * (x−m)  +  m_T * y  +  m_T * x  −  m_T * m

    the centered dot product is approximated by the integer
    dot product scaled by the step sizes, plus the sum-of-components
    correction terms stored in QuantizationResult.

    for EUCLIDEAN the 'additional_correction' is norm2(x−m) (stored per doc)
    and the Euclidean distance is recovered analogously.

    """
    qa, qb = query_result.lower_interval, query_result.upper_interval
    da, db = doc_result.lower_interval,   doc_result.upper_interval

    q_nsteps = float((1 << _infer_bits(query_quantized)) - 1)
    d_nsteps = float((1 << _infer_bits(doc_quantized))   - 1)

    q_step = (qb - qa) / q_nsteps
    d_step = (db - da) / d_nsteps

    dim           = len(query_quantized)
    int_dot       = int(np.dot(query_quantized.astype(np.int64),
                               doc_quantized.astype(np.int64)))
    centered_approx = (
        qa * da * dim
        + qa * d_step * doc_result.quantized_component_sum
        + da * q_step * query_result.quantized_component_sum
        + q_step * d_step * int_dot
    )

    if similarity == Similarity.EUCLIDEAN:
        return centered_approx
    else:
        m_dot_m = float(np.dot(centroid, centroid))
        return centered_approx + query_result.additional_correction + doc_result.additional_correction - m_dot_m


def _clamp(x: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, x))


def _welford_mean_var(v: np.ndarray) -> tuple[float, float]:
    """
    online Welford algorithm for mean and variance.
    """
    mean = 0.0
    var  = 0.0
    for i, xi in enumerate(v):
        delta = xi - mean
        mean += delta / (i + 1)
        var  += delta * (xi - mean)
    var /= len(v)
    return mean, var


def _quantize_vector(
    centered: np.ndarray,
    a:        float,
    b:        float,
    step:     float,
    n_steps:  float,
) -> tuple[np.ndarray, int]:
    clamped  = np.clip(centered, a, b)
    q        = np.round((clamped - a) / step).astype(np.int32)
    q        = np.clip(q, 0, int(n_steps))   # guard against fp rounding
    return q, int(q.sum())


def _infer_bits(quantized: np.ndarray) -> int:
    mx = int(quantized.max())
    for bits in range(1, 9):
        if mx < (1 << bits):
            return bits
    return 8


def compute_centroid(vectors: np.ndarray) -> np.ndarray:
    """
    compute the segment centroid as the component-wise mean.
    """
    return vectors.mean(axis=0).astype(np.float32)


def build_osq_index(
    vectors:    np.ndarray,          # float32, shape (N, d)
    bits:       int = 1,
    similarity: Similarity = Similarity.DOT_PRODUCT,
    lambda_:    float = DEFAULT_LAMBDA,
    iters:      int   = DEFAULT_ITERS,
) -> tuple[np.ndarray, list[QuantizationResult], np.ndarray]:
    """
    quantize an entire corpus of vectors into an OSQ index.
    """
    centroid  = compute_centroid(vectors)
    quantizer = OptimizedScalarQuantizer(similarity, lambda_, iters)
    N, d      = vectors.shape
    quantized_corpus = np.empty((N, d), dtype=np.int32)
    results: list[QuantizationResult] = []

    for i, vec in enumerate(vectors):
        result = quantizer.scalar_quantize(
            vec.astype(np.float64), bits, centroid.astype(np.float64),
            destination=quantized_corpus[i],
        )
        results.append(result)

    return quantized_corpus, results, centroid


def osq_search(
    query:           np.ndarray,           # float32, shape (d,)
    quantized_corpus:np.ndarray,           # int32, shape (N, d)
    corpus_results:  list[QuantizationResult],
    centroid:        np.ndarray,           # float32, shape (d,)
    bits:            int,
    similarity:      Similarity = Similarity.DOT_PRODUCT,
    lambda_:         float = DEFAULT_LAMBDA,
    iters:           int   = DEFAULT_ITERS,
    top_k:           int   = 10,
) -> list[tuple[int, float]]:
    """
    brute-force search over an OSQ index.
    """
    quantizer = OptimizedScalarQuantizer(similarity, lambda_, iters)
    d = len(query)
    q_dest = np.empty(d, dtype=np.int32)
    q_result = quantizer.scalar_quantize(
        query.astype(np.float64), bits, centroid.astype(np.float64), destination=q_dest,
    )

    qa, qb      = q_result.lower_interval, q_result.upper_interval
    q_nsteps    = float((1 << bits) - 1)
    q_step      = (qb - qa) / q_nsteps

    m_dot_m = float(np.dot(centroid, centroid))
    scores  = []

    for i, (doc_q, doc_r) in enumerate(zip(quantized_corpus, corpus_results)):
        da, db      = doc_r.lower_interval, doc_r.upper_interval
        d_nsteps    = float((1 << bits) - 1)
        d_step      = (db - da) / d_nsteps

        dim         = d
        int_dot     = int(np.dot(q_dest.astype(np.int64), doc_q.astype(np.int64)))

        centered_approx = (
            qa * da * dim
            + qa * d_step * doc_r.quantized_component_sum
            + da * q_step * q_result.quantized_component_sum
            + q_step * d_step * int_dot
        )

        if similarity == Similarity.EUCLIDEAN:
            score = centered_approx
        else:
            score = (centered_approx
                     + q_result.additional_correction
                     + doc_r.additional_correction
                     - m_dot_m)

        scores.append((i, score))

    scores.sort(key=lambda x: -x[1])
    return scores[:top_k]


# "demo"
if __name__ == "__main__":
    rng = np.random.default_rng(42)
    d   = 128
    N   = 1000

    # simulate l2-normalised embeddings 
    raw = rng.standard_normal((N, d)).astype(np.float32)
    raw /= np.linalg.norm(raw, axis=1, keepdims=True)
    query_raw = rng.standard_normal(d).astype(np.float32)
    query_raw /= np.linalg.norm(query_raw)

    centroid  = compute_centroid(raw)
    quantizer = OptimizedScalarQuantizer(Similarity.DOT_PRODUCT)

    # show scalar_quantize results
    vec   = raw[0].copy()
    dest  = np.empty(d, dtype=np.int32)
    result = quantizer.scalar_quantize(vec.astype(np.float64), bits=1,
                                       centroid=centroid.astype(np.float64),
                                       destination=dest)
    print(f"[1-bit] interval=[{result.lower_interval:.4f}, {result.upper_interval:.4f}]  "
          f"qsum={result.quantized_component_sum}  "
          f"correction={result.additional_correction:.4f}")

    recon = dequantize(dest, 1, result.lower_interval, result.upper_interval, centroid)
    cos_sim = float(np.dot(recon, vec) / (np.linalg.norm(recon) * np.linalg.norm(vec)))
    print(f"cosine(original, dequantized) = {cos_sim:.4f}")

    for bits in [1, 2, 4, 7]:
        dst2 = np.empty(d, dtype=np.int32)
        r2   = quantizer.scalar_quantize(vec.astype(np.float64), bits=bits,
                                         centroid=centroid.astype(np.float64), destination=dst2)
        recon2 = dequantize(dst2, bits, r2.lower_interval, r2.upper_interval, centroid)
        cos = float(np.dot(recon2, vec) / (np.linalg.norm(recon2) * np.linalg.norm(vec)))
        mem_ratio = 32 / bits
        print(f"[{bits}-bit] cosine(orig, recon)={cos:.4f}   compression={mem_ratio:.0f}x")

    # show multi_scalar_quantize 
    dests, results = quantizer.multi_scalar_quantize(vec.astype(np.float64),
                                                     [1, 2, 4],
                                                     centroid.astype(np.float64))
    for bits, r in zip([1, 2, 4], results):
        print(f"bits={bits}  a={r.lower_interval:.4f}  b={r.upper_interval:.4f}")

    # e2e bf recall test ---
    true_dots = (raw @ query_raw)
    true_top10 = set(np.argsort(-true_dots)[:10].tolist())

    for bits in [1, 2, 4]:
        corpus_q, corpus_r, seg_centroid = build_osq_index(raw, bits=bits,
                                                            similarity=Similarity.DOT_PRODUCT)
        approx = osq_search(query_raw, corpus_q, corpus_r, seg_centroid, bits=bits,
                            similarity=Similarity.DOT_PRODUCT, top_k=10)
        approx_top10 = {idx for idx, _ in approx}
        recall = len(true_top10 & approx_top10) / 10
        print(f"bits={bits}  recall@10={recall:.1%}  "
              f"(memory: {bits/32:.1%} of float32)")