Extract sub-array: staggered indexing

I was wondering whether there is an obvious way of improving this code. I tried a few things (numpy.ix_ and numpy.take) but I could not get them to work.

For a given array of indeces idx_array of size (nZ, nA, nP) I want to get the sub-array of the array A of size (nZ, nA, nW, nP) as follows

for iz in range(nZ):
    for ip in range(nP):
         for ia in range(nA):
             Anew[iz, ia, ip] = A[iz, ia, idx_array[iz, ia, ip], ip]

Thank you in advance for your help.

It looks like you are trying to slice the arrays, if so then these may help: a. b.
If not then can you explain a little more background of what you are trying to do?

Thank you for your reply.

I guess I was hoping there was a built-in function that would allow me not to go through the for-loop and just do the extraction in one pop. Something like

Anew = np.extract(A, idxZ, idxA, idxW, idxP)

where idxZ are the full set of iz indices, idxA are the full set of ia indices, idxW are the full set of idx_array[iz, ia, ip] indices, and idxP are the full set of ip indices.

I hope that makes it clearer. Thank you again!

mmm not so much, you have not defined indices so I cant understand if these are lists or not. If they are numbers then I don’t see why they cant’t be a list of numbers or you can convert them to a list of numbers. Also, its hard to picture what you mean as you have given very little info: background, data, examples.
But I was thinking that Dash maybe could do what you want since it handles multidimensional data. But I was wrong, however, I noticed it doesn’t do slicing with multiple axes but check this link. Because Dash uses numpy, numpy doesn’t do slicing with lists in multiple axes.
But there is a workaround via this link.
Your worse case should be slicing three times but I would check the timing between

for i in range(j) ...

you are using above and slicing (use Python stdlib ‘timeit’).