Python vs. Julia ODE Solver

General Usage Performance
14

There are some other things that I have fiddled around:

  1. rand(Float64, (N, 1))
    These are more clearly written as rand(Float64,N) (the Float64 is redundant there, but ok).
    The thing is that that was generating an two-dimensional array instead of a vector, and that didn’t allow the broadcasting of the result of the first of the four lines there.

  2. The remaining allocations come from, as far as I could see:
    (a) -s which can be written as -1 .* s
    (b) Two matrix-vector multiplications: p[21]*s and p[20]*sin.(p[23][:, 1].*t)

To solve those maybe someone has better advice, but I think you have to use an inplace matrix multiplication routine as mul! from LinearAlgebra. But for that you need to store auxiliary arrays to store the results.

I have done that in the code bellow (added some auxiliary arrays to the parameters). That code does not allocate anything in that function. It is possible that I have done something wrong, but if I didn’t the code runs in:

julia> @time solution = solve(prob, lsoda(), saveat=0.1, reltol=1e-8, abstol=1e-8);
  0.135555 seconds (166.55 k allocations: 16.639 MiB)

for N=20

and for N=500 it runs in:

 71.583810 seconds (945.29 k allocations: 149.516 MiB, 0.11% gc time)

Is that good?

In the code below I also added some views, but I think those didn’t make much of a difference. To check.

using OrdinaryDiffEq, BenchmarkTools, LSODA, Plots, Profile
using LinearAlgebra

N = 20

parameters = (2, N,
              -500, 0.2, 0.0, 100.0, 0.7, 0.0, 15.0, 0.24, 0.02, 1700.0, 8.0, 0.19, -0.5,
              25.0, 6.0, 12.0, 10.0,
              2.0*1.4 .* (rand(Float64,N,2) .- 0.5), #20 
              0.5 .* (rand(Float64,N,N) .- 0.8), #21
              0.14 .* (rand(Float64,N) .- 0.5), #22
              rand(Float64,2,4),  #23
              zeros(N), #24
              zeros(N), #25 
              zeros(2)) #26

function forward_RHS(du, u, p, t)
    v_s = view(u, 1:p[2])
    u_s = view(u, p[2]+1:2*p[2])
    w_s = view(u, 2*p[2]+1:3*p[2])
    s = view(u, 3*p[2]+1:4*p[2])

    p[26] .= t .* view(p[23],:,1)
    p[26] .= sin.(p[26])
    mul!(p[24],p[20],p[26])

    du[1:p[2]] .= ((p[3].*(v_s .- p[4]).^2.0 .+ p[5]).*(v_s .- 1.0)
                  .- p[12].*u_s.*(v_s .- p[14]) .- p[13].*w_s.*(v_s .- p[15])
                  .+ p[22]
                  .+ p[24]
                  .+ mul!(p[25],p[21],s))./p[16]
    du[p[2]+1:2*p[2]] .= (exp.(p[6].*(v_s .- p[7])) .+ p[8] .- u_s)./p[17]
    du[2*p[2]+1:3*p[2]] .= (p[9].*(v_s .- p[10]).^2.0 .+ p[11] .- w_s)./p[18]
    du[3*p[2]+1:4*p[2]] .= (-1 .* s .+ (v_s .> 0.25).*(v_s .< 0.3).*(view(du,1:p[2]) .> 0.0).*view(du,1:p[2])./0.05)./p[19]

end

prob = ODEProblem(forward_RHS, zeros(4*N), (0.0, 500.0), parameters)
@time solution = solve(prob, lsoda(), saveat=0.1, reltol=1e-8, abstol=1e-8)

One note: In Julia you do not need to write everything in vectorized form. Maybe you are comfortable with that, but I would have split those operations in a lot of lines to makes things more clear and debug errors easier.

ps: There is an auxiliary vector of dimension 2 which could be substituted by an StaticArray, but that I think is for another moment :-).

7 Likes