Merge pull request #179 from JuliaControl/diff_interface_linearize

added: `linearize!` now uses `DifferentiationInterface.jl`

Author: franckgaga

src/ModelPredictiveControl.jl

@@ -1,6 +1,6 @@
 module ModelPredictiveControl
 
-using PrecompileTools # TODO: remove this dep if possible (with Cache of DI.jl)
+using PrecompileTools 
 using LinearAlgebra
 using Random: randn
 
@@ -13,7 +13,7 @@ using DifferentiationInterface: Constant, Cache
 using SparseConnectivityTracer: TracerSparsityDetector
 using SparseMatrixColorings: GreedyColoringAlgorithm, sparsity_pattern
 
-import ForwardDiff #TODO: delete this after `linearize!` and `ExtendedKalmanFilter` are updated
+import ForwardDiff
 
 import ControlSystemsBase
 import ControlSystemsBase: ss, tf, delay
@@ -25,7 +25,7 @@ import JuMP
 import JuMP: MOIU, MOI, GenericModel, Model, optimizer_with_attributes, register
 import JuMP: @variable, @operator, @constraint, @objective
 
-import PreallocationTools: DiffCache, get_tmp
+import PreallocationTools: DiffCache, get_tmp # TODO: remove this dep if possible (with Cache of DI.jl)
 
 import OSQP, Ipopt
 

src/controller/linmpc.jl

@@ -135,7 +135,7 @@ arguments. This controller allocates memory at each time step for the optimizati
 
 # Arguments
 - `model::LinModel` : model used for controller predictions and state estimations.
-- `Hp=10+nk`: prediction horizon ``H_p``, `nk` is the number of delays in `model`.
+- `Hp=10+nk` : prediction horizon ``H_p``, `nk` is the number of delays in `model`.
 - `Hc=2` : control horizon ``H_c``.
 - `Mwt=fill(1.0,model.ny)` : main diagonal of ``\mathbf{M}`` weight matrix (vector).
 - `Nwt=fill(0.1,model.nu)` : main diagonal of ``\mathbf{N}`` weight matrix (vector).

src/general.jl

@@ -6,31 +6,6 @@ const DEFAULT_LWT = 0.0
 const DEFAULT_CWT = 1e5
 const DEFAULT_EWT = 0.0
 
-"Abstract type for all differentiation buffers."
-abstract type DifferentiationBuffer end
-
-function Base.show(io::IO, buffer::DifferentiationBuffer) 
-    return print(io, "DifferentiationBuffer with a $(typeof(buffer.config).name.name)")
-end
-
-"Struct with both function and configuration for ForwardDiff Jacobian."
-struct JacobianBuffer{FT<:Function, CT<:ForwardDiff.JacobianConfig} <: DifferentiationBuffer
-    f!::FT
-    config::CT
-end
-
-"Create a JacobianBuffer with in-place 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 get_jacobian!(A, buffer::JacobianBuffer, y, x)
-    return ForwardDiff.jacobian!(A, buffer.f!, y, x, buffer.config)
-end
-
-"Init a differentiation result matrix as dense or sparse matrix, as required by `backend`."
-init_diffmat(T, backend::AbstractADType, _  , nx , ny) = Matrix{T}(undef, ny, nx)
-init_diffmat(T, backend::AutoSparse    ,prep , _ , _ ) = similar(sparsity_pattern(prep), T)
-
 "Termination status that means 'no solution available'."
 const ERROR_STATUSES = (
     JuMP.INFEASIBLE, JuMP.DUAL_INFEASIBLE, JuMP.LOCALLY_INFEASIBLE, 
@@ -81,6 +56,10 @@ function limit_solve_time(optim::GenericModel, Ts)
     end
 end
 
+"Init a differentiation result matrix as dense or sparse matrix, as required by `backend`."
+init_diffmat(T, backend::AbstractADType, _  , nx , ny) = Matrix{T}(undef, ny, nx)
+init_diffmat(T, backend::AutoSparse    ,prep , _ , _ ) = similar(sparsity_pattern(prep), T)
+
 "Verify that x and y elements are different using `!==`."
 isdifferent(x, y) = any(xi !== yi for (xi, yi) in zip(x, y))
 

src/model/linearization.jl

@@ -1,81 +1,59 @@
-"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
+"""
+    get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p, backend) -> linfunc!
 
-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
+Return the `linfunc!` function that computes the Jacobians of `f!` and `h!` functions.
 
-"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!
+The function has the following signature: 
+```
+    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 `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, 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)
+    strict  = Val(true)
+    xnext = zeros(NT, nx)
+    y = zeros(NT, ny)
+    x = zeros(NT, nx)
+    u = zeros(NT, nu)
+    d = zeros(NT, nd)
+    cst_x = Constant(x)
+    cst_u = Constant(u)
+    cst_d = Constant(d)
+    A_prep  = prepare_jacobian(f_x!, xnext, backend, x, cst_u, cst_d; strict)
+    Bu_prep = prepare_jacobian(f_u!, xnext, backend, u, cst_x, cst_d; strict)
+    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, 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)
+        jacobian!(f_d!, xnext, Bd, Bd_prep, backend, d, cst_x, cst_u)
+        jacobian!(h_x!, y,     C,  C_prep,  backend, x, cst_d)
+        jacobian!(h_d!, y,     Dd, Dd_prep, backend, d, cst_x)
+        return nothing
+    end
+    return linfunc!
 end
 
+
 @doc raw"""
     linearize(model::SimModel; x=model.x0+model.xop, u=model.uop, d=model.dop) -> linmodel
 
 Linearize `model` at the operating points `x`, `u`, `d` and return the [`LinModel`](@ref).
 
 The arguments `x`, `u` and `d` are the linearization points for the state ``\mathbf{x}``,
 manipulated input ``\mathbf{u}`` and measured disturbance ``\mathbf{d}``, respectively (not
-necessarily an equilibrium, details in Extended Help). The Jacobians of ``\mathbf{f}`` and 
-``\mathbf{h}`` functions are automatically computed with [`ForwardDiff`](@extref ForwardDiff).
+necessarily an equilibrium, details in Extended Help). By default, [`ForwardDiff`](@extref ForwardDiff)
+automatically computes the Jacobians of ``\mathbf{f}`` and ``\mathbf{h}`` functions. Modify
+the `jacobian` keyword argument at the construction of `model` to swap the backend.
 
 !!! warning
     See Extended Help if you get an error like:    
