用Numpy矢量化操作替换For循环

时间:2017-03-30 04:01:34

标签: python numpy

在下面的代码中,我有一个简单的for循环,我希望用更快的矢量化numpy操作替换它。

import numpy as np

b = np.array([9,8100,-60,7], dtype=np.float64)
a = np.array([584,-11,23,79,1001,0,-19], dtype=np.float64)
m = 3
n = b.shape[0]
l = n-m+1
k = a.shape[0]-m+1

QT = np.array([-85224., 181461., 580047., 8108811., 10149.])
QT_first = QT.copy()

out = [None] * l
for i in range(1, l):
    QT[1:] = QT[:k-1] - b[i-1]*a[:k-1] + b[i-1+m]*a[-(k-1):]
    QT[0] = QT_first[i]

    # Update: Use this QT to do something with the ith element of array x
    # As i updates in each iteration, QT changes
    out[i] = np.argmin((QT + b_mean[i] * m) / (b_stddev[i] * m * a_stddev))

return out

我需要一个足够通用的解决方案来承担更长的阵列。请注意,QT取决于mb的长度,并且始终会提供。

如何用numpy向量化操作替换for循环以使其更快?

更新

我修改了原始代码,以更清楚地说明为什么卷积无法解决我的问题。 Convolve只给出了最终的QT,但实际上我需要使用中间QT值进行另一次计算,然后再更新for循环的迭代。

2 个答案:

答案 0 :(得分:5)

我没有详细研究过这一点,以确切知道它在做什么- b[i-1]*a[:k-1] + b[i-1+m]*a[-(k-1):]。但我怀疑它可以在没有i循环的情况下编写。某种-b[slice1]*a[slice2] + b[slice3]*a[slice4]

留下了是否

的问题
for i in range(1, l):
    QT[1:] = QT[:k-1] + c[i]
    QT[0] = QT_first[i]

本质上是串行和迭代的。你可能比我更了解发生了什么。我必须做一些手动循环来了解正在发生的事情。

答案 1 :(得分:5)

修改请注意,以下解决方案并未解决实际上意味着什么(仅限他们所写的内容; - ))。

fftconvolve似乎有点帮助(选择卷积解决方案):

k, n, m = 100, 100, 10
loopy                 0.47971359 ms
opt                   0.18304241 ms
k, n, m = 1000, 800, 20
loopy                 6.44813690 ms
opt                   0.33353859 ms
k, n, m = 10000, 10000, 1000
loopy               313.31501430 ms
opt                   2.61437489 ms

代码:

import numpy as np
from scipy import signal

import types
from timeit import timeit

b = np.array([9,8100,-60,7], dtype=np.float64)
a = np.array([584,-11,23,79,1001,0,-19], dtype=np.float64)
m = 3
n = b.shape[0]
l = n-m+1
k = a.shape[0]-m+1

QT = np.array([-85224., 181461., 580047., 8108811., 10149.])

def mock_data(k, n, m):
    QT = np.random.random((k,))
    b =  np.random.random((n,))
    a = np.random.random((k+m-1,))
    return QT, a, b, m

def f_loopy(QT, a, b, m):
    k, l = a.size - m + 1, b.size - m + 1
    QT, QT_first = QT.copy(), QT

    for i in range(1, l):
        QT[1:] = QT[:k-1] - b[i-1]*a[:k-1] + b[i-1+m]*a[-(k-1):]
        QT[0] = QT_first[i]
    return QT

def f_opt(QT, a, b, m):
    k, l = a.size - m + 1, b.size - m + 1
    ab1 = signal.fftconvolve(b[l-2::-1], a[:k-1], mode='full')[:QT.size-1]
    ab2 = signal.fftconvolve(b[l-2+m:m-1:-1], a[1-k:], mode='full')[:QT.size-1]
    return np.r_[0, ab2 - ab1] + np.r_[QT[l-1::-1], QT[1:QT.size+1-l]]

for k, n, m in [(100, 100, 10),
                (1000, 800, 20),
                (10000, 10000, 1000)]:
    data = mock_data(k, n, m)
    ref = f_loopy(*data)
    print(f'k, n, m = {k}, {n}, {m}')
    for name, func in list(globals().items()):
        if not name.startswith('f_') or not isinstance(func, types.FunctionType):
            continue
        try:
            assert np.allclose(ref, func(*data))
            print("{:16s}{:16.8f} ms".format(name[2:], timeit(
                'f(*data)', globals={'f':func, 'data':data}, number=10)*100))
        except:
            print("{:16s} apparently failed".format(name[2:]))