From cc37063e16865439b45cd842049d36f3087c732c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 13 Aug 2026 10:05:00 +0200 Subject: [PATCH 1/6] Add QPBlockData --- docs/src/submodules/Nonlinear/reference.md | 2 + src/Nonlinear/Nonlinear.jl | 2 + src/Nonlinear/qp_block_data.jl | 739 +++++++++++++++++++++ 3 files changed, 743 insertions(+) create mode 100644 src/Nonlinear/qp_block_data.jl diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index 842db49713..ab48feb1dd 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -72,6 +72,8 @@ Nonlinear.AbstractAutomaticDifferentiation Nonlinear.ExprGraphOnly Nonlinear.SparseReverseMode Nonlinear.SymbolicMode +Nonlinear.QPBlockData + ``` ## Data-structure diff --git a/src/Nonlinear/Nonlinear.jl b/src/Nonlinear/Nonlinear.jl index 9a5e4bd8d9..e057052029 100644 --- a/src/Nonlinear/Nonlinear.jl +++ b/src/Nonlinear/Nonlinear.jl @@ -42,4 +42,6 @@ include("evaluator.jl") include("ReverseAD/ReverseAD.jl") include("SymbolicAD/SymbolicAD.jl") +include("qp_block_data.jl") + end # module diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl new file mode 100644 index 0000000000..1b16647d47 --- /dev/null +++ b/src/Nonlinear/qp_block_data.jl @@ -0,0 +1,739 @@ +# 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`. +# +# 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, + _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 the storage behind [`ModelWithQuad`](@ref); it is not typically used +directly. + +## 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. +""" +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::Dict{Int64,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[], + Dict{Int64,T}(), + ) + end +end + +_is_parameter(v::MOI.VariableIndex, p::Dict) = haskey(p, v.value) + +function _value(v::MOI.VariableIndex, x, p::Dict) + return _is_parameter(v, p) ? p[v.value] : x[v.value] +end + +function _eval_function( + f::MOI.ScalarQuadraticFunction{T}, + x::AbstractVector{T}, + p::Dict{Int64,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::Dict{Int64,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::Dict{Int64,T}, +)::Nothing where {T} + for term in f.affine_terms + if !_is_parameter(term.variable, p) + ∇f[term.variable.value] += term.coefficient + end + end + for term in f.quadratic_terms + if !_is_parameter(term.variable_1, p) + 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) + 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::Dict{Int64,T}, +)::Nothing where {T} + for term in f.terms + if !_is_parameter(term.variable, p) + ∇f[term.variable.value] += term.coefficient + end + end + return +end + +function _append_sparse_gradient_structure!( + f::MOI.ScalarQuadraticFunction, + J, + row, + p::Dict, +) + for term in f.affine_terms + if !_is_parameter(term.variable, p) + push!(J, (row, term.variable.value)) + end + end + for term in f.quadratic_terms + if !_is_parameter(term.variable_1, p) + push!(J, (row, term.variable_1.value)) + end + if term.variable_1 != term.variable_2 && + !_is_parameter(term.variable_2, p) + push!(J, (row, term.variable_2.value)) + end + end + return +end + +function _append_sparse_gradient_structure!( + f::MOI.ScalarAffineFunction, + J, + row, + p::Dict, +) + for term in f.terms + if !_is_parameter(term.variable, p) + 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::Dict{Int64,T}, +)::Int where {T} + i = 0 + for term in f.affine_terms + if !_is_parameter(term.variable, p) + i += 1 + ∇f[i] = term.coefficient + end + end + for term in f.quadratic_terms + if !_is_parameter(term.variable_1, p) + 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) + 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::Dict{Int64,T}, +)::Int where {T} + i = 0 + for term in f.terms + if !_is_parameter(term.variable, p) + i += 1 + ∇f[i] = term.coefficient + end + end + return i +end + +function _append_sparse_hessian_structure!( + f::MOI.ScalarQuadraticFunction, + H, + p::Dict, +) + for term in f.quadratic_terms + if _is_parameter(term.variable_1, p) || + _is_parameter(term.variable_2, p) + continue + end + push!(H, (term.variable_1.value, term.variable_2.value)) + end + return +end + +function _append_sparse_hessian_structure!( + ::MOI.ScalarAffineFunction, + H, + ::Dict, +) + return nothing +end + +function _eval_sparse_hessian( + ∇²f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + σ::T, + p::Dict{Int64,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) + 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::Dict{Int64,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]) + elseif block.bound_type[row] == _kBoundTypeLessThan + return MOI.LessThan(block.g_U[row]) + elseif block.bound_type[row] == _kBoundTypeGreaterThan + return MOI.GreaterThan(block.g_L[row]) + else + @assert block.bound_type[row] == _kBoundTypeInterval + return MOI.Interval(block.g_L[row], block.g_U[row]) + 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 + +# Returns the number of entries written to `J`. +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 i +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 + +# Returns the number of entries written to `H`. +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 i +end + +# The product evaluators below ACCUMULATE into their output vector, so that +# they compose with the products of the other layers. Zero the output before +# the first call. + +function _eval_Jv_product( + f::MOI.ScalarAffineFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Dict{Int64,T}, + i::Int, +)::Nothing where {T} + for term in f.terms + if !_is_parameter(term.variable, p) + y[i] += term.coefficient * w[term.variable.value] + end + end + return +end + +function _eval_Jv_product( + f::MOI.ScalarQuadraticFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Dict{Int64,T}, + i::Int, +)::Nothing where {T} + for term in f.affine_terms + if !_is_parameter(term.variable, p) + y[i] += term.coefficient * w[term.variable.value] + end + end + for term in f.quadratic_terms + if !_is_parameter(term.variable_1, p) + 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) + v = _value(term.variable_1, x, p) + y[i] += term.coefficient * v * w[term.variable_2.value] + end + end + return +end + +function _eval_Jtv_product( + f::MOI.ScalarAffineFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Dict{Int64,T}, + i::Int, +)::Nothing where {T} + for term in f.terms + if !_is_parameter(term.variable, p) + y[term.variable.value] += term.coefficient * w[i] + end + end + return +end + +function _eval_Jtv_product( + f::MOI.ScalarQuadraticFunction{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, + p::Dict{Int64,T}, + i::Int, +)::Nothing where {T} + for term in f.affine_terms + if !_is_parameter(term.variable, p) + y[term.variable.value] += term.coefficient * w[i] + end + end + for term in f.quadratic_terms + if !_is_parameter(term.variable_1, p) + 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) + v = _value(term.variable_1, x, p) + y[term.variable_2.value] += term.coefficient * v * w[i] + end + end + return +end + +function _eval_Hv_product( + f::MOI.ScalarQuadraticFunction{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + λ::T, + p::Dict{Int64,T}, +)::Nothing where {T} + for term in f.quadratic_terms + if _is_parameter(term.variable_1, p) || + _is_parameter(term.variable_2, p) + 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 _eval_Hv_product( + ::MOI.ScalarAffineFunction{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + λ::T, + p::Dict{Int64,T}, +) where {T} + return nothing +end + +function MOI.eval_constraint_jacobian_product( + block::QPBlockData{T}, + y::AbstractVector{T}, + x::AbstractVector{T}, + w::AbstractVector{T}, +) where {T} + for (i, constraint) in enumerate(block.constraints) + _eval_Jv_product(constraint, y, x, w, block.parameters, i) + end + return +end + +function MOI.eval_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) + _eval_Jtv_product(constraint, y, x, w, block.parameters, i) + end + return +end + +function MOI.eval_hessian_lagrangian_product( + block::QPBlockData{T}, + H::AbstractVector{T}, + x::AbstractVector{T}, + v::AbstractVector{T}, + σ::T, + μ::AbstractVector{T}, +) where {T} + _eval_Hv_product(block.objective, H, x, v, σ, block.parameters) + for (i, constraint) in enumerate(block.constraints) + _eval_Hv_product(constraint, H, x, v, μ[i], block.parameters) + end + return +end From 876fa345214864d8ec9d28889baa88052a1d1f7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 13 Aug 2026 08:16:53 +0000 Subject: [PATCH 2/6] Fix reference to a type of another PR in the QPBlockData docstring --- src/Nonlinear/qp_block_data.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index 1b16647d47..2384e53532 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -66,8 +66,9 @@ 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 the storage behind [`ModelWithQuad`](@ref); it is not typically used -directly. +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 From 1a3b7018a0fc46df07a627c61ea5732a96f46151 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 13 Aug 2026 10:57:02 +0200 Subject: [PATCH 3/6] Add type assert --- src/Nonlinear/qp_block_data.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index 2384e53532..0623049095 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -431,14 +431,14 @@ function MOI.get( ) where {T,F,S} row = c.value if block.bound_type[row] == _kBoundTypeEqualTo - return MOI.EqualTo(block.g_L[row]) + return MOI.EqualTo(block.g_L[row])::S elseif block.bound_type[row] == _kBoundTypeLessThan - return MOI.LessThan(block.g_U[row]) + return MOI.LessThan(block.g_U[row])::S elseif block.bound_type[row] == _kBoundTypeGreaterThan - return MOI.GreaterThan(block.g_L[row]) + 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]) + return MOI.Interval(block.g_L[row], block.g_U[row])::S end end From d939dd4b61d4863a26d1bb9ed3237f4b0b52e613 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 13 Aug 2026 10:37:04 +0000 Subject: [PATCH 4/6] Return nothing from the evaluator methods of QPBlockData The generic functions MOI.eval_constraint_jacobian and MOI.eval_hessian_lagrangian are documented to return nothing: the caller constructs the sparsity pattern, so it knows how many entries are written. Returning the count invited callers to rely on a return value that no other evaluator provides. --- src/Nonlinear/qp_block_data.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index 0623049095..bc528b43ca 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -538,7 +538,6 @@ function MOI.jacobian_structure(block::QPBlockData) return J end -# Returns the number of entries written to `J`. function MOI.eval_constraint_jacobian( block::QPBlockData{T}, J::AbstractVector{T}, @@ -549,7 +548,7 @@ function MOI.eval_constraint_jacobian( ∇f = view(J, (i+1):length(J)) i += _eval_sparse_gradient(∇f, constraint, x, block.parameters) end - return i + return end function MOI.hessian_lagrangian_structure(block::QPBlockData) @@ -561,7 +560,6 @@ function MOI.hessian_lagrangian_structure(block::QPBlockData) return H end -# Returns the number of entries written to `H`. function MOI.eval_hessian_lagrangian( block::QPBlockData{T}, H::AbstractVector{T}, @@ -574,7 +572,7 @@ function MOI.eval_hessian_lagrangian( ∇²f = view(H, (i+1):length(H)) i += _eval_sparse_hessian(∇²f, constraint, μ[row], block.parameters) end - return i + return end # The product evaluators below ACCUMULATE into their output vector, so that From 6725299a61c66b362b6f80e8fbda0ba16ecaaa25 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 13 Aug 2026 12:34:28 +0000 Subject: [PATCH 5/6] Rename the product functions of QPBlockData to add_... The products accumulate into their output vector, so that the contributions of several blocks (the QP block, oracle constraints, and an MOI.AbstractNLPEvaluator) can be composed into the same output. This is incompatible with the contract of MOI.eval_constraint_jacobian_product and friends, which store the result, so the functions are renamed add_constraint_jacobian_product, add_constraint_jacobian_transpose_product, and add_hessian_lagrangian_product instead of overloading the MOI generic functions: QPBlockData is not an MOI.AbstractNLPEvaluator, so it does not have to define the same interface as evaluators. The private helpers are renamed _eval_... to _add_... accordingly. --- docs/src/submodules/Nonlinear/reference.md | 3 + src/Nonlinear/qp_block_data.jl | 90 ++++++++++++++++++---- 2 files changed, 76 insertions(+), 17 deletions(-) diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index ab48feb1dd..468f44cdfe 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -73,6 +73,9 @@ Nonlinear.ExprGraphOnly Nonlinear.SparseReverseMode Nonlinear.SymbolicMode Nonlinear.QPBlockData +Nonlinear.add_constraint_jacobian_product +Nonlinear.add_constraint_jacobian_transpose_product +Nonlinear.add_hessian_lagrangian_product ``` diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index bc528b43ca..85e6ad8ce4 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -575,11 +575,15 @@ function MOI.eval_hessian_lagrangian( return end -# The product evaluators below ACCUMULATE into their output vector, so that -# they compose with the products of the other layers. Zero the output before -# the first call. - -function _eval_Jv_product( +# 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}, @@ -595,7 +599,7 @@ function _eval_Jv_product( return end -function _eval_Jv_product( +function _add_Jv_product( f::MOI.ScalarQuadraticFunction{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -622,7 +626,7 @@ function _eval_Jv_product( return end -function _eval_Jtv_product( +function _add_Jtv_product( f::MOI.ScalarAffineFunction{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -638,7 +642,7 @@ function _eval_Jtv_product( return end -function _eval_Jtv_product( +function _add_Jtv_product( f::MOI.ScalarQuadraticFunction{T}, y::AbstractVector{T}, x::AbstractVector{T}, @@ -665,7 +669,7 @@ function _eval_Jtv_product( return end -function _eval_Hv_product( +function _add_Hv_product( f::MOI.ScalarQuadraticFunction{T}, H::AbstractVector{T}, x::AbstractVector{T}, @@ -687,7 +691,7 @@ function _eval_Hv_product( return end -function _eval_Hv_product( +function _add_Hv_product( ::MOI.ScalarAffineFunction{T}, H::AbstractVector{T}, x::AbstractVector{T}, @@ -698,31 +702,83 @@ function _eval_Hv_product( return nothing end -function MOI.eval_constraint_jacobian_product( +# 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) - _eval_Jv_product(constraint, y, x, w, block.parameters, i) + _add_Jv_product(constraint, y, x, w, block.parameters, i) end return end -function MOI.eval_constraint_jacobian_transpose_product( +""" + 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) - _eval_Jtv_product(constraint, y, x, w, block.parameters, i) + _add_Jtv_product(constraint, y, x, w, block.parameters, i) end return end -function MOI.eval_hessian_lagrangian_product( +""" + 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}, @@ -730,9 +786,9 @@ function MOI.eval_hessian_lagrangian_product( σ::T, μ::AbstractVector{T}, ) where {T} - _eval_Hv_product(block.objective, H, x, v, σ, block.parameters) + _add_Hv_product(block.objective, H, x, v, σ, block.parameters) for (i, constraint) in enumerate(block.constraints) - _eval_Hv_product(constraint, H, x, v, μ[i], block.parameters) + _add_Hv_product(constraint, H, x, v, μ[i], block.parameters) end return end From ff34b9936c7c723569e4d4453944c4aed94a25e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 21 Aug 2026 07:16:03 +0200 Subject: [PATCH 6/6] Add ModelWithQuad and privatize the product functions of QPBlockData (#3053) * Add ModelWithQuad and privatize the product functions of QPBlockData Making a first release with public add_constraint_jacobian_product, add_constraint_jacobian_transpose_product, and add_hessian_lagrangian_product commits MOI to three new function names. Instead, rename them _add_... and add ModelWithQuad/EvaluatorWithQuad: the evaluator implements the standard MOI.AbstractNLPEvaluator interface (including the product callbacks, composed inner-first so that store-semantics inner implementations such as ReverseAD do not clobber the QP contribution), so consumers only rely on existing generic functions and the QPBlockData internals stay private. ModelWithQuad routes affine and quadratic objectives and constraints to a QPBlockData and everything else to an inner model, forwards the MOI attribute queries of the QP block, and tracks where the objective lives. EvaluatorWithQuad remaps the variables of the QP block to their consecutive index in ordered_variables during MOI.initialize (variables absent from ordered_variables are parameters and keep their index, with the parameters dictionary aliased so per-solve value syncs are visible), and MOI.NLPBlockData(evaluator) assembles the combined constraint bounds. * Remove the unused nothing-inner support of ModelWithQuad No solver passes `inner = nothing`: the consumers all wrap an eager inner `Nonlinear.Model`, so drop the `set_objective(::Nothing, ::Nothing)` hook and the docstring sentence advertising it. A solver that manages its own nonlinear storage can still define `set_objective` for its inner type. * Remove MOI.set for the objective of ModelWithQuad Nonlinear.set_objective is the single way to set the objective: it does everything the MOI.set method did, and it also accepts nothing to clear the objective, which the MOI attribute interface cannot express. The MOI getters stay since the function API has no counterpart for them. * Move the variables and the parameter convention into ModelWithQuad ModelWithQuad now owns a MOI.Utilities.VariablesContainer and implements MOI.add_variable, guaranteeing variable indices 1:n like MOI.Utilities.MatrixOfConstraints, so the EvaluatorWithQuad no longer remaps the QP block. Parameters are added with MOI.add_constrained_variable: their indices are offset by _PARAMETER_OFFSET (moved here from Ipopt), QPBlockData is back to the simple offset-based _is_parameter instead of the parameters dictionary, and its parameters vector aliases the parameter storage of the inner Nonlinear.Model, so solvers no longer sync parameter values before a solve. * Always alias the parameter storage of the inner model Sharing only when the inner model is a Nonlinear.Model silently left qp.parameters empty for any other inner model type. Assume instead that the inner model exposes its parameter values as parameters::Vector{T}, like Nonlinear.Model does, and always alias it; an inner model without that field fails loudly at construction. * Define _is_parameter on the affine and quadratic terms They belong next to the convention; Ipopt and MadNLP defined them on their local alias of the function, which pirates it. * Move the parameter substitution into ModelWithQuad The substitution of the offset parameter indices by ParameterIndex before the inner model parses a function is index arithmetic tied to the parameter convention of the layer, so it belongs here: the generic add_constraint and set_objective now perform it, and the solvers just forward. * Use an enum for the objective sink of ModelWithQuad * Implement NumberOfVariables and ListOfVariableIndices * Return ListOfVariableIndices in the order of creation Its docstring requires the order in which the variables were added, with the parameters interleaved, so ModelWithQuad records that order instead of appending the parameters after the variables. * Fail if the inner evaluator does not implement the structure queries An inner evaluator with constraint rows must support :Jac (and :Hess when the Hessian is queried), so do not silently skip it in jacobian_structure and hessian_lagrangian_structure. * Rely on the store contract of the product callbacks MOI.eval_constraint_jacobian_transpose_product and MOI.eval_hessian_lagrangian_product store their result, zeroing the output, so calling the inner evaluator first makes pre-filling the output redundant. The Jacobian product only zeroes the rows of the QP block, into which the block accumulates; the inner evaluator stores its own rows. * Fix unresolvable references in the documentation _PARAMETER_OFFSET and _is_parameter are private and not part of any docs block, so the docstrings of QPBlockData and ModelWithQuad cannot link to them with @ref. * Cover all lines of model_with_quad.jl and qp_block_data.jl Add tests for the gradients of affine and quadratic objectives (affine and off-diagonal terms), the nonlinear objective gradient, the VariableIndex objective function type, the ConstraintSet getters and constraint-index filtering of every bound set, the Hessian and product evaluations of a parameterized QP block, nonlinear functions mentioning the parameter and the variable directly, and a parameter-free affine subfunction that the substitution leaves as is. * Cover the quadratic parameter substitution Add nonlinear constraints with quadratic subfunctions, with and without a parameter, and query ListOfSupportedNonlinearOperators through the layer. These methods were never compiled by the tests, which local coverage reports as non-executable lines but codecov reports as uncovered. --- docs/src/submodules/Nonlinear/reference.md | 5 +- src/Nonlinear/Nonlinear.jl | 1 + src/Nonlinear/model_with_quad.jl | 559 +++++++++++++++++++++ src/Nonlinear/qp_block_data.jl | 154 +++--- test/Nonlinear/test_model_with_quad.jl | 466 +++++++++++++++++ 5 files changed, 1111 insertions(+), 74 deletions(-) create mode 100644 src/Nonlinear/model_with_quad.jl create mode 100644 test/Nonlinear/test_model_with_quad.jl 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..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 index 85e6ad8ce4..6b8a8ccfa2 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`; 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}} @@ -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..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()