test: more test fixes

Author: vyudu

ext/MTKJuMPControlExt.jl

@@ -30,7 +30,7 @@ struct InfiniteOptControlProblem{uType, tType, isinplace, P, F, K} <:
     model::InfiniteModel
     kwargs::K
 
-    function InfiniteOptControlProblem(f, u0, tspan, p, model; kwargs...)
+    function InfiniteOptControlProblem(f, u0, tspan, p, model, kwargs...)
         new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f),
             typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
     end
@@ -58,7 +58,7 @@ 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))
-    f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     pmap = MTK.todict(pmap)
@@ -84,7 +84,7 @@ function MTK.InfiniteOptControlProblem(sys::ODESystem, u0map, tspan, pmap;
         guesses = Dict(), kwargs...)
     MTK.warn_overdetermined(sys, u0map)
     _u0map = has_alg_eqs(sys) ? u0map : merge(Dict(u0map), Dict(guesses))
-    f, u0, p = MTK.process_SciMLProblem(ControlFunction, sys, _u0map, pmap;
+    f, u0, p = MTK.process_SciMLProblem(ODEInputFunction, sys, _u0map, pmap;
         t = tspan !== nothing ? tspan[1] : tspan, kwargs...)
 
     pmap = MTK.todict(pmap)
@@ -116,16 +116,17 @@ function init_model(sys, tspan, steps, u0map, pmap, u0; is_free_t = false)
 
     @infinite_parameter(model, t in [tspan[1], tspan[2]], num_supports=steps)
     @variable(model, U[i = 1:length(states)], Infinite(t), start=u0[i])
-    c0 = [pmap[c] for c in ctrls]
+    c0 = MTK.value.([pmap[c] for c in ctrls])
     @variable(model, V[i = 1:length(ctrls)], Infinite(t), start=c0[i])
 
     set_jump_bounds!(model, sys, pmap)
     add_jump_cost_function!(model, sys, (tspan[1], tspan[2]), pmap; is_free_t)
     add_user_constraints!(model, sys, pmap; is_free_t)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(states)])
+    u0map = Dict([MTK.default_toterm(MTK.value(k)) => v for (k, v) in u0map])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(states)) :
-              [stidxmap[k] for (k, v) in u0map]
+        [stidxmap[MTK.default_toterm(k)] for (k, v) in u0map]
     add_initial_constraints!(model, u0, u0_idxs, tspan[1])
     return model
 end
@@ -190,7 +191,10 @@ function add_user_constraints!(model::InfiniteModel, sys, pmap; is_free_t = fals
         end
     end
 
-    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints)
+    auxmap = Dict([u => MTK.default_toterm(MTK.value(u)) for u in unknowns(conssys)])
+    jconstraints = substitute_jump_vars(model, sys, pmap, jconstraints; auxmap)
+    
+    # Substitute to-term'd variables
     for (i, cons) in enumerate(jconstraints)
         if cons isa Equation
             @constraint(model, cons.lhs - cons.rhs==0, base_name="user[$i]")
@@ -207,25 +211,28 @@ function add_initial_constraints!(model::InfiniteModel, u0, u0_idxs, ts)
     @constraint(model, initial[i in u0_idxs], U[i](ts)==u0[i])
 end
 
-function substitute_jump_vars(model, sys, pmap, exprs)
+function substitute_jump_vars(model, sys, pmap, exprs; auxmap = Dict())
     iv = MTK.get_iv(sys)
     sts = unknowns(sys)
     cts = MTK.unbound_inputs(sys)
     U = model[:U]
     V = model[:V]
+    exprs = map(c -> Symbolics.fixpoint_sub(c, auxmap), exprs)
+
     # 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)
