[ENH] KNN n_jobs and updated kneighbours method

aeon/classification/distance_based/_time_series_neighbors.py

@@ -5,16 +5,18 @@
 distances in aeon.distances.
 """
 
+import numbers
+from typing import Optional
+
 __maintainer__ = []
 __all__ = ["KNeighborsTimeSeriesClassifier"]
 
-from typing import Callable, Optional, Union
+from typing import Callable, Union
 
 import numpy as np
-from joblib import Parallel, delayed
 
 from aeon.classification.base import BaseClassifier
-from aeon.distances import get_distance_function
+from aeon.distances import pairwise_distance
 from aeon.utils.validation import check_n_jobs
 
 WEIGHTS_SUPPORTED = ["uniform", "distance"]
@@ -47,21 +49,17 @@ class KNeighborsTimeSeriesClassifier(BaseClassifier):
         n_timepoints)` as input and returns a ``float``.
     distance_params : dict, default = None
         Dictionary for metric parameters for the case that distance is a ``str``.
-    n_jobs : int, default = 1
-        The number of parallel jobs to run for neighbors search.
-        ``None`` means 1 unless in a :obj:``joblib.parallel_backend`` context.
-        ``-1`` means using all processors.
-    parallel_backend : str, ParallelBackendBase instance or None, default=None
-        Specify the parallelisation backend implementation in joblib, if None
-        a ‘prefer’ value of “threads” is used by default. Valid options are
-        “loky”, “multiprocessing”, “threading” or a custom backend.
-        See the joblib Parallel documentation for more details.
+    n_jobs : int, default=1
+        The number of jobs to run in parallel. If -1, then the number of jobs is set
+        to the number of CPU cores. If 1, then the function is executed in a single
+        thread. If greater than 1, then the function is executed in parallel.
 
     Raises
     ------
     ValueError
         If ``weights`` is not among the supported values.
         See the ``weights`` parameter description for valid options.
+        Dictionary for metric parameters for the case that distance is a str.
 
     Examples
     --------
@@ -90,13 +88,11 @@ def __init__(
         n_neighbors: int = 1,
         weights: Union[str, Callable] = "uniform",
         n_jobs: int = 1,
-        parallel_backend: str = None,
     ) -> None:
         self.distance = distance
         self.distance_params = distance_params
         self.n_neighbors = n_neighbors
         self.n_jobs = n_jobs
-        self.parallel_backend = parallel_backend
 
         self._distance_params = distance_params
         if self._distance_params is None:
@@ -124,7 +120,6 @@ def _fit(self, X, y):
         y : array-like, shape = (n_cases)
             The class labels.
         """
-        self.metric_ = get_distance_function(method=self.distance)
         self.X_ = X
         _, self.y_ = np.unique(y, return_inverse=True)
         self._n_jobs = check_n_jobs(self.n_jobs)
@@ -148,10 +143,26 @@ def _predict_proba(self, X):
             The class probabilities of the input samples. Classes are ordered
             by lexicographic order.
         """
-        preds = Parallel(n_jobs=self._n_jobs, backend=self.parallel_backend)(
-            delayed(self._proba_row)(x) for x in X
-        )
-        return np.array(preds)
+        preds = np.zeros((len(X), len(self.classes_)))
+        for i in range(len(X)):
+            neigh_dist, neigh_ind = self.kneighbors(X[i : i + 1])
+            neigh_dist = neigh_dist[0]
+            neigh_ind = neigh_ind[0]
+
+            if self.weights == "distance":
+                weights = 1 / (neigh_dist + np.finfo(float).eps)
+            elif self.weights == "uniform":
+                weights = np.repeat(1.0, len(neigh_ind))
+            else:
+                raise Exception(f"Invalid kNN weights: {self.weights}")
+
+            for id, w in zip(neigh_ind, weights):
+                predicted_class = self.y_[id]
+                preds[i, predicted_class] += w
+
+            preds[i] = preds[i] / np.sum(preds[i])
+
+        return preds
 
     def _predict(self, X):
         """
@@ -170,70 +181,136 @@ def _predict(self, X):
         y : array of shape (n_cases)
             Class labels for each data sample.
         """
-        preds = Parallel(n_jobs=self._n_jobs, backend=self.parallel_backend)(
-            delayed(self._predict_row)(x) for x in X
+        self._check_is_fitted()
+
+        neigh_ind = self._kneighbors(
+            X, n_neighbors=1, return_distance=False, query_is_train=False
         )
-        return np.array(preds, dtype=self.classes_.dtype)
+        indexes = neigh_ind[:, 0]
+        return self.classes_[self.y_[indexes]]
 
-    def _proba_row(self, x):
-        scores = self._predict_scores(x)
-        return scores / np.sum(scores)
+    def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
+        """Find the K-neighbors of a point.
 
-    def _predict_row(self, x):
-        scores = self._predict_scores(x)
-        return self.classes_[np.argmax(scores)]
+        Returns indices of and distances to the neighbors of each point.
 
