Explicitly declaring the ODE Jacobian being less efficiently

General Usage Performance
1

One of the problems that I am solving, without explicitly declaring the jacobian, was being solved in, roughly, 800 seconds. It isn’t that fast, the system is really stiff, but I am trying to optimize it as far as I can. Then I decided the put the jacobian by hand following this guide Solving Stiff Equations · DifferentialEquations.jl (sciml.ai), however, the time simulation has increase to, roughly, 1000 seconds.

I verified with ModelingToolkit.jl if I was declaring everything right, apparently yes, but unfortunately as my jacobian has explicitly time dependency, I had some difficulties using ModelingToolkit.jl (does it only accepts non explicitly time dependent functions?)

Worth mention that both simulations (with jacobian declaration and without) at least give the same numerical solution.

I would appreciate any help towards my code optimization. I will post here the part of the code that I think is the most important, without declaring the parameters/functions.

#--------------------
function oscillator!(du,u,p,τ)
A, ϕa, B, ϕb = u
du[1] = -(1/2)(κ/ξ)A + sqrt(κₑ/ξ)(s_func(τ/ξ)/sqrt(ξ))cos(ϕa)
du[2] = (Δ_func(τ/ξ)/ξ) - (2g₀/ξ)Bcos(ϕb) - sqrt(κₑ/ξ)(s_func(τ/ξ)/sqrt(ξ))(sin(ϕa)/A)
du[3] = -(1/2)(Γₘ/ξ)B - (g₀/ξ)(A^2)(sin(ϕb))
du[4] = -(Ωₘ/ξ) - (g₀/ξ)(A^2)*(cos(ϕb)/B)
end

function oscillator_jac(J,u,p,τ)
A, ϕa, B, ϕb = u
J[1,1] = -(1/2)(κ/ξ)
J[2,1] = sqrt(κₑ/ξ)(s_func(τ/ξ)/sqrt(ξ))(sin(ϕa)/(AA))
J[3,1] = -(2g₀/ξ)A(sin(ϕb))
J[4,1] = -(2g₀/ξ)A(cos(ϕb)/B)
J[1,2] = -sqrt(κₑ/ξ)(s_func(τ/ξ)/sqrt(ξ))sin(ϕa)
J[2,2] = -sqrt(κₑ/ξ)(s_func(τ/ξ)/sqrt(ξ))(cos(ϕa)/A)
J[3,2] = 0
J[4,2] = 0
J[1,3] = 0
J[2,3] = -(2g₀/ξ)cos(ϕb)
J[3,3] = -(1/2)(Γₘ/ξ)
J[4,3] = (g₀/ξ)(A^2)(cos(ϕb)/(BB))
J[1,4] = 0
J[2,4] = (2g₀/ξ)Bsin(ϕb)
J[3,4] = -(g₀/ξ)(A^2)(cos(ϕb))
J[4,4] = (g₀/ξ)(A^2)*(sin(ϕb)/B)
nothing
end

initial conditions

A₀ = 0.1
ϕa₀ = 2πrand(1)[1]
B₀ = 220
ϕb₀ = 2πrand(1)[1]
u₀ = [A₀,ϕa₀,B₀,ϕb₀]

time span

τ_span = (0, ξ*t_final)

the system

ODEfunc = ODEFunction(oscillator!, jac=oscillator_jac)
ODEprob = ODEProblem(ODEfunc, u₀, τ_span)

solving

@time sol = solve(ODEprob, CVODE_BDF(), abstol=1e-6, reltol=1e-6, maxiters=Int(100e6))
#--------------------

2

If you fence your code with backticks ‘```’ it will format nicely in your message.

There are some constants that you have not defined in this example (ξ, g₀, κₑ and several others). How are these defined in your real code? The code should be modified to pass these in as part of p, or they should be explicitly declared const if they are in the global scope (see here). A runable example would be helpful.

ModelingToolkit should be able to create time dependent jacobians from a system of equations, but modelingtoolkitize can have trouble automatically deriving those equations from your code. Can you show what you tried and how it failed?

Do you have some reason to use the CVODE_BDF solver? I have found that DifferentialEquations built-in heuristics often do a good job of automatically selecting a solver, in your case I would try

sol = solve(ODEprob; alg_hints=[:stiff], abstol=1e-6, reltol=1e-6, maxiters=Int(100e6))
4 Likes
3

If you fence your code with backticks ‘```’ it will format nicely in your message.