@@ -181,8 +159,7 @@ function linearize!(
     u0 .= u .- nonlinmodel.uop
     d0 .= d .- nonlinmodel.dop
     # --- compute the Jacobians at linearization points ---
-    xnext0::Vector{NT}, y0::Vector{NT} = linmodel.buffer.x, linmodel.buffer.y
-    get_jacobians!(linmodel, xnext0, y0, nonlinmodel, x0, u0, d0)
+    linearize_core!(linmodel, nonlinmodel, x0, u0, d0)
     # --- compute the nonlinear model output at operating points ---
     xnext0, y0 = linmodel.buffer.x, linmodel.buffer.y
     h!(y0, nonlinmodel, x0, d0, model.p)
@@ -203,22 +180,20 @@ function linearize!(
     return linmodel
 end
 
-"Compute the 5 Jacobians of `model` at the linearization point and write them in `linmodel`."
-function get_jacobians!(linmodel::LinModel, xnext0, y0, model::SimModel, x0, u0, d0)
-    linbuffer = model.linbuffer # a LinearizationBuffer object
-    linbuffer.x .= x0
-    linbuffer.u .= u0
-    linbuffer.d .= d0
-    get_jacobian!(linmodel.A,  linbuffer.buffer_f_at_u_d, xnext0, x0)
-    get_jacobian!(linmodel.Bu, linbuffer.buffer_f_at_x_d, xnext0, u0)
-    get_jacobian!(linmodel.Bd, linbuffer.buffer_f_at_x_u, xnext0, d0)
-    get_jacobian!(linmodel.C,  linbuffer.buffer_h_at_d, y0, x0)
-    get_jacobian!(linmodel.Dd, linbuffer.buffer_h_at_x, y0, d0)
+"Call `linfunc!` function to compute the Jacobians of `model` at the linearization point."
+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
+    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
 
 "Copy the state-space matrices of `model` to `linmodel` if `model` is already linear."
-function get_jacobians!(linmodel::LinModel, _ , _ , model::LinModel, _ , _ , _)
+function linearize_core!(linmodel::LinModel, model::LinModel, _ , _ , _ )
     linmodel.A  .= model.A
     linmodel.Bu .= model.Bu
     linmodel.C  .= model.C

src/model/nonlinmodel.jl

@@ -1,5 +1,11 @@
 struct NonLinModel{
-    NT<:Real, F<:Function, H<:Function, PT<:Any, DS<:DiffSolver, LB<:LinearizationBuffer
+    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}
-    linbuffer::LB
+    jacobian::JB
+    linfunc!::LF
     buffer::SimModelBuffer{NT}
     function NonLinModel{NT}(
-        f!::F, h!::H, Ts, nu, nx, ny, nd, p::PT, solver::DS, linbuffer::LB
-    ) where {NT<:Real, F<:Function, H<:Function, PT<:Any, DS<:DiffSolver, LB<:LinearizationBuffer}
+        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, LB}(
+        return new{NT, F, H, PT, DS, JB, LF}(
             x0, 
             f!, h!,
             p,
@@ -48,15 +63,15 @@ struct NonLinModel{
             nu, nx, ny, nd, 
             uop, yop, dop, xop, fop,
             uname, yname, dname, xname,
-            linbuffer,
+            jacobian, linfunc!,
             buffer
         )
     end
 end
 
 @doc raw"""
-    NonLinModel{NT}(f::Function,  h::Function,  Ts, nu, nx, ny, nd=0; p=[], solver=RungeKutta(4))
-    NonLinModel{NT}(f!::Function, h!::Function, Ts, nu, nx, ny, nd=0; p=[], solver=RungeKutta(4))
+    NonLinModel{NT}(f::Function,  h::Function,  Ts, nu, nx, ny, nd=0; <keyword arguments>)
+    NonLinModel{NT}(f!::Function, h!::Function, Ts, nu, nx, ny, nd=0; <keyword arguments>)
 
 Construct a nonlinear model from state-space functions `f`/`f!` and `h`/`h!`.
 
@@ -71,24 +86,20 @@ functions are defined as:
 ```
 where ``\mathbf{x}``, ``\mathbf{y}``, ``\mathbf{u}``, ``\mathbf{d}`` and ``\mathbf{p}`` are
 respectively the state, output, manipulated input, measured disturbance and parameter
-vectors. If the dynamics is a function of the time, simply add a measured disturbance
-defined as ``d(t) = t``. The functions can be implemented in two possible ways:
+vectors. As a mather of fact, the parameter argument `p` can be any Julia objects but use a
+mutable type if you want to change them later e.g.: a vector. If the dynamics is a function
+of the time, simply add a measured disturbance defined as ``d(t) = t``. The functions can be
+implemented in two possible ways:
 
 1. **Non-mutating functions** (out-of-place): define them as `f(x, u, d, p) -> ẋ` and
    `h(x, d, p) -> y`. This syntax is simple and intuitive but it allocates more memory.
 2. **Mutating functions** (in-place): define them as `f!(ẋ, x, u, d, p) -> nothing` and
    `h!(y, x, d, p) -> nothing`. This syntax reduces the allocations and potentially the 
    computational burden as well.
 
-`Ts` is the sampling time in second. `nu`, `nx`, `ny` and `nd` are the respective number of 
-manipulated inputs, states, outputs and measured disturbances. The keyword argument `p`
-is the parameters of the model passed to the two functions. It can be any Julia objects but
-use a mutable type if you want to change them later e.g.: a vector.
-
 !!! tip
     Replace the `d` or `p` argument with `_` in your functions if not needed (see Examples below).
     
-A 4th order [`RungeKutta`](@ref) solver discretizes the differential equations by default. 
 The rest of the documentation assumes discrete dynamics since all models end up in this 
 form. The optional parameter `NT` explicitly set the number type of vectors (default to 
 `Float64`).
@@ -100,6 +111,20 @@ form. The optional parameter `NT` explicitly set the number type of vectors (def
 
 See also [`LinModel`](@ref).
 
+# Arguments
+- `f::Function` or `f!`: state function.
+- `h::Function` or `h!`: output function.
+- `Ts`: sampling time of in second.
+- `nu`: number of manipulated inputs.
+- `nx`: number of states.
+- `ny`: number of outputs.
+- `nd=0`: number of measured disturbances.
+- `p=[]`: parameters of the model (any type).
+- `solver=RungeKutta(4)`: a [`DiffSolver`](@ref) object for the discretization of continuous
+  dynamics, use `nothing` for discrete-time models (default to 4th order [`RungeKutta`](@ref)).
+- `jacobian=AutoForwardDiff()`: an `AbstractADType` backend when [`linearize`](@ref) is
+   called, see [`DifferentiationInterface` doc](@extref DifferentiationInterface List).
+
 # Examples
 ```jldoctest
 julia> f!(ẋ, x, u, _ , p) = (ẋ .= p*x .+ u; nothing);
@@ -143,20 +168,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)
-    linbuffer = LinearizationBuffer(NT, f!, h!, nu, nx, ny, nd, p)
-    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver, linbuffer)
+    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."