改进矢量化的Julia代码或如何对它们进行矢量化处理？

Question

我正在尝试将一些python代码转换为julia。它们旨在解决一些微分方程。由于它们包含方程式，以提高可读性，因此我给出了如下图像。

我还没有儒略语的资格。这是上述方程式的茱莉亚代码。但是它们非常慢。可以加快代码执行哪些改进？我试图对其进行向量化处理，以为它会加快速度，但是我没有成功。另外，我应该注意D1和D2实际上更密集，为便于说明，我仅在代码中以对角线形式给出它们。您的帮助将不胜感激。

更新的代码：

using SparseArrays
using BenchmarkTools
using LinearAlgebra
using Profile


function Rhs(W,D1,D2)
    u = @views W[:,1]; v = @views W[:,2]
    ut = -0.5*(D1*u) .+ im*(0.5*D2*u .+ (abs2.(u) + 2/3*abs2.(v)).*u)
    vt = 0.5*(D1*v) .+ im*(0.5*D2*v .+ (abs2.(v) + 2/3*abs2.(u)).*v)

    return [ut vt]
end


function RK4(U, dt, Rhs,D1,D2)
    k1 = Rhs(U,D1,D2);
    k2 = Rhs(U.+dt/2*k1,D1,D2);
    k3 = Rhs(U.+dt/2*k2,D1,D2);
    k4 = Rhs(U.+dt*k3,D1,D2);

    return @. U + dt*(k1 + 2*k2 + 2*k3 + k4)/6;
end


Uexact(X, T) = 1.0954*sech.(sqrt(2)*(X.-T)).*exp.(im*(0.5*X .+ 0.625*T))
Vexact(X, T) = 1.0954*sech.(sqrt(2)*(X.-T)).*exp.(im*(1.5*X .+ 0.625*T))


function main(N,xl,xr,dt,t0,tf) 
    D2 = Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1])
    D1 = Tridiagonal([0.0 for i in 1:N-1],[-1.0 for i in 1:N],[1.0 for i in 1:N-1])


    D1 = sparse(D1);  
    D2 = sparse(D2); 
    x = range(xl,xr,length=N);
    bp = [1,N]; 
    NN = Int(round((tf-t0)/dt))

    # Initial cond.
    u = Uexact(x,0.0); 
    v = Vexact(x,0.0); 
    W = [u v];
    WW = zeros(Complex,NN,2*N);

    # ODE integration    
    for i in 1:NN
        Wn = RK4(W,dt,Rhs,D1,D2)         

        u = @views Wn[:,1]; v = @views Wn[:,2]
        u[bp] .= 0.0
        v[bp] .= 0.0

        W = [u v];   
        WW[i,:] = W[:]; 
    end

end
@benchmark main(300,-20.0,60.0,0.0001,0.0,1.0)

Answer 1

您的问题在RHS中。取一个矢量，将其分成两部分，对这两个部分应用相同的运算，然后将它们组合在一起。这种方法会多次复制数据，并且没有真正的原因就涉及更改数据形状。如果改为将W设为2xn矩阵，则代码将更快，并且RHS将是更简单的操作，不需要分配。