fix: use dt in the constraints

Author: vyudu

Project.toml

@@ -31,7 +31,6 @@ 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"
@@ -120,7 +119,6 @@ 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"

ext/MTKJuMPControlExt.jl

@@ -5,7 +5,7 @@ using DiffEqDevTools, DiffEqBase, SciMLBase
 using LinearAlgebra
 const MTK = ModelingToolkit
 
-struct JuMPControlProblem{uType, tType, P, F, K}
+struct JuMPControlProblem{uType, tType, isinplace, P, F, K} <: SciMLBase.AbstractODEProblem{uType, tType, isinplace}
        f::F
        u0::uType
        tspan::tType
@@ -14,7 +14,7 @@ struct JuMPControlProblem{uType, tType, P, F, K}
        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)
+           new{typeof(u0), typeof(tspan), SciMLBase.isinplace(f), typeof(p), typeof(f), typeof(kwargs)}(f, u0, tspan, p, model, kwargs)
        end
 end
 
@@ -50,9 +50,9 @@ function MTK.JuMPControlProblem(sys::ODESystem, u0map, tspan, pmap; dt = error("
         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)
+    @infinite_parameter(model, t in [ts, te], num_supports = length(steps))
+    @variable(model, U[i = 1:length(states)], Infinite(t))
+    @variable(model, V[1:length(ctrls)], Infinite(t))
 
     add_jump_cost_function!(model, sys)
     add_user_constraints!(model, sys)
@@ -136,7 +136,8 @@ 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])
+    U = model[:U]
+    @constraint(model, initial[i in u0_idxs], U[i](ts) == u0[i])
 end
 
 is_explicit(tableau) = tableau isa DiffEqDevTools.ExplicitRKTableau
@@ -148,6 +149,7 @@ function add_solve_constraints!(prob, tableau)
     model = prob.model
     f = prob.f
     p = prob.p
+    t = model[:t]
     tsteps = supports(model[:t])
     pop!(tsteps)
     dt = tsteps[2] - tsteps[1]
@@ -163,21 +165,21 @@ function add_solve_constraints!(prob, tableau)
                 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))
+            ΔU = dt*sum([α[i] * K[i] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n](τ) + ΔU[n] == U[n](τ + dt), base_name = "solve_time_$τ")
             empty!(K)
         end
     else
-        @variable(model, K[1:length(a), 1:nᵤ], Infinite(t), start = tsteps[1])
+        @variable(model, K[1:length(α), 1:nᵤ], Infinite(t), start = tsteps[1])
         for τ in tsteps
-            ΔUs = [A * K(τ)]
+            Δ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))
+                ΔU = ΔUs[i, :]
+                Uₙ = [U[j] + ΔU[j]*dt for j in 1:nᵤ]
+                @constraint(model, [j in 1:nᵤ], K[i, j] == f(Uₙ, p, τ + h*dt)[j], DomainRestrictions(t => τ), base_name = "solve_K($τ)")
             end
-            ΔU = sum([α[i] * K[i] for i in 1:length(α)])
-            @constraint(model, U(τ) + dot(α, K(τ)) == U(τ + dt))
+            ΔU = dt*sum([α[i] * K[i, :] for i in 1:length(α)])
+            @constraint(model, [n = 1:nᵤ], U[n] + ΔU[n] == U[n](τ + dt), DomainRestrictions(t => τ), base_name = "solve_U($τ)")
         end
     end
 end
