CODEWITHSUNDEEP
pythonnumpy
Rufus
Rufus

Reputation: 5566

Iterate over columns of array as column vectors

Is there a way to iterate over the columns of a 2D numpy array such that the iterators remain column vectors?

i.e.

>>> A = np.arange(9).reshape((3,3))
[[0 1 2]
 [3 4 5]
 [6 7 8]]
>>> np.hstack([a in some_way_of_iterating(A)])
[[0 1 2]
 [3 4 5]
 [6 7 8]]

This is useful, for example, when I want to pass the column vectors into a function that transforms the individual column vector without having to clutter stuff with reshapes

Upvotes: 2

Views: 82

Answers (3)

Grayrigel
Grayrigel

Reputation: 3594

An ugly but another possible way with index and transpose:

np.hstack([A[:,i][np.newaxis].T for i in range(len(A.T))])

I am using np.newaxis to facilitate the transpose. Based on @hpaulj suggestion this can be significantly cleaned off:

np.hstack([A[:,i,np.newaxis] for i in range(A.shape[1])])

Output:

array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

Upvotes: 0

hpaulj
hpaulj

Reputation: 231395

In [39]: A = np.arange(1,10).reshape(3,3)
In [40]: A
Out[40]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

Iteration on an array operates on the first dimension. It's much like iterating on a nested list - but slower. And like the list case it too reduces the dimension.

You could iterate on the range, and use advanced indexing, [i] to maintain the 2d, "column vector" shape:

In [41]: [A[:,[i]] for i in range(3)]
Out[41]: 
[array([[1],
        [4],
        [7]]),
 array([[2],
        [5],
        [8]]),
 array([[3],
        [6],
        [9]])]

Or iterate on the transpose - but this still requires some form of reshape. I prefer the None/newaxis syntax.

In [42]: [a[:,None] for a in A.T]
Out[42]: 
[array([[1],
        [4],
        [7]]),
 array([[2],
        [5],
        [8]]),
 array([[3],
        [6],
        [9]])]

Indexing and reshape can be combined with:

In [43]: A[:,0,None]
Out[43]: 
array([[1],
       [4],
       [7]])

Or with slicing:

In [44]: A[:,1:2]
Out[44]: 
array([[2],
       [5],
       [8]])

There is a difference that may matter. A[:,[i]] makes a copy, A[:,i,None] is a view.

This may be the time to reread the basic numpy indexing docs.

https://numpy.org/doc/stable/reference/arrays.indexing.html

Upvotes: 0

Georgina Skibinski
Georgina Skibinski

Reputation: 13387

How about simple transpose:

B = np.hstack([a.reshape(-1,1) for a in A.T])

You require .reshape(-1,1) to get size of n X 1 instead of just n

Upvotes: 2

Related Questions