got basic node copy_to working

jalving · Jan 9, 2024 · 936130b · 936130b
1 parent dbdeb39
commit 936130b
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 193 deletions.
diff --git a/src/Plasmo.jl b/src/Plasmo.jl
@@ -12,7 +12,7 @@ const MOI = MathOptInterface
 using Reexport
 @reexport using JuMP
 
-export OptiGraph, graph_backend
+export OptiGraph, graph_backend, graph_index
 
 abstract type AbstractOptiGraph <: JuMP.AbstractModel end
 

diff --git a/src/moi_aggregate.jl b/src/moi_aggregate.jl
@@ -19,41 +19,51 @@ end
 function _append_node_to_backend!(graph::OptiGraph, node::OptiNode)
     src = graph_backend(node)
     dest = graph_backend(graph)
-
-    # node_variables = all_variables(node)
-    # vis_src = graph_index.(node_variables) # variable indices on src graph
-
-    vis_src = MOI.get(src, node, MOI.ListOfConstraintIndices())
-
-    # TODO: get variable constraints specifically for this node
-    # variable_constraints = Any[
-    #     MOI.get(src, MOI.ListOfConstraintIndices{F,S}()) for
-    #     (F, S) in MOI.get(src, MOI.ListOfConstraintTypesPresent()) if
-    #     MOIU._is_variable_function(F)
-    # ]
-
     index_map = MOIU.IndexMap()
 
-    # copy node variables
-    _copy_node_variables(dest, index_map, node_variables)
-    # println(index_map)
+    # copy node variables and variable attributes
+    _copy_node_variables(dest, node, index_map)
 
-    # copy variable attributes (e.g. VariablePrimalStart(), VariableName())
-    MOI.Utilities.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src)
+    # copy constraints and constraint attributes
+    # NOTE: for now, we split between variable and non-variable, but they do the same thing
+    # eventually, we might try doing something more similar to MOI `default_copy_to` where
+    # we try to constraint variables on creation.
+    all_constraint_types = MOI.get(node, MOI.ListOfConstraintTypesPresent())
+    variable_constraint_types = filter(all_constraint_types) do (F, S)
+        return MOIU._is_variable_function(F)
+    end
+    _copy_node_constraints(
+        dest, 
+        node,
+        index_map, 
+        variable_constraint_types
+    )
+
+    # # copy non-variable constraints
+    nonvariable_constraint_types = filter(all_constraint_types) do (F, S)
+        return !MOIU._is_variable_function(F)
+    end
+    _copy_node_constraints(
+        dest,
+        node,
+        index_map,
+        nonvariable_constraint_types
+    )
 
-    _copy_variable_node_constraints(dest.moi_backend, src.moi_backend, index_map, vis_src)
+    # TODO: pass non-objective attributes (use MOI Filter?)
 
-    # TODO: constraints
-    # MOI.Utilities._pass_constraints(dest.moi_backend, src.moi_backend, index_map, constraints_not_added)
+    return
 end
 
-# TODO: MOI functions to get variables and constraints on nodes
-
+# TODO: update graph backend mappings
 function _copy_node_variables(
-    dest::GraphMOIBackend, 
-    index_map::MOIU.IndexMap,
-    node_variables::Vector{NodeVariableRef}
+    dest::GraphMOIBackend,
+    node::OptiNode,
+    index_map::MOIU.IndexMap
 )
+    src = graph_backend(node)
+    node_variables = all_variables(node)
+
     # map existing variables in index_map
     existing_vars = intersect(node_variables, keys(dest.node_to_graph_map.var_map))
     for var in existing_vars
@@ -69,168 +79,33 @@ function _copy_node_variables(
         dest_graph_index = _add_variable_to_backend(dest, var)
         index_map[src_graph_index] = dest_graph_index
     end
-    return
-end
-
-# TODO
 
-function _copy_variable_node_constraints(dest::GraphMOIBackend, src::OptiNode)
-    for cis in variable_constraints
-        _copy_constraints(dest, src, index_map, cis)
-    end
+    # pass variable attributes
+    vis_src = graph_index.(node_variables)
+    MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src)
+    return
 end
 
