Buckets:

DERBACH
/

latent-rubric

Files

xet

DERBACH/latent-rubric / tests /thread_statistics /test_composition.py

DERBACH

about 1 month ago

download

raw

10.7 kB

"""Deterministic unit/algorithm checks, not research experiments."""

import itertools
import json

import numpy as np
import pytest

from analysis.thread_statistics import composition_point_difference, paired_composition_bootstrap
from eval.metrics import composition


def example():
    # Three pairs belong to A, one to B. Ties depend on rubric, not model.
    a = np.zeros((2, 4, 2))
    b = np.array([[[1, 0], [1, np.nan], [1, np.inf], [0, 1]]] * 2)
    valid = np.array([[[1, 1], [1, 0], [1, 0], [1, 1]]] * 2, dtype=bool)
    return a, b, valid, valid.copy(), np.array(["A", "A", "A", "B"])


def call(data=None, **kwargs):
    data = example() if data is None else data
    defaults = dict(seed_ids_a=[7, 11], seed_ids_b=[7, 11], n_bootstrap=300)
    defaults.update(kwargs)
    return paired_composition_bootstrap(*data, **defaults)


def test_rubric_first_point_repeated_prompt_pairs_and_ties():
    result = call(n_bootstrap=10_000, return_bootstrap_distribution=True)
    assert result["point_a"] == 0
    assert result["point_b"] == result["difference"] == 0.625
    assert result["bootstrap_mean"] != result["difference"]
    assert result["ci95"] == [0.5, 0.625]
    assert set(result["bootstrap_distribution"]) == {0.5, 0.625}
    assert result["counts"]["pairs_per_prompt"] == [3, 1]
    assert result["counts"]["valid_pairs_per_rubric"] == [4, 2]
    assert result["counts"]["human_tie_pair_rubric_decisions"] == 2
    assert result["counts"]["valid_model_decisions_per_method"] == 12
    assert result["bootstrap_samples"] == 10_000
    assert result["bootstrap_undefined_samples"] == 0
    json.dumps(result, allow_nan=False)
    # Equal prompts after per-pair averaging is a different quantity (2/3).
    _, b, v, _, _ = example()
    per_pair = np.nanmean(np.where(v, b, np.nan), axis=2)
    old = (per_pair[:, :3].mean(axis=1) + per_pair[:, 3]) / 2
    assert old.mean() == pytest.approx(2 / 3)


def test_exact_point_matches_existing_composition_with_64_rubrics():
    rng = np.random.default_rng(205)
    attributes = rng.integers(0, 4, size=(16, 4)).astype(float)
    attributes[1] = attributes[0]  # All-rubric human tie for this pair.
    pairs = np.array([(i, j) for start in range(0, 16, 4) for i, j in itertools.combinations(range(start, start + 4), 2)])
    a, b, valid, ma, mb = [], [], [], [], []
    for _ in range(3):
        m_a, e_a = composition(rng.normal(size=attributes.shape), attributes, pairs)
        m_b, e_b = composition(attributes + rng.normal(size=attributes.shape), attributes, pairs)
        a.append(e_a["pair_correct"])
        b.append(e_b["pair_correct"])
        valid.append(e_a["pair_valid"])
        np.testing.assert_array_equal(e_a["pair_valid"], e_b["pair_valid"])
        ma.append(m_a["composition_accuracy"])
        mb.append(m_b["composition_accuracy"])
    point = composition_point_difference(a, b, valid, valid, seed_ids_a=[1, 2, 3], seed_ids_b=[1, 2, 3])
    # Array reduction layout can differ by a final floating-point bit.
    assert point["point_a"] == pytest.approx(np.mean(ma), abs=5e-16, rel=0)
    assert point["point_b"] == pytest.approx(np.mean(mb), abs=5e-16, rel=0)
    assert point["difference"] == pytest.approx(np.mean(mb) - np.mean(ma), abs=5e-16, rel=0)
    assert point["counts"]["rubrics"] == 64


def test_zero_model_difference_is_incorrect_not_excluded():
    a, b, va, vb, pids = example()
    b[:] = 0
    result = call((a, b, va, vb, pids))
    assert result["point_b"] == 0
    assert result["counts"]["valid_pairs_per_rubric"] == [4, 2]


