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


# ----------------------------------------------------------- monotonicity


def test_monotone_alignment_strength_in_concentration():
    """As W concentrates more mass on the low-λ subspace, the
    alignment ratio should strictly increase.

    Pattern from fox's `test_monotone_in_window_length` /
    `test_monotone_in_gradient_fraction`: scaling one knob while
    holding others fixed verifies the function responds in the
    expected direction. Catches sign errors + drops.
    """
    dim_d = 256
    H = _diagonal_h(dim_d)
    seed = 7

    # Three concentrations: spread (32 rows), tighter (16 rows),
    # tightest (8 rows). Tighter concentration on low-λ → larger
    # ratio. (Frobenius-norm-normalized so all three have same
    # ‖W‖_F.)
    rng = np.random.default_rng(seed)
    base_signal = rng.standard_normal((32, 256))

    def _make_w(low_k):
        W = np.zeros((dim_d, 256))
        # Reuse same total energy across concentrations.
        scale = (32 / low_k) ** 0.5
        W[:low_k, :] = base_signal[:low_k, :] * scale
        return W

    r_spread = measure_alignment(_make_w(32), hessian_from_dense(H), k_top=32, k_bot=32)
    r_medium = measure_alignment(_make_w(16), hessian_from_dense(H), k_top=32, k_bot=32)
    r_tight = measure_alignment(_make_w(8), hessian_from_dense(H), k_top=32, k_bot=32)

    assert r_spread.ratio < r_medium.ratio, (
        f"32-row vs 16-row: spread ratio {r_spread.ratio:.3f} should "
        f"be less than medium {r_medium.ratio:.3f}"
    )
    assert r_medium.ratio < r_tight.ratio, (
        f"16-row vs 8-row: medium ratio {r_medium.ratio:.3f} should "
        f"be less than tight {r_tight.ratio:.3f}"
    )


def test_monotone_alignment_in_dim_h():
    """Doubling W's column count (with the same row distribution)
    should leave the alignment ratio approximately invariant — the
    ratio is row-distribution-driven, not column-count-driven.

    Sanity invariant: A(W, H) is normalized by ‖W‖_F², so scaling
    column count (more samples of same row distribution) shouldn't
    move the ratio by more than sampling noise."""
    dim_d = 128
    H = _diagonal_h(dim_d)
    rng = np.random.default_rng(11)

    # Same row distribution, different number of columns.
    pattern = np.zeros(dim_d)
    pattern[:16] = 1.0
    W_128cols = rng.standard_normal((dim_d, 128)) * pattern[:, None]
    W_512cols = rng.standard_normal((dim_d, 512)) * pattern[:, None]

    r_128 = measure_alignment(W_128cols, hessian_from_dense(H), k_top=8, k_bot=8)
    r_512 = measure_alignment(W_512cols, hessian_from_dense(H), k_top=8, k_bot=8)

    # Tolerance is wide because column count affects sample noise.
    assert r_128.verdict == r_512.verdict, (
        f"verdict shouldn't change with column count: "
        f"128 cols → {r_128.verdict}, 512 cols → {r_512.verdict}"
    )


# ----------------------------------------------------------- hand-formula


def test_uniform_baseline_matches_analytical_formula():
    """The isotropic baseline ``a_uniform_bot`` in the
    AlignmentReport is computed analytically per #000034 §2.1
    derivation: ``E[A_bot(W_uniform, H)] = (1/dim_d) Σ 1/(λ_j+ε)``
    over the bot-k subset. Hand-compute and assert exact agreement.

    Pattern from fox's `test_b1_exact_formula` / `test_b2_exact_formula`:
    don't trust the implementation — derive the math in the test
    file from the spec, run the function, compare. Catches
    algorithm drift even when KAT regression passes.
    """
    dim_d = 64
    H = _diagonal_h(dim_d)  # eigenvalues 1..64 along canonical basis
    epsilon = 1e-6
    W = _uniform_w(dim_d, 32, seed=21)
    k = 8

    r = measure_alignment(W, hessian_from_dense(H), k_top=k, k_bot=k, epsilon=epsilon)

    # Hand-computed expected: bot-k eigenvalues are 1, 2, ..., k.
    expected_a_uniform_bot = (1.0 / dim_d) * sum(
        1.0 / (lam + epsilon) for lam in range(1, k + 1)
    )
    assert r.a_uniform == pytest.approx(expected_a_uniform_bot, abs=1e-9), (
        f"a_uniform mismatch: function {r.a_uniform}, "
        f"hand-computed {expected_a_uniform_bot}"
    )


def test_full_spectrum_a_top_plus_a_bot_covers_full_isotropic_baseline():
    """Closure invariant. When k_top + k_bot covers the full
    spectrum (k_top = k_bot = dim_d / 2), the function's a_top +
    a_bot should equal the full-spectrum alignment for any W.

    Pattern from fox's `test_total_equals_sum_of_three_contributions`:
    no missing term, no double-counting. Here the sum across all
    eigenvectors should reconstruct the full alignment.
    """
    dim_d = 32
    H = _diagonal_h(dim_d)
    W = _uniform_w(dim_d, 16, seed=33)
    k = dim_d // 2  # split the spectrum exactly

    r = measure_alignment(W, hessian_from_dense(H), k_top=k, k_bot=k)

    # Hand-compute full-spectrum A directly via dense decomposition.
    eigvals, eigvecs = np.linalg.eigh(H)
    proj = W.T @ eigvecs
    col_sq = np.sum(proj ** 2, axis=0)  # (dim_d,)
    weighted = col_sq / (eigvals + 1e-6)
    w_fro_sq = float(np.linalg.norm(W, "fro") ** 2)
    full_a = float(np.sum(weighted) / w_fro_sq)

    assert r.a_top + r.a_bot == pytest.approx(full_a, rel=1e-6), (
        f"a_top + a_bot ({r.a_top} + {r.a_bot} = "
        f"{r.a_top + r.a_bot}) should equal full-spectrum A "
        f"({full_a})"
    )


def test_eigenvalue_ordering_top_dominates_bot():
    """Closure invariant: top-k eigenvalues from Lanczos must all
    be ≥ all bot-k eigenvalues from Lanczos. Catches a bug where
    eigsh's 'LA'/'SA' modes returned overlapping ranges (which
    would happen on near-degenerate spectra without proper
    selection logic)."""
    dim_d = 64
    H = _diagonal_h(dim_d)
    W = _uniform_w(dim_d, 32, seed=44)
    r = measure_alignment(W, hessian_from_dense(H), k_top=8, k_bot=8)
    # Both lists are unordered Lanczos output; pick min/max.
    top_min = min(r.eigenvalues_top)
    bot_max = max(r.eigenvalues_bot)
    assert top_min >= bot_max, (
        f"top-k min {top_min} should be ≥ bot-k max {bot_max}"
    )


# ----------------------------------------------------------- validation (parametrized)


@pytest.mark.parametrize("bad_w_shape", [
    (16,),                        # 1-D
    (16, 16, 16),                 # 3-D
    (),                           # 0-D scalar
])
def test_rejects_wrong_dim_w(bad_w_shape):
    """Pattern from fox: parametrize over the shape-error cone."""
    H = _diagonal_h(16)
    W = np.zeros(bad_w_shape) if bad_w_shape else np.array(0.0)
    with pytest.raises(ValueError, match="2-D"):
        measure_alignment(W, hessian_from_dense(H), k_top=2, k_bot=2)


@pytest.mark.parametrize("bad_epsilon", [0, -1e-6, -1.0])
def test_rejects_non_positive_epsilon(bad_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=bad_epsilon)


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_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"