fix aggregate; add value (graph,expr)

src/Plasmo.jl

@@ -60,7 +60,11 @@ neighborhood, expand,
 
 # macros
 
-@optinode, @nodevariables, @linkconstraint
+@optinode, @nodevariables, @linkconstraint,
+
+# other functions
+
+set_jump_model
 
 include("core_types.jl")
 

src/aggregate.jl

@@ -6,7 +6,7 @@ The reference of the combined model can be obtained by indexing the map with the
 reference of the corresponding original optinode.
 """
 struct GraphReferenceMap
-    #map variables and from original optigraph to new aggregate node or graph
+    # map variables and from original optigraph to new aggregate node or graph
     var_map::OrderedDict{NodeVariableRef,NodeVariableRef}
     con_map::OrderedDict{JuMP.ConstraintRef,JuMP.ConstraintRef}  
 end
@@ -58,13 +58,16 @@ end
 
 Aggregate an optigraph into a graph containing a single optinode.
 """
-function aggregate(graph::OptiGraph)
+function aggregate(graph::OptiGraph; name=gensym())
     # aggregate into a graph containing a single optinode
-    new_graph = OptiGraph()
+    new_graph = OptiGraph(;name=name)
     new_node, ref_map = _copy_graph_elements_to!(new_graph, graph)
     # copy other model attributes (such as the objective function)
     # setup index_map from the ref_map
     _copy_attributes_to!(new_graph, graph, ref_map)
+
+    # copy objective function to node
+    JuMP.set_objective(new_node, objective_sense(new_graph), objective_function(new_graph))
     return new_node, ref_map
 end
 
