Lorentz96をJuliaでとく その３

前回まで

グローバル変数は使わないようにしたコードは作りましたが，今回はできるだけ早くします。

コード

高速化の方針

• メモリアロケーションができるだけ起きないようにする。
  • そのために，一時変数も最初に作ったものを使い回す。
• パラメータはParametersパッケージでやりとりする。

using Parameters

@with_kw struct Param{T1,T2}
    dt::T1 = 0.01
    tend::T1 = 146.0
    nstep::T2 = Int(tend/dt)
    
    F::T2 = 8
    N::T2 = 40
end

function Lorentz96(para, x, f)
    @unpack N, F = para
    #f = zeros(N)
    
    for i=3:N-1
        f[i] = (x[i+1]-x[i-2])x[i-1] - x[i] + F
    end
    
    # 周期境界
    f[1] = (x[2]-x[N-1])x[N] - x[1] + F
    f[2] = (x[3]-x[N])x[1] - x[2] + F
    f[N] = (x[1]-x[N-2])x[N-1] - x[N] + F

end

function RK4(para, dt, x, k1, k2, k3, k4, f)
    
    #k1, k2, k3, k4 = zeros(N), zeros(N), zeros(N), zeros(N)

    Lorentz96(para, x, f)
    k1 .= f
    Lorentz96(para, x .+ k1 .* dt .* 0.5, f)
    k2 .= f
    Lorentz96(para, x .+ k2 .* dt .* 0.5, f)
    k3 .= f
    Lorentz96(para, x .+ k3 .* dt, f)
    k4 .= f
    return x .+ dt .*(k1 .+ 2 .* k2 .+ 2 .* k3 .+ k4)./6.
    
end

function main(para)
    @unpack dt, nstep, N, F = para

    X = zeros(N, nstep)
    t = zeros(nstep)
    k1, k2, k3, k4 = zeros(N), zeros(N), zeros(N), zeros(N)
    f = zeros(N)
    
    # initialize
    for i=1:N
        X[i, 1] = F
    end
    X[20, 1] = 1.001*F

    for i=2:nstep
        X[:,i] .=  RK4(para, dt, view(X,:,i-1), k1, k2, k3, k4, f)
        t[i] = dt*(i-1)
    end

    return X, t
end

実行

para = Param()

@time   X, t = main(para)

ベンチマーク結果

using BenchmarkTools
para = Param()
@benchmark main(para)

平均7.7msで，適当に書いたFortranより早くなりました。

BenchmarkTools.Trial: 645 samples with 1 evaluation.
 Range (min … max):  7.013 ms …   9.304 ms  ┊ GC (min … max): 0.00% … 17.44%
 Time  (median):     7.384 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   7.755 ms ± 629.416 μs  ┊ GC (mean ± σ):  4.99% ±  5.20%

    ▁▄▆█▂▄                                                     
  ▃▇██████▆█▆▄▄▃▂▂▂▃▁▁▁▂▃▃▄▄▆▄▃▂▁▁▂▃▃▄▄▆█▇▇▆▆▄▄▅▄▃▃▄▂▃▄▃▄▄▃▂▂ ▄
  7.01 ms         Histogram: frequency by time        9.03 ms <

 Memory estimate: 26.85 MiB, allocs estimate: 58406.