Source code for knncolle.find_distance

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

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




@singledispatch
def find_distance(
    X: Index,
    num_neighbors: Union[int, Sequence],
    num_threads: int = 1,
    subset: Optional[Sequence] = None, 
    **kwargs
) -> numpy.ndarray:
    """Find the distance to the k-th closest point for each observation.

    Args:
        X:
            A prebuilt search index.

        num_neighbors:
            Number of nearest neighbors at which to compute the distance from
            each observation in ``X``, i.e., k. 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
            neighbor at which to compute the distance 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 neighbor 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 compute the distances. All indices should be non-negative
            and less than the total number of observations.

        kwargs:
            Additional arguments to pass to specific methods.

    Returns:
        A NumPy array of length equal to the number of observations in
        ``X`` (or ``subset``, if provided) containing the distance to the
        ``num_neighbors``-th point for each observation.
    """
    raise NotImplementedError("no available method for '" + str(type(X)) + "'")



@find_distance.register
def _find_distance_generic(
    X: GenericIndex,
    num_neighbors: Union[int, Sequence],
    num_threads: int = 1,
    subset: Optional[Sequence] = None,
    **kwargs
) -> numpy.ndarray:
    num_neighbors, force_variable = process_num_neighbors(num_neighbors)
    return lib.generic_find_knn(
        X.ptr, 
        num_neighbors,
        force_variable,
        process_subset(subset), 
        num_threads, 
        True,
        False,
        False
    )