@pytest.mark.parametrize("resample_seeds", [True, False])
def test_paired_identical_models_have_zero_entire_distribution(resample_seeds):
    a, b, va, vb, pids = example()
    result = call((b, b.copy(), va, vb, pids), resample_seeds=resample_seeds, n_bootstrap=10_000, return_bootstrap_distribution=True)
    assert result["difference"] == 0
    assert result["ci95"] == [0.0, 0.0]
    assert result["bootstrap_distribution"] == [0.0] * 10_000


@pytest.mark.parametrize("resample_seeds", [True, False])
def test_every_draw_matches_literal_cluster_and_seed_replication(resample_seeds):
    rng = np.random.default_rng(18)
    a = rng.integers(0, 2, size=(3, 9, 4))
    b = rng.integers(0, 2, size=a.shape)
    v = rng.integers(0, 2, size=(9, 4)).astype(bool)
    # At least one valid pair per rubric in each of three unequal clusters.
    v[[0, 2, 5]] = True
    v = np.broadcast_to(v, a.shape)
    pids = np.repeat(["p0", "p1", "p2"], [2, 3, 4])
    count, random_seed = 257, 810
    got = paired_composition_bootstrap(a, b, v, v, pids, seed_ids_a=[2, 4, 6], seed_ids_b=[2, 4, 6], n_bootstrap=count, random_seed=random_seed, batch_size=11, resample_seeds=resample_seeds, return_bootstrap_distribution=True)
    ss, ps = np.random.SeedSequence(random_seed).spawn(2)
    sg, pg = np.random.default_rng(ss), np.random.default_rng(ps)
    expected = []
    for _ in range(count):
        prompt_counts = pg.multinomial(3, [1 / 3] * 3)
        seed_counts = sg.multinomial(3, [1 / 3] * 3) if resample_seeds else [1] * 3
        ix = np.concatenate([np.tile(np.flatnonzero(pids == prompt), prompt_counts[i]) for i, prompt in enumerate(np.unique(pids))])
        si = np.repeat(np.arange(3), seed_counts)
        accuracies = []
        for array in (a, b):
            by_seed = []
            for seed in si:
                by_rubric = [array[seed, ix, r][v[seed, ix, r]].mean() for r in range(4)]
                by_seed.append(np.mean(by_rubric))
            accuracies.append(np.mean(by_seed))
        expected.append(accuracies[1] - accuracies[0])
    np.testing.assert_allclose(got["bootstrap_distribution"], expected, rtol=0, atol=2e-16)
    np.testing.assert_allclose(got["ci95"], np.quantile(expected, [0.025, 0.975]), atol=2e-16)


def test_batch_size_does_not_change_random_draws_or_results():
    small = call(batch_size=1, return_bootstrap_distribution=True)
    large = call(batch_size=173, return_bootstrap_distribution=True)
    assert small == large


def test_swapping_models_negates_point_and_interval():
    a, b, va, vb, pids = example()
    forward = call()
    reverse = call((b, a, vb, va, pids))
    assert reverse["difference"] == -forward["difference"]
    assert reverse["ci95"] == [-forward["ci95"][1], -forward["ci95"][0]]


def test_observed_all_tie_rubric_omitted_and_disclosed():
    a, b, va, vb, pids = example()
    va[:, :, 1] = vb[:, :, 1] = False
    result = call((a, b, va, vb, pids))
    assert result["point_b"] == 0.75
    assert result["counts"]["excluded_all_tie_rubric_indices"] == [1]
    assert result["counts"]["active_rubrics"] == 1


def test_empty_resampled_rubric_withholds_ci_without_redraw_or_dropping():
    a, b, va, vb, pids = example()
    va[:, -1, 1] = vb[:, -1, 1] = False  # Rubric 1 supported only in A.
    result = call((a, b, va, vb, pids), n_bootstrap=10_000, return_bootstrap_distribution=True)
    assert result["difference"] == 0.375
    assert result["ci95"] is None
    assert result["interval_status"] == "undefined_rubric_in_bootstrap"
    assert result["bootstrap_samples"] == 10_000
    assert 0 < result["bootstrap_undefined_samples"] < 10_000
    assert result["bootstrap_distribution"].count(None) == result["bootstrap_undefined_samples"]
    json.dumps(result, allow_nan=False)