-    def _predict_scores(self, x):
-        scores = np.zeros(len(self.classes_))
-        idx, weights = self._kneighbors(x)
-        for id, weight in zip(idx, weights):
-            predicted_class = self.y_[id]
-            scores[predicted_class] += weight
-        return scores
+        Parameters
+        ----------
+        X : 3D np.ndarray of shape = (n_cases, n_channels, n_timepoints) or list of
+        shape [n_cases] of 2D arrays shape (n_channels,n_timepoints_i)
+            The query point or points.
+            If not provided, neighbors of each indexed point are returned.
+            In this case, the query point is not considered its own neighbor.
+        n_neighbors : int, default=None
+            Number of neighbors required for each sample. The default is the value
+            passed to the constructor.
+        return_distance : bool, default=True
+            Whether or not to return the distances.
 
-    def _kneighbors(self, X):
+        Returns
+        -------
+        neigh_dist : ndarray of shape (n_queries, n_neighbors)
+            Array representing the distances to points, only present if
+            return_distance=True.
+        neigh_ind : ndarray of shape (n_queries, n_neighbors)
+            Indices of the nearest points in the population matrix.
         """
-        Find the K-neighbors of a point.
+        self._check_is_fitted()
+
+        # Input validation
+        if n_neighbors is None:
+            n_neighbors = self.n_neighbors
+        elif not isinstance(n_neighbors, numbers.Integral):
+            raise TypeError(
+                f"n_neighbors does not take {type(n_neighbors)} value, enter integer "
+                f"value"
+            )
+        elif n_neighbors <= 0:
+            raise ValueError(f"Expected n_neighbors > 0. Got {n_neighbors}")
+
+        if not isinstance(return_distance, bool):
+            raise TypeError(
+                f"return_distance must be a boolean, got {type(return_distance)}"
+            )
+
+        # Preprocess X if provided
+        query_is_train = X is None
+        if query_is_train:
+            X = self.X_
+        else:
+            X = self._preprocess_collection(X, store_metadata=False)
+            self._check_shape(X)
+
+        # Validate n_neighbors against data size
+        n_samples_fit = len(self.X_)
+        if query_is_train:
+            if not (n_neighbors < n_samples_fit):
+                raise ValueError(
+                    "Expected n_neighbors < n_samples_fit, but "
+                    f"n_neighbors = {n_neighbors}, n_samples_fit = {n_samples_fit}, "
+                    f"n_samples = {len(X)}"
+                )
+        else:
+            if not (n_neighbors <= n_samples_fit):
+                raise ValueError(
+                    "Expected n_neighbors <= n_samples_fit, but "
+                    f"n_neighbors = {n_neighbors}, n_samples_fit = {n_samples_fit}, "
+                    f"n_samples = {len(X)}"
+                )
 
-        Returns indices and weights of each point.
+        return self._kneighbors(X, n_neighbors, return_distance, query_is_train)
+
+    def _kneighbors(self, X, n_neighbors, return_distance, query_is_train):
+        """Find the K-neighbors of a point.
+
+        Returns indices of and distances to the neighbors of each point.
 
         Parameters
         ----------
-        X : np.ndarray
-            A single time series instance if shape = ``(n_channels, n_timepoints)``
+        X : 3D np.ndarray of shape = (n_cases, n_channels, n_timepoints) or list of
+        shape [n_cases] of 2D arrays shape (n_channels,n_timepoints_i)
+            The query point or points.
+        n_neighbors : int
+            Number of neighbors required for each sample.
+        return_distance : bool
+            Whether or not to return the distances.
+        query_is_train : bool
+            Whether the query points are from the training set.
 
         Returns
         -------
-        ind : array
+        neigh_dist : ndarray of shape (n_queries, n_neighbors)
+            Array representing the distances to points, only present if
+            return_distance=True.
+        neigh_ind : ndarray of shape (n_queries, n_neighbors)
             Indices of the nearest points in the population matrix.
