Repository: wandb/wandb

Issue:
[CLI]: ROC Curve clips remaining points
### Describe the bug

`wandb.plot.roc_curve()` will clip the remaining points in an illogical way when the data exceeds `chart_limit = wandb.Table.MAX_ROWS`. This is problematic particularly because it forces the user to use less data for evaluation, as opposed to subsampling the true/false positive rates after computing them with all data.

### Additional Files

_No response_

### Environment

N/A

### Additional Context

I already have the fix, making a PR now


Discussion:

Resolving patch:
diff --git a/wandb/plot/pr_curve.py b/wandb/plot/pr_curve.py
--- a/wandb/plot/pr_curve.py
+++ b/wandb/plot/pr_curve.py
@@ -3,10 +3,14 @@
 from wandb.plots.utils import test_missing, test_types
 
 
-chart_limit = wandb.Table.MAX_ROWS
-
-
-def pr_curve(y_true=None, y_probas=None, labels=None, classes_to_plot=None, title=None):
+def pr_curve(
+    y_true=None,
+    y_probas=None,
+    labels=None,
+    classes_to_plot=None,
+    interp_size=21,
+    title=None,
+):
     """
     Computes the tradeoff between precision and recall for different thresholds.
     A high area under the curve represents both high recall and high precision,
@@ -17,12 +21,18 @@ def pr_curve(y_true=None, y_probas=None, labels=None, classes_to_plot=None, titl
     PR curve is useful when the classes are very imbalanced.
 
     Arguments:
-        y_true (arr): Test set labels.
-        y_probas (arr): Test set predicted probabilities.
+        y_true (arr): true sparse labels
+        y_probas (arr): Target scores, can either be probability estimates, confidence
+                         values, or non-thresholded measure of decisions.
+                         shape: (*y_true.shape, num_classes)
         labels (list): Named labels for target variable (y). Makes plots easier to
-            read by replacing target values with corresponding index.
-            For example labels = ["dog", "cat", "owl"] all 0s are
-            replaced by "dog", 1s by "cat".
+                        read by replacing target values with corresponding index.
+                        For example labels = ['dog', 'cat', 'owl'] all 0s are
+                        replaced by 'dog', 1s by 'cat'.
+        classes_to_plot (list): unique values of y_true to include in the plot
+        interp_size (int): the recall values will be fixed to `interp_size` points
+                            uniform on [0, 1] and the precision will be interpolated for
+                            these recall values.
 
     Returns:
         Nothing. To see plots, go to your W&B run page then expand the 'media' tab
@@ -37,77 +47,87 @@ def pr_curve(y_true=None, y_probas=None, labels=None, classes_to_plot=None, titl
         "numpy",
         required="roc requires the numpy library, install with `pip install numpy`",
     )
-    preprocessing = util.get_module(
-        "sklearn.preprocessing",
-        "roc requires the scikit preprocessing submodule, install with `pip install scikit-learn`",
+    pd = util.get_module(
+        "pandas",
+        required="roc requires the pandas library, install with `pip install pandas`",
     )
-
-    metrics = util.get_module(
+    sklearn_metrics = util.get_module(
         "sklearn.metrics",
-        "roc requires the scikit metrics submodule, install with `pip install scikit-learn`",
+        "roc requires the scikit library, install with `pip install scikit-learn`",
+    )
+    sklearn_utils = util.get_module(
+        "sklearn.utils",
+        "roc requires the scikit library, install with `pip install scikit-learn`",
     )
 
+    def _step(x):
+        y = np.array(x)
+        for i in range(1, len(y)):
+            y[i] = max(y[i], y[i - 1])
+        return y
+
     y_true = np.array(y_true)
     y_probas = np.array(y_probas)
 
-    if test_missing(y_true=y_true, y_probas=y_probas) and test_types(
-        y_true=y_true, y_probas=y_probas
-    ):
-        classes = np.unique(y_true)
-        probas = y_probas
-
-        if classes_to_plot is None:
-            classes_to_plot = classes
-
-        binarized_y_true = preprocessing.label_binarize(y_true, classes=classes)
-        if len(classes) == 2:
-            binarized_y_true = np.hstack((1 - binarized_y_true, binarized_y_true))
-
-        pr_curves = {}
-        indices_to_plot = np.in1d(classes, classes_to_plot)
-        for i, to_plot in enumerate(indices_to_plot):
-            if to_plot:
-                precision, recall, _ = metrics.precision_recall_curve(
-                    y_true, probas[:, i], pos_label=classes[i]
-                )
-
-                samples = 20
-                sample_precision = []
-                sample_recall = []
-                for k in range(samples):
-                    sample_precision.append(
-                        precision[int(len(precision) * k / samples)]
-                    )
-                    sample_recall.append(recall[int(len(recall) * k / samples)])
-
-                pr_curves[classes[i]] = (sample_precision, sample_recall)
+    if not test_missing(y_true=y_true, y_probas=y_probas):
+        return
+    if not test_types(y_true=y_true, y_probas=y_probas):
+        return
+
+    classes = np.unique(y_true)
+    if classes_to_plot is None:
+        classes_to_plot = classes
+
+    precision = dict()
+    interp_recall = np.linspace(0, 1, interp_size)[::-1]
+    indices_to_plot = np.where(np.isin(classes, classes_to_plot))[0]
+    for i in indices_to_plot:
+        if labels is not None and (
+            isinstance(classes[i], int) or isinstance(classes[0], np.integer)
+        ):
+            class_label = labels[classes[i]]
+        else:
+            class_label = classes[i]
+
+        cur_precision, cur_recall, _ = sklearn_metrics.precision_recall_curve(
+            y_true, y_probas[:, i], pos_label=classes[i]
+        )
+        # smooth the precision (monotonically increasing)
+        cur_precision = _step(cur_precision)
+
+        # reverse order so that recall in ascending
+        cur_precision = cur_precision[::-1]
+        cur_recall = cur_recall[::-1]
+        indices = np.searchsorted(cur_recall, interp_recall, side="left")
+        precision[class_label] = cur_precision[indices]
+
+    df = pd.DataFrame(
+        {
+            "class": np.hstack([[k] * len(v) for k, v in precision.items()]),
+            "precision": np.hstack(list(precision.values())),
+            "recall": np.tile(interp_recall, len(precision)),
+        }
+    )
+    df = df.round(3)
 
-        data = []
-        count = 0
-        for class_name in pr_curves.keys():
-            precision, recall = pr_curves[class_name]
-            for p, r in zip(precision, recall):
-                # if class_names are ints and labels are set
-                if labels is not None and (
-                    isinstance(class_name, int) or isinstance(class_name, np.integer)
-                ):
-                    class_name = labels[class_name]
-                # if class_names are ints and labels are not set
-                # or, if class_names have something other than ints
-                # (string, float, date) - user class_names
-                data.append([class_name, round(p, 3), round(r, 3)])
-                count += 1
-                if count >= chart_limit:
-                    wandb.termwarn(
-                        "wandb uses only the first %d datapoints to create the plots."
-                        % wandb.Table.MAX_ROWS
-                    )
-                    break
-        table = wandb.Table(columns=["class", "precision", "recall"], data=data)
-        title = title or "Precision v. Recall"
-        return wandb.plot_table(
-            "wandb/area-under-curve/v0",
-            table,
-            {"x": "recall", "y": "precision", "class": "class"},
-            {"title": title},
+    if len(df) > wandb.Table.MAX_ROWS:
+        wandb.termwarn(
+            "wandb uses only %d data points to create the plots." % wandb.Table.MAX_ROWS
         )
+        # different sampling could be applied, possibly to ensure endpoints are kept
+        df = sklearn_utils.resample(
+            df,
+            replace=False,
+            n_samples=wandb.Table.MAX_ROWS,
+            random_state=42,
+            stratify=df["class"],
+        ).sort_values(["precision", "recall", "class"])
+
+    table = wandb.Table(dataframe=df)
+    title = title or "Precision v. Recall"
+    return wandb.plot_table(
+        "wandb/area-under-curve/v0",
+        table,
+        {"x": "recall", "y": "precision", "class": "class"},
+        {"title": title},
+    )
diff --git a/wandb/plot/roc_curve.py b/wandb/plot/roc_curve.py
--- a/wandb/plot/roc_curve.py
+++ b/wandb/plot/roc_curve.py
@@ -2,8 +2,6 @@
 from wandb import util
 from wandb.plots.utils import test_missing, test_types
 
-chart_limit = wandb.Table.MAX_ROWS
-
 
 def roc_curve(
     y_true=None, y_probas=None, labels=None, classes_to_plot=None, title=None
@@ -13,12 +11,15 @@ def roc_curve(
     ROC curve.
 
     Arguments:
-        y_true (arr): Test set labels.
-        y_probas (arr): Test set predicted probabilities.
-        labels (list): Named labels for target varible (y). Makes plots easier to
+        y_true (arr): true sparse labels
+        y_probas (arr): Target scores, can either be probability estimates, confidence
+                         values, or non-thresholded measure of decisions.
+                         shape: (*y_true.shape, num_classes)
+        labels (list): Named labels for target variable (y). Makes plots easier to
                         read by replacing target values with corresponding index.
-                        For example labels= ['dog', 'cat', 'owl'] all 0s are
+                        For example labels = ['dog', 'cat', 'owl'] all 0s are
                         replaced by 'dog', 1s by 'cat'.
+        classes_to_plot (list): unique values of y_true to include in the plot
 
     Returns:
         Nothing. To see plots, go to your W&B run page then expand the 'media' tab
@@ -33,66 +34,77 @@ def roc_curve(
         "numpy",
         required="roc requires the numpy library, install with `pip install numpy`",
     )
-    util.get_module(
-        "sklearn",
-        required="roc requires the scikit library, install with `pip install scikit-learn`",
+    pd = util.get_module(
+        "pandas",
+        required="roc requires the pandas library, install with `pip install pandas`",
+    )
+    sklearn_metrics = util.get_module(
+        "sklearn.metrics",
+        "roc requires the scikit library, install with `pip install scikit-learn`",
     )
-    from sklearn.metrics import roc_curve
+    sklearn_utils = util.get_module(
+        "sklearn.utils",
+        "roc requires the scikit library, install with `pip install scikit-learn`",
+    )
+
+    y_true = np.array(y_true)
+    y_probas = np.array(y_probas)
 
-    if test_missing(y_true=y_true, y_probas=y_probas) and test_types(
-        y_true=y_true, y_probas=y_probas
-    ):
-        y_true = np.array(y_true)
-        y_probas = np.array(y_probas)
-        classes = np.unique(y_true)
-        probas = y_probas
+    if not test_missing(y_true=y_true, y_probas=y_probas):
+        return
+    if not test_types(y_true=y_true, y_probas=y_probas):
+        return
 
-        if classes_to_plot is None:
-            classes_to_plot = classes
+    classes = np.unique(y_true)
+    if classes_to_plot is None:
+        classes_to_plot = classes
 
-        fpr_dict = dict()
-        tpr_dict = dict()
+    fpr = dict()
+    tpr = dict()
+    indices_to_plot = np.where(np.isin(classes, classes_to_plot))[0]
+    for i in indices_to_plot:
+        if labels is not None and (
+            isinstance(classes[i], int) or isinstance(classes[0], np.integer)
+        ):
+            class_label = labels[classes[i]]
+        else:
+            class_label = classes[i]
 
-        indices_to_plot = np.in1d(classes, classes_to_plot)
+        fpr[class_label], tpr[class_label], _ = sklearn_metrics.roc_curve(
+            y_true, y_probas[..., i], pos_label=classes[i]
+        )
 
-        data = []
-        count = 0
+    df = pd.DataFrame(
+        {
+            "class": np.hstack([[k] * len(v) for k, v in fpr.items()]),
+            "fpr": np.hstack(list(fpr.values())),
+            "tpr": np.hstack(list(tpr.values())),
+        }
+    )
+    df = df.round(3)
 
-        for i, to_plot in enumerate(indices_to_plot):
-            fpr_dict[i], tpr_dict[i], _ = roc_curve(
-                y_true, probas[:, i], pos_label=classes[i]
-            )
-            if to_plot:
-                for j in range(len(fpr_dict[i])):
-                    if labels is not None and (
-                        isinstance(classes[i], int)
-                        or isinstance(classes[0], np.integer)
-                    ):
-                        class_dict = labels[classes[i]]
-                    else:
-                        class_dict = classes[i]
-                    fpr = [
-                        class_dict,
-                        round(fpr_dict[i][j], 3),
-                        round(tpr_dict[i][j], 3),
-                    ]
-                    data.append(fpr)
-                    count += 1
-                    if count >= chart_limit:
-                        wandb.termwarn(
-                            "wandb uses only the first %d datapoints to create the plots."
-                            % wandb.Table.MAX_ROWS
-                        )
-                        break
-        table = wandb.Table(columns=["class", "fpr", "tpr"], data=data)
-        title = title or "ROC"
-        return wandb.plot_table(
-            "wandb/area-under-curve/v0",
-            table,
-            {"x": "fpr", "y": "tpr", "class": "class"},
-            {
-                "title": title,
-                "x-axis-title": "False positive rate",
-                "y-axis-title": "True positive rate",
-            },
+    if len(df) > wandb.Table.MAX_ROWS:
+        wandb.termwarn(
+            "wandb uses only %d data points to create the plots." % wandb.Table.MAX_ROWS
         )
+        # different sampling could be applied, possibly to ensure endpoints are kept
+        df = sklearn_utils.resample(
+            df,
+            replace=False,
+            n_samples=wandb.Table.MAX_ROWS,
+            random_state=42,
+            stratify=df["class"],
+        ).sort_values(["fpr", "tpr", "class"])
+
+    table = wandb.Table(dataframe=df)
+    title = title or "ROC"
+    return wandb.plot_table(
+        "wandb/area-under-curve/v0",
+        table,
+        {"x": "fpr", "y": "tpr", "class": "class"},
+        {
+            "title": title,
+            "x-axis-title": "False positive rate",
+            "y-axis-title": "True positive rate",
+        },
+    )


Test patch:
diff --git a/tests/unit_tests/test_plot.py b/tests/unit_tests/test_plot.py
--- a/tests/unit_tests/test_plot.py
+++ b/tests/unit_tests/test_plot.py
@@ -34,11 +34,8 @@ def test_roc(dummy_classifier, wandb_init_run):
     (nb, x_train, y_train, x_test, y_test, y_pred, y_probas) = dummy_classifier
     custom_chart_no_title = roc_curve(y_test, y_probas)
     assert custom_chart_no_title.string_fields["title"] == "ROC"
-    assert custom_chart_no_title.table.data[0] == [
-        0,
-        0.0,
-        0.0,
-    ], custom_chart_no_title.table.data[0]
+    first_row = custom_chart_no_title.table.data[0]
+    assert first_row == [0, 0.0, 0.0], first_row
     custom_chart_with_title = roc_curve(y_test, y_probas, title="New title")
     assert custom_chart_with_title.string_fields["title"] == "New title"
 
@@ -47,12 +44,8 @@ def test_pr(dummy_classifier, wandb_init_run):
     (nb, x_train, y_train, x_test, y_test, y_pred, y_probas) = dummy_classifier
     custom_chart_no_title = pr_curve(y_test, y_probas)
     assert custom_chart_no_title.string_fields["title"] == "Precision v. Recall"
-    # note: see https://github.com/wandb/client/pull/3735/ for context
-    sklearn_version = parse_version(sklearn.__version__)
-    pr = [0, 1.0, 1.0] if sklearn_version < parse_version("1.1") else [0, 0.5, 1.0]
-    assert custom_chart_no_title.table.data[0] == pr, custom_chart_no_title.table.data[
-        0
-    ]
+    first_row = custom_chart_no_title.table.data[0]
+    assert first_row == [0, 1.0, 1.0], first_row
     custom_chart_with_title = pr_curve(y_test, y_probas, title="New title")
     assert custom_chart_with_title.string_fields["title"] == "New title"