numpy 向量化而不是 for 循环

我用 Python 写了一些代码，运行良好，但速度很慢；我认为是由于 for 循环。我希望可以使用 numpy 命令加速以下操作。让我定义目标。

假设我有一个 2D numpy 数组all_CMs尺寸row x col。例如考虑一个6x11数组（见下图）。

1. 我想计算所有行的平均值，即sumⱼ aᵢⱼ 生成一个数组。这当然可以轻松完成。 （我称这个值为CM_tilde)

2. 现在，为了each row我想计算一些选定值的平均值，即通过计算它们的总和并将其除以所有列的数量来计算低于特定阈值的所有值（N）。如果该值高于此定义的阈值，CM_tilde添加值（整行的平均值）。这个值称为CM

3. 随后，CM从行中的每个元素中减去值

除此之外，我想要一个 numpy 数组或列表，其中所有这些CM列出了值。

如图：

以下代码可以工作，但速度非常慢（特别是当数组变大时）

CM_tilde = np.mean(data, axis=1)
N = data.shape[1]
data_cm = np.zeros(( data.shape[0], data.shape[1], data.shape[2] ))
all_CMs = np.zeros(( data.shape[0], data.shape[2]))
for frame in range(data.shape[2]):
    for row in range(data.shape[0]):
        CM=0
        for col in range(data.shape[1]):
            if data[row, col, frame] < (CM_tilde[row, frame]+threshold):
               CM += data[row, col, frame]
            else:
               CM += CM_tilde[row, frame]
        CM = CM/N
        all_CMs[row, frame] = CM
        # calculate CM corrected value
        for col in range(data.shape[1]):
            data_cm[row, col, frame] = data[row, col, frame] - CM
    print "frame: ", frame
return data_cm, all_CMs

有任何想法吗？

将您正在做的事情矢量化非常容易：

import numpy as np

#generate dummy data
nrows=6
ncols=11
nframes=3
threshold=0.3
data=np.random.rand(nrows,ncols,nframes)

CM_tilde = np.mean(data, axis=1)
N = data.shape[1]

all_CMs2 = np.mean(np.where(data < (CM_tilde[:,None,:]+threshold),data,CM_tilde[:,None,:]),axis=1)
data_cm2 = data - all_CMs2[:,None,:]

将此与您的原件进行比较：

In [684]: (data_cm==data_cm2).all()
Out[684]: True

In [685]: (all_CMs==all_CMs2).all()
Out[685]: True

逻辑是我们使用大小的数组[nrows,ncols,nframes]同时地。主要技巧是利用Python的广播，通过转动CM_tilde大小的[nrows,nframes] into CM_tilde[:,None,:]大小的[nrows,1,nframes]。然后，Python 将为每一列使用相同的值，因为这是此修改后的单一维度CM_tilde.

通过使用np.where我们选择（基于threshold) 是否要获取对应的值data，或者，再次，广播值CM_tilde。一个新的用途np.mean允许我们计算all_CMs2.

在最后一步中，我们通过直接减去这个新的来利用广播all_CMs2从相应的元素data.

通过查看临时变量的隐式索引，可能有助于以这种方式矢量化代码。我的意思是你的临时变量CM生活在一个循环中[nrows,nframes]，并且其值在每次迭代时都会重置。这意味着CM实际上是一个数量CM[row,frame]（后来显式分配给二维数组all_CMs），从这里很容易看出，您可以通过总结适当的CMtmp[row,col,frames]沿其列尺寸的数量。如果有帮助，您可以命名np.where(...)部分作为CMtmp为此目的，然后计算np.mean(CMtmp,axis=1)从那。显然，结果相同，但可能更透明。