Skip to content

jaxknn.tree¤

The raw kd-tree. Distances are squared, and there is no input validation, scaling or size bucketing; use jaxknn.knn unless you need a custom query.

jaxknn.tree.build_tree(points: Float[Array, 'n d'], /, *, leaf_size: int = 16) -> Tree ¤

Build a tree over points (n >= 0 rows, any dimension d >= 1).

jaxknn.tree.Tree ¤

A built tree. Arrays are in tree (permuted) order unless noted.

__abstractclassvars__ = frozenset() class-attribute ¤

Build an immutable unordered collection of unique elements.

__abstractmethods__ = frozenset() class-attribute ¤

Build an immutable unordered collection of unique elements.

__abstractvars__ = frozenset() class-attribute ¤

Build an immutable unordered collection of unique elements.

__annotations_cache__ = {'points': <class 'jaxtyping.Float[Array, 'n_pad d']'>, 'perm': <class 'jaxtyping.Int[Array, 'n_pad']'>, 'valid': <class 'jaxtyping.Bool[Array, 'n_pad']'>, 'cell_lo': tuple[jaxtyping.Float[Array, '_nodes d'], ...], 'cell_hi': tuple[jaxtyping.Float[Array, '_nodes d'], ...], 'split_dim': tuple[jaxtyping.Int[Array, '_nodes'], ...], 'split_val': tuple[jaxtyping.Float[Array, '_nodes'], ...], 'n': <class 'int'>, 'depth': <class 'int'>, 'leaf_size': <class 'int'>, 'n_pad': <class 'int'>} class-attribute ¤

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_fields__ = {'points': Field(name='points',type=<class 'jaxtyping.Float[Array, 'n_pad d']'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'perm': Field(name='perm',type=<class 'jaxtyping.Int[Array, 'n_pad']'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'valid': Field(name='valid',type=<class 'jaxtyping.Bool[Array, 'n_pad']'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'cell_lo': Field(name='cell_lo',type=tuple[jaxtyping.Float[Array, '_nodes d'], ...],default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'cell_hi': Field(name='cell_hi',type=tuple[jaxtyping.Float[Array, '_nodes d'], ...],default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'split_dim': Field(name='split_dim',type=tuple[jaxtyping.Int[Array, '_nodes'], ...],default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'split_val': Field(name='split_val',type=tuple[jaxtyping.Float[Array, '_nodes'], ...],default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=False,doc=None,_field_type=_FIELD), 'n': Field(name='n',type=<class 'int'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'static': True}),kw_only=False,doc=None,_field_type=_FIELD), 'depth': Field(name='depth',type=<class 'int'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'static': True}),kw_only=False,doc=None,_field_type=_FIELD), 'leaf_size': Field(name='leaf_size',type=<class 'int'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'static': True}),kw_only=False,doc=None,_field_type=_FIELD), 'n_pad': Field(name='n_pad',type=<class 'int'>,default=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,default_factory=<dataclasses._MISSING_TYPE object at 0x7f5518f0d7f0>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({'static': True}),kw_only=False,doc=None,_field_type=_FIELD)} class-attribute ¤

dict() -> new empty dictionary dict(mapping) -> new dictionary initialized from a mapping object's (key, value) pairs dict(iterable) -> new dictionary initialized as if via: d = {} for k, v in iterable: d[k] = v dict(**kwargs) -> new dictionary initialized with the name=value pairs in the keyword argument list. For example: dict(one=1, two=2)

__dataclass_params__ = _DataclassParams(init=True,repr=False,eq=False,order=False,unsafe_hash=False,frozen=False,match_args=True,kw_only=False,slots=False,weakref_slot=False) class-attribute ¤
__doc__ = 'A built tree. Arrays are in tree (permuted) order unless noted.' class-attribute ¤

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__firstlineno__ = 25 class-attribute ¤

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.int(). For floating-point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by '+' or '-' and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer iteral.

int('0b100', base=0) 4

__match_args__ = ('points', 'perm', 'valid', 'cell_lo', 'cell_hi', 'split_dim', 'split_val', 'n', 'depth', 'leaf_size', 'n_pad') class-attribute ¤

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

__module__ = 'jaxknn.tree._build' class-attribute ¤

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to 'utf-8'. errors defaults to 'strict'.

__static_attributes__ = () class-attribute ¤

Built-in immutable sequence.

If no argument is given, the constructor returns an empty tuple. If iterable is specified the tuple is initialized from iterable's items.

If the argument is a tuple, the return value is the same object.

n_leaves property ¤

Number of leaves, 2**depth.

__init__(points: Float[Array, 'n_pad d'], perm: Int[Array, n_pad], valid: Bool[Array, n_pad], cell_lo: tuple[jaxtyping.Float[Array, '_nodes d'], ...], cell_hi: tuple[jaxtyping.Float[Array, '_nodes d'], ...], split_dim: tuple[jaxtyping.Int[Array, _nodes], ...], split_val: tuple[jaxtyping.Float[Array, _nodes], ...], n: int, depth: int, leaf_size: int, n_pad: int) -> None ¤

Initialize self. See help(type(self)) for accurate signature.

jaxknn.tree.knn(tree: Tree, queries: Float[Array, 'm d'], /, k: int, *, frontier: int = 16, exclude: tuple[jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex, jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex] | None = None) -> tuple[jaxtyping.Int[Array, 'm k'], jaxtyping.Float[Array, 'm k']] ¤

Exact k nearest tree points to each query (no self-exclusion).

exclude=(point_labels, query_labels) (integers in [0, n)) skips every point whose label equals the query's, exactly: subtrees holding only that label are pruned whole. Returns (indices into the tree's original points, sq_dist), rows sorted by distance; missing neighbours are index tree.n with distance inf.

jaxknn.tree.all_knn(points: Float[Array, 'n d'], /, k: int, *, leaf_size: int = 16, frontier: int = 16) -> tuple[jaxtyping.Int[Array, 'n k'], jaxtyping.Float[Array, 'n k']] ¤

Exact k nearest neighbours of every point among the others.

Self is excluded by index. Returns (indices, sq_dist) in input order, each row sorted by distance; missing neighbours (fewer than k other points) are index n with distance inf.

jaxknn.tree.locate_leaves(tree: Tree, queries: Float[Array, 'm d']) -> Int[Array, m] ¤

Return the leaf each query falls in, by descending the split planes.

jaxknn.tree.node_labels(tree: Tree, labels: Int[Array, n]) -> tuple[jax.Array, tuple[jax.Array, ...]] ¤

Per-row labels and every node's homogeneous label, for exclude.

Rows: labels in tree order, tree.n + 1 (EMPTY) on padding. Nodes (per level 0..depth): the label all their points share, tree.n (MIXED) if they differ, EMPTY if they hold no points. Labels must be in [0, n).

jaxknn.tree.brute_knn(points: Float[Array, 'n d'], /, k: int, *, queries: jaxtyping.Float[Array, 'm d'] | None = None, chunk: int = 1024) -> tuple[jaxtyping.Int[Array, 'm k'], jaxtyping.Float[Array, 'm k']] ¤

Exact kNN comparing every query with every point, chunk queries at a time.

Same contract as all_knn (queries=None: self excluded by index) and knn (queries given: no exclusion): (indices, sq_dist), rows sorted, missing neighbours index n with distance inf.