|
| 1 | +module ConvexOptimization |
| 2 | + |
| 3 | +using Reexport |
| 4 | +@reexport using SciMLBase |
| 5 | +using SciMLBase: ConvexOptimizationProblem, ConvexOptimizationSolution, |
| 6 | + OptimizationFunction, AbstractOptimizationCache, AbstractOptimizationAlgorithm, |
| 7 | + NullParameters, ReturnCode |
| 8 | +import MathOptInterface as MOI |
| 9 | +import Symbolics |
| 10 | +using Symbolics: variable, unwrap, linear_expansion |
| 11 | +import SymbolicAnalysis |
| 12 | +using SymbolicAnalysis: analyze |
| 13 | +using LinearAlgebra |
| 14 | + |
| 15 | +""" |
| 16 | + ConeConstraint(g, set) |
| 17 | +
|
| 18 | +One convex cone constraint of a [`ConvexOptimizationProblem`](@ref). `g(u, p)` |
| 19 | +returns the affine map whose image must lie in the MathOptInterface vector cone |
| 20 | +`set` (`MOI.Zeros`, `MOI.Nonnegatives`, `MOI.Nonpositives`, `MOI.SecondOrderCone`, |
| 21 | +…). The output length of `g` must equal `MOI.dimension(set)`. |
| 22 | +
|
| 23 | +The backend traces `g` on its own symbolic variables, so each `ConeConstraint` |
| 24 | +maps to exactly one MOI constraint and therefore one entry of |
| 25 | +`ConvexOptimizationSolution.dual`, in the order the constraints are given. |
| 26 | +Because the cone is named explicitly, the returned dual is already expressed in |
| 27 | +the user's variables (no sign remap): `>=` → `MOI.Nonnegatives`, `<=` → |
| 28 | +`MOI.Nonpositives`, `==` → `MOI.Zeros`. |
| 29 | +""" |
| 30 | +struct ConeConstraint{G, S <: MOI.AbstractVectorSet} |
| 31 | + g::G |
| 32 | + set::S |
| 33 | +end |
| 34 | + |
| 35 | +abstract type AbstractConvexOptAlgorithm <: AbstractOptimizationAlgorithm end |
| 36 | + |
| 37 | +""" |
| 38 | + ConvexMOI(optimizer_constructor = Clarabel.Optimizer) |
| 39 | +
|
| 40 | +Conic backend: certify convexity with SymbolicAnalysis, lower the (affine) |
| 41 | +objective and each `ConeConstraint` to a MathOptInterface cone, and solve with |
| 42 | +`optimizer_constructor`. |
| 43 | +""" |
| 44 | +struct ConvexMOI{O} <: AbstractConvexOptAlgorithm |
| 45 | + optimizer_constructor::O |
| 46 | +end |
| 47 | + |
| 48 | +SciMLBase.allowsbounds(::AbstractConvexOptAlgorithm) = true |
| 49 | +SciMLBase.allowsconstraints(::AbstractConvexOptAlgorithm) = true |
| 50 | + |
| 51 | +# Must be <: AbstractOptimizationCache (build_convex_solution requires it) and |
| 52 | +# carry real `f`/`p` fields for the solution's SymbolicIndexingInterface glue. |
| 53 | +# No `reinit_cache` field (that would reroute getproperty(:u0/:p)). |
| 54 | +struct ConvexOptimizationCache{F, U, P, A, AR, MOD, XV, CR} <: AbstractOptimizationCache |
| 55 | + f::F |
| 56 | + u0::U |
| 57 | + p::P |
| 58 | + alg::A |
| 59 | + analysis::AR |
| 60 | + model::MOD # lowered MOI model |
| 61 | + xvars::XV # Vector{MOI.VariableIndex}: user u -> MOI variables |
| 62 | + conrefs::CR # Vector{MOI.ConstraintIndex}, 1:1 with prob.constraints |
| 63 | +end |
| 64 | + |
| 65 | +# `solve(prob, alg)` routes through CommonSolve: solve = solve! ∘ init. Neither |
| 66 | +# `init` nor `solve!` is inherited here (no OptimizationBase in the dep tree), so |
| 67 | +# both thin methods are defined explicitly. |
| 68 | +function SciMLBase.init( |
| 69 | + prob::ConvexOptimizationProblem, |
| 70 | + alg::AbstractConvexOptAlgorithm, args...; kwargs... |
| 71 | + ) |
| 72 | + return SciMLBase.__init(prob, alg, args...; prob.kwargs..., kwargs...) |
| 73 | +end |
| 74 | +SciMLBase.solve!(cache::ConvexOptimizationCache) = SciMLBase.__solve(cache) |
| 75 | + |
| 76 | +function SciMLBase.__init( |
| 77 | + prob::ConvexOptimizationProblem, |
| 78 | + alg::AbstractConvexOptAlgorithm, args...; kwargs... |
| 79 | + ) |
| 80 | + analysis = certify_convex(prob) |
| 81 | + model, xvars, conrefs = lower_to_moi(prob, alg) |
| 82 | + return ConvexOptimizationCache( |
| 83 | + prob.f, prob.u0, prob.p, alg, analysis, model, xvars, conrefs |
| 84 | + ) |
| 85 | +end |
| 86 | + |
| 87 | +function SciMLBase.__solve(cache::ConvexOptimizationCache) |
| 88 | + model = cache.model |
| 89 | + MOI.optimize!(model) |
| 90 | + ret = _moi_status_to_retcode(MOI.get(model, MOI.TerminationStatus())) |
| 91 | + if MOI.get(model, MOI.ResultCount()) >= 1 |
| 92 | + u = MOI.get(model, MOI.VariablePrimal(), cache.xvars) |
| 93 | + objective = MOI.get(model, MOI.ObjectiveValue()) |
| 94 | + else |
| 95 | + u = fill(NaN, length(cache.xvars)) |
| 96 | + objective = NaN |
| 97 | + end |
| 98 | + dual = if MOI.get(model, MOI.DualStatus()) == MOI.NO_SOLUTION |
| 99 | + nothing |
| 100 | + else |
| 101 | + [MOI.get(model, MOI.ConstraintDual(), c) for c in cache.conrefs] |
| 102 | + end |
| 103 | + return SciMLBase.build_convex_solution( |
| 104 | + cache, cache.alg, u, objective; |
| 105 | + dual = dual, retcode = ret, original = model, |
| 106 | + stats = SciMLBase.OptimizationStats() |
| 107 | + ) |
| 108 | +end |
| 109 | + |
| 110 | +function certify_convex(prob::ConvexOptimizationProblem) |
| 111 | + vars, params = _symbolic_vars(prob) |
| 112 | + obj = unwrap(_scalar(prob.f.f(vars, params))) |
| 113 | + obj_res = analyze(obj) |
| 114 | + ok = prob.sense === SciMLBase.MaxSense ? |
| 115 | + obj_res.curvature in (SymbolicAnalysis.Concave, SymbolicAnalysis.Affine) : |
| 116 | + obj_res.curvature in (SymbolicAnalysis.Convex, SymbolicAnalysis.Affine) |
| 117 | + ok || error( |
| 118 | + "Objective is not certified convex for $(prob.sense): curvature = " * |
| 119 | + "$(obj_res.curvature). Route to a general OptimizationProblem/NLP solver." |
| 120 | + ) |
| 121 | + cons_res = _certify_constraints(prob, vars, params) |
| 122 | + return (; objective = obj_res, constraints = cons_res) |
| 123 | +end |
| 124 | + |
| 125 | +# MVP: constraints are affine-in-cone, so every output component must be Affine. |
| 126 | +function _certify_constraints(prob, vars, params) |
| 127 | + prob.constraints === nothing && return nothing |
| 128 | + res = [] |
| 129 | + for con in prob.constraints |
| 130 | + cres = analyze.(unwrap.(_asvec(con.g(vars, params)))) |
| 131 | + all(r -> r.curvature == SymbolicAnalysis.Affine, cres) || error( |
| 132 | + "This backend supports affine-in-cone constraints only; got " * |
| 133 | + "curvatures $(getproperty.(cres, :curvature)) for cone $(con.set)." |
| 134 | + ) |
| 135 | + push!(res, cres) |
| 136 | + end |
| 137 | + return res |
| 138 | +end |
| 139 | + |
| 140 | +function lower_to_moi(prob::ConvexOptimizationProblem, alg::ConvexMOI) |
| 141 | + model = MOI.instantiate(alg.optimizer_constructor; with_bridge_type = Float64) |
| 142 | + MOI.set(model, MOI.Silent(), true) |
| 143 | + n = length(prob.u0) |
| 144 | + x = MOI.add_variables(model, n) |
| 145 | + vars, params = _symbolic_vars(prob) |
| 146 | + |
| 147 | + if prob.lb !== nothing |
| 148 | + for i in 1:n |
| 149 | + prob.lb[i] > -Inf && |
| 150 | + MOI.add_constraint(model, x[i], MOI.GreaterThan(Float64(prob.lb[i]))) |
| 151 | + prob.ub[i] < Inf && |
| 152 | + MOI.add_constraint(model, x[i], MOI.LessThan(Float64(prob.ub[i]))) |
| 153 | + end |
| 154 | + end |
| 155 | + |
| 156 | + conrefs = MOI.ConstraintIndex[] |
| 157 | + if prob.constraints !== nothing |
| 158 | + for con in prob.constraints |
| 159 | + gvals = _asvec(con.g(vars, params)) |
| 160 | + A, b, islin = linear_expansion(gvals, vars) # gvals == A*vars + b |
| 161 | + islin || error("Constraint $(con.set) is not affine in the variables.") |
| 162 | + f = _affine_to_vaf(_tofloat.(A), _tofloat.(b), x) |
| 163 | + push!(conrefs, MOI.add_constraint(model, f, con.set)) |
| 164 | + end |
| 165 | + end |
| 166 | + |
| 167 | + objexpr = _asvec(_scalar(prob.f.f(vars, params))) |
| 168 | + Ao, bo, olin = linear_expansion(objexpr, vars) |
| 169 | + olin || error("This backend requires an affine objective.") |
| 170 | + c = vec(_tofloat.(Ao)) |
| 171 | + d = _tofloat(only(bo)) |
| 172 | + saterms = [MOI.ScalarAffineTerm(c[j], x[j]) for j in 1:n if !iszero(c[j])] |
| 173 | + MOI.set( |
| 174 | + model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), |
| 175 | + MOI.ScalarAffineFunction(saterms, d) |
| 176 | + ) |
| 177 | + MOI.set( |
| 178 | + model, MOI.ObjectiveSense(), |
| 179 | + prob.sense === SciMLBase.MaxSense ? MOI.MAX_SENSE : MOI.MIN_SENSE |
| 180 | + ) |
| 181 | + return model, x, conrefs |
| 182 | +end |
| 183 | + |
| 184 | +function _symbolic_vars(prob) |
| 185 | + vars = [variable(:x, i) for i in 1:length(prob.u0)] |
| 186 | + params = prob.p isa NullParameters ? Float64[] : |
| 187 | + [variable(:α, i) for i in eachindex(prob.p)] |
| 188 | + return vars, params |
| 189 | +end |
| 190 | + |
| 191 | +_asvec(v::AbstractVector) = v |
| 192 | +_asvec(v) = [v] |
| 193 | +_scalar(v::AbstractVector) = only(v) |
| 194 | +_scalar(v) = v |
| 195 | + |
| 196 | +_tofloat(x) = Float64(Symbolics.value(x)) |
| 197 | + |
| 198 | +function _affine_to_vaf(A::AbstractMatrix, b::AbstractVector, x) |
| 199 | + terms = MOI.VectorAffineTerm{Float64}[] |
| 200 | + m, n = size(A) |
| 201 | + for i in 1:m, j in 1:n |
| 202 | + iszero(A[i, j]) && continue |
| 203 | + push!(terms, MOI.VectorAffineTerm(i, MOI.ScalarAffineTerm(A[i, j], x[j]))) |
| 204 | + end |
| 205 | + return MOI.VectorAffineFunction(terms, collect(float.(b))) |
| 206 | +end |
| 207 | + |
| 208 | +function _moi_status_to_retcode(s::MOI.TerminationStatusCode) |
| 209 | + s in ( |
| 210 | + MOI.OPTIMAL, MOI.LOCALLY_SOLVED, MOI.ALMOST_OPTIMAL, |
| 211 | + MOI.ALMOST_LOCALLY_SOLVED, |
| 212 | + ) && return ReturnCode.Success |
| 213 | + s in ( |
| 214 | + MOI.INFEASIBLE, MOI.DUAL_INFEASIBLE, MOI.LOCALLY_INFEASIBLE, |
| 215 | + MOI.INFEASIBLE_OR_UNBOUNDED, |
| 216 | + ) && return ReturnCode.Infeasible |
| 217 | + s == MOI.TIME_LIMIT && return ReturnCode.MaxTime |
| 218 | + s in (MOI.ITERATION_LIMIT, MOI.NODE_LIMIT, MOI.SLOW_PROGRESS) && |
| 219 | + return ReturnCode.MaxIters |
| 220 | + s in (MOI.NUMERICAL_ERROR, MOI.INVALID_MODEL, MOI.OTHER_ERROR) && |
| 221 | + return ReturnCode.Failure |
| 222 | + return ReturnCode.Default |
| 223 | +end |
| 224 | + |
| 225 | +export ConvexMOI, ConeConstraint |
| 226 | + |
| 227 | +end # module |
0 commit comments