-        ws : array
-            Array representing the weights of each neighbor.
         """
-        distances = np.array(
-            [
-                self.metric_(X, self.X_[j], **self._distance_params)
-                for j in range(len(self.X_))
-            ]
+        distances = pairwise_distance(
+            X,
+            None if query_is_train else self.X_,
+            method=self.distance,
+            n_jobs=self.n_jobs,
+            **self._distance_params,
         )
 
-        # Find indices of k nearest neighbors using partitioning:
-        # [0..k-1], [k], [k+1..n-1]
-        # They might not be ordered within themselves,
-        # but it is not necessary and partitioning is
-        # O(n) while sorting is O(nlogn)
-        closest_idx = np.argpartition(distances, self.n_neighbors)
-        closest_idx = closest_idx[: self.n_neighbors]
-
-        if self.weights == "distance":
-            ws = distances[closest_idx]
-            # Using epsilon ~= 0 to avoid division by zero
-            ws = 1 / (ws + np.finfo(float).eps)
-        elif self.weights == "uniform":
-            ws = np.repeat(1.0, self.n_neighbors)
-        else:
-            raise Exception(f"Invalid kNN weights: {self.weights}")
-
-        return closest_idx, ws
+        # If querying the training set, exclude self by setting diag to +inf
+        if query_is_train:
+            np.fill_diagonal(distances, np.inf)
+
+        k = n_neighbors
+        # 1) partial select smallest k
+        idx_part = np.argpartition(distances, kth=k - 1, axis=1)[:, :k]
+        # 2) sort those k by (distance, index)
+        row_idx = np.arange(distances.shape[0])[:, None]
+        part_d = distances[row_idx, idx_part]
+        # argsort by distance, then by index for ties (lexsort uses last key as primary)
+        order = np.lexsort((idx_part, part_d), axis=1)
+        neigh_ind = idx_part[row_idx, order]
+
+        if return_distance:
+            neigh_dist = distances[row_idx, neigh_ind]
+            return neigh_dist, neigh_ind
+        return neigh_ind
 
     @classmethod
     def _get_test_params(

aeon/classification/distance_based/tests/test_time_series_neighbors.py

@@ -1,5 +1,6 @@
 """Tests for KNeighborsTimeSeriesClassifier."""
 
+import numpy as np
 import pytest
 
 from aeon.classification.distance_based import KNeighborsTimeSeriesClassifier
@@ -42,12 +43,9 @@
 @pytest.mark.parametrize("distance_key", distance_functions)
 def test_knn_on_unit_test(distance_key):
     """Test function for elastic knn, to be reinstated soon."""
-    # load arrowhead data for unit tests
     X_train, y_train = load_unit_test(split="train")
     X_test, y_test = load_unit_test(split="test")
-    knn = KNeighborsTimeSeriesClassifier(
-        distance=distance_key,
-    )
+    knn = KNeighborsTimeSeriesClassifier(distance=distance_key)
     knn.fit(X_train, y_train)
     pred = knn.predict(X_test)
     correct = 0
@@ -77,3 +75,62 @@ def test_knn_bounding_matrix(distance_key):
         if pred[j] == y_test[j]:
             correct = correct + 1
     assert correct == expected_correct_window[distance_key]
+
+
+@pytest.mark.parametrize("distance_key", distance_functions)
+def test_knn_kneighbors(distance_key):
+    """Test knn kneighbors with comprehensive validation."""
+    X_train, y_train = load_unit_test(split="train")
+    X_test, y_test = load_unit_test(split="test")
+
+    knn = KNeighborsTimeSeriesClassifier(distance=distance_key)
+    knn.fit(X_train, y_train)
+
+    # Test basic kneighbors functionality
+    dists, ind = knn.kneighbors(X_test, n_neighbors=3)
+    assert isinstance(dists, np.ndarray)
+    assert isinstance(ind, np.ndarray)
+    assert dists.shape == (X_test.shape[0], 3)
+    assert ind.shape == (X_test.shape[0], 3)
+
+    # Test that distances are non-negative
+    assert np.all(dists >= 0)
+
+    # Test that indices are within valid range
+    assert np.all(ind >= 0)
+    assert np.all(ind < len(X_train))
+
+    # Test that distances are sorted (closest first)
+    assert np.all(dists[:, 0] <= dists[:, 1])
+    assert np.all(dists[:, 1] <= dists[:, 2])
+
+    # Test using kneighbors results to make predictions manually
+    # This validates that the kneighbors method returns correct neighbor indices
+    indexes = ind[:, 0]
+    classes, y = np.unique(y_train, return_inverse=True)
+    pred = classes[y[indexes]]
+    correct = 0
+    for j in range(0, len(pred)):
+        if pred[j] == y_test[j]:
+            correct = correct + 1
+    assert correct == expected_correct_window[distance_key]
+
+    # Test kneighbors with different n_neighbors values
+    dists_2, ind_2 = knn.kneighbors(X_test, n_neighbors=2)
+    assert dists_2.shape == (X_test.shape[0], 2)
+    assert ind_2.shape == (X_test.shape[0], 2)
+
+    # Test kneighbors without returning distances
+    ind_only = knn.kneighbors(X_test, n_neighbors=3, return_distance=False)
+    assert isinstance(ind_only, np.ndarray)
+    assert ind_only.shape == (X_test.shape[0], 3)
+    # Should return same indices as when return_distance=True
+    np.testing.assert_array_equal(ind_only, ind)
+
+    # Test kneighbors on training data (should exclude self)
+    train_dists, train_ind = knn.kneighbors(n_neighbors=2)
+    assert train_dists.shape == (len(X_train), 2)
+    assert train_ind.shape == (len(X_train), 2)
+    # Each point should not be its own neighbor (diagonal should be excluded)
+    for i in range(len(X_train)):
+        assert i not in train_ind[i]