Python 3.x Python向量化与Julia for循环

在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循环，但最终使用了一次comprehension，因为它比另一个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个回路）

55.075毫秒（2176212次分配：61.20兆字节）

---

也就是说，Julia代码要慢得多。我绝对肯定我在这里做错了什么，因为拨款不应该那么大。我试着尽可能地优化它，但几天前我开始学习Julia，不能再进一步了。任何关于如何提高性能的提示，我们都将不胜感激

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

一旦这两个问题都得到解决，认识到问题中的线性代数既可以让代码更简洁，也可以让julia调用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。手动对乘法进行编码可以得到这样的结果

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

---

为了便于比较，我想添加一个优化的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计时（来自Oscar Smith的函数）

---

进步很好！过一会儿我会研究你的解决方案。我真的希望Julia能在这里超越Python。我把它降到了1.6msDon't do this:

length=size（-x0:dx:x0dx）[1]

。执行此操作：

N=length（-x0:dx:x0dx）

<代码>长度是一个内置函数，最好先使用它，然后像这样扭转它。因此，您再次使Python比Julia更快！我又一次困惑了，翻译我的代码是否值得。你能运行我的代码吗？我用的是一台4年的笔记本电脑，所以我的cpu肯定不是最强的。对于我来说，使用eqivilent设置（Fastmath=False，parallel=False）运行测试需要763毫秒，而我最近的版本需要227ms。我不久前安装了Julia，并做了一些尝试。当然，我必须再试一次。我的问题是找到一些非常好的例子来编写代码。事实上，我确实遇到了与提问者相同的问题。我在击败numba时遇到困难的部分原因是，目前存在一个julia性能错误，其中real*complex数组比它应该的慢，因为没有向量{real}*矩阵{complex}的方法，所以它提升了向量。

%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

@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