View Source Scholar.Neighbors.KDTree (Scholar v0.3.0)
Implements a k-d tree, a space-partitioning data structure for organizing points in a k-dimensional space.
It can be used to predict the K-Nearest Neighbors of a given input.
This is implemented as one-dimensional tensor with indices pointed to highest
dimension of the given tensor. Traversal starts by calling root/0 and then
accessing the left_child/1 and right_child/1. The tree is left-balanced.
Each level traverses over the last axis of tensor, the index for a level can be
computed as: rem(level, Nx.axis_size(tensor, -1)).
References
Summary
Functions
Builds a KDTree.
Returns the index of the left child of i.
Returns the number of resulting levels in a KDTree for tensor.
Returns the parent of child i.
Predict the K nearest neighbors of x_predict in KDTree.
Returns the index of the right child of i.
Returns the root index.
Functions
Builds a KDTree.
Examples
iex> tree = Scholar.Neighbors.KDTree.fit(Nx.iota({5, 2}))
iex> tree.data
Nx.tensor(
[
[0, 1],
[2, 3],
[4, 5],
[6, 7],
[8, 9]
]
)
iex> tree.levels
3
iex> tree.indices
Nx.u32([3, 1, 4, 0, 2])Returns the index of the left child of i.
It is your responsibility to guarantee the result is not greater than the leading axis of the tensor.
Examples
iex> Scholar.Neighbors.KDTree.left_child(0)
1
iex> Scholar.Neighbors.KDTree.left_child(1)
3
iex> Scholar.Neighbors.KDTree.left_child(Nx.u32(3))
#Nx.Tensor<
u32
7
>Returns the number of resulting levels in a KDTree for tensor.
Examples
iex> Scholar.Neighbors.KDTree.levels(Nx.iota({10, 3}))
4Returns the parent of child i.
It is your responsibility to guarantee the result is positive.
Examples
iex> Scholar.Neighbors.KDTree.parent(1)
0
iex> Scholar.Neighbors.KDTree.parent(2)
0
iex> Scholar.Neighbors.KDTree.parent(Nx.u32(3))
#Nx.Tensor<
u32
1
>Predict the K nearest neighbors of x_predict in KDTree.
Examples
iex> x = Nx.iota({10, 2})
iex> x_predict = Nx.tensor([[2, 5], [1, 9], [6, 4]])
iex> kdtree = Scholar.Neighbors.KDTree.fit(x, num_neighbors: 3)
iex> {indices, distances} = Scholar.Neighbors.KDTree.predict(kdtree, x_predict)
iex> indices
#Nx.Tensor<
s64[3][3]
[
[2, 1, 0],
[2, 3, 1],
[2, 3, 1]
]
>
iex> distances
#Nx.Tensor<
f32[3][3]
[
[2.0, 2.0, 4.4721360206604],
[5.0, 5.385164737701416, 6.082762718200684],
[2.2360680103302, 3.0, 4.123105525970459]
]
>
iex> x = Nx.iota({10, 2})
iex> x_predict = Nx.tensor([[2, 5], [1, 9], [6, 4]])
iex> kdtree = Scholar.Neighbors.KDTree.fit(x, num_neighbors: 3, metric: {:minkowski, 1})
iex> {indices, distances} = Scholar.Neighbors.KDTree.predict(kdtree, x_predict)
iex> indices
#Nx.Tensor<
s64[3][3]
[
[2, 1, 0],
[4, 3, 1],
[2, 3, 1]
]
>
iex> distances
#Nx.Tensor<
f32[3][3]
[
[2.0, 2.0, 6.0],
[7.0, 7.0, 7.0],
[3.0, 3.0, 5.0]
]
>Returns the index of the right child of i.
It is your responsibility to guarantee the result is not greater than the leading axis of the tensor.
Examples
iex> Scholar.Neighbors.KDTree.right_child(0)
2
iex> Scholar.Neighbors.KDTree.right_child(1)
4
iex> Scholar.Neighbors.KDTree.right_child(Nx.u32(3))
#Nx.Tensor<
u32
8
>Returns the root index.
Examples
iex> Scholar.Neighbors.KDTree.root()
0