Handle constraints inside modelingtoolkitize of OptimizationSystem

src/ModelingToolkit.jl

@@ -152,6 +152,7 @@ include("systems/nonlinear/modelingtoolkitize.jl")
 
 include("systems/optimization/constraints_system.jl")
 include("systems/optimization/optimizationsystem.jl")
+include("systems/optimization/modelingtoolkitize.jl")
 
 include("systems/pde/pdesystem.jl")
 

src/systems/diffeqs/modelingtoolkitize.jl

@@ -217,26 +217,3 @@ function modelingtoolkitize(prob::DiffEqBase.SDEProblem; kwargs...)
 
     de
 end
-
-"""
-$(TYPEDSIGNATURES)
-
-Generate `OptimizationSystem`, dependent variables, and parameters from an `OptimizationProblem`.
-"""
-function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem; kwargs...)
-    if prob.p isa Tuple || prob.p isa NamedTuple
-        p = [x for x in prob.p]
-    else
-        p = prob.p
-    end
-
-    vars = reshape([variable(:x, i) for i in eachindex(prob.u0)], size(prob.u0))
-    params = p isa DiffEqBase.NullParameters ? [] :
-             reshape([variable(:α, i) for i in eachindex(p)], size(Array(p)))
-
-    eqs = prob.f(vars, params)
-    de = OptimizationSystem(eqs, vec(vars), vec(params);
-                            name = gensym(:MTKizedOpt),
-                            kwargs...)
-    de
-end

src/systems/optimization/modelingtoolkitize.jl