Good to know about the backticks haha

Do you have some reason to use the CVODE_BDF solver?

Actually yes… When I tried using other stiff solvers like QNDF, KenCarp etc, the solution didn’t converges, even for smaller tolerances. However, I would like to use others solvers because CVODE_BDF can’t handle complex-valued ODE, that’s why I always open complex functions in polar coordinate, going from a N-dim complex space to a 2N-dim real space where CVODE_BDF do its job.

The code should be modified to pass these in as part of p

I need to understand better how this works…

There are some constants that you have not defined in this example;
Can you show what you tried and how it failed?

The code is a bit too long, that’s why I haven’t copied it here haha. Nevertheless, here is a runnable example

const ħ = 1.054571817e-34 # [J.s]
const c = 299792458 # [m/s]
const kb = 1.380649e-23 # [J/K]

# Inputs
λ = 1560.0e-9 # bare optical cavity ressonance wavelength [m]
P₀ = 425e-6 # optical power that reaches the cavity (not necessarily, and probably not, the same of the pump laser) [W]
ω₀ = 2*π*c/λ # optical angular frequency [rad/s]
Q₀ = 165000 # optical quality factor
κ = ω₀/Q₀ # optical loss [rad/s]
g₀ = -2*π*16200 # vacuum single-photon optomechanical coupling [rad/s]
η = 0.75 # κe/κ 
κₑ = η*κ # external coupling [rad/s]
κ₀ = κ - κₑ # intrinsic optical loss [rad/s]
Ωₘ = 2*π*31.86e6 # natural angular mechanical oscillator frequency [rad/s]
Qₘ = 1255 # mechanical quality factor
Γₘ = Ωₘ/Qₘ # mechanical loss [rad/s]
T = 300 # temperature [K]
nth_a = (kb*T)/(ħ*ω₀) # thermal photons number
nth_b = (kb*T)/(ħ*Ωₘ) # thermal phonons number

# Rescaling parameter
ξ = Ωₘ

# Optical Dynamic Parameters
Δi_x = 8.0*κ # initial optical detuning with bistability correction
Δf_x = 0.35*κ # final optical detuning with bistability correction
#-----------------------------------------------------------------------
# Static Corrections (i.e., Optomechanical Bistability)
function Δi_correction!(F, Δi)
    F[1] = Δi[1] - Δi_x + (2*g₀*g₀*κₑ/Ωₘ)*(P₀/(ħ*(Δi[1]+ω₀)))/(κ^2/4 + Δi_x^2)
end
function Δf_correction!(F, Δf)
    F[1] = Δf[1] - Δf_x + (2*g₀*g₀*κₑ/Ωₘ)*(P₀/(ħ*(Δf[1]+ω₀)))/(κ^2/4 + Δf_x^2)
end

Δi_guess = Δi_x - (2*g₀*g₀*κₑ/Ωₘ)*(P₀/(ħ*(Δi_x+ω₀)))/(κ^2/4 + Δi_x^2)
Δf_guess = Δf_x - (2*g₀*g₀*κₑ/Ωₘ)*(P₀/(ħ*(Δf_x+ω₀)))/(κ^2/4 + Δf_x^2)

Δi = (nlsolve(Δi_correction!, [Δi_guess], show_trace=false, ftol=1E-10, xtol=1E-12)).zero[1]
Δf = (nlsolve(Δf_correction!, [Δf_guess], show_trace=false, ftol=1E-10, xtol=1E-12)).zero[1]
#-----------------------------------------------------------------------
vΔ = 0.01*κ^2/(2*π) # optical detuning sweep velocity
vΔ = ifelse(Δf-Δi>0, abs(vΔ), -abs(vΔ)) # this corrects the signal of the sweep velocity, if necessary

