added: custom estimator for MHE covariance at arrival (now default to UKF)

Author: franckgaga

src/estimator/kalman.jl

@@ -163,7 +163,7 @@ This syntax allows nonzero off-diagonal elements in ``\mathbf{Q̂, R̂}``.
 """
 function SteadyKalmanFilter(model::SM, i_ym, nint_u, nint_ym, Q̂, R̂) where {NT<:Real, SM<:LinModel{NT}}
     Q̂, R̂ = to_mat(Q̂), to_mat(R̂)
-    return SteadyKalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, Q̂ , R̂)
+    return SteadyKalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, Q̂, R̂)
 end
 
 
@@ -315,7 +315,7 @@ This syntax allows nonzero off-diagonal elements in ``\mathbf{P̂}_{-1}(0), \mat
 """
 function KalmanFilter(model::SM, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂) where {NT<:Real, SM<:LinModel{NT}}
     P̂0, Q̂, R̂ = to_mat(P̂0), to_mat(Q̂), to_mat(R̂)
-    return KalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, P̂0, Q̂ , R̂)
+    return KalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
 end
 
 @doc raw"""
@@ -478,14 +478,14 @@ function UnscentedKalmanFilter(
 end
 
 @doc raw"""
-    UnscentedKalmanFilter(model, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, α, β, κ)
+    UnscentedKalmanFilter(model, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, α=1e-3, β=2, κ=0)
 
 Construct the estimator from the augmented covariance matrices `P̂0`, `Q̂` and `R̂`.
 
 This syntax allows nonzero off-diagonal elements in ``\mathbf{P̂}_{-1}(0), \mathbf{Q̂, R̂}``.
 """
 function UnscentedKalmanFilter(
-    model::SM, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, α, β, κ
+    model::SM, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, α=1e-3, β=2, κ=0
 ) where {NT<:Real, SM<:SimModel{NT}}
     P̂0, Q̂, R̂ = to_mat(P̂0), to_mat(Q̂), to_mat(R̂)
     return UnscentedKalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, P̂0, Q̂ , R̂, α, β, κ)
@@ -742,7 +742,7 @@ This syntax allows nonzero off-diagonal elements in ``\mathbf{P̂}_{-1}(0), \mat
 """
 function ExtendedKalmanFilter(model::SM, i_ym, nint_u, nint_ym,P̂0, Q̂, R̂) where {NT<:Real, SM<:SimModel{NT}}
     P̂0, Q̂, R̂ = to_mat(P̂0), to_mat(Q̂), to_mat(R̂)
