From 4e34d46ff57bb2f1635d3b9dd03899267cfa6e7f Mon Sep 17 00:00:00 2001 From: jalving Date: Sun, 12 May 2024 19:48:07 -0700 Subject: [PATCH] wip: more progress on JuMP/MOI methods --- src/Plasmo.jl | 2 +- src/{ => backends}/moi_backend.jl | 117 +++++++++++++++++++++++++----- src/core_types.jl | 3 +- src/node_variables.jl | 7 +- src/optinode.jl | 106 +++++++++++---------------- 5 files changed, 144 insertions(+), 91 deletions(-) rename src/{ => backends}/moi_backend.jl (85%) diff --git a/src/Plasmo.jl b/src/Plasmo.jl index 37cd0fe..afcbe44 100644 --- a/src/Plasmo.jl +++ b/src/Plasmo.jl @@ -38,7 +38,7 @@ include("optinode.jl") include("optiedge.jl") -include("moi_backend.jl") +include("backends/moi_backend.jl") include("aggregate.jl") diff --git a/src/moi_backend.jl b/src/backends/moi_backend.jl similarity index 85% rename from src/moi_backend.jl rename to src/backends/moi_backend.jl index 0154b5b..c29a25e 100644 --- a/src/moi_backend.jl +++ b/src/backends/moi_backend.jl @@ -81,6 +81,8 @@ mutable struct GraphMOIBackend <: MOI.AbstractOptimizer # map of nodes and edges to variables and constraints. node_variables::OrderedDict{OptiNode,Vector{MOI.VariableIndex}} element_constraints::OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}} + element_attributes::OrderedDict{Tuple{OptiElement,MOI.AbstractModelAttribute},Any} + operator_map::OrderedDict{Tuple{OptiElement,Symbol},Symbol} end """ @@ -98,10 +100,20 @@ function GraphMOIBackend(graph::OptiGraph) ElementToGraphMap(), GraphToElementMap(), OrderedDict{OptiNode,Vector{MOI.VariableIndex}}(), - OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}() + OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}(), + OrderedDict{Tuple{OptiElement,MOI.AbstractModelAttribute},Any}(), + OrderedDict{Tuple{OptiElement,Symbol},Symbol}() ) end +function graph_index(gb::GraphMOIBackend, nvref::NodeVariableRef) + return gb.element_to_graph_map[nvref] +end + +function graph_operator(gb::GraphMOIBackend, element::OptiElement, name::Symbol) + return gb.operator_map[(element,name)] +end + function _add_node(gb::GraphMOIBackend, node::OptiNode) if !haskey(gb.node_variables, node) gb.node_variables[node] = MOI.VariableIndex[] @@ -125,37 +137,104 @@ function JuMP.backend(gb::GraphMOIBackend) return gb.moi_backend end -# MOI Methods +### MOI Methods + +# graph attributes -function MOI.get(gb::GraphMOIBackend, attr::MOI.AnyAttribute) +function MOI.get(gb::GraphMOIBackend, attr::AT) where + AT <: Union{MOI.AbstractModelAttribute, MOI.AbstractOptimizerAttribute} return MOI.get(gb.moi_backend, attr) end -function MOI.get(gb::GraphMOIBackend, attr::MOI.AnyAttribute, ref::ConstraintRef) - graph_index = gb.element_to_graph_map[ref] - return MOI.get(gb.moi_backend, attr, graph_index) +function MOI.set(gb::GraphMOIBackend, attr::AT, args...) where + AT <: Union{MOI.AbstractModelAttribute, MOI.AbstractOptimizerAttribute} + MOI.set(gb.moi_backend, attr, args...) + return end -function MOI.get(gb::GraphMOIBackend, attr::MOI.AnyAttribute, ref::NodeVariableRef) - graph_index = gb.element_to_graph_map[ref] - return MOI.get(gb.moi_backend, attr, graph_index) +# element attributes + +# function MOI.get(gb::GraphMOIBackend, attr::AT, element::OptiElement) where +# AT <: MOI.AbstractModelAttribute +# return gb.element_attributes[(element,attr)] +# end + +# function MOI.set(gb::GraphMOIBackend, attr::AT, element::OptiElement, args...) where +# AT <: MOI.AbstractModelAttribute +# MOI.set(gb.moi_backend, attr, args...) +# gb.element_attributes[(element,attr)] = tuple(args...) +# return +# end +function MOI.set( + gb::GraphMOIBackend, + attr::MOI.UserDefinedFunction, + node::OptiNode, + args... +) + registered_name = Symbol(node.label, ".", attr.name) + MOI.set( + gb.moi_backend, + MOI.UserDefinedFunction(registered_name, attr.arity), + args... + ) + gb.element_attributes[(node,attr)] = tuple(args...) + gb.operator_map[(node,attr.name)] = registered_name end function MOI.get(gb::GraphMOIBackend, attr::MOI.NumberOfVariables, node::OptiNode) return length(gb.node_variables[node]) end -function MOI.get(gb::GraphMOIBackend, attr::MOI.NumberOfConstraints{F,S}, node::OptiNode) where{F<:MOI.AbstractFunction,S<:MOI.AbstractSet} - cons = MOI.get(gb.moi_backend, MOI.ListOfConstraintIndices{F,S}()) - refs = [gb.graph_to_element_map[con] for con in cons] - return length(filter((cref) -> cref.model == node, refs)) +function MOI.get( + gb::GraphMOIBackend, + attr::MOI.ListOfConstraintTypesPresent, + element::OptiElement +) + cons = gb.element_constraints[element] + con_types = unique(typeof.(cons)) + type_tuple = [(type.parameters[1],type.parameters[2]) for type in con_types] + return type_tuple end -function MOI.set(gb::GraphMOIBackend, attr::MOI.AnyAttribute, args...) - MOI.set(gb.moi_backend, attr, args...) +function MOI.get( + gb::GraphMOIBackend, + attr::MOI.NumberOfConstraints{F,S}, + element::OptiElement +) where{F<:MOI.AbstractFunction,S<:MOI.AbstractSet} + return length(gb.element_constraints[element]) +end + +# variable attributes + +function MOI.get(gb::GraphMOIBackend, attr::AT, nvref::NodeVariableRef) where + AT <: MOI.AbstractVariableAttribute + graph_index = gb.element_to_graph_map[nvref] + return MOI.get(gb.moi_backend, attr, graph_index) +end + +function MOI.set(gb::GraphMOIBackend, attr::AT, nvref::NodeVariableRef, args...) where + AT <: MOI.AbstractVariableAttribute + graph_index = gb.element_to_graph_map[nvref] + MOI.set(gb.moi_backend, attr, graph_index, args...) return end +# constraint attributes + +function MOI.get(gb::GraphMOIBackend, attr::AT, cref::ConstraintRef) where + AT <: MOI.AbstractConstraintAttribute + graph_index = gb.element_to_graph_map[cref] + return MOI.get(gb.moi_backend, attr, graph_index) +end + +function MOI.set(gb::GraphMOIBackend, attr::AT, cref::ConstraintRef, args...) where + AT <: MOI.AbstractConstraintAttribute + graph_index = gb.element_to_graph_map[nvref] + MOI.set(gb.moi_backend, attr, graph_index, args...) +end + +# delete + function MOI.delete(gb::GraphMOIBackend, nvref::NodeVariableRef) MOI.delete(gb.moi_backend, gb.element_to_graph_map[nvref]) delete!(gb.graph_to_element_map.var_map, gb.element_to_graph_map[nvref]) @@ -170,6 +249,8 @@ function MOI.delete(gb::GraphMOIBackend, cref::ConstraintRef) return end +# is_valid + function MOI.is_valid(gb::GraphMOIBackend, vi::MOI.VariableIndex) return MOI.is_valid(gb.moi_backend, vi) end @@ -178,6 +259,8 @@ function MOI.is_valid(gb::GraphMOIBackend, ci::MOI.ConstraintIndex) return MOI.is_valid(gb.moi_backend, ci) end +# optimize! + function MOI.optimize!(gb::GraphMOIBackend) MOI.optimize!(gb.moi_backend) return @@ -185,10 +268,6 @@ end ### Variables and Constraints -function graph_index(gb::GraphMOIBackend, nvref::NodeVariableRef) - return gb.element_to_graph_map[nvref] -end - function _add_variable_to_backend( graph_backend::GraphMOIBackend, vref::NodeVariableRef diff --git a/src/core_types.jl b/src/core_types.jl index 632ae13..e53c045 100644 --- a/src/core_types.jl +++ b/src/core_types.jl @@ -49,5 +49,4 @@ end const OptiElement = Union{OptiNode,OptiEdge} -const OptiObject = Union{OptiNode, OptiEdge, OptiGraph} - +const OptiObject = Union{OptiNode, OptiEdge, OptiGraph} \ No newline at end of file diff --git a/src/node_variables.jl b/src/node_variables.jl index e3a7fa2..a49dccf 100644 --- a/src/node_variables.jl +++ b/src/node_variables.jl @@ -30,7 +30,7 @@ function MOI.get( attr::MOI.AbstractVariableAttribute, nvref::NodeVariableRef ) - return MOI.get(graph_backend(node), attr, graph_index) + return MOI.get(graph_backend(node), attr, nvref) end function MOI.set( @@ -41,8 +41,7 @@ function MOI.set( ) for graph in containing_optigraphs(node) gb = graph_backend(graph) - graph_index = gb.element_to_graph_map[nvref] - MOI.set(gb, attr, graph_index, args...) + MOI.set(gb, attr, nvref, args...) end return end @@ -179,7 +178,7 @@ end ### node variable bounds function JuMP.has_lower_bound(nvref::NodeVariableRef) - return _moi_nv_has_lower_bound(graph_backend(nvref), nvref) + return _moi_nv_has_lower_bound(nvref) end function JuMP.set_lower_bound(nvref::NodeVariableRef, lower::Number) diff --git a/src/optinode.jl b/src/optinode.jl index d10332f..b591417 100644 --- a/src/optinode.jl +++ b/src/optinode.jl @@ -15,15 +15,27 @@ function Base.getindex(node::OptiNode, name::Symbol) return source_graph(node).node_obj_dict[t] end -function JuMP.num_variables(node::OptiNode) - return MOI.get(graph_backend(node), MOI.NumberOfVariables(), node) - #return length(graph_backend(node).node_variables[node]) +""" + source_graph(node::OptiNode) + +Return the optigraph that contains the optinode. This is the optigraph that +defined said node and stores node object dictionary data. +""" +function source_graph(node::OptiNode) + return node.source_graph.x end -function JuMP.all_variables(node::OptiNode) - gb = graph_backend(node) - graph_indices = gb.node_variables[node] - return getindex.(Ref(gb.graph_to_element_map), graph_indices) +function containing_optigraphs(node::OptiNode) + source = source_graph(node) + graphs = [source] + if haskey(source.node_to_graphs, node) + graphs = [graphs; source.node_to_graphs[node]] + end + return graphs +end + +function containing_backends(node::OptiNode) + return graph_backend.(containing_optigraphs(node)) end """ @@ -44,28 +56,7 @@ function graph_backend(node::OptiNode) return graph_backend(source_graph(node)) end -""" - source_graph(node::OptiNode) -Return the optigraph that contains the optinode. This is the optigraph that -defined said node and stores node object dictionary data. -""" -function source_graph(node::OptiNode) - return node.source_graph.x -end - -function containing_optigraphs(node::OptiNode) - source = source_graph(node) - graphs = [source] - if haskey(source.node_to_graphs, node) - graphs = [graphs; source.node_to_graphs[node]] - end - return graphs -end - -function containing_backends(node::OptiNode) - return graph_backend.(containing_optigraphs(node)) -end """ Filter the object dictionary for values that belong to node. Keep in mind that @@ -88,6 +79,17 @@ end ### JuMP Methods +function JuMP.num_variables(node::OptiNode) + return MOI.get(graph_backend(node), MOI.NumberOfVariables(), node) + #return length(graph_backend(node).node_variables[node]) +end + +function JuMP.all_variables(node::OptiNode) + gb = graph_backend(node) + graph_indices = gb.node_variables[node] + return getindex.(Ref(gb.graph_to_element_map), graph_indices) +end + function JuMP.delete(node::OptiNode, cref::ConstraintRef) if node !== JuMP.owner_model(cref) error( @@ -150,9 +152,12 @@ function JuMP.add_nonlinear_operator( "hesssian provided)", ) end - name = Symbol(node.label, ".", name) - MOI.set(graph_backend(node), MOI.UserDefinedFunction(name, dim), tuple(f, args...)) - return JuMP.NonlinearOperator(f, name) + #registered_name = Symbol(node.label, ".", name) + #MOI.set(node, MOI.UserDefinedFunction(registered_name, dim), tuple(f, args...)) + #return JuMP.NonlinearOperator(f, registered_name) + MOI.set(node, MOI.UserDefinedFunction(name, dim), tuple(f, args...)) + registered_name = graph_operator(graph_backend(node), node, name) + return JuMP.NonlinearOperator(f, registered_name) end function _set_dirty(node::OptiNode) @@ -212,48 +217,20 @@ end # TODO: store objective functions on nodes and query as node attributes -function MOI.get(node::OptiNode, attr::MOI.AnyAttribute) - return MOI.get(graph_backend(node), attr) -end - # function MOI.get(node::OptiNode, attr::MOI.UserDefinedFunction) # return MOI.get(graph_backend(node), attr) # end -# TODO: consider caching constraint types in graph backend versus using unique to filter -function MOI.get(node::OptiNode, attr::MOI.ListOfConstraintTypesPresent) - cons = graph_backend(node).element_constraints[node] - con_types = unique(typeof.(cons)) - type_tuple = [(type.parameters[1],type.parameters[2]) for type in con_types] - return type_tuple +function MOI.get(node::OptiNode, attr::MOI.AnyAttribute) + return MOI.get(graph_backend(node), attr, node) 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).element_constraints[node] - if (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S) - push!(con_inds, con) - end +function MOI.set(node::OptiNode, attr::MOI.AnyAttribute, args...) + for graph in containing_optigraphs(node) + MOI.set(graph_backend(node), attr, node, args...) end - return con_inds 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).element_constraints[node] -# if (typeof(con).parameters[1] == F && typeof(con).parameters[2] == S) -# push!(con_inds, con) -# end -# end -# return con_inds -# end - function MOI.get( node::OptiNode, attr::MOI.AbstractConstraintAttribute, @@ -270,7 +247,6 @@ function MOI.set( ) for graph in containing_optigraphs(JuMP.owner_model(cref)) gb = graph_backend(graph) - graph_index = gb.element_to_graph_map[cref] MOI.set(gb, attr, graph_index, args...) end return