wip: writing copy functions for aggregate

src/aggregate.jl

@@ -6,10 +6,8 @@ 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 in original optigraph to optinode
+    #map variables and from original optigraph to new aggregate node or graph
     var_map::OrderedDict{NodeVariableRef,NodeVariableRef}
-
-    #map constraints in original optigraph to optinode
     con_map::OrderedDict{JuMP.ConstraintRef,JuMP.ConstraintRef}  
 end
 
@@ -61,17 +59,36 @@ end
 Aggregate an optigraph into a graph containing a single optinode.
 """
 function aggregate(graph::OptiGraph)
-    # aggregate into a graph containing a single node.
+    # aggregate into a graph containing a single optinode
     new_graph = OptiGraph()
     new_node, ref_map = _copy_graph_elements_to!(new_graph, graph)
-
-    # pass model attributes
-    # copy other model attributes (including objective function)
+    # copy other model attributes (such as the objective function)
     # setup index_map from the ref_map
     _copy_attributes_to!(new_graph, graph, ref_map)
     return new_node, ref_map
 end
 
+"""
+    Aggregate `source_graph` into an optinode within new graph. Return new node and mapping.
+"""
+function _copy_graph_elements_to!(new_graph::OptiGraph, source_graph::OptiGraph)
+    new_node = add_node(
+        new_graph; 
+        label=Symbol(new_graph.label, Symbol(".n"),length(new_graph.optinodes)+1)
+    )
+    ref_map = GraphReferenceMap()
+
+    for node in all_nodes(source_graph)
+        _copy_node_to!(new_node, node, ref_map)
+    end
+
+    for edge in all_edges(source_graph)
+        _copy_edge_to!(new_node, edge, ref_map)
+    end
+
+    return new_node, ref_map
+end
+
 """
     Copy graph attributes from optigraph `source_graph` to the new optigraph `new_graph`.
 """
@@ -88,6 +105,7 @@ function _copy_attributes_to!(
     # the variables or constraints. we just want to pass the attributes that would be 
     # exposed such as the 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)
             index_map[graph_index(source_graph, source_vref)] = graph_index(dest_vref)
         end
@@ -103,117 +121,138 @@ end
 
 
 """
-    Aggregate `source_graph` into an optinode within new graph. Return new node and mapping.
+    Copy (append) the `source_node` backend model into `new_node`.
 """
-function _copy_graph_elements_to!(new_graph::OptiGraph, source_graph::OptiGraph)
-    new_node = add_node(new_graph)
-    ref_map = GraphReferenceMap()
+function _copy_node_to!(
+    new_node::OptiNode,
+    source_node::OptiNode, 
+    ref_map::GraphReferenceMap
+)
+    src = graph_backend(source_node)
+    dest = graph_backend(new_node)
+    index_map = MOIU.IndexMap()
 
-    for node in all_nodes(source_graph)
-        _copy_node_to!(new_node, node, ref_map)
-    end
+    _copy_variables!(new_node, source_node, ref_map, index_map)
 
-    for edge in all_edges(source_graph)
-        _copy_edge_to!(new_node, edge, ref_map)
-    end
+    _copy_constraints!(new_node, source_node, ref_map, index_map)
 
-    return new_node, ref_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
 
 """