-    return ExtendedKalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, P̂0, Q̂ , R̂)
+    return ExtendedKalmanFilter{NT, SM}(model, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
 end
 
 
@@ -759,7 +759,8 @@ substitutions ``\mathbf{Â = F̂}(k)`` and ``\mathbf{Ĉ^m = Ĥ^m}(k)``:
                            [\mathbf{Ĥ^m}(k)\mathbf{P̂}_{k-1}(k)\mathbf{Ĥ^m}'(k) + \mathbf{R̂}]^{-1}    \\
     \mathbf{K̂}(k)       &= \mathbf{F̂}(k) \mathbf{M̂}(k)                                    \\
     \mathbf{ŷ^m}(k)     &= \mathbf{ĥ^m}\Big( \mathbf{x̂}_{k-1}(k), \mathbf{d}(k) \Big)     \\
-    \mathbf{x̂}_{k}(k+1) &= \mathbf{f̂}\Big( \mathbf{x̂}_{k-1}(k), \mathbf{u}(k), \mathbf{d}(k) \Big)
+    \mathbf{x̂}_{k}(k+1) &= \math, covestim::CE,
+    )    bf{f̂}\Big( \mathbf{x̂}_{k-1}(k), \mathbf{u}(k), \mathbf{d}(k) \Big)
                            + \mathbf{K̂}(k)[\mathbf{y^m}(k) - \mathbf{ŷ^m}(k)]             \\
     \mathbf{P̂}_{k}(k+1) &= \mathbf{F̂}(k)[\mathbf{P̂}_{k-1}(k)
                            - \mathbf{M̂}(k)\mathbf{Ĥ^m}(k)\mathbf{P̂}_{k-1}(k)]\mathbf{F̂}'(k) 
@@ -770,7 +771,8 @@ substitutions ``\mathbf{Â = F̂}(k)`` and ``\mathbf{Ĉ^m = Ĥ^m}(k)``:
 automatically computes the Jacobians:
 ```math
 \begin{aligned}
-    \mathbf{F̂}(k) &= \left. \frac{∂\mathbf{f̂}(\mathbf{x̂}, \mathbf{u}, \mathbf{d})}{∂\mathbf{x̂}} \right|_{\mathbf{x̂ = x̂}_{k-1}(k),\, \mathbf{u = u}(k),\, \mathbf{d = d}(k)}  \\
+    \mathbf{F̂}(k) &= \left. \fra, covestim::CE,
+    )    c{∂\mathbf{f̂}(\mathbf{x̂}, \mathbf{u}, \mathbf{d})}{∂\mathbf{x̂}} \right|_{\mathbf{x̂ = x̂}_{k-1}(k),\, \mathbf{u = u}(k),\, \mathbf{d = d}(k)}  \\
     \mathbf{Ĥ}(k) &= \left. \frac{∂\mathbf{ĥ}(\mathbf{x̂}, \mathbf{d})}{∂\mathbf{x̂}}             \right|_{\mathbf{x̂ = x̂}_{k-1}(k),\, \mathbf{d = d}(k)}
 \end{aligned}
 ```

src/estimator/mhe/construct.jl

@@ -43,14 +43,16 @@ end
 
 struct MovingHorizonEstimator{
     NT<:Real, 
-    SM<:SimModel, 
-    JM<:JuMP.GenericModel
+    SM<:SimModel,
+    JM<:JuMP.GenericModel,
+    CE<:StateEstimator,
 } <: StateEstimator{NT}
     model::SM
     # note: `NT` and the number type `JNT` in `JuMP.GenericModel{JNT}` can be
     # different since solvers that support non-Float64 are scarce.
     optim::JM
     con::EstimatorConstraint{NT}
+    covestim::CE
     Z̃::Vector{NT}
     lastu0::Vector{NT}
     x̂::Vector{NT}
@@ -95,17 +97,16 @@ struct MovingHorizonEstimator{
     x̂arr_old::Vector{NT}
     P̂arr_old::Hermitian{NT, Matrix{NT}}
     Nk::Vector{Int}
-    function MovingHorizonEstimator{NT, SM, JM}(
-        model::SM, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim::JM
-    ) where {NT<:Real, SM<:SimModel{NT}, JM<:JuMP.GenericModel}
+    function MovingHorizonEstimator{NT, SM, JM, CE}(
+        model::SM, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim::JM, covestim::CE
+    ) where {NT<:Real, SM<:SimModel{NT}, JM<:JuMP.GenericModel, CE<:StateEstimator{NT}}
         nu, nd = model.nu, model.nd
         He < 1  && throw(ArgumentError("Estimation horizon He should be ≥ 1"))
         Cwt < 0 && throw(ArgumentError("Cwt weight should be ≥ 0"))
         nym, nyu = validate_ym(model, i_ym)
         As, Cs_u, Cs_y, nint_u, nint_ym = init_estimstoch(model, i_ym, nint_u, nint_ym)
         nxs = size(As, 1)
         nx̂  = model.nx + nxs
-        nŵ = nx̂
         Â, B̂u, Ĉ, B̂d, D̂d = augment_model(model, As, Cs_u, Cs_y)
         validate_kfcov(nym, nx̂, Q̂, R̂, P̂0)
         lastu0 = zeros(NT, model.nu)
@@ -129,8 +130,8 @@ struct MovingHorizonEstimator{
         x̂arr_old = zeros(NT, nx̂)
         P̂arr_old = copy(P̂0)
         Nk = [0]
-        estim = new{NT, SM, JM}(
-            model, optim, con, 
+        estim = new{NT, SM, JM, CE}(
+            model, optim, con, covestim,  
             Z̃, lastu0, x̂, 
             He,
             i_ym, nx̂, nym, nyu, nxs, 
@@ -188,7 +189,7 @@ The vectors ``\mathbf{Ŵ}`` and ``\mathbf{V̂}`` encompass the estimated proces
 ``\mathbf{ŵ}(k-j)`` and sensor noise ``\mathbf{v̂}(k-j)`` from ``j=N_k-1`` to ``0``. The 
 Extended Help defines the two vectors and the scalar ``ϵ``. See [`UnscentedKalmanFilter`](@ref)
 for details on the augmented process model and ``\mathbf{R̂}, \mathbf{Q̂}`` covariances. The
-matrix ``\mathbf{P̂}_{k-N_k}(k-N_k+1)`` is estimated with an [`ExtendedKalmanFilter`](@ref).
+covariance ``\mathbf{P̂}_{k-N_k}(k-N_k+1)`` is estimated with an [`ExtendedKalmanFilter`](@ref).
 
 !!! warning
     See the Extended Help if you get an error like:    
@@ -262,7 +263,7 @@ MovingHorizonEstimator estimator with a sample time Ts = 10.0 s, Ipopt optimizer
 """
 function MovingHorizonEstimator(
     model::SM;
-    He::Union{Int, Nothing}=nothing,
+    He::Union{Int, Nothing} = nothing,
     i_ym::IntRangeOrVector = 1:model.ny,
     σP0::Vector = fill(1/model.nx, model.nx),
     σQ ::Vector = fill(1/model.nx, model.nx),
@@ -280,33 +281,40 @@ function MovingHorizonEstimator(
     P̂0 = Hermitian(diagm(NT[σP0; σP0int_u; σP0int_ym].^2), :L)
     Q̂  = Hermitian(diagm(NT[σQ;  σQint_u;  σQint_ym ].^2), :L)
     R̂  = Hermitian(diagm(NT[σR;].^2), :L)
-    isnothing(He) && throw(ArgumentError("Estimation horizon He must be explicitly specified"))        
-    return MovingHorizonEstimator{NT, SM, JM}(
-        model, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim
-    )
+    isnothing(He) && throw(ArgumentError("Estimation horizon He must be explicitly specified")) 
+    return MovingHorizonEstimator(model, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim)
 end
 
-"Return a `JuMP.Model` with OSQP optimizer if `model` is a [`LinModel`](@ref)."
 default_optim_mhe(::LinModel) = JuMP.Model(DEFAULT_LINMHE_OPTIMIZER, add_bridges=false)
-"Else, return it with Ipopt optimizer."
 default_optim_mhe(::SimModel) = JuMP.Model(DEFAULT_NONLINMHE_OPTIMIZER, add_bridges=false)
 
 @doc raw"""
-    MovingHorizonEstimator(model, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim)
+    MovingHorizonEstimator(model, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim[, covestim])
 
 Construct the estimator from the augmented covariance matrices `P̂0`, `Q̂` and `R̂`.
 
 This syntax allows nonzero off-diagonal elements in ``\mathbf{P̂}_{-1}(0), \mathbf{Q̂, R̂}``.
+The final argument `covestim` also allows specifying a custom [`StateEstimator`](@ref) 
+object for the estimation of covariance at the arrival ``\mathbf{P̂}_{k-N_k}(k-N_k+1)``. The
+supported types are [`KalmanFilter`](@ref), [`UnscentedKalmanFilter`](@ref) and 
+[`ExtendedKalmanFilter`](@ref).
 """
 function MovingHorizonEstimator(
-    model::SM, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim::JM
-) where {NT<:Real, SM<:SimModel{NT}, JM<:JuMP.GenericModel}
+    model::SM, He, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂, Cwt, optim::JM, 
+    covestim::CE = default_covestim_mhe(model, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
+) where {NT<:Real, SM<:SimModel{NT}, JM<:JuMP.GenericModel, CE<:StateEstimator{NT}}
     P̂0, Q̂, R̂ = to_mat(P̂0), to_mat(Q̂), to_mat(R̂)
-    return MovingHorizonEstimator{NT, SM, JM}(
-        model, He, i_ym, nint_u, nint_ym, P̂0, Q̂ , R̂, Cwt, optim
+    return MovingHorizonEstimator{NT, SM, JM, CE}(
+        model, He, i_ym, nint_u, nint_ym, P̂0, Q̂ , R̂, Cwt, optim, covestim
     )
 end
 
+function default_covestim_mhe(model::LinModel, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
+    return KalmanFilter(model, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
+end
+function default_covestim_mhe(model::SimModel, i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
+    return UnscentedKalmanFilter(model,  i_ym, nint_u, nint_ym, P̂0, Q̂, R̂)
+end
 
 @doc raw"""
     setconstraint!(estim::MovingHorizonEstimator; <keyword arguments>) -> estim

src/estimator/mhe/execute.jl

@@ -78,11 +78,7 @@ function update_estimate!(estim::MovingHorizonEstimator{NT}, u, ym, d) where NT<
     estim.Ŵ[1:nŵ*Nk] .= @views estim.Z̃[nx̃+1:nx̃+nŵ*Nk] # update Ŵ with optimum for warm-start
     V̂, X̂ = predict!(V̂, X̂, ŷ, estim, model, estim.Z̃)
     x̂ .= X̂[end-nx̂+1:end]
-    if Nk == estim.He
-        uarr, ymarr, darr = estim.U[1:model.nu], estim.Ym[1:nym], estim.D[1:model.nd]
-        update_cov!(estim, model, estim.x̂arr_old, estim.P̂arr_old, uarr, ymarr, darr)
-        estim.invP̄.data .= Hermitian(inv(estim.P̂arr_old), :L)
-    end
+    Nk == estim.He && update_cov!(estim::MovingHorizonEstimator)
     return nothing
 end
 
@@ -323,20 +319,18 @@ function trunc_bounds(estim::MovingHorizonEstimator{NT}, Bmin, Bmax, n) where NT
     return Bmin_t, Bmax_t
 end
 
-"Update the covariance `P̂` with the `KalmanFilter` if `model` is a `LinModel`."
-function update_cov!(estim::MovingHorizonEstimator, ::LinModel, _ , P̂, u, ym, d) 
-    return update_estimate_kf!(estim, u, ym, d, estim.Â, estim.Ĉ[estim.i_ym, :], P̂)
-end
-"Update it with the `ExtendedKalmanFilter` if model is not a `LinModel`."
-function update_cov!(estim::MovingHorizonEstimator, model::SimModel, x̂, P̂, u, ym, d) 
-    # TODO: also support UnscentedKalmanFilter
-    x̂next, ŷ = similar(estim.x̂), similar(model.yop)
-    F̂ = ForwardDiff.jacobian((x̂next, x̂) -> f̂!(x̂next, estim, estim.model, x̂, u, d), x̂next, x̂)
-    Ĥ = ForwardDiff.jacobian((ŷ, x̂)     -> ĥ!(ŷ, estim, estim.model, x̂, d), ŷ, x̂)
-    return update_estimate_kf!(estim, u, ym, d, F̂, Ĥ[estim.i_ym, :], P̂)
+"Update the covariance estimate at arrival using `covestim` [`StateEstimator`](@ref)."
+function update_cov!(estim::MovingHorizonEstimator)
+    nu, nd, nym = estim.model.nu, estim.model.nd, estim.nym
+    uarr, ymarr, darr = @views estim.U[1:nu], estim.Ym[1:nym], estim.D[1:nd]
+    estim.covestim.x̂      .= estim.x̂arr_old
+    estim.covestim.P̂.data .= estim.P̂arr_old # .data is necessary for Hermitian
+    update_estimate!(estim.covestim, uarr, ymarr, darr)
+    estim.P̂arr_old.data   .= estim.covestim.P̂
+    estim.invP̄.data       .= Hermitian(inv(estim.P̂arr_old), :L)
+    return nothing
 end
 
-
 """
     obj_nonlinprog!( _ , estim::MovingHorizonEstimator, ::LinModel, _ , Z̃)