Question

As part of a larger function, I’m writing some code to generate a vector/matrix (depending on the input) containing the mean value of each column of the input vector/matrix ‘x’. These values are stored in a vector/matrix of the same shape as the input vector.

My preliminary solution for it to work on both a 1-D and matrix arrays is very(!) messy:

# 'x' is of type array and can be a vector or matrix.
import scipy as sp
shp = sp.shape(x)
x_mean = sp.array(sp.zeros(sp.shape(x)))

try: # if input is a matrix
    shp_range = range(shp[1])
    for d in shp_range:
        x_mean[:,d] = sp.mean(x[:,d])*sp.ones(sp.shape(z))
except IndexError: # error occurs if the input is a vector
    z = sp.zeros((shp[0],))
    x_mean = sp.mean(x)*sp.ones(sp.shape(z))       

Coming from a MATLAB background, this is what it would look like in MATLAB:

[R,C] = size(x);
for d = 1:C,
    xmean(:,d) = zeros(R,1) + mean(x(:,d));
end

This works on both vectors as well as matrices without errors.

My question is, how can I make my python code work on input of both vector and matrix format without the (ugly) try/except block?

Thanks!

Answer 1

First A quick note about broadcasting in numpy. Broadcasting was kinda confusing to me when I switched over from matlab to python, but once I took the time to understand it I realized how useful it could be. To learn more about broadcasting take a look at http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html,

Because of broadcasting an (m,) array in numpy (what you’re calling a vector) is essentially equivelent to an (1, m) array or (1, 1, m) array and so on. It seems like you want to have an (m,) array behave like a (m, 1) array. I believe this happens sometimes, especially in the linalg module, but if you’re going to do it you should know that you’re breaking the numpy convention.

With that warning there’s the code:

import scipy as sp
def my_mean(x):
    if x.ndim == 1:
        x = x[:, sp.newaxis]
    m = sp.empty(x.shape)
    m[:] = x.mean(0)
    return sp.squeeze(m)

and an example:

In [6]: x = sp.arange(30).reshape(5,6)

In [7]: x
Out[7]: 
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [24, 25, 26, 27, 28, 29]])

In [8]: my_mean(x)
Out[8]: 
array([[ 12.,  13.,  14.,  15.,  16.,  17.],
       [ 12.,  13.,  14.,  15.,  16.,  17.],
       [ 12.,  13.,  14.,  15.,  16.,  17.],
       [ 12.,  13.,  14.,  15.,  16.,  17.],
       [ 12.,  13.,  14.,  15.,  16.,  17.]])

In [9]: my_mean(x[0])
Out[9]: array([ 2.5,  2.5,  2.5,  2.5,  2.5,  2.5])

This is faster than using tile, the timing is bellow:

In [1]: import scipy as sp

In [2]: x = sp.arange(30).reshape(5,6)

In [3]: m = x.mean(0)

In [5]: timeit m_2d = sp.empty(x.shape); m_2d[:] = m
100000 loops, best of 3: 2.58 us per loop

In [6]: timeit m_2d = sp.tile(m, (len(x), 1))
100000 loops, best of 3: 13.3 us per loop