Skip to content

Commit 3d4f917

Browse files
franckgagafranckgaga
committed
added: similar refactor for nonlinear inequality constraints
1 parent 38531a8 commit 3d4f917

File tree

3 files changed

+186
-149
lines changed

3 files changed

+186
-149
lines changed

‎src/controller/nonlinmpc.jl‎

Lines changed: 39 additions & 142 deletions
Original file line numberDiff line numberDiff line change
@@ -505,10 +505,10 @@ function init_optimization!(mpc::NonLinMPC, model::SimModel, optim)
505505
JuMP.set_attribute(optim, "nlp_scaling_max_gradient", 10.0/C)
506506
end
507507
end
508-
Jfunc, ∇Jfunc!, gfuncs, geqfuncs, ∇geqfuncs! = get_optim_functions(mpc, optim)
508+
Jfunc, ∇Jfunc!, gfuncs, ∇gfuncs!, geqfuncs, ∇geqfuncs! = get_optim_functions(mpc, optim)
509509
@operator(optim, J, nZ̃, Jfunc, ∇Jfunc!)
510510
@objective(optim, Min, J(Z̃var...))
511-
init_nonlincon!(mpc, model, transcription, gfuncs, geqfuncs, ∇geqfuncs!)
511+
init_nonlincon!(mpc, model, transcription, gfuncs, ∇gfuncs!, geqfuncs, ∇geqfuncs!)
512512
set_nonlincon!(mpc, model, optim)
513513
return nothing
514514
end
@@ -569,8 +569,15 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
569569
end
570570
return nothing
571571
end
572+
# --------------------- cache for the AD functions -----------------------------------
573+
Z̃arg_vec = Vector{JNT}(undef, nZ̃)
574+
g_vec = Vector{JNT}(undef, ng)
575+
∇g = Matrix{JNT}(undef, ng, nZ̃) # Jacobian of g
576+
geq_vec = Vector{JNT}(undef, neq)
577+
∇geq = Matrix{JNT}(undef, neq, nZ̃) # Jacobian of geq
572578
# --------------------- objective function --------------------------------------------
573-
function Jfunc(Z̃arg::Vararg{T, N}) where {N, T<:Real} # force specialization with Vararg, see https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing
579+
# force specialization with Vararg, see https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing
580+
function Jfunc(Z̃arg::Vararg{T, N}) where {N, T<:Real}
574581
update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
575582
ΔŨ = get_tmp(ΔŨ_cache, T)
576583
Ue, Ŷe = get_tmp(Ue_cache, T), get_tmp(Ŷe_cache, T)
@@ -582,27 +589,36 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
582589
ΔŨ = get_tmp(ΔŨ_cache, T)
583590
Ue, Ŷe = get_tmp(Ue_cache, T), get_tmp(Ŷe_cache, T)
584591
U0, Ŷ0 = get_tmp(U0_cache, T), get_tmp(Ŷ0_cache, T)
585-
return obj_nonlinprog!(Ŷ0, U0, mpc, model, Ue, Ŷe, ΔŨ)::T
592+
return obj_nonlinprog!(Ŷ0, U0, mpc, model, Ue, Ŷe, ΔŨ)
586593
end
587-
Z̃vec = Vector{JNT}(undef, nZ̃)
588-
∇Jbuffer = GradientBuffer(Jfunc_vec, Z̃vec)
594+
∇J_buffer = GradientBuffer(Jfunc_vec, Z̃arg_vec)
589595
function ∇Jfunc!(∇J, Z̃arg::Vararg{T, N}) where {N, T<:Real}
590-
Z̃vec .= Z̃arg
591-
return gradient!(∇J, ∇Jbuffer, Z̃vec)
596+
Z̃arg_vec .= Z̃arg
597+
return gradient!(∇J, ∇J_buffer, Z̃arg_vec)
592598
end
593599
# --------------------- inequality constraint functions -------------------------------
594-
# TODO:re-déplacer en haut:
595-
# TODO: modifier pour combiner en 1 seule fonction comme pour geqfuncs
596-
function gfunc_i(i, Z̃arg::NTuple{N, T}) where {N, T<:Real}
597-
update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
598-
g = get_tmp(g_cache, T)
599-
return g[i]::T
600-
end
601600
gfuncs = Vector{Function}(undef, ng)
602-
for i in 1:ng
603-
# this is another syntax for anonymous function, allowing parameters T and N:
604-
gfuncs[i] = function (Z̃arg::Vararg{T, N}) where {N, T<:Real}
605-
return gfunc_i(i, Z̃arg)
601+
for i in eachindex(gfuncs)
602+
func_i = function (Z̃arg::Vararg{T, N}) where {N, T<:Real} # see comment above
603+
update_simulations!(Z̃arg, get_tmp(Z̃_cache, T))
604+
g = get_tmp(g_cache, T)
605+
return g[i]::T
606+
end
607+
gfuncs[i] = func_i
608+
end
609+
function gfunc_vec!(g, Z̃vec::AbstractVector{T}) where T<:Real
610+
update_simulations!(Z̃vec, get_tmp(Z̃_cache, T))
611+
g .= get_tmp(g_cache, T)
612+
return g
613+
end
614+
∇g_buffer = JacobianBuffer(gfunc_vec!, g_vec, Z̃arg_vec)
615+
∇gfuncs! = Vector{Function}(undef, ng)
616+
for i in eachindex(∇gfuncs!)
617+
∇gfuncs![i] = function (∇g_i, Z̃arg::Vararg{T, N}) where {N, T<:Real}
618+
Z̃arg_vec .= Z̃arg
619+
jacobian!(∇g, ∇g_buffer, g_vec, Z̃arg_vec)
620+
∇g_i .= @views ∇g[i, :]
621+
return ∇g_i
606622
end
607623
end
608624
# --------------------- equality constraint functions ---------------------------------
@@ -620,136 +636,17 @@ function get_optim_functions(mpc::NonLinMPC, ::JuMP.GenericModel{JNT}) where JNT
620636
geq .= get_tmp(geq_cache, T)
621637
return geq
622638
end
623-
∇geq = Matrix{JNT}(undef, neq, nZ̃) # Jacobian of geq
624-
geqvec = Vector{JNT}(undef, neq)
639+
∇geq_buffer = JacobianBuffer(geqfunc_vec!, geq_vec, Z̃arg_vec)
625640
∇geqfuncs! = Vector{Function}(undef, neq)
626641
for i in eachindex(∇geqfuncs!)
627642
∇geqfuncs![i] = function (∇geq_i, Z̃arg::Vararg{T, N}) where {N, T<:Real}
628-
Z̃vec .= Z̃arg
629-
ForwardDiff.jacobian!(∇geq, geqfunc_vec!, geqvec, Z̃vec)
643+
Z̃arg_vec .= Z̃arg
644+
jacobian!(∇geq, ∇geq_buffer, geq_vec, Z̃arg_vec)
630645
∇geq_i .= @views ∇geq[i, :]
631646
return ∇geq_i
632647
end
633648
end
634-
return Jfunc, ∇Jfunc!, gfuncs, geqfuncs, ∇geqfuncs!
635-
end
636-
637-
"""
638-
init_nonlincon!(
639-
mpc::NonLinMPC, model::LinModel, ::TranscriptionMethod, gfuncs, geqfuncs, ∇geqfuncs!
640-
)
641-
642-
Init nonlinear constraints for [`LinModel`](@ref).
643-
644-
The only nonlinear constraints are the custom inequality constraints `gc`.
645-
"""
646-
function init_nonlincon!(
647-
mpc::NonLinMPC, ::LinModel, ::TranscriptionMethod,
648-
gfuncs::Vector{<:Function}, _ , _
649-
)
650-
optim, con = mpc.optim, mpc.con
651-
nZ̃ = length(mpc.Z̃)
652-
if length(con.i_g) ≠ 0
653-
i_base = 0
654-
for i in 1:con.nc
655-
name = Symbol("g_c_$i")
656-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
657-
end
658-
end
659-
return nothing
660-
end
661-
662-
"""
663-
init_nonlincon!(
664-
mpc::NonLinMPC, model::NonLinModel, ::MultipleShooting, gfuncs, geqfuncs, ∇geqfuncs!
665-
)
666-
667-
Init nonlinear constraints for [`NonLinModel`](@ref) and [`MultipleShooting`](@ref).
668-
669-
The nonlinear constraints are the output prediction `Ŷ` bounds, the custom inequality
670-
constraints `gc` and all the nonlinear equality constraints `geq`.
671-
"""
672-
function init_nonlincon!(
673-
mpc::NonLinMPC, ::NonLinModel, ::MultipleShooting,
674-
gfuncs::Vector{<:Function}, geqfuncs::Vector{<:Function}, ∇geqfuncs!::Vector{<:Function}
675-
)
676-
optim, con = mpc.optim, mpc.con
677-
ny, nx̂, Hp, nZ̃ = mpc.estim.model.ny, mpc.estim.nx̂, mpc.Hp, length(mpc.Z̃)
678-
# --- nonlinear inequality constraints ---
679-
if length(con.i_g) ≠ 0
680-
i_base = 0
681-
for i in eachindex(con.Y0min)
682-
name = Symbol("g_Y0min_$i")
683-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
684-
end
685-
i_base = 1Hp*ny
686-
for i in eachindex(con.Y0max)
687-
name = Symbol("g_Y0max_$i")
688-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
689-
end
690-
i_base = 2Hp*ny
691-
for i in 1:con.nc
692-
name = Symbol("g_c_$i")
693-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
694-
end
695-
end
696-
# --- nonlinear equality constraints ---
697-
Z̃var = optim[:Z̃var]
698-
for i in eachindex(geqfuncs)
699-
name = Symbol("geq_$i")
700-
geqfunc_i = JuMP.add_nonlinear_operator(optim, nZ̃, geqfuncs[i], ∇geqfuncs![i]; name)
701-
optim[name] = geqfunc_i
702-
# set with @constrains here instead of set_nonlincon!, since the number of nonlinear
703-
# equality constraints is known and constant (±Inf are impossible):
704-
@constraint(optim, geqfunc_i(Z̃var...) == 0)
705-
end
706-
return nothing
707-
end
708-
709-
"""
710-
init_nonlincon!(
711-
mpc::NonLinMPC, model::NonLinModel, ::SingleShooting, gfuncs, geqfuncs, ∇geqfuncs!
712-
)
713-
714-
Init nonlinear constraints for [`NonLinModel`](@ref) and [`SingleShooting`](@ref).
715-
716-
The nonlinear constraints are the custom inequality constraints `gc`, the output
717-
prediction `Ŷ` bounds and the terminal state `x̂end` bounds.
718-
"""
719-
function init_nonlincon!(
720-
mpc::NonLinMPC, ::NonLinModel, ::SingleShooting,
721-
gfuncs::Vector{<:Function}, _ , _
722-
)
723-
optim, con = mpc.optim, mpc.con
724-
ny, nx̂, Hp, nZ̃ = mpc.estim.model.ny, mpc.estim.nx̂, mpc.Hp, length(mpc.Z̃)
725-
if length(con.i_g) ≠ 0
726-
i_base = 0
727-
for i in eachindex(con.Y0min)
728-
name = Symbol("g_Y0min_$i")
729-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
730-
end
731-
i_base = 1Hp*ny
732-
for i in eachindex(con.Y0max)
733-
name = Symbol("g_Y0max_$i")
734-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
735-
end
736-
i_base = 2Hp*ny
737-
for i in eachindex(con.x̂0min)
738-
name = Symbol("g_x̂0min_$i")
739-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
740-
end
741-
i_base = 2Hp*ny + nx̂
742-
for i in eachindex(con.x̂0max)
743-
name = Symbol("g_x̂0max_$i")
744-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
745-
end
746-
i_base = 2Hp*ny + 2nx̂
747-
for i in 1:con.nc
748-
name = Symbol("g_c_$i")
749-
optim[name] = JuMP.add_nonlinear_operator(optim, nZ̃, gfuncs[i_base+i]; name)
750-
end
751-
end
752-
return nothing
649+
return Jfunc, ∇Jfunc!, gfuncs, ∇gfuncs!, geqfuncs, ∇geqfuncs!
753650
end
754651

