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