Merge pull request #165 from JuliaControl/new_jacbuffer

added: `linearize!` is allocation-free once again

Author: franckgaga

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelPredictiveControl"
 uuid = "61f9bdb8-6ae4-484a-811f-bbf86720c31c"
 authors = ["Francis Gagnon"]
-version = "1.3.4"
+version = "1.3.5"
 
 [deps]
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"

src/general.jl

@@ -6,6 +6,29 @@ const DEFAULT_LWT = 0.0
 const DEFAULT_CWT = 1e5
 const DEFAULT_EWT = 0.0
 
+"Abstract type for all differentiation buffers."
+abstract type DifferentiationBuffer end
+
+"Struct with both function and configuration for ForwardDiff differentiation."
+struct JacobianBuffer{FT<:Function, CT<:ForwardDiff.JacobianConfig} <: DifferentiationBuffer
+    f!::FT
+    config::CT
+end
+
+function Base.show(io::IO, buffer::DifferentiationBuffer) 
+    return print(io, "DifferentiationBuffer with a $(typeof(buffer.config).name.name)")
+end
+
+"Create a JacobianBuffer with function `f!`, output `y` and input `x`."
+JacobianBuffer(f!, y, x) = JacobianBuffer(f!, ForwardDiff.JacobianConfig(f!, y, x))
+
+"Compute in-place and return the Jacobian matrix of `buffer.f!` at `x`."
+function jacobian!(
+    A, buffer::JacobianBuffer, y, x
+)
+    return ForwardDiff.jacobian!(A, buffer.f!, y, x, buffer.config)
+end
+
 "Termination status that means 'no solution available'."
 const ERROR_STATUSES = (
     JuMP.INFEASIBLE, JuMP.DUAL_INFEASIBLE, JuMP.LOCALLY_INFEASIBLE, 

src/model/linearization.jl

@@ -1,9 +1,71 @@
-"""
-    LinModel(model::NonLinModel; x=model.x0+model.xop, u=model.uop, d=model.dop)
+"A linearization buffer for the [`linearize`](@ref) function."
+struct LinearizationBuffer{
+    NT <:Real,
+    JB_FUD <:JacobianBuffer,
+    JB_FXD <:JacobianBuffer,
+    JB_FXU <:JacobianBuffer,
+    JB_HD  <:JacobianBuffer,
+    JB_HX  <:JacobianBuffer
+}
+    x::Vector{NT}
+    u::Vector{NT}
+    d::Vector{NT}
+    buffer_f_at_u_d::JB_FUD
+    buffer_f_at_x_d::JB_FXD
+    buffer_f_at_x_u::JB_FXU
+    buffer_h_at_d  ::JB_HD
+    buffer_h_at_x  ::JB_HX
+    function LinearizationBuffer(
+        x::Vector{NT}, 
+        u::Vector{NT}, 
+        d::Vector{NT},
+        buffer_f_at_u_d::JB_FUD, 
+        buffer_f_at_x_d::JB_FXD, 
+        buffer_f_at_x_u::JB_FXU, 
+        buffer_h_at_d::JB_HD, 
+        buffer_h_at_x::JB_HX
+    ) where {NT<:Real, JB_FUD, JB_FXD, JB_FXU, JB_HD, JB_HX}
+        return new{NT, JB_FUD, JB_FXD, JB_FXU, JB_HD, JB_HX}(
+            x, u, d, 
+            buffer_f_at_u_d, 
+            buffer_f_at_x_d, 
+            buffer_f_at_x_u, 
+            buffer_h_at_d, 
+            buffer_h_at_x
+        )
+    end
+    
+end
 
-Call [`linearize(model; x, u, d)`](@ref) and return the resulting linear model.
-"""
-LinModel(model::NonLinModel; kwargs...) = linearize(model; kwargs...)
+Base.show(io::IO, buffer::LinearizationBuffer) = print(io, "LinearizationBuffer object")
+
+function LinearizationBuffer(NT, f!, h!, nu, nx, ny, nd, p)
+    x, u, d, f_at_u_d!, f_at_x_d!, f_at_x_u!, h_at_d!, h_at_x! = get_linearize_funcs(
+        NT, f!, h!, nu, nx, ny, nd, p
+    )
+    xnext, y = Vector{NT}(undef, nx), Vector{NT}(undef, ny) # TODO: to replace ?
+    return LinearizationBuffer(
+        x, u, d, 
+        JacobianBuffer(f_at_u_d!, xnext, x),
+        JacobianBuffer(f_at_x_d!, xnext, u),
+        JacobianBuffer(f_at_x_u!, xnext, d),
+        JacobianBuffer(h_at_d!, y, x),
+        JacobianBuffer(h_at_x!, y, d)
+    )
+end
+
+"Get the linearization functions for [`NonLinModel`](@ref) `f!` and `h!` functions."
+function get_linearize_funcs(NT, f!, h!, nu, nx, ny, nd, p)
+    x = Vector{NT}(undef, nx)
+    u = Vector{NT}(undef, nu)
+    d = Vector{NT}(undef, nd)
+    f_at_u_d!(xnext, x) = f!(xnext, x, u, d, p)
+    f_at_x_d!(xnext, u) = f!(xnext, x, u, d, p)
+    f_at_x_u!(xnext, d) = f!(xnext, x, u, d, p)
+    h_at_d!(y, x)       = h!(y, x, d, p)
+    h_at_x!(y, d)       = h!(y, x, d, p)
+    return x, u, d, f_at_u_d!, f_at_x_d!, f_at_x_u!, h_at_d!, h_at_x!
+end
 
 @doc raw"""
     linearize(model::SimModel; x=model.x0+model.xop, u=model.uop, d=model.dop) -> linmodel
@@ -92,8 +154,8 @@ end
 
 Linearize `model` and store the result in `linmodel` (in-place).
 
-The keyword arguments are identical to [`linearize`](@ref). The code allocates a small
-amount of memory to compute the Jacobians.
+The keyword arguments are identical to [`linearize`](@ref). The code is allocation-free if
+`model` simulations does not allocate.
 
 # Examples
 ```jldoctest
@@ -112,26 +174,22 @@ function linearize!(
     linmodel::LinModel{NT}, model::SimModel; x=model.x0+model.xop, u=model.uop, d=model.dop
 ) where NT<:Real
     nonlinmodel = model
-    buffer = nonlinmodel.buffer
+    buffer, linbuffer = nonlinmodel.buffer, nonlinmodel.linbuffer
     # --- remove the operating points of the nonlinear model (typically zeros) ---
     x0, u0, d0 = buffer.x, buffer.u, buffer.d
     u0 .= u .- nonlinmodel.uop
     d0 .= d .- nonlinmodel.dop
     x0 .= x .- nonlinmodel.xop
     # --- compute the Jacobians at linearization points ---
-    A::Matrix{NT}, Bu::Matrix{NT}, Bd::Matrix{NT}  = linmodel.A, linmodel.Bu, linmodel.Bd
-    C::Matrix{NT}, Dd::Matrix{NT} = linmodel.C, linmodel.Dd
     xnext0::Vector{NT}, y0::Vector{NT} = linmodel.buffer.x, linmodel.buffer.y
-    myf_x0!(xnext0, x0) = f!(xnext0, nonlinmodel, x0, u0, d0, model.p)
-    myf_u0!(xnext0, u0) = f!(xnext0, nonlinmodel, x0, u0, d0, model.p)
-    myf_d0!(xnext0, d0) = f!(xnext0, nonlinmodel, x0, u0, d0, model.p)
-    myh_x0!(y0, x0) = h!(y0, nonlinmodel, x0, d0, model.p)
-    myh_d0!(y0, d0) = h!(y0, nonlinmodel, x0, d0, model.p)
-    ForwardDiff.jacobian!(A,  myf_x0!, xnext0, x0)
-    ForwardDiff.jacobian!(Bu, myf_u0!, xnext0, u0)
-    ForwardDiff.jacobian!(Bd, myf_d0!, xnext0, d0)
-    ForwardDiff.jacobian!(C,  myh_x0!, y0, x0)
-    ForwardDiff.jacobian!(Dd, myh_d0!, y0, d0)
+    linbuffer.x .= x0
+    linbuffer.u .= u0
+    linbuffer.d .= d0
+    jacobian!(linmodel.A,  linbuffer.buffer_f_at_u_d, xnext0, x0)
+    jacobian!(linmodel.Bu, linbuffer.buffer_f_at_x_d, xnext0, u0)
+    jacobian!(linmodel.Bd, linbuffer.buffer_f_at_x_u, xnext0, d0)
+    jacobian!(linmodel.C,  linbuffer.buffer_h_at_d, y0, x0)
+    jacobian!(linmodel.Dd, linbuffer.buffer_h_at_x, y0, d0)
     # --- compute the nonlinear model output at operating points ---
     xnext0, y0 = linmodel.buffer.x, linmodel.buffer.y
     h!(y0, nonlinmodel, x0, d0, model.p)

src/model/nonlinmodel.jl

@@ -1,5 +1,5 @@
 struct NonLinModel{
-    NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver
+    NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver, LB<:LinearizationBuffer
 } <: SimModel{NT}
     x0::Vector{NT}
     f!::F
@@ -21,10 +21,11 @@ struct NonLinModel{
     yname::Vector{String}
     dname::Vector{String}
     xname::Vector{String}
+    linbuffer::LB
     buffer::SimModelBuffer{NT}
     function NonLinModel{NT}(
-        f!::F, h!::H, Ts, nu, nx, ny, nd, p::P, solver::DS
-    ) where {NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver}
+        f!::F, h!::H, Ts, nu, nx, ny, nd, p::P, solver::DS, linbuffer::LB
+    ) where {NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver, LB<:LinearizationBuffer}
         Ts > 0 || error("Sampling time Ts must be positive")
         uop = zeros(NT, nu)
         yop = zeros(NT, ny)
@@ -38,7 +39,7 @@ struct NonLinModel{
         x0 = zeros(NT, nx)
         t  = zeros(NT, 1)
         buffer = SimModelBuffer{NT}(nu, nx, ny, nd)
-        return new{NT, F, H, P, DS}(
+        return new{NT, F, H, P, DS, LB}(
             x0, 
             f!, h!,
             p,
@@ -47,6 +48,7 @@ struct NonLinModel{
             nu, nx, ny, nd, 
             uop, yop, dop, xop, fop,
             uname, yname, dname, xname,
+            linbuffer,
             buffer
         )
     end
@@ -145,7 +147,8 @@ function NonLinModel{NT}(
     isnothing(solver) && (solver=EmptySolver())
     f!, h! = get_mutating_functions(NT, f, h)
     f!, h! = get_solver_functions(NT, solver, f!, h!, Ts, nu, nx, ny, nd)
-    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver)
+    linbuffer = LinearizationBuffer(NT, f!, h!, nu, nx, ny, nd, p)
+    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver, linbuffer)
 end
 
 function NonLinModel(
@@ -226,6 +229,13 @@ end
 "Do nothing if `model` is a [`NonLinModel`](@ref)."
 steadystate!(::SimModel, _ , _ ) = nothing
 
+"""
+    LinModel(model::NonLinModel; x=model.x0+model.xop, u=model.uop, d=model.dop)
+
+Call [`linearize(model; x, u, d)`](@ref) and return the resulting linear model.
+"""
+LinModel(model::NonLinModel; kwargs...) = linearize(model; kwargs...)
+
 "Call `model.f!(xnext0, x0, u0, d0, p)` for [`NonLinModel`](@ref)."
 f!(xnext0, model::NonLinModel, x0, u0, d0, p) = model.f!(xnext0, x0, u0, d0, p)
 

src/sim_model.jl

@@ -29,13 +29,13 @@ struct SimModelBuffer{NT<:Real}
 end
 
 @doc raw"""
-    SimModelBuffer{NT}(nu::Int, nx::Int, ny::Int, nd::Int)
+    SimModelBuffer{NT}(nu::Int, nx::Int, ny::Int, nd::Int, linearization=nothing)
 
 Create a buffer for `SimModel` objects for inputs, states, outputs, and disturbances.
 
 The buffer is used to store intermediate results during simulation without allocating.
 """
-function SimModelBuffer{NT}(nu::Int, nx::Int, ny::Int, nd::Int) where {NT <: Real}
+function SimModelBuffer{NT}(nu::Int, nx::Int, ny::Int, nd::Int, ) where {NT<:Real}
     u = Vector{NT}(undef, nu)
     x = Vector{NT}(undef, nx)
     y = Vector{NT}(undef, ny)
@@ -371,5 +371,5 @@ to_mat(A::Real, dims...) = fill(A, dims)
 
 include("model/linmodel.jl")
 include("model/solver.jl")
-include("model/nonlinmodel.jl")
-include("model/linearization.jl")
+include("model/linearization.jl")
+include("model/nonlinmodel.jl")

test/1_test_sim_model.jl

@@ -285,6 +285,8 @@ end
     @test linmodel1.Bd ≈ linmodel2.Bd
     @test linmodel1.C  ≈ linmodel2.C
     @test linmodel1.Dd ≈ linmodel2.Dd 
+    @test repr(nonlinmodel1.linbuffer) == "LinearizationBuffer object"
+    @test repr(nonlinmodel1.linbuffer.buffer_f_at_u_d) == "DifferentiationBuffer with a JacobianConfig"
 
     f1!(ẋ, x, u, d, _) = (ẋ .= x.^5 + u.^4 + d.^3; nothing)
     h1!(y, x, d, _) = (y .= x.^2 + d; nothing)

test/4_test_plot_sim.jl

@@ -2,7 +2,8 @@
     using .SetupMPCtests, ControlSystemsBase, LinearAlgebra
     model = LinModel(sys, Ts, i_d=[3])
     res = sim!(model, 15)
-    display(res)
+
+    @test repr(res) == "Simulation results of LinModel with 15 time steps."
     @test isa(res.obj, LinModel)
     @test length(res.T_data) == 15
     @test res.U_data[:, 1] ≈ model.uop .+ 1