如何减少此嵌套for循环的运行时？

Question

我正在编写一个代码，将转换应用于一组xyz坐标。旧xyz坐标保留在阵列fea中，新xyz坐标指定给新fea。npref是fea的较小“基本”版本（也包含xyz坐标），W是定义平移的矩阵

代码如下：

@njit(cache=True, nogil=True)
def outerLoop(npref, fea, newFea, W):
n = len(W)
for i in range(1000):

    aTerms = W[n - 4] + fea[i][0] * W[n - 3] + fea[i][1] * W[n - 2] + fea[i][2] * W[n - 1]

    total = zeros((3,))

    for j in range(len(npref)):
        dist = sqrt(sum(square(npref[j] - fea[i])))
        u = square(dist) * sqrt(log(square(dist)))
        total += u * W[j]

    newFea[i] = total + aTerms
return newFea

目前，我正在调试它，我将达到1000，但在实践中，这将不得不运行到大约130万。我得到一个大约4秒的运行时间，对于实际应用程序来说大约是一个小时，我希望运行时间减少到1秒以下，这样完整的代码可以在15分钟内运行

Answer 1

我找到的解决方案是使用numba的内置并行=真正的装饰器。这将运行时间缩短到了可接受的20分钟。下面是最终结果的代码块

 """
@njit(cache=True, nogil=True, fastmath=True, parallel=True)
def outerLoop(npref, fea, newFea, W):
    n = len(W)
    W4, W3, W2, W1 = W[n - 4], W[n - 3], W[n - 2], W[n - 1]
    for i in prange(len(fea)):

        aTerms = W4 + fea[i][0] * W3 + fea[i][1] * W2 + fea[i][2] * W1

        total = zeros((3,))

        for j in prange(len(npref)):
            dist = sum(square(npref[j] - fea[i]))
            u = dist * log(dist)
            total += u * W[j]

       
        newFea[i] = total + aTerms
    return newFea