wip: simplifying allocations in DiffSolvers

Author: franckgaga

src/estimator/execute.jl

@@ -18,7 +18,7 @@ function remove_op!(estim::StateEstimator, ym, d, u=nothing)
 end
 
 @doc raw"""
-    f̂!(x̂next0, û0, estim::StateEstimator, model::SimModel, x̂0, u0, d0) -> nothing
+    f̂!(x̂0next, û0, x0i, estim::StateEstimator, model::SimModel, x̂0, u0, d0) -> nothing
 
 Mutating state function ``\mathbf{f̂}`` of the augmented model.
 
@@ -30,39 +30,40 @@ function returns the next state of the augmented model, defined as:
     \mathbf{ŷ_0}(k)   &= \mathbf{ĥ}\Big(\mathbf{x̂_0}(k), \mathbf{d_0}(k)\Big) 
 \end{aligned}
 ```
-where ``\mathbf{x̂_0}(k+1)`` is stored in `x̂next0` argument. The method mutates `x̂next0` and
-`û0` in place, the latter stores the input vector of the augmented model 
-``\mathbf{u_0 + ŷ_{s_u}}``. The model parameter vector `model.p` is not included in the 
-function signature for conciseness.
+where ``\mathbf{x̂_0}(k+1)`` is stored in `x̂0next` argument. The method mutates `x̂0next`, 
+`û0` and `x0i` in place. The argument `û0` is the input vector of the augmented model, 
+computed by ``\mathbf{û_0 = u_0 + ŷ_{s_u}}``. The argument `x0i` is used to store the
+intermediate stage values of `model.solver` (when applicable). The model parameter vector
+`model.p` is not included in the function signature for conciseness.
 """
-function f̂!(x̂next0, û0, estim::StateEstimator, model::SimModel, x̂0, u0, d0)
-    return f̂!(x̂next0, û0, model, estim.As, estim.Cs_u, x̂0, u0, d0)
+function f̂!(x̂0next, û0, x0i, estim::StateEstimator, model::SimModel, x̂0, u0, d0)
+    return f̂!(x̂0next, û0, x0i, model, estim.As, estim.Cs_u, x̂0, u0, d0)
 end
 
 """
-    f̂!(x̂next0, _ , estim::StateEstimator, model::LinModel, x̂0, u0, d0) -> nothing
+    f̂!(x̂0next, _ , _ , estim::StateEstimator, model::LinModel, x̂0, u0, d0) -> nothing
 
 Use the augmented model matrices if `model` is a [`LinModel`](@ref).
 """
-function f̂!(x̂next0, _ , estim::StateEstimator, ::LinModel, x̂0, u0, d0)
-    mul!(x̂next0, estim.Â,  x̂0)
-    mul!(x̂next0, estim.B̂u, u0, 1, 1)
-    mul!(x̂next0, estim.B̂d, d0, 1, 1)
+function f̂!(x̂0next, _ , _ , estim::StateEstimator, ::LinModel, x̂0, u0, d0)
+    mul!(x̂0next, estim.Â,  x̂0)
+    mul!(x̂0next, estim.B̂u, u0, 1, 1)
+    mul!(x̂0next, estim.B̂d, d0, 1, 1)
     return nothing
 end
 
 """
-    f̂!(x̂next0, û0, model::SimModel, As, Cs_u, x̂0, u0, d0)
+    f̂!(x̂0next, û0, x0i, model::SimModel, As, Cs_u, x̂0, u0, d0)
 
 Same than [`f̂!`](@ref) for [`SimModel`](@ref) but without the `estim` argument.
 """
-function f̂!(x̂next0, û0, model::SimModel, As, Cs_u, x̂0, u0, d0)
+function f̂!(x̂0next, û0, x0i, model::SimModel, As, Cs_u, x̂0, u0, d0)
     # `@views` macro avoid copies with matrix slice operator e.g. [a:b]
     @views x̂d, x̂s = x̂0[1:model.nx], x̂0[model.nx+1:end]
