AD interface compatibile with both ArrayDiff and ExaModels

docs/src/submodules/Nonlinear/reference.md

@@ -74,6 +74,27 @@ Nonlinear.SparseReverseMode
 Nonlinear.SymbolicMode
 ```
 
+## Linearity
+
+```@docs
+Nonlinear.Linearity
+Nonlinear.num_constraints
+Nonlinear.constraint_linearity
+Nonlinear.objective_linearity
+Nonlinear.constraint_bounds
+```
+
+## Model layers
+
+```@docs
+Nonlinear.model
+Nonlinear.ModelWithQuad
+Nonlinear.ModelWithOracles
+Nonlinear.EvaluatorWithQuad
+Nonlinear.EvaluatorWithOracles
+Nonlinear.QPBlockData
+```
+
 ## Data-structure
 
 ```@docs

src/Nonlinear/Nonlinear.jl

@@ -42,4 +42,10 @@ include("evaluator.jl")
 include("ReverseAD/ReverseAD.jl")
 include("SymbolicAD/SymbolicAD.jl")
 
+# Depends on the `Linearity` enum defined in `ReverseAD`, so must be included
+# after it.
+include("linearity.jl")
+include("qp_block_data.jl")
+include("layers.jl")
+
 end  # module

src/Nonlinear/ReverseAD/graph_tools.jl

@@ -33,22 +33,38 @@
     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 @@
             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 @@
                 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

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,

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