diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index 468f44cdfe..f5c3746069 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -73,9 +73,8 @@ Nonlinear.ExprGraphOnly Nonlinear.SparseReverseMode Nonlinear.SymbolicMode Nonlinear.QPBlockData -Nonlinear.add_constraint_jacobian_product -Nonlinear.add_constraint_jacobian_transpose_product -Nonlinear.add_hessian_lagrangian_product +Nonlinear.ModelWithQuad +Nonlinear.EvaluatorWithQuad ``` diff --git a/src/Nonlinear/Nonlinear.jl b/src/Nonlinear/Nonlinear.jl index e057052029..2475fba33c 100644 --- a/src/Nonlinear/Nonlinear.jl +++ b/src/Nonlinear/Nonlinear.jl @@ -43,5 +43,6 @@ include("ReverseAD/ReverseAD.jl") include("SymbolicAD/SymbolicAD.jl") include("qp_block_data.jl") +include("model_with_quad.jl") end # module diff --git a/src/Nonlinear/model_with_quad.jl b/src/Nonlinear/model_with_quad.jl new file mode 100644 index 0000000000..8176616abf --- /dev/null +++ b/src/Nonlinear/model_with_quad.jl @@ -0,0 +1,545 @@ +# Copyright (c) 2017: Miles Lubin and contributors +# Copyright (c) 2017: Google Inc. +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. + +""" + ModelWithQuad{T,M}( + qp::QPBlockData{T}, + inner::M; + objective_sink::Symbol = :none, + ) where {T,M} + +A model layer that owns the variables of the model, stores affine and +quadratic objectives and constraints in a [`QPBlockData`](@ref), and forwards +everything else to the `inner` model, typically a [`Model`](@ref). + +`ModelWithQuad(inner)` and `ModelWithQuad{T}(inner)` create an empty +[`QPBlockData`](@ref), with `T` defaulting to `Float64`. + +Add variables with `MOI.add_variable`: the layer guarantees that the variable +indices are `1:n`, like `MOI.Utilities.MatrixOfConstraints`. Add parameters +with `MOI.add_constrained_variable(model, ::MOI.Parameter)`: parameters get +indices offset by [`_PARAMETER_OFFSET`](@ref), and their values are stored in +the inner model through [`add_parameter`](@ref). The inner model must expose +that storage as `parameters::Vector{T}`, like [`Model`](@ref) does: +`qp.parameters` aliases it, so a parameter update is visible to both blocks. + +Add constraints with [`add_constraint`](@ref) or `MOI.add_constraint`, and +set the objective with [`set_objective`](@ref): affine and quadratic +functions are routed to the QP block, everything else to the inner model. +`objective_sink` records where the objective currently lives (`:none`, +`:quad` or `:inner`). + +Create the corresponding evaluator, [`EvaluatorWithQuad`](@ref), with +`Evaluator(model, backend)`, or construct it directly from an inner +`MOI.AbstractNLPEvaluator`. The rows of the QP block come first, followed by +the rows of the inner evaluator. +""" +mutable struct ModelWithQuad{T,M} + variables::MOI.Utilities.VariablesContainer{T} + qp::QPBlockData{T} + inner::M + objective_sink::Symbol # :none, :quad or :inner + + function ModelWithQuad{T}( + qp::QPBlockData{T}, + inner::M; + objective_sink::Symbol = :none, + ) where {T,M} + model = new{T,M}( + MOI.Utilities.VariablesContainer{T}(), + qp, + inner, + objective_sink, + ) + # The QP block reads the parameter values from the storage of the + # inner model, which must expose them as `parameters::Vector{T}`, + # like [`Model`](@ref) does. + model.qp.parameters = inner.parameters + return model + end +end + +function ModelWithQuad{T}(inner) where {T} + return ModelWithQuad{T}(QPBlockData{T}(), inner) +end + +ModelWithQuad(inner) = ModelWithQuad{Float64}(inner) + +# The variables and the parameters. + +MOI.add_variable(model::ModelWithQuad) = MOI.add_variable(model.variables) + +function MOI.add_constrained_variable( + model::ModelWithQuad{T}, + set::MOI.Parameter{T}, +) where {T} + p = add_parameter(model.inner, set.value) + x = MOI.VariableIndex(_PARAMETER_OFFSET + p.value) + ci = MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}(x.value) + return x, ci +end + +function MOI.is_valid(model::ModelWithQuad, x::MOI.VariableIndex) + if _is_parameter(x) + return 1 <= x.value - _PARAMETER_OFFSET <= length(model.qp.parameters) + end + return MOI.is_valid(model.variables, x) +end + +function MOI.is_valid( + model::ModelWithQuad{T}, + ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}, +) where {T} + return MOI.is_valid(model, MOI.VariableIndex(ci.value)) +end + +function MOI.get( + model::ModelWithQuad{T}, + ::MOI.NumberOfConstraints{MOI.VariableIndex,MOI.Parameter{T}}, +) where {T} + return length(model.qp.parameters) +end + +function MOI.get( + model::ModelWithQuad{T}, + ::MOI.ListOfConstraintIndices{F,S}, +) where {T,F<:MOI.VariableIndex,S<:MOI.Parameter{T}} + n = length(model.qp.parameters) + return MOI.ConstraintIndex{F,S}.(_PARAMETER_OFFSET .+ (1:n)) +end + +function MOI.get( + model::ModelWithQuad{T}, + ::MOI.ConstraintFunction, + ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}, +) where {T} + return MOI.VariableIndex(ci.value) +end + +function MOI.get( + model::ModelWithQuad{T}, + ::MOI.ConstraintSet, + ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}, +) where {T} + return MOI.Parameter(model.qp.parameters[ci.value-_PARAMETER_OFFSET]) +end + +function MOI.set( + model::ModelWithQuad{T}, + ::MOI.ConstraintSet, + ci::MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}, + set::MOI.Parameter{T}, +) where {T} + model.qp.parameters[ci.value-_PARAMETER_OFFSET] = set.value + return +end + +""" + Base.length(model::ModelWithQuad) + +The number of affine and quadratic constraint rows of `model`, which come +before the rows of the inner model in the corresponding evaluator. +""" +Base.length(model::ModelWithQuad) = length(model.qp) + +# Replace the parameters of `f`, encoded as `MOI.VariableIndex`es offset by +# [`_PARAMETER_OFFSET`](@ref), by the corresponding [`ParameterIndex`](@ref), +# which the inner model understands. An affine or quadratic function that +# contains a parameter is converted to `MOI.ScalarNonlinearFunction`, because +# the inner model parses such functions with their variable indices verbatim. +_replace_parameters(f) = f + +function _replace_parameters(f::MOI.VariableIndex) + if _is_parameter(f) + return ParameterIndex(f.value - _PARAMETER_OFFSET) + end + return f +end + +function _replace_parameters(f::MOI.ScalarAffineFunction) + if any(_is_parameter, f.terms) + return _replace_parameters(convert(MOI.ScalarNonlinearFunction, f)) + end + return f +end + +function _replace_parameters(f::MOI.ScalarQuadraticFunction) + if any(_is_parameter, f.affine_terms) || + any(_is_parameter, f.quadratic_terms) + return _replace_parameters(convert(MOI.ScalarNonlinearFunction, f)) + end + return f +end + +function _replace_parameters(f::MOI.ScalarNonlinearFunction) + for (i, arg) in enumerate(f.args) + f.args[i] = _replace_parameters(arg) + end + return f +end + +# Methods forwarded to the inner model. + +function add_parameter(model::ModelWithQuad, value::Real) + return add_parameter(model.inner, value) +end + +add_expression(model::ModelWithQuad, expr) = add_expression(model.inner, expr) + +Base.getindex(model::ModelWithQuad, index::ExpressionIndex) = model.inner[index] + +function register_operator( + model::ModelWithQuad, + op::Symbol, + nargs::Int, + f::Function..., +) + return register_operator(model.inner, op, nargs, f...) +end + +function MOI.is_valid(model::ModelWithQuad, index::ConstraintIndex) + return MOI.is_valid(model.inner, index) +end + +function MOI.get( + model::ModelWithQuad, + attr::MOI.ListOfSupportedNonlinearOperators, +) + return MOI.get(model.inner, attr) +end + +# The objective. + +function set_objective( + model::ModelWithQuad{T}, + obj::Union{ + MOI.VariableIndex, + MOI.ScalarAffineFunction{T}, + MOI.ScalarQuadraticFunction{T}, + }, +) where {T} + MOI.set(model.qp, MOI.ObjectiveFunction{typeof(obj)}(), obj) + set_objective(model.inner, nothing) + model.objective_sink = :quad + return +end + +function set_objective(model::ModelWithQuad{T}, obj) where {T} + F = MOI.ScalarAffineFunction{T} + MOI.set(model.qp, MOI.ObjectiveFunction{F}(), zero(F)) + if !isempty(model.qp.parameters) + obj = _replace_parameters(obj) + end + set_objective(model.inner, obj) + model.objective_sink = obj === nothing ? :none : :inner + return +end + +function MOI.get(model::ModelWithQuad, attr::MOI.ObjectiveFunctionType) + return MOI.get(model.qp, attr) +end + +function MOI.get(model::ModelWithQuad, attr::MOI.ObjectiveFunction{F}) where {F} + return MOI.get(model.qp, attr) +end + +# The affine and quadratic constraints. The MOI attribute methods are +# forwarded to the QP block, which implements them. + +const _QPFunction{T} = + Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}} + +const _QPSet{T} = + Union{MOI.LessThan{T},MOI.GreaterThan{T},MOI.EqualTo{T},MOI.Interval{T}} + +function add_constraint( + model::ModelWithQuad{T}, + func::_QPFunction{T}, + set::_QPSet{T}, +) where {T} + return MOI.add_constraint(model.qp, func, set) +end + +function add_constraint(model::ModelWithQuad, func, set) + if !isempty(model.qp.parameters) + func = _replace_parameters(func) + end + return add_constraint(model.inner, func, set) +end + +function MOI.add_constraint( + model::ModelWithQuad{T}, + func::_QPFunction{T}, + set::_QPSet{T}, +) where {T} + return MOI.add_constraint(model.qp, func, set) +end + +function MOI.get(model::ModelWithQuad, attr::MOI.ListOfConstraintTypesPresent) + return MOI.get(model.qp, attr) +end + +function MOI.is_valid( + model::ModelWithQuad{T}, + ci::MOI.ConstraintIndex{F,S}, +) where {T,F<:_QPFunction{T},S<:_QPSet{T}} + return MOI.is_valid(model.qp, ci) +end + +function MOI.get( + model::ModelWithQuad{T}, + attr::Union{MOI.ListOfConstraintIndices{F,S},MOI.NumberOfConstraints{F,S}}, +) where {T,F<:_QPFunction{T},S<:_QPSet{T}} + return MOI.get(model.qp, attr) +end + +function MOI.get( + model::ModelWithQuad{T}, + attr::Union{ + MOI.ConstraintFunction, + MOI.ConstraintSet, + MOI.ConstraintDualStart, + }, + ci::MOI.ConstraintIndex{F,S}, +) where {T,F<:_QPFunction{T},S<:_QPSet{T}} + return MOI.get(model.qp, attr, ci) +end + +function MOI.set( + model::ModelWithQuad{T}, + attr::MOI.ConstraintSet, + ci::MOI.ConstraintIndex{F,S}, + set::S, +) where {T,F<:_QPFunction{T},S<:_QPSet{T}} + return MOI.set(model.qp, attr, ci, set) +end + +function MOI.set( + model::ModelWithQuad{T}, + attr::MOI.ConstraintDualStart, + ci::MOI.ConstraintIndex{F,S}, + value, +) where {T,F<:_QPFunction{T},S<:_QPSet{T}} + return MOI.set(model.qp, attr, ci, value) +end + +""" + EvaluatorWithQuad( + model::ModelWithQuad, + inner::MOI.AbstractNLPEvaluator, + ) <: MOI.AbstractNLPEvaluator + +The evaluator of a [`ModelWithQuad`](@ref) layer. It implements the +[`MOI.AbstractNLPEvaluator`](@ref) interface: the rows of the QP block come +first, followed by the rows of `inner`, and the Jacobian and Hessian product +callbacks compose the contributions of the two blocks. + +Create it with `Evaluator(model::ModelWithQuad, backend)`, which recursively +creates the evaluator of the inner model, or construct it directly from an +existing inner evaluator. + +The QP block is evaluated as stored: [`ModelWithQuad`](@ref) owns the +variables of the model, so their indices are the columns `1:n` and no +remapping is needed. +""" +mutable struct EvaluatorWithQuad{T,M,E<:MOI.AbstractNLPEvaluator} <: + MOI.AbstractNLPEvaluator + model::ModelWithQuad{T,M} + inner::E + # The number of entries of the Jacobian and of the Hessian of the + # Lagrangian of the QP block, computed during `MOI.initialize`. + qp_nnzj::Int + qp_nnzh::Int + + function EvaluatorWithQuad( + model::ModelWithQuad{T,M}, + inner::E, + ) where {T,M,E<:MOI.AbstractNLPEvaluator} + return new{T,M,E}(model, inner, 0, 0) + end +end + +function Evaluator( + model::ModelWithQuad, + backend::AbstractAutomaticDifferentiation, +) + vars = MOI.get(model.variables, MOI.ListOfVariableIndices()) + inner = Evaluator(model.inner, backend, vars) + return EvaluatorWithQuad(model, inner) +end + +function MOI.features_available(d::EvaluatorWithQuad) + features = MOI.features_available(d.inner) + return filter(f -> f in (:Grad, :Jac, :JacVec, :Hess, :HessVec), features) +end + +function MOI.initialize(d::EvaluatorWithQuad, features::Vector{Symbol}) + d.qp_nnzj = length(MOI.jacobian_structure(d.model.qp)) + d.qp_nnzh = length(MOI.hessian_lagrangian_structure(d.model.qp)) + MOI.initialize(d.inner, features) + return +end + +function MOI.eval_objective(d::EvaluatorWithQuad{T}, x) where {T} + sink = d.model.objective_sink + if sink == :quad + return MOI.eval_objective(d.model.qp, x) + elseif sink == :inner + return MOI.eval_objective(d.inner, x) + else + return zero(T) + end +end + +function MOI.eval_objective_gradient(d::EvaluatorWithQuad{T}, grad, x) where {T} + sink = d.model.objective_sink + if sink == :quad + MOI.eval_objective_gradient(d.model.qp, grad, x) + elseif sink == :inner + MOI.eval_objective_gradient(d.inner, grad, x) + else + grad .= zero(T) + end + return +end + +function MOI.eval_constraint(d::EvaluatorWithQuad, g, x) + m = length(d.model.qp) + MOI.eval_constraint(d.model.qp, view(g, 1:m), x) + MOI.eval_constraint(d.inner, view(g, (m+1):length(g)), x) + return +end + +function MOI.jacobian_structure(d::EvaluatorWithQuad) + J = MOI.jacobian_structure(d.model.qp) + offset = length(d.model.qp) + # An evaluator is only required to implement `jacobian_structure` if it + # supports `:Jac`. If the inner evaluator does not (it then must not have + # any rows for the stack to be usable), append nothing. + if :Jac in MOI.features_available(d.inner) + for (row, col) in MOI.jacobian_structure(d.inner) + push!(J, (row + offset, col)) + end + end + return J +end + +function MOI.eval_constraint_jacobian(d::EvaluatorWithQuad, J, x) + MOI.eval_constraint_jacobian(d.model.qp, J, x) + MOI.eval_constraint_jacobian(d.inner, view(J, (d.qp_nnzj+1):length(J)), x) + return +end + +function MOI.hessian_lagrangian_structure(d::EvaluatorWithQuad) + H = MOI.hessian_lagrangian_structure(d.model.qp) + if :Hess in MOI.features_available(d.inner) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + end + return H +end + +function MOI.eval_hessian_lagrangian(d::EvaluatorWithQuad, H, x, σ, μ) + m = length(d.model.qp) + # If the objective is not in the QP block, `d.model.qp.objective` is zero, so + # passing `σ` is harmless; and vice versa for the inner evaluator. + MOI.eval_hessian_lagrangian(d.model.qp, H, x, σ, view(μ, 1:m)) + MOI.eval_hessian_lagrangian( + d.inner, + view(H, (d.qp_nnzh+1):length(H)), + x, + σ, + view(μ, (m+1):length(μ)), + ) + return +end + +# The rows of the two blocks are disjoint, so zero everything and let each +# block write its own rows. +function MOI.eval_constraint_jacobian_product(d::EvaluatorWithQuad, y, x, w) + fill!(y, zero(eltype(y))) + m = length(d.model.qp) + MOI.eval_constraint_jacobian_product( + d.inner, + view(y, (m+1):length(y)), + x, + w, + ) + _add_constraint_jacobian_product(d.model.qp, y, x, w) + return +end + +# Both blocks accumulate into the same variable-dimensional output. Call the +# inner evaluator FIRST because implementations are allowed to overwrite the +# output, and accumulate the QP block afterwards. +function MOI.eval_constraint_jacobian_transpose_product( + d::EvaluatorWithQuad, + y, + x, + w, +) + fill!(y, zero(eltype(y))) + m = length(d.model.qp) + MOI.eval_constraint_jacobian_transpose_product( + d.inner, + y, + x, + view(w, (m+1):length(w)), + ) + _add_constraint_jacobian_transpose_product(d.model.qp, y, x, view(w, 1:m)) + return +end + +function MOI.eval_hessian_lagrangian_product( + d::EvaluatorWithQuad, + H, + x, + v, + σ, + μ, +) + fill!(H, zero(eltype(H))) + m = length(d.model.qp) + MOI.eval_hessian_lagrangian_product( + d.inner, + H, + x, + v, + σ, + view(μ, (m+1):length(μ)), + ) + _add_hessian_lagrangian_product(d.model.qp, H, x, v, σ, view(μ, 1:m)) + return +end + +# The lower and upper bounds of each constraint row, in the row order of the +# evaluator. Solvers that use their own inner evaluator type can add a method +# for it so that `MOI.NLPBlockData(::EvaluatorWithQuad)` works. +function _constraint_bounds(evaluator::Evaluator) + return MOI.NLPBoundsPair[ + _bound(c.set) for (_, c) in evaluator.model.constraints + ] +end + +function _constraint_bounds(d::EvaluatorWithQuad) + bounds = MOI.NLPBoundsPair[ + MOI.NLPBoundsPair(l, u) for + (l, u) in zip(d.model.qp.g_L, d.model.qp.g_U) + ] + return append!(bounds, _constraint_bounds(d.inner)) +end + +_has_objective(d::Evaluator) = d.model.objective !== nothing + +function _has_objective(d::EvaluatorWithQuad) + if d.model.objective_sink == :quad + return true + end + return _has_objective(d.inner) +end + +function MOI.NLPBlockData(d::EvaluatorWithQuad) + return MOI.NLPBlockData(_constraint_bounds(d), d, _has_objective(d)) +end diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index 85e6ad8ce4..a7c5c7ebe7 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -4,12 +4,6 @@ # in the LICENSE.md file or at https://opensource.org/licenses/MIT. # This file is adapted from `Ipopt.jl/ext/IpoptMathOptInterfaceExt/utils.jl`. -# -# Unlike the Ipopt version, a variable is treated as a parameter if and only -# if its index is a key of the `parameters` dictionary, instead of an -# index-offset convention. Parameters must therefore be registered in -# `parameters` before any structure query, but their values may be updated -# freely between function evaluations. @enum( _FunctionType, @@ -72,11 +66,12 @@ the solver through the same callbacks as an [`MOI.AbstractNLPEvaluator`](@ref) ## Parameters -A variable is treated as a parameter if and only if its index is a key of the -`parameters` dictionary, which maps the raw `MOI.VariableIndex` value of the -parameter to its current value. Register every parameter in `parameters` -before querying any structure; the values may be updated freely between -function evaluations. +A variable is treated as a parameter if and only if its index is offset by +[`_PARAMETER_OFFSET`](@ref); see [`_is_parameter`](@ref). The value of the +parameter `x` is `parameters[x.value - _PARAMETER_OFFSET]`, following the +indexing of [`ParameterIndex`](@ref), so that `parameters` can alias the +parameter storage of a [`Model`](@ref). The values may be updated freely +between function evaluations. """ mutable struct QPBlockData{T} objective::Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}} @@ -89,7 +84,7 @@ mutable struct QPBlockData{T} mult_g::Vector{Union{Nothing,T}} function_type::Vector{_FunctionType} bound_type::Vector{_BoundType} - parameters::Dict{Int64,T} + parameters::Vector{T} function QPBlockData{T}() where {T} return new( @@ -101,21 +96,46 @@ mutable struct QPBlockData{T} Union{Nothing,T}[], _FunctionType[], _BoundType[], - Dict{Int64,T}(), + T[], ) end end -_is_parameter(v::MOI.VariableIndex, p::Dict) = haskey(p, v.value) +""" + _PARAMETER_OFFSET + +The offset of the `MOI.VariableIndex` value of a parameter: the variable +`x` is a parameter if and only if `x.value >= _PARAMETER_OFFSET`, and +`x.value - _PARAMETER_OFFSET` is the value of the corresponding +[`ParameterIndex`](@ref). +""" +const _PARAMETER_OFFSET = 0x00f0000000000000 + +""" + _is_parameter(x::MOI.VariableIndex) -function _value(v::MOI.VariableIndex, x, p::Dict) - return _is_parameter(v, p) ? p[v.value] : x[v.value] +Return whether `x` is a parameter, following the [`_PARAMETER_OFFSET`](@ref) +convention. +""" +_is_parameter(x::MOI.VariableIndex) = x.value >= _PARAMETER_OFFSET + +_is_parameter(term::MOI.ScalarAffineTerm) = _is_parameter(term.variable) + +function _is_parameter(term::MOI.ScalarQuadraticTerm) + return _is_parameter(term.variable_1) || _is_parameter(term.variable_2) +end + +function _value(v::MOI.VariableIndex, x, p::Vector) + if _is_parameter(v) + return p[v.value-_PARAMETER_OFFSET] + end + return x[v.value] end function _eval_function( f::MOI.ScalarQuadraticFunction{T}, x::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, )::T where {T} y = f.constant for term in f.affine_terms @@ -136,7 +156,7 @@ end function _eval_function( f::MOI.ScalarAffineFunction{T}, x::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, )::T where {T} y = f.constant for term in f.terms @@ -149,20 +169,19 @@ function _eval_dense_gradient( ∇f::AbstractVector{T}, f::MOI.ScalarQuadraticFunction{T}, x::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, )::Nothing where {T} for term in f.affine_terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) ∇f[term.variable.value] += term.coefficient end end for term in f.quadratic_terms - if !_is_parameter(term.variable_1, p) + if !_is_parameter(term.variable_1) v = _value(term.variable_2, x, p) ∇f[term.variable_1.value] += term.coefficient * v end - if term.variable_1 != term.variable_2 && - !_is_parameter(term.variable_2, p) + if term.variable_1 != term.variable_2 && !_is_parameter(term.variable_2) v = _value(term.variable_1, x, p) ∇f[term.variable_2.value] += term.coefficient * v end @@ -174,10 +193,10 @@ function _eval_dense_gradient( ∇f::AbstractVector{T}, f::MOI.ScalarAffineFunction{T}, x::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, )::Nothing where {T} for term in f.terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) ∇f[term.variable.value] += term.coefficient end end @@ -188,19 +207,18 @@ function _append_sparse_gradient_structure!( f::MOI.ScalarQuadraticFunction, J, row, - p::Dict, + p::Vector, ) for term in f.affine_terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) push!(J, (row, term.variable.value)) end end for term in f.quadratic_terms - if !_is_parameter(term.variable_1, p) + if !_is_parameter(term.variable_1) push!(J, (row, term.variable_1.value)) end - if term.variable_1 != term.variable_2 && - !_is_parameter(term.variable_2, p) + if term.variable_1 != term.variable_2 && !_is_parameter(term.variable_2) push!(J, (row, term.variable_2.value)) end end @@ -211,10 +229,10 @@ function _append_sparse_gradient_structure!( f::MOI.ScalarAffineFunction, J, row, - p::Dict, + p::Vector, ) for term in f.terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) push!(J, (row, term.variable.value)) end end @@ -225,23 +243,22 @@ function _eval_sparse_gradient( ∇f::AbstractVector{T}, f::MOI.ScalarQuadraticFunction{T}, x::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, )::Int where {T} i = 0 for term in f.affine_terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) i += 1 ∇f[i] = term.coefficient end end for term in f.quadratic_terms - if !_is_parameter(term.variable_1, p) + if !_is_parameter(term.variable_1) v = _value(term.variable_2, x, p) i += 1 ∇f[i] = term.coefficient * v end - if term.variable_1 != term.variable_2 && - !_is_parameter(term.variable_2, p) + if term.variable_1 != term.variable_2 && !_is_parameter(term.variable_2) v = _value(term.variable_1, x, p) i += 1 ∇f[i] = term.coefficient * v @@ -254,11 +271,11 @@ function _eval_sparse_gradient( ∇f::AbstractVector{T}, f::MOI.ScalarAffineFunction{T}, x::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, )::Int where {T} i = 0 for term in f.terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) i += 1 ∇f[i] = term.coefficient end @@ -269,11 +286,10 @@ end function _append_sparse_hessian_structure!( f::MOI.ScalarQuadraticFunction, H, - p::Dict, + p::Vector, ) for term in f.quadratic_terms - if _is_parameter(term.variable_1, p) || - _is_parameter(term.variable_2, p) + if _is_parameter(term.variable_1) || _is_parameter(term.variable_2) continue end push!(H, (term.variable_1.value, term.variable_2.value)) @@ -284,7 +300,7 @@ end function _append_sparse_hessian_structure!( ::MOI.ScalarAffineFunction, H, - ::Dict, + ::Vector, ) return nothing end @@ -293,12 +309,11 @@ function _eval_sparse_hessian( ∇²f::AbstractVector{T}, f::MOI.ScalarQuadraticFunction{T}, σ::T, - p::Dict{Int64,T}, + p::Vector{T}, )::Int where {T} i = 0 for term in f.quadratic_terms - if _is_parameter(term.variable_1, p) || - _is_parameter(term.variable_2, p) + if _is_parameter(term.variable_1) || _is_parameter(term.variable_2) continue end i += 1 @@ -311,7 +326,7 @@ function _eval_sparse_hessian( ∇²f::AbstractVector{T}, f::MOI.ScalarAffineFunction{T}, σ::T, - p::Dict{Int64,T}, + p::Vector{T}, )::Int where {T} return 0 end @@ -588,11 +603,11 @@ function _add_Jv_product( y::AbstractVector{T}, x::AbstractVector{T}, w::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, i::Int, )::Nothing where {T} for term in f.terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) y[i] += term.coefficient * w[term.variable.value] end end @@ -604,21 +619,20 @@ function _add_Jv_product( y::AbstractVector{T}, x::AbstractVector{T}, w::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, i::Int, )::Nothing where {T} for term in f.affine_terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) y[i] += term.coefficient * w[term.variable.value] end end for term in f.quadratic_terms - if !_is_parameter(term.variable_1, p) + if !_is_parameter(term.variable_1) v = _value(term.variable_2, x, p) y[i] += term.coefficient * v * w[term.variable_1.value] end - if term.variable_1 != term.variable_2 && - !_is_parameter(term.variable_2, p) + if term.variable_1 != term.variable_2 && !_is_parameter(term.variable_2) v = _value(term.variable_1, x, p) y[i] += term.coefficient * v * w[term.variable_2.value] end @@ -631,11 +645,11 @@ function _add_Jtv_product( y::AbstractVector{T}, x::AbstractVector{T}, w::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, i::Int, )::Nothing where {T} for term in f.terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) y[term.variable.value] += term.coefficient * w[i] end end @@ -647,21 +661,20 @@ function _add_Jtv_product( y::AbstractVector{T}, x::AbstractVector{T}, w::AbstractVector{T}, - p::Dict{Int64,T}, + p::Vector{T}, i::Int, )::Nothing where {T} for term in f.affine_terms - if !_is_parameter(term.variable, p) + if !_is_parameter(term.variable) y[term.variable.value] += term.coefficient * w[i] end end for term in f.quadratic_terms - if !_is_parameter(term.variable_1, p) + if !_is_parameter(term.variable_1) v = _value(term.variable_2, x, p) y[term.variable_1.value] += term.coefficient * v * w[i] end - if term.variable_1 != term.variable_2 && - !_is_parameter(term.variable_2, p) + if term.variable_1 != term.variable_2 && !_is_parameter(term.variable_2) v = _value(term.variable_1, x, p) y[term.variable_2.value] += term.coefficient * v * w[i] end @@ -675,11 +688,10 @@ function _add_Hv_product( x::AbstractVector{T}, v::AbstractVector{T}, λ::T, - p::Dict{Int64,T}, + p::Vector{T}, )::Nothing where {T} for term in f.quadratic_terms - if _is_parameter(term.variable_1, p) || - _is_parameter(term.variable_2, p) + if _is_parameter(term.variable_1) || _is_parameter(term.variable_2) continue end i, j = term.variable_1.value, term.variable_2.value @@ -697,7 +709,7 @@ function _add_Hv_product( x::AbstractVector{T}, v::AbstractVector{T}, λ::T, - p::Dict{Int64,T}, + p::Vector{T}, ) where {T} return nothing end @@ -705,7 +717,7 @@ end # These are used to add the QP contribution on top of the NL contribution. """ - add_constraint_jacobian_product( + _add_constraint_jacobian_product( block::QPBlockData{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -720,7 +732,7 @@ accumulates into `y` instead of storing the result, so that the contributions of several blocks can be composed: the caller is responsible for zeroing `y` before the first contribution. """ -function add_constraint_jacobian_product( +function _add_constraint_jacobian_product( block::QPBlockData{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -733,7 +745,7 @@ function add_constraint_jacobian_product( end """ - add_constraint_jacobian_transpose_product( + _add_constraint_jacobian_transpose_product( block::QPBlockData{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -748,7 +760,7 @@ function accumulates into `y` instead of storing the result, so that the contributions of several blocks can be composed: the caller is responsible for zeroing `y` before the first contribution. """ -function add_constraint_jacobian_transpose_product( +function _add_constraint_jacobian_transpose_product( block::QPBlockData{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -761,7 +773,7 @@ function add_constraint_jacobian_transpose_product( end """ - add_hessian_lagrangian_product( + _add_hessian_lagrangian_product( block::QPBlockData{T}, H::AbstractVector{T}, x::AbstractVector{T}, @@ -778,7 +790,7 @@ accumulates into `H` instead of storing the result, so that the contributions of several blocks can be composed: the caller is responsible for zeroing `H` before the first contribution. """ -function add_hessian_lagrangian_product( +function _add_hessian_lagrangian_product( block::QPBlockData{T}, H::AbstractVector{T}, x::AbstractVector{T}, diff --git a/test/Nonlinear/test_model_with_quad.jl b/test/Nonlinear/test_model_with_quad.jl new file mode 100644 index 0000000000..d061a1b1c8 --- /dev/null +++ b/test/Nonlinear/test_model_with_quad.jl @@ -0,0 +1,330 @@ +# Copyright (c) 2017: Miles Lubin and contributors +# Copyright (c) 2017: Google Inc. +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. + +module TestNonlinearModelWithQuad + +using Test +import MathOptInterface as MOI + +import MathOptInterface.Nonlinear + +function runtests() + for name in names(@__MODULE__; all = true) + if startswith("$(name)", "test_") + @testset "$(name)" begin + getfield(@__MODULE__, name)() + end + end + end + return +end + +# A model with, in row order: +# row 1 (quad layer, linear): 2x + 3y <= 4 +# row 2 (quad layer, quadratic): x^2 + xy + y in [0, 1] +# row 3 (inner nlp): sin(x) <= 0.5 +# and the objective x^2 in the quad layer. +function _test_model() + model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + x = MOI.add_variable(model) + y = MOI.add_variable(model) + @test (x, y) == (MOI.VariableIndex(1), MOI.VariableIndex(2)) + @test MOI.is_valid(model, x) && !MOI.is_valid(model, MOI.VariableIndex(3)) + Nonlinear.set_objective( + model, + MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ), + ) + c1 = MOI.add_constraint( + model, + MOI.ScalarAffineFunction( + [MOI.ScalarAffineTerm(2.0, x), MOI.ScalarAffineTerm(3.0, y)], + 0.0, + ), + MOI.LessThan(4.0), + ) + @test c1 isa MOI.ConstraintIndex{ + MOI.ScalarAffineFunction{Float64}, + MOI.LessThan{Float64}, + } + c2 = Nonlinear.add_constraint( + model, + MOI.ScalarQuadraticFunction( + [ + MOI.ScalarQuadraticTerm(2.0, x, x), + MOI.ScalarQuadraticTerm(1.0, x, y), + ], + [MOI.ScalarAffineTerm(1.0, y)], + 0.0, + ), + MOI.Interval(0.0, 1.0), + ) + c3 = Nonlinear.add_constraint(model, :(sin($x)), MOI.LessThan(0.5)) + @test c3 isa Nonlinear.ConstraintIndex + @test length(model) == 2 + return model, x, y +end + +function test_evaluator_with_quad() + model, x, y = _test_model() + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + @test d isa Nonlinear.EvaluatorWithQuad + @test d.inner isa Nonlinear.Evaluator + @test MOI.features_available(d) == [:Grad, :Jac, :JacVec, :Hess, :HessVec] + MOI.initialize(d, [:Grad, :Jac, :Hess]) + xv = [1.0, 2.0] # x = 1, y = 2 + @test MOI.eval_objective(d, xv) == 1.0 + grad = fill(NaN, 2) + MOI.eval_objective_gradient(d, grad, xv) + @test grad == [2.0, 0.0] + g = fill(NaN, 3) + MOI.eval_constraint(d, g, xv) + @test g ≈ [8.0, 5.0, sin(1.0)] + # Jacobian: accumulate the sparse entries into a dense matrix. + J_structure = MOI.jacobian_structure(d) + J_values = fill(NaN, length(J_structure)) + MOI.eval_constraint_jacobian(d, J_values, xv) + J = zeros(3, 2) + for ((row, col), value) in zip(J_structure, J_values) + J[row, col] += value + end + @test J ≈ [ + 2.0 3.0 + 4.0 2.0 + cos(1.0) 0.0 + ] + # Hessian of the Lagrangian: accumulate into a dense matrix. + H_structure = MOI.hessian_lagrangian_structure(d) + σ, μ = 2.0, [10.0, 100.0, 1_000.0] + H_values = fill(NaN, length(H_structure)) + MOI.eval_hessian_lagrangian(d, H_values, xv, σ, μ) + H = zeros(2, 2) + for ((row, col), value) in zip(H_structure, H_values) + H[row, col] += value + end + # σ * ∇²(x^2) + μ₂ * ∇²(x^2 + xy) + μ₃ * ∇²(sin(x)) + @test H[1, 1] ≈ 2σ + 2 * μ[2] - sin(1.0) * μ[3] + @test H[1, 2] + H[2, 1] ≈ μ[2] + @test H[2, 2] ≈ 0.0 + block = MOI.NLPBlockData(d) + @test block.has_objective + @test block.constraint_bounds == [ + MOI.NLPBoundsPair(-Inf, 4.0), + MOI.NLPBoundsPair(0.0, 1.0), + MOI.NLPBoundsPair(-Inf, 0.5), + ] + return +end + +function test_evaluator_products() + model, x, y = _test_model() + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac, :JacVec, :Hess, :HessVec]) + xv = [1.0, 2.0] + # Dense Jacobian from the sparse callback, as the reference. + J_structure = MOI.jacobian_structure(d) + J_values = fill(NaN, length(J_structure)) + MOI.eval_constraint_jacobian(d, J_values, xv) + J = zeros(3, 2) + for ((row, col), value) in zip(J_structure, J_values) + J[row, col] += value + end + w = [1.0, -2.0] + Jv = fill(NaN, 3) + MOI.eval_constraint_jacobian_product(d, Jv, xv, w) + @test Jv ≈ J * w + u = [1.0, -1.0, 2.0] + Jtv = fill(NaN, 2) + MOI.eval_constraint_jacobian_transpose_product(d, Jtv, xv, u) + @test Jtv ≈ J' * u + # Dense Hessian of the Lagrangian, as the reference. + H_structure = MOI.hessian_lagrangian_structure(d) + σ, μ = 2.0, [10.0, 100.0, 1_000.0] + H_values = fill(NaN, length(H_structure)) + MOI.eval_hessian_lagrangian(d, H_values, xv, σ, μ) + H = zeros(2, 2) + for ((row, col), value) in zip(H_structure, H_values) + H[row, col] += value + if row != col + H[col, row] += value + end + end + v = [1.0, -3.0] + Hv = fill(NaN, 2) + MOI.eval_hessian_lagrangian_product(d, Hv, xv, v, σ, μ) + @test Hv ≈ H * v + return +end + +function test_objective_sink_switching() + model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + x = MOI.add_variable(model) + @test model.objective_sink == :none + f = MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ) + Nonlinear.set_objective(model, f) + @test model.objective_sink == :quad + @test MOI.get(model, MOI.ObjectiveFunctionType()) == + MOI.ScalarQuadraticFunction{Float64} + @test MOI.get(model, MOI.ObjectiveFunction{typeof(f)}()) ≈ f + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test MOI.eval_objective(d, [3.0]) == 9.0 + @test MOI.NLPBlockData(d).has_objective + # Switch to a nonlinear objective: the quadratic objective must be + # cleared, including its Hessian entries. + Nonlinear.set_objective(model, :(sin($x))) + @test model.objective_sink == :inner + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test MOI.eval_objective(d, [3.0]) == sin(3.0) + @test MOI.NLPBlockData(d).has_objective + H_structure = MOI.hessian_lagrangian_structure(d) + H = fill(NaN, length(H_structure)) + MOI.eval_hessian_lagrangian(d, H, [3.0], 1.0, Float64[]) + @test sum(H) ≈ -sin(3.0) + # Switch to a linear objective, and then remove it. + g = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(2.0, x)], 1.0) + Nonlinear.set_objective(model, g) + @test model.objective_sink == :quad + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + @test MOI.eval_objective(d, [3.0]) == 7.0 + Nonlinear.set_objective(model, nothing) + @test model.objective_sink == :none + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + @test MOI.eval_objective(d, [3.0]) == 0.0 + grad = fill(NaN, 1) + MOI.eval_objective_gradient(d, grad, [3.0]) + @test grad == [0.0] + @test !MOI.NLPBlockData(d).has_objective + return +end + +function test_quad_parameters() + model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + x = MOI.add_variable(model) + p, cp = MOI.add_constrained_variable(model, MOI.Parameter(5.0)) + @test p.value == Nonlinear._PARAMETER_OFFSET + 1 + @test MOI.is_valid(model, p) && MOI.is_valid(model, cp) + @test MOI.get(model, MOI.ConstraintFunction(), cp) == p + @test MOI.get(model, MOI.ConstraintSet(), cp) == MOI.Parameter(5.0) + F, S = MOI.VariableIndex, MOI.Parameter{Float64} + @test MOI.get(model, MOI.NumberOfConstraints{F,S}()) == 1 + @test MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) == [cp] + # The value is stored in the inner model, aliased by the QP block. + @test model.qp.parameters === model.inner.parameters + Nonlinear.add_constraint( + model, + MOI.ScalarAffineFunction( + [MOI.ScalarAffineTerm(2.0, x), MOI.ScalarAffineTerm(3.0, p)], + 0.0, + ), + MOI.LessThan(10.0), + ) + Nonlinear.add_constraint( + model, + MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(1.0, p, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ), + MOI.LessThan(10.0), + ) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + g = fill(NaN, 2) + MOI.eval_constraint(d, g, [1.0]) + @test g == [2.0 * 1.0 + 3.0 * 5.0, 5.0 * 1.0] + # Parameters never appear in the Jacobian or Hessian structure. + @test MOI.jacobian_structure(d) == [(1, 1), (2, 1)] + J = fill(NaN, 2) + MOI.eval_constraint_jacobian(d, J, [1.0]) + @test J == [2.0, 5.0] + @test isempty(MOI.hessian_lagrangian_structure(d)) + # Updating the parameter value must be visible without re-initializing. + MOI.set(model, MOI.ConstraintSet(), cp, MOI.Parameter(7.0)) + MOI.eval_constraint(d, g, [1.0]) + @test g == [2.0 * 1.0 + 3.0 * 7.0, 7.0 * 1.0] + # A nonlinear constraint with the parameter in an embedded affine + # subfunction: the layer substitutes the parameter before the inner model + # parses the function. + aff = MOI.ScalarAffineFunction( + [MOI.ScalarAffineTerm(3.0, p), MOI.ScalarAffineTerm(1.0, x)], + 0.0, + ) + snf = MOI.ScalarNonlinearFunction(:sqrt, Any[aff]) + Nonlinear.add_constraint(model, snf, MOI.LessThan(10.0)) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + g = fill(NaN, 3) + MOI.eval_constraint(d, g, [1.0]) + @test g ≈ [2.0 + 3.0 * 7.0, 7.0, sqrt(3.0 * 7.0 + 1.0)] + return +end + +function test_attribute_forwarding() + model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + x = MOI.add_variable(model) + F, S = MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64} + f = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0) + ci = MOI.add_constraint(model, f, MOI.GreaterThan(1.0)) + @test MOI.is_valid(model, ci) + @test !MOI.is_valid(model, typeof(ci)(ci.value + 1)) + @test MOI.get(model, MOI.NumberOfConstraints{F,S}()) == 1 + @test MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) == [ci] + @test (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent()) + @test MOI.get(model, MOI.ConstraintFunction(), ci) ≈ f + @test MOI.get(model, MOI.ConstraintSet(), ci) == MOI.GreaterThan(1.0) + MOI.set(model, MOI.ConstraintSet(), ci, MOI.GreaterThan(2.0)) + @test MOI.get(model, MOI.ConstraintSet(), ci) == MOI.GreaterThan(2.0) + @test MOI.get(model, MOI.ConstraintDualStart(), ci) === nothing + MOI.set(model, MOI.ConstraintDualStart(), ci, 1.5) + @test MOI.get(model, MOI.ConstraintDualStart(), ci) == 1.5 + # Nonlinear-model forwarding + p = Nonlinear.add_parameter(model, 2.0) + @test p isa Nonlinear.ParameterIndex + ex = Nonlinear.add_expression(model, :($p * $x)) + @test model[ex] isa Nonlinear.Expression + Nonlinear.register_operator(model, :my_square, 1, z -> z^2) + c = Nonlinear.add_constraint(model, :(my_square($ex)), MOI.LessThan(1.0)) + @test MOI.is_valid(model, c) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + g = fill(NaN, 2) + MOI.eval_constraint(d, g, [3.0]) + @test g == [3.0, 36.0] + return +end + +function test_quad_only_with_empty_inner() + model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + x = MOI.add_variable(model) + Nonlinear.add_constraint( + model, + MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0), + MOI.GreaterThan(1.0), + ) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + g = fill(NaN, 1) + MOI.eval_constraint(d, g, [1.5]) + @test g == [1.5] + @test isempty(MOI.hessian_lagrangian_structure(d)) + @test MOI.NLPBlockData(d).constraint_bounds == [MOI.NLPBoundsPair(1.0, Inf)] + return +end + +end # module + +TestNonlinearModelWithQuad.runtests()