feat: implementation for CasADiProblem

Author: vyudu

Project.toml

@@ -65,7 +65,9 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [weakdeps]
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
+CasADi = "c49709b8-5c63-11e9-2fb2-69db5844192f"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
@@ -75,6 +77,7 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 
 [extensions]
 MTKBifurcationKitExt = "BifurcationKit"
+MTKCasADiDynamicOptExt = ["CasADi", "DataInterpolations", "DiffEqDevTools"]
 MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
@@ -89,6 +92,7 @@ BifurcationKit = "0.4"
 BlockArrays = "1.1"
 BoundaryValueDiffEqAscher = "1.1.0"
 BoundaryValueDiffEqMIRK = "1.4.0"
+CasADi = "1.0.0"
 ChainRulesCore = "1"
 Combinatorics = "1"
 CommonSolve = "0.2.4"

ext/MTKCasADiDynamicOptExt.jl

@@ -3,6 +3,7 @@ using ModelingToolkit
 using CasADi
 using DiffEqDevTools, DiffEqBase
 using DataInterpolations
+using UnPack
 const MTK = MOdelingToolkit
 
 struct CasADiDynamicOptProblem{uType, tType, isinplace, P, F, K} <:
@@ -22,16 +23,15 @@ end
 
 struct CasADiModel
     opti::Opti
-    U::MX
-    V::MX
-end
-
-struct TimedMX
+    U::AbstractInterpolation
+    V::AbstractInterpolation
+    tₛ::Union{Number, MX}
 end
 
 function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
         dt = nothing,
-        steps = nothing, 
+        steps = nothing,
+        interpolation_method::AbstractInterpolation = LinearInterpolation,
         guesses = Dict(), kwargs...) 
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
@@ -43,73 +43,228 @@ function MTK.CasADiDynamicOptProblem(sys::ODESystem, u0map, tspan, pmap;
     model = init_model()
 end
 
-function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t)
+function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false, interp_type::AbstractInterpolation)
     ctrls = MTK.unbound_inputs(sys)
     states = unknowns(sys)
-    model = CasADi.Opti()
+    opti = CasADi.Opti()
+
+    if is_free_t
+        (ts_sym, te_sym) = tspan
+        MTK.symbolic_type(ts_sym) !== MTK.NotSymbolic() &&
+            error("Free initial time problems are not currently supported.")
+        tₛ = variable!(opti)
+        tsteps = LinRange(0, 1, steps)
+    else
+        tₛ = 1
+        tsteps = LinRange(tspan[1], tspan[2], steps)
+    end
 
-    U = CasADi.variable!(model, length(states), steps)
-    V = CasADi.variable!(model, length(ctrls), steps)
+    U = CasADi.variable!(opti, length(states), steps)
+    V = CasADi.variable!(opti, length(ctrls), steps)
+
+    U_interp = interp_type(Matrix(U), tsteps)
+    V_interp = interp_type(Matrix(V), tsteps)
+
+    CasADiModel(opti, U_interp, V_interp, tₛ)
 end
 
-function add_initial_constraints!()
-    
+function set_casadi_bounds!(model, sys, pmap)
+    @unpack opti, U, V = model
+    for (i, u) in enumerate(unknowns(sys))
+        if MTK.hasbounds(u)
+            lo, hi = MTK.getbounds(u)
+            subject_to!(opti, lo <= U[i, :] <= hi)
+        end
+    end
+    for (i, v) in enumerate(MTK.unbound_inputs(sys))
+        if MTK.hasbounds(v)
+            lo, hi = MTK.getbounds(v)
+            subject_to!(opti, lo <= V[i, :] <= hi)
+        end
+    end
+end
+
+function add_initial_constraints!(model::CasADiModel, u0, u0_idxs, ts)
+    @unpack opti, U = model
+    for i in u0_idxs
+        subject_to!(opti, U.u[i, 1] == u0[i])
+    end
 end
 
 function add_user_constraints!(model::CasADiModel, sys, pmap; is_free_t = false)