-    @views x̂d_next, x̂s_next = x̂next0[1:model.nx], x̂next0[model.nx+1:end]
-    mul!(û0, Cs_u, x̂s)
+    @views x̂d_next, x̂s_next = x̂0next[1:model.nx], x̂0next[model.nx+1:end]
+    mul!(û0, Cs_u, x̂s) # ŷs_u = Cs_u * x̂s
     û0 .+= u0
-    f!(x̂d_next, model, x̂d, û0, d0, model.p)
+    f!(x̂d_next, x0i, model, x̂d, û0, d0, model.p)
     mul!(x̂s_next, As, x̂s)
     return nothing
 end
@@ -96,7 +97,7 @@ function ĥ!(ŷ0, model::SimModel, Cs_y, x̂0, d0)
     # `@views` macro avoid copies with matrix slice operator e.g. [a:b]
     @views x̂d, x̂s = x̂0[1:model.nx], x̂0[model.nx+1:end]
     h!(ŷ0, model, x̂d, d0, model.p)
-    mul!(ŷ0, Cs_y, x̂s, 1, 1)
+    mul!(ŷ0, Cs_y, x̂s, 1, 1) # ŷ0 = ŷ0 + Cs_y*x̂s
     return nothing
 end
 

src/estimator/internal_model.jl

@@ -168,20 +168,25 @@ function matrices_internalmodel(model::SimModel{NT}) where NT<:Real
 end
 
 @doc raw"""
-    f̂!(x̂next0, _ , estim::InternalModel, model::NonLinModel, x̂0, u0, d0)
+    f̂!(x̂0next, _ , x̂0i, estim::InternalModel, model::NonLinModel, x̂0, u0, d0)
 
 State function ``\mathbf{f̂}`` of [`InternalModel`](@ref) for [`NonLinModel`](@ref).
 
-It calls `model.f!(x̂next0, x̂0, u0 ,d0, model.p)` since this estimator does not augment the states.
+It calls `model.solver_f!(x̂0next, x̂0i, x̂0, u0 ,d0, model.p)` directly since this estimator
+does not augment the states.
 """
-f̂!(x̂next0, _, ::InternalModel, model::NonLinModel, x̂0, u0, d0) = model.f!(x̂next0, x̂0, u0, d0, model.p)
+function f̂!(x̂0next, _ , x̂0i, ::InternalModel, model::NonLinModel, x̂0, u0, d0)
+    return model.solver_f!(x̂0next, x̂0i, x̂0, u0, d0, model.p)
+end
 
 @doc raw"""
     ĥ!(ŷ0, estim::InternalModel, model::NonLinModel, x̂0, d0)
 
-Output function ``\mathbf{ĥ}`` of [`InternalModel`](@ref), it calls `model.h!`.
+Output function ``\mathbf{ĥ}`` of [`InternalModel`](@ref), it calls `model.solver_h!`.
 """
-ĥ!(x̂next0, ::InternalModel, model::NonLinModel, x̂0, d0) = model.h!(x̂next0, x̂0, d0, model.p)
+function ĥ!(ŷ0, ::InternalModel, model::NonLinModel, x̂0, d0)
+    return model.solver_h!(ŷ0, x̂0, d0, model.p)
+end
 
 
 @doc raw"""

src/model/linearization.jl

@@ -1,7 +1,9 @@
 """
-    get_linearization_func(NT, f!, h!, nu, nx, ny, nd, p, backend) -> linfunc!
+    get_linearization_func(
+        NT, solver_f!, solver_h!, nu, nx, ny, nd, ns, p, solver, backend
+    ) -> linfunc!
 
-Return the `linfunc!` function that computes the Jacobians of `f!` and `h!` functions.
+Return `linfunc!` function that computes Jacobians of `solver_f!` and `solver_h!` functions.
 
 The function has the following signature: 
 ```
@@ -11,31 +13,33 @@ and it should modifies in-place all the arguments before `backend`. The `backend
 is an `AbstractADType` object 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)
