I have a numpy.array
, called grid
, with shape:
grid.shape = [N, M_1, M_2, ..., M_N]
The values of N, M_1, M_2, ..., M_N are known only after initialization.
For this example, let's say N=3 and M_1 = 20, M_2 = 17, M_3 = 9:
grid = np.arange(3*20*17*9).reshape(3, 20, 17, 9)
I am trying to loop over this array, as follows:
for indices, val in np.ndenumerate(grid[0]):print indices_some_func_with_N_arguments(*grid[:, indices])
At the first iteration, indices = (0, 0, 0) and:
grid[:, indices] # array with shape 3,3,17,9
whereas I want it to be an array of three elements only, much like:
grid[:, indices[0], indices[1], indices[2]] # array([ 0, 3060, 6120])
which however I cannot implement like in the above line, because I don't know a-priori what is the length of indices
.
I am using python 2.7, but a version-agnostic implementation is welcome :-)
I think you want something like this:
In [134]: x=np.arange(24).reshape(4,3,2)In [135]: x
Out[135]:
array([[[ 0, 1],[ 2, 3],[ 4, 5]],[[ 6, 7],[ 8, 9],[10, 11]],[[12, 13],[14, 15],[16, 17]],[[18, 19],[20, 21],[22, 23]]])In [136]: for i,j in np.ndindex(x[0].shape):...: print(i,j,x[:,i,j])...:
(0, 0, array([ 0, 6, 12, 18]))
(0, 1, array([ 1, 7, 13, 19]))
(1, 0, array([ 2, 8, 14, 20]))
(1, 1, array([ 3, 9, 15, 21]))
(2, 0, array([ 4, 10, 16, 22]))
(2, 1, array([ 5, 11, 17, 23]))
where the 1st line is:
In [142]: x[:,0,0]
Out[142]: array([ 0, 6, 12, 18])
Unpacking the index tuple as i,j
and using that in x[:,i,j]
is the simplest way of doing this index. But to generalize it to other numbers of dimensions I'll have to play with tuples a bit. x[i,j]
is the same as x[(i,j)]
.
In [147]: for ind in np.ndindex(x.shape[1:]):...: print(ind,x[(slice(None),)+ind])...:
((0, 0), array([ 0, 6, 12, 18]))
((0, 1), array([ 1, 7, 13, 19]))
...
with enumerate
:
for ind,val in np.ndenumerate(x[0]):print(ind,x[(slice(None),)+ind])