Merge pull request #160 from JuliaControl/type_stability_nonlinmodel

debug: revert allocation-free `linearize!` to avoid `NonLinModel` type-instability

Author: franckgaga

Project.toml

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

src/model/linearization.jl

@@ -1,31 +1,3 @@
-
-function jacobianA!(A, jb::JacobianBuffer, model::SimModel, x, u, d)
-    jb.x .= x; jb.u .= u; jb.d .= d
-    return ForwardDiff.jacobian!(A, jb.f_x!, jb.xnext, jb.x, jb.f_x_cfg)
-end
-jacobianA!( _ , _ , model::LinModel, _ , _ , _ ) = model.A
-function jacobianBu!(Bu, jb::JacobianBuffer, model::SimModel, x, u, d)
-    jb.x .= x; jb.u .= u; jb.d .= d
-    return ForwardDiff.jacobian!(Bu, jb.f_u!, jb.xnext, jb.u, jb.f_u_cfg)
-end
-jacobianBu!( _ , _ , model::LinModel, _ , _ , _ ) = model.Bu
-function jacobianBd!(Bd, jb::JacobianBuffer, model::SimModel, x, u, d)
-    jb.x .= x; jb.u .= u; jb.d .= d
-    return ForwardDiff.jacobian!(Bd, jb.f_d!, jb.xnext, jb.d, jb.f_d_cfg)
-end
-jacobianBd!( _ , _ , model::LinModel, _ , _ , _ ) = model.Bd
-function jacobianC!(C, jb::JacobianBuffer, model::SimModel, x, d)
-    jb.x .= x; jb.d .= d
-    return ForwardDiff.jacobian!(C, jb.h_x!, jb.y, jb.x, jb.h_x_cfg)
-end
-jacobianC!( _ , _ , model::LinModel, _ , _ ) = model.C
-function jacobianDd!(Dd, jb::JacobianBuffer, model::SimModel, x, d)
-    jb.x .= x; jb.d .= d
-    return ForwardDiff.jacobian!(Dd, jb.h_d!, jb.y, jb.d, jb.h_d_cfg)
-end
-jacobianDd!( _ , _ , model::LinModel, _ , _ ) = model.Dd
-
-
 """
     LinModel(model::NonLinModel; x=model.x0+model.xop, u=model.uop, d=model.dop)
 
@@ -120,8 +92,8 @@ end
 
 Linearize `model` and store the result in `linmodel` (in-place).
 
-The keyword arguments are identical to [`linearize`](@ref). The code is allocation-free if
-`model` simulations does not allocate.
+The keyword arguments are identical to [`linearize`](@ref). The code allocates a small
+amount of memory to compute the Jacobians.
 
 # Examples
 ```jldoctest
@@ -137,8 +109,8 @@ julia> linearize!(linmodel, model, x=[20.0], u=[0.0]); linmodel.A
 ```
 """
 function linearize!(
-    linmodel::LinModel, model::SimModel; x=model.x0+model.xop, u=model.uop, d=model.dop
-)
+    linmodel::LinModel{NT}, model::SimModel; x=model.x0+model.xop, u=model.uop, d=model.dop
+) where NT<:Real
     nonlinmodel = model
     buffer = nonlinmodel.buffer
     # --- remove the operating points of the nonlinear model (typically zeros) ---
@@ -147,12 +119,19 @@ function linearize!(
     d0 .= d .- nonlinmodel.dop
     x0 .= x .- nonlinmodel.xop
     # --- compute the Jacobians at linearization points ---
-    jb = nonlinmodel.buffer.jacobian
-    jacobianA!(linmodel.A,   jb, nonlinmodel, x0, u0, d0)
-    jacobianBu!(linmodel.Bu, jb, nonlinmodel, x0, u0, d0)
-    jacobianBd!(linmodel.Bd, jb, nonlinmodel, x0, u0, d0)
-    jacobianC!(linmodel.C,   jb, nonlinmodel, x0, d0)
-    jacobianDd!(linmodel.Dd, jb, nonlinmodel, x0, d0)
+    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)
     # --- 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, SMB<:SimModelBuffer
+    NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver
 } <: SimModel{NT}
     x0::Vector{NT}
     f!::F
