Insane amount of space allocated solving system of stochastic differential equations

Question

first time asking a question here. I previously used a simple MATLAB script to model 90 Hopf oscillators, coupled through a matrix, with randn noise, with a simple Euler step integration. I wanted to upgrade this, so I got into Julia, seems to have many exciting properties.

This is the system of equations I'm solving

I'm kinda lost. I started using differentialequations.jl (stochastic solver) , arrived to a solution, and found myself with a benchmark that tells me that solving 200 seconds occupies like 4 Gb!!! (2.5 Gb with alg_hints=[:stiff]) (I haven't fixed dt, previously I used dt=0.1)

function Shopf(du,u,p,t)

  du[1:90,1]=(p[1:90,1]-u[1:90,1].^2.0-u[1:90,2].^2.0).*u[1:90,1]-p[1:90,2].*u[1:90,2] + 0.5*(-p[: , end].*u[:,1]+p[:,4:end-1] *u[:,1])
  du[1:90,2]=(p[1:90,1]-u[1:90,1].^2.0-u[1:90,2].^2.0).*u[1:90,1]+p[1:90,2].*u[1:90,1] + 0.5*(-p[: , end].*u[:,2]+p[:,4:end-1] *u[:,2])

end


function σ_Shopf(du,u,p,t)

du[1:90,1]=0.04*ones(90,1)
du[1:90,2]=0.04*ones(90,1)


end


#initial condition
u0=-0.1*ones(90,2);
#initial time
t0=0.0;
#final time
tend=200.0;
#setting parameter matrix
p0=[0.1 , 2*pi*0.04]
push!(p0,-p0[2])
p=p0'.*ones(90,3);
SC=SC;
p=[p SC]
p=[p sum(SC,dims=2)]
#


#col 1 :alpha
#col 2-3 : [w0 -w0]

#col 3-93 : coupling matrix
#col 94: col-wise sum of coupling matrix


@benchmark solve(prob_sde_Shopf,nlsolver=Rosenbrock23(),alg_hints=[:stiff])

BenchmarkTools.Trial: memory estimate: 2.30 GiB

allocs estimate: 722769

---

minimum time: 859.224 ms (13.24% GC)

median time: 942.707 ms (13.10% GC)

mean time: 975.430 ms (12.99% GC)

maximum time: 1.223 s (13.00% GC)

---

samples: 6

evals/sample: 1

Any thoughts? I'm checking out several solutions, but none of them reduce the amount of memory to a reasonable amount. Thanks in advance.

Answer 1

You are creating a staggering number of temporary arrays. Every slice creates a temporary. You put in a dot here and there, but you have to dot everything to get fused broadcasting. Instead, you can just use the @. macro which will do it for you. Also, using @views will make sure that slices don't copy:

function Shopf(du, u, p, t)
    @. du[1:90, 1] = @views (p[1:90, 1] - u[1:90, 1]^2 - u[1:90, 2]^2) * u[1:90, 1] -
        p[1:90, 2] * u[1:90,2] + 0.5 * (-p[:, end] * u[:, 1] + p[:, 4:end-1] * u[:,1])
    @. du[1:90, 2] = @views (p[1:90, 1] - u[1:90, 1]^2 - u[1:90, 2]^2) * u[1:90, 1] + 
        p[1:90, 2] * u[1:90,1] + 0.5 * (-p[:, end] * u[:, 2] + p[:, 4:end-1] * u[:,2])  
end

Also, dont write x^2.0, use x^2, the former is a slow float power, while the latter is a fast x * x. In fact, try to use integers wherever you can, in multiplications, additions, etc.

Here's another thing

function σ_Shopf(du,u,p,t)

du[1:90,1]=0.04*ones(90,1)
du[1:90,2]=0.04*ones(90,1)


end

No need to create two temporary arrays on the right side of the assignment. Just write this:

function σ_Shopf(du, u, p, t)
    du[1:90, 1:2] .= 0.04
end

Faster and simpler. Note, that I haven't tested this, so please fix any typos.

(Finally, please use indentation and put spaces around operators, it makes your code much nicer to read.)

Update: I don't really know what your code is supposed to do, what with the strange indices, but here is a possible improvement that just uses loops (which I think is actually cleaner, and will let you make further optimizations):

The operation producing A is a matrix product, so you cannot avoid allocations there, unless you can pass in a cache array to work on, using mul!. Aside from that, you should have no allocations below.

function shopf!(du, u, p, t)
    A = @view p[:, 4:end-1] * u
    # mul!(A, view(p, 4:end-1), u)  # in-place matrix product
    for i in axes(u, 1)
        val = (p[i, 1] - u[i, 1]^2 - u[i, 2]^2) * u[i, 1]  # don't calculate this twice
        du[i, 1] = val - (p[i, 2] * u[i, 2]) - (0.5 * p[i, end] * u[i, 1]) + 
            (0.5 * A[i, 1])
        du[i, 2] = val + (p[i, 2] * u[i, 1]) - (0.5 * p[i, end] * u[i, 2]) + 
            (0.5 * A[i, 2])
    end
end

After this, you can add various optimizations, @inbounds if you are sure about the array sizes, multithreading, @simd or even @avx from the LoopVectorization experimental package.