@@ -194,25 +196,46 @@ 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.
+
+Returns a JuMPControlSolution, which contains both the model and the ODE solution.
 """
 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])
+
+    # Unregister current solver constraints
+    for con in all_constraints(model)
+        if occursin("solve", JuMP.name(con))
+            unregister(model, Symbol(JuMP.name(con)))
+            delete(model, con)
+        end
+    end
+    for var in all_variables(model)
+        @show JuMP.name(var)
+        if occursin("K", JuMP.name(var))
+            unregister(model, Symbol(JuMP.name(var)))
+            delete(model, var)
+        end
+    end
     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))
+    tstatus = termination_status(model)
+    pstatus = primal_status(model)
+    !has_values(model) && error("Model not solvable; please report this to github.com/SciML/ModelingToolkit.jl.")
+
+    U_vals = value.(model[: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)
+
+    if !(pstatus === FEASIBLE_POINT && (tstatus === OPTIMAL || tstatus === LOCALLY_SOLVED || tstatus === ALMOST_OPTIMAL || tstatus === ALMOST_LOCALLY_SOLVED))
         sol = SciMLBase.solution_new_retcode(sol, SciMLBase.ReturnCode.ConvergenceFailure)
-        JuMPControlSolution(model, sol)
     end
+    JuMPControlSolution(model, sol)
 end
 
 end

test/extensions/Project.toml

@@ -14,10 +14,10 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Nemo = "2edaba10-b0f1-5616-af89-8c11ac63239a"
 NonlinearSolveHomotopyContinuation = "2ac3b008-d579-4536-8c91-a1a5998c2f8b"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
+SimpleDiffEq = "05bca326-078c-5bf0-a5bf-ce7c7982d7fd"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

test/extensions/ad.jl

@@ -3,7 +3,7 @@ using ModelingToolkit: t_nounits as t, D_nounits as D
 using Zygote
 using SymbolicIndexingInterface
 using SciMLStructures
-using OrdinaryDiffEq
+using OrdinaryDiffEqTsit5
 using NonlinearSolve
 using SciMLSensitivity
 using ForwardDiff

test/extensions/jump_control.jl

@@ -2,7 +2,8 @@ using ModelingToolkit
 using JuMP, InfiniteOpt
 using DiffEqDevTools, DiffEqBase
 using SimpleDiffEq
-using HiGHS
+using OrdinaryDiffEqSDIRK
+using Ipopt 
 const M = ModelingToolkit
 
 @testset "ODE Solution, no cost" begin
@@ -22,36 +23,33 @@ const M = ModelingToolkit
     # Test explicit method.
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap, dt = 0.01)
     @test num_constraints(jprob.model) == 2 # initials
-    jsol = solve(jprob, Ipopt.Optimizer, :Tsitouras5)
+    jsol = solve(jprob, Ipopt.Optimizer, :RK4)
     oprob = ODEProblem(sys, u0map, tspan, parammap)
-    osol = solve(oprob, SimpleTsit5(), adaptive = false)
+    osol = solve(oprob, SimpleRK4(), dt = 0.01)
     @test jsol.sol.u ≈ osol.u
 
     # Implicit method.
-    jsol2 = solve(prob, Ipopt.Optimizer, :RK4)
-    osol2 = solve(oprob, SimpleRK4(), adaptive = false)
-    @test jsol2.sol.u ≈ osol2.u
+    jsol2 = solve(jprob, Ipopt.Optimizer, :ImplicitEuler)
+    osol2 = solve(oprob, ImplicitEuler(), dt = 0.01, adaptive = false)
+    @test ≈(jsol2.sol.u, osol2.u, rtol = 0.001)
 
     # With a constraint
-    @parameters α=1.5 β=1.0 γ=3.0 δ=1.0
-    @variables x(..) y(..)
-
-    eqs = [D(x(t)) ~ α * x(t) - β * x(t) * y(t),
-        D(y(t)) ~ -γ * y(t) + δ * x(t) * y(t)]
-
-    u0map = []
-    tspan = (0.0, 1.0)
+    u0map = Pair[]
     guess = [x(t) => 4.0, y(t) => 2.0]
     constr = [x(0.6) ~ 3.5, x(0.3) ~ 7.0]
     @mtkbuild lksys = ODESystem(eqs, t; constraints = constr)
 
     jprob = JuMPControlProblem(sys, u0map, tspan, parammap; guesses, dt = 0.01)
     @test num_constraints(jprob.model) == 2 == num_variables(jprob.model) == 2
-    jsol = solve(prob, HiGHS.Optimizer, :Tsitouras5)
+    jsol = solve(prob, Ipopt.Optimizer, :Tsitouras5)
     sol = jsol.sol
     @test sol(0.6)[1] ≈ 3.5
     @test sol(0.3)[1] ≈ 7.0
 end
 
 @testset "Optimal control problem" begin
+    # Investing
+    
+   
+    # Bang-bang control
 end