Remove need for the MOI fork

.github/workflows/ci.yml

@@ -25,13 +25,6 @@ jobs:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v1
-      - name: MOI
-        shell: julia --project=@. {0}
-        run: |
-          using Pkg
-          Pkg.add([
-              PackageSpec(name="MathOptInterface", rev="bl/arraydiff"),
-          ])
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         with:

Project.toml

@@ -9,6 +9,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 
 [compat]
 DataStructures = "0.18, 0.19"

src/ArrayDiff.jl

@@ -10,25 +10,10 @@ import ForwardDiff
 import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 import SparseArrays
+import OrderedCollections: OrderedDict
 
-"""
-    Mode() <: AbstractAutomaticDifferentiation
-
-Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
-"""
 struct Mode <: MOI.Nonlinear.AbstractAutomaticDifferentiation end
 
-function MOI.Nonlinear.Evaluator(
-    model::MOI.Nonlinear.Model,
-    ::Mode,
-    ordered_variables::Vector{MOI.VariableIndex},
-)
-    return MOI.Nonlinear.Evaluator(
-        model,
-        NLPEvaluator(model, ordered_variables),
-    )
-end
-
 # Override basic math functions to return NaN instead of throwing errors.
 # This is what NLP solvers expect, and sometimes the results aren't needed
 # anyway, because the code may compute derivatives wrt constants.
@@ -57,5 +42,23 @@ include("utils.jl")
 include("reverse_mode.jl")
 include("forward_over_reverse.jl")
 include("mathoptinterface_api.jl")
+include("operators.jl")
+include("model.jl")
+include("parse.jl")
+include("evaluator.jl")
+
+"""
+    Mode() <: AbstractAutomaticDifferentiation
+
+Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
+"""
+
+function Evaluator(
+    model::ArrayDiff.Model,
+    ::Mode,
+    ordered_variables::Vector{MOI.VariableIndex},
+)
+    return Evaluator(model, NLPEvaluator(model, ordered_variables))
+end
 
 end  # module

src/evaluator.jl