-    
+    @unpack opti, U, V, tₛ = model
+
+    iv = get_iv(sys)
+    conssys = MTK.get_constraintsystem(sys)
+    jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
+    (isnothing(jconstraints) || isempty(jconstraints)) && return nothing
+
+    stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
+    pidxmap = Dict([v => i for (i, v) in enumerate(ps)])
+    cons_unknowns = map(MTK.default_toterm, unknowns(conssys))
+    for st in cons_unknowns
+        x = operation(st)
+        t = only(argments(st))       
+        idx = stidxmap[x(iv)]
+
+        jconstraints = map(c -> Symbolics.substitute(c, Dict(x(t) => U(t)[idx])), jconstraints)
+    end
+    jconstraints = substitute_casadi_vars(model, sys, pmap, jconstraints)
+
+    for (i, cons) in enumerate(jconstraints)
+        if cons isa Equation
+            subject_to!(opti, cons.lhs - cons.rhs==0)
+        elseif cons.relational_op === Symbolics.geq
+            subject_to!(model, cons.lhs - cons.rhs≥0)
+        else
+            subject_to!(model, cons.lhs - cons.rhs≤0)
+        end
+    end
 end
 
-function add_cost_function!(model)
+function add_cost_function!(model::CasADiModel, sys, tspan, pmap)
+    @unpack opti, U, V, tₛ = model
+    jcosts = MTK.get_costs(sys)
+    consolidate = MTK.get_consolidate(sys)
+
+    if isnothing(jcosts) || isempty(jcosts)
+        minimize!(opti, 0)
+        return
+    end
+    stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
+    pidxmap = Dict([v => i for (i, v) in enumerate(ps)])
 
+    for i in 1:length(jcosts)
+        vars = vars(jcosts[i])
+        for st in vars
+            x = operation(st)
+            t = only(arguments(st))
+            t isa Union{Num, MTK.Symbolic} && continue
+            idx = stidxmap[x(iv)]
+            jcosts[i] = Symbolics.substitute(jcosts[i], Dict(x(t) => U(t)[idx]))
+        end
+    end
+    jcosts = substitute_casadi_vars(model::CasADiModel, sys, pmap, jcosts; auxmap)
+    jcosts = map(
+        c -> Symbolics.substitute(c, MTK.∫() => Symbolics.Integral(iv in tspan)), jcosts)
+
+    dt = U.t[2] - U.t[1]
+    intmap = Dict()
+    for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral)
+        op = MTK.operation(int)
+        arg = only(arguments(MTK.value(int)))
+        lo, hi = (op.domain.domain.left, op.domain.domain.right)
+        (lo, hi) !== tspan && error("Non-whole interval bounds for integrals are not currently supported.")
+        intmap[int] = dt * tₛ * sum(arg)
+    end
+    jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
+    minimize!(opti, consolidate(jcosts))
+end
+
+function substitute_casadi_vars(model::CasADiModel, sys, pmap, exprs; auxmap = Dict())
+    @unpack opti, U, V = model
+    iv = MTK.get_iv(sys)
+    sts = unknowns(sys)
+    cts = MTK.unbound_inputs(sys)
+
+    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
+    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
+
+    exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
+    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
+
+    # for variables like x(t)
+    whole_interval_map = Dict([[v => U.u[i, :] for (i, v) in enumerate(sts)];
+                               [v => V.u[i, :] for (i, v) in enumerate(cts)]])
+    exprs = map(c -> Symbolics.fixpoint_sub(c, whole_interval_map), exprs)
+    exprs
 end
 
 function add_solve_constraints!(prob, tableau; is_free_t)
-    A = tableau.A
-    α = tableau.α
-    c = tableau.c
-    model = prob.model
-    f = prob.f
-    p = prob.p
+    @unpack A, α, c = tableau
+    @unpack model, f, p = prob 
+    @unpack opti, U, V, tₛ = model
 
-    opti = model.opti
-    t = model[:t]
-    tsteps = supports(t)
-    tmax = tsteps[end]
-    pop!(tsteps)
-    tₛ = is_free_t ? model[:tf] : 1
+    tsteps = U.t 
     dt = tsteps[2] - tsteps[1]
 
-    U = model.U
-    V = model.V
     nᵤ = length(U)
     nᵥ = length(V)
 
     if is_explicit(tableau)
         K = Any[]
-        for τ in tsteps
+        for k in 1:length(tsteps)-1
             for (i, h) in enumerate(c)
                 ΔU = sum([A[i, j] * K[j] for j in 1:(i - 1)], init = zeros(nᵤ))
