added: JacobianBuffer to get rid of the allocation with jacobian!

Author: franckgaga

src/model/linearization.jl

@@ -1,3 +1,31 @@
+
+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)
 

src/model/linmodel.jl

@@ -21,7 +21,7 @@ struct LinModel{NT<:Real} <: SimModel{NT}
     yname::Vector{String}
     dname::Vector{String}
     xname::Vector{String}
-    buffer::SimModelBuffer{NT}
+    buffer::SimModelBuffer{NT, Nothing}
     function LinModel{NT}(A, Bu, C, Bd, Dd, Ts) where {NT<:Real}
         A, Bu = to_mat(A, 1, 1), to_mat(Bu, 1, 1)
         nu, nx = size(Bu, 2), size(A, 2)

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, SMB<:SimModelBuffer
 } <: SimModel{NT}
     x0::Vector{NT}
     f!::F
@@ -21,10 +21,10 @@ struct NonLinModel{
     yname::Vector{String}
     dname::Vector{String}
     xname::Vector{String}
-    buffer::SimModelBuffer{NT}
+    buffer::SMB
     function NonLinModel{NT, F, H, P, DS}(
-        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, buffer::SMB
+    ) where {NT<:Real, F<:Function, H<:Function, P<:Any, DS<:DiffSolver, SMB<:SimModelBuffer}
         Ts > 0 || error("Sampling time Ts must be positive")
         uop = zeros(NT, nu)
         yop = zeros(NT, ny)
@@ -37,8 +37,7 @@ struct NonLinModel{
         xname = ["\$x_{$i}\$" for i in 1:nx]
         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, SMB}(
             x0, 
             f!, h!,
             p,
@@ -146,7 +145,9 @@ function NonLinModel{NT}(
     f!, h! = get_mutating_functions(NT, f, h)
     f!, h! = get_solver_functions(NT, solver, f!, h!, Ts, nu, nx, ny, nd)
     F, H, P, DS = get_types(f!, h!, p, solver)
-    return NonLinModel{NT, F, H, P, DS}(f!, h!, Ts, nu, nx, ny, nd, p, solver)
+    jacobian = JacobianBuffer{NT}(f!, h!, nu, nx, ny, nd)
+    buffer = SimModelBuffer{NT}(nu, nx, ny, nd, jacobian)
+    return NonLinModel{NT, F, H, P, DS}(f!, h!, Ts, nu, nx, ny, nd, p, solver, buffer)
 end
 
 function NonLinModel(

src/sim_model.jl

@@ -20,12 +20,73 @@ julia> y = model()
 """
 abstract type SimModel{NT<:Real} end
 
-struct SimModelBuffer{NT<:Real}
+struct JacobianBuffer{
+    NT<:Real, 
+    FX<:Function,
+    FU<:Function,
+    FD<:Function,
+    HX<:Function,
+    HU<:Function,
+    CFX<:ForwardDiff.JacobianConfig,
+    CFU<:ForwardDiff.JacobianConfig,
+    CFD<:ForwardDiff.JacobianConfig,
+    CHU<:ForwardDiff.JacobianConfig,
+    CHX<:ForwardDiff.JacobianConfig
+}
+    xnext::Vector{NT}
+    x::Vector{NT}
+    y::Vector{NT}
+    u::Vector{NT}
+    d::Vector{NT}
+    f_x!::FX
+    f_u!::FU
+    f_d!::FD
+    h_x!::HX
+    h_d!::HU
+    f_x_cfg::CFX
+    f_u_cfg::CFU
+    f_d_cfg::CFD
+    h_x_cfg::CHX
+    h_d_cfg::CHU
+end
+
+"""
+    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, nx::Int, nu::Int, ny::Int, nd::Int
+) where NT <: Real 
+    xnext = Vector{NT}(undef, nx)
+    x = Vector{NT}(undef, nx)
+    y = Vector{NT}(undef, ny)
+    u = Vector{NT}(undef, nu)
+    d = Vector{NT}(undef, nd)
+    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,      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
+    )
+end
+
+struct SimModelBuffer{NT<:Real, JB<:Union{JacobianBuffer, Nothing}}
     u::Vector{NT}
     x::Vector{NT}
     y::Vector{NT}
     d::Vector{NT}
     empty::Vector{NT}
+    jacobian::JB
 end
 
 @doc raw"""
@@ -35,13 +96,15 @@ Create a buffer for `SimModel` objects for inputs, states, outputs, and disturba
 
 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, jacobian::JB=nothing
+) where {NT <: Real, JB <: Union{JacobianBuffer, Nothing}}
     u = Vector{NT}(undef, nu)
     x = Vector{NT}(undef, nx)
     y = Vector{NT}(undef, ny)
     d = Vector{NT}(undef, nd)
     empty = Vector{NT}(undef, 0)
-    return SimModelBuffer{NT}(u, x, y, d, empty)
+    return SimModelBuffer{NT, JB}(u, x, y, d, empty, jacobian)
 end