valor-lite 0.36.6__py3-none-any.whl → 0.37.5__py3-none-any.whl

valor_lite/classification/computation.py

@@ -1,4 +1,4 @@
-from enum import IntFlag, auto
+from enum import IntFlag
 
 import numpy as np
 from numpy.typing import NDArray
@@ -6,256 +6,141 @@ from numpy.typing import NDArray
 import valor_lite.classification.numpy_compatibility as npc
 
 
-def compute_label_metadata(
-    ids: NDArray[np.int32],
+def compute_rocauc(
+    rocauc: NDArray[np.float64],
+    array: NDArray[np.float64],
+    gt_count_per_label: NDArray[np.uint64],
+    pd_count_per_label: NDArray[np.uint64],
     n_labels: int,
-) -> NDArray[np.int32]:
+    prev: NDArray[np.uint64],
+) -> tuple[NDArray[np.float64], NDArray[np.uint64]]:
     """
-    Computes label metadata returning a count of annotations per label.
+    Compute ROCAUC.
 
     Parameters
     ----------
-    detailed_pairs : NDArray[np.int32]
-        Detailed annotation pairings with shape (n_pairs, 3).
-            Index 0 - Datum Index
-            Index 1 - GroundTruth Label Index
-            Index 2 - Prediction Label Index
+    rocauc : NDArray[np.float64]
+        The running ROCAUC.
+    array : NDArray[np.float64]
+        An sorted array of ROCAUC intermediate values with shape (n_pairs, 3).
+            Index 0 - Prediction Label Index
+            Index 1 - Cumulative FP
+            Index 2 - Cumulative TP
+    gt_count_per_label : NDArray[np.uint64]
+        The number of ground truth occurences per label.
+    pd_count_per_label : NDArray[np.uint64]
+        The number of prediction occurences per label.
     n_labels : int
-        The total number of unique labels.
+        The total number of unqiue labels.
+    prev : NDArray[np.uint64]
+        The previous cumulative sum for FP's and TP's. Used as intermediate in chunking operations.
 
     Returns
     -------
-    NDArray[np.int32]
-        The label metadata array with shape (n_labels, 2).
-            Index 0 - Ground truth label count
-            Index 1 - Prediction label count
-    """
-    label_metadata = np.zeros((n_labels, 2), dtype=np.int32)
-    ground_truth_pairs = ids[:, (0, 1)]
-    ground_truth_pairs = ground_truth_pairs[ground_truth_pairs[:, 1] >= 0]
-    unique_pairs = np.unique(ground_truth_pairs, axis=0)
-    label_indices, unique_counts = np.unique(
-        unique_pairs[:, 1], return_counts=True
-    )
-    label_metadata[label_indices.astype(np.int32), 0] = unique_counts
-
-    prediction_pairs = ids[:, (0, 2)]
-    prediction_pairs = prediction_pairs[prediction_pairs[:, 1] >= 0]
-    unique_pairs = np.unique(prediction_pairs, axis=0)
-    label_indices, unique_counts = np.unique(
-        unique_pairs[:, 1], return_counts=True
-    )
-    label_metadata[label_indices.astype(np.int32), 1] = unique_counts
-
-    return label_metadata
-
-
-def filter_cache(
-    detailed_pairs: NDArray[np.float64],
-    datum_mask: NDArray[np.bool_],
-    valid_label_indices: NDArray[np.int32] | None,
-    n_labels: int,
-) -> tuple[NDArray[np.float64], NDArray[np.int32]]:
-    # filter by datum
-    detailed_pairs = detailed_pairs[datum_mask].copy()
-
-    n_rows = detailed_pairs.shape[0]
-    mask_invalid_groundtruths = np.zeros(n_rows, dtype=np.bool_)
-    mask_invalid_predictions = np.zeros_like(mask_invalid_groundtruths)
-
-    # filter labels
-    if valid_label_indices is not None:
-        mask_invalid_groundtruths[
-            ~np.isin(detailed_pairs[:, 1], valid_label_indices)
-        ] = True
-        mask_invalid_predictions[
-            ~np.isin(detailed_pairs[:, 2], valid_label_indices)
-        ] = True
-
-    # filter cache
-    if mask_invalid_groundtruths.any():
-        invalid_groundtruth_indices = np.where(mask_invalid_groundtruths)[0]
-        detailed_pairs[invalid_groundtruth_indices[:, None], 1] = np.array(
-            [[-1.0]]
-        )
-
-    if mask_invalid_predictions.any():
-        invalid_prediction_indices = np.where(mask_invalid_predictions)[0]
-        detailed_pairs[
-            invalid_prediction_indices[:, None], (2, 3, 4)
-        ] = np.array([[-1.0, -1.0, -1.0]])
-
-    # filter null pairs
-    mask_null_pairs = np.all(
-        np.isclose(
-            detailed_pairs[:, 1:5],
-            np.array([-1.0, -1.0, -1.0, -1.0]),
-        ),
-        axis=1,
-    )
-    detailed_pairs = detailed_pairs[~mask_null_pairs]
-
-    detailed_pairs = np.unique(detailed_pairs, axis=0)
-    indices = np.lexsort(
-        (
-            detailed_pairs[:, 1],  # ground truth
-            detailed_pairs[:, 2],  # prediction
-            -detailed_pairs[:, 3],  # score
-        )
-    )
-    detailed_pairs = detailed_pairs[indices]
-    label_metadata = compute_label_metadata(
-        ids=detailed_pairs[:, :3].astype(np.int32),
-        n_labels=n_labels,
-    )
-    return detailed_pairs, label_metadata
-
-
-def _compute_rocauc(
-    data: NDArray[np.float64],
-    label_metadata: NDArray[np.int32],
-    n_datums: int,
-    n_labels: int,
-    mask_matching_labels: NDArray[np.bool_],
-    pd_labels: NDArray[np.int32],
-) -> tuple[NDArray[np.float64], float]:
-    """
-    Compute ROCAUC and mean ROCAUC.
+    NDArray[np.float64]
+        ROCAUC.
+    NDArray[np.uint64]
+        The previous cumulative sum for FP's. Used as intermediate in chunking operations.
+    NDArray[np.uint64]
+        The previous cumulative sum for TP's. Used as intermediate in chunking operations.
     """
