added: DI backend as NonLinModel parameter in the struct

Author: franckgaga

src/model/linearization.jl

@@ -1,22 +1,22 @@
 """
-    get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p) -> linfunc!
+    get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p, backend) -> linfunc!
 
 Return the `linfunc!` function that computes the Jacobians of `f!` and `h!` functions.
 
 The function has the following signature: 
 ```
-    linfunc!(xnext, y, A, Bu, C, Bd, Dd, x, u, d, cst_x, cst_u, cst_d) -> nothing
+    linfunc!(xnext, y, A, Bu, C, Bd, Dd, backend, x, u, d, cst_x, cst_u, cst_d) -> nothing
 ```
-and it should modifies in-place all the arguments before `x`. The `cst_x`, `cst_u`, `cst_d` 
-and are `DifferentiationInterface.Constant` objects with the linearization points.
+and it should modifies in-place all the arguments before `backend`. The `backend` argument
+is an `AbstractADType` backend from `DifferentiationInterface`. The `cst_x`, `cst_u` and 
+`cst_d` are `DifferentiationInterface.Constant` objects with the linearization points.
 """
-function get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p)
+function get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p, backend)
     f_x!(xnext, x, u, d) = f!(xnext, x, u, d, p)
     f_u!(xnext, u, x, d) = f!(xnext, x, u, d, p)
     f_d!(xnext, d, x, u) = f!(xnext, x, u, d, p)
     h_x!(y, x, d) = h!(y, x, d, p)
     h_d!(y, d, x) = h!(y, x, d, p)
-    backend = AutoForwardDiff()
     strict  = Val(true)
     xnext = zeros(NT, nx)
     y = zeros(NT, ny)
@@ -31,7 +31,7 @@ function get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p)
     Bd_prep = prepare_jacobian(f_d!, xnext, backend, d, cst_x, cst_u; strict)
     C_prep  = prepare_jacobian(h_x!, y,     backend, x, cst_d       ; strict)
     Dd_prep = prepare_jacobian(h_d!, y,     backend, d, cst_x       ; strict)
-    function linfunc!(xnext, y, A, Bu, C, Bd, Dd, x, u, d, cst_x, cst_u, cst_d)
+    function linfunc!(xnext, y, A, Bu, C, Bd, Dd, backend, x, u, d, cst_x, cst_u, cst_d)
         # all the arguments before `x` are mutated in this function
         jacobian!(f_x!, xnext, A,  A_prep,  backend, x, cst_u, cst_d)
         jacobian!(f_u!, xnext, Bu, Bu_prep, backend, u, cst_x, cst_d)
@@ -183,10 +183,11 @@ end
 function linearize_core!(linmodel::LinModel, model::SimModel, x0, u0, d0)
     xnext0, y0 = linmodel.buffer.x, linmodel.buffer.y
     A, Bu, C, Bd, Dd = linmodel.A, linmodel.Bu, linmodel.C, linmodel.Bd, linmodel.Dd
-    cx = Constant(x0)
-    cu = Constant(u0)
-    cd = Constant(d0)
-    model.linfunc!(xnext0, y0, A, Bu, C, Bd, Dd, x0, u0, d0, cx, cu, cd)
+    cst_x = Constant(x0)
+    cst_u = Constant(u0)
+    cst_d = Constant(d0)
+    backend = model.jacobian
+    model.linfunc!(xnext0, y0, A, Bu, C, Bd, Dd, backend, x0, u0, d0, cst_x, cst_u, cst_d)
     return nothing
 end
 

src/model/nonlinmodel.jl

@@ -1,5 +1,11 @@
 struct NonLinModel{
-    NT<:Real, F<:Function, H<:Function, PT<:Any, DS<:DiffSolver, LF<:Function
+    NT<:Real, 
+    F<:Function, 
+    H<:Function, 
+    PT<:Any, 
+    DS<:DiffSolver, 
+    JB<:AbstractADType,
+    LF<:Function
 } <: SimModel{NT}
     x0::Vector{NT}
     f!::F
@@ -21,11 +27,20 @@ struct NonLinModel{
     yname::Vector{String}
     dname::Vector{String}
     xname::Vector{String}
+    jacobian::JB
     linfunc!::LF
     buffer::SimModelBuffer{NT}
     function NonLinModel{NT}(
-        f!::F, h!::H, Ts, nu, nx, ny, nd, p::PT, solver::DS, linfunc!::LF
-    ) where {NT<:Real, F<:Function, H<:Function, PT<:Any, DS<:DiffSolver, LF<:Function}
+        f!::F, h!::H, Ts, nu, nx, ny, nd, p::PT, solver::DS, jacobian::JB, linfunc!::LF
+    ) where {
+            NT<:Real, 
+            F<:Function, 
+            H<:Function, 
+            PT<:Any, 
+            DS<:DiffSolver, 
+            JB<:AbstractADType,
+            LF<:Function
+        }
         Ts > 0 || error("Sampling time Ts must be positive")
         uop = zeros(NT, nu)
         yop = zeros(NT, ny)
@@ -39,7 +54,7 @@ struct NonLinModel{
         x0 = zeros(NT, nx)
         t  = zeros(NT, 1)
         buffer = SimModelBuffer{NT}(nu, nx, ny, nd)
-        return new{NT, F, H, PT, DS, LF}(
+        return new{NT, F, H, PT, DS, JB, LF}(
             x0, 
             f!, h!,
             p,
@@ -48,7 +63,7 @@ struct NonLinModel{
             nu, nx, ny, nd, 
             uop, yop, dop, xop, fop,
             uname, yname, dname, xname,
-            linfunc!,
+            jacobian, linfunc!,
             buffer
         )
     end
@@ -143,20 +158,20 @@ NonLinModel with a sample time Ts = 2.0 s, empty solver and:
 """
 function NonLinModel{NT}(
     f::Function, h::Function, Ts::Real, nu::Int, nx::Int, ny::Int, nd::Int=0;
-    p=NT[], solver=RungeKutta(4)
+    p=NT[], solver=RungeKutta(4), jacobian=AutoForwardDiff()
 ) where {NT<:Real}
     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)
-    linfunc! = get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p)
-    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver, linfunc!)
+    linfunc! = get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p, jacobian)
+    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver, jacobian, linfunc!)
 end
 
 function NonLinModel(
     f::Function, h::Function, Ts::Real, nu::Int, nx::Int, ny::Int, nd::Int=0;
-    p=Float64[], solver=RungeKutta(4)
+    p=Float64[], solver=RungeKutta(4), jacobian=AutoForwardDiff()
 )
-    return NonLinModel{Float64}(f, h, Ts, nu, nx, ny, nd; p, solver)
+    return NonLinModel{Float64}(f, h, Ts, nu, nx, ny, nd; p, solver, jacobian)
 end
 
 "Get the mutating functions `f!` and `h!` from the provided functions in argument."