[ENH] Add capability:predict tag for Clustering

aeon/clustering/base.py

@@ -24,6 +24,7 @@ class BaseClusterer(ClusterMixin, BaseCollectionEstimator):
 
     _tags = {
         "fit_is_empty": False,
+        "capability:predict": True,
     }
 
     @abstractmethod
@@ -35,6 +36,9 @@ def __init__(self):
     def fit(self, X, y=None) -> BaseCollectionEstimator:
         """Fit time series clusterer to training data.
 
+        Clusters the collection ``X`` and stores the cluster index of each time
+        series in the ``labels_`` attribute.
+
         Parameters
         ----------
         X : 3D np.ndarray (any number of channels, equal length series)
@@ -62,6 +66,11 @@ def fit(self, X, y=None) -> BaseCollectionEstimator:
     def predict(self, X) -> np.ndarray:
         """Predict the closest cluster each sample in X belongs to.
 
+        Only clusterers with the ``capability:predict`` tag set to True support
+        assigning previously unseen cases to clusters. Transductive clusterers
+        (tag set to False) can only cluster the collection they are fitted on;
+        use ``fit_predict(X)`` or inspect ``labels_`` after ``fit(X)`` instead.
+
         Parameters
         ----------
         X : 3D np.ndarray
@@ -79,8 +88,15 @@ def predict(self, X) -> np.ndarray:
         np.array
             shape ``(n_cases,)``, index of the cluster to which each time series in X
             belongs.
+
+        Raises
+        ------
+        NotImplementedError
+            If the ``capability:predict`` tag is False, as the clusterer does not
+            support out-of-sample prediction.
         """
         self._check_is_fitted()
+        self._check_predict_capability()
         X = self._preprocess_collection(X, store_metadata=False)
         self._check_shape(X)
         return self._predict(X)
@@ -111,17 +127,28 @@ def predict_proba(self, X) -> np.ndarray:
             1st dimension indices correspond to instance indices in X
             2nd dimension indices correspond to possible labels (integers)
             (i, j)-th entry is predictive probability that i-th instance is of class j
+
+        Raises
+        ------
+        NotImplementedError
+            If the ``capability:predict`` tag is False, as the clusterer does not
+            support out-of-sample prediction.
         """
         self._check_is_fitted()
+        self._check_predict_capability()
         X = self._preprocess_collection(X, store_metadata=False)
         self._check_shape(X)
         return self._predict_proba(X)
 
     @final
     def fit_predict(self, X, y=None) -> np.ndarray:
-        """Compute cluster centers and predict cluster index for each time series.
+        """Fit the clusterer to X and return the cluster index of each time series.
 
-        Convenience method; equivalent of calling fit(X) followed by predict(X)
+        Convenience method; equivalent of calling ``fit(X)`` and returning
+        ``labels_``, the cluster indices assigned to the fitted collection. It
+        does not call ``predict(X)``, so it is valid for all clusterers,
+        including transductive ones that do not support out-of-sample
+        prediction (``capability:predict`` tag set to False).
 
         Parameters
         ----------
@@ -136,9 +163,18 @@ def fit_predict(self, X, y=None) -> np.ndarray:
         np.ndarray (1d array of shape (n_cases,))
             Index of the cluster each time series in X belongs to.
         """
-        self.fit(X)
+        self.fit(X, y)
         return self.labels_
 
+    def _check_predict_capability(self):
+        """Raise an error if the clusterer cannot predict on unseen data."""
+        if not self.get_tag("capability:predict"):
+            raise NotImplementedError(
+                f"{self.__class__.__name__} does not support out-of-sample "
+                "prediction. Use fit_predict(X) to cluster a collection, or "
+                "inspect labels_ after fit(X)."
+            )
+
     def _predict_proba(self, X) -> np.ndarray:
         """Predicts labels probabilities for sequences in X.
 

aeon/clustering/tests/test_base.py

