OBLITERATUS/tests/test_analysis_display_contracts.py

"""Numerical and presentation contracts for analysis results on the CPU."""

from types import SimpleNamespace as NS

import matplotlib.pyplot as plt
import numpy as np
import pytest
import torch

from obliteratus.analysis import visualization as plots
from obliteratus.analysis.spectral_certification import SpectralCertifier
from obliteratus.analysis.tuned_lens import RefusalTunedLens, TunedLensProbe

pytestmark = pytest.mark.cpu


@pytest.fixture(autouse=True)
def close_figures():
    yield
    plt.close("all")


@pytest.mark.parametrize("save", [False, True])
def test_token_spectrum_preserves_order_signs_and_layer_selection(tmp_path, monkeypatch, save):
    shown = []
    monkeypatch.setattr(plt, "show", lambda: shown.append(True))
    layer = NS(
        layer_idx=7,
        top_promoted=[("yes", 3.0), ("perhaps", 1.0)],
        top_suppressed=[("no", -4.0), ("never", -2.0)],
        refusal_compliance_gap=2.5,
        refusal_specificity=0.75,
    )
    result = NS(per_layer={7: layer}, strongest_refusal_layer=7)
    path = tmp_path / "tokens.png" if save else None
    fig = plots.plot_logit_lens_spectrum(
        result, layer_idx=7 if save else None, output_path=path, title="Spectrum" if save else None
    )
    ax = fig.axes[0]
    assert [patch.get_width() for patch in ax.patches] == [-2, -4, 3, 1]
    assert [label.get_text() for label in ax.get_yticklabels()] == [
        "'never'",
        "'no'",
        "'yes'",
        "'perhaps'",
    ]
    assert "2.5000" in ax.texts[0].get_text()
    assert ax.get_title() == ("Spectrum" if save else "Logit Lens — Layer 7")
    if save:
        assert path.read_bytes().startswith(b"\x89PNG")
        assert not shown
    else:
        assert shown == [True]
    assert plots.plot_logit_lens_spectrum(result, layer_idx=99) is None


@pytest.mark.parametrize("save", [False, True])
def test_pareto_averages_scores_and_preserves_reference_coordinates(tmp_path, monkeypatch, save):
    monkeypatch.setattr(plt, "show", lambda: None)
    path = tmp_path / "pareto.png" if save else None
    fig = plots.plot_capability_safety_pareto(
        {"a": NS(score=0.2), "b": NS(score=0.8)},
        0.3,
        other_points=[(0.9, 0.6, "Baseline")],
        output_path=path,
    )
    assert np.asarray(fig.axes[0].collections[0].get_offsets()).tolist() == [[0.3, 0.5]]
    assert np.asarray(fig.axes[0].collections[1].get_offsets()).tolist() == [[0.9, 0.6]]
    assert "Baseline" in [text.get_text() for text in fig.axes[0].texts]
    if save:
        assert path.read_bytes().startswith(b"\x89PNG")
    empty = plots.plot_capability_safety_pareto({}, 0.0)
    assert np.asarray(empty.axes[0].collections[0].get_offsets()).tolist() == [[0.0, 0.0]]


def test_topology_missing_means_have_zero_strength_and_squeezes_direction(monkeypatch):
    monkeypatch.setattr(plt, "show", lambda: None)
    fig = plots.plot_refusal_topology(
        {4: torch.tensor([[3.0, 0.0]]), 8: torch.tensor([0.0, 1.0])},
        {4: torch.tensor([2.0, 7.0])},
        {4: torch.tensor([0.0, 7.0])},
        [4],
    )
    assert [p.get_height() for p in fig.axes[0].patches] == [2.0, 0.0]
    assert [t.get_text() for t in fig.axes[0].get_xticklabels()] == ["4", "8"]


@pytest.mark.parametrize(
    "path", ["embed_out", "output", "transformer.wte", "model.embed_tokens", "gpt_neox.embed_in"]
)
def test_tuned_lens_finds_supported_output_or_tied_embeddings(path):
    weight = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
    node = NS(weight=weight)
    for part in reversed(path.split(".")):
        node = NS(**{part: node})
    torch.testing.assert_close(RefusalTunedLens()._get_unembedding_matrix(node), weight)


def test_tuned_lens_missing_output_is_actionable():
    with pytest.raises(RuntimeError, match="Cannot locate unembedding"):
        RefusalTunedLens()._get_unembedding_matrix(NS())


def test_tuned_lens_group_filters_invalid_and_unencodable_tokens():
    def encode(text, **kwargs):
        if text == "error":
            raise ValueError("unknown token")
        return {"valid": [1, 0], "empty": [], "negative": [-1], "outside": [3]}[text]

    boosts = RefusalTunedLens()._get_token_group_boosts(
        torch.tensor([2.0, 5.0, 9.0]),
        NS(encode=encode),
        ["error", "empty", "negative", "outside", "valid"],
    )
    assert boosts == [5.0]


