+import numpy as np
+
+from napf import _napf as core # noqa: F401
+
+np2napf_dtypes = {
+ "int32": "i",
+ "int64": "l",
+ "float32": "f",
+ "float64": "d",
+}
+
+
+
+
+
+
+
+
[docs]
+
def enforce_contiguous(array, dtype=None):
+
"""
+
If input is an instance / subclass of np.ndarray, this will check
+
if they are configuous. If so, returns same object, else turns makes it
+
contiguous and returns.
+
+
Parameters
+
----------
+
array: array-like
+
+
Returns
+
-------
+
contiguous_array: array-like
+
"""
+
if isinstance(array, np.ndarray):
+
if array.flags["C_CONTIGUOUS"] and (
+
dtype is None or dtype is array.dtype
+
):
+
return array
+
return np.ascontiguousarray(array, dtype=dtype)
+
+
return array
+
[docs]
+
def core_class_str_and_data(tree_data, metric):
+
"""
+
Returns class name of current setting.
+
Also checks if it is valid dtype.
+
+
Parameters
+
-----------
+
tree_data: (n, dim) np.ndarray
+
metric: int or str
+
+
Returns
+
--------
+
core_class_str: str
+
core_class_data: (n, dim) np.ndarray
+
"""
+
arr = enforce_contiguous(
+
tree_data, dtype=getattr(tree_data, "dtype", None)
+
)
+
dtypestr = str(arr.dtype)
+
if dtypestr not in np2napf_dtypes:
+
raise TypeError(f"Sorry, `napf` does not support ({dtypestr}) dtypes.")
+
+
if arr.ndim != 2:
+
raise ValueError(
+
"Please make sure your input array is 2D! "
+
"Given array's dim is {arr.ndim}."
+
)
+
+
# extract info
+
data_t = np2napf_dtypes[dtypestr]
+
metric = validate_metric_input(metric)
+
+
return f"KDT{data_t}L{metric}", arr
+
[docs]
+
class KDT:
+
"""
+
`napf` is implemented as template, thus, there are separate classes
+
for each {data_type, dim, metric}.
+
Currently following combinations are supported:
+
data_type: {double, int}
+
dim: {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
+
19, 20}
+
metric: {L1, L2}
+
+
Given tree_data, creates corresponding core kdt class.
+
Tree is initialized using `newtree()`.
+
+
+
Parameters
+
-----------
+
tree_data: (n, dim) np.ndarray
+
Default is None. {double, float, int, long}
+
metric: int or str
+
Default is 2 and distance will be a squared euklidian distance.
+
Valid options are {1, l1, L1, 2, l2, L2}.
+
leaf_size: int
+
nthread: int
+
Default thread count for all multi-thread-
+
+
Returns
+
--------
+
core_obj: KDT{data_t}D{dim}L{metric}
+
"""
+
+
__slots__ = (
+
"_core_tree",
+
"_nthread",
+
"_dtype",
+
)
+
+
def __init__(self, tree_data, metric=2, leaf_size=10, nthread=1):
+
"""
+
Init
+
"""
+
self.newtree(tree_data, metric, leaf_size, nthread)
+
self.nthread = nthread
+
+
@property
+
def nthread(self):
+
"""
+
Returns saved default value for nthread.
+
+
Parameters
+
-----------
+
None
+
+
Returns
+
--------
+
nthread: int
+
"""
+
return self._nthread
+
+
@nthread.setter
+
def nthread(self, nthread_):
+
"""
+
Sets default value for nthread.
+
+
Parameters
+
-----------
+
nthread_: int
+
+
Returns
+
--------
+
None
+
"""
+
self._nthread = nthread_
+
+
@property
+
def core_tree(self):
+
"""
+
Returns initialized core tree, if there's any.
+
+
Parameters
+
-----------
+
None
+
+
Returns
+
--------
+
core_tree: napf.core.KDT*
+
"""
+
if hasattr(self, "_core_tree"):
+
return self._core_tree
+
+
else:
+
return None
+
+
@property
+
def tree_data(self):
+
"""
+
Returns data used to initialize core tree. Read only
+
+
Parameters
+
-----------
+
None
+
+
Returns
+
--------
+
tree_data: (n, d) np.ndarray
+
{double, float, int, long}
+
"""
+
coretree = self.core_tree
+
if coretree is not None:
+
return coretree.tree_data
+
+
else:
+
return None
+
+
@property
+
def dtype(self):
+
"""
+
Returns dtype of current tree
+
"""
+
return self._dtype
+
+
+
[docs]
+
def newtree(self, tree_data, metric=2, leaf_size=10, nthread=1):
+
"""
+
Given 2D array-like tree_data, it:
+
1. makes sure data is a contiguous array
+
2. builds a corresponding kdt.
+
+
Note that core kdt objects also have `newtree()` function with the same
+
parameters, and build a new kdt each time it is called.
+
+
Parameters
+
-----------
+
tree_data: (n, d) np.ndarray
+
{double, float, int, long}
+
+
"""
+
core_cls, tdata = core_class_str_and_data(
+
np.ascontiguousarray(tree_data), metric
+
) # checks and raises error
+
# we can call newtree() function of the core class,
+
# if _core_tree already exists.
+
# However, creating a new kdt should not add significant overhead.
+
self._core_tree = eval(f"core.{core_cls}(tdata, leaf_size, nthread)")
+
self._dtype = tdata.dtype
+
+
+
+
[docs]
+
def knn_search(self, queries, kneighbors, nthread=None):
+
"""
+
k-nearest-neighbor search.
+
+
Parameters
+
-----------
+
queires: (m, d) np.ndarray
+
Data type will be casted to the same type as `tree_data`.
+
kneighbors: int
+
nthread: int
+
Default is None and will use self.nthread.
+
+
Returns
+
--------
+
ids_and_distances: tuple
+
((m, kneighbors) np.ndarray - double dists,)
+
(m, kneighbors) np.ndarray - uint ids)
+
"""
+
if nthread is None:
+
nthread = self.nthread
+
+
return self.core_tree.knn_search(
+
enforce_contiguous(queries, self.dtype), kneighbors, nthread
+
)
+
+
+
+
[docs]
+
def query(self, queries, nthread=None):
+
"""
+
scipy-like KDTree query call.
+
Same as `knn_search(queries, 1, nthreads)`
+
+
Parameters
+
-----------
+
queries: (m, d) np.ndarray
+
Data type will be casted to the same type as `tree_data`.
+
nthread: int
+
Default is None and will use self.nthread.
+
+
Returns
+
--------
+
ids_and_distances: tuple
+
((m, 1) np.ndarray - double dists,)
+
(m, 1) np.ndarray - uint ids)
+
"""
+
if nthread is None:
+
nthread = self.nthread
+
+
return self.core_tree.query(
+
enforce_contiguous(queries, self.dtype), nthread
+
)
+
+
+
+
[docs]
+
def radius_search(self, queries, radius, return_sorted, nthread=None):
+
"""
+
Searches for neighbors in given radius.
+
+
Parameters
+
-----------
+
queries: (m, d) np.ndarray
+
Data type will be casted to the same type as `tree_data`.
+
radius: float
+
return_sorted: bool
+
nthread: int
+
Default is None and will use self.nthread
+
+
Returns
+
--------
+
ids_and_distances: tuple
+
((m, 1) np.ndarray - uint ids,
+
(m, 1) np.ndarray - double dists)
+
"""
+
if nthread is None:
+
nthread = self.nthread
+
+
return self.core_tree.radius_search(
+
enforce_contiguous(queries, self.dtype),
+
radius,
+
return_sorted,
+
nthread,
+
)
+
+
+
+
[docs]
+
def query_ball_point(self, queries, radius, return_sorted, nthread=None):
+
"""
+
scipy-like KDTree query_ball_point call.
+
+
Parameters
+
----------
+
queries: (m, d) np.ndarray
+
Data type will be casted to the same type as `tree_data`.
+
radius: float
+
return_sorted: bool
+
nthread: int
+
Default is None and will use self.nthread
+
+
Returns
+
-------
+
ids: list
+
list of np.array
+
"""
+
if nthread is None:
+
nthread = self.nthread
+
+
return self.core_tree.query_ball_point(
+
enforce_contiguous(queries, self.dtype),
+
radius,
+
return_sorted,
+
nthread,
+
)
+
+
+
+
[docs]
+
def radii_search(self, queries, radii, return_sorted, nthread=None):
+
"""
+
Similar to `radius_search`, but you can specify radius for each query.
+
+
Parameters
+
-----------
+
queries: (m, d) np.ndarray
+
Data type will be casted to the same type as `tree_data`.
+
radii: (m,) np.ndarray
+
return_sorted: bool
+
nthread: int
+
Default is None and will use self.nthread
+
+
Returns
+
--------
+
ids_and_distances: tuple
+
((m, 1) np.ndarray - uint ids,
+
(m, 1) np.ndarray - double dists)
+
"""
+
# input size check
+
if len(queries) != len(radii):
+
raise ValueError(
+
f"Input size mismatch between queires ({len(queries)}) "
+
f" and radii ({len(radii)})."
+
"They should be the same."
+
)
+
+
if nthread is None:
+
nthread = self.nthread
+
+
return self.core_tree.radii_search(
+
enforce_contiguous(queries, self.dtype),
+
enforce_contiguous(radii, self.dtype),
+
return_sorted,
+
nthread,
+
)
+
+
+
+
[docs]
+
def unique_data_and_inverse(
+
self,
+
radius,
+
return_unique=True,
+
return_intersection=True,
+
nthread=None,
+
):
+
"""
+
Finds unique tree data with in given radius tolerance.
+
+
+
Parameters
+
----------
+
radius: float
+
return_unique: bool
+
Default is True. Otherwise, will be an empty array return.
+
return_intersection: bool
+
Default is True, Otherwise, will be an empty UIntVectorVector return.
+
nthread:int
+
+
Returns
+
-------
+
unique_data: np.ndarray
+
Empty if return_unique is False.
+
Same as kdt.tree_data[unique_ids].
+
unique_ids: np.ndarray
+
Indices of unique entries from tree data.
+
First occurance is considered unique.
+
inverse_ids: np.ndarray
+
Indices to reconstruct original tree_data with
+
unique_data. kdt.tree_data == unique_data[inverse_ids]
+
intersection: UIntVectorVector
+
Empty if return_intersection is False.
+
Intersection of each data with respect to all the others.
+
"""
+
if nthread is None:
+
nthread = self.nthread
+
+
(
+
original_inverse,
+
intersection,
+
) = self.core_tree.tree_data_unique_inverse(
+
radius, return_intersection, nthread
+
)
+
+
unique_ids, inverse_ids = np.unique(
+
original_inverse, return_inverse=True
+
)
+
+
if return_unique:
+
return (
+
self.core_tree.tree_data[unique_ids],
+
unique_ids,
+
inverse_ids,
+
intersection,
+
)
+
else:
+
return np.array(), unique_ids, inverse_ids, intersection
+