support direct mode for optigraphs

src/Plasmo.jl

@@ -29,6 +29,7 @@ export OptiGraph,
     OptiNode,
     OptiEdge,
     NodeVariableRef,
+    direct_moi_graph,
     graph_backend,
     graph_index,
     add_node,

src/backends/moi_backend.jl

@@ -104,12 +104,10 @@ end
 Initialize an empty backend given an optigraph.
 By default we use a `CachingOptimizer` to store the underlying optimizer just like JuMP.
 """
-function GraphMOIBackend(graph::OptiGraph)
-    inner = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}())
-    cache = MOI.Utilities.CachingOptimizer(inner, MOI.Utilities.AUTOMATIC)
+function GraphMOIBackend(graph::OptiGraph, backend::MOI.ModelLike)
     return GraphMOIBackend(
         graph,
-        cache,
+        backend,
         ElementToGraphMap(),
         GraphToElementMap(),
         OrderedDict{OptiNode,Vector{MOI.VariableIndex}}(),
@@ -119,10 +117,72 @@ function GraphMOIBackend(graph::OptiGraph)
     )
 end
 
-function graph_index(
-    backend::GraphMOIBackend, ref::RT
-) where {RT<:Union{NodeVariableRef,ConstraintRef}}
-    return backend.element_to_graph_map[ref]
+function cached_moi_backend(graph::OptiGraph)
+    inner = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}())
+    cache = MOI.Utilities.CachingOptimizer(inner, MOI.Utilities.AUTOMATIC)
+    return GraphMOIBackend(graph, cache)
+end
+
+function direct_moi_backend(graph::OptiGraph, backend::MOI.ModelLike;)
+    @assert MOI.is_empty(backend)
+    return GraphMOIBackend(graph, backend)
+end
+
+function direct_moi_backend(graph::OptiGraph, factory::MOI.OptimizerWithAttributes)
+    optimizer = MOI.instantiate(factory)
+    return direct_moi_backend(graph, optimizer)
+end
+
+# NOTE: _moi_mode adapted from JuMP.jl
+# https://github.com/jump-dev/JuMP.jl/blob/301d46e81cb66c74c6e22cd89fb89ced740f157b/src/JuMP.jl#L571-L575
+_moi_mode(::MOI.ModelLike) = DIRECT
+function _moi_mode(model::MOIU.CachingOptimizer)
+    return model.mode == MOIU.AUTOMATIC ? AUTOMATIC : MANUAL
+end
+
+function _moi_mode(backend::GraphMOIBackend)
+    return _moi_mode(backend.moi_backend)
+end
+
+function JuMP.error_if_direct_mode(backend::GraphMOIBackend, func::Symbol)
+    if _moi_mode(backend) == DIRECT
+        error("The `$func` function is not supported in DIRECT mode.")
+    end
+    return nothing
+end
+
+# MOI Utilities
+
+function MOIU.state(backend::GraphMOIBackend)
+    return MOIU.state(JuMP.backend(backend))
+end
+
+function MOIU.reset_optimizer(
+    backend::GraphMOIBackend, optimizer::MOI.AbstractOptimizer, ::Bool=true
+)
+    JuMP.error_if_direct_mode(backend, :reset_optimizer)
+    MOIU.reset_optimizer(JuMP.backend(backend), optimizer)
+    return nothing
+end
+
+function MOIU.reset_optimizer(backend::GraphMOIBackend)
+    JuMP.error_if_direct_mode(backend, :reset_optimizer)
+    if MOI.Utilities.state(JuMP.backend(backend)) == MOI.Utilities.ATTACHED_OPTIMIZER
+        MOIU.reset_optimizer(JuMP.backend(backend))
+    end
+    return nothing
+end
+
+function MOIU.drop_optimizer(backend::GraphMOIBackend)
+    JuMP.error_if_direct_mode(backend, :drop_optimizer)
+    MOIU.drop_optimizer(JuMP.backend(backend))
+    return nothing
+end
+
+function MOIU.attach_optimizer(backend::GraphMOIBackend)
+    JuMP.error_if_direct_mode(backend, :attach_optimizer)
+    MOIU.attach_optimizer(JuMP.backend(backend))
+    return nothing
 end
 
 # JuMP Methods
@@ -147,6 +207,21 @@ function JuMP.constraint_ref_with_index(backend::GraphMOIBackend, idx::MOI.Index
     return backend.graph_to_element_map[idx]
 end
 
+"""
+    graph_index(
+        backend::GraphMOIBackend, 
+        ref::RT
+    ) where {RT<:Union{NodeVariableRef,ConstraintRef}}
+
+Return the actual variable or constraint index of the backend model that corresponds to the
+local index of a node or edge.
+"""
+function graph_index(
+    backend::GraphMOIBackend, ref::RT
+) where {RT<:Union{NodeVariableRef,ConstraintRef}}
+    return backend.element_to_graph_map[ref]
+end
+
 """
     graph_operator(backend::GraphMOIBackend, element::OptiElement, name::Symbol)
 
