Source code for knncolle.find_neighbors

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_threshold, process_subset




@dataclass
class FindNeighborsResults:
    """Results of :py:func:`~knncolle.find_neighbors.find_neighbors`.

    ``index`` and ``distance`` are lists where each element corresponds to an
    observation in ``X``. Each element is a NumPy array containing the indices
    of (for ``index``) or distances to (for ``distance``) the neighbors of the
    corresponding observation within the specified threshold distance. For each
    observation, 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 length of ``index`` and ``distance`` 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[list]
    distance: Optional[list]





@singledispatch
def find_neighbors(
    X: Index,
    threshold: Union[float, Sequence],
    num_threads: int = 1,
    subset: Optional[Sequence] = None, 
    get_index: bool = True,
    get_distance: bool = True,
    **kwargs
) -> FindNeighborsResults:
    """Find all neighbors within a certain distance for each observation.

    Args:
        X:
            A prebuilt search index.

        threshold:
            Distance threshold at which to identify neighbors for each
            observation in ``X``. 

            Alternatively, this may be a sequence of non-negative floats of
            length equal to the number of observations in ``X``, specifying
            the distance threshold to search for each observation.

            If ``subset`` is supplied and ``threshold`` is a sequence, it
            should have length equal to ``subset`` instead, and should specify
            the distance threshold 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 neighbor search.
    """
    raise NotImplementedError("no available method for '" + str(type(X)) + "'")



@find_neighbors.register
def _find_neighbors_generic(
    X: GenericIndex,
    threshold: Union[int, Sequence],
    num_threads: int = 1,
    subset: Optional[Sequence] = None,
    get_index: bool = True,
    get_distance: bool = True,
    **kwargs
) -> FindNeighborsResults:
    idx, dist = lib.generic_find_all(
        X.ptr, 
        process_subset(subset), 
        process_threshold(threshold),
        num_threads, 
        get_index,
        get_distance
    )
    return FindNeighborsResults(index = idx, distance = dist)