def test_all_tie_prompt_is_retained_in_sampling_frame():
    a, b, va, vb, pids = example()
    va[:, -1] = vb[:, -1] = False
    result = call((a, b, va, vb, pids))
    assert result["counts"]["prompts"] == 2
    assert result["counts"]["prompts_with_valid_pairs"] == 1
    assert result["bootstrap_undefined_samples"] > 0


def test_single_prompt_withholds_prompt_ci():
    a, b, va, vb, pids = example()
    result = call((a, b, va, vb, ["A"] * 4))
    assert result["difference"] == 0.625
    assert result["ci95"] is None
    assert result["interval_status"] == "insufficient_prompt_clusters"


def test_single_seed_discloses_conditional_uncertainty():
    a, b, va, vb, pids = example()
    result = call((a[:1], b[:1], va[:1], vb[:1], pids), seed_ids_a=[11], seed_ids_b=[11])
    assert result["ci95"] is not None
    assert any("Only one paired seed" in note for note in result["assumptions"])


@pytest.mark.parametrize("kwargs, match", [
    ({"seed_ids_a": [7, 11], "seed_ids_b": [11, 7]}, "seed axis alignment"),
    ({"seed_ids_a": [7, 7], "seed_ids_b": [7, 7]}, "unique IDs"),
    ({"seed_ids_a": [7], "seed_ids_b": [7]}, "length 2"),
    ({"seed_ids_a": None, "seed_ids_b": None}, "required"),
    ({"seed_ids_a": [7, 11], "seed_ids_b": None}, "provide both"),
    ({"seed_ids_a": [7, "11"], "seed_ids_b": [7, "11"]}, "must not mix"),
    ({"rubric_ids_a": ["r0", "r1"], "rubric_ids_b": ["r1", "r0"]}, "rubric axis alignment"),
    ({"pair_ids_a": [0, 1, 2, 3], "pair_ids_b": [1, 0, 2, 3]}, "pair axis alignment"),
    ({"pair_ids_a": [0, 1, 2, 3]}, "provide both"),
    ({"pair_prompt_ids_b": ["A", "B", "A", "A"]}, "prompt axis alignment"),
    ({"n_bootstrap": 1}, "n_bootstrap"),
    ({"n_bootstrap": 3.5}, "n_bootstrap"),
    ({"batch_size": 0}, "batch_size"),
    ({"random_seed": -1}, "random_seed"),
    ({"resample_seeds": "yes"}, "resample_seeds"),
])
def test_axis_seed_alignment_and_options_rejected(kwargs, match):
    with pytest.raises(ValueError, match=match):
        call(**kwargs)


@pytest.mark.parametrize("failure", ["transposed", "unequal_shape", "empty", "different_seed_mask", "different_model_mask", "nan_on_valid", "fractional_correct", "nan_valid", "all_ties", "short_prompts", "missing_prompt"])
def test_bad_arrays_are_not_silently_repaired(failure):
    a, b, va, vb, pids = example()
    if failure == "transposed":
        a, b, va, vb = [x.transpose(1, 0, 2) for x in (a, b, va, vb)]
    elif failure == "unequal_shape":
        b = b[:1]
    elif failure == "empty":
        a, b, va, vb = [x[:, :0] for x in (a, b, va, vb)]
    elif failure == "different_seed_mask":
        va[0, 0, 0] = vb[0, 0, 0] = False
    elif failure == "different_model_mask":
        vb[0, 0, 0] = False
    elif failure == "nan_on_valid":
        b[0, 0, 0] = np.nan
    elif failure == "fractional_correct":
        b[0, 0, 0] = 0.5
    elif failure == "nan_valid":
        va = va.astype(float)
        va[0, 0, 0] = np.nan
    elif failure == "all_ties":
        va[:] = vb[:] = False
    elif failure == "short_prompts":
        pids = pids[:3]
    elif failure == "missing_prompt":
        pids = ["A", None, "A", "B"]
    with pytest.raises(ValueError):
        call((a, b, va, vb, pids))


def test_inputs_are_not_mutated():
    data = example()
    copies = [x.copy() for x in data]
    call(data)
    for after, before in zip(data, copies):
        np.testing.assert_array_equal(after, before)

