custom nonlinear constraint starting to work yay!

franckgaga · franckgaga · commit 4c645f3605ad · 2024-11-21T16:03:26.000-05:00
diff --git a/src/controller/construct.jl b/src/controller/construct.jl
@@ -148,7 +148,8 @@ function setconstraint!(
     C_Δumin = C_Deltaumin, C_Δumax = C_Deltaumax,
 )
     model, con, optim = mpc.estim.model, mpc.con, mpc.optim
-    nu, ny, nx̂, Hp, Hc, nϵ = model.nu, model.ny, mpc.estim.nx̂, mpc.Hp, mpc.Hc, mpc.nϵ
+    nu, ny, nx̂, Hp, Hc = model.nu, model.ny, mpc.estim.nx̂, mpc.Hp, mpc.Hc
+    nϵ, nc = mpc.nϵ, con.nc
     notSolvedYet = (JuMP.termination_status(optim) == JuMP.OPTIMIZE_NOT_CALLED)
     if isnothing(Umin) && !isnothing(umin)
         size(umin) == (nu,) || throw(ArgumentError("umin size must be $((nu,))"))
@@ -277,7 +278,8 @@ function setconstraint!(
     i_Ymin,  i_Ymax  = .!isinf.(con.Y0min), .!isinf.(con.Y0max)
     i_x̂min,  i_x̂max  = .!isinf.(con.x̂0min), .!isinf.(con.x̂0max)
     if notSolvedYet
-        con.i_b[:], con.i_g[:], con.A[:] = init_matconstraint_mpc(model,
+        con.i_b[:], con.i_g[:], con.A[:] = init_matconstraint_mpc(
+            model, nc,
             i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, 
             i_Ymin, i_Ymax, i_x̂min, i_x̂max,
             con.A_Umin, con.A_Umax, con.A_ΔŨmin, con.A_ΔŨmax, 
@@ -291,7 +293,8 @@ function setconstraint!(
         @constraint(optim, linconstraint, A*ΔŨvar .≤ b)
         setnonlincon!(mpc, model, optim)
     else
-        i_b, i_g = init_matconstraint_mpc(model, 
+        i_b, i_g = init_matconstraint_mpc(
+            model, nc,
             i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, 
             i_Ymin, i_Ymax, i_x̂min, i_x̂max
         )
@@ -304,8 +307,10 @@ end
 
 
 @doc raw"""
-    init_matconstraint_mpc(model::LinModel,
-        i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max, args...
+    init_matconstraint_mpc(
+        model::LinModel, nc::Int,
+        i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max, 
+        args...
     ) -> i_b, i_g, A
 
 Init `i_b`, `i_g` and `A` matrices for the linear and nonlinear inequality constraints.
@@ -317,17 +322,20 @@ The linear and nonlinear inequality constraints are respectively defined as:
     \mathbf{g(ΔŨ)} &≤ \mathbf{0}
 \end{aligned}
 ```
-`i_b` is a `BitVector` including the indices of ``\mathbf{b}`` that are finite numbers. 
-`i_g` is a similar vector but for the indices of ``\mathbf{g}`` (empty if `model` is a 
-[`LinModel`](@ref)). The method also returns the ``\mathbf{A}`` matrix if `args` is
-provided. In such a case, `args`  needs to contain all the inequality constraint matrices: 
+The argument `nc` is the number of custom nonlinear constraints in ``\mathbf{g_c}``. `i_b` 
+is a `BitVector` including the indices of ``\mathbf{b}`` that are finite numbers. `i_g` is a
+similar vector but for the indices of ``\mathbf{g}``. The method also returns the 
+``\mathbf{A}`` matrix if `args` is provided. In such a case, `args`  needs to contain all
+the inequality constraint matrices: 
 `A_Umin, A_Umax, A_ΔŨmin, A_ΔŨmax, A_Ymin, A_Ymax, A_x̂min, A_x̂max`.
 """
-function init_matconstraint_mpc(::LinModel{NT}, 
-    i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max, args...
+function init_matconstraint_mpc(
+    ::LinModel{NT}, nc::Int,
+    i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max, 
+    args...
 ) where {NT<:Real}
     i_b = [i_Umin; i_Umax; i_ΔŨmin; i_ΔŨmax; i_Ymin; i_Ymax; i_x̂min; i_x̂max]
-    i_g = BitVector()
+    i_g = trues(nc)
     if isempty(args)
         A = nothing
     else
@@ -338,11 +346,13 @@ function init_matconstraint_mpc(::LinModel{NT},
 end
 
 "Init `i_b, A` without outputs and terminal constraints if `model` is not a [`LinModel`](@ref)."
-function init_matconstraint_mpc(::SimModel{NT},
-    i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max, args...
+function init_matconstraint_mpc(
+    ::SimModel{NT}, nc::Int,
+    i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max, 
+    args...
 ) where {NT<:Real}
     i_b = [i_Umin; i_Umax; i_ΔŨmin; i_ΔŨmax]
-    i_g = [i_Ymin; i_Ymax; i_x̂min; i_x̂max]
+    i_g = [i_Ymin; i_Ymax; i_x̂min;  i_x̂max; trues(nc)]
     if isempty(args)
         A = nothing
     else
@@ -667,7 +677,7 @@ function init_defaultcon_mpc(
     i_Ymin,  i_Ymax  = .!isinf.(Y0min),  .!isinf.(Y0max)
     i_x̂min,  i_x̂max  = .!isinf.(x̂0min),  .!isinf.(x̂0max)
     i_b, i_g, A = init_matconstraint_mpc(
-        model, 
+        model, nc,
         i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, i_Ymin, i_Ymax, i_x̂min, i_x̂max,
         A_Umin, A_Umax, A_ΔŨmin, A_ΔŨmax, A_Ymin, A_Ymax, A_x̂max, A_x̂min
     )
diff --git a/src/controller/nonlinmpc.jl b/src/controller/nonlinmpc.jl
@@ -384,7 +384,7 @@ function get_mutating_gc(NT, gc)
 end
 
 function test_custom_functions(JE, gc!, uop; Uop, dop, Dop, ΔŨ, p)
-    # TODO: contunue here (important to guide the users, sim! can be used on NonLinModel
+    # TODO: contunue here (important to guide the user, sim! can be used on NonLinModel,
     # but there is no similar function for the custom functions of NonLinMPC)
     Ue = [Uop; uop]
     D̂e = [dop; Dop]
@@ -458,6 +458,11 @@ function init_optimization!(mpc::NonLinMPC, model::SimModel, optim)
             name = Symbol("g_x̂0max_$i")
             optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_x̂min+i]; name)
         end
+        i_end_x̂max = 2Hp*ny + 2nx̂
+        for i in 1:con.nc
+            name = Symbol("g_c_$i")
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_x̂max+i]; name)
+        end
     end
     return nothing
 end
@@ -471,20 +476,21 @@ Inspired from: [User-defined operators with vector outputs](https://jump.dev/JuM
 """
 function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT<:Real
     model = mpc.estim.model
-    nu, ny, nx̂, Hp = model.nu, model.ny, mpc.estim.nx̂, mpc.Hp
-    ng, nΔŨ, nU, nŶ = length(mpc.con.i_g), length(mpc.ΔŨ), Hp*nu, Hp*ny
+    nu, ny, nx̂, nϵ, Hp = model.nu, model.ny, mpc.estim.nx̂, mpc.nϵ, mpc.Hp
+    ng, nc, nΔŨ, nU, nŶ = length(mpc.con.i_g), mpc.con.nc, length(mpc.ΔŨ), Hp*nu, Hp*ny
     nUe, nŶe = nU + nu, nŶ + ny
-    Nc = nΔŨ + 3
-    ΔŨ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nΔŨ), Nc)
-    Ŷe_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶe), Nc)
-    Ue_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nUe), Nc)
-    Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, nŶ),  Nc)
-    Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, nU),  Nc)
-    x̂0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nx̂),  Nc)
-    x̂0next_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂),  Nc)
-    u0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu),  Nc)
-    û0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu),  Nc)
-    g_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, ng),  Nc)
+    Ncache = nΔŨ + 3
+    ΔŨ_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nΔŨ), Ncache)
+    Ŷe_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nŶe), Ncache)
+    Ue_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nUe), Ncache)
+    Ȳ_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, nŶ),  Ncache)
+    Ū_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, nU),  Ncache)
+    x̂0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nx̂),  Ncache)
+    x̂0next_cache::DiffCache{Vector{JNT}, Vector{JNT}} = DiffCache(zeros(JNT, nx̂),  Ncache)
+    u0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu),  Ncache)
+    û0_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nu),  Ncache)
+    g_cache::DiffCache{Vector{JNT}, Vector{JNT}}      = DiffCache(zeros(JNT, ng),  Ncache)
+    gc_cache::DiffCache{Vector{JNT}, Vector{JNT}}     = DiffCache(zeros(JNT, nc),  Ncache)
     function Jfunc(ΔŨtup::T...) where T<:Real
         ΔŨ1 = ΔŨtup[begin]
         ΔŨ, g = get_tmp(ΔŨ_cache, ΔŨ1), get_tmp(g_cache, ΔŨ1) 
@@ -495,10 +501,12 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
         Ȳ,  Ū      = get_tmp(Ȳ_cache, ΔŨ1),  get_tmp(Ū_cache, ΔŨ1)
         x̂0, x̂0next = get_tmp(x̂0_cache, ΔŨ1), get_tmp(x̂0next_cache, ΔŨ1)
         u0, û0     = get_tmp(u0_cache, ΔŨ1), get_tmp(û0_cache, ΔŨ1)
-        g          = get_tmp(g_cache, ΔŨ1)
+        g,  gc     = get_tmp(g_cache, ΔŨ1),  get_tmp(gc_cache, ΔŨ1)
         Ŷ0, x̂0end  = predict!(Ȳ, x̂0, x̂0next, u0, û0, mpc, model, ΔŨ)
         Ŷe, Ue     = extended_predictions!(Ŷe, Ue, Ū, mpc, model, Ŷ0, ΔŨ)
-        g          = con_nonlinprog!(g, mpc, model, x̂0end, Ŷ0, ΔŨ)
+        ϵ = (nϵ == 1) ? ΔŨ[end] : zero(JNT) # ϵ = 0 if nϵ == 0 (meaning no relaxation)
+        mpc.con.gc!(gc, Ue, Ŷe, mpc.D̂e, mpc.p, ϵ)
+        g = con_nonlinprog!(g, mpc, model, x̂0end, gc, Ŷ0, ΔŨ, ϵ)
         return obj_nonlinprog!(Ȳ, Ū, mpc, model, Ŷe, Ue, ΔŨ)::T
     end
     function gfunc_i(i, ΔŨtup::NTuple{N, T}) where {N, T<:Real}
@@ -512,10 +520,12 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
             Ȳ,  Ū      = get_tmp(Ȳ_cache, ΔŨ1),  get_tmp(Ū_cache, ΔŨ1)
             x̂0, x̂0next = get_tmp(x̂0_cache, ΔŨ1), get_tmp(x̂0next_cache, ΔŨ1)
             u0, û0     = get_tmp(u0_cache, ΔŨ1), get_tmp(û0_cache, ΔŨ1)
-            g          = get_tmp(g_cache, ΔŨ1)
+            g,  gc     = get_tmp(g_cache, ΔŨ1),  get_tmp(gc_cache, ΔŨ1)
             Ŷ0, x̂0end  = predict!(Ȳ, x̂0, x̂0next, u0, û0, mpc, model, ΔŨ)
             Ŷe, Ue     = extended_predictions!(Ŷe, Ue, Ū, mpc, model, Ŷ0, ΔŨ)
-            g          = con_nonlinprog!(g, mpc, model, x̂0end, Ŷ0, ΔŨ)
+            ϵ = (nϵ == 1) ? ΔŨ[end] : zero(JNT) # ϵ = 0 if nϵ == 0 (meaning no relaxation)
+            mpc.con.gc!(gc, Ue, Ŷe, mpc.D̂e, mpc.p, ϵ)
+            g = con_nonlinprog!(g, mpc, model, x̂0end, gc, Ŷ0, ΔŨ, ϵ)
         end
         return g[i]::T
     end
@@ -568,19 +578,29 @@ function setnonlincon!(
         gfunc_i = optim[Symbol("g_x̂0max_$(i)")]
         @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
+    for i in 1:con.nc
+        gfunc_i = optim[Symbol("g_c_$i")]
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
+    end
+    return nothing
+end
+
+# TODO: MODIF THE FOLLOWING METHOD!
+function setnonlincon!(
+    mpc::NonLinMPC, ::LinModel, optim::JuMP.GenericModel{JNT}
+) where JNT<:Real
     return nothing
 end
 
 """
-    con_nonlinprog!(g, mpc::NonLinMPC, model::SimModel, x̂end, Ŷ, ΔŨ) -> g
+    con_nonlinprog!(g, mpc::NonLinMPC, model::SimModel, x̂end, gc, Ŷ0, ΔŨ, ϵ) -> g
 
 Nonlinear constrains for [`NonLinMPC`](@ref) when `model` is not a [`LinModel`](@ref).
 
-The method mutates the `g` vector in argument and returns it.
+The method mutates the `g` and `gc` vectors in argument.
 """
-function con_nonlinprog!(g, mpc::NonLinMPC, ::SimModel, x̂0end, Ŷ0, ΔŨ)
-    nx̂, nŶ = mpc.estim.nx̂, length(Ŷ0)
-    ϵ = mpc.nϵ == 1 ? ΔŨ[end] : 0 # ϵ = 0 if Cwt=Inf (meaning: no relaxation)
+function con_nonlinprog!(g, mpc::NonLinMPC, ::SimModel, x̂0end, gc, Ŷ0, ΔŨ, ϵ)
+    nx̂, nŶ = length(x̂0end), length(Ŷ0)
     for i in eachindex(g)
         mpc.con.i_g[i] || continue
         if i ≤ nŶ
@@ -592,16 +612,20 @@ function con_nonlinprog!(g, mpc::NonLinMPC, ::SimModel, x̂0end, Ŷ0, ΔŨ)
         elseif i ≤ 2nŶ + nx̂
             j = i - 2nŶ
             g[i] = (mpc.con.x̂0min[j] - x̂0end[j])  - ϵ*mpc.con.c_x̂min[j]
-        else
+        elseif i ≤ 2nŶ + 2nx̂
             j = i - 2nŶ - nx̂
             g[i] = (x̂0end[j] - mpc.con.x̂0max[j])  - ϵ*mpc.con.c_x̂max[j]
+        else
+            j = i - 2nŶ - 2nx̂
+            g[i] = gc[j]
         end
     end
     return g
 end
 
+#TODO: MODIF THE FOLLOWING METHOD!
 "No nonlinear constraints if `model` is a [`LinModel`](@ref), return `g` unchanged."
-con_nonlinprog!(g, ::NonLinMPC, ::LinModel, _ , _ , _ ) = g
+con_nonlinprog!(g, ::NonLinMPC, ::LinModel, _ , _ , _ , _ , _ ) = g
 
 "Evaluate the economic term `E*JE` of the objective function for [`NonLinMPC`](@ref)."
 function obj_econ!(Ue, Ŷe, mpc::NonLinMPC, model::SimModel)
