Merge pull request #51 from JuliaControl/cleanup

franckgaga · web-flow · commit 3228f293238a · 2024-04-11T18:31:16.000-04:00
Namespace, `NonLinMPC` and `MovingHorizonEsitimator` cleanup
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ModelPredictiveControl"
 uuid = "61f9bdb8-6ae4-484a-811f-bbf86720c31c"
 authors = ["Francis Gagnon"]
-version = "0.20.1"
+version = "0.20.2"
 
 [deps]
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"
diff --git a/src/controller/linmpc.jl b/src/controller/linmpc.jl
@@ -76,7 +76,7 @@ struct LinMPC{
             d0, D̂0, D̂E,
             Ŷop, Dop,
         )
-        init_optimization!(mpc, optim)
+        init_optimization!(mpc, model, optim)
         return mpc
     end
 end
@@ -235,16 +235,16 @@ function LinMPC(
 end
 
 """
-    init_optimization!(mpc::LinMPC, optim::JuMP.GenericModel)
+    init_optimization!(mpc::LinMPC, model::LinModel, optim)
 
 Init the quadratic optimization for [`LinMPC`](@ref) controllers.
 """
-function init_optimization!(mpc::LinMPC, optim::JuMP.GenericModel)
+function init_optimization!(mpc::LinMPC, model::LinModel, optim)
     # --- variables and linear constraints ---
     con = mpc.con
     nΔŨ = length(mpc.ΔŨ)
     JuMP.set_silent(optim)
-    limit_solve_time(mpc.optim, mpc.estim.model.Ts)
+    limit_solve_time(mpc.optim, model.Ts)
     @variable(optim, ΔŨvar[1:nΔŨ])
     A = con.A[con.i_b, :]
     b = con.b[con.i_b]
diff --git a/src/controller/nonlinmpc.jl b/src/controller/nonlinmpc.jl
@@ -77,7 +77,7 @@ struct NonLinMPC{
             d0, D̂0, D̂E,
             Ŷop, Dop,
         )
-        init_optimization!(mpc, optim)
+        init_optimization!(mpc, model, optim)
         return mpc
     end
 end
@@ -277,11 +277,11 @@ function addinfo!(info, mpc::NonLinMPC)
 end
 
 """
-    init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel)
+    init_optimization!(mpc::NonLinMPC, model::SimModel, optim)
 
 Init the nonlinear optimization for [`NonLinMPC`](@ref) controllers.
 """
-function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where JNT<:Real
+function init_optimization!(mpc::NonLinMPC, model::SimModel, optim)
     # --- variables and linear constraints ---
     C, con = mpc.C, mpc.con
     nΔŨ = length(mpc.ΔŨ)
@@ -300,65 +300,11 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
             JuMP.set_attribute(optim, "nlp_scaling_max_gradient", 10.0/C)
         end
     end
-    model = mpc.estim.model
-    nu, ny, nx̂, Hp, ng = model.nu, model.ny, mpc.estim.nx̂, mpc.Hp, length(con.i_g)
-    # inspired from https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/tips_and_tricks/#User-defined-operators-with-vector-outputs
-    Jfunc, gfunc = let mpc=mpc, model=model, ng=ng, nΔŨ=nΔŨ, nŶ=Hp*ny, nx̂=nx̂, nu=nu, nU=Hp*nu
-        Nc = nΔŨ + 3
-        last_ΔŨtup_float, last_ΔŨtup_dual = nothing, nothing
-        ΔŨ_cache::DiffCache{Vector{JNT}, Vector{JNT}}    = DiffCache(zeros(JNT, nΔŨ), Nc)
-        Ŷ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ), Nc)
-        U_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU), Nc)
-        g_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, ng), Nc)
-        x̂_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nx̂), Nc)
-        x̂next_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), Nc)
-        u_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu), Nc)
-        û_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu), Nc)
-        Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ), Nc)
-        Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU), Nc)
-        function Jfunc(ΔŨtup::T...)::T where {T <: Real}
-            ΔŨ1 = ΔŨtup[begin]
-            if T == JNT
-                last_ΔŨtup_float = ΔŨtup
-            else
-                last_ΔŨtup_dual  = ΔŨtup
-            end
-            ΔŨ, Ŷ = get_tmp(ΔŨ_cache, ΔŨ1), get_tmp(Ŷ_cache, ΔŨ1)
-            x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
-            u, û = get_tmp(u_cache, ΔŨ1), get_tmp(û_cache, ΔŨ1)
-            ΔŨ .= ΔŨtup
-            Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, û, mpc, model, ΔŨ)
-            g = get_tmp(g_cache, ΔŨ1)
-            g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
-            U, Ȳ, Ū = get_tmp(U_cache, ΔŨ1), get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
-            return obj_nonlinprog!(U, Ȳ, Ū, mpc, model, Ŷ, ΔŨ)::T
-        end
-        function gfunc_i(i, ΔŨtup::NTuple{N, T})::T where {N, T <:Real}
-            ΔŨ1 = ΔŨtup[begin]
-            g = get_tmp(g_cache, ΔŨ1)
-            if T == JNT
-                isnewvalue = (ΔŨtup !== last_ΔŨtup_float)
-                isnewvalue && (last_ΔŨtup_float = ΔŨtup)
-            else
-                isnewvalue = (ΔŨtup !== last_ΔŨtup_dual)
-                isnewvalue && (last_ΔŨtup_dual = ΔŨtup)
-            end
-            if isnewvalue
-                ΔŨ, Ŷ = get_tmp(ΔŨ_cache, ΔŨ1), get_tmp(Ŷ_cache, ΔŨ1)
-                x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
-                u, û = get_tmp(u_cache, ΔŨ1), get_tmp(û_cache, ΔŨ1)
-                ΔŨ .= ΔŨtup
-                Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, û, mpc, model, ΔŨ)
-                g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
-            end
-            return g[i]::T
-        end
-        gfunc = [(ΔŨ...) -> gfunc_i(i, ΔŨ) for i in 1:ng]
-        (Jfunc, gfunc)
-    end
+    Jfunc, gfunc = get_optim_functions(mpc, mpc.optim)
     register(optim, :Jfunc, nΔŨ, Jfunc, autodiff=true)
     @NLobjective(optim, Min, Jfunc(ΔŨvar...))
