Hi everyone! While Julia-evangelizing (as usual) in my research group, in order to encourage them to move to Julia I offered to convert some of their Python codes for solving a network of Kuramoto oscillators with delay (with a different delays for each connection). To my surprise, the Python version (which uses JiTCDDE) is much more performant (and scales much better with the system size) than what I could achieve in Julia so far. I must be committing a serious mistake here but I’m not seeing it. I would be really grateful if anyone could help me with this.
To have a sense of the performances:
Julia:
(median execution times measured with BenchmarkTools.jl)
Python:
(this is discarding the compilation time)
At first I tried to implement it with NetworkDynamics.jl which seemed appropriate for this case but I didn’t find the way to set different lags for each edge (does anyone know if this is possible?) so I moved to plain DifferentialEquations.jl :
Julia code
using DelayDiffEq
using Distributions
using Random
function f!(dθ, θ, h, p, t)
ω, A = p
n = length(θ)
lags = reshape(lag, n,n)
@inbounds for j in 1:n
coupling = 0.0
@inbounds for i in 1:n
coupling += A[i,j]*sin(h(p, t-lags[i,j]; idxs=i) - θ[j])
end
dθ[j] = ω[j] + coupling
end
nothing
end
n = 10
Random.seed!(1)
ω = rand(n)
A = rand(n,n)
const lag = rand(n*n)
θ₀ = rand(Uniform(0, 2π), n)
p = (ω, A)
h(p, t; idxs=nothing) = typeof(idxs) <: Number ? 1.0 : rand(Uniform(0, 2π), n)
prob = DDEProblem(f!, θ₀, h, (0.0, 1.0), p, constant_lags=lag)
solve(prob, MethodOfSteps(BS3()), saveat=0.01, rtol=0, atol=1e-5, force_dtmin=true)
(I made the diffeq function inplace because the goal is to simulate 180 oscillators)
Status `~/Documents/julia_kuramoto/Project.toml`
[6e4b80f9] BenchmarkTools v1.2.0
[bcd4f6db] DelayDiffEq v5.31.1
[31c24e10] Distributions v0.25.19
[9a3f8284] Random
Python code:
import os
from tqdm import tqdm
import numpy as np
from numpy import pi, random, max
os.environ["CC"] = "clang"
# https://jitcdde.readthedocs.io/en/stable/
from jitcdde import jitcdde, y, t
from symengine import sin
import symengine
import sympy
n = 10
ω = random.rand(n)
A = random.rand(n, n)
τ = random.rand(n, n)
def kuramotos():
for i in range(n):
yield ω[i] + sum(
A[j, i]*sin(y(j, t-τ[i, j])-y(i))
for j in range(n)
)
DDE = jitcdde(kuramotos, n=n, verbose=True)
DDE.compile_C(simplify=False, do_cse=False, chunk_size=150)
DDE.set_integration_parameters(rtol=0, atol=1e-5)
DDE.constant_past(random.uniform(0, 2*np.pi, n), time=0.0)
DDE.integrate_blindly(max(τ), 0.1)
output = []
for time in tqdm(DDE.t + arange(0, 1, 0.01)):
output.append([*DDE.integrate(time)])
. And 