added: 🟨📐 TrapezoidalCollocation 🟨📐 for MovingHorizonEstimator

Project.toml

@@ -1,6 +1,6 @@
 name = "ModelPredictiveControl"
 uuid = "61f9bdb8-6ae4-484a-811f-bbf86720c31c"
-version = "2.10.1"
+version = "2.11.0"
 authors = ["Francis Gagnon"]
 
 [deps]

src/controller/nonlinmpc.jl

@@ -584,7 +584,7 @@ function addinfo!(info, mpc::NonLinMPC{NT}) where NT<:Real
     hess = mpc.hessian
     transcription = mpc.transcription
     nu, ny, nx̂, nϵ = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ
-    nk = get_nk_mpc(model, transcription)
+    nk = get_nk(model, transcription)
     Hp, Hc = mpc.Hp, mpc.Hc
     i_g = findall(mpc.con.i_g) # convert to non-logical indices for non-allocating @views
     ng, ngi = length(mpc.con.i_g), sum(mpc.con.i_g)
@@ -823,7 +823,7 @@ function get_nonlinobj_op(mpc::NonLinMPC, optim::JuMP.GenericModel{JNT}) where J
     transcription = mpc.transcription
     grad, hess = mpc.gradient, mpc.hessian
     nu, ny, nx̂, nϵ = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ
-    nk = get_nk_mpc(model, transcription)
+    nk = get_nk(model, transcription)
     Hp, Hc = mpc.Hp, mpc.Hc
     ng = length(mpc.con.i_g)
     nc, neq = mpc.con.nc, mpc.con.neq