@@ -0,0 +1,63 @@
+"""
+$(TYPEDSIGNATURES)
+
+Generate `OptimizationSystem`, dependent variables, and parameters from an `OptimizationProblem`.
+"""
+function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem; kwargs...)
+    num_cons = isnothing(prob.lcons) ? 0 : length(prob.lcons)
+    if prob.p isa Tuple || prob.p isa NamedTuple
+        p = [x for x in prob.p]
+    else
+        p = prob.p
+    end
+
+    vars = ArrayInterfaceCore.restructure(prob.u0,
+                                          [variable(:x, i) for i in eachindex(prob.u0)])
+    params = p isa DiffEqBase.NullParameters ? [] :
+             ArrayInterfaceCore.restructure(p, [variable(:α, i) for i in eachindex(p)])
+
+    eqs = prob.f(vars, params)
+
+    if DiffEqBase.isinplace(prob) && !isnothing(prob.f.cons)
+        lhs = Array{Num}(undef, num_cons)
+        prob.f.cons(lhs, vars, params)
+        cons = Union{Equation, Inequality}[]
+
+        if !isnothing(prob.lcons)
+            for i in 1:num_cons
+                !isinf(prob.lcons[i]) && prob.lcons[i] !== prob.ucons[i] &&
+                    push!(cons, prob.lcons[i] ≲ lhs[i])
+                if !isinf(prob.ucons[i])
+                    prob.lcons[i] == prob.ucons[i] ? push!(cons, lhs[i] ~ prob.ucons[i]) :
+                    push!(cons, lhs[i] ≲ prob.ucons[i])
+                end
+            end
+        end
+
+        if !isnothing(prob.ucons)
+            for i in 1:num_cons
+                if !isinf(prob.ucons[i])
+                    prob.lcons[i] == prob.ucons[i] ? push!(cons, lhs[i] ~ prob.ucons[i]) :
+                    push!(cons, lhs[i] ≲ prob.ucons[i])
+                end
+            end
+        end
+
+        if (isnothing(prob.lcons) || all(isinf.(prob.lcons))) &&
+           (isnothing(prob.ucons) || all(isinf.(prob.ucons)))
+            throw(ArgumentError("Constraints passed have no proper bounds defined.
+            Ensure you pass equal bounds (the scalar that the constraint should evaluate to) for equality constraints
+            or pass the lower and upper bounds for inequality constraints."))
+        end
+    elseif !isnothing(prob.f.cons)
+        cons = prob.f.cons(vars, params)
+    else
+        cons = []
+    end
+
+    de = OptimizationSystem(eqs, vec(vars), vec(toparam.(params));
+                            name = gensym(:MTKizedOpt),
+                            constraints = cons,
+                            kwargs...)
+    de
+end

src/systems/optimization/optimizationsystem.jl

@@ -80,7 +80,6 @@ function OptimizationSystem(op, states, ps;
                             metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
-
     constraints = value.(scalarize(constraints))
     states′ = value.(scalarize(states))
     ps′ = value.(scalarize(ps))
@@ -199,7 +198,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,u0map,
 
 Generates an OptimizationProblem from an OptimizationSystem and allows for automatically
 symbolically calculating numerical enhancements.
-                
+
 Certain solvers require setting `cons_j`, `cons_h` to `true` for constrained-optimization problems.
 """
 function DiffEqBase.OptimizationProblem(sys::OptimizationSystem, args...; kwargs...)
@@ -211,7 +210,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                              grad = false,
                                              hess = false, sparse = false,
                                              cons_j = false, cons_h = false,
-                                             checkbounds = false,
+                                             cons_sparse = false, checkbounds = false,
                                              linenumbers = true, parallel = SerialForm(),
                                              use_union = false,
                                              kwargs...) where {iip}
@@ -313,12 +312,14 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                                  linenumbers = linenumbers,
                                                  expression = Val{false})
         if cons_j
-            _cons_j = generate_jacobian(cons_sys; expression = Val{false}, sparse = sparse)[2]
+            _cons_j = generate_jacobian(cons_sys; expression = Val{false},
+                                        sparse = cons_sparse)[2]
         else
             _cons_j = nothing
         end
         if cons_h
-            _cons_h = generate_hessian(cons_sys; expression = Val{false}, sparse = sparse)[2]
+            _cons_h = generate_hessian(cons_sys; expression = Val{false},
+                                       sparse = cons_sparse)[2]
         else
             _cons_h = nothing
         end
@@ -339,7 +340,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
             ucons = haskey(kwargs, :ucons)
         end
 
-        if sparse
+        if cons_sparse
             cons_jac_prototype = jacobian_sparsity(cons_sys)
             cons_hess_prototype = hessian_sparsity(cons_sys)
         else
@@ -507,7 +508,7 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0map,
             lcons = lcons_
             ucons = ucons_
         else # use the user supplied constraints bounds
-            haskey(kwargs, :lcons) && haskey(kwargs, :ucons) &&
+            !haskey(kwargs, :lcons) && !haskey(kwargs, :ucons) &&
                 throw(ArgumentError("Expected both `ucons` and `lcons` to be supplied"))
             haskey(kwargs, :lcons) && length(kwargs[:lcons]) != length(cstr) &&
                 throw(ArgumentError("Expected `lcons` to be of the same length as the vector of constraints"))
@@ -598,12 +599,12 @@ function structural_simplify(sys::OptimizationSystem; kwargs...)
     obs = observed(snlsys)
     subs = Dict(eq.lhs => eq.rhs for eq in observed(snlsys))
     seqs = equations(snlsys)
-    sizehint!(icons, length(icons) + length(seqs))
-    for eq in seqs
-        push!(icons, substitute(eq, subs))
+    cons_simplified = Array{eltype(cons), 1}(undef, length(icons) + length(seqs))
+    for (i, eq) in enumerate(Iterators.flatten((seqs, icons)))
+        cons_simplified[i] = substitute(eq, subs)
     end
     newsts = setdiff(states(sys), keys(subs))
-    @set! sys.constraints = icons
+    @set! sys.constraints = cons_simplified
     @set! sys.observed = [observed(sys); obs]
     @set! sys.op = substitute(equations(sys), subs)
     @set! sys.states = newsts

test/optimizationsystem.jl

@@ -77,7 +77,8 @@ end
     @parameters a b
     loss = (a - x)^2 + b * z^2
     cons = [1.0 ~ x^2 + y^2
-            z ~ y - x^2]
+            z ~ y - x^2
+            z^2 + y^2 ≲ 1.0]
     @named sys = OptimizationSystem(loss, [x, y, z], [a, b], constraints = cons)
     sys = structural_simplify(sys)
     prob = OptimizationProblem(sys, [x => 0.0, y => 0.0, z => 0.0], [a => 1.0, b => 1.0],
@@ -246,3 +247,25 @@ end
     @test prob.lb == [-Inf, 0.0]
     @test prob.ub == [Inf, Inf]
 end
+
+@testset "modelingtoolkitize" begin
+    @variables x₁ x₂
+    @parameters α₁ α₂
+    loss = (α₁ - x₁)^2 + α₂ * (x₂ - x₁^2)^2
+    cons = [
+        x₁^2 + x₂^2 ≲ 1.0,
+    ]
+    sys1 = OptimizationSystem(loss, [x₁, x₂], [α₁, α₂], name = :sys1, constraints = cons)
+
+    prob1 = OptimizationProblem(sys1, [x₁ => 0.0, x₂ => 0.0], [α₁ => 1.0, α₂ => 100.0],
+                                grad = true, hess = true, cons_j = true, cons_h = true)
+
+    sys2 = modelingtoolkitize(prob1)
+    prob2 = OptimizationProblem(sys2, [x₁ => 0.0, x₂ => 0.0], [α₁ => 1.0, α₂ => 100.0],
+                                grad = true, hess = true, cons_j = true, cons_h = true)
+
+    sol1 = Optimization.solve(prob1, Ipopt.Optimizer())
+    sol2 = Optimization.solve(prob2, Ipopt.Optimizer())
+
+    @test sol1.u ≈ sol2.u
+end