debug: add JuMP prefix where needed

Author: franckgaga

src/ModelPredictiveControl.jl

@@ -16,7 +16,7 @@ import ControlSystemsBase: iscontinuous, isdiscrete, sminreal, minreal, c2d, d2c
 
 import JuMP
 import JuMP: MOIU, MOI, GenericModel, Model, optimizer_with_attributes, register
-import JuMP: @variable, @constraint, @objective, @NLconstraint, @NLobjective
+import JuMP: @variable, @operator, @constraint, @objective
 
 import PreallocationTools: DiffCache, get_tmp
 

src/controller/nonlinmpc.jl

@@ -319,22 +319,22 @@ function init_optimization!(mpc::NonLinMPC, model::SimModel, optim)
     if length(con.i_g) ≠ 0
         for i in eachindex(con.Y0min)
             name = Symbol("g_Y0min_$i")
-            optim[name] = add_nonlinear_operator(optim, nΔŨ, gfunc[i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i]; name)
         end
         i_end_Ymin = 1Hp*ny
         for i in eachindex(con.Y0max)
             name = Symbol("g_Y0max_$i")
-            optim[name] = add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_Ymin+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_Ymin+i]; name)
         end
         i_end_Ymax = 2Hp*ny
         for i in eachindex(con.x̂0min)
             name = Symbol("g_x̂0min_$i")
-            optim[name] = add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_Ymax+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_Ymax+i]; name)
         end
         i_end_x̂min = 2Hp*ny + nx̂
         for i in eachindex(con.x̂0max)
             name = Symbol("g_x̂0max_$i")
-            optim[name] = add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_x̂min+i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_x̂min+i]; name)
         end
     end
     return nothing
@@ -402,23 +402,23 @@ function setnonlincon!(
 ) where JNT<:Real
     ΔŨvar = optim[:ΔŨvar]
     con = mpc.con
-    nonlin_constraints = all_constraints(optim, NonlinearExpr, MOI.LessThan{JNT})
-    map(con_ref -> JuMP.delete(optim, con_ref), JuMP.nonlin_constraints)
+    nonlin_constraints = JuMP.all_constraints(optim, JuMP.NonlinearExpr, MOI.LessThan{JNT})
+    map(con_ref -> JuMP.delete(optim, con_ref), nonlin_constraints)
     for i in findall(.!isinf.(con.Y0min))
         gfunc_i = optim[Symbol("g_Y0min_$(i)")]
-        JuMP.@constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     for i in findall(.!isinf.(con.Y0max))
         gfunc_i = optim[Symbol("g_Y0max_$(i)")]
-        JuMP.@constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     for i in findall(.!isinf.(con.x̂0min))
         gfunc_i = optim[Symbol("g_x̂0min_$(i)")]
-        JuMP.@constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     for i in findall(.!isinf.(con.x̂0max))
         gfunc_i = optim[Symbol("g_x̂0max_$(i)")]
-        JuMP.@constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     return nothing
 end

src/estimator/mhe/construct.jl

@@ -606,15 +606,15 @@ function setnonlincon!(
 ) where JNT<:Real
     optim, con = estim.optim, estim.con
     Z̃var = optim[:Z̃var]
-    nonlin_constraints = all_constraints(optim, NonlinearExpr, MOI.LessThan{JNT})
-    map(con_ref -> JuMP.delete(optim, con_ref), JuMP.nonlin_constraints)
+    nonlin_constraints = JuMP.all_constraints(optim, JuMP.NonlinearExpr, MOI.LessThan{JNT})
+    map(con_ref -> JuMP.delete(optim, con_ref), nonlin_constraints)
     for i in findall(.!isinf.(con.X̂0min))
         gfunc_i = optim[Symbol("g_X̂0min_$(i)")]
-        JuMP.@constraint(optim, gfunc_i(Z̃var...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.X̂0max))
         gfunc_i = optim[Symbol("g_X̂0max_$(i)")]
-        JuMP.@constraint(optim, gfunc_i(Z̃var...) <= 0)
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.V̂min))
         gfunc_i = optim[Symbol("g_V̂min_$(i)")]
