got aggregate_to_depth working

src/aggregate.jl

@@ -9,6 +9,17 @@ struct GraphReferenceMap
     # map variables and from original optigraph to new aggregate node or graph
     var_map::OrderedDict{NodeVariableRef,NodeVariableRef}
     con_map::OrderedDict{JuMP.ConstraintRef,JuMP.ConstraintRef}
+    node_map::OrderedDict{OptiNode,OptiNode}
+    edge_map::OrderedDict{OptiEdge,OptiEdge}
+end
+
+function GraphReferenceMap()
+    return GraphReferenceMap(
+        OrderedDict{NodeVariableRef,NodeVariableRef}(),
+        OrderedDict{JuMP.ConstraintRef,JuMP.ConstraintRef}(),
+        OrderedDict{OptiNode,OptiNode}(),
+        OrderedDict{OptiEdge,OptiEdge}(),
+    )
 end
 
 function Base.getindex(ref_map::GraphReferenceMap, vref::NodeVariableRef)
@@ -40,15 +51,12 @@ function Base.setindex!(
     return ref_map.var_map[node_vref] = graph_vref
 end
 
-function GraphReferenceMap()
-    return GraphReferenceMap(
-        OrderedDict{NodeVariableRef,NodeVariableRef}(),
-        OrderedDict{JuMP.ConstraintRef,JuMP.ConstraintRef}(),
-    )
-end
 function Base.merge!(ref_map1::GraphReferenceMap, ref_map2::GraphReferenceMap)
     merge!(ref_map1.var_map, ref_map2.var_map)
-    return merge!(ref_map1.con_map, ref_map2.con_map)
+    merge!(ref_map1.con_map, ref_map2.con_map)
+    merge!(ref_map1.node_map, ref_map2.node_map)
+    merge!(ref_map1.edge_map, ref_map2.edge_map)
+    return nothing
 end
 
 ### Aggregate Functions
@@ -127,8 +135,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)
+    MOIU.pass_attributes(dest, src, index_map)
     return nothing
 end
 