-                Uₙ = [U[i](τ) + ΔU[i] * dt for i in 1:nᵤ]
-                Vₙ = [V[i](τ) for i in 1:nᵥ]
+                Uₙ = U.u[:, k] + ΔU*dt
+                Vₙ = V.u[:, k]
                 Kₙ = tₛ * f(Uₙ, Vₙ, p, τ + h * dt) # scale the time
                 push!(K, Kₙ)
             end
             ΔU = dt * sum([α[i] * K[i] for i in 1:length(α)])
-            subject_to!(model, U[n](τ) + ΔU[n]==U[n](τ + dt))
-            empty!(K)
+            subject_to!(opti, U.u[:, k] + ΔU == U.u[:, k+1])
         end
     else
+        ΔU_tot = dt * (K' * α)
+        for k in 1:length(tsteps)-1
+            Kᵢ = variable!(opti, length(α), nᵤ)
+            ΔUs = A * Kᵢ # the stepsize at each stage of the implicit method
+            for (i, h) in enumerate(c)
+                ΔU = @view ΔUs[i, :]
+                Uₙ = U.u[:,k] + ΔU
+                Vₙ = V.u[:,k]
+                subject_to!(opti, K[i,:] == tₛ * f(Uₙ, Vₙ, p, τ + h*dt))
+            end
+            ΔU_tot = dt*(Kᵢ'*α)
+            subject_to!(opti, U.u[:, k] + ΔU_tot == U.u[:,k+1])
+        end
     end
 end
 
-function DiffEqBase.solve(prob::CasADiDynamicOptProblem, solver::Union{String, Symbol}, ode_solver::Union{String, Symbol}; silent = false)
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
+
+"""
+    solve(prob::CasADiDynamicOptProblem, casadi_solver, ode_solver; plugin_options, solver_options) 
+
+`plugin_options` and `solver_options` get propagated to the Opti object in CasADi.
+"""
+function DiffEqBase.solve(prob::CasADiDynamicOptProblem, solver::Union{String, Symbol}, ode_solver::Union{String, Symbol}; plugin_options::Dict = Dict(), solver_options::Dict = Dict(), silent = false)
     model = prob.model
     opti = model.opti
 
-    solver!(opti, solver)
-    sol = solve(opti)
-    DynamicOptSolution(model, sol, input_sol)
+    solver!(opti, solver, plugin_options, solver_options)
+    add_casadi_solve_constraints!(prob, tableau)
+    solver!(cmodel, "$solver", plugin_options, solver_options)
+
+    failed = false
+    try
+        sol = solve(opti)
+        value_getter = x -> CasADi.value(sol, x)
+    catch ErrorException
+        value_getter = x -> CasADi.debug_value(opti, x)
+        failed = true
+        continue
+    end
+
+    ts = value_getter(tₛ) * U.t
+    U_vals = value_getter(U)
+    U_vals = [[U_vals[i][j] for i in 1:length(U_vals)] for j in 1:length(ts)]
+    sol = DiffEqBase.build_solution(prob, ode_solver, ts, U_vals)
+
+    input_sol = nothing
+    if !isempty(V)
+        V_vals = value_getter(V)
+        V_vals = [[V_vals[i][j] for i in 1:length(V_vals)] for j in 1:length(ts)]
+        input_sol = DiffEqBase.build_solution(prob, ode_solver, ts, V_vals)
+    end
+
+    if failed
+        sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
+        !isnothing(input_sol) && (input_sol = SciMLBase.solution_new_retcode(
+            input_sol, SciMLBase.ReturnCode.ConvergenceFailure))
+    end
+
+    DynamicOptSolution(cmodel, sol, input_sol)
 end
 
 end

test/downstream/dynamic_opt_systems.jl

@@ -1,9 +1,27 @@
-function lotkavolterra()
-    
-end
+function build_lotkavolterra(; with_constraint = false)
+    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
+    @variables x(..) y(..)
+    t = M.t_nounits
+    D = M.D_nounits
 
-function ()
-    
+    eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
+        D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
+
+    tspan = (0.0, 1.0)
+    parammap = [α => 1.5, β => 1.0, γ => 3.0, δ => 1.0]
+
+    if with_constraint
+        constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0]
+        guess = [x(t) => 4.0, y(t) => 2.0]
+        u0map = Pair[]
+    else
+        constr = nothing
+        guess = Pair[]
+        u0map = [x(t) => 4.0, y(t) => 2.0]
+    end
+
+    @mtkbuild sys = ODESystem(eqs, t; constraints = constr)
+    sys, u0map, tspan, parammap, guess
 end
 
 function rocket_fft()