partial: add free final time and bounds-handling

Author: vyudu

ext/MTKJuMPControlExt.jl

@@ -55,8 +55,11 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
+    @show _u0map
     f, u0, p = MTK.process_SciMLProblem(ODEFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
+
+    (f_i, f_o) = generate_control_function(sys)
     model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
 
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
@@ -86,13 +89,14 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
 end
 
 function init_model(sys, tsteps, u0map, u0)
-    ctrls = controls(sys)
+    ctrls = MTK.unbound_inputs(sys)
     states = unknowns(sys)
     model = InfiniteModel()
     @infinite_parameter(model, t in [tsteps[1], tsteps[end]], num_supports=length(tsteps))
     @variable(model, U[i = 1:length(states)], Infinite(t))
     @variable(model, V[1:length(ctrls)], Infinite(t))
 
+    set_bounds!(model, sys)
     add_jump_cost_function!(model, sys)
     add_user_constraints!(model, sys)
 
@@ -103,6 +107,22 @@ function init_model(sys, tsteps, u0map, u0)
     return model
 end
 
+function set_bounds!(model, sys)
+    U = model[:U]
+    for (i, u) in enumerate(unknowns(sys))
+        lo, hi = MTK.getbounds(u)
+        set_lower_bound(U[i], lo)
+        set_upper_bound(U[i], hi)
+    end
+
+    V = model[:V]
+    for (i, v) in enumerate(MTK.unbound_inputs(sys))
+        lo, hi = MTK.getbounds(v)
+        set_lower_bound(V[i], lo)
+        set_upper_bound(V[i], hi)
+    end
+end
+
 function add_jump_cost_function!(model::InfiniteModel, sys)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
@@ -113,7 +133,7 @@ function add_jump_cost_function!(model::InfiniteModel, sys)
     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))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
 
     for st in unknowns(sys)
         x = operation(st)
@@ -123,7 +143,7 @@ function add_jump_cost_function!(model::InfiniteModel, sys)
         jcosts = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jcosts)
     end
 
-    for ct in controls(sys)
+    for ct in MTK.unbound_inputs(sys)
         p = operation(ct)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
@@ -141,7 +161,7 @@ function add_user_constraints!(model::InfiniteModel, sys)
 
     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))])
+    cidxmap = Dict([v => i for (i, v) in enumerate(MTK.unbound_inputs(sys))])
 
     for st in unknowns(conssys)
         x = operation(st)
@@ -151,7 +171,7 @@ function add_user_constraints!(model::InfiniteModel, sys)
         jconstraints = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jconstraints)
     end
 
-    for ct in controls(sys)
+    for ct in MTK.unbound_inputs(sys)
         p = operation(ct)
         t = only(arguments(ct))
         idx = cidxmap[p(iv)]
@@ -185,9 +205,8 @@ function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
     V = model[:V]
 
     stmap = Dict([v => U[i] for (i, v) in enumerate(unknowns(sys))])
-    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(controls(sys))])
+    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))])
     submap = merge(stmap, ctrlmap, Dict(pmap))
-    @show submap
 
     # Differential equations
     diff_eqs = diff_equations(sys)

src/systems/optimal_control_interface.jl

@@ -19,3 +19,9 @@ function warn_overdetermined(sys, u0map)
             @warn "The control problem 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
 end
+
+"""
+IntegralNorm. When applied to an expression.
+"""
+struct IntegralNorm end
+

src/systems/problem_utils.jl

