CODEWITHSUNDEEP
pythonnumpy
normanius
normanius

Reputation: 9782

How can I ensure a numpy array to be either a 2D row- or column vector?

Is there a numpy function to ensure a 1D- or 2D- array to be either a column or row vector?

For example, I have either one of the following vectors/lists. What is the easiest way to convert any of the input into a column vector?

x1 = np.array(range(5))
x2 = x1[np.newaxis, :]
x3 = x1[:, np.newaxis]

def ensureCol1D(x):
    # The input is either a 0D list or 1D.
    assert(len(x.shape)==1 or (len(x.shape)==2 and 1 in x.shape))
    x = np.atleast_2d(x)
    n = x.size
    print(x.shape, n)
    return x if x.shape[0] == n else x.T

assert(ensureCol1D(x1).shape == (x1.size, 1))
assert(ensureCol1D(x2).shape == (x2.size, 1))
assert(ensureCol1D(x3).shape == (x3.size, 1))

Instead of writing my own function ensureCol1D, is there something similar already available in numpy that ensures a vector to be column?

Upvotes: 5

Views: 1674

Answers (1)

Primusa
Primusa

Reputation: 13498

Your question is essentially how to convert an array into a "column", a column being a 2D array with a row length of 1. This can be done with ndarray.reshape(-1, 1).

This means that you reshape your array to have a row length of one, and let numpy infer the number of rows / column length. 

x1 = np.array(range(5))
print(x1.reshape(-1, 1))

Output:

array([[0],
   [1],
   [2],
   [3],
   [4]])

You get the same output when reshaping x2 and x3. Additionally this also works for n-dimensional arrays:

x = np.random.rand(1, 2, 3)
print(x.reshape(-1, 1).shape)

Output:

(6, 1)

Finally the only thing missing here is that you make some assertions to ensure that arrays that cannot be converted are not converted incorrectly. The main check you're making is that the number of non-one integers in the shape is less than or equal to one. This can be done with:

assert sum(i != 1 for i in x1.shape) <= 1

This check along with .reshape let's you apply your logic on all numpy arrays.

Upvotes: 5

Related Questions