+function get_linearization_func(NT, solver_f!, solver_h!, nu, nx, ny, nd, p, solver, backend)
+    f_x!(xnext, x, xi, u, d) = solver_f!(xnext, xi, x, u, d, p)
+    f_u!(xnext, u, xi, x, d) = solver_f!(xnext, xi, x, u, d, p)
+    f_d!(xnext, d, xi, x, u) = solver_f!(xnext, xi, x, u, d, p)
+    h_x!(y, x, d) = solver_h!(y, x, d, p)
+    h_d!(y, d, x) = solver_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)
+    xi = zeros(NT, nx*(solver.ni+1))
+    cache_xi = Cache(xi)
     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)
+    A_prep  = prepare_jacobian(f_x!, xnext, backend, x, cache_xi, cst_u, cst_d; strict)
+    Bu_prep = prepare_jacobian(f_u!, xnext, backend, u, cache_xi, cst_x, cst_d; strict)
+    Bd_prep = prepare_jacobian(f_d!, xnext, backend, d, cache_xi, 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 `backend` 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!(f_x!, xnext, A,  A_prep,  backend, x, cache_xi, cst_u, cst_d)
+        jacobian!(f_u!, xnext, Bu, Bu_prep, backend, u, cache_xi, cst_x, cst_d)
+        jacobian!(f_d!, xnext, Bd, Bd_prep, backend, d, cache_xi, 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
@@ -154,21 +158,21 @@ function linearize!(
     nonlinmodel = model
     buffer = nonlinmodel.buffer
     # --- remove the operating points of the nonlinear model (typically zeros) ---
-    x0, u0, d0 = buffer.x, buffer.u, buffer.d
+    x0, u0, d0, x0i = buffer.x, buffer.u, buffer.d, buffer.xi
     x0 .= x .- nonlinmodel.xop
     u0 .= u .- nonlinmodel.uop
     d0 .= d .- nonlinmodel.dop
     # --- compute the Jacobians at linearization points ---
     linearize_core!(linmodel, nonlinmodel, x0, u0, d0)
     # --- compute the nonlinear model output at operating points ---
-    xnext0, y0 = linmodel.buffer.x, linmodel.buffer.y
+    x0next, y0 = linmodel.buffer.x, linmodel.buffer.y
     h!(y0, nonlinmodel, x0, d0, model.p)
     y  = y0
     y .= y0 .+ nonlinmodel.yop
     # --- compute the nonlinear model next state at operating points ---
-    f!(xnext0, nonlinmodel, x0, u0, d0, model.p)
-    xnext  = xnext0
-    xnext .= xnext0 .+ nonlinmodel.fop .- nonlinmodel.xop
+    f!(x0next, x0i, nonlinmodel, x0, u0, d0, model.p)
+    xnext  = x0next
+    xnext .= x0next .+ nonlinmodel.fop .- nonlinmodel.xop
     # --- modify the linear model operating points ---
     linmodel.uop .= u
     linmodel.yop .= y
@@ -182,13 +186,13 @@ end
 
 "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
+    x0next, 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)
+    model.linfunc!(x0next, y0, A, Bu, C, Bd, Dd, backend, x0, u0, d0, cst_x, cst_u, cst_d)
     return nothing
 end
 

src/model/linmodel.jl

@@ -258,20 +258,20 @@ function steadystate!(model::LinModel, u0, d0)
 end
 
 """
-    f!(xnext0, model::LinModel, x0, u0, d0, p) -> nothing
+    f!(x0next, _ , model::LinModel, x0, u0, d0, _ ) -> nothing
 
-Evaluate `xnext0 = A*x0 + Bu*u0 + Bd*d0` in-place when `model` is a [`LinModel`](@ref).
+Evaluate `x0next = A*x0 + Bu*u0 + Bd*d0` in-place when `model` is a [`LinModel`](@ref).
 """
-function f!(xnext0, model::LinModel, x0, u0, d0, _ )
-    mul!(xnext0, model.A,  x0)
-    mul!(xnext0, model.Bu, u0, 1, 1)
-    mul!(xnext0, model.Bd, d0, 1, 1)
+function f!(x0next, _ , model::LinModel, x0, u0, d0, _ )
+    mul!(x0next, model.A,  x0)
+    mul!(x0next, model.Bu, u0, 1, 1)
+    mul!(x0next, model.Bd, d0, 1, 1)
     return nothing
 end
 
 
 """
-    h!(y0, model::LinModel, x0, d0, p) -> nothing
+    h!(y0, model::LinModel, x0, d0, _ ) -> nothing
 
 Evaluate `y0 = C*x0 + Dd*d0` in-place when `model` is a [`LinModel`](@ref).
 """

src/model/nonlinmodel.jl

@@ -8,8 +8,8 @@ struct NonLinModel{
     LF<:Function
 } <: SimModel{NT}
     x0::Vector{NT}
-    f!::F
-    h!::H
+    solver_f!::F
+    solver_h!::H
     p::PT
     solver::DS
     Ts::NT
@@ -31,7 +31,8 @@ struct NonLinModel{
     linfunc!::LF
     buffer::SimModelBuffer{NT}
     function NonLinModel{NT}(
-        f!::F, h!::H, Ts, nu, nx, ny, nd, p::PT, solver::DS, jacobian::JB, linfunc!::LF
+        solver_f!::F, solver_h!::H, Ts, nu, nx, ny, nd, 
+        p::PT, solver::DS, jacobian::JB, linfunc!::LF
     ) where {
             NT<:Real, 
             F<:Function, 
@@ -53,10 +54,10 @@ 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, solver.ns)
+        buffer = SimModelBuffer{NT}(nu, nx, ny, nd, solver.ni)
         return new{NT, F, H, PT, DS, JB, LF}(
             x0, 
-            f!, h!,
+            solver_f!, solver_h!,
             p,
             solver, 
             Ts, t,
@@ -172,9 +173,13 @@ function NonLinModel{NT}(
 ) 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, jacobian)
-    return NonLinModel{NT}(f!, h!, Ts, nu, nx, ny, nd, p, solver, jacobian, linfunc!)
+    solver_f!, solver_h! = get_solver_functions(NT, solver, f!, h!, Ts, nu, nx, ny, nd)
+    linfunc! = get_linearization_func(
+        NT, solver_f!, solver_h!, nu, nx, ny, nd, p, solver, jacobian
+    )
+    return NonLinModel{NT}(
+        solver_f!, solver_h!, Ts, nu, nx, ny, nd, p, solver, jacobian, linfunc!
+    )
 end
 
 function NonLinModel(
@@ -262,11 +267,11 @@ 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, model.buffer.K, x0, u0, d0, p)
+"Call `model.solver_f!(x0next, x0i, x0, u0, d0, p)` for [`NonLinModel`](@ref)."
+f!(x0next, x0i, model::NonLinModel, x0, u0, d0, p) = model.solver_f!(x0next, x0i, x0, u0, d0, p)
 
-"Call `model.h!(y0, x0, d0, p)` for [`NonLinModel`](@ref)."
-h!(y0, model::NonLinModel, x0, d0, p) = model.h!(y0, x0, d0, p)
+"Call `model.solver_h!(y0, x0, d0, p)` for [`NonLinModel`](@ref)."
+h!(y0, model::NonLinModel, x0, d0, p) = model.solver_h!(y0, x0, d0, p)
 
 detailstr(model::NonLinModel) = ", $(typeof(model.solver).name.name)($(model.solver.order)) solver"
 detailstr(::NonLinModel{<:Real, <:Function, <:Function, <:Any, <:EmptySolver}) = ", empty solver"