Any possible way to stop ODE solver (with DifferentialEquations.jl)? - julia

I'm trying to solve an ODE problem (with Julia), that can stop early when a specific condition is satisfied.
Let's say I have a Lorenz system as below
using DifferentialEquations
function lorenz!(du,u)
du[1] = 10.0*(u[2]-u[1]);
du[2] = u[1]*(28.0-u[3]) - u[2];
du[3] = u[1]*u[2] - (8/3)*u[3];
end
u0 = [1.0;0.0;0.0]
tspan = (0.0, 100.0)
prob = ODEProblem(lorenz!,u0, tspan);
sol = solve(prob);
And, For example, I want to stop the ODE solver when u[3] is higher than 10, like below.
sol = solve(prob, stopcondition = u[3]>10);
But I'm not sure that there is a possible way to stop ODE solver with a given condition.
Any relevant comments would be thankful :)

Yes, use the terminate!(integrator) functionality within the event handling system. That would look like this here:
condition(u,t,integrator) = u[3] - 10 # Is zero when u[3] = 10
affect!(integrator) = terminate!(integrator)
cb = ContinuousCallback(condition,affect!)
sol = solve(prob, callback = cb);

Related

Improve plots quality in log scale in Julia

I have a function solving two ODEs and joining the solution which gives a really nice plot in linear scales, but which has a high drop in quality in log scale depending on the parameters I use. In the code below, I plot the solution for two sets of parameters, in which you can see the first set is not smooth, while the second one is kind of okay. If I try to obtain a smoother visualisation using the saveat option in the second ODEs, solclass = solve(probclass,Tsit5(),saveat=0.001), I get an error when plotting the second set: ArgumentError: range step cannot be zero. Is there a way to obtain smooth linear-log other than manually changing the saveat option? Also, I have tried using a few other backends, but they gave an error at ploting the solution.
using DifferentialEquations, Plots, RecursiveArrayTools
function alpha_of_phi!(s2,d,a0,ϕ₀)
# α in the quantum phase
function quantum!(dv,v,p,ϕ)
s2,d=p
α = v[1]
dv[1] = (ϕ*s2*sin(2*d*α)+2*d*sinh(s2*α*ϕ))/(-α*s2*sin(2*d*α)+2*d*cos(2*d*α)+2*d*cosh(s2*α*ϕ))
end
# When dα/dϕ = 1, we reach the classical regime, and stop the integration
condition(v,ϕ,integrator) = (ϕ*s2*sin(2*d*v[1])+2*d*sinh(s2*v[1]*ϕ))/(-v[1]*s2*sin(2*d*v[1])+2*d*cos(2*d*v[1])+2*d*cosh(s2*v[1]*ϕ))==1.0
affect!(integrator) = terminate!(integrator)
cb = DiscreteCallback(condition,affect!)
# Initial Condition at the bounce
α₀ = [a0]
classspan = (0,ϕ₀)
probquant = ODEProblem(quantum!,α₀,classspan,(s2,d))
solquant = solve(probquant,Tsit5(),callback=cb)
# α in the classical phase
function classic!(du,u,p,ϕ)
αc = u[1]
dαc = u[2]
du[1] = dαc
du[2] = 3*(-dαc^3/sqrt(2)+dαc^2+dαc/sqrt(2)-1)
end
# IC retrieved from end of quantum phase
init = [last(solquant);1.0]
classspan = (last(solquant.t),ϕ₀)
probclass = ODEProblem(classic!,init,classspan)
solclass = solve(probclass,Tsit5())
# α(ϕ) for ϕ>0
solu = append!(solquant[1,:],solclass[1,:]) # α
solt = append!(solquant.t,solclass.t) # ϕ
# α(ϕ) for ϕ<0
soloppu = reverse(solu)
soloppt = -reverse(solt)
pop!(soloppu)
pop!(soloppt)
# Join the two solutions
soltotu = append!(soloppu,solu)
soltott = append!(soloppt,solt)
soltot = DiffEqArray(soltotu,soltott)
end
plot(alpha_of_phi!(10000.0,-0.0009,0.0074847,2.0),yaxis=:log)
plot!(alpha_of_phi!(16.0,-0.1,0.00001,2.0))
```
If you were plotting a solution returned by solve directly, then the Plots recipes for DifferentialEquations enable an optional keyword argument for plot entitled plotdensity, which would let you choose the number of points plotted, and thus smoothness, as described in the docs, e.g.:
plot(sol,plotdensity=10000)
However, this keyword appears to require a solution object, rather than a DiffEqArray. Consequently, your best bet will indeed be manually setting saveat. For this approach, saveat = 0.01 would seem to be plenty to obtain fully smooth lines. However, this still gives the "range step cannot be zero" error you describe.
While I have no deep understanding of the system you are solving, an inspection of the results revealed duplicate timesteps in the results for alpha_of_phi!(16.0,-0.1,0.00001,2.0) run without saveat, suggesting that the classical simulation was being run with over a range of no time. In other words, this hints that last(solquant.t) may well be equal to or greater than ϕ₀ with these parameters, resulting in a timespan of zero. If so, this will quite understandably fail when you request to saveat some finite time within that timespan (last(solquant.t), ϕ₀).
Consequently, working on this hypothesis, if we just rewrite your function to check for this condition
using DifferentialEquations, Plots
function alpha_of_phi!(s2,d,a0,ϕ₀)
# α in the quantum phase
function quantum!(dv,v,p,ϕ)
s2,d=p
α = v[1]
dv[1] = (ϕ*s2*sin(2*d*α)+2*d*sinh(s2*α*ϕ))/(-α*s2*sin(2*d*α)+2*d*cos(2*d*α)+2*d*cosh(s2*α*ϕ))
end
# When dα/dϕ = 1, we reach the classical regime, and stop the integration
condition(v,ϕ,integrator) = (ϕ*s2*sin(2*d*v[1])+2*d*sinh(s2*v[1]*ϕ))/(-v[1]*s2*sin(2*d*v[1])+2*d*cos(2*d*v[1])+2*d*cosh(s2*v[1]*ϕ))==1.0
affect!(integrator) = terminate!(integrator)
cb = DiscreteCallback(condition,affect!)
# Initial Condition at the bounce
α₀ = [a0]
classspan = (0,ϕ₀)
probquant = ODEProblem(quantum!,α₀,classspan,(s2,d))
solquant = solve(probquant,Tsit5(),callback=cb,saveat=0.01)
# α in the classical phase
function classic!(du,u,p,ϕ)
αc = u[1]
dαc = u[2]
du[1] = dαc
du[2] = 3*(-dαc^3/sqrt(2)+dαc^2+dαc/sqrt(2)-1)
end
if last(solquant.t) < ϕ₀
# IC retrieved from end of quantum phase
init = [last(solquant);1.0]
classspan = (last(solquant.t),ϕ₀)
probclass = ODEProblem(classic!,init,classspan)
solclass = solve(probclass,Tsit5(),saveat=0.01)
# α(ϕ) for ϕ>0
solu = append!(solquant[1,:],solclass[1,:]) # α
solt = append!(solquant.t,solclass.t) # ϕ
else
solu = solquant[1,:] # α
solt = solquant.t # ϕ
end
# α(ϕ) for ϕ<0
soloppu = reverse(solu)
soloppt = -reverse(solt)
pop!(soloppu)
pop!(soloppt)
# Join the two solutions
soltotu = append!(soloppu,solu)
soltott = append!(soloppt,solt)
soltot = DiffEqArray(soltotu,soltott)
end
plot(alpha_of_phi!(10000.0,-0.0009,0.0074847,2.0),yaxis=:log)
plot!(alpha_of_phi!(16.0,-0.1,0.00001,2.0))
then we would seem to be in business!

What's the NeuralNetDiffEq.jl equivalent of this piece of code for solving ODE in Julia?

I am trying to make documentation for NeuraNetDiffEq.jl. I find an ODE solution with Flux.jl from this amazing tutorial here https://mitmath.github.io/18S096SciML/lecture2/ml .
using Flux
using DifferentialEquations
using LinearAlgebra
using Plots
using Statistics
NNODE = Chain(x -> [x],
Dense(1, 32, tanh),
Dense(32, 1),
first)
NNODE(1.0)
g(t) = t * NNODE(t) + 1f0
ϵ = sqrt(eps(Float32))
loss() = mean(abs2(((g(t + ϵ) - g(t)) / ϵ) - cos(2π * t)) for t in 0:1f-2:1f0)
opt = Flux.Descent(0.01)
data = Iterators.repeated((), 5000)
iter = 0
cb = function () # callback function to observe training
global iter += 1
if iter % 500 == 0
display(loss())
end
end
display(loss())
Flux.train!(loss, Flux.params(NNODE), data, opt; cb=cb)
t = 0:0.001:1.0
plot(t,g.(t),label="NN")
plot!(t,1.0 .+ sin.(2π .* t) / 2π, label="True")
I am having trouble understanding the parameters involved to invoke training process for NueralNetDiffEq.jl as in:
function DiffEqBase.solve(
prob::DiffEqBase.AbstractODEProblem,
alg::NeuralNetDiffEqAlgorithm,
args...;
dt,
timeseries_errors = true,
save_everystep=true,
adaptive=false,
abstol = 1f-6,
verbose = false,
maxiters = 100)
What would be a valid input for alg parameter? What would be the equivalent code in NeuralNetDiffEq.jl for the above ODE example?

How to train a Neural ODE to predict Lotka Voltera Time Series in Julia?

I want to decouple the ODE from which a time series data is generated and a Neural Network embedded in an ODE which is trying to learn the structure of this data. In other terms, I want to replicate the example on time-series extrapolation provided in https://julialang.org/blog/2019/01/fluxdiffeq/, but with a different underlying function, i.e. I am using Lotka-Voltera to generate the data.
My workflow in Julia is the following (Note that I am rather new to Julia, but I hope it's clear.):
train_size = 32
tspan_train = (0.0f0,4.00f0)
u0 = [1.0,1.0]
p = [1.5,1.0,3.0,1.0]
function lotka_volterra(du,u,p,t)
x, y = u
α, β, δ, γ = p
du[1] = dx = α*x - β*x*y
du[2] = dy = -δ*y + γ*x*y
end
t_train = range(tspan_train[1],tspan_train[2],length = train_size)
prob = ODEProblem(lotka_volterra, u0, tspan_train,p)
ode_data_train = Array(solve(prob, Tsit5(),saveat=t_train))
function create_neural_ode(solver, tspan, t_saveat)
dudt = Chain(
Dense(2,50,tanh),
Dense(50,2))
ps = Flux.params(dudt)
n_ode = NeuralODE(dudt, tspan, solver, saveat = t_saveat, reltol=1e-7, abstol=1e-9)
n_ode
end
function predict_n_ode(ps)
n_ode(u0,ps)
end
function loss_n_ode(ps)
pred = predict_n_ode(ps)
loss = sum(abs2, ode_data_train .- pred)
loss,pred
end
n_ode = create_neural_ode(Tsit5(), tspan_train, t_train)
final_p = Any[]
losses = []
cb = function(p,loss,pred)
display(loss)
display(p)
push!(final_p, copy(p))
push!(losses,loss)
pl = scatter(t_train, ode_data_train[1,:],label="data")
scatter!(pl,t_train,pred[1,:],label="prediction")
display(plot(pl))
end
sol = DiffEqFlux.sciml_train!(loss_n_ode, n_ode.p, ADAM(0.05), cb = cb, maxiters = 100)
# Plot and save training results
x = 1:100
plot_to_save = plot(x,losses,title=solver_name,label="loss")
plot(x,losses,title=solver_name, label="loss")
xlabel!("Epochs")
However I can observe that my NN is not learning much, it stagnates and the loss stays at around 155 with Euler and Tsit5, and behaves a bit better with RK4 (loss 142).
I would be very thankful if someone points out if I'm doing an error in my implementation or if this behaviour is expected.
Increasing the number for maxiters = to 300 helped achieving better fits, but the training is extremely unstable.

Stochastic differential equation sensitivity analysis with specified noise

I am trying to calculate the gradient of a functional of a stochastic differential equation (SDE) solution given a specific realization of the noise. I can successfully calculate these gradients if I leave the noise unspecified, as shown in DiffEqFlux.jl: Using Other Differential Equations. I can also successfully obtain the solution to my SDE for a specific noise realization, like shown in DifferentialEquations.jl: NoiseWrapper Example. When I try and put the two together, though, the code returns an error.
Here is a minimal working example adapted from the two separate examples referenced above:
using StochasticDiffEq, DiffEqBase, DiffEqNoiseProcess, DiffEqSensitivity, Zygote
function lotka_volterra(du,u,p,t)
x, y = u
α, β, δ, γ = p
du[1] = dx = α*x - β*x*y
du[2] = dy = -δ*y + γ*x*y
end
function lotka_volterra_noise(du,u,p,t)
du[1] = 0.1u[1]
du[2] = 0.1u[2]
end
dt = 1//2^(4)
u0 = [1.0,1.0]
p = [2.2, 1.0, 2.0, 0.4]
prob1 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p)
sol1 = solve(prob1,EM(),dt=dt,save_noise=true)
W2 = NoiseWrapper(sol1.W)
prob2 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p,noise=W2)
sol2 = solve(prob2,EM(),dt=dt)
function predict_sde1(p)
Array(concrete_solve(remake(prob1,p=p),EM(),dt=dt,sensealg=ForwardDiffSensitivity(),saveat=0.1))
end
loss_sde1(p) = sum(abs2,x-1 for x in predict_sde1(p))
loss_sde1(p)
# This gradient is successfully calculated
Zygote.gradient(loss_sde1,p)
function predict_sde2(p)
W2 = NoiseWrapper(sol1.W)
Array(concrete_solve(remake(prob2,p=p,noise=W2),EM(),dt=dt,sensealg=ForwardDiffSensitivity(),saveat=0.1))
end
loss_sde2(p) = sum(abs2,x-1 for x in predict_sde2(p))
# This loss is successfully calculated
loss_sde2(p)
# This gradient calculation raises and error
Zygote.gradient(loss_sde2,p)
The error I get at the end of running this code is
TypeError: in setfield!, expected Float64, got ForwardDiff.Dual{Nothing,Float64,4}
Stacktrace:
[1] setproperty! at ./Base.jl:21 [inlined]
...
followed by an interminable conclusion to the stacktrace (I can post it if you think it would be helpful, but since it's longer than the rest of this question I'd rather not clutter things up off the bat).
Is calculating gradients for SDE problems with specified noise realizations not currently supported, or am I just not making the appropriate function calls? I could easily believe the latter, since it was a bit of a struggle just to get to the point where the working parts of the above code worked, but I couldn't find any clue as to what I had incorrectly supplied after stepping through this code with the Juno debugger.
As a StackOverflow solution, you can use ForwardDiffSensitivity(convert_tspan=false) to work around this. Working code:
using StochasticDiffEq, DiffEqBase, DiffEqNoiseProcess, DiffEqSensitivity, Zygote
function lotka_volterra(du,u,p,t)
x, y = u
α, β, δ, γ = p
du[1] = dx = α*x - β*x*y
du[2] = dy = -δ*y + γ*x*y
end
function lotka_volterra_noise(du,u,p,t)
du[1] = 0.1u[1]
du[2] = 0.1u[2]
end
dt = 1//2^(4)
u0 = [1.0,1.0]
p = [2.2, 1.0, 2.0, 0.4]
prob1 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p)
sol1 = solve(prob1,EM(),dt=dt,save_noise=true)
W2 = NoiseWrapper(sol1.W)
prob2 = SDEProblem(lotka_volterra,lotka_volterra_noise,u0,(0.0,10.0),p,noise=W2)
sol2 = solve(prob2,EM(),dt=dt)
function predict_sde1(p)
Array(concrete_solve(remake(prob1,p=p),EM(),dt=dt,sensealg=ForwardDiffSensitivity(convert_tspan=false),saveat=0.1))
end
loss_sde1(p) = sum(abs2,x-1 for x in predict_sde1(p))
loss_sde1(p)
# This gradient is successfully calculated
Zygote.gradient(loss_sde1,p)
function predict_sde2(p)
Array(concrete_solve(prob2,EM(),prob2.u0,p,dt=dt,sensealg=ForwardDiffSensitivity(convert_tspan=false),saveat=0.1))
end
loss_sde2(p) = sum(abs2,x-1 for x in predict_sde2(p))
# This loss is successfully calculated
loss_sde2(p)
# This gradient calculation raises and error
Zygote.gradient(loss_sde2,p)
As a developer... this isn't a nice solution and our default should be better here. I'll work on this. You can track the development here https://github.com/JuliaDiffEq/DiffEqSensitivity.jl/issues/204. It'll probably get solved in an hour or so.
Edit: The fix is released and your original code works.

Output the index of the event in the solution when using VectorContinuousCallback in Julia/DifferentialEquations

I have a dynamical system for which many events could occur. I want to terminate the integration of the trajectory on an event, but I also want to known which event has been activated.
The workaround I found is to use a global variable to save the event index in the affect! function. Here is a modified version of the bouncing ball example:
using DifferentialEquations
function f(du,u,p,t)
du[1] = u[2]
du[2] = -p
du[3] = u[4]
du[4] = 0.0
end
function condition(out,u,t,integrator) # Event when event_f(u,t) == 0
out[1] = u[1]
out[2] = (u[3] - 10.0)u[3]
end
event_idx = [0, ] # global variable
function affect!(integrator, idx)
event_idx[1] = idx
terminate!(integrator)
end
cb = VectorContinuousCallback(condition,affect!,2)
u0 = [50.0,0.0,0.0,2.0]
tspan = (0.0,15.0)
p = 9.8
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,Tsit5(),callback=cb,dt=1e-3,adaptive=false)
println(sol.retcode)
println(event_idx) # [1]
Is there a better solution for doing this?
You can use integrator.p as a cache, and change around that value, and use sol.prob.p at the end. But yeah, this or a Ref is a fine way to do it.

Resources