Merge pull request #113 from DimitriAlston/master

Refactor of JuMP's nonlinear API

.github/workflows/documentation.yml

@@ -12,8 +12,8 @@ jobs:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@latest
         with:
-          # Build documentation on Julia 1.3
-          version: '1.5'
+          # Build documentation on Julia 1.6
+          version: '1.6'
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy

Project.toml

@@ -1,7 +1,7 @@
 name = "EAGO"
 uuid = "bb8be931-2a91-5aca-9f87-79e1cb69959a"
 authors = ["Matthew Wilhelm <[email protected]>"]
-version = "0.7.3"
+version = "0.8.0"
 
 [deps]
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
@@ -36,7 +36,7 @@ ForwardDiff = "~0.10"
 IntervalArithmetic = "~0.20"
 IntervalContractors = "~0.4"
 Ipopt = "~1"
-JuMP = "1.0.0 - 1.1.1"
+JuMP = "1.11"
 MINLPTests = "0.5.2"
 MathOptInterface = "~1"
 McCormick = "0.13"

docs/Project.toml

@@ -10,4 +10,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 Documenter = "^0.26"
-JuMP = "1.0.0 - 1.1.1"
+JuMP = "1.11"

src/EAGO.jl

@@ -13,19 +13,14 @@ __precompile__()
 
 module EAGO
 
+    using MathOptInterface
     import MathOptInterface
-
+    import MathOptInterface.Nonlinear: DEFAULT_MULTIVARIATE_OPERATORS, DEFAULT_UNIVARIATE_OPERATORS
     using Reexport, Requires, Cassette, IntervalArithmetic, DocStringExtensions,
           FastRounding, SpecialFunctions, Ipopt, Cbc, Printf, PrettyTables
 
     using JuMP
     import JuMP
-    import JuMP._Derivatives: operators, NodeData
-    using JuMP._Derivatives: univariate_operators,
-                             univariate_operator_to_id
-    import JuMP: _SubexpressionStorage
-    import JuMP._Derivatives: NodeType, UserOperatorRegistry
-    const JuMPOpReg = JuMP._Derivatives.UserOperatorRegistry
 
     using DataStructures: OrderedDict, BinaryMinMaxHeap, popmin!, popmax!, top
     using SparseArrays: SparseMatrixCSC, spzeros, rowvals, nzrange, nonzeros, sparse, findnz
@@ -54,6 +49,8 @@ module EAGO
     const MOI = MathOptInterface
     const MOIU = MOI.Utilities
     const MOIB = MOI.Bridges
+    const MOINL = MOI.Nonlinear
+    const MOIRAD = MOINL.ReverseAD
 
     const SAF = MOI.ScalarAffineFunction{Float64}
     const SQF = MOI.ScalarQuadraticFunction{Float64}

src/eago_optimizer/functions/nonlinear/auxiliary_variables.jl

@@ -1,7 +1,7 @@
 
 function _not_EAGO_error!(m::JuMP.Model)
     if JuMP.solver_name(m) !== "EAGO: Easy Advanced Global Optimization"
-        error("Solve attached to model must be EAGO.Optimizer")
+        error("Solver attached to model must be EAGO.Optimizer")
     end
 end
 

src/eago_optimizer/functions/nonlinear/composite_relax/forward_propagation.jl

@@ -313,10 +313,10 @@ function fprop!(t::Relax, v::Val{USERN}, g::DAT, b::RelaxCache{V,N,T}, k::Int) w
         i += 1
     end
     if anysets
-        z = MOI.eval_objective(mv, set_input)::MC{N,T}
+        z = mv.f(set_input)::MC{N,T}
         b[k] = cut(z, set(b, k), b.ic.v, zero(Float64), sparsity(g,k), b.cut, b.post)
     else
-        b[k] = MOI.eval_objective(mv, num_input)
+        b[k] = mv.f(num_input)
     end
 end
 

src/eago_optimizer/functions/nonlinear/graph/graph.jl

@@ -20,27 +20,43 @@ function _variable_count(g::AbstractDG)::Int
     error("Variable count not defined for graph type = $(typeof(g))")
 end
 