@@ -150,13 +157,16 @@ function _copy_node_to!(
 
     _copy_variables!(new_node, source_node, ref_map, index_map)
 
-    return _copy_constraints!(new_node, source_node, ref_map, index_map)
+    _copy_constraints!(new_node, source_node, ref_map, index_map)
+
+    return nothing
 end
 
 """
     _copy_node_to!(new_graph::OptiGraph, source_node::OptiNode)
 
-Copy an optinode to a new optigraph.
+Copy an optinode to a new optigraph. Creates the new node and then calls the 
+underlying method to copy to a node.
 """
 function _copy_node_to!(new_graph::OptiGraph, source_node::OptiNode)
     ref_map = GraphReferenceMap()
@@ -187,7 +197,7 @@ function _copy_variables!(
 
     # pass variable attributes
     vis_src = graph_index.(source_variables)
-    MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src)
+    MOIU.pass_attributes(dest, src, index_map, vis_src)
 
     # set new variable names
     for nvref in all_variables(source_node)
@@ -217,19 +227,13 @@ function _copy_constraints!(
     dest = graph_backend(new_node)
 
     # copy each constraint by iterating through each type
-    constraint_types = MOI.get(src.moi_backend, MOI.ListOfConstraintTypesPresent())
+    constraint_types = JuMP.list_of_constraint_types(source_element)
     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)
+        for src_cref in JuMP.all_constraints(source_element, F, S)
+            # get source constraint data
+            src_func = MOI.get(source_element, MOI.ConstraintFunction(), src_cref)
+            src_set = MOI.get(source_element, MOI.ConstraintSet(), src_cref)
             new_shape = src_cref.shape
 
             # create a new ConstraintRef
@@ -243,11 +247,66 @@ function _copy_constraints!(
             dest_index = MOI.add_constraint(dest, new_cref, new_func, src_set)
 
             # update index_map and ref_map
-            index_map_FS[ci] = dest_index
+            index_map_FS[graph_index(src_cref)] = dest_index
+            ref_map[src_cref] = new_cref
+        end
+        # pass constraint attributes
+        F_moi = JuMP.moi_function_type(F)
+        cis_src = MOI.get(source_element, MOI.ListOfConstraintIndices{F_moi,S}())
+        MOIU.pass_attributes(dest, src, index_map_FS, cis_src)
+    end
+end
+
+"""
+    _copy_constraints!(
+        new_edge::OptiEdge,
+        source_edge::OptiEdge,
+        ref_map::GraphReferenceMap,
+        index_map::MOIU.IndexMap,
+    )
+
+Copy the constraints from a node or edge into a new node.
+"""
+function _copy_constraints!(
+    new_edge::OptiEdge,
+    source_edge::OptiEdge,
+    ref_map::GraphReferenceMap,
+    index_map::MOIU.IndexMap,
+)
+    source_element = source_edge
+
+    # get relevant backends
+    src = graph_backend(source_element)
+    dest = graph_backend(new_edge)
+
+    # copy each constraint by iterating through each type
+    constraint_types = JuMP.list_of_constraint_types(source_element)
+    for (F, S) in constraint_types
+        index_map_FS = index_map[F, S]
+        for src_cref in JuMP.all_constraints(source_element, F, S)
+            # get source constraint data
+            src_func = MOI.get(source_element, MOI.ConstraintFunction(), src_cref)
+            src_set = MOI.get(source_element, MOI.ConstraintSet(), src_cref)
+            new_shape = src_cref.shape
+
+            # create a new ConstraintRef
+            new_constraint_index = next_constraint_index(
+                new_edge, typeof(src_func), typeof(src_set)
+            )::MOI.ConstraintIndex{typeof(src_func),typeof(src_set)}
+            new_cref = ConstraintRef(new_edge, new_constraint_index, new_shape)
+
+            # create new MOI function
+            new_func = MOIU.map_indices(index_map, src_func)
+            dest_index = MOI.add_constraint(dest, new_cref, new_func, src_set)
+
+            # update index_map and ref_map
+            index_map_FS[graph_index(src_cref)] = 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)
+        F_moi = JuMP.moi_function_type(F)
+        cis_src = MOI.get(source_element, MOI.ListOfConstraintIndices{F_moi,S}())
+        MOIU.pass_attributes(dest, src, index_map_FS, cis_src)
     end
 end
 
@@ -287,9 +346,21 @@ end
 function _copy_edge_to!(
     new_graph::OptiGraph, source_edge::OptiEdge, ref_map::GraphReferenceMap
 )
-    # create a new edge
+    source_nodes = all_nodes(source_edge)
+    new_nodes = Base.getindex.(Ref(ref_map.node_map), source_nodes)
+    new_edge = add_edge(new_graph, new_nodes...)
 
-    # copy constraints from source_edge
+    # setup variable index map
+    index_map = MOIU.IndexMap()
+    vars = all_variables(source_edge)
+    for var in vars
+        source_index = graph_index(graph_backend(source_edge), var)
+        dest_index = graph_index(graph_backend(new_edge), ref_map[var])
+        index_map[source_index] = dest_index
+    end
+
+    # copy constraints from source_edge to new edge
+    return _copy_constraints!(new_edge, source_edge, ref_map, index_map)
 end
 
 """
@@ -301,12 +372,12 @@ how many levels of subgraphs remain in the new aggregated optigraph. For example
 Return a new aggregated optigraph and reference map that maps elements from the old 
 optigraph to the new aggregate optigraph.
 """
-function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
+function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0; name=gensym())
     if num_subgraphs(graph) == 0
         error("`aggregate_to_depth` requires the graph to contain subgraphs.")
     end
 
-    root_optigraph = OptiGraph()
+    root_optigraph = OptiGraph(; name=name)
     ref_map = GraphReferenceMap()
     subgraph_dict = Dict(graph => root_optigraph)
 
@@ -352,7 +423,7 @@ function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
         merge!(ref_map, sub_ref_map)
     end
 
-    #now copy nodes and edges going back up the tree
+    # now copy nodes and edges going back up the tree
     for rgraph in reverse(all_parents)
         nodes = local_nodes(rgraph)
         edges = local_edges(rgraph)
@@ -362,11 +433,14 @@ function aggregate_to_depth(graph::OptiGraph, max_depth::Int64=0)
         for node in nodes
             new_node, node_ref_map = _copy_node_to!(new_graph, node)
             merge!(ref_map, node_ref_map)
+            ref_map.node_map[node] = new_node
         end
 
         # copy optiedges
         for edge in edges
             _copy_edge_to!(new_graph, edge, ref_map)
+            # new_edge, edge_ref_map = 
+            # ref_map[edge] = new_edge
         end
     end
 
@@ -384,8 +458,11 @@ function _aggregate_subgraphs!(new_graph::OptiGraph, source_graph::OptiGraph)
     # aggregate each subgraph into a node in new_graph
     nodes, ref_maps = [], []
     for subgraph in local_subgraphs(source_graph)
