Skip to main content

A kdtree implementation for numba.

Project description

Numba-Kdtree - a fast KDTree implementation for numba

numba-kdtree provides a fast KDTree implementation for numba CPU functions based on scipys ckdtree.

The provided KDTree class is usable from both python and numba nopython functions calling directly into C code for performance critical sections utilizing a ctypes-like interface.

Once the KDTree class is compiled by numba, it is just as fast as the original scipy version.

Note: Currently only a basic subset of the original ckdtree interface is implemented.

Installation

Using pip

pip install numba-kdtree

From source

git clone https://github.com/mortacious/numba-kdtree.git
cd numba-kdtree
python -m pip install

Usage

import numpy as np
from numba_kdtree import KDTree
data = np.random.random(3_000_000).reshape(-1, 3)
kdtree = KDTree(data, leafsize=10)

# compute the nearest neighbors of the first 100 points
distances, indices = kdtree.query(data[:100], k=30)

# compute all points in a 0.5 radius around the first 100 points and return it's sorted indices
indices = kdtree.query_radius(data[:100], r=0.5, return_sorted=True)

The KDTree can be freely passed into a numba nopython function or constructed directly within it:

import numpy as np
from numba import njit
from numba_kdtree import KDTree

def numba_function_with_kdtree(kdtree, data):
    for i in range(data.shape[0]):
        distances, indices = kdtree.query(data[i], k=30)
        #<Use the computed neighbors
        
data = np.random.random(3_000_000).reshape(-1, 3)
kdtree = KDTree(data, leafsize=10)

numba_function_with_kdtree(kdtree)
import numpy as np
from numba import njit
from numba_kdtree import KDTree

def numba_function_constructing_kdtree(kdtree, data):
    kdtree = KDTree(data, leafsize=10)
    for i in range(data.shape[0]):
        distances, indices = kdtree.query(data[i], k=30)
        #<Use the computed neighbors
        
data = np.random.random(10_000_000).reshape(-1, 10)

numba_function_constructing_kdtree(data)

Additionally, the KDTree object can also be serialized via pickle:

import pickle

data = np.random.random(3_000_000).reshape(-1, 3)
kdtree = KDTree(data, leafsize=10) 

# pass the tree through pickle
# Note: This also copies the data array preserving it's integrity
serialized_tree = pickle.dumps(kdtree)

# The copied data array is now owned by the restored tree
restored_kdtree = pickle.loads(serialized_tree)

k = 10
# query the old tree
dd, ii, nn = kdtree.query_parallel(data[:100], k=k)

# query the new tree
dd_r, ii_r, nn_r = restored_kdtree.query_parallel(data[:100], k=k)

TODOs

  • Implement the full scipy ckdtree interface

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

numba_kdtree-0.6.0.tar.gz (31.7 kB view details)

Uploaded Source

Built Distributions

numba_kdtree-0.6.0-cp313-cp313-win_amd64.whl (46.2 kB view details)

Uploaded CPython 3.13 Windows x86-64

numba_kdtree-0.6.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (481.3 kB view details)

Uploaded CPython 3.13 manylinux: glibc 2.17+ x86-64

numba_kdtree-0.6.0-cp312-cp312-win_amd64.whl (46.2 kB view details)

Uploaded CPython 3.12 Windows x86-64

numba_kdtree-0.6.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (481.3 kB view details)

Uploaded CPython 3.12 manylinux: glibc 2.17+ x86-64

numba_kdtree-0.6.0-cp311-cp311-win_amd64.whl (46.2 kB view details)

Uploaded CPython 3.11 Windows x86-64

numba_kdtree-0.6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (481.1 kB view details)

Uploaded CPython 3.11 manylinux: glibc 2.17+ x86-64

numba_kdtree-0.6.0-cp310-cp310-win_amd64.whl (46.2 kB view details)

Uploaded CPython 3.10 Windows x86-64

numba_kdtree-0.6.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (481.1 kB view details)

Uploaded CPython 3.10 manylinux: glibc 2.17+ x86-64

numba_kdtree-0.6.0-cp39-cp39-win_amd64.whl (46.2 kB view details)

Uploaded CPython 3.9 Windows x86-64

numba_kdtree-0.6.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (481.1 kB view details)

