返回正确维度数组的 Pythonic 方法

Question

我想创建一个函数，如果给定的话返回一个 numpy 数组，或者如果给定的话返回一个多维 numpy 数组。例如:

import numpy as np;
def running_average(data,windowSize):
    dShape = np.shape(data);
    if(len(dShape)==1):
        raOut = np.zeros(len(data));
        rSum = 0.0;
        for row,value in enumerate(data):
            if row<windowSize:
                rSum+=float(value);
            else:
                rSum=rSum-data[row-windowSize]+value;
            raOut[row]=rSum/windowSize;
    else:
        raOut = np.zeros(dShape);
        for col in xrange(dShape[1]):
            rSum=0.0;
            for row,value in enumerate(data[:,col]):
                if row<windowSize:
                    rSum+=float(value);
                else:
                    rSum=rSum-data[row-windowSize,col]+value;
                raOut[row,col]=rSum/windowSize;

    return raOut;

但是必须要有一个好的测试，这样我就不必在 if 和 else 语句中重复自己。

我是 python 新手，首选方法是什么？

Answer 1

怎么样:

def running_avg(data, ws):
    tmp = np.cumsum(data, axis=-1, dtype='float')
    ra = (tmp[..., ws:] - tmp[..., :-ws]) / ws
    return ra

这将在最后一个轴上取平均值，如果你想变得非常聪明，你可以让函数接受一个轴参数并在任意轴上取平均值。

更新

我相信这个版本与您上面的代码一致。

def running_avg(data, ws):
    ra = np.cumsum(data, axis=-1, dtype='float') / ws
    ra[..., ws:] = ra[..., ws:] - ra[..., :-ws]
    return ra

对于更一般的问题，使用 numpys 内置函数(例如 cumsum)会有所帮助，因为它们已经这样做了，但是如果您确实必须循环，则可以使用 A = np.zeros(A.shape) 来获得相同的数组shape 作为输入，然后使用 A[..., i] 始终对最后一个维度进行操作，或使用 A[..., i, :] 始终对倒数第二个维度进行操作，依此类推。此外，有时人们会执行 data = np.roll(data, axis) 将轴移动到开头，然后您使用 A[i] 在第一个维度上操作并在需要时将轴移回。

更新 2:

我只记得为什么以下是一个非常糟糕的主意(至少在这种情况下):

ra[..., ws:] -= ra[..., :-ws]

你应该改用这个:

ra[..., ws:] = ra[..., ws:] - ra[..., :-ws]