Blame for analog.py

"""Geometry-preserving weight analogs for Q2_K bilinear factorization.

At 2 bits per weight (4 levels ``q in {0, 1, 2, 3}`` and 4-bit sub-block
scales ``Ls, Lm in {0..15}``), raw weight blocks suffer from two structural
bottlenecks:

1. **Intra-sub-block outlier collapse**: a single heavy-tailed weight stretches
   the sub-block span ``xmax - xmin``, forcing 13-15 of the 16 weights into
   one or two quantization bins.
2. **Inter-sub-block scale crushing**: a single high-variance sub-block inflates
   the shared superblock scale ``d``, collapsing ``Ls_j`` of the remaining
   15 sub-blocks toward small integers with severe rounding loss.

Rather than factoring raw weights alone, V2 constructs a bank of
**geometry-preserving weight analogs** ``x_tilde`` of the eligible weights:

* :func:`subblock_isomoment_analog` maps each 16-weight sub-block through an
  importance-weighted saturating transport blended with the optimal 4-level
  quantile spacing, then orthogonally projects back onto the exact weighted
  mean ``mu_j`` and weighted variance ``sigma_j^2`` of the original sub-block.
* :func:`vesica_gram_analog` operates in the D6 / Vesica fold basis
  ``(x_p + x_{p+128}, x_p - x_{p+128})``, regularizing outliers in the whitened
  eigen-basis and coloring back so the exact 2x2 weighted Gram matrix and means
  of the two half-blocks are preserved.
* :func:`subscale_lattice_analog` tempers cross-sub-block dispersion toward the
  4-bit ``{0..15}`` scale lattice while preserving every sub-block's unit
  direction, mean, and total block weighted energy.

The discrete Q2_K lattice coordinates ``(Ls, Lm, q)`` discovered on each analog
are then projected back onto the **true weights** ``x`` via coupled bilinear
least-squares, guaranteeing that the final emitted block is strictly optimal
under the true objective.
"""

from __future__ import annotations

import numpy as np

from .q2k import NS, QK_K

_Q2_IDEAL_QUANTILES = np.linspace(-1.35, 1.35, 16, dtype=np.float64)


def subblock_isomoment_analog(
    x: np.ndarray,
    w: np.ndarray | None = None,
    tau: float = 1.65,
    beta: float = 0.25,
    quantile_blend: float = 0.30,
    enforce_nonpositive_min: bool = False,
) -> np.ndarray:
    """Create a sub-block iso-moment analog preserving weighted mean & variance.

    Invariants per 16-weight sub-block ``j`` (when ``enforce_nonpositive_min``
    is ``False``):
      * ``sum_k w_{j,k} x_tilde_{j,k} == sum_k w_{j,k} x_{j,k}``
      * ``sum_k w_{j,k} (x_tilde_{j,k} - mu_j)^2 == sum_k w_{j,k} (x_{j,k} - mu_j)^2``
      * strict rank-order preservation with respect to ``x_{j,k}`` (when ``beta == 0``)
    """
    xb = np.asarray(x, dtype=np.float64).reshape(NS, 16)
    wb = (
        np.ones((NS, 16), dtype=np.float64)
        if w is None
        else np.asarray(w, dtype=np.float64).reshape(NS, 16)
    )
    wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30)
    mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum
    centered = xb - mu
    var = np.sum(wb * centered * centered, axis=1, keepdims=True) / wsum
    sigma = np.sqrt(np.maximum(var, 1.0e-30))
    active = (sigma[:, 0] >= 1.0e-15)

    z = centered / sigma
    w_mean = wsum / 16.0
    w_rel = np.power(np.maximum(wb / w_mean, 1.0e-12), beta) if beta != 0.0 else 1.0
    z_comp = tau * np.tanh((z * w_rel) / max(tau, 1.0e-6))

    if quantile_blend > 0.0:
        order = np.argsort(z, axis=1, kind="stable")
        z_quant = np.empty_like(z)
        rows = np.arange(NS)[:, None]
        z_quant[rows, order] = _Q2_IDEAL_QUANTILES[None, :]
        z_analog = (1.0 - quantile_blend) * z_comp + quantile_blend * z_quant
    else:
        z_analog = z_comp

    mu_a = np.sum(wb * z_analog, axis=1, keepdims=True) / wsum
    z_analog = z_analog - mu_a
    var_a = np.sum(wb * z_analog * z_analog, axis=1, keepdims=True) / wsum
    scale = np.where(var_a > 1.0e-30, np.sqrt(var / np.maximum(var_a, 1.0e-30)), 0.0)
    out = np.where(active[:, None], mu + z_analog * scale, xb)

    if enforce_nonpositive_min:
        orig_min = np.min(xb, axis=1)
        new_min = np.min(out, axis=1)
        shift = np.where((orig_min <= 0.0) & (new_min > 0.0), new_min, 0.0)
        out = out - shift[:, None]

    return out.reshape(QK_K).astype(np.float32)


