How To Mix Numpy Slices To List Of Indices?

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'

Solution 1:

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])

Solution 2:

You can add slice(None) to the index tuple manually:

>>> grid.shape
(3, 20, 17, 9)
>>> indices
(19, 16, 8)
>>> grid[:,19,16,8]array([3059, 6119, 9179])
>>> grid[(slice(None),) + indices]array([3059, 6119, 9179])

See here in the documentation for more.

Solution 3:

I believe what you are looking for is grid[1:][grid[0]].

grid = np.array([
        [0, 2, 1],  # N
        [1, 9, 3, 6],  # M_1
        [7, 8, 2, 5, 0, 8, 3],  # M_2
        [4, 8]  # M_3
    ])

np.array([grid[a[0] + 1][n] for a, n in np.ndenumerate(grid[0])])
# array([1, 2, 8])