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

Question

在下面的代码中，我有一个简单的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取决于m和b的长度，并将始终提供。

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

更新

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

Answer 1

我没有详细研究这个以确切地知道它在做什么- 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]

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

Answer 2

编辑请注意，下面的解决方案没有解决OP实际意味着什么（只有他们写的，最初的;-)）。

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