Breadth-first search

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

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

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}
 

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