reduce allocation NonLinMPC

franckgaga · franckgaga · commit 8874af3b5361 · 2024-01-28T17:28:17.000-05:00
diff --git a/src/controller/execute.jl b/src/controller/execute.jl
@@ -104,11 +104,11 @@ 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̂, 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̃)
+    Ŷ, x̂, u     = similar(mpc.Ŷop), similar(mpc.estim.x̂), similar(mpc.estim.lastu0)
+    Ŷ, x̂end     = predict!(Ŷ, x̂, u, mpc, mpc.estim.model, mpc.ΔŨ)
     U           = mpc.S̃*mpc.ΔŨ + mpc.T*(mpc.estim.lastu0 + mpc.estim.model.uop)
     Ȳ, Ū        = similar(Ŷ), similar(U)
-    J           = obj_nonlinprog!(Ȳ, Ū, mpc, mpc.estim.model, Ŷ, mpc.ΔŨ)
+    J           = obj_nonlinprog!(Ȳ, Ū, u, mpc, mpc.estim.model, Ŷ, mpc.ΔŨ)
     info[:ΔU]   = mpc.ΔŨ[1:mpc.Hc*mpc.estim.model.nu]
     info[:ϵ]    = isinf(mpc.C) ? NaN : mpc.ΔŨ[end]
     info[:J]    = J
@@ -284,20 +284,20 @@ function predict!(
 end
 
 @doc raw"""
-    predict!(Ŷ, x̂, u0, mpc::PredictiveController, model::SimModel, ΔŨ) -> Ŷ, x̂end
+    predict!(Ŷ, x̂, u, mpc::PredictiveController, model::SimModel, ΔŨ) -> Ŷ, x̂end
 
 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)``.
+The method mutates `Ŷ`, `x̂` and `u` arguments.
 """
 function predict!(
-    Ŷ, x̂, u0, mpc::PredictiveController, model::SimModel, ΔŨ::Vector{NT}
+    Ŷ, x̂, u, 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
+    u0 = u
     x̂  .= mpc.estim.x̂
     u0 .= mpc.estim.lastu0
-    d0 = @views mpc.d0[1:end]
+    d0  = @views mpc.d0[1:end]
     for j=1:Hp
         if j ≤ Hc
             u0 .+= @views ΔŨ[(1 + nu*(j-1)):(nu*j)]
@@ -312,21 +312,23 @@ function predict!(
 end
 
 """
-    obj_nonlinprog!(_ , _ , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ)
+    obj_nonlinprog!( _ , _ , u , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ)
 
 Nonlinear programming objective function when `model` is a [`LinModel`](@ref).
 
 The function is called by the nonlinear optimizer of [`NonLinMPC`](@ref) controllers. It can
 also be called on any [`PredictiveController`](@ref)s to evaluate the objective function `J`
-at specific input increments `ΔŨ` and predictions `Ŷ` values. This method does not mutate
-its argument.
+at specific input increments `ΔŨ` and predictions `Ŷ` values. This method mutate `u` 
+argument.
 """
 function obj_nonlinprog!(
-    _ , _ , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ::Vector{NT}
+    _ , _ , u , mpc::PredictiveController, model::LinModel, Ŷ, ΔŨ::Vector{NT}
 ) where {NT<:Real}
     J = obj_quadprog(ΔŨ, mpc.H̃, mpc.q̃) + mpc.p[]
     if !iszero(mpc.E)
-        U = mpc.S̃*ΔŨ + mpc.T*(mpc.estim.lastu0 + model.uop)
+        lastu  = u
+        lastu .= mpc.estim.lastu0 .+ model.uop
+        U = mpc.S̃*ΔŨ + mpc.T*lastu
         UE = [U; U[(end - model.nu + 1):end]]
         ŶE = [mpc.ŷ; Ŷ]
         J += mpc.E*mpc.JE(UE, ŶE, mpc.D̂E)
@@ -335,14 +337,14 @@ function obj_nonlinprog!(
 end
 
 """
-    obj_nonlinprog!(Ȳ, Ū. mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ)
+    obj_nonlinprog!(Ȳ, Ū, u, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ)
 
 Nonlinear programming objective function when `model` is not a [`LinModel`](@ref). The
 function `dot(x, A, x)` is a performant way of calculating `x'*A*x`. This method mutates
 `Ȳ` and `Ū` vector arguments (output and input setpoint tracking error, respectively).
 """
 function obj_nonlinprog!(
-    Ȳ, Ū, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ::Vector{NT}
+    Ȳ, Ū, u, mpc::PredictiveController, model::SimModel, Ŷ, ΔŨ::Vector{NT}
 ) where {NT<:Real}
     # --- output setpoint tracking term ---
     Ȳ  .= mpc.R̂y .- Ŷ
@@ -351,7 +353,9 @@ function obj_nonlinprog!(
     JΔŨ = dot(ΔŨ, mpc.Ñ_Hc, ΔŨ)
     # --- input over prediction horizon ---
     if !mpc.noR̂u || !iszero(mpc.E)
-        U = mpc.S̃*ΔŨ + mpc.T*(mpc.estim.lastu0 + model.uop)
+        lastu  = u
+        lastu .= mpc.estim.lastu0 .+ model.uop
+        U = mpc.S̃*ΔŨ + mpc.T*lastu
     end
     # --- input setpoint tracking term ---
     if !mpc.noR̂u
diff --git a/src/controller/nonlinmpc.jl b/src/controller/nonlinmpc.jl
@@ -301,61 +301,58 @@ function init_optimization!(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where
     # 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
         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)
-        u0_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nu), nΔŨ + 3)
-        Ȳ_cache ::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŶ), nΔŨ + 3)
-        Ū_cache ::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nU), 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)
+        u_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nu), nΔŨ + 3)
+        Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŶ), nΔŨ + 3)
+        Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nU), nΔŨ + 3)
         function Jfunc(ΔŨtup::JNT...)