def vesica_gram_analog(
    x: np.ndarray,
    w: np.ndarray | None = None,
    tau: float = 1.9,
) -> np.ndarray:
    """Create a D6 / Vesica fold analog preserving the 2x2 half-block Gram matrix.

    Decomposes ``x`` into symmetric ``v = 0.5*(x[:128] + x[128:])`` and
    antisymmetric ``u = 0.5*(x[:128] - x[128:])`` modes under pair weights
    ``wp = 0.5*(w[:128] + w[128:])``.  Whitens ``(v, u)``, applies smooth
    hyperbolic saturation in the principal axes, and colors back with the exact
    2x2 covariance square-root so ``Gram(v_tilde, u_tilde) == Gram(v, u)``.
    """
    x_f64 = np.asarray(x, dtype=np.float64).reshape(QK_K)
    w_f64 = (
        np.ones(QK_K, dtype=np.float64)
        if w is None
        else np.asarray(w, dtype=np.float64).reshape(QK_K)
    )
    half = QK_K // 2
    x0, x1 = x_f64[:half], x_f64[half:]
    wp = np.maximum(0.5 * (w_f64[:half] + w_f64[half:]), 1.0e-30)
    wsum = float(np.sum(wp))

    v = 0.5 * (x0 + x1)
    u = 0.5 * (x0 - x1)
    mv = float(np.sum(wp * v) / wsum)
    mu = float(np.sum(wp * u) / wsum)
    Y = np.vstack([v - mv, u - mu])  # (2, 128)

    G = (Y * wp[None, :]) @ Y.T / wsum  # 2x2 weighted covariance
    evals, evecs = np.linalg.eigh(G)
    if float(np.max(evals)) < 1.0e-28:
        return x_f64.astype(np.float32)

    pos = evals > 1.0e-28
    inv_sqrt = np.where(pos, 1.0 / np.sqrt(np.maximum(evals, 1.0e-30)), 0.0)
    W_white = (evecs * inv_sqrt[None, :]) @ evecs.T
    Z = W_white @ Y  # whitened (2, 128)

    Z_comp = tau * np.tanh(Z / max(tau, 1.0e-6))
    Z_mean = np.sum(Z_comp * wp[None, :], axis=1, keepdims=True) / wsum
    Z_comp = Z_comp - Z_mean

    G_z = (Z_comp * wp[None, :]) @ Z_comp.T / wsum
    ez, Vz = np.linalg.eigh(G_z)
    pos_z = ez > 1.0e-28
    if not np.all(pos_z):
        return x_f64.astype(np.float32)
    inv_sqrt_z = np.where(pos_z, 1.0 / np.sqrt(np.maximum(ez, 1.0e-30)), 0.0)
    Z_orth = (Vz * inv_sqrt_z[None, :]) @ Vz.T @ Z_comp

    sqrt_G = (evecs * np.sqrt(np.maximum(evals, 0.0))[None, :]) @ evecs.T
    Y_tilde = sqrt_G @ Z_orth
    v_new = mv + Y_tilde[0]
    u_new = mu + Y_tilde[1]
    return np.concatenate([v_new + u_new, v_new - u_new]).astype(np.float32)


def subscale_lattice_analog(
    x: np.ndarray,
    w: np.ndarray | None = None,
    gamma: float = 0.65,
) -> np.ndarray:
    """Create a cross-sub-block analog whose sub-block dispersions fit 4-bit scales.

    Preserves each sub-block's mean ``mu_j`` and unit direction
    ``(x_j - mu_j) / sigma_j``, while tempering extreme sub-block standard
    deviations ``sigma_j`` toward the superblock root-mean-square dispersion
    and rescaling to preserve total centered weighted energy.
    """
    xb = np.asarray(x, dtype=np.float64).reshape(NS, 16)
    wb = (
        np.ones((NS, 16), dtype=np.float64)
        if w is None
        else np.asarray(w, dtype=np.float64).reshape(NS, 16)
    )
    wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30)
    mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum
    centered = xb - mu
    energy_j = np.sum(wb * centered * centered, axis=1, keepdims=True)
    total_energy = float(np.sum(energy_j))
    if total_energy <= 1.0e-30:
        return xb.reshape(QK_K).astype(np.float32)

    sigma = np.sqrt(np.maximum(energy_j / wsum, 1.0e-30))
    sigma_ref = float(np.sqrt(total_energy / max(float(np.sum(wsum)), 1.0e-30)))
    target_sigma = np.power(np.maximum(sigma, 1.0e-30), gamma) * (
        sigma_ref ** (1.0 - gamma)
    )
    scaled = np.where(sigma > 1.0e-15, centered * (target_sigma / sigma), 0.0)
    new_energy = float(np.sum(wb * scaled * scaled))
    if new_energy > 1.0e-30:
        scaled = scaled * np.sqrt(total_energy / new_energy)
    out = mu + scaled
    return out.reshape(QK_K).astype(np.float32)


def build_block_analogs(
    x: np.ndarray,
    w: np.ndarray | None = None,
) -> list[np.ndarray]:
    """Return the portfolio of geometry-preserving weight analogs for a block."""
    a1 = subblock_isomoment_analog(
        x, w, tau=1.60, beta=0.25, quantile_blend=0.30, enforce_nonpositive_min=True
    )
    a2 = subblock_isomoment_analog(
        x, w, tau=2.15, beta=0.15, quantile_blend=0.15, enforce_nonpositive_min=True
    )
    a3 = vesica_gram_analog(x, w, tau=1.85)
    a4 = subscale_lattice_analog(a1, w, gamma=0.70)
    return [a1, a2, a3, a4]