5

我有一个带坐标的 3d 网格

x = linspace(0, Lx, Nx)
y = linspace(0, Ly, Ny)
z = linspace(0, Lz, Nz)

我需要在某个位置（x0，y0，z0）的某个半径R内索引点（即x [i]，y [j]，z [k]）。N_i 可以很大。我可以做一个简单的循环来找到我需要的东西

points=[]
i0,j0,k0 = floor( (x0,y0,z0)/grid_spacing )
Nr = (i0,j0,k0)/grid_spacing + 2
for i in range(i0-Nr, i0+Nr):
    for j in range(j0-Nr, j0+Nr):
        for k in range(k0-Nr, k0+Nr):
            if norm(array([i,j,k])*grid_spacing - (x0,y0,k0)) < cutoff:
                points.append((i,j,k))

但这很慢。有没有更自然/更快的方法来使用 numpy 进行这种类型的操作？

4

1 回答 1

3

这个怎么样：

import scipy.spatial as sp
x = np.linspace(0, Lx, Nx)
y = np.linspace(0, Ly, Ny)
z = np.linspace(0, Lz, Nz)

#Manipulate x,y,z here to obtain the dimensions you are looking for

center=np.array([x0,y0,z0])

#First mask the obvious points- may actually slow down your calculation depending.
x=x[abs(x-x0)<cutoff]
y=y[abs(y-y0)<cutoff]
z=z[abs(z-z0)<cutoff]


#Generate grid of points
X,Y,Z=np.meshgrid(x,y,z)
data=np.vstack((X.ravel(),Y.ravel(),Z.ravel())).T

distance=sp.distance.cdist(data,center.reshape(1,-1)).ravel()
points_in_sphere=data[distance<cutoff]

您应该能够执行以下操作，而不是最后两行：

tree=sp.cKDTree(data)
mask=tree.query_ball_point(center,cutoff)
points_in_sphere=data[mask]

如果您不想调用空间：

distance=np.power(np.sum(np.power(data-center,2),axis=1),.5)
points_in_sphere=data[distance<cutoff]
于 2013-07-15T20:44:28.750 回答