fastmatch.knn_scikit

 1import numpy as np
 2from sklearn.neighbors import NearestNeighbors
 3
 4
 5class ScikitNearestNeighbors:
 6    """
 7    A wrapper for scikit-learn's NearestNeighbors to conform to the API
 8    used by the other backends in this package.
 9    """
10
11    def __init__(self, **kwargs):
12        """
13        Initializes the scikit-learn NearestNeighbors model.
14        Any keyword arguments are passed directly to the scikit-learn constructor.
15        """
16        self.model = NearestNeighbors(**kwargs)
17
18    def fit(self, X: np.ndarray):
19        """
20        Fits the NearestNeighbors model to the provided data.
21
22        Args:
23            X (np.ndarray): Array of shape (N, M) to fit the model on.
24
25        Returns:
26            self: The fitted object.
27        """
28        self.model.fit(X)
29        return self
30
31    def kneighbors(self, X: np.ndarray, n_neighbors: int = 1):
32        """
33        Finds the k-nearest neighbors for the samples in X.
34
35        Args:
36            X (np.ndarray): Array of shape (N, M) to find neighbors for.
37            n_neighbors (int): The number of neighbors to find.
38
39        Returns:
40            (distances, indices): Tuple of arrays containing the distances
41                                  and indices of the nearest neighbors.
42        """
43        # n_neighbors in the method call overrides the one in the constructor
44        distances, indices = self.model.kneighbors(X, n_neighbors=n_neighbors)
45        return distances, indices