File size: 2,613 Bytes
0cef53c | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import itertools
import math
import pytest
import torch
from gnn4colliders.tasks import (
BinaryClassificationTask,
MulticlassClassificationTask,
macro_scores,
per_class_weighted_mean,
weighted_binary_auc,
)
def _pairwise_auc(labels, scores, weights):
total = correct = 0.0
for i, j in itertools.product(range(len(labels)), repeat=2):
if labels[i] and not labels[j]:
pair = weights[i] * weights[j]
total += pair
if scores[i] > scores[j]:
correct += pair
elif scores[i] == scores[j]:
correct += 0.5 * pair
return correct / total
def test_sorted_auc_matches_pairwise_definition_with_ties():
generator = torch.Generator().manual_seed(0)
labels = torch.rand(60, generator=generator) > 0.4
scores = torch.randint(0, 8, (60,), generator=generator).float()
weights = torch.rand(60, generator=generator)
assert weighted_binary_auc(labels, scores, weights) == pytest.approx(
_pairwise_auc(labels.tolist(), scores.tolist(), weights.tolist())
)
assert math.isnan(weighted_binary_auc(labels & False, scores, weights))
def test_loss_is_the_mean_of_per_class_weighted_means():
loss = torch.tensor([1.0, 2.0, 3.0, 4.0])
labels = torch.tensor([0, 0, 2, 2])
weights = torch.tensor([1.0, 3.0, 1.0, 1.0])
expected = ((1 * 1 + 3 * 2) / 4 + (3 + 4) / 2) / 2
assert float(per_class_weighted_mean(loss, labels, weights, 3)) == pytest.approx(
expected
)
def test_macro_scores():
scores = macro_scores(torch.tensor([0, 0, 1, 1]), torch.tensor([0, 1, 1, 1]), 2)
assert scores["precision"] == pytest.approx((1 + 2 / 3) / 2)
assert scores["recall"] == pytest.approx((0.5 + 1) / 2)
def test_multiclass_metrics_exclude_non_positive_weights():
task = MulticlassClassificationTask()
logits = torch.tensor([[2.0, 0.0], [0.0, 2.0], [2.0, 0.0]])
labels = torch.tensor([0, 1, 1])
metrics = task.metrics(logits, labels, torch.tensor([1.0, 1.0, -1.0]))
assert metrics["auc"] == 1.0
assert metrics["accuracy"] == pytest.approx(2 / 3)
def test_binary_task_balances_classes_and_uses_absolute_weights():
task = BinaryClassificationTask(absolute_weights=True)
logits = torch.tensor([[0.0], [0.0], [0.0]])
labels = torch.tensor([1, 0, 0])
loss = task.loss(logits, labels, torch.tensor([-1.0, 1.0, 2.0]))
assert float(loss) == pytest.approx(float(torch.log(torch.tensor(2.0))))
assert task.metrics(torch.tensor([3.0, -1.0, 0.0]), labels, torch.ones(3))[
"auc"
] == pytest.approx(1.0)
|