@@ -640,6 +640,7 @@ function maybe_build_initialization_problem(
         t = zero(floatT)
     end
 
+    @show u0map
     initializeprob = ModelingToolkit.InitializationProblem{true, SciMLBase.FullSpecialize}(
         sys, t, u0map, pmap; guesses, kwargs...)
     if state_values(initializeprob) !== nothing

test/extensions/jump_control.jl

@@ -1,26 +1,27 @@
 using ModelingToolkit
-using JuMP, InfiniteOpt
+import JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
 using OrdinaryDiffEqSDIRK
 using Ipopt
 using BenchmarkTools
+using CairoMakie
 const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
     # Test solving without anything attached.
     @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
-    M.@variables x(..) y(..)
+    @variables x(..) y(..)
     t = M.t_nounits
     D = M.D_nounits
 
     eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
         D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
 
+    @mtkbuild sys = ODESystem(eqs, t)
     tspan = (0.0, 1.0)
     u0map = [x(t) => 4.0, y(t) => 2.0]
     parammap = [α => 1.5, β => 1.0, γ => 3.0, δ => 1.0]
-    @mtkbuild sys = ODESystem(eqs, t)
 
     # Test explicit method.
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
@@ -58,27 +59,70 @@ const M = ModelingToolkit
     sol = isol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
+
+    # Test whole-interval constraints
+    constr = [x(t) > 3, y(t) > 4]
+    @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
+    iprob = InfiniteOptControlProblem(
+        lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    isol = @btime solve(
+        $iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
+    sol = isol.sol
+    @test all(u -> u .> [3, 4], sol.u)
+
+    jprob = JuMPControlProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
+    jsol = @btime solve($jprob, Ipopt.Optimizer, :RadauIA3, silent = true) # 12.190 s, 9.68 GiB
+    sol = jsol.sol
+    @test all(u -> u .> [3, 4], sol.u)
 end
 
-#@testset "Optimal control: bees" begin
-#    # Example from Lawrence Evans' notes
-#    M.@variables w(..) q(..) 
-#    M.@parameters α(t) [bounds = [0, 1]] b c μ s ν
-#
-#    tspan = (0, 4)
-#    eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
-#           D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
-#    costs = [-q(tspan[2])]
-#   
-#    @mtkbuild beesys = ODESystem(eqs, t; costs)
-#    u0map = [w(0) => 40, q(0) => 2]
-#    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1]
-#
-#    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
-#    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
-#    control_sol = jsol.control_sol
-#    # Bang-bang control
-#end
+@testset "Linear systems" begin
+    function is_bangbang(input_sol, lbounds, ubounds)
+        bangbang = true
+        for v in 1:length(input_sol.u[1])
+            all(i -> i[v] ≈ bounds[v] || i[v] ≈ bounds[u], input_sol.u) || (bangbang = false)
+        end
+        bangbang
+    end
+
+    # Double integrator
+    @variables x(..) [bounds = (0., 0.25)] v(..)
+    @variables u(t) [bounds = (-1., 1.), input = true]
+    constr = [v(1.0) ~ 0.0]
+    cost = [-x(1.0)] # Optimize the final distance.
+    @named block = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u], t)
+    block, input_idxs = structural_simplify(block, ([u],[]))
+
+    u0map = [x(t) => 0., v(t) => 0.]
+    tspan = (0., 1.)
+    parammap = [u => 0.]
+    jprob = JuMPControlProblem(block, u0map, tspan, parammap; dt = 0.01)
+    jsol = solve(jprob, Ipopt.Optimizer, :Verner8)
+    # Linear systems have bang-bang controls
+    @test is_bangbang(jsol.input_sol, [-1.], [1.])
+    # Test reached final position.
+    @test jsol.sol.u[end][1] ≈ 0.25
+
+    # Cart-pole system
+
+    # Bee example (from Lawrence Evans' notes)
+    M.@variables w(..) q(..) 
+    M.@parameters α(t) [bounds = [0, 1]] b c μ s ν
+
+    tspan = (0, 4)
+    eqs = [D(w(t)) ~ -μ*w(t) + b*s*α*w(t),
+           D(q(t)) ~ -ν*q(t) + c*(1 - α)*s*w(t)]
+    costs = [-q(tspan[2])]
+   
+    @mtkbuild beesys = ODESystem(eqs, t; costs)
+    u0map = [w(0) => 40, q(0) => 2]
+    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1]
+
+    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
+    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    control_sol = jsol.control_sol
+    # Bang-bang control
+end
 #
 #@testset "Constrained optimal control problems" begin
 #end