reduce allocation for NonLinMPC

franckgaga · franckgaga · commit 4462cd774378 · 2024-01-25T17:44:39.000-05:00
diff --git a/src/controller/execute.jl b/src/controller/execute.jl
@@ -104,8 +104,8 @@ julia> round.(getinfo(mpc)[:Ŷ], digits=3)
 """
 function getinfo(mpc::PredictiveController{NT}) where NT<:Real
     info = Dict{Symbol, Union{JuMP._SolutionSummary, Vector{NT}, NT}}()
-    Ŷ, x̂    = similar(mpc.Ŷop), similar(mpc.estim.x̂)
-    Ŷ, x̂end = predict!(Ŷ, x̂, mpc, mpc.estim.model, mpc.ΔŨ)
+    Ŷ, x̂, u0    = similar(mpc.Ŷop), similar(mpc.estim.x̂), similar(mpc.estim.lastu0)
+    Ŷ, x̂end     = predict!(Ŷ, x̂, u0, mpc, mpc.estim.model, mpc.ΔŨ)
     info[:ΔU]   = mpc.ΔŨ[1:mpc.Hc*mpc.estim.model.nu]
     info[:ϵ]    = isinf(mpc.C) ? NaN : mpc.ΔŨ[end]
     info[:J]    = obj_nonlinprog(mpc, mpc.estim.model, Ŷ, mpc.ΔŨ)
@@ -263,14 +263,16 @@ function linconstraint!(mpc::PredictiveController, model::SimModel)
 end
 
 @doc raw"""
-    predict!(Ŷ, x̂, mpc::PredictiveController, model::LinModel, ΔŨ) -> Ŷ, x̂end
+    predict!(Ŷ, x̂, _ , mpc::PredictiveController, model::LinModel, ΔŨ) -> Ŷ, x̂end
 
 Compute the predictions `Ŷ` and terminal states `x̂end` if model is a [`LinModel`](@ref).
 
 The method mutates `Ŷ` and `x̂` vector arguments. The `x̂end` vector is used for
 the terminal constraints applied on ``\mathbf{x̂}_{k-1}(k+H_p)``.
 """
-function predict!(Ŷ, x̂, mpc::PredictiveController, ::LinModel, ΔŨ::Vector{NT}) where {NT<:Real}
+function predict!(
+    Ŷ, x̂, _ , mpc::PredictiveController, ::LinModel, ΔŨ::Vector{NT}
+) where {NT<:Real}
     # in-place operations to reduce allocations :
     Ŷ .= mul!(Ŷ, mpc.Ẽ, ΔŨ) .+ mpc.F
     x̂ .= mul!(x̂, mpc.con.ẽx̂, ΔŨ) .+ mpc.con.fx̂
@@ -279,14 +281,19 @@ function predict!(Ŷ, x̂, mpc::PredictiveController, ::LinModel, ΔŨ::Vector
 end
 
 @doc raw"""
-    predict!(Ŷ, x̂, mpc::PredictiveController, model::SimModel, ΔŨ) -> Ŷ, x̂end
+    predict!(Ŷ, x̂, u0, mpc::PredictiveController, model::SimModel, ΔŨ) -> Ŷ, x̂end
 
