45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
"""Tests for the evaluation metrics on small known fixtures."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from protean_research.evaluation.metrics import (
|
|
classification_metrics,
|
|
confusion_matrix,
|
|
rare_class_metrics,
|
|
)
|
|
|
|
|
|
def test_confusion_matrix_counts():
|
|
matrix = confusion_matrix([0, 0, 1, 1], [0, 1, 1, 1], num_classes=2)
|
|
assert matrix.tolist() == [[1, 1], [0, 2]]
|
|
|
|
|
|
def test_classification_metrics_known_fixture():
|
|
y_true = [0, 0, 1, 1, 2, 2]
|
|
y_pred = [0, 0, 1, 2, 2, 1]
|
|
metrics = classification_metrics(y_true, y_pred, num_classes=3)
|
|
assert metrics["accuracy"] == pytest.approx(4 / 6)
|
|
assert metrics["macro_f1"] == pytest.approx((1.0 + 0.5 + 0.5) / 3)
|
|
per_class = {row["class_id"]: row for row in metrics["per_class"]}
|
|
assert per_class[0]["f1"] == pytest.approx(1.0)
|
|
assert per_class[1]["precision"] == pytest.approx(0.5)
|
|
assert per_class[2]["recall"] == pytest.approx(0.5)
|
|
|
|
|
|
def test_rare_class_metrics_restricts_to_given_classes():
|
|
y_true = [0, 0, 1, 1, 2, 2]
|
|
y_pred = [0, 0, 1, 2, 2, 1]
|
|
rare = rare_class_metrics(y_true, y_pred, num_classes=3, rare_class_ids=[1, 2])
|
|
assert rare["rare_class_macro_f1"] == pytest.approx(0.5)
|
|
assert rare["rare_class_recall"] == pytest.approx(0.5)
|
|
assert rare["rare_support"] == 4
|
|
|
|
|
|
def test_rare_class_metrics_skips_absent_classes():
|
|
rare = rare_class_metrics([0, 0], [0, 0], num_classes=3, rare_class_ids=[2])
|
|
assert rare["rare_class_macro_f1"] == 0.0
|
|
assert rare["rare_support"] == 0
|