python - 一段 Julia 和 Python 代码的优化建议

Question

我有一个模拟程序，在程序中我有一个函数。我已经意识到该功能消耗了大部分模拟时间。所以，我首先尝试优化功能。函数如下

Julia 1.1 版:

function fun_jul(M,ksi,xi,x)

  F(n,x) = sin(n*pi*(x+1)/2)*cos(n*pi*(x+1)/2);
  K = length(ksi);
  Z = zeros(length(x),K);
  for n in 1:M
    for k in 1:K
        for l in 1:length(x)
          Z[l,k] +=  (1-(n/(M+1))^2)^xi*F(n,ksi[k])*F(n,x[l]);
        end
    end
  end

return Z

end

我也用python+numba重写了上面的函数，对比如下

Python+数字

import numpy as np
from numba import prange, jit    
@jit(nopython=True, parallel=True)
def fun_py(M,ksi,xi,x):

    K = len(ksi);
    F = lambda nn,xx: np.sin(nn*np.pi*(xx+1)/2)*np.cos(nn*np.pi*(xx+1)/2);

    Z = np.zeros((len(x),K));
    for n in range(1,M+1):
        for k in prange(0,K):
            Z[:,k] += (1-(n/(M+1))**2)**xi*F(n,ksi[k])*F(n,x);


    return Z

但 Julia 代码非常慢，这是我的结果:

Julia 结果:

using BenchmarkTools
N=400; a=-0.5; b=0.5; x=range(a,b,length=N); cc=x; M = 2*N+100; xi = M/40;
@benchmark fun_jul(M,cc,xi,x)

BenchmarkTools.Trial: 
  memory estimate:  1.22 MiB
  allocs estimate:  2
  --------------
  minimum time:     25.039 s (0.00% GC)
  median time:      25.039 s (0.00% GC)
  mean time:        25.039 s (0.00% GC)
  maximum time:     25.039 s (0.00% GC)
  --------------
  samples:          1
  evals/sample:     1

Python 结果:

N=400;a = -0.5;b = 0.5;x = np.linspace(a,b,N);cc = x;M = 2*N + 100;xi = M/40;

%timeit fun_py(M,cc,xi,x);
1.2 s ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如果能帮助改进 julia 和 python+numba 的代码，我们将不胜感激。

已更新

基于@Przemyslaw Szufel 的回答和其他帖子，我改进了 numba 和 julia 代码。现在两者都是并行的。这是时间

Python+Numba 时间:

@jit(nopython=True, parallel=True)
def fun_py(M,ksi,xi,x):

    K = len(ksi);
    F = lambda nn,xx: np.sin(nn*np.pi*(xx+1)/2)*np.cos(nn*np.pi*(xx+1)/2);

Z = np.zeros((K,len(x)));
for n in range(1,M+1):
    pw = (1-(n/(M+1))**2)**xi; f=F(n,x)
    for k in prange(0,K):
        Z[k,:] = Z[k,:] + pw*F(n,ksi[k])*f;


    return Z


N=1000; a=-0.5; b=0.5; x=np.linspace(a,b,N); cc=x; M = 2*N+100; xi = M/40;

%timeit fun_py(M,cc,xi,x);
733 ms ± 13.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Julia 时代

N=1000; a=-0.5; b=0.5; x=range(a,b,length=N); cc=x; M = 2*N+100; xi = M/40;
@benchmark fun_jul2(M,cc,xi,x)

BenchmarkTools.Trial: 
  memory estimate:  40.31 MiB
  allocs estimate:  6302
  --------------
  minimum time:     705.470 ms (0.17% GC)
  median time:      726.403 ms (0.17% GC)
  mean time:        729.032 ms (1.68% GC)
  maximum time:     765.426 ms (5.27% GC)
  --------------
  samples:          7
  evals/sample:     1

Answer 1

我使用以下代码在单线程(而不是我的机器上的 28 秒)上减少到 300 毫秒。

您正在为 Numba 使用多线程。在 Julia 中，您应该使用并行处理(Julia 的多线程支持是实验性的)。看起来你的代码正在做某种参数扫描——这样的代码很容易并行化，但它通常需要对你的计算过程进行一些调整。

代码如下:

function fun_jul2(M,ksi,xi,x)
  F(n,x) = sin(n*pi*(x+1))/2;
  K = length(ksi);
  L = length(x);
  Z = zeros(length(x),K);
  for n in 1:M
    F_im1= [F(n,ksi[k]) for k in 1:K]
    F_im2 = [F(n,x[l]) for l in 1:L]
    pow = (1-(n/(M+1))^2)^xi    
    for k in 1:K
       for l in 1:L    
          Z[l,k] += pow*F_im1[k]*F_im2[l];
      end
    end
  end
  Z
end

julia> fun_jul2(M,cc,xi,x) ≈ fun_jul(M,cc,xi,x)
true

julia> @time fun_jul2(M,cc,xi,x);
  0.305269 seconds (1.81 k allocations: 6.934 MiB, 1.60% gc time)

** 编辑:使用 DNF 建议的多线程和入站:

function fun_jul3(M,ksi,xi,x)
  F(n,x) = sin(n*pi*(x+1))/2;
  K = length(ksi);
  L = length(x);
  Z = zeros(length(x),K);
  for n in 1:M
    F_im1= [F(n,ksi[k]) for k in 1:K]
    F_im2 = [F(n,x[l]) for l in 1:L]
    pow = (1-(n/(M+1))^2)^xi    
    Threads.@threads for k in 1:K
       for l in 1:L    
          @inbounds Z[l,k] += pow*F_im1[k]*F_im2[l];
      end
    end
  end
  Z
end

现在是运行时间(记得在启动 Julia 之前运行 set JULIA_NUM_THREADS=4 或等效的 Linux):

julia>  fun_jul2(M,cc,xi,x) ≈ fun_jul3(M,cc,xi,x)
true

julia> @time fun_jul3(M,cc,xi,x);
  0.051470 seconds (2.71 k allocations: 6.989 MiB)

您还可以尝试进一步试验 F_im1 和 F_im2 的并行计算。