# RF Parameters
ε = 0.03 # modulation depth
ρ = 2 # winding number, i.e., the harmonic that you are are pumping (for example, ρ = 2 means Ωd around 2Ω₀)
Ω₀_guess = (2*π)*3.2054910037748326e7 # what is, approximatly, the self-sustained frequency of the oscillator?
Ωⁱd = ρ*Ω₀_guess - 4.8*Γₘ # initial modulation frequency
Ωᶠd = ρ*Ω₀_guess + 9.1*Γₘ # final modulation frequency
vΩd = 0.275*Γₘ^2/(2*π) # drive frequency sweep velocity
vΩd = ifelse(Ωᶠd-Ωⁱd>0, abs(vΩd), -abs(vΩd)) # this corrects the signal of the sweep velocity, if necessary

# Time/Delay/Waiting Setup
t_vΔ = (Δf-Δi)/vΔ # time until the optical sweep reaches Δf
t_Δf = 5/(Γₘ/(2*π)) # waiting time after t_vΔ to turn on the modulation depth ε
t_ε = 5/(Γₘ/(2*π)) # waiting time after t_Δf to turn on vΩd
t_vΩd = (Ωᶠd - Ωⁱd)/vΩd # waiting time after t_ε to reach Ωᶠd
t_Ωᶠd = 1/(Γₘ/(2*π)) # waiting time after t_vΩd, here the experiment has already finished
t_final = t_vΔ + t_Δf + t_ε + t_vΩd + t_Ωᶠd

# Optical Detuning function
Δ_func(t) = ifelse(t<=t_vΔ, Δi + vΔ*t, Δf)

# Modulation Depth function
ε_func(t) = ifelse(t<=t_vΔ+t_Δf, 0, ε)

# RF drive function
Θ₁ = 0.5*vΩd*(t_vΔ+t_Δf+t_ε)^2 - vΩd*(t_vΔ+t_Δf+t_ε)*(t_vΔ+t_Δf+t_ε)
Θ₂ = Ωⁱd*(t_vΔ+t_Δf+t_ε+t_vΩd) + 0.5*vΩd*(t_vΔ+t_Δf+t_ε+t_vΩd)^2 - vΩd*(t_vΔ+t_Δf+t_ε)*(t_vΔ+t_Δf+t_ε+t_vΩd) - Θ₁
Θd_func(t) = @. ifelse(t<=t_vΔ+t_Δf+t_ε+t_vΩd, ifelse(t<=t_vΔ+t_Δf+t_ε, Ωⁱd*t, Ωⁱd*t + 0.5*vΩd*t^2 - vΩd*(t_vΔ+t_Δf+t_ε)*t - Θ₁), Ωᶠd*(t - (t_vΔ+t_Δf+t_ε+t_vΩd)) + Θ₂)
Ωd_func(t) = @. ForwardDiff.derivative(Θd_func, t)
fd_func(t) = Ωd_func.(t)/(2*π)

# Optical Pump functions
s₀_func(t) = sqrt(P₀/(ħ*(Δ_func(t)+ω₀)))
s_func(t) = s₀_func(t)*sqrt(1+ε_func(t)*sin(Θd_func(t)));

# Amplitude and Phase Model
#--------------------
function amplitude_phase_optomechanics!(du,u,p,τ)
    A,ϕa,B,ϕb = u
    du[1] = -(1/2)*(κ/ξ)*A + sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*cos(ϕa)
    du[2] = (Δ_func(τ/ξ)/ξ) - (2*g₀/ξ)*B*cos(ϕb) - sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(sin(ϕa)/A)
    du[3] = -(1/2)*(Γₘ/ξ)*B - (g₀/ξ)*(A^2)*(sin(ϕb))
    du[4] = -(Ωₘ/ξ) - (g₀/ξ)*(A^2)*(cos(ϕb)/B)
end