-    if ng ≠ 0
+    ny, nx̂, Hp = model.ny, mpc.estim.nx̂, mpc.Hp
+    if length(con.i_g) ≠ 0
         for i in eachindex(con.Ymin)
             sym = Symbol("g_Ymin_$i")
             register(optim, sym, nΔŨ, gfunc[i], autodiff=true)
@@ -382,6 +328,62 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
     return nothing
 end
 
+"""
+    get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel) -> Jfunc, gfunc
+
+Get the objective `Jfunc` and constraints `gfunc` functions for [`NonLinMPC`](@ref).
+
+Inspired from: [User-defined operators with vector outputs](https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/tips_and_tricks/#User-defined-operators-with-vector-outputs)
+"""
+function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT<:Real
+    model = mpc.estim.model
+    nu, ny, nx̂, Hp = model.nu, model.ny, mpc.estim.nx̂, mpc.Hp
+    ng, nΔŨ, nU, nŶ = length(mpc.con.i_g), length(mpc.ΔŨ), Hp*nu, Hp*ny
+    Nc = nΔŨ + 3
+    ΔŨ_cache::DiffCache{Vector{JNT}, Vector{JNT}}    = DiffCache(zeros(JNT, nΔŨ), Nc)
+    Ŷ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ), Nc)
+    U_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU), Nc)
+    g_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, ng), Nc)
+    x̂_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nx̂), Nc)
+    x̂next_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), Nc)
+    u_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu), Nc)
+    û_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu), Nc)
+    Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶ), Nc)
+    Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nU), Nc)
+    function Jfunc(ΔŨtup::T...) where T<:Real
+        ΔŨ1 = ΔŨtup[begin]
+        ΔŨ, g = get_tmp(ΔŨ_cache, ΔŨ1), get_tmp(g_cache, ΔŨ1) 
+        for i in eachindex(ΔŨtup)
+            ΔŨ[i] = ΔŨtup[i] # ΔŨ .= ΔŨtup seems to produce a type instability
+        end 
+        Ŷ = get_tmp(Ŷ_cache, ΔŨ1)
+        x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
+        u, û = get_tmp(u_cache, ΔŨ1), get_tmp(û_cache, ΔŨ1)
+        Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, û, mpc, model, ΔŨ)
+        g = get_tmp(g_cache, ΔŨ1)
+        g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
+        U, Ȳ, Ū = get_tmp(U_cache, ΔŨ1), get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
+        return obj_nonlinprog!(U, Ȳ, Ū, mpc, model, Ŷ, ΔŨ)::T
+    end
+    function gfunc_i(i, ΔŨtup::NTuple{N, T}) where {N, T<:Real}
+        ΔŨ1 = ΔŨtup[begin]
+        ΔŨ, g = get_tmp(ΔŨ_cache, ΔŨ1), get_tmp(g_cache, ΔŨ1)     
+        if any(new !== old for (new, old) in zip(ΔŨtup, ΔŨ)) # new ΔŨtup, update predictions:
+            for i in eachindex(ΔŨtup)
+                ΔŨ[i] = ΔŨtup[i] # ΔŨ .= ΔŨtup seems to produce a type instability
+            end
+            Ŷ = get_tmp(Ŷ_cache, ΔŨ1)
+            x̂, x̂next = get_tmp(x̂_cache, ΔŨ1), get_tmp(x̂next_cache, ΔŨ1)
+            u, û = get_tmp(u_cache, ΔŨ1), get_tmp(û_cache, ΔŨ1)
+            Ŷ, x̂end = predict!(Ŷ, x̂, x̂next, u, û, mpc, model, ΔŨ)
+            g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
+        end
+        return g[i]::T
+    end
+    gfunc = [(ΔŨ...) -> gfunc_i(i, ΔŨ) for i in 1:ng]
+    return Jfunc, gfunc
+end
+
 "Set the nonlinear constraints on the output predictions `Ŷ` and terminal states `x̂end`."
 function setnonlincon!(mpc::NonLinMPC, ::NonLinModel)
     optim = mpc.optim
