Numpy multidimensional indexing for np.ufunc.at and np.ix_

问题

I would like to know how I can take index from an array and multiply with another array. I have two 4d arrays and one 2d index array:

base = np.ones((2, 3, 5, 5))
to_multiply = np.arange(120).reshape(2, 3, 4, 5)
index = np.array([[0, 2, 4, 2], [0, 3, 3, 2]])

The row index of the index array corresponds to the 1st dimension of base and to_multiply, and the value of the index array corresponds to the 3rd dimension of base. I want to take the slice from base according to the index and multiply with to_multiply.

Using for loops and np.multiply.at (because I may have same index) I can achieve it by:

for i, x in enumerate(index):
    np.multiply.at(base[i, :, :, :], np.s_[:, x, :], to_multiply[i, :, :, :])

The correctness of above can be validated by:

to_multiply[0, 0, :, 0]
array([ 0,  5, 10, 15])

base[0, 0, :, 0]
array([ 0.,  1., 75.,  1., 10.])

However, I would like to know if there is one-line solution using np.multiply.at and np.ix_

I tried to use np.ix_ but I'm very confused about it because in this case it's multidimensional.

回答1:

Can't be done with ix_. From its docs:

This function takes N 1-D sequences and returns N outputs with N dimensions each,

Your index is 2d.

However, we can do the equivalent 'by-hand':

In [196]: np.multiply.at(base1, (np.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:]), to_
     ...: multiply)                                                                                  
In [197]: np.allclose(base,base1)                                                                    
Out[197]: True

The goal was to make 3 arrays that broadcast together to match to_multiply (except for the last size 5 dimension).

That is a (2,1,1), (1,3,1) and (2,1,4) => (2,3,4)

In [199]: np.broadcast_arrays(np.arange(2)[:,None,None],np.arange(3)[:,None],index[:,None,:])        
Out[199]: 
[array([[[0, 0, 0, 0],
         [0, 0, 0, 0],
         [0, 0, 0, 0]],
 
        [[1, 1, 1, 1],
         [1, 1, 1, 1],
         [1, 1, 1, 1]]]),
 array([[[0, 0, 0, 0],
         [1, 1, 1, 1],
         [2, 2, 2, 2]],
 
        [[0, 0, 0, 0],
         [1, 1, 1, 1],
         [2, 2, 2, 2]]]),
 array([[[0, 2, 4, 2],
         [0, 2, 4, 2],
         [0, 2, 4, 2]],
 
        [[0, 3, 3, 2],
         [0, 3, 3, 2],
         [0, 3, 3, 2]]])]

While I had a general idea of where I wanted to go, I had to try quite a number of ideas first.

来源：https://stackoverflow.com/questions/63143788/numpy-multidimensional-indexing-for-np-ufunc-at-and-np-ix