from abc import abstractmethod from typing import Any, Literal import numpy as np from scipy.sparse import csr_matrix, spmatrix def sqeuclidean_row_norms(X: np.ndarray | csr_matrix, num_threads: int) -> np.ndarray: """Compute the squared euclidean norm of the rows of X in parallel. Parameters ---------- X : ndarray or CSR matrix of shape (n_samples, n_features) Input data. Must be c-contiguous. num_threads : int The number of OpenMP threads to use. Returns ------- sqeuclidean_row_norms : ndarray of shape (n_samples,) Arrays containing the squared euclidean norm of each row of X. """ class BaseDistancesReductionDispatcher: """Abstract base dispatcher for pairwise distance computation & reduction. Each dispatcher extending the base :class:`BaseDistancesReductionDispatcher` dispatcher must implement the :meth:`compute` classmethod. """ @classmethod def valid_metrics(cls) -> list[str]: ... @classmethod def is_usable_for(cls, X: np.ndarray | spmatrix, Y: np.ndarray | spmatrix, metric: str = "euclidean") -> bool: """Return True if the dispatcher can be used for the given parameters. Parameters ---------- X : {ndarray, sparse matrix} of shape (n_samples_X, n_features) Input data. Y : {ndarray, sparse matrix} of shape (n_samples_Y, n_features) Input data. metric : str, default='euclidean' The distance metric to use. For a list of available metrics, see the documentation of :class:`~sklearn.metrics.DistanceMetric`. Returns ------- True if the dispatcher can be used, else False. """ @classmethod @abstractmethod def compute( cls, X: np.ndarray | csr_matrix, Y: np.ndarray | csr_matrix, **kwargs, ): ... class ArgKmin(BaseDistancesReductionDispatcher): """Compute the argkmin of row vectors of X on the ones of Y. For each row vector of X, computes the indices of k first the rows vectors of Y with the smallest distances. ArgKmin is typically used to perform bruteforce k-nearest neighbors queries. This class is not meant to be instanciated, one should only use its :meth:`compute` classmethod which handles allocation and deallocation consistently. """ @classmethod def compute( cls, X: np.ndarray | csr_matrix, Y: np.ndarray | csr_matrix, k: int, metric: str = "euclidean", chunk_size: int | None = None, metric_kwargs: dict[str, Any] | None = None, strategy: Literal["auto", "parallel_on_X", "parallel_on_Y"] | None = None, return_distance: bool = False, ) -> np.ndarray | tuple[np.ndarray, np.ndarray]: """Compute the argkmin reduction. Parameters ---------- X : ndarray or CSR matrix of shape (n_samples_X, n_features) Input data. Y : ndarray or CSR matrix of shape (n_samples_Y, n_features) Input data. k : int The k for the argkmin reduction. metric : str, default='euclidean' The distance metric to use for argkmin. For a list of available metrics, see the documentation of :class:`~sklearn.metrics.DistanceMetric`. chunk_size : int, default=None, The number of vectors per chunk. If None (default) looks-up in scikit-learn configuration for `pairwise_dist_chunk_size`, and use 256 if it is not set. metric_kwargs : dict, default=None Keyword arguments to pass to specified metric function. strategy : str, {'auto', 'parallel_on_X', 'parallel_on_Y'}, default=None The chunking strategy defining which dataset parallelization are made on. For both strategies the computations happens with two nested loops, respectively on chunks of X and chunks of Y. Strategies differs on which loop (outer or inner) is made to run in parallel with the Cython `prange` construct: - 'parallel_on_X' dispatches chunks of X uniformly on threads. Each thread then iterates on all the chunks of Y. This strategy is embarrassingly parallel and comes with no datastructures synchronisation. - 'parallel_on_Y' dispatches chunks of Y uniformly on threads. Each thread processes all the chunks of X in turn. This strategy is a sequence of embarrassingly parallel subtasks (the inner loop on Y chunks) with intermediate datastructures synchronisation at each iteration of the sequential outer loop on X chunks. - 'auto' relies on a simple heuristic to choose between 'parallel_on_X' and 'parallel_on_Y': when `X.shape[0]` is large enough, 'parallel_on_X' is usually the most efficient strategy. When `X.shape[0]` is small but `Y.shape[0]` is large, 'parallel_on_Y' brings more opportunity for parallelism and is therefore more efficient - None (default) looks-up in scikit-learn configuration for `pairwise_dist_parallel_strategy`, and use 'auto' if it is not set. return_distance : boolean, default=False Return distances between each X vector and its argkmin if set to True. Returns ------- If return_distance=False: - argkmin_indices : ndarray of shape (n_samples_X, k) Indices of the argkmin for each vector in X. If return_distance=True: - argkmin_distances : ndarray of shape (n_samples_X, k) Distances to the argkmin for each vector in X. - argkmin_indices : ndarray of shape (n_samples_X, k) Indices of the argkmin for each vector in X. Notes ----- This classmethod inspects the arguments values to dispatch to the dtype-specialized implementation of :class:`ArgKmin`. This allows decoupling the API entirely from the implementation details whilst maintaining RAII: all temporarily allocated datastructures necessary for the concrete implementation are therefore freed when this classmethod returns. """ class RadiusNeighbors(BaseDistancesReductionDispatcher): """Compute radius-based neighbors for two sets of vectors. For each row-vector X[i] of the queries X, find all the indices j of row-vectors in Y such that: dist(X[i], Y[j]) <= radius The distance function `dist` depends on the values of the `metric` and `metric_kwargs` parameters. This class is not meant to be instanciated, one should only use its :meth:`compute` classmethod which handles allocation and deallocation consistently. """ @classmethod def compute( cls, X: np.ndarray | csr_matrix, Y: np.ndarray | csr_matrix, radius: float, metric: str = "euclidean", chunk_size: int | None = None, metric_kwargs: dict[str, Any] | None = None, strategy: Literal["auto", "parallel_on_X", "parallel_on_Y"] | None = None, return_distance: bool = False, sort_results: bool = False, ): """Return the results of the reduction for the given arguments. Parameters ---------- X : ndarray or CSR matrix of shape (n_samples_X, n_features) Input data. Y : ndarray or CSR matrix of shape (n_samples_Y, n_features) Input data. radius : float The radius defining the neighborhood. metric : str, default='euclidean' The distance metric to use. For a list of available metrics, see the documentation of :class:`~sklearn.metrics.DistanceMetric`. chunk_size : int, default=None, The number of vectors per chunk. If None (default) looks-up in scikit-learn configuration for `pairwise_dist_chunk_size`, and use 256 if it is not set. metric_kwargs : dict, default=None Keyword arguments to pass to specified metric function. strategy : str, {'auto', 'parallel_on_X', 'parallel_on_Y'}, default=None The chunking strategy defining which dataset parallelization are made on. For both strategies the computations happens with two nested loops, respectively on chunks of X and chunks of Y. Strategies differs on which loop (outer or inner) is made to run in parallel with the Cython `prange` construct: - 'parallel_on_X' dispatches chunks of X uniformly on threads. Each thread then iterates on all the chunks of Y. This strategy is embarrassingly parallel and comes with no datastructures synchronisation. - 'parallel_on_Y' dispatches chunks of Y uniformly on threads. Each thread processes all the chunks of X in turn. This strategy is a sequence of embarrassingly parallel subtasks (the inner loop on Y chunks) with intermediate datastructures synchronisation at each iteration of the sequential outer loop on X chunks. - 'auto' relies on a simple heuristic to choose between 'parallel_on_X' and 'parallel_on_Y': when `X.shape[0]` is large enough, 'parallel_on_X' is usually the most efficient strategy. When `X.shape[0]` is small but `Y.shape[0]` is large, 'parallel_on_Y' brings more opportunity for parallelism and is therefore more efficient despite the synchronization step at each iteration of the outer loop on chunks of `X`. - None (default) looks-up in scikit-learn configuration for `pairwise_dist_parallel_strategy`, and use 'auto' if it is not set. return_distance : boolean, default=False Return distances between each X vector and its neighbors if set to True. sort_results : boolean, default=False Sort results with respect to distances between each X vector and its neighbors if set to True. Returns ------- If return_distance=False: - neighbors_indices : ndarray of n_samples_X ndarray Indices of the neighbors for each vector in X. If return_distance=True: - neighbors_indices : ndarray of n_samples_X ndarray Indices of the neighbors for each vector in X. - neighbors_distances : ndarray of n_samples_X ndarray Distances to the neighbors for each vector in X. Notes ----- This classmethod inspects the arguments values to dispatch to the dtype-specialized implementation of :class:`RadiusNeighbors`. This allows decoupling the API entirely from the implementation details whilst maintaining RAII: all temporarily allocated datastructures necessary for the concrete implementation are therefore freed when this classmethod returns. """