How to optimize the julia differential equation solution

Question

I am running the following Julia code to solve a set of differential equations, but it takes too much time, and sometimes my machine gets crushed. Is there any solution to reduce the running time and memory usage? I am not too familiar with Julia. So a help would be very much appreciated.

using DelimitedFiles
using LinearAlgebra
using Random
using PyPlot
using BenchmarkTools
using SparseArrays
using Distributions
using DifferentialEquations



N=10000;
Random.seed!(123)
d=Cauchy()
omega=rand(d,N);

function kuramoto(du, u, pp, t)

    u1 = @view u[1:N] #### θ
    du1 = @view du[1:N]  #### dθ
    u2 = @view u[N+1:2*N]  ####### λ
    du2 = @view du[N+1:2*N]  ####### dλ

    a=1
    b=3
    α=0.005
    β=0.002
    λ0=pp

    z1 = Array{Complex{Float64},1}(undef, N)
    z1=mean(cis,u1)
    z1c=conj(z1)

    z2 = Array{Complex{Float64},1}(undef, N)
    z2=mean(cis,2*u1)




    ####### equ of motion
    @. du1  = omega + u2 *(a*( imag(z1  * exp((-1im) * (u1)) )) + b*(imag(z2 * z1c * exp((-1im) * (u1)) )))
    @. du2 = α *(λ0-u2)- β * (abs(z1))

    return nothing
end;


# setting up time steps and integration intervals

dt = 0.01 # time step
dts = 0.1 # save time
ti = 0.0
tt = 1000.0
tf = 20000.0
nt = Int(div(tt,dts))
nf = Int(div(tf,dts))

tspan = (ti, tf); # time interval
t=range(0,stop=tf-tt,length=nf-nt+1);

pp=2.05
Random.seed!(123)
u0=[rand(N)*2*pi;pp*ones(N)];
du = similar(u0);


prob = ODEProblem(kuramoto, u0, tspan, pp)
sol = solve(prob, RK4(), reltol=1e-4, saveat=dts,progress=true);



function global_order(_u)
    _R_global = zeros(nf-nt+1)

    for n in nt:nf
        _re = mean(j -> cos(_u[j,n]), 1:N)
        _im = mean(j -> sin(_u[j,n]), 1:N)
        _R_global[n-nt+1] = sqrt(_re^2 + _im^2)
    end

    return (_R_global)
end;


r1= global_order(sol[1:N,:]);


ppp=[t r1];
writedlm("k2=3_cauchy_time_order_hoi_α=0.005_β=0.002_λ0=$pp.txt",ppp)


λ_avg=mean(sol[N+1:2*N,:],dims=1)[nt:nf];
qqq=[λ_avg r1];
writedlm("k2=3_cauchy_λ_avg,R_α=0.005_β=0.002_λ0=$pp.txt",qqq)


clf()
plot(t,r1,c="blue")
ylim([0,1])
xlabel("time")
ylabel("order parameter R")
title("λ=$pp")
gcf()

I have to solve it with the RK4 method only.

Answer 1

Just calculate it out. You're saving ((20000.0 - 1000.0) / 0.1) = 190000.0 steps of a 2 * 10000 = 20000 length array. This means the solution is ((20000.0 - 1000.0) / 0.1) * (2 * 10000) = 3.8e9 64-bit numbers, or 30.4 GB. Then you slice half of it here in global_order(sol[1:N,:]);. so that slice will make an extra 15 GB. If you don't have 45 GB of RAM, your computer will crash at this point. So yeah, not really surprising: don't tell it to generate 45 GB of vectors if it cannot be stored. That's also going to be horrible for performance.

The solution of course is just to not save so much. Use the saving callback so that you just save the time series for _R_global array. Saving a size 1 object every time step instead of 20000 is going to be a huge difference.

For your code, this would look like:

saved_values = SavedValues(Float64, Float64)
function saver(u,t,integrator)
    _re = mean(cos.(u))
    _im = mean(sin.(u))
    out = sqrt(_re^2 + _im^2)
end
cb = SavingCallback(saver, saved_values)

and then change the solve call to:

sol = solve(prob, RK4(), reltol=1e-4, saveat=dts,progress=true,callback = cb);

then after solving, saved_values.u will be the array of values you want.

P.S.

Also, the number of steps 20000.0 - 1000.0 = 190000 looks odd: almost surely that's a typo meant to be 1000 to 2000 not 20000?