Uploaded CPython 3.9 manylinux: glibc 2.17+ x86-64

File details

Details for the file numba_kdtree-0.6.0.tar.gz.

File metadata

  • Download URL: numba_kdtree-0.6.0.tar.gz
  • Upload date:
  • Size: 31.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.9.21

File hashes

Hashes for numba_kdtree-0.6.0.tar.gz
Algorithm Hash digest
SHA256 53ba8bb82cc820b06b771018698468cad5a9356bc9d8bfd7df7332c1887bd251
MD5 f5bcec192a9bcae69c803ee955c2e091
BLAKE2b-256 4862d99cb142b118144635fae85368e2e586c52e4ec9d49d38aaf2176cd0ddcf

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp313-cp313-win_amd64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 ede5d08abaa069a4fe24059ef278118ae2afecaf23ab57fc6858cc334e81cf6c
MD5 5c2b79d495724ba423d8b2b066a67c62
BLAKE2b-256 0248559aa77dbfbc738f3f2997f6452303a274294031f17577d19f7d1b134005

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 a2b31ffd14cfa5a50f3107ca6ecdac3bc1892cc3cc8208e053ff41aaf8687476
MD5 9f8da4765b55a3b8b63d032769a60895
BLAKE2b-256 3f052a24018fe50f4e593b6795eb8040728f013af37606326e1a7a42861dd823

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp312-cp312-win_amd64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 c3537d4bd35934ff9716e4161d715e2036ddaa03080a83ea22b8365a89860b3c
MD5 ee3d5cd69e9f3f18497deb68714a0984
BLAKE2b-256 a09c2f93702257f79fdb7dd3f3a55ac87bdd5a69d651df3829a835582b3ec4e6

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 93f8bb9cf5965b5b6a24e0b9f665df4837cfecf025d9e87d25b3d0b59e96a9fe
MD5 6510dc83dbf972a9898f41087683e554
BLAKE2b-256 3af357c096b81ba13cc6bdd4f2d2e827da7f1c41b25c3c53e76e97e402b611d2

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp311-cp311-win_amd64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 029f1cfa815fa5e58a8ee47f02c0926f3d54ef44ab50cb68884e5928b22a1e63
MD5 17b9cfc556c84090fb5c8934ba87f8a3
BLAKE2b-256 da25902e307bff7faa716a59eda062a861173428df02989b640adbea6a838c25

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 4fe41e287392a901b90c8ab341cd21ff6285571f1175a368860a966a20b425f3
MD5 ddd141dc208ad0edd6e2956c83049318
BLAKE2b-256 426dc2d892b8fc5f74603bb467330652983af5269c0d4d8f0b0ec57a1710a8d5

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp310-cp310-win_amd64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 49b7d13f92d859e022c06c91ab38c04c6bba1ad891aaaa1646199600ccb57913
MD5 21437f6e9ed6f68aaeb2b4e9c8b9eae0
BLAKE2b-256 a20f92deda143a75f9ace20ab435dfdc22360cdc4aeb8fa892351f6370515336

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 b2b62bb23c0db17c908e5bf1968b71340411ef062d9e961bdc58f20fe57572c9
MD5 8c04827c4e24cbc983a8c2dca37a7eba
BLAKE2b-256 7ce053604ec230c656e50971260c5ee7836082219b8abdfd99974c1d60dff9a3

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp39-cp39-win_amd64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp39-cp39-win_amd64.whl
Algorithm Hash digest
SHA256 1f7c221f502fb06016e3a76c49a4243f6426e93f289eaba805f62e9151b331c0
MD5 92e4eda88bfff0fc9fd12d6dfc065e33
BLAKE2b-256 f3c8d2411728e4cc7373ff2e8fe9fd10009901f9b499b9181a4b17ccf1bffc91

See more details on using hashes here.

File details

Details for the file numba_kdtree-0.6.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.

File metadata

File hashes

Hashes for numba_kdtree-0.6.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
Algorithm Hash digest
SHA256 a668a123a552bd8cbd4d31c57b149c5ba9721247736e65dde02c731164e9a455
MD5 eba8e9c7df53164cc53d3121daffdf9c
BLAKE2b-256 01849aeff3848f2ab9e65c07701ae1d48d2d377dd7f31508d770b714276c9284

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page