-    positive_count = label_metadata[:, 0]
-    negative_count = label_metadata[:, 1] - label_metadata[:, 0]
+    pd_labels = array[:, 0]
+    cumulative_fp = array[:, 1]
+    cumulative_tp = array[:, 2]
 
-    true_positives = np.zeros((n_labels, n_datums), dtype=np.int32)
-    false_positives = np.zeros_like(true_positives)
-    scores = np.zeros_like(true_positives, dtype=np.float64)
+    positive_count = gt_count_per_label
+    negative_count = pd_count_per_label - gt_count_per_label
 
     for label_idx in range(n_labels):
-        if label_metadata[label_idx, 1] == 0:
-            continue
-
         mask_pds = pd_labels == label_idx
-        true_positives[label_idx] = mask_matching_labels[mask_pds]
-        false_positives[label_idx] = ~mask_matching_labels[mask_pds]
-        scores[label_idx] = data[mask_pds, 3]
-
-    cumulative_fp = np.cumsum(false_positives, axis=1)
-    cumulative_tp = np.cumsum(true_positives, axis=1)
+        n_masked_pds = mask_pds.sum()
+        if pd_count_per_label[label_idx] == 0 or n_masked_pds == 0:
+            continue
 
-    fpr = np.zeros_like(true_positives, dtype=np.float64)
-    np.divide(
-        cumulative_fp,
-        negative_count[:, np.newaxis],
-        where=negative_count[:, np.newaxis] > 1e-9,
-        out=fpr,
-    )
-    tpr = np.zeros_like(true_positives, dtype=np.float64)
-    np.divide(
-        cumulative_tp,
-        positive_count[:, np.newaxis],
-        where=positive_count[:, np.newaxis] > 1e-9,
-        out=tpr,
-    )
+        fps = cumulative_fp[mask_pds]
+        tps = cumulative_tp[mask_pds]
+        if prev[label_idx, 0] > 0 or prev[label_idx, 1] > 0:
+            fps = np.r_[prev[label_idx, 0], fps]
+            tps = np.r_[prev[label_idx, 1], tps]
 