def test_tuned_lens_affine_direction_and_report_have_known_oracle():
    lens = RefusalTunedLens(top_k=1)
    tokenizer = NS(
        decode=lambda ids: str(ids[0]), encode=lambda text, **kw: [0] if text == "sorry" else []
    )
    model = NS(
        lm_head=NS(weight=torch.eye(2)),
        model=NS(norm=NS(weight=torch.tensor([2.0, 3.0]), bias=torch.tensor([1.0, -1.0]))),
    )
    probe = TunedLensProbe(5, torch.eye(2), torch.tensor([100.0, 100.0]), 0.0)
    result = lens.analyze_all_layers(
        {5: torch.tensor([[1.0, 0.0]]), 6: torch.ones(2)}, {5: probe}, model, tokenizer
    )
    layer = result.per_layer[5]
    assert layer.top_promoted == [("0", 3.0)]
    assert layer.top_suppressed == [("1", -1.0)]
    assert layer.refusal_compliance_gap == 3.0
    assert layer.correction_magnitude == 0.0
    report = lens.format_report(result)
    assert "Strongest refusal layer: 5" in report
    assert "Mean refusal-compliance gap: 3.0000" in report
    assert "Top promoted:" in report
    assert "Top suppressed:" in report
    empty = lens.analyze_all_layers({6: torch.ones(2)}, {}, model, tokenizer)
    assert empty.per_layer == {}
    assert empty.mean_refusal_compliance_gap == 0.0
    assert "No layers analyzed." in lens.format_report(empty)


def test_tuned_lens_comparison_intersects_layers_and_detects_reversed_ranking():
    result = NS(
        per_layer={
            1: NS(refusal_compliance_gap=1.0),
            2: NS(refusal_compliance_gap=2.0),
            3: NS(refusal_compliance_gap=3.0),
        }
    )
    assert RefusalTunedLens.compare_with_logit_lens(result, {1: 1.0}) == 1.0
    assert (
        RefusalTunedLens.compare_with_logit_lens(result, {1: 3.0, 2: 2.0, 3: 1.0, 4: 100.0}) == -1.0
    )


def test_spectral_diagonal_oracle_and_degenerate_condition_estimates():
    certifier = SpectralCertifier()
    covariance = torch.diag(torch.tensor([0.0, 2.0, 8.0]))
    result = certifier._eigenvalue_analysis(covariance, bbp_threshold=4.0, mp_upper=2.0)
    torch.testing.assert_close(result.eigenvalues, torch.tensor([8.0, 2.0, 0.0]))
    torch.testing.assert_close(
        covariance @ result.eigenvectors, result.eigenvectors @ torch.diag(result.eigenvalues)
    )
    assert result.above_threshold == [0]
    assert result.in_bulk == [1]
    assert result.signal_subspace_dim == 1
    assert certifier._estimate_condition_number(covariance) == 4.0
    assert certifier._estimate_condition_number(torch.zeros(2, 2)) == 1.0
    assert certifier._estimate_noise_variance(covariance, n=2, d=3) == 2.0
    ratio = (1 - (3 / 12) ** 0.5) ** 2 + (3 / 12) ** (1 / 3)
    assert certifier._estimate_noise_variance(covariance, n=12, d=3) == pytest.approx(2 / ratio)
    assert certifier._estimate_noise_variance(torch.zeros(2, 2), n=2, d=2) == 1e-10
    for spectrum in [
        torch.tensor([]),
        torch.tensor([1.0]),
        torch.tensor([1.0, float("nan")]),
        torch.tensor([-1.0, 2.0]),
    ]:
        assert certifier._estimate_condition_number_from_spectrum(spectrum) == 1.0
    assert certifier._estimate_condition_number_from_spectrum(torch.tensor([1.0, 1e8])) == 1e6


def test_spectral_linear_algebra_failure_returns_conservative_fallbacks(monkeypatch):
    def fail(*args, **kwargs):
        raise RuntimeError("eigensolver failed")

    monkeypatch.setattr(torch.linalg, "eigvalsh", fail)
    monkeypatch.setattr(torch.linalg, "eigh", fail)
    certifier = SpectralCertifier()
    assert certifier._estimate_noise_variance(torch.eye(2), 2, 2) == 1.0
    assert certifier._estimate_condition_number(torch.eye(2)) == 1.0
    result = certifier._eigenvalue_analysis(torch.eye(2), 1.0, 1.0)
    assert result.signal_subspace_dim == 0
    assert result.above_threshold == []
    assert result.in_bulk == []
    torch.testing.assert_close(result.eigenvalues, torch.zeros(1))