@@ -90,7 +93,9 @@ function _copy_graph_elements_to!(new_graph::OptiGraph, source_graph::OptiGraph)
 end
 
 """
-    Copy graph attributes from optigraph `source_graph` to the new optigraph `new_graph`.
+    Copy graph attributes from `source_graph` to the new optigraph `new_graph`.
+
+These attributes include all model and optimizer attributes
 """
 function _copy_attributes_to!(
     new_graph::OptiGraph, 
@@ -103,7 +108,7 @@ function _copy_attributes_to!(
 
     # NOTE: we use an if statement because the source graph does not necessarily have all 
     # the variables or constraints. we just want to pass the attributes that would be 
-    # exposed such as the objective function.
+    # exposed such as the graph objective function.
     for (source_vref, dest_vref) in ref_map.var_map
         # TODO: use outer method here; don't access data members directly 
         if source_vref in keys(src.element_to_graph_map.var_map)
@@ -115,6 +120,7 @@ function _copy_attributes_to!(
             index_map[graph_index(source_graph, source_cref)] = graph_index(dest_cref)
         end
     end
+    # TODO: avoid using direct reference to moi_backend
     MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map)
     return
 end
@@ -135,17 +141,10 @@ function _copy_node_to!(
     _copy_variables!(new_node, source_node, ref_map, index_map)
 
     _copy_constraints!(new_node, source_node, ref_map, index_map)
-
-    # update reference map
-    for (source_index, dest_index) in index_map.con_map
-        source_cref = src.graph_to_element_map.con_map[source_index]
-        dest_cref = dest.graph_to_element_map.con_map[dest_index]
-        ref_map[source_cref] = dest_cref
-    end
 end
 
 """
-    Copy an optinode to a new optigraph
+    Copy an optinode to a new optigraph.
 """
 function _copy_node_to!(new_graph::OptiGraph, source_node::OptiNode)
     ref_map = GraphReferenceMap()
@@ -160,6 +159,7 @@ function _copy_variables!(
     ref_map::GraphReferenceMap,
     index_map::MOIU.IndexMap
 )
+    # get relevant backends
     src = graph_backend(source_node)
     dest = graph_backend(new_node)
 
@@ -177,37 +177,44 @@ function _copy_variables!(
     vis_src = graph_index.(source_variables)
     MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src)
 
-    # TODO: set new variable names (otherwise they can be confusing)
-    # if set_new_names
-    #     for nvref in all_variables(new_node)
-    #         JuMP.set_name(nvref, "$(new_node.label)")
-    #     end
-    # end
-
+    # set new variable names
+    for nvref in all_variables(source_node)
+        new_variable = ref_map[nvref]
+        JuMP.set_name(new_variable, "$(new_node.label)"*"."*JuMP.name(nvref))
+    end
 end
 
+"""
+    Copy the constraints from a node or edge into a new node.
+"""
 function _copy_constraints!(
     new_node::OptiNode, 
     source_element::OptiElement, 
     ref_map::GraphReferenceMap,
     index_map::MOIU.IndexMap
 )
+    # get relevant backends
     src = graph_backend(source_element)
     dest = graph_backend(new_node)
+
+    # copy each constraint by iterating through each type
     constraint_types = MOI.get(src.moi_backend, MOI.ListOfConstraintTypesPresent())
     for (F, S) in constraint_types
         cis_src = MOI.get(source_element, MOI.ListOfConstraintIndices{F,S}())
         index_map_FS = index_map[F,S]
         for ci in cis_src
+            # TODO: use references to elements instead so we don't have to hardcode backend
             src_func = MOI.get(src.moi_backend, MOI.ConstraintFunction(), ci)
             src_set = MOI.get(src.moi_backend, MOI.ConstraintSet(), ci)
+            # src_func = MOI.get(source_element, MOI.ConstraintFunction(), ci)
+            # src_set = MOI.get(source_element, MOI.ConstraintSet(), ci)
 
             # get source cref to lookup constraint shape
             src_cref = JuMP.constraint_ref_with_index(src, ci)
             new_shape = src_cref.shape
 
-            # create new ConstraintRef
-            new_constraint_index = Plasmo.next_constraint_index(
+            # create a new ConstraintRef
+            new_constraint_index = next_constraint_index(
                 new_node, 
                 typeof(src_func), 
                 typeof(src_set)
@@ -216,63 +223,59 @@ function _copy_constraints!(
 
             # create new MOI function
             new_func = MOIU.map_indices(index_map, src_func)
-            dest_index = Plasmo._add_element_constraint_to_backend(
+            dest_index = _add_element_constraint_to_backend(
                 dest,
                 new_cref,
                 new_func,
                 src_set
             )
+
+            # update index_map and ref_map
             index_map_FS[ci] = dest_index
             ref_map[src_cref] = new_cref
         end
+        # pass constraint attributes
+        MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map_FS, cis_src)
     end
 end
 
 """
     Aggregate an optiedge `source_edge` into new optinode `new_node`.
 """
 function _copy_edge_to!(
-    new_graph::OptiGraph, 
+    new_node::OptiNode, 
     source_edge::OptiEdge,
     ref_map::GraphReferenceMap
 )
     src = graph_backend(source_edge)
-    dest = graph_backend(new_graph)
-    index_map = MOIU.IndexMap()
-
-    # populate index_map
+    dest = graph_backend(new_node)
+
+    # setup variable index map
     index_map = MOIU.IndexMap()
     vars = all_variables(source_edge)
     for var in vars
         source_index = graph_index(src, var)
         dest_index = graph_index(dest, ref_map[var])
-        #source_index = src.element_to_graph_map[var]
-        #dest_index = dest.element_to_graph_map.var_map[ref_map[var]]
         index_map[source_index] = dest_index
     end
 
     # copy constraints
     _copy_constraints!(new_node, source_edge, ref_map, index_map)
-
-    # update ref_map constraints
-    for (source_index, dest_index) in index_map.con_map
-        source_cref = src.graph_to_element_map.con_map[source_index]
-        dest_cref = dest.graph_to_element_map.con_map[dest_index]
-        ref_map[source_cref] = dest_cref
-    end
-
     return 
 end
 
+# TODO
 """
     Aggregate an optiedge `source_edge` into new optinode `new_node`.
 """
 function _copy_edge_to!(
-    new_node::OptiNode, 
+    new_graph::OptiGraph, 
     source_edge::OptiEdge,
     ref_map::GraphReferenceMap
 )
-    _copy_edge_to!(source_graph(new_node), source_edge, ref_map) 
+    # create a new edge
+
+    # copy constraints from source_edge
 end
 
 function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
@@ -330,8 +333,8 @@ function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
 
         # copy optinodes
         for node in nodes
-            new_node, ref_map = _copy_node_to!(new_graph, node)
-            merge!(ref_map, ref_map)
+            new_node, node_ref_map = _copy_node_to!(new_graph, node)
+            merge!(ref_map, node_ref_map)
         end
 
         # copy optiedges
@@ -343,7 +346,6 @@ function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
     _copy_attributes_to!(root_optigraph, graph, ref_map)
 
     return root_optigraph, ref_map
-
 end
 
 """

src/backends/moi_backend.jl

@@ -149,7 +149,6 @@ function JuMP.backend(backend::GraphMOIBackend)
 end
 
 
-
 ### MOI Methods
 
 # graph attributes

src/core_types.jl

@@ -38,6 +38,9 @@ struct NodeIndex
     value::Symbol
 end
 
+# NOTE: a node or edge could technically have their `source_graph` 
+# changed using `apply_partition!`. that is why we use a RefValue here.
+
 struct OptiNode <: AbstractNode
     source_graph::Base.RefValue{<:OptiGraph}
     idx::NodeIndex

src/optielement.jl

@@ -4,7 +4,7 @@
 # MOI Methods
 #
 
-function MOI.get(element::OptiElement, attr::MOI.AnyAttribute)
+function MOI.get(element::OptiElement, attr::MOI.AnyAttribute, args...)
     return MOI.get(graph_backend(element), attr, element)
 end
 
@@ -144,7 +144,4 @@ function JuMP.all_constraints(
         append!(constraints, JuMP.all_constraints(element, function_type, set_type))
     end
     return constraints
-end
-
-
-
+end

src/optigraph.jl

@@ -35,6 +35,8 @@ function Base.getindex(graph::OptiGraph, idx::Int)
     return graph.optinodes[idx]
 end
 
+Base.broadcastable(graph::OptiGraph) = Ref(graph)
+
 # TODO: parameterize on numerical precision like JuMP Models do
 JuMP.value_type(::Type{OptiGraph})  = Float64
 
@@ -507,6 +509,36 @@ function JuMP.value(graph::OptiGraph, nvref::NodeVariableRef; result::Int = 1)
     return MOI.get(graph_backend(graph), MOI.VariablePrimal(result), nvref)
 end
 
+function JuMP.value(graph::OptiGraph, expr::JuMP.GenericAffExpr; result::Int = 1)
+    return JuMP.value(expr) do x
+        return JuMP.value(graph, x; result = result)
+    end
+end
+
+function JuMP.value(graph::OptiGraph, expr::JuMP.GenericQuadExpr; result::Int = 1)
+    return JuMP.value(expr) do x
+        return JuMP.value(graph, x; result = result)
+    end
+end
+
+function JuMP.value(graph::OptiGraph, expr::GenericNonlinearExpr; result::Int = 1)
+    return value(expr) do x
+        return value(graph, x; result = result)
+    end
+end
+
+### Expression values
+
+# function JuMP.value(var_value::Function, ex::GenericAffExpr{T,V}) where {T,V}
+#     S = Base.promote_op(var_value, V)
+#     U = Base.promote_op(*, T, S)
+#     ret = convert(U, ex.constant)
+#     for (var, coef) in ex.terms
+#         ret += coef * var_value(var)
+#     end
+#     return ret
+# end
+
 ### Constraints
 
 function JuMP.add_constraint(

src/optimizer_interface.jl

@@ -164,12 +164,6 @@ function JuMP.optimize!(
     return
 end
 
-function JuMP.set_optimizer(node::OptiNode; kwargs...)
-    error("Optinodes currently do not support setting an optimizer. Optimization is now 
-        handled through the optigraph. If you wish to optimize a single optinode, 
-        you can create a new optigraph using `assemble_optigraph(node)`.")
-end
-
 #
 # status results
 #
@@ -224,7 +218,7 @@ function JuMP.dual_status(graph::OptiGraph; result::Int = 1)
 end
 
 #
-# Optinode Interface
+# Optinode optimizer
 #
 
 function JuMP.set_optimizer(
@@ -238,8 +232,11 @@ function JuMP.set_optimizer(
     else
         node_graph = source_graph(node).node_graphs[node]
     end
+
+    # set objective on node graph
+    JuMP.set_objective(node_graph, objective_sense(node), objective_function(node))
     JuMP.set_optimizer(node_graph, optimizer_constructor; add_bridges=add_bridges)
-    return
+    return node_graph
 end
 
 function JuMP.optimize!(node::OptiNode; kwargs...)