feat: working linear control problems

Author: vyudu

ext/MTKJuMPControlExt.jl

@@ -58,7 +58,7 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, pmap, u0)
     JuMPControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
@@ -80,22 +80,23 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
     f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
-    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, u0)
+    model = init_model(sys, tspan[1]:dt:tspan[2], u0map, pmap, u0)
     add_infopt_solve_constraints!(model, sys, pmap)
     InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
 end
 
-function init_model(sys, tsteps, u0map, u0)
+function init_model(sys, tsteps, u0map, pmap, u0)
     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)
+    add_jump_cost_function!(model, sys, (tsteps[1], tsteps[2]), pmap)
+    add_user_constraints!(model, sys, pmap)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
@@ -120,63 +121,35 @@ function set_bounds!(model, sys)
     end
 end
 
-function add_jump_cost_function!(model::InfiniteModel, sys)
+function add_jump_cost_function!(model::InfiniteModel, sys, tspan, pmap)
     jcosts = MTK.get_costs(sys)
     consolidate = MTK.get_consolidate(sys)
     if isnothing(jcosts) || isempty(jcosts)
         @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(MTK.unbound_inputs(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 = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jcosts)
-    end
+    jcosts = substitute_jump_vars(model, sys, pmap, jcosts)
 
-    for ct in MTK.unbound_inputs(sys)
-        p = operation(ct)
-        t = only(arguments(ct))
-        idx = cidxmap[p(iv)]
-        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jcosts = map(c -> Symbolics.substitute(c, Dict(p(t) => subval)), jcosts)
+    # Substitute integral
+    iv = MTK.get_iv(sys)
+    jcosts = map(c -> Symbolics.substitute(c, ∫ => Symbolics.Integral(iv in tspan)), jcosts)
+    intmap = Dict()
+    
+    for int in MTK.collect_applied_operators(jcosts, Symbolics.Integral)
+        arg = only(arguments(MTK.value(int)))
+        lower_bound, upper_bound = (int.domain.domain.left, int.domain.domain.right)
+        intmap[int] = InfiniteOpt.∫(arg, iv; lower_bound, upper_bound)
     end
-
+    jcosts = map(c -> Symbolics.substitute(c, intmap), jcosts)
     @objective(model, Min, consolidate(jcosts))
 end
 
-function add_user_constraints!(model::InfiniteModel, sys)
+function add_user_constraints!(model::InfiniteModel, sys, pmap)
     conssys = MTK.get_constraintsystem(sys)
     jconstraints = isnothing(conssys) ? nothing : MTK.get_constraints(conssys)
     (isnothing(jconstraints) || isempty(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(MTK.unbound_inputs(sys))])
-
-    for st in unknowns(conssys)
-        x = operation(st)
-        t = only(arguments(st))
-        idx = stidxmap[x(iv)]
-        subval = isequal(t, iv) ? model[:U][idx] : model[:U][idx](t)
-        jconstraints = map(c -> Symbolics.substitute(c, Dict(x(t) => subval)), jconstraints)
-    end
-
-    for ct in MTK.unbound_inputs(sys)
-        p = operation(ct)
-        t = only(arguments(ct))
-        idx = cidxmap[p(iv)]
-        subval = isequal(t, iv) ? model[:V][idx] : model[:V][idx](t)
-        jconstraints = map(
-            c -> Symbolics.substitute(jconstraints, Dict(p(t) => subval)), jconstriants)
-    end
-
+    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints)
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
             @constraint(model, cons.lhs - cons.rhs==0, base_name="user[$i]")
@@ -193,31 +166,41 @@ function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
-is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
-
-function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
+function substitute_jump_vars(model, sys, pmap, exprs)
     iv = MTK.get_iv(sys)
-    t = model[:t]
+    sts = unknowns(sys)
+    cts = MTK.unbound_inputs(sys)
     U = model[:U]
     V = model[:V]