@@ -1082,28 +1082,90 @@ function init_optimization!(
     if length(con.i_g) ≠ 0
         for i in eachindex(con.X̂0min)
             name = Symbol("g_X̂0min_$i")
-            optim[name] = add_nonlinear_operator(optim, nZ̃, gfunc[i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(
+                optim, nZ̃, gfunc[i]; name
+            )
         end
         i_end_X̂min = nX̂
         for i in eachindex(con.X̂0max)
             name = Symbol("g_X̂0max_$i")
-            optim[name] = add_nonlinear_operator(optim, nZ̃, gfunc[i_end_X̂min + i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(
+                optim, nZ̃, gfunc[i_end_X̂min + i]; name
+            )
         end
         i_end_X̂max = 2*nX̂
         for i in eachindex(con.V̂min)
             name = Symbol("g_V̂min_$i")
-            optim[name] = add_nonlinear_operator(optim, nZ̃, gfunc[i_end_X̂max + i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(
+                optim, nZ̃, gfunc[i_end_X̂max + i]; name
+            )
         end
         i_end_V̂min = 2*nX̂ + nV̂
         for i in eachindex(con.V̂max)
             name = Symbol("g_V̂max_$i")
-            optim[name] = add_nonlinear_operator(optim, nZ̃, gfunc[i_end_V̂min + i]; name)
+            optim[name] = JuMP.add_nonlinear_operator(
+                optim, nZ̃, gfunc[i_end_V̂min + i]; name
+            )
         end
     end
     return nothing
 end
 
 
+"""
+    get_optim_functions(estim::MovingHorizonEstimator, ::JuMP.GenericModel) -> Jfunc, gfunc
+
+Get the objective `Jfunc` and constraints `gfunc` functions for [`MovingHorizonEstimator`](@ref).
+
+Inspired from: [User-defined operators with vector outputs](https://jump.dev/JuMP.jl/stable/tutorials/nonlinear/tips_and_tricks/#User-defined-operators-with-vector-outputs)
+"""
+function get_optim_functions(
+    estim::MovingHorizonEstimator, ::JuMP.GenericModel{JNT}
+) where {JNT <: Real}
+    model, con = estim.model, estim.con
+    nx̂, nym, nŷ, nu, He = estim.nx̂, estim.nym, model.ny, model.nu, estim.He
+    nV̂, nX̂, ng, nZ̃ = He*nym, He*nx̂, length(con.i_g), length(estim.Z̃)
+    Nc = nZ̃ + 3
+    Z̃_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nZ̃), Nc)
+    V̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nV̂), Nc)
+    g_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, ng), Nc)
+    X̂_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nX̂), Nc)
+    x̄_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂), Nc)
+    û_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nu), Nc)
+    ŷ_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nŷ), Nc)
+    function Jfunc(Z̃tup::T...)::T where {T <: Real}
+        Z̃1 = Z̃tup[begin]
+        Z̃, g = get_tmp(Z̃_cache, Z̃1), get_tmp(g_cache, Z̃1)
+        for i in eachindex(Z̃tup)
+            Z̃[i] = Z̃tup[i] # Z̃ .= Z̃tup seems to produce a type instability
+        end
+        V̂, X̂ = get_tmp(V̂_cache, Z̃1), get_tmp(X̂_cache, Z̃1)
+        û, ŷ = get_tmp(û_cache, Z̃1), get_tmp(ŷ_cache, Z̃1)
+        V̂, X̂ = predict!(V̂, X̂, û, ŷ, estim, model, Z̃)
+        g = get_tmp(g_cache, Z̃1)
+        g = con_nonlinprog!(g, estim, model, X̂, V̂, Z̃)
+        x̄ = get_tmp(x̄_cache, Z̃1)
+        return obj_nonlinprog!(x̄, estim, model, V̂, Z̃)::T
+    end
+    function gfunc_i(i, Z̃tup::NTuple{N, T})::T where {N, T <:Real}
+        Z̃1 = Z̃tup[begin]
+        Z̃, g = get_tmp(Z̃_cache, Z̃1), get_tmp(g_cache, Z̃1)
+        if any(new !== old for (new, old) in zip(Z̃tup, Z̃)) # new Z̃tup, update predictions:
+            V̂, X̂ = get_tmp(V̂_cache, Z̃1), get_tmp(X̂_cache, Z̃1)
+            û, ŷ = get_tmp(û_cache, Z̃1), get_tmp(ŷ_cache, Z̃1)
+            for i in eachindex(Z̃tup)
+                Z̃[i] = Z̃tup[i] # Z̃ .= Z̃tup seems to produce a type instability
+            end
+            V̂, X̂ = predict!(V̂, X̂, û, ŷ, estim, model, Z̃)
+            g = con_nonlinprog!(g, estim, model, X̂, V̂, Z̃)
+        end
+        return g[i]
+    end
+    gfunc = [(Z̃...) -> gfunc_i(i, Z̃) for i in 1:ng]
+    return Jfunc, gfunc
+end
+
+
 """
     get_optim_functions(estim::MovingHorizonEstimator, ::JuMP.GenericModel) -> Jfunc, gfunc
 

test/test_predictive_control.jl

@@ -507,11 +507,11 @@ end
     u = moveinput!(nmpc4, [0], d, R̂u=fill(12, nmpc4.Hp))
     @test u ≈ [12] atol=5e-2
     nmpc5 = setconstraint!(NonLinMPC(nonlinmodel, Hp=15, Cwt=Inf), ymin=[1])
-    g_Ymin_end = nmpc5.optim[:g_Ymin_15].func
+    g_Y0min_end = nmpc5.optim[:g_Y0min_15].func
     # test gfunc_i(i,::NTuple{N, Float64}):
-    @test g_Ymin_end(20.0, 10.0) ≤ 0.0 
+    @test g_Y0min_end(20.0, 10.0) ≤ 0.0 
     # test gfunc_i(i,::NTuple{N, ForwardDiff.Dual}) : 
-    @test ForwardDiff.gradient(vec->g_Ymin_end(vec...), [20.0, 10.0]) ≈ [-5, -5] atol=1e-3
+    @test ForwardDiff.gradient(vec->g_Y0min_end(vec...), [20.0, 10.0]) ≈ [-5, -5] atol=1e-3
     linmodel3 = LinModel{Float32}(0.5*ones(1,1), ones(1,1), ones(1,1), zeros(1,0), zeros(1,0), 1.0)
     nmpc6  = NonLinMPC(linmodel3, Hp=10)
     @test moveinput!(nmpc6, [0]) ≈ [0.0]