How to use Numba for a spatial.cKDTree script ?

[HerveGlz]

Hi all,

I am searching for a way to use Numba so that I can improve a KDTree processing script with scipy.
But it seems that it doesn't work when I apply the procedure from numba-scipy. I wonder if the spatial.cKDTree / spatial.KDTree is implemented within numba-scipy... Is it ?

Is there someone who could give me a hand on it ?
Or have I to compile my script with Cython ?

Thanks in advance,
Warm regards,
Hervé

[stuartarchibald]

Thanks for the request. numba-scipy only supports a few things at present, scipy.spatial isn't one of those, if you do write support please consider contributing it to this project. Thanks!

[jackd]

@HerveGlz this isn't really related to numba-scipy, but a little while ago I ported part of scikit-learn's KDTree implementation - you can find it here. It's a bit messy, but depending on what usage you want and how keen you are to look through the examples/benchmark scripts it may be enough to cover what you want.

[HerveGlz]

@jackd Thanks for your message. If I may, your link seems to be dead... Will it be possible to update it ?

[jackd]

@HerveGlz my bad, had it private. Let me know if you still can't see it.

[HerveGlz]

@jackd Thanks a lot Jack ! The link works perfectly. I will try to implement your solution in my script and will get back to you.

[sgbaird]

Another resource to look at: https://github.com/lmcinnes/pynndescent
UMAP uses this under the hood if I remember correctly, and many metrics (including custom) are supported and fast.

[promitmoitra]

https://github.com/mortacious/numba-kdtree
The generation of the tree has to be done outside a numba compiled function, but the generated tree can be passed to, and parsed within, a numba compiled function.