function to get number of columns in a NumPy array that returns 1 if it is a 1D array

Question

I have defined operations on 3xN NumPy arrays, and I want to loop over each column of the array.

I tried:

for i in range(nparray.shape[1]):

However, if nparray.ndim == 1, this fails.

Is there a clean way to ascertain the number of columns of a NumPy array, for example, to get 1 if it is a 1D array (like MATLAB's size operation does)?

Otherwise, I have implemented:

if nparray.ndim == 1:
    num_points = 1
else:
    num_points = nparray.shape[1]

for i in range(num_points):

Answer 1

If you're just looking for something less verbose, you could do this:

num_points = np.atleast_2d(nparray).shape[1]

That will, of course, make a new temporary array just to take its shape, which is a little silly… but it'll be pretty cheap, because it's just a view of the same memory.

However, I think your explicit code is more readable, except that I might do it with a try:

try:
    num_points = nparray.shape[1]
except IndexError:
    num_points = 1

If you're doing this repeatedly, whatever you do, you should wrap it in a function. For example:

def num_points(arr, axis):
    try:
        return arr.shape[axis]
    except IndexError:
        return 1

Then all you have to write is:

for i in range(num_points(nparray, 1)):

And of course it means you can change things everywhere by just editing one place, e.g.,:

def num_points(arr, axis):
    return nparray[:,...,np.newaxis].shape[1]

Answer 2

If you want to keep the one-liner, how about using conditional expressions:

for i in range(nparray.shape[1] if nparray.ndim > 1 else 1):
    pass

Answer 3

By default, to iterate a np.array means to iterate over the rows. If you have to iterate over columns, just iterate through the transposed array:

>>> a2=array(range(12)).reshape((3,4))
>>> for col in a2.T:
    print col

[0 4 8]
[1 5 9]
[ 2  6 10]
[ 3  7 11]

What's the intended behavior of an array array([1,2,3]), it is treated as having one column or having 3 cols? It is confusing that you mentioned that the arrays are all 3XN arrays, which means this should be the intended behavior, as it should be treated as having just 1 column:

>>> a1=array(range(3))
>>> for col in a1.reshape((3,-1)).T:
print col

[0 1 2]

So, a general solution: for col in your_array.reshape((3,-1)).T: #do something

Answer 4

I think the easiest way is to use the len function:

for i in range(len(nparray)):
   ...

Why? Because if the number of nparray is a like a one dimensional vector, len will return the number of elements. In your case, the number of columns.

nparray = numpy.ones(10)
print(len(nparray))
Out: 10

If nparray is like a matrix, the number of columns will be returned.

nparray = numpy.ones((10, 5))
print(len(nparray))
Out: 10

If you have a list of numpy arrays with different sizes, just use len inside a loop based on your list.