diff --git a/src/MathProgIncidence.jl b/src/MathProgIncidence.jl index 350fdf5..6fb2355 100644 --- a/src/MathProgIncidence.jl +++ b/src/MathProgIncidence.jl @@ -32,6 +32,7 @@ include("incidence_graph.jl") include("maximum_matching.jl") include("dulmage_mendelsohn.jl") include("block_triangularize.jl") +include("bfs.jl") include("interface.jl") # Methods to get incidence matrices (as SparseMatrixCSC) from diff --git a/src/bfs.jl b/src/bfs.jl new file mode 100644 index 0000000..cb3a2b0 --- /dev/null +++ b/src/bfs.jl @@ -0,0 +1,66 @@ +# ___________________________________________________________________________ +# +# MathProgIncidence.jl: Math Programming Incidence Graph Analysis +# Copyright (c) 2023. Triad National Security, LLC. All rights reserved. +# +# This program was produced under U.S. Government contract 89233218CNA000001 +# for Los Alamos National Laboratory (LANL), which is operated by Triad +# National Security, LLC for the U.S. Department of Energy/National Nuclear +# Security Administration. All rights in the program are reserved by Triad +# National Security, LLC, and the U.S. Department of Energy/National Nuclear +# Security Administration. The Government is granted for itself and others +# acting on its behalf a nonexclusive, paid-up, irrevocable worldwide license +# in this material to reproduce, prepare derivative works, distribute copies +# to the public, perform publicly and display publicly, and to permit others +# to do so. +# +# This software is distributed under the 3-clause BSD license. +# ___________________________________________________________________________ + +import Graphs + +function _limited_bfs(adjacency_list::Vector{Vector{Int}}, root::Int; depth::Int = 1) + # This is the queue + nodes = Int[root] + idx = 1 + # I use a dict rather than an array here to avoid constructing an O(n) + # array for every (potentially) small BFS in a (potentially) large graph + parents = Dict(root => 0) + # We differ from traditional BFS by imposing a depth limit. We store the index + # (in the queue) corresponding to the start of the current "level". + maxdepth = depth + current_depth = 0 + levelstart = 1 + while idx <= length(nodes) + node = nodes[idx] + parent_idx = idx + idx += 1 + for other in adjacency_list[node] + if other ∉ keys(parents) + if parent_idx >= levelstart + # When we first encounter the child of a node in the current level, + # we advance (a) the depth and (b) the level start index. + # The level starts at the child node we are about to add. We do the + # checking here so we can break before adding the child if doing + # so would violate the depth limit. + levelstart = length(nodes) + 1 + current_depth += 1 + if current_depth > maxdepth + # Break before adding this child to the queue + break + end + end + parents[other] = parent_idx + push!(nodes, other) + end + end + # Once we exceed the depth limit, we must stop processing children + # of future nodes. + if current_depth > maxdepth + break + end + end + parents = map(n -> parents[n], nodes) + dag = Graphs.tree(parents) # Returns a DiGraph + return nodes, dag +end diff --git a/src/identify_variables.jl b/src/identify_variables.jl index 1911b9e..1e77125 100644 --- a/src/identify_variables.jl +++ b/src/identify_variables.jl @@ -199,7 +199,7 @@ function identify_unique_variables( end """ - identify_unique_variables(fcn)::Vector{JuMP.VariableIndex} + identify_unique_variables(fcn)::Vector{MOI.VariableIndex} Return the variables that appear in the provided MathOptInterface function. @@ -224,7 +224,7 @@ function identify_unique_variables( return _filter_duplicates(variables) end -# This method is used to handle function-in-set constraints. +# This method is used to handle variable-in-set constraints. function identify_unique_variables( var::MOI.VariableIndex )::Vector{MOI.VariableIndex} @@ -248,6 +248,39 @@ function identify_unique_variables( )) end +function identify_unique_variables(expr::JuMP.AffExpr)::Vector{JuMP.VariableRef} + return _filter_duplicates(collect(keys(expr.terms))) +end + +function identify_unique_variables(expr::JuMP.QuadExpr)::Vector{JuMP.VariableRef} + vars = collect(keys(expr.aff.terms)) + for term in keys(expr.terms) + push!(vars, term.a) + push!(vars, term.b) + end + return _filter_duplicates(vars) +end + +function identify_unique_variables(expr::JuMP.VariableRef)::Vector{JuMP.VariableRef} + return [expr] +end + +function identify_unique_variables(expr::T)::Vector{JuMP.VariableRef} where {T<:Real} + return [] +end + +function identify_unique_variables(expr::JuMP.NonlinearExpr)::Vector{JuMP.VariableRef} + vars = JuMP.VariableRef[] + for arg in expr.args + # NOTE: This could be made more efficient by avoiding the construction + # of many small vectors. However, this method is probably not critical + # for overall performance as of 2025-10-31. + append!(vars, identify_unique_variables(arg)) + end + return _filter_duplicates(vars) +end + + """ _identify_variables(fcn)::Vector{JuMP.VariableIndex} diff --git a/src/interface.jl b/src/interface.jl index 2b4662f..ae2e87e 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -713,3 +713,114 @@ end block_triangularize(model::JuMP.Model) = block_triangularize(IncidenceGraphInterface(model)) block_triangularize(matrix::SparseMatrixCSC) = block_triangularize(IncidenceGraphInterface(matrix)) block_triangularize(matrix::Matrix) = block_triangularize(IncidenceGraphInterface(matrix)) + +const SubtreeNodeType = Union{ + JuMP.VariableRef, + JuMP.ConstraintRef, + Int, + # This is starting to get a little excessive. + # Should I just use Vector{Any}? + JuMP.AffExpr, + JuMP.QuadExpr, + JuMP.NonlinearExpr, +} + +struct IncidenceSubtree{T} + _nodes::Vector{SubtreeNodeType} + _dag::Graphs.DiGraph{T} +end + +function _collect_lines!(lines::Vector, tree::IncidenceSubtree; node = 1, indent = "") + indent = (indent == "" ? "├ " : join(["│ ", indent])) + orig_indent = indent + neighbors = Graphs.neighbors(tree._dag, node) + for i in neighbors + child = tree._nodes[i] + if i == last(neighbors) + indent = replace(indent, "├" => "└") + end + push!(lines, "$indent$child") + # This could be handled better... + indent = orig_indent + _collect_lines!(lines, tree; node = i, indent) + end + return +end + +function Base.show(io::IO, tree::IncidenceSubtree) + root = tree._nodes[1] + lines = ["$root"] + _collect_lines!(lines, tree) + for line in lines + println(io, line) + end + return +end + +function limited_bfs( + igraph::IncidenceGraphInterface, + root::Union{JuMP.VariableRef,<:JuMP.ConstraintRef,Int}; + depth::Int = 1, +) + root = if root in keys(igraph._var_node_map) + igraph._var_node_map[root] + elseif root in keys(igraph._con_node_map) + igraph._con_node_map[root] + else + error("root is not a node in this graph") + end + adjlist = Graphs.SimpleGraphs.adj(igraph._graph) + nodes, dag = _limited_bfs(adjlist, root; depth) + nodes = convert(Vector{SubtreeNodeType}, map(n -> igraph._nodes[n], nodes)) + return IncidenceSubtree(nodes, dag) +end + +function limited_bfs( + root::Union{JuMP.VariableRef,<:JuMP.ConstraintRef}; + depth::Int = 1, + kwds..., +) + igraph = IncidenceGraphInterface(root.model; kwds...) + return limited_bfs(igraph, root; depth) +end + +function limited_bfs( + igraph::IncidenceGraphInterface, + root::Union{JuMP.AffExpr,JuMP.QuadExpr,JuMP.NonlinearExpr}; + depth::Int = 1, +) + nnodes = length(igraph._nodes) + # Here, I assume the nodes are contiguous integers + expr_node = nnodes + 1 + # Get the expression's neighbors. These must be variables + expr_neighbors = identify_unique_variables(root) + expr_neighbors = map(v -> igraph._var_node_map[v], expr_neighbors) + g = igraph._graph + adjlist = map(n -> Graphs.neighbors(g, n), 1:nnodes) + # Add the new node's edges to our adjacency list + push!(adjlist, expr_neighbors) + subtree_nodes, dag = _limited_bfs(adjlist, expr_node; depth) + # We know that root (or its corresponding integer, expr_node) is the first + # node in subtree_nodes. We need to use this because we can't look it up + # in igraph._nodes. + jump_subtree_nodes = SubtreeNodeType[root] + append!( + jump_subtree_nodes, + map(n -> igraph._nodes[n], subtree_nodes[2:end]), + ) + return IncidenceSubtree(jump_subtree_nodes, dag) +end + +function limited_bfs( + root::Union{JuMP.AffExpr, JuMP.QuadExpr, JuMP.NonlinearExpr}; + depth::Int = 1, + kwds..., +) + expr_neighbors = identify_unique_variables(root) + if isempty(expr_neighbors) + return IncidenceSubtree(SubtreeNodeType[root], Graphs.DiGraph(1)) + else + igraph = IncidenceGraphInterface(expr_neighbors[1].model; kwds...) + return limited_bfs(igraph, root; depth) + end +end