Hello, I’ve been working on a project which relies on being able to swiftly solve something like the following system:
using DifferentialEquations
using StaticArrays
using Einsum
using Random
using Graphs
function wattsstrogatzmatrix(size, neighbors, rewiring_prob)
g = watts_strogatz(size, 2*neighbors, rewiring_prob)
coupling_matrix = adjacency_matrix(g)
coupling_matrix = Matrix(coupling_matrix)
coupling_matrix = coupling_matrix - Diagonal(vec(sum(coupling_matrix, dims=2))) # Ensure zero row sum
return -coupling_matrix
end
function fhn_eom(x, params)
a = params[1]
eps = params[2]
dx = (x[1] - (x[1]^3)/3 - x[2])/eps
dy = x[1] + a
return SVector{2}(dx, dy)
end
function bmatrix(phi, eps)
return -[cos(phi)/eps sin(phi)/eps; -sin(phi) cos(phi)]
#return SArray{Tuple{2,2}}(-cos(phi)/eps, sin(phi)/eps, -sin(phi)/eps, -cos(phi))
end
function coupled_fhn_eom!(dx, x, a, eps, coupling_strength, coupling_matrix, b)
N = length(coupling_matrix[1, :])
eachneuron = reshape(x, (2, N))
coupling_terms = b * eachneuron
@einsum coupling[i, j] := coupling_matrix[i, r] * coupling_terms[j, r]
eom_params = SVector{2}(a, eps)
for i in range(1, N)
thisneuron = @view eachneuron[:, i]
thiscoupling = @view coupling[i, :]
dx_i = fhn_eom(thisneuron, eom_params) .+ coupling_strength .* thiscoupling
# dx_i = fhn_eom(eachneuron[:, i], eom_params) .+ coupling_strength .* coupling[i, :]
dx[2*i-1:2*i] = dx_i
end
nothing
end
N = 175
eps = 0.05
a = 0.5
b = bmatrix(pi/2-0.1, eps)
σ = 0.0506
G = wattsstrogatzmatrix(N, 3, 1);#0.232)
x_0 = zeros(2*N)
x_0[2 .* (1:N) .- 1] = rand(N) .* 2 .* a .- a
x_0[2 .* (1:N)] = rand(N) .* 2 .* (-a + a^3 / 3) .- (-a + a^3 / 3)
prob = ODEProblem((dx, x, params, t) -> coupled_fhn_eom!(dx, x, params[1], params[2], params[3], G, b), x_0, (0.0, 1000.0), [a, eps, σ])
sol = solve(prob);
I’m fairly new to Julia, and tried optimizing it as much as I could with some tips I found online. I’ve had some success, the first versions were much slower than this. However I’ve a hard time knowing how I could optimize this further, with more Julia-specific optimizations, which I’m keen to learn.
In my pc, this is taking close to 7 seconds (measured using btime).
Any tips would be appreciated!
