New jump nlp syntax, merge into main?

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelPredictiveControl"
 uuid = "61f9bdb8-6ae4-484a-811f-bbf86720c31c"
 authors = ["Francis Gagnon"]
-version = "0.21.3"
+version = "0.22.0"
 
 [deps]
 ControlSystemsBase = "aaaaaaaa-a6ca-5380-bf3e-84a91bcd477e"

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/construct.jl

@@ -256,7 +256,7 @@ function setconstraint!(
         JuMP.delete(optim, optim[:linconstraint])
         JuMP.unregister(optim, :linconstraint)
         @constraint(optim, linconstraint, A*ΔŨvar .≤ b)
-        setnonlincon!(mpc, model)
+        setnonlincon!(mpc, model, optim)
     else
         i_b, i_g = init_matconstraint_mpc(model, 
             i_Umin, i_Umax, i_ΔŨmin, i_ΔŨmax, 
@@ -320,7 +320,7 @@ function init_matconstraint_mpc(::SimModel{NT},
 end
 
 "By default, there is no nonlinear constraint, thus do nothing."
-setnonlincon!(::PredictiveController, ::SimModel) = nothing
+setnonlincon!(::PredictiveController, ::SimModel, ::JuMP.GenericModel) = nothing
 
 """
     default_Hp(model::LinModel)

src/controller/nonlinmpc.jl

@@ -313,28 +313,28 @@ function init_optimization!(mpc::NonLinMPC, model::SimModel, optim)
         end
     end
     Jfunc, gfunc = get_optim_functions(mpc, mpc.optim)
-    register(optim, :Jfunc, nΔŨ, Jfunc, autodiff=true)
-    @NLobjective(optim, Min, Jfunc(ΔŨvar...))
+    @operator(optim, J, nΔŨ, Jfunc)
+    @objective(optim, Min, J(ΔŨvar...))
     ny, nx̂, Hp = model.ny, mpc.estim.nx̂, mpc.Hp
     if length(con.i_g) ≠ 0
         for i in eachindex(con.Y0min)
-            sym = Symbol("g_Y0min_$i")
-            register(optim, sym, nΔŨ, gfunc[i], autodiff=true)
+            name = Symbol("g_Y0min_$i")
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i]; name)
         end
         i_end_Ymin = 1Hp*ny
         for i in eachindex(con.Y0max)
-            sym = Symbol("g_Y0max_$i")
-            register(optim, sym, nΔŨ, gfunc[i_end_Ymin+i], autodiff=true)
+            name = Symbol("g_Y0max_$i")
+            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)
-            sym = Symbol("g_x̂0min_$i")
-            register(optim, sym, nΔŨ, gfunc[i_end_Ymax+i], autodiff=true)
+            name = Symbol("g_x̂0min_$i")
+            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)
-            sym = Symbol("g_x̂0max_$i")
-            register(optim, sym, nΔŨ, gfunc[i_end_x̂min+i], autodiff=true)
+            name = Symbol("g_x̂0max_$i")
+            optim[name] = JuMP.add_nonlinear_operator(optim, nΔŨ, gfunc[i_end_x̂min+i]; name)
         end
     end
     return nothing
@@ -397,26 +397,28 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
 end
 
 "Set the nonlinear constraints on the output predictions `Ŷ` and terminal states `x̂end`."
-function setnonlincon!(mpc::NonLinMPC, ::NonLinModel)
-    optim = mpc.optim
+function setnonlincon!(
+    mpc::NonLinMPC, ::NonLinModel, optim::JuMP.GenericModel{JNT}
+) where JNT<:Real
     ΔŨvar = optim[:ΔŨvar]
     con = mpc.con
-    map(con -> JuMP.delete(optim, con), JuMP.all_nonlinear_constraints(optim))
+    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))
-        f_sym = Symbol("g_Y0min_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(ΔŨvar...)) <= 0))
+        gfunc_i = optim[Symbol("g_Y0min_$(i)")]
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     for i in findall(.!isinf.(con.Y0max))
-        f_sym = Symbol("g_Y0max_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(ΔŨvar...)) <= 0))
+        gfunc_i = optim[Symbol("g_Y0max_$(i)")]
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     for i in findall(.!isinf.(con.x̂0min))
-        f_sym = Symbol("g_x̂0min_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(ΔŨvar...)) <= 0))
+        gfunc_i = optim[Symbol("g_x̂0min_$(i)")]
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     for i in findall(.!isinf.(con.x̂0max))
-        f_sym = Symbol("g_x̂0max_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(ΔŨvar...)) <= 0))
+        gfunc_i = optim[Symbol("g_x̂0max_$(i)")]
+        @constraint(optim, gfunc_i(ΔŨvar...) <= 0)
     end
     return nothing
 end

src/estimator/mhe/construct.jl

@@ -536,7 +536,7 @@ function setconstraint!(
         JuMP.delete(optim, optim[:linconstraint])
         JuMP.unregister(optim, :linconstraint)
         @constraint(optim, linconstraint, A*Z̃var .≤ b)
-        setnonlincon!(estim, model)
+        setnonlincon!(estim, model, optim)
     else
         i_b, i_g = init_matconstraint_mhe(model, 
             i_x̃min, i_x̃max, i_X̂min, i_X̂max, i_Ŵmin, i_Ŵmax, i_V̂min, i_V̂max
@@ -598,28 +598,31 @@ function init_matconstraint_mhe(::SimModel{NT},
 end
 
 "By default, no nonlinear constraints in the MHE, thus return nothing."
-setnonlincon!(::MovingHorizonEstimator, ::SimModel) = nothing
+setnonlincon!(::MovingHorizonEstimator, ::SimModel, ::JuMP.GenericModel) = nothing
 
 "Set the nonlinear constraints on the output predictions `Ŷ` and terminal states `x̂end`."
-function setnonlincon!(estim::MovingHorizonEstimator, ::NonLinModel)
+function setnonlincon!(
+    estim::MovingHorizonEstimator, ::NonLinModel, optim::JuMP.GenericModel{JNT}
+) where JNT<:Real
     optim, con = estim.optim, estim.con
     Z̃var = optim[:Z̃var]
-    map(con -> JuMP.delete(optim, con), JuMP.all_nonlinear_constraints(optim))
+    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))
-        f_sym = Symbol("g_X̂0min_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(Z̃var...)) <= 0))
+        gfunc_i = optim[Symbol("g_X̂0min_$(i)")]
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.X̂0max))
-        f_sym = Symbol("g_X̂0max_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(Z̃var...)) <= 0))
+        gfunc_i = optim[Symbol("g_X̂0max_$(i)")]
+        @constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.V̂min))
-        f_sym = Symbol("g_V̂min_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(Z̃var...)) <= 0))
+        gfunc_i = optim[Symbol("g_V̂min_$(i)")]
+        JuMP.@constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     for i in findall(.!isinf.(con.V̂max))
-        f_sym = Symbol("g_V̂max_$(i)")
-        JuMP.add_nonlinear_constraint(optim, :($(f_sym)($(Z̃var...)) <= 0))
+        gfunc_i = optim[Symbol("g_V̂max_$(i)")]
+        JuMP.@constraint(optim, gfunc_i(Z̃var...) <= 0)
     end
     return nothing
 end
@@ -1073,28 +1076,36 @@ function init_optimization!(
         end
     end
     Jfunc, gfunc = get_optim_functions(estim, optim)
-    register(optim, :Jfunc, nZ̃, Jfunc, autodiff=true)
-    @NLobjective(optim, Min, Jfunc(Z̃var...))
+    @operator(optim, J, nZ̃, Jfunc)
+    @objective(optim, Min, J(Z̃var...))
     nV̂, nX̂ = estim.He*estim.nym, estim.He*estim.nx̂
     if length(con.i_g) ≠ 0
         for i in eachindex(con.X̂0min)
-            sym = Symbol("g_X̂0min_$i")
-            register(optim, sym, nZ̃, gfunc[i], autodiff=true)
+            name = Symbol("g_X̂0min_$i")
+            optim[name] = JuMP.add_nonlinear_operator(
+                optim, nZ̃, gfunc[i]; name
+            )
         end
         i_end_X̂min = nX̂
         for i in eachindex(con.X̂0max)
-            sym = Symbol("g_X̂0max_$i")
-            register(optim, sym, nZ̃, gfunc[i_end_X̂min+i], autodiff=true)
+            name = Symbol("g_X̂0max_$i")
+            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)
-            sym = Symbol("g_V̂min_$i")
-            register(optim, sym, nZ̃, gfunc[i_end_X̂max+i], autodiff=true)
+            name = Symbol("g_V̂min_$i")
+            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)
-            sym = Symbol("g_V̂max_$i")
-            register(optim, sym, nZ̃, gfunc[i_end_V̂min+i], autodiff=true)
+            name = Symbol("g_V̂max_$i")
+            optim[name] = JuMP.add_nonlinear_operator(
+                optim, nZ̃, gfunc[i_end_V̂min + i]; name
+            )
         end
     end
     return nothing

src/estimator/mhe/execute.jl

@@ -564,4 +564,7 @@ function setmodel_estimator!(
 end
 
 "Called by plots recipes for the estimated states constraints."
-getX̂con(estim::MovingHorizonEstimator, _ ) = estim.con.X̂0min+estim.X̂op, estim.con.X̂0max+estim.X̂op
+getX̂con(estim::MovingHorizonEstimator, _ ) = estim.con.X̂0min+estim.X̂op, estim.con.X̂0max+estim.X̂op
+
+"No nonlinear constraints if `model` is a [`LinModel`](@ref), return `g` unchanged."
+con_nonlinprog!(g, ::MovingHorizonEstimator, ::LinModel, _ , _ , _ ) = g

test/test_predictive_control.jl

@@ -439,8 +439,10 @@ end
     nmpc7 = NonLinMPC(nonlinmodel, Hp=15, Ewt=1e-3, JE=(UE,ŶE,D̂E) -> UE.*ŶE.*D̂E)
     @test nmpc7.E == 1e-3
     @test nmpc7.JE([1,2],[3,4],[4,6]) == [12, 48]
-    nmpc8 = NonLinMPC(nonlinmodel, Hp=15, optim=JuMP.Model(OSQP.MathOptInterfaceOSQP.Optimizer))
-    @test solver_name(nmpc8.optim) == "OSQP"
+    optim = JuMP.Model(optimizer_with_attributes(Ipopt.Optimizer, "nlp_scaling_max_gradient"=>1.0))
+    nmpc8 = NonLinMPC(nonlinmodel, Hp=15, optim=optim)
+    @test solver_name(nmpc8.optim) == "Ipopt"
+    @test get_attribute(nmpc8.optim, "nlp_scaling_max_gradient") == 1.0
     im = InternalModel(nonlinmodel)
     nmpc9 = NonLinMPC(im, Hp=15)
     @test isa(nmpc9.estim, InternalModel)
@@ -504,11 +506,12 @@ end
     nmpc4 = NonLinMPC(nonlinmodel, Hp=15, Mwt=[0], Nwt=[0], Lwt=[1])
     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), ymax=[1])
-    g_Ymax_end = nmpc5.optim.nlp_model.operators.registered_multivariate_operators[end].f
-    @test g_Ymax_end((1.0, 1.0)) ≤ 0.0 # test gfunc_i(i,::NTuple{N, Float64})
+    nmpc5 = setconstraint!(NonLinMPC(nonlinmodel, Hp=15, Cwt=Inf), ymin=[1])
+    g_Y0min_end = nmpc5.optim[:g_Y0min_15].func
+    # test gfunc_i(i,::NTuple{N, Float64}):
+    @test g_Y0min_end(20.0, 10.0) ≤ 0.0 
     # test gfunc_i(i,::NTuple{N, ForwardDiff.Dual}) : 
-    @test ForwardDiff.gradient(g_Ymax_end, [1.0, 1.0]) ≈ [0.0, 0.0]
+    @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]

test/test_state_estim.jl

@@ -666,14 +666,12 @@ end
     @test mhe9.Q̂ ≈ I(6)
     @test mhe9.R̂ ≈ I(2)
 
-    optim = Model(Ipopt.Optimizer)
+    optim = JuMP.Model(optimizer_with_attributes(Ipopt.Optimizer, "nlp_scaling_max_gradient"=>1.0))
     covestim = ExtendedKalmanFilter(nonlinmodel, 1:2, 0, [1, 1], I_6, I_6, I_2)
     mhe10 = MovingHorizonEstimator(
         nonlinmodel, 5, 1:2, 0, [1, 1], I_6, I_6, I_2, Inf, optim, covestim
     )
-
-    mhe11 = MovingHorizonEstimator(nonlinmodel, He=5, optim=Model(OSQP.Optimizer))
-    @test solver_name(mhe11.optim) == "OSQP"
+    @test solver_name(mhe10.optim) == "Ipopt"
 
     mhe12 = MovingHorizonEstimator(nonlinmodel, He=5, Cwt=1e3)
     @test size(mhe12.Ẽ, 2) == 6*mhe12.nx̂ + 1
@@ -735,15 +733,14 @@ end
     x̂ = updatestate!(mhe3, [0], [0])
     @test x̂ ≈ [0, 0] atol=1e-3
     @test isa(x̂, Vector{Float32})
+
+    mhe4 = setconstraint!(MovingHorizonEstimator(nonlinmodel, He=1, nint_ym=0), v̂max=[50,50])
+    g_V̂max_end = mhe4.optim[:g_V̂max_2].func
+     # test gfunc_i(i,::NTuple{N, Float64})
+    @test g_V̂max_end(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) ≤ 0.0 
+    # test gfunc_i(i,::NTuple{N, ForwardDiff.Dual}) : 
+    @test ForwardDiff.gradient(vec->g_V̂max_end(vec...), zeros(8)) ≈ zeros(8)
 
-    mhe4 = setconstraint!(MovingHorizonEstimator(nonlinmodel, He=1, nint_ym=0), x̂max=[50,50,50,50])
-    g_X̂max_end = mhe4.optim.nlp_model.operators.registered_multivariate_operators[end].f
-    # test gfunc_i(i,::NTuple{N, Float64}):
-    @test g_X̂max_end(
-        (1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0)) ≤ 0.0 
-    # test gfunc_i(i,::NTuple{N, ForwardDiff.Dual}): 
-    @test ForwardDiff.gradient(
-        g_X̂max_end, [1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0]) ≈ [0, 0, 0, 0, 0, 0, 0, 0]
     Q̂ = diagm([1/4, 1/4, 1/4, 1/4].^2) 
     R̂ = diagm([1, 1].^2)
     optim = Model(Ipopt.Optimizer)