diff --git a/src/estimator/kalman.jl b/src/estimator/kalman.jl
@@ -368,9 +368,12 @@ struct UnscentedKalmanFilter{NT<:Real, SM<:SimModel} <: StateEstimator{NT}
     B̂d::Matrix{NT}
     D̂d::Matrix{NT}
     P̂0::Hermitian{NT, Matrix{NT}}
-    Q̂::Hermitian{NT, Matrix{NT}}
-    R̂::Hermitian{NT, Matrix{NT}}
+    Q̂ ::Hermitian{NT, Matrix{NT}}
+    R̂ ::Hermitian{NT, Matrix{NT}}
     K̂::Matrix{NT}
+    M̂::Hermitian{NT, Matrix{NT}}
+    X̂ ::Matrix{NT}
+    Ŷm::Matrix{NT}
     nσ::Int 
     γ::NT
     m̂::Vector{NT}
@@ -391,14 +394,16 @@ struct UnscentedKalmanFilter{NT<:Real, SM<:SimModel} <: StateEstimator{NT}
         P̂0 = Hermitian(P̂0, :L)
         P̂ = copy(P̂0)
         K̂ = zeros(NT, nx̂, nym)
+        M̂ = Hermitian(zeros(NT, nym, nym), :L)
+        X̂, Ŷm = zeros(NT, nx̂, nσ), zeros(NT, nym, nσ)
         return new{NT, SM}(
             model,
             lastu0, x̂, P̂, 
             i_ym, nx̂, nym, nyu, nxs, 
             As, Cs_u, Cs_y, nint_u, nint_ym,
             Â, B̂u, Ĉ, B̂d, D̂d,
             P̂0, Q̂, R̂,
-            K̂,
+            K̂, M̂, X̂, Ŷm,
             nσ, γ, m̂, Ŝ
         )
     end
