From 6f1f64217998578083d1960841ecef8dc59a2a6a Mon Sep 17 00:00:00 2001 From: jalving Date: Sun, 7 Jan 2024 20:56:45 -0800 Subject: [PATCH] working on moi aggregate --- src/Plasmo.jl | 12 +- src/aggregate.jl | 50 ---- src/dev3.jl | 5 +- src/jump_interop.jl | 44 ---- src/moi_aggregate.jl | 235 ++++++++++++++++++ ...{graph_backend.jl => moi_graph_backend.jl} | 92 ++++++- src/optigraph.jl | 1 + src/optinode.jl | 24 +- 8 files changed, 349 insertions(+), 114 deletions(-) delete mode 100644 src/aggregate.jl create mode 100644 src/moi_aggregate.jl rename src/{graph_backend.jl => moi_graph_backend.jl} (81%) diff --git a/src/Plasmo.jl b/src/Plasmo.jl index 13b56b9..3d76258 100644 --- a/src/Plasmo.jl +++ b/src/Plasmo.jl @@ -12,13 +12,7 @@ const MOI = MathOptInterface using Reexport @reexport using JuMP -export OptiGraph - -# import JuMP: AbstractModel, AbstractConstraint, AbstractJuMPScalar, ConstraintRef -# import Base: ==, show, print, string, getindex, copy -# import LightGraphs: AbstractGraph, AbstractEdge, Graph -# import DataStructures.OrderedDict -# import Base: ==, string, print, show +export OptiGraph, graph_backend abstract type AbstractOptiGraph <: JuMP.AbstractModel end @@ -28,7 +22,9 @@ include("optiedge.jl") include("optigraph.jl") -include("graph_backend.jl") +include("moi_graph_backend.jl") + +include("moi_aggregate.jl") include("optimizer_interface.jl") diff --git a/src/aggregate.jl b/src/aggregate.jl deleted file mode 100644 index e16fe44..0000000 --- a/src/aggregate.jl +++ /dev/null @@ -1,50 +0,0 @@ -""" - aggregate_backends!(graph::OptiGraph) - -Aggregate the moi backends from each subgraph within `graph` to create a single backend. -""" -function aggregate_backends!(graph::OptiGraph) - dest = JuMP.backend(graph) - -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 \ No newline at end of file diff --git a/src/dev3.jl b/src/dev3.jl index ce0cfe5..b846ab2 100644 --- a/src/dev3.jl +++ b/src/dev3.jl @@ -19,6 +19,7 @@ n2 = Plasmo.add_node(sg1) @variable(n2, y >= 2) @constraint(n2, ref2, n2[:x] + n2[:y] <= 4) + # linking constraint on subgraph1 edge1 = Plasmo.add_edge(sg1, n1, n2) @constraint(edge1, ref_edge_1, n1[:x] == n2[:x]) @@ -72,4 +73,6 @@ optimize!(sg2) # @show value(n2[:y]) # @show value(n4[:y]) -# println(objective_value(graph)) \ No newline at end of file +# println(objective_value(graph)) + +# Plasmo._append_node_to_backend!(graph, n1) \ No newline at end of file diff --git a/src/jump_interop.jl b/src/jump_interop.jl index 35eb978..245966b 100644 --- a/src/jump_interop.jl +++ b/src/jump_interop.jl @@ -1,48 +1,4 @@ ### directly copied functions from JuMP -# TODO: attribute file -function _moi_constrain_node_variable( - gb::GraphMOIBackend, - index, - info, - ::Type{T}, -) where {T} - if info.has_lb - _moi_add_constraint( - gb.moi_backend, - index, - MOI.GreaterThan{T}(info.lower_bound), - ) - end - if info.has_ub - _moi_add_constraint( - gb.moi_backend, - index, - MOI.LessThan{T}(info.upper_bound), - ) - end - if info.has_fix - _moi_add_constraint( - gb.moi_backend, - index, - MOI.EqualTo{T}(info.fixed_value), - ) - end - if info.binary - _moi_add_constraint(gb.moi_backend, index, MOI.ZeroOne()) - end - if info.integer - _moi_add_constraint(gb.moi_backend, index, MOI.Integer()) - end - if info.has_start && info.start !== nothing - MOI.set( - gb.moi_backend, - MOI.VariablePrimalStart(), - index, - convert(T, info.start), - ) - end -end - # copied from: https://github.com/jump-dev/JuMP.jl/blob/0df25a9185ceede762af533bc965c9374c97450c/src/constraints.jl function _moi_add_constraint( model::MOI.ModelLike, diff --git a/src/moi_aggregate.jl b/src/moi_aggregate.jl new file mode 100644 index 0000000..81f4a32 --- /dev/null +++ b/src/moi_aggregate.jl @@ -0,0 +1,235 @@ +""" + aggregate_backends!(graph::OptiGraph) + +Aggregate the moi backends from each subgraph within `graph` to create a single backend. +""" +function aggregate_backends!(graph::OptiGraph) + for subgraph in get_subgraphs(graph) + _aggregate_subgraph_nodes!(subgraph) + _aggregate_subgraph_edges!(subgraph) + end +end + +function _aggregate_subgraph_nodes!(graph::OptiGraph) + for node in all_nodes(graph) + _append_node_to_backend!(graph, node) + end +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 + + # 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 variable attributes (e.g. VariablePrimalStart(), VariableName()) + MOI.Utilities.pass_attributes(dest.moi_backend, src.moi_backend, index_map, vis_src) + + _copy_variable_node_constraints(dest.moi_backend, src.moi_backend, index_map, vis_src) + + # TODO: constraints + # MOI.Utilities._pass_constraints(dest.moi_backend, src.moi_backend, index_map, constraints_not_added) +end + +# TODO: MOI functions to get variables and constraints on nodes + +function _copy_node_variables( + dest::GraphMOIBackend, + index_map::MOIU.IndexMap, + node_variables::Vector{NodeVariableRef} +) + # map existing variables in index_map + existing_vars = intersect(node_variables, keys(dest.node_to_graph_map.var_map)) + for var in existing_vars + src_graph_index = graph_index(var) + dest_graph_index = dest.node_to_graph_map[var] + index_map[src_graph_index] = dest_graph_index + end + + # create and add new variables + vars_to_add = setdiff(node_variables, keys(dest.node_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) + index_map[src_graph_index] = dest_graph_index + end + return +end + +function _copy_variable_node_constraints(dest::GraphMOIBackend, src::OptiNode) + for cis in variable_constraints + _copy_constraints(dest, src, index_map, cis) + end +end + +function _copy_nonvariable_node_constraints( + dest::MOI.ModelLike, + src::MOI.ModelLike, + index_map::MOI.Utilities.IndexMap +) + + all_constraint_types = MOI.get(src, MOI.ListOfConstraintTypesPresent()) + nonvariable_constraint_types = filter(all_constraint_types) do (F, S) + return !_is_variable_function(F) + 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, + index_map, + MOI.get(src, 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 \ No newline at end of file diff --git a/src/graph_backend.jl b/src/moi_graph_backend.jl similarity index 81% rename from src/graph_backend.jl rename to src/moi_graph_backend.jl index 4655f75..0680220 100644 --- a/src/graph_backend.jl +++ b/src/moi_graph_backend.jl @@ -12,6 +12,18 @@ function NodeToGraphMap() ) end +# mutable struct NodeToGraphMap +# var_map::OrderedDict{NodeVariableRef,MOI.VariableIndex} +# con_map::OrderedDict{ConstraintRef,MOI.ConstraintIndex} +# end +# function NodeToGraphMap() +# return NodeToGraphMap( +# OrderedDict{NodeVariableRef,MOI.VariableIndex}(), +# OrderedDict{ConstraintRef,MOI.ConstraintIndex}(), +# ) +# end + + function Base.setindex!(n2g_map::NodeToGraphMap, idx::MOI.VariableIndex, vref::NodeVariableRef) n2g_map.var_map[vref] = idx return @@ -68,19 +80,20 @@ end # try to support Direct, Manual, and Automatic modes on an optigraph. mutable struct GraphMOIBackend <: MOI.AbstractOptimizer optigraph::AbstractOptiGraph - # TODO: legacy nlp model + # TODO (maybe): legacy nlp model # nlp_model::MOI.Nonlinear.Model moi_backend::MOI.AbstractOptimizer node_to_graph_map::NodeToGraphMap graph_to_node_map::GraphToNodeMap + node_variables::OrderedDict{OptiNode,Vector{MOI.VariableIndex}} + node_constraints::OrderedDict{OptiNode,Vector{MOI.ConstraintIndex}} end """ GraphMOIBackend() -Initialize an empty optigraph backend. Contains a model_cache that can be used to set -`MOI.AbstractModelAttribute`s and `MOI.AbstractOptimizerAttribute`s. By default we -use a `CachingOptimizer` to store the underlying optimizer. +Initialize an empty optigraph backend that uses MOI. +By default we use a `CachingOptimizer` to store the underlying optimizer just like JuMP. """ function GraphMOIBackend(optigraph::AbstractOptiGraph) inner = MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()) @@ -89,10 +102,17 @@ function GraphMOIBackend(optigraph::AbstractOptiGraph) optigraph, cache, NodeToGraphMap(), - GraphToNodeMap() + GraphToNodeMap(), + OrderedDict{OptiNode,Vector{MOI.VariableIndex}}(), + OrderedDict{OptiNode,Vector{MOI.ConstraintIndex}}() ) end +function add_node(gb::GraphMOIBackend, node::OptiNode) + gb.node_variables[node] = MOI.VariableIndex[] + gb.node_constraints[node] = MOI.ConstraintIndex[] +end + # JuMP Extension function JuMP.backend(gb::GraphMOIBackend) @@ -135,8 +155,10 @@ 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) - _moi_constrain_node_variable( + push!(graph_backend(graph).node_variables[node], graph_var_index) + _moi_constrain_node_variable( graph_backend(graph), + node, graph_var_index, v.info, Float64 @@ -145,10 +167,68 @@ function _moi_add_node_variable( return vref end +function _moi_constrain_node_variable( + gb::GraphMOIBackend, + node::OptiNode, + index, + info, + ::Type{T}, +) where {T} + #TODO: set local node constraint indices + if info.has_lb + # next_node_index = next_constraint_index(node, MOI.VariableIndex, MOI.GreaterThan{T}) + con = _moi_add_constraint( + gb.moi_backend, + index, + MOI.GreaterThan{T}(info.lower_bound), + ) + push!(gb.node_constraints[node], con) + #push!(gb.node_constraints[node], next_node_index) + # gb.node_to_graph_map.con_map[next_node_index] = con + # gb.graph_to_node_map.con_map[con] = next_node_index + end + if info.has_ub + con = _moi_add_constraint( + gb.moi_backend, + index, + MOI.LessThan{T}(info.upper_bound), + ) + push!(gb.node_constraints[node], con) + end + if info.has_fix + con = _moi_add_constraint( + gb.moi_backend, + index, + MOI.EqualTo{T}(info.fixed_value), + ) + push!(gb.node_constraints[node], con) + end + if info.binary + con = _moi_add_constraint(gb.moi_backend, index, MOI.ZeroOne()) + push!(gb.node_constraints[node], con) + end + if info.integer + con = _moi_add_constraint(gb.moi_backend, index, MOI.Integer()) + push!(gb.node_constraints[node], con) + end + if info.has_start && info.start !== nothing + MOI.set( + gb.moi_backend, + MOI.VariablePrimalStart(), + index, + convert(T, info.start), + ) + end +end + function _add_variable_to_backend( graph_backend::GraphMOIBackend, vref::NodeVariableRef ) + # return if variable already in backend + vref in keys(graph_backend.node_to_graph_map.var_map) && return + + # add variable, track index 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 diff --git a/src/optigraph.jl b/src/optigraph.jl index 58040c8..caca828 100644 --- a/src/optigraph.jl +++ b/src/optigraph.jl @@ -113,6 +113,7 @@ function add_node( node_index = NodeIndex(length(graph.optinodes)+1) optinode = OptiNode{OptiGraph}(graph, node_index, label) push!(graph.optinodes, optinode) + add_node(graph.backend, optinode) return optinode end diff --git a/src/optinode.jl b/src/optinode.jl index 920200c..4a866a7 100644 --- a/src/optinode.jl +++ b/src/optinode.jl @@ -85,13 +85,26 @@ function JuMP.backend(node::OptiNode) return JuMP.backend(graph_backend(node)) 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)) + ) +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))) +end + function JuMP.num_constraints( node::OptiNode, ::Type{F}, ::Type{S} )::Int64 where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet} g2n = graph_backend(node).graph_to_node_map - cons = MOI.get(JuMP.backend(node),MOI.ListOfConstraintIndices{F,S}()) + cons = MOI.get(JuMP.backend(node), MOI.ListOfConstraintIndices{F,S}()) + println(cons) refs = [g2n[con] for con in cons] return length(filter((cref) -> cref.model == node, refs)) end @@ -110,6 +123,10 @@ Base.print(io::IO, vref::NodeVariableRef) = Base.print(io, Base.string(vref)) Base.show(io::IO, vref::NodeVariableRef) = Base.print(io, vref) Base.broadcastable(vref::NodeVariableRef) = Ref(vref) +function graph_index(vref::NodeVariableRef) + return graph_backend(vref.node).node_to_graph_map[vref] +end + """ JuMP.add_variable(node::OptiNode, v::JuMP.AbstractVariable, name::String="") @@ -151,10 +168,7 @@ function JuMP.set_name(vref::NodeVariableRef, s::String) return 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))) -end + ### Node Constraints