Iterate numpy array and keep dimensions

Question

I would like to iterate through a 2D numpy array (especially using nditer) and keep the number of dimensions of the iterated array. In doing this I would like to avoid any operation on the iterated elements. The solution should be set only once either on the iterated array or the iterator.

import numpy as np

X = np.random.randn(5, 2)

for row in X:
    print(row.shape)

>>> (2,)
(2,)
(2,)
(2,)
(2,)

However I would like to have:

>> (1, 2)
(1, 2)
(1, 2)
(1, 2)
(1, 2)

I would like to find a solution that can be applied only once and not for each row like row = row.reshape(1,-1).

Answer 1

Iterating on an array is like iterating on a list of lists - it returns the elements as indexed on the first dimension:

In [48]: X = np.random.randn(5, 2)                                           
In [49]: X[0,:]                                                              
Out[49]: array([0.59964924, 0.46057338])
In [50]: for row in X:print(row)                                             
[0.59964924 0.46057338]
[1.09308258 0.06495922]
[ 0.98928476 -1.07894574]
[-1.31303644 -0.34589506]
[0.31475676 0.3003112 ]

You could iterate on indices, and use a list index:

In [51]: for i in range(X.shape[0]): print(X[[i],:])                         
[[0.59964924 0.46057338]]
[[1.09308258 0.06495922]]
[[ 0.98928476 -1.07894574]]
[[-1.31303644 -0.34589506]]
[[0.31475676 0.3003112 ]]

X[[1]], X[1:2], X[1][None,:], X[None,1] all do the same.

I don't think there's a way of incorporating those directly into a

for ... in X:

expression.

nditer can be awkward to use. Normally it iterates at the element level, not at the 'row' level, giving us a 0d array. And it isn't any faster than a for iteration. So I don't think it's useful here.

====

The suggest link, Numpy index slice without losing dimension information, inspired me to try:

In [57]: for row in X[:,None]: print(row)                                    
[[0.59964924 0.46057338]]
[[1.09308258 0.06495922]]
[[ 0.98928476 -1.07894574]]
[[-1.31303644 -0.34589506]]
[[0.31475676 0.3003112 ]]

In effect I'm turning X into a (5,1,2) array, so iteration on the first dimension produces (1,2) elements. Instead of preserving a dimension, I'm adding one.