Indexing a spherical subset of 3d mesh data in numpy

Question

I have a 3d grid with coordinates

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

and I need to index the points (ie x [i], y [j], z [k]) within a certain radius R of the position (x0, y0, z0). N_i can be quite large. I can do a simple loop to find what I need

 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)) 

but it is rather slow. Is there a more natural / faster way to do this type of operation with numpy?

Answer 1

How about this:

 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] 

Instead of the last two lines you can:

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

If you do not want to name spatial:

 distance=np.power(np.sum(np.power(data-center,2),axis=1),.5) points_in_sphere=data[distance<cutoff]