755652
"""

‎src/controller/transcription.jl‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,148 @@ function init_defectmat(
508508
return Eŝ, Gŝ, Jŝ, Kŝ, Vŝ, Bŝ
509509
end
510510

511+
"""
512+
init_nonlincon!(
513+
mpc::PredictiveController, model::LinModel, transcription::TranscriptionMethod,
514+
gfuncs , ∇gfuncs!,
515+
geqfuncs, ∇geqfuncs!
516+
)
517+
518+
Init nonlinear constraints for [`LinModel`](@ref) for all [`TranscriptionMethod`](@ref).
519+
520+
The only nonlinear constraints are the custom inequality constraints `gc`.
521+
"""
522+
function init_nonlincon!(
523+
mpc::PredictiveController, ::LinModel, ::TranscriptionMethod, gfuncs, ∇gfuncs!, _ , _
524+
)
525+
optim, con = mpc.optim, mpc.con
526+
nZ̃ = length(mpc.Z̃)
527+
if length(con.i_g) ≠ 0
528+
i_base = 0
529+
for i in 1:con.nc
530+
name = Symbol("g_c_$i")
531+
optim[name] = JuMP.add_nonlinear_operator(
532+
optim, nZ̃, gfuncs[i_base+i]; name
533+
)
534+
end
535+
end
536+
return nothing
537+
end
538+
539+
"""
540+
init_nonlincon!(
541+
mpc::PredictiveController, model::NonLinModel, transcription::MultipleShooting,
542+
gfuncs, ∇gfuncs!,
543+
geqfuncs, ∇geqfuncs!
544+
)
545+
546+
Init nonlinear constraints for [`NonLinModel`](@ref) and [`MultipleShooting`](@ref).
547+
548+
The nonlinear constraints are the output prediction `Ŷ` bounds, the custom inequality
549+
constraints `gc` and all the nonlinear equality constraints `geq`.
550+
"""
551+
function init_nonlincon!(
552+
mpc::PredictiveController, ::NonLinModel, ::MultipleShooting,
553+
gfuncs, ∇gfuncs!,
554+
geqfuncs, ∇geqfuncs!
555+
)
556+
optim, con = mpc.optim, mpc.con
557+
ny, nx̂, Hp, nZ̃ = mpc.estim.model.ny, mpc.estim.nx̂, mpc.Hp, length(mpc.Z̃)
558+
# --- nonlinear inequality constraints ---
559+
if length(con.i_g) ≠ 0
560+
i_base = 0
561+
for i in eachindex(con.Y0min)
562+
name = Symbol("g_Y0min_$i")
563+
optim[name] = JuMP.add_nonlinear_operator(
564+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
565+
)
566+
end
567+
i_base = 1Hp*ny
568+
for i in eachindex(con.Y0max)
569+
name = Symbol("g_Y0max_$i")
570+
optim[name] = JuMP.add_nonlinear_operator(
571+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
572+
)
573+
end
574+
i_base = 2Hp*ny
575+
for i in 1:con.nc
576+
name = Symbol("g_c_$i")
577+
optim[name] = JuMP.add_nonlinear_operator(
578+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
579+
)
580+
end
581+
end
582+
# --- nonlinear equality constraints ---
583+
Z̃var = optim[:Z̃var]
584+
for i in eachindex(geqfuncs)
585+
name = Symbol("geq_$i")
586+
geqfunc_i = optim[name] = JuMP.add_nonlinear_operator(
587+
optim, nZ̃, geqfuncs[i], ∇geqfuncs![i]; name
588+
)
589+
# set with @constrains here instead of set_nonlincon!, since the number of nonlinear
590+
# equality constraints is known and constant (±Inf are impossible):
591+
@constraint(optim, geqfunc_i(Z̃var...) == 0)
592+
end
593+
return nothing
594+
end
595+
596+
"""
597+
init_nonlincon!(
598+
mpc::PredictiveController, model::NonLinModel, ::SingleShooting,
599+
gfuncs, ∇gfuncs!,
600+
geqfuncs, ∇geqfuncs!
601+
)
602+
603+
Init nonlinear constraints for [`NonLinModel`](@ref) and [`SingleShooting`](@ref).
604+
605+
The nonlinear constraints are the custom inequality constraints `gc`, the output
606+
prediction `Ŷ` bounds and the terminal state `x̂end` bounds.
607+
"""
608+
function init_nonlincon!(
609+
mpc::PredictiveController, ::NonLinModel, ::SingleShooting, gfuncs, ∇gfuncs!, _ , _
610+
)
611+
optim, con = mpc.optim, mpc.con
612+
ny, nx̂, Hp, nZ̃ = mpc.estim.model.ny, mpc.estim.nx̂, mpc.Hp, length(mpc.Z̃)
613+
if length(con.i_g) ≠ 0
614+
i_base = 0
615+
for i in eachindex(con.Y0min)
616+
name = Symbol("g_Y0min_$i")
617+
optim[name] = JuMP.add_nonlinear_operator(
618+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
619+
)
620+
end
621+
i_base = 1Hp*ny
622+
for i in eachindex(con.Y0max)
623+
name = Symbol("g_Y0max_$i")
624+
optim[name] = JuMP.add_nonlinear_operator(
625+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
626+
)
627+
end
628+
i_base = 2Hp*ny
629+
for i in eachindex(con.x̂0min)
630+
name = Symbol("g_x̂0min_$i")
631+
optim[name] = JuMP.add_nonlinear_operator(
632+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
633+
)
634+
end
635+
i_base = 2Hp*ny + nx̂
636+
for i in eachindex(con.x̂0max)
637+
name = Symbol("g_x̂0max_$i")
638+
optim[name] = JuMP.add_nonlinear_operator(
639+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
640+
)
641+
end
642+
i_base = 2Hp*ny + 2nx̂
643+
for i in 1:con.nc
644+
name = Symbol("g_c_$i")
645+
optim[name] = JuMP.add_nonlinear_operator(
646+
optim, nZ̃, gfuncs[i_base+i], ∇gfuncs![i_base+i]; name
647+
)
648+
end
649+
end
650+
return nothing
651+
end
652+
511653
@doc raw"""
512654
linconstrainteq!(
513655
mpc::PredictiveController, model::LinModel, transcription::MultipleShooting

0 commit comments

Comments
 (0)