+    exprs = map(c -> Symbolics.fixpoint_sub(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 = map(c -> Symbolics.fixpoint_sub(c, fixed_t_map), exprs)
+
+    exprs = map(c -> Symbolics.fixpoint_sub(c, Dict(pmap)), exprs)
     exprs
 end
 
@@ -255,8 +262,10 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
     model = prob.model
     f = prob.f
     p = prob.p
+
     t = model[:t]
-    tsteps = supports(model[:t])
+    tsteps = supports(t)
+    tmax = tsteps[end]
     pop!(tsteps)
     tₛ = is_free_t ? model[:tf] : 1
     dt = tsteps[2] - tsteps[1]
@@ -280,6 +289,7 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
                 base_name="solve_time_$τ")
             empty!(K)
         end
+        @show num_variables(model)
     else
         @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start=tsteps[1])
         ΔUs = A * K
@@ -288,10 +298,10 @@ function add_jump_solve_constraints!(prob, tableau; is_free_t = false)
             for (i, h) in enumerate(c)
                 ΔU = @view ΔUs[i, :]
                 Uₙ = U + ΔU * dt
-                @constraint(model, [j = 1:nᵤ], K[i, j](τ)==tₛ * f(Uₙ, V, p, τ + h * dt)[j],
-                    DomainRestrictions(t => τ + h * dt), base_name="solve_K($τ)")
+                @constraint(model, [j = 1:nᵤ], K[i, j](τ)==(tₛ * f(Uₙ, V, p, τ + h * dt)[j]),
+                            DomainRestrictions(t => min(τ + h * dt, tmax)), base_name="solve_K($τ)")
             end
-            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](τ + dt),
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU_tot[n]==U[n](min(τ + dt, tmax)),
                 DomainRestrictions(t => τ), base_name="solve_U($τ)")
         end
     end
@@ -323,6 +333,7 @@ function DiffEqBase.solve(
     unregister(model, :K)
     for var in all_variables(model)
         if occursin("K", JuMP.name(var))
+            unregister(model, Symbol(JuMP.name(var)))
             delete(model, var)
         end
     end

src/systems/diffeqs/odesystem.jl

@@ -406,7 +406,7 @@ function ODESystem(eqs, iv; constraints = Equation[], costs = Num[], kwargs...)
     end
     costs = wrap.(costs)
 
-    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars, consvars))),
+    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars))),
         collect(new_ps); constraintsystem, costs, kwargs...)
 end
 
