Repository: scikit-activeml/scikit-activeml (language: python)

Issue:
FR: Allow limiting training data in Subsampling Wrapper
### Describe the workflow you want to enable

Currently, the Subsampling Wrapper forwards all unlabeled data to the wrapped query strategy and the subset is purely defined by the `query`'s  candidate parameter. However, strategies that utilize unlabeled data for things like clustering might benefit from excluding the unlabeled data that is not included in the subsample itself.

### Describe your proposed solution

Add a boolean flag that allows for excluding the non-subsampled unlabeled data from `X` and `Y`. this flag should be turned off by default.

### Describe alternatives you've considered, if relevant

_No response_

### Additional context

_No response_


Discussion / hints:

Fix patch:
diff --git a/skactiveml/pool/_wrapper.py b/skactiveml/pool/_wrapper.py
index 44169490..551a8e7d 100644
--- a/skactiveml/pool/_wrapper.py
+++ b/skactiveml/pool/_wrapper.py
@@ -3,6 +3,7 @@
     MISSING_LABEL,
     check_random_state,
     is_labeled,
+    labeled_indices,
     unlabeled_indices,
     check_scalar,
     simple_batch,
@@ -26,11 +27,17 @@ class SubSamplingWrapper(SingleAnnotatorPoolQueryStrategy):
     query_strategy : skactiveml.base.SingleAnnotatorPoolQueryStrategy
         The strategy used for computing the utilities of the candidate
         sub-sample.
-    max_candidates : int or float
-         Determines the number of candidates. If `max_candidates` is an
-         integer, `max_candidates` is the maximum number of candidates whose
-         utilities are computed. If `max_candidates` is a float,
-         `max_candidates` is the fraction of the original number of candidates.
+    max_candidates : int or float, default=0.1
+        Determines the number of candidates. If `max_candidates` is an
+        integer, `max_candidates` is the maximum number of candidates whose
+        utilities are computed. If `max_candidates` is a float,
+        `max_candidates` is the fraction of the original number of candidates.
+    exclude_non_subsample : bool, default=False
+        - If `True`, unlabeled candidates in `X` and `y` are excluded which
+          are not part of the subsample.  If `candidates` is an array-like of
+          shape `(n_candidates, n_features)`, all unlabeled data will be
+          removed from `X` and `y`.
+        - If `False`, `X` and `y` stay the same.
     missing_label : scalar or string or np.nan or None, default=np.nan
         Value to represent a missing label.
     random_state : int or np.random.RandomState
@@ -40,7 +47,8 @@ class SubSamplingWrapper(SingleAnnotatorPoolQueryStrategy):
     def __init__(
         self,
         query_strategy=None,
-        max_candidates=None,
+        max_candidates=0.1,
+        exclude_non_subsample=False,
         missing_label=MISSING_LABEL,
         random_state=None,
     ):
@@ -49,6 +57,7 @@ def __init__(
         )
         self.query_strategy = query_strategy
         self.max_candidates = max_candidates
