diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index 842db49713..f5c3746069 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -72,6 +72,10 @@ Nonlinear.AbstractAutomaticDifferentiation Nonlinear.ExprGraphOnly Nonlinear.SparseReverseMode Nonlinear.SymbolicMode +Nonlinear.QPBlockData +Nonlinear.ModelWithQuad +Nonlinear.EvaluatorWithQuad + ``` ## Data-structure diff --git a/src/Nonlinear/Nonlinear.jl b/src/Nonlinear/Nonlinear.jl index 9a5e4bd8d9..2475fba33c 100644 --- a/src/Nonlinear/Nonlinear.jl +++ b/src/Nonlinear/Nonlinear.jl @@ -42,4 +42,7 @@ include("evaluator.jl") 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..1bd21b2774 --- /dev/null +++ b/src/Nonlinear/model_with_quad.jl @@ -0,0 +1,559 @@ +# 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. + +# Where the objective of a `ModelWithQuad` currently lives. +@enum(_ObjectiveSink, _NONE, _QUAD, _INNER) + +""" + ModelWithQuad{T,M}( + qp::QPBlockData{T}, + inner::M; + objective_sink::_ObjectiveSink = _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`, 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} + # The variables and the parameters, in the order they were added, as + # `MOI.ListOfVariableIndices` requires. + list_of_variable_indices::Vector{MOI.VariableIndex} + qp::QPBlockData{T} + inner::M + objective_sink::_ObjectiveSink + + function ModelWithQuad{T}( + qp::QPBlockData{T}, + inner::M; + objective_sink::_ObjectiveSink = _NONE, + ) where {T,M} + model = new{T,M}( + MOI.Utilities.VariablesContainer{T}(), + MOI.VariableIndex[], + 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. + +function MOI.add_variable(model::ModelWithQuad) + x = MOI.add_variable(model.variables) + push!(model.list_of_variable_indices, x) + return x +end + +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) + push!(model.list_of_variable_indices, x) + ci = MOI.ConstraintIndex{MOI.VariableIndex,MOI.Parameter{T}}(x.value) + return x, ci +end + +function MOI.get(model::ModelWithQuad, ::MOI.NumberOfVariables) + return length(model.list_of_variable_indices) +end + +function MOI.get(model::ModelWithQuad, ::MOI.ListOfVariableIndices) + return model.list_of_variable_indices +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) + for (row, col) in MOI.jacobian_structure(d.inner) + push!(J, (row + offset, col)) + 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) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + 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: the inner evaluator stores its +# rows, and the QP block accumulates into its rows, which must be zeroed +# first. +function MOI.eval_constraint_jacobian_product(d::EvaluatorWithQuad, y, x, w) + m = length(d.model.qp) + fill!(view(y, 1:m), zero(eltype(y))) + 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 contribute to the same variable-dimensional output. +# `MOI.eval_constraint_jacobian_transpose_product` is called first as it +# zeroes the output before accumulating, then the QP block accumulates. +function MOI.eval_constraint_jacobian_transpose_product( + d::EvaluatorWithQuad, + y, + x, + w, +) + 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 + +# `MOI.eval_hessian_lagrangian_product` is called first as it zeroes the +# output before accumulating, then the QP block accumulates. +function MOI.eval_hessian_lagrangian_product( + d::EvaluatorWithQuad, + H, + x, + v, + σ, + μ, +) + 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 new file mode 100644 index 0000000000..6b8a8ccfa2 --- /dev/null +++ b/src/Nonlinear/qp_block_data.jl @@ -0,0 +1,806 @@ +# Copyright (c) 2013: Iain Dunning, Miles Lubin, and contributors +# +# 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. + +# This file is adapted from `Ipopt.jl/ext/IpoptMathOptInterfaceExt/utils.jl`. + +@enum( + _FunctionType, + _kFunctionTypeVariableIndex, + _kFunctionTypeScalarAffine, + _kFunctionTypeScalarQuadratic, +) + +function _function_type_to_func(::Type{T}, k::_FunctionType) where {T} + if k == _kFunctionTypeVariableIndex + return MOI.VariableIndex + elseif k == _kFunctionTypeScalarAffine + return MOI.ScalarAffineFunction{T} + else + @assert k == _kFunctionTypeScalarQuadratic + return MOI.ScalarQuadraticFunction{T} + end +end + +_function_info(::MOI.VariableIndex) = _kFunctionTypeVariableIndex +_function_info(::MOI.ScalarAffineFunction) = _kFunctionTypeScalarAffine +_function_info(::MOI.ScalarQuadraticFunction) = _kFunctionTypeScalarQuadratic + +@enum( + _BoundType, + _kBoundTypeLessThan, + _kBoundTypeGreaterThan, + _kBoundTypeEqualTo, + _kBoundTypeInterval, +) + +_set_info(s::MOI.LessThan) = _kBoundTypeLessThan, -Inf, s.upper +_set_info(s::MOI.GreaterThan) = _kBoundTypeGreaterThan, s.lower, Inf +_set_info(s::MOI.EqualTo) = _kBoundTypeEqualTo, s.value, s.value +_set_info(s::MOI.Interval) = _kBoundTypeInterval, s.lower, s.upper + +function _bound_type_to_set(::Type{T}, k::_BoundType) where {T} + if k == _kBoundTypeEqualTo + return MOI.EqualTo{T} + elseif k == _kBoundTypeLessThan + return MOI.LessThan{T} + elseif k == _kBoundTypeGreaterThan + return MOI.GreaterThan{T} + else + @assert k == _kBoundTypeInterval + return MOI.Interval{T} + end +end + +""" + QPBlockData{T}() + +A data structure holding an affine or quadratic objective and a block of +affine and quadratic constraints, together with methods to evaluate them +following the [`MOI.AbstractNLPEvaluator`](@ref) callback conventions. + +This is a helper for solvers that pass affine and quadratic constraints to +the solver through the same callbacks as an [`MOI.AbstractNLPEvaluator`](@ref) +(for example, Ipopt and MadNLP). + +## Parameters + +A variable is treated as a parameter if and only if its index is offset by +`_PARAMETER_OFFSET`; see `_is_parameter`. 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}} + objective_function_type::_FunctionType + constraints::Vector{ + Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}, + } + g_L::Vector{T} + g_U::Vector{T} + mult_g::Vector{Union{Nothing,T}} + function_type::Vector{_FunctionType} + bound_type::Vector{_BoundType} + parameters::Vector{T} + + function QPBlockData{T}() where {T} + return new( + zero(MOI.ScalarQuadraticFunction{T}), + _kFunctionTypeScalarAffine, + Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}[], + T[], + T[], + Union{Nothing,T}[], + _FunctionType[], + _BoundType[], + T[], + ) + end +end + +""" + _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) + +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::Vector{T}, +)::T where {T} + y = f.constant + for term in f.affine_terms + y += term.coefficient * _value(term.variable, x, p) + end + for term in f.quadratic_terms + v1 = _value(term.variable_1, x, p) + v2 = _value(term.variable_2, x, p) + if term.variable_1 == term.variable_2 + y += term.coefficient * v1 * v2 / 2 + else + y += term.coefficient * v1 * v2 + end + end + return y +end + +function _eval_function( + f::MOI.ScalarAffineFunction{T}, + x::AbstractVector{T}, + p::Vector{T}, +)::T where {T} + y = f.constant + for term in f.terms + y += term.coefficient * _value(term.variable, x, p) + end + return y +end + +function _eval_dense_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + x::AbstractVector{T}, + p::Vector{T}, +)::Nothing where {T} + for term in f.affine_terms + 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) + 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) + v = _value(term.variable_1, x, p) + ∇f[term.variable_2.value] += term.coefficient * v + end + end + return +end + +function _eval_dense_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarAffineFunction{T}, + x::AbstractVector{T}, + p::Vector{T}, +)::Nothing where {T} + for term in f.terms + if !_is_parameter(term.variable) + ∇f[term.variable.value] += term.coefficient + end + end + return +end + +function _append_sparse_gradient_structure!( + f::MOI.ScalarQuadraticFunction, + J, + row, + p::Vector, +) + for term in f.affine_terms + 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) + push!(J, (row, term.variable_1.value)) + end + if term.variable_1 != term.variable_2 && !_is_parameter(term.variable_2) + push!(J, (row, term.variable_2.value)) + end + end + return +end + +function _append_sparse_gradient_structure!( + f::MOI.ScalarAffineFunction, + J, + row, + p::Vector, +) + for term in f.terms + if !_is_parameter(term.variable) + push!(J, (row, term.variable.value)) + end + end + return +end + +function _eval_sparse_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + x::AbstractVector{T}, + p::Vector{T}, +)::Int where {T} + i = 0 + for term in f.affine_terms + 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) + 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) + v = _value(term.variable_1, x, p) + i += 1 + ∇f[i] = term.coefficient * v + end + end + return i +end + +function _eval_sparse_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarAffineFunction{T}, + x::AbstractVector{T}, + p::Vector{T}, +)::Int where {T} + i = 0 + for term in f.terms + if !_is_parameter(term.variable) + i += 1 + ∇f[i] = term.coefficient + end + end + return i +end + +function _append_sparse_hessian_structure!( + f::MOI.ScalarQuadraticFunction, + H, + p::Vector, +) + for term in f.quadratic_terms + if _is_parameter(term.variable_1) || _is_parameter(term.variable_2) + continue + end + push!(H, (term.variable_1.value, term.variable_2.value)) + end + return +end + +function _append_sparse_hessian_structure!( + ::MOI.ScalarAffineFunction, + H, + ::Vector, +) + return nothing +end + +function _eval_sparse_hessian( + ∇²f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + σ::T, + p::Vector{T}, +)::Int where {T} + i = 0 + for term in f.quadratic_terms + if _is_parameter(term.variable_1) || _is_parameter(term.variable_2) + continue + end + i += 1 + ∇²f[i] = term.coefficient * σ + end + return i +end + +function _eval_sparse_hessian( + ∇²f::AbstractVector{T}, + f::MOI.ScalarAffineFunction{T}, + σ::T, + p::Vector{T}, +)::Int where {T} + return 0 +end + +Base.length(block::QPBlockData) = length(block.bound_type) + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ObjectiveFunction{F}, + f::F, +) where {T,F<:Union{MOI.VariableIndex,MOI.ScalarAffineFunction{T}}} + block.objective = convert(MOI.ScalarAffineFunction{T}, f) + block.objective_function_type = _function_info(f) + return +end + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}}, + f::MOI.ScalarQuadraticFunction{T}, +) where {T} + block.objective = f + block.objective_function_type = _function_info(f) + return +end + +function MOI.get(block::QPBlockData{T}, ::MOI.ObjectiveFunctionType) where {T} + return _function_type_to_func(T, block.objective_function_type) +end + +function MOI.get(block::QPBlockData{T}, ::MOI.ObjectiveFunction{F}) where {T,F} + return convert(F, block.objective) +end + +function MOI.get( + block::QPBlockData{T}, + ::MOI.ListOfConstraintTypesPresent, +) where {T} + constraints = Set{Tuple{Type,Type}}() + for i in 1:length(block) + F = _function_type_to_func(T, block.function_type[i]) + S = _bound_type_to_set(T, block.bound_type[i]) + push!(constraints, (F, S)) + end + return collect(constraints) +end + +function MOI.is_valid( + block::QPBlockData{T}, + ci::MOI.ConstraintIndex{F,S}, +) where { + T, + F<:Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}, + S<:Union{MOI.LessThan{T},MOI.GreaterThan{T},MOI.EqualTo{T},MOI.Interval{T}}, +} + return 1 <= ci.value <= length(block) +end + +function MOI.get( + block::QPBlockData{T}, + ::MOI.ListOfConstraintIndices{F,S}, +) where { + T, + F<:Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}, + S<:Union{MOI.LessThan{T},MOI.GreaterThan{T},MOI.EqualTo{T},MOI.Interval{T}}, +} + ret = MOI.ConstraintIndex{F,S}[] + for i in 1:length(block) + if _bound_type_to_set(T, block.bound_type[i]) != S + continue + elseif _function_type_to_func(T, block.function_type[i]) != F + continue + end + push!(ret, MOI.ConstraintIndex{F,S}(i)) + end + return ret +end + +function MOI.get( + block::QPBlockData{T}, + ::MOI.NumberOfConstraints{F,S}, +) where { + T, + F<:Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}, + S<:Union{MOI.LessThan{T},MOI.GreaterThan{T},MOI.EqualTo{T},MOI.Interval{T}}, +} + return length(MOI.get(block, MOI.ListOfConstraintIndices{F,S}())) +end + +function MOI.add_constraint( + block::QPBlockData{T}, + f::Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}, + s::Union{MOI.LessThan{T},MOI.GreaterThan{T},MOI.EqualTo{T},MOI.Interval{T}}, +) where {T} + push!(block.constraints, f) + bound_type, l, u = _set_info(s) + push!(block.g_L, l) + push!(block.g_U, u) + push!(block.mult_g, nothing) + push!(block.bound_type, bound_type) + push!(block.function_type, _function_info(f)) + return MOI.ConstraintIndex{typeof(f),typeof(s)}(length(block.bound_type)) +end + +function MOI.get( + block::QPBlockData{T}, + ::MOI.ConstraintFunction, + c::MOI.ConstraintIndex{F,S}, +) where {T,F,S} + return convert(F, block.constraints[c.value]) +end + +function MOI.get( + block::QPBlockData{T}, + ::MOI.ConstraintSet, + c::MOI.ConstraintIndex{F,S}, +) where {T,F,S} + row = c.value + if block.bound_type[row] == _kBoundTypeEqualTo + return MOI.EqualTo(block.g_L[row])::S + elseif block.bound_type[row] == _kBoundTypeLessThan + return MOI.LessThan(block.g_U[row])::S + elseif block.bound_type[row] == _kBoundTypeGreaterThan + return MOI.GreaterThan(block.g_L[row])::S + else + @assert block.bound_type[row] == _kBoundTypeInterval + return MOI.Interval(block.g_L[row], block.g_U[row])::S + end +end + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ConstraintSet, + c::MOI.ConstraintIndex{F,MOI.LessThan{T}}, + set::MOI.LessThan{T}, +) where {T,F} + block.g_U[c.value] = set.upper + return +end + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ConstraintSet, + c::MOI.ConstraintIndex{F,MOI.GreaterThan{T}}, + set::MOI.GreaterThan{T}, +) where {T,F} + block.g_L[c.value] = set.lower + return +end + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ConstraintSet, + c::MOI.ConstraintIndex{F,MOI.EqualTo{T}}, + set::MOI.EqualTo{T}, +) where {T,F} + block.g_L[c.value] = set.value + block.g_U[c.value] = set.value + return +end + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ConstraintSet, + c::MOI.ConstraintIndex{F,MOI.Interval{T}}, + set::MOI.Interval{T}, +) where {T,F} + block.g_L[c.value] = set.lower + block.g_U[c.value] = set.upper + return +end + +function MOI.get( + block::QPBlockData{T}, + ::MOI.ConstraintDualStart, + c::MOI.ConstraintIndex{F,S}, +) where {T,F,S} + return block.mult_g[c.value] +end + +function MOI.set( + block::QPBlockData{T}, + ::MOI.ConstraintDualStart, + c::MOI.ConstraintIndex{F,S}, + value, +) where {T,F,S} + block.mult_g[c.value] = value + return +end + +function MOI.eval_objective( + block::QPBlockData{T}, + x::AbstractVector{T}, +) where {T} + return _eval_function(block.objective, x, block.parameters) +end + +function MOI.eval_objective_gradient( + block::QPBlockData{T}, + ∇f::AbstractVector{T}, + x::AbstractVector{T}, +) where {T} + ∇f .= zero(T) + _eval_dense_gradient(∇f, block.objective, x, block.parameters) + return +end + +function MOI.eval_constraint( + block::QPBlockData{T}, + g::AbstractVector{T}, + x::AbstractVector{T}, +) where {T} + for (i, constraint) in enumerate(block.constraints) + g[i] = _eval_function(constraint, x, block.parameters) + end + return +end + +function MOI.jacobian_structure(block::QPBlockData) + J = Tuple{Int,Int}[] + for (row, constraint) in enumerate(block.constraints) + _append_sparse_gradient_structure!(constraint, J, row, block.parameters) + end + return J +end + +function MOI.eval_constraint_jacobian( + block::QPBlockData{T}, + J::AbstractVector{T}, + x::AbstractVector{T}, +) where {T} + i = 0 + for constraint in block.constraints + ∇f = view(J, (i+1):length(J)) + i += _eval_sparse_gradient(∇f, constraint, x, block.parameters) + end + return +end + +function MOI.hessian_lagrangian_structure(block::QPBlockData) + H = Tuple{Int,Int}[] + _append_sparse_hessian_structure!(block.objective, H, block.parameters) + for constraint in block.constraints + _append_sparse_hessian_structure!(constraint, H, block.parameters) + end + return H +end + +function MOI.eval_hessian_lagrangian( + block::QPBlockData{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + σ::T, + μ::AbstractVector{T}, +) where {T} + i = _eval_sparse_hessian(H, block.objective, σ, block.parameters) + for (row, constraint) in enumerate(block.constraints) + ∇²f = view(H, (i+1):length(H)) + i += _eval_sparse_hessian(∇²f, constraint, μ[row], block.parameters) + end + return +end + +# The product functions below ACCUMULATE into their output vector, so that +# the contributions of several blocks (for example, the QP block, the +# vector-nonlinear-oracle constraints, and an `MOI.AbstractNLPEvaluator`) can +# be composed into the same output. This is why they are not methods of the +# corresponding `MOI.eval_...` functions, whose contract is to store the +# result: `QPBlockData` is not an `MOI.AbstractNLPEvaluator`, so it does not +# have to define the same interface as evaluators. + +function _add_Jv_product( + f::MOI.ScalarAffineFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Vector{T}, + i::Int, +)::Nothing where {T} + for term in f.terms + if !_is_parameter(term.variable) + y[i] += term.coefficient * w[term.variable.value] + end + end + return +end + +function _add_Jv_product( + f::MOI.ScalarQuadraticFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Vector{T}, + i::Int, +)::Nothing where {T} + for term in f.affine_terms + 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) + 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) + v = _value(term.variable_1, x, p) + y[i] += term.coefficient * v * w[term.variable_2.value] + end + end + return +end + +function _add_Jtv_product( + f::MOI.ScalarAffineFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Vector{T}, + i::Int, +)::Nothing where {T} + for term in f.terms + if !_is_parameter(term.variable) + y[term.variable.value] += term.coefficient * w[i] + end + end + return +end + +function _add_Jtv_product( + f::MOI.ScalarQuadraticFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Vector{T}, + i::Int, +)::Nothing where {T} + for term in f.affine_terms + 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) + 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) + v = _value(term.variable_1, x, p) + y[term.variable_2.value] += term.coefficient * v * w[i] + end + end + return +end + +function _add_Hv_product( + f::MOI.ScalarQuadraticFunction{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + λ::T, + p::Vector{T}, +)::Nothing where {T} + for term in f.quadratic_terms + if _is_parameter(term.variable_1) || _is_parameter(term.variable_2) + continue + end + i, j = term.variable_1.value, term.variable_2.value + H[i] += λ * term.coefficient * v[j] + if i != j + H[j] += λ * term.coefficient * v[i] + end + end + return +end + +function _add_Hv_product( + ::MOI.ScalarAffineFunction{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + λ::T, + p::Vector{T}, +) where {T} + return nothing +end + +# These are used to add the QP contribution on top of the NL contribution. + +""" + _add_constraint_jacobian_product( + block::QPBlockData{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + )::Nothing where {T} + +Add to `y` the product of the Jacobian of the constraints of `block` at `x` +with `w`. + +Unlike [`MOI.eval_constraint_jacobian_product`](@ref), this 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_product( + block::QPBlockData{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, +) where {T} + for (i, constraint) in enumerate(block.constraints) + _add_Jv_product(constraint, y, x, w, block.parameters, i) + end + return +end + +""" + _add_constraint_jacobian_transpose_product( + block::QPBlockData{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + )::Nothing where {T} + +Add to `y` the product of the transpose of the Jacobian of the constraints of +`block` at `x` with `w`. + +Unlike [`MOI.eval_constraint_jacobian_transpose_product`](@ref), this +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( + block::QPBlockData{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, +) where {T} + for (i, constraint) in enumerate(block.constraints) + _add_Jtv_product(constraint, y, x, w, block.parameters, i) + end + return +end + +""" + _add_hessian_lagrangian_product( + block::QPBlockData{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + σ::T, + μ::AbstractVector{T}, + )::Nothing where {T} + +Add to `H` the product of the Hessian of the Lagrangian of `block` at `x`, +with objective weight `σ` and constraint weights `μ`, with `v`. + +Unlike [`MOI.eval_hessian_lagrangian_product`](@ref), this function +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( + block::QPBlockData{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + σ::T, + μ::AbstractVector{T}, +) where {T} + _add_Hv_product(block.objective, H, x, v, σ, block.parameters) + for (i, constraint) in enumerate(block.constraints) + _add_Hv_product(constraint, H, x, v, μ[i], block.parameters) + end + return +end diff --git a/test/Nonlinear/test_model_with_quad.jl b/test/Nonlinear/test_model_with_quad.jl new file mode 100644 index 0000000000..f9e4a193ef --- /dev/null +++ b/test/Nonlinear/test_model_with_quad.jl @@ -0,0 +1,466 @@ +# 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 == Nonlinear._NONE + f = MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ) + Nonlinear.set_objective(model, f) + @test model.objective_sink == Nonlinear._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 == Nonlinear._INNER + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test MOI.eval_objective(d, [3.0]) == sin(3.0) + grad = fill(NaN, 1) + MOI.eval_objective_gradient(d, grad, [3.0]) + @test grad ≈ [cos(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 == Nonlinear._QUAD + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + @test MOI.eval_objective(d, [3.0]) == 7.0 + grad = fill(NaN, 1) + MOI.eval_objective_gradient(d, grad, [3.0]) + @test grad == [2.0] + Nonlinear.set_objective(model, nothing) + @test model.objective_sink == Nonlinear._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 + # `ListOfVariableIndices` is in the order of creation, parameters + # included. + @test MOI.get(model, MOI.NumberOfVariables()) == 2 + @test MOI.get(model, MOI.ListOfVariableIndices()) == [x, p] + let model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + q, _ = MOI.add_constrained_variable(model, MOI.Parameter(1.0)) + z = MOI.add_variable(model) + @test MOI.get(model, MOI.ListOfVariableIndices()) == [q, z] + end + 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] + # The Hessian and the products skip the terms with a parameter: the + # second constraint, `p * x`, has no entry. + @test isempty(MOI.hessian_lagrangian_structure(d)) + H = Float64[] + MOI.eval_hessian_lagrangian(d, H, [1.0], 1.0, [1.0, 1.0]) + Jv = fill(NaN, 2) + MOI.eval_constraint_jacobian_product(d, Jv, [1.0], [1.5]) + @test Jv == [2.0 * 1.5, 7.0 * 1.5] + Jtv = fill(NaN, 1) + MOI.eval_constraint_jacobian_transpose_product(d, Jtv, [1.0], [1.0, 1.0]) + @test Jtv == [2.0 + 7.0] + Hv = fill(NaN, 1) + MOI.eval_hessian_lagrangian_product(d, Hv, [1.0], [1.5], 1.0, [1.0, 1.0]) + @test Hv == [0.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)) + # A nonlinear constraint with a parameter-free affine subfunction, which + # the substitution leaves as is. + Nonlinear.add_constraint( + model, + MOI.ScalarNonlinearFunction( + :sqrt, + Any[MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0)], + ), + MOI.LessThan(10.0), + ) + # Nonlinear constraints with quadratic subfunctions: with a parameter + # (converted to `ScalarNonlinearFunction`) and without (left as is). + q_p = MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, p, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ) + Nonlinear.add_constraint( + model, + MOI.ScalarNonlinearFunction(:sqrt, Any[q_p]), + MOI.LessThan(30.0), + ) + q_x = MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ) + Nonlinear.add_constraint( + model, + MOI.ScalarNonlinearFunction(:sqrt, Any[q_x]), + MOI.LessThan(30.0), + ) + # A nonlinear constraint and objective mentioning the parameter and the + # variable directly. + Nonlinear.add_constraint( + model, + MOI.ScalarNonlinearFunction(:+, Any[x, p]), + MOI.LessThan(20.0), + ) + Nonlinear.set_objective(model, MOI.ScalarNonlinearFunction(:*, Any[p, x])) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + g = fill(NaN, 7) + MOI.eval_constraint(d, g, [1.0]) + @test g ≈ [ + 2.0 + 3.0 * 7.0, + 7.0, + sqrt(3.0 * 7.0 + 1.0), + 1.0, + sqrt(2.0 * 7.0), + 1.0, + 1.0 + 7.0, + ] + @test MOI.eval_objective(d, [1.5]) == 7.0 * 1.5 + 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) + ops = MOI.get(model, MOI.ListOfSupportedNonlinearOperators()) + @test :my_square in ops + 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_qp_attribute_types() + model = Nonlinear.ModelWithQuad(Nonlinear.Model()) + x = MOI.add_variable(model) + y = MOI.add_variable(model) + Nonlinear.set_objective(model, x) + @test MOI.get(model, MOI.ObjectiveFunctionType()) == MOI.VariableIndex + @test MOI.get(model, MOI.ObjectiveFunction{MOI.VariableIndex}()) == x + F = MOI.ScalarAffineFunction{Float64} + f = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0) + q = MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, y)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ) + c1 = MOI.add_constraint(model, f, MOI.LessThan(1.0)) + c2 = MOI.add_constraint(model, f, MOI.EqualTo(2.0)) + c3 = MOI.add_constraint(model, f, MOI.Interval(3.0, 4.0)) + c4 = MOI.add_constraint(model, q, MOI.GreaterThan(5.0)) + @test MOI.get(model, MOI.ConstraintSet(), c1) == MOI.LessThan(1.0) + @test MOI.get(model, MOI.ConstraintSet(), c2) == MOI.EqualTo(2.0) + @test MOI.get(model, MOI.ConstraintSet(), c3) == MOI.Interval(3.0, 4.0) + @test MOI.get(model, MOI.ConstraintSet(), c4) == MOI.GreaterThan(5.0) + for (S, ci) in [ + (MOI.LessThan{Float64}, c1), + (MOI.EqualTo{Float64}, c2), + (MOI.Interval{Float64}, c3), + ] + @test MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) == [ci] + @test MOI.get(model, MOI.NumberOfConstraints{F,S}()) == 1 + end + Q = MOI.ScalarQuadraticFunction{Float64} + S = MOI.GreaterThan{Float64} + c5 = MOI.add_constraint(model, f, MOI.GreaterThan(6.0)) + @test MOI.get(model, MOI.ListOfConstraintIndices{Q,S}()) == [c4] + @test MOI.get(model, MOI.ListOfConstraintIndices{F,S}()) == [c5] + # The gradient of a quadratic objective with affine and off-diagonal + # terms. + g = MOI.ScalarQuadraticFunction( + [ + MOI.ScalarQuadraticTerm(2.0, x, x), + MOI.ScalarQuadraticTerm(1.0, x, y), + ], + [MOI.ScalarAffineTerm(3.0, x)], + 0.0, + ) + Nonlinear.set_objective(model, g) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode()) + MOI.initialize(d, [:Grad, :Jac]) + xv = [1.0, 2.0] + @test MOI.eval_objective(d, xv) == 1.0 + 2.0 + 3.0 + grad = fill(NaN, 2) + MOI.eval_objective_gradient(d, grad, xv) + @test grad == [2.0 * 1.0 + 2.0 + 3.0, 1.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()