I am working with numpy indexing in a for loop and wanted to know if there is a way to do it without the loop. I have tried searching through stackoverflow but couldn't quite find an answer to my specific use case (so I apologize if this question has already been asked and answered).

Essentially I have a 3d numpy array of zeros that I am trying to fill on specific indices based on multi-dimensional indexing arrays. Here is a simplified version of my current implementation.

import numpy as np

# Some shape parameters (for the example)
x = 3
y = 2
h = 10
w = 12

# Create the arrays 
arr = np.zeros((x,h,w))
index_arr1 = np.array( # has shape (x,y,4)
      [[[9, 0, 7, 1],
        [1, 3, 6, 8]],

       [[8, 1, 2, 8],
        [0, 1, 0, 9]],

       [[0, 8, 7, 8],
        [0, 4, 5, 9]]])
index_arr2 = np.array([4, 8]) # has shape (y,)
value_arr = np.ones((x,y))

# Loop through the indices to fill 'arr'
for i in range(y):
    for j in range(x):
        arr[j, index_arr1[j,i], index_arr2[i]] = value_arr[j,i]

Is there a way to remove the loop to make the code faster and more efficient while still having the same output?

Thanks in advance for any help :)

CodePudding user response：

Replacing the ranges with arange:

In [59]: I = np.arange(y); J = np.arange(x)[:,None];

In [60]: new = np.zeros((x,h,w))
In [62]: new[J[:,None], index_arr1[J,I], index_arr2[I,None]] = value_arr[J,I,None]

In [63]: np.allclose(new, arr)
Out[63]: True

The main challenge was that the middle index was:

In [64]: index_arr1[J,I].shape
Out[64]: (3, 2, 4)

The 1st index then has to be (3,1,1), and the last (2,1) to broadcast with this. And value has to be (3,2,1). I had to try several things before I got the indexing right. With patience you probably could have do so as well.

put/take_along_axis were intended to take care of many of these details, but I haven't used them enough to apply them to your complicated cases.