+    # for variables like x(t)
+    whole_interval_map = Dict([[v => U[i] for (i, v) in enumerate(sts)]; [v => V[i] for (i, v) in enumerate(cts)]])
+    exprs = map(c -> Symbolics.substitute(c, whole_interval_map), exprs)
+
+    # for variables like x(1.0)
+    x_ops = [MTK.operation(MTK.unwrap(st)) for st in sts]
+    c_ops = [MTK.operation(MTK.unwrap(ct)) for ct in cts]
+    fixed_t_map = Dict([[x_ops[i] => U[i] for i in 1:length(U)]; [c_ops[i] => V[i] for i in 1:length(V)]])
+    exprs = map(c -> Symbolics.substitute(c, fixed_t_map), exprs)
+
+    exprs = map(c -> Symbolics.substitute(c, Dict(pmap)), exprs)
+    exprs
+end
 
-    stmap = Dict([v => U[i] for (i, v) in enumerate(unknowns(sys))])
-    ctrlmap = Dict([v => V[i] for (i, v) in enumerate(MTK.unbound_inputs(sys))])
-    submap = merge(stmap, ctrlmap, Dict(pmap))
+is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
 
+function add_infopt_solve_constraints!(model::InfiniteModel, sys, pmap)
     # Differential equations
-    diff_eqs = diff_equations(sys)
-    D = Differential(iv)
+    U = model[:U]
+    t = model[:t]
+    D = Differential(MTK.get_iv(sys))
     diffsubmap = Dict([D(U[i]) => ∂(U[i], t) for i in 1:length(U)])
-    for u in unknowns(sys)
-        diff_eqs = map(e -> Symbolics.substitute(e, submap), diff_eqs)
-        diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
-    end
+
+    diff_eqs = substitute_jump_vars(model, sys, pmap, diff_equations(sys))
+    diff_eqs = map(e -> Symbolics.substitute(e, diffsubmap), diff_eqs)
     @constraint(model, D[i = 1:length(diff_eqs)], diff_eqs[i].lhs==diff_eqs[i].rhs)
 
     # Algebraic equations
-    alg_eqs = alg_equations(sys)
-    alg_eqs = map(e -> Symbolics.substitute(e, submap), alg_eqs)
+    alg_eqs = substitute_jump_vars(model, sys, pmap, alg_equations(sys))
     @constraint(model, A[i = 1:length(alg_eqs)], alg_eqs[i].lhs==alg_eqs[i].rhs)
 end
 
@@ -306,9 +289,10 @@ end
 `derivative_method` kwarg refers to the method used by InfiniteOpt to compute derivatives. The list of possible options can be found at https://infiniteopt.github.io/InfiniteOpt.jl/stable/guide/derivative/. Defaults to FiniteDifference(Backward()).
 """
 function DiffEqBase.solve(prob::InfiniteOptControlProblem, jump_solver;
-        derivative_method = InfiniteOpt.FiniteDifference(Backward()))
+        derivative_method = InfiniteOpt.FiniteDifference(Backward()), silent = false)
+    model = prob.model
     silent && set_silent(model)
-    set_derivative_method(prob.model[:t], derivative_method)
+    set_derivative_method(model[:t], derivative_method)
     _solve(prob, jump_solver, derivative_method)
 end
 

src/systems/optimal_control_interface.jl

@@ -134,9 +134,15 @@ function SciMLBase.ControlFunction{false}(sys::AbstractODESystem, args...;
 end
 
 """
-IntegralNorm. When applied to an expression.
+IntegralNorm. When applied to an expression in a cost
+function, assumes that the integration variable is the
+iv of the system, and assumes that the bounds are the
+tspan.
+Equivalent to Integral(t in tspan) in Symbolics.
 """
-struct IntegralNorm end
+struct ∫ <: Symbolics.Operator end
+∫(x) = ∫()(x)
+Base.show(io::IO, x::∫) = print(io, "∫")
 
 """
 $(SIGNATURES)

src/systems/problem_utils.jl

@@ -640,7 +640,6 @@ 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

@@ -77,10 +77,10 @@ const M = ModelingToolkit
 end
 
 @testset "Linear systems" begin
-    function is_bangbang(input_sol, lbounds, ubounds)
+    function is_bangbang(input_sol, lbounds, ubounds, rtol = 1e-4)
         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)
