generalisation between vector and matrix, matrix and tensor with numpy ndarray

Question

I found an interesting thing when comparing MATLAB and numpy.

MATLAB:
x = [1, 2]
n = size(X, 2) 
% n = 1

Python:
x = np.array([1, 2])
n = x.shape[1]
# error

The question is: how to handle input which may be both ndarray with shape (n,) and ndarray with shape (n, m).

e.g.

def my_summation(X):
    """
    X : ndarray
        each column of X is an observation.
    """
    # my solution for ndarray shape (n,)
    # if X.ndim == 1:
    #     X = X.reshape((-1, 1))

    num_of_sample = X.shape[1]
    sum = np.zeros(X.shape[0])
    for i in range(num_of_sample):
        sum = sum + X[:, i]

    return sum


a = np.array([[1, 2], [3, 4]])
b = np.array([1, 2])

print my_summation(a)
print my_summation(b)

My solution is forcing ndarray shape (n,) to be shape (n, 1).

The summation is used as an example. What I want is to find an elegant way to handle the possibility of matrix with only one observation(vector) and matrix with more than one observation using ndarray.

Does anyone have better solutions?

For a 1D case, since you are summing along axis=1 and you wanted (n) to become (n,1), so for a 1D array case the output from my_summation(X) would be same as input X, right? — Divakar
@mose there is nothing wrong with reshaping. I just want to find an elegant way to handle the possibility of matrix with only one observation and matrix with more than one observation using ndarray. — user2262504
@Divakar in 1D case, the summation is still X with shape (n,). By the way, what is the pros and cons of returning sum as (n,) versus (n,1)? — user2262504
@user2262504 I don't see any con in that, just a different shape occupying the same number of elements. — Divakar

Rory Yorke Rory Yorke · Accepted Answer · 2015-10-10T07:57:51

I recently learned about numpy.atleast_2d from the Python control toolbox. You also don't need a for-loop for summation, rather use numpy.sum.

import numpy as np


def my_summation(X):
    """
    X : ndarray
        each column of X is an observation.
    """
    # my solution for ndarray shape (n,)
    # if X.ndim == 1:
    #     X = X.reshape((-1, 1))
    X = np.atleast_2d(X)
    return np.sum(X, axis=1)


a = np.array([[1, 2], [3, 4]])
b = np.array([1, 2])

print my_summation(a)
print my_summation(b)

gives

[3 7]
[3]

generalisation between vector and matrix, matrix and tensor with numpy ndarray

2 Answers