-            ΔŨtud1 = ΔŨtup[begin]
-            Ŷ = get_tmp(Ŷ_cache, ΔŨtud1)
+            ΔŨ1 = ΔŨtup[begin]
+            Ŷ, u = get_tmp(Ŷ_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1)
             ΔŨ = collect(ΔŨtup)
             if ΔŨtup !== last_ΔŨtup_float
-                x̂, u0  = get_tmp(x̂_cache, ΔŨtud1), get_tmp(u0_cache, ΔŨtud1)
-                g  = get_tmp(g_cache, ΔŨtup[1])
-                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
+                x̂, g = get_tmp(x̂_cache, ΔŨ1), get_tmp(g_cache, ΔŨ1)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_float = ΔŨtup
             end
-            Ȳ, Ū = get_tmp(Ȳ_cache, ΔŨtud1), get_tmp(Ū_cache, ΔŨtud1)
-            return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ŷ, ΔŨ)
+            Ȳ, Ū = get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
+            return obj_nonlinprog!(Ȳ, Ū, u, mpc, model, Ŷ, ΔŨ)
         end
         function Jfunc(ΔŨtup::ForwardDiff.Dual...)
-            ΔŨtud1 = ΔŨtup[begin]
-            Ŷ = get_tmp(Ŷ_cache, ΔŨtud1)
+            ΔŨ1 = ΔŨtup[begin]
+            Ŷ, u = get_tmp(Ŷ_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1)
             ΔŨ = collect(ΔŨtup)
             if ΔŨtup !== last_ΔŨtup_dual
-                x̂, u0  = get_tmp(x̂_cache, ΔŨtud1), get_tmp(u0_cache, ΔŨtud1)
-                g = get_tmp(g_cache, ΔŨtud1)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
+                x̂, g = get_tmp(x̂_cache, ΔŨ1), get_tmp(g_cache, ΔŨ1)
+                g = get_tmp(g_cache, ΔŨ1)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_dual = ΔŨtup
             end
-            Ȳ, Ū = get_tmp(Ȳ_cache, ΔŨtud1), get_tmp(Ū_cache, ΔŨtud1)
-            return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ŷ, ΔŨ)
+            Ȳ, Ū = get_tmp(Ȳ_cache, ΔŨ1), get_tmp(Ū_cache, ΔŨ1)
+            return obj_nonlinprog!(Ȳ, Ū, u, mpc, model, Ŷ, ΔŨ)
         end
         function gfunc_i(i, ΔŨtup::NTuple{N, JNT}) where N
-            ΔŨtud1 = ΔŨtup[begin]
-            g = get_tmp(g_cache, ΔŨtud1)
+            ΔŨ1 = ΔŨtup[begin]
+            g = get_tmp(g_cache, ΔŨ1)
             if ΔŨtup !== last_ΔŨtup_float
-                x̂, u0  = get_tmp(x̂_cache, ΔŨtud1), get_tmp(u0_cache, ΔŨtud1)
-                Ŷ = get_tmp(Ŷ_cache, ΔŨtud1)
+                Ŷ, u, x̂ = get_tmp(Ŷ_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1), get_tmp(x̂_cache, ΔŨ1)
                 ΔŨ = collect(ΔŨtup)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_float = ΔŨtup
             end
             return g[i]
         end 
         function gfunc_i(i, ΔŨtup::NTuple{N, ForwardDiff.Dual}) where N
-            ΔŨtud1 = ΔŨtup[begin]
-            g = get_tmp(g_cache, ΔŨtud1)
+            ΔŨ1 = ΔŨtup[begin]
+            g = get_tmp(g_cache, ΔŨ1)
             if ΔŨtup !== last_ΔŨtup_dual
-                x̂, u0  = get_tmp(x̂_cache, ΔŨtud1), get_tmp(u0_cache, ΔŨtud1)
-                Ŷ = get_tmp(Ŷ_cache, ΔŨtud1)
+                Ŷ, u, x̂ = get_tmp(Ŷ_cache, ΔŨ1), get_tmp(u_cache, ΔŨ1), get_tmp(x̂_cache, ΔŨ1)
                 ΔŨ = collect(ΔŨtup)
-                Ŷ, x̂end = predict!(Ŷ, x̂, u0, mpc, model, ΔŨ)
+                Ŷ, x̂end = predict!(Ŷ, x̂, u, mpc, model, ΔŨ)
                 g = con_nonlinprog!(g, mpc, model, x̂end, Ŷ, ΔŨ)
                 last_ΔŨtup_dual = ΔŨtup
             end
