Merge pull request #1660 from SciML/optprobconsexpr

Constraints and expression handling in `OptimizationProblem`

Author: ChrisRackauckas

Project.toml

@@ -79,9 +79,13 @@ Unitful = "1.1"
 julia = "1.6"
 
 [extras]
+AmplNLWriter = "7c4d4715-977e-5154-bfe0-e096adeac482"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+OptimizationMOI = "fd9f6733-72f4-499f-8506-86b2bdd0dea1"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -93,4 +97,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BenchmarkTools", "ForwardDiff", "Optimization", "OptimizationOptimJL", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]
+test = ["AmplNLWriter", "BenchmarkTools", "ForwardDiff", "Ipopt", "Ipopt_jll", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]

src/systems/control/controlsystem.jl

@@ -223,6 +223,6 @@ function runge_kutta_discretize(sys::ControlSystem, dt, tspan;
                       reduce(vcat, reduce(vcat, control_timeseries)))
 
     OptimizationSystem(reduce(+, losses, init = 0), opt_states, ps,
-                       equality_constraints = equalities, name = nameof(sys),
+                       constraints = equalities, name = nameof(sys),
                        checks = false)
 end

src/systems/optimization/optimizationsystem.jl

@@ -26,8 +26,7 @@ struct OptimizationSystem <: AbstractTimeIndependentSystem
     """Array variables."""
     var_to_name::Any
     observed::Vector{Equation}
-    equality_constraints::Vector{Equation}
-    inequality_constraints::Vector
+    constraints::Vector
     """
     Name: the name of the system.  These are required to have unique names.
     """
