in-place NonLinModel seems to work

franckgaga · franckgaga · commit 21bd5895ee91 · 2024-03-02T14:28:42.000-05:00
diff --git a/src/controller/execute.jl b/src/controller/execute.jl
@@ -325,15 +325,18 @@ function predict!(
     nu, ny, nd, Hp, Hc = model.nu, model.ny, model.nd, mpc.Hp, mpc.Hc
     u0 = u
     x̂  .= mpc.estim.x̂
+    x̂next = similar(x̂) # TODO: avoid this allocation if possible
     u0 .= mpc.estim.lastu0
     d0  = @views mpc.d0[1:end]
     for j=1:Hp
         if j ≤ Hc
             u0 .+= @views ΔŨ[(1 + nu*(j-1)):(nu*j)]
         end
-        x̂[:]  = f̂(mpc.estim, model, x̂, u0, d0)
-        d0    = @views mpc.D̂0[(1 + nd*(j-1)):(nd*j)]
-        Ŷ[(1 + ny*(j-1)):(ny*j)] = ĥ(mpc.estim, model, x̂, d0)
+        f̂!(x̂next, mpc.estim, model, x̂, u0, d0)
+        x̂ .= x̂next
+        d0 = @views mpc.D̂0[(1 + nd*(j-1)):(nd*j)]
+        ŷ  = @views Ŷ[(1 + ny*(j-1)):(ny*j)]
+        ĥ!(ŷ, mpc.estim, model, x̂, d0)
     end
     Ŷ .= Ŷ .+ mpc.Ŷop # Ŷop = Ŷs + Yop, and Ŷs=0 if mpc.estim is not an InternalModel
     x̂end = x̂
diff --git a/src/estimator/execute.jl b/src/estimator/execute.jl
@@ -15,9 +15,9 @@ function remove_op!(estim::StateEstimator, u, ym, d)
 end
 
 @doc raw"""
-    f̂(estim::StateEstimator, model::SimModel, x̂, u, d)
+    f̂!(x̂next, estim::StateEstimator, model::SimModel, x̂, u, d) -> x̂next
 
-State function ``\mathbf{f̂}`` of the augmented model.
+Mutating state function ``\mathbf{f̂}`` of the augmented model.
 
 By introducing an augmented state vector ``\mathbf{x̂}`` like in [`augment_model`](@ref), the
 function returns the next state of the augmented model, defined as:
@@ -26,28 +26,48 @@ function returns the next state of the augmented model, defined as:
     \mathbf{x̂}(k+1) &= \mathbf{f̂}\Big(\mathbf{x̂}(k), \mathbf{u}(k), \mathbf{d}(k)\Big) \\
     \mathbf{ŷ}(k)   &= \mathbf{ĥ}\Big(\mathbf{x̂}(k), \mathbf{d}(k)\Big) 
 \end{aligned}
+where ``\mathbf{x̂}(k+1)`` is stored in `x̂next` argument.
 ```
 """
-function f̂(estim::StateEstimator, model::SimModel, x̂, u, d)
+function f̂!(x̂next, estim::StateEstimator, model::SimModel, x̂, u, d)
     # `@views` macro avoid copies with matrix slice operator e.g. [a:b]
     @views x̂d, x̂s = x̂[1:model.nx], x̂[model.nx+1:end]
-    return [f(model, x̂d, u + estim.Cs_u*x̂s, d); estim.As*x̂s]
+    @views x̂d_next, x̂s_next = x̂next[1:model.nx], x̂next[model.nx+1:end]
+    T = promote_type(eltype(x̂), eltype(u), eltype(d))
+    u_us = Vector{T}(undef, model.nu) # TODO: avoid this allocation if possible
+    u_us .= u .+ mul!(u_us, estim.Cs_u, x̂s)
+    f!(x̂d_next, model, x̂d, u_us, d)
+    mul!(x̂s_next, estim.As, x̂s)
+    return x̂next
 end
 "Use the augmented model matrices if `model` is a [`LinModel`](@ref)."
-f̂(estim::StateEstimator, ::LinModel, x̂, u, d) = estim.Â * x̂ + estim.B̂u * u + estim.B̂d * d
+function f̂!(x̂next, estim::StateEstimator, ::LinModel, x̂, u, d)
+    x̂next .= 0
+    mul!(x̂next, estim.Â,  x̂, 1, 1)
+    mul!(x̂next, estim.B̂u, u, 1, 1)
+    mul!(x̂next, estim.B̂d, d, 1, 1)
+    return x̂next
+end
 
 @doc raw"""
-    ĥ(estim::StateEstimator, model::SimModel, x̂, d)
+    ĥ!(ŷ, estim::StateEstimator, model::SimModel, x̂, d) -> ŷ
 
-Output function ``\mathbf{ĥ}`` of the augmented model, see [`f̂`](@ref) for details.
+Mutating output function ``\mathbf{ĥ}`` of the augmented model, see [`f̂!`](@ref).
 """
-function ĥ(estim::StateEstimator, model::SimModel, x̂, d)
+function ĥ!(ŷ, estim::StateEstimator, model::SimModel, x̂, d)
     # `@views` macro avoid copies with matrix slice operator e.g. [a:b]
     @views x̂d, x̂s = x̂[1:model.nx], x̂[model.nx+1:end]
-    return h(model, x̂d, d) + estim.Cs_y*x̂s
+    h!(ŷ, model, x̂d, d)
+    mul!(ŷ, estim.Cs_y, x̂s, 1, 1)
+    return ŷ
 end
 "Use the augmented model matrices if `model` is a [`LinModel`](@ref)."
-ĥ(estim::StateEstimator, ::LinModel, x̂, d) = estim.Ĉ * x̂ + estim.D̂d * d
+function ĥ!(ŷ, estim::StateEstimator, ::LinModel, x̂, d)
+    ŷ .= 0
+    mul!(ŷ, estim.Ĉ,  x̂, 1, 1)
+    mul!(ŷ, estim.D̂d, d, 1, 1)
+    return ŷ
+end
 
 
 @doc raw"""
@@ -147,8 +167,12 @@ julia> ŷ = evaloutput(kf)
  20.0
 ```
 """
-function evaloutput(estim::StateEstimator, d=empty(estim.x̂)) 
-    return ĥ(estim, estim.model, estim.x̂, d - estim.model.dop) + estim.model.yop
+function evaloutput(estim::StateEstimator{NT}, d=empty(estim.x̂)) where NT <: Real
+    validate_args(estim.model, d)
+    ŷ = Vector{NT}(undef, estim.model.ny)
+    ĥ!(ŷ, estim, estim.model, estim.x̂, d - estim.model.dop)
+    ŷ .+= estim.model.yop
+    return ŷ
 end
 
 "Functor allowing callable `StateEstimator` object as an alias for `evaloutput`."
@@ -186,7 +210,7 @@ updatestate!(::StateEstimator, _ ) = throw(ArgumentError("missing measured outpu
 Check `u`, `ym` and `d` sizes against `estim` dimensions.
 """
 function validate_args(estim::StateEstimator, u, ym, d)
-    validate_args(estim.model, u, d)
+    validate_args(estim.model, d, u)
     nym = estim.nym
     size(ym) ≠ (nym,) && throw(DimensionMismatch("ym size $(size(ym)) ≠ meas. output size ($nym,)"))
 end
diff --git a/src/estimator/internal_model.jl b/src/estimator/internal_model.jl
@@ -144,20 +144,20 @@ function matrices_internalmodel(model::SimModel{NT}) where NT<:Real
 end
 
 @doc raw"""
-    f̂(::InternalModel, model::NonLinModel, x̂, u, d)
+    f̂!(x̂next, ::InternalModel, model::NonLinModel, x̂, u, d)
 
 State function ``\mathbf{f̂}`` of [`InternalModel`](@ref) for [`NonLinModel`](@ref).
 
-It calls `model.f(x̂, u ,d)` since this estimator does not augment the states.
+It calls `model.f!(x̂next, x̂, u ,d)` since this estimator does not augment the states.
 """
-f̂(::InternalModel, model::NonLinModel, x̂, u, d) = model.f(x̂, u, d)
+f̂!(x̂next, ::InternalModel, model::NonLinModel, x̂, u, d) = model.f!(x̂next, x̂, u, d)
 
 @doc raw"""
-    ĥ(::InternalModel, model::NonLinModel, x̂, d)
+    ĥ!(ŷ, ::InternalModel, model::NonLinModel, x̂, d)
 
-Output function ``\mathbf{ĥ}`` of [`InternalModel`](@ref), it calls `model.h`.
+Output function ``\mathbf{ĥ}`` of [`InternalModel`](@ref), it calls `model.h!`.
 """
-ĥ(::InternalModel, model::NonLinModel, x̂, d) = model.h(x̂, d)
+ĥ!(x̂next, ::InternalModel, model::NonLinModel, x̂, d) = model.h!(x̂next, x̂, d)
 
 
 @doc raw"""
@@ -218,12 +218,17 @@ function update_estimate!(
     model = estim.model
     x̂d, x̂s = estim.x̂d, estim.x̂s
     # -------------- deterministic model ---------------------
-    ŷd = h(model, x̂d, d)
-    x̂d[:] = f(model, x̂d, u, d) # this also updates estim.x̂ (they are the same object)
+    ŷd, x̂dnext = Vector{NT}(undef, model.ny), Vector{NT}(undef, model.nx)
+    h!(ŷd, model, x̂d, d)
+    f!(x̂dnext, model, x̂d, u, d) 
+    x̂d .= x̂dnext # this also updates estim.x̂ (they are the same object)
     # --------------- stochastic model -----------------------
+    x̂snext = Vector{NT}(undef, estim.nxs)
     ŷs = zeros(NT, model.ny)
     ŷs[estim.i_ym] = ym - ŷd[estim.i_ym]   # ŷs=0 for unmeasured outputs
-    x̂s[:] = estim.Âs*x̂s + estim.B̂s*ŷs
+    mul!(x̂snext, estim.Âs, x̂s)
+    mul!(x̂snext, estim.B̂s, ŷs, 1, 1)
+    x̂s .= x̂snext
     return nothing
 end
 
@@ -247,11 +252,12 @@ This estimator does not augment the state vector, thus ``\mathbf{x̂ = x̂_d}``.
 """
 function init_estimate!(estim::InternalModel, model::LinModel{NT}, u, ym, d) where NT<:Real
     x̂d, x̂s = estim.x̂d, estim.x̂s
-    x̂d[:] = (I - model.A)\(model.Bu*u + model.Bd*d)
-    ŷd = h(model, x̂d, d)
+    x̂d .= (I - model.A)\(model.Bu*u + model.Bd*d)
+    ŷd = Vector{NT}(undef, model.ny)
+    h!(ŷd, model, x̂d, d)
     ŷs = zeros(NT, model.ny)
     ŷs[estim.i_ym] = ym - ŷd[estim.i_ym]  # ŷs=0 for unmeasured outputs
-    x̂s[:] = (I-estim.Âs)\estim.B̂s*ŷs
+    x̂s .= (I-estim.Âs)\estim.B̂s*ŷs
     return nothing
 end
 
diff --git a/src/estimator/kalman.jl b/src/estimator/kalman.jl
@@ -569,6 +569,7 @@ function update_estimate!(estim::UnscentedKalmanFilter{NT}, u, ym, d) where NT<:
     # --- initialize matrices ---
     X̂  = Matrix{NT}(undef, nx̂, nσ)
     ŷm = Vector{NT}(undef, nym)
+    ŷ  = Vector{NT}(undef, estim.model.ny)
     Ŷm = Matrix{NT}(undef, nym, nσ)
     sqrt_P̂ = LowerTriangular{NT, Matrix{NT}}(Matrix{NT}(undef, nx̂, nx̂))
     # --- correction step ---
@@ -578,7 +579,8 @@ function update_estimate!(estim::UnscentedKalmanFilter{NT}, u, ym, d) where NT<:
     X̂[:, 2:nx̂+1]   .+= γ_sqrt_P̂
     X̂[:, nx̂+2:end] .-= γ_sqrt_P̂
     for j in axes(Ŷm, 2)
-        Ŷm[:, j] = @views ĥ(estim, estim.model, X̂[:, j], d)[estim.i_ym]
+        @views ĥ!(ŷ, estim, estim.model, X̂[:, j], d)
+        @views Ŷm[:, j] .= ŷ[estim.i_ym]
     end
     mul!(ŷm, Ŷm, m̂)
     X̄, Ȳm = X̂, Ŷm
@@ -600,7 +602,7 @@ function update_estimate!(estim::UnscentedKalmanFilter{NT}, u, ym, d) where NT<:
     X̂_cor[:, nx̂+2:end] .-= γ_sqrt_P̂_cor
     X̂_next = X̂_cor
     for j in axes(X̂_next, 2)
-        X̂_next[:, j] = @views f̂(estim, estim.model, X̂_cor[:, j], u, d)
+        @views f̂!(X̂_next[:, j], estim, estim.model, X̂_cor[:, j], u, d)
     end
     x̂_next = mul!(x̂, X̂_next, m̂)
     X̄_next = X̂_next
@@ -757,9 +759,13 @@ automatically computes the Jacobians:
 ```
 The matrix ``\mathbf{Ĥ^m}`` is the rows of ``\mathbf{Ĥ}`` that are measured outputs.
 """
-function update_estimate!(estim::ExtendedKalmanFilter, u, ym, d=empty(estim.x̂))
-    F̂ = ForwardDiff.jacobian(x̂ -> f̂(estim, estim.model, x̂, u, d), estim.x̂)
-    Ĥ = ForwardDiff.jacobian(x̂ -> ĥ(estim, estim.model, x̂, d), estim.x̂)
+function update_estimate!(
+    estim::ExtendedKalmanFilter{NT}, u, ym, d=empty(estim.x̂)
+) where NT<:Real
+    model = estim.model
+    x̂next, ŷ = Vector{NT}(undef, estim.nx̂), Vector{NT}(undef, model.ny)
+    F̂ = ForwardDiff.jacobian((x̂next, x̂) -> f̂!(x̂next, estim, model, x̂, u, d), x̂next, estim.x̂)
+    Ĥ = ForwardDiff.jacobian((ŷ, x̂)     -> ĥ!(ŷ, estim, model, x̂, d), ŷ, estim.x̂)
     return update_estimate_kf!(estim, u, ym, d, F̂, Ĥ[estim.i_ym, :], estim.P̂, estim.x̂)
 end
 
@@ -790,7 +796,7 @@ function validate_kfcov(nym, nx̂, Q̂, R̂, P̂0=nothing)
 end
 
 """
-    update_estimate_kf!(estim, u, ym, d, Â, Ĉm, P̂, x̂=nothing)
+    update_estimate_kf!(estim::StateEstimator, u, ym, d, Â, Ĉm, P̂, x̂=nothing)
 
 Update time-varying/extended Kalman Filter estimates with augmented `Â` and `Ĉm` matrices.
 
@@ -801,16 +807,22 @@ The implementation uses in-place operations and explicit factorization to reduce
 allocations. See e.g. [`KalmanFilter`](@ref) docstring for the equations. If `isnothing(x̂)`,
 only the covariance `P̂` is updated.
 """
-function update_estimate_kf!(estim, u, ym, d, Â, Ĉm, P̂, x̂=nothing)
+function update_estimate_kf!(
+    estim::StateEstimator{NT}, u, ym, d, Â, Ĉm, P̂, x̂=nothing
+) where NT<:Real
     Q̂, R̂, M̂ = estim.Q̂, estim.R̂, estim.M̂
     mul!(M̂, P̂, Ĉm')
     rdiv!(M̂, cholesky!(Hermitian(Ĉm * P̂ * Ĉm' .+ R̂)))
     if !isnothing(x̂)
         mul!(estim.K̂, Â, M̂)
-        ŷm = @views ĥ(estim, estim.model, x̂, d)[estim.i_ym]
+        x̂next, ŷ = Vector{NT}(undef, estim.nx̂), Vector{NT}(undef, estim.model.ny)
+        ĥ!(ŷ, estim, estim.model, x̂, d)
+        ŷm = @views ŷ[estim.i_ym]
         v̂  = ŷm
         v̂ .= ym .- ŷm
-        x̂ .= f̂(estim, estim.model, x̂, u, d) .+ mul!(x̂, estim.K̂, v̂)
+        f̂!(x̂next, estim, estim.model, x̂, u, d)
+        mul!(x̂next, estim.K̂, v̂, 1, 1)
+        estim.x̂ .= x̂next
     end
     P̂.data .= Â * (P̂ .- M̂ * Ĉm * P̂) * Â' .+ Q̂ # .data is necessary for Hermitians
     return nothing
diff --git a/src/model/linmodel.jl b/src/model/linmodel.jl
@@ -217,23 +217,34 @@ disturbances ``\mathbf{d_0 = d - d_{op}}``. The Moore-Penrose pseudo-inverse com
 ``\mathbf{(I - A)^{-1}}`` to support integrating `model` (integrator states will be 0).
 """
 function steadystate!(model::LinModel, u, d)
-    model.x[:] = pinv(I - model.A)*(model.Bu*(u - model.uop) + model.Bd*(d - model.dop))
+    model.x .= pinv(I - model.A)*(model.Bu*(u - model.uop) + model.Bd*(d - model.dop))
     return nothing
 end
 
 """
-    f(model::LinModel, x, u, d)
+    f!(xnext, model::LinModel, x, u, d) -> xnext
 
-Evaluate ``\\mathbf{A x + B_u u + B_d d}`` when `model` is a [`LinModel`](@ref).
+Evaluate `xnext = A*x + Bu*u + Bd*d` in-place when `model` is a [`LinModel`](@ref).
 """
-f(model::LinModel, x, u, d) = model.A * x + model.Bu * u + model.Bd * d
+function f!(xnext, model::LinModel, x, u, d)
+    xnext .= 0
+    mul!(xnext, model.A,  x, 1, 1)
+    mul!(xnext, model.Bu, u, 1, 1)
+    mul!(xnext, model.Bd, d, 1, 1)
+    return xnext
+end
 
 
 """
-    h(model::LinModel, x, d)
+    h!(y, model::LinModel, x, d) -> y
 
-Evaluate ``\\mathbf{C x + D_d d}`` when `model` is a [`LinModel`](@ref).
+Evaluate `y = C*x + Dd*d` in-place when `model` is a [`LinModel`](@ref).
 """
-h(model::LinModel, x, d) = model.C*x + model.Dd*d
+function h!(y, model::LinModel, x, d)
+    y .= 0
+    mul!(y, model.C,  x, 1, 1)
+    mul!(y, model.Dd, d, 1, 1)
+    return y
+end
 
 typestr(model::LinModel) = "linear"
diff --git a/src/model/nonlinmodel.jl b/src/model/nonlinmodel.jl
@@ -122,13 +122,10 @@ end
 steadystate!(::SimModel, _ , _ ) = nothing
 
 
-
-
-
-"Call ``\\mathbf{f!(x, u, d)}`` with `model.f!` function for [`NonLinModel`](@ref)."
+"Call `f!(xnext, x, u, d)` with `model.f!` method for [`NonLinModel`](@ref)."
 f!(xnext, model::NonLinModel, x, u, d) = model.f!(xnext, x, u, d)
 
-"Call ``\\mathbf{h!(y, x, d)}`` with `model.h` function for [`NonLinModel`](@ref)."
+"Call `h!(y, x, d)` with `model.h` method for [`NonLinModel`](@ref)."
 h!(y, model::NonLinModel, x, d) = model.h!(y, x, d)
 
 typestr(model::NonLinModel) = "nonlinear"
diff --git a/src/sim_model.jl b/src/sim_model.jl
@@ -123,7 +123,7 @@ true
 
 """
 function initstate!(model::SimModel, u, d=empty(model.x))
-    validate_args(model::SimModel, u, d)
+    validate_args(model::SimModel, d, u)
     steadystate!(model, u, d)
     return model.x
 end
@@ -142,9 +142,11 @@ julia> x = updatestate!(model, [1])
  1.0
 ```
 """
-function updatestate!(model::SimModel, u, d=empty(model.x))
+function updatestate!(model::SimModel{NT}, u, d=empty(model.x)) where NT <: Real
     validate_args(model::SimModel, d, u)
-    f!(model.x, model, model.x, u .- model.uop, d .- model.dop)
+    xnext = Vector{NT}(undef, model.nx)
+    f!(xnext, model, model.x, u - model.uop, d - model.dop)
+    model.x .= xnext
     return model.x
 end
 
@@ -167,7 +169,8 @@ julia> y = evaloutput(model)
 function evaloutput(model::SimModel{NT}, d=empty(model.x)) where NT <: Real
     validate_args(model, d)
     y = Vector{NT}(undef, model.ny)
-    h!(y, model, model.x, d .- model.dop) .+ model.yop
+    h!(y, model, model.x, d - model.dop)
+    y .+= model.yop
     return y
 end
 
