wip: multiple shooting transcription

franckgaga · franckgaga · commit 3db0f9612a00 · 2025-02-23T17:44:52.000-05:00
diff --git a/src/controller/construct.jl b/src/controller/construct.jl
@@ -379,10 +379,10 @@ function setconstraint!(
         )
         A = con.A[con.i_b, :]
         b = con.b[con.i_b]
-        ΔŨvar::Vector{JuMP.VariableRef} = optim[:ΔŨvar]
+        Z̃var::Vector{JuMP.VariableRef} = optim[:Z̃var]
         JuMP.delete(optim, optim[:linconstraint])
         JuMP.unregister(optim, :linconstraint)
-        @constraint(optim, linconstraint, A*ΔŨvar .≤ b)
+        @constraint(optim, linconstraint, A*Z̃var .≤ b)
         set_nonlincon!(mpc, model, optim)
     else
         i_b, i_g = init_matconstraint_mpc(
diff --git a/src/controller/execute.jl b/src/controller/execute.jl
@@ -67,7 +67,7 @@ function moveinput!(
     validate_args(mpc, ry, d, D̂, R̂y, R̂u)
     initpred!(mpc, mpc.estim.model, d, D̂, R̂y, R̂u)
     linconstraint!(mpc, mpc.estim.model)
-    linconstrainteq!(mpc, mpc.transcription)
+    linconstrainteq!(mpc, mpc.estim.model, mpc.transcription)
     Z̃ = optim_objective!(mpc)
     return getinput(mpc, Z̃)
 end
@@ -268,10 +268,11 @@ predictstoch!(Ŷs, mpc::PredictiveController, ::StateEstimator) = (Ŷs .= 0; n
 @doc raw"""
     linconstraint!(mpc::PredictiveController, model::LinModel)
 
-Set `b` vector for the linear model inequality constraints (``\mathbf{A ΔŨ ≤ b}``).
+Set `b` vector for the linear model inequality constraints (``\mathbf{A Z̃ ≤ b}``).
 
-Also init ``\mathbf{f_x̂} = \mathbf{g_x̂ d}(k) + \mathbf{j_x̂ D̂} + \mathbf{k_x̂ x̂_0}(k) + \mathbf{v_x̂ u}(k-1) + \mathbf{b_x̂}``
-vector for the terminal constraints, see [`init_predmat`](@ref).
+Also init ``\mathbf{f_x̂} = \mathbf{g_x̂ d_0}(k) + \mathbf{j_x̂ D̂_0} + \mathbf{k_x̂ x̂_0}(k) + 
+\mathbf{v_x̂ u_0}(k-1) + \mathbf{b_x̂}`` vector for the terminal constraints, see
+[`init_predmat`](@ref).
 """
 function linconstraint!(mpc::PredictiveController, model::LinModel)
     nU, nΔŨ, nY = length(mpc.con.U0min), length(mpc.con.ΔŨmin), length(mpc.con.Y0min)
@@ -324,22 +325,6 @@ function linconstraint!(mpc::PredictiveController, ::SimModel)
     return nothing
 end
 
-linconstrainteq!(mpc::PredictiveController, transcription::SingleShooting) = nothing
-function linconstrainteq!(mpc::PredictiveController, transcription::MultipleShooting)
-    nx̂, Fŝ = mpc.estim.nx̂, mpc.con.Fŝ
-    Fŝ .= mpc.con.Bŝ
-    mul!(Fŝ, mpc.con.Kŝ, mpc.estim.x̂0, 1, 1)
-    mul!(Fŝ, mpc.con.Vŝ, mpc.estim.lastu0, 1, 1)
-    if mpc.estim.model.nd ≠ 0
-        mul!(Fŝ, mpc.con.Gŝ, mpc.d0, 1, 1)
-        mul!(Fŝ, mpc.con.Jŝ, mpc.D̂0, 1, 1)
-    end
-    mpc.con.beq .= @. -Fŝ
-    linconeq = mpc.optim[:linconstrainteq]
-    JuMP.set_normalized_rhs(linconeq, mpc.con.beq)
-    return nothing
-end
-
 @doc raw"""
     predict!(Ŷ0, x̂0, _, _, _, mpc::PredictiveController, model::LinModel, ΔŨ) -> Ŷ0, x̂0end
 
diff --git a/src/controller/explicitmpc.jl b/src/controller/explicitmpc.jl
@@ -184,18 +184,18 @@ end
 linconstraint!(::ExplicitMPC, ::LinModel) = nothing
 
 @doc raw"""
-    optim_objective!(mpc::ExplicitMPC) -> ΔŨ
+    optim_objective!(mpc::ExplicitMPC) -> Z̃
 
 Analytically solve the optimization problem for [`ExplicitMPC`](@ref).
 
-The solution is ``\mathbf{ΔŨ = - H̃^{-1} q̃}``, see [`init_quadprog`](@ref).
+The solution is ``\mathbf{Z̃ = - H̃^{-1} q̃}``, see [`init_quadprog`](@ref).
 """
-optim_objective!(mpc::ExplicitMPC) = lmul!(-1, ldiv!(mpc.ΔŨ, mpc.H̃_chol, mpc.q̃))
+optim_objective!(mpc::ExplicitMPC) = lmul!(-1, ldiv!(mpc.Z̃, mpc.H̃_chol, mpc.q̃))
 
 "Compute the predictions but not the terminal states if `mpc` is an [`ExplicitMPC`](@ref)."
-function predict!(Ŷ, x̂, _ , _ , _ , mpc::ExplicitMPC, ::LinModel, ΔŨ)
+function predict!(Ŷ, x̂, _ , _ , _ , mpc::ExplicitMPC, ::LinModel, Z̃)
     # in-place operations to reduce allocations :
-    Ŷ .= mul!(Ŷ, mpc.Ẽ, ΔŨ) .+ mpc.F
+    Ŷ .= mul!(Ŷ, mpc.Ẽ, Z̃) .+ mpc.F
     x̂ .= NaN
     return Ŷ, x̂
 end
diff --git a/src/controller/nonlinmpc.jl b/src/controller/nonlinmpc.jl
@@ -524,10 +524,10 @@ Inspired from: [User-defined operators with vector outputs](https://jump.dev/JuM
 function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT<:Real
     model = mpc.estim.model
     nu, ny, nx̂, nϵ, Hp = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ, mpc.Hp
-    ng, nc, nΔŨ, nU, nŶ = length(mpc.con.i_g), mpc.con.nc, length(mpc.ΔŨ), Hp*nu, Hp*ny
+    ng, nc, nZ̃, nU, nŶ = length(mpc.con.i_g), mpc.con.nc, length(mpc.Z̃), Hp*nu, Hp*ny
     nUe, nŶe = nU + nu, nŶ + ny
-    Ncache = nΔŨ + 3
-    ΔŨ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nΔŨ), Ncache)
+    Ncache = nZ̃ + 3 
+    Z̃_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, nZ̃),  Ncache)
     Ŷe_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶe), Ncache)
     Ue_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nUe), Ncache)
     Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, nŶ),  Ncache)
@@ -538,91 +538,91 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
     û0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu),  Ncache)
     g_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, ng),  Ncache)
     gc_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nc),  Ncache)
-    function update_simulations!(ΔŨ, ΔŨtup::NTuple{N, T}) where {N, T<:Real}
-        if any(new !== old for (new, old) in zip(ΔŨtup, ΔŨ)) # new ΔŨtup, update predictions:
-            ΔŨ1 = ΔŨtup[begin]
-            for i in eachindex(ΔŨtup)
-                ΔŨ[i] = ΔŨtup[i] # ΔŨ .= ΔŨtup seems to produce a type instability
+    function update_simulations!(Z̃, Z̃tup::NTuple{N, T}) where {N, T<:Real}
+        if any(new !== old for (new, old) in zip(Z̃tup, Z̃)) # new Z̃tup, update predictions:
+            Z̃1 = Z̃tup[begin]
+            for i in eachindex(Z̃tup)
+                Z̃[i] = Z̃tup[i] # Z̃ .= Z̃tup seems to produce a type instability
             end
-            Ŷe, Ue     = get_tmp(Ŷe_cache, ΔŨ1), get_tmp(Ue_cache, ΔŨ1)
-            Ȳ,  Ū      = get_tmp(Ȳ_cache, ΔŨ1),  get_tmp(Ū_cache, ΔŨ1)
-            x̂0, x̂0next = get_tmp(x̂0_cache, ΔŨ1), get_tmp(x̂0next_cache, ΔŨ1)
-            u0, û0     = get_tmp(u0_cache, ΔŨ1), get_tmp(û0_cache, ΔŨ1)
-            gc, g      = get_tmp(gc_cache, ΔŨ1), get_tmp(g_cache, ΔŨ1)
-            Ŷ0, x̂0end  = predict!(Ȳ, x̂0, x̂0next, u0, û0, mpc, model, ΔŨ)
-            Ŷe, Ue     = extended_predictions!(Ŷe, Ue, Ū, mpc, model, Ŷ0, ΔŨ)
-            ϵ = (nϵ ≠ 0) ? ΔŨ[end] : zero(T) # ϵ = 0 if nϵ == 0 (meaning no relaxation)
+            Ŷe, Ue     = get_tmp(Ŷe_cache, Z̃1), get_tmp(Ue_cache, Z̃1)
+            Ȳ,  Ū      = get_tmp(Ȳ_cache, Z̃1),  get_tmp(Ū_cache, Z̃1)
+            x̂0, x̂0next = get_tmp(x̂0_cache, Z̃1), get_tmp(x̂0next_cache, Z̃1)
+            u0, û0     = get_tmp(u0_cache, Z̃1), get_tmp(û0_cache, Z̃1)
+            gc, g      = get_tmp(gc_cache, Z̃1), get_tmp(g_cache, Z̃1)
+            Ŷ0, x̂0end  = predict!(Ȳ, x̂0, x̂0next, u0, û0, mpc, model, Z̃)
+            Ŷe, Ue     = extended_predictions!(Ŷe, Ue, Ū, mpc, model, Ŷ0, Z̃)
+            ϵ = (nϵ ≠ 0) ? Z̃[end] : zero(T) # ϵ = 0 if nϵ == 0 (meaning no relaxation)
             gc = con_custom!(gc, mpc, Ue, Ŷe, ϵ)
             g  = con_nonlinprog!(g, mpc, model, x̂0end, Ŷ0, gc, ϵ)
         end
         return nothing
     end
-    function Jfunc(ΔŨtup::Vararg{T, N}) where {N, T<:Real}
-        ΔŨ1 = ΔŨtup[begin]
-        ΔŨ = get_tmp(ΔŨ_cache, ΔŨ1)
-        update_simulations!(ΔŨ, ΔŨtup)
-        Ȳ,  Ū  = get_tmp(Ȳ_cache, ΔŨ1),  get_tmp(Ū_cache, ΔŨ1)
-        Ŷe, Ue = get_tmp(Ŷe_cache, ΔŨ1), get_tmp(Ue_cache, ΔŨ1)
-        return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ue, Ŷe, ΔŨ)::T
+    function Jfunc(Z̃tup::Vararg{T, N}) where {N, T<:Real}
+        Z̃1 = Z̃tup[begin]
+        Z̃ = get_tmp(Z̃_cache, Z̃1)
+        update_simulations!(Z̃, Z̃tup)
+        Ȳ,  Ū  = get_tmp(Ȳ_cache, Z̃1),  get_tmp(Ū_cache, Z̃1)
+        Ŷe, Ue = get_tmp(Ŷe_cache, Z̃1), get_tmp(Ue_cache, Z̃1)
+        return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ue, Ŷe, Z̃)::T
     end
-    function gfunc_i(i, ΔŨtup::NTuple{N, T}) where {N, T<:Real}
-        ΔŨ1 = ΔŨtup[begin]
-        ΔŨ = get_tmp(ΔŨ_cache, ΔŨ1)
-        update_simulations!(ΔŨ, ΔŨtup)
-        g = get_tmp(g_cache, ΔŨ1)
+    function gfunc_i(i, Z̃tup::NTuple{N, T}) where {N, T<:Real}
+        Z̃1 = Z̃tup[begin]
+        Z̃ = get_tmp(Z̃_cache, Z̃1)
+        update_simulations!(Z̃, Z̃tup)
+        g = get_tmp(g_cache, Z̃1)
         return g[i]::T
     end
     gfuncs = Vector{Function}(undef, ng)
     for i in 1:ng
         # this is another syntax for anonymous function, allowing parameters T and N:
-        gfuncs[i] = function (ΔŨtup::Vararg{T, N}) where {N, T<:Real}
-            return gfunc_i(i, ΔŨtup)
+        gfuncs[i] = function (Z̃tup::Vararg{T, N}) where {N, T<:Real}
+            return gfunc_i(i, Z̃tup)
         end
     end
     return Jfunc, gfuncs
 end
 
 function init_nonlincon!(mpc::NonLinMPC, ::LinModel, gfuncs::Vector{<:Function}) 
     optim, con = mpc.optim, mpc.con
-    nΔŨ = length(mpc.ΔŨ)
+    nZ̃ = length(mpc.Z̃)
     if length(con.i_g) ≠ 0
         i_base = 0
         for i in 1:con.nc
             name = Symbol("g_c_$i")
-            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfuncs[i_base+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
         end
     end
     return nothing
 end
 
 function init_nonlincon!(mpc::NonLinMPC, ::NonLinModel, gfuncs::Vector{<:Function}) 
     optim, con = mpc.optim, mpc.con
-    ny, nx̂, Hp, nΔŨ = mpc.estim.model.ny, mpc.estim.nx̂, mpc.Hp, length(mpc.ΔŨ)
+    ny, nx̂, Hp, nZ̃ = mpc.estim.model.ny, mpc.estim.nx̂, mpc.Hp, length(mpc.Z̃)
     if length(con.i_g) ≠ 0
         i_base = 0
         for i in eachindex(con.Y0min)
             name = Symbol("g_Y0min_$i")
-            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfuncs[i_base+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
         end
         i_base = 1Hp*ny
         for i in eachindex(con.Y0max)
             name = Symbol("g_Y0max_$i")
-            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfuncs[i_base+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
         end
         i_base = 2Hp*ny
         for i in eachindex(con.x̂0min)
             name = Symbol("g_x̂0min_$i")
-            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfuncs[i_base+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
         end
         i_base = 2Hp*ny + nx̂
         for i in eachindex(con.x̂0max)
             name = Symbol("g_x̂0max_$i")
-            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfuncs[i_base+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
         end
         i_base = 2Hp*ny + 2nx̂
         for i in 1:con.nc
             name = Symbol("g_c_$i")
-            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfuncs[i_base+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
         end
     end
     return nothing
@@ -636,13 +636,13 @@ Set the custom nonlinear inequality constraints for `LinModel`.
 function set_nonlincon!(
     mpc::NonLinMPC, ::LinModel, optim::JuMP.GenericModel{JNT}
 ) where JNT<:Real
-    ΔŨvar = optim[:ΔŨvar]
+    Z̃var = optim[:Z̃var]
     con = mpc.con
     nonlin_constraints = JuMP.all_constraints(optim, JuMP.NonlinearExpr, MOI.LessThan{JNT})
     map(con_ref -> JuMP.delete(optim, con_ref), nonlin_constraints)
     for i in 1:con.nc
         gfunc_i = optim[Symbol("g_c_$i")]
-        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     return nothing
 end
@@ -655,29 +655,29 @@ Also set output prediction `Ŷ` and terminal state `x̂end` constraints when no
 function set_nonlincon!(
     mpc::NonLinMPC, ::SimModel, optim::JuMP.GenericModel{JNT}
 ) where JNT<:Real
-    ΔŨvar = optim[:ΔŨvar]
+    Z̃var = optim[:Z̃var]
     con = mpc.con
     nonlin_constraints = JuMP.all_constraints(optim, JuMP.NonlinearExpr, MOI.LessThan{JNT})
     map(con_ref -> JuMP.delete(optim, con_ref), nonlin_constraints)
     for i in findall(.!isinf.(con.Y0min))
         gfunc_i = optim[Symbol("g_Y0min_$(i)")]
-        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.Y0max))
         gfunc_i = optim[Symbol("g_Y0max_$(i)")]
-        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.x̂0min))
         gfunc_i = optim[Symbol("g_x̂0min_$(i)")]
-        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.x̂0max))
         gfunc_i = optim[Symbol("g_x̂0max_$(i)")]
-        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in 1:con.nc
         gfunc_i = optim[Symbol("g_c_$i")]
-        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     return nothing
 end
diff --git a/src/controller/transcription.jl b/src/controller/transcription.jl
@@ -508,6 +508,33 @@ function init_defectmat(
     return Eŝ, Gŝ, Jŝ, Kŝ, Vŝ, Bŝ
 end
 
+@doc raw"""
+    linconstrainteq!(
+        mpc::PredictiveController, model::LinModel, transcription::MultipleShooting
+    )
+
+Set `beq` vector for the linear model equality constraints (``\mathbf{Aeq Z̃ = beq}``).
+
+Also init ``\mathbf{F_ŝ} = \mathbf{G_ŝ d_0}(k) + \mathbf{J_ŝ D̂_0} + \mathbf{K_ŝ x̂_0}(k) + 
+\mathbf{V_ŝ u_0}(k-1) + \mathbf{B_ŝ}``, see [`init_defectmat`](@ref).
+"""
+function linconstrainteq!(mpc::PredictiveController, model::LinModel, ::MultipleShooting)
+    Fŝ  = mpc.con.Fŝ
+    Fŝ .= mpc.con.Bŝ
+    mul!(Fŝ, mpc.con.Kŝ, mpc.estim.x̂0, 1, 1)
+    mul!(Fŝ, mpc.con.Vŝ, mpc.estim.lastu0, 1, 1)
+    if model.nd ≠ 0
+        mul!(Fŝ, mpc.con.Gŝ, mpc.d0, 1, 1)
+        mul!(Fŝ, mpc.con.Jŝ, mpc.D̂0, 1, 1)
+    end
+    mpc.con.beq .= @. -Fŝ
+    linconeq = mpc.optim[:linconstrainteq]
+    JuMP.set_normalized_rhs(linconeq, mpc.con.beq)
+    return nothing
+end
+linconstrainteq!(::PredictiveController, ::SimModel, ::SingleShooting) = nothing
+linconstrainteq!(::PredictiveController, ::SimModel, ::MultipleShooting) = nothing
+
 @doc raw"""
     set_warmstart!(mpc::PredictiveController, transcription::SingleShooting, Z̃var) -> Z̃0
 