function amplitude_phase_optomechanics_jac(J,u,p,τ)
    A,ϕa,B,ϕb = u 
    J[1,1] = -(1/2)*(κ/ξ)
    J[2,1] = sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(sin(ϕa)/(A*A))
    J[3,1] = -(2*g₀/ξ)*A*(sin(ϕb))
    J[4,1] = -(2*g₀/ξ)*A*(cos(ϕb)/B) 
    J[1,2] = -sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*sin(ϕa)
    J[2,2] = -sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(cos(ϕa)/A)
    J[3,2] = 0
    J[4,2] = 0
    J[1,3] = 0
    J[2,3] = -(2*g₀/ξ)*cos(ϕb)
    J[3,3] = -(1/2)*(Γₘ/ξ)
    J[4,3] = (g₀/ξ)*(A^2)*(cos(ϕb)/(B*B))
    J[1,4] = 0
    J[2,4] = (2*g₀/ξ)*B*sin(ϕb)
    J[3,4] = -(g₀/ξ)*(A^2)*(cos(ϕb))
    J[4,4] = (g₀/ξ)*(A^2)*(sin(ϕb)/B)
    nothing
end

A₀ = sqrt(nth_a)
ϕa₀ = 2*π*rand(1)[1]
B₀ = sqrt(nth_b)
ϕb₀ = 2*π*rand(1)[1]
#--------------------

τ_span = (0, ξ*t_final)
u₀ = [A₀,ϕa₀,B₀,ϕb₀]
ODEfunc = ODEFunction(amplitude_phase_optomechanics!)
ODEprob = ODEProblem(ODEfunc, u₀, τ_span)

# Solver method
sol_abstol = 1e-6
sol_reltol = 1e-6
@time sol = solve(ODEprob, CVODE_BDF(), abstol=sol_abstol, reltol=sol_reltol, maxiters=Int(100e6))

I also would appreciate any hint/good practice that I am not doing in my code so far

4

I tried this

sys = modelingtoolkitize(ODEprob)
jac = eval(ModelingToolkit.generate_jacobian(sys)[2])

and the error was “TypeError: non-boolean (Num) used in boolean context”

5

You’re using a lot of non-constant global variables and reference them all throughout your code. This is not good for performance. You’re supposed to use the p argument of the differential equation function to pass in your parameter values, so you don’t have to use globals.

Please consult the documentation for DifferentialEquations as well as the manual section on Performance Tips for more information.

Not related to performance and mostly a style thing, but you can just use rand() here instead. There’s no need to allocate a single element array, just to index into it directly.

3 Likes
6

This error is probably from ifelse. I would add

using IfElse

And then replace all occurences of ifelse with IfElse.ifelse.

ifelse is in Core and can’t be overloaded, so the IfElse package defines a new ifelse that can be overloaded. ModelingToolkit then has overloads for this new ifelse to help with extracting the equations.

1 Like
7

For making the parameters nicer to use, consider Parameters.jl. You can construct a struct with all your ‘inputs’, pass that in as p, and then use @unpack to get the values you need in each function. This will mean you need to modify the signatures of some of your intermediate functions, but it should be worth it for the performance improvement.

8

Before I check Parameters.jl, following the tutorial Ordinary Differential Equations · DifferentialEquations.jl (sciml.ai) I think this should fix the global variable issue:

κ  = 1
κₑ = 2
g₀ = 3
Ωₘ = 4
Γₘ = 5
ξ  = 6

function amplitude_phase_optomechanics!(du,u,p,τ)
    A,ϕa,B,ϕb = u
    κ,ξ,κₑ,g₀,Γₘ,Ωₘ = p
    du[1] = -(1/2)*(κ/ξ)*A + sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*cos(ϕa)
    du[2] = (Δ_func(τ/ξ)/ξ) - (2*g₀/ξ)*B*cos(ϕb) - sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(sin(ϕa)/A)
    du[3] = -(1/2)*(Γₘ/ξ)*B - (g₀/ξ)*(A^2)*(sin(ϕb))
    du[4] = -(Ωₘ/ξ) - (g₀/ξ)*(A^2)*(cos(ϕb)/B)
end
ODEfunc = ODEFunction(amplitude_phase_optomechanics!)

A₀ = 0.1
ϕa₀ = 2*π*rand()
B₀ = 200
ϕb₀ = 2*π*rand()
u₀ = [A₀,ϕa₀,B₀,ϕb₀]
τ_span = (0, ξ*10)
p=[κ,ξ,κₑ,g₀,Γₘ,Ωₘ]
ODEprob = ODEProblem(ODEfunc, u₀, τ_span, p)