-# added id field from JuMP UserOperatorRegistry, expect more extensive changes in future.
+# added id field to MOI OperatorRegistry
 struct OperatorRegistry
-    multivariate_id::Vector{Symbol}
-    multivariate_operator_to_id::Dict{Symbol,Int}
-    multivariate_operator_evaluator::Vector{MOI.AbstractNLPEvaluator}
+    univariate_operators::Vector{Symbol}
     univariate_operator_id::Vector{Symbol}
     univariate_operator_to_id::Dict{Symbol,Int}
-    univariate_operator_f::Vector{Any}
-    univariate_operator_fprime::Vector{Any}
-    univariate_operator_fprimeprime::Vector{Any}
+    univariate_user_operator_start::Int
+    registered_univariate_operators::Vector{MOINL._UnivariateOperator}
+    multivariate_operators::Vector{Symbol}
+    multivariate_id::Vector{Symbol}
+    multivariate_operator_to_id::Dict{Symbol,Int}
+    multivariate_user_operator_start::Int
+    registered_multivariate_operators::Vector{MOINL._MultivariateOperator}
+    logic_operators::Vector{Symbol}
+    logic_operator_id::Vector{Symbol}
+    logic_operator_to_id::Dict{Symbol,Int}
+    comparison_operators::Vector{Symbol}
+    comparison_operator_id::Vector{Symbol}
+    comparison_operator_to_id::Dict{Symbol,Int}
 end
 function OperatorRegistry()
     return OperatorRegistry(
+        Symbol[],
         Symbol[],
         Dict{Symbol,Int}(),
-        MOI.AbstractNLPEvaluator[],
+        0,
+        MOINL._UnivariateOperator[],
+        Symbol[],
         Symbol[],
         Dict{Symbol,Int}(),
-        [],
-        [],
-        [],
+        0,
+        MOINL._MultivariateOperator[],
+        Symbol[],
+        Symbol[],
+        Dict{Symbol,Int}(),
+        Symbol[],
+        Symbol[],
+        Dict{Symbol,Int}()
     )
 end
 

src/eago_optimizer/functions/nonlinear/graph/graphs/directed_tree.jl

@@ -224,12 +224,12 @@ dep_subexpr_count(g::DAT)                = length(g.dependent_subexpressions)
 sparsity(g::DAT, i)                      = g.sparsity
 rev_sparsity(g::DAT, i::Int, k::Int)     = g.rev_sparsity[i]
 
-user_univariate_operator(g::DAT, i) = g.user_operators.univariate_operator_f[i]
-user_multivariate_operator(g::DAT, i) = g.user_operators.multivariate_operator_evaluator[i]
+user_univariate_operator(g::DAT, i) = g.user_operators.registered_univariate_operators[i].f
+user_multivariate_operator(g::DAT, i) = g.user_operators.registered_multivariate_operators[i]
 
 function DirectedTree(aux_info, d, op::OperatorRegistry, sub_sparsity::Dict{Int,Vector{Int}}, subexpr_linearity, parameter_values, is_sub, subexpr_indx)
-
-    nd = copy(d.nd)
+    
+    nd = copy(d.nodes)
     adj = copy(d.adj)
     const_values = copy(d.const_values)
 