-function _copy_nonvariable_node_constraints(
-    dest::MOI.ModelLike,
-    src::MOI.ModelLike,
-    index_map::MOI.Utilities.IndexMap
+# TODO: update graph backend mappings
+function _copy_node_constraints(
+    dest::GraphMOIBackend, 
+    node::OptiNode,
+    index_map::MOIU.IndexMap,
+    constraint_types
 )
-
-    all_constraint_types = MOI.get(src, MOI.ListOfConstraintTypesPresent())
-    nonvariable_constraint_types = filter(all_constraint_types) do (F, S)
-        return !_is_variable_function(F)
+    src = graph_backend(node)
+    for (F, S) in constraint_types
+        cis_src = MOI.get(node, MOI.ListOfConstraintIndices{F,S}())
+        MOIU._copy_constraints(dest.moi_backend, src.moi_backend, index_map, cis_src)
     end
-
-    pass_nonvariable_constraints(
-        dest,
-        src,
-        index_map,
-        nonvariable_constraint_types,
-    )
-
-    # pass constraint attributes
-    for (F, S) in all_constraint_types
-        pass_attributes(
-            dest,
-            src,
+
+    for (F, S) in constraint_types
+        MOIU.pass_attributes(
+            dest.moi_backend,
+            src.moi_backend,
             index_map,
-            MOI.get(src, MOI.ListOfConstraintIndices{F,S}()),
+            MOI.get(node, MOI.ListOfConstraintIndices{F,S}()),
         )
     end
-
     return
-end
-
-# function MOIU.pass_attributes(
-#     dest::GraphMOIBackend, 
-#     src::GraphMOIBackend,
-#     index_map::MOIU.IndexMap
-# )
-#     MOIU.pass_attributes(dest.moi_backend, src.moi_backend, index_map)
-# end
-
-# function MOI.Utilities.pass_nonvariable_constraints(
-#     dest::GraphMOIBackend,
-#     src::OptiNode,
-#     index_map::MOIU.IndexMap,
-#     constraint_types,
-# )
-#   for (F, S) in constraint_types
-#       cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
-#       _copy_constraints(dest, src, index_map, cis_src)
-#     end
-
-#     return
-# end
-
-# function pass_nonvariable_constraints_fallback(
-#     dest::MOI.ModelLike,
-#     src::MOI.ModelLike,
-#     index_map::IndexMap,
-#     constraint_types,
-# )
-#     for (F, S) in constraint_types
-#         cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
-#         _copy_constraints(dest, src, index_map, cis_src)
-#     end
-#     return
-# end
-
-# function _copy_constraints(
-#     dest::MOI.ModelLike,
-#     src::MOI.ModelLike,
-#     index_map,
-#     index_map_FS,
-#     cis_src::Vector{<:MOI.ConstraintIndex},
-# )
-#     for ci in cis_src
-#         f = MOI.get(src, MOI.ConstraintFunction(), ci)
-#         s = MOI.get(src, MOI.ConstraintSet(), ci)
-#         index_map_FS[ci] =
-#             MOI.add_constraint(dest, map_indices(index_map, f), s)
-#     end
-#     return
-# end
-
-# function _copy_constraints(
-#     dest::MOI.ModelLike,
-#     src::MOI.ModelLike,
-#     index_map,
-#     cis_src::Vector{MOI.ConstraintIndex{F,S}},
-# ) where {F,S}
-#     return _copy_constraints(dest, src, index_map, index_map[F, S], cis_src)
-# end
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-### Helpful utilities
-
-# """
-#     append_to_backend!(dest::MOI.ModelLike, src::MOI.ModelLike)
-
-# Copy the underylying model from `src` into `dest`, but ignore attributes
-# such as objective function and objective sense
-# """
-# function append_to_backend!(dest::MOI.ModelLike, src::MOI.ModelLike)
-#     vis_src = MOI.get(src, MOI.ListOfVariableIndices())   #returns vector of MOI.VariableIndex
-#     index_map = MOIU.IndexMap()
-
-
-#     # has_nlp = MOI.NLPBlock() in MOI.get(src, MOI.ListOfModelAttributesSet())
-#     # constraints_not_added = if has_nlp
-#     constraints_not_added = Any[
-#         MOI.get(src, MOI.ListOfConstraintIndices{F,S}()) for
-#         (F, S) in MOI.get(src, MOI.ListOfConstraintTypesPresent()) if
-#         MOIU._is_variable_function(F)
-#     ]
-#     # else
-#     #     Any[
-#     #         MOIU._try_constrain_variables_on_creation(dest, src, index_map, S)
-#     #         for S in MOIU.sorted_variable_sets_by_cost(dest, src)
-#     #     ]
-#     # end
-
-#     # Copy free variables into graph optimizer
-#     MOI.Utilities._copy_free_variables(dest, index_map, vis_src)
-
-#     # Copy variable attributes (e.g. name, and VariablePrimalStart())
-#     MOI.Utilities.pass_attributes(dest, src, index_map, vis_src)
-
-#     # Normally this copies ObjectiveSense() and ObjectiveFunction(), but we don't want to do that here
-#     # MOI.Utilities.pass_attributes(dest, src, idxmap)
-
-#     MOI.Utilities._pass_constraints(dest, src, index_map, constraints_not_added)
-
-#     return index_map    #return an idxmap for each source model
-# end
+end
diff --git a/src/moi_graph_backend.jl b/src/moi_graph_backend.jl
@@ -144,7 +144,7 @@ function _moi_add_node_variable(
     # add variable to all containing optigraphs
     for graph in containing_optigraphs(node)
         graph_var_index = _add_variable_to_backend(graph_backend(graph), vref)
-        push!(graph_backend(graph).node_variables[node], graph_var_index) # TODO: move
+        #push!(graph_backend(graph).node_variables[node], graph_var_index) # TODO: move
         _moi_constrain_node_variable(
             graph_backend(graph),
             vref,
@@ -204,6 +204,11 @@ function _add_variable_to_backend(
     graph_var_index = MOI.add_variable(graph_backend.moi_backend)
     graph_backend.node_to_graph_map[vref] = graph_var_index
     graph_backend.graph_to_node_map[graph_var_index] = vref
+
+    if !haskey(graph_backend.node_variables, vref.node)
+        graph_backend.node_variables[vref.node] = MOI.VariableIndex[]
+    end
+    push!(graph_backend.node_variables[vref.node], graph_var_index)
     return graph_var_index
 end
 
@@ -235,8 +240,6 @@ function _moi_add_node_constraint(
     cref = ConstraintRef(node, constraint_index, JuMP.shape(con))
 
     for graph in containing_optigraphs(node)
-        # moi_func_graph = deepcopy(moi_func)
-
         # update func variable indices
         moi_func_graph = _create_graph_moi_func(graph_backend(graph), moi_func, jump_func)
 

diff --git a/src/optinode.jl b/src/optinode.jl
@@ -87,16 +87,16 @@ end
 
 # TODO: determine if caching node references is possible without dict-of-dicts
 function JuMP.all_variables(node::OptiNode)
-    return collect(
-        filter(var -> var.node == node, keys(graph_backend(node).node_to_graph_map.var_map))
-    )
+    gb = graph_backend(node)
+    graph_indices = gb.node_variables[node]
+    return getindex.(Ref(gb.graph_to_node_map), graph_indices)
 end
 
 function JuMP.num_variables(node::OptiNode)
-    n2g = graph_backend(node).node_to_graph_map
-    return length(filter((vref) -> vref.node == node, keys(n2g.var_map)))
+    return length(graph_backend(node).node_variables[node])
 end
 
+# TODO: update
 function JuMP.num_constraints(
     node::OptiNode,
     ::Type{F}, 
@@ -167,8 +167,6 @@ function JuMP.set_name(vref::NodeVariableRef, s::String)
     return
 end
 
-
-
 ### Node Constraints
 
 """
@@ -182,4 +180,32 @@ function JuMP.add_constraint(
     con = JuMP.model_convert(node, con)
     cref = _moi_add_node_constraint(node, con)
     return cref
+end
+
+### MOI extensions for optinode
+# TODO: cache constraint types in graph backend versus using unique and filters?
+
+function MOI.get(node::OptiNode, attr::MOI.ListOfConstraintTypesPresent)
+    cons = graph_backend(node).node_constraints[node]
+    con_types = unique(typeof.(cons))
+    type_tuple = [(type.parameters[1],type.parameters[2]) for type in con_types]  
+    return type_tuple
+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).node_constraints[node]
+        if (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S)
+            push!(con_inds, con)
+        end
+    end
+    return con_inds
+    # return collect(filter(
+    #     con -> (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S),
+    #     graph_backend(node).node_constraints[node]
+    # ))
 end