How to index multidimensional array multiple times in a vectorized way numpy?

Question

I am trying to index a multidimensional array (4-dimensions) in numpy. The shape of the array is of (125,125,125,3). I have 3 separate 2D lists of index arrays. The lists are of size (N,4), (M,4), and (1,4) respectively. The 3 separate lists represent the rows, columns, and depth values in the 4D array that I am trying to index. For example consider the following:

ix = [[0,1,2,3],
     [3,4,5,6]]
iy = [[2,3,4,5],
     [5,6,7,8]]
iz = [[1,2,3,4]]

weights.shape = (125,125,125,3)

I want to index weights with every possible combination of row, column, and depth index arrays in ix,iy, and iz. For example, if I take the first row in each of the index matrices, that means I want to select rows [0,1,2,3], columns [2,3,4,5], and depth values [1,2,3,4] in weights. I always want to select all elements in the 4th dimension of weights. This means that I am essentially selecting a (4,4,4,3) slice of weights.

Right now, I have implemented this by indexing with loops using the following code

w = np.empty(shape=(X,Y,Z,4,4,4,weights.ndim-1))
for i in range(X):
    for j in range(Y):
        w_ij = np.ix_(ix[i,:], iy[j,:], iz[0,:])
        w[i,j,0,:,:,:,:] = weights[w_ij[0], w_ij[1], w_ij[2], :]

My final goal is to construct the multidimensional array w that is of shape (N,M,1,4,4,4,3) as fast as possible. This part of the code is going to run multiple times, so if there is a vectorized way of doing this with built-in numpy functions, that would be ideal.

Please let me know if there are any clarifying questions. This is my first time asking a question on stack overflow, so I apologize if anything is unclear or confusing!

Answer 1

You can use indexing with broadcasting to achieve this.

import numpy as np

weights = np.random.rand(125, 125, 125, 3)

ix = np.array([[0,1,2,3], [3,4,5,6]])
iy = np.array([[2,3,4,5], [5,6,7,8]])
iz = np.array([[1,2,3,4]])

X = len(ix)
Y = len(iy)
Z = len(iz)

def compute1(weights):
    w = np.empty(shape=(X, Y, Z, 4, 4, 4, weights.ndim-1))
    for i in range(X):
        for j in range(Y):
            w_ij = np.ix_(ix[i,:], iy[j,:], iz[0,:])
            w[i,j,0,:,:,:,:] = weights[w_ij[0], w_ij[1], w_ij[2], :]
    return w

def compute2(weights):
    return weights[ix[:, None, None, :, None, None], iy[None, :, None, None, :, None], iz[None, None, :, None, None, :]]

print(np.allclose(compute1(weights), compute2(weights)))

Gives True.

Bench-marking -

%timeit compute1(weights)
%timeit compute2(weights)

Gives -

36.7 µs ± 897 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
6.28 µs ± 62.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

As you can see, the broadcasting solution is around 6x faster for data of this size.