+        for v in 1:length(input_sol.u[1]) - 1
+            all(i -> ≈(i[v], bounds[v]; rtol) || ≈(i[v], bounds[u]; rtol), input_sol.u) || (bangbang = false)
         end
         bangbang
     end
@@ -91,8 +91,8 @@ end
     @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)
+    cost = [-x(1.0)] # Maximize the final distance.
+    @named block = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u], t; costs = cost, constraints = constr)
     block, input_idxs = structural_simplify(block, ([u],[]))
 
     u0map = [x(t) => 0., v(t) => 0.]
@@ -103,28 +103,85 @@ end
     # 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
+    @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
 
-    # Cart-pole system
+    iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
+    isol = solve(iprob, Ipopt.Optimizer; silent = true)
+    @test is_bangbang(isol.input_sol, [-1.], [1.])
+    @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
 
-    # Bee example (from Lawrence Evans' notes)
-    @variables w(..) q(..) 
-    @parameters α(t) [bounds = [0, 1]] b c μ s ν
+    ###################
+    ### Bee example ###
+    ###################
+    # From Lawrence Evans' notes
+    @variables w(..) q(..) α(t) [input = true, bounds = (0, 1)]
+    @parameters 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]
+    @named beesys = ODESystem(eqs, t; costs)
+    beesys, input_idxs = structural_simplify(beesys, ([α],[]))
+    u0map = [w(t) => 40, q(t) => 2]
+    pmap = [b => 1, c => 1, μ => 1, s => 1, ν => 1, α => 1]
 
-    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap)
+    jprob = JuMPControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
     jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
-    control_sol = jsol.control_sol
-    # Bang-bang control
+    @test is_bangbang(jsol.input_sol, [0.], [1.])
+    iprob = InfiniteOptControlProblem(beesys, u0map, tspan, pmap, dt = 0.01)
+    isol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    @test is_bangbang(isol.input_sol, [0.], [1.])
 end
-#
+
+@testset "Rocket launch" begin
+    t = M.t_nounits
+    D = M.D_nounits
+
+    @variables h(..) v(..) m(..) T(..) [input = true, bounds = (0, tₘ)]
+    @parameters h_c m₀ h₀ g₀ D_c c Tₘ
+    @parameters tf
+    drag(h, v) = D_c * v^2 * exp(-h_c * (h - h₀) / h₀)
+    gravity(h) = g₀ * (h₀ / h)
+
+    eqs = [D(h(t)) ~ v(t), 
+           D(v(t)) ~ (T(t) - drag(h(t), v(t))) / m(t) - gravity(t),
+           D(m(t)) ~ -T(t) / c]
+
+    costs = [-h(tf)]
+    constraints = [T(tf) ~ 0]
+    @named rocket = ODESystem(eqs, t; costs, constraints)
+    @test tf ∈ Set(parameters(rocket))
+
+    u0map = [h(t) => h₀, m(t) => m₀, v(t) => 0]
+    pmap = [g₀ => 1, m₀ => 1.0, h_c => 500, c => 0.5*√(g₀*h₀), D_C => 0.5 * 620 * m₀/g₀, Tₘ => 3.5*g₀*m₀]
+    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
+    @test jsol.sol.u[end][1] ≈ 1.012
+end
+
+@testset "Free final time problem" begin
+    t = M.t_nounits
+    D = M.D_nounits
+
+    @variables x(..) u(..) [input = true, bounds = (0,1)]
+    @parameters tf
+    eqs = [D(x(t)) ~ -2 + 0.5*u]
+
+    # Integral cost function
+    costs = [∫(x-u), x(tf)]
+    consolidate(u) = u[1] + u[2]
+    jprob = JuMPControlProblem(rocket, u0map, (0, tf), pmap)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
+    @test jsol.sol.t[end] ≈ 10.0
+    iprob = InfiniteOptControlProblem(rocket, u0map, (0, tf), pmap)
+    isol = solve(iprob, Ipopt.Optimizer, :RadauIA3)
+    @test isol.sol.t[end] ≈ 10.0
+end
+
+@testset "Cart-pole problem" begin
+end
+
 #@testset "Constrained optimal control problems" begin
 #end