wip

Author: czaloom

src/valor_lite/cache/compute.py

@@ -1,7 +1,7 @@
 import heapq
 import tempfile
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Any
 
 import pyarrow as pa
 
@@ -16,12 +16,13 @@ def _merge(
     batch_size: int,
     sorting: list[tuple[str, str]],
     columns: list[str] | None = None,
-    table_sort_override: Callable[[pa.Table], pa.Table] | None = None,
+    sort_override: Callable[[pa.Table], pa.Table] | None = None,
+    merge_override: Callable[[pa.RecordBatch, Any], tuple[pa.RecordBatch | None, Any]] | None = None,
 ):
     """Merge locally sorted cache fragments."""
     for tbl in source.iterate_tables(columns=columns):
-        if table_sort_override is not None:
-            sorted_tbl = table_sort_override(tbl)
+        if sort_override is not None:
+            sorted_tbl = sort_override(tbl)
         else:
             sorted_tbl = tbl.sort_by(sorting)
         intermediate_sink.write_table(sorted_tbl)
@@ -57,12 +58,18 @@ def create_sort_key(
         if batches[batch_idx] is not None and len(batches[batch_idx]) > 0:
             heapq.heappush(heap, create_sort_key(batches, batch_idx, 0))
 
+    prev_state = None
     while heap:
         row = heapq.heappop(heap)
         batch_idx = row[-2]
         row_idx = row[-1]
         row_table = batches[batch_idx].slice(row_idx, 1)
-        sink.write_batch(row_table)
+        if merge_override is not None:
+            batch = merge_override(row_table, prev_state)
+            if batch is not None:
+                sink.write_batch(row_table)    
+        else:
+            sink.write_batch(row_table)
         row_idx += 1
         if row_idx < len(batches[batch_idx]):
             heapq.heappush(

src/valor_lite/classification/computation.py

@@ -14,8 +14,6 @@ def compute_rocauc(
     n_labels: int,
     accumulated_tp: NDArray[np.uint64],
     accumulated_fp: NDArray[np.uint64],
-    prev_fpr: NDArray[np.float64],
-    prev_tpr: NDArray[np.float64],
 ) -> tuple[NDArray[np.float64], NDArray[np.uint64], NDArray[np.uint64]]:
     """
     Compute ROCAUC.
@@ -54,98 +52,81 @@ def compute_rocauc(
     positive_count = gt_count_per_label
     negative_count = pd_count_per_label - gt_count_per_label
 
+    print()
     for label_idx in range(n_labels):
         mask_pds = pd_labels == label_idx
         n_masked_pds = mask_pds.sum()
         if pd_count_per_label[label_idx] == 0 or n_masked_pds == 0:
             continue
 
         true_positives = mask_matching_labels[mask_pds]
-        false_positives = ~mask_matching_labels[mask_pds]
         tp_scores = scores[mask_pds]
 
-        cumulative_fp = np.cumsum(false_positives) + accumulated_fp[label_idx]
-        cumulative_tp = np.cumsum(true_positives) + accumulated_tp[label_idx]
+        distinct_score_indices = np.where(np.diff(tp_scores))[0]
+        indices = np.r_[distinct_score_indices, n_masked_pds - 1]
+        cumulative_tp = np.cumsum(true_positives, dtype=np.uint64)[indices]
+        cumulative_fp = indices + 1 - cumulative_tp
+
+        cumulative_tp += accumulated_tp[label_idx]
+        cumulative_fp += accumulated_fp[label_idx]
+
+        cumulative_tp = np.concatenate([accumulated_tp[label_idx:label_idx+1], cumulative_tp])
+        cumulative_fp = np.concatenate([accumulated_fp[label_idx:label_idx+1], cumulative_fp])
 
-        accumulated_fp[label_idx] = cumulative_fp[-1]
         accumulated_tp[label_idx] = cumulative_tp[-1]
+        accumulated_fp[label_idx] = cumulative_fp[-1]
 
-        fpr = np.zeros(n_masked_pds, dtype=np.float64)
+        fpr = np.zeros_like(cumulative_fp, dtype=np.float64)
         np.divide(
             cumulative_fp,
             negative_count[label_idx],
             where=negative_count[label_idx] > 0,
             out=fpr,
         )
-        tpr = np.zeros(n_masked_pds, dtype=np.float64)
+        tpr = np.zeros_like(cumulative_tp, dtype=np.float64)
         np.divide(
             cumulative_tp,
             positive_count[label_idx],
             where=positive_count[label_idx] > 0,
             out=tpr,
         )
 
-        if prev_fpr[label_idx] > -0.5 and prev_tpr[label_idx] > -0.5:
-            fpr = np.concatenate([prev_fpr[label_idx:label_idx+1], fpr])
-            tpr = np.concatenate([prev_tpr[label_idx:label_idx+1], tpr])
-
-        # sort by -tpr, -score
-        indices = np.lexsort((-tpr, -tp_scores))
-        fpr = fpr[indices]
-        tpr = tpr[indices]
-
-        sfpr = fpr.copy()
-        stpr = tpr.copy()
+        # # sort by -tpr, -score
+        # indices = np.lexsort((-tpr, -tp_scores))
+        # fpr = fpr[indices]
+        # tpr = tpr[indices]
 
         # running max of tpr
         np.maximum.accumulate(tpr, out=tpr)
 
-
-        prev_fpr[label_idx] = fpr[-1]
-        prev_tpr[label_idx] = tpr[-1]
-
         # compute rocauc
         rocauc[label_idx] += npc.trapezoid(x=fpr, y=tpr, axis=0)
 
-        print()
-        print(label_idx, rocauc[label_idx])
-        print("====")
-        print(
-            f"{'FP':4}",
-            f"{'TP':4}",
-            f"{'CFP':4}",
-            f"{'CTP':4}",
-            f"{'FPR':4}",
-            f"{'TPR':4}",
-            f"{'SFPR':4}",
-            f"{'STPR':4}",
-            f"{'SCO':4}",
-        )
-        for f, t, af, at, fr, tr, sf, st, s in zip(
-            false_positives,
-            true_positives,
-            cumulative_fp,
-            cumulative_tp,
-            fpr, 
-            tpr, 
-            sfpr,
-            stpr,
-            tp_scores,
-        ):
+        if label_idx == 3:
+            print(rocauc[label_idx])
             print(
-                f"{f:.2f}", 
-                f"{t:.2f}",
-                f"{af:.2f}", 
-                f"{at:.2f}", 
-                f"{fr:.2f}", 
-                f"{tr:.2f}", 
-                f"{sf:.2f}", 
-                f"{st:.2f}", 
-                f"{s:.2f}",
+                f"{'CFP':4}",
+                f"{'CTP':4}",
+                f"{'FPR':4}",
+                f"{'TPR':4}",
+                # f"{'SCO':4}",
             )
+            for af, at, fr, tr in zip(
+                cumulative_fp,
+                cumulative_tp,
+                fpr, 
+                tpr, 
+                # tp_scores,
+            ):
+                print(
+                    f"{af:.2f}", 
+                    f"{at:.2f}", 
+                    f"{fr:.2f}", 
+                    f"{tr:.2f}", 
+                    # f"{s:.2f}",
+                )
 
-
-    return rocauc, accumulated_fp, accumulated_tp, prev_fpr, prev_tpr
+    return rocauc, accumulated_fp, accumulated_tp
 
 
 def compute_counts(

src/valor_lite/classification/evaluator.py

@@ -294,17 +294,14 @@ def compute_rocauc(self) -> dict[MetricType, list[Metric]]:
             #         tpr[plabel]
             #     else:
             #         cumulative_fp[plabel] += 1
-            print("loop", loopid)
-            rocauc, accumulated_fp, accumulated_tp, prev_fpr, prev_tpr = compute_rocauc(
+            rocauc, accumulated_fp, accumulated_tp = compute_rocauc(
                 rocauc=rocauc,
                 array=array,
                 gt_count_per_label=self._label_counts[:, 0],
                 pd_count_per_label=self._label_counts[:, 1],
                 n_labels=self.info.number_of_labels,
                 accumulated_fp=accumulated_fp,
                 accumulated_tp=accumulated_tp,
-                prev_fpr=prev_fpr,
-                prev_tpr=prev_tpr,
             )
 
         mean_rocauc = rocauc.mean()

src/valor_lite/classification/loader.py

@@ -1,6 +1,7 @@
 import json
 from pathlib import Path
 
+import pyarrow as pa
 import numpy as np
 from pyarrow import DataType
 from tqdm import tqdm
@@ -229,14 +230,24 @@ def finalize(
 
         # post-process into sorted writer
         reader = self._writer.to_reader()
+
+        n_labels = len(self._index_to_label)
+
+        def accumulate(batch: pa.RecordBatch, prev: np.ndarray | None) -> pa.RecordBatch:
+            pd_label_id = batch["pd_label_id"].as_py()
+            matched = batch["match"][0].as_py()
+            if prev is None:
+                prev = np.zeros(n_labels, dtype=np.uint64)
+                return None, ()
+
         sort(
             source=reader,
             sink=self._rocauc_writer,
             batch_size=batch_size,
             sorting=[
                 ("score", "descending"),
-                ("match", "descending"),
-                ("pd_label_id", "ascending"),
+                # ("match", "descending"),
+                # ("pd_label_id", "ascending"),
             ],
             columns=[
                 "pd_label_id",

tests/classification/conftest.py

@@ -8,7 +8,7 @@
 @pytest.fixture(
     params=[
         ("persistent", 10_000, 100_000),
-        ("persistent", 1, 1),
+        ("persistent", 2, 2),
         ("memory", 10_000, 0),
         ("memory", 1, 0),
     ],

tests/classification/test_rocauc.py

@@ -60,9 +60,11 @@ def test_rocauc_with_color_example(
 
     loader.add_data(classifications_color_example)
     evaluator = loader.finalize()
+    print(evaluator._index_to_label)
 
     metrics = evaluator.compute_rocauc()
 
+
     # test ROCAUC
     actual_metrics = [m.to_dict() for m in metrics[MetricType.ROCAUC]]
     expected_metrics = [