JuMP optimizing parametric DifferentialEquation

Hello,
I am trying to do the following: I have to optimize some parameters of an ODE, using JuMP (I have an objective function that depends on the solution ode the ODE). Now the problem is that the ODE itself has some other parameters (that I don’t want to optimize) cooked into a struct (using Parameters.jl).
So my idea has been that of defining an objective function for JuMP via JuMP.register that takes the parameter that I need to optimize and the constant one and bash them into the real objective function (joining them together as a tuple).

What happens is that Julia complains that I am trying to convert Dual numbers into Float64. I guess that is an issue when I extract the parameter that I need to optimize and the constant one (stored in a struct) inside the derivative in the ODE.

I have a minimal (non-)working example which is basically the adaptation of Using JuMP with DiffEqParamEstim · DiffEqParamEstim.jl
by taking an ODE with additional parameter as in my case (also usign Ipopt because I like it more than NLopt).

Thanks in advance for the help!

using DifferentialEquations, JuMP, Ipopt, Parameters

@with_kw struct parameters
    extra = 1.0
end
param = parameters()

function g(du,u,p_join,t)
    p,param = p_join
    σ,ρ,β = p
    x,y,z = u
    du[1] = dx = σ*(y-x)*param.extra
    du[2] = dy = x*(ρ-z) - y
    du[3] = dz = x*y - β*z

end

u0 = [1.0;0.0;0.0]
t = 0.0:0.01:1.0
tspan = (0.0,1.0)
model_ode(p_) = ODEProblem(g, u0, tspan,p_)
solve_model(mp_) = DifferentialEquations.solve(model_ode(mp_), Tsit5(),saveat=0.01)

#test that the ODE works
solve_model(([8,15,10/3],param)) 

function loss_objective(mp_)
    mp_join_ = (mp_,param)
    sol = solve_model(mp_join_)
    return sol[1,end]
end
juobj(args...) = loss_objective(args)

jumodel  = Model(Ipopt.Optimizer)

JuMP.register(jumodel, :juobj, 3, juobj, autodiff=true)
JuMP.@variables jumodel begin
    σ,(start=8)
    ρ,(start=25.0)
    β,(start=10/3)
end
@NLobjective(jumodel, Min, juobj(σ, ρ, β))
JuMP.unset_silent(jumodel)

JuMP.optimize!(jumodel)

output (up to relatively helpful lines):

julia> JuMP.optimize!(jumodel)
This is Ipopt version 3.13.4, running with linear solver mumps.
NOTE: Other linear solvers might be more efficient (see Ipopt documentation).

Number of nonzeros in equality constraint Jacobian...:        0
Number of nonzeros in inequality constraint Jacobian.:        0
Number of nonzeros in Lagrangian Hessian.............:        0

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#138#140"{typeof(juobj)}, Float64}, Float64, 3})
Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  (::Type{T})(::T) where T<:Number at boot.jl:760
  (::Type{T})(::VectorizationBase.Double{T}) where T<:Union{Float16, Float32, Float64, VectorizationBase.Vec{var"#s33", var"#s32"} where {var"#s33", var"#s32"<:Union{Float16, Float32, Float64}}, VectorizationBase.VecUnroll{var"#s31", var"#s30", var"#s29", V} where {var"#s31", var"#s30", var"#s29"<:Union{Float16, Float32, Float64}, V<:Union{Bool, Float16, Float32, Float64, Int16, Int32, Int64, Int8, UInt16, UInt32, UInt64, UInt8, SIMDTypes.Bit, VectorizationBase.AbstractSIMD{var"#s30", var"#s29"}}}} at /home/cricci/.julia/packages/VectorizationBase/TQLhj/src/special/double.jl:84
  ...
Stacktrace:
  [1] convert(#unused#::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#138#140"{typeof(juobj)}, Float64}, Float64, 3})
    @ Base ./number.jl:7
  [2] setindex!(A::Vector{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#138#140"{typeof(juobj)}, Float64}, Float64, 3}, i1::Int64)
    @ Base ./array.jl:839
  [3] g(du::Vector{Float64}, u::Vector{Float64}, p_join::Tuple{Tuple{ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#138#140"{typeof(juobj)}, Float64}, Float64, 3}, ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#138#140"{typeof(juobj)}, Float64}, Float64, 3}, ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#138#140"{typeof(juobj)}, Float64}, Float64, 3}}, parameters}, t::Float64)
    @ Main ~/actualpath/JumpDiffEq.jl:7

JuMP is not a good choice for this sort of problem. See https://jump.dev/JuMP.jl/stable/background/should_i_use/#Black-box,-derivative-free,-or-unconstrained-optimization.

Since this is unconstrained, you’re better off using Optim.jl or GalacticOptim.jl for this.

Your error arises because your ODE solve needs to support generic number types. I assume the DifferentialEquations documentation has some advice on this somewhere.

Thanks!
Sorry, this is by bad. My example here is unconstrained optimization, while the problem I actually have to solve is constrained (I just wanted to provide a short working example)

Your error arises because your ODE solve needs to support generic number types. I assume the DifferentialEquations documentation has some advice on this somewhere.

I completely agree. It seems to me that while without the extra parameters DifferentialEquations is able to detect the type of u or du when p is a Dual number. While when I add the extra parameter then It doesn’t and just assumes Float64 that doesn’t work then.

I assume you’ll need something like

function model_ode(p_::AbstractVector{T}) where {T}
    u = [one(T), zero(T), zero(T)]
    return ODEProblem(g, u, tspan, p_)
end

Passing in u0 which is a Vector{Float64} won’t work.

Wow, how can I have been so dumb not to think of It. Thank you very much sir for you help.
I will post here the pieces of code that need to be changed, in the spirit of what you suggested, in order to make it work.

function model_ode(p_::Tuple{AbstractVector{T}, parameters}) where {T}
    u0 = [one(T), zero(T), zero(T)]
    return ODEProblem(g, u0, tspan, p_)
end

function loss_objective(mp_)
    mp_join_ = (collect(mp_),param)
    sol = solve_model(mp_join_)
    return sol[1,end]
end
1 Like