tests/test_core.py

#!/usr/bin/env python3
"""Deterministic unit tests for the reproduction's pure-numeric core.
Avoids torch / sentence-transformers imports (native-instability on this box);
those paths are exercised by the experiment drivers + stored outputs."""
import sys, os, json
import numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))

from exp_claim2_diversity import greedy_diverse, avg_pairwise_sim
from exp_claim2b_draft_length import alpha_from_meanAL

def test_alpha_inversion_roundtrip():
    for a in [0.3, 0.5, 0.72, 0.9]:
        for k in [4, 8]:
            E = a * (1 - a**k) / (1 - a)          # forward E[AL]
            a_hat = alpha_from_meanAL(E, k)        # invert
            assert abs(a_hat - a) < 1e-4, (a, k, a_hat)

def test_greedy_more_diverse_than_max_similarity():
    rng = np.random.default_rng(0)
    E = rng.normal(size=(60, 16))
    E /= np.linalg.norm(E, axis=1, keepdims=True)
    k = 12
    gidx = greedy_diverse(E, k, seed=0)
    # a deliberately low-diversity subset: take one anchor's nearest neighbours
    anchor = 0
    sims = E @ E[anchor]
    low_idx = list(np.argsort(-sims)[:k])
    assert avg_pairwise_sim(E, gidx) < avg_pairwise_sim(E, low_idx)

def test_greedy_beats_random_on_average():
    rng = np.random.default_rng(1)
    E = rng.normal(size=(80, 24)); E /= np.linalg.norm(E, axis=1, keepdims=True)
    k = 20
    g = avg_pairwise_sim(E, greedy_diverse(E, k, seed=0))
    rand = np.mean([avg_pairwise_sim(E, list(rng.choice(80, k, replace=False))) for _ in range(30)])
    assert g < rand

def test_roofline_kstar_monotone_nonincreasing():
    # standard SD throughput model must give non-increasing k* as batch grows
    alpha = 0.72
    M_t, C_t, M_d, C_d = 1.0, 0.01, 0.2, 0.002
    egen = lambda k: (1 - alpha**(k+1)) / (1 - alpha)
    ks = range(1, 13)
    def kstar(b):
        Ld = max(M_d, C_d*b)
        return max(ks, key=lambda k: egen(k)/(k*Ld + max(M_t, C_t*b*(k+1))))
    batches = [1, 2, 4, 8, 16, 32, 64, 128, 256]
    stars = [kstar(b) for b in batches]
    assert all(stars[i] >= stars[i+1] for i in range(len(stars)-1)), stars
    assert stars[0] > stars[-1]   # strict overall decrease

def test_stored_outputs_present_and_consistent():
    base = os.path.join(os.path.dirname(__file__), "..", "outputs")
    div = json.load(open(os.path.join(base, "claim2_diversity.json"), encoding="utf-8"))
    assert div["overall"]["n_better"] == div["overall"]["n_total"]      # greedy<random in every category
    assert div["overall"]["greedy_avg"] < div["overall"]["random_avg"]
    c1 = json.load(open(os.path.join(base, "claim1_throughput_bias.json"), encoding="utf-8"))
    assert c1["config"]["n_real_rounds"] > 100 and c1["config"]["n_synth_rounds"] > 100

if __name__ == "__main__":
    import pytest
    raise SystemExit(pytest.main([__file__, "-v"]))