编写针对向量和矩阵的函数

Question

在一个更大的功能中，我正在写一些代码，用来生成一个向量或矩阵（取决于输入），这个向量/矩阵包含输入向量/矩阵 'x' 每一列的平均值。这些平均值会被存储在一个和输入向量形状相同的向量/矩阵里。

我目前的解决方案可以处理一维数组和矩阵，但代码写得非常(!)乱：

# '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))       

我之前是用MATLAB的，这在MATLAB中看起来是这样的：

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

这个方法可以在向量和矩阵上都正常工作，没有错误。

我的问题是，如何才能让我的Python代码在处理向量和矩阵格式的输入时，不用（丑陋的）try/except块呢？

谢谢！

Answer 1

在计算平均值的时候，其实你不需要区分向量和矩阵。如果你使用Numpy的axis参数，它会根据你是用向量还是矩阵来自动进行计算。对于向量，它会沿着向量计算；对于矩阵，它会沿着每一列计算。然后，生成输出的时候，你可以用传统的列表推导式，虽然对于特别大的矩阵来说，这可能会有点慢：

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

用列表推导式生成输出会比较慢，因为它需要分配两次内存——一次是为列表，另一次是为数组。如果你想要更快，可以使用np.repeat或np.tile，但它们在处理向量输入时会有点奇怪——输出会变成一个嵌套的矩阵，每一行里都有一个长度为1的向量。如果速度比优雅更重要，你可以把最后一行替换成这个：

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

顺便提一下，你的Matlab代码在处理行向量和列向量时表现是不同的（试着用x和x'运行一下）。

Answer 2

首先，关于numpy中的广播，有一点要快速说明一下。当我从matlab转到python时，广播让我有点困惑，但一旦我花时间去理解它，我意识到它是多么有用。如果想了解更多关于广播的内容，可以看看这个链接：http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html。

因为广播，在numpy中一个形状为(m,)的数组（你可以理解为一个向量）实际上等同于一个形状为(1, m)的数组或者(1, 1, m)的数组，依此类推。看起来你想让一个(m,)的数组表现得像一个(m, 1)的数组。我相信这种情况有时会发生，特别是在linalg模块中，但如果你要这样做，你需要知道这会打破numpy的常规用法。

在这个警告之后，下面是代码：

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)

还有一个例子：

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])

这个方法比使用tile要快，下面是时间的对比：

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