feat: solver tableau debugging

Author: vyudu

Project.toml

@@ -31,6 +31,7 @@ FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
@@ -66,8 +67,10 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 
 [extensions]
@@ -76,7 +79,7 @@ MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
-MTKJuMP = "JuMP"
+MTKJuMPControlExt = ["JuMP", "DiffEqDevTools"]
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
@@ -98,6 +101,7 @@ DeepDiffs = "1"
 DelayDiffEq = "5.50"
 DiffEqBase = "6.165.1"
 DiffEqCallbacks = "2.16, 3, 4"
+DiffEqDevTools = "2.48.0"
 DiffEqNoiseProcess = "5"
 DiffRules = "0.1, 1.0"
 DifferentiationInterface = "0.6.47"
@@ -116,6 +120,8 @@ FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
 InfiniteOpt = "0.5"
 InteractiveUtils = "1"
+Ipopt = "1.8.0"
+JuMP = "1.25.0"
 JuliaFormatter = "1.0.47"
 JumpProcesses = "9.13.1"
 LabelledArrays = "1.3"

ext/MTKJuMPControlExt.jl

@@ -0,0 +1,218 @@
+module MTKJuMPControlExt
+using ModelingToolkit
+using JuMP, InfiniteOpt
+using DiffEqDevTools, DiffEqBase, SciMLBase
+using LinearAlgebra
+const MTK = ModelingToolkit
+
+struct JuMPControlProblem{uType, tType, P, F, K}
+       f::F
+       u0::uType
+       tspan::tType
+       p::P
+       model::InfiniteModel
+       kwargs::K
+
+       function JuMPControlProblem(f, u0, tspan, p, model; kwargs...) 
+           new{typeof(u0), typeof(tspan), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
+       end
+end
+
+"""
+    JuMPControlProblem(sys::ODESystem, u0, tspan, p; dt)
+
+Convert an ODESystem representing an optimal control system into a JuMP model
+for solving using optimization. Must provide `dt` for determining the length 
+of the interpolation arrays.
+
+The optimization variables:
+- a vector-of-vectors U representing the unknowns as an interpolation array
+- a vector-of-vectors V representing the controls as an interpolation array
+
+The constraints are:
+- The set of user constraints passed to the ODESystem via `constraints`
+- The solver constraints that encode the time-stepping used by the solver
+"""
+function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("dt must be provided for JuMPControlProblem."), kwargs...)
+    ts = tspan[1]
+    te = tspan[2]
+    steps = ts:dt:te
+    ctrls = controls(sys)
+    states = unknowns(sys)
+    constraintsys = MTK.get_constraintsystem(sys)
+
+    if !isnothing(constraintsys)
+        (length(constraints(constraintsys)) + length(u0map) > length(sts)) &&
+            @warn "The JuMPControlProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The solvers will default to doing a nonlinear least-squares optimization."
+    end
+
+    f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, u0map, pmap;
+        t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
+
+    model = InfiniteModel()
+    @infinite_parameter(model, t in [ts, te], num_supports = length(steps), derivative_method = OrthogonalCollocation(2))
+    @variable(model, U[1:length(states)], Infinite(t), start = ts)
+    @variable(model, V[1:length(ctrls)], Infinite(t), start = ts)
+
+    add_jump_cost_function!(model, sys)
+    add_user_constraints!(model, sys)
+
+    stidxmap = Dict([v => i for (i, v) in enumerate(states)])
+    u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) : [stidxmap[k] for (k, v) in u0map]
+    add_initial_constraints!(model, u0, u0_idxs, tspan)
+
+    JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
+end
+
+function add_jump_cost_function!(model, sys)
+    jcosts = MTK.get_costs(sys)
+    consolidate = MTK.get_consolidate(sys)
+    if isnothing(consolidate)
+        @objective(model, Min, 0)
+        return
+    end
+    iv = MTK.get_iv(sys)
+
+    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
+
+    for st in unknowns(sys)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
+        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => subval))
+    end
+    
+    for ct in controls(sys)
+        p = operation(ct)
+        t = only(arguments(ct))
+        idx = cidxmap[p(iv)]
+        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
+        jcosts = Symbolics.substitute(jcosts, Dict(x(t) => subval))
+    end
+    
+    @objective(model, Min, consolidate(jcosts))
+end
+
+function add_user_constraints!(model, sys)
+    jconstraints = if !(csys = MTK.get_constraintsystem(sys) isa Nothing)
+        MTK.get_constraints(csys)
+    else
+        nothing
+    end
+    isnothing(jconstraints) && return nothing
+
+    iv = MTK.get_iv(sys)
+    stidxmap = Dict([v => i for (i, v) in enumerate(unknowns(sys))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(controls(sys))])
+
+    for st in unknowns(sys)
+        x = operation(st)
+        t = only(arguments(st))
+        idx = stidxmap[x(iv)]
+        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
+        jconstraints = Symbolics.substitute(jconstraints, Dict(x(t) => subval))
+    end
+
+    for ct in controls(sys)
+        p = operation(ct)
+        t = only(arguments(ct))
+        idx = cidxmap[p(iv)]
+        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
+        jconstraints = Symbolics.substitute(jconstraints, Dict(p(t) => subval))
+    end
+
+    for (i, cons) in enumerate(jconstraints)
+        if cons isa Equation
+            @constraint(model, user[i], cons.lhs - cons.rhs == 0)
+        elseif cons.relational_op === Symbolics.geq 
+            @constraint(model, user[i], cons.lhs - cons.rhs ≥ 0)
+        else
+            @constraint(model, user[i], cons.lhs - cons.rhs ≤ 0)
+        end
+    end
+end
+
+function add_initial_constraints!(model, u0, u0_idxs, tspan)
+    ts = tspan[1]
+    @constraint(model, init_u0_idx[i in u0_idxs], model[:U][i](ts) == u0[i])
+end
+
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
+
+function add_solve_constraints!(prob, tableau)
+    A = tableau.A
+    α = tableau.α
+    c = tableau.c
+    model = prob.model
+    f = prob.f
+    p = prob.p
+    tsteps = supports(model[:t])
+    pop!(tsteps)
+    dt = tsteps[2] - tsteps[1]
+
+    U = model[:U]
+    nᵤ = length(U)
+    if is_explicit(tableau)
+        K = Any[]
+        for τ in tsteps
+            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ᵤ]
+                Kₙ = f(Uₙ, p, τ + h*dt)
+                push!(K, Kₙ)
+            end
+            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt))
+            empty!(K)
+        end
+    else
+        @variable(model, K[1:length(a), 1:nᵤ], Infinite(t), start = tsteps[1])
+        for τ in tsteps
+            ΔUs = [A * K(τ)]
+            for (i, h) in enumerate(c)
+                ΔU = ΔUs[i]
+                Uₙ = [U[j](τ) + ΔU[j](τ)*dt for j in 1:nᵤ]
+                @constraint(model, K[i](τ) == f(Uₙ, p, τ + h*dt))
+            end
+            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, U(τ) + dot(α, K(τ)) == U(τ + dt))
+        end
+    end
+end
+
+"""
+"""
+struct JuMPControlSolution
+    model::InfiniteModel
+    sol::ODESolution
+end
+
+"""
+Solve JuMPControlProblem. Arguments:
+- prob: a JumpControlProblem
+- jump_solver: a LP solver such as HiGHS
+- ode_solver: Takes in a symbol representing the solver. Acceptable solvers may be found at https://docs.sciml.ai/DiffEqDevDocs/stable/internals/tableaus/. Note that the symbol may be different than the typical name of the solver, e.g. :Tsitouras5 rather than Tsit5.
+"""
+function DiffEqBase.solve(prob::JuMPControlProblem, jump_solver, ode_solver::Symbol)
+    model = prob.model
+    tableau_getter = Symbol(:construct, ode_solver)
+    @eval tableau = $tableau_getter()
+    ts = supports(model[:t])
+    add_solve_constraints!(prob, tableau)
+
+    set_optimizer(model, jump_solver)
+    optimize!(model)
+
+    if is_solved_and_feasible(model)
+        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
+        JuMPControlSolution(model, sol)
+    else
+        sol = DiffEqBase.build_solution(prob, ode_solver, ts, value.(U))
+        sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
+        JuMPControlSolution(model, sol)
+    end
+end
+
+end