orca-sdk 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

orca_sdk/_shared/metrics.py

@@ -8,7 +8,8 @@ IMPORTANT:
 
 """
 
-from dataclasses import dataclass
+import logging
+from dataclasses import dataclass, field
 from typing import Any, Literal, Sequence, TypedDict, cast
 
 import numpy as np
@@ -20,7 +21,9 @@ from numpy.typing import NDArray
 def softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
     shifted = logits - np.max(logits, axis=axis, keepdims=True)
     exps = np.exp(shifted)
-    return exps / np.sum(exps, axis=axis, keepdims=True)
+    sums = np.sum(exps, axis=axis, keepdims=True)
+    # Guard against division by zero (can happen if all logits are -inf or NaN)
+    return exps / np.where(sums > 0, sums, 1.0)
 
 
 # We don't want to depend on transformers just for the eval_pred type in orca_sdk
@@ -240,6 +243,12 @@ class ClassificationMetrics:
     roc_curve: ROCCurve | None = None
     """Receiver operating characteristic curve"""
 
+    confusion_matrix: list[list[int]] | None = None
+    """Confusion matrix where confusion_matrix[i][j] is the count of samples with true label i predicted as label j"""
+
+    warnings: list[str] = field(default_factory=list)
+    """Human-readable warnings about skipped or adjusted metrics"""
+
     def __repr__(self) -> str:
         return (
             "ClassificationMetrics({\n"
@@ -300,7 +309,9 @@ def convert_logits_to_probabilities(logits: NDArray[np.float32]) -> NDArray[np.f
             probabilities = cast(NDArray[np.float32], softmax(logits))
         elif not np.allclose(logits.sum(-1, keepdims=True), 1.0):
             # Rows don't sum to 1: normalize to probabilities
-            probabilities = cast(NDArray[np.float32], logits / logits.sum(-1, keepdims=True))
+            row_sums = logits.sum(-1, keepdims=True)
+            # Guard against division by zero (can happen if all values in a row are 0 or NaN)
+            probabilities = cast(NDArray[np.float32], logits / np.where(row_sums > 0, row_sums, 1.0))
         else:
             # Already normalized probabilities
             probabilities = logits
@@ -317,7 +328,9 @@ def calculate_classification_metrics(
     average: Literal["micro", "macro", "weighted", "binary"] | None = None,
     multi_class: Literal["ovr", "ovo"] = "ovr",
     include_curves: bool = False,
+    include_confusion_matrix: bool = False,
 ) -> ClassificationMetrics:
+    warnings: list[str] = []
     references = np.array(expected_labels)
 
     # Convert to numpy array, handling None values
@@ -338,6 +351,7 @@ def calculate_classification_metrics(
             pr_auc=None,
             pr_curve=None,
             roc_curve=None,
+            confusion_matrix=None,
         )
 
     # Convert logits to probabilities
@@ -347,41 +361,102 @@ def calculate_classification_metrics(
     predictions[np.isnan(probabilities).all(axis=-1)] = -1  # set predictions to -1 for all nan logits
 
     num_classes_references = len(set(references))
-    num_classes_predictions = len(set(predictions))
+    num_classes_predictions = probabilities.shape[1]  # Number of probability columns (model's known classes)
     num_none_predictions = np.isnan(probabilities).all(axis=-1).sum()
-    coverage = 1 - num_none_predictions / len(probabilities)
+    coverage = 1 - (num_none_predictions / len(probabilities) if len(probabilities) > 0 else 0)
+    if num_none_predictions > 0:
+        warnings.append(f"Some predictions were missing (coverage={coverage:.3f}); loss and AUC metrics were skipped.")
 
     if average is None:
         average = "binary" if num_classes_references == 2 and num_none_predictions == 0 else "weighted"
 
     accuracy = sklearn.metrics.accuracy_score(references, predictions)
     f1 = sklearn.metrics.f1_score(references, predictions, average=average)
+
+    # Check for unknown classes early (before log_loss)
+    classes_in_references = np.unique(references)
+    has_unknown_classes = np.max(classes_in_references) >= num_classes_predictions
+    if has_unknown_classes:
+        logging.warning(
+            f"Test labels contain classes not in the model's predictions. "
+            f"Model has {num_classes_predictions} classes (0 - {num_classes_predictions - 1}), "
+            f"but test labels contain class {np.max(classes_in_references)}. "
+            f"ROC AUC and PR AUC cannot be calculated."
+        )
+        warnings.append("y_true contains classes unknown to the model; loss and AUC metrics were skipped.")
+
     # Ensure sklearn sees the full class set corresponding to probability columns
     # to avoid errors when y_true does not contain all classes.
+    # Skip log_loss if there are unknown classes (would cause ValueError)
     loss = (
         sklearn.metrics.log_loss(
             references,
             probabilities,
             labels=list(range(probabilities.shape[1])),
         )
-        if num_none_predictions == 0
+        if num_none_predictions == 0 and not has_unknown_classes
         else None
     )
 
-    if num_classes_references == num_classes_predictions and num_none_predictions == 0:
-        # special case for binary classification: https://github.com/scikit-learn/scikit-learn/issues/20186
-        if num_classes_references == 2:
-            # Use probabilities[:, 1] which is guaranteed to be 2D
-            probabilities_positive = probabilities[:, 1]
-            roc_auc = sklearn.metrics.roc_auc_score(references, probabilities_positive)
-            roc_curve = calculate_roc_curve(references, probabilities_positive) if include_curves else None
-            pr_auc = sklearn.metrics.average_precision_score(references, probabilities_positive)
-            pr_curve = calculate_pr_curve(references, probabilities_positive) if include_curves else None
-        else:
-            roc_auc = sklearn.metrics.roc_auc_score(references, probabilities, multi_class=multi_class)
+    # Calculate ROC AUC with filtering for class mismatch
+    if num_none_predictions == 0:
+        # Check if y_true contains classes not in the model (unknown classes)
+        if has_unknown_classes:
+            # Unknown classes present - can't calculate meaningful ROC AUC
+            logging.warning(
+                "Cannot calculate ROC AUC and PR AUC: test labels contain classes not in the model's predictions."
+            )
+            if "y_true contains classes unknown to the model" not in " ".join(warnings):
+                warnings.append("y_true contains classes unknown to the model; loss and AUC metrics were skipped.")
+            roc_auc = None
+            pr_auc = None
+            pr_curve = None
             roc_curve = None
+        elif len(classes_in_references) < 2:
+            # Need at least 2 classes for ROC AUC
+            logging.warning(
+                f"Cannot calculate ROC AUC and PR AUC: need at least 2 classes, but only {len(classes_in_references)} class(es) found in test labels."
+            )
+            roc_auc = None
             pr_auc = None
             pr_curve = None
+            roc_curve = None
+            warnings.append("ROC AUC requires at least 2 classes; metric was skipped.")
+        else:
+            # Filter probabilities to only classes present in references
+            if len(classes_in_references) < num_classes_predictions:
+                # Subset and renormalize probabilities
+                probabilities_filtered = probabilities[:, classes_in_references]
+                # Safe renormalization: guard against zero denominators
+                row_sums = probabilities_filtered.sum(axis=1, keepdims=True)
+                probabilities_filtered = probabilities_filtered / np.where(row_sums > 0, row_sums, 1.0)
+
+                # Remap references to filtered indices
+                class_mapping = {cls: idx for idx, cls in enumerate(classes_in_references)}
+                references_remapped = np.array([class_mapping[y] for y in references])
+                warnings.append(
+                    f"ROC AUC computed only on classes present in y_true: {classes_in_references.tolist()}."
+                )
+            else:
+                # All classes present, no filtering needed
+                probabilities_filtered = probabilities
+                references_remapped = references
+
+            # special case for binary classification: https://github.com/scikit-learn/scikit-learn/issues/20186
+            if len(classes_in_references) == 2:
+                # Use probabilities[:, 1] which is guaranteed to be 2D
+                probabilities_positive = cast(NDArray[np.float32], probabilities_filtered[:, 1].astype(np.float32))
+                roc_auc = sklearn.metrics.roc_auc_score(references_remapped, probabilities_positive)
+                roc_curve = calculate_roc_curve(references_remapped, probabilities_positive) if include_curves else None
+                pr_auc = sklearn.metrics.average_precision_score(references_remapped, probabilities_positive)
+                pr_curve = calculate_pr_curve(references_remapped, probabilities_positive) if include_curves else None
+            else:
+                roc_auc = sklearn.metrics.roc_auc_score(
+                    references_remapped, probabilities_filtered, multi_class=multi_class
+                )
+                roc_curve = None
+                pr_auc = None
+                pr_curve = None
     else:
         roc_auc = None
         pr_auc = None
@@ -391,6 +466,31 @@ def calculate_classification_metrics(
     # Calculate anomaly score statistics
     anomaly_score_mean, anomaly_score_median, anomaly_score_variance = calculate_anomaly_score_stats(anomaly_scores)
 
+    # Calculate confusion matrix if requested
+    confusion_matrix: list[list[int]] | None = None
+    if include_confusion_matrix:
+        # Get the number of classes from the probabilities shape
+        num_classes = probabilities.shape[1]
+        labels = list(range(num_classes))
+        # Filter out NaN predictions (which are set to -1) before computing confusion matrix
+        valid_mask = predictions != -1
+        num_filtered = (~valid_mask).sum()
+        if num_filtered > 0:
+            warning_msg = (
+                f"Confusion matrix computation: filtered out {num_filtered} samples with NaN predictions "
+                f"({num_filtered}/{len(predictions)} = {num_filtered / len(predictions):.1%})"
+            )
+            logging.warning(warning_msg)
+            warnings.append(warning_msg)
+
+        if np.any(valid_mask):
+            # Compute confusion matrix with explicit labels to ensure consistent shape
+            cm = sklearn.metrics.confusion_matrix(references[valid_mask], predictions[valid_mask], labels=labels)
+        else:
+            # No valid predictions; return an all-zero confusion matrix
+            cm = np.zeros((num_classes, num_classes), dtype=int)
+        confusion_matrix = cast(list[list[int]], cm.tolist())
+
     return ClassificationMetrics(
         coverage=coverage,
         accuracy=float(accuracy),
@@ -403,6 +503,8 @@ def calculate_classification_metrics(
         pr_auc=float(pr_auc) if pr_auc is not None else None,
         pr_curve=pr_curve,
         roc_curve=roc_curve,
+        confusion_matrix=confusion_matrix,
+        warnings=warnings,
     )
 
 
@@ -503,7 +605,7 @@ def calculate_regression_metrics(
     # Filter out NaN values from predictions (expected_scores are already validated to be non-NaN)
     valid_mask = ~np.isnan(predictions)
     num_none_predictions = (~valid_mask).sum()
-    coverage = 1 - num_none_predictions / len(predictions)
+    coverage = 1 - (num_none_predictions / len(predictions) if len(predictions) > 0 else 0)
     if num_none_predictions > 0:
         references = references[valid_mask]
         predictions = predictions[valid_mask]

orca_sdk/_shared/metrics_test.py

@@ -364,3 +364,207 @@ def test_regression_metrics_all_predictions_nan():
     assert metrics.anomaly_score_mean is None
     assert metrics.anomaly_score_median is None
     assert metrics.anomaly_score_variance is None
+
+
+def test_roc_auc_handles_missing_classes_in_y_true():
+    """Test that ROC AUC is calculated with filtering when test set has fewer classes than model predictions."""
+    # Model trained on classes [0, 1, 2], but test set only has [0, 1]
+    y_true = np.array([0, 1, 0, 1])
+    y_score = np.array(
+        [
+            [0.7, 0.2, 0.1],  # Predicts class 0
+            [0.1, 0.8, 0.1],  # Predicts class 1
+            [0.6, 0.3, 0.1],  # Predicts class 0
+            [0.2, 0.7, 0.1],  # Predicts class 1
+        ]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score)
+
+    # Should calculate ROC AUC by filtering to classes [0, 1]
+    assert metrics.roc_auc is not None
+    assert metrics.roc_auc == 1.0
+    assert any("computed only on classes present" in w for w in metrics.warnings)
+    # Other metrics should still work
+    assert metrics.accuracy == 1.0
+    assert metrics.f1_score == 1.0
+    assert metrics.loss is not None
+
+
+def test_roc_auc_with_all_classes_present():
+    """Test that ROC AUC works when all classes are present in test set."""
+    # Model trained on classes [0, 1, 2], test set has all three
+    y_true = np.array([0, 1, 2, 0, 1, 2])
+    y_score = np.array(
+        [
+            [0.9, 0.05, 0.05],  # Predicts class 0
+            [0.1, 0.8, 0.1],  # Predicts class 1
+            [0.1, 0.1, 0.8],  # Predicts class 2
+            [0.7, 0.2, 0.1],  # Predicts class 0
+            [0.2, 0.7, 0.1],  # Predicts class 1
+            [0.1, 0.2, 0.7],  # Predicts class 2
+        ]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score, multi_class="ovr")
+
+    # ROC AUC should be calculated when all classes present
+    assert metrics.roc_auc is not None
+    assert metrics.accuracy == 1.0
+    assert metrics.f1_score == 1.0
+
+
+def test_roc_auc_handles_subset_of_many_classes():
+    """Test ROC AUC where model knows 15 classes, test has 10."""
+    # Simulate the actual error scenario from the bug report
+    num_model_classes = 15
+    num_test_classes = 10
+    num_samples = 50
+
+    # Test set only uses classes 0-9
+    y_true = np.random.randint(0, num_test_classes, size=num_samples)
+
+    # Model produces predictions for all 15 classes
+    y_score = np.random.rand(num_samples, num_model_classes)
+    y_score = y_score / y_score.sum(axis=1, keepdims=True)  # Normalize to probabilities
+
+    metrics = calculate_classification_metrics(y_true, y_score, multi_class="ovr")
+
+    # Should calculate ROC AUC by filtering to classes 0-9
+    assert metrics.roc_auc is not None
+    assert 0.0 <= metrics.roc_auc <= 1.0
+    assert any("computed only on classes present" in w for w in metrics.warnings)
+    # Other metrics should still work
+    assert metrics.accuracy is not None
+    assert metrics.f1_score is not None
+    assert metrics.loss is not None
+
+
+def test_roc_auc_handles_unknown_classes_in_y_true():
+    """Test that metrics handle when y_true contains classes not in y_score."""
+    # Model trained on classes [0, 1, 2], but test set has class 3
+    y_true = np.array([0, 1, 2, 3])
+    y_score = np.array(
+        [
+            [0.7, 0.2, 0.1],
+            [0.1, 0.8, 0.1],
+            [0.1, 0.1, 0.8],
+            [0.3, 0.4, 0.3],  # Unknown class 3
+        ]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score)
+
+    # Should skip ROC AUC and loss when unknown classes present
+    assert metrics.roc_auc is None
+    assert metrics.loss is None  # Loss also skipped to avoid ValueError
+    assert any("unknown" in w for w in metrics.warnings)
+    # Other metrics should still work (they handle extra classes)
+    assert metrics.accuracy is not None
+    assert metrics.f1_score is not None
+
+
+def test_roc_auc_handles_zero_probability_on_present_classes():
+    """Test ROC AUC when a sample has zero probability on all present classes (edge case for renormalization)."""
+    # Model trained on classes [0, 1, 2, 3], test set only has [0, 1, 2]
+    # One sample has ALL probability mass on excluded class 3 (zero on [0, 1, 2])
+    y_true = np.array([0, 1, 2, 0, 1, 2])
+    y_score = np.array(
+        [
+            [0.7, 0.2, 0.08, 0.02],
+            [0.1, 0.8, 0.08, 0.02],
+            [0.1, 0.1, 0.78, 0.02],
+            [0.6, 0.3, 0.08, 0.02],
+            [0.0, 0.0, 0.0, 1.0],  # zero denominator
+            [0.1, 0.1, 0.78, 0.02],
+        ]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score, multi_class="ovr")
+
+    # Should still calculate ROC AUC despite zero-denominator case
+    # The safe renormalization should prevent NaN/inf
+    assert metrics.roc_auc is not None
+    assert not np.isnan(metrics.roc_auc)
+    assert not np.isinf(metrics.roc_auc)
+    assert any("computed only on classes present" in w for w in metrics.warnings)
+    assert metrics.accuracy is not None
+    assert metrics.f1_score is not None
+    assert metrics.loss is not None
+
+
+def test_roc_auc_skipped_for_single_class():
+    """Test that ROC AUC is skipped when only one class is present in y_true."""
+    # Model trained on classes [0, 1, 2], but test set only has class 0
+    y_true = np.array([0, 0, 0, 0])
+    y_score = np.array(
+        [
+            [0.9, 0.05, 0.05],
+            [0.8, 0.1, 0.1],
+            [0.85, 0.1, 0.05],
+            [0.9, 0.05, 0.05],
+        ]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score)
+
+    # ROC AUC requires at least 2 classes
+    assert metrics.roc_auc is None
+    assert metrics.accuracy == 1.0
+    assert metrics.loss is not None
+    assert any("requires at least 2 classes" in w for w in metrics.warnings)
+
+
+# Confusion Matrix Tests
+def test_confusion_matrix_binary_classification():
+    y_true = np.array([0, 1, 1, 0, 1])
+    y_score = np.array([0.1, 0.9, 0.8, 0.3, 0.2])
+
+    metrics = calculate_classification_metrics(y_true, y_score, include_confusion_matrix=True)
+
+    assert metrics.confusion_matrix is not None
+    expected_cm = sklearn.metrics.confusion_matrix(y_true, [0, 1, 1, 0, 0], labels=[0, 1])
+    assert metrics.confusion_matrix == expected_cm.tolist()
+
+
+def test_confusion_matrix_multiclass():
+    y_true = np.array([0, 1, 2, 0, 1, 2])
+    y_score = np.array(
+        [[0.9, 0.1, 0.0], [0.1, 0.8, 0.1], [0.0, 0.1, 0.9], [0.7, 0.2, 0.1], [0.2, 0.7, 0.1], [0.1, 0.1, 0.8]]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score, include_confusion_matrix=True)
+
+    assert metrics.confusion_matrix is not None
+    # All predictions correct
+    assert metrics.confusion_matrix == [[2, 0, 0], [0, 2, 0], [0, 0, 2]]
+
+
+def test_confusion_matrix_with_misclassifications():
+    y_true = np.array([0, 1, 2, 0, 1, 2])
+    y_score = np.array(
+        [[0.9, 0.1, 0.0], [0.1, 0.8, 0.1], [0.1, 0.8, 0.1], [0.2, 0.7, 0.1], [0.2, 0.7, 0.1], [0.1, 0.1, 0.8]]
+    )
+
+    metrics = calculate_classification_metrics(y_true, y_score, include_confusion_matrix=True)
+
+    assert metrics.confusion_matrix is not None
+    # Class 0: 1 correct (index 0), 1 predicted as class 1 (index 3)
+    # Class 1: 2 correct (indices 1, 4)
+    # Class 2: 1 predicted as class 1 (index 2), 1 correct (index 5)
+    assert metrics.confusion_matrix == [[1, 1, 0], [0, 2, 0], [0, 1, 1]]
+
+
+def test_confusion_matrix_handles_nan_logits():
+    logits = np.array([[np.nan, np.nan], [np.nan, np.nan], [0.1, 0.9], [0.2, 0.8]])
+    expected_labels = [0, 1, 0, 1]
+    metrics = calculate_classification_metrics(expected_labels, logits, include_confusion_matrix=True)
+
+    # NaN predictions are set to -1, so they won't match any true label
+    # Only the last 2 predictions are valid: pred=[1, 1], true=[0, 1]
+    assert metrics.confusion_matrix is not None
+    # With NaN handling, predictions become [-1, -1, 1, 1]
+    # Only position 3 is correct (true=1, pred=1)
+    # Position 2 is wrong (true=0, pred=1)
+    assert len(metrics.confusion_matrix) == 2  # 2 classes
+    assert len(metrics.confusion_matrix[0]) == 2