-        node, ref_map = _copy_graph_elements_to!(new_graph, subgraph)
-        push!(nodes, node)
+        new_node, ref_map = _copy_graph_elements_to!(new_graph, subgraph)
+        for node in all_nodes(subgraph)
+            ref_map.node_map[node] = new_node
+        end
+        push!(nodes, new_node)
         push!(ref_maps, ref_map)
     end
     return nodes, ref_maps

src/backends/moi_backend.jl

@@ -252,7 +252,6 @@ end
 function MOI.get(
     backend::GraphMOIBackend, attr::MOI.NumberOfConstraints{F,S}, element::OptiElement
 ) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
-    # filter F and S
     n_cons = 0
     for con in backend.element_constraints[element]
         if (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S)
@@ -702,6 +701,7 @@ function _copy_subgraph_edges!(backend::GraphMOIBackend, subgraph::OptiGraph)
     end
 end
 
+# TODO: re-implement MOIU._is_variable_function
 function _copy_node_to_backend!(backend::GraphMOIBackend, node::OptiNode)
     dest = backend
     index_map = MOIU.IndexMap()
@@ -768,7 +768,6 @@ function _copy_node_variables(
     vars_to_add = setdiff(node_variables, keys(dest.element_to_graph_map.var_map))
     for var in vars_to_add
         src_graph_index = graph_index(var)
-        #dest_graph_index = _add_variable_to_backend(dest, var)
         dest_graph_index = MOI.add_variable(dest, var)
         index_map[src_graph_index] = dest_graph_index
     end
@@ -833,3 +832,30 @@ function _copy_element_constraints(
     end
     return nothing
 end
+
+#
+# MOIU Utility Methods
+#
+function MOIU.pass_attributes(
+    dest::GraphMOIBackend, src::GraphMOIBackend, index_map::MOIU.IndexMap
+)
+    return MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map)
+end
+
+function MOIU.pass_attributes(
+    dest::GraphMOIBackend,
+    src::GraphMOIBackend,
+    index_map::MOIU.IndexMap,
+    vis_src::Vector{MOI.VariableIndex},
+)
+    return MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src)
+end
+
+function MOIU.pass_attributes(
+    dest::GraphMOIBackend,
+    src::GraphMOIBackend,
+    index_map::MOIU.IndexMap,
+    cis_src::Vector{MOI.ConstraintIndex{F,S}},
+) where {F,S}
+    return MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map, cis_src)
+end

test/test_aggregate.jl

@@ -20,6 +20,18 @@ function _create_test_optigraph()
     return graph
 end
 
+function _create_nested_test_optigraph()
+    graph = OptiGraph()
+    for _ in 1:4
+        add_subgraph(graph, _create_test_optigraph())
+    end
+    subs = local_subgraphs(graph)
+    for i in 1:3
+        @linkconstraint(graph, subs[i + 1][1][:x] == subs[i][10][:x])
+    end
+    return graph
+end
+
 function _create_test_model()
     model = Model()
     @variable(model, x[1:10] >= 0)
@@ -64,22 +76,16 @@ function test_set_model()
     @test value.(all_variables(m)) == value.(graph, all_variables(n1))
 end
 
-# TODO
-# function test_aggregate_to_depth()
-#     graph = _create_test_optigraph_w_subgraphs()
-#     new_graph, ref = aggregate(graph, 0)
-#     @test num_all_nodes(new_graph) == 5
-#     @test num_all_linkconstraints(new_graph) == 4
-
-#     aggregate!(graph, 0)
-#     @test num_all_nodes(graph) == 5
-#     @test num_all_linkconstraints(graph) == 4
-
-#     #TODO: more checks
-#     graph = _create_test_optigraph_w_recursive_subgraphs()
-#     new_graph, ref = aggregate(graph, 1)
-#     @test num_all_subgraphs(new_graph) == 5
-# end
+function test_aggregate_to_depth()
+    graph = _create_nested_test_optigraph()
+    agg_graph, ref_map = aggregate_to_depth(graph, 0; name=:agg_graph)
+    @test num_nodes(agg_graph) == 4
+    @test num_edges(agg_graph) == 3
+    @test num_subgraphs(agg_graph) == 0
+    @test num_variables(agg_graph) == 80
+    @test num_constraints(agg_graph) == 319
+    @test num_local_link_constraints(agg_graph) == 3
+end
 
 function run_tests()
     for name in names(@__MODULE__; all=true)