+        self.exclude_non_subsample = exclude_non_subsample
 
     @match_signature("query_strategy", "query")
     def query(
@@ -65,25 +74,24 @@ def query(
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
-           Training data set, usually complete, i.e. including the labeled and
-           unlabeled samples.
+            Training data set, usually complete, i.e. including the labeled and
+            unlabeled samples.
         y : array-like of shape (n_samples)
-           Labels of the training data set (possibly including unlabeled ones
-           indicated by self.MISSING_LABEL).
+            Labels of the training data set (possibly including unlabeled ones
+            indicated by self.MISSING_LABEL).
         candidates : None or array-like of shape (n_candidates), dtype=int or
-           array-like of shape (n_candidates, n_features),
-           optional (default=None)
-           If candidates is None, the unlabeled samples from (X,y) are
-           considered as candidates.
-           If candidates is of shape (n_candidates) and of type int,
-           candidates is considered as the indices of the samples in (X,y).
-           If candidates is of shape (n_candidates, n_features), the
-           candidates are directly given in candidates (not necessarily
-           contained in X). This is not supported by all query strategies.
-        batch_size : int, optional (default=1)
-           The number of samples to be selected in one AL cycle.
-        return_utilities : bool, optional (default=False)
-           If true, also return the utilities based on the query strategy.
+            array-like of shape (n_candidates, n_features), default=None
+            If candidates is None, the unlabeled samples from (X,y) are
+            considered as candidates.
+            If candidates is of shape (n_candidates) and of type int,
+            candidates is considered as the indices of the samples in (X,y).
+            If candidates is of shape (n_candidates, n_features), the
+            candidates are directly given in candidates (not necessarily
+            contained in X). This is not supported by all query strategies.
+        batch_size : int, default=1
+            The number of samples to be selected in one AL cycle.
+        return_utilities : bool, default=False
+            If true, also return the utilities based on the query strategy.
         **query_kwargs : dict-like
             Further keyword arguments are passed to the `query` method of the
             `query_strategy` object.
@@ -91,23 +99,23 @@ def query(
         Returns
         -------
         query_indices : numpy.ndarray of shape (batch_size)
-           The query_indices indicate for which candidate sample a label is
-           to queried, e.g., `query_indices[0]` indicates the first selected
-           sample.
-           If candidates is None or of shape (n_candidates), the indexing
-           refers to samples in X.
-           If candidates is of shape (n_candidates, n_features), the indexing
-           refers to samples in candidates.
+            The query_indices indicate for which candidate sample a label is
+            to queried, e.g., `query_indices[0]` indicates the first selected
+            sample.
+            If candidates is None or of shape (n_candidates), the indexing
+            refers to samples in X.
+            If candidates is of shape (n_candidates, n_features), the indexing
+            refers to samples in candidates.
         utilities : numpy.ndarray of shape (batch_size, n_samples) or
-           numpy.ndarray of shape (batch_size, n_candidates)
-           The utilities of samples after each selected sample of the batch,
-           e.g., `utilities[0]` indicates the utilities used for selecting
-           the first sample (with index `query_indices[0]`) of the batch.
-           Utilities for labeled samples will be set to np.nan.
-           If candidates is None or of shape (n_candidates), the indexing
-           refers to samples in X.
-           If candidates is of shape (n_candidates, n_features), the indexing
-           refers to samples in candidates.
+            numpy.ndarray of shape (batch_size, n_candidates)
+            The utilities of samples after each selected sample of the batch,
+            e.g., `utilities[0]` indicates the utilities used for selecting
+            the first sample (with index `query_indices[0]`) of the batch.
+            Utilities for labeled samples will be set to np.nan.
+            If candidates is None or of shape (n_candidates), the indexing
+            refers to samples in X.
+            If candidates is of shape (n_candidates, n_features), the indexing
+            refers to samples in candidates.
         """
 
         X, y, candidates, batch_size, return_utilities = self._validate_data(
@@ -120,6 +128,7 @@ def query(
                 f"`query_strategy` is of type `{type(self.query_strategy)}` "
                 f"but must be of type `SingleAnnotatorPoolQueryStrategy`."
             )
+        check_scalar(self.exclude_non_subsample, "exclude_non_subsample", bool)
         seed_multiplier = (
             int(is_labeled(y, missing_label=self.missing_label_).sum()) + 1
         )
@@ -149,65 +158,128 @@ def query(
             )
         random_state = check_random_state(self.random_state, seed_multiplier)
 
+        # subsampling with no explicit provided candidates
         if candidates is None:
             candidate_indices = unlabeled_indices(
                 y=y, missing_label=self.missing_label_
             )
+            # transform max_candidates to int if a ratio is given
             if isinstance(max_candidates, float):
                 max_candidates = ceil(
                     len(candidate_indices) * self.max_candidates
                 )
             max_candidates = min(max_candidates, len(candidate_indices))
+            # subsample new candidates
             new_candidates = random_state.choice(
                 a=candidate_indices, size=max_candidates, replace=False
             )
+        # subsampling with privided explicit candidates
         else:
+            # transform max_candidates to int if a ratio is given
             if isinstance(max_candidates, float):
                 max_candidates = ceil(len(candidates) * self.max_candidates)
             max_candidates = min(max_candidates, len(candidates))
             if candidates.ndim == 1:
                 candidate_indices = candidates
+                # subsample new candidates
                 new_candidates = random_state.choice(
                     a=candidates, size=max_candidates, replace=False
                 )
             else:
                 candidate_indices = range(len(candidates))
+                # subsample new candidates
                 new_candidate_indices = random_state.choice(
                     a=candidate_indices, size=max_candidates, replace=False
                 )
                 new_candidates = candidates[new_candidate_indices]
 
+        # check if to exclude unlabeled non-candidate training data
+        if self.exclude_non_subsample:
+            all_labeled = labeled_indices(
+                y=y, missing_label=self.missing_label_
+            )
+            if candidates is not None and candidates.ndim > 1:
+                subset_and_labeled_indices = all_labeled
+            else:
+                subset_and_labeled_indices = np.concatenate(
+                    [all_labeled, new_candidates]
+                )
+            subset_and_labeled_indices = np.sort(subset_and_labeled_indices)
+
+            new_X = X[subset_and_labeled_indices]
+            new_y = y[subset_and_labeled_indices]
+            # for explicitely provided candidates recalculate candidate indices
+            # that are passed to the wrapped query strategy
+            if candidates is None or candidates.ndim == 1:
+                new_candidates = unlabeled_indices(
+                    y=new_y, missing_label=self.missing_label_
+                )
+        else:
+            new_X = X
+            new_y = y
+
         qs_output = self.query_strategy.query(
-            X=X,
-            y=y,
+            X=new_X,
+            y=new_y,
             candidates=new_candidates,
             batch_size=batch_size,
             return_utilities=return_utilities,
             **query_kwargs,
         )
 
-        if not return_utilities:
-            if candidates is not None and candidates.ndim > 1:
-                return new_candidate_indices[qs_output]
-            return qs_output
-
-        queried_indices, utilities = qs_output
+        # unpack result of query strategy if needed
+        queried_indices = qs_output
+        utilities = None
+        if return_utilities:
+            queried_indices, utilities = qs_output
 
-        if candidates is None or candidates.ndim == 1:
-            new_utilities = np.full(
-                shape=(batch_size, len(X)), fill_value=np.nan
-            )
-            new_utilities[:, candidate_indices] = -np.inf
-            new_utilities[:, new_candidates] = utilities[:, new_candidates]
+        # retransform queried indices and utilities as if no training data was
+        # removed
+        if self.exclude_non_subsample and (
+            candidates is None or candidates.ndim == 1
+        ):
+            # transform to original candidate indices
+            queried_indices = subset_and_labeled_indices[queried_indices]
+            # transform to original utilities shape
+            if utilities is not None:
+                new_utilities = np.full(
+                    shape=(batch_size, len(X)), fill_value=np.nan
+                )
+                transformed_new_candidates = subset_and_labeled_indices[
+                    new_candidates
+                ]
+                new_utilities[:, transformed_new_candidates] = utilities[
+                    :, new_candidates
+                ]
+                utilities = new_utilities
+                new_candidates = transformed_new_candidates
+
+        # transform indices if candidates was provided in the shape of
+        # (n_candidates, n_features)
+        if candidates is not None and candidates.ndim > 1:
+            new_queried_indices = new_candidate_indices[queried_indices]
         else:
-            new_utilities = np.full(
-                shape=(batch_size, len(candidates)), fill_value=np.nan
-            )
-            new_utilities[:, candidate_indices] = -np.inf
-            new_utilities[:, new_candidate_indices] = utilities
-            queried_indices = new_candidate_indices[queried_indices]
+            new_queried_indices = queried_indices
 
-        return queried_indices, new_utilities
+        # transform utilities from subsampled shape to original utilities shape
+        if return_utilities:
+            if candidates is None or candidates.ndim == 1:
+                new_utilities = np.full(
+                    shape=(batch_size, len(X)), fill_value=np.nan
+                )
+                new_utilities[:, candidate_indices] = -np.inf
+                new_utilities[:, new_candidates] = utilities[:, new_candidates]
+            else:
+                new_utilities = np.full(
+                    shape=(batch_size, len(candidates)), fill_value=np.nan
+                )
+                new_utilities[:, candidate_indices] = -np.inf
+                new_utilities[:, new_candidate_indices] = utilities
+
+        if return_utilities:
+            return new_queried_indices, new_utilities
+        else:
+            return new_queried_indices
 
 
 class ParallelUtilityEstimationWrapper(SingleAnnotatorPoolQueryStrategy):
@@ -267,26 +339,25 @@ def query(
         Parameters
         ----------
         X : array-like of shape (n_samples, n_features)
-           Training data set, usually complete, i.e. including the labeled and
-           unlabeled samples.
+            Training data set, usually complete, i.e. including the labeled and
+            unlabeled samples.
         y : array-like of shape (n_samples)
-           Labels of the training data set (possibly including unlabeled ones
-           indicated by self.MISSING_LABEL).
+            Labels of the training data set (possibly including unlabeled ones
+            indicated by self.MISSING_LABEL).
         candidates : None or array-like of shape (n_candidates), dtype=int or
-           array-like of shape (n_candidates, n_features),
-           optional (default=None)
-           If candidates is None, the unlabeled samples from (X,y) are
-           considered as candidates.
-           If candidates is of shape (n_candidates) and of type int,
-           candidates is considered as the indices of the samples in (X,y).
-           If candidates is of shape (n_candidates, n_features), the
-           candidates are directly given in candidates (not necessarily
-           contained in X). This is not supported by all query strategies.
-        batch_size : int, optional (default=1)
-           The number of samples to be selected in one AL cycle. For this
-           wrapper, only `batch_size=1` is supported.
-        return_utilities : bool, optional (default=False)
-           If true, also return the utilities based on the query strategy.
+            array-like of shape (n_candidates, n_features), (default=None)
+            If candidates is None, the unlabeled samples from (X,y) are
+            considered as candidates.
+            If candidates is of shape (n_candidates) and of type int,
+            candidates is considered as the indices of the samples in (X,y).
+            If candidates is of shape (n_candidates, n_features), the
+            candidates are directly given in candidates (not necessarily
+            contained in X). This is not supported by all query strategies.
+        batch_size : int, default=1
+            The number of samples to be selected in one AL cycle. For this
+            wrapper, only `batch_size=1` is supported.
+        return_utilities : bool, default=False
+            If true, also return the utilities based on the query strategy.
         **query_kwargs : dict-like
             Further keyword arguments are passed to the `query` method of the
             `query_strategy` object.
@@ -294,23 +365,23 @@ def query(
         Returns
         -------
         query_indices : numpy.ndarray of shape (batch_size)
-           The query_indices indicate for which candidate sample a label is
-           to queried, e.g., `query_indices[0]` indicates the first selected
-           sample.
-           If candidates is None or of shape (n_candidates), the indexing
-           refers to samples in X.
-           If candidates is of shape (n_candidates, n_features), the indexing
-           refers to samples in candidates.
+            The query_indices indicate for which candidate sample a label is
+            to queried, e.g., `query_indices[0]` indicates the first selected
+            sample.
+            If candidates is None or of shape (n_candidates), the indexing
+            refers to samples in X.
+            If candidates is of shape (n_candidates, n_features), the indexing
+            refers to samples in candidates.
         utilities : numpy.ndarray of shape (batch_size, n_samples) or
-           numpy.ndarray of shape (batch_size, n_candidates)
-           The utilities of samples after each selected sample of the batch,
-           e.g., `utilities[0]` indicates the utilities used for selecting
-           the first sample (with index `query_indices[0]`) of the batch.
-           Utilities for labeled samples will be set to np.nan.
-           If candidates is None or of shape (n_candidates), the indexing
-           refers to samples in X.
-           If candidates is of shape (n_candidates, n_features), the indexing
-           refers to samples in candidates.
+            numpy.ndarray of shape (batch_size, n_candidates)
+            The utilities of samples after each selected sample of the batch,
+            e.g., `utilities[0]` indicates the utilities used for selecting
+            the first sample (with index `query_indices[0]`) of the batch.
+            Utilities for labeled samples will be set to np.nan.
+            If candidates is None or of shape (n_candidates), the indexing
+            refers to samples in X.
+            If candidates is of shape (n_candidates, n_features), the indexing
+            refers to samples in candidates.
         """
 
         X, y, candidates, batch_size, return_utilities = self._validate_data(


Test patch:
diff --git a/skactiveml/pool/tests/test_wrapper.py b/skactiveml/pool/tests/test_wrapper.py
index d2c78d4e..fb11a7dd 100644
--- a/skactiveml/pool/tests/test_wrapper.py
+++ b/skactiveml/pool/tests/test_wrapper.py
@@ -19,7 +19,7 @@
 from skactiveml.tests.template_query_strategy import (
     TemplateSingleAnnotatorPoolQueryStrategy,
 )
-from skactiveml.utils import MISSING_LABEL
+from skactiveml.utils import MISSING_LABEL, unlabeled_indices
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
 
 
@@ -71,6 +71,7 @@ def setUp(self):
         init_default_params = {
             "query_strategy": QueryByCommittee(random_state=0),
             "max_candidates": 10,
+            "exclude_non_subsample": False,
             "random_state": 0,
             "missing_label": MISSING_LABEL,
         }
@@ -104,6 +105,18 @@ def test_init_param_query_strategy(self, test_cases=None):
         ]
         self._test_param("init", "query_strategy", test_cases)
 
+    def test_init_param_exclude_non_subsample(self, test_cases=None):
+        test_cases = [] if test_cases is None else test_cases
+        test_cases += [
+            (True, None),
+            (False, None),
+            (0, TypeError),
+            (1, TypeError),
+            ("1.2", TypeError),
+            (DummyNonQueryStrategy(), TypeError),
+        ]
+        self._test_param("init", "exclude_non_subsample", test_cases)
+
     def test_query_param_query_kwargs(self, test_cases=None):
         test_cases = [] if test_cases is None else test_cases
         test_cases += [(2, TypeError), (True, None), ("Hello", TypeError)]
@@ -112,6 +125,7 @@ def test_query_param_query_kwargs(self, test_cases=None):
         self._test_param("query", "fit_ensemble", test_cases)
 
     def test_query(self):
+        # check consistency with wrapped and non-wrapped query strategy
         qs = deepcopy(self.init_default_params["query_strategy"])
         qs_sub = self.qs_class(**deepcopy(self.init_default_params))
 
@@ -136,6 +150,49 @@ def test_query(self):
             q_sub = qs_sub.query(**query_params)
             self.assertEqual(len(q_sub), 1)
 
+        # check consistency of exclude_non_subsample with varying candidates
+        # and batch_sizes for classification and regression
+        for query_params in [
+            self.query_default_params_clf,
+            self.query_default_params_reg,
+        ]:
+            init_params_base = deepcopy(self.init_default_params)
+            init_params_base.pop("exclude_non_subsample")
+            candidate_indices = unlabeled_indices(
+                query_params["y"], init_params_base["missing_label"]
+            )
+            candidates_list = [
+                None,
+                candidate_indices,
+                query_params["X"][candidate_indices],
+            ]
+            for batch_size in [1, 3]:
+                for candidates in candidates_list:
+                    if query_params is not None:
+                        query_params_base = deepcopy(query_params)
+                        query_params_base["return_utilities"] = True
+                        query_params_base["candidates"] = candidates
+                        query_params_base["batch_size"] = batch_size
+                        qs_false = SubSamplingWrapper(
+                            exclude_non_subsample=False, **init_params_base
+                        )
+                        qs_true = SubSamplingWrapper(
+                            exclude_non_subsample=True, **init_params_base
+                        )
+                        query_indices_false, utilities_false = qs_false.query(
+                            **query_params_base
+                        )
+                        query_indices_true, utilities_true = qs_true.query(
+                            **query_params_base
+                        )
+
+                        np.testing.assert_array_equal(
+                            query_indices_false, query_indices_true
+                        )
+                        np.testing.assert_array_equal(
+                            utilities_false, utilities_true
+                        )
+
         us = UncertaintySampling()
         qs_us = SubSamplingWrapper(us)
         sig_qs_us = inspect.signature(qs_us.query).parameters