metrics.py

"""
Classification metrics for MAVNet's two heads.

The original paper mixed regression metrics (EVA, RMSE) with classification
metrics (F-measure, accuracy) for what is fundamentally a discrete action
space. Since both heads here are classification (softmax movement + binary
junction), we only report accuracy/precision/recall/F1 - no EVA/RMSE.
"""

from dataclasses import dataclass

import torch


@dataclass
class ClassificationMetrics:
    accuracy: float
    precision: float
    recall: float
    f1: float


def binary_classification_metrics(preds: torch.Tensor, targets: torch.Tensor) -> ClassificationMetrics:
    """preds, targets: 1D tensors of 0/1 (already thresholded)."""
    preds = preds.bool()
    targets = targets.bool()

    tp = (preds & targets).sum().item()
    fp = (preds & ~targets).sum().item()
    fn = (~preds & targets).sum().item()
    tn = (~preds & ~targets).sum().item()

    accuracy = (tp + tn) / max(1, tp + fp + fn + tn)
    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    f1 = 2 * precision * recall / max(1e-8, precision + recall)

    return ClassificationMetrics(accuracy, precision, recall, f1)


def multiclass_accuracy(logits: torch.Tensor, targets: torch.Tensor) -> float:
    preds = logits.argmax(dim=-1)
    return (preds == targets).float().mean().item()


def evaluate_batch(outputs: dict, movement_labels: torch.Tensor, junction_labels: torch.Tensor) -> dict:
    movement_acc = multiclass_accuracy(outputs["movement_logits"], movement_labels)

    junction_preds = (torch.sigmoid(outputs["junction_logit"]) > 0.5).long()
    junction_metrics = binary_classification_metrics(junction_preds, junction_labels.long())

    return {
        "movement_accuracy": movement_acc,
        "junction_accuracy": junction_metrics.accuracy,
        "junction_precision": junction_metrics.precision,
        "junction_recall": junction_metrics.recall,
        "junction_f1": junction_metrics.f1,
    }


if __name__ == "__main__":
    logits = torch.tensor([[2.0, 0.1, 0.0, 0.0], [0.0, 0.0, 3.0, 0.0]])
    targets = torch.tensor([0, 2])
    print("multiclass_accuracy:", multiclass_accuracy(logits, targets))

    preds = torch.tensor([1, 0, 1, 1])
    binary_targets = torch.tensor([1, 0, 0, 1])
    print(binary_classification_metrics(preds, binary_targets))