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 毫秒 ± 247 微秒（平均值 ± 标准偏差，7 次运行，每次 100 次循环）

@btime test(10,.1,10)

55.075 毫秒（2176212 次分配：61.20 MiB）

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

Answer 1

只是为了比较，我想添加一个优化的 Python (Numba) 解决方案。 Oscar Smith 的解决方案看起来仍然有点慢，但这也可能是处理器速度较慢的结果。

一些比较，优化的 Numba 与优化的 Julia 代码对于学习如何在 Julia 中编写高效代码非常有用。

代码

import numpy as np
import numba as nb

def Test_orig(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

@nb.njit(fastmath=True,parallel=True)
def Test_nb(x0,dx,n):
    q = np.arange(-x0,x0,dx)
    B = np.exp(-1j*q)
    z = np.zeros(shape=(n,q.shape[0]), dtype=np.complex128)

    TMP_1 = np.empty(shape=(q.shape[0],q.shape[0]), dtype=np.complex128)
    for i in nb.prange(q.shape[0]):
        for j in range(q.shape[0]):
            TMP_1[i,j]=np.exp(-(q[i]-q[j])**2)

    TMP_2 = np.empty(shape=(q.shape[0],q.shape[0]), dtype=np.complex128)
    for s in nb.prange(1,n):
        for j in range(q.shape[0]):
            TMP_2[s,j]=np.exp(-1j*s*q[j])

    for s in nb.prange(1,n):
        for i in range(q.shape[0]):
            sum=0.j
            for j in range(q.shape[0]):
                sum+=TMP_1[i,j]*TMP_2[s,j]
            z[s,i]=sum*B[i]*dx
    return z

时间

%timeit res_1=Test_orig(10,.1,10)
2.64 ms ± 276 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit res_2=Test_nb(10,.1,10)
192 µs ± 7.06 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#@nb.njit(fastmath=True,parallel=False,cache=True) -> single threaded, caching possible
472 µs ± 19.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
#@nb.njit(fastmath=False,parallel=False,cache=True) -> without fastmath compiler flag
643 µs ± 10.9 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

res_1=Test_orig(10,.1,10)
res_2=Test_nb(10,.1,10)
np.allclose(res_1,res_2)
True

编辑：Julia Timings（来自 Oscar Smith 的功能）

@benchmark test(10,.1,10)
BenchmarkTools.Trial:
  memory estimate:  1.16 MiB
  allocs estimate:  1014
  --------------
  minimum time:     643.072 μs (0.00% GC)
  median time:      662.869 μs (0.00% GC)
  mean time:        729.127 μs (4.15% GC)
  maximum time:     34.947 ms (97.25% GC)
  --------------
  samples:          6841
  evals/sample:     1

Answer 2

最初有几件事阻碍了这一点：第一个也是最大的问题是循环理解，当您正确诊断时（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