是否可以在Python / Numpy中向量化这个三重for循环?

时间:2012-10-25 09:31:42

标签: python numpy vectorization

我正在尝试加速我的代码,目前需要花费一个多小时才能在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 = 75 xdim = 25

2 个答案:

答案 0 :(得分:4)

我同意,这很棘手,因为所有四个方面的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。但上面应该比原来快......

答案 1 :(得分:2)

这看起来像Cython的理想情况。我建议在Cython中编写该函数,它可能要快几百倍。