@@ -645,7 +720,6 @@ end
 function _add_backend_variables(backend::GraphMOIBackend, vars::Vector{NodeVariableRef})
     vars_to_add = setdiff(vars, keys(backend.element_to_graph_map.var_map))
     for var in vars_to_add
-        # _add_variable_to_backend(backend, var)
         MOI.add_variable(backend, var)
     end
     return nothing

src/optigraph.jl

@@ -3,7 +3,7 @@
 #  License, v. 2.0. If a copy of the MPL was not distributed with this
 #  file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
-function OptiGraph(; name::Symbol=Symbol(:g, gensym()))
+function _initialize_optigraph(name::Symbol)
     graph = OptiGraph(
         name,
         OrderedSet{OptiNode}(),
@@ -18,9 +18,22 @@ function OptiGraph(; name::Symbol=Symbol(:g, gensym()))
         Set{Any}(),
         false,
     )
+    return graph
+end
+
+function OptiGraph(; name::Symbol=Symbol(:g, gensym()))
+    graph = _initialize_optigraph(name)
+    # default is to use a CachingOptimizer backend
+    graph.backend = cached_moi_backend(graph)
+    return graph
+end
 
-    # default is MOI backend
-    graph.backend = GraphMOIBackend(graph)
+function direct_moi_graph(
+    backend::Union{MOI.ModelLike,MOI.OptimizerWithAttributes};
+    name::Symbol=Symbol(:g, gensym()),
+)
+    graph = _initialize_optigraph(name)
+    graph.backend = direct_moi_backend(graph, backend)
     return graph
 end
 
