Get 26 nearest neighbors of a point in 3D space

Say you have a point in 3D space with coordinate (2,2,2). How can you vectorize the operation with either numpy (I was thinking of using meshgrid, I just have not been able to get it to work) or scipy to find the 26 nearest neighbors in 3D space? There are 26 neighbors because I am considering the point as a cube, and thus the neighbors would be the 6 neighbors along the cube faces 8 neighbors along the cube corners 12 neighbors connected to cube edges.

So for point (2,2,2), how can I get the following coordinates:

(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 1), (1, 2, 2), (1, 2, 3), (1, 3, 1), (1, 3, 2), (1, 3, 3), (2, 1, 1), (2, 1, 2), (2, 1, 3), (2, 2, 1), (2, 2, 3), (2, 3, 1), (2, 3, 2), (2, 3, 3), (3, 1, 1), (3, 1, 2), (3, 1, 3), (3, 2, 1), (3, 2, 2), (3, 2, 3), (3, 3, 1), (3, 3, 2), (3, 3, 3)

I have already implemented this with a triple for loop, which works. However, speed is critical for my system and thus I need to vectorize this operation in order for my system not to fail. The triple for loop is as follows:

neighbors = [] # initialize the empty neighbor list

# Find the 26 neighboring voxels' coordinates
for i in [-1, 0, 1]: # i coordinate
    for j in [-1, 0, 1]: # j coordinate
        for k in [-1, 0, 1]: # k coordinate
            if (i ==0 and j ==0 and k ==0): # if at the same point
                pass # skip the current point
            else:
                neighbors.append((self.current_point[0] i,self.current_point[1] j,self.current_point[2] k)) # add the neighbor to the neighbors list

This is my first post to StackOverflow, so apologies if I missed anything. This is for a path planning algorithm I wish to put on a drone, and thus time is critical so that I don't hit a wall or something.

CodePudding user response：

You could create a motion directions in easy way using itertools.product:

from itertools import product 
motion = np.array(list(product(m, repeat=3)))
motion
>>> array([[-1, -1, -1],
           [-1, -1,  0],
           [-1, -1,  1],
           [-1,  0, -1],
           [-1,  0,  0],
           [-1,  0,  1],
           [-1,  1, -1],
           [-1,  1,  0],
           [-1,  1,  1],
           [ 0, -1, -1],
           [ 0, -1,  0],
           [ 0, -1,  1],
           [ 0,  0, -1],
           [ 0,  0,  0],
           [ 0,  0,  1],
           [ 0,  1, -1],
           [ 0,  1,  0],
           [ 0,  1,  1],
           [ 1, -1, -1],
           [ 1, -1,  0],
           [ 1, -1,  1],
           [ 1,  0, -1],
           [ 1,  0,  0],
           [ 1,  0,  1],
           [ 1,  1, -1],
           [ 1,  1,  0],
           [ 1,  1,  1]])

This is quite easy but slower than pure numpy:

m = [-1, 0, 1]
motion = np.stack(np.meshgrid(m, m, m), axis=-1).reshape(-1, 3)

or:

motion = np.transpose(np.indices((3,3,3)) - 1).reshape(-1, 3)

Then remove origin (0, 0, 0) (it's in the middle):

motion = np.delete(motion, 13, axis=0)

And then you're able to capture all the surrounding points:

motion   [[2, 2, 2]]
>>> array([[1, 1, 1],
           [1, 1, 2],
           [1, 1, 3],
           [1, 2, 1],
           [1, 2, 2],
           [1, 2, 3],
           [1, 3, 1],
           [1, 3, 2],
           [1, 3, 3],
           [2, 1, 1],
           [2, 1, 2],
           [2, 1, 3],
           [2, 2, 1],
           [2, 2, 3],
           [2, 3, 1],
           [2, 3, 2],
           [2, 3, 3],
           [3, 1, 1],
           [3, 1, 2],
           [3, 1, 3],
           [3, 2, 1],
           [3, 2, 2],
           [3, 2, 3],
           [3, 3, 1],
           [3, 3, 2],
           [3, 3, 3]])

CodePudding user response：

If you already have the coordinates that you are comparing against as a numpy array, say it is x, then you can calculate the euclidean distance between (2, 2, 2) and x with

distances = np.power(x - (2, 2, 2), 2).sum(axis=1)

Now you just want the indices of the 26 smallest, which you can do with np.argpartition

indices = np.argpartition(distances, 26)[:26]

Now, to get the actual elements:

elements = x[indices, :]

As an example, if x is just the set of points that you provided, then

distances = np.power(x - (2, 2, 2), 2).sum(axis=1)
indices = np.argpartition(distances, 2)[:2]
elements = x[indices, :]

This returns

array([[1, 2, 2],
       [2, 1, 2]])

As expected.

Note that this highly rated answer makes the claim that this function argpartition runs in linear time, as opposed to nlog(n) time that a regular sort would take to run.