Verwenden von Bäumen zum Beschleunigen der nächsten Nachbarsuche auf 3D-Array mit periodischen Randbedingungen

Question

Mit einem Code aus der Antwort auf this Frage kann ich eine Brute-Force-NN-Suche auf einem 3D-Array unter Berücksichtigung der periodischen Randbedingungen durchführen. Der Code gibt dann den Index des nächsten Nachbarn zurück und dies für alle Nachbarn.

import numpy as np 
from multiprocessing.dummy import Pool as ThreadPool 

N = 10000 # Num of objects 
# create random point positions 
coords = np.random.random((3, N)).transpose() 

def NN(point): 
    dist = np.abs(np.subtract(coords, point)) # Distance between point and all neighbors xyz values 
    dist = np.where(dist > 0.5, dist - 1, dist) # checking if distance is closer if it wraps around 
    return (np.square(dist)).sum(axis=-1).argsort()[1] # Calc distance and find index of nearest neighbour 

# multi threading for speed increase 
pool = ThreadPool(12) 
match = pool.map(NN, coords) 
pool.close() 
pool.join() 

Für N ~ 50000 wird es sehr langsam wie erwartet.

Ich mag würde wissen, wie ich dies mit Bäumen wie sklearn.BallTree implementieren würde oder scipy.spacial.cKDTree und möchte, dies zu tun, ohne den Raum 8 weitere Male dupliziert als here vorgeschlagen.

Answer 1

Weder sklearn.BallTree noch scipy.spatial.cKDTree können leicht modifiziert werden, um periodische Randbedingungen zu behandeln. Periodische Grenzen erfordern grundsätzlich andere Annahmen als die in diesen Implementierungen verwendeten.

Sie sollten eine alternative Implementierung in Betracht ziehen, z. B. periodic_kdtree, beachten Sie jedoch, dass dies eine (etwas veraltete) Python-Implementierung ist und nicht annähernd so schnell wie die von Ihnen erwähnte Cython/C++ - Implementierung sein wird.