-    Aggregate an optinode `source_node` into new optinode `new_node`.
+    Copy an optinode to a new optigraph
 """
-function _copy_node_to!(
+function _copy_node_to!(new_graph::OptiGraph, source_node::OptiNode)
+    ref_map = GraphReferenceMap()
+    new_node = add_node(new_graph)
+    _copy_node_to!(new_node, source_node, ref_map)
+    return new_node, ref_map
+end
+
+function _copy_variables!(
     new_node::OptiNode, 
     source_node::OptiNode, 
-    ref_map::GraphReferenceMap
+    ref_map::GraphReferenceMap,
+    index_map::MOIU.IndexMap
 )
     src = graph_backend(source_node)
     dest = graph_backend(new_node)
-    index_map = MOIU.IndexMap()
-
-    # create new variable references
+
+    # create new variables
     source_variables = all_variables(source_node)
-    new_vars = NodeVariableRef[]
-    for vref in source_variables
+    for nvref in source_variables
         new_variable_index = next_variable_index(new_node)
         new_vref = NodeVariableRef(new_node, new_variable_index)
         _add_variable_to_backend(dest, new_vref)
-        index_map[graph_index(vref)] = graph_index(new_vref)
-        ref_map[vref] = new_vref
+        index_map[graph_index(nvref)] = graph_index(new_vref)
+        ref_map[nvref] = new_vref
     end
 
     # pass variable attributes
     vis_src = graph_index.(source_variables)
     MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src)
 
-    # copy constraints
-    constraint_types = MOI.get(source_node, MOI.ListOfConstraintTypesPresent())
-    _copy_element_constraints(
-        dest,
-        source_node,
-        index_map,
-        constraint_types
-    )
+    # 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
 
-    # update aggregation 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
 
-"""
-    Aggregate an optiedge `source_edge` into new optinode `new_node`.
-"""
-function _copy_edge_to!(
+function _copy_constraints!(
     new_node::OptiNode, 
-    source_edge::OptiEdge,
-    ref_map::GraphReferenceMap
+    source_element::OptiElement, 
+    ref_map::GraphReferenceMap,
+    index_map::MOIU.IndexMap
 )
-    _copy_edge_to!(source_graph(new_node), source_edge, ref_map) 
-end
-
-"""
-    Copy an optinode to a new optigraph
-"""
-function _copy_node_to!(new_graph::OptiGraph, source_node::OptiNode)
-    ref_map = GraphReferenceMap()
-    new_node = add_node(new_graph)
-    _add_to_aggregate_node!(new_node, source_node)
-    return new_node, ref_map
+    src = graph_backend(source_element)
+    dest = graph_backend(new_node)
+    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
+            src_func = MOI.get(src.moi_backend, MOI.ConstraintFunction(), ci)
+            src_set = MOI.get(src.moi_backend, 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(
+                new_node, 
+                typeof(src_func), 
+                typeof(src_set)
+            )::MOI.ConstraintIndex{typeof(src_func),typeof(src_set)}
+            new_cref = ConstraintRef(new_node, new_constraint_index, new_shape);
+
+            # create new MOI function
+            new_func = MOIU.map_indices(index_map, src_func)
+            dest_index = Plasmo._add_element_constraint_to_backend(
+                dest,
+                new_cref,
+                new_func,
+                src_set
+            )
+            index_map_FS[ci] = dest_index
+            ref_map[src_cref] = new_cref
+        end
+    end
 end
 
 """
     Aggregate an optiedge `source_edge` into new optinode `new_node`.
 """
 function _copy_edge_to!(
-    graph::OptiGraph, 
+    new_graph::OptiGraph, 
     source_edge::OptiEdge,
     ref_map::GraphReferenceMap
 )
     src = graph_backend(source_edge)
-    dest = graph_backend(graph)
+    dest = graph_backend(new_graph)
     index_map = MOIU.IndexMap()
 
-    # populate edge index_map
+    # populate index_map
     index_map = MOIU.IndexMap()
     vars = all_variables(source_edge)
     for var in vars
-        source_index = src.element_to_graph_map[var]
-        dest_index = dest.element_to_graph_map.var_map[ref_map[var]]
+        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
-    constraint_types = MOI.get(source_edge, MOI.ListOfConstraintTypesPresent())
-    _copy_element_constraints(
-        dest,
-        source_edge,
-        index_map,
-        constraint_types
-    )
+    _copy_constraints!(new_node, source_edge, ref_map, index_map)
 
     # update ref_map constraints
     for (source_index, dest_index) in index_map.con_map
@@ -225,6 +264,17 @@ function _copy_edge_to!(
     return 
 end
 
+"""
+    Aggregate an optiedge `source_edge` into new optinode `new_node`.
+"""
+function _copy_edge_to!(
+    new_node::OptiNode, 
+    source_edge::OptiEdge,
+    ref_map::GraphReferenceMap
+)
+    _copy_edge_to!(source_graph(new_node), source_edge, ref_map) 
+end
+
 function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
     root_optigraph = OptiGraph()
     ref_map = GraphReferenceMap()

src/backends/moi_backend.jl

@@ -591,6 +591,8 @@ function _add_backend_variables(
 end
 
 ### Aggregate MOI backends
+# Note that these methods do not create copies of nodes or edges; they create model 
+# data for new backends. The nodes and edges will then reference data for multiple backends. 
 
 """
     _copy_subgraph_backends!(graph::OptiGraph)
@@ -788,15 +790,17 @@ function _copy_element_constraints(
 )
     src = graph_backend(element)
     for ci in cis_src
-        f = MOI.get(src.moi_backend, MOI.ConstraintFunction(), ci)
-        s = MOI.get(src.moi_backend, MOI.ConstraintSet(), ci)
+        func = MOI.get(src.moi_backend, MOI.ConstraintFunction(), ci)
+        set = MOI.get(src.moi_backend, MOI.ConstraintSet(), ci)
         cref = src.graph_to_element_map[ci]
+
+        # avoid creating duplicate constraints if destination has this reference already.
         cref in keys(dest.element_to_graph_map.con_map) && return
         dest_index = _add_element_constraint_to_backend(
             dest,
             cref,
-            MOIU.map_indices(index_map, f), 
-            s
+            MOIU.map_indices(index_map, func),
+            set
         )
         index_map_FS[ci] = dest_index
     end

src/node_variables.jl

@@ -146,6 +146,11 @@ function JuMP.is_valid(node::OptiNode, nvref::NodeVariableRef)
            MOI.is_valid(graph_backend(node), nvref)
 end
 
+function JuMP.is_valid(node::OptiNode, cref::ConstraintRef)
+    return node === JuMP.owner_model(cref) &&
+           MOI.is_valid(graph_backend(node), cref)
+end
+
 function JuMP.owner_model(nvref::NodeVariableRef)
     return nvref.node
 end

src/optinode.jl

@@ -182,7 +182,7 @@ function _moi_add_node_constraint(
         _add_element_constraint_to_backend(
             graph_backend(graph), 
             cref, 
-            moi_func_graph, 
+            moi_func_graph,
             moi_set
         )
     end
@@ -240,7 +240,7 @@ function set_model(node::OptiNode, model::JuMP.Model)
     #_copy_model_to!(node, model)
 end
 
-# TODO
+# TODO: copy JuMP models to optinodes
 # function _copy_model_to!(node::OptiNode, model::JuMP.Model)
 #     if !(num_variables(node) == 0 && num_constraints(node) == 0)
 #         error("An optinode must be empty to set a JuMP Model.")