From 354fc7d221315a3fc8670bbe0d9aed5e4c21952f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 8 Apr 2026 22:12:16 +0200 Subject: [PATCH 01/12] Add nonlinear_model function --- src/Nonlinear/model.jl | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/Nonlinear/model.jl b/src/Nonlinear/model.jl index b570a0dcc3..023640625e 100644 --- a/src/Nonlinear/model.jl +++ b/src/Nonlinear/model.jl @@ -77,6 +77,16 @@ function set_objective(model::Model, ::Nothing) return end +""" + nonlinear_model(backend::AbstractAutomaticDifferentiation) + +Return a new nonlinear model appropriate for the given AD `backend`. + +The default returns `Model()`. Custom AD backends can override this +to return their own model type. +""" +nonlinear_model(::AbstractAutomaticDifferentiation) = Model() + """ add_expression(model::Model, expr)::ExpressionIndex From 46e51031856e05a62bd3f1f21f46e8fed3bc7479 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Tue, 21 Apr 2026 10:34:35 +0200 Subject: [PATCH 02/12] Renam --- src/Nonlinear/model.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Nonlinear/model.jl b/src/Nonlinear/model.jl index 023640625e..8577de2f34 100644 --- a/src/Nonlinear/model.jl +++ b/src/Nonlinear/model.jl @@ -78,14 +78,14 @@ function set_objective(model::Model, ::Nothing) end """ - nonlinear_model(backend::AbstractAutomaticDifferentiation) + model(backend::AbstractAutomaticDifferentiation) Return a new nonlinear model appropriate for the given AD `backend`. The default returns `Model()`. Custom AD backends can override this to return their own model type. """ -nonlinear_model(::AbstractAutomaticDifferentiation) = Model() +model(::AbstractAutomaticDifferentiation) = Model() """ add_expression(model::Model, expr)::ExpressionIndex From 459c20a2484333cb94294ab0f13b70293e54087a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 17:34:20 +0000 Subject: [PATCH 03/12] Add QUADRATIC linearity and public linearity queries to Nonlinear --- docs/src/submodules/Nonlinear/reference.md | 9 ++ src/Nonlinear/Nonlinear.jl | 4 + src/Nonlinear/ReverseAD/graph_tools.jl | 124 +++++++++------ .../ReverseAD/mathoptinterface_api.jl | 2 +- src/Nonlinear/ReverseAD/types.jl | 21 ++- src/Nonlinear/linearity.jl | 128 ++++++++++++++++ test/Nonlinear/test_ReverseAD.jl | 141 +++++++++++++++++- 7 files changed, 369 insertions(+), 60 deletions(-) create mode 100644 src/Nonlinear/linearity.jl diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index 842db49713..9f10b404f8 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -74,6 +74,15 @@ Nonlinear.SparseReverseMode Nonlinear.SymbolicMode ``` +## Linearity + +```@docs +Nonlinear.Linearity +Nonlinear.num_constraints +Nonlinear.constraint_linearity +Nonlinear.objective_linearity +``` + ## Data-structure ```@docs diff --git a/src/Nonlinear/Nonlinear.jl b/src/Nonlinear/Nonlinear.jl index 9a5e4bd8d9..32ceb5182f 100644 --- a/src/Nonlinear/Nonlinear.jl +++ b/src/Nonlinear/Nonlinear.jl @@ -42,4 +42,8 @@ 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") + 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/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/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 From 7172c9be06575be716b46da20c2d4e95f086f609 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 18:47:07 +0000 Subject: [PATCH 04/12] Add ModelWithQuad and ModelWithOracles layers to Nonlinear --- docs/src/submodules/Nonlinear/reference.md | 12 + src/Nonlinear/Nonlinear.jl | 2 + src/Nonlinear/layers.jl | 548 +++++++++++++++++++++ src/Nonlinear/model.jl | 16 +- src/Nonlinear/qp_block_data.jl | 508 +++++++++++++++++++ test/Nonlinear/test_layers.jl | 286 +++++++++++ 6 files changed, 1369 insertions(+), 3 deletions(-) create mode 100644 src/Nonlinear/layers.jl create mode 100644 src/Nonlinear/qp_block_data.jl create mode 100644 test/Nonlinear/test_layers.jl diff --git a/docs/src/submodules/Nonlinear/reference.md b/docs/src/submodules/Nonlinear/reference.md index 9f10b404f8..e8584a0ea2 100644 --- a/docs/src/submodules/Nonlinear/reference.md +++ b/docs/src/submodules/Nonlinear/reference.md @@ -81,6 +81,18 @@ 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 diff --git a/src/Nonlinear/Nonlinear.jl b/src/Nonlinear/Nonlinear.jl index 32ceb5182f..ee062767ca 100644 --- a/src/Nonlinear/Nonlinear.jl +++ b/src/Nonlinear/Nonlinear.jl @@ -45,5 +45,7 @@ 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/layers.jl b/src/Nonlinear/layers.jl new file mode 100644 index 0000000000..298e940188 --- /dev/null +++ b/src/Nonlinear/layers.jl @@ -0,0 +1,548 @@ +# 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 + + function ModelWithOracles{T}(inner::M) where {T,M} + return new{T,M}( + Tuple{MOI.VectorOfVariables,MOI.VectorNonlinearOracle{T}}[], + inner, + ) + end +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 containing parameters must not be added to this layer: convert them +to [`MOI.ScalarNonlinearFunction`](@ref) so that they are forwarded to the +inner model, where parameters are first-class. +""" +mutable struct ModelWithQuad{T,M} + qp::QPBlockData{T} + inner::M + objective_sink::Symbol # :none, :quad or :inner + + function ModelWithQuad{T}(inner::M) where {T,M} + return new{T,M}(QPBlockData{T}(), inner, :none) + end +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 + +# 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, :Hess), 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) + ) + fmap = Base.Fix1(getindex, index_map) + 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) + 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) + 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) + 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) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + 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 + +""" + 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, :Hess), features) + 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 + for (row, col) in MOI.jacobian_structure(d.inner) + push!(J, (row + row_offset, col)) + 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 + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + 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 + +# 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 + +function _inner_constraint_linearity(inner::MOI.AbstractNLPEvaluator) + linearity = constraint_linearity(inner) + if linearity === nothing + return fill(NONLINEAR, num_constraints(inner)) + end + return linearity +end + +function constraint_linearity(d::EvaluatorWithQuad) + linearity = Linearity[ + ft == _kFunctionTypeScalarQuadratic ? QUADRATIC : LINEAR for + ft in d.model.qp.function_type + ] + return vcat(linearity, _inner_constraint_linearity(d.inner)) +end + +function constraint_linearity(d::EvaluatorWithOracles) + linearity = fill(NONLINEAR, _num_rows(d)) + return vcat(linearity, _inner_constraint_linearity(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 + +_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/model.jl b/src/Nonlinear/model.jl index 8577de2f34..5b3ed65c48 100644 --- a/src/Nonlinear/model.jl +++ b/src/Nonlinear/model.jl @@ -82,10 +82,20 @@ end Return a new nonlinear model appropriate for the given AD `backend`. -The default returns `Model()`. Custom AD backends can override this -to return their own model type. +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())`. """ -model(::AbstractAutomaticDifferentiation) = Model() +function model(::AbstractAutomaticDifferentiation) + return ModelWithQuad(ModelWithOracles(Model())) +end """ 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..a2349013d3 --- /dev/null +++ b/src/Nonlinear/qp_block_data.jl @@ -0,0 +1,508 @@ +# 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, `QPBlockData` does not support parameters: +# functions containing parameters must be converted to +# `MOI.ScalarNonlinearFunction` and routed to the inner nonlinear model, where +# parameters are first-class. + +@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. + +The functions must not contain parameters: convert functions with parameters +to [`MOI.ScalarNonlinearFunction`](@ref) and add them to the inner nonlinear +model instead. +""" +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} + + 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[], + ) + end +end + +function _eval_function( + f::MOI.ScalarQuadraticFunction{T}, + x::AbstractVector{T}, +)::T where {T} + y = f.constant + for term in f.affine_terms + y += term.coefficient * x[term.variable.value] + end + for term in f.quadratic_terms + v1 = x[term.variable_1.value] + v2 = x[term.variable_2.value] + 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}, +)::T where {T} + y = f.constant + for term in f.terms + y += term.coefficient * x[term.variable.value] + end + return y +end + +function _eval_dense_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + x::AbstractVector{T}, +)::Nothing where {T} + for term in f.affine_terms + ∇f[term.variable.value] += term.coefficient + end + for term in f.quadratic_terms + ∇f[term.variable_1.value] += term.coefficient * x[term.variable_2.value] + if term.variable_1 != term.variable_2 + ∇f[term.variable_2.value] += + term.coefficient * x[term.variable_1.value] + end + end + return +end + +function _eval_dense_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarAffineFunction{T}, + x::AbstractVector{T}, +)::Nothing where {T} + for term in f.terms + ∇f[term.variable.value] += term.coefficient + end + return +end + +function _append_sparse_gradient_structure!( + f::MOI.ScalarQuadraticFunction, + J, + row, +) + for term in f.affine_terms + push!(J, (row, term.variable.value)) + end + for term in f.quadratic_terms + push!(J, (row, term.variable_1.value)) + if term.variable_1 != term.variable_2 + push!(J, (row, term.variable_2.value)) + end + end + return +end + +function _append_sparse_gradient_structure!(f::MOI.ScalarAffineFunction, J, row) + for term in f.terms + push!(J, (row, term.variable.value)) + end + return +end + +function _eval_sparse_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + x::AbstractVector{T}, +)::Int where {T} + i = 0 + for term in f.affine_terms + i += 1 + ∇f[i] = term.coefficient + end + for term in f.quadratic_terms + i += 1 + ∇f[i] = term.coefficient * x[term.variable_2.value] + if term.variable_1 != term.variable_2 + i += 1 + ∇f[i] = term.coefficient * x[term.variable_1.value] + end + end + return i +end + +function _eval_sparse_gradient( + ∇f::AbstractVector{T}, + f::MOI.ScalarAffineFunction{T}, + x::AbstractVector{T}, +)::Int where {T} + i = 0 + for term in f.terms + i += 1 + ∇f[i] = term.coefficient + end + return i +end + +function _append_sparse_hessian_structure!(f::MOI.ScalarQuadraticFunction, H) + for term in f.quadratic_terms + push!(H, (term.variable_1.value, term.variable_2.value)) + end + return +end + +_append_sparse_hessian_structure!(::MOI.ScalarAffineFunction, H) = nothing + +function _eval_sparse_hessian( + ∇²f::AbstractVector{T}, + f::MOI.ScalarQuadraticFunction{T}, + σ::T, +)::Int where {T} + i = 0 + for term in f.quadratic_terms + i += 1 + ∇²f[i] = term.coefficient * σ + end + return i +end + +function _eval_sparse_hessian( + ∇²f::AbstractVector{T}, + f::MOI.ScalarAffineFunction{T}, + σ::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) +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) + 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) + 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) + 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) + end + return i +end + +function MOI.hessian_lagrangian_structure(block::QPBlockData) + H = Tuple{Int,Int}[] + _append_sparse_hessian_structure!(block.objective, H) + for constraint in block.constraints + _append_sparse_hessian_structure!(constraint, H) + 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, σ) + for (row, constraint) in enumerate(block.constraints) + ∇²f = view(H, (i+1):length(H)) + i += _eval_sparse_hessian(∇²f, constraint, μ[row]) + end + return i +end diff --git a/test/Nonlinear/test_layers.jl b/test/Nonlinear/test_layers.jl new file mode 100644 index 0000000000..22c1ef29d0 --- /dev/null +++ b/test/Nonlinear/test_layers.jl @@ -0,0 +1,286 @@ +# 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_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() From f68f598aec25403e5929a3c967ab5f1025e5bfed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 19:21:02 +0000 Subject: [PATCH 05/12] Support parameters in QPBlockData and wrap existing solver storage --- src/Nonlinear/layers.jl | 46 +++++++--- src/Nonlinear/qp_block_data.jl | 151 ++++++++++++++++++++++++--------- test/Nonlinear/test_layers.jl | 41 +++++++++ 3 files changed, 186 insertions(+), 52 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index 298e940188..a84b5831b4 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -25,14 +25,23 @@ mutable struct ModelWithOracles{T,M} } inner::M - function ModelWithOracles{T}(inner::M) where {T,M} - return new{T,M}( - Tuple{MOI.VectorOfVariables,MOI.VectorNonlinearOracle{T}}[], - inner, - ) + # 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) """ @@ -46,20 +55,30 @@ A model layer that stores affine and quadratic objectives and constraints in a When wrapped in an [`Evaluator`](@ref), the layer's rows come first, followed by the rows of the inner evaluator. -Functions containing parameters must not be added to this layer: convert them -to [`MOI.ScalarNonlinearFunction`](@ref) so that they are forwarded to the -inner model, where parameters are first-class. +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 - function ModelWithQuad{T}(inner::M) where {T,M} - return new{T,M}(QPBlockData{T}(), inner, :none) + # 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} @@ -197,7 +216,9 @@ function MOI.initialize( index_map = Dict{MOI.VariableIndex,MOI.VariableIndex}( x => MOI.VariableIndex(i) for (i, x) in enumerate(d.ordered_variables) ) - fmap = Base.Fix1(getindex, index_map) + # 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) @@ -210,6 +231,9 @@ function MOI.initialize( 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 diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index a2349013d3..5196c4a3c5 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -5,10 +5,11 @@ # This file is adapted from `Ipopt.jl/ext/IpoptMathOptInterfaceExt/utils.jl`. # -# Unlike the Ipopt version, `QPBlockData` does not support parameters: -# functions containing parameters must be converted to -# `MOI.ScalarNonlinearFunction` and routed to the inner nonlinear model, where -# parameters are first-class. +# 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, @@ -68,9 +69,13 @@ following the [`MOI.AbstractNLPEvaluator`](@ref) callback conventions. This is the storage behind [`ModelWithQuad`](@ref); it is not typically used directly. -The functions must not contain parameters: convert functions with parameters -to [`MOI.ScalarNonlinearFunction`](@ref) and add them to the inner nonlinear -model instead. +## 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}} @@ -83,6 +88,7 @@ mutable struct QPBlockData{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( @@ -94,21 +100,29 @@ mutable struct QPBlockData{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 * x[term.variable.value] + y += term.coefficient * _value(term.variable, x, p) end for term in f.quadratic_terms - v1 = x[term.variable_1.value] - v2 = x[term.variable_2.value] + 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 @@ -121,10 +135,11 @@ 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 * x[term.variable.value] + y += term.coefficient * _value(term.variable, x, p) end return y end @@ -133,15 +148,22 @@ 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 - ∇f[term.variable.value] += term.coefficient + if !_is_parameter(term.variable, p) + ∇f[term.variable.value] += term.coefficient + end end for term in f.quadratic_terms - ∇f[term.variable_1.value] += term.coefficient * x[term.variable_2.value] - if term.variable_1 != term.variable_2 - ∇f[term.variable_2.value] += - term.coefficient * x[term.variable_1.value] + 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 @@ -151,9 +173,12 @@ 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 - ∇f[term.variable.value] += term.coefficient + if !_is_parameter(term.variable, p) + ∇f[term.variable.value] += term.coefficient + end end return end @@ -162,22 +187,35 @@ function _append_sparse_gradient_structure!( f::MOI.ScalarQuadraticFunction, J, row, + p::Dict, ) for term in f.affine_terms - push!(J, (row, term.variable.value)) + if !_is_parameter(term.variable, p) + push!(J, (row, term.variable.value)) + end end for term in f.quadratic_terms - push!(J, (row, term.variable_1.value)) - if term.variable_1 != term.variable_2 + 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) +function _append_sparse_gradient_structure!( + f::MOI.ScalarAffineFunction, + J, + row, + p::Dict, +) for term in f.terms - push!(J, (row, term.variable.value)) + if !_is_parameter(term.variable, p) + push!(J, (row, term.variable.value)) + end end return end @@ -186,18 +224,26 @@ 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 - i += 1 - ∇f[i] = term.coefficient + if !_is_parameter(term.variable, p) + i += 1 + ∇f[i] = term.coefficient + end end for term in f.quadratic_terms - i += 1 - ∇f[i] = term.coefficient * x[term.variable_2.value] - if term.variable_1 != term.variable_2 + 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 * x[term.variable_1.value] + ∇f[i] = term.coefficient * v end end return i @@ -207,31 +253,53 @@ 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 - i += 1 - ∇f[i] = term.coefficient + 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) +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 -_append_sparse_hessian_structure!(::MOI.ScalarAffineFunction, H) = nothing +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 @@ -242,6 +310,7 @@ function _eval_sparse_hessian( ∇²f::AbstractVector{T}, f::MOI.ScalarAffineFunction{T}, σ::T, + p::Dict{Int64,T}, )::Int where {T} return 0 end @@ -436,7 +505,7 @@ function MOI.eval_objective( block::QPBlockData{T}, x::AbstractVector{T}, ) where {T} - return _eval_function(block.objective, x) + return _eval_function(block.objective, x, block.parameters) end function MOI.eval_objective_gradient( @@ -445,7 +514,7 @@ function MOI.eval_objective_gradient( x::AbstractVector{T}, ) where {T} ∇f .= zero(T) - _eval_dense_gradient(∇f, block.objective, x) + _eval_dense_gradient(∇f, block.objective, x, block.parameters) return end @@ -455,7 +524,7 @@ function MOI.eval_constraint( x::AbstractVector{T}, ) where {T} for (i, constraint) in enumerate(block.constraints) - g[i] = _eval_function(constraint, x) + g[i] = _eval_function(constraint, x, block.parameters) end return end @@ -463,7 +532,7 @@ 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) + _append_sparse_gradient_structure!(constraint, J, row, block.parameters) end return J end @@ -477,16 +546,16 @@ function MOI.eval_constraint_jacobian( i = 0 for constraint in block.constraints ∇f = view(J, (i+1):length(J)) - i += _eval_sparse_gradient(∇f, constraint, x) + 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) + _append_sparse_hessian_structure!(block.objective, H, block.parameters) for constraint in block.constraints - _append_sparse_hessian_structure!(constraint, H) + _append_sparse_hessian_structure!(constraint, H, block.parameters) end return H end @@ -499,10 +568,10 @@ function MOI.eval_hessian_lagrangian( σ::T, μ::AbstractVector{T}, ) where {T} - i = _eval_sparse_hessian(H, block.objective, σ) + 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]) + i += _eval_sparse_hessian(∇²f, constraint, μ[row], block.parameters) end return i end diff --git a/test/Nonlinear/test_layers.jl b/test/Nonlinear/test_layers.jl index 22c1ef29d0..09f8aa2fe9 100644 --- a/test/Nonlinear/test_layers.jl +++ b/test/Nonlinear/test_layers.jl @@ -261,6 +261,47 @@ function test_forwarding_through_layers() 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()) From 06e9194e7b51045bb17c307b25d2c3adc71c3c2c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 19:25:28 +0000 Subject: [PATCH 06/12] Return nothing from layer constraint_linearity when inner is unknown --- src/Nonlinear/layers.jl | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index a84b5831b4..66753612af 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -487,25 +487,37 @@ function num_constraints(d::EvaluatorWithOracles) return _num_rows(d) + num_constraints(d.inner) 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 fill(NONLINEAR, num_constraints(inner)) + if linearity !== nothing + return linearity + elseif !applicable(num_constraints, inner) + return nothing end - return linearity + return fill(NONLINEAR, num_constraints(inner)) 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_constraint_linearity(d.inner)) + return vcat(linearity, inner) end function constraint_linearity(d::EvaluatorWithOracles) - linearity = fill(NONLINEAR, _num_rows(d)) - return vcat(linearity, _inner_constraint_linearity(d.inner)) + 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) From d278d81b636b9278de19a855e3ad629d8b7c4609 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 19:37:19 +0000 Subject: [PATCH 07/12] Add Jacobian and Hessian product callbacks to the Nonlinear layers From MadNLP's copy (parameter-aware, accumulating into the output) --- src/Nonlinear/layers.jl | 111 +++++++++++++++++++++- src/Nonlinear/qp_block_data.jl | 162 +++++++++++++++++++++++++++++++++ test/Nonlinear/test_layers.jl | 72 +++++++++++++++ 3 files changed, 343 insertions(+), 2 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index 66753612af..aef595f481 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -206,7 +206,7 @@ end function MOI.features_available(d::EvaluatorWithQuad) features = MOI.features_available(d.inner) - return filter(f -> f in (:Grad, :Jac, :Hess), features) + return filter(f -> f in (:Grad, :Jac, :JacVec, :Hess, :HessVec), features) end function MOI.initialize( @@ -305,6 +305,64 @@ function MOI.eval_hessian_lagrangian(d::EvaluatorWithQuad, H, x, σ, μ) 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, @@ -364,7 +422,12 @@ end function MOI.features_available(d::EvaluatorWithOracles) features = MOI.features_available(d.inner) - features = filter(f -> f in (:Grad, :Jac, :Hess), features) + 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 @@ -477,6 +540,50 @@ function MOI.eval_hessian_lagrangian(d::EvaluatorWithOracles, H, x, σ, μ) return end +# The product callbacks are only supported when the layer has no oracle +# constraints (see `MOI.features_available`), in which case the layer is a +# pass-through. + +function _check_no_oracles(d::EvaluatorWithOracles) + if !isempty(d.model.constraints) + error( + "The product callbacks are not supported in the presence of " * + "VectorNonlinearOracle constraints.", + ) + end + return +end + +function MOI.eval_constraint_jacobian_product(d::EvaluatorWithOracles, y, x, w) + _check_no_oracles(d) + MOI.eval_constraint_jacobian_product(d.inner, y, x, w) + return +end + +function MOI.eval_constraint_jacobian_transpose_product( + d::EvaluatorWithOracles, + y, + x, + w, +) + _check_no_oracles(d) + MOI.eval_constraint_jacobian_transpose_product(d.inner, y, x, w) + return +end + +function MOI.eval_hessian_lagrangian_product( + d::EvaluatorWithOracles, + H, + x, + v, + σ, + μ, +) + _check_no_oracles(d) + MOI.eval_hessian_lagrangian_product(d.inner, H, x, v, σ, μ) + return +end + # Linearity and row queries function num_constraints(d::EvaluatorWithQuad) diff --git a/src/Nonlinear/qp_block_data.jl b/src/Nonlinear/qp_block_data.jl index 5196c4a3c5..1b16647d47 100644 --- a/src/Nonlinear/qp_block_data.jl +++ b/src/Nonlinear/qp_block_data.jl @@ -575,3 +575,165 @@ function MOI.eval_hessian_lagrangian( 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_layers.jl b/test/Nonlinear/test_layers.jl index 09f8aa2fe9..b1cd39195a 100644 --- a/test/Nonlinear/test_layers.jl +++ b/test/Nonlinear/test_layers.jl @@ -261,6 +261,78 @@ function test_forwarding_through_layers() 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()) From 13a7687d0d5653707006bdf4674bee94b86acc2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 19:45:42 +0000 Subject: [PATCH 08/12] Recurse gracefully in _try_num_constraints through layer stacks The applicable() guard in _inner_constraint_linearity did not recurse: num_constraints on a layer is applicable but throws when the innermost evaluator (for example, a legacy NLPBlock evaluator) does not implement the query. Add _try_num_constraints, which returns nothing in that case at any depth of the stack. --- src/Nonlinear/layers.jl | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index aef595f481..99c18d2026 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -594,6 +594,25 @@ 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. @@ -601,10 +620,12 @@ function _inner_constraint_linearity(inner::MOI.AbstractNLPEvaluator) linearity = constraint_linearity(inner) if linearity !== nothing return linearity - elseif !applicable(num_constraints, inner) + end + n = _try_num_constraints(inner) + if n === nothing return nothing end - return fill(NONLINEAR, num_constraints(inner)) + return fill(NONLINEAR, n) end function constraint_linearity(d::EvaluatorWithQuad) From 88111a0fbe6242cdd6e60ae81d1fcc70d104ae0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 19:52:38 +0000 Subject: [PATCH 09/12] Gate inner structure queries on the inner evaluator's features An evaluator is only required to implement jacobian_structure and hessian_lagrangian_structure when it supports :Jac and :Hess. Solvers such as Ipopt query the structures eagerly even for models whose legacy NLPBlock evaluator supports neither (for example, an objective-only evaluator with feature :Grad), so the layers must not call into the inner evaluator in that case. --- src/Nonlinear/layers.jl | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index 99c18d2026..d5e664675b 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -272,8 +272,13 @@ end function MOI.jacobian_structure(d::EvaluatorWithQuad) J = MOI.jacobian_structure(d.qp) offset = length(d.qp) - for (row, col) in MOI.jacobian_structure(d.inner) - push!(J, (row + offset, col)) + # 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 @@ -286,7 +291,9 @@ end function MOI.hessian_lagrangian_structure(d::EvaluatorWithQuad) H = MOI.hessian_lagrangian_structure(d.qp) - append!(H, MOI.hessian_lagrangian_structure(d.inner)) + if :Hess in MOI.features_available(d.inner) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + end return H end @@ -490,8 +497,10 @@ function MOI.jacobian_structure(d::EvaluatorWithOracles) end row_offset += s.output_dimension end - for (row, col) in MOI.jacobian_structure(d.inner) - push!(J, (row + row_offset, col)) + 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 @@ -516,7 +525,9 @@ function MOI.hessian_lagrangian_structure(d::EvaluatorWithOracles) push!(H, (columns[i], columns[j])) end end - append!(H, MOI.hessian_lagrangian_structure(d.inner)) + if :Hess in MOI.features_available(d.inner) + append!(H, MOI.hessian_lagrangian_structure(d.inner)) + end return H end From 96d49299b7bc95af38423d7a9e88669b2fc21563 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 20:18:48 +0000 Subject: [PATCH 10/12] Add Nonlinear.exploits_structure trait and _has_objective fallback exploits_structure(backend) returns true when the model returned by model(backend) natively accepts ScalarAffine/ScalarQuadratic functions (no ModelWithQuad layer), so consumers such as NLPModelsJuMP know to pass all functions unparsed. Defined in MOI.Nonlinear because a trait spanning two package extensions (an AD backend and a consumer) has no reliable home in either: extension load order is unspecified. _has_objective falls back to true for unknown evaluators, which is conservative: the solver then calls eval_objective, and evaluators without an objective return zero. --- src/Nonlinear/layers.jl | 4 ++++ src/Nonlinear/model.jl | 17 +++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index d5e664675b..78d54cfe1e 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -707,6 +707,10 @@ function constraint_bounds(d::EvaluatorWithOracles) 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) diff --git a/src/Nonlinear/model.jl b/src/Nonlinear/model.jl index 5b3ed65c48..68cd2110db 100644 --- a/src/Nonlinear/model.jl +++ b/src/Nonlinear/model.jl @@ -97,6 +97,23 @@ 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 From 3737ff0ac976c51aa2d3ed0044ee8b46131508ed Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 20:32:05 +0000 Subject: [PATCH 11/12] Answer product callbacks with oracles via materialized derivatives Solvers such as MadNLP call the Jacobian-transpose product for dual computations even when :JacVec is not advertised, so erroring in the presence of oracle constraints is not an option. The oracle layer now materializes each oracle's Jacobian (or Hessian) to answer the product callbacks, which can be slow on large oracles; the :JacVec/:HessVec features remain removed from features_available so that solvers do not rely on them for performance. Also forward MOI.is_valid for constraint indices and MOI.ListOfSupportedNonlinearOperators through the model layers (needed by NLopt), and add the Nonlinear.exploits_structure trait. --- src/Nonlinear/layers.jl | 98 +++++++++++++++++++++++++++++++++-------- 1 file changed, 79 insertions(+), 19 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index 78d54cfe1e..bb0f006ff4 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -102,6 +102,12 @@ function register_operator( return register_operator(model.inner, op, nargs, f...) end +MOI.is_valid(model::_LayerModel, index::ConstraintIndex) = + MOI.is_valid(model.inner, index) + +MOI.get(model::_LayerModel, attr::MOI.ListOfSupportedNonlinearOperators) = + MOI.get(model.inner, attr) + # ModelWithQuad function set_objective( @@ -551,23 +557,31 @@ function MOI.eval_hessian_lagrangian(d::EvaluatorWithOracles, H, x, σ, μ) return end -# The product callbacks are only supported when the layer has no oracle -# constraints (see `MOI.features_available`), in which case the layer is a -# pass-through. - -function _check_no_oracles(d::EvaluatorWithOracles) - if !isempty(d.model.constraints) - error( - "The product callbacks are not supported in the presence of " * - "VectorNonlinearOracle constraints.", - ) - end - 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) - _check_no_oracles(d) - MOI.eval_constraint_jacobian_product(d.inner, 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 @@ -577,8 +591,25 @@ function MOI.eval_constraint_jacobian_transpose_product( x, w, ) - _check_no_oracles(d) - MOI.eval_constraint_jacobian_transpose_product(d.inner, 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 @@ -590,8 +621,37 @@ function MOI.eval_hessian_lagrangian_product( σ, μ, ) - _check_no_oracles(d) - MOI.eval_hessian_lagrangian_product(d.inner, 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 From e7048ed714da84d3932a43742172d634d2d81d98 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Wed, 12 Aug 2026 20:40:24 +0000 Subject: [PATCH 12/12] Formatting --- src/Nonlinear/layers.jl | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/Nonlinear/layers.jl b/src/Nonlinear/layers.jl index bb0f006ff4..e213eb9b96 100644 --- a/src/Nonlinear/layers.jl +++ b/src/Nonlinear/layers.jl @@ -102,11 +102,16 @@ function register_operator( return register_operator(model.inner, op, nargs, f...) end -MOI.is_valid(model::_LayerModel, index::ConstraintIndex) = - MOI.is_valid(model.inner, index) +function MOI.is_valid(model::_LayerModel, index::ConstraintIndex) + return MOI.is_valid(model.inner, index) +end -MOI.get(model::_LayerModel, attr::MOI.ListOfSupportedNonlinearOperators) = - MOI.get(model.inner, attr) +function MOI.get( + model::_LayerModel, + attr::MOI.ListOfSupportedNonlinearOperators, +) + return MOI.get(model.inner, attr) +end # ModelWithQuad