Repository: cleanlab/cleanlab

Issue:
extend non-iid issue check in Datalab
Currently Datalab's non-iid issue type is only detected based on `features`:
https://docs.cleanlab.ai/master/cleanlab/datalab/guide/issue_type_description.html#non-iid-issue

If the user does not input `features`, but does provide` pred_probs`, run this same check based on the `pred_probs` instead (just treat them as features). 

Note that if the user did provide `features` or there was already a KNN graph constructed in Datalab, the results should be returned as they currently are, not using the `pred_probs` at all!

Reference:  https://github.com/cleanlab/cleanlab/blob/master/cleanlab/datalab/internal/issue_manager/noniid.py  


Discussion:
@jwmueller can you assign this issue to me?

Resolving patch:
diff --git a/cleanlab/datalab/internal/issue_finder.py b/cleanlab/datalab/internal/issue_finder.py
--- a/cleanlab/datalab/internal/issue_finder.py
+++ b/cleanlab/datalab/internal/issue_finder.py
@@ -199,7 +199,7 @@ def _resolve_required_args(self, pred_probs, features, knn_graph):
             "label": {"pred_probs": pred_probs, "features": features},
             "outlier": {"pred_probs": pred_probs, "features": features, "knn_graph": knn_graph},
             "near_duplicate": {"features": features, "knn_graph": knn_graph},
