Writing functions to work on vectors and matrices

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!

To be fair, matlab doesn't have 1D arrays at all. A "1D" vector in matlab _is_ 2D. — Joe Kington, Dec 16 '11 at 22:20
If you are going to use the 2D/1D output for something like computing a norm, then you don't need a 2D output, and can get away with 1D array in every case by broadcasting. — cyborg, Dec 17 '11 at 10:07

mtrw · Answer 1 · 2011-12-16T22:48:29.960

You don't need to distinguish between vectors and matrices for the mean calculation itself - if you use the axis parameter Numpy will perform the calculation along the vector (for vectors) or columns (for matrices). And then to construct the output, you can use a good old-fashioned list comprehension, although it might be a bit slow for huge matrices:

import numpy as np
m = np.mean(x,axis=0) # For vector x, calculate the mean. For matrix x, calculate the means of the columns
x_mean = np.array([m for k in x]) # replace elements for vectors or rows for matrices

Creating the output with a list comprehension is slow because it has to allocate memory twice - once for the list and once for the array. Using np.repeat or np.tile would be faster, but acts funny for vector inputs - the output will be a nested matrix with a 1-long vector in each row. If speed matters more than elegance you can replace the last line with this if:

if len(x.shape) == 1:
    x_mean = m*np.ones(len(x))
else:
    x_mean = np.tile(m, (x.shape[1],1))

By the way, your Matlab code behaves differently for row vectors and column vectors (try running it with x and x').

Good point. I was rushing through and should have take more time to explain things. Your answer does a very good job of explaining it. (As well as actually being correct, unlike mine :) ) — Joe Kington, Dec 16 '11 at 22:28
@JoeKington - Thanks! I added an alternative based on your original answer for speed. — mtrw, Dec 16 '11 at 22:49
Sorry about the difference in the MATLAB code - I was just trying to illustrate the point for column vectors (I don't expect any row vectors in my program). Thanks for the clarification though! — rapidfyre, Dec 17 '11 at 14:10

score 3 · Accepted Answer · answered Dec 17 '11 at 06:08

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

Writing functions to work on vectors and matrices

2 Answers2