Remove need for the MOI fork

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,6 +10,7 @@ import ForwardDiff
 import MathOptInterface as MOI
 const Nonlinear = MOI.Nonlinear
 import SparseArrays
+import OrderedCollections: OrderedDict
 
 """
     Mode() <: AbstractAutomaticDifferentiation

src/MOI_Nonlinear_fork.jl

@@ -1,19 +1,118 @@
 # Inspired by MathOptInterface/src/Nonlinear/parse_expression.jl
 
+const DEFAULT_MULTIVARIATE_OPERATORS = [
+    :+,
+    :-,
+    :*,
+    :^,
+    :/,
+    :ifelse,
+    :atan,
+    :min,
+    :max,
+    :vect,
+    :dot,
+    :hcat,
+    :vcat,
+    :norm,
+    :sum,
+    :row,
+]
+
+struct OperatorRegistry
+    # NODE_CALL_UNIVARIATE
+    univariate_operators::Vector{Symbol}
+    univariate_operator_to_id::Dict{Symbol,Int}
+    univariate_user_operator_start::Int
+    registered_univariate_operators::Vector{MOI.Nonlinear._UnivariateOperator}
+    # NODE_CALL_MULTIVARIATE
+    multivariate_operators::Vector{Symbol}
+    multivariate_operator_to_id::Dict{Symbol,Int}
+    multivariate_user_operator_start::Int
+    registered_multivariate_operators::Vector{
+        MOI.Nonlinear._MultivariateOperator,
+    }
+    # NODE_LOGIC
+    logic_operators::Vector{Symbol}
+    logic_operator_to_id::Dict{Symbol,Int}
+    # NODE_COMPARISON
+    comparison_operators::Vector{Symbol}
+    comparison_operator_to_id::Dict{Symbol,Int}
+    function OperatorRegistry()
+        univariate_operators = copy(DEFAULT_UNIVARIATE_OPERATORS)
+        multivariate_operators = copy(DEFAULT_MULTIVARIATE_OPERATORS)
+        logic_operators = [:&&, :||]
+        comparison_operators = [:<=, :(==), :>=, :<, :>]
+        return new(
+            # NODE_CALL_UNIVARIATE
+            univariate_operators,
+            Dict{Symbol,Int}(
+                op => i for (i, op) in enumerate(univariate_operators)
+            ),
+            length(univariate_operators),
+            _UnivariateOperator[],
+            # NODE_CALL
+            multivariate_operators,
+            Dict{Symbol,Int}(
+                op => i for (i, op) in enumerate(multivariate_operators)
+            ),
+            length(multivariate_operators),
+            _MultivariateOperator[],
+            # NODE_LOGIC
+            logic_operators,
+            Dict{Symbol,Int}(op => i for (i, op) in enumerate(logic_operators)),
+            # NODE_COMPARISON
+            comparison_operators,
+            Dict{Symbol,Int}(
+                op => i for (i, op) in enumerate(comparison_operators)
+            ),
+        )
+    end
+end
+
+"""
+    Model()
+
+The core datastructure for representing a nonlinear optimization problem.
+
+It has the following fields:
+ * `objective::Union{Nothing,Expression}` : holds the nonlinear objective
+   function, if one exists, otherwise `nothing`.
+ * `expressions::Vector{Expression}` : a vector of expressions in the model.
+ * `constraints::OrderedDict{ConstraintIndex,Constraint}` : a map from
+   [`ConstraintIndex`](@ref) to the corresponding [`Constraint`](@ref). An
+   `OrderedDict` is used instead of a `Vector` to support constraint deletion.
+ * `parameters::Vector{Float64}` : holds the current values of the parameters.
+ * `operators::OperatorRegistry` : stores the operators used in the model.
+"""
+mutable struct Model
+    objective::Union{Nothing,MOI.Nonlinear.Expression}
+    expressions::Vector{MOI.Nonlinear.Expression}
+    constraints::OrderedDict{
+        MOI.Nonlinear.ConstraintIndex,
+        MOI.Nonlinear.Constraint,
+    }
+    parameters::Vector{Float64}
+    operators::OperatorRegistry
+    # This is a private field, used only to increment the ConstraintIndex.
+    last_constraint_index::Int64
+    function Model()
+        model = MOI.Nonlinear.Model()
+        ops = [:vect, :dot, :hcat, :vcat, :norm, :sum, :row]
+        start = length(model.operators.multivariate_operators)
+        append!(model.operators.multivariate_operators, ops)
+        for (i, op) in enumerate(ops)
+            model.operators.multivariate_operator_to_id[op] = start + i
+        end
+        return model
+    end
+end
+
 function set_objective(model::MOI.Nonlinear.Model, obj)
     model.objective = parse_expression(model, obj)
     return
 end
 
-function Model()
-    model = MOI.Nonlinear.Model()
-    append!(
-        model.operators.multivariate_operators,
-        [:vect, :dot, :hcat, :vcat, :norm, :sum, :row],
-    )
-    return model
-end
-
 function parse_expression(data::MOI.Nonlinear.Model, input)
     expr = MOI.Nonlinear.Expression()
     parse_expression(data, expr, input, -1)

src/reverse_mode.jl

@@ -430,9 +430,8 @@ function _reverse_eval(f::_SubexpressionStorage)
         node = f.nodes[k]
         children_indices = SparseArrays.nzrange(f.adj, k)
         if node.type == MOI.Nonlinear.NODE_CALL_MULTIVARIATE
-            if node.index in
-               eachindex(MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS)
-                op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index]
+            if node.index in eachindex(DEFAULT_MULTIVARIATE_OPERATORS)
+                op = DEFAULT_MULTIVARIATE_OPERATORS[node.index]
                 if op == :vect
                     @assert _eachindex(f.sizes, k) ==
                             eachindex(children_indices)

src/sizes.jl

@@ -163,14 +163,11 @@ function _infer_sizes(
         children_indices = SparseArrays.nzrange(adj, k)
         N = length(children_indices)
         if node.type == Nonlinear.NODE_CALL_MULTIVARIATE
-            if !(
-                node.index in
-                eachindex(MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS)
-            )
+            if !(node.index in eachindex(DEFAULT_MULTIVARIATE_OPERATORS))
                 # TODO user-defined operators
                 continue
             end
-            op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index]
+            op = DEFAULT_MULTIVARIATE_OPERATORS[node.index]
             if op == :vect
                 _assert_scalar_children(
                     sizes,

test/Project.toml

@@ -4,3 +4,4 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"