diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index 842db49713..e8584a0ea2 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -74,6 +74,27 @@ Nonlinear.SparseReverseMode Nonlinear.SymbolicMode ``` +## Linearity + +```@docs +Nonlinear.Linearity +Nonlinear.num_constraints +Nonlinear.constraint_linearity +Nonlinear.objective_linearity +Nonlinear.constraint_bounds +``` + +## Model layers + +```@docs +Nonlinear.model +Nonlinear.ModelWithQuad +Nonlinear.ModelWithOracles +Nonlinear.EvaluatorWithQuad +Nonlinear.EvaluatorWithOracles +Nonlinear.QPBlockData +``` + ## Data-structure ```@docs diff --git a/src/Nonlinear/Nonlinear.jl b/src/Nonlinear/Nonlinear.jl index 9a5e4bd8d9..ee062767ca 100644 --- a/src/Nonlinear/Nonlinear.jl +++ b/src/Nonlinear/Nonlinear.jl @@ -42,4 +42,10 @@ include("evaluator.jl") include("ReverseAD/ReverseAD.jl") include("SymbolicAD/SymbolicAD.jl") +# Depends on the `Linearity` enum defined in `ReverseAD`, so must be included +# after it. +include("linearity.jl") +include("qp_block_data.jl") +include("layers.jl") + end # module diff --git a/src/Nonlinear/ReverseAD/graph_tools.jl b/src/Nonlinear/ReverseAD/graph_tools.jl index 837fbf23b4..6292ad35cf 100644 --- a/src/Nonlinear/ReverseAD/graph_tools.jl +++ b/src/Nonlinear/ReverseAD/graph_tools.jl @@ -33,22 +33,38 @@ function _replace_moi_variables( return new_nodes end -@enum(Linearity, CONSTANT, LINEAR, PIECEWISE_LINEAR, NONLINEAR) +@enum(Linearity, CONSTANT, LINEAR, PIECEWISE_LINEAR, QUADRATIC, NONLINEAR) """ _classify_linearity( nodes::Vector{Nonlinear.Node}, adj::SparseArrays.SparseMatrixCSC, subexpression_linearity::Vector{Linearity}, + const_values::Vector{Float64}, ) -Classify the nodes in a tree as constant, linear, or nonlinear with respect to -the input. +Classify the nodes in a tree as constant, linear, piecewise linear, quadratic, +or nonlinear with respect to the input. + +The classification is conservative: a node may be classified less strictly +than the tightest class that applies (for example, an expression that +simplifies to an affine function may be classified as `NONLINEAR`), but never +more strictly. For fixed values of the parameters, each class guarantees: + + * `CONSTANT`: the value does not depend on the input + * `LINEAR`: the value is an affine function of the input; the gradient is + constant and the Hessian is zero + * `PIECEWISE_LINEAR`: the gradient is piecewise constant and the Hessian is + zero almost everywhere + * `QUADRATIC`: the gradient is an affine function of the input and the + Hessian is constant + * `NONLINEAR`: no guarantee """ function _classify_linearity( nodes::Vector{Nonlinear.Node}, adj::SparseArrays.SparseMatrixCSC, subexpression_linearity::Vector{Linearity}, + const_values::Vector{Float64}, ) linearity = Array{Linearity}(undef, length(nodes)) children_arr = SparseArrays.rowvals(adj) @@ -68,46 +84,32 @@ function _classify_linearity( continue end children_idx = SparseArrays.nzrange(adj, k) - num_constant_children, any_nonlinear = 0, false + num_constant, num_piecewise, num_quadratic = 0, 0, 0 + worst = CONSTANT for r in children_idx - if linearity[children_arr[r]] == NONLINEAR - any_nonlinear = true - break - elseif linearity[children_arr[r]] == CONSTANT - num_constant_children += 1 + child = linearity[children_arr[r]] + worst = max(worst, child) + if child == CONSTANT + num_constant += 1 + elseif child == PIECEWISE_LINEAR + num_piecewise += 1 + elseif child == QUADRATIC + num_quadratic += 1 end end - if any_nonlinear - # If any children are nonlinear, then we're nonlinear... - linearity[k] = NONLINEAR - # ...except in the case of ifelse. If the operands are linear then - # we're piecewise linear. - op = get( - Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS, - node.index, - nothing, - ) - if ( - node.type == Nonlinear.NODE_CALL_MULTIVARIATE && - op == :ifelse && - linearity[children_arr[children_idx[2]]] == LINEAR && - linearity[children_arr[children_idx[3]]] == LINEAR - ) - linearity[k] = PIECEWISE_LINEAR - end - continue - elseif num_constant_children == length(children_idx) + if worst == CONSTANT # If all children are constant, then we're constant. linearity[k] = CONSTANT continue end - # By this point, some children are constant and some are linear, so if - # the operator is nonlinear, then we're nonlinear. if node.type == Nonlinear.NODE_CALL_UNIVARIATE op = get(Nonlinear.DEFAULT_UNIVARIATE_OPERATORS, node.index, nothing) if op == :+ || op == :- - linearity[k] = LINEAR + # Unary plus and minus preserve the linearity of the child. + linearity[k] = worst + elseif op == :abs && worst <= PIECEWISE_LINEAR + linearity[k] = PIECEWISE_LINEAR else linearity[k] = NONLINEAR end @@ -117,26 +119,60 @@ function _classify_linearity( node.index, nothing, ) - if op == :+ - linearity[k] = LINEAR - elseif op == :- - linearity[k] = LINEAR + if op == :+ || op == :- + if num_quadratic > 0 && num_piecewise > 0 + # The sum of a quadratic term and a piecewise linear term + # is piecewise quadratic: neither class applies. + linearity[k] = NONLINEAR + else + linearity[k] = worst + end elseif op == :* - # Multiplication is linear if there is one non-constant term. - one_op = num_constant_children == length(children_idx) - 1 - linearity[k] = one_op ? LINEAR : NONLINEAR + num_nonconstant = length(children_idx) - num_constant + if num_nonconstant == 1 + # Multiplication by constants preserves the linearity of + # the non-constant term. + linearity[k] = worst + elseif num_nonconstant == 2 && worst == LINEAR + # The product of two linear terms is quadratic. + linearity[k] = QUADRATIC + else + linearity[k] = NONLINEAR + end elseif op == :^ - linearity[k] = NONLINEAR + expo = nodes[children_arr[children_idx[2]]] + if expo.type == Nonlinear.NODE_VALUE && + const_values[expo.index] == 2.0 && + linearity[children_arr[children_idx[1]]] == LINEAR + # A linear term squared is quadratic. We do not attempt to + # classify other constant exponents; note in particular + # that a NODE_PARAMETER exponent is CONSTANT but its value + # may change between solves, so it must stay NONLINEAR. + linearity[k] = QUADRATIC + else + linearity[k] = NONLINEAR + end elseif op == :/ if linearity[children_arr[children_idx[2]]] == CONSTANT - # If the denominator is constant, we're linear. - linearity[k] = LINEAR + # If the denominator is constant, division preserves the + # linearity of the numerator. + linearity[k] = linearity[children_arr[children_idx[1]]] else linearity[k] = NONLINEAR end elseif op == :ifelse - linearity[k] = NONLINEAR - else # User-defined functions + # The condition is typically NONLINEAR (comparisons are), but + # if both branches are at worst piecewise linear, so are we. + if linearity[children_arr[children_idx[2]]] <= + PIECEWISE_LINEAR && + linearity[children_arr[children_idx[3]]] <= PIECEWISE_LINEAR + linearity[k] = PIECEWISE_LINEAR + else + linearity[k] = NONLINEAR + end + elseif (op == :min || op == :max) && worst <= PIECEWISE_LINEAR + linearity[k] = PIECEWISE_LINEAR + else # Other operators and user-defined functions linearity[k] = NONLINEAR end elseif node.type == Nonlinear.NODE_LOGIC diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index f67d9caf8e..0faf740962 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -69,7 +69,6 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) d.data.expressions[k], d.subexpression_linearity, moi_index_to_consecutive_index, - d.want_hess, ) d.subexpressions[k] = subex d.subexpression_linearity[k] = subex.linearity @@ -88,6 +87,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol}) subex.nodes, subex.adj, d.subexpression_linearity, + subex.const_values, ) edgelist = _compute_hessian_sparsity( subex.nodes, diff --git a/src/Nonlinear/ReverseAD/types.jl b/src/Nonlinear/ReverseAD/types.jl index ce3e807b94..e3c9cbd4c5 100644 --- a/src/Nonlinear/ReverseAD/types.jl +++ b/src/Nonlinear/ReverseAD/types.jl @@ -18,17 +18,17 @@ struct _SubexpressionStorage expr::Nonlinear.Expression, subexpression_linearity, moi_index_to_consecutive_index, - want_hess::Bool, ) nodes = _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index) adj = Nonlinear.adjacency_matrix(nodes) N = length(nodes) - linearity = if want_hess - _classify_linearity(nodes, adj, subexpression_linearity)[1] - else - NONLINEAR - end + linearity = _classify_linearity( + nodes, + adj, + subexpression_linearity, + expr.values, + )[1] return new( nodes, adj, @@ -85,8 +85,13 @@ struct _FunctionStorage end grad_sparsity = sort!(collect(coloring_storage)) empty!(coloring_storage) + linearity = _classify_linearity( + nodes, + adj, + subexpression_linearity, + const_values, + ) if want_hess - linearity = _classify_linearity(nodes, adj, subexpression_linearity) edgelist = _compute_hessian_sparsity( nodes, adj, @@ -128,7 +133,7 @@ struct _FunctionStorage Int[], Coloring.RecoveryInfo(), Array{Float64}(undef, 0, 0), - NONLINEAR, + linearity[1], dependent_subexpressions, ) end diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl new file mode 100644 index 0000000000..e213eb9b96 --- /dev/null +++ b/src/Nonlinear/layers.jl @@ -0,0 +1,792 @@ +# 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. + +# The evaluator layers in this file are adapted from +# `Ipopt.jl/ext/IpoptMathOptInterfaceExt`. + +""" + ModelWithOracles{T,M}(inner::M) where {T,M} + +A model layer that stores `MOI.VectorOfVariables`-in- +[`MOI.VectorNonlinearOracle`](@ref) constraints and forwards everything else +to the `inner` model. + +`ModelWithOracles(inner)` defaults `T` to `Float64`. + +When wrapped in an [`Evaluator`](@ref), the layer's rows come first, followed +by the rows of the inner evaluator. +""" +mutable struct ModelWithOracles{T,M} + constraints::Vector{ + Tuple{MOI.VectorOfVariables,MOI.VectorNonlinearOracle{T}}, + } + inner::M + + # This constructor wraps an existing vector of constraints so that + # solvers with their own storage can assemble a layer without copying. + function ModelWithOracles{T}( + constraints::Vector{ + Tuple{MOI.VectorOfVariables,MOI.VectorNonlinearOracle{T}}, + }, + inner::M, + ) where {T,M} + return new{T,M}(constraints, inner) + end +end + +function ModelWithOracles{T}(inner) where {T} + constraints = Tuple{MOI.VectorOfVariables,MOI.VectorNonlinearOracle{T}}[] + return ModelWithOracles{T}(constraints, inner) +end + +ModelWithOracles(inner) = ModelWithOracles{Float64}(inner) + +""" + ModelWithQuad{T,M}(inner::M) where {T,M} + +A model layer that stores affine and quadratic objectives and constraints in a +[`QPBlockData`](@ref) and forwards everything else to the `inner` model. + +`ModelWithQuad(inner)` defaults `T` to `Float64`. + +When wrapped in an [`Evaluator`](@ref), the layer's rows come first, followed +by the rows of the inner evaluator. + +Functions added to this layer may contain parameters, following the +convention documented in [`QPBlockData`](@ref): a variable is a parameter if +and only if its index is a key of `qp.parameters`. +""" +mutable struct ModelWithQuad{T,M} + qp::QPBlockData{T} + inner::M + objective_sink::Symbol # :none, :quad or :inner + + # This constructor wraps an existing QPBlockData so that solvers with + # their own storage can assemble a layer without copying. + function ModelWithQuad{T}( + qp::QPBlockData{T}, + inner::M; + objective_sink::Symbol = :none, + ) where {T,M} + return new{T,M}(qp, inner, objective_sink) + end +end + +function ModelWithQuad{T}(inner) where {T} + return ModelWithQuad{T}(QPBlockData{T}(), inner) +end + +ModelWithQuad(inner) = ModelWithQuad{Float64}(inner) + +const _LayerModel = Union{ModelWithQuad,ModelWithOracles} + +# Forwarded methods common to all layers. + +function add_parameter(model::_LayerModel, value::Real) + return add_parameter(model.inner, value) +end + +add_expression(model::_LayerModel, expr) = add_expression(model.inner, expr) + +Base.getindex(model::_LayerModel, index::ExpressionIndex) = model.inner[index] + +function register_operator( + model::_LayerModel, + op::Symbol, + nargs::Int, + f::Function..., +) + return register_operator(model.inner, op, nargs, f...) +end + +function MOI.is_valid(model::_LayerModel, index::ConstraintIndex) + return MOI.is_valid(model.inner, index) +end + +function MOI.get( + model::_LayerModel, + attr::MOI.ListOfSupportedNonlinearOperators, +) + return MOI.get(model.inner, attr) +end + +# ModelWithQuad + +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)) + set_objective(model.inner, obj) + model.objective_sink = obj === nothing ? :none : :inner + return +end + +function add_constraint( + model::ModelWithQuad{T}, + func::Union{MOI.ScalarAffineFunction{T},MOI.ScalarQuadraticFunction{T}}, + set::Union{ + MOI.LessThan{T}, + MOI.GreaterThan{T}, + MOI.EqualTo{T}, + MOI.Interval{T}, + }, +) where {T} + return MOI.add_constraint(model.qp, func, set) +end + +function add_constraint(model::ModelWithQuad, func, set) + return add_constraint(model.inner, func, set) +end + +# ModelWithOracles + +set_objective(model::ModelWithOracles, obj) = set_objective(model.inner, obj) + +function add_constraint( + model::ModelWithOracles{T}, + func::MOI.VectorOfVariables, + set::MOI.VectorNonlinearOracle{T}, +) where {T} + if length(func.variables) != set.input_dimension + throw(DimensionMismatch()) + end + push!(model.constraints, (func, set)) + F, S = MOI.VectorOfVariables, MOI.VectorNonlinearOracle{T} + return MOI.ConstraintIndex{F,S}(length(model.constraints)) +end + +function add_constraint(model::ModelWithOracles, func, set) + return add_constraint(model.inner, func, set) +end + +""" + EvaluatorWithQuad( + model::ModelWithQuad, + inner::MOI.AbstractNLPEvaluator, + ordered_variables::Vector{MOI.VariableIndex}, + ) <: MOI.AbstractNLPEvaluator + +The evaluator of a [`ModelWithQuad`](@ref) layer. The rows of the +[`QPBlockData`](@ref) come first, followed by the rows of `inner`. + +Create it with `Evaluator(model::ModelWithQuad, backend, ordered_variables)`, +which recursively creates the evaluator of the inner model. +""" +mutable struct EvaluatorWithQuad{T,M,E<:MOI.AbstractNLPEvaluator} <: + MOI.AbstractNLPEvaluator + model::ModelWithQuad{T,M} + inner::E + ordered_variables::Vector{MOI.VariableIndex} + # A copy of `model.qp` with the variables mapped to their consecutive + # 1-based index in `ordered_variables`. Rebuilt during `MOI.initialize`. + qp::QPBlockData{T} + + function EvaluatorWithQuad( + model::ModelWithQuad{T,M}, + inner::E, + ordered_variables::Vector{MOI.VariableIndex}, + ) where {T,M,E<:MOI.AbstractNLPEvaluator} + return new{T,M,E}(model, inner, ordered_variables, QPBlockData{T}()) + end +end + +function Evaluator( + model::ModelWithQuad, + backend::AbstractAutomaticDifferentiation, + ordered_variables::Vector{MOI.VariableIndex}, +) + inner = Evaluator(model.inner, backend, ordered_variables) + return EvaluatorWithQuad(model, inner, ordered_variables) +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{T}, + features::Vector{Symbol}, +) where {T} + index_map = Dict{MOI.VariableIndex,MOI.VariableIndex}( + x => MOI.VariableIndex(i) for (i, x) in enumerate(d.ordered_variables) + ) + # Variables absent from `ordered_variables` are parameters: they keep + # their index, which the `parameters` dictionary uses as key. + fmap = v::MOI.VariableIndex -> get(index_map, v, v) + src = d.model.qp + qp = QPBlockData{T}() + qp.objective = MOI.Utilities.map_indices(fmap, src.objective) + qp.objective_function_type = src.objective_function_type + for f in src.constraints + push!(qp.constraints, MOI.Utilities.map_indices(fmap, f)) + end + append!(qp.g_L, src.g_L) + append!(qp.g_U, src.g_U) + append!(qp.mult_g, src.mult_g) + append!(qp.function_type, src.function_type) + append!(qp.bound_type, src.bound_type) + # Alias, do not copy: parameter values updated in the model between + # solves must be visible to the evaluator without re-initializing. + qp.parameters = src.parameters + d.qp = 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.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.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.qp) + MOI.eval_constraint(d.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.qp) + offset = length(d.qp) + # An evaluator is only required to implement `jacobian_structure` if it + # supports `:Jac`. If the inner evaluator does not (it then must not have + # any rows for the stack to be usable), append nothing. + if :Jac in MOI.features_available(d.inner) + for (row, col) in MOI.jacobian_structure(d.inner) + push!(J, (row + offset, col)) + end + end + return J +end + +function MOI.eval_constraint_jacobian(d::EvaluatorWithQuad, J, x) + nnz = MOI.eval_constraint_jacobian(d.qp, J, x) + MOI.eval_constraint_jacobian(d.inner, view(J, (nnz+1):length(J)), x) + return +end + +function MOI.hessian_lagrangian_structure(d::EvaluatorWithQuad) + H = MOI.hessian_lagrangian_structure(d.qp) + if :Hess in MOI.features_available(d.inner) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + end + return H +end + +function MOI.eval_hessian_lagrangian(d::EvaluatorWithQuad, H, x, σ, μ) + m = length(d.qp) + # If the objective is not in the QP block, `d.qp.objective` is zero, so + # passing `σ` is harmless; and vice versa for the inner evaluator. + nnz = MOI.eval_hessian_lagrangian(d.qp, H, x, σ, view(μ, 1:m)) + MOI.eval_hessian_lagrangian( + d.inner, + view(H, (nnz+1):length(H)), + x, + σ, + view(μ, (m+1):length(μ)), + ) + return +end + +# The rows of the two blocks are disjoint, so zero everything and let each +# block write its own rows. +function MOI.eval_constraint_jacobian_product(d::EvaluatorWithQuad, y, x, w) + fill!(y, zero(eltype(y))) + m = length(d.qp) + MOI.eval_constraint_jacobian_product( + d.inner, + view(y, (m+1):length(y)), + x, + w, + ) + MOI.eval_constraint_jacobian_product(d.qp, y, x, w) + return +end + +# Both blocks accumulate into the same variable-dimensional output. Call the +# inner evaluator FIRST because implementations are allowed to overwrite the +# output, and accumulate the QP block afterwards. +function MOI.eval_constraint_jacobian_transpose_product( + d::EvaluatorWithQuad, + y, + x, + w, +) + fill!(y, zero(eltype(y))) + m = length(d.qp) + MOI.eval_constraint_jacobian_transpose_product( + d.inner, + y, + x, + view(w, (m+1):length(w)), + ) + MOI.eval_constraint_jacobian_transpose_product(d.qp, y, x, view(w, 1:m)) + return +end + +function MOI.eval_hessian_lagrangian_product( + d::EvaluatorWithQuad, + H, + x, + v, + σ, + μ, +) + fill!(H, zero(eltype(H))) + m = length(d.qp) + MOI.eval_hessian_lagrangian_product( + d.inner, + H, + x, + v, + σ, + view(μ, (m+1):length(μ)), + ) + MOI.eval_hessian_lagrangian_product(d.qp, H, x, v, σ, view(μ, 1:m)) + return +end + +""" + EvaluatorWithOracles( + model::ModelWithOracles, + inner::MOI.AbstractNLPEvaluator, + ordered_variables::Vector{MOI.VariableIndex}, + ) <: MOI.AbstractNLPEvaluator + +The evaluator of a [`ModelWithOracles`](@ref) layer. The rows of the oracles +come first, in the order they were added, followed by the rows of `inner`. + +Create it with `Evaluator(model::ModelWithOracles, backend, +ordered_variables)`, which recursively creates the evaluator of the inner +model. + +If an oracle does not implement `eval_hessian_lagrangian`, the `:Hess` feature +is removed from [`MOI.features_available`](@ref). +""" +mutable struct EvaluatorWithOracles{T,M,E<:MOI.AbstractNLPEvaluator} <: + MOI.AbstractNLPEvaluator + model::ModelWithOracles{T,M} + inner::E + ordered_variables::Vector{MOI.VariableIndex} + # For each oracle, the consecutive 1-based index of each of its input + # variables. Rebuilt during `MOI.initialize`. + columns::Vector{Vector{Int}} + # For each oracle, a buffer to gather its input variables into. The public + # `MOI.VectorNonlinearOracle` has no scratch storage of its own. + x_buffer::Vector{Vector{T}} + + function EvaluatorWithOracles( + model::ModelWithOracles{T,M}, + inner::E, + ordered_variables::Vector{MOI.VariableIndex}, + ) where {T,M,E<:MOI.AbstractNLPEvaluator} + return new{T,M,E}( + model, + inner, + ordered_variables, + Vector{Int}[], + Vector{T}[], + ) + end +end + +function Evaluator( + model::ModelWithOracles, + backend::AbstractAutomaticDifferentiation, + ordered_variables::Vector{MOI.VariableIndex}, +) + inner = Evaluator(model.inner, backend, ordered_variables) + return EvaluatorWithOracles(model, inner, ordered_variables) +end + +function _num_rows(d::EvaluatorWithOracles) + return sum(s.output_dimension for (_, s) in d.model.constraints; init = 0) +end + +function MOI.features_available(d::EvaluatorWithOracles) + features = MOI.features_available(d.inner) + features = + filter(f -> f in (:Grad, :Jac, :JacVec, :Hess, :HessVec), features) + if !isempty(d.model.constraints) + # The oracles have no product callbacks. + filter!(f -> !(f in (:JacVec, :HessVec)), features) + end + no_hessian = any(d.model.constraints) do (_, s) + return s.eval_hessian_lagrangian === nothing + end + if no_hessian + filter!(f -> f != :Hess, features) + end + return features +end + +function MOI.initialize( + d::EvaluatorWithOracles{T}, + features::Vector{Symbol}, +) where {T} + index_map = Dict{MOI.VariableIndex,Int}( + x => i for (i, x) in enumerate(d.ordered_variables) + ) + empty!(d.columns) + empty!(d.x_buffer) + for (f, s) in d.model.constraints + push!(d.columns, [index_map[x] for x in f.variables]) + push!(d.x_buffer, zeros(T, s.input_dimension)) + end + MOI.initialize(d.inner, features) + return +end + +function _gather!(d::EvaluatorWithOracles, k::Int, x) + xk = d.x_buffer[k] + for (j, col) in enumerate(d.columns[k]) + xk[j] = x[col] + end + return xk +end + +MOI.eval_objective(d::EvaluatorWithOracles, x) = MOI.eval_objective(d.inner, x) + +function MOI.eval_objective_gradient(d::EvaluatorWithOracles, grad, x) + MOI.eval_objective_gradient(d.inner, grad, x) + return +end + +function MOI.eval_constraint(d::EvaluatorWithOracles, g, x) + offset = 0 + for (k, (_, s)) in enumerate(d.model.constraints) + xk = _gather!(d, k, x) + s.eval_f(view(g, offset .+ (1:s.output_dimension)), xk) + offset += s.output_dimension + end + MOI.eval_constraint(d.inner, view(g, (offset+1):length(g)), x) + return +end + +function MOI.jacobian_structure(d::EvaluatorWithOracles) + J = Tuple{Int,Int}[] + row_offset = 0 + for (k, (_, s)) in enumerate(d.model.constraints) + columns = d.columns[k] + for (i, j) in s.jacobian_structure + push!(J, (row_offset + i, columns[j])) + end + row_offset += s.output_dimension + end + if :Jac in MOI.features_available(d.inner) + for (row, col) in MOI.jacobian_structure(d.inner) + push!(J, (row + row_offset, col)) + end + end + return J +end + +function MOI.eval_constraint_jacobian(d::EvaluatorWithOracles, J, x) + offset = 0 + for (k, (_, s)) in enumerate(d.model.constraints) + xk = _gather!(d, k, x) + nnz = length(s.jacobian_structure) + s.eval_jacobian(view(J, offset .+ (1:nnz)), xk) + offset += nnz + end + MOI.eval_constraint_jacobian(d.inner, view(J, (offset+1):length(J)), x) + return +end + +function MOI.hessian_lagrangian_structure(d::EvaluatorWithOracles) + H = Tuple{Int,Int}[] + for (k, (_, s)) in enumerate(d.model.constraints) + columns = d.columns[k] + for (i, j) in s.hessian_lagrangian_structure + push!(H, (columns[i], columns[j])) + end + end + if :Hess in MOI.features_available(d.inner) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + end + return H +end + +function MOI.eval_hessian_lagrangian(d::EvaluatorWithOracles, H, x, σ, μ) + offset, μ_offset = 0, 0 + for (k, (_, s)) in enumerate(d.model.constraints) + xk = _gather!(d, k, x) + nnz = length(s.hessian_lagrangian_structure) + μk = view(μ, μ_offset .+ (1:s.output_dimension)) + s.eval_hessian_lagrangian(view(H, offset .+ (1:nnz)), xk, μk) + offset += nnz + μ_offset += s.output_dimension + end + MOI.eval_hessian_lagrangian( + d.inner, + view(H, (offset+1):length(H)), + x, + σ, + view(μ, (μ_offset+1):length(μ)), + ) + return +end + +# The oracles have no product callbacks, so the layer removes :JacVec and +# :HessVec from `MOI.features_available` when it has oracle constraints. +# The products are still answered if called (some solvers use the transpose +# product for dual computations regardless), by materializing each oracle's +# Jacobian or Hessian. This can be slow on large oracles. + +function MOI.eval_constraint_jacobian_product(d::EvaluatorWithOracles, y, x, w) + fill!(y, zero(eltype(y))) + offset = 0 + for (k, (_, s)) in enumerate(d.model.constraints) + xk = _gather!(d, k, x) + J_val = zeros(eltype(y), length(s.jacobian_structure)) + s.eval_jacobian(J_val, xk) + columns = d.columns[k] + for ((r, c), v) in zip(s.jacobian_structure, J_val) + y[offset+r] += v * w[columns[c]] + end + offset += s.output_dimension + end + MOI.eval_constraint_jacobian_product( + d.inner, + view(y, (offset+1):length(y)), + x, + w, + ) + return +end + +function MOI.eval_constraint_jacobian_transpose_product( + d::EvaluatorWithOracles, + y, + x, + w, +) + # Inner first: implementations are allowed to overwrite `y`. + offset = _num_rows(d) + MOI.eval_constraint_jacobian_transpose_product( + d.inner, + y, + x, + view(w, (offset+1):length(w)), + ) + row_offset = 0 + for (k, (_, s)) in enumerate(d.model.constraints) + xk = _gather!(d, k, x) + J_val = zeros(eltype(y), length(s.jacobian_structure)) + s.eval_jacobian(J_val, xk) + columns = d.columns[k] + for ((r, c), v) in zip(s.jacobian_structure, J_val) + y[columns[c]] += v * w[row_offset+r] + end + row_offset += s.output_dimension + end + return +end + +function MOI.eval_hessian_lagrangian_product( + d::EvaluatorWithOracles, + H, + x, + v, + σ, + μ, +) + μ_offset = _num_rows(d) + MOI.eval_hessian_lagrangian_product( + d.inner, + H, + x, + v, + σ, + view(μ, (μ_offset+1):length(μ)), + ) + row_offset = 0 + for (k, (_, s)) in enumerate(d.model.constraints) + if s.eval_hessian_lagrangian === nothing + error( + "The Hessian-vector product is not available because a " * + "VectorNonlinearOracle does not implement " * + "`eval_hessian_lagrangian`.", + ) + end + xk = _gather!(d, k, x) + H_val = zeros(eltype(H), length(s.hessian_lagrangian_structure)) + μk = view(μ, row_offset .+ (1:s.output_dimension)) + s.eval_hessian_lagrangian(H_val, xk, μk) + columns = d.columns[k] + for ((i, j), h) in zip(s.hessian_lagrangian_structure, H_val) + H[columns[i]] += h * v[columns[j]] + if i != j + H[columns[j]] += h * v[columns[i]] + end + end + row_offset += s.output_dimension + end + return +end + +# Linearity and row queries + +function num_constraints(d::EvaluatorWithQuad) + return length(d.model.qp) + num_constraints(d.inner) +end + +function num_constraints(d::EvaluatorWithOracles) + return _num_rows(d) + num_constraints(d.inner) +end + +# Like `num_constraints`, but returns `nothing` instead of erroring when an +# evaluator at the bottom of the stack does not implement the query. +function _try_num_constraints(ev::MOI.AbstractNLPEvaluator) + if !applicable(num_constraints, ev) + return nothing + end + return num_constraints(ev) +end + +function _try_num_constraints(d::EvaluatorWithQuad) + n = _try_num_constraints(d.inner) + return n === nothing ? nothing : length(d.model.qp) + n +end + +function _try_num_constraints(d::EvaluatorWithOracles) + n = _try_num_constraints(d.inner) + return n === nothing ? nothing : _num_rows(d) + n +end + +# Returns `nothing` if the inner evaluator implements neither +# `constraint_linearity` nor `num_constraints`, in which case the layer +# cannot describe its rows either. +function _inner_constraint_linearity(inner::MOI.AbstractNLPEvaluator) + linearity = constraint_linearity(inner) + if linearity !== nothing + return linearity + end + n = _try_num_constraints(inner) + if n === nothing + return nothing + end + return fill(NONLINEAR, n) +end + +function constraint_linearity(d::EvaluatorWithQuad) + inner = _inner_constraint_linearity(d.inner) + if inner === nothing + return nothing + end + linearity = Linearity[ + ft == _kFunctionTypeScalarQuadratic ? QUADRATIC : LINEAR for + ft in d.model.qp.function_type + ] + return vcat(linearity, inner) +end + +function constraint_linearity(d::EvaluatorWithOracles) + inner = _inner_constraint_linearity(d.inner) + if inner === nothing + return nothing + end + return vcat(fill(NONLINEAR, _num_rows(d)), inner) +end + +function objective_linearity(d::EvaluatorWithQuad) + sink = d.model.objective_sink + if sink == :quad + if d.model.qp.objective_function_type == _kFunctionTypeScalarQuadratic + return QUADRATIC + end + return LINEAR + elseif sink == :inner + return objective_linearity(d.inner) + else + return CONSTANT + end +end + +objective_linearity(d::EvaluatorWithOracles) = objective_linearity(d.inner) + +""" + constraint_bounds( + evaluator::MOI.AbstractNLPEvaluator, + )::Vector{MOI.NLPBoundsPair} + +Return the lower and upper bounds of each constraint in `evaluator`, aligned +with the rows of [`MOI.eval_constraint`](@ref). +""" +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 vcat(bounds, constraint_bounds(d.inner)) +end + +function constraint_bounds(d::EvaluatorWithOracles) + bounds = MOI.NLPBoundsPair[] + for (_, s) in d.model.constraints + for (l, u) in zip(s.l, s.u) + push!(bounds, MOI.NLPBoundsPair(l, u)) + end + end + return vcat(bounds, constraint_bounds(d.inner)) +end + +# `true` is the conservative default: the solver will then call +# `eval_objective`, and evaluators without an objective return zero. +_has_objective(::MOI.AbstractNLPEvaluator) = true + +_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 + +_has_objective(d::EvaluatorWithOracles) = _has_objective(d.inner) + +function MOI.NLPBlockData(d::Union{EvaluatorWithQuad,EvaluatorWithOracles}) + return MOI.NLPBlockData(constraint_bounds(d), d, _has_objective(d)) +end diff --git a/src/Nonlinear/linearity.jl b/src/Nonlinear/linearity.jl new file mode 100644 index 0000000000..b94dd02cc3 --- /dev/null +++ b/src/Nonlinear/linearity.jl @@ -0,0 +1,128 @@ +# 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. + +""" + Linearity + +An enum describing the linearity of an expression with respect to the decision +variables, for fixed values of the parameters. + +The classification is conservative: an expression may be classified less +strictly than the tightest class that applies (for example, an expression that +simplifies to an affine function may be classified as `NONLINEAR`), but never +more strictly. + +## Values + + * `CONSTANT`: the value does not depend on the decision variables + * `LINEAR`: the value is an affine function of the decision variables; the + gradient is constant and the Hessian is zero + * `PIECEWISE_LINEAR`: the gradient is piecewise constant and the Hessian is + zero almost everywhere + * `QUADRATIC`: the gradient is an affine function of the decision variables + and the Hessian is constant + * `NONLINEAR`: no guarantee + +Because parameters are held fixed by this classification, consumers that cache +constant derivatives (for example, the Jacobian coefficients of `LINEAR` rows) +must invalidate their cache when parameter values change. +""" +const Linearity = ReverseAD.Linearity + +const CONSTANT = ReverseAD.CONSTANT +const LINEAR = ReverseAD.LINEAR +const PIECEWISE_LINEAR = ReverseAD.PIECEWISE_LINEAR +const QUADRATIC = ReverseAD.QUADRATIC +const NONLINEAR = ReverseAD.NONLINEAR + +""" + num_constraints(evaluator::MOI.AbstractNLPEvaluator)::Int + +Return the number of constraints in `evaluator`, that is, the length of the +vector `g` filled by [`MOI.eval_constraint`](@ref). + +## Implementation + +There is no default fallback: evaluators opt in to this query by adding a +method. +""" +function num_constraints end + +num_constraints(evaluator::Evaluator) = length(evaluator.model.constraints) + +num_constraints(d::ReverseAD.NLPEvaluator) = length(d.data.constraints) + +""" + constraint_linearity( + evaluator::MOI.AbstractNLPEvaluator, + )::Union{Nothing,Vector{Linearity}} + +Return a vector of the [`Linearity`](@ref) of each constraint in `evaluator`, +aligned with the rows of [`MOI.eval_constraint`](@ref), or `nothing` if the +evaluator does not implement this query. + +Callers must treat `nothing` as if every row were `NONLINEAR`. + +The length of a non-`nothing` return value is [`num_constraints`](@ref). + +## Initialize + +Before querying this function, you must call [`MOI.initialize`](@ref). +""" +constraint_linearity(::MOI.AbstractNLPEvaluator) = nothing + +function constraint_linearity(evaluator::Evaluator) + if evaluator.backend === nothing + return nothing + end + return constraint_linearity(evaluator.backend) +end + +function constraint_linearity(d::ReverseAD.NLPEvaluator) + if !isdefined(d, :constraints) + error( + "Unable to query constraint_linearity because MOI.initialize " * + "has not been called.", + ) + end + return Linearity[c.linearity for c in d.constraints] +end + +""" + objective_linearity(evaluator::MOI.AbstractNLPEvaluator)::Linearity + +Return the [`Linearity`](@ref) of the objective function in `evaluator`. + +The default fallback returns `NONLINEAR`, which is always a valid (if +conservative) answer. If the evaluator has no objective, return `CONSTANT`. + +## Initialize + +Before querying this function, you must call [`MOI.initialize`](@ref). +""" +objective_linearity(::MOI.AbstractNLPEvaluator) = NONLINEAR + +function objective_linearity(evaluator::Evaluator) + if evaluator.model.objective === nothing + return CONSTANT + elseif evaluator.backend === nothing + return NONLINEAR + end + return objective_linearity(evaluator.backend) +end + +function objective_linearity(d::ReverseAD.NLPEvaluator) + if !isdefined(d, :objective) + error( + "Unable to query objective_linearity because MOI.initialize " * + "has not been called.", + ) + end + if d.objective === nothing + return CONSTANT + end + return something(d.objective).linearity +end diff --git a/src/Nonlinear/model.jl b/src/Nonlinear/model.jl index b570a0dcc3..68cd2110db 100644 --- a/src/Nonlinear/model.jl +++ b/src/Nonlinear/model.jl @@ -77,6 +77,43 @@ function set_objective(model::Model, ::Nothing) return end +""" + model(backend::AbstractAutomaticDifferentiation) + +Return a new nonlinear model appropriate for the given AD `backend`. + +The default returns `ModelWithQuad(ModelWithOracles(Model()))`, that is, a +[`Model`](@ref) wrapped in the [`ModelWithOracles`](@ref) and +[`ModelWithQuad`](@ref) layers, so that the backend receives only the +constraints it can differentiate, and affine, quadratic, and +[`MOI.VectorNonlinearOracle`](@ref) constraints are handled by the layers. + +Custom AD backends can override this method to return their own model type, +wrapped in the layers they do not handle themselves. For example, a backend +that exploits the structure of affine and quadratic constraints but cannot +evaluate oracles would return `ModelWithOracles(CustomModel())`. +""" +function model(::AbstractAutomaticDifferentiation) + return ModelWithQuad(ModelWithOracles(Model())) +end + +""" + exploits_structure(backend::AbstractAutomaticDifferentiation)::Bool + +Return `true` if the model returned by [`model`](@ref)`(backend)` natively +accepts [`MOI.ScalarAffineFunction`](@ref) and +[`MOI.ScalarQuadraticFunction`](@ref) objectives and constraints, that is, if +the backend exploits the structure of affine and quadratic functions itself +instead of relying on the [`ModelWithQuad`](@ref) layer. + +Consumers that parse affine and quadratic constraints into their own data +structures (for example, NLPModelsJuMP) should pass them unparsed to the +model of a backend for which this function returns `true`. + +The default is `false`. +""" +exploits_structure(::AbstractAutomaticDifferentiation) = false + """ add_expression(model::Model, expr)::ExpressionIndex 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 diff --git a/test/Nonlinear/test_ReverseAD.jl b/test/Nonlinear/test_ReverseAD.jl index f1a6cc4fb0..0809a9191a 100644 --- a/test/Nonlinear/test_ReverseAD.jl +++ b/test/Nonlinear/test_ReverseAD.jl @@ -554,7 +554,12 @@ function test_linearity() expr = model[ex] adj = Nonlinear.adjacency_matrix(expr.nodes) nodes = ReverseAD._replace_moi_variables(expr.nodes, variables) - ret = ReverseAD._classify_linearity(nodes, adj, ReverseAD.Linearity[]) + ret = ReverseAD._classify_linearity( + nodes, + adj, + ReverseAD.Linearity[], + expr.values, + ) @test ret[1] == test_value indexed_set = Coloring.IndexedSet(100) edge_list = ReverseAD._compute_hessian_sparsity( @@ -564,7 +569,13 @@ function test_linearity() Set{Tuple{Int,Int}}[], Vector{Int}[], ) - if ret[1] != ReverseAD.NONLINEAR + if ret[1] <= ReverseAD.LINEAR + # Constant or linear: the Hessian is zero, so there must not be + # any Hessian edges. Note that this does not hold for + # PIECEWISE_LINEAR: the Hessian is zero almost everywhere, but + # the operator-driven structural sparsity may still contain + # entries (for example, `abs` registers a self-interaction), and + # their values evaluate to zero. @test length(edge_list) == 0 elseif length(IJ) > 0 @test IJ == edge_list @@ -584,7 +595,7 @@ function test_linearity() [1, 2], ) _test_linearity(:(3 * 4 * ($x + $y)), ReverseAD.LINEAR) - _test_linearity(:($z * $y), ReverseAD.NONLINEAR, Set([(3, 2)]), [2, 3]) + _test_linearity(:($z * $y), ReverseAD.QUADRATIC, Set([(3, 2)]), [2, 3]) _test_linearity(:(3 + 4), ReverseAD.CONSTANT) _test_linearity(:(sin(3) + $x), ReverseAD.LINEAR) _test_linearity( @@ -635,6 +646,33 @@ function test_linearity() Set([(3, 3), (3, 2), (3, 1)]), [1, 2, 3], ) + # QUADRATIC + _test_linearity(:($x^2), ReverseAD.QUADRATIC, Set([(1, 1)]), [1]) + _test_linearity(:($x^2.0), ReverseAD.QUADRATIC) + _test_linearity(:(($x + 3 * $y)^2), ReverseAD.QUADRATIC) + _test_linearity(:(2 * $x^2), ReverseAD.QUADRATIC) + _test_linearity(:($x^2 + 3 * $x + 1), ReverseAD.QUADRATIC) + _test_linearity(:(-($x^2)), ReverseAD.QUADRATIC) + _test_linearity(:($x^2 / 4), ReverseAD.QUADRATIC) + _test_linearity(:($x * $y / 2), ReverseAD.QUADRATIC) + _test_linearity(:($x^3), ReverseAD.NONLINEAR) + _test_linearity(:($x^$x), ReverseAD.NONLINEAR) + _test_linearity(:(2^$x), ReverseAD.NONLINEAR) + _test_linearity(:($x^2 * $y), ReverseAD.NONLINEAR) + _test_linearity(:($x * $y * $z), ReverseAD.NONLINEAR) + _test_linearity(:(sin($x)^2), ReverseAD.NONLINEAR) + # PIECEWISE_LINEAR + _test_linearity(:(abs($x)), ReverseAD.PIECEWISE_LINEAR) + _test_linearity(:(abs(2 * $x + $y) / 2), ReverseAD.PIECEWISE_LINEAR) + _test_linearity(:(abs($x^2)), ReverseAD.NONLINEAR) + _test_linearity(:(min($x, $y)), ReverseAD.PIECEWISE_LINEAR) + _test_linearity(:(max($x, 2 * $y, 1)), ReverseAD.PIECEWISE_LINEAR) + _test_linearity(:(min($x^2, $y)), ReverseAD.NONLINEAR) + _test_linearity(:($x + abs($y)), ReverseAD.PIECEWISE_LINEAR) + _test_linearity(:(ifelse($x <= 1, 2.0, $y)), ReverseAD.PIECEWISE_LINEAR) + # The sum of a quadratic and a piecewise linear term is piecewise + # quadratic: neither class applies. + _test_linearity(:($x^2 + abs($y)), ReverseAD.NONLINEAR) return end @@ -645,10 +683,99 @@ function test_linearity_no_hess() Nonlinear.set_objective(model, ex) evaluator = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) MOI.initialize(evaluator, [:Grad, :Jac]) - # We initialized without the need for the hessian so - # the linearity shouldn't be computed. - @test only(evaluator.backend.subexpressions).linearity == - ReverseAD.NONLINEAR + # The linearity is computed even when we initialize without :Hess, so + # that `Nonlinear.constraint_linearity` can be queried. + @test only(evaluator.backend.subexpressions).linearity == ReverseAD.LINEAR + return +end + +function test_linearity_queries() + x = MOI.VariableIndex(1) + y = MOI.VariableIndex(2) + model = Nonlinear.Model() + Nonlinear.set_objective(model, :($x^2 + $y)) + Nonlinear.add_constraint(model, :($x + $y), MOI.LessThan(1.0)) + Nonlinear.add_constraint(model, :($x * $y), MOI.LessThan(1.0)) + Nonlinear.add_constraint(model, :(sin($x)), MOI.LessThan(1.0)) + Nonlinear.add_constraint(model, :(abs($x)), MOI.LessThan(1.0)) + Nonlinear.add_constraint( + model, + :(ifelse($x <= 1, $x, $y)), + MOI.LessThan(1.0), + ) + Nonlinear.add_constraint(model, :($x / 2), MOI.LessThan(1.0)) + for features in ([:Grad, :Jac], [:Grad, :Jac, :Hess]) + evaluator = + Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y]) + MOI.initialize(evaluator, features) + @test Nonlinear.num_constraints(evaluator) == 6 + @test Nonlinear.objective_linearity(evaluator) == Nonlinear.QUADRATIC + @test Nonlinear.constraint_linearity(evaluator) == [ + Nonlinear.LINEAR, + Nonlinear.QUADRATIC, + Nonlinear.NONLINEAR, + Nonlinear.PIECEWISE_LINEAR, + Nonlinear.PIECEWISE_LINEAR, + Nonlinear.LINEAR, + ] + end + # An evaluator without a backend does not implement the linearity queries, + # so it returns the documented conservative fallbacks. + evaluator = Nonlinear.Evaluator(model) + MOI.initialize(evaluator, Symbol[]) + @test Nonlinear.num_constraints(evaluator) == 6 + @test Nonlinear.constraint_linearity(evaluator) === nothing + @test Nonlinear.objective_linearity(evaluator) == Nonlinear.NONLINEAR + return +end + +function test_linearity_queries_before_initialize() + x = MOI.VariableIndex(1) + model = Nonlinear.Model() + Nonlinear.set_objective(model, :($x^2)) + Nonlinear.add_constraint(model, :($x + 1.0), MOI.LessThan(1.0)) + evaluator = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + @test Nonlinear.num_constraints(evaluator) == 1 + @test_throws( + ErrorException( + "Unable to query constraint_linearity because MOI.initialize " * + "has not been called.", + ), + Nonlinear.constraint_linearity(evaluator.backend), + ) + @test_throws( + ErrorException( + "Unable to query objective_linearity because MOI.initialize " * + "has not been called.", + ), + Nonlinear.objective_linearity(evaluator.backend), + ) + return +end + +function test_linearity_queries_no_objective() + x = MOI.VariableIndex(1) + model = Nonlinear.Model() + Nonlinear.add_constraint(model, :($x^2), MOI.LessThan(1.0)) + evaluator = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(evaluator, [:Grad, :Jac]) + @test Nonlinear.objective_linearity(evaluator) == Nonlinear.CONSTANT + @test Nonlinear.constraint_linearity(evaluator) == [Nonlinear.QUADRATIC] + return +end + +function test_linearity_queries_subexpression() + x = MOI.VariableIndex(1) + y = MOI.VariableIndex(2) + model = Nonlinear.Model() + ex = Nonlinear.add_expression(model, :($x^2 + $y)) + Nonlinear.add_constraint(model, :($ex + 1.0), MOI.LessThan(1.0)) + Nonlinear.add_constraint(model, :(sin($ex)), MOI.LessThan(1.0)) + evaluator = + Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y]) + MOI.initialize(evaluator, [:Grad, :Jac]) + @test Nonlinear.constraint_linearity(evaluator) == + [Nonlinear.QUADRATIC, Nonlinear.NONLINEAR] return end diff --git a/test/Nonlinear/test_layers.jl b/test/Nonlinear/test_layers.jl new file mode 100644 index 0000000000..b1cd39195a --- /dev/null +++ b/test/Nonlinear/test_layers.jl @@ -0,0 +1,399 @@ +# 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 TestNonlinearLayers + +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 + +function _squared_oracle() + return MOI.VectorNonlinearOracle(; + dimension = 1, + l = [0.0], + u = [1.0], + eval_f = (ret, x) -> (ret[1] = x[1]^2), + jacobian_structure = [(1, 1)], + eval_jacobian = (ret, x) -> (ret[1] = 2.0 * x[1]), + hessian_lagrangian_structure = [(1, 1)], + eval_hessian_lagrangian = (ret, x, μ) -> (ret[1] = 2.0 * μ[1]), + ) +end + +# A stacked 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 (oracle layer): x^2 in [0, 1] +# row 4 (inner nlp): sin(x) <= 0.5 +# and the objective x^2 in the quad layer. +function _test_model(x, y) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + @test model isa Nonlinear.ModelWithQuad + @test model.inner isa Nonlinear.ModelWithOracles + @test model.inner.inner isa Nonlinear.Model + Nonlinear.set_objective( + model, + MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ), + ) + c1 = Nonlinear.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, + MOI.VectorOfVariables([x]), + _squared_oracle(), + ) + @test c3 isa MOI.ConstraintIndex{ + MOI.VectorOfVariables, + MOI.VectorNonlinearOracle{Float64}, + } + c4 = Nonlinear.add_constraint(model, :(sin($x)), MOI.LessThan(0.5)) + @test c4 isa Nonlinear.ConstraintIndex + return model +end + +function test_evaluator_with_stack() + # Use non-consecutive variable indices to test that the layers remap the + # variables to their consecutive index in `ordered_variables`. + x, y = MOI.VariableIndex(5), MOI.VariableIndex(9) + model = _test_model(x, y) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y]) + @test d isa Nonlinear.EvaluatorWithQuad + @test d.inner isa Nonlinear.EvaluatorWithOracles + @test d.inner.inner isa Nonlinear.Evaluator + @test MOI.features_available(d) == [:Grad, :Jac, :Hess] + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test Nonlinear.num_constraints(d) == 4 + @test Nonlinear.constraint_bounds(d) == [ + MOI.NLPBoundsPair(-Inf, 4.0), + MOI.NLPBoundsPair(0.0, 1.0), + MOI.NLPBoundsPair(0.0, 1.0), + MOI.NLPBoundsPair(-Inf, 0.5), + ] + @test Nonlinear.constraint_linearity(d) == [ + Nonlinear.LINEAR, + Nonlinear.QUADRATIC, + Nonlinear.NONLINEAR, + Nonlinear.NONLINEAR, + ] + @test Nonlinear.objective_linearity(d) == Nonlinear.QUADRATIC + 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, 4) + MOI.eval_constraint(d, g, xv) + @test g ≈ [8.0, 5.0, 1.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(4, 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 + 2.0 0.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, 10_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) + μ₃ * ∇²(x^2) + μ₄ * ∇²(sin(x)) + @test H[1, 1] ≈ 2σ + 2 * μ[2] + 2 * μ[3] - sin(1.0) * μ[4] + @test H[1, 2] + H[2, 1] ≈ μ[2] + @test H[2, 2] ≈ 0.0 + block = MOI.NLPBlockData(d) + @test block.has_objective + @test length(block.constraint_bounds) == 4 + return +end + +function test_objective_sink_switching() + x = MOI.VariableIndex(1) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + @test model.objective_sink == :none + f = MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ) + Nonlinear.set_objective(model, f) + @test model.objective_sink == :quad + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test MOI.eval_objective(d, [3.0]) == 9.0 + @test Nonlinear.objective_linearity(d) == Nonlinear.QUADRATIC + @test MOI.NLPBlockData(d).has_objective + # Switch to a nonlinear objective: the quadratic objective must be + # cleared, including its Hessian entries. + Nonlinear.set_objective(model, :(sin($x))) + @test model.objective_sink == :inner + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test MOI.eval_objective(d, [3.0]) == sin(3.0) + @test Nonlinear.objective_linearity(d) == Nonlinear.NONLINEAR + @test MOI.NLPBlockData(d).has_objective + H_structure = MOI.hessian_lagrangian_structure(d) + H = fill(NaN, length(H_structure)) + MOI.eval_hessian_lagrangian(d, H, [3.0], 1.0, Float64[]) + @test sum(H) ≈ -sin(3.0) + # Switch to a linear objective, and then remove it. + g = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(2.0, x)], 1.0) + Nonlinear.set_objective(model, g) + @test model.objective_sink == :quad + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(d, [:Grad, :Jac]) + @test MOI.eval_objective(d, [3.0]) == 7.0 + @test Nonlinear.objective_linearity(d) == Nonlinear.LINEAR + Nonlinear.set_objective(model, nothing) + @test model.objective_sink == :none + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + 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 Nonlinear.objective_linearity(d) == Nonlinear.CONSTANT + @test !MOI.NLPBlockData(d).has_objective + return +end + +function test_oracle_without_hessian() + x = MOI.VariableIndex(1) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + set = MOI.VectorNonlinearOracle(; + dimension = 1, + l = [0.0], + u = [1.0], + eval_f = (ret, x) -> (ret[1] = x[1]^2), + jacobian_structure = [(1, 1)], + eval_jacobian = (ret, x) -> (ret[1] = 2.0 * x[1]), + ) + Nonlinear.add_constraint(model, MOI.VectorOfVariables([x]), set) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + @test MOI.features_available(d) == [:Grad, :Jac] + MOI.initialize(d, [:Grad, :Jac]) + g = fill(NaN, 1) + MOI.eval_constraint(d, g, [2.0]) + @test g == [4.0] + return +end + +function test_oracle_dimension_mismatch() + x, y = MOI.VariableIndex(1), MOI.VariableIndex(2) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + @test_throws( + DimensionMismatch, + Nonlinear.add_constraint( + model, + MOI.VectorOfVariables([x, y]), + _squared_oracle(), + ), + ) + return +end + +function test_forwarding_through_layers() + x = MOI.VariableIndex(1) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + p = Nonlinear.add_parameter(model, 2.0) + @test p isa Nonlinear.ParameterIndex + ex = Nonlinear.add_expression(model, :($p * $x)) + @test ex isa Nonlinear.ExpressionIndex + @test model[ex] isa Nonlinear.Expression + Nonlinear.register_operator(model, :my_square, 1, z -> z^2) + Nonlinear.add_constraint(model, :(my_square($ex)), MOI.LessThan(1.0)) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(d, [:Grad, :Jac]) + g = fill(NaN, 1) + MOI.eval_constraint(d, g, [3.0]) + @test g == [36.0] + return +end + +function test_stack_products() + x, y = MOI.VariableIndex(1), MOI.VariableIndex(2) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + Nonlinear.set_objective( + model, + MOI.ScalarQuadraticFunction( + [MOI.ScalarQuadraticTerm(2.0, x, x)], + MOI.ScalarAffineTerm{Float64}[], + 0.0, + ), + ) + Nonlinear.add_constraint( + model, + MOI.ScalarAffineFunction( + [MOI.ScalarAffineTerm(2.0, x), MOI.ScalarAffineTerm(3.0, y)], + 0.0, + ), + MOI.LessThan(4.0), + ) + 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), + ) + Nonlinear.add_constraint(model, :(sin($x)), MOI.LessThan(0.5)) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x, y]) + # No oracle constraints, so the product features pass through. + @test MOI.features_available(d) == [:Grad, :Jac, :JacVec, :Hess, :HessVec] + 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_quad_layer_parameters() + x, p = MOI.VariableIndex(1), MOI.VariableIndex(42) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + # Register the parameter before any structure query. Its value may be + # updated between evaluations. + model.qp.parameters[p.value] = 5.0 + 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(), [x]) + 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. + model.qp.parameters[p.value] = 7.0 + MOI.eval_constraint(d, g, [1.0]) + @test g == [2.0 * 1.0 + 3.0 * 7.0, 7.0 * 1.0] + return +end + +function test_quad_layer_row_order_with_empty_inner() + x = MOI.VariableIndex(1) + model = Nonlinear.model(Nonlinear.SparseReverseMode()) + Nonlinear.add_constraint( + model, + MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0), + MOI.GreaterThan(1.0), + ) + d = Nonlinear.Evaluator(model, Nonlinear.SparseReverseMode(), [x]) + MOI.initialize(d, [:Grad, :Jac, :Hess]) + @test Nonlinear.num_constraints(d) == 1 + @test Nonlinear.constraint_linearity(d) == [Nonlinear.LINEAR] + @test Nonlinear.constraint_bounds(d) == [MOI.NLPBoundsPair(1.0, Inf)] + g = fill(NaN, 1) + MOI.eval_constraint(d, g, [1.5]) + @test g == [1.5] + @test isempty(MOI.hessian_lagrangian_structure(d)) + return +end + +end # module + +TestNonlinearLayers.runtests()