@@ -0,0 +1,156 @@
+# Largely inspired by MathOptInterface/src/Nonlinear/parse.jl
+# Most functions have been copy-pasted and slightly modified to adapt to small changes in OperatorRegistry and Model.
+
+function MOI.initialize(evaluator::Evaluator, features::Vector{Symbol})
+    start_time = time()
+    empty!(evaluator.ordered_constraints)
+    evaluator.eval_objective_timer = 0.0
+    evaluator.eval_objective_gradient_timer = 0.0
+    evaluator.eval_constraint_timer = 0.0
+    evaluator.eval_constraint_gradient_timer = 0.0
+    evaluator.eval_constraint_jacobian_timer = 0.0
+    evaluator.eval_hessian_objective_timer = 0.0
+    evaluator.eval_hessian_constraint_timer = 0.0
+    evaluator.eval_hessian_lagrangian_timer = 0.0
+    append!(evaluator.ordered_constraints, keys(evaluator.model.constraints))
+    # Every backend supports :ExprGraph, so don't forward it.
+    filter!(f -> f != :ExprGraph, features)
+    if evaluator.backend !== nothing
+        MOI.initialize(evaluator.backend, features)
+    elseif !isempty(features)
+        @assert evaluator.backend === nothing  # ==> ExprGraphOnly used
+        error(
+            "Unable to initialize `Nonlinear.Evaluator` because the " *
+            "following features are not supported: $features",
+        )
+    end
+    evaluator.initialize_timer = time() - start_time
+    return
+end
+
+function MOI.eval_objective(evaluator::Evaluator, x)
+    start = time()
+    obj = MOI.eval_objective(evaluator.backend, x)
+    evaluator.eval_objective_timer += time() - start
+    return obj
+end
+
+function MOI.eval_objective_gradient(evaluator::Evaluator, g, x)
+    start = time()
+    MOI.eval_objective_gradient(evaluator.backend, g, x)
+    evaluator.eval_objective_gradient_timer += time() - start
+    return
+end
+
+function MOI.eval_constraint_gradient(evaluator::Evaluator, ∇g, x, i)
+    start = time()
+    MOI.eval_constraint_gradient(evaluator.backend, ∇g, x, i)
+    evaluator.eval_constraint_gradient_timer += time() - start
+    return
+end
+
+function MOI.constraint_gradient_structure(evaluator::Evaluator, i)
+    return MOI.constraint_gradient_structure(evaluator.backend, i)
+end
+
+function MOI.eval_constraint(evaluator::Evaluator, g, x)
+    start = time()
+    MOI.eval_constraint(evaluator.backend, g, x)
+    evaluator.eval_constraint_timer += time() - start
+    return
+end
+
+function MOI.jacobian_structure(evaluator::Evaluator)
+    return MOI.jacobian_structure(evaluator.backend)
+end
+
+function MOI.eval_constraint_jacobian(evaluator::Evaluator, J, x)
+    start = time()
+    MOI.eval_constraint_jacobian(evaluator.backend, J, x)
+    evaluator.eval_constraint_jacobian_timer += time() - start
+    return
+end
+
+function MOI.hessian_objective_structure(evaluator::Evaluator)
+    return MOI.hessian_objective_structure(evaluator.backend)
+end
+
+function MOI.hessian_constraint_structure(evaluator::Evaluator, i)
+    return MOI.hessian_constraint_structure(evaluator.backend, i)
+end
+
+function MOI.hessian_lagrangian_structure(evaluator::Evaluator)
+    return MOI.hessian_lagrangian_structure(evaluator.backend)
+end
+
+function MOI.eval_hessian_objective(evaluator::Evaluator, H, x)
+    start = time()
+    MOI.eval_hessian_objective(evaluator.backend, H, x)
+    evaluator.eval_hessian_objective_timer += time() - start
+    return
+end
+
+function MOI.eval_hessian_constraint(evaluator::Evaluator, H, x, i)
+    start = time()
+    MOI.eval_hessian_constraint(evaluator.backend, H, x, i)
+    evaluator.eval_hessian_constraint_timer += time() - start
+    return
+end
+
+function MOI.eval_hessian_lagrangian(evaluator::Evaluator, H, x, σ, μ)
+    start = time()
+    MOI.eval_hessian_lagrangian(evaluator.backend, H, x, σ, μ)
+    evaluator.eval_hessian_lagrangian_timer += time() - start
+    return
+end
+
+function MOI.eval_constraint_jacobian_product(evaluator::Evaluator, y, x, w)
+    start = time()
+    MOI.eval_constraint_jacobian_product(evaluator.backend, y, x, w)
+    evaluator.eval_constraint_jacobian_timer += time() - start
+    return
+end
+
+function MOI.eval_constraint_jacobian_transpose_product(
+    evaluator::Evaluator,
+    y,
+    x,
+    w,
+)
+    start = time()
+    MOI.eval_constraint_jacobian_transpose_product(evaluator.backend, y, x, w)
+    evaluator.eval_constraint_jacobian_timer += time() - start
+    return
+end
+
+function MOI.eval_hessian_lagrangian_product(
+    evaluator::Evaluator,
+    H,
+    x,
+    v,
+    σ,
+    μ,
+)
+    start = time()
+    MOI.eval_hessian_lagrangian_product(evaluator.backend, H, x, v, σ, μ)
+    evaluator.eval_hessian_lagrangian_timer += time() - start
+    return
+end
+
+function eval_univariate_hessian(
+    registry::OperatorRegistry,
+    id::Integer,
+    x::T,
+) where {T}
+    if id <= registry.univariate_user_operator_start
+        ret = Nonlinear._eval_univariate_2nd_deriv(id, x)
+        if ret === nothing
+            op = registry.univariate_operators[id]
+            error("Hessian is not defined for operator $op")
+        end
+        return ret::T
+    end
+    offset = id - registry.univariate_user_operator_start
+    operator = registry.registered_univariate_operators[offset]
+    return eval_univariate_hessian(operator, x)
+end