@@ -1,65 +1,189 @@
 """Unit tests for clustering base class functionality."""
 
 import numpy as np
-import numpy.random
 import pytest
+from sklearn.exceptions import NotFittedError
 
 from aeon.clustering.base import BaseClusterer
-from aeon.testing.mock_estimators import MockCluster
+from aeon.testing.mock_estimators import MockCluster, MockTransductiveCluster
 
 
 def test_correct_input():
-    """Tests errors raised with wrong inputs: X and/or y."""
+    """Test fit rejects invalid X and accepts a valid 2D array.
+
+    Exercises the input validation in ``BaseClusterer.fit``: a list of strings,
+    a dict, and a 2D list of ints each raise a ``TypeError`` with a specific
+    message, while a valid 2D numpy array fits and can then predict.
+    """
     dummy = MockCluster()
 
     X = ["list", "of", "invalid", "test", "strings"]
-    msg1 = r"ERROR passed a list containing <class 'str'>"
-    with pytest.raises(TypeError, match=msg1):
+    msg = r"ERROR passed a list containing <class 'str'>"
+    with pytest.raises(TypeError, match=msg):
         dummy.fit(X)
 
-    # dict X
-    X = {
-        0: "invalid",
-        1: "input",
-        2: "dict",
-    }
-    msg2 = r"ERROR passed input of type <class 'dict'>"
-    with pytest.raises(TypeError, match=msg2):
+    X = {0: "invalid", 1: "input", 2: "dict"}
+    with pytest.raises(TypeError, match=r"ERROR passed input of type <class 'dict'>"):
         dummy.fit(X)
 
-    # 2d list of int X
     X = [[1, 1, 1], [1, 3, 4]]
-    msg3 = r"lists should either 2D numpy arrays or pd.DataFrames"
-    with pytest.raises(TypeError, match=msg3):
+    with pytest.raises(TypeError, match=r"lists should either 2D numpy arrays"):
         dummy.fit(X)
 
-    # correct X
     X = np.random.randn(5, 5)
     dummy.fit(X)
-    assert (dummy.predict(X)).shape == (5,)
-    assert (dummy.predict_proba(X)).shape == (5,)
+    assert dummy.predict(X).shape == (5,)
+    assert dummy.predict_proba(X).shape == (5,)
+
 
+class _DefaultProbaClusterer(BaseClusterer):
+    """Clusterer that assigns every case to cluster 0.
 
-class _TestClusterer(BaseClusterer):
-    """Clusterer for testing base class fit/predict/predict_proba."""
+    It overrides only ``_fit`` and ``_predict`` and deliberately does not
+    override ``_predict_proba``, so it exercises the ``BaseClusterer`` default
+    ``_predict_proba`` (which one-hot encodes the ``_predict`` output).
+    """
 
     def __init__(self):
         super().__init__()
 
     def _fit(self, X, y=None):
-        """Fit dummy."""
+        self.labels_ = np.zeros(len(X), dtype=int)
         return self
 
     def _predict(self, X):
-        """Predict dummy."""
-        return np.zeros(shape=(len(X),), dtype=int)
+        return np.zeros(len(X), dtype=int)
+
 
+def test_base_predict_proba_is_one_hot_of_predict():
+    """Test the base default ``_predict_proba`` one-hot encodes ``_predict``.
 
-def test_base_clusterer():
-    """Test with no clusters."""
-    clst = _TestClusterer()
+    A clusterer that does not override ``_predict_proba`` falls back to the
+    base implementation: run ``_predict`` and put probability 1 on the assigned
+    cluster. With a single cluster this is an ``(n_cases, 1)`` column of ones.
+    """
     X = np.random.random(size=(10, 1, 20))