Xet Storage Details

Size:: 10.7 kB
Xet hash:: 47fdba351ef5a48eaf3e6e3b34bcbca2cd1bcf183c4cbc168b791263dc35cf57

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Deterministic unit/algorithm checks, not research experiments."""

	import itertools
	import json

	import numpy as np
	import pytest

	from analysis.thread_statistics import composition_point_difference, paired_composition_bootstrap
	from eval.metrics import composition


	def example():
	# Three pairs belong to A, one to B. Ties depend on rubric, not model.
	a = np.zeros((2, 4, 2))
	b = np.array([[[1, 0], [1, np.nan], [1, np.inf], [0, 1]]] * 2)
	valid = np.array([[[1, 1], [1, 0], [1, 0], [1, 1]]] * 2, dtype=bool)
	return a, b, valid, valid.copy(), np.array(["A", "A", "A", "B"])


	def call(data=None, **kwargs):
	data = example() if data is None else data
	defaults = dict(seed_ids_a=[7, 11], seed_ids_b=[7, 11], n_bootstrap=300)
	defaults.update(kwargs)
	return paired_composition_bootstrap(data, *defaults)


	def test_rubric_first_point_repeated_prompt_pairs_and_ties():
	result = call(n_bootstrap=10_000, return_bootstrap_distribution=True)
	assert result["point_a"] == 0
	assert result["point_b"] == result["difference"] == 0.625
	assert result["bootstrap_mean"] != result["difference"]
	assert result["ci95"] == [0.5, 0.625]
	assert set(result["bootstrap_distribution"]) == {0.5, 0.625}
	assert result["counts"]["pairs_per_prompt"] == [3, 1]
	assert result["counts"]["valid_pairs_per_rubric"] == [4, 2]
	assert result["counts"]["human_tie_pair_rubric_decisions"] == 2
	assert result["counts"]["valid_model_decisions_per_method"] == 12
	assert result["bootstrap_samples"] == 10_000
	assert result["bootstrap_undefined_samples"] == 0
	json.dumps(result, allow_nan=False)
	# Equal prompts after per-pair averaging is a different quantity (2/3).
	_, b, v, _, _ = example()
	per_pair = np.nanmean(np.where(v, b, np.nan), axis=2)
	old = (per_pair[:, :3].mean(axis=1) + per_pair[:, 3]) / 2
	assert old.mean() == pytest.approx(2 / 3)


	def test_exact_point_matches_existing_composition_with_64_rubrics():
	rng = np.random.default_rng(205)
	attributes = rng.integers(0, 4, size=(16, 4)).astype(float)
	attributes[1] = attributes[0] # All-rubric human tie for this pair.
	pairs = np.array([(i, j) for start in range(0, 16, 4) for i, j in itertools.combinations(range(start, start + 4), 2)])
	a, b, valid, ma, mb = [], [], [], [], []
	for _ in range(3):
	m_a, e_a = composition(rng.normal(size=attributes.shape), attributes, pairs)
	m_b, e_b = composition(attributes + rng.normal(size=attributes.shape), attributes, pairs)
	a.append(e_a["pair_correct"])
	b.append(e_b["pair_correct"])
	valid.append(e_a["pair_valid"])
	np.testing.assert_array_equal(e_a["pair_valid"], e_b["pair_valid"])
	ma.append(m_a["composition_accuracy"])
	mb.append(m_b["composition_accuracy"])
	point = composition_point_difference(a, b, valid, valid, seed_ids_a=[1, 2, 3], seed_ids_b=[1, 2, 3])
	# Array reduction layout can differ by a final floating-point bit.
	assert point["point_a"] == pytest.approx(np.mean(ma), abs=5e-16, rel=0)
	assert point["point_b"] == pytest.approx(np.mean(mb), abs=5e-16, rel=0)
	assert point["difference"] == pytest.approx(np.mean(mb) - np.mean(ma), abs=5e-16, rel=0)
	assert point["counts"]["rubrics"] == 64


	def test_zero_model_difference_is_incorrect_not_excluded():
	a, b, va, vb, pids = example()
	b[:] = 0
	result = call((a, b, va, vb, pids))
	assert result["point_b"] == 0
	assert result["counts"]["valid_pairs_per_rubric"] == [4, 2]


	@pytest.mark.parametrize("resample_seeds", [True, False])
	def test_paired_identical_models_have_zero_entire_distribution(resample_seeds):
	a, b, va, vb, pids = example()
	result = call((b, b.copy(), va, vb, pids), resample_seeds=resample_seeds, n_bootstrap=10_000, return_bootstrap_distribution=True)
	assert result["difference"] == 0
	assert result["ci95"] == [0.0, 0.0]
	assert result["bootstrap_distribution"] == [0.0] * 10_000


	@pytest.mark.parametrize("resample_seeds", [True, False])
	def test_every_draw_matches_literal_cluster_and_seed_replication(resample_seeds):
	rng = np.random.default_rng(18)
	a = rng.integers(0, 2, size=(3, 9, 4))
	b = rng.integers(0, 2, size=a.shape)
	v = rng.integers(0, 2, size=(9, 4)).astype(bool)
	# At least one valid pair per rubric in each of three unequal clusters.
	v[[0, 2, 5]] = True
	v = np.broadcast_to(v, a.shape)
	pids = np.repeat(["p0", "p1", "p2"], [2, 3, 4])
	count, random_seed = 257, 810
	got = paired_composition_bootstrap(a, b, v, v, pids, seed_ids_a=[2, 4, 6], seed_ids_b=[2, 4, 6], n_bootstrap=count, random_seed=random_seed, batch_size=11, resample_seeds=resample_seeds, return_bootstrap_distribution=True)
	ss, ps = np.random.SeedSequence(random_seed).spawn(2)
	sg, pg = np.random.default_rng(ss), np.random.default_rng(ps)
	expected = []
	for _ in range(count):
	prompt_counts = pg.multinomial(3, [1 / 3] * 3)
	seed_counts = sg.multinomial(3, [1 / 3] * 3) if resample_seeds else [1] * 3
	ix = np.concatenate([np.tile(np.flatnonzero(pids == prompt), prompt_counts[i]) for i, prompt in enumerate(np.unique(pids))])
	si = np.repeat(np.arange(3), seed_counts)
	accuracies = []
	for array in (a, b):
	by_seed = []
	for seed in si:
	by_rubric = [array[seed, ix, r][v[seed, ix, r]].mean() for r in range(4)]
	by_seed.append(np.mean(by_rubric))
	accuracies.append(np.mean(by_seed))
	expected.append(accuracies[1] - accuracies[0])
	np.testing.assert_allclose(got["bootstrap_distribution"], expected, rtol=0, atol=2e-16)
	np.testing.assert_allclose(got["ci95"], np.quantile(expected, [0.025, 0.975]), atol=2e-16)


	def test_batch_size_does_not_change_random_draws_or_results():
	small = call(batch_size=1, return_bootstrap_distribution=True)
	large = call(batch_size=173, return_bootstrap_distribution=True)
	assert small == large


	def test_swapping_models_negates_point_and_interval():
	a, b, va, vb, pids = example()
	forward = call()
	reverse = call((b, a, vb, va, pids))
	assert reverse["difference"] == -forward["difference"]
	assert reverse["ci95"] == [-forward["ci95"][1], -forward["ci95"][0]]


	def test_observed_all_tie_rubric_omitted_and_disclosed():
	a, b, va, vb, pids = example()
	va[:, :, 1] = vb[:, :, 1] = False
	result = call((a, b, va, vb, pids))
	assert result["point_b"] == 0.75
	assert result["counts"]["excluded_all_tie_rubric_indices"] == [1]
	assert result["counts"]["active_rubrics"] == 1


	def test_empty_resampled_rubric_withholds_ci_without_redraw_or_dropping():
	a, b, va, vb, pids = example()
	va[:, -1, 1] = vb[:, -1, 1] = False # Rubric 1 supported only in A.
	result = call((a, b, va, vb, pids), n_bootstrap=10_000, return_bootstrap_distribution=True)
	assert result["difference"] == 0.375
	assert result["ci95"] is None
	assert result["interval_status"] == "undefined_rubric_in_bootstrap"
	assert result["bootstrap_samples"] == 10_000
	assert 0 < result["bootstrap_undefined_samples"] < 10_000
	assert result["bootstrap_distribution"].count(None) == result["bootstrap_undefined_samples"]
	json.dumps(result, allow_nan=False)


	def test_all_tie_prompt_is_retained_in_sampling_frame():
	a, b, va, vb, pids = example()
	va[:, -1] = vb[:, -1] = False
	result = call((a, b, va, vb, pids))
	assert result["counts"]["prompts"] == 2
	assert result["counts"]["prompts_with_valid_pairs"] == 1
	assert result["bootstrap_undefined_samples"] > 0


	def test_single_prompt_withholds_prompt_ci():
	a, b, va, vb, pids = example()
	result = call((a, b, va, vb, ["A"] * 4))
	assert result["difference"] == 0.625
	assert result["ci95"] is None
	assert result["interval_status"] == "insufficient_prompt_clusters"


	def test_single_seed_discloses_conditional_uncertainty():
	a, b, va, vb, pids = example()
	result = call((a[:1], b[:1], va[:1], vb[:1], pids), seed_ids_a=[11], seed_ids_b=[11])
	assert result["ci95"] is not None
	assert any("Only one paired seed" in note for note in result["assumptions"])


	@pytest.mark.parametrize("kwargs, match", [
	({"seed_ids_a": [7, 11], "seed_ids_b": [11, 7]}, "seed axis alignment"),
	({"seed_ids_a": [7, 7], "seed_ids_b": [7, 7]}, "unique IDs"),
	({"seed_ids_a": [7], "seed_ids_b": [7]}, "length 2"),
	({"seed_ids_a": None, "seed_ids_b": None}, "required"),
	({"seed_ids_a": [7, 11], "seed_ids_b": None}, "provide both"),
	({"seed_ids_a": [7, "11"], "seed_ids_b": [7, "11"]}, "must not mix"),
	({"rubric_ids_a": ["r0", "r1"], "rubric_ids_b": ["r1", "r0"]}, "rubric axis alignment"),
	({"pair_ids_a": [0, 1, 2, 3], "pair_ids_b": [1, 0, 2, 3]}, "pair axis alignment"),
	({"pair_ids_a": [0, 1, 2, 3]}, "provide both"),
	({"pair_prompt_ids_b": ["A", "B", "A", "A"]}, "prompt axis alignment"),
	({"n_bootstrap": 1}, "n_bootstrap"),
	({"n_bootstrap": 3.5}, "n_bootstrap"),
	({"batch_size": 0}, "batch_size"),
	({"random_seed": -1}, "random_seed"),
	({"resample_seeds": "yes"}, "resample_seeds"),
	])
	def test_axis_seed_alignment_and_options_rejected(kwargs, match):
	with pytest.raises(ValueError, match=match):
	call(**kwargs)


	@pytest.mark.parametrize("failure", ["transposed", "unequal_shape", "empty", "different_seed_mask", "different_model_mask", "nan_on_valid", "fractional_correct", "nan_valid", "all_ties", "short_prompts", "missing_prompt"])
	def test_bad_arrays_are_not_silently_repaired(failure):
	a, b, va, vb, pids = example()
	if failure == "transposed":
	a, b, va, vb = [x.transpose(1, 0, 2) for x in (a, b, va, vb)]
	elif failure == "unequal_shape":
	b = b[:1]
	elif failure == "empty":
	a, b, va, vb = [x[:, :0] for x in (a, b, va, vb)]
	elif failure == "different_seed_mask":
	va[0, 0, 0] = vb[0, 0, 0] = False
	elif failure == "different_model_mask":
	vb[0, 0, 0] = False
	elif failure == "nan_on_valid":
	b[0, 0, 0] = np.nan
	elif failure == "fractional_correct":
	b[0, 0, 0] = 0.5
	elif failure == "nan_valid":
	va = va.astype(float)
	va[0, 0, 0] = np.nan
	elif failure == "all_ties":
	va[:] = vb[:] = False
	elif failure == "short_prompts":
	pids = pids[:3]
	elif failure == "missing_prompt":
	pids = ["A", None, "A", "B"]
	with pytest.raises(ValueError):
	call((a, b, va, vb, pids))


	def test_inputs_are_not_mutated():
	data = example()
	copies = [x.copy() for x in data]
	call(data)
	for after, before in zip(data, copies):
	np.testing.assert_array_equal(after, before)