| 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) |
|
|
