I want to extract a complete slice from a 3D numpy array using ndeumerate or something similar.
arr = np.random.rand(4, 3, 3)
I want to extract all possible arr[:, x, y] where x, y range from 0 to 2
I want to extract a complete slice from a 3D numpy array using ndeumerate or something similar.
arr = np.random.rand(4, 3, 3)
I want to extract all possible arr[:, x, y] where x, y range from 0 to 2
ndindex
is a convenient way of generating the indices corresponding to a shape:
In [33]: arr = np.arange(36).reshape(4,3,3)
In [34]: for xy in np.ndindex((3,3)):...: print(xy, arr[:,xy[0],xy[1]])...:
(0, 0) [ 0 9 18 27]
(0, 1) [ 1 10 19 28]
(0, 2) [ 2 11 20 29]
(1, 0) [ 3 12 21 30]
(1, 1) [ 4 13 22 31]
(1, 2) [ 5 14 23 32]
(2, 0) [ 6 15 24 33]
(2, 1) [ 7 16 25 34]
(2, 2) [ 8 17 26 35]
It uses nditer
, but doesn't have any speed advantages over a nested pair of for loops.
In [35]: for x in range(3):...: for y in range(3):...: print((x,y), arr[:,x,y])
ndenumerate
uses arr.flat
as the iterator, but using it to
In [38]: for xy, _ in np.ndenumerate(arr[0,:,:]):...: print(xy, arr[:,xy[0],xy[1]])
does the same thing, iterating on the elements of a 3x3 subarray. As with ndindex
it generates the indices. The element won't be the size 4 array that you want, so I ignored that.
A different approach is to flatten the later axes, transpose, and then just iterate on the (new) first axis:
In [43]: list(arr.reshape(4,-1).T)
Out[43]:
[array([ 0, 9, 18, 27]),array([ 1, 10, 19, 28]),array([ 2, 11, 20, 29]),array([ 3, 12, 21, 30]),array([ 4, 13, 22, 31]),array([ 5, 14, 23, 32]),array([ 6, 15, 24, 33]),array([ 7, 16, 25, 34]),array([ 8, 17, 26, 35])]
or with the print as before:
In [45]: for a in arr.reshape(4,-1).T:print(a)