-    # sort by -tpr, -score
-    indices = np.lexsort((-tpr, -scores), axis=1)
-    fpr = np.take_along_axis(fpr, indices, axis=1)
-    tpr = np.take_along_axis(tpr, indices, axis=1)
+        prev[label_idx, 0] = fps[-1]
+        prev[label_idx, 1] = tps[-1]
 
-    # running max of tpr
-    np.maximum.accumulate(tpr, axis=1, out=tpr)
+        if fps.size == 1:
+            continue
 
-    # compute rocauc
-    rocauc = npc.trapezoid(x=fpr, y=tpr, axis=1)
+        fpr = np.zeros_like(fps, dtype=np.float64)
+        np.divide(
+            fps,
+            negative_count[label_idx],
+            where=negative_count[label_idx] > 0,
+            out=fpr,
+        )
+        tpr = np.zeros_like(tps, dtype=np.float64)
+        np.divide(
+            tps,
+            positive_count[label_idx],
+            where=positive_count[label_idx] > 0,
+            out=tpr,
+        )
 
-    # compute mean rocauc
-    mean_rocauc = rocauc.mean()
+        # compute rocauc
+        rocauc[label_idx] += npc.trapezoid(x=fpr, y=tpr, axis=0)
 
-    return rocauc, mean_rocauc  # type: ignore[reportReturnType]
+    return rocauc, prev
 
 
-def compute_precision_recall_rocauc(
-    detailed_pairs: NDArray[np.float64],
-    label_metadata: NDArray[np.int32],
+def compute_counts(
+    ids: NDArray[np.int64],
+    scores: NDArray[np.float64],
+    winners: NDArray[np.bool_],
     score_thresholds: NDArray[np.float64],
     hardmax: bool,
-    n_datums: int,
-) -> tuple[
-    NDArray[np.int32],
-    NDArray[np.float64],
-    NDArray[np.float64],
-    NDArray[np.float64],
-    NDArray[np.float64],
-    NDArray[np.float64],
-    float,
-]:
+    n_labels: int,
+) -> NDArray[np.uint64]:
     """
-    Computes classification metrics.
+    Computes counts of TP, FP and FN's per label.
 
     Parameters
     ----------
-    detailed_pairs : NDArray[np.float64]
-        A sorted array of classification pairs with shape (n_pairs, 5).
+    ids : NDArray[np.int64]
+        A sorted array of classification pairs with shape (n_pairs, 3).
             Index 0 - Datum Index
             Index 1 - GroundTruth Label Index
             Index 2 - Prediction Label Index
-            Index 3 - Score
-            Index 4 - Hard-Max Score
-    label_metadata : NDArray[np.int32]
-        An array containing metadata related to labels with shape (n_labels, 2).
-            Index 0 - GroundTruth Label Count
-            Index 1 - Prediction Label Count
+    scores : NDArray[np.float64]
+        A sorted array of classification scores with shape (n_pairs,).
+    winner : NDArray[np.bool_]
+        Marks predictions with highest score over a datum.
     score_thresholds : NDArray[np.float64]
         A 1-D array contains score thresholds to compute metrics over.
     hardmax : bool
         Option to only allow a single positive prediction.
-    n_datums : int
-        The number of datums being operated over.
+    n_labels : int
+        The total number of unqiue labels.
 
     Returns
     -------
     NDArray[np.int32]
         TP, FP, FN, TN counts.
-    NDArray[np.float64]
-        Precision.
-    NDArray[np.float64]
-        Recall.
-    NDArray[np.float64]
-        Accuracy
-    NDArray[np.float64]
-        F1 Score
-    NDArray[np.float64]
-        ROCAUC.
-    float
-        mROCAUC.
     """
