How can I always have numpy.ndarray.shape return a two valued tuple?

Question

I'm trying to get the values of (nRows, nCols) from a 2D Matrix but when it's a single row (ie x = np.array([1, 2, 3, 4])), x.shape will return (4,) and so my statement of (nRows, nCols) = x.shape returns "ValueError: need more than 1 value to unpack"

Any suggestions on how I can make this statement more adaptable? It's for a function that is used in many programs and should work with both single row and multi-row matices. Thanks!

Answer 1

You could create a function that returns a tuple of rows and columns like this:

def rowsCols(a):
    if len(a.shape) > 1:
        rows = a.shape[0]
        cols = a.shape[1]
    else:
        rows = a.shape[0]
        cols = 0
    return (rows, cols)

where a is the array you input to the function. Here's an example of using the function:

import numpy as np

x = np.array([1,2,3])

y = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])

def rowsCols(a):
    if len(a.shape) > 1:
        rows = a.shape[0]
        cols = a.shape[1]
    else:
        rows = a.shape[0]
        cols = 0
    return (rows, cols)

(nRows, nCols) = rowsCols(x)

print('rows {} and columns {}'.format(nRows, nCols))

(nRows, nCols) = rowsCols(y)

print('rows {} and columns {}'.format(nRows, nCols))

This prints rows 3 and columns 0 then rows 4 and columns 3 . Alternatively, you can use the atleast_2d function for a more concise approach:

(r, c) = np.atleast_2d(x).shape

print('rows {}, cols {}'.format(r, c))

(r, c) = np.atleast_2d(y).shape

print('rows {}, cols {}'.format(r, c))

Which prints rows 1, cols 3 and rows 4, cols 3 .

Answer 2

If your function uses

 (nRows, nCols) = x.shape 

it probably also indexes or iterates on x with the assumption that it has nRows rows, eg

 x[0,:]
 for row in x:
    # do something with the row

Common practice is to reshape x (as needed) so it has at least 1 row. In other words, change the shape from (n,) to (1,n) .

 x = np.atleast_2d(x)

does this nicely. Inside a function, such a change to x won't affect x outside it. This way you can treat x as 2d through out your function, rather than constantly looking to see whether it is 1d v 2d.

Python: How can I force 1-element NumPy arrays to be two-dimensional?

is one of many previous SO questions that asks about treating 1d arrays as 2d.