pythonnumpy

How to use Numpy advanced indexing with a grid to select certain values from a multidimensional array?

I have a simulation of particles moving on a grid, created with numpy.meshgrid, and I want to determine the distance of every point on the grid to the closest particle. I've already found the closest particles, and stored them in a 2d array closest which has the same dimensions as the grid and contains the index of the particle closest to the corresponding point. I also have the 3d array distances with dimensions (number of parties, x size, y size) and which holds every point's distance to every particle.

I need to do calculations with the distance of each point to the closest particle.

I made the mistake of using numpy advanced indexing without fully understanding it. I tried

distances[closest[0], closest[1]]

with the intention of returning a matrix of distances to the closest particles. This was fast and produced the right shape of output, but is incorrect. What did I do wrong?

For now I have fixed the issue with explicit for loops, but since this is in my code's innermost loop I would like to make better use of numpy's speed. This code works:

for x in range(0, xSize):
    for y in range(0, ySize):
        closestDistances[x, y] = distances[closest[x, y], x, y]

Thanks in advance!

Solution

  • Check this out:

    N = 10
X = 5
Y = 4

distances = np.random.rand(N,X,Y)
closest = np.random.randint(N, size=(X,Y))

# your reference slow implementation
closestDistances = np.zeros((X,Y))
for x in range(X):
    for y in range(Y):
        closestDistances[x, y] = distances[closest[x, y], x, y]

# fast vectorized numpy implementation
x,y = np.meshgrid(np.arange(X), np.arange(Y), indexing='ij')
closestDistances_2 = distances[closest[x,y],x,y]

# they are the same:
np.all(closestDistances_2 == closestDistances) # True