@@ -769,7 +769,9 @@ function process_constraint_system(
     constraintps = OrderedSet()
     for cons in constraints
         collect_vars!(constraintsts, constraintps, cons, iv)
+        union!(constraintsts, collect_applied_operators(cons, Differential))
     end
+    @show constraintsts
 
     # Validate the states.
     validate_vars_and_find_ps!(constraintsts, constraintps, sts, iv)
@@ -811,6 +813,9 @@ function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
         elseif length(arguments(var)) > 1
             throw(ArgumentError("Too many arguments for variable $var."))
         elseif length(arguments(var)) == 1
+            if iscall(var) && operation(var) isa Differential
+                var = only(arguments(var))
+            end
             arg = only(arguments(var))
             operation(var)(iv) ∈ sts ||
                 throw(ArgumentError("Variable $var is not a variable of the ODESystem. Called variables must be variables of the ODESystem."))

src/systems/systems.jl

@@ -163,6 +163,11 @@ function __structural_simplify(
     end
 end
 
+function toterm_auxsystems(system::ODESystem)
+    constraints = system.constraintsystem.constraints
+
+end
+
 """
     $(TYPEDSIGNATURES)
 

test/dynamic_optimization/Project.toml

@@ -1,5 +1,14 @@
 [deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
-Multibody = "e1cad5d1-98ef-44f9-a79a-9ca4547f95b9"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
+OrdinaryDiffEqFIRK = "5960d6e9-dd7a-4743-88e7-cf307b64f125"
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
+SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"

test/dynamic_optimization/jump_control.jl

@@ -2,7 +2,7 @@ using ModelingToolkit
 import JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
-using OrdinaryDiffEqSDIRK
+using OrdinaryDiffEqSDIRK, OrdinaryDiffEqVerner, OrdinaryDiffEqTsit5, OrdinaryDiffEqFIRK
 using Ipopt
 using BenchmarkTools
 using CairoMakie
@@ -100,8 +100,8 @@ end
     constr = [v(1.0) ~ 0.0]
     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], []))
+                             [D(x(t)) ~ v(t), D(v(t)) ~ u(t)], t; costs = cost, constraints = constr)
+    block, input_idxs = structural_simplify(block, ([u(t)], []))
 
     u0map = [x(t) => 0.0, v(t) => 0.0]
     tspan = (0.0, 1.0)
@@ -113,19 +113,19 @@ end
     # Test reached final position.
     @test ≈(jsol.sol.u[end][1], 0.25, rtol = 1e-5)
     # Test dynamics
-    @parameters (u_interp::LinearInterpolation)(..)
-    block_ode = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u_interp(t)], t)
-    spline = ctrl_to_spline(jsol.input_sol, LinearInterpolation)
-    oprob = ODEProblem(block, u0map, tspan, [u_interp => spline])
-    osol = solve(oprob, Vern8())
-    @test jsol.sol.u ≈ osol.u
+    @parameters (u_interp::ConstantInterpolation)(..)
+    @mtkbuild block_ode = ODESystem([D(x(t)) ~ v(t), D(v(t)) ~ u_interp(t)], t)
+    spline = ctrl_to_spline(jsol.input_sol, ConstantInterpolation)
+    oprob = ODEProblem(block_ode, u0map, tspan, [u_interp => spline])
+    osol = solve(oprob, Vern8(), dt = 0.01, adaptive = false)
+    @test ≈(jsol.sol.u, osol.u, rtol = 0.05)
 
     iprob = InfiniteOptControlProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
     @test ≈(isol.sol.u[end][1], 0.25, rtol = 1e-5)
-    osol = solve(oprob, ImplicitEuler())
-    @test isol.sol.u ≈ osol.u
+    osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
+    @test ≈(isol.sol.u, osol.u, rtol = 0.05)
 
     ###################
     ### Bee example ###
@@ -154,10 +154,12 @@ end
     @parameters (α_interp::LinearInterpolation)(..)
     eqs = [D(w(t)) ~ -μ * w(t) + b * s * α_interp(t) * w(t),
            D(q(t)) ~ -ν * q(t) + c * (1 - α_interp(t)) * s * w(t)]
-    beesys_ode = ODESystem(eqs, t)
-    oprob = ODEProblem(beesys_ode, u0map, tspan, [α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation)])
-    osol = solve(oprob, Tsit5())
-    @test osol.u ≈ jsol.sol.u
+    @mtkbuild beesys_ode = ODESystem(eqs, t)
+    oprob = ODEProblem(beesys_ode, u0map, tspan, merge(Dict(pmap), Dict(α_interp => ctrl_to_spline(jsol.input_sol, LinearInterpolation))))
+    osol = solve(oprob, Tsit5(); dt = 0.01, adaptive = false)
+    @test ≈(osol.u, jsol.sol.u, rtol = 0.01)
+    osol2 = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
+    @test ≈(osol2.u, isol.sol.u, rtol = 0.01)
 end
 
 @testset "Rocket launch" begin
@@ -175,27 +177,31 @@ end
 
     (ts, te) = (0.0, 0.2)
     costs = [-h(te)]
-    constraints = [T(te) ~ 0]
-    @named rocket = ODESystem(eqs, t; costs, constraints)
+    cons = [T(te) ~ 0]
+    @named rocket = ODESystem(eqs, t; costs, constraints = cons)
     rocket, input_idxs = structural_simplify(rocket, ([T(t)], []))
 
     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₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
-    jsol = solve(jprob, Ipopt.Optimizer, :RadauIA3)
+    jsol = solve(jprob, Ipopt.Optimizer, :RadauIIA3)
     @test jsol.sol.u[end][1] > 1.012
+
+    iprob = InfiniteOptControlProblem(rocket, u0map, (ts, te), pmap; dt = 0.005, cse = false)
+    isol = solve(iprob, Ipopt.Optimizer, derivative_method = OrthogonalCollocation(3))
+    @test isol.sol.u[end][1] > 1.012
 
     # Test solution
     @parameters (T_interp::CubicSpline)(..)
     eqs = [D(h(t)) ~ v(t),
         D(v(t)) ~ (T_interp(t) - drag(h(t), v(t))) / m(t) - gravity(h(t)),
         D(m(t)) ~ -T_interp(t) / c]
-    rocket_ode = ODESystem(eqs, t)
-    interpmap = Dict(T_interp => ctrl_to_spline(jsol.inputsol, CubicSpline))
-    oprob = ODEProblem(rocket_ode, u0map, tspan, merge(pmap, interpmap))
-    osol = solve(oprob, RadauIA3())
+    @mtkbuild rocket_ode = ODESystem(eqs, t)
+    interpmap = Dict(T_interp => ctrl_to_spline(jsol.input_sol, CubicSpline))
+    oprob = ODEProblem(rocket_ode, u0map, (ts, te), merge(Dict(pmap), interpmap))
+    osol = solve(oprob, RadauIIA3())
     @test jsol.sol.u ≈ osol.u
 end
 
@@ -223,10 +229,44 @@ end
     @test isapprox(isol.sol.t[end], 10.0, rtol = 1e-3)
 end
 
-using JuliaSimCompiler
-using Multibody.PlanarMechanics
-
 @testset "Cart-pole problem" begin
+    # gravity, length, moment of Inertia, drag coeff
+    @parameters g l mₚ mₖ
+    @variables x(..) θ(..) u(t) [input = true, bounds = (-10, 10)]
+
+    s = sin(θ(t))
+    c = cos(θ(t))
+    H = [mₖ+mₚ  mₚ*l*c
+         mₚ*l*c mₚ*l^2]
+    C = [0 -mₚ*D(θ(t))*l*s
+         0  0]
+    qd = [D(x(t)), D(θ(t))]
+    G = [0, mₚ*g*l*s]
+    B = [1, 0]
+
+    tf = 5
+    rhss = -H \ Vector(C*qd + G - B*u)
+    eqs = [D(D(x(t))) ~ rhss[1], D(D(θ(t))) ~ rhss[2]]
+    cons = [θ(tf) ~ π, x(tf) ~ 0, D(θ(tf)) ~ 0, D(x(tf)) ~ 0]
+    costs = [∫(u^2)]
+    tspan = (0, tf)
+
+    @named cartpole = ODESystem(eqs, t; costs, constraints = cons)
+    cartpole, input_idxs = structural_simplify(cartpole, ([u], []))
+
+    u0map = [D(x(t)) => 0., D(θ(t)) => 0., θ(t) => 0., x(t) => 0.]
+    pmap = [mₖ => 1., mₚ => 0.2, l => 0.5, g => 9.81, u => 0]
+    jprob = JuMPControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
+    jsol = solve(jprob, Ipopt.Optimizer, :RK4)
+    @test jsol.sol.u[end] ≈ [π, 0, 0, 0]
+
+    iprob = InfiniteOptControlProblem(cartpole, u0map, tspan, pmap; dt = 0.04)
+    isol = solve(iprob, Ipopt.Optimizer)
+    @test isol.sol.u[end] ≈ [π, 0, 0, 0]
+end
+
+# RC Circuit
+@testset "MTK Components" begin
 end
 
 #@testset "Constrained optimal control problems" begin

test/runtests.jl

@@ -22,6 +22,12 @@ function activate_downstream_env()
     Pkg.instantiate()
 end
 
+function activate_dynamic_optimization_env()
+    Pkg.activate("dynamic_optimization")
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+end
+
 @time begin
     if GROUP == "All" || GROUP == "InterfaceI"
         @testset "InterfaceI" begin
@@ -143,4 +149,9 @@ end
         @safetestset "InfiniteOpt Extension Test" include("extensions/test_infiniteopt.jl")
         @safetestset "JuMPControl Extension Test" include("extensions/jump_control.jl")
     end
+
+    if GROUP == "All" || GROUP == "Dynamic Optimization"
+        activate_dynamic_optimization_env()
+        @safetestset "JuMP Collocation Solvers" include("dynamic_optimization/jump_control")
+    end
 end