@@ -949,7 +949,7 @@ function get_nonlincon_oracle(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JN
     transcription = mpc.transcription
     jac, hess = mpc.jacobian, mpc.hessian
     nu, ny, nx̂, nϵ = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ
-    nk = get_nk_mpc(model, transcription)
+    nk = get_nk(model, transcription)
     Hp, Hc = mpc.Hp, mpc.Hc
     i_g = findall(mpc.con.i_g) # convert to non-logical indices for non-allocating @views
     ng, ngi = length(mpc.con.i_g), sum(mpc.con.i_g)

src/controller/transcription.jl

@@ -9,10 +9,6 @@ function get_nZ_mpc(estim::StateEstimator, transcription::OrthogonalCollocation,
     return estim.model.nu*Hc + estim.nx̂*Hp + estim.model.nx*transcription.no*Hp
 end
 
-"Get length of the `k` vector with all the solver intermediate steps or all the collocation pts."
-get_nk_mpc(model::SimModel, ::ShootingMethod) = model.nk
-get_nk_mpc(model::SimModel, transcription::CollocationMethod) = model.nx*transcription.no
-
 @doc raw"""
     init_predmat(
         model::LinModel, estim, transcription::SingleShooting, Hp, Hc, nb
@@ -561,7 +557,7 @@ function init_defectmat_orthocolloc(
     Hp, Hc, Co, λo, As, nxs
 ) where {NT<:Real}
     nu, nx, nd, nx̂ = model.nu, model.nx, model.nd, estim.nx̂
-    nk = get_nk_mpc(model, transcription)
+    nk = get_nk(model, transcription)
     λo_I = λo*I(nx)
     # --- current state estimates x̂0 ---
     KS = zeros(NT, nx̂*Hp, nx̂)
@@ -1049,7 +1045,7 @@ function set_warmstart_mpc!(
 )
     nu, nx̂ = mpc.estim.model.nu, mpc.estim.nx̂
     Hp, Hc, Z̃s  = mpc.Hp, mpc.Hc, mpc.buffer.Z̃
-    nk = get_nk_mpc(mpc.estim.model, transcription)
+    nk = get_nk(mpc.estim.model, transcription)
     nΔU, nX̂, nK = nu*Hc, nx̂*Hp, nk*Hp
     # --- input increments ΔU ---
     Z̃s[1:(nΔU-nu)]       .= @views mpc.Z̃[(nu+1):(nΔU)]
@@ -1403,7 +1399,7 @@ function con_nonlinprogeq!(
     nΔU, nX̂ = nu*Hc, nx̂*Hp
     f_threads = transcription.f_threads
     Ts, p = model.Ts, model.p
-    nk = get_nk_mpc(model, transcription)
+    nk = get_nk(model, transcription)
     D̂0 = mpc.D̂0
     X̂0_Z̃ = @views Z̃[(nΔU+1):(nΔU+nX̂)]
     Û0 = disturbedinput!(Û0, mpc.estim, mpc.estim.x̂0, X̂0_Z̃, U0)
@@ -1498,7 +1494,7 @@ function con_nonlinprogeq!(
     f_threads = transcription.f_threads
     p = model.p
     Mo, no, τ =  mpc.Mo, transcription.no, transcription.τ
-    nk = get_nk_mpc(model, transcription)
+    nk = get_nk(model, transcription)
     D̂0 = mpc.D̂0
     X̂0_Z̃, K_Z̃ = @views Z̃[(nΔU+1):(nΔU+nX̂)], Z̃[(nΔU+nX̂+1):(nΔU+nX̂+nk*Hp)]
     D̂temp = mpc.buffer.D̂

src/estimator/mhe/construct.jl

@@ -170,6 +170,10 @@ struct MovingHorizonEstimator{
         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)
+        validate_transcription(model, transcription)
+        if transcription isa OrthogonalCollocation 
+            error("OrthogonalCollocation is not supported for the MHE for now.")
+        end
         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
@@ -329,7 +333,7 @@ at each time step for the optimization.
    (details in Extended Help).
 - `nc=0` : number of custom nonlinear inequality constraints.
 - `p=model.p` : ``\mathbf{g_c}`` functions parameter ``\mathbf{p}`` (any type).
-- `transcription=SingleShooting()` : [`SingleShooting`](@ref) or [`MultipleShooting`](@ref).
+- `transcription=SingleShooting()` : a [`TranscriptionMethod`](@ref) for the optimization.
 - `optim=default_optim_mhe(model,nc)` : a [`JuMP.Model`](@extref) object with a quadratic or
    nonlinear optimizer for solving (default to [`Ipopt`](https://github.com/jump-dev/Ipopt.jl),
    or [`OSQP`](https://osqp.org/docs/parsers/jump.html) if `model` is a [`LinModel`](@ref)).
@@ -539,7 +543,7 @@ function MovingHorizonEstimator(
     gc ::Function = gc!,
     nc ::Int = 0,
     p = model.p,
-    transcription::ShootingMethod = DEFAULT_MHE_TRANSCRIPTION,
+    transcription::TranscriptionMethod = DEFAULT_MHE_TRANSCRIPTION,
     optim::JM = default_optim_mhe(model, nc),
     gradient::AbstractADType = DEFAULT_GRADIENT,
     jacobian::AbstractADType = default_jacobian(transcription),
@@ -593,7 +597,7 @@ function MovingHorizonEstimator(
     gc ::Function = gc!,
     nc = 0,
     p = model.p,
-    transcription::ShootingMethod = DEFAULT_MHE_TRANSCRIPTION,
+    transcription::TranscriptionMethod = DEFAULT_MHE_TRANSCRIPTION,
     optim::JM = default_optim_mhe(model, nc),
     gradient::AbstractADType = DEFAULT_GRADIENT,
     jacobian::AbstractADType = default_jacobian(transcription),
@@ -1416,7 +1420,8 @@ function get_nonlinobj_op(
 ) where JNT<:Real
     model, con = estim.model, estim.con
     grad, hess = estim.gradient, estim.hessian
-    nx̂, nym, nŷ, nu, nk = estim.nx̂, estim.nym, model.ny, model.nu, model.nk
+    nx̂, nym, nŷ, nu = estim.nx̂, estim.nym, model.ny, model.nu
+    nk = get_nk(model, estim.transcription)
     He = estim.He
     nc, neq, ng = con.nc, con.neq, length(con.i_g)
     nŴ, nV̂, nX̂, ng, nZ̃ = He*nx̂, He*nym, He*nx̂, length(con.i_g), length(estim.Z̃)
@@ -1531,7 +1536,8 @@ function get_nonlincon_oracle(
     # ----------- common cache for all functions  ----------------------------------------
     model, con = estim.model, estim.con
     jac, hess = estim.jacobian, estim.hessian
-    nx̂, nym, nŷ, nu, nk = estim.nx̂, estim.nym, model.ny, model.nu, model.nk
+    nx̂, nym, nŷ, nu = estim.nx̂, estim.nym, model.ny, model.nu
+    nk = get_nk(model, estim.transcription)
     He = estim.He
     nc, neq, ng = con.nc, con.neq, length(con.i_g)
     i_g = findall(con.i_g) # convert to non-logical indices for non-allocating @views

src/estimator/mhe/execute.jl

@@ -213,7 +213,8 @@ function addinfo!(info, estim::MovingHorizonEstimator{NT}, model::SimModel) wher
     # --- objective derivatives ---
     optim, con = estim.optim, estim.con
     hess = estim.hessian
-    nx̂, nym, nŷ, nu, nk, nc = estim.nx̂, estim.nym, model.ny, model.nu, model.nk, con.nc
+    nx̂, nym, nŷ, nu, nc = estim.nx̂, estim.nym, model.ny, model.nu, con.nc
+    nk = get_nk(model, estim.transcription)
     He = estim.He
     nc, neq, ng = con.nc, con.neq, length(con.i_g)
     i_g = findall(con.i_g) # convert to non-logical indices for non-allocating @views

src/estimator/mhe/transcription.jl

@@ -914,7 +914,7 @@ function linconstrainteq!(
         JuMP.delete(optim, optim[:linconstrainteq_temp])
         JuMP.unregister(optim, :linconstrainteq_temp)
     end
-    if estim.Nk[] < estim.He
+    if Nk < estim.He
         if haskey(optim, :linconstrainteq)
             JuMP.delete(optim, optim[:linconstrainteq])
             JuMP.unregister(optim, :linconstrainteq)
@@ -1002,10 +1002,10 @@ end
 
 @doc raw"""
     set_warmstart_mhe!(
-        estim::MovingHorizonEstimator, transcription::MultipleShooting, Z̃var
+        estim::MovingHorizonEstimator, transcription::TranscriptionMethod, Z̃var
     ) -> Z̃s
 
-Do the same but based on a [`MultipleShooting`](@ref) transcription.
+Do the same but for other transcription [`TranscriptionMethod`](@ref).
 
 If supported by `estim.optim`, it warm-starts the solver at:
 ```math
@@ -1035,7 +1035,7 @@ last control period ``k-1``, expressed as a deviation from the operating point
 of vectors ``\mathbf{0_x̂}`` and ``\mathbf{0_ŵ}``.
 """
 function set_warmstart_mhe!(
-    estim::MovingHorizonEstimator{NT}, transcription::MultipleShooting, Z̃var
+    estim::MovingHorizonEstimator{NT}, transcription::TranscriptionMethod, Z̃var
 ) where NT<:Real
     model, buffer = estim.model, estim.buffer
     nu, nk = model.nu, model.nk
@@ -1342,13 +1342,14 @@ end
 
 Nonlinear MHE equality constrains for [`NonLinModel`](@ref) and [`MultipleShooting`](@ref).
 
-The method mutates the `geq`, `X̂0`, `Û0` and `K` vectors in argument. The defects of the
-estimated states are computed with:
+The method mutates the `geq`, `X̂0`, `Û0` and `K` vectors in argument. By introducing 
+the integer ``ℓ = k - N_k + p`` to shorten the notation, the defects of the estimated states
+are computed with:
 ```math
-\mathbf{ŝ}(k+j+1) = \mathbf{f̂}\Big(\mathbf{x̂_0}(k+j), \mathbf{u_0}(k+j), \mathbf{d_0}(k+j)\Big) 
-                      - \mathbf{x̂_0}(k+j+1)
+\mathbf{ŝ}(ℓ+j+1) = \mathbf{f̂}\Big(\mathbf{x̂_0}(ℓ+j), \mathbf{u_0}(ℓ+j), \mathbf{d_0}(ℓ+j)\Big) 
+                           + \mathbf{ŵ}(ℓ+j) - \mathbf{x̂_0}(ℓ+j+1)
 ```
-for ``j = 0, 1, ... , H_p-1`` and in which the augmented state vectors ``\mathbf{x̂_0}`` are 
+for ``j = 0, 1, ... , N_k-1`` and in which the augmented state vectors ``\mathbf{x̂_0}`` are 
 extracted from the decision variable `Z̃`. The function ``\mathbf{f̂}`` is defined at [`f̂!`](@ref).
 """
 function con_nonlinprogeq_mhe!(
@@ -1385,5 +1386,85 @@ function con_nonlinprogeq_mhe!(
     Nk < He && (geq[nx̂*Nk+1:end] .= 0)
     return geq
 end
+
+@doc raw"""
+    con_nonlinprogeq_mhe!(
+        geq, X̂0, Û0, K̇,
+        estim::MovingHorizonEstimator, model::NonLinModel, ::TrapezoidalCollocation, 
+        x̂0arr, Ŵ, Z̃
+    ) -> geq
+
+Nonlinear MHE equality constrains for [`NonLinModel`](@ref) and [`TrapezoidalCollocation`](@ref).
+
+By introducing the integer ``ℓ = k - N_k + p`` to shorten the notation, the deterministic
+state defects are computed with:
+```math
+\mathbf{ŝ_d}(ℓ+j+1) = \mathbf{x̂_d}(ℓ+j) + 0.5 T_s [\mathbf{k̇}_1(ℓ+j) + \mathbf{k̇}_2(ℓ+j)] 
+                       + \mathbf{ŵ_d}(ℓ+j) - \mathbf{x̂_d}(ℓ+j+1)                                              
+```
+for ``j = 0, 1, ... , N_k-1``, and in which ``\mathbf{x̂_d}`` and ``\mathbf{ŵ_d}`` are the
+deterministic state and process noise estimates, respectively, extracted from the decision
+variable `Z̃`. The ``\mathbf{k̇}`` coefficients are  evaluated from the continuous-time
+function `model.f!` and:
+```math
+\begin{aligned}
+\mathbf{k̇}_1(ℓ+j) &= \mathbf{f}\Big(\mathbf{x̂_d}(ℓ+j),   \mathbf{û_0}(ℓ+j),   \mathbf{d̂_0}(ℓ+j),   \mathbf{p}\Big) \\
+\mathbf{k̇}_2(ℓ+j) &= \mathbf{f}\Big(\mathbf{x̂_d}(ℓ+j+1), \mathbf{û_0}(ℓ+j+h), \mathbf{d̂_0}(ℓ+j+1), \mathbf{p}\Big) 
+\end{aligned}
+```
+in which ``h`` is the hold order `transcription.h` and the disturbed input ``\mathbf{û_0}``
+is defined in [`f̂!`](@ref) documentation.
+"""
+function con_nonlinprogeq_mhe!(
+    geq, _ , Û0, K̇,
+    estim::MovingHorizonEstimator, model::NonLinModel, transcription::TrapezoidalCollocation, 
+    x̂0arr, Ŵ, Z̃
+)
+    nu, nx, nd, h = model.nu, model.nx, model.nd, transcription.h
+    nx̂, nxs, nŵ, He = estim.nx̂, estim.nxs, estim.nx̂, estim.He
+    Nk = estim.Nk[]
+    f_threads = transcription.f_threads
+    Ts = model.Ts
+    nk = get_nk(model, transcription)
+    nw = nŵ - nxs
+    nx̃ = estim.nε + nx̂
+    p = estim.direct ? 0 : 1
+    X̂0_Z̃ = @views Z̃[(nx̃+1):(nx̃+nx̂*He)]
+    Û0 = disturbedinput!(Û0, estim, x̂0arr, X̂0_Z̃, estim.U0)
+    @threadsif f_threads for j=1:Nk
+        if j < 2
+            x̂d_Z̃ = @views x̂0arr[1:nx]
+        else
+            x̂d_Z̃ = @views X̂0_Z̃[(1 + nx̂*(j-2)):(nx̂*(j-2) + nx)]
+        end
+        d0       = @views   estim.D0[(1 + nd*(j+p-1)):(nd*(j+p))]
+        û0       = @views         Û0[(1 + nu*(j-1)):(nu*j)]
+        k̇        = @views          K̇[(1 + nk*(j-1)):(nk*j)]
+        ŵd       = @views          Ŵ[(1 + nŵ*(j-1)):(nŵ*(j-1) + nw)]
+        x̂dnext_Z̃ = @views       X̂0_Z̃[(1 + nx̂*(j-1)):(nx̂*(j-1) + nx)]
+        sdnext   = @views        geq[(1 + nx*(j-1)):(nx*j)]
+        k̇1, k̇2   = @views          k̇[1:nx], k̇[nx+1:2*nx]    
+        d0next   = @views   estim.D0[(1 + nd*(j+p)):(nd*(j+p+1))]
+        if f_threads || h < 1 || j < 2
+            # we need to recompute k1 with multi-threading, even with h==1, since the 
+            # last iteration (j-1) may not be executed (iterations are re-orderable)
+            model.f!(k̇1, x̂d_Z̃, û0, d0, model.p)
+        else
+            k̇1 .= @views K̇[(1 + nk*(j-1)-nx):(nk*(j-1))] # k2 of of the last iter. j-1
+        end
+        if h < 1
+            model.f!(k̇2, x̂dnext_Z̃, û0, d0next, model.p)
+        else
+            # special case: û0(k+p) ≈ û0(k+p-1), since û0(k+p) is not available at k!
+            û0next = @views j ≥ Nk ? û0 : Û0[(1 + nu*j):(nu*(j+1))]
+            model.f!(k̇2, x̂dnext_Z̃, û0next, d0next, model.p)
+        end
+        sdnext  .= @. x̂d_Z̃ - x̂dnext_Z̃ + 0.5*Ts*(k̇1 + k̇2)
+        sdnext .+= ŵd
+    end
+    Nk < He && (geq[nx̂*Nk+1:end] .= 0)
+    return geq
+end
+
 "No nonlinear eq. const. for other cases e.g. [`SingleShooting`](@ref), returns `geq` unchanged."
 con_nonlinprogeq_mhe!(geq,_,_,_,::MovingHorizonEstimator, ::SimModel, ::TranscriptionMethod, _,_,_) = geq

src/transcription.jl

@@ -136,9 +136,11 @@ equality constraint function and by using the implicit trapezoidal rule. It can
 moderately stiff systems and is A-stable. See Extended Help for more details.
 
 !!! warning
-    The built-in [`StateEstimator`](@ref) will still use the `solver` provided at the
-    construction of the [`NonLinModel`](@ref) to estimate the plant states, not the 
-    trapezoidal rule (see `supersample` option of  [`RungeKutta`](@ref) for stiff systems).
+    Except if you construct your MPC with a [`MovingHorizonEstimator`](@ref) based on a
+    `TrapezoidalCollocation` transcription, the built-in [`StateEstimator`](@ref) will
+    still use the `solver` provided at the construction of the [`NonLinModel`](@ref) to
+    estimate the plant states, not the trapezoidal rule (see `supersample` option of
+    [`RungeKutta`](@ref) for stiff systems).
 
 Sparse optimizers like `Ipopt` and sparse Jacobian computations are recommended for this
 transcription method.
@@ -149,7 +151,7 @@ transcription method.
     as described in [`ModelPredictiveControl.init_estimstoch`](@ref). Collocation methods
     require continuous-time dynamics. Because of this, the stochastic states are transcribed
     separately using a [`MultipleShooting`](@ref) method. See [`con_nonlinprogeq!`](@ref)
-    for more details.
+    and [`con_nonlinprogeq_mhe!`](@ref) for more details.
 """
 struct TrapezoidalCollocation <: CollocationMethod
     h::Int
@@ -200,7 +202,8 @@ Gauss-Legendre quadrature, respectively. See [`MultipleShooting`](@ref) docstrin
 descriptions of `f_threads` and `h_threads` keywords. This transcription computes the
 predictions by enforcing the collocation and continuity constraints at the collocation
 points. It is efficient for highly stiff systems, but generally more expensive than the
-other methods for non-stiff systems. See Extended Help for more details.
+other methods for non-stiff systems. This transcription is not supported by the
+[`MovingHorizonEstimator`](@ref) for now. See Extended Help for more details.
 
 !!! warning
     The built-in [`StateEstimator`](@ref) will still use the `solver` provided at the
@@ -352,3 +355,7 @@ function validate_transcription(::NonLinModel{<:Real, <:EmptySolver}, ::Collocat
     return nothing
 end
 validate_transcription(::SimModel, ::TranscriptionMethod) = nothing
+
+"Get length of the `k` vector with all the solver intermediate steps or all the collocation pts."
+get_nk(model::SimModel, ::ShootingMethod) = model.nk
+get_nk(model::SimModel, transcription::CollocationMethod) = model.nx*transcription.no