src/forward_over_reverse.jl

@@ -324,7 +324,7 @@ function _forward_eval_ϵ(
                         d.user_output_buffer,
                         n_children,
                     )
-                    has_hessian = Nonlinear.eval_multivariate_hessian(
+                    has_hessian = eval_multivariate_hessian(
                         d.data.operators,
                         d.data.operators.multivariate_operators[node.index],
                         H,
@@ -351,7 +351,7 @@ function _forward_eval_ϵ(
                 end
             elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE
                 @inbounds child_idx = children_arr[ex.adj.colptr[k]]
-                f′′ = Nonlinear.eval_univariate_hessian(
+                f′′ = eval_univariate_hessian(
                     d.data.operators,
                     node.index,
                     ex.forward_storage[child_idx],

src/model.jl

@@ -0,0 +1,85 @@
+# Largely inspired by MathOptInterface/src/Nonlinear/model.jl
+# Most functions have been copy-pasted and slightly modified to adapt to small changes in OperatorRegistry and Model.
+
+function set_objective(model::Model, obj)
+    model.objective = parse_expression(model, obj)
+    return
+end
+
+function add_constraint(
+    model::Model,
+    func,
+    set::Union{
+        MOI.GreaterThan{Float64},
+        MOI.LessThan{Float64},
+        MOI.Interval{Float64},
+        MOI.EqualTo{Float64},
+    },
+)
+    f = parse_expression(model, func)
+    model.last_constraint_index += 1
+    index = MOI.Nonlinear.ConstraintIndex(model.last_constraint_index)
+    model.constraints[index] = MOI.Nonlinear.Constraint(f, set)
+    return index
+end
+
+function add_parameter(model::Model, value::Float64)
+    push!(model.parameters, value)
+    return MOI.Nonlinear.ParameterIndex(length(model.parameters))
+end
+
+function add_expression(model::Model, expr)
+    push!(model.expressions, parse_expression(model, expr))
+    return Nonlinear.ExpressionIndex(length(model.expressions))
+end
+
+function Base.getindex(model::Model, index::Nonlinear.ExpressionIndex)
+    return model.expressions[index.value]
+end
+
+function register_operator(model::Model, op::Symbol, nargs::Int, f::Function...)
+    return register_operator(model.operators, op, nargs, f...)
+end
+
+function register_operator(
+    registry::OperatorRegistry,
+    op::Symbol,
+    nargs::Int,
+    f::Function...,
+)
+    if nargs == 1
+        if haskey(registry.univariate_operator_to_id, op)
+            error("Operator $op is already registered.")
+        elseif haskey(registry.multivariate_operator_to_id, op)
+            error("Operator $op is already registered.")
+        end
+        operator = Nonlinear._UnivariateOperator(op, f...)
+        push!(registry.univariate_operators, op)
+        push!(registry.registered_univariate_operators, operator)
+        registry.univariate_operator_to_id[op] =
+            length(registry.univariate_operators)
+    else
+        if haskey(registry.multivariate_operator_to_id, op)
+            error("Operator $op is already registered.")
+        elseif haskey(registry.univariate_operator_to_id, op)
+            error("Operator $op is already registered.")
+        end
+        operator = Nonlinear._MultivariateOperator{nargs}(op, f...)
+        push!(registry.multivariate_operators, op)
+        push!(registry.registered_multivariate_operators, operator)
+        registry.multivariate_operator_to_id[op] =
+            length(registry.multivariate_operators)
+    end
+    return
+end
+
+function features_available(evaluator::Evaluator)
+    features = Symbol[]
+    if evaluator.backend !== nothing
+        append!(features, MOI.features_available(evaluator.backend))
+    end
+    if !(:ExprGraph in features)
+        push!(features, :ExprGraph)
+    end
+    return features
+end