Or am I missing something? We can notice that I have also declared explicitly time dependent functions inside the ODE (s_func(τ/ξ) and Δ_func(τ/ξ)), and I guess that should be good to put them also inside the parameter ‘p’, but how? I tried to follow the example from here Ordinary Differential Equations · DifferentialEquations.jl (sciml.ai) but something is not good. I tried this:

function amplitude_phase_optomechanics!(du,u,p,τ)
    A,ϕa,B,ϕb = u
    κ,ξ,κₑ,g₀,Γₘ,Ωₘ,s_p(τ),Δ_p(τ) = p
    du[1] = -(1/2)*(κ/ξ)*A + sqrt(κₑ/ξ)*(s_p(τ)/sqrt(ξ))*cos(ϕa)
    du[2] = (Δ_p(τ)/ξ) - (2*g₀/ξ)*B*cos(ϕb) - sqrt(κₑ/ξ)*(s_p(τ)/sqrt(ξ))*(sin(ϕa)/A)
    du[3] = -(1/2)*(Γₘ/ξ)*B - (g₀/ξ)*(A^2)*(sin(ϕb))
    du[4] = -(Ωₘ/ξ) - (g₀/ξ)*(A^2)*(cos(ϕb)/B)
end

A₀ = sqrt(nth_a)
ϕa₀ = 2*π*rand()
B₀ = sqrt(nth_b)
ϕb₀ = 2*π*rand()
amplitude_phase_u₀ = [A₀,ϕa₀,B₀,ϕb₀]
#--------------------
ODEfunc = ODEFunction(amplitude_phase_optomechanics!)

A₀ = 0.1
ϕa₀ = 2*π*rand()
B₀ = 200
ϕb₀ = 2*π*rand()
u₀ = [A₀,ϕa₀,B₀,ϕb₀]
τ_span = (0, ξ*10)
p=[κ,ξ,κₑ,g₀,Γₘ,Ωₘ,τ->s_func(τ/ξ),τ->Δ_func(τ/ξ)]
ODEprob = ODEProblem(ODEfunc, u₀, τ_span, p)

