[MNT, ENH, DOC] Rework similarity search

aeon/similarity_search/collection/__init__.py

@@ -1,6 +1,9 @@
 """Similarity search for time series collection."""
 
-__all__ = ["BaseCollectionSimilaritySearch", "RandomProjectionIndexANN"]
+__all__ = [
+    "BaseCollectionSimilaritySearch",
+    "RandomProjectionIndexANN",
+]
 
 from aeon.similarity_search.collection._base import BaseCollectionSimilaritySearch
 from aeon.similarity_search.collection.neighbors._rp_cosine_lsh import (

aeon/similarity_search/collection/_base.py

@@ -57,9 +57,8 @@ def fit(
         """
         self.reset()
         X = self._preprocess_collection(X)
-        # Store minimum number of n_timepoints for unequal length collections
-        self.n_channels_ = X[0].shape[0]
-        self.n_cases_ = len(X)
+        self.n_channels_ = self.metadata_["n_channels"]
+        self.n_cases_ = self.metadata_["n_cases"]
         self._fit(X, y=y)
         self.is_fitted = True
         return self
@@ -89,9 +88,7 @@ def predict(self, X, **kwargs):
 
         """
         self._check_is_fitted()
-        if X[0].ndim == 1:
-            X = X[np.newaxis, :, :]
-        X = self._preprocess_collection(X)
+        X = self._preprocess_collection(X, store_metadata=False)
         self._check_predict_series_format(X)
         indexes, distances = self._predict(X, **kwargs)
         return indexes, distances
@@ -105,16 +102,6 @@ def _check_predict_series_format(self, X):
         X : np.ndarray, shape = (n_channels, n_timepoints)
             A series to be used in predict.
         """
-        if isinstance(X, np.ndarray):
-            if X[0].ndim != 2:
-                raise TypeError(
-                    "A np.ndarray given in predict must be 3D"
-                    f"(n_channels, n_timepoints) but found {X.ndim}D."
-                )
-        else:
-            raise TypeError(
-                "Expected a 3D np.ndarray in predict but found" f" {type(X)}."
-            )
         if self.n_channels_ != X[0].shape[0]:
             raise ValueError(
                 f"Expected X to have {self.n_channels_} channels but"

aeon/similarity_search/collection/neighbors/_rp_cosine_lsh.py

@@ -140,6 +140,7 @@ class RandomProjectionIndexANN(BaseCollectionSimilaritySearch):
 
     _tags = {
         "capability:unequal_length": False,
+        "capability:multivariate": True,
         "capability:multithreading": True,
     }
 

aeon/similarity_search/collection/tests/test_base.py

@@ -2,55 +2,18 @@
 
 __maintainer__ = ["baraline"]
 
-import pytest
-
 from aeon.testing.mock_estimators._mock_similarity_searchers import (
     MockCollectionSimilaritySearch,
 )
-from aeon.testing.testing_data import (
-    make_example_1d_numpy,
-    make_example_2d_numpy_series,
-    make_example_3d_numpy,
-)
+from aeon.testing.testing_data import FULL_TEST_DATA_DICT, _get_datatypes_for_estimator
 
 
 def test_input_shape_fit_predict_collection():
     """Test input shapes."""
     estimator = MockCollectionSimilaritySearch()
+    datatypes = _get_datatypes_for_estimator(estimator)
     # dummy data to pass to fit when testing predict/predict_proba
-    X_3D_uni = make_example_3d_numpy(n_channels=1, return_y=False)
-    X_3D_multi = make_example_3d_numpy(n_channels=2, return_y=False)
-    X_2D_uni = make_example_2d_numpy_series(n_channels=1)
-    X_2D_multi = make_example_2d_numpy_series(n_channels=2)
-    X_1D = make_example_1d_numpy()
-
-    # 2D are converted to 3D
-    valid_inputs_fit = [
-        X_3D_uni,
-        X_3D_multi,
-        X_2D_uni,
-        X_2D_multi,
-    ]
-    # Valid inputs
-    for _input in valid_inputs_fit:
-        estimator.fit(_input)
-
-    with pytest.raises(ValueError):
-        estimator.fit(X_1D)
-
-    estimator_multi = MockCollectionSimilaritySearch().fit(X_3D_multi)
-    estimator_uni = MockCollectionSimilaritySearch().fit(X_3D_uni)
-
-    estimator_uni.predict(X_2D_uni)
-    estimator_uni.predict(X_3D_uni)
-    estimator_multi.predict(X_2D_multi)
-    estimator_multi.predict(X_3D_multi)
-
-    with pytest.raises(ValueError):
-        estimator_uni.predict(X_2D_multi)
-    with pytest.raises(ValueError):
-        estimator_multi.predict(X_2D_uni)
-    with pytest.raises(ValueError):
-        estimator_multi.predict(X_3D_uni)
-    with pytest.raises(ValueError):
-        estimator_uni.predict(X_3D_multi)
+    for datatype in datatypes:
+        X_train, y_train = FULL_TEST_DATA_DICT[datatype]["train"]
+        X_test, y_test = FULL_TEST_DATA_DICT[datatype]["test"]
+        estimator.fit(X_train, y_train).predict(X_test)

aeon/similarity_search/series/__init__.py

@@ -1,8 +1,15 @@
 """Similarity search for series."""
 
-__all__ = ["BaseSeriesSimilaritySearch", "MassSNN", "StompMotif", "DummySNN"]
+__all__ = [
+    "BaseSeriesSimilaritySearch",
+    "MassSNN",
+    "StompMotif",
+    "DummySNN",
+]
 
-from aeon.similarity_search.series._base import BaseSeriesSimilaritySearch
+from aeon.similarity_search.series._base import (
+    BaseSeriesSimilaritySearch,
+)
 from aeon.similarity_search.series.motifs._stomp import StompMotif
 from aeon.similarity_search.series.neighbors._dummy import DummySNN
 from aeon.similarity_search.series.neighbors._mass import MassSNN

aeon/similarity_search/series/_base.py

@@ -1,5 +1,8 @@
 """Base similiarity search for series."""
 
+__maintainer__ = ["baraline"]
+__all__ = ["BaseSeriesSimilaritySearch"]
+
 from abc import abstractmethod
 from typing import final
 
@@ -53,9 +56,9 @@ def fit(
         """
         self.reset()
         X = self._preprocess_series(X, self.axis, True)
-        # Store minimum number of n_timepoints for unequal length collections
-        self.n_channels_ = X.shape[0]
-        self.n_timepoints_ = X.shape[1]
+        self.n_channels_ = self.metadata_["n_channels"]
+        timepoint_idx = 1 if self.axis == 1 else 0
+        self.n_timepoints_ = X.shape[timepoint_idx]
         self.X_ = X
         self._fit(X, y=y)
         self.is_fitted = True
@@ -69,7 +72,7 @@ def _fit(
     ): ...
 
     @final
-    def predict(self, X=None, **kwargs):
+    def predict(self, X, **kwargs):
         """
         Predict function.
 
@@ -79,7 +82,7 @@ def predict(self, X=None, **kwargs):
             Series to predict on.
         kwargs : dict, optional
             Additional keyword argument as dict or individual keywords args
-            to pass to use.
+            to pass to the estimator.
 
         Returns
         -------
@@ -90,41 +93,14 @@ def predict(self, X=None, **kwargs):
 
         """
         self._check_is_fitted()
-        if X is not None:
-            X = self._preprocess_series(X, self.axis, False)
-            self._check_predict_series_format(X)
-        else:
-            X = self.X_
+        X = self._preprocess_series(X, self.axis, False)
+        self._check_predict_series_format(X)
         indexes, distances = self._predict(X, **kwargs)
         return indexes, distances
 
     @abstractmethod
     def _predict(self, X, **kwargs): ...
 
-    def _check_X_index(self, X_index: int):
-        """
-        Check wheter a X_index parameter is correctly formated and is admissible.
-
-        Parameters
-        ----------
-        X_index : int
-            Index of a timestamp in X_.
-
-        """
-        if X_index is not None:
-            if not isinstance(X_index, int):
-                raise TypeError("Expected an integer for X_index but got {X_index}")
-
-            max_timepoints = self.n_timepoints_
-            if hasattr(self, "length"):
-                max_timepoints -= self.length
-            if X_index >= max_timepoints or X_index < 0:
-                raise ValueError(
-                    "The value of X_index cannot exced the number "
-                    "of timepoint in series given during fit. Expected a value "
-                    f"between [0, {max_timepoints - 1}] but got {X_index}"
-                )
-
     def _check_predict_series_format(self, X):
         """
         Check wheter a series X is correctly formated regarding series given in fit.
@@ -135,8 +111,9 @@ def _check_predict_series_format(self, X):
             A series to be used in predict.
 
         """
-        if self.n_channels_ != X.shape[0]:
+        channel_idx = 0 if self.axis == 1 else 1
+        if self.n_channels_ != X.shape[channel_idx]:
             raise ValueError(
                 f"Expected X to have {self.n_channels_} channels but"
-                f" got {X.shape[0]} channels."
+                f" got {X.shape[channel_idx]} channels."
             )

aeon/similarity_search/series/_commons.py

@@ -9,6 +9,33 @@
 from aeon.utils.numba.general import AEON_NUMBA_STD_THRESHOLD
 
 
+def _check_X_index(X_index: int, n_timepoints: int, length: int):
+    """
+    Check wheter a X_index parameter is correctly formated and is admissible.
+
+    Parameters
+    ----------
+    X_index : int
+        Index of a timestamp in X_.
+     n_timepoints: int
+         Number of timepoints in the serie X_
+     length: int
+         Length parameter of the estimator
+
+    """
+    if X_index is not None:
+        if not isinstance(X_index, int):
+            raise TypeError("Expected an integer for X_index but got {X_index}")
+
+        max_timepoints = n_timepoints - length
+        if X_index >= max_timepoints or X_index < 0:
+            raise ValueError(
+                "The value of X_index cannot exced the number "
+                "of timepoint in series given during fit. Expected a value "
+                f"between [0, {max_timepoints - 1}] but got {X_index}"
+            )
+
+
 def fft_sliding_dot_product(X, q):
     """
     Use FFT convolution to calculate the sliding window dot product.
@@ -161,7 +188,10 @@ def _extract_top_k_motifs(MP, IP, k, allow_trivial_matches, exclusion_size):
     idx, _ = _extract_top_k_from_dist_profile(
         criterion, k, np.inf, allow_trivial_matches, exclusion_size
     )
-    return [MP[i] for i in idx], [IP[i] for i in idx]
+    return (
+        [IP[i] for i in idx],
+        [MP[i] for i in idx],
+    )
 
 
 def _extract_top_r_motifs(MP, IP, k, allow_trivial_matches, exclusion_size):
@@ -175,7 +205,7 @@ def _extract_top_r_motifs(MP, IP, k, allow_trivial_matches, exclusion_size):
         allow_trivial_matches,
         exclusion_size,
     )
-    return [MP[i] for i in idx], [IP[i] for i in idx]
+    return [IP[i] for i in idx], [MP[i] for i in idx]
 
 
 @njit(cache=True, fastmath=True)