是否可以在 Python / Numpy 中对这个三重循环进行矢量化？答案

【问题标题】：Is vectorizing this triple for loop in Python / Numpy possible?是否可以在 Python / Numpy 中对这个三重循环进行矢量化？
【发布时间】：2012-10-25 09:31:42
【问题描述】：

我正在尝试加快我的代码，目前在 Python / Numpy 中运行需要一个多小时。大部分计算时间发生在下面粘贴的函数中。

我正在尝试对 Z 进行矢量化，但我发现三重 for 循环相当困难。我可以在某处实现numpy.diff 函数吗？看看：

def MyFESolver(KK,D,r,Z):
    global tdim
    global xdim
    global q1
    global q2
    for k in range(1,tdim):
        for i in range(1,xdim-1):
            for j in range (1,xdim-1):
                Z[k,i,j]=Z[k-1,i,j]+r*q1*Z[k-1,i,j]*(KK-Z[k-1,i,j])+D*q2*(Z[k-1,i-1,j]-4*Z[k-1,i,j]+Z[k-1,i+1,j]+Z[k-1,i,j-1]+Z[k-1,i,j+1])
    return Z

tdim = 75xdim = 25

【问题讨论】：

试试Z[1:, 1:-1, 1:-1] = Z[:-1, 1:-1, 1:-1] + r*q1*Z[:-1, 1:-1, 1:-1]*(KK-Z[:-1, 1:-1, 1:-1]) + D*q2*(Z[:-1,:-2,1:-1] - 4*Z[:-1, 1:-1, 1:-1] + Z[:-1, 2:, 1:-1] + Z[:-1, 1:-1, :-2] + Z[:-1, 1:-1, 2:])这行，而不是你的三重行。
无关，但很重要：您可能需要重新检查 global 关键字的使用。在你的情况下它是没用的。

标签： python numpy vectorization

【解决方案1】：

我同意，这很棘手，因为四个边上的 BC 破坏了刚度矩阵的简单结构。您可以像这样摆脱空间循环：

from pylab import *
from scipy.sparse.lil import lil_matrix
tdim = 3;     xdim = 4;  r = 1.0;  q1, q2 = .05, .05; KK= 1.0; D = .5  #random values
Z = ones((tdim, xdim, xdim))
#Iterate in time
for k in range(1,tdim):
    Z_prev = Z[k-1,:,:] #may need to flatten
    Z_up = Z_prev[1:-1,2:]
    Z_down = Z_prev[1:-1,:-2]

    Z_left = Z_prev[:-2,1:-1]
    Z_right = Z_prev[2:,1:-1]

    centre_term  = (q1*r*(Z_prev[1:-1,1:-1] + KK) - 4*D*q2)* Z_prev[1:-1,1:-1] 

    Z[k,1:-1,1:-1]= Z_prev[1:-1,1:-1]+ centre_term + q2*(Z_up+Z_left+Z_right+Z_down)

但我不认为你可以摆脱时间循环......

我认为的表达方式：

Z_up = Z_prev[1:-1,2:]

在 numpy 中创建一个副本，而您想要的是一个视图 - 如果您能弄清楚如何做到这一点 - 它应该会更快（多少？）

最后，我同意其他回答者的观点——根据经验，这种循环最好在 C 中完成，然后包装到 numpy.不过上面的应该比原来的要快...

【讨论】：

是的，对，除了时间循环之外的一切都可以矢量化。 Numpy 总是为切片操作创建视图。

【解决方案2】：

这看起来是 Cython 的理想案例。我建议在 Cython 中编写该函数，它可能会快数百倍。

【讨论】：