-
-    n_labels = label_metadata.shape[0]
     n_scores = score_thresholds.shape[0]
+    counts = np.zeros((n_scores, n_labels, 4), dtype=np.uint64)
+    if ids.size == 0:
+        return counts
 
-    pd_labels = detailed_pairs[:, 2].astype(int)
+    gt_labels = ids[:, 1]
+    pd_labels = ids[:, 2]
 
-    mask_matching_labels = np.isclose(
-        detailed_pairs[:, 1], detailed_pairs[:, 2]
-    )
-    mask_score_nonzero = ~np.isclose(detailed_pairs[:, 3], 0.0)
-    mask_hardmax = detailed_pairs[:, 4] > 0.5
-
-    # calculate ROCAUC
-    rocauc, mean_rocauc = _compute_rocauc(
-        data=detailed_pairs,
-        label_metadata=label_metadata,
-        n_datums=n_datums,
-        n_labels=n_labels,
-        mask_matching_labels=mask_matching_labels,
-        pd_labels=pd_labels,
-    )
+    mask_matching_labels = np.isclose(gt_labels, pd_labels)
+    mask_score_nonzero = ~np.isclose(scores, 0.0)
+    mask_hardmax = winners > 0.5
+    mask_valid_gts = gt_labels >= 0
+    mask_valid_pds = pd_labels >= 0
 
     # calculate metrics at various score thresholds
-    counts = np.zeros((n_scores, n_labels, 4), dtype=np.int32)
     for score_idx in range(n_scores):
-        mask_score_threshold = (
-            detailed_pairs[:, 3] >= score_thresholds[score_idx]
-        )
+        mask_score_threshold = scores >= score_thresholds[score_idx]
         mask_score = mask_score_nonzero & mask_score_threshold
 
         if hardmax:
