wip: more progress on JuMP/MOI methods

jalving · May 13, 2024 · 4e34d46 · 4e34d46
1 parent 6cf17a1
commit 4e34d46
Show file tree

Hide file tree

Showing 5 changed files with 144 additions and 91 deletions.
diff --git a/src/Plasmo.jl b/src/Plasmo.jl
@@ -38,7 +38,7 @@ include("optinode.jl")
 
 include("optiedge.jl")
 
-include("moi_backend.jl")
+include("backends/moi_backend.jl")
 
 include("aggregate.jl")
 

diff --git a/src/moi_backend.jl → src/backends/moi_backend.jl b/src/moi_backend.jl → src/backends/moi_backend.jl
@@ -81,6 +81,8 @@ mutable struct GraphMOIBackend <: MOI.AbstractOptimizer
     # map of nodes and edges to variables and constraints.
     node_variables::OrderedDict{OptiNode,Vector{MOI.VariableIndex}}
     element_constraints::OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}
+    element_attributes::OrderedDict{Tuple{OptiElement,MOI.AbstractModelAttribute},Any}
+    operator_map::OrderedDict{Tuple{OptiElement,Symbol},Symbol}
 end
 
 """
@@ -98,10 +100,20 @@ function GraphMOIBackend(graph::OptiGraph)
         ElementToGraphMap(),
         GraphToElementMap(),
         OrderedDict{OptiNode,Vector{MOI.VariableIndex}}(),
-        OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}()
+        OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}(),
+        OrderedDict{Tuple{OptiElement,MOI.AbstractModelAttribute},Any}(),
+        OrderedDict{Tuple{OptiElement,Symbol},Symbol}()
     )
 end
 
+function graph_index(gb::GraphMOIBackend, nvref::NodeVariableRef)
+    return gb.element_to_graph_map[nvref]
+end
+
+function graph_operator(gb::GraphMOIBackend, element::OptiElement, name::Symbol)
+    return gb.operator_map[(element,name)]
+end
+
 function _add_node(gb::GraphMOIBackend, node::OptiNode)
     if !haskey(gb.node_variables, node)
         gb.node_variables[node] = MOI.VariableIndex[]
@@ -125,37 +137,104 @@ function JuMP.backend(gb::GraphMOIBackend)
     return gb.moi_backend
 end
 
-# MOI Methods
+### MOI Methods
+
+# graph attributes
 
-function MOI.get(gb::GraphMOIBackend, attr::MOI.AnyAttribute)
+function MOI.get(gb::GraphMOIBackend, attr::AT) where 
+    AT <: Union{MOI.AbstractModelAttribute, MOI.AbstractOptimizerAttribute}
     return MOI.get(gb.moi_backend, attr)
 end
 
-function MOI.get(gb::GraphMOIBackend, attr::MOI.AnyAttribute, ref::ConstraintRef)
-    graph_index = gb.element_to_graph_map[ref]
-    return MOI.get(gb.moi_backend, attr, graph_index)
+function MOI.set(gb::GraphMOIBackend, attr::AT, args...) where
+    AT <: Union{MOI.AbstractModelAttribute, MOI.AbstractOptimizerAttribute}
+    MOI.set(gb.moi_backend, attr, args...)
+    return
 end
 
-function MOI.get(gb::GraphMOIBackend, attr::MOI.AnyAttribute, ref::NodeVariableRef)
-    graph_index = gb.element_to_graph_map[ref]
-    return MOI.get(gb.moi_backend, attr, graph_index)
+# element attributes
+
+# function MOI.get(gb::GraphMOIBackend, attr::AT, element::OptiElement) where
+#     AT <: MOI.AbstractModelAttribute
+#     return gb.element_attributes[(element,attr)]
+# end
+
+# function MOI.set(gb::GraphMOIBackend, attr::AT, element::OptiElement, args...) where
+#     AT <: MOI.AbstractModelAttribute
+#     MOI.set(gb.moi_backend, attr, args...)
+#     gb.element_attributes[(element,attr)] = tuple(args...)
+#     return
+# end
+function MOI.set(
+    gb::GraphMOIBackend,
+    attr::MOI.UserDefinedFunction,
+    node::OptiNode,
+    args...
+)
+    registered_name = Symbol(node.label, ".", attr.name)
+    MOI.set(
+        gb.moi_backend, 
+        MOI.UserDefinedFunction(registered_name, attr.arity),
+        args...
+    )
+    gb.element_attributes[(node,attr)] = tuple(args...)
+    gb.operator_map[(node,attr.name)] = registered_name
 end
 
 function MOI.get(gb::GraphMOIBackend, attr::MOI.NumberOfVariables, node::OptiNode)
     return length(gb.node_variables[node])
 end
 
-function MOI.get(gb::GraphMOIBackend, attr::MOI.NumberOfConstraints{F,S}, node::OptiNode) where{F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
-    cons = MOI.get(gb.moi_backend, MOI.ListOfConstraintIndices{F,S}())
-    refs = [gb.graph_to_element_map[con] for con in cons]
-    return length(filter((cref) -> cref.model == node, refs))
+function MOI.get(
+    gb::GraphMOIBackend, 
+    attr::MOI.ListOfConstraintTypesPresent, 
+    element::OptiElement
+)
+    cons = gb.element_constraints[element]
+    con_types = unique(typeof.(cons))
+    type_tuple = [(type.parameters[1],type.parameters[2]) for type in con_types]  
+    return type_tuple
 end
 
-function MOI.set(gb::GraphMOIBackend, attr::MOI.AnyAttribute, args...)
-    MOI.set(gb.moi_backend, attr, args...)
+function MOI.get(
+    gb::GraphMOIBackend, 
+    attr::MOI.NumberOfConstraints{F,S}, 
+    element::OptiElement
+) where{F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
+    return length(gb.element_constraints[element])
+end
+
+# variable attributes
+
+function MOI.get(gb::GraphMOIBackend, attr::AT, nvref::NodeVariableRef) where 
+    AT <: MOI.AbstractVariableAttribute
+    graph_index = gb.element_to_graph_map[nvref]
+    return MOI.get(gb.moi_backend, attr, graph_index)
+end
+
+function MOI.set(gb::GraphMOIBackend, attr::AT, nvref::NodeVariableRef, args...) where 
+    AT <: MOI.AbstractVariableAttribute
+    graph_index = gb.element_to_graph_map[nvref]
+    MOI.set(gb.moi_backend, attr, graph_index, args...)
     return
 end
 
+# constraint attributes
+
+function MOI.get(gb::GraphMOIBackend, attr::AT, cref::ConstraintRef) where
+    AT <: MOI.AbstractConstraintAttribute
+    graph_index = gb.element_to_graph_map[cref]
+    return MOI.get(gb.moi_backend, attr, graph_index)
+end
+
+function MOI.set(gb::GraphMOIBackend, attr::AT, cref::ConstraintRef, args...) where 
+    AT <: MOI.AbstractConstraintAttribute
+    graph_index = gb.element_to_graph_map[nvref]
+    MOI.set(gb.moi_backend, attr, graph_index, args...)
+end
+
+# delete
+
 function MOI.delete(gb::GraphMOIBackend, nvref::NodeVariableRef)
     MOI.delete(gb.moi_backend, gb.element_to_graph_map[nvref])
     delete!(gb.graph_to_element_map.var_map, gb.element_to_graph_map[nvref])
@@ -170,6 +249,8 @@ function MOI.delete(gb::GraphMOIBackend, cref::ConstraintRef)
     return
 end
 
+# is_valid
+
 function MOI.is_valid(gb::GraphMOIBackend, vi::MOI.VariableIndex)
     return MOI.is_valid(gb.moi_backend, vi)
 end
@@ -178,17 +259,15 @@ function MOI.is_valid(gb::GraphMOIBackend, ci::MOI.ConstraintIndex)
     return MOI.is_valid(gb.moi_backend, ci)
 end
 
+# optimize!
+
 function MOI.optimize!(gb::GraphMOIBackend)
     MOI.optimize!(gb.moi_backend)
     return
 end
 
 ### Variables and Constraints
 
-function graph_index(gb::GraphMOIBackend, nvref::NodeVariableRef)
-    return gb.element_to_graph_map[nvref]
-end
-
 function _add_variable_to_backend(
     graph_backend::GraphMOIBackend,
     vref::NodeVariableRef

diff --git a/src/core_types.jl b/src/core_types.jl
@@ -49,5 +49,4 @@ end
 
 const OptiElement = Union{OptiNode,OptiEdge}
 
-const OptiObject = Union{OptiNode, OptiEdge, OptiGraph}
-
+const OptiObject = Union{OptiNode, OptiEdge, OptiGraph}
diff --git a/src/node_variables.jl b/src/node_variables.jl
@@ -30,7 +30,7 @@ function MOI.get(
     attr::MOI.AbstractVariableAttribute,
     nvref::NodeVariableRef
 )
-    return MOI.get(graph_backend(node), attr, graph_index)
+    return MOI.get(graph_backend(node), attr, nvref)
 end
 
 function MOI.set(
@@ -41,8 +41,7 @@ function MOI.set(
 )
     for graph in containing_optigraphs(node)
         gb = graph_backend(graph)
-        graph_index = gb.element_to_graph_map[nvref]
-        MOI.set(gb, attr, graph_index, args...)
+        MOI.set(gb, attr, nvref, args...)
     end
     return
 end
@@ -179,7 +178,7 @@ end
 ### node variable bounds
 
 function JuMP.has_lower_bound(nvref::NodeVariableRef)
-    return _moi_nv_has_lower_bound(graph_backend(nvref), nvref)
+    return _moi_nv_has_lower_bound(nvref)
 end
 
 function JuMP.set_lower_bound(nvref::NodeVariableRef, lower::Number)

diff --git a/src/optinode.jl b/src/optinode.jl
@@ -15,15 +15,27 @@ function Base.getindex(node::OptiNode, name::Symbol)
     return source_graph(node).node_obj_dict[t]
 end
 
-function JuMP.num_variables(node::OptiNode)
-    return MOI.get(graph_backend(node), MOI.NumberOfVariables(), node)
-    #return length(graph_backend(node).node_variables[node])
+"""
+    source_graph(node::OptiNode)
+
+Return the optigraph that contains the optinode. This is the optigraph that 
+defined said node and stores node object dictionary data.
+"""
+function source_graph(node::OptiNode)
+    return node.source_graph.x
 end
 
-function JuMP.all_variables(node::OptiNode)
-    gb = graph_backend(node)
-    graph_indices = gb.node_variables[node]
-    return getindex.(Ref(gb.graph_to_element_map), graph_indices)
+function containing_optigraphs(node::OptiNode)
+    source = source_graph(node)
+    graphs = [source]
+    if haskey(source.node_to_graphs, node)
+        graphs = [graphs; source.node_to_graphs[node]]
+    end
+    return graphs
+end
+
+function containing_backends(node::OptiNode)
+    return graph_backend.(containing_optigraphs(node))
 end
 
 """
