python-3.x - 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 per loop (mean ± std. dev. of 7 runs, 100 loops each)

@btime test(10,.1,10)

55.075 ms (2176212 allocations: 61.20 MiB)

也就是说，Julia 代码要慢得多。我绝对确定我在这里做错了什么，因为分配不应该那么大。我尝试尽可能地优化它，但我几天前开始学习 Julia，但无法走得更远。非常感谢任何有关如何提高性能的提示。

Answer 1

最初有几个因素阻碍了这一点:第一个也是最大的一个是循环理解，它在您正确诊断时分配了大量资源(16 毫秒)。一旦解决了这个问题，最大的问题是它重复计算相同复数指数的方式(1.6 毫秒)。

解决这两个问题后，认识到问题中的线性代数可以使代码更清晰，并且还允许 Julia 调用 blas 进行更有效的矩阵乘法。 (900 微秒)。这是最新的代码，比等效的 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 尝试，这里还有一项尝试，通过解决将实数和复数数组相乘会将实数转换为复数的错误，将时间降至 725μs。手动编码乘法会产生这个结果。

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