Finding nearest neighbors within a distance and taking the average of those neighbors using cKDTree

Question

I'm using python scripting to read in two large (millions of points) point clouds as arrays ("A" and "B").

I need to find the nearest "B" neighbors of the points in "A", but within 5 cm of each point in "A". I also want to average the neighbors within the 5 cm radius of the points in "A."

Is there a way to do this using cKTree all at once, including the averaging?

Answer 1

I'm not sure about what do you want to do, but If I understand you correctly you can follow this steps:

# this are just random arrays for test
A = 20 * np.random.rand(1000, 3)
B = 20 * np.random.rand(1000, 3)

Compute a cKDTree for each point cloud

tree_A = cKDTree(A)
tree_B = cKDTree(B)

Find the points in A that are at most at 5 units from each point in B:

# faster than loop + query_ball_point
neighbourhood = tree_B.query_ball_tree(tree_A, 5)

Compute the mean over all of those groups of points:

means = np.zeros_like(A)
for i in range(len(neighbourhood)):
    means[i] = A[neighbourhood[i]].mean(0)