diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl index 19f3334005..c50803498f 100644 --- a/src/ModelingToolkit.jl +++ b/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") diff --git a/src/systems/diffeqs/modelingtoolkitize.jl b/src/systems/diffeqs/modelingtoolkitize.jl index 107a538dde..edfecdc705 100644 --- a/src/systems/diffeqs/modelingtoolkitize.jl +++ b/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 diff --git a/src/systems/optimization/modelingtoolkitize.jl b/src/systems/optimization/modelingtoolkitize.jl new file mode 100644 index 0000000000..715989a047 --- /dev/null +++ b/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 diff --git a/src/systems/optimization/optimizationsystem.jl b/src/systems/optimization/optimizationsystem.jl index e07d716f41..a5699c8648 100644 --- a/src/systems/optimization/optimizationsystem.jl +++ b/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 diff --git a/test/optimizationsystem.jl b/test/optimizationsystem.jl index a09b2bffde..cd27d6e1b4 100644 --- a/test/optimizationsystem.jl +++ b/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