Python矢量化与Julia for循环

Question

在Julia中有很多关于速度矢量化和去矢量化代码的文章，但是我的问题是我在Python和Julia中运行的一个非常简单的代码片段来比较性能。我有一个巨大的Python代码，当且仅当此类功能明显加速时，我才会将其转录为Julia。

我正在Jupyter笔记本中运行所有内容，并且我知道Julia在.jl文件中运行得更快；但是，与从终端运行相比，这种情况的差异可以忽略不计。

Python代码使用广播和矢量化：

def test(x0,dx,i):
   q = np.arange(-x0,x0,dx)  
   z = np.zeros(shape=(i,len(q)), dtype=np.complex128)
   B = np.exp(-1j*q)

   for s in range(1,i):
      A = np.exp(-(q-q[:,np.newaxis])**2)*np.exp(-1j*s*q)
      z[s] = B*np.sum(A,axis=1)*dx

   return z

在Julia翻译中，我尝试使用for循环，但最终一次使用了理解，因为它比另一个for循环要快：

function test(x0,dx,n)
    z = zeros(ComplexF64, (n, Int(2*x0/dx + 1)))
    B = exp.(-1im*(-x0:dx:x0-dx))

    for s in 1:n-1
        for (i,q) in enumerate(-x0:dx:x0-dx)    
            A = [exp(-(q-Q)^2)exp(-1im*s*Q) for Q in -x0:dx:x0-dx]
            z[s,i] = B[i]*sum(A)*dx
        end       
    end   
    z
end

结果是

%timeit test(10,.1,10)

每个循环2.81 ms±247 µs（平均±标准偏差，共运行7次，每个循环100个）

@btime test(10,.1,10)

55.075毫秒（2176212分配：61.20 MiB）

也就是说，Julia代码要慢得多。我绝对可以确定我在这里做错了，因为分配的空间不会那么大。我尽了最大努力来优化它，但是几天前我开始学习Julia，并且走得更远。任何有关如何提高性能的提示都将受到高度赞赏。

Answer 1

最初有几件事阻碍了这一点：第一个也是最大的问题是循环理解，即在正确诊断后分配一吨（16毫秒）。一旦解决，最大的问题就是重复计算相同复杂指数（1.6 ms）的方式。

解决了这两个问题后，认识到问题中的线性代数既可以使代码更简洁，也可以让茱莉亚召唤blas进行更有效的矩阵乘法。（900 s）。这是最新的代码，比等效的numpy好约3倍

using LinearAlgebra
function test(x0,dx,n)
    Q = collect(-x0:dx:x0-dx)
    A = complex.(exp.(-(Q.-transpose(Q)).^2))
    B = exp.(-im.*transpose(1:n-1).*Q)
    z = Matrix{ComplexF64}(undef, n-1, length(Q))
    for (i, q) in enumerate(Q)
        total = transpose(@view A[:, i]) * B
        z[:, i] = dx*exp(-q*im) .* total
    end
    z
end

作为对numba尝试的响应，这是另一种尝试，它通过解决一个bug来降低到725？s，该bug中乘以实数和复数数组会将实数转换为复数。手动编码乘法即可得出结果。

function test(x0,dx,n)
    Q = collect(-x0:dx:x0-dx)
    A = exp.(-(Q.-transpose(Q)).^2)
    B = exp.(-im.*transpose(1:n-1).*Q)
    rB = real.(B)
    iB = imag.(B)
    z = Matrix{ComplexF64}(undef, n-1, length(Q))
    for (i, q) in enumerate(Q)
        At = transpose(@view A[:, i])
        total =  (At * rB) .+ (At * iB).*im
        z[:, i] = dx*exp(-q*im) .* total
    end
    z
end