Skip to content

Commit

Permalink
got backend aggregation fully working
Browse files Browse the repository at this point in the history
  • Loading branch information
@jalving
jalving committed Jan 13, 2024
1 parent 92a16a5 commit 4d0b6da
Show file tree
Hide file tree
Showing 6 changed files with 250 additions and 196 deletions.
58 changes: 33 additions & 25 deletions src/moi_aggregate.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,20 +5,23 @@ Aggregate the moi backends from each subgraph within `graph` to create a single
"""
function aggregate_backends!(graph::OptiGraph)
for subgraph in get_subgraphs(graph)
_aggregate_subgraph_nodes!(graph, subgraph)
# _aggregate_subgraph_edges!(subgraph)
_copy_subgraph_nodes!(graph, subgraph)
_copy_subgraph_edges!(graph, subgraph)

# TODO: pass non-objective graph attributes (use an MOI Filter?)

end
end

function _aggregate_subgraph_nodes!(graph::OptiGraph, subgraph::OptiGraph)
function _copy_subgraph_nodes!(graph::OptiGraph, subgraph::OptiGraph)
for node in all_nodes(subgraph)
_append_node_to_backend!(graph, node)
end
end

function _aggregate_subgraph_edges!(graph::OptiGraph, subgraph::OptiGraph)
for edges in all_edges(subgraph)
_append_edge_to_backend!(graph, node)
function _copy_subgraph_edges!(graph::OptiGraph, subgraph::OptiGraph)
for edge in all_edges(subgraph)
_append_edge_to_backend!(graph, edge)
end
end

Expand All @@ -38,7 +41,7 @@ function _append_node_to_backend!(graph::OptiGraph, node::OptiNode)
variable_constraint_types = filter(all_constraint_types) do (F, S)
return MOIU._is_variable_function(F)
end
_copy_node_constraints(
_copy_element_constraints(
dest,
node,
index_map,
Expand All @@ -49,15 +52,12 @@ function _append_node_to_backend!(graph::OptiGraph, node::OptiNode)
nonvariable_constraint_types = filter(all_constraint_types) do (F, S)
return !MOIU._is_variable_function(F)
end
_copy_node_constraints(
_copy_element_constraints(
dest,
node,
index_map,
nonvariable_constraint_types
)

# TODO: pass non-objective graph attributes (use an MOI Filter?)

return
end

Expand Down Expand Up @@ -91,52 +91,53 @@ function _copy_node_variables(
return
end

function _copy_node_constraints(
function _copy_element_constraints(
dest::GraphMOIBackend,
node::OptiNode,
element::OptiElement,
index_map::MOIU.IndexMap,
constraint_types
)
for (F, S) in constraint_types
cis_src = MOI.get(node, MOI.ListOfConstraintIndices{F,S}())
_copy_node_constraints(dest, node, index_map, cis_src)
cis_src = MOI.get(element, MOI.ListOfConstraintIndices{F,S}())
_copy_element_constraints(dest, element, index_map, cis_src)
end

src = graph_backend(node)
# pass constraint attributes
src = graph_backend(element)
for (F, S) in constraint_types
MOIU.pass_attributes(
dest.moi_backend,
src.moi_backend,
index_map,
MOI.get(node, MOI.ListOfConstraintIndices{F,S}()),
MOI.get(element, MOI.ListOfConstraintIndices{F,S}()),
)
end
return
end

function _copy_node_constraints(
function _copy_element_constraints(
dest::GraphMOIBackend,
node::OptiNode,
element::OptiElement,
index_map::MOIU.IndexMap,
cis_src::Vector{MOI.ConstraintIndex{F,S}}
) where {F,S}
return _copy_node_constraints(dest, node, index_map, index_map[F, S], cis_src)
return _copy_element_constraints(dest, element, index_map, index_map[F, S], cis_src)
end

function _copy_node_constraints(
function _copy_element_constraints(
dest::GraphMOIBackend,
node::OptiNode,
element::OptiElement,
index_map::MOIU.IndexMap,
index_map_FS,
cis_src::Vector{<:MOI.ConstraintIndex},
)
src = graph_backend(node)
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)
cref = src.graph_to_element_map[ci]
cref in keys(dest.element_to_graph_map.con_map) && return
dest_index = _add_node_constraint_to_backend(
dest_index = _add_element_constraint_to_backend(
dest,
cref,
MOIU.map_indices(index_map, f),
Expand All @@ -150,10 +151,17 @@ end
function _append_edge_to_backend!(graph::OptiGraph, edge::OptiEdge)
src = graph_backend(edge)
dest = graph_backend(graph)

# populate index map with node data for src -- > dest
index_map = MOIU.IndexMap()
vars = all_variables(edge)
for var in vars
index_map[src.element_to_graph_map[var]] = dest.element_to_graph_map[var]
end

# copy the constraints
constraint_types = MOI.get(edge, MOI.ListOfConstraintTypesPresent())
_copy_edge_constraints(
_copy_element_constraints(
dest,
edge,
index_map,
Expand Down
127 changes: 58 additions & 69 deletions src/moi_graph_backend.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,8 @@ mutable struct GraphMOIBackend <: MOI.AbstractOptimizer

# map of variables and constraints on nodes and edges to graph backend indices
node_variables::OrderedDict{OptiNode,Vector{MOI.VariableIndex}}
node_constraints::OrderedDict{OptiNode,Vector{MOI.ConstraintIndex}}
edge_constraints::OrderedDict{OptiEdge,Vector{MOI.ConstraintIndex}}
element_constraints::OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}
#edge_constraints::OrderedDict{OptiEdge,Vector{MOI.ConstraintIndex}}

# TODO (maybe): legacy JuMP nonlinear support
# nlp_model::MOI.Nonlinear.Model
Expand All @@ -97,14 +97,18 @@ function GraphMOIBackend(optigraph::AbstractOptiGraph)
ElementToGraphMap(),
GraphToElementMap(),
OrderedDict{OptiNode,Vector{MOI.VariableIndex}}(),
OrderedDict{OptiNode,Vector{MOI.ConstraintIndex}}(),
OrderedDict{OptiEdge,Vector{MOI.ConstraintIndex}}()
OrderedDict{OptiElement,Vector{MOI.ConstraintIndex}}(),
#OrderedDict{OptiEdge,Vector{MOI.ConstraintIndex}}()
)
end

function add_node(gb::GraphMOIBackend, node::OptiNode)
gb.node_variables[node] = MOI.VariableIndex[]
gb.node_constraints[node] = MOI.ConstraintIndex[]
gb.element_constraints[node] = MOI.ConstraintIndex[]
end

function add_edge(gb::GraphMOIBackend, edge::OptiEdge)
gb.element_constraints[edge] = MOI.ConstraintIndex[]
end

# JuMP Extension
Expand Down Expand Up @@ -218,7 +222,7 @@ function _add_variable_to_backend(
return graph_var_index
end

### Node Constraints
### Constraints

function next_constraint_index(
node::OptiNode,
Expand Down Expand Up @@ -250,39 +254,70 @@ function _moi_add_node_constraint(
moi_func_graph = _create_graph_moi_func(graph_backend(graph), moi_func, jump_func)

# add to optinode backend
_add_node_constraint_to_backend(graph_backend(graph), cref, moi_func_graph, moi_set)
# _add_node_constraint_to_backend(graph_backend(graph), cref, moi_func_graph, moi_set)
_add_element_constraint_to_backend(graph_backend(graph), cref, moi_func_graph, moi_set)
end
return cref
end

function next_constraint_index(
edge::OptiEdge,
::Type{F},
::Type{S}
)::MOI.ConstraintIndex{F,S} where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
index = num_constraints(edge, F, S)
return MOI.ConstraintIndex{F,S}(index + 1)
end

function _moi_add_edge_constraint(
edge::OptiEdge,
con::JuMP.AbstractConstraint
)
# get moi function and set
jump_func = JuMP.jump_function(con)
moi_func = JuMP.moi_function(con)
moi_set = JuMP.moi_set(con)

# create constraint index and reference
constraint_index = next_constraint_index(
edge,
typeof(moi_func),
typeof(moi_set)
)::MOI.ConstraintIndex{typeof(moi_func),typeof(moi_set)}
cref = ConstraintRef(edge, constraint_index, JuMP.shape(con))

# update graph backends
for graph in containing_optigraphs(edge)
# add backend variables if linking across optigraphs
_add_backend_variables(graph_backend(graph), jump_func)

# update the moi function variable indices
moi_func_graph = _create_graph_moi_func(graph_backend(graph), moi_func, jump_func)

# add the constraint to the backend
_add_element_constraint_to_backend(graph_backend(graph), cref, moi_func_graph, moi_set)
end

return cref
end

function _add_node_constraint_to_backend(
function _add_element_constraint_to_backend(
graph_backend::GraphMOIBackend,
cref::ConstraintRef,
func::F,
set::S
) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
cref in keys(graph_backend.element_to_graph_map.con_map) && return
if !haskey(graph_backend.node_constraints, cref.model)
graph_backend.node_constraints[cref.model] = MOI.ConstraintIndex[]
if !haskey(graph_backend.element_constraints, cref.model)
graph_backend.element_constraints[cref.model] = MOI.ConstraintIndex[]
end
graph_con_index = MOI.add_constraint(graph_backend.moi_backend, func, set)
graph_backend.element_to_graph_map[cref] = graph_con_index
graph_backend.graph_to_element_map[graph_con_index] = cref
push!(graph_backend.node_constraints[cref.model], graph_con_index)
push!(graph_backend.element_constraints[cref.model], graph_con_index)
return graph_con_index
end

### Edge Constraints

function next_constraint_index(
edge::OptiEdge,
::Type{F},
::Type{S}
)::MOI.ConstraintIndex{F,S} where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
index = num_constraints(edge, F, S)
return MOI.ConstraintIndex{F,S}(index + 1)
end

### Graph MOI Utilities

"""
Expand Down Expand Up @@ -384,9 +419,7 @@ function _update_nonlinear_func!(
return
end

### add variables to a backend for purpose of creating linking constraints or objectives
# across subgraphs

# add variables to a backend for linking across subgraphs
function _add_backend_variables(
backend::GraphMOIBackend,
jump_func::JuMP.GenericAffExpr
Expand Down Expand Up @@ -433,50 +466,6 @@ function _add_backend_variables(
return
end

function _moi_add_edge_constraint(
edge::OptiEdge,
con::JuMP.AbstractConstraint
)
# get moi function and set
jump_func = JuMP.jump_function(con)
moi_func = JuMP.moi_function(con)
moi_set = JuMP.moi_set(con)

# create constraint index and reference
constraint_index = next_constraint_index(
edge,
typeof(moi_func),
typeof(moi_set)
)::MOI.ConstraintIndex{typeof(moi_func),typeof(moi_set)}
cref = ConstraintRef(edge, constraint_index, JuMP.shape(con))

# update graph backends
for graph in containing_optigraphs(edge)
# add backend variables if linking across optigraphs
_add_backend_variables(graph_backend(graph), jump_func)

# update the moi function variable indices
moi_func_graph = _create_graph_moi_func(graph_backend(graph), moi_func, jump_func)

# add the constraint to the backend
_add_edge_constraint_to_backend(graph_backend(graph), cref, moi_func_graph, moi_set)
end

return cref
end

function _add_edge_constraint_to_backend(
graph_backend::GraphMOIBackend,
cref::ConstraintRef,
func::F,
set::S
) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
graph_con_index = MOI.add_constraint(graph_backend.moi_backend, func, set)
graph_backend.element_to_graph_map[cref] = graph_con_index
graph_backend.graph_to_element_map[graph_con_index] = cref
return graph_con_index
end

### Objective Function

function _moi_set_objective_function(
Expand Down
Loading

0 comments on commit 4d0b6da

Please sign in to comment.