File size: 1,501 Bytes
8b37c3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | from __future__ import annotations
import numpy as np
import pandas as pd
from src.evaluation.protocol import (
battery_cluster_bootstrap,
grouped_train_val_test_folds,
holm_adjust,
paired_battery_wilcoxon,
)
def test_grouped_folds_keep_batteries_disjoint():
groups = np.repeat([f"B{i}" for i in range(10)], 3)
folds = list(grouped_train_val_test_folds(groups, n_splits=5))
assert len(folds) == 5
for train, val, test in folds:
assert not set(groups[train]) & set(groups[val])
assert not set(groups[train]) & set(groups[test])
assert not set(groups[val]) & set(groups[test])
def test_cluster_bootstrap_and_holm_outputs():
y = np.arange(12, dtype=float)
pred = y + 0.2
groups = np.repeat(["A", "B", "C"], 4)
ci = battery_cluster_bootstrap(y, pred, groups, n_bootstrap=50)
assert {"metric", "estimate", "ci_low", "ci_high"} <= set(ci)
adjusted = holm_adjust([0.01, 0.04, 0.03])
assert np.all((adjusted >= 0) & (adjusted <= 1))
def test_paired_wilcoxon_uses_per_battery_scores():
rows = []
for battery, a, b in zip(["A", "B", "C", "D"], [1, 1, 1, 1], [2, 3, 2, 4]):
rows.extend([
{"battery_id": battery, "model": "reference", "mae": a},
{"battery_id": battery, "model": "candidate", "mae": b},
])
result = paired_battery_wilcoxon(pd.DataFrame(rows), reference_model="reference")
assert result.loc[0, "n_batteries"] == 4
assert "p_value_holm" in result
|