@@ -565,18 +570,18 @@ noise, respectively.
      ISBN9780470045343.
 """
 function update_estimate!(estim::UnscentedKalmanFilter{NT}, u, ym, d) where NT<:Real
-    x̂, P̂, Q̂, R̂, K̂ = estim.x̂, estim.P̂, estim.Q̂, estim.R̂, estim.K̂
+    x̂, P̂, Q̂, R̂, K̂, M̂ = estim.x̂, estim.P̂, estim.Q̂, estim.R̂, estim.K̂, estim.M̂
+    X̂, Ŷm = estim.X̂, estim.Ŷm
     nym, nx̂, nσ = estim.nym, estim.nx̂, estim.nσ
     γ, m̂, Ŝ = estim.γ, estim.m̂, estim.Ŝ
     # --- initialize matrices ---
-    X̂, X̂_next = Matrix{NT}(undef, nx̂, nσ), Matrix{NT}(undef, nx̂, nσ)
+    X̂_next = Matrix{NT}(undef, nx̂, nσ)
     û  = Vector{NT}(undef, estim.model.nu)
     ŷm = Vector{NT}(undef, nym)
     ŷ  = Vector{NT}(undef, estim.model.ny)
-    Ŷm = Matrix{NT}(undef, nym, nσ)
     sqrt_P̂ = LowerTriangular{NT, Matrix{NT}}(Matrix{NT}(undef, nx̂, nx̂))
     # --- correction step ---
-    sqrt_P̂.data .= cholesky(P̂).L
+    sqrt_P̂ .= cholesky(P̂).L
     γ_sqrt_P̂ = lmul!(γ, sqrt_P̂)
     X̂ .= x̂
     X̂[:, 2:nx̂+1]   .+= γ_sqrt_P̂
@@ -589,7 +594,7 @@ function update_estimate!(estim::UnscentedKalmanFilter{NT}, u, ym, d) where NT<:
     X̄, Ȳm = X̂, Ŷm
     X̄  .= X̂  .- x̂
     Ȳm .= Ŷm .- ŷm
-    M̂ = Hermitian(Ȳm * Ŝ * Ȳm' + R̂, :L)
+    M̂ .= Hermitian(Ȳm * Ŝ * Ȳm' .+ R̂)
     mul!(K̂, X̄, lmul!(Ŝ, Ȳm'))
     rdiv!(K̂, cholesky(M̂))
     v̂ = ŷm
@@ -598,19 +603,18 @@ function update_estimate!(estim::UnscentedKalmanFilter{NT}, u, ym, d) where NT<:
     P̂_cor = P̂ - Hermitian(K̂ * M̂ * K̂', :L)
     # --- prediction step ---
     X̂_cor, sqrt_P̂_cor = X̂, sqrt_P̂
-    sqrt_P̂_cor.data .= cholesky(P̂_cor).L
+    sqrt_P̂_cor .= cholesky(P̂_cor).L
     γ_sqrt_P̂_cor = lmul!(γ, sqrt_P̂_cor)
     X̂_cor .= x̂_cor
     X̂_cor[:, 2:nx̂+1]   .+= γ_sqrt_P̂_cor
     X̂_cor[:, nx̂+2:end] .-= γ_sqrt_P̂_cor
-    X̂_next = similar(X̂_cor)
     for j in axes(X̂_next, 2)
         @views f̂!(X̂_next[:, j], û, estim, estim.model, X̂_cor[:, j], u, d)
     end
     x̂_next = mul!(x̂, X̂_next, m̂)
     X̄_next = X̂_next
     X̄_next .= X̂_next .- x̂_next
-    P̂.data .= X̄_next * Ŝ * X̄_next' .+ Q̂ # .data is necessary for Hermitians
+    P̂ .= Hermitian(X̄_next * Ŝ * X̄_next' .+ Q̂)
     return nothing
 end
 
@@ -781,7 +785,7 @@ end
 function init_estimate_cov!(
     estim::Union{KalmanFilter, UnscentedKalmanFilter, ExtendedKalmanFilter}, _ , _ , _
 ) 
-    estim.P̂.data .= estim.P̂0 # .data is necessary for Hermitians
+    estim.P̂ .= estim.P̂0
     return nothing
 end
 
@@ -818,6 +822,7 @@ function update_estimate_kf!(estim::StateEstimator{NT}, u, ym, d, Â, Ĉm, P̂
     Q̂, R̂, M̂, K̂ = estim.Q̂, estim.R̂, estim.M̂, estim.K̂
     nx̂, nu, ny = estim.nx̂, estim.model.nu, estim.model.ny
     x̂next, û, ŷ = Vector{NT}(undef, nx̂), Vector{NT}(undef, nu), Vector{NT}(undef, ny)
+    P̂next = similar(P̂)
     mul!(M̂, P̂.data, Ĉm') # the ".data" weirdly removes a type instability in mul!
     rdiv!(M̂, cholesky!(Hermitian(Ĉm * P̂ * Ĉm' .+ R̂)))
     mul!(K̂, Â, M̂)
@@ -828,6 +833,6 @@ function update_estimate_kf!(estim::StateEstimator{NT}, u, ym, d, Â, Ĉm, P̂
     f̂!(x̂next, û, estim, estim.model, x̂, u, d)
     mul!(x̂next, K̂, v̂, 1, 1)
     estim.x̂ .= x̂next
-    P̂.data .= Â * (P̂ .- M̂ * Ĉm * P̂) * Â' .+ Q̂ # .data is necessary for Hermitians
+    P̂ .= Hermitian(Â * (P̂ .- M̂ * Ĉm * P̂) * Â' .+ Q̂)
     return nothing
 end
diff --git a/src/estimator/mhe/construct.jl b/src/estimator/mhe/construct.jl
diff --git a/src/estimator/mhe/execute.jl b/src/estimator/mhe/execute.jl