-            "non_iid": {"features": features, "knn_graph": knn_graph},
+            "non_iid": {"pred_probs": pred_probs, "features": features, "knn_graph": knn_graph},
         }
 
         args_dict = {
diff --git a/cleanlab/datalab/internal/issue_manager/noniid.py b/cleanlab/datalab/internal/issue_manager/noniid.py
--- a/cleanlab/datalab/internal/issue_manager/noniid.py
+++ b/cleanlab/datalab/internal/issue_manager/noniid.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union, cast, Tuple
 import warnings
 import itertools
 
@@ -124,38 +124,103 @@ def __init__(
         self.seed = seed
         self.significance_threshold = significance_threshold
 
-    def find_issues(self, features: Optional[npt.NDArray] = None, **kwargs) -> None:
+        # TODO: Temporary flag introduced to decide on storing knn graphs based on pred_probs.
+        # Revisit and finalize the implementation.
+        self._skip_storing_knn_graph_for_pred_probs: bool = False
+
+    @staticmethod
+    def _determine_features(
+        features: Optional[npt.NDArray],
+        pred_probs: Optional[np.ndarray],
+    ) -> npt.NDArray:
+        """
+        Determines the feature array to be used for the non-IID check. Prioritizing the original features array over pred_probs.
+
+        Parameters
+        ----------
+        features :
+            Original feature array or None.
+
+        pred_probs :
+            Predicted probabilities array or None.
+
+        Returns
+        -------
+        features_to_use :
+            Either the original feature array or the predicted probabilities array,
+            intended to be used for the non-IID check.
+
+        Raises
+        ------
+        ValueError :
+            If both `features` and `pred_probs` are None.
+        """
+        if features is not None:
+            return features
+
+        if pred_probs is not None:
+            return pred_probs
+
+        raise ValueError(
+            "If a knn_graph is not provided, either 'features' or 'pred_probs' must be provided to fit a new knn."
+        )
+
+    def _select_features_and_setup_knn(
+        self,
+        features: Optional[npt.NDArray],
+        pred_probs: Optional[np.ndarray],
+        knn_graph: Optional[csr_matrix],
+        metric_changes: bool,
+    ) -> Tuple[Optional[NearestNeighbors], npt.NDArray]:
+        """
+        Selects features (or pred_probs if features are None) and sets up a NearestNeighbors object if needed.
+
+        # Add type-hints and document the arguments.
+        """
+        if features is None and pred_probs is not None:
+            self._skip_storing_knn_graph_for_pred_probs = True
+        features_to_use = self._determine_features(features, pred_probs)
+
+        if self.metric is None:
+            self.metric = "cosine" if features_to_use.shape[1] > 3 else "euclidean"
+
+        if knn_graph is not None and not metric_changes:
+            return None, features_to_use
+
+        knn = NearestNeighbors(n_neighbors=self.k, metric=self.metric)
+
+        if self.metric != knn.metric:
+            warnings.warn(
+                f"Metric {self.metric} does not match metric {knn.metric} used to fit knn. "
+                "Most likely an existing NearestNeighbors object was passed in, but a different "
+                "metric was specified."
+            )
+        self.metric = knn.metric
+
+        try:
+            check_is_fitted(knn)
+        except NotFittedError:
+            knn.fit(features_to_use)
+
+        return knn, features_to_use
+
+    def find_issues(
+        self,
+        features: Optional[npt.NDArray] = None,
+        pred_probs: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> None:
         knn_graph = self._process_knn_graph_from_inputs(kwargs)
         old_knn_metric = self.datalab.get_info("statistics").get("knn_metric")
         metric_changes = self.metric and self.metric != old_knn_metric
-
-        knn = None  # Won't be used if knn_graph is not None
+        knn, features_used = self._select_features_and_setup_knn(
+            features, pred_probs, knn_graph, metric_changes
+        )
 
         if knn_graph is None or metric_changes:
-            if features is None:
-                raise ValueError(
-                    "If a knn_graph is not provided, features must be provided to fit a new knn."
-                )
-
-            if self.metric is None:
-                self.metric = "cosine" if features.shape[1] > 3 else "euclidean"
-            knn = NearestNeighbors(n_neighbors=self.k, metric=self.metric)
-
-            if self.metric and self.metric != knn.metric:
-                warnings.warn(
-                    f"Metric {self.metric} does not match metric {knn.metric} used to fit knn. "
-                    "Most likely an existing NearestNeighbors object was passed in, but a different "
-                    "metric was specified."
-                )
-            self.metric = knn.metric
-
-            try:
-                check_is_fitted(knn)
-            except NotFittedError:
-                knn.fit(features)
-
             self.neighbor_index_choices = self._get_neighbors(knn=knn)
         else:
+            self._skip_storing_knn_graph_for_pred_probs = False
             self.neighbor_index_choices = self._get_neighbors(knn_graph=knn_graph)
 
         self.num_neighbors = self.k
@@ -234,6 +299,8 @@ def collect_info(
     def _build_statistics_dictionary(self, knn_graph: csr_matrix) -> Dict[str, Dict[str, Any]]:
         statistics_dict: Dict[str, Dict[str, Any]] = {"statistics": {}}
 
+        if self._skip_storing_knn_graph_for_pred_probs:
+            return statistics_dict
         # Add the knn graph as a statistic if necessary
         graph_key = "weighted_knn_graph"
         old_knn_graph = self.datalab.get_info("statistics").get(graph_key, None)


Test patch:
diff --git a/tests/datalab/issue_manager/test_noniid.py b/tests/datalab/issue_manager/test_noniid.py
--- a/tests/datalab/issue_manager/test_noniid.py
+++ b/tests/datalab/issue_manager/test_noniid.py
@@ -60,6 +60,13 @@ def embeddings(self, lab):
         embeddings_array = np.arange(lab.get_info("statistics")["num_examples"] * 10).reshape(-1, 1)
         return embeddings_array
 
+    @pytest.fixture
+    def pred_probs(self, lab):
+        pred_probs_array = (
+            np.arange(lab.get_info("statistics")["num_examples"] * 10).reshape(-1, 1)
+        ) / len(np.arange(lab.get_info("statistics")["num_examples"] * 10).reshape(-1, 1))
+        return pred_probs_array
+
     @pytest.fixture
     def issue_manager(self, lab):
         return NonIIDIssueManager(
@@ -116,7 +123,49 @@ def test_find_issues(self, issue_manager, embeddings):
             issues_perm["is_non_iid_issue"] == expected_permuted_issue_mask
         ), "Issue mask should be correct"
         assert summary_perm["issue_type"][0] == "non_iid"
-        # ensure score is large, cannot easily ensure precise value because random seed has different effects on different OS:
+        # ensure score is large, cannot easily ensure precise value because random seed has different effects on
+        # different OS:
+        assert summary_perm["score"][0] > 0.05
+        assert info_perm.get("p-value", None) is not None, "Should have p-value"
+        assert summary_perm["score"][0] == pytest.approx(expected=info_perm["p-value"], abs=1e-7)
+
+    def test_find_issues_using_pred_probs(self, issue_manager, pred_probs):
+        np.random.seed(SEED)
+        issue_manager.find_issues(pred_probs=pred_probs)
+        issues_sort, summary_sort, info_sort = (
+            issue_manager.issues,
+            issue_manager.summary,
+            issue_manager.info,
+        )
+        expected_sorted_issue_mask = np.array([False] * 46 + [True] + [False] * 3)
+        assert np.all(
+            issues_sort["is_non_iid_issue"] == expected_sorted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_sort["issue_type"][0] == "non_iid"
+        assert summary_sort["score"][0] == pytest.approx(expected=0.0, abs=1e-7)
+        assert info_sort.get("p-value", None) is not None, "Should have p-value"
+        assert summary_sort["score"][0] == pytest.approx(expected=info_sort["p-value"], abs=1e-7)
+
+        permutation = np.random.permutation(len(pred_probs))
+        new_issue_manager = NonIIDIssueManager(
+            datalab=issue_manager.datalab,
+            metric="euclidean",
+            k=10,
+        )
+
+        new_issue_manager.find_issues(pred_probs=pred_probs[permutation])
+        issues_perm, summary_perm, info_perm = (
+            new_issue_manager.issues,
+            new_issue_manager.summary,
+            new_issue_manager.info,
+        )
+        expected_permuted_issue_mask = np.array([False] * len(pred_probs))
+        assert np.all(
+            issues_perm["is_non_iid_issue"] == expected_permuted_issue_mask
+        ), "Issue mask should be correct"
+        assert summary_perm["issue_type"][0] == "non_iid"
+        # ensure score is large, cannot easily ensure precise value because random seed has different effects on
+        # different OS:
         assert summary_perm["score"][0] > 0.05
         assert info_perm.get("p-value", None) is not None, "Should have p-value"
         assert summary_perm["score"][0] == pytest.approx(expected=info_perm["p-value"], abs=1e-7)
@@ -136,12 +185,38 @@ def test_report(self, issue_manager, embeddings):
             "Number of examples with this issue:"
         ) in report
 
+        issue_manager.find_issues(features=embeddings)
         report = issue_manager.report(
             issues=issue_manager.issues,
             summary=issue_manager.summary,
             info=issue_manager.info,
             verbosity=3,
         )
+
+        assert "Additional Information: " in report
+
+    def test_report_using_pred_probs(self, issue_manager, pred_probs):
+        np.random.seed(SEED)
+        issue_manager.find_issues(pred_probs=pred_probs)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+        )
+
+        assert (
+            "---------------------- non_iid issues ----------------------\n\n"
+            "Number of examples with this issue:"
+        ) in report
+
+        issue_manager.find_issues(pred_probs=pred_probs)
+        report = issue_manager.report(
+            issues=issue_manager.issues,
+            summary=issue_manager.summary,
+            info=issue_manager.info,
+            verbosity=3,
+        )
+
         assert "Additional Information: " in report
 
     def test_collect_info(self, issue_manager, embeddings):
@@ -157,6 +232,18 @@ def test_collect_info(self, issue_manager, embeddings):
         assert info["metric"] == "euclidean"
         assert info["k"] == 10
 
+    def test_collect_info_using_pred_probs(self, issue_manager, pred_probs):
+        """Test some values in the info dict.
+
+        Mainly focused on the nearest neighbor info.
+        """
+        issue_manager.find_issues(pred_probs=pred_probs)
+        info = issue_manager.info
+
+        assert info["p-value"] == 0
+        assert info["metric"] == "euclidean"
+        assert info["k"] == 10
+
     @pytest.mark.parametrize(
         "seed",
         [
@@ -221,3 +308,32 @@ def generate_data_iid():
             assert p_value == p_value2
         else:
             assert p_value != p_value2
+
+        # using pred_probs
+        # normalizing pred_probs (0 to 1)
+        pred_probs = embeddings / (np.max(embeddings) - np.min(embeddings))
+        if seed == "default":
+            issue_manager = NonIIDIssueManager(
+                datalab=lab,
+                metric="euclidean",
+                k=10,
+            )
+        else:
+            issue_manager = NonIIDIssueManager(
+                datalab=lab,
+                metric="euclidean",
+                k=10,
+                seed=seed,
+            )
+        issue_manager.find_issues(pred_probs=pred_probs)
+        p_value = issue_manager.info["p-value"]
+
+        # Run again with the same seed
+        issue_manager.find_issues(pred_probs=pred_probs)
+        p_value2 = issue_manager.info["p-value"]
+
+        assert p_value > 0.0
+        if seed is not None or seed == "default":
+            assert p_value == p_value2
+        else:
+            assert p_value != p_value2
diff --git a/tests/datalab/test_datalab.py b/tests/datalab/test_datalab.py
--- a/tests/datalab/test_datalab.py
+++ b/tests/datalab/test_datalab.py
@@ -776,6 +776,16 @@ def test_find_non_iid_issues(self, random_embeddings):
         assert summary["score"].values[0] > 0.05
         assert lab.get_issues()["is_non_iid_issue"].sum() == 0
 
+    def test_find_non_iid_issues_using_pred_probs(self, random_embeddings):
+        data = {"labels": [0, 1, 0]}
+        lab = Datalab(data=data, label_name="labels")
+        pred_probs = random_embeddings / random_embeddings.sum(axis=1, keepdims=True)
+        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
+        summary = lab.get_issue_summary()
+        assert ["non_iid"] == summary["issue_type"].values
+        assert summary["score"].values[0] > 0.05
+        assert lab.get_issues()["is_non_iid_issue"].sum() == 0
+
     def test_find_non_iid_issues_sorted(self, sorted_embeddings):
         data = {"labels": [0, 1, 0]}
         lab = Datalab(data=data, label_name="labels")
@@ -785,6 +795,16 @@ def test_find_non_iid_issues_sorted(self, sorted_embeddings):
         assert summary["score"].values[0] == 0
         assert lab.get_issues()["is_non_iid_issue"].sum() == 1
 
+    def test_find_non_iid_issues_sorted_using_pred_probs(self, sorted_embeddings):
+        data = {"labels": [0, 1, 0]}
+        lab = Datalab(data=data, label_name="labels")
+        pred_probs = sorted_embeddings / sorted_embeddings.sum(axis=1, keepdims=True)
+        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
+        summary = lab.get_issue_summary()
+        assert ["non_iid"] == summary["issue_type"].values
+        assert summary["score"].values[0] == 0
+        assert lab.get_issues()["is_non_iid_issue"].sum() == 1
+
     def test_incremental_search(self, sorted_embeddings):
         data = {"labels": [0, 1, 0]}
         lab = Datalab(data=data, label_name="labels")
@@ -799,6 +819,21 @@ def test_incremental_search(self, sorted_embeddings):
         assert non_iid_summary["score"].values[0] == 0
         assert non_iid_summary["num_issues"].values[0] == 1
 
+    def test_incremental_search_using_pred_probs(self, sorted_embeddings):
+        data = {"labels": [0, 1, 0]}
+        lab = Datalab(data=data, label_name="labels")
+        pred_probs = sorted_embeddings / sorted_embeddings.sum(axis=1, keepdims=True)
+        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
+        summary = lab.get_issue_summary()
+        assert len(summary) == 1
+        lab.find_issues(pred_probs=pred_probs, issue_types={"non_iid": {}})
+        summary = lab.get_issue_summary()
+        assert len(summary) == 1
+        assert "non_iid" in summary["issue_type"].values
+        non_iid_summary = lab.get_issue_summary("non_iid")
+        assert non_iid_summary["score"].values[0] == 0
+        assert non_iid_summary["num_issues"].values[0] == 1
+
 
 class TestDatalabFindLabelIssues:
     @pytest.fixture
@@ -1042,7 +1077,7 @@ def test_init(self, lab, features):
     def test_find_issues(self, lab, features, pred_probs):
         lab = Datalab(data={"X": features})
         lab.find_issues(pred_probs=pred_probs)
-        assert lab.issues.empty
+        assert set(lab.issues.columns) == {"is_non_iid_issue", "non_iid_score"}
 
         lab = Datalab(data={"X": features})
         lab.find_issues(features=features)