@@ -21,10 +21,10 @@ struct NonLinModel{
     yname::Vector{String}
     dname::Vector{String}
     xname::Vector{String}
-    buffer::SMB
+    buffer::SimModelBuffer{NT}
     function NonLinModel{NT}(
-        f!::F, h!::H, Ts, nu, nx, ny, nd, p::P, solver::DS, buffer::SMB
-    ) where {NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver, SMB<:SimModelBuffer}
+        f!::F, h!::H, Ts, nu, nx, ny, nd, p::P, solver::DS
+    ) where {NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver}
         Ts > 0 || error("Sampling time Ts must be positive")
         uop = zeros(NT, nu)
         yop = zeros(NT, ny)
@@ -37,7 +37,8 @@ struct NonLinModel{
         xname = ["\$x_{$i}\$" for i in 1:nx]
         x0 = zeros(NT, nx)
         t  = zeros(NT, 1)
-        return new{NT, F, H, P, DS, SMB}(
+        buffer = SimModelBuffer{NT}(nu, nx, ny, nd)
+        return new{NT, F, H, P, DS}(
             x0, 
             f!, h!,
             p,
@@ -144,9 +145,7 @@ 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)
-    jacobian = JacobianBuffer{NT}(f!, h!, nu, nx, ny, nd, p)
-    buffer = SimModelBuffer{NT}(nu, nx, ny, nd, jacobian)
-    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver, buffer)
+    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver)
 end
 
 function NonLinModel(

src/sim_model.jl

@@ -34,9 +34,9 @@ struct JacobianBuffer{
     CHX<:ForwardDiff.JacobianConfig
 }
     xnext::Vector{NT}
-    u::Vector{NT}
     x::Vector{NT}
     y::Vector{NT}
+    u::Vector{NT}
     d::Vector{NT}
     f_x!::FX
     f_u!::FU
@@ -48,34 +48,33 @@ struct JacobianBuffer{
     f_d_cfg::CFD
     h_x_cfg::CHX
     h_d_cfg::CHU
-    
 end
 
 """
-    JacobianBuffer(NT::DataType, f!::Function, h!::Function, nu, nx, ny, nd)
+    JacobianBuffer(NT::DataType, f!::Function, h!::Function, nx, nu, ny, nd)
 
 Buffer object for calling `ForwardDiff.jacobian!` on `SimModel` without any allocation.
 """
 function JacobianBuffer{NT}(
-    f!::Function, h!::Function, nu::Int, nx::Int, ny::Int, nd::Int, p
+    f!::Function, h!::Function, nx::Int, nu::Int, ny::Int, nd::Int
 ) where NT <: Real 
     xnext = Vector{NT}(undef, nx)
-    u = Vector{NT}(undef, nu)
     x = Vector{NT}(undef, nx)
     y = Vector{NT}(undef, ny)
+    u = Vector{NT}(undef, nu)
     d = Vector{NT}(undef, nd)
-    f_x!(xnext, x) = f!(xnext, x, u, d, p)
-    f_u!(xnext, u) = f!(xnext, x, u, d, p)
-    f_d!(xnext, d) = f!(xnext, x, u, d, p)
-    h_x!(y, x) = h!(y, x, d, p)
-    h_d!(y, d) = h!(y, x, d, p)
+    f_x!(y, x) = f!(y, x, u, d)
+    f_u!(y, u) = f!(y, x, u, d)
+    f_d!(y, d) = f!(y, x, u, d)
+    h_x!(y, x) = h!(y, x, d)
+    h_d!(y, d) = h!(y, x, d)
     f_x_cfg = ForwardDiff.JacobianConfig(f_x!, xnext, x)
     f_u_cfg = ForwardDiff.JacobianConfig(f_u!, xnext, u)
     f_d_cfg = ForwardDiff.JacobianConfig(f_d!, xnext, d)
     h_x_cfg = ForwardDiff.JacobianConfig(h_x!, y, x)
     h_d_cfg = ForwardDiff.JacobianConfig(h_d!, y, d)
     return JacobianBuffer(
-        xnext,      u,          x,          y,           d,
+        xnext,      x,          y,          u,          d,
         f_x!,       f_u!,       f_d!,       h_x!,       h_d!, 
         f_x_cfg,    f_u_cfg,    f_d_cfg,    h_x_cfg,    h_d_cfg
     )