Source code for knncolle.find_knn

from functools import singledispatch
from typing import Sequence, Optional, Union
from dataclasses import dataclass
import numpy

from .classes import Index, GenericIndex
from . import lib_knncolle as lib
from ._utils import process_num_neighbors, process_subset




@dataclass
class FindKnnResults:
    """Results of :py:func:`~knncolle.find_knn.find_knn`.

    If ``num_neighbors`` is an integer, ``index`` and ``distance`` are both
    matrices. Each row corresponds to an observation in ``X`` and each column
    corresponds to one of its neighbors. ``index`` contains the indices of the
    nearest neighbors while ``distance`` contains the distance to those
    neighbors. In each row, neighbors are guaranteed to be sorted in order of
    increasing distance. Each row of ``index`` is guaranteed to not contain the
    index of the corresponding observation.

    If ``num_neighbors`` is a sequence, ``index`` and ``distance`` are lists
    instead. Each list element corresponds to an observation in ``X`` and is a
    NumPy array containing the indices (for ``index``) or distances (for
    ``distance``) to the requested number of neighbors for that observation.
    For each observation, the neighbors are guaranteed to be sorted in order of
    increasing distance. Each element of ``index`` is guaranteed to not contain
    the index of the corresponding observation.

    If ``get_index = False``, ``index`` is set to None.

    If ``get_distance = False``, ``distance`` is set to None.

    If ``subset`` is provided, the number of rows in ``index`` and ``distance``
    (if ``num_neighbors`` is an integer) or their length (otherwise) is instead
    equal to the length of the subset. Each row or list entry corresponds to
    one of the observations in the subset.
    """
    index: Optional[numpy.ndarray]
    distance: Optional[numpy.ndarray]





@singledispatch
def find_knn(
    X: Index,
    num_neighbors: Union[int, Sequence],
    num_threads: int = 1,
    subset: Optional[Sequence] = None, 
    get_index: bool = True,
    get_distance: bool = True,
    **kwargs
) -> FindKnnResults:
    """Find the k-nearest neighbors for each observation.

    Args:
        X:
            A prebuilt search index.

        num_neighbors:
            Number of nearest neighbors to identify for each observation in
            ``X``. This is automatically capped at the number of observations
            minus 1.

            Alternatively, this may be a sequence of non-negative integers of
            length equal to the number of observations in ``X``, specifying the
            number of neighbors to find for each observation.

            If ``subset`` is supplied and ``num_neighbors`` is a sequence, it
            should have length equal to ``subset`` instead, and should specify
            the number of neighbors for each observation in the subset.

        num_threads:
            Number of threads to use for the search.

        subset:
            Sequence of integers containing the indices of the observations for
            which to identify neighbors. All indices should be non-negative
            and less than the total number of observations.

        get_index:
            Whether to report the indices of each nearest neighbor.

        get_distance:
            Whether to report the distances to each nearest neighbor.

        kwargs:
            Additional arguments to pass to specific methods.

    Returns:
        Results of the nearest-neighbor search.
    """
    raise NotImplementedError("no available method for '" + str(type(X)) + "'")



@find_knn.register
def _find_knn_generic(
    X: GenericIndex,
    num_neighbors: Union[int, Sequence],
    num_threads: int = 1,
    subset: Optional[Sequence] = None,
    get_index: bool = True,
    get_distance: bool = True,
    **kwargs
) -> FindKnnResults:
    num_neighbors, force_variable = process_num_neighbors(num_neighbors)
    idx, dist = lib.generic_find_knn(
        X.ptr, 
        num_neighbors,
        force_variable,
        process_subset(subset), 
        num_threads, 
        False,
        get_index,
        get_distance
    )
    return FindKnnResults(index = idx, distance = dist)