@@ -41,24 +40,22 @@ struct OptimizationSystem <: AbstractTimeIndependentSystem
     parameters are not supplied in `ODEProblem`.
     """
     defaults::Dict
-    function OptimizationSystem(op, states, ps, var_to_name, observed, equality_constraints,
-                                inequality_constraints, name, systems, defaults;
+    function OptimizationSystem(op, states, ps, var_to_name, observed,
+                                constraints, name, systems, defaults;
                                 checks::Bool = true)
         if checks
             check_units(op)
             check_units(observed)
-            check_units(equality_constraints)
-            all_dimensionless([states; ps]) || check_units(inequality_constraints)
+            all_dimensionless([states; ps]) || check_units(constraints)
         end
-        new(op, states, ps, var_to_name, observed, equality_constraints,
-            inequality_constraints, name, systems, defaults)
+        new(op, states, ps, var_to_name, observed,
+            constraints, name, systems, defaults)
     end
 end
 
 function OptimizationSystem(op, states, ps;
                             observed = [],
-                            equality_constraints = Equation[],
-                            inequality_constraints = [],
+                            constraints = [],
                             default_u0 = Dict(),
                             default_p = Dict(),
                             defaults = _merge(Dict(default_u0), Dict(default_p)),
@@ -85,7 +82,7 @@ function OptimizationSystem(op, states, ps;
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
     OptimizationSystem(value(op), states, ps, var_to_name,
                        observed,
-                       equality_constraints, inequality_constraints,
+                       constraints,
                        name, systems, defaults; checks = checks)
 end
 
@@ -133,6 +130,13 @@ function DiffEqBase.OptimizationProblem(sys::OptimizationSystem, args...; kwargs
     DiffEqBase.OptimizationProblem{true}(sys::OptimizationSystem, args...; kwargs...)
 end
 
+function rep_pars_vals!(e::Expr, p)
+    rep_pars_vals!.(e.args, Ref(p))
+    replace!(e.args, p...)
+end
+
+function rep_pars_vals!(e, p) end
+
 """
 ```julia
 function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,u0map,
@@ -160,9 +164,26 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
     dvs = states(sys)
     ps = parameters(sys)
 
+    defs = defaults(sys)
+    defs = mergedefaults(defs, parammap, ps)
+    defs = mergedefaults(defs, u0map, dvs)
+
+    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
+    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
+    lb = varmap_to_vars(lb, dvs; check = false, tofloat = false, use_union)
+    ub = varmap_to_vars(ub, dvs; check = false, tofloat = false, use_union)
+
     f = generate_function(sys, checkbounds = checkbounds, linenumbers = linenumbers,
                           expression = Val{false})
 
+    obj_expr = toexpr(equations(sys))
+    pairs_arr = p isa SciMLBase.NullParameters ?
+                [Symbol(_s) => Expr(:ref, :x, i) for (i, _s) in enumerate(dvs)] :
+                [
+        [Symbol(_s) => Expr(:ref, :x, i) for (i, _s) in enumerate(dvs)]...,
+        [Symbol(_p) => p[i] for (i, _p) in enumerate(ps)]...,
+    ]
+    rep_pars_vals!(obj_expr, pairs_arr)
     if grad
         grad_oop, grad_iip = generate_gradient(sys, checkbounds = checkbounds,
                                                linenumbers = linenumbers,
@@ -190,20 +211,46 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         hess_prototype = nothing
     end
 
-    _f = DiffEqBase.OptimizationFunction{iip}(f,
-                                              SciMLBase.NoAD();
-                                              grad = _grad,
-                                              hess = _hess,
-                                              hess_prototype = hess_prototype)
+    if length(sys.constraints) > 0
+        @named cons_sys = NonlinearSystem(sys.constraints, dvs, ps)
+        cons = generate_function(cons_sys, checkbounds = checkbounds,
+                                 linenumbers = linenumbers,
+                                 expression = Val{false})[1]
+        cons_j = generate_jacobian(cons_sys; expression = Val{false}, sparse = sparse)[2]
+        cons_h = generate_hessian(cons_sys; expression = Val{false}, sparse = sparse)[2]
+
+        cons_expr = toexpr(equations(cons_sys))
+        rep_pars_vals!.(cons_expr, Ref(pairs_arr))
+
+        if sparse
+            cons_jac_prototype = jacobian_sparsity(cons_sys)
+            cons_hess_prototype = hessian_sparsity(cons_sys)
+        else
+            cons_jac_prototype = nothing
+            cons_hess_prototype = nothing
+        end
 
-    defs = defaults(sys)
-    defs = mergedefaults(defs, parammap, ps)
-    defs = mergedefaults(defs, u0map, dvs)
+        _f = DiffEqBase.OptimizationFunction{iip}(f,
+                                                  SciMLBase.NoAD();
+                                                  grad = _grad,
+                                                  hess = _hess,
+                                                  hess_prototype = hess_prototype,
+                                                  cons = cons,
+                                                  cons_j = cons_j,
+                                                  cons_h = cons_h,
+                                                  cons_jac_prototype = cons_jac_prototype,
+                                                  cons_hess_prototype = cons_hess_prototype,
+                                                  expr = obj_expr,
+                                                  cons_expr = cons_expr)
+    else
+        _f = DiffEqBase.OptimizationFunction{iip}(f,
+                                                  SciMLBase.NoAD();
+                                                  grad = _grad,
+                                                  hess = _hess,
+                                                  hess_prototype = hess_prototype,
+                                                  expr = obj_expr)
+    end
 
-    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
-    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
-    lb = varmap_to_vars(lb, dvs; check = false, tofloat = false, use_union)
-    ub = varmap_to_vars(ub, dvs; check = false, tofloat = false, use_union)
     OptimizationProblem{iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
 end
 
@@ -272,18 +319,74 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
     p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
     lb = varmap_to_vars(lb, dvs; check = false, tofloat = false, use_union)
     ub = varmap_to_vars(ub, dvs; check = false, tofloat = false, use_union)
-    quote
-        f = $f
-        p = $p
-        u0 = $u0
-        grad = $_grad
-        hess = $_hess
-        lb = $lb
-        ub = $ub
-        _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
-                                       grad = grad,
-                                       hess = hess,
-                                       hess_prototype = hess_prototype)
-        OptimizationProblem{$iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
+
+    obj_expr = toexpr(equations(sys))
+    pairs_arr = p isa SciMLBase.NullParameters ?
+                [Symbol(_s) => Expr(:ref, :x, i) for (i, _s) in enumerate(dvs)] :
+                [
+        [Symbol(_s) => Expr(:ref, :x, i) for (i, _s) in enumerate(dvs)]...,
+        [Symbol(_p) => p[i] for (i, _p) in enumerate(ps)]...,
+    ]
+    rep_pars_vals!(obj_expr, pairs_arr)
+
+    if length(sys.constraints) > 0
+        @named cons_sys = NonlinearSystem(sys.constraints, dvs, ps)
+        cons = generate_function(cons_sys, checkbounds = checkbounds,
+                                 linenumbers = linenumbers,
+                                 expression = Val{false})[1]
+        cons_j = generate_jacobian(cons_sys; expression = Val{false}, sparse = sparse)[2]
+
+        cons_h = generate_hessian(cons_sys; expression = Val{false}, sparse = sparse)[2]
+
+        cons_expr = toexpr(equations(cons_sys))
+        rep_pars_vals!.(cons_expr, Ref(pairs_arr))
+
+        if sparse
+            cons_jac_prototype = jacobian_sparsity(cons_sys)
+            cons_hess_prototype = hessian_sparsity(cons_sys)
+        else
+            cons_jac_prototype = nothing
+            cons_hess_prototype = nothing
+        end
+        quote
+            f = $f
+            p = $p
+            u0 = $u0
+            grad = $_grad
+            hess = $_hess
+            lb = $lb
+            ub = $ub
+            cons = $cons
+            cons_j = $cons_j
+            cons_h = $cons_h
+            _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
+                                           grad = grad,
+                                           hess = hess,
+                                           hess_prototype = hess_prototype,
+                                           cons = cons,
+                                           cons_j = cons_j,
+                                           cons_h = cons_h,
+                                           cons_jac_prototype = cons_jac_prototype,
+                                           cons_hess_prototype = cons_hess_prototype,
+                                           expr = obj_expr,
+                                           cons_expr = cons_expr)
+            OptimizationProblem{$iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
+        end
+    else
+        quote
+            f = $f
+            p = $p
+            u0 = $u0
+            grad = $_grad
+            hess = $_hess
+            lb = $lb
+            ub = $ub
+            _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
+                                           grad = grad,
+                                           hess = hess,
+                                           hess_prototype = hess_prototype,
+                                           expr = obj_expr)
+            OptimizationProblem{$iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
+        end
     end
 end

test/controlsystem.jl

@@ -16,6 +16,7 @@ sys = runge_kutta_discretize(sys, dt, tspan)
 u0 = rand(length(states(sys))) # guess for the state values
 prob = OptimizationProblem(sys, u0, [0.1, 0.1], grad = true)
 sol = solve(prob, BFGS())
+@test prob.f(sol.minimizer, prob.p) < prob.f(u0, prob.p)
 
 # issue #819
 @testset "Combined system name collisions" begin

test/optimizationsystem.jl

@@ -1,10 +1,13 @@
-using ModelingToolkit, SparseArrays, Test, Optimization, OptimizationOptimJL
+using ModelingToolkit, SparseArrays, Test, Optimization, OptimizationOptimJL,
+      OptimizationMOI, Ipopt, AmplNLWriter, Ipopt_jll
 
 @variables x y
 @parameters a b
 loss = (a - x)^2 + b * (y - x^2)^2
 sys1 = OptimizationSystem(loss, [x, y], [a, b], name = :sys1)
-sys2 = OptimizationSystem(loss, [x, y], [a, b], name = :sys2)
+
+cons2 = [x^2 + y^2 ~ 0, y * sin(x) - x ~ 0]
+sys2 = OptimizationSystem(loss, [x, y], [a, b], name = :sys2, constraints = cons2)
 
 @variables z
 @parameters β
@@ -47,6 +50,32 @@ sol = solve(prob, BFGS(initial_stepnorm = 0.0001), allow_f_increases = true)
 sol = solve(prob2, BFGS(initial_stepnorm = 0.0001), allow_f_increases = true)
 @test sol.minimum < -1e9
 
+#inequality constraint, the bounds for constraints lcons !== ucons
+prob = OptimizationProblem(sys2, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0],
+                           lcons = [-1.0, -1.0], ucons = [500.0, 500.0], grad = true,
+                           hess = true)
+sol = solve(prob, IPNewton(), allow_f_increases = true)
+@test sol.minimum < 1.0
+sol = solve(prob, Ipopt.Optimizer())
+@test sol.minimum < 1.0
+sol = solve(prob, AmplNLWriter.Optimizer(Ipopt_jll.amplexe))
+@test sol.minimum < 1.0
+
+#equality constraint, lcons == ucons
+cons2 = [0.0 ~ x^2 + y^2]
+sys2 = OptimizationSystem(loss, [x, y], [a, b], name = :sys2, constraints = cons2)
+prob = OptimizationProblem(sys2, [x => 0.0, y => 0.0], [a => 1.0, b => 1.0], lcons = [1.0],
+                           ucons = [1.0], grad = true, hess = true)
+sol = solve(prob, IPNewton())
+@test sol.minimum < 1.0
+@test prob.f.cons(sol.minimizer, [1.0, 1.0]) ≈ [1.0]
+sol = solve(prob, Ipopt.Optimizer())
+@test sol.minimum < 1.0
+@test prob.f.cons(sol.minimizer, [1.0, 1.0]) ≈ [1.0]
+sol = solve(prob, AmplNLWriter.Optimizer(Ipopt_jll.amplexe))
+@test sol.minimum < 1.0
+@test prob.f.cons(sol.minimizer, [1.0, 1.0]) ≈ [1.0]
+
 rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
 x0 = zeros(2)
 _p = [1.0, 100.0]