@@ -44,28 +56,7 @@ function graph_backend(node::OptiNode)
     return graph_backend(source_graph(node))
 end
 
-"""
-    source_graph(node::OptiNode)
 
-Return the optigraph that contains the optinode. This is the optigraph that 
-defined said node and stores node object dictionary data.
-"""
-function source_graph(node::OptiNode)
-    return node.source_graph.x
-end
-
-function containing_optigraphs(node::OptiNode)
-    source = source_graph(node)
-    graphs = [source]
-    if haskey(source.node_to_graphs, node)
-        graphs = [graphs; source.node_to_graphs[node]]
-    end
-    return graphs
-end
-
-function containing_backends(node::OptiNode)
-    return graph_backend.(containing_optigraphs(node))
-end
 
 """
     Filter the object dictionary for values that belong to node. Keep in mind that 
@@ -88,6 +79,17 @@ end
 
 ### JuMP Methods
 
+function JuMP.num_variables(node::OptiNode)
+    return MOI.get(graph_backend(node), MOI.NumberOfVariables(), node)
+    #return length(graph_backend(node).node_variables[node])
+end
+
+function JuMP.all_variables(node::OptiNode)
+    gb = graph_backend(node)
+    graph_indices = gb.node_variables[node]
+    return getindex.(Ref(gb.graph_to_element_map), graph_indices)
+end
+
 function JuMP.delete(node::OptiNode, cref::ConstraintRef)
     if node !== JuMP.owner_model(cref)
         error(
@@ -150,9 +152,12 @@ function JuMP.add_nonlinear_operator(
             "hesssian provided)",
         )
     end
-    name = Symbol(node.label, ".", name)
-    MOI.set(graph_backend(node), MOI.UserDefinedFunction(name, dim), tuple(f, args...))
-    return JuMP.NonlinearOperator(f, name)
+    #registered_name = Symbol(node.label, ".", name)
+    #MOI.set(node, MOI.UserDefinedFunction(registered_name, dim), tuple(f, args...))
+    #return JuMP.NonlinearOperator(f, registered_name)
+    MOI.set(node, MOI.UserDefinedFunction(name, dim), tuple(f, args...))
+    registered_name = graph_operator(graph_backend(node), node, name)
+    return JuMP.NonlinearOperator(f, registered_name)
 end
 
 function _set_dirty(node::OptiNode)
@@ -212,48 +217,20 @@ end
 
 # TODO: store objective functions on nodes and query as node attributes
 
-function MOI.get(node::OptiNode, attr::MOI.AnyAttribute)
-    return MOI.get(graph_backend(node), attr)
-end
-
 # function MOI.get(node::OptiNode, attr::MOI.UserDefinedFunction)
 #     return MOI.get(graph_backend(node), attr)
 # end
 
-# TODO: consider caching constraint types in graph backend versus using unique to filter
-function MOI.get(node::OptiNode, attr::MOI.ListOfConstraintTypesPresent)
-    cons = graph_backend(node).element_constraints[node]
-    con_types = unique(typeof.(cons))
-    type_tuple = [(type.parameters[1],type.parameters[2]) for type in con_types]  
-    return type_tuple
+function MOI.get(node::OptiNode, attr::MOI.AnyAttribute)
+    return MOI.get(graph_backend(node), attr, node)
 end
 
-function MOI.get(
-    node::OptiNode, 
-    attr::MOI.ListOfConstraintIndices{F,S}
-) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
-    con_inds = MOI.ConstraintIndex{F,S}[]
-    for con in graph_backend(node).element_constraints[node]
-        if (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S)
-            push!(con_inds, con)
-        end
+function MOI.set(node::OptiNode, attr::MOI.AnyAttribute, args...)
+    for graph in containing_optigraphs(node)
+        MOI.set(graph_backend(node), attr, node, args...)
     end
-    return con_inds
 end
 
-# function MOI.get(
-#     node::OptiNode, 
-#     attr::MOI.ListOfConstraintIndices{F,S}
-# ) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
-#     con_inds = MOI.ConstraintIndex{F,S}[]
-#     for con in graph_backend(node).element_constraints[node]
-#         if (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S)
-#             push!(con_inds, con)
-#         end
-#     end
-#     return con_inds
-# end
-
 function MOI.get(
     node::OptiNode, 
     attr::MOI.AbstractConstraintAttribute,
@@ -270,7 +247,6 @@ function MOI.set(
 )
     for graph in containing_optigraphs(JuMP.owner_model(cref))
         gb = graph_backend(graph)
-        graph_index = gb.element_to_graph_map[cref]
         MOI.set(gb, attr, graph_index, args...)
     end
     return