python - Finding k nearest neighbors in 3d numpy array

Question

So I'm trying to find the k nearest neighbors in a pyvista numpy array from an example mesh. With the neighbors received, I want to implement some region growing in my 3d model.

But unfortunaley I receive some weird output, which you can see in the following picture. It seems like I'm missing something on the KDTree implementation. I was following the answer on a similar question: https://stackoverflow.com/a/2486341/9812286

import numpy as np 
from sklearn.neighbors import KDTree

import pyvista as pv

from pyvista import examples

# Example dataset with normals
mesh = examples.load_random_hills()

smooth = mesh

NDIM = 3
X = smooth.points
point = X[5000]

tree = KDTree(X, leaf_size=X.shape[0]+1)
# ind = tree.query_radius([point], r=10) # indices of neighbors within distance 0.3
distances, ind = tree.query([point], k=1000)

p = pv.Plotter()
p.add_mesh(smooth)

ids = np.arange(smooth.n_points)[ind[0]]
top = smooth.extract_cells(ids)
random_color = np.random.random(3)
p.add_mesh(top, color=random_color)

p.show()

question from:https://stackoverflow.com/questions/65891163/finding-k-nearest-neighbors-in-3d-numpy-array

Answer 1

You're almost there :) The problem is that you are using the points in the mesh to build the tree, but then extracting cells. Of course these are unrelated in the sense that indices for points will give you nonsense when applied as indices of cells.

Either you have to extract_points:

import numpy as np 
from sklearn.neighbors import KDTree

import pyvista as pv

from pyvista import examples

# Example dataset with normals
mesh = examples.load_random_hills()

smooth = mesh

NDIM = 3
X = smooth.points
point = X[5000]

tree = KDTree(X, leaf_size=X.shape[0]+1)
# ind = tree.query_radius([point], r=10) # indices of neighbors within distance 0.3
distances, ind = tree.query([point], k=1000)

p = pv.Plotter()
p.add_mesh(smooth)

ids = np.arange(smooth.n_points)[ind[0]]
top = smooth.extract_points(ids)  # changed here!
random_color = np.random.random(3)
p.add_mesh(top, color=random_color)

p.show()

Or you have to work with cell centers to begin with:

import numpy as np 
from sklearn.neighbors import KDTree

import pyvista as pv

from pyvista import examples

# Example dataset with normals
mesh = examples.load_random_hills()

smooth = mesh

NDIM = 3
X = smooth.cell_centers().points  # changed here!
point = X[5000]

tree = KDTree(X, leaf_size=X.shape[0]+1)
# ind = tree.query_radius([point], r=10) # indices of neighbors within distance 0.3
distances, ind = tree.query([point], k=1000)

p = pv.Plotter()
p.add_mesh(smooth)

ids = np.arange(smooth.n_points)[ind[0]]
top = smooth.extract_cells(ids)
random_color = np.random.random(3)
p.add_mesh(top, color=random_color)

p.show()

As you can see, the two results differ, since index 5000 (which we used for the reference point) means something else when indexing points or when indexing cells.