+    clst = _DefaultProbaClusterer().fit(X)
+    proba = clst._predict_proba(X)
+    assert proba.shape == (10, 1)
+    assert np.all(proba == 1.0)
+
+
+class _NClustersClusterer(BaseClusterer):
+    """Clusterer exposing an ``n_clusters`` attribute, assigning cases round-robin.
+
+    The base ``_predict_proba`` sizes its probability matrix from ``n_clusters``
+    when that attribute is present, and only falls back to ``max(_predict) + 1``
+    when it is ``None``. This mock lets a test drive both branches; ``_predict``
+    returns three distinct clusters (labels 0, 1, 2).
+    """
+
+    def __init__(self, n_clusters):
+        self.n_clusters = n_clusters
+        super().__init__()
+
+    def _fit(self, X, y=None):
+        self.labels_ = np.arange(len(X)) % 3
+        return self
+
+    def _predict(self, X):
+        return np.arange(len(X)) % 3
+
+
+def test_base_predict_proba_width_from_n_clusters_attribute():
+    """Test base ``_predict_proba`` matrix width follows the ``n_clusters`` attribute.
+
+    When a clusterer defines ``n_clusters``, the probability matrix has exactly
+    that many columns (even if wider than the number of clusters actually
+    predicted). When ``n_clusters`` is ``None`` the base falls back to
+    ``max(_predict) + 1`` (here 3, for labels 0, 1, 2). Either way each row is a
+    one-hot of the assigned cluster.
+    """
+    X = np.random.random(size=(9, 1, 20))
+    expected = np.arange(9) % 3
+
+    # explicit n_clusters -> that many columns (line: n_clusters = self.n_clusters)
+    proba = _NClustersClusterer(n_clusters=4).fit(X)._predict_proba(X)
+    assert proba.shape == (9, 4)
+    assert np.all(proba.sum(axis=1) == 1)
+    assert np.array_equal(np.argmax(proba, axis=1), expected)
+
+    # n_clusters is None -> fall back to max(_predict) + 1 == 3
+    proba = _NClustersClusterer(n_clusters=None).fit(X)._predict_proba(X)
+    assert proba.shape == (9, 3)
+    assert np.array_equal(np.argmax(proba, axis=1), expected)
+
+
+def test_fit_predict_returns_labels_from_its_own_fit():
+    """Test ``fit_predict(X)`` returns ``labels_`` from the fit it performs.
+
+    ``fit_predict`` is defined as ``fit(X)`` followed by returning ``labels_``
+    (the assignment of the fitted collection); it must not return ``predict(X)``.
+    ``MockTransductiveCluster`` has a known ``labels_`` pattern (alternating 0/1)
+    and a ``_predict`` that raises, so matching that pattern shows ``fit_predict``
+    returned ``labels_`` rather than routing through ``predict``. The check is
+    repeated on ``MockCluster`` to confirm the same for a predict-capable
+    clusterer.
+    """
+    X = np.random.random(size=(10, 1, 20))
+
+    labels = MockTransductiveCluster().fit_predict(X)
+    assert np.array_equal(labels, np.arange(10) % 2)
+
+    clst = MockCluster()
+    assert np.array_equal(clst.fit_predict(X), clst.labels_)
+
+
+def test_fit_predict_valid_without_predict_capability():
+    """Test ``fit_predict`` works for a clusterer without out-of-sample predict.
+
+    Because ``fit_predict`` never routes through ``_predict``, it is valid even
+    when ``capability:predict`` is False. ``MockTransductiveCluster._predict``
+    raises, so ``fit_predict`` completing normally (rather than erroring) proves
+    ``_predict`` was never called.
+    """
+    X = np.random.random(size=(10, 1, 20))
+    clst = MockTransductiveCluster()
+    assert not clst.get_tag("capability:predict")
+
+    labels = clst.fit_predict(X)
+    assert labels.shape == (10,)
+
+
+def test_predict_without_capability_raises():
+    """Test ``predict``/``predict_proba`` raise for a transductive clusterer.
+
+    When ``capability:predict`` is False, ``predict`` and ``predict_proba`` must
+    raise ``NotFittedError`` before fit and ``NotImplementedError`` after fit
+    (out-of-sample prediction is unsupported). Only the distinctive part of the
+    message is matched, so the test is robust to wording of the accompanying
+    guidance text.
+    """
+    X = np.random.random(size=(10, 1, 20))
+    clst = MockTransductiveCluster()
+    assert not clst.get_tag("capability:predict")
+
+    with pytest.raises(NotFittedError, match="has not been fitted"):
+        clst.predict(X)
+
+    clst.fit(X)
+    msg = "does not support out-of-sample prediction"
+    with pytest.raises(NotImplementedError, match=msg):
+        clst.predict(X)
+    with pytest.raises(NotImplementedError, match=msg):
+        clst.predict_proba(X)
+
+
+def test_predict_with_capability():
+    """Test ``predict``/``predict_proba`` work when ``capability:predict`` is True.
+
+    The default, non-transductive case: after ``fit`` both ``predict`` and
+    ``predict_proba`` return one row per case without raising.
+    """
+    X = np.random.random(size=(10, 1, 20))
+    clst = MockCluster()
+    assert clst.get_tag("capability:predict")
+
     clst.fit(X)