@@ -266,8 +151,11 @@ def compute_precision_recall_rocauc(
         mask_fn = (mask_matching_labels & ~mask_score) | mask_fp
         mask_tn = ~mask_matching_labels & ~mask_score
 
-        fn = np.unique(detailed_pairs[mask_fn][:, [0, 1]].astype(int), axis=0)
-        tn = np.unique(detailed_pairs[mask_tn][:, [0, 2]].astype(int), axis=0)
+        mask_fn &= mask_valid_gts
+        mask_fp &= mask_valid_pds
+
+        fn = np.unique(ids[mask_fn][:, [0, 1]].astype(int), axis=0)
+        tn = np.unique(ids[mask_tn][:, [0, 2]].astype(int), axis=0)
 
         counts[score_idx, :, 0] = np.bincount(
             pd_labels[mask_tp], minlength=n_labels
@@ -278,29 +166,45 @@ def compute_precision_recall_rocauc(
         counts[score_idx, :, 2] = np.bincount(fn[:, 1], minlength=n_labels)
         counts[score_idx, :, 3] = np.bincount(tn[:, 1], minlength=n_labels)
 
-    recall = np.zeros((n_scores, n_labels), dtype=np.float64)
-    np.divide(
-        counts[:, :, 0],
-        (counts[:, :, 0] + counts[:, :, 2]),
-        where=(counts[:, :, 0] + counts[:, :, 2]) > 1e-9,
-        out=recall,
-    )
+    return counts
+
 
-    precision = np.zeros_like(recall)
+def compute_precision(counts: NDArray[np.uint64]) -> NDArray[np.float64]:
+    """
+    Compute precision metric using result of compute_counts.
+    """
+    n_scores, n_labels, _ = counts.shape
+    precision = np.zeros((n_scores, n_labels), dtype=np.float64)
     np.divide(
         counts[:, :, 0],
         (counts[:, :, 0] + counts[:, :, 1]),
-        where=(counts[:, :, 0] + counts[:, :, 1]) > 1e-9,
+        where=(counts[:, :, 0] + counts[:, :, 1]) > 0,
         out=precision,
     )
+    return precision
 
-    accuracy = np.zeros(n_scores, dtype=np.float64)
+
+def compute_recall(counts: NDArray[np.uint64]) -> NDArray[np.float64]:
+    """
+    Compute recall metric using result of compute_counts.
+    """
+    n_scores, n_labels, _ = counts.shape
+    recall = np.zeros((n_scores, n_labels), dtype=np.float64)
     np.divide(
-        counts[:, :, 0].sum(axis=1),
-        float(n_datums),
-        out=accuracy,
+        counts[:, :, 0],
+        (counts[:, :, 0] + counts[:, :, 2]),
+        where=(counts[:, :, 0] + counts[:, :, 2]) > 0,
+        out=recall,
     )
+    return recall
+
 
+def compute_f1_score(
+    precision: NDArray[np.float64], recall: NDArray[np.float64]
+) -> NDArray[np.float64]:
+    """
+    Compute f1 metric using result of compute_precision and compute_recall.
+    """
     f1_score = np.zeros_like(recall)
     np.divide(
         (2 * precision * recall),
@@ -308,68 +212,86 @@ def compute_precision_recall_rocauc(
         where=(precision + recall) > 1e-9,
         out=f1_score,
     )
+    return f1_score
 
-    return (
-        counts,
-        precision,
-        recall,
-        accuracy,
-        f1_score,
-        rocauc,
-        mean_rocauc,
+
+def compute_accuracy(
+    counts: NDArray[np.uint64], n_datums: int
+) -> NDArray[np.float64]:
+    """
+    Compute accuracy metric using result of compute_counts.
+    """
+    n_scores, _, _ = counts.shape
+    accuracy = np.zeros(n_scores, dtype=np.float64)
+    if n_datums == 0:
+        return accuracy
+    np.divide(
+        counts[:, :, 0].sum(axis=1),
+        n_datums,
+        out=accuracy,
     )
+    return accuracy
 
 
 class PairClassification(IntFlag):
-    TP = auto()
-    FP_FN_MISCLF = auto()
-    FN_UNMATCHED = auto()
+    TP = 1 << 0
+    FP_FN_MISCLF = 1 << 1
+    FN_UNMATCHED = 1 << 2
 
 
-def compute_confusion_matrix(
-    detailed_pairs: NDArray[np.float64],
+def compute_pair_classifications(
+    ids: NDArray[np.int64],
+    scores: NDArray[np.float64],
+    winners: NDArray[np.bool_],
     score_thresholds: NDArray[np.float64],
     hardmax: bool,
-) -> NDArray[np.uint8]:
+) -> tuple[NDArray[np.bool_], NDArray[np.bool_], NDArray[np.bool_]]:
     """
-    Compute detailed confusion matrix.
+    Classifiy ID pairs as TP, FP or FN.
 
     Parameters
     ----------
-    detailed_pairs : NDArray[np.float64]
-        A 2-D sorted array summarizing the IOU calculations of one or more pairs with shape (n_pairs, 5).
+    ids : NDArray[np.int64]
+        A sorted array of classification pairs with shape (n_pairs, 3).
             Index 0 - Datum Index
             Index 1 - GroundTruth Label Index
             Index 2 - Prediction Label Index
-            Index 3 - Score
-            Index 4 - Hard Max Score
-    iou_thresholds : NDArray[np.float64]
-        A 1-D array containing IOU thresholds.
+    scores : NDArray[np.float64]
+        A sorted array of classification scores with shape (n_pairs,).
+    winner : NDArray[np.bool_]
+        Marks predictions with highest score over a datum.
     score_thresholds : NDArray[np.float64]
         A 1-D array containing score thresholds.
+    hardmax : bool
+        Option to only allow a single positive prediction.
 
     Returns
     -------
-    NDArray[uint8]
-        Row-wise classification of pairs.
+    NDArray[bool]
+        True-positive mask.
+    NDArray[bool]
+        Misclassification FP, FN mask.
+    NDArray[bool]
+        Unmatched FN mask.
     """
-    n_pairs = detailed_pairs.shape[0]
+    n_pairs = ids.shape[0]
     n_scores = score_thresholds.shape[0]
 
+    gt_labels = ids[:, 1]
+    pd_labels = ids[:, 2]
+    groundtruths = ids[:, [0, 1]]
+
     pair_classifications = np.zeros(
         (n_scores, n_pairs),
         dtype=np.uint8,
     )
 
-    mask_label_match = np.isclose(detailed_pairs[:, 1], detailed_pairs[:, 2])
-    mask_score = detailed_pairs[:, 3] > 1e-9
-
-    groundtruths = detailed_pairs[:, [0, 1]].astype(int)
-
+    mask_label_match = np.isclose(gt_labels, pd_labels)
+    mask_score = scores > 1e-9
     for score_idx in range(n_scores):
-        mask_score &= detailed_pairs[:, 3] >= score_thresholds[score_idx]
+        mask_score &= scores >= score_thresholds[score_idx]
         if hardmax:
-            mask_score &= detailed_pairs[:, 4] > 0.5
+            mask_score &= winners
 
         mask_true_positives = mask_label_match & mask_score
         mask_misclassifications = ~mask_label_match & mask_score
@@ -393,4 +315,64 @@ def compute_confusion_matrix(
             score_idx, mask_unmatched_groundtruths
         ] |= np.uint8(PairClassification.FN_UNMATCHED)
 
-    return pair_classifications
+    mask_tp = np.bitwise_and(pair_classifications, PairClassification.TP) > 0
+    mask_fp_fn_misclf = (
+        np.bitwise_and(pair_classifications, PairClassification.FP_FN_MISCLF)
+        > 0
+    )
+    mask_fn_unmatched = (
+        np.bitwise_and(pair_classifications, PairClassification.FN_UNMATCHED)
+        > 0
+    )
+
+    return (
+        mask_tp,
+        mask_fp_fn_misclf,
+        mask_fn_unmatched,
+    )
+
+
+def compute_confusion_matrix(
+    ids: NDArray[np.int64],
+    mask_tp: NDArray[np.bool_],
+    mask_fp_fn_misclf: NDArray[np.bool_],
+    mask_fn_unmatched: NDArray[np.bool_],
+    score_thresholds: NDArray[np.float64],
+    n_labels: int,
+):
+    """
+    Compute confusion matrix using output of compute_pair_classifications.
+    """
+    n_scores = score_thresholds.size
+
+    # initialize arrays
+    confusion_matrices = np.zeros(
+        (n_scores, n_labels, n_labels), dtype=np.uint64
+    )
+    unmatched_groundtruths = np.zeros((n_scores, n_labels), dtype=np.uint64)
+
+    mask_matched = mask_tp | mask_fp_fn_misclf
+    for score_idx in range(n_scores):
+        # matched annotations
+        unique_pairs = np.unique(
+            ids[np.ix_(mask_matched[score_idx], (0, 1, 2))],  # type: ignore - numpy ix_ typing
+            axis=0,
+        )
+        unique_labels, unique_label_counts = np.unique(
+            unique_pairs[:, (1, 2)], axis=0, return_counts=True
+        )
+        confusion_matrices[
+            score_idx, unique_labels[:, 0], unique_labels[:, 1]
+        ] = unique_label_counts
+
+        # unmatched groundtruths
+        unique_pairs = np.unique(
+            ids[np.ix_(mask_fn_unmatched[score_idx], (0, 1))],  # type: ignore - numpy ix_ typing
+            axis=0,
+        )
+        unique_labels, unique_label_counts = np.unique(
+            unique_pairs[:, 1], return_counts=True
+        )
+        unmatched_groundtruths[score_idx, unique_labels] = unique_label_counts
+
+    return confusion_matrices, unmatched_groundtruths