Update _safe_updates to use global dict. Adds tests for plots

src/datagnosis/plugins/core/plugin.py

@@ -25,6 +25,20 @@
 from datagnosis.utils.constants import DEVICE
 from datagnosis.utils.reproducibility import clear_cache, enable_reproducible_results
 
+# The complete list of parameters that can be passed to the _updates method of any of the plugins.
+UPDATE_PARAMS: Dict[str, Any] = {
+    "y_pred": None,
+    "y_batch": None,
+    "sample_ids": None,
+    "net": None,
+    "device": None,
+    "logits": None,
+    "targets": None,
+    "probs": None,
+    "indices": None,
+    "data_uncert_class": None,
+}
+
 
 # Base class for Hardness Classification Methods (HCMs)
 class Plugin(metaclass=ABCMeta):
@@ -89,6 +103,8 @@ def __init__(
             log.debug("Fixing seed for reproducibility.")
             enable_reproducible_results(0)
         log.debug(f"Initialized parent plugin for {self.name()}")
+        # Update UPDATE_PARAMS
+        UPDATE_PARAMS["device"] = self.device
 
     @staticmethod
     @abstractmethod
@@ -194,11 +210,15 @@ def fit(
                         indices,
                         batch_idx,
                     ) = utils.load_update_values_from_cache(update_values_cache_file)
+
                     if self.update_point == "mid-epoch":
                         log.debug("Updating scores mid-epoch")
-                        self._safe_update(
-                            y_pred=outputs, y_batch=true_label, sample_ids=indices
-                        )
+                        # Update UPDATE_PARAMS before calling _safe_update() for mid-epoch
+                        UPDATE_PARAMS["y_pred"] = outputs
+                        UPDATE_PARAMS["y_batch"] = true_label
+                        UPDATE_PARAMS["sample_ids"] = indices
+                        UPDATE_PARAMS["net"] = self.model
+                        self._safe_update(**UPDATE_PARAMS)
                 else:
                     if (
                         use_caches_if_exist
@@ -220,9 +240,12 @@ def fit(
 
                     if self.update_point == "mid-epoch":
                         log.debug("Updating scores mid-epoch")
-                        self._safe_update(
-                            y_pred=outputs, y_batch=true_label, sample_ids=indices
-                        )
+                        # Update UPDATE_PARAMS before calling _safe_update() for mid-epoch
+                        UPDATE_PARAMS["y_pred"] = outputs
+                        UPDATE_PARAMS["y_batch"] = true_label
+                        UPDATE_PARAMS["sample_ids"] = indices
+                        UPDATE_PARAMS["net"] = self.model
+                        self._safe_update(**UPDATE_PARAMS)
                     loss = self.criterion(outputs, true_label)
                     loss.backward()
                     self.optimizer.step()
@@ -286,25 +309,23 @@ def fit(
 
             if self.update_point == "per-epoch":
                 log.debug(f"Updating plugin after epoch {epoch+1}")
-                self._safe_update(
-                    net=self.model,
-                    device=self.device,
-                    logits=logits,
-                    targets=targets,
-                    probs=probs,
-                    indices=indices,
-                )
+                # Update UPDATE_PARAMS before calling _safe_update() for per-epoch
+                UPDATE_PARAMS["net"] = self.model
+                UPDATE_PARAMS["logits"] = logits
+                UPDATE_PARAMS["targets"] = targets
+                UPDATE_PARAMS["probs"] = probs
+                UPDATE_PARAMS["indices"] = indices
+                self._safe_update(**UPDATE_PARAMS)
 
         if self.update_point == "post-epoch":
             log.debug("Updating plugin after training")
-            self._safe_update(
-                net=self.model,
-                data_uncert_class=self.data_uncert_class,
-                device=self.device,
-                logits=logits,
-                targets=targets,
-                probs=probs,
-            )
+            # Update UPDATE_PARAMS before calling _safe_update() for per-epoch
+            UPDATE_PARAMS["net"] = self.model
+            UPDATE_PARAMS["logits"] = logits
+            UPDATE_PARAMS["targets"] = targets
+            UPDATE_PARAMS["probs"] = probs
+            UPDATE_PARAMS["data_uncert_class"] = self.data_uncert_class
+            self._safe_update(**UPDATE_PARAMS)
         self.has_been_fit = True
         self.compute_scores()
         log.debug("Plugin fit complete and scores computed.")
@@ -447,7 +468,6 @@ def _extract_datapoints_by_top_n(
             log.info(extracted)
             extracted = sorted(extracted)
             log.info(extracted)
-        print("top_n", self.dataloader_unshuffled.dataset[extracted])
         return (
             self.dataloader_unshuffled.dataset[extracted],  # pyright: ignore
             extraction_scores[extracted],

tests/plugins/core/test_plugin.py

@@ -172,19 +172,19 @@ def test_mock_plugin_extract_all_methods() -> None:
         [0.10, 0.16, 0.15, 0.14, 0.13, 0.12, 0.11, 0.14, 0.12, 0.06]
     ), np.asarray([0.01, 0.16, 0.15, 0.14, 0.13, 0.12, 0.11, 0.14, 0.12, 0.12])
 
-    # # extract datapoints
-    # # extracted have format: ((features, Labels, Indices), scores)
-    # extract_indices = [0, 1, 5]
-    # extracted = plugin.extract_datapoints(method="index", indices=extract_indices)
-    # assert isinstance(extracted, tuple)
-    # assert isinstance(extracted[0][0], torch.Tensor)
-    # assert isinstance(extracted[0][1], torch.Tensor)
-    # assert isinstance(extracted[0][2], List)
-    # assert isinstance(extracted[1], np.ndarray)
-    # assert extracted[0][0].shape[0] == len(extract_indices)
-    # assert extracted[0][1].shape[0] == len(extract_indices)
-    # assert len(extracted[0][2]) == len(extract_indices)
-    # assert extracted[1].shape[0] == len(extract_indices)
+    # extract datapoints
+    # extracted have format: ((features, Labels, Indices), scores)
+    extract_indices = [0, 1, 5]
+    extracted = plugin.extract_datapoints(method="index", indices=extract_indices)
+    assert isinstance(extracted, tuple)
+    assert isinstance(extracted[0][0], torch.Tensor)
+    assert isinstance(extracted[0][1], torch.Tensor)
+    assert isinstance(extracted[0][2], List)
+    assert isinstance(extracted[1], np.ndarray)
+    assert extracted[0][0].shape[0] == len(extract_indices)
+    assert extracted[0][1].shape[0] == len(extract_indices)
+    assert len(extracted[0][2]) == len(extract_indices)
+    assert extracted[1].shape[0] == len(extract_indices)
 
     extracted = plugin.extract_datapoints(method="top_n", n=3, sort_by_index=True)
     assert isinstance(extracted, tuple)

tests/plugins/generic/test_aum.py

@@ -73,3 +73,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert len(scores) == len(y)
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_conf_agree.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score <= 1.0 for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_confident_learning.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score <= 1.0 for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_data_iq.py

@@ -75,3 +75,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert scores[0].dtype in [np.float32, np.float64]
     assert scores[1].dtype in [np.float32, np.float64]
     assert all([0.0 <= score <= 1.0 for score in scores[0]])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_data_maps.py

@@ -75,3 +75,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert scores[0].dtype in [np.float32, np.float64]
     assert scores[1].dtype in [np.float32, np.float64]
     assert all([0.0 <= score <= 1.0 for score in scores[0]])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_el2n.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_forgetting.py

@@ -77,3 +77,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score <= 1.0 for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_grand.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_large_loss.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_prototypicality.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/generic/test_vog.py

@@ -74,3 +74,17 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
     assert all([0.0 <= score <= 1.0 for score in scores])
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y = load_iris(return_X_y=True, as_frame=True)
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)

tests/plugins/images/test_allsh.py

@@ -129,9 +129,24 @@ def test_plugin_scores(test_plugin: Plugin) -> None:
         datahandler=datahander,
         use_caches_if_exist=False,
         workspace="test_workspace",
-        epochs=2,
     )
     scores = test_plugin.scores
     assert len(scores) == len(y)
     assert isinstance(scores, np.ndarray)
     assert scores.dtype in [np.float32, np.float64]
+
+
+@pytest.mark.parametrize(
+    "test_plugin", generate_fixtures(plugin_name, plugin, plugin_args_mnist_lenet)
+)
+def test_plugin_plots(test_plugin: Plugin) -> None:
+    X, y, _, _ = load_mnist()
+    X = X[:100]
+    y = y[:100]
+    datahander = DataHandler(X, y, batch_size=32)  # pyright: ignore
+    test_plugin.fit(
+        datahandler=datahander,
+        use_caches_if_exist=False,
+        workspace="test_workspace",
+    )
+    test_plugin.plot_scores(show=False)