@@ -243,7 +243,7 @@ function DirectedTree(aux_info, d, op::OperatorRegistry, sub_sparsity::Dict{Int,
         rev_sparsity[s] = i
     end
 
-    nodes = _convert_node_list(aux_info, d.nd, op)
+    nodes = _convert_node_list(aux_info, d.nodes, op)
     lin = linearity(nd, adj, subexpr_linearity)
     DirectedTree(nodes = nodes,
                     variables = rev_sparsity,

src/eago_optimizer/functions/nonlinear/graph/node.jl

@@ -77,7 +77,7 @@ node_is_class(::Variable, n::Node) = node_class(n) == VARIABLE
 node_is_class(::Parameter, n::Node) = node_class(n) == PARAMETER
 node_is_class(::Constant, n::Node) = node_class(n) == CONSTANT
 
-mv_eago_not_jump = setdiff(JuMP._Derivatives.operators,
+mv_eago_not_jump = setdiff(MOINL.DEFAULT_MULTIVARIATE_OPERATORS,
                            union(Symbol[k for k in keys(REV_BIVARIATE_ATOM_DICT)],
                                  Symbol[k for k in keys(REV_NARITY_ATOM_DICT)]))
 eago_mv_switch = quote end
@@ -100,8 +100,14 @@ end
         return Node(Val(true), Val(DIV), v[c])
     elseif i == 6
         error("If-else currently unsupported...")
-    elseif i >= JuMP._Derivatives.USER_OPERATOR_ID_START
-        i_mv = i - JuMP._Derivatives.USER_OPERATOR_ID_START + 1
+    elseif i == 7
+        return Node(Val(true), Val(ATAN), v[c])
+    elseif i == 8
+        return Node(Val(true), Val(MIN), v[c])
+    elseif i == 9
+        return Node(Val(true), Val(MAX), v[c])
+    elseif i >= length(DEFAULT_MULTIVARIATE_OPERATORS) + 1
+        i_mv = i - length(DEFAULT_MULTIVARIATE_OPERATORS)
         d = op.multivariate_id[i_mv]
         $eago_mv_switch
         return Node(Val(true), Val(USERN), i_mv, v[c])
@@ -126,15 +132,15 @@ end
 
 indx_JuMP = Int[]
 indx_EAGO = AtomType[]
-for k in univariate_operators
-    if haskey(REV_UNIVARIATE_ATOM_DICT, k)
-        k_EAGO = REV_UNIVARIATE_ATOM_DICT[k]
-        push!(indx_JuMP, univariate_operator_to_id[k])
+for (k,j) in enumerate(MOINL.DEFAULT_UNIVARIATE_OPERATORS)
+    if haskey(REV_UNIVARIATE_ATOM_DICT, j)
+        k_EAGO = REV_UNIVARIATE_ATOM_DICT[j]
+        push!(indx_JuMP, k)
         push!(indx_EAGO, k_EAGO)
     end
 end
 
-uni_eago_not_jump = setdiff(univariate_operators, Symbol[k for k in keys(REV_UNIVARIATE_ATOM_DICT)])
+uni_eago_not_jump = setdiff(MOINL.DEFAULT_UNIVARIATE_OPERATORS, Symbol[k for k in keys(REV_UNIVARIATE_ATOM_DICT)])
 uni_eago_not_jump = push!(uni_eago_not_jump, :-)
 eago_uni_switch = quote end
 for s in uni_eago_not_jump
@@ -146,8 +152,8 @@ end
 atom_switch = binary_switch_typ(indx_JuMP, indx_EAGO)
 @eval function _create_call_node_uni(i::Int, v, c::UnitRange{Int}, op::OperatorRegistry)
 
-    if i >= JuMP._Derivatives.USER_UNIVAR_OPERATOR_ID_START
-        j = i - JuMP._Derivatives.USER_UNIVAR_OPERATOR_ID_START + 1
+    if i >= length(DEFAULT_UNIVARIATE_OPERATORS) + 1
+        j = i - length(DEFAULT_UNIVARIATE_OPERATORS)
         dop = op.univariate_operator_id[j]
         d = op.univariate_operator_to_id[dop]
         $eago_uni_switch

src/eago_optimizer/functions/nonlinear/graph/utilities.jl

@@ -16,45 +16,46 @@ function binary_switch(ids; is_forward = true)
     end
 end
 
-function Node(aux_info, d::JuMP._Derivatives.NodeData, child_vec, c::UnitRange{Int}, op)
-    nt = d.nodetype
+function Node(aux_info, d::MOINL.Node, child_vec, c::UnitRange{Int}, op)
+    nt = d.type
     i = d.index
-    (nt == JuMP._Derivatives.CALL)          && return _create_call_node(i, child_vec, c, op)
-    (nt == JuMP._Derivatives.CALLUNIVAR)    && return _create_call_node_uni(i, child_vec, c, op)
-    (nt == JuMP._Derivatives.VALUE)         && return Node(Constant(), i)
-    (nt == JuMP._Derivatives.PARAMETER)     && return Node(Parameter(), i)
-    (nt == JuMP._Derivatives.SUBEXPRESSION) && return Node(Subexpression(), i)
-    (nt == JuMP._Derivatives.VARIABLE)      && return !is_auxiliary_variable(aux_info, i) ? Node(Variable(), i) : Node(Select(), i)
-    (nt == JuMP._Derivatives.LOGIC)         && error("Unable to load JuMP expression. Logical operators not currently supported.")
-    (nt == JuMP._Derivatives.COMPARISON)    && error("Unable to load JuMP expression. Comparisons not currently supported.")
+    (nt == MOINL.NODE_CALL_MULTIVARIATE)    && return _create_call_node(i, child_vec, c, op)
+    (nt == MOINL.NODE_CALL_UNIVARIATE)      && return _create_call_node_uni(i, child_vec, c, op)
+    (nt == MOINL.NODE_VALUE)                && return Node(Constant(), i)
+    (nt == MOINL.NODE_PARAMETER)            && return Node(Parameter(), i)
+    (nt == MOINL.NODE_SUBEXPRESSION)        && return Node(Subexpression(), i)
+    (nt == MOINL.NODE_VARIABLE)             && return !is_auxiliary_variable(aux_info, i) ? Node(Variable(), i) : Node(Select(), i)
+    (nt == MOINL.NODE_MOI_VARIABLE)         && return !is_auxiliary_variable(aux_info, i) ? Node(Variable(), i) : Node(Select(), i)
+    (nt == MOINL.NODE_LOGIC)                && error("Unable to load JuMP expression. Logical operators not currently supported.")
+    (nt == MOINL.NODE_COMPARISON)           && error("Unable to load JuMP expression. Comparisons not currently supported.")
     error("Node type = $nt not expected from JuMP.")
 end
 
-function _convert_node_list(aux_info, x::Vector{JuMP._Derivatives.NodeData}, op)
+function _convert_node_list(aux_info, x::Vector{MOINL.Node}, op)
     y = Vector{Node}(undef, length(x))
-    adj = JuMP._Derivatives.adjmat(x)
+    adj = MOINL.adjacency_matrix(x)
     child_vec = rowvals(adj)
     for i in eachindex(x)
         y[i] = Node(aux_info, x[i], child_vec, nzrange(adj, i), op)
     end
     return y
 end
 
-# Access gradient sparsity of JuMP storage.
-sparsity(d::JuMP._FunctionStorage) = d.grad_sparsity
-sparsity(d::JuMP._SubexpressionStorage) = d.sparsity
+# Access gradient sparsity of MOI storage.
+sparsity(d::MOIRAD._FunctionStorage) = d.grad_sparsity
+#sparsity(d::MOIRAD._SubexpressionStorage) = d.sparsity
 
-# Compute gradient sparsity from JuMP storage.
-function _compute_sparsity(d::JuMP._FunctionStorage, sparse_dict::Dict{Int,Vector{Int}}, is_sub, subexpr_indx)
+# Compute gradient sparsity from MOI storage.
+function _compute_sparsity(d::MOIRAD._FunctionStorage, sparse_dict::Dict{Int,Vector{Int}}, is_sub, subexpr_indx)
     dep_subexpression = Int[]
     variable_dict = Dict{Int,Bool}()
-    for n in d.nd
-        if n.nodetype == JuMP._Derivatives.VARIABLE
+    for n in d.nodes
+        if n.type == MOINL.NODE_VARIABLE
             if !haskey(variable_dict, n.index)
                 variable_dict[n.index] = true
             end
         end
-        if n.nodetype == JuMP._Derivatives.SUBEXPRESSION
+        if n.type == MOINL.NODE_SUBEXPRESSION
             push!(dep_subexpression, n.index)
         end
     end
@@ -71,16 +72,16 @@ function _compute_sparsity(d::JuMP._FunctionStorage, sparse_dict::Dict{Int,Vecto
     end
     sparsity, dep_subexpression
 end
-function _compute_sparsity(d::JuMP._SubexpressionStorage, sparse_dict::Dict{Int,Vector{Int}}, is_sub, subexpr_indx)
+function _compute_sparsity(d::MOIRAD._SubexpressionStorage, sparse_dict::Dict{Int,Vector{Int}}, is_sub, subexpr_indx)
     dep_subexpression = Int[]
     variable_dict = Dict{Int,Bool}()
-    for n in d.nd
-        if n.nodetype == JuMP._Derivatives.VARIABLE
+    for n in d.nodes
+        if n.type == MOINL.NODE_VARIABLE
             if !haskey(variable_dict, n.index)
                 variable_dict[n.index] = true
             end
         end
-        if n.nodetype == JuMP._Derivatives.SUBEXPRESSION
+        if n.type == MOINL.NODE_SUBEXPRESSION
             push!(dep_subexpression, n.index)
         end
     end
@@ -98,25 +99,34 @@ function _compute_sparsity(d::JuMP._SubexpressionStorage, sparse_dict::Dict{Int,
     sparsity, dep_subexpression
 end
 
-function linearity(d::JuMP._Derivatives.Linearity)
-    (d == JuMP._Derivatives.LINEAR)            && return LIN_LINEAR
-    (d == JuMP._Derivatives.PIECEWISE_LINEAR)  && return LIN_PIECEWISE_LINEAR
-    (d == JuMP._Derivatives.NONLINEAR)         && return LIN_NONLINEAR
-    LIN_CONSTANT                               # assumes d is then JuMP._Derivatives.CONSTANT
+function linearity(d::MOIRAD.Linearity)
+    (d == MOIRAD.LINEAR)            && return LIN_LINEAR
+    (d == MOIRAD.PIECEWISE_LINEAR)  && return LIN_PIECEWISE_LINEAR
+    (d == MOIRAD.NONLINEAR)         && return LIN_NONLINEAR
+    LIN_CONSTANT                   # assumes d is then MOINL.CONSTANT
 end
-function linearity(nd::Vector{JuMP._Derivatives.NodeData}, adj::SparseMatrixCSC{Bool,Int}, d::Vector{JuMP._Derivatives.Linearity})
-    x = JuMP._Derivatives.classify_linearity(nd, adj, d)
+
+function linearity(nd::Vector{MOINL.Node}, adj::SparseMatrixCSC{Bool,Int}, d::Vector{MOIRAD.Linearity})
+    x = MOIRAD._classify_linearity(nd, adj, d)
     linearity.(x)
 end
 
-function OperatorRegistry(d::JuMP._Derivatives.UserOperatorRegistry)
-    mv_id = collect(keys(d.multivariate_operator_to_id))
-    mv_operator_to_id = d.multivariate_operator_to_id
-    mv_operator_evaluator = d.multivariate_operator_evaluator
+function OperatorRegistry(d::MOINL.OperatorRegistry)
+    u_operators = d.univariate_operators
     u_operator_id = collect(keys(d.univariate_operator_to_id))
     u_to_id = d.univariate_operator_to_id
-    u_to_f = d.univariate_operator_f
-    u_fprime = d.univariate_operator_fprime
-    u_fprimeprime = d.univariate_operator_fprimeprime
-    OperatorRegistry(mv_id, mv_operator_to_id, mv_operator_evaluator, u_operator_id, u_to_id, u_to_f, u_fprime, u_fprimeprime)
+    u_operator_start = d.univariate_user_operator_start
+    r_u_operators = d.registered_univariate_operators
+    mv_operators = d.multivariate_operators
+    mv_operator_id = collect(keys(d.multivariate_operator_to_id))
+    mv_to_id = d.multivariate_operator_to_id
+    mv_operator_start = d.multivariate_user_operator_start
+    r_mv_operators = d.registered_multivariate_operators
+    l_operators = d.logic_operators
+    l_operator_id = collect(keys(d.logic_operator_to_id))
+    l_to_id = d.logic_operator_to_id
+    c_operators = d.comparison_operators
+    c_operator_id = collect(keys(d.comparison_operator_to_id))
+    c_to_id = d.comparison_operator_to_id
+    OperatorRegistry(u_operators, u_operator_id, u_to_id, u_operator_start, r_u_operators, mv_operators, mv_operator_id, mv_to_id, mv_operator_start, r_mv_operators, l_operators, l_operator_id, l_to_id, c_operators, c_operator_id, c_to_id)
 end