arborist/tests/test_phi_alignment_probe.py

"""Tests for bench.scripts.phi_alignment_probe (ticket #000034 §3.2).

Coverage matches the ticket's acceptance criteria:

- Determinism: same (W, H, seed) → identical ratio + verdict.
- Verdict thresholds: engineered cases (aligned / uniform / anti)
  land in their expected bucket; boundary cases land cleanly.
- Lanczos convergence: top-k eigenvalues from scipy.eigsh match
  a known dense decomposition on a synthetic Hessian.
- Module export shape: AlignmentReport dataclass round-trips via
  ``asdict``; ``PROBE_VERSION`` + thresholds exported.
- Input validation: shape mismatches, bad k, bad epsilon all raise.
- KAT regression: pinned (seed, dim, k, class) → expected_verdict
  triples replay against the synthetic-checkpoints.jsonl fixture.
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

# Module imports skip if numpy/scipy aren't installed (the [hessian]
# extras gate per #000034 §3 + pyproject.toml).
np = pytest.importorskip("numpy")
pytest.importorskip("scipy.sparse.linalg")

from bench.scripts.phi_alignment_probe import (
    ANTI_ALIGNED_RATIO_CEILING,
    DEFAULT_K_TOP,
    DEFAULT_K_BOT,
    PROBE_VERSION,
    STRUCTURAL_ALIGNMENT_RATIO_FLOOR,
    AlignmentReport,
    hessian_from_dense,
    measure_alignment,
)


# ----------------------------------------------------------- helpers


def _diagonal_h(dim_d: int, scale: float = 1.0):
    """Diagonal Hessian with eigenvalues 1..dim_d * scale.
    Eigenvectors are the canonical basis."""
    return np.diag(np.arange(1, dim_d + 1).astype(float) * scale)


def _aligned_w(dim_d: int, n_cols: int, low_k: int, seed: int):
    """W with mass concentrated on the bottom-low_k canonical-basis
    rows — i.e. the low-eigenvalue subspace of _diagonal_h."""
    rng = np.random.default_rng(seed)
    W = np.zeros((dim_d, n_cols))
    W[:low_k, :] = rng.standard_normal((low_k, n_cols))
    return W


def _uniform_w(dim_d: int, n_cols: int, seed: int):
    rng = np.random.default_rng(seed)
    return rng.standard_normal((dim_d, n_cols))


def _anti_w(dim_d: int, n_cols: int, high_k: int, seed: int):
    """W with mass concentrated on the top-high_k canonical-basis
    rows — i.e. the high-eigenvalue subspace."""
    rng = np.random.default_rng(seed)
    W = np.zeros((dim_d, n_cols))
    W[-high_k:, :] = rng.standard_normal((high_k, n_cols))
    return W


# ----------------------------------------------------------- determinism


def test_measure_alignment_deterministic_same_inputs():
    """Same inputs produce the same verdict + ratio within Lanczos
    float tolerance.

    Strict bit-equality fails because scipy.eigsh's Lanczos
    iteration uses a randomized starting vector internally; output
    differs in the ~15th decimal across runs even with identical
    user inputs. The verdict + alignment-class signal is stable —
    that's what callers depend on.
    """
    dim_d = 64
    W = _uniform_w(dim_d, 256, seed=0)
    H = _diagonal_h(dim_d)
    r1 = measure_alignment(W, hessian_from_dense(H), k_top=8, k_bot=8)
    r2 = measure_alignment(W, hessian_from_dense(H), k_top=8, k_bot=8)
    assert r1.verdict == r2.verdict
    assert r1.ratio == pytest.approx(r2.ratio, rel=1e-9)
    assert sorted(r1.eigenvalues_top) == pytest.approx(
        sorted(r2.eigenvalues_top), rel=1e-6
    )
    assert sorted(r1.eigenvalues_bot) == pytest.approx(
        sorted(r2.eigenvalues_bot), rel=1e-6
    )


# ----------------------------------------------------------- verdicts


def test_aligned_w_returns_structural_alignment():
    """W concentrated on low-λ subspace must produce
    STRUCTURAL_ALIGNMENT (ratio > 1.5)."""
    W = _aligned_w(256, 256, low_k=32, seed=42)
    H = _diagonal_h(256)
    r = measure_alignment(W, hessian_from_dense(H), k_top=32, k_bot=32)
    assert r.verdict == "STRUCTURAL_ALIGNMENT"
    assert r.ratio > STRUCTURAL_ALIGNMENT_RATIO_FLOOR


def test_uniform_w_returns_no_alignment():
    """W with isotropically-distributed entries must produce
    NO_ALIGNMENT (ratio near 1.0)."""
    W = _uniform_w(256, 256, seed=42)
    H = _diagonal_h(256)
    r = measure_alignment(W, hessian_from_dense(H), k_top=32, k_bot=32)
    assert r.verdict == "NO_ALIGNMENT"
    assert ANTI_ALIGNED_RATIO_CEILING <= r.ratio <= STRUCTURAL_ALIGNMENT_RATIO_FLOOR


def test_anti_aligned_w_returns_anti_aligned():
    """W concentrated on high-λ subspace must produce
    ANTI_ALIGNED (ratio < 0.7)."""
    W = _anti_w(256, 256, high_k=32, seed=42)
    H = _diagonal_h(256)
    r = measure_alignment(W, hessian_from_dense(H), k_top=32, k_bot=32)
    assert r.verdict == "ANTI_ALIGNED"
    assert r.ratio < ANTI_ALIGNED_RATIO_CEILING


# ----------------------------------------------------------- Lanczos convergence


def test_lanczos_top_k_matches_dense_decomposition():
    """Synthetic diagonal Hessian: top-k Lanczos eigenvalues must
    match the analytical top-k eigenvalues exactly (within tol)."""
    dim_d = 128
    H = _diagonal_h(dim_d)
    W = _uniform_w(dim_d, 256, seed=0)
    k = 16
    r = measure_alignment(W, hessian_from_dense(H), k_top=k, k_bot=k)

    # Diagonal H has eigenvalues 1..dim_d, so top-k = dim_d-k+1..dim_d
    expected_top = list(range(dim_d - k + 1, dim_d + 1))
    expected_bot = list(range(1, k + 1))

    # Lanczos returns eigenvalues in algebraic order; sort to match
    # canonical ordering. Allow small numerical tolerance.
    observed_top = sorted(r.eigenvalues_top)
    observed_bot = sorted(r.eigenvalues_bot)
    assert all(
        abs(o - e) < 1e-6 for o, e in zip(observed_top, expected_top)
    ), f"top: expected {expected_top}, got {observed_top}"
    assert all(
        abs(o - e) < 1e-6 for o, e in zip(observed_bot, expected_bot)
    ), f"bot: expected {expected_bot}, got {observed_bot}"


# ----------------------------------------------------------- module shape


def test_alignment_report_round_trips_via_asdict():
    W = _uniform_w(64, 32, seed=0)
    H = _diagonal_h(64)
    r = measure_alignment(W, hessian_from_dense(H), k_top=4, k_bot=4)
    d = r.as_dict()
    assert isinstance(d, dict)
    assert d["verdict"] in {
        "STRUCTURAL_ALIGNMENT", "NO_ALIGNMENT", "ANTI_ALIGNED",
    }
    assert isinstance(d["eigenvalues_top"], list)  # not tuple
    assert isinstance(d["eigenvalues_bot"], list)
    assert d["probe_version"] == PROBE_VERSION
    # JSON-serializable
    json.dumps(d)


def test_module_exports_thresholds_and_version():
    assert PROBE_VERSION == "phi-alignment-v1-lanczos"
    assert STRUCTURAL_ALIGNMENT_RATIO_FLOOR == 1.5
    assert ANTI_ALIGNED_RATIO_CEILING == 0.7
    assert DEFAULT_K_TOP == 100
    assert DEFAULT_K_BOT == 100


# ----------------------------------------------------------- validation


def test_rejects_non_2d_w():
    H = _diagonal_h(16)
    with pytest.raises(ValueError, match="2-D"):
        measure_alignment(np.zeros(16), hessian_from_dense(H), k_top=2, k_bot=2)


def test_rejects_dim_d_mismatch():
    W = np.zeros((16, 256))
    H = _diagonal_h(32)  # different dim
    with pytest.raises(ValueError, match="dim_d"):
        measure_alignment(W, hessian_from_dense(H), dim_d=32, k_top=4, k_bot=4)


def test_rejects_k_too_large():
    W = _uniform_w(8, 8, seed=0)
    H = _diagonal_h(8)
    with pytest.raises(ValueError, match="dim_d"):
        measure_alignment(W, hessian_from_dense(H), k_top=10, k_bot=10)


def test_rejects_zero_epsilon():
    W = _uniform_w(16, 32, seed=0)
    H = _diagonal_h(16)
    with pytest.raises(ValueError, match="epsilon"):
        measure_alignment(W, hessian_from_dense(H), k_top=2, k_bot=2, epsilon=0)


def test_rejects_zero_norm_w():
    W = np.zeros((16, 32))
    H = _diagonal_h(16)
    with pytest.raises(ValueError, match="Frobenius"):
        measure_alignment(W, hessian_from_dense(H), k_top=2, k_bot=2)


def test_hessian_from_dense_rejects_non_square():
    with pytest.raises(ValueError, match="square"):
        hessian_from_dense(np.zeros((4, 8)))


# ----------------------------------------------------------- KAT regression


KAT_FIXTURE = (
    Path(__file__).parent.parent
    / "bench" / "fixtures" / "phi-alignment"
    / "synthetic-checkpoints.jsonl"
)


@pytest.mark.skipif(
    not KAT_FIXTURE.exists(),
    reason="KAT fixture not yet generated",
)
def test_phi_alignment_known_answer_tests():
    """Replay the pinned synthetic-checkpoints.jsonl fixture; every
    entry's verdict must match what the probe produces today.

    Algorithm change (probe rewrite, threshold tuning, etc.) MUST
    bump PROBE_VERSION + emit a new fixture file under
    bench/fixtures/phi-alignment/ — old runs replay against old data.
    """
    cls_to_w_factory = {
        "aligned": lambda dim_d, n, low_k, seed: _aligned_w(dim_d, n, low_k, seed),
        "uniform": lambda dim_d, n, _low_k, seed: _uniform_w(dim_d, n, seed),
        "anti": lambda dim_d, n, high_k, seed: _anti_w(dim_d, n, high_k, seed),
    }
    n_kats = 0
    for line in KAT_FIXTURE.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        kat = json.loads(line)
        n_kats += 1
        cls = kat["class"]
        dim_d = int(kat["dim_d"])
        n_cols = int(kat["n_cols"])
        seed = int(kat["seed"])
        k_top = int(kat["k_top"])
        k_bot = int(kat["k_bot"])
        # All synthetic KATs use low_k = high_k = k_top = 32.
        W = cls_to_w_factory[cls](dim_d, n_cols, k_top, seed)
        H = _diagonal_h(dim_d)
        r = measure_alignment(
            W, hessian_from_dense(H),
            k_top=k_top, k_bot=k_bot,
        )
        assert r.verdict == kat["expected_verdict"], (
            f"KAT mismatch for {kat['label']!r}: "
            f"expected {kat['expected_verdict']}, got {r.verdict} "
            f"(ratio={r.ratio:.4f})"
        )
        # Float tolerance for the ratio match (Lanczos has minor
        # numerical noise across scipy versions).
        observed = round(r.ratio, 4)
        recorded = float(kat["observed_ratio"])
        assert abs(observed - recorded) < 0.05, (
            f"ratio drift for {kat['label']!r}: recorded {recorded}, "
            f"observed {observed}"
        )

    assert n_kats >= 30, f"KAT fixture seems incomplete: only {n_kats} entries"