使用scipy的kdtree寻找最近邻点

最近邻

from scipy import spatial
import numpy as np
import meshio
mesh = meshio.read("D:/Dev/nonNewton/SPH/data/models/raindrop.ply")
pos = mesh.points
tree = spatial.KDTree(mesh.points)

nearest_dist = np.zeros(shape=pos.shape[0])
nearest_indx = np.zeros(shape=pos.shape[0], dtype=np.int32)
for i in range(pos.shape[0]):
    nearest = tree.query(pos[i],k=2)
    nearest_dist[i] = nearest[0][1]
    nearest_indx[i] = nearest[1][1]

avg_distance = np.mean(nearest_dist)
max_distance = np.max(nearest_dist)
min_distance = np.min(nearest_dist)
print("total points: ", pos.shape[0])
print("max_distance: ", max_distance)
print("min_distance: ", min_distance)
print("avg_distance: ", avg_distance)

关键的API就两个
一个是tree = spatial.KDTree(mesh.points)输入点云
一个是nearest = tree.query(pos[i],k=2) 寻找最近的两个点。因为第一个点必定是自己所以找俩。返回的是两个np array。第一个是所有的最小距离，第二个是所有的最近点index

固定半径内的邻居

from scipy import spatial
import numpy as np
import meshio
mesh = meshio.read("D:/Dev/nonNewton/SPH/data/models/raindrop.ply")
pos = mesh.points
tree = spatial.KDTree(mesh.points)

nearest_dist = np.zeros(shape=pos.shape[0])
nearest_indx = np.zeros(shape=pos.shape[0], dtype=np.int32)
neighbor_indices = np.ones(shape=(pos.shape[0], 50), dtype=np.int32) * (-1)
num_neighbors = np.zeros(shape=pos.shape[0], dtype=np.int32)
for i in range(pos.shape[0]):
    nearest = tree.query(pos[i],k=2)
    nearest_dist[i] = nearest[0][1]
    nearest_indx[i] = nearest[1][1]

    neighbors_  = tree.query_ball_point(pos[i], r=0.015)
    num_neighbors[i] = len(neighbors_)
    for j in range(len(neighbors_)):
        neighbor_indices[i,j] = neighbors_[j]

nearest_dist
# nearest_indx
mean = np.mean(nearest_dist)
mean
num_neighbors

改为tree.query_ball_point(pos[i], r=0.015)即可