but this happened “MethodError: no method matching /(::var”#17#19", ::Float64)". The way I defined these function were

Δ_func(t) = ifelse(t<=t_vΔ, Δi + vΔ*t, Δf)

similar for the s_func(τ/ξ).

9

It should not be necessary to pass the functions in like that, and you are still treating ξ as a global variable since you include it in the definition of the anonymous function (maybe a typo?).

A first step would be to modify those functions to accept p as an argument as well and extract the values they needed.

Looking at your code more carefully, I think you should consider using events rather than step functions to provide the driving functions. Each of the functions (Δ_func, ε_func, Θd_func,...) should be replaced by a value in p (or a zero-derivative entry in u), and an event to modify it at the correct time.

1 Like
10

I have used continuous functions instead of events because it was just what I was used to manipulate, mainly because I use Wolfram Mathematica a lot.

I don’t know how to use events in Julia right now. Do you think the performance would increase significantly?

11

That’s a good question, I am not sure of the performance implications. It does seem to be more accurate though since events will root-find and solve right at the transition point, where the transition might be missed by the adaptive stepping.

I think this is the kind of thing Mathematica might do automatically, it has had a lot longer to add features like that.

1 Like
12

I’ve done most of the suggestions: declared const in most of the parameters at the global scope; start using Parameters.jl; and the ifelse → IfElse.ifelse also have done well with ModelingToolkit.

I am surprised that now my code runs easily x4 fasters, and explicitly declaring the Jacobian now indeed improves performance (by almost 20%).

One last thing that I want to try is StaticArrays.jl. I was following Optimizing DiffEq Code (sciml.ai) but I keep getting an error.

The part of the code that I declared StaticArray

function amplitude_phase_optomechanics(u,p,τ)
    A,ϕa,B,ϕb = u
    dA = -(1/2)*(κ/ξ)*A + sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*cos(ϕa)
    dϕa = (Δ_func(τ/ξ)/ξ) - (2*g₀/ξ)*B*cos(ϕb) - sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(sin(ϕa)/A)
    dB = -(1/2)*(Γₘ/ξ)*B - (g₀/ξ)*(A^2)*(sin(ϕb))
    dϕb = -(Ωₘ/ξ) - (g₀/ξ)*(A^2)*(cos(ϕb)/B)
    @SVector [dA,dϕa,dB,dϕb]
end
amplitude_phase_M = [1. 0  0  0
                     0  1. 0  0
                     0  0  1. 0
                     0  0  0  1.]

function amplitude_phase_optomechanics_jac(J,u,p,τ)
    A,ϕa,B,ϕb = u 
    J[1,1] = -(1/2)*(κ/ξ)
    J[2,1] = sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(sin(ϕa)/(A*A))
    J[3,1] = -(2*g₀/ξ)*A*(sin(ϕb))
    J[4,1] = -(2*g₀/ξ)*A*(cos(ϕb)/B) 
    J[1,2] = -sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*sin(ϕa)
    J[2,2] = -sqrt(κₑ/ξ)*(s_func(τ/ξ)/sqrt(ξ))*(cos(ϕa)/A)
    J[3,2] = 0
    J[4,2] = 0
    J[1,3] = 0
    J[2,3] = -(2*g₀/ξ)*cos(ϕb)
    J[3,3] = -(1/2)*(Γₘ/ξ)
    J[4,3] = (g₀/ξ)*(A^2)*(cos(ϕb)/(B*B))
    J[1,4] = 0
    J[2,4] = (2*g₀/ξ)*B*sin(ϕb)
    J[3,4] = -(g₀/ξ)*(A^2)*(cos(ϕb))
    J[4,4] = (g₀/ξ)*(A^2)*(sin(ϕb)/B)
end

A₀ = sqrt(nth_a)
ϕa₀ = 2*π*rand()
B₀ = sqrt(nth_b)
ϕb₀ = 2*π*rand()
#--------------------
ODEfunc = ODEFunction(amplitude_phase_optomechanics, jac=amplitude_phase_optomechanics_jac)

τ_span = (0, ξ*t_final)
u₀ = @SVector [A₀,ϕa₀,B₀,ϕb₀]
ODEprob = ODEProblem(ODEfunc, u₀, τ_span)

sol_abstol = 1e-6
sol_reltol = 1e-6
@time sol = solve(ODEprob, CVODE_BDF(), 
    abstol=sol_abstol, reltol=sol_reltol, maxiters=Int(100e6))

And the error is this

setindex!(::SVector{4, Float64}, value, ::Int) is not defined.
13

An SVector is an immutable Vector of static size, i.e. you can’t write to it. You probably want a regular Vector for u₀.

While declaring globals as const helps in some regard, it’s best to put all of your code into functions (including parameters! Passing them around via p to custom inner functions is most idiomatic and will lead to the best performance. It’s just not done in the examples because they’re quick to execute anyway and performance is not an issue there).

14

Is that error from somewhere in your code or from inside DifferentialEquations? I don’t usually use StaticVector, but it is documented as working in the tutorial. If it is in your code, check out Setfield.jl. If it is in DifferentialEquations, then it is either a bug or a limitation of CVODE_BDF, might be worth filing an issue if it looks like a bug.

15

I’ve just done what Sukera suggested You probably want a regular Vector for u₀. and everything work just fine. But indeed in the tutorial they let u₀ as a @SVector aswell.

16

I noticed you call s_func with the same parameters multiple times, you might consider making a temporary variable and storing the value of s_func for one timestep there, i.e. s_func_val = s_func(τ/ξ) and filling in s_func_val where previously you evaluted s_func(τ/ξ). That way you only require one function evaluation for each timestep, and depending on the evaluation time of s_func(τ/ξ) this can result in a speedup.

1 Like
17

CVODE_BDF is a solver in C, not Julia, so it cannot support static arrays. Use Rodas4 or Rodas5 are usually better if it’s static array sized anyways.

1 Like