@@ -1425,5 +1438,5 @@ function _set_objective_coefficient(
 end
 
 function JuMP.unregister(graph::OptiGraph, key::Symbol)
-    delete!(object_dictionary(graph), key)
+    return delete!(object_dictionary(graph), key)
 end

src/optimizer_interface.jl

@@ -65,25 +65,41 @@ function JuMP.set_attributes(destination::Union{OptiGraph,NodeVariableRef}, pair
     return nothing
 end
 
+function JuMP.time_limit_sec(graph::OptiGraph)
+    return MOI.get(graph, MOI.TimeLimitSec())
+end
+
 #
 # set optimizer
 #
 
-# NOTE: _moi_mode adapted from JuMP.jl
-# https://github.com/jump-dev/JuMP.jl/blob/301d46e81cb66c74c6e22cd89fb89ced740f157b/src/JuMP.jl#L571-L575
-_moi_mode(::MOI.ModelLike) = DIRECT
-function _moi_mode(model::MOIU.CachingOptimizer)
-    return model.mode == MOIU.AUTOMATIC ? AUTOMATIC : MANUAL
+function JuMP.mode(graph::OptiGraph)
+    return _moi_mode(JuMP.backend(graph))
 end
 
-function JuMP.mode(graph::OptiGraph)
-    return _moi_mode(JuMP.backend(graph_backend(graph)))
+function MOIU.state(graph)
+    return MOIU.state(JuMP.backend(graph))
 end
 
-function JuMP.error_if_direct_mode(graph::OptiGraph, func::Symbol)
-    if JuMP.mode(graph) == DIRECT
-        error("The `$func` function is not supported in DIRECT mode.")
-    end
+function MOIU.reset_optimizer(
+    graph::OptiGraph, optimizer::MOI.AbstractOptimizer, ::Bool=true
+)
+    MOIU.reset_optimizer(JuMP.backend(graph), optimizer)
+    return nothing
+end
+
+function MOIU.reset_optimizer(graph::OptiGraph)
+    MOIU.reset_optimizer(JuMP.backend(graph))
+    return nothing
+end
+
+function MOIU.drop_optimizer(graph::OptiGraph)
+    MOIU.drop_optimizer(JuMP.backend(graph))
+    return nothing
+end
+
+function MOIU.attach_optimizer(graph::OptiGraph)
+    MOIU.attach_optimizer(JuMP.backend(graph))
     return nothing
 end
 
@@ -99,7 +115,7 @@ Set the optimizer on `graph` by passing an `optimizer_constructor`.
 function JuMP.set_optimizer(
     graph::OptiGraph, JuMP.@nospecialize(optimizer_constructor); add_bridges::Bool=true
 )
-    JuMP.error_if_direct_mode(graph, :set_optimizer)
+    JuMP.error_if_direct_mode(JuMP.backend(graph), :set_optimizer)
     if add_bridges
         optimizer = MOI.instantiate(optimizer_constructor)#; with_bridge_type = T)
         for BT in graph.bridge_types

src/optinode.jl

@@ -122,7 +122,7 @@ function JuMP.all_variables(node::OptiNode)
 end
 
 function JuMP.unregister(node::OptiNode, key::Symbol)
-    delete!(object_dictionary(node), (node, key))
+    return delete!(object_dictionary(node), (node, key))
 end
 
 ### Duals

test/test_optigraph.jl

@@ -20,7 +20,48 @@ function test_simple_graph()
     @linkconstraint(graph, nodes[1][:x] + nodes[2][:x] == 4)
     @objective(graph, Max, nodes[1][:x] + 2 * nodes[2][:x])
 
+    @test MOIU.state(graph) == MOIU.NO_OPTIMIZER
     set_optimizer(graph, HiGHS.Optimizer)
+    @test MOIU.state(graph) == MOIU.EMPTY_OPTIMIZER
+    @suppress optimize!(graph)
+    @test MOIU.state(graph) == MOIU.ATTACHED_OPTIMIZER
+
+    @test objective_value(graph) == 7.0
+    @test value(nodes[1][:x]) == 1.0
+    @test value(nodes[2][:x]) == 3.0
+    @test value(nodes[1][:x] + nodes[2][:x]) == value(graph, nodes[1][:x] + nodes[2][:x])
+    @test value(nodes[1][:x]^2 + nodes[2][:x]^2) ==
+        value(graph, nodes[1][:x]^2 + nodes[2][:x]^2)
+    @test value(nodes[1][:x]^3 + nodes[2][:x]^3) ==
+        value(graph, nodes[1][:x]^3 + nodes[2][:x]^3)
+
+    @test JuMP.termination_status(graph) == MOI.OPTIMAL
+    @test JuMP.primal_status(graph) == MOI.FEASIBLE_POINT
+    @test JuMP.dual_status(graph) == MOI.FEASIBLE_POINT
+    @test JuMP.result_count(graph) == 1
+    @test JuMP.raw_status(graph) == "kHighsModelStatusOptimal"
+
+    constraints = all_constraints(graph)
+    @test JuMP.dual(constraints[1]) == 1.0
+    @test JuMP.dual(constraints[2]) == 0.0
+    @test JuMP.dual(constraints[3]) == -2.0
+
+    MOIU.reset_optimizer(graph)
+    @test MOIU.state(graph) == MOIU.EMPTY_OPTIMIZER
+    MOIU.attach_optimizer(graph)
+    @test MOIU.state(graph) == MOIU.ATTACHED_OPTIMIZER
+    MOIU.drop_optimizer(graph)
+    @test MOIU.state(graph) == MOIU.NO_OPTIMIZER
+end
+
+function test_direct_moi_graph()
+    graph = direct_moi_graph(HiGHS.Optimizer())
+    @optinode(graph, nodes[1:2])
+
+    @variable(nodes[1], x >= 1)
+    @variable(nodes[2], x >= 2)
+    @linkconstraint(graph, nodes[1][:x] + nodes[2][:x] == 4)
+    @objective(graph, Max, nodes[1][:x] + 2 * nodes[2][:x])
     @suppress optimize!(graph)
 
     @test objective_value(graph) == 7.0