-Compute both vectors if `model` is not a [`LinModel`](@ref).
+Compute both vectors if `model` is not a [`LinModel`](@ref). 
+    
+The method mutates `Ŷ`, `x̂` and `u0` arguments. The latter is the manipulated input without
+the operating points ``\mathbf{u_0}(k) = \mathbf{u}(k) - \mathbf{u_{op}}(k)``.
 """
-function predict!(Ŷ, x̂, mpc::PredictiveController, model::SimModel, ΔŨ::Vector{NT}) where {NT<:Real}
+function predict!(
+    Ŷ, x̂, u0, mpc::PredictiveController, model::SimModel, ΔŨ::Vector{NT}
+) where {NT<:Real}
     nu, ny, nd, Hp, Hc = model.nu, model.ny, model.nd, mpc.Hp, mpc.Hc
-    x̂ .= mpc.estim.x̂
-    u0::Vector{NT} = copy(mpc.estim.lastu0)
+    x̂  .= mpc.estim.x̂
+    u0 .= mpc.estim.lastu0
     d0 = @views mpc.d0[1:end]
     for j=1:Hp
         if j ≤ Hc
diff --git a/src/controller/nonlinmpc.jl b/src/controller/nonlinmpc.jl
@@ -297,20 +297,22 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
         end
     end
     model = mpc.estim.model
-    ny, nx̂, Hp, ng = model.ny, mpc.estim.nx̂, mpc.Hp, length(con.i_g)
+    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̂
+    Jfunc, gfunc = let mpc=mpc, model=model, ng=ng, nΔŨ=nΔŨ , nŶ=Hp*ny, nx̂=nx̂, nu=nu
         last_ΔŨtup_float, last_ΔŨtup_dual = nothing, nothing
-        Ŷ_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŶ), nΔŨ + 3)
-        g_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, ng), nΔŨ + 3)
-        x̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), nΔŨ + 3)
+        Ŷ_cache ::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŶ), nΔŨ + 3)
+        g_cache ::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, ng), nΔŨ + 3)
+        x̂_cache ::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), nΔŨ + 3)
+        u0_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nu), nΔŨ + 3)
         function Jfunc(ΔŨtup::JNT...)
             Ŷ = get_tmp(Ŷ_cache, ΔŨtup[1])
             ΔŨ = collect(ΔŨtup)
             if ΔŨtup !== last_ΔŨtup_float
-                x̂ = get_tmp(x̂_cache, ΔŨtup[1])
-                g = get_tmp(g_cache, ΔŨtup[1])
-                Ŷ, x̂end = predict!(Ŷ, x̂, mpc, model, ΔŨ)
+                x̂  = get_tmp(x̂_cache, ΔŨtup[1])
+                u0 = get_tmp(u0_cache, ΔŨtup[1])
+                g  = get_tmp(g_cache, ΔŨtup[1])
+                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_float = ΔŨtup
             end
@@ -321,8 +323,9 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
             ΔŨ = collect(ΔŨtup)
             if ΔŨtup !== last_ΔŨtup_dual
                 x̂ = get_tmp(x̂_cache, ΔŨtup[1])
+                u0 = get_tmp(u0_cache, ΔŨtup[1])
                 g = get_tmp(g_cache, ΔŨtup[1])
-                Ŷ, x̂end = predict!(Ŷ, x̂, mpc, model, ΔŨ)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_dual = ΔŨtup
             end
@@ -332,9 +335,10 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
             g = get_tmp(g_cache, ΔŨtup[1])
             if ΔŨtup !== last_ΔŨtup_float
                 x̂ = get_tmp(x̂_cache, ΔŨtup[1])
+                u0 = get_tmp(u0_cache, ΔŨtup[1])
                 Ŷ = get_tmp(Ŷ_cache, ΔŨtup[1])
                 ΔŨ = collect(ΔŨtup)
-                Ŷ, x̂end = predict!(Ŷ, x̂, mpc, model, ΔŨ)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_float = ΔŨtup
             end
@@ -344,9 +348,10 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
             g = get_tmp(g_cache, ΔŨtup[1])
             if ΔŨtup !== last_ΔŨtup_dual
                 x̂ = get_tmp(x̂_cache, ΔŨtup[1])
+                u0 = get_tmp(u0_cache, ΔŨtup[1])
                 Ŷ = get_tmp(Ŷ_cache, ΔŨtup[1])
                 ΔŨ = collect(ΔŨtup)
-                Ŷ, x̂end = predict!(Ŷ, x̂, mpc, model, ΔŨ)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_dual = ΔŨtup
             end
