arborist/bench/scripts/phi_alignment_probe.py

"""φ_linear Hessian-alignment probe (ticket #000034 Phase 1a).

Measures whether the v7 reference frozen-seed linear projection
``W`` (dim_d × 256, mapping SHA-256-digest space to the dim_h
parameter-vector space) has structural alignment with the
low-eigenvalue subspace of the typical training-loss Hessian
``H(Θ)``.

Why this matters — from #000018 §5.1:

    The v7 reference choice (frozen-seed Gaussian linear projection)
    is **not obviously independent of parameter geometry**. If W's
    column space aligns with low-eigenvalue directions of the loss
    Hessian, the L2 anchor pull provides cheap parameter movement
    in directions correlated with SHA-256 buckets. Cheap = "reachable
    within a per-step gradient budget" = exploitable by a T2
    gradient-shaping adversary.

This module ships only the probe + verdict logic. Phase 1a uses
synthetic Hessians (engineered eigenstructure + known W placement)
to validate the probe end-to-end. Phase 1b runs the same probe
against a real v7 checkpoint when one becomes available — that's
when the verdict actually informs the M1 (φ_PRG) vs M2 (anchor
nonce) priority decision.

Construction — the alignment score (#000034 §2.1):

    A(W, H) = Σ_i ⟨W·e_i, v_i⟩² / (λ_i + ε)   /   ‖W‖_F²

where the sum runs over the columns ``W·e_i`` of W and the
eigenvectors ``v_i`` of H, ``λ_i`` are the matching eigenvalues,
and ε is a small regularizer to avoid division by 0 on flat
directions. Under random-oracle modeling A is uniform across the
spectrum; concentration in the low-λ tail is the alignment
signature.

Verdict thresholds (#000034 §3.3):

- ``A_low / A_uniform > 1.5`` → ``STRUCTURAL_ALIGNMENT`` (M1
  priority).
- ``A_low / A_uniform`` near ``1.0`` → ``NO_ALIGNMENT`` (M2
  alone sufficient).
- ``A_low / A_uniform < 0.7`` → ``ANTI_ALIGNED`` (φ_linear is
  actually safe vs the loss landscape).

This module imports numpy + scipy at function call time (not
module load time) so the rest of arborist can import it without
the [hessian] extras installed; tests skip via
``pytest.importorskip`` when absent.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Callable

# Module-level marker so callers can detect probe version drift.
PROBE_VERSION = "phi-alignment-v1-lanczos"

# Verdict thresholds per #000034 §3.3.
STRUCTURAL_ALIGNMENT_RATIO_FLOOR = 1.5
ANTI_ALIGNED_RATIO_CEILING = 0.7

# Default Lanczos hyperparameters per #000034 §2.2 recommendation.
DEFAULT_K_TOP = 100
DEFAULT_K_BOT = 100
DEFAULT_EPSILON = 1e-6


@dataclass(frozen=True)
class AlignmentReport:
    """Output of :func:`measure_alignment`. Frozen so a caller's
    JSON serialization is stable; round-trips via ``asdict``.
    """
    verdict: str  # STRUCTURAL_ALIGNMENT | NO_ALIGNMENT | ANTI_ALIGNED
    ratio: float  # A(W, H_low) / A(W, H_uniform)
    a_top: float  # alignment score concentrated on top-k (high λ)
    a_bot: float  # alignment score concentrated on bot-k (low λ)
    a_uniform: float  # baseline assuming uniform alignment
    eigenvalues_top: tuple[float, ...]
    eigenvalues_bot: tuple[float, ...]
    w_frobenius_norm_squared: float
    k_top: int
    k_bot: int
    epsilon: float
    probe_version: str = field(default_factory=lambda: PROBE_VERSION)

    def as_dict(self) -> dict:
        d = asdict(self)
        # tuples → lists for JSON friendliness
        d["eigenvalues_top"] = list(self.eigenvalues_top)
        d["eigenvalues_bot"] = list(self.eigenvalues_bot)
        return d


def measure_alignment(
    W,
    hessian_eval: Callable,
    *,
    dim_d: int | None = None,
    k_top: int = DEFAULT_K_TOP,
    k_bot: int = DEFAULT_K_BOT,
    epsilon: float = DEFAULT_EPSILON,
) -> AlignmentReport:
    """Run the alignment probe on ``W`` against the Hessian operator
    ``hessian_eval``.

    Parameters
    ----------
    W
        Numpy array of shape ``(dim_d, 256)`` per the v7 reference
        spec — the frozen-seed linear projection from SHA-256-digest
        space to parameter-vector space.
    hessian_eval
        Callable ``x → H @ x`` representing the symmetric Hessian
        ``H(Θ)`` as a linear operator. The probe never materializes
        the full Hessian; Lanczos only needs matrix-vector products.
    dim_d
        Parameter-space dimension. Defaults to ``W.shape[0]``.
    k_top, k_bot
        Number of top and bottom eigenpairs to extract. ``k_top +
        k_bot`` should be ≤ dim_d; per #000034 §2.2,
        100/100 is empirically tractable for ~10⁸-parameter
        checkpoints.
    epsilon
        Regularizer inside the alignment score's
        ``1 / (λ_i + ε)`` term to avoid division by 0 on flat
        directions of the loss landscape.

    Returns
    -------
    AlignmentReport
        Per #000034 §3.3 the verdict is
        ``STRUCTURAL_ALIGNMENT`` / ``NO_ALIGNMENT`` /
        ``ANTI_ALIGNED``.

    Raises
    ------
    ValueError
        If ``W`` is not 2-D, if its shape doesn't agree with
        ``dim_d``, or if Lanczos can't extract the requested number
        of eigenpairs.
    """
    import numpy as np
    from scipy.sparse.linalg import LinearOperator, eigsh

    # Validation — fail loud on shape errors rather than producing a
    # silent zero alignment.
    if W.ndim != 2:
        raise ValueError(
            f"W must be 2-D (dim_d × 256); got shape {W.shape!r}"
        )
    if dim_d is None:
        dim_d = W.shape[0]
    if W.shape[0] != dim_d:
        raise ValueError(
            f"W.shape[0] = {W.shape[0]} but dim_d = {dim_d}"
        )
    if k_top + k_bot > dim_d:
        raise ValueError(
            f"k_top ({k_top}) + k_bot ({k_bot}) > dim_d ({dim_d}); "
            "Lanczos cannot extract more eigenpairs than the matrix dim"
        )
    if epsilon <= 0:
        raise ValueError(f"epsilon must be positive; got {epsilon!r}")

    # Wrap the user closure in a scipy LinearOperator so eigsh can
    # call it via matvec. The operator is symmetric by construction
    # (Hessian) so eigsh's symmetric path applies.
    op = LinearOperator(
        shape=(dim_d, dim_d),
        matvec=hessian_eval,
        dtype=np.float64,
    )

    # Lanczos top-k via 'LA' (largest algebraic). Bottom-k via 'SA'
    # (smallest algebraic). Both directions are needed — the
    # alignment score weights inversely (1/(λ+ε)), so the bottom-k
    # dominates for any non-flat spectrum but we report top-k to
    # detect anti-alignment too.
    eigvals_top, eigvecs_top = eigsh(op, k=k_top, which="LA")
    eigvals_bot, eigvecs_bot = eigsh(op, k=k_bot, which="SA")

    # Frobenius normalizer.
    w_fro_sq = float(np.linalg.norm(W, "fro") ** 2)
    if w_fro_sq == 0.0:
        raise ValueError("W has zero Frobenius norm; degenerate input")

    # Alignment score on a given (eigvals, eigvecs) pair.
    # ⟨W·e_i, v_j⟩ = (W^T v_j)_i, so the inner products across all
    # 256 columns of W are the entries of (W^T @ v_j). Squared, then
    # weighted by 1/(λ_j + ε) and summed over j, then normalized.
    def _score(eigvals: "np.ndarray", eigvecs: "np.ndarray") -> float:
        # eigvecs.shape = (dim_d, k); each column is a v_j.
        proj = W.T @ eigvecs  # (256, k); col j is W^T v_j
        # Squared norms per column = Σ_i ⟨W·e_i, v_j⟩²
        col_sq = np.sum(proj ** 2, axis=0)  # (k,)
        weighted = col_sq / (eigvals + epsilon)
        return float(np.sum(weighted) / w_fro_sq)

    a_top = _score(eigvals_top, eigvecs_top)
    a_bot = _score(eigvals_bot, eigvecs_bot)

    # Isotropic baseline (#000034 §2.1, derived 2026-05-10).
    #
    # Under random-oracle modeling W's columns are isotropically
    # distributed Gaussians; for each unit-norm eigenvector v_j,
    # E[‖W^T v_j‖²] = 256 σ² and ‖W‖_F² = dim_d × 256 σ², so
    # E[‖W^T v_j‖² / ‖W‖_F²] = 1/dim_d. Thus:
    #
    #     E[A_k(W_uniform, H)] = (1/dim_d) Σ_{j in k-subset} 1/(λ_j+ε)
    #
    # This is the per-eigenvalue-subset baseline a uniformly-
    # distributed W would produce; the verdict ratio is the
    # observed alignment on the BOTTOM-k normalized by the
    # bottom-k's isotropic baseline. A previous formulation
    # (mean(a_top, a_bot)) was wrong because the 1/(λ+ε) weighting
    # makes a_bot mechanically large for any W shape — so the
    # ratio came out > 1 even for genuinely uniform W. The
    # analytical baseline corrects this.
    a_uniform_top = float(np.sum(1.0 / (eigvals_top + epsilon)) / dim_d)
    a_uniform_bot = float(np.sum(1.0 / (eigvals_bot + epsilon)) / dim_d)

    if a_uniform_bot == 0.0:
        # Degenerate — bottom-k baseline is exactly zero (would
        # require infinite eigenvalues, which Lanczos can't return).
        # Defensive: treat as NO_ALIGNMENT.
        ratio = 1.0
    else:
        ratio = a_bot / a_uniform_bot

    if ratio > STRUCTURAL_ALIGNMENT_RATIO_FLOOR:
        verdict = "STRUCTURAL_ALIGNMENT"
    elif ratio < ANTI_ALIGNED_RATIO_CEILING:
        verdict = "ANTI_ALIGNED"
    else:
        verdict = "NO_ALIGNMENT"

    return AlignmentReport(
        verdict=verdict,
        ratio=ratio,
        a_top=a_top,
        a_bot=a_bot,
        a_uniform=a_uniform_bot,  # bottom-k isotropic baseline
        eigenvalues_top=tuple(float(x) for x in eigvals_top),
        eigenvalues_bot=tuple(float(x) for x in eigvals_bot),
        w_frobenius_norm_squared=w_fro_sq,
        k_top=k_top,
        k_bot=k_bot,
        epsilon=epsilon,
    )


def hessian_from_dense(H):
    """Helper for unit tests + KAT generation: turn a dense
    symmetric matrix into a ``hessian_eval`` closure.

    Production callsites (real v7 checkpoints) should pass an
    HVP closure instead of materializing the Hessian; this helper
    is for the synthetic Phase 1a path only.
    """
    import numpy as np

    H = np.asarray(H, dtype=np.float64)
    if H.ndim != 2 or H.shape[0] != H.shape[1]:
        raise ValueError(f"H must be square 2-D; got shape {H.shape!r}")
    return lambda x: H @ x


__all__ = [
    "PROBE_VERSION",
    "STRUCTURAL_ALIGNMENT_RATIO_FLOOR",
    "ANTI_ALIGNED_RATIO_CEILING",
    "DEFAULT_K_TOP",
    "DEFAULT_K_BOT",
    "DEFAULT_EPSILON",
    "AlignmentReport",
    "measure_alignment",
    "hessian_from_dense",
]