-    assert clst.is_fitted
-    preds = clst._predict_proba(X)
-    assert preds.shape == (10, 1)
+    assert clst.predict(X).shape == (10,)
+    assert clst.predict_proba(X).shape == (10,)

aeon/testing/estimator_checking/_yield_clustering_checks.py

@@ -60,8 +60,9 @@ def check_clusterer_tags_consistent(estimator_class):
         X = np.random.random((10, 2, 10))
         inst = estimator_class._create_test_instance(parameter_set="default")
         inst.fit(X)
-        inst.predict(X)
-        inst.predict_proba(X)
+        if estimator_class.get_class_tag("capability:predict"):
+            inst.predict(X)
+            inst.predict_proba(X)
 
 
 def check_clusterer_does_not_override_final_methods(estimator_class):
@@ -117,16 +118,41 @@ def check_clustering_random_state_deep_learning(estimator, datatype):
 def check_clusterer_output(estimator, datatype):
     """Test clusterer outputs the correct data types and values.
 
-    Test predict produces a np.array or pd.Series with only values seen in the train
-    data, and that predict_proba probability estimates add up to one.
+    Test fit sets labels_ and fit_predict returns it. For clusterers which support
+    out-of-sample prediction, test predict produces a np.array or pd.Series with only
+    values seen in the train data, and that predict_proba probability estimates add
+    up to one. For transductive clusterers (capability:predict tag set to False),
+    test predict and predict_proba raise a NotImplementedError instead.
     """
+    import pytest
+
     estimator = _clone_estimator(estimator)
 
-    # run fit and predict
+    # run fit and check labels_
     data = FULL_TEST_DATA_DICT[datatype]["train"][0]
     estimator.fit(data)
     assert hasattr(estimator, "labels_")
     assert isinstance(estimator.labels_, np.ndarray)
+    assert estimator.labels_.shape == (get_n_cases(data),)
+
+    if not estimator.get_tag("capability:predict"):
+        # transductive clusterers cannot assign clusters to unseen data
+        with pytest.raises(
+            NotImplementedError, match="does not support out-of-sample prediction"
+        ):
+            estimator.predict(data)
+        with pytest.raises(
+            NotImplementedError, match="does not support out-of-sample prediction"
+        ):
+            estimator.predict_proba(data)
+
+        # fit_predict must still work, returning the labels_ of a fit on X
+        labels = estimator.fit_predict(data)
+        assert isinstance(labels, np.ndarray)
+        assert labels.shape == (get_n_cases(data),)
+        assert np.array_equal(labels, estimator.labels_)
+        return
+
     assert np.array_equal(estimator.labels_, estimator.predict(data))
 
     y